rustytorch/crates/tooling/rtx-eval/src/core.rs

//! Core benchmarking framework for RTX-Eval

use crate::error::{RTXEvalError, RTXEvalResult};
use crate::metrics::MetricsEngine;
use anyhow::Result;
use async_trait::async_trait;
use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::time::{Duration, Instant};
use tracing::{error, info, warn};

/// Core benchmark trait that all benchmarks must implement
#[async_trait]
pub trait Benchmark: Send + Sync + std::fmt::Debug {
    /// Unique identifier for this benchmark
    fn name(&self) -> &str;

    /// Description of what this benchmark measures
    fn description(&self) -> &str;

    /// Category this benchmark belongs to
    fn category(&self) -> BenchmarkCategory;

    /// Expected runtime for this benchmark
    fn estimated_duration(&self) -> Duration;

    /// Run the benchmark and return results
    async fn run(&mut self, config: &BenchmarkConfig) -> RTXEvalResult<BenchmarkResult>;

    /// Validate benchmark prerequisites
    async fn validate_prerequisites(&self) -> RTXEvalResult<()>;

    /// Setup required resources
    async fn setup(&mut self) -> RTXEvalResult<()>;

    /// Cleanup resources
    async fn cleanup(&mut self) -> RTXEvalResult<()>;
}

/// Benchmark configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkConfig {
    pub use_gpu: bool,
    pub precision: PrecisionMode,
    pub batch_size: usize,
    pub max_sequence_length: usize,
    pub timeout: Duration,
    pub num_samples: Option<usize>,
    pub model_path: Option<String>,
    pub dataset_path: Option<String>,
    pub output_path: String,
}

/// Precision modes for benchmarking
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum PrecisionMode {
    FP16,
    FP32,
    FP64,
    Mixed,
}

/// Benchmark categories
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum BenchmarkCategory {
    Language,
    Vision,
    Multimodal,
    Scientific,
    Performance,
    Robustness,
    Fairness,
}

/// Individual benchmark result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkResult {
    pub benchmark_name: String,
    pub category: BenchmarkCategory,
    pub start_time: chrono::DateTime<chrono::Utc>,
    pub duration: Duration,
    pub success: bool,
    pub metrics: HashMap<String, f64>,
    pub metadata: HashMap<String, String>,
    pub error_message: Option<String>,
}

/// Collection of benchmark results
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkResults {
    pub category: BenchmarkCategory,
    pub results: Vec<BenchmarkResult>,
    pub summary: BenchmarkSummary,
    pub timestamp: chrono::DateTime<chrono::Utc>,
}

/// Summary statistics for a group of benchmarks
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkSummary {
    pub total_benchmarks: usize,
    pub successful_benchmarks: usize,
    pub failed_benchmarks: usize,
    pub total_duration: Duration,
    pub average_accuracy: f64,
    pub performance_score: f64,
}

/// Main benchmark suite orchestrator
#[derive(Debug)]
pub struct BenchmarkSuite {
    benchmarks: DashMap<BenchmarkCategory, Vec<Box<dyn Benchmark>>>,
    config: BenchmarkConfig,
    metrics_engine: Arc<MetricsEngine>,
}

impl BenchmarkSuite {
    /// Create a new benchmark suite
    pub fn new(eval_config: &crate::EvalConfig) -> Result<Self> {
        let config = BenchmarkConfig {
            use_gpu: eval_config.use_gpu,
            precision: eval_config.precision,
            batch_size: 32,
            max_sequence_length: 2048,
            timeout: eval_config.timeout,
            num_samples: None,
            model_path: None,
            dataset_path: None,
            output_path: eval_config.output_dir.clone(),
        };

        let metrics_engine = Arc::new(MetricsEngine::new(eval_config)?);

        let suite = Self {
            benchmarks: DashMap::new(),
            config,
            metrics_engine,
        };

        Ok(suite)
    }

    /// Register a benchmark with the suite
    pub fn register_benchmark(&self, benchmark: Box<dyn Benchmark>) {
        let category = benchmark.category();
        self.benchmarks.entry(category).or_default().push(benchmark);
    }

    /// Run all benchmarks in a category
    pub async fn run_category(
        &mut self,
        category: BenchmarkCategory,
    ) -> RTXEvalResult<BenchmarkResults> {
        info!("Running {:?} benchmark category", category);

        let mut results = Vec::new();
        let start_time = Instant::now();

        // Get benchmarks for this category
        let mut benchmarks =
            self.benchmarks
                .get_mut(&category)
                .ok_or_else(|| RTXEvalError::ConfigError {
                    message: format!("No benchmarks registered for category {category:?}"),
                })?;

        let total_benchmarks = benchmarks.len();
        let mut successful = 0;
        let mut failed = 0;

        // Run each benchmark
        for benchmark in benchmarks.iter_mut() {
            info!("Running benchmark: {}", benchmark.name());

            match self.run_single_benchmark(benchmark).await {
                Ok(result) => {
                    if result.success {
                        successful += 1;
                    } else {
                        failed += 1;
                    }
                    results.push(result);
                }
                Err(e) => {
                    error!("Benchmark {} failed: {}", benchmark.name(), e);
                    failed += 1;

                    // Create error result
                    let error_result = BenchmarkResult {
                        benchmark_name: benchmark.name().to_string(),
                        category,
                        start_time: chrono::Utc::now(),
                        duration: Duration::from_secs(0),
                        success: false,
                        metrics: HashMap::new(),
                        metadata: HashMap::new(),
                        error_message: Some(e.to_string()),
                    };
                    results.push(error_result);
                }
            }
        }

        let total_duration = start_time.elapsed();

        // Calculate summary statistics
        let average_accuracy = self.calculate_average_accuracy(&results);
        let performance_score = self.calculate_performance_score(&results);

        let summary = BenchmarkSummary {
            total_benchmarks,
            successful_benchmarks: successful,
            failed_benchmarks: failed,
            total_duration,
            average_accuracy,
            performance_score,
        };

        let benchmark_results = BenchmarkResults {
            category,
            results,
            summary,
            timestamp: chrono::Utc::now(),
        };

        info!(
            "Completed {:?} benchmarks: {}/{} successful in {:.2}s",
            category,
            successful,
            total_benchmarks,
            total_duration.as_secs_f64()
        );

        Ok(benchmark_results)
    }

    /// Run a single benchmark with full lifecycle management
    async fn run_single_benchmark(
        &self,
        benchmark: &mut Box<dyn Benchmark>,
    ) -> RTXEvalResult<BenchmarkResult> {
        let start_time = chrono::Utc::now();
        let timer = Instant::now();

        // Validate prerequisites
        benchmark.validate_prerequisites().await?;

        // Setup benchmark
        benchmark.setup().await?;

        // Run the actual benchmark with timeout
        let result = tokio::time::timeout(self.config.timeout, benchmark.run(&self.config)).await;

        // Cleanup regardless of success/failure
        if let Err(cleanup_err) = benchmark.cleanup().await {
            warn!(
                "Cleanup failed for benchmark {}: {}",
                benchmark.name(),
                cleanup_err
            );
        }

        // Process result
        match result {
            Ok(Ok(mut benchmark_result)) => {
                benchmark_result.start_time = start_time;
                benchmark_result.duration = timer.elapsed();
                Ok(benchmark_result)
            }
            Ok(Err(e)) => Err(RTXEvalError::BenchmarkFailed {
                message: format!("Benchmark {} failed: {}", benchmark.name(), e),
            }),
            Err(_) => Err(RTXEvalError::TimeoutError {
                benchmark: benchmark.name().to_string(),
                timeout_secs: self.config.timeout.as_secs(),
            }),
        }
    }

    /// Calculate average accuracy across results
    fn calculate_average_accuracy(&self, results: &[BenchmarkResult]) -> f64 {
        let successful_results: Vec<_> = results.iter().filter(|r| r.success).collect();

        if successful_results.is_empty() {
            return 0.0;
        }

        let total_accuracy: f64 = successful_results
            .iter()
            .filter_map(|r| r.metrics.get("accuracy"))
            .sum();

        total_accuracy / successful_results.len() as f64
    }

    /// Calculate overall performance score
    fn calculate_performance_score(&self, results: &[BenchmarkResult]) -> f64 {
        let successful_results: Vec<_> = results.iter().filter(|r| r.success).collect();

        if successful_results.is_empty() {
            return 0.0;
        }

        // Weighted score based on accuracy, throughput, and efficiency
        let mut total_score = 0.0_f64;
        let mut total_weight = 0.0_f64;

        for result in successful_results {
            let accuracy = result.metrics.get("accuracy").copied().unwrap_or(0.0);
            let throughput = result.metrics.get("throughput").copied().unwrap_or(0.0);
            let efficiency = result.metrics.get("efficiency").copied().unwrap_or(0.0);

            let score = (accuracy * 0.5) + (throughput * 0.3) + (efficiency * 0.2);
            total_score += score;
            total_weight += 1.0;
        }

        if total_weight > 0.0 {
            total_score / total_weight
        } else {
            0.0
        }
    }

    /// Run language benchmarks
    pub async fn run_language_suite(&mut self) -> RTXEvalResult<BenchmarkResults> {
        self.run_category(BenchmarkCategory::Language).await
    }

    /// Run vision benchmarks  
    pub async fn run_vision_suite(&mut self) -> RTXEvalResult<BenchmarkResults> {
        self.run_category(BenchmarkCategory::Vision).await
    }

    /// Run multimodal benchmarks
    pub async fn run_multimodal_suite(&mut self) -> RTXEvalResult<BenchmarkResults> {
        self.run_category(BenchmarkCategory::Multimodal).await
    }

    /// Run scientific benchmarks
    pub async fn run_scientific_suite(&mut self) -> RTXEvalResult<BenchmarkResults> {
        self.run_category(BenchmarkCategory::Scientific).await
    }

    /// Run performance benchmarks
    pub async fn run_performance_suite(&mut self) -> RTXEvalResult<BenchmarkResults> {
        self.run_category(BenchmarkCategory::Performance).await
    }

    /// Run robustness benchmarks
    pub async fn run_robustness_suite(&mut self) -> RTXEvalResult<BenchmarkResults> {
        self.run_category(BenchmarkCategory::Robustness).await
    }
}

impl Default for BenchmarkConfig {
    fn default() -> Self {
        Self {
            use_gpu: true,
            precision: PrecisionMode::FP32,
            batch_size: 32,
            max_sequence_length: 2048,
            timeout: Duration::from_secs(3600),
            num_samples: None,
            model_path: None,
            dataset_path: None,
            output_path: "./benchmark_results".to_owned(),
        }
    }
}

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

    #[derive(Debug)]
    struct MockBenchmark {
        name: String,
        category: BenchmarkCategory,
        should_fail: bool,
    }

    impl MockBenchmark {
        fn new(name: &str, category: BenchmarkCategory, should_fail: bool) -> Self {
            Self {
                name: name.to_string(),
                category,
                should_fail,
            }
        }
    }

    #[async_trait]
    impl Benchmark for MockBenchmark {
        fn name(&self) -> &str {
            &self.name
        }

        fn description(&self) -> &str {
            "Mock benchmark for testing"
        }

        fn category(&self) -> BenchmarkCategory {
            self.category
        }

        fn estimated_duration(&self) -> Duration {
            Duration::from_millis(100)
        }

        async fn run(&mut self, _config: &BenchmarkConfig) -> RTXEvalResult<BenchmarkResult> {
            if self.should_fail {
                return Err(RTXEvalError::BenchmarkFailed {
                    message: "Mock failure".to_owned(),
                });
            }

            let mut metrics = HashMap::new();
            metrics.insert("accuracy".to_owned(), 0.95);
            metrics.insert("throughput".to_owned(), 1000.0);

            Ok(BenchmarkResult {
                benchmark_name: self.name.clone(),
                category: self.category,
                start_time: chrono::Utc::now(),
                duration: Duration::from_millis(100),
                success: true,
                metrics,
                metadata: HashMap::new(),
                error_message: None,
            })
        }

        async fn validate_prerequisites(&self) -> RTXEvalResult<()> {
            Ok(())
        }

        async fn setup(&mut self) -> RTXEvalResult<()> {
            Ok(())
        }

        async fn cleanup(&mut self) -> RTXEvalResult<()> {
            Ok(())
        }
    }

    #[tokio::test]
    async fn test_benchmark_suite_creation() {
        let config = EvalConfig::default();
        let suite = BenchmarkSuite::new(&config);
        assert!(suite.is_ok());
    }

    #[tokio::test]
    async fn test_benchmark_registration_and_execution() {
        let config = EvalConfig::default();
        let mut suite = BenchmarkSuite::new(&config).unwrap();

        // Register a mock benchmark
        let benchmark = Box::new(MockBenchmark::new(
            "test_benchmark",
            BenchmarkCategory::Language,
            false,
        ));
        suite.register_benchmark(benchmark);

        // Run the category
        let results = suite.run_language_suite().await;
        assert!(results.is_ok());

        let results = results.unwrap();
        assert_eq!(results.results.len(), 1);
        assert!(results.results[0].success);
        assert_eq!(results.summary.successful_benchmarks, 1);
    }

    #[tokio::test]
    async fn test_failed_benchmark_handling() {
        let config = EvalConfig::default();
        let mut suite = BenchmarkSuite::new(&config).unwrap();

        // Register a failing benchmark
        let benchmark = Box::new(MockBenchmark::new(
            "failing_benchmark",
            BenchmarkCategory::Language,
            true,
        ));
        suite.register_benchmark(benchmark);

        // Run the category
        let results = suite.run_language_suite().await;
        assert!(results.is_ok());

        let results = results.unwrap();
        assert_eq!(results.results.len(), 1);
        assert!(!results.results[0].success);
        assert_eq!(results.summary.failed_benchmarks, 1);
    }
}