//! A/B Testing Framework for ML Model Deployments //! //! This module provides statistical A/B testing capabilities for comparing //! model performance with traffic splitting and significance testing. use std::collections::HashMap; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant, SystemTime}; use async_trait::async_trait; use serde::{Deserialize, Serialize}; use tokio::sync::{RwLock, Mutex}; use uuid::Uuid; use crate::{MLError, MLResult, ModelType, Features, ModelPrediction, MLModel}; use super::{ModelVersion, DeploymentStatus}; /// A/B test configuration #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ABTestConfig { /// Test name/identifier pub test_name: String, /// Description of the test pub description: String, /// Control group percentage (0.0 to 1.0) pub control_percentage: f32, /// Treatment group percentage (0.0 to 1.0) pub treatment_percentage: f32, /// Minimum sample size per group pub min_sample_size: u32, /// Statistical significance threshold (e.g., 0.05 for 95% confidence) pub significance_threshold: f64, /// Test duration limit pub max_duration: Duration, /// Metrics to track for comparison pub tracked_metrics: Vec, /// Early stopping criteria pub early_stopping: Option, /// Traffic splitting strategy pub splitting_strategy: TrafficSplittingStrategy, } impl Default for ABTestConfig { fn default() -> Self { Self { test_name: "default_test".to_string(), description: "Default A/B test configuration".to_string(), control_percentage: 0.5, treatment_percentage: 0.5, min_sample_size: 1000, significance_threshold: 0.05, max_duration: Duration::from_hours(24), tracked_metrics: vec!["latency".to_string(), "accuracy".to_string(), "error_rate".to_string()], early_stopping: Some(EarlyStoppingConfig::default()), splitting_strategy: TrafficSplittingStrategy::HashBased, } } } /// Early stopping configuration for A/B tests #[derive(Debug, Clone, Serialize, Deserialize)] pub struct EarlyStoppingConfig { /// Check interval for early stopping pub check_interval: Duration, /// Minimum samples before early stopping is considered pub min_samples_for_early_stop: u32, /// Maximum degradation allowed before stopping (0.0 to 1.0) pub max_degradation_threshold: f64, /// Statistical power threshold for early stopping pub statistical_power_threshold: f64, } impl Default for EarlyStoppingConfig { fn default() -> Self { Self { check_interval: Duration::from_minutes(5), min_samples_for_early_stop: 100, max_degradation_threshold: 0.1, // 10% degradation statistical_power_threshold: 0.8, // 80% power } } } /// Traffic splitting strategies #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum TrafficSplittingStrategy { /// Hash-based splitting using feature hash HashBased, /// Random splitting Random, /// Round-robin splitting RoundRobin, /// Weighted random splitting WeightedRandom, } /// A/B test status #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub enum ABTestStatus { /// Test is being set up Initializing, /// Test is running Running, /// Test is paused Paused, /// Test completed successfully Completed, /// Test was stopped early due to significance StoppedEarly, /// Test was stopped due to degradation StoppedDegradation, /// Test failed Failed, } /// A/B test experiment #[derive(Debug, Clone)] pub struct ABTestExperiment { /// Unique experiment ID pub experiment_id: Uuid, /// Test configuration pub config: ABTestConfig, /// Control model pub control_model: Arc, /// Treatment model pub treatment_model: Arc, /// Test status pub status: ABTestStatus, /// Start time pub start_time: SystemTime, /// End time (if completed) pub end_time: Option, /// Traffic splitter pub traffic_splitter: Arc, /// Metrics collector pub metrics_collector: Arc, } impl ABTestExperiment { /// Create new A/B test experiment pub fn new( config: ABTestConfig, control_model: Arc, treatment_model: Arc, ) -> Self { let experiment_id = Uuid::new_v4(); let traffic_splitter = Arc::new(TrafficSplitter::new( config.control_percentage, config.treatment_percentage, config.splitting_strategy, )); let metrics_collector = Arc::new(ABTestMetricsCollector::new( experiment_id, config.tracked_metrics.clone(), )); Self { experiment_id, config, control_model, treatment_model, status: ABTestStatus::Initializing, start_time: SystemTime::now(), end_time: None, traffic_splitter, metrics_collector, } } /// Start the A/B test pub async fn start(&mut self) -> MLResult<()> { if self.status != ABTestStatus::Initializing { return Err(MLError::ValidationError { message: "Test can only be started from Initializing status".to_string(), }); } self.status = ABTestStatus::Running; self.start_time = SystemTime::now(); tracing::info!( "Started A/B test {} with control model {} and treatment model {}", self.config.test_name, self.control_model.name(), self.treatment_model.name() ); Ok(()) } /// Perform prediction with A/B test traffic splitting pub async fn predict(&self, features: &Features) -> MLResult { if self.status != ABTestStatus::Running { return Err(MLError::ValidationError { message: "Test is not running".to_string(), }); } let start_time = Instant::now(); // Determine which model to use let group = self.traffic_splitter.assign_group(features); let (model, group_name) = match group { TestGroup::Control => (&self.control_model, "control"), TestGroup::Treatment => (&self.treatment_model, "treatment"), }; // Make prediction let prediction = model.predict(features).await?; let latency = start_time.elapsed(); // Record metrics self.metrics_collector.record_prediction( group, &prediction, latency, ).await?; Ok(ABTestPrediction { prediction, assigned_group: group, model_name: model.name().to_string(), latency, experiment_id: self.experiment_id, }) } /// Get current test results pub async fn get_results(&self) -> ABTestResults { self.metrics_collector.get_results().await } /// Check if test should be stopped early pub async fn check_early_stopping(&mut self) -> MLResult { if let Some(ref early_config) = self.config.early_stopping { let results = self.get_results().await; // Check minimum samples if results.control_metrics.sample_count < early_config.min_samples_for_early_stop || results.treatment_metrics.sample_count < early_config.min_samples_for_early_stop { return Ok(false); } // Check for statistical significance if let Some(significance) = self.calculate_statistical_significance(&results).await? { if significance.p_value < self.config.significance_threshold { self.status = ABTestStatus::StoppedEarly; self.end_time = Some(SystemTime::now()); tracing::info!( "A/B test {} stopped early due to statistical significance (p-value: {})", self.config.test_name, significance.p_value ); return Ok(true); } } // Check for degradation if self.check_degradation_threshold(&results, early_config.max_degradation_threshold).await? { self.status = ABTestStatus::StoppedDegradation; self.end_time = Some(SystemTime::now()); tracing::warn!( "A/B test {} stopped due to performance degradation", self.config.test_name ); return Ok(true); } } Ok(false) } /// Calculate statistical significance between groups async fn calculate_statistical_significance(&self, results: &ABTestResults) -> MLResult> { // For latency comparison (continuous metric) if let (Some(control_latency), Some(treatment_latency)) = (results.control_metrics.avg_latency, results.treatment_metrics.avg_latency) { let t_stat = self.calculate_t_statistic( control_latency, treatment_latency, results.control_metrics.latency_std_dev.unwrap_or(0.0), results.treatment_metrics.latency_std_dev.unwrap_or(0.0), results.control_metrics.sample_count as f64, results.treatment_metrics.sample_count as f64, ); let p_value = self.calculate_p_value(t_stat, (results.control_metrics.sample_count + results.treatment_metrics.sample_count - 2) as f64); return Ok(Some(StatisticalSignificance { metric_name: "latency".to_string(), t_statistic: t_stat, p_value, confidence_interval: self.calculate_confidence_interval( control_latency, treatment_latency, results.control_metrics.latency_std_dev.unwrap_or(0.0), results.treatment_metrics.latency_std_dev.unwrap_or(0.0), results.control_metrics.sample_count as f64, results.treatment_metrics.sample_count as f64, ), effect_size: (treatment_latency - control_latency) / control_latency, })); } Ok(None) } /// Calculate t-statistic for two-sample t-test fn calculate_t_statistic(&self, mean1: f64, mean2: f64, std1: f64, std2: f64, n1: f64, n2: f64) -> f64 { let pooled_variance = ((n1 - 1.0) * std1.powi(2) + (n2 - 1.0) * std2.powi(2)) / (n1 + n2 - 2.0); let standard_error = (pooled_variance * (1.0 / n1 + 1.0 / n2)).sqrt(); if standard_error == 0.0 { 0.0 } else { (mean1 - mean2) / standard_error } } /// Calculate p-value from t-statistic (simplified approximation) fn calculate_p_value(&self, t_stat: f64, degrees_of_freedom: f64) -> f64 { // Simplified p-value calculation using normal approximation // In production, you'd use a proper statistical library let abs_t = t_stat.abs(); if degrees_of_freedom > 30.0 { // Normal approximation for large samples 2.0 * (1.0 - self.normal_cdf(abs_t)) } else { // Conservative estimate for small samples if abs_t > 2.0 { 0.05 } else { 0.1 } } } /// Normal cumulative distribution function (approximation) fn normal_cdf(&self, x: f64) -> f64 { 0.5 * (1.0 + erf(x / 2.0_f64.sqrt())) } /// Calculate confidence interval fn calculate_confidence_interval(&self, mean1: f64, mean2: f64, std1: f64, std2: f64, n1: f64, n2: f64) -> (f64, f64) { let diff = mean2 - mean1; let pooled_variance = ((n1 - 1.0) * std1.powi(2) + (n2 - 1.0) * std2.powi(2)) / (n1 + n2 - 2.0); let standard_error = (pooled_variance * (1.0 / n1 + 1.0 / n2)).sqrt(); let margin_of_error = 1.96 * standard_error; // 95% confidence (diff - margin_of_error, diff + margin_of_error) } /// Check if treatment shows significant degradation async fn check_degradation_threshold(&self, results: &ABTestResults, threshold: f64) -> MLResult { // Check latency degradation if let (Some(control_latency), Some(treatment_latency)) = (results.control_metrics.avg_latency, results.treatment_metrics.avg_latency) { let degradation = (treatment_latency - control_latency) / control_latency; if degradation > threshold { return Ok(true); } } // Check error rate degradation if let (Some(control_error), Some(treatment_error)) = (results.control_metrics.error_rate, results.treatment_metrics.error_rate) { let degradation = (treatment_error - control_error) / control_error.max(0.001); // Avoid division by zero if degradation > threshold { return Ok(true); } } Ok(false) } /// Stop the test pub async fn stop(&mut self) -> MLResult { if self.status != ABTestStatus::Running { return Err(MLError::ValidationError { message: "Test is not running".to_string(), }); } self.status = ABTestStatus::Completed; self.end_time = Some(SystemTime::now()); let results = self.get_results().await; tracing::info!( "A/B test {} completed with {} control samples and {} treatment samples", self.config.test_name, results.control_metrics.sample_count, results.treatment_metrics.sample_count ); Ok(results) } } /// Test group assignment #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum TestGroup { /// Control group (original model) Control, /// Treatment group (new model) Treatment, } /// Traffic splitter for A/B testing pub struct TrafficSplitter { /// Control group percentage control_percentage: f32, /// Treatment group percentage treatment_percentage: f32, /// Splitting strategy strategy: TrafficSplittingStrategy, /// Round-robin counter round_robin_counter: AtomicU64, } impl TrafficSplitter { /// Create new traffic splitter pub fn new( control_percentage: f32, treatment_percentage: f32, strategy: TrafficSplittingStrategy, ) -> Self { Self { control_percentage, treatment_percentage, strategy, round_robin_counter: AtomicU64::new(0), } } /// Assign test group for given features pub fn assign_group(&self, features: &Features) -> TestGroup { match self.strategy { TrafficSplittingStrategy::HashBased => self.hash_based_assignment(features), TrafficSplittingStrategy::Random => self.random_assignment(), TrafficSplittingStrategy::RoundRobin => self.round_robin_assignment(), TrafficSplittingStrategy::WeightedRandom => self.weighted_random_assignment(), } } /// Hash-based assignment using feature hash fn hash_based_assignment(&self, features: &Features) -> TestGroup { use std::collections::hash_map::DefaultHasher; use std::hash::{Hash, Hasher}; let mut hasher = DefaultHasher::new(); // Hash feature values for value in &features.values { value.to_bits().hash(&mut hasher); } // Hash timestamp for additional randomness features.timestamp.hash(&mut hasher); let hash_value = hasher.finish(); let normalized = (hash_value % 1000) as f32 / 1000.0; if normalized < self.control_percentage { TestGroup::Control } else { TestGroup::Treatment } } /// Random assignment fn random_assignment(&self) -> TestGroup { let random_value: f32 = rand::random(); if random_value < self.control_percentage { TestGroup::Control } else { TestGroup::Treatment } } /// Round-robin assignment fn round_robin_assignment(&self) -> TestGroup { let counter = self.round_robin_counter.fetch_add(1, Ordering::Relaxed); let normalized = (counter % 1000) as f32 / 1000.0; if normalized < self.control_percentage { TestGroup::Control } else { TestGroup::Treatment } } /// Weighted random assignment fn weighted_random_assignment(&self) -> TestGroup { // Similar to random but with more sophisticated weighting self.random_assignment() } } /// A/B test prediction result #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ABTestPrediction { /// Underlying model prediction pub prediction: ModelPrediction, /// Assigned test group pub assigned_group: TestGroup, /// Model name used pub model_name: String, /// Prediction latency pub latency: Duration, /// Experiment ID pub experiment_id: Uuid, } /// Metrics collector for A/B tests pub struct ABTestMetricsCollector { /// Experiment ID experiment_id: Uuid, /// Tracked metrics tracked_metrics: Vec, /// Control group metrics control_metrics: Arc>, /// Treatment group metrics treatment_metrics: Arc>, } /// Metrics for a test group #[derive(Debug, Clone, Default)] struct GroupMetrics { /// Total predictions sample_count: u32, /// Total errors error_count: u32, /// Latency measurements latencies: Vec, /// Accuracy scores accuracy_scores: Vec, /// Custom metrics custom_metrics: HashMap>, } impl ABTestMetricsCollector { /// Create new metrics collector pub fn new(experiment_id: Uuid, tracked_metrics: Vec) -> Self { Self { experiment_id, tracked_metrics, control_metrics: Arc::new(Mutex::new(GroupMetrics::default())), treatment_metrics: Arc::new(Mutex::new(GroupMetrics::default())), } } /// Record prediction metrics pub async fn record_prediction( &self, group: TestGroup, prediction: &ModelPrediction, latency: Duration, ) -> MLResult<()> { let metrics = match group { TestGroup::Control => &self.control_metrics, TestGroup::Treatment => &self.treatment_metrics, }; let mut group_metrics = metrics.lock().await; group_metrics.sample_count += 1; group_metrics.latencies.push(latency.as_micros() as f64); group_metrics.accuracy_scores.push(prediction.confidence); Ok(()) } /// Record error pub async fn record_error(&self, group: TestGroup) -> MLResult<()> { let metrics = match group { TestGroup::Control => &self.control_metrics, TestGroup::Treatment => &self.treatment_metrics, }; let mut group_metrics = metrics.lock().await; group_metrics.error_count += 1; Ok(()) } /// Get test results pub async fn get_results(&self) -> ABTestResults { let control_metrics = self.control_metrics.lock().await; let treatment_metrics = self.treatment_metrics.lock().await; ABTestResults { experiment_id: self.experiment_id, control_metrics: GroupMetricsSummary::from_group_metrics(&control_metrics), treatment_metrics: GroupMetricsSummary::from_group_metrics(&treatment_metrics), test_duration: SystemTime::now(), } } } /// Summary of group metrics #[derive(Debug, Clone, Serialize, Deserialize)] pub struct GroupMetricsSummary { /// Sample count pub sample_count: u32, /// Error count pub error_count: u32, /// Error rate pub error_rate: Option, /// Average latency pub avg_latency: Option, /// Latency standard deviation pub latency_std_dev: Option, /// Average accuracy pub avg_accuracy: Option, /// Accuracy standard deviation pub accuracy_std_dev: Option, /// Custom metrics pub custom_metrics: HashMap, } impl GroupMetricsSummary { /// Create summary from group metrics fn from_group_metrics(metrics: &GroupMetrics) -> Self { let error_rate = if metrics.sample_count > 0 { Some(metrics.error_count as f64 / metrics.sample_count as f64) } else { None }; let (avg_latency, latency_std_dev) = if !metrics.latencies.is_empty() { let avg = metrics.latencies.iter().sum::() / metrics.latencies.len() as f64; let variance = metrics.latencies.iter() .map(|x| (x - avg).powi(2)) .sum::() / metrics.latencies.len() as f64; (Some(avg), Some(variance.sqrt())) } else { (None, None) }; let (avg_accuracy, accuracy_std_dev) = if !metrics.accuracy_scores.is_empty() { let avg = metrics.accuracy_scores.iter().sum::() / metrics.accuracy_scores.len() as f64; let variance = metrics.accuracy_scores.iter() .map(|x| (x - avg).powi(2)) .sum::() / metrics.accuracy_scores.len() as f64; (Some(avg), Some(variance.sqrt())) } else { (None, None) }; Self { sample_count: metrics.sample_count, error_count: metrics.error_count, error_rate, avg_latency, latency_std_dev, avg_accuracy, accuracy_std_dev, custom_metrics: HashMap::new(), // Would be populated from metrics.custom_metrics } } } /// A/B test results #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ABTestResults { /// Experiment ID pub experiment_id: Uuid, /// Control group results pub control_metrics: GroupMetricsSummary, /// Treatment group results pub treatment_metrics: GroupMetricsSummary, /// Test duration pub test_duration: SystemTime, } /// Statistical significance result #[derive(Debug, Clone, Serialize, Deserialize)] pub struct StatisticalSignificance { /// Metric name pub metric_name: String, /// T-statistic pub t_statistic: f64, /// P-value pub p_value: f64, /// Confidence interval (lower, upper) pub confidence_interval: (f64, f64), /// Effect size pub effect_size: f64, } /// Error function approximation for normal CDF fn erf(x: f64) -> f64 { // Abramowitz and Stegun approximation let a1 = 0.254829592; let a2 = -0.284496736; let a3 = 1.421413741; let a4 = -1.453152027; let a5 = 1.061405429; let p = 0.3275911; let sign = if x < 0.0 { -1.0 } else { 1.0 }; let x = x.abs(); let t = 1.0 / (1.0 + p * x); let y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * (-x * x).exp(); sign * y } #[cfg(test)] mod tests { use super::*; use crate::model_factory; #[test] fn test_traffic_splitter_creation() { let splitter = TrafficSplitter::new(0.5, 0.5, TrafficSplittingStrategy::Random); // Test that assignments are roughly balanced over many iterations let mut control_count = 0; let mut treatment_count = 0; for _ in 0..1000 { let features = Features::new(vec![1.0, 2.0], vec!["f1".to_string(), "f2".to_string()]); match splitter.assign_group(&features) { TestGroup::Control => control_count += 1, TestGroup::Treatment => treatment_count += 1, } } // Should be roughly balanced (within 20% of expected) let total = control_count + treatment_count; let control_ratio = control_count as f64 / total as f64; assert!(control_ratio > 0.3 && control_ratio < 0.7); } #[test] fn test_group_metrics_summary() { let mut metrics = GroupMetrics::default(); metrics.sample_count = 100; metrics.error_count = 5; metrics.latencies = vec![100.0, 200.0, 150.0, 175.0, 125.0]; metrics.accuracy_scores = vec![0.8, 0.85, 0.9, 0.75, 0.88]; let summary = GroupMetricsSummary::from_group_metrics(&metrics); assert_eq!(summary.sample_count, 100); assert_eq!(summary.error_count, 5); assert_eq!(summary.error_rate, Some(0.05)); assert!(summary.avg_latency.is_some()); assert!(summary.avg_accuracy.is_some()); } #[tokio::test] async fn test_ab_test_experiment_creation() { let control_model = Arc::from(model_factory::create_dqn_wrapper().unwrap()); let treatment_model = Arc::from(model_factory::create_dqn_wrapper().unwrap()); let config = ABTestConfig::default(); let experiment = ABTestExperiment::new(config, control_model, treatment_model); assert_eq!(experiment.status, ABTestStatus::Initializing); assert!(experiment.end_time.is_none()); } #[test] fn test_statistical_calculations() { let experiment = ABTestExperiment::new( ABTestConfig::default(), Arc::from(model_factory::create_dqn_wrapper().unwrap()), Arc::from(model_factory::create_dqn_wrapper().unwrap()), ); // Test t-statistic calculation let t_stat = experiment.calculate_t_statistic(100.0, 110.0, 10.0, 12.0, 50.0, 50.0); assert!(t_stat.is_finite()); // Test confidence interval calculation let (lower, upper) = experiment.calculate_confidence_interval(100.0, 110.0, 10.0, 12.0, 50.0, 50.0); assert!(lower < upper); } }