diff --git a/eoc/eoc-rs/crates/eoc-route-learned/src/bandit.rs b/eoc/eoc-rs/crates/eoc-route-learned/src/bandit.rs index 4d3c093d6..cae5f7214 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/src/bandit.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/src/bandit.rs @@ -3,10 +3,20 @@ //! Two algorithms ship here: //! //! - **LinUCB** (Li et al. 2010) — contextual bandit with linear reward -//! models. Each arm (stage) keeps `A = X^T X + I` and `b = X^T r`. Pick -//! the arm that maximises `θ_a · x + α · √(x^T A^{-1} x)`. +//! models. Each arm (stage) keeps `A = X^T X + I` and `b = X^T r`, where +//! `r` is 1 when the stage resolved the query and 0 when it did not. //! - **Thompson Sampling** — Beta-Bernoulli on per-arm success rates. Cheap //! and surprisingly competitive when contexts are noisy. +//! +//! Both learn each stage's chance of resolving a query, and the cost of each +//! stage from the joules `observe` reports. Both pick the stage that is +//! cheapest to answer from in expectation +//! ([`cheapest_in_expectation`](crate::router::cheapest_in_expectation)), +//! not the arm with the best success rate: the neural stage always has the +//! best rate, and picking it skips every cheaper stage that would have +//! answered. Exploration enters through the success estimate, optimistic for +//! LinUCB (`θ_a · x + α · √(x^T A^{-1} x)`) and sampled for Thompson, so an +//! untried cheap stage looks likely to resolve and gets tried. use std::collections::{BTreeMap, HashMap}; @@ -18,7 +28,10 @@ use rand_distr::{Beta, Distribution}; use serde::{Deserialize, Serialize}; use crate::error::{Error, Result}; -use crate::router::{ALL_STAGES, LearnedRouter, RouterState, StagePrediction}; +use crate::router::{ + ALL_STAGES, LearnedRouter, RouterState, StagePrediction, blend_joules, cheapest_in_expectation, + default_stage_joules, stage_index, +}; /// Per-arm state for [`LinUcbRouter`]. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -52,7 +65,8 @@ impl LinUcbArm { } } - fn ucb(&self, x: &DVector, alpha: f32) -> f32 { + /// The mean reward estimate `θ · x` and its upper confidence bound. + fn estimate(&self, x: &DVector, alpha: f32) -> (f32, f32) { let theta = self.theta(); let mean = theta.dot(x); let a_inv = self @@ -61,7 +75,7 @@ impl LinUcbArm { .try_inverse() .unwrap_or_else(|| DMatrix::::identity(self.dim, self.dim)); let bonus_sq = (x.transpose() * &a_inv * x)[(0, 0)].max(0.0); - mean + alpha * bonus_sq.sqrt() + (mean, mean + alpha * bonus_sq.sqrt()) } fn update(&mut self, x: &DVector, reward: f32) { @@ -71,6 +85,10 @@ impl LinUcbArm { } /// LinUCB contextual bandit over the four cascade stages. +/// +/// The reward model is linear in the context with no intercept, as in Li et +/// al.: put a constant feature in the context for a stage's base rate to be +/// learnable. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct LinUcbRouter { /// Context (embedding) dimension. @@ -81,6 +99,9 @@ pub struct LinUcbRouter { pub ridge: f32, /// Per-stage arm state. pub arms: HashMap, + /// Per-stage cost estimate in microjoules, learned from `observe`. + #[serde(default = "default_stage_joules")] + pub joules: HashMap, } impl LinUcbRouter { @@ -95,11 +116,14 @@ impl LinUcbRouter { alpha, ridge, arms, + joules: default_stage_joules(), } } - /// Argmax UCB across stages. - pub fn pick(&self, embedding: &[f32]) -> Result<(Stage, f32)> { + /// Each stage's mean estimate and upper confidence bound that it resolves + /// the query, in [`ALL_STAGES`] order. A missing arm counts as a certain + /// miss. + fn estimates(&self, embedding: &[f32]) -> Result<([f32; 4], [f32; 4])> { if embedding.len() != self.dim { return Err(Error::DimensionMismatch { expected: self.dim, @@ -107,20 +131,27 @@ impl LinUcbRouter { }); } let x = DVector::::from_row_slice(embedding); - let mut best = Stage::Cache; - let mut best_score = f32::NEG_INFINITY; - for s in ALL_STAGES { - let arm = match self.arms.get(&s) { - Some(a) => a, - None => continue, - }; - let u = arm.ucb(&x, self.alpha); - if u > best_score { - best_score = u; - best = s; + let mut means = [0.0_f32; 4]; + let mut bounds = [0.0_f32; 4]; + for (i, s) in ALL_STAGES.iter().enumerate() { + if let Some(arm) = self.arms.get(s) { + (means[i], bounds[i]) = arm.estimate(&x, self.alpha); } } - Ok((best, best_score)) + Ok((means, bounds)) + } + + /// The stage cheapest to answer from in expectation, with every stage's + /// chance taken at its upper confidence bound, and that stage's bound. + pub fn pick(&self, embedding: &[f32]) -> Result<(Stage, f32)> { + let (_, bounds) = self.estimates(embedding)?; + let stage = cheapest_in_expectation(&bounds, &self.joules); + Ok((stage, bounds[stage_index(stage)])) + } + + /// Record an observed stage cost — EMA with coefficient 0.1. + pub fn observe_cost(&mut self, stage: Stage, microjoules: u64) { + blend_joules(&mut self.joules, stage, microjoules); } /// Update the chosen arm with the observed reward. @@ -145,13 +176,14 @@ impl LearnedRouter for LinUcbRouter { let Some(emb) = query.embedding.as_ref() else { return StagePrediction::new(Stage::Cache, 0.0); }; - match self.pick(emb) { - Ok((stage, score)) => { - // Squash UCB score into [0,1] via logistic for "confidence". - let conf = 1.0 / (1.0 + (-score).exp()); - StagePrediction::new(stage, conf) + match (self.pick(emb), self.estimates(emb)) { + // The confidence is the mean estimate, not the optimistic bound + // the pick used: the bound is how much the router is willing to + // explore, not how likely the stage is to resolve. + (Ok((stage, _)), Ok((means, _))) => { + StagePrediction::new(stage, means[stage_index(stage)]) } - Err(_) => StagePrediction::new(Stage::Cache, 0.0), + _ => StagePrediction::new(Stage::Cache, 0.0), } } @@ -160,8 +192,9 @@ impl LearnedRouter for LinUcbRouter { query: &Query, chosen_stage: Stage, was_correct: bool, - _joule_cost: u64, + joule_cost: u64, ) { + self.observe_cost(chosen_stage, joule_cost); if let Some(emb) = query.embedding.as_ref() { let reward = if was_correct { 1.0 } else { 0.0 }; let _ = self.update(emb, chosen_stage, reward); @@ -224,6 +257,9 @@ pub struct ThompsonSamplingRouter { pub arms: HashMap, /// RNG seed for reproducibility — incremented every sample. pub seed: u64, + /// Per-stage cost estimate in microjoules, learned from `observe`. + #[serde(default = "default_stage_joules")] + pub joules: HashMap, } impl ThompsonSamplingRouter { @@ -233,30 +269,43 @@ impl ThompsonSamplingRouter { for s in ALL_STAGES { arms.insert(s, BetaArm::new()); } - Self { arms, seed } + Self { + arms, + seed, + joules: default_stage_joules(), + } } - /// Sample once from each arm; return the argmax. Iteration order is the - /// fixed cheapest-first stage order, so the result is deterministic for - /// a given `seed`. + /// Sample each stage's chance of resolving once, and return the stage + /// cheapest to answer from in expectation under those samples, with its + /// sample. Sampling runs in the fixed cheapest-first stage order, so the + /// result is deterministic for a given `seed`. A missing arm counts as a + /// certain miss. pub fn pick(&self) -> (Stage, f32) { let mut rng = rand::rngs::StdRng::seed_from_u64(self.seed); - let mut best = Stage::Cache; - let mut best_v = f32::NEG_INFINITY; - for s in ALL_STAGES { - let arm = match self.arms.get(&s) { - Some(a) => a, - None => continue, + let mut samples = [0.0_f32; 4]; + for (i, s) in ALL_STAGES.iter().enumerate() { + let Some(arm) = self.arms.get(s) else { + continue; }; let dist = Beta::new(arm.alpha as f64, arm.beta as f64) .unwrap_or_else(|_| Beta::new(1.0, 1.0).expect("beta(1,1) ok")); - let v = dist.sample(&mut rng) as f32; - if v > best_v { - best_v = v; - best = s; - } + samples[i] = dist.sample(&mut rng) as f32; } - (best, best_v) + let stage = cheapest_in_expectation(&samples, &self.joules); + (stage, samples[stage_index(stage)]) + } + + /// The posterior mean chance that `stage` resolves a query. + pub fn mean(&self, stage: Stage) -> f32 { + self.arms + .get(&stage) + .map_or(0.0, |a| a.alpha / (a.alpha + a.beta)) + } + + /// Record an observed stage cost — EMA with coefficient 0.1. + pub fn observe_cost(&mut self, stage: Stage, microjoules: u64) { + blend_joules(&mut self.joules, stage, microjoules); } /// Update Beta posterior for `stage`. @@ -275,8 +324,11 @@ impl ThompsonSamplingRouter { #[async_trait] impl LearnedRouter for ThompsonSamplingRouter { async fn route(&self, _query: &Query) -> StagePrediction { - let (stage, v) = self.pick(); - StagePrediction::new(stage, v.clamp(0.0, 1.0)) + // The confidence is the posterior mean, not the sample the pick used: + // a sample is a draw for exploration, and would make the threshold + // policy's skip decision a coin toss. + let (stage, _) = self.pick(); + StagePrediction::new(stage, self.mean(stage)) } async fn observe( @@ -284,8 +336,9 @@ impl LearnedRouter for ThompsonSamplingRouter { _query: &Query, chosen_stage: Stage, was_correct: bool, - _joule_cost: u64, + joule_cost: u64, ) { + self.observe_cost(chosen_stage, joule_cost); self.update(chosen_stage, was_correct); } @@ -332,4 +385,39 @@ mod tests { assert_eq!(arm.alpha, 2.0); assert_eq!(arm.beta, 1.0); } + + /// `observe` is how a stage's cost reaches the router; it was ignored. + #[tokio::test] + async fn observing_a_stage_learns_its_cost() { + let q = Query::new("q").with_embedding(vec![1.0, 0.0]); + let mut lin = LinUcbRouter::new(2, 0.1, 1.0); + let mut ts = ThompsonSamplingRouter::new(0); + lin.observe(&q, Stage::Kv, false, 1_050_000).await; + ts.observe(&q, Stage::Kv, false, 1_050_000).await; + // A tenth of the way from the default 50 000. + assert_eq!(lin.joules[&Stage::Kv], 150_000); + assert_eq!(ts.joules[&Stage::Kv], 150_000); + assert_eq!(ts.arms[&Stage::Kv].beta, 2.0); + } + + /// The confidence handed to the threshold policy is the router's belief + /// that the stage resolves, not the bound or sample it explored with: an + /// untried router believes nothing yet. + #[tokio::test] + async fn confidence_is_the_mean_belief_not_the_exploration_bound() { + let q = Query::new("q").with_embedding(vec![1.0, 0.0]); + let lin = LinUcbRouter::new(2, 1.0, 1.0); + let (stage, bound) = lin.pick(&[1.0, 0.0]).unwrap(); + assert!(bound > 0.5, "an untried arm explores optimistically"); + let pred = lin.route(&q).await; + assert_eq!((pred.recommended, pred.confidence), (stage, 0.0)); + + let mut ts = ThompsonSamplingRouter::new(3); + for _ in 0..3 { + ts.update(Stage::Cache, true); + } + let pred = ts.route(&q).await; + assert_eq!(pred.confidence, ts.mean(pred.recommended)); + assert_eq!(ts.mean(Stage::Cache), 0.8); + } } diff --git a/eoc/eoc-rs/crates/eoc-route-learned/src/classifier.rs b/eoc/eoc-rs/crates/eoc-route-learned/src/classifier.rs index 48e432fe4..8883172ce 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/src/classifier.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/src/classifier.rs @@ -21,19 +21,10 @@ use rand_distr::{Distribution, Normal}; use serde::{Deserialize, Serialize}; use crate::error::{Error, Result}; -use crate::router::{ALL_STAGES, LearnedRouter, RouterState, StagePrediction}; +use crate::router::{ALL_STAGES, LearnedRouter, RouterState, StagePrediction, stage_index}; const NUM_STAGES: usize = 4; -fn stage_index(s: Stage) -> usize { - match s { - Stage::Cache => 0, - Stage::Kv => 1, - Stage::Graph => 2, - Stage::Neural => 3, - } -} - fn stage_from_index(i: usize) -> Stage { ALL_STAGES[i.min(3)] } @@ -179,8 +170,16 @@ impl LearnedRouter for LogRegRouter { was_correct: bool, _joule_cost: u64, ) { - // Only train on successful resolutions — we want P(stage = correct | x). - if was_correct && let Some(emb) = query.embedding.as_ref() { + // The classifier predicts the cheapest stage that resolves a query. + // A stage that resolved it shows that stage is enough, not that it is + // the cheapest that is, so it is evidence only when it is no costlier + // than the stage already predicted. Otherwise a correct Neural answer, + // which every query gets, would teach it to predict Neural for all. + if was_correct + && let Some(emb) = query.embedding.as_ref() + && let Ok((predicted, _, _)) = self.predict(emb) + && stage_index(chosen_stage) <= stage_index(predicted) + { let _ = self.sgd_step(emb, chosen_stage, 0.01); } } @@ -228,4 +227,57 @@ mod tests { let r = LogRegRouter::new(8, 1e-3, 1); assert!(r.predict_probs(&[0.0; 7]).is_err()); } + + /// Neural answers every query correctly. Reinforcing it whenever it did + /// drifted the classifier to predict Neural for everything; only a + /// correct stage no costlier than the prediction may move it. + #[tokio::test] + async fn a_correct_costlier_stage_does_not_move_the_prediction() { + let emb = vec![1.0, -0.5, 0.25]; + let q = Query::new("q").with_embedding(emb.clone()); + let mut r = LogRegRouter::new(3, 0.0, 9); + for _ in 0..200 { + r.sgd_step(&emb, Stage::Kv, 0.1).unwrap(); + } + assert_eq!(r.predict(&emb).unwrap().0, Stage::Kv); + + for _ in 0..500 { + r.observe(&q, Stage::Neural, true, 50_000_000).await; + } + assert_eq!( + r.predict(&emb).unwrap().0, + Stage::Kv, + "Neural being enough is no evidence Kv is not" + ); + + for _ in 0..500 { + r.observe(&q, Stage::Cache, true, 1_000).await; + } + assert_eq!( + r.predict(&emb).unwrap().0, + Stage::Cache, + "a cheaper stage that resolved it is" + ); + } + + /// The predicted stage resolving the query is evidence for it. + #[tokio::test] + async fn the_predicted_stage_resolving_reinforces_it() { + let emb = vec![1.0, -0.5, 0.25]; + let q = Query::new("q").with_embedding(emb.clone()); + let mut r = LogRegRouter::new(3, 0.0, 9); + for _ in 0..20 { + r.sgd_step(&emb, Stage::Kv, 0.1).unwrap(); + } + let (stage, before, _) = r.predict(&emb).unwrap(); + assert_eq!(stage, Stage::Kv); + for _ in 0..50 { + r.observe(&q, Stage::Kv, true, 50_000).await; + } + let (_, after, _) = r.predict(&emb).unwrap(); + assert!( + after > before, + "confidence in Kv went from {before} to {after}" + ); + } } diff --git a/eoc/eoc-rs/crates/eoc-route-learned/src/evaluation.rs b/eoc/eoc-rs/crates/eoc-route-learned/src/evaluation.rs index ba31b8485..7d3f27681 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/src/evaluation.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/src/evaluation.rs @@ -1,13 +1,25 @@ //! Routing-quality metrics. //! -//! `evaluate` consumes a `LearnedRouter` and a held-out test set, returning -//! `RouterMetrics`: +//! `evaluate` replays a router's recommendations over a held-out set of +//! per-stage outcomes, returning `RouterMetrics`: //! -//! - **accuracy** — fraction of examples where the predicted stage matched -//! the recorded successful stage. -//! - **joules_per_correct** — total observed joules / correct predictions. -//! - **lift** — joules saved vs. always running the full cascade (sum of all -//! stage costs). +//! - **accuracy** — fraction of scored queries the dispatch answered +//! correctly. +//! - **total_microjoules** — what the dispatches the router chose spend. +//! - **full_cascade_microjoules** — what the plain cascade walk spends on the +//! same queries. +//! - **lift** — the walk's joules minus the router's: what routing saved, or +//! cost when negative. +//! - **joules_per_correct** — total joules / correct answers, absent when +//! none was correct. +//! +//! Every joule figure is priced from the recorded per-stage costs, run the way +//! [`LearnedCascade`](crate::inference::LearnedCascade) dispatches: Kv or +//! Graph, falling through to Neural on a miss; Neural alone; or, for Cache, +//! the full walk. A row that did not succeed is a miss. A query whose +//! dispatch or walk would run a stage the set has no row for is not scored: +//! it is counted in `unscored` and enters no figure, rather than being priced +//! at zero. //! //! Latency metrics are deferred to the caller (we don't run the cascade //! here, just predictions). @@ -17,76 +29,238 @@ use std::collections::HashMap; use eoc_core::{Query, Stage}; use serde::{Deserialize, Serialize}; -use crate::router::LearnedRouter; -use crate::training::Example; +use crate::router::{LearnedRouter, stage_index}; +use crate::training::{Example, query_key}; /// Summary statistics produced by [`evaluate`]. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct RouterMetrics { - /// Number of test examples scored. + /// Number of queries scored. pub n: usize, - /// Number of correct (predicted == actual successful) predictions. + /// Queries left out because their dispatch or the walk would run a stage + /// the test set has no row for. + pub unscored: usize, + /// Scored queries whose dispatch ended on a correct answer. pub correct: usize, /// `correct / n`. pub accuracy: f32, - /// Sum of observed joule cost. + /// Joules the router's dispatches spend on the scored queries. pub total_microjoules: u64, - /// `total_microjoules / max(correct, 1)`. - pub joules_per_correct: u64, - /// Per-stage prediction count. + /// `total_microjoules / correct`; `None` when no answer was correct. + pub joules_per_correct: Option, + /// Per-stage recommendation count, over every query routed. pub stage_counts: HashMap, - /// Joules saved vs. running the full cascade for every example - /// (sum of all stage costs in the dataset). + /// Joules the plain cascade walk spends on the scored queries. + pub full_cascade_microjoules: u64, + /// `full_cascade_microjoules − total_microjoules`: what routing saved + /// over the walk, negative when it cost more. pub lift_microjoules: i64, } -/// Evaluate a router on a labelled test set. +/// One stage's recorded outcome on a query: whether it resolved it, and what +/// it spent. +type Outcome = Option<(bool, u64)>; + +/// What skipping to `stage` spends on a query and whether it ends correct, +/// or `None` when it would run a stage with no row. +fn dispatch(stage: Stage, rows: &[Outcome; 4]) -> Option<(bool, u64)> { + let neural = rows[stage_index(Stage::Neural)]; + match stage { + Stage::Cache => walk(rows), + Stage::Neural => neural, + Stage::Kv | Stage::Graph => match rows[stage_index(stage)]? { + (true, uj) => Some((true, uj)), + (false, uj) => neural.map(|(ok, n)| (ok, uj.saturating_add(n))), + }, + } +} + +/// The plain cascade: every stage in order until one resolves the query, +/// Neural last. +fn walk(rows: &[Outcome; 4]) -> Option<(bool, u64)> { + let mut spent = 0u64; + for (i, row) in rows.iter().enumerate() { + let (ok, uj) = (*row)?; + spent = spent.saturating_add(uj); + if ok || i == rows.len() - 1 { + return Some((ok, spent)); + } + } + unreachable!("the walk ends at Neural") +} + +/// Evaluate a router on a labelled test set of per-stage outcomes. /// -/// Each example is treated as a (embedding → routed_stage) prediction. If the -/// example's `success` is true and the router picks the same stage, we count -/// it correct. If the example was a failure, we ignore correctness but still -/// account for joules. +/// Rows with the same embedding are one query, and each row is one stage's +/// outcome on it; when a stage has more than one row for a query, the first +/// counts. See the module docs for how each recommendation is priced. pub async fn evaluate(router: &R, test_set: &[Example]) -> RouterMetrics { + let mut index: HashMap, usize> = HashMap::new(); + let mut queries: Vec<(&[f32], [Outcome; 4])> = Vec::new(); + for ex in test_set { + let i = *index.entry(query_key(&ex.embedding)).or_insert_with(|| { + queries.push((&ex.embedding, [None; 4])); + queries.len() - 1 + }); + let slot = &mut queries[i].1[stage_index(ex.stage)]; + if slot.is_none() { + *slot = Some((ex.success, ex.cost_microjoules)); + } + } + + let mut n = 0usize; + let mut unscored = 0usize; let mut correct = 0usize; let mut total_microjoules: u64 = 0; let mut full_cascade_microjoules: u64 = 0; let mut stage_counts: HashMap = HashMap::new(); - for ex in test_set { - let q = Query::new("").with_embedding(ex.embedding.clone()); - let prediction = router.route(&q).await; - *stage_counts.entry(prediction.recommended).or_insert(0) += 1; - if ex.success && prediction.recommended == ex.stage { - correct += 1; - } - // Saved joules: only pay for the recommended stage, not all four. - total_microjoules = total_microjoules.saturating_add(ex.cost_microjoules); - // "Full-cascade" baseline pays every stage's average cost — we use - // 4x the recorded cost as a conservative proxy when per-stage costs - // aren't broken out in the test set. - full_cascade_microjoules = full_cascade_microjoules.saturating_add(ex.cost_microjoules * 4); + for (embedding, rows) in &queries { + let q = Query::new("").with_embedding(embedding.to_vec()); + let recommended = router.route(&q).await.recommended; + *stage_counts.entry(recommended).or_insert(0) += 1; + let (Some((ok, uj)), Some((_, walk_uj))) = (dispatch(recommended, rows), walk(rows)) else { + unscored += 1; + continue; + }; + n += 1; + correct += usize::from(ok); + total_microjoules = total_microjoules.saturating_add(uj); + full_cascade_microjoules = full_cascade_microjoules.saturating_add(walk_uj); } - let n = test_set.len(); let accuracy = if n == 0 { 0.0 } else { correct as f32 / n as f32 }; - let joules_per_correct = if correct == 0 { - 0 - } else { - total_microjoules / correct as u64 - }; + let joules_per_correct = (correct > 0).then(|| total_microjoules / correct as u64); let lift_microjoules = full_cascade_microjoules as i64 - total_microjoules as i64; RouterMetrics { n, + unscored, correct, accuracy, total_microjoules, joules_per_correct, stage_counts, + full_cascade_microjoules, lift_microjoules, } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::router::{RouterState, StagePrediction}; + use async_trait::async_trait; + + /// Recommends one stage for every query. + struct Always(Stage); + + #[async_trait] + impl LearnedRouter for Always { + async fn route(&self, _q: &Query) -> StagePrediction { + StagePrediction::new(self.0, 1.0) + } + async fn observe(&mut self, _q: &Query, _s: Stage, _ok: bool, _j: u64) {} + fn export_state(&self) -> RouterState { + unimplemented!() + } + fn import_state(_s: RouterState) -> crate::Result { + unimplemented!() + } + } + + fn row(q: f32, stage: Stage, success: bool, uj: u64) -> Example { + Example::new(vec![q], stage, success, uj) + } + + /// Query 1 is resolved first by Graph. Query 2 has rows only for Cache + /// (a hit) and Neural. + fn set() -> Vec { + vec![ + row(1.0, Stage::Cache, false, 1), + row(1.0, Stage::Kv, false, 10), + row(1.0, Stage::Graph, true, 100), + row(1.0, Stage::Neural, true, 1_000), + row(2.0, Stage::Cache, true, 1), + row(2.0, Stage::Neural, true, 1_000), + ] + } + + #[tokio::test] + async fn the_routers_choice_decides_the_joules() { + let m = evaluate(&Always(Stage::Graph), &set()).await; + // Query 1: Graph resolves it for 100; the walk spends 1 + 10 + 100. + // Query 2: no Graph row, so it is not scored. + assert_eq!((m.n, m.unscored, m.correct), (1, 1, 1)); + assert_eq!( + (m.total_microjoules, m.full_cascade_microjoules), + (100, 111) + ); + assert_eq!(m.lift_microjoules, 11); + assert_eq!(m.joules_per_correct, Some(100)); + assert_eq!(m.stage_counts[&Stage::Graph], 2); + + let m = evaluate(&Always(Stage::Kv), &set()).await; + // Query 1: Kv misses for 10 and Neural answers for 1 000. + assert_eq!((m.n, m.total_microjoules), (1, 1_010)); + assert_eq!(m.lift_microjoules, 111 - 1_010); + + let m = evaluate(&Always(Stage::Cache), &set()).await; + // Both walks: 111, and 1 for query 2's cache hit. + assert_eq!((m.n, m.total_microjoules, m.lift_microjoules), (2, 112, 0)); + + let m = evaluate(&Always(Stage::Neural), &set()).await; + assert_eq!((m.n, m.total_microjoules), (2, 2_000)); + } + + /// The old metric charged each row's own cost whatever the router chose, + /// and priced the walk at four times it, so every router showed a lift of + /// three times its spend. One row per query prices no walk, and no + /// dispatch except to the stage the row is for. + #[tokio::test] + async fn one_row_per_query_scores_only_what_it_can_price() { + let rows = vec![ + row(1.0, Stage::Kv, true, 50_000), + row(2.0, Stage::Neural, true, 50_000_000), + ]; + let m = evaluate(&Always(Stage::Kv), &rows).await; + assert_eq!((m.n, m.unscored), (0, 2)); + assert_eq!( + ( + m.total_microjoules, + m.full_cascade_microjoules, + m.lift_microjoules + ), + (0, 0, 0) + ); + assert_eq!(m.joules_per_correct, None); + } + + #[tokio::test] + async fn no_correct_answer_has_no_joules_per_correct() { + let rows = vec![ + row(1.0, Stage::Cache, false, 1), + row(1.0, Stage::Kv, false, 10), + row(1.0, Stage::Graph, false, 100), + row(1.0, Stage::Neural, false, 1_000), + ]; + let m = evaluate(&Always(Stage::Neural), &rows).await; + assert_eq!((m.n, m.correct, m.total_microjoules), (1, 0, 1_000)); + assert_eq!( + m.joules_per_correct, None, + "zero would read as the best possible figure" + ); + } + + #[tokio::test] + async fn the_first_row_for_a_stage_counts() { + let mut rows = set(); + rows.push(row(1.0, Stage::Graph, false, 7)); + let m = evaluate(&Always(Stage::Graph), &rows).await; + assert_eq!((m.correct, m.total_microjoules), (1, 100)); + } +} diff --git a/eoc/eoc-rs/crates/eoc-route-learned/src/lib.rs b/eoc/eoc-rs/crates/eoc-route-learned/src/lib.rs index 71c3c6b2e..7ca105129 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/src/lib.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/src/lib.rs @@ -4,8 +4,11 @@ //! strict order. That's optimal when most queries land in the cheap stages. //! For workloads where many queries unavoidably need the neural stage, walking //! the cheaper stages first is wasted joules. A *learned router* predicts the -//! cheapest stage that is likely to satisfy the query and lets us **skip -//! ahead** when the prediction is confident. +//! stage the cascade answers from most cheaply in expectation +//! ([`router::expected_joules`]) and lets us **skip ahead** when the +//! prediction is confident. Cheapest in expectation is not likeliest to +//! succeed: Neural succeeds on everything and costs the most, and a router +//! that ranked stages by success alone learned to send every query there. //! //! This crate ships three orthogonal router families plus a threshold policy //! and an end-to-end cascade extension: diff --git a/eoc/eoc-rs/crates/eoc-route-learned/src/matrix_factorization.rs b/eoc/eoc-rs/crates/eoc-route-learned/src/matrix_factorization.rs index 3667a5ddf..7907db673 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/src/matrix_factorization.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/src/matrix_factorization.rs @@ -30,24 +30,14 @@ use rand_distr::{Distribution, Normal}; use serde::{Deserialize, Serialize}; use crate::error::{Error, Result}; -use crate::router::{ALL_STAGES, LearnedRouter, RouterState, StagePrediction}; +use crate::router::{ + LearnedRouter, RouterState, StagePrediction, blend_joules, cheapest_in_expectation, + default_stage_joules, stage_index, +}; /// Default latent dimension. pub const DEFAULT_LATENT_DIM: usize = 32; -fn stage_index(s: Stage) -> usize { - match s { - Stage::Cache => 0, - Stage::Kv => 1, - Stage::Graph => 2, - Stage::Neural => 3, - } -} - -fn stage_from_index(i: usize) -> Stage { - ALL_STAGES[i.min(3)] -} - /// Matrix-factorisation router (embedding-conditioned). #[derive(Debug, Clone, Serialize, Deserialize)] pub struct MfRouter { @@ -91,7 +81,7 @@ impl MfRouter { stage_factors, stage_bias, global_bias: 0.0, - joule_ema: default_joule_ema(), + joule_ema: default_stage_joules(), } } @@ -105,11 +95,11 @@ impl MfRouter { self.embedding_dim } - /// Forward pass: predict success probability per stage. + /// Forward pass: the probability that each stage resolves the query. /// - /// Returns logits *and* softmax-normalised probabilities. The softmax is - /// taken over per-stage sigmoid outputs so the recommended-stage - /// probability stays in [0, 1] when the router is uncertain. + /// One sigmoid per stage, not a softmax: the stages are not alternatives + /// to one another (a query Kv resolves, Neural resolves too), so the + /// probabilities need not sum to one. pub fn predict_probs(&self, embedding: &[f32]) -> Result<[f32; 4]> { if embedding.len() != self.embedding_dim { return Err(Error::DimensionMismatch { @@ -129,11 +119,13 @@ impl MfRouter { Ok(probs) } - /// Pick the recommended stage and its confidence. + /// Pick the stage cheapest to answer from in expectation + /// ([`cheapest_in_expectation`], priced at [`Self::joule_estimates`]). + /// The confidence is that stage's probability of resolving the query. pub fn predict(&self, embedding: &[f32]) -> Result<(Stage, f32, [f32; 4])> { let probs = self.predict_probs(embedding)?; - let (idx, conf) = argmax(&probs); - Ok((stage_from_index(idx), conf, probs)) + let stage = cheapest_in_expectation(&probs, &self.joule_ema); + Ok((stage, probs[stage_index(stage)], probs)) } /// One SGD step on a single example. @@ -196,9 +188,7 @@ impl MfRouter { /// Record an observed joule cost — EMA with coefficient 0.1. pub fn observe_cost(&mut self, stage: Stage, microjoules: u64) { - let prev = self.joule_ema.get(&stage).copied().unwrap_or(microjoules) as f64; - let blended = 0.9 * prev + 0.1 * microjoules as f64; - self.joule_ema.insert(stage, blended as u64); + blend_joules(&mut self.joule_ema, stage, microjoules); } /// Current per-stage joule estimate (EMA). @@ -207,16 +197,6 @@ impl MfRouter { } } -fn default_joule_ema() -> HashMap { - let mut m = HashMap::new(); - // Conservative defaults, ordered cheap → expensive. - m.insert(Stage::Cache, 1_000); - m.insert(Stage::Kv, 50_000); - m.insert(Stage::Graph, 500_000); - m.insert(Stage::Neural, 50_000_000); - m -} - fn sigmoid(x: f32) -> f32 { if x >= 0.0 { let z = (-x).exp(); @@ -227,18 +207,6 @@ fn sigmoid(x: f32) -> f32 { } } -fn argmax(xs: &[f32; 4]) -> (usize, f32) { - let mut best = 0; - let mut bv = xs[0]; - for (i, v) in xs.iter().enumerate().skip(1) { - if *v > bv { - best = i; - bv = *v; - } - } - (best, bv) -} - #[async_trait] impl LearnedRouter for MfRouter { async fn route(&self, query: &Query) -> StagePrediction { @@ -323,6 +291,15 @@ mod tests { } } + #[test] + fn the_confidence_is_the_chosen_stages_chance() { + let r = MfRouter::new(4, 4, 7); + let emb = [0.5, -0.5, 0.25, -0.25]; + let (stage, conf, probs) = r.predict(&emb).unwrap(); + assert_eq!(conf, probs[stage_index(stage)]); + assert!(conf < 1.0); + } + #[test] fn sgd_reduces_loss() { let mut r = MfRouter::new(4, 4, 7); diff --git a/eoc/eoc-rs/crates/eoc-route-learned/src/router.rs b/eoc/eoc-rs/crates/eoc-route-learned/src/router.rs index 9ea875328..b856eccb2 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/src/router.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/src/router.rs @@ -63,10 +63,20 @@ pub struct RouterState { /// A learned router predicts the best cascade stage from a query. #[async_trait] pub trait LearnedRouter: Send + Sync { - /// Predict the best stage for `query`. + /// Predict the best stage for `query`: the one [`LearnedCascade`] answers + /// from most cheaply in expectation (see [`expected_joules`]), not the one + /// most likely to succeed. The neural stage is the most likely to succeed + /// and the most expensive, so a router ranking by success alone learns to + /// skip every cheap stage. + /// + /// [`LearnedCascade`]: crate::inference::LearnedCascade async fn route(&self, query: &Query) -> StagePrediction; - /// Online update — observe what actually happened. + /// Online update with the outcome of one stage that ran: whether + /// `chosen_stage` resolved `query` correctly, and the microjoules that + /// stage itself spent. When a dispatch ran several stages (a full walk, + /// or a miss that fell through to Neural), report each of them: a stage + /// that missed is evidence too. async fn observe( &mut self, query: &Query, @@ -86,3 +96,186 @@ pub trait LearnedRouter: Send + Sync { /// The four cascade stages, in cheapest-first order, useful for iteration. pub const ALL_STAGES: [Stage; 4] = [Stage::Cache, Stage::Kv, Stage::Graph, Stage::Neural]; + +/// A stage's position in [`ALL_STAGES`]. +pub(crate) fn stage_index(s: Stage) -> usize { + match s { + Stage::Cache => 0, + Stage::Kv => 1, + Stage::Graph => 2, + Stage::Neural => 3, + } +} + +/// Per-stage cost assumed before any is observed, in microjoules. Conservative +/// defaults, ordered cheap → expensive. +pub fn default_stage_joules() -> HashMap { + ALL_STAGES.iter().map(|s| (*s, default_cost(*s))).collect() +} + +fn default_cost(stage: Stage) -> u64 { + match stage { + Stage::Cache => 1_000, + Stage::Kv => 50_000, + Stage::Graph => 500_000, + Stage::Neural => 50_000_000, + } +} + +/// Fold one observed stage cost into a per-stage estimate — EMA with +/// coefficient 0.1. +pub fn blend_joules(joules: &mut HashMap, stage: Stage, microjoules: u64) { + let prev = joules.get(&stage).copied().unwrap_or(microjoules) as f64; + let blended = 0.9 * prev + 0.1 * microjoules as f64; + joules.insert(stage, blended as u64); +} + +/// Expected microjoules for [`LearnedCascade`] to answer a query it was told +/// to skip to `stage`, given each stage's probability of resolving the query +/// (`resolves`, in [`ALL_STAGES`] order) and each stage's cost (`joules`). +/// +/// It follows the dispatch exactly: +/// +/// ```text +/// Neural c_N +/// Kv, Graph c_s + (1 − p_s) · c_N a miss falls through to Neural +/// Cache c_C + (1 − p_C)(c_K + (1 − p_K)(c_G + (1 − p_G) · c_N)) +/// the full walk +/// ``` +/// +/// Neural always answers, so its own probability does not enter. The walk +/// treats the stages' misses as independent, because the routers estimate +/// each stage's chance on its own rather than given that a cheaper one +/// missed. A probability outside [0, 1] is clamped and NaN counts as a miss, +/// so a broken estimate prices a stage dearer, never cheaper. A stage with no +/// cost estimate is priced at [`default_stage_joules`]. +/// +/// [`LearnedCascade`]: crate::inference::LearnedCascade +pub fn expected_joules(stage: Stage, resolves: &[f32; 4], joules: &HashMap) -> f64 { + let cost = |s: Stage| joules.get(&s).copied().unwrap_or_else(|| default_cost(s)) as f64; + let miss = |i: usize| { + let p = resolves[i]; + if p.is_nan() { + 1.0 + } else { + 1.0 - f64::from(p.clamp(0.0, 1.0)) + } + }; + let neural = cost(Stage::Neural); + match stage { + Stage::Neural => neural, + Stage::Kv => cost(Stage::Kv) + miss(1) * neural, + Stage::Graph => cost(Stage::Graph) + miss(2) * neural, + Stage::Cache => { + let graph_on = cost(Stage::Graph) + miss(2) * neural; + let kv_on = cost(Stage::Kv) + miss(1) * graph_on; + cost(Stage::Cache) + miss(0) * kv_on + } + } +} + +/// The stage [`LearnedCascade`] answers from most cheaply in expectation +/// ([`expected_joules`]). A tie goes to the cheaper stage. +/// +/// [`LearnedCascade`]: crate::inference::LearnedCascade +pub fn cheapest_in_expectation(resolves: &[f32; 4], joules: &HashMap) -> Stage { + let mut best = ALL_STAGES[0]; + let mut best_uj = expected_joules(best, resolves, joules); + for s in &ALL_STAGES[1..] { + let uj = expected_joules(*s, resolves, joules); + if uj < best_uj { + best = *s; + best_uj = uj; + } + } + best +} + +#[cfg(test)] +mod tests { + use super::*; + + fn costs() -> HashMap { + let mut m = HashMap::new(); + m.insert(Stage::Cache, 1); + m.insert(Stage::Kv, 10); + m.insert(Stage::Graph, 100); + m.insert(Stage::Neural, 1_000); + m + } + + #[test] + fn each_dispatch_is_priced_as_the_cascade_runs_it() { + let p = [0.5, 0.25, 0.75, 0.0]; + let c = costs(); + assert_eq!(expected_joules(Stage::Neural, &p, &c), 1_000.0); + assert_eq!(expected_joules(Stage::Kv, &p, &c), 10.0 + 0.75 * 1_000.0); + assert_eq!( + expected_joules(Stage::Graph, &p, &c), + 100.0 + 0.25 * 1_000.0 + ); + let walk = 1.0 + 0.5 * (10.0 + 0.75 * (100.0 + 0.25 * 1_000.0)); + assert_eq!(expected_joules(Stage::Cache, &p, &c), walk); + } + + /// The defect this rule replaces: Neural is the likeliest to succeed, so + /// an argmax over success picked it, at a thousand times Graph's cost. + #[test] + fn the_likeliest_stage_is_not_the_cheapest_to_answer_from() { + let c = costs(); + let p = [0.0, 0.0, 0.95, 1.0]; + assert_eq!(cheapest_in_expectation(&p, &c), Stage::Graph); + // Hopeless cheap stages: skipping them all is cheapest. + assert_eq!( + cheapest_in_expectation(&[0.0, 0.0, 0.0, 1.0], &c), + Stage::Neural + ); + // A likely cache hit: walk from the top. + assert_eq!( + cheapest_in_expectation(&[0.9, 0.0, 0.0, 1.0], &c), + Stage::Cache + ); + } + + #[test] + fn a_broken_estimate_never_prices_a_stage_cheaper() { + let c = costs(); + let nan = [f32::NAN; 4]; + assert_eq!(expected_joules(Stage::Kv, &nan, &c), 1_010.0); + assert_eq!(expected_joules(Stage::Kv, &[0.0, 7.0, 0.0, 0.0], &c), 10.0); + assert_eq!( + expected_joules(Stage::Kv, &[0.0, -3.0, 0.0, 0.0], &c), + 1_010.0 + ); + } + + #[test] + fn a_tie_goes_to_the_cheaper_stage() { + let mut c = costs(); + c.insert(Stage::Graph, 10); + // Kv and Graph both cost 10 and both certainly resolve. + assert_eq!( + cheapest_in_expectation(&[0.0, 1.0, 1.0, 1.0], &c), + Stage::Kv + ); + } + + #[test] + fn a_stage_without_an_estimate_is_priced_at_the_default() { + let p = [0.0; 4]; + assert_eq!( + expected_joules(Stage::Neural, &p, &HashMap::new()), + 50_000_000.0 + ); + } + + #[test] + fn blending_moves_a_tenth_of_the_way() { + let mut j = costs(); + blend_joules(&mut j, Stage::Neural, 2_000); + assert_eq!(j[&Stage::Neural], 1_100); + let mut empty = HashMap::new(); + blend_joules(&mut empty, Stage::Kv, 40); + assert_eq!(empty[&Stage::Kv], 40); + } +} diff --git a/eoc/eoc-rs/crates/eoc-route-learned/src/training.rs b/eoc/eoc-rs/crates/eoc-route-learned/src/training.rs index c0d8dfd46..3ac727f28 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/src/training.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/src/training.rs @@ -5,6 +5,8 @@ //! `train_logreg` (logistic regression). Both consume `&[Example]` and emit a //! ready-to-serve router. +use std::collections::HashMap; +use std::collections::hash_map::Entry; use std::path::Path; use eoc_core::Stage; @@ -15,6 +17,7 @@ use serde::{Deserialize, Serialize}; use crate::classifier::LogRegRouter; use crate::error::{Error, Result}; use crate::matrix_factorization::{DEFAULT_LATENT_DIM, MfRouter}; +use crate::router::stage_index; /// One labeled training example. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -141,16 +144,46 @@ pub fn train_mf_with_dim( router } -/// Train a logistic-regression router. Only successful examples shape the -/// classifier — failed-stage examples are skipped because we want -/// `P(stage = correct | x)`. +/// The key that makes rows one query: their embedding, bit for bit. +pub(crate) fn query_key(embedding: &[f32]) -> Vec { + embedding.iter().map(|v| v.to_bits()).collect() +} + +/// Each query's cheapest successful row, in the order queries first appear. +/// Rows with the same embedding are one query; a query no stage resolved has +/// no row here. +pub(crate) fn cheapest_successes(examples: &[Example]) -> Vec<&Example> { + let mut index: HashMap, usize> = HashMap::new(); + let mut out: Vec<&Example> = Vec::new(); + for e in examples.iter().filter(|e| e.success) { + match index.entry(query_key(&e.embedding)) { + Entry::Occupied(slot) => { + let kept = &mut out[*slot.get()]; + if stage_index(e.stage) < stage_index(kept.stage) { + *kept = e; + } + } + Entry::Vacant(slot) => { + slot.insert(out.len()); + out.push(e); + } + } + } + out +} + +/// Train a logistic-regression router to predict the cheapest stage that +/// resolves a query. Each query is labelled with its cheapest successful row +/// ([`Example`]s with the same embedding are one query); failed rows label +/// nothing. Labelling every successful row instead would teach it the stage +/// that most often succeeds, and Neural succeeds on everything. pub fn train_logreg(examples: &[Example], config: TrainingConfig) -> LogRegRouter { let dim = first_embedding_dim(examples).unwrap_or(1); let mut router = LogRegRouter::new(dim, config.l2, config.seed); if examples.is_empty() { return router; } - let positive: Vec<&Example> = examples.iter().filter(|e| e.success).collect(); + let positive = cheapest_successes(examples); if positive.is_empty() { return router; } @@ -182,4 +215,50 @@ mod tests { let r = train_logreg(&[], TrainingConfig::default()); assert_eq!(r.embedding_dim(), 1); } + + #[test] + fn a_query_is_labelled_with_its_cheapest_success() { + let a = vec![1.0, 0.0]; + let b = vec![0.0, 1.0]; + let rows = vec![ + Example::new(a.clone(), Stage::Neural, true, 50_000_000), + Example::new(a.clone(), Stage::Cache, false, 1_000), + Example::new(b.clone(), Stage::Graph, false, 500_000), + Example::new(a.clone(), Stage::Graph, true, 500_000), + Example::new(a.clone(), Stage::Kv, true, 50_000), + Example::new(b.clone(), Stage::Neural, true, 50_000_000), + Example::new(vec![9.0, 9.0], Stage::Kv, false, 50_000), + ]; + let kept: Vec<(Vec, Stage)> = cheapest_successes(&rows) + .into_iter() + .map(|e| (e.embedding.clone(), e.stage)) + .collect(); + assert_eq!(kept, vec![(a, Stage::Kv), (b, Stage::Neural)]); + } + + /// Every query here is resolved by Kv, its cheapest stage, and by Graph + /// and Neural; Neural ran twice. Trained on every successful row, the + /// classifier learned Neural, the stage that most often succeeds. + #[test] + fn logreg_learns_the_cheapest_stage_not_the_most_frequent_success() { + let mut rows = Vec::new(); + for i in 0..40 { + let emb = vec![1.0, i as f32 / 40.0]; + rows.push(Example::new(emb.clone(), Stage::Kv, true, 50_000)); + rows.push(Example::new(emb.clone(), Stage::Graph, true, 500_000)); + rows.push(Example::new(emb.clone(), Stage::Neural, true, 50_000_000)); + rows.push(Example::new(emb, Stage::Neural, true, 50_000_000)); + } + let cfg = TrainingConfig { + epochs: 30, + learning_rate: 0.1, + l2: 0.0, + ..TrainingConfig::default() + }; + let r = train_logreg(&rows, cfg); + for i in 0..40 { + let emb = [1.0, i as f32 / 40.0]; + assert_eq!(r.predict(&emb).unwrap().0, Stage::Kv); + } + } } diff --git a/eoc/eoc-rs/crates/eoc-route-learned/tests/linucb_bandit.rs b/eoc/eoc-rs/crates/eoc-route-learned/tests/linucb_bandit.rs index 76b5a94d3..5118306c3 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/tests/linucb_bandit.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/tests/linucb_bandit.rs @@ -1,68 +1,193 @@ -//! 100-round LinUCB synthetic bandit test. +//! Synthetic bandit tests: LinUCB and Thompson sampling. //! -//! Four arms with known mean rewards. Context is a fixed-length embedding -//! that *informs* but doesn't determine reward (Gaussian noise). After ~50 -//! rounds the router should pull the best arm more than any other. +//! Each stage resolves a query with a known probability. Each round the +//! router picks a stage, the pick is run the way `LearnedCascade` dispatches +//! it (Kv or Graph falling through to Neural on a miss, Neural alone, or for +//! Cache the full walk), and every stage that ran reports its outcome and +//! cost. The router must learn the dispatch that spends least, not the stage +//! likeliest to succeed: before, both bandits converged on Neural. use eoc_core::Stage; -use eoc_route_learned::bandit::LinUcbRouter; -use rand::SeedableRng; +use eoc_route_learned::bandit::{LinUcbRouter, ThompsonSamplingRouter}; +use eoc_route_learned::router::ALL_STAGES; +use rand::{RngExt, SeedableRng}; use rand_distr::{Distribution, Normal}; use std::collections::HashMap; const DIM: usize = 4; -fn arm_mean(stage: Stage) -> f32 { +fn cost(stage: Stage) -> u64 { match stage { - Stage::Cache => 0.10, - Stage::Kv => 0.20, - Stage::Graph => 0.30, - // Best arm: Neural. - Stage::Neural => 0.80, + Stage::Cache => 1_000, + Stage::Kv => 50_000, + Stage::Graph => 500_000, + Stage::Neural => 50_000_000, } } -#[test] -fn linucb_converges_to_best_arm() { - // Lower alpha → less exploration, more exploitation. Ridge=0.1 so the - // estimated theta moves quickly. - let mut router = LinUcbRouter::new(DIM, 0.1, 0.1); - let mut rng = rand::rngs::StdRng::seed_from_u64(0xBEEF); - let ctx_dist = Normal::new(0.0_f32, 1.0).expect("valid normal"); - let reward_noise = Normal::new(0.0_f32, 0.02).expect("valid normal"); - - let mut pulls: HashMap = HashMap::new(); - let mut tail_pulls: HashMap = HashMap::new(); - - for round in 0..100 { - let ctx: Vec = (0..DIM).map(|_| ctx_dist.sample(&mut rng)).collect(); - let (pick, _score) = router.pick(&ctx).expect("pick"); - *pulls.entry(pick).or_insert(0) += 1; - if round >= 50 { - *tail_pulls.entry(pick).or_insert(0) += 1; +/// The stages a dispatch to `stage` runs, each with whether it resolved. +fn dispatch(stage: Stage, resolves: &HashMap) -> Vec<(Stage, bool)> { + let order: &[Stage] = match stage { + Stage::Cache => &ALL_STAGES, + Stage::Kv => &[Stage::Kv, Stage::Neural], + Stage::Graph => &[Stage::Graph, Stage::Neural], + Stage::Neural => &[Stage::Neural], + }; + let mut ran = Vec::new(); + for s in order { + ran.push((*s, resolves[s])); + if resolves[s] { + break; } - let reward = (arm_mean(pick) + reward_noise.sample(&mut rng)).clamp(0.0, 1.0); - router.update(&ctx, pick, reward).expect("update"); } + ran +} - println!("LinUCB pulls (all 100 rounds): {pulls:?}"); - println!("LinUCB pulls (rounds 50-99): {tail_pulls:?}"); - - let neural_tail = *tail_pulls.get(&Stage::Neural).unwrap_or(&0); - let best_other = [Stage::Cache, Stage::Kv, Stage::Graph] - .iter() - .map(|s| *tail_pulls.get(s).unwrap_or(&0)) - .max() - .unwrap_or(0); - assert!( - neural_tail >= best_other, - "neural ({neural_tail}) should be pulled at least as often as any other ({best_other}) after round 50" - ); - // And neural should dominate over the full 100 rounds, since its mean - // reward is 0.8 vs ≤0.3 for every other arm. - let neural_total = *pulls.get(&Stage::Neural).unwrap_or(&0); - assert!( - neural_total >= 40, - "neural should be pulled ≥40 times over 100 rounds, got {neural_total}" - ); +/// Pick counts over the last `tail` of `rounds`, and the joules spent there. +struct Run { + tail_picks: HashMap, + tail_microjoules: u64, +} + +fn run( + rounds: usize, + tail: usize, + chance: [f64; 4], + mut pick: impl FnMut(&[f32]) -> Stage, + mut learn: impl FnMut(&[f32], Stage, bool, u64), +) -> Run { + let mut rng = rand::rngs::StdRng::seed_from_u64(0xBEEF); + let noise = Normal::new(0.0_f32, 1.0).expect("valid normal"); + let mut out = Run { + tail_picks: HashMap::new(), + tail_microjoules: 0, + }; + for round in 0..rounds { + // A constant feature first, so LinUCB can learn each stage's base rate. + let mut ctx = vec![1.0_f32]; + ctx.extend((1..DIM).map(|_| 0.1 * noise.sample(&mut rng))); + let resolves: HashMap = ALL_STAGES + .iter() + .zip(chance) + .map(|(s, p)| (*s, rng.random_bool(p))) + .collect(); + let stage = pick(&ctx); + for (s, ok) in dispatch(stage, &resolves) { + learn(&ctx, s, ok, cost(s)); + if round >= rounds - tail { + out.tail_microjoules += cost(s); + } + } + if round >= rounds - tail { + *out.tail_picks.entry(stage).or_insert(0) += 1; + } + } + out +} + +fn linucb(rounds: usize, tail: usize, chance: [f64; 4]) -> Run { + let router = std::cell::RefCell::new(LinUcbRouter::new(DIM, 0.1, 1.0)); + run( + rounds, + tail, + chance, + |ctx| router.borrow().pick(ctx).expect("pick").0, + |ctx, s, ok, uj| { + let mut r = router.borrow_mut(); + r.update(ctx, s, if ok { 1.0 } else { 0.0 }) + .expect("update"); + r.observe_cost(s, uj); + }, + ) +} + +fn thompson(rounds: usize, tail: usize, chance: [f64; 4]) -> Run { + let router = std::cell::RefCell::new(ThompsonSamplingRouter::new(7)); + run( + rounds, + tail, + chance, + |_| router.borrow().pick().0, + |_, s, ok, uj| { + let mut r = router.borrow_mut(); + r.update(s, ok); + r.observe_cost(s, uj); + }, + ) +} + +fn share(run: &Run, stage: Stage, tail: usize) -> f32 { + *run.tail_picks.get(&stage).unwrap_or(&0) as f32 / tail as f32 +} + +/// Graph resolves 95% of queries and nothing cheaper resolves any: skip to +/// Graph, at about 3 million µJ a query, not Neural at 50 million. +#[test] +fn bandits_skip_to_the_stage_that_resolves() { + let chance = [0.0, 0.0, 0.95, 1.0]; + for (name, r) in [ + ("linucb", linucb(3000, 500, chance)), + ("thompson", thompson(3000, 500, chance)), + ] { + println!("{name}: {:?}, {} µJ", r.tail_picks, r.tail_microjoules); + assert!( + share(&r, Stage::Graph, 500) > 0.8, + "{name} picks Graph only {}", + share(&r, Stage::Graph, 500) + ); + assert!( + share(&r, Stage::Neural, 500) < 0.05, + "{name} still skips to Neural" + ); + assert!( + r.tail_microjoules / 500 < 5_000_000, + "{name} spends {} µJ a query", + r.tail_microjoules / 500 + ); + } +} + +/// Nothing short of Neural resolves anything, so going straight to Neural, +/// at 50 million µJ, is cheapest; the full walk costs 50.551 million. The +/// bandits come within 1% of Neural. They still probe Kv now and then: a +/// probe costs a thousandth of a Neural call and would save the whole call +/// if it hit, so it is worth trying until the bound on Kv's chance falls +/// below 0.1%. +#[test] +fn bandits_come_within_a_percent_of_neural_when_nothing_cheaper_resolves() { + let chance = [0.0, 0.0, 0.0, 1.0]; + for (name, r) in [ + ("linucb", linucb(4000, 500, chance)), + ("thompson", thompson(4000, 500, chance)), + ] { + println!("{name}: {:?}, {} µJ", r.tail_picks, r.tail_microjoules); + let per_query = r.tail_microjoules / 500; + assert!( + per_query < 50_500_000, + "{name} spends {per_query} µJ a query" + ); + } +} + +/// The old test's arms. Neural is likeliest to succeed (0.8), and the old +/// bandits learned to pick it every time. But the cheap stages resolve often +/// enough that walking them first costs least, so the walk is the pick. +#[test] +fn bandits_walk_when_cheap_stages_often_resolve() { + let chance = [0.1, 0.2, 0.3, 0.8]; + for (name, r) in [ + ("linucb", linucb(1000, 500, chance)), + ("thompson", thompson(1000, 500, chance)), + ] { + println!("{name}: {:?}, {} µJ", r.tail_picks, r.tail_microjoules); + assert!( + share(&r, Stage::Cache, 500) > 0.8, + "{name} walks only {}", + share(&r, Stage::Cache, 500) + ); + assert!( + share(&r, Stage::Neural, 500) < 0.05, + "{name} skips to Neural" + ); + } } diff --git a/eoc/eoc-rs/crates/eoc-route-learned/tests/mf_synthetic.rs b/eoc/eoc-rs/crates/eoc-route-learned/tests/mf_synthetic.rs index 1338e996e..c6a9b3efb 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/tests/mf_synthetic.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/tests/mf_synthetic.rs @@ -1,73 +1,204 @@ -//! Synthetic-data accuracy test for the MF router. +//! Synthetic test for the MF router. //! -//! Generates 1000 examples where the "best stage" is a deterministic -//! function of the embedding (first two dims define quadrants → stage). The -//! MF router should learn this assignment with >80% accuracy. +//! Each query has one row per stage. Each cheap stage has its own coverage: +//! the chance it resolves a query is a logistic function of one feature of +//! the embedding (Cache's centred at 0.5, Kv's at 0, Graph's at −0.5), and +//! Neural resolves everything. The router must learn to skip to the stage +//! that is cheapest to answer from in expectation, not to the one likeliest +//! to succeed, which is Neural. +//! +//! The coverage is soft, as real coverage is, and of the family the router +//! models, so its learned chances can be right. They have to be: with the +//! default costs a wrong skip pays a Neural call, a thousand times Graph, +//! while a right one saves at most a hundredth of one. A router sure of +//! itself on noise-free boundaries skips wrongly at their edges and spends +//! more than the plain walk. +//! +//! With the default costs the plain walk is itself within a fraction of a +//! percent of the best choices, so there is little for routing to win, and +//! the router lands within 1% of the best. Where Kv and Graph cost a real +//! share of a Neural call, the walk is 4–5% above the best, but at this +//! training budget MF's chances are off by 5–9% a stage, and the router lands +//! at about the walk's cost rather than capturing that. Training longer made +//! it beat the walk by 2% while its chances got worse, so that is a property +//! of the estimator, not of the rule; the test holds it to not losing. -use eoc_core::Stage; +use std::collections::HashMap; + +use async_trait::async_trait; +use eoc_core::{Query, Stage}; +use eoc_route_learned::evaluation::evaluate; use eoc_route_learned::matrix_factorization::MfRouter; +use eoc_route_learned::router::{ + ALL_STAGES, LearnedRouter, RouterState, StagePrediction, cheapest_in_expectation, + default_stage_joules, expected_joules, +}; use eoc_route_learned::training::{Example, TrainingConfig, train_mf}; -use rand::SeedableRng; +use rand::{RngExt, SeedableRng}; use rand_distr::{Distribution, Normal}; const DIM: usize = 8; -fn quadrant_stage(emb: &[f32]) -> Stage { - let (a, b) = (emb[0], emb[1]); - match (a >= 0.0, b >= 0.0) { - (true, true) => Stage::Cache, - (true, false) => Stage::Kv, - (false, true) => Stage::Graph, - (false, false) => Stage::Neural, - } +/// Kv and Graph at a real share of a Neural call. +fn dear_cheap_stages() -> HashMap { + let mut c = default_stage_joules(); + c.insert(Stage::Kv, 2_000_000); + c.insert(Stage::Graph, 10_000_000); + c } -fn synthesise(n: usize, seed: u64) -> Vec { +/// The true chance that each stage resolves the query with this embedding, +/// in `ALL_STAGES` order. +fn chances(emb: &[f32]) -> [f32; 4] { + let s = |x: f32| 1.0 / (1.0 + (-4.0 * x).exp()); + [s(emb[0] - 0.5), s(emb[1]), s(emb[2] + 0.5), 1.0] +} + +fn synthesise(n: usize, seed: u64, costs: &HashMap) -> Vec { let mut rng = rand::rngs::StdRng::seed_from_u64(seed); let dist = Normal::new(0.0_f32, 1.0).expect("valid normal"); - let mut out = Vec::with_capacity(n); + let mut out = Vec::with_capacity(n * 4); for _ in 0..n { - let mut emb = vec![0.0_f32; DIM]; - for v in emb.iter_mut() { - *v = dist.sample(&mut rng); + let emb: Vec = (0..DIM).map(|_| dist.sample(&mut rng)).collect(); + let p = chances(&emb); + for (i, stage) in ALL_STAGES.into_iter().enumerate() { + let success = rng.random_bool(f64::from(p[i])); + out.push(Example::new(emb.clone(), stage, success, costs[&stage])); } - let stage = quadrant_stage(&emb); - // Pure labels — the synthetic test focuses on classification accuracy, - // not noise-handling. - let success = true; - // cost varies per stage so cost estimates are meaningful. - let cost = match stage { - Stage::Cache => 1_000, - Stage::Kv => 50_000, - Stage::Graph => 500_000, - Stage::Neural => 50_000_000, - }; - out.push(Example::new(emb, stage, success, cost)); } out } -#[test] -fn mf_router_accuracy_above_80() { - let train = synthesise(1000, 17); - let test = synthesise(200, 42); - +fn router(costs: &HashMap) -> MfRouter { let cfg = TrainingConfig { epochs: 30, learning_rate: 0.1, l2: 1e-5, ..TrainingConfig::default() }; - let router: MfRouter = train_mf(&train, cfg); + train_mf(&synthesise(1000, 17, costs), cfg) +} - let mut correct = 0usize; - for ex in &test { - let (stage, _conf, _probs) = router.predict(&ex.embedding).expect("predict"); - if stage == ex.stage { - correct += 1; +/// The rule the routers used before: the stage likeliest to succeed. +struct LikeliestToSucceed<'a>(&'a MfRouter); + +#[async_trait] +impl LearnedRouter for LikeliestToSucceed<'_> { + async fn route(&self, q: &Query) -> StagePrediction { + let probs = self.0.predict_probs(q.embedding.as_ref().unwrap()).unwrap(); + let best = (0..4).fold(0, |b, i| if probs[i] > probs[b] { i } else { b }); + StagePrediction::new(ALL_STAGES[best], probs[best]) + } + async fn observe(&mut self, _q: &Query, _s: Stage, _ok: bool, _j: u64) {} + fn export_state(&self) -> RouterState { + unimplemented!() + } + fn import_state(_s: RouterState) -> eoc_route_learned::Result { + unimplemented!() + } +} + +/// Expected µJ over held-out queries, priced at their true chances: the best +/// choices, the router's, the plain walk's and the likeliest-stage rule's; +/// and the share of queries where the router's choice is the best one. +struct Priced { + best: f64, + routed: f64, + walk: f64, + likeliest: f64, + agree: f64, +} + +fn priced(costs: &HashMap) -> Priced { + let router = router(costs); + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let dist = Normal::new(0.0_f32, 1.0).expect("valid normal"); + let mut out = Priced { + best: 0.0, + routed: 0.0, + walk: 0.0, + likeliest: 0.0, + agree: 0.0, + }; + let mut fit = [0.0_f64; 4]; + let n = 2000; + for _ in 0..n { + let emb: Vec = (0..DIM).map(|_| dist.sample(&mut rng)).collect(); + let p = chances(&emb); + let (choice, _, learned) = router.predict(&emb).expect("predict"); + let ideal = cheapest_in_expectation(&p, costs); + let top = (0..4).fold(0, |b, i| if learned[i] > learned[b] { i } else { b }); + for i in 0..4 { + fit[i] += f64::from((learned[i] - p[i]).abs()) / f64::from(n); } + out.agree += f64::from(u8::from(choice == ideal)) / f64::from(n); + out.routed += expected_joules(choice, &p, costs); + out.walk += expected_joules(Stage::Cache, &p, costs); + out.best += expected_joules(ideal, &p, costs); + out.likeliest += expected_joules(ALL_STAGES[top], &p, costs); } - let acc = correct as f32 / test.len() as f32; - println!("MF synthetic accuracy: {:.3}", acc); - assert!(acc > 0.8, "MF accuracy {acc} below 0.8 threshold"); + println!( + "expected µJ: best {:.0}, routed {:.0}, walk {:.0}, likeliest {:.0}; agrees with best {:.3}", + out.best, out.routed, out.walk, out.likeliest, out.agree + ); + println!("mean |learned − true| chance per stage: {fit:.3?}"); + out +} + +#[test] +fn with_the_default_costs_the_router_lands_within_a_percent_of_the_best() { + let r = priced(&default_stage_joules()); + assert!( + r.agree > 0.9, + "the router's choice is the best for only {} of queries", + r.agree + ); + assert!( + r.routed <= r.best * 1.01, + "routed {:.0} µJ is over 1% above the best {:.0}", + r.routed, + r.best + ); + assert!( + r.routed * 2.0 < r.likeliest, + "routed {:.0} µJ is not far below the likeliest-stage rule's {:.0}", + r.routed, + r.likeliest + ); +} + +#[test] +fn where_cheap_stages_are_dear_the_router_does_not_lose_to_the_walk() { + let r = priced(&dear_cheap_stages()); + assert!( + r.routed <= r.walk * 1.01, + "routed {:.0} µJ is over 1% above the walk's {:.0}", + r.routed, + r.walk + ); + assert!( + r.routed * 2.0 < r.likeliest, + "routed {:.0} µJ is not far below the likeliest-stage rule's {:.0}", + r.routed, + r.likeliest + ); +} + +/// The same, replayed over sampled outcomes by `evaluate`. +#[tokio::test] +async fn mf_router_spends_far_less_than_the_likeliest_stage_rule() { + let costs = default_stage_joules(); + let router = router(&costs); + let test = synthesise(500, 42, &costs); + let routed = evaluate(&router, &test).await; + let likeliest = evaluate(&LikeliestToSucceed(&router), &test).await; + println!("routed: {routed:?}"); + println!("likeliest: {likeliest:?}"); + assert_eq!((routed.n, routed.unscored), (500, 0)); + assert!( + routed.total_microjoules * 2 < likeliest.total_microjoules, + "routed {} µJ is not far below the likeliest-stage rule's {}", + routed.total_microjoules, + likeliest.total_microjoules + ); } diff --git a/eoc/eoc-rs/crates/eoc-route-learned/tests/state_roundtrip.rs b/eoc/eoc-rs/crates/eoc-route-learned/tests/state_roundtrip.rs index fe0b56677..04e840d35 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/tests/state_roundtrip.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/tests/state_roundtrip.rs @@ -46,3 +46,45 @@ fn rejects_wrong_algorithm_tag() { s.algorithm = "wrong".into(); assert!(MfRouter::import_state(s).is_err()); } + +/// The bandits learn each stage's cost. A learned cost survives export, and +/// a state saved before they learned costs imports with the defaults rather +/// than failing. +#[test] +fn bandit_stage_costs_roundtrip_and_older_states_import() { + use eoc_core::Stage; + use eoc_route_learned::router::default_stage_joules; + + fn without_joules(state: &mut eoc_route_learned::router::RouterState) { + let mut v: serde_json::Value = serde_json::from_slice(&state.weights_blob).unwrap(); + v.as_object_mut() + .unwrap() + .remove("joules") + .expect("the field is exported"); + state.weights_blob = serde_json::to_vec(&v).unwrap(); + } + + let mut lin = LinUcbRouter::new(4, 1.0, 1.0); + lin.observe_cost(Stage::Neural, 10_000); + let restored = LinUcbRouter::import_state(lin.export_state()).expect("import"); + assert_eq!(restored.joules, lin.joules); + let mut old = lin.export_state(); + without_joules(&mut old); + assert_eq!( + LinUcbRouter::import_state(old).expect("import").joules, + default_stage_joules() + ); + + let mut ts = ThompsonSamplingRouter::new(3); + ts.observe_cost(Stage::Kv, 7); + let restored = ThompsonSamplingRouter::import_state(ts.export_state()).expect("import"); + assert_eq!(restored.joules, ts.joules); + let mut old = ts.export_state(); + without_joules(&mut old); + assert_eq!( + ThompsonSamplingRouter::import_state(old) + .expect("import") + .joules, + default_stage_joules() + ); +}