[Fix] Reduce power of two to constant time (#28228)

Signed-off-by: JiangJiaWei1103 <waynechuang97@gmail.com>

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);
+                let i = rng.gen_range(0..len);
-                chosen.sort_by_key(|w| w.active_load());
+                let mut j = rng.gen_range(0..len - 1);
-                Some(chosen[0].clone())
+                if j >= i {
+                    j += 1;
+                }
+                Some(std::cmp::min_by_key(&workers[i], &workers[j], |w| w.active_load()).clone())
             }
         }
     }

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:?}"
+    );
+}