The crate says a learned router "predicts the cheapest stage that is likely to satisfy the query". Three of its four routers recommended the stage likeliest to succeed: MF and both bandits took an argmax over success, and the bandits ignored the joule cost `observe` passed them. Neural succeeds on everything and costs the most, so they learned to send every query there. The synthetic LinUCB test pinned exactly that as the goal. - router: `expected_joules` prices a skip the way LearnedCascade dispatches it (Kv or Graph falling through to Neural, Neural alone, Cache as the full walk), and `cheapest_in_expectation` picks the least. A broken probability prices a stage dearer, never cheaper. `observe` is the outcome of one stage that ran, with that stage's own joules. - MF picks by that rule over its per-stage chances and learned costs. - LinUCB and Thompson learn each stage's cost and pick by the rule over optimistic and sampled chances. Confidence is the mean belief, not the bound or sample used to explore. Older saved states import with the default costs. - train_logreg labels each query with its cheapest successful stage; the classifier's online update no longer reinforces a correct stage costlier than the one it predicts. - evaluate replays each recommendation over per-stage outcomes. It charged every row's own cost whatever the router chose and priced the walk at 4x that, so every router showed a lift of 3x its spend. A query it cannot price is counted unscored, not charged zero, and joules_per_correct is None when nothing was correct. Co-Authored-By: Claude Opus 5.5 <[email protected]>
90 lines
3.1 KiB
Rust
90 lines
3.1 KiB
Rust
//! State import/export round-trip across all router families.
|
|
|
|
use eoc_route_learned::bandit::{LinUcbRouter, ThompsonSamplingRouter};
|
|
use eoc_route_learned::classifier::LogRegRouter;
|
|
use eoc_route_learned::matrix_factorization::MfRouter;
|
|
use eoc_route_learned::router::LearnedRouter;
|
|
|
|
#[test]
|
|
fn mf_roundtrip() {
|
|
let router = MfRouter::new(8, 4, 11);
|
|
let state = router.export_state();
|
|
let restored = MfRouter::import_state(state).expect("import");
|
|
assert_eq!(router.latent_dim(), restored.latent_dim());
|
|
assert_eq!(router.embedding_dim(), restored.embedding_dim());
|
|
}
|
|
|
|
#[test]
|
|
fn logreg_roundtrip() {
|
|
let router = LogRegRouter::new(8, 1e-4, 11);
|
|
let state = router.export_state();
|
|
let restored = LogRegRouter::import_state(state).expect("import");
|
|
assert_eq!(router.embedding_dim(), restored.embedding_dim());
|
|
}
|
|
|
|
#[test]
|
|
fn linucb_roundtrip() {
|
|
let router = LinUcbRouter::new(8, 1.0, 1.0);
|
|
let state = router.export_state();
|
|
let restored = LinUcbRouter::import_state(state).expect("import");
|
|
assert_eq!(router.dim, restored.dim);
|
|
assert_eq!(router.alpha, restored.alpha);
|
|
}
|
|
|
|
#[test]
|
|
fn thompson_roundtrip() {
|
|
let router = ThompsonSamplingRouter::new(7);
|
|
let state = router.export_state();
|
|
let restored = ThompsonSamplingRouter::import_state(state).expect("import");
|
|
assert_eq!(router.seed, restored.seed);
|
|
}
|
|
|
|
#[test]
|
|
fn rejects_wrong_algorithm_tag() {
|
|
let r = MfRouter::new(4, 4, 0);
|
|
let mut s = r.export_state();
|
|
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()
|
|
);
|
|
}
|