diff --git a/experimental/sgl-router/src/policies/power_of_two.rs b/experimental/sgl-router/src/policies/power_of_two.rs index a6c4aacf6..554f70511 100644 --- a/experimental/sgl-router/src/policies/power_of_two.rs +++ b/experimental/sgl-router/src/policies/power_of_two.rs @@ -3,7 +3,7 @@ use crate::policies::{Policy, SelectionContext}; use crate::workers::Worker; -use rand::seq::IteratorRandom; +use rand::Rng; use std::sync::Arc; #[derive(Debug, Default)] @@ -20,11 +20,14 @@ impl Policy for PowerOfTwoChoicesPolicy { match workers.len() { 0 => None, 1 => Some(workers[0].clone()), - _ => { + len => { let mut rng = rand::thread_rng(); - let mut chosen = workers.iter().choose_multiple(&mut rng, 2); - chosen.sort_by_key(|w| w.active_load()); - Some(chosen[0].clone()) + let i = rng.gen_range(0..len); + let mut j = rng.gen_range(0..len - 1); + if j >= i { + j += 1; + } + Some(std::cmp::min_by_key(&workers[i], &workers[j], |w| w.active_load()).clone()) } } } diff --git a/experimental/sgl-router/tests/component/policies/power_of_two.rs b/experimental/sgl-router/tests/component/policies/power_of_two.rs index 21178e372..183f9f658 100644 --- a/experimental/sgl-router/tests/component/policies/power_of_two.rs +++ b/experimental/sgl-router/tests/component/policies/power_of_two.rs @@ -71,3 +71,29 @@ fn single_worker_returns_it() { let ctx = SelectionContext::new(&model_id, None); assert_eq!(p.select(&ws, &ctx).unwrap().id.0, "only"); } + +#[test] +fn all_workers_reachable() { + // Ensure all workers are selectable under equal load to prevent an + // off-by-one error from skipping any worker during sampling + let workers = vec![ + worker("a"), + worker("b"), + worker("c"), + worker("d"), + worker("e"), + ]; + let p = PowerOfTwoChoicesPolicy::new(); + let model_id = ModelId("m".into()); + let ctx = SelectionContext::new(&model_id, None); + let mut seen = std::collections::HashSet::new(); + for _ in 0..1000 { + let w = p.select(&workers, &ctx).unwrap(); + seen.insert(w.id.0.clone()); + } + assert_eq!( + seen.len(), + workers.len(), + "every worker should be reachable, saw {seen:?}" + ); +}