//! Prediction validation and model poisoning detection //! //! Provides statistical bounds checking and outlier detection to identify //! poisoned or adversarial models in the ML inference pipeline. use std::collections::VecDeque; use std::time::{Duration, Instant}; use serde::{Deserialize, Serialize}; use tokio::sync::RwLock; use tracing::{debug, warn}; use ml_core::MLError; /// Prediction validator with statistical bounds checking /// /// Detects model poisoning through: /// 1. Range validation (predictions must be in [-1.0, 1.0]) /// 2. Z-score outlier detection (>3sigma from mean) /// 3. Confidence sanity checks /// 4. Extreme prediction rate limiting /// /// # Performance /// - Validation: ~5us per prediction /// - Statistics update: ~2us pub struct PredictionValidator { /// Historical prediction statistics (rolling window) stats: RwLock, /// Configuration config: ValidationConfig, /// Extreme prediction tracking (for rate limiting) extreme_tracker: RwLock, } impl std::fmt::Debug for PredictionValidator { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("PredictionValidator") .field("config", &self.config) .field("stats", &format_args!(">")) .field("extreme_tracker", &format_args!(">")) .finish() } } /// Configuration for prediction validation #[derive(Debug, Clone)] pub struct ValidationConfig { /// Z-score threshold for outlier detection (default: 3.0) pub z_score_threshold: f64, /// Minimum confidence threshold (default: 0.5) pub confidence_threshold: f64, /// Extreme prediction threshold (default: 0.9) pub extreme_prediction_threshold: f64, /// Maximum allowed extreme prediction rate (default: 0.1 = 10%) pub max_extreme_rate: f64, /// Window duration for rate limiting (default: 60 seconds) pub window_duration: Duration, /// Bootstrap sample count (default: 1000) pub bootstrap_samples: u64, /// Exponential moving average weight (default: 0.05) pub ema_alpha: f64, } impl Default for ValidationConfig { fn default() -> Self { Self { z_score_threshold: 3.0, confidence_threshold: 0.5, extreme_prediction_threshold: 0.9, max_extreme_rate: 0.1, window_duration: Duration::from_secs(60), bootstrap_samples: 1000, ema_alpha: 0.05, } } } /// Historical prediction statistics #[derive(Debug, Clone)] struct PredictionStatistics { /// Exponential moving average of predictions mean: f64, /// Standard deviation (Welford's online algorithm) std_dev: f64, /// Variance (for std_dev calculation) variance: f64, /// Total number of samples processed sample_count: u64, /// Last update timestamp last_updated: Instant, /// Sum of squared differences (for variance) m2: f64, } impl PredictionStatistics { fn new() -> Self { Self { mean: 0.0, std_dev: 0.3, // Conservative bootstrap value variance: 0.09, sample_count: 0, last_updated: Instant::now(), m2: 0.0, } } /// Update statistics with new prediction (Welford's algorithm) fn update(&mut self, prediction: f64, ema_alpha: f64, bootstrap_samples: u64) { self.sample_count += 1; if self.sample_count <= bootstrap_samples { // Bootstrap phase: Use Welford's algorithm for stable variance let delta = prediction - self.mean; self.mean += delta / self.sample_count as f64; let delta2 = prediction - self.mean; self.m2 += delta * delta2; if self.sample_count > 1 { self.variance = self.m2 / (self.sample_count - 1) as f64; self.std_dev = self.variance.sqrt().max(0.01); // Minimum 0.01 to avoid division by zero } } else { // Production phase: Use exponential moving average self.mean = (1.0 - ema_alpha) * self.mean + ema_alpha * prediction; // Update variance with EMA let delta = prediction - self.mean; self.variance = (1.0 - ema_alpha) * self.variance + ema_alpha * delta * delta; self.std_dev = self.variance.sqrt().max(0.01); } self.last_updated = Instant::now(); } } /// Extreme prediction tracker for rate limiting #[derive(Debug)] struct ExtremeTracker { /// Recent extreme prediction timestamps timestamps: VecDeque, } impl ExtremeTracker { fn new() -> Self { Self { timestamps: VecDeque::new(), } } /// Record an extreme prediction fn record_extreme(&mut self, now: Instant, window_duration: Duration) { // Remove old timestamps outside window while let Some(oldest) = self.timestamps.front() { if now.duration_since(*oldest) > window_duration { self.timestamps.pop_front(); } else { break; } } self.timestamps.push_back(now); } /// Calculate current extreme prediction rate fn calculate_rate(&self, window_duration: Duration) -> f64 { let now = Instant::now(); let count = self .timestamps .iter() .filter(|&&ts| now.duration_since(ts) <= window_duration) .count(); // Estimate total predictions based on typical inference rate (~1000/sec) let window_secs = window_duration.as_secs_f64(); let estimated_total = (window_secs * 1000.0).max(1.0); count as f64 / estimated_total } } /// Validated prediction result #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ValidatedPrediction { /// Original prediction value pub value: f64, /// Whether prediction is an outlier pub is_outlier: bool, /// Z-score (standard deviations from mean) pub z_score: f64, /// Whether to override with ensemble fallback pub should_override: bool, /// Validation flags pub validation_flags: Vec, } /// Validation flags for different failure modes #[derive(Debug, Clone, Serialize, Deserialize)] pub enum ValidationFlag { /// Prediction outside valid range [-1.0, 1.0] OutOfBounds { value: f64 }, /// Prediction is statistical outlier Outlier { z_score: f64 }, /// Confidence below threshold LowConfidence { confidence: f64 }, /// Extreme prediction rate exceeded ExtremeRateLimit { rate: f64 }, } impl PredictionValidator { /// Create a new prediction validator with default configuration pub fn new() -> Self { Self::with_config(ValidationConfig::default()) } /// Create a new prediction validator with custom configuration pub fn with_config(config: ValidationConfig) -> Self { Self { stats: RwLock::new(PredictionStatistics::new()), config, extreme_tracker: RwLock::new(ExtremeTracker::new()), } } /// Validate prediction with multi-layer checks /// /// # Arguments /// * `prediction` - Prediction value to validate /// * `confidence` - Model confidence (0.0-1.0) /// * `model_id` - Model identifier for logging /// /// # Returns /// ValidatedPrediction with flags and z-score /// /// # Errors /// Returns error if prediction is invalid (out of bounds, extreme rate exceeded) /// /// # Performance /// ~5us per call (includes statistics update) pub async fn validate( &self, prediction: f64, confidence: f64, model_id: &str, ) -> Result { let mut flags = Vec::new(); // Layer 1: Range check if !(-1.0..=1.0).contains(&prediction) { warn!( "Prediction out of bounds for model {}: {} (valid range: [-1.0, 1.0])", model_id, prediction ); flags.push(ValidationFlag::OutOfBounds { value: prediction }); return Err(MLError::ValidationError { message: format!( "Prediction {} out of bounds [-1.0, 1.0] for model {}", prediction, model_id ), }); } // Layer 2: Z-score outlier detection let stats = self.stats.read().await; let z_score = if stats.std_dev > 0.001 { (prediction - stats.mean) / stats.std_dev } else { 0.0 }; drop(stats); // Release read lock if z_score.abs() > self.config.z_score_threshold { warn!( "Outlier prediction detected for model {}: {} (z-score: {:.2}, mean: {:.3}, std: {:.3})", model_id, prediction, z_score, self.stats.read().await.mean, self.stats.read().await.std_dev ); flags.push(ValidationFlag::Outlier { z_score }); } // Layer 3: Confidence sanity check if confidence < self.config.confidence_threshold { warn!( "Low confidence prediction for model {}: {:.3}", model_id, confidence ); flags.push(ValidationFlag::LowConfidence { confidence }); } // Layer 4: Rate limit extreme predictions let is_extreme = prediction.abs() > self.config.extreme_prediction_threshold; if is_extreme { let now = Instant::now(); let mut tracker = self.extreme_tracker.write().await; tracker.record_extreme(now, self.config.window_duration); let extreme_rate = tracker.calculate_rate(self.config.window_duration); if extreme_rate > self.config.max_extreme_rate { warn!( "Extreme prediction rate exceeded for model {}: {:.3} (threshold: {:.3})", model_id, extreme_rate, self.config.max_extreme_rate ); flags.push(ValidationFlag::ExtremeRateLimit { rate: extreme_rate }); return Err(MLError::ValidationError { message: format!( "Too many extreme predictions from model {} - possible poisoning (rate: {:.3}, threshold: {:.3})", model_id, extreme_rate, self.config.max_extreme_rate ), }); } } // Update statistics (after validation to avoid corrupting stats with invalid data) self.update_statistics(prediction).await; let is_outlier = !flags.is_empty(); let should_override = z_score.abs() > self.config.z_score_threshold; debug!( "Validated prediction for {}: value={:.3}, z={:.2}, outlier={}, override={}", model_id, prediction, z_score, is_outlier, should_override ); Ok(ValidatedPrediction { value: prediction, is_outlier, z_score, should_override, validation_flags: flags, }) } /// Update prediction statistics with new value pub async fn update_statistics(&self, prediction: f64) { let mut stats = self.stats.write().await; stats.update( prediction, self.config.ema_alpha, self.config.bootstrap_samples, ); } /// Get current prediction statistics pub async fn get_statistics(&self) -> PredictionStats { let stats = self.stats.read().await; PredictionStats { mean: stats.mean, std_dev: stats.std_dev, sample_count: stats.sample_count, } } /// Reset statistics (useful for testing or model retraining) pub async fn reset_statistics(&self) { let mut stats = self.stats.write().await; *stats = PredictionStatistics::new(); let mut tracker = self.extreme_tracker.write().await; tracker.timestamps.clear(); } } impl Default for PredictionValidator { fn default() -> Self { Self::new() } } /// Public statistics snapshot #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PredictionStats { pub mean: f64, pub std_dev: f64, pub sample_count: u64, } #[cfg(test)] #[allow( clippy::let_underscore_must_use, clippy::assertions_on_result_states, clippy::field_reassign_with_default )] mod tests { use super::*; #[tokio::test] async fn test_validate_normal_prediction() { let validator = PredictionValidator::new(); // Add some bootstrap samples for i in 0..100 { let _ = validator .validate(0.0 + (i as f64 / 1000.0), 0.8, "test_model") .await; } // Validate normal prediction let result = validator.validate(0.05, 0.8, "test_model").await; assert!(result.is_ok()); let validated = result.unwrap(); assert!(!validated.should_override); assert!(validated.validation_flags.is_empty()); } #[tokio::test] async fn test_validate_out_of_bounds() { let validator = PredictionValidator::new(); // Test upper bound let result = validator.validate(1.5, 0.8, "test_model").await; assert!(result.is_err()); // Test lower bound let result = validator.validate(-1.5, 0.8, "test_model").await; assert!(result.is_err()); } #[tokio::test] async fn test_validate_outlier() { let validator = PredictionValidator::new(); // Build statistics with normal predictions around 0.0 for _ in 0..1000 { let _ = validator.validate(0.0, 0.8, "test_model").await; } // Inject extreme outlier let result = validator.validate(0.95, 0.8, "test_model").await.unwrap(); assert!(result.is_outlier); assert!(result.z_score.abs() > 3.0); assert!(result.should_override); } #[tokio::test] async fn test_low_confidence_flag() { let validator = PredictionValidator::new(); let result = validator.validate(0.5, 0.3, "test_model").await.unwrap(); // Should flag low confidence but not reject let has_low_conf = result .validation_flags .iter() .any(|f| matches!(f, ValidationFlag::LowConfidence { .. })); assert!(has_low_conf); } #[tokio::test] async fn test_extreme_rate_limiting() { let mut config = ValidationConfig::default(); config.max_extreme_rate = 0.05; // 5% config.window_duration = Duration::from_secs(1); let validator = PredictionValidator::with_config(config); // Inject many extreme predictions quickly for _ in 0..100 { let result = validator.validate(0.95, 0.8, "test_model").await; // Should eventually hit rate limit if result.is_err() { // Verify it's a rate limit error let err = result.unwrap_err(); assert!(err.to_string().contains("extreme predictions")); return; } tokio::time::sleep(Duration::from_millis(1)).await; } panic!("Rate limit should have been triggered"); } #[tokio::test] async fn test_statistics_update() { let validator = PredictionValidator::new(); // Add predictions for i in 0..100 { let value = (i as f64 / 100.0) - 0.5; // Range: -0.5 to 0.5 validator.update_statistics(value).await; } let stats = validator.get_statistics().await; assert!(stats.sample_count == 100); assert!(stats.mean.abs() < 0.1); // Should be close to 0 assert!(stats.std_dev > 0.0); } #[tokio::test] async fn test_reset_statistics() { let validator = PredictionValidator::new(); // Add some data for _ in 0..100 { validator.update_statistics(0.5).await; } let stats_before = validator.get_statistics().await; assert!(stats_before.sample_count == 100); // Reset validator.reset_statistics().await; let stats_after = validator.get_statistics().await; assert_eq!(stats_after.sample_count, 0); assert_eq!(stats_after.mean, 0.0); } #[tokio::test] async fn test_bootstrap_phase() { let mut config = ValidationConfig::default(); config.bootstrap_samples = 10; let validator = PredictionValidator::with_config(config); // Add bootstrap samples for i in 0..10 { validator.update_statistics(i as f64 / 10.0).await; } let stats = validator.get_statistics().await; assert!(stats.std_dev > 0.0); } }