use std::collections::VecDeque;
use parking_lot::Mutex;
use crate::lifecycle::registry::AccuracyMetrics;
#[inline]
fn approx_eq(a: f64, b: f64) -> bool {
let diff = (a - b).abs();
let magnitude = a.abs().max(b.abs());
diff <= 1e-9 * magnitude.max(1e-3)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum PromotionCriteria {
LowerMape,
LowerRmse,
LowerBoth,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ABTestConfig {
pub challenger_traffic_pct: f64,
pub evaluation_windows: usize,
pub criteria: PromotionCriteria,
}
impl Default for ABTestConfig {
fn default() -> Self {
Self {
challenger_traffic_pct: 0.10,
evaluation_windows: 10,
criteria: PromotionCriteria::LowerMape,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ABTestResult {
Promote,
Keep,
Inconclusive,
}
#[derive(Debug, Clone)]
struct PredictionRecord {
actual: f64,
champion_prediction: f64,
challenger_prediction: f64,
}
#[derive(Debug)]
pub struct ABTestEvaluator {
config: ABTestConfig,
inner: Mutex<ABTestInner>,
}
#[derive(Debug)]
struct ABTestInner {
records: VecDeque<PredictionRecord>,
consecutive_wins: usize,
last_eval_count: usize,
}
impl ABTestEvaluator {
pub fn new(config: ABTestConfig) -> Self {
Self {
config,
inner: Mutex::new(ABTestInner {
records: VecDeque::new(),
consecutive_wins: 0,
last_eval_count: 0,
}),
}
}
pub fn should_use_challenger(&self, request_id: u64) -> bool {
let frac = (request_id % 10000) as f64 / 10000.0;
frac < self.config.challenger_traffic_pct
}
pub fn log_prediction(&self, actual: f64, champion_pred: f64, challenger_pred: f64) {
let mut inner = self.inner.lock();
inner.records.push_back(PredictionRecord {
actual,
champion_prediction: champion_pred,
challenger_prediction: challenger_pred,
});
let max_records = self.config.evaluation_windows * 100;
while inner.records.len() > max_records {
inner.records.pop_front();
}
}
pub fn evaluate(&self) -> ABTestResult {
let mut inner = self.inner.lock();
if inner.records.len() < self.config.evaluation_windows {
return ABTestResult::Inconclusive;
}
let current_count = inner.records.len();
if current_count == inner.last_eval_count {
return if inner.consecutive_wins >= 3 {
ABTestResult::Promote
} else if current_count >= self.config.evaluation_windows * 5
&& inner.consecutive_wins == 0
{
ABTestResult::Keep
} else {
ABTestResult::Inconclusive
};
}
inner.last_eval_count = current_count;
let window = self.config.evaluation_windows;
let recent = inner.records.iter().rev().take(window);
let actuals: Vec<f64> = recent.clone().map(|r| r.actual).collect();
let champ_preds: Vec<f64> = recent.clone().map(|r| r.champion_prediction).collect();
let chal_preds: Vec<f64> = recent.map(|r| r.challenger_prediction).collect();
let champ_metrics = AccuracyMetrics::compute(&actuals, &champ_preds);
let chal_metrics = AccuracyMetrics::compute(&actuals, &chal_preds);
metrics::gauge!("chronix_model_ab_test_champion_mape").set(champ_metrics.mape);
metrics::gauge!("chronix_model_ab_test_challenger_mape").set(chal_metrics.mape);
let challenger_wins = match self.config.criteria {
PromotionCriteria::LowerMape => {
!chal_metrics.mape.is_nan()
&& !champ_metrics.mape.is_nan()
&& chal_metrics.mape < champ_metrics.mape
}
PromotionCriteria::LowerRmse => {
!chal_metrics.rmse.is_nan()
&& !champ_metrics.rmse.is_nan()
&& chal_metrics.rmse < champ_metrics.rmse
}
PromotionCriteria::LowerBoth => {
chal_metrics.rmse < champ_metrics.rmse
&& !chal_metrics.mape.is_nan()
&& !champ_metrics.mape.is_nan()
&& chal_metrics.mape < champ_metrics.mape
}
};
let is_tie = match self.config.criteria {
PromotionCriteria::LowerMape => {
!chal_metrics.mape.is_nan()
&& !champ_metrics.mape.is_nan()
&& approx_eq(chal_metrics.mape, champ_metrics.mape)
}
PromotionCriteria::LowerRmse => {
!chal_metrics.rmse.is_nan()
&& !champ_metrics.rmse.is_nan()
&& approx_eq(chal_metrics.rmse, champ_metrics.rmse)
}
PromotionCriteria::LowerBoth => {
!chal_metrics.mape.is_nan()
&& !champ_metrics.mape.is_nan()
&& approx_eq(chal_metrics.mape, champ_metrics.mape)
&& !chal_metrics.rmse.is_nan()
&& !champ_metrics.rmse.is_nan()
&& approx_eq(chal_metrics.rmse, champ_metrics.rmse)
}
};
if challenger_wins {
inner.consecutive_wins += 1;
} else if is_tie {
} else {
inner.consecutive_wins = 0;
}
if inner.consecutive_wins >= 3 {
ABTestResult::Promote
} else if inner.records.len() >= self.config.evaluation_windows * 5
&& inner.consecutive_wins == 0
{
ABTestResult::Keep
} else {
ABTestResult::Inconclusive
}
}
pub fn champion_metrics(&self) -> Option<AccuracyMetrics> {
let inner = self.inner.lock();
if inner.records.is_empty() {
return None;
}
let actuals: Vec<f64> = inner.records.iter().map(|r| r.actual).collect();
let preds: Vec<f64> = inner
.records
.iter()
.map(|r| r.champion_prediction)
.collect();
Some(AccuracyMetrics::compute(&actuals, &preds))
}
pub fn challenger_metrics(&self) -> Option<AccuracyMetrics> {
let inner = self.inner.lock();
if inner.records.is_empty() {
return None;
}
let actuals: Vec<f64> = inner.records.iter().map(|r| r.actual).collect();
let preds: Vec<f64> = inner
.records
.iter()
.map(|r| r.challenger_prediction)
.collect();
Some(AccuracyMetrics::compute(&actuals, &preds))
}
pub fn record_count(&self) -> usize {
self.inner.lock().records.len()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_routing_fraction() {
let eval = ABTestEvaluator::new(ABTestConfig {
challenger_traffic_pct: 0.10,
..Default::default()
});
let mut challenger_count = 0;
for id in 0..10000 {
if eval.should_use_challenger(id) {
challenger_count += 1;
}
}
assert!((900..=1100).contains(&challenger_count));
}
#[test]
fn test_inconclusive_with_few_records() {
let eval = ABTestEvaluator::new(ABTestConfig {
evaluation_windows: 10,
..Default::default()
});
for i in 0..5 {
eval.log_prediction(100.0, 110.0 + i as f64, 105.0 + i as f64);
}
assert_eq!(eval.evaluate(), ABTestResult::Inconclusive);
}
#[test]
fn test_promote_when_challenger_consistently_better() {
let eval = ABTestEvaluator::new(ABTestConfig {
evaluation_windows: 5,
criteria: PromotionCriteria::LowerMape,
..Default::default()
});
for i in 0..50 {
let actual = 100.0 + (i as f64 * 0.1).sin();
let champ_pred = actual + 10.0; let chal_pred = actual + 1.0; eval.log_prediction(actual, champ_pred, chal_pred);
}
for batch in 0..3 {
for i in 0..5 {
let actual = 100.0 + ((50 + batch * 5 + i) as f64 * 0.1).sin();
let champ_pred = actual + 10.0;
let chal_pred = actual + 1.0;
eval.log_prediction(actual, champ_pred, chal_pred);
}
eval.evaluate();
}
let result = eval.evaluate();
assert_eq!(result, ABTestResult::Promote);
}
#[test]
fn test_keep_when_champion_wins() {
let eval = ABTestEvaluator::new(ABTestConfig {
evaluation_windows: 5,
criteria: PromotionCriteria::LowerRmse,
..Default::default()
});
for i in 0..50 {
let actual = 100.0 + (i as f64 * 0.1).sin();
let champ_pred = actual + 0.5;
let chal_pred = actual + 20.0;
eval.log_prediction(actual, champ_pred, chal_pred);
}
let result = eval.evaluate();
assert_eq!(result, ABTestResult::Keep);
}
#[test]
fn test_metrics_available() {
let eval = ABTestEvaluator::new(Default::default());
eval.log_prediction(100.0, 110.0, 105.0);
eval.log_prediction(200.0, 190.0, 195.0);
let cm = eval.champion_metrics().unwrap();
let clm = eval.challenger_metrics().unwrap();
assert!(cm.mae > 0.0);
assert!(clm.mae > 0.0);
assert!(clm.mae < cm.mae); }
#[test]
fn test_record_count() {
let eval = ABTestEvaluator::new(Default::default());
assert_eq!(eval.record_count(), 0);
eval.log_prediction(1.0, 1.0, 1.0);
assert_eq!(eval.record_count(), 1);
}
#[test]
fn test_ab_test_tie_preserves_streak() {
let eval = ABTestEvaluator::new(ABTestConfig {
evaluation_windows: 5,
criteria: PromotionCriteria::LowerMape,
..Default::default()
});
for _ in 0..10 {
let actual = 100.0;
let champ_pred = 110.0; let chal_pred = 101.0; eval.log_prediction(actual, champ_pred, chal_pred);
}
let r1 = eval.evaluate();
assert_eq!(r1, ABTestResult::Inconclusive);
for _ in 0..5 {
eval.log_prediction(100.0, 110.0, 101.0);
}
let r2 = eval.evaluate();
assert_eq!(r2, ABTestResult::Inconclusive);
for _ in 0..5 {
eval.log_prediction(100.0, 105.0, 105.0);
}
let r3 = eval.evaluate();
assert_eq!(r3, ABTestResult::Inconclusive);
for _ in 0..5 {
eval.log_prediction(100.0, 110.0, 101.0);
}
let r4 = eval.evaluate();
assert_eq!(r4, ABTestResult::Promote);
}
}