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

//! Burn inference session management
//!
//! Provides a high-level API for running inference with Burn models.

use crate::backend::{BackendConfig, BurnBackend};
use crate::error::{BurnError, Result};
use crate::model::BurnModel;
use crate::tensor_bridge::{burn_to_rtx, rtx_to_burn};
use rtx_tensor::{Device, Tensor};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;
use tracing::{debug, info};

/// Configuration for a Burn session
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BurnConfig {
    /// Backend configuration
    pub backend: BackendConfig,
    /// Enable model caching
    pub cache_models: bool,
    /// Maximum cached models
    pub max_cached_models: usize,
    /// Output device for results
    pub output_device: Device,
    /// Enable profiling
    pub enable_profiling: bool,
}

impl Default for BurnConfig {
    fn default() -> Self {
        Self {
            backend: BackendConfig::default(),
            cache_models: true,
            max_cached_models: 10,
            output_device: Device::Cpu,
            enable_profiling: false,
        }
    }
}

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

    /// Set the device index
    pub fn with_device(mut self, index: usize) -> Self {
        self.backend.device_index = index;
        self
    }

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

    /// Enable profiling
    pub fn with_profiling(mut self) -> Self {
        self.enable_profiling = true;
        self
    }
}

/// Statistics for a Burn 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,
    /// Cache hits
    pub cache_hits: u64,
    /// Cache misses
    pub cache_misses: 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
        }
    }

    /// Cache hit rate
    pub fn cache_hit_rate(&self) -> f64 {
        let total = self.cache_hits + self.cache_misses;
        if total == 0 {
            0.0
        } else {
            self.cache_hits as f64 / total as f64
        }
    }
}

/// Burn inference session
///
/// Manages model loading, caching, and inference execution.
pub struct BurnSession {
    /// Session configuration
    config: BurnConfig,
    /// Loaded models cache
    models: HashMap<String, BurnModel>,
    /// Session statistics
    stats: SessionStats,
}

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

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

        // Initialize seed if provided
        if let Some(seed) = config.backend.seed {
            debug!("Setting seed: {}", seed);
            // Burn seed initialization would go here
        }

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

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

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

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

        self.stats.cache_misses += 1;

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

        // Cache if enabled
        if self.config.cache_models {
            // Evict oldest if at capacity
            if self.models.len() >= self.config.max_cached_models
                && let Some(key) = self.models.keys().next().cloned()
            {
                debug!("Evicting cached model: {}", key);
                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())
    }

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

        // Convert inputs to Burn format
        let mut burn_inputs = HashMap::new();
        for (name, tensor) in inputs {
            let (data, shape) = rtx_to_burn(tensor)?;
            burn_inputs.insert(name.clone(), (data, shape));
        }

        // Run inference (simulated for now - real implementation would use Burn's runtime)
        let burn_outputs = model.forward(burn_inputs)?;

        // Convert outputs back to rtx tensors
        let mut outputs = HashMap::new();
        for (name, (data, shape)) in burn_outputs {
            let tensor = burn_to_rtx(data, shape, &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: &BurnModel, 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(|| BurnError::Inference("No output tensor found".to_string()))
    }

    /// Get session configuration
    pub fn config(&self) -> &BurnConfig {
        &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) -> BurnBackend {
        self.config.backend.backend
    }
}

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

    #[test]
    fn test_config_default() {
        let config = BurnConfig::default();
        assert!(config.cache_models);
        assert_eq!(config.max_cached_models, 10);
    }

    #[test]
    fn test_config_builder() {
        let config = BurnConfig::default()
            .with_backend(BurnBackend::NdArray)
            .with_device(1)
            .with_profiling();

        assert_eq!(config.backend.backend, BurnBackend::NdArray);
        assert_eq!(config.backend.device_index, 1);
        assert!(config.enable_profiling);
    }

    #[test]
    fn test_session_stats() {
        let mut stats = SessionStats::default();
        stats.total_inferences = 10;
        stats.total_inference_ms = 100.0;
        stats.cache_hits = 8;
        stats.cache_misses = 2;

        assert!((stats.avg_inference_ms() - 10.0).abs() < 0.001);
        assert!((stats.cache_hit_rate() - 0.8).abs() < 0.001);
    }

    #[test]
    fn test_session_creation() {
        // NdArray backend should always be available with the default feature
        let config = BurnConfig::default().with_backend(BurnBackend::NdArray);
        let session = BurnSession::new(config);
        // May fail if ndarray feature is not enabled, that's ok for test
        if let Ok(session) = session {
            assert_eq!(session.backend(), BurnBackend::NdArray);
        }
    }
}