rustytorch/crates/integration/rtx-candle/src/session.rs

//! Candle inference session management

use crate::backend::{CandleBackend, CandleDevice};
use crate::error::{CandleError, Result};
use crate::model::CandleModel;
use crate::tensor_bridge::{candle_to_rtx, rtx_to_candle};
use rtx_tensor::{Device, Tensor};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;
use tracing::{debug, info};

/// Configuration for a Candle session
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CandleConfig {
    /// Device to use
    pub device: CandleDevice,
    /// Enable model caching
    pub cache_models: bool,
    /// Maximum cached models
    pub max_cached_models: usize,
    /// Output device for results
    pub output_device: Device,
    /// Use half precision (fp16)
    pub use_half: bool,
    /// Enable flash attention if available
    pub flash_attention: bool,
}

impl Default for CandleConfig {
    fn default() -> Self {
        Self {
            device: CandleDevice::default(),
            cache_models: true,
            max_cached_models: 10,
            output_device: Device::Cpu,
            use_half: false,
            flash_attention: true,
        }
    }
}

impl CandleConfig {
    /// Set the backend
    pub fn with_backend(mut self, backend: CandleBackend) -> Self {
        self.device.backend = backend;
        self
    }

    /// Set the device ordinal
    pub fn with_device_ordinal(mut self, ordinal: usize) -> Self {
        self.device.ordinal = ordinal;
        self
    }

    /// Set the output device
    pub fn with_output_device(mut self, device: Device) -> Self {
        self.output_device = device;
        self
    }

    /// Enable half precision
    pub fn with_half_precision(mut self) -> Self {
        self.use_half = true;
        self
    }

    /// Disable flash attention
    pub fn without_flash_attention(mut self) -> Self {
        self.flash_attention = false;
        self
    }
}

/// Statistics for a Candle session
#[derive(Debug, Clone, Default)]
pub struct SessionStats {
    /// Total inferences run
    pub total_inferences: u64,
    /// Total inference time in milliseconds
    pub total_inference_ms: f64,
    /// Models currently cached
    pub cached_models: usize,
    /// Total tokens processed (for language models)
    pub total_tokens: u64,
}

impl SessionStats {
    /// Average inference time
    pub fn avg_inference_ms(&self) -> f64 {
        if self.total_inferences == 0 {
            0.0
        } else {
            self.total_inference_ms / self.total_inferences as f64
        }
    }

    /// Tokens per second
    pub fn tokens_per_second(&self) -> f64 {
        if self.total_inference_ms == 0.0 {
            0.0
        } else {
            (self.total_tokens as f64) / (self.total_inference_ms / 1000.0)
        }
    }
}

/// Candle inference session
pub struct CandleSession {
    /// Session configuration
    config: CandleConfig,
    /// Loaded models cache
    models: HashMap<String, CandleModel>,
    /// Session statistics
    stats: SessionStats,
}

impl CandleSession {
    /// Create a new session with the given configuration
    pub fn new(config: CandleConfig) -> Result<Self> {
        info!(
            "Creating Candle session with {} backend",
            config.device.backend.name()
        );

        // Validate backend availability
        if !config.device.backend.is_available() {
            return Err(CandleError::BackendUnavailable(
                config.device.backend.name().to_string(),
            ));
        }

        Ok(Self {
            config,
            models: HashMap::new(),
            stats: SessionStats::default(),
        })
    }

    /// Create a session with default configuration
    pub fn default_session() -> Result<Self> {
        Self::new(CandleConfig::default())
    }

    /// Load a model from a file path
    pub fn load_model(&mut self, path: impl AsRef<Path>) -> Result<&CandleModel> {
        let path = path.as_ref();
        let path_str = path.to_string_lossy().to_string();

        // Check cache
        if self.config.cache_models && self.models.contains_key(&path_str) {
            debug!("Model cache hit: {}", path_str);
            return Ok(self.models.get(&path_str).unwrap());
        }

        // Load the model
        info!("Loading model from: {}", path.display());
        let model = CandleModel::load(path, &self.config)?;

        // Cache if enabled
        if self.config.cache_models {
            if self.models.len() >= self.config.max_cached_models
                && let Some(key) = self.models.keys().next().cloned()
            {
                self.models.remove(&key);
            }
            self.models.insert(path_str.clone(), model);
            self.stats.cached_models = self.models.len();
        } else {
            self.models.insert(path_str.clone(), model);
        }

        Ok(self.models.get(&path_str).unwrap())
    }

    /// Load a model from HuggingFace Hub
    #[cfg(feature = "transformers")]
    pub fn load_from_hub(&mut self, model_id: &str) -> Result<&CandleModel> {
        use crate::hub::download_model;

        let path = download_model(model_id, &Default::default())?;
        self.load_model(path)
    }

    /// Run inference on a model
    pub fn run(
        &mut self,
        model: &CandleModel,
        inputs: HashMap<String, &Tensor>,
    ) -> Result<HashMap<String, Tensor>> {
        let start = std::time::Instant::now();

        // Convert inputs to Candle format
        let mut candle_inputs = HashMap::new();
        for (name, tensor) in inputs {
            let candle_tensor = rtx_to_candle(tensor, &self.config.device)?;
            candle_inputs.insert(name.clone(), candle_tensor);
        }

        // Run inference
        let candle_outputs = model.forward(candle_inputs)?;

        // Convert outputs back to rtx tensors
        let mut outputs = HashMap::new();
        for (name, candle_tensor) in candle_outputs {
            let tensor = candle_to_rtx(&candle_tensor, &self.config.output_device)?;
            outputs.insert(name, tensor);
        }

        // Update stats
        let elapsed = start.elapsed().as_secs_f64() * 1000.0;
        self.stats.total_inferences += 1;
        self.stats.total_inference_ms += elapsed;

        debug!("Inference completed in {:.2}ms", elapsed);
        Ok(outputs)
    }

    /// Run inference with a single input/output
    pub fn run_simple(&mut self, model: &CandleModel, input: &Tensor) -> Result<Tensor> {
        let inputs = HashMap::from([("input".to_string(), input)]);
        let mut outputs = self.run(model, inputs)?;

        outputs
            .remove("output")
            .ok_or_else(|| CandleError::Inference("No output tensor found".to_string()))
    }

    /// Get session configuration
    pub fn config(&self) -> &CandleConfig {
        &self.config
    }

    /// Get session statistics
    pub fn stats(&self) -> &SessionStats {
        &self.stats
    }

    /// Clear the model cache
    pub fn clear_cache(&mut self) {
        self.models.clear();
        self.stats.cached_models = 0;
        info!("Model cache cleared");
    }

    /// Get the backend being used
    pub fn backend(&self) -> CandleBackend {
        self.config.device.backend
    }

    /// Check if using GPU
    pub fn is_gpu(&self) -> bool {
        self.config.device.is_gpu()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_config_default() {
        let config = CandleConfig::default();
        assert!(config.cache_models);
        assert!(!config.use_half);
    }

    #[test]
    fn test_config_builder() {
        let config = CandleConfig::default()
            .with_backend(CandleBackend::Cpu)
            .with_half_precision()
            .without_flash_attention();

        assert_eq!(config.device.backend, CandleBackend::Cpu);
        assert!(config.use_half);
        assert!(!config.flash_attention);
    }

    #[test]
    fn test_session_stats() {
        let mut stats = SessionStats::default();
        stats.total_inferences = 10;
        stats.total_inference_ms = 1000.0;
        stats.total_tokens = 1000;

        assert!((stats.avg_inference_ms() - 100.0).abs() < 0.001);
        assert!((stats.tokens_per_second() - 1000.0).abs() < 0.001);
    }

    #[test]
    fn test_session_creation() {
        let config = CandleConfig::default().with_backend(CandleBackend::Cpu);
        let session = CandleSession::new(config);
        assert!(session.is_ok());

        let session = session.unwrap();
        assert_eq!(session.backend(), CandleBackend::Cpu);
    }
}