[Router] Add load-aware prefill admission and bounded policy proposals (#37843)

Signed-off-by: Vincent Gao <vincentbo@linux.alibaba.com>
Co-authored-by: Kangyan Zhou <zky314343421@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>

experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs

@@ -26,6 +26,7 @@ use sgl_router::config::{ActiveLoadConfig, ProxyConfig};
 
 use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
 use sgl_router::policies::cache_aware_zmq::CacheAwareZmqPolicy;
+use sgl_router::policies::engine_load::EngineLoadTable;
 use sgl_router::policies::kv_events::{compute_block_hashes, discovery::EventConfig, KvEventIndex};
 use sgl_router::policies::{Policy, SelectionContext};
 use sgl_router::tokenizer::TokenizerRegistry;
@@ -72,6 +73,7 @@ async fn zmq_indexer_routes_to_publishing_worker_e2e() {
             circuit_breaker: None,
             cache_aware: None,
             sticky: None,
+            affinity: None,
             fused: None,
             eligibility: None,
         },
@@ -116,6 +118,7 @@ async fn zmq_indexer_routes_to_publishing_worker_e2e() {
         kv_index.tree(),
         Arc::clone(&tokenizers),
         block_size_oracle,
+        EngineLoadTable::new(),
     );
 
     // 5. Register two workers. They share `127.0.0.1` so both
@@ -129,6 +132,8 @@ async fn zmq_indexer_routes_to_publishing_worker_e2e() {
         topic: String::new(),
         block_size,
         dp_size: 1,
+        load_port_base: None,
+        load_topic: None,
         is_bigram: false,
     };
     kv_index.add_worker(url_a, Some(preresolved.clone())).await;

experimental/sgl-router/tests/component/policies/fused_score.rs

@@ -12,14 +12,17 @@
 //! and the idlest at 1.0, so every assertion below holds either way.
 
 use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
+use sgl_router::policies::engine_load::{EngineLoadSnapshot, EngineWorkerLoad};
 use sgl_router::policies::kv_events::{
     compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId,
 };
 use sgl_router::policies::load_based::LoadBasedPolicy;
-use sgl_router::policies::scoring::{prefix_cache::PrefixCachePolicy, FusedScorePolicy};
+use sgl_router::policies::scoring::{
+    prefix_cache::PrefixCachePolicy, FusedScorePolicy, ScorePolicy,
+};
 use sgl_router::policies::{Policy, SelectionContext};
 use sgl_router::workers::Worker;
-use std::sync::Arc;
+use std::{collections::HashMap, sync::Arc, time::Instant};
 
 const BLOCK: usize = 4;
 
@@ -73,3 +76,85 @@ fn the_weight_override_steers_a_two_term_fusion_past_either_term_alone() {
         assert_eq!(got.id, want.id, "--fuse load_based={load_weight}");
     }
 }
+
+#[test]
+fn fused_load_based_term_uses_the_request_snapshot() {
+    let ws = vec![worker("w0"), worker("w1")];
+    // Local counters changed after the request snapshot and prefer w0.
+    let _after_snapshot: Vec<_> = (0..10).map(|_| ws[1].load_guard()).collect();
+    let snapshot = EngineLoadSnapshot::from_workers(
+        29,
+        HashMap::from([
+            (
+                ws[0].url.clone(),
+                EngineWorkerLoad {
+                    num_running_reqs: 50,
+                    num_waiting_reqs: 0,
+                    num_tokens: 0,
+                    max_total_num_tokens: 0,
+                    captured_at: Instant::now(),
+                },
+            ),
+            (
+                ws[1].url.clone(),
+                EngineWorkerLoad {
+                    num_running_reqs: 1,
+                    num_waiting_reqs: 0,
+                    num_tokens: 0,
+                    max_total_num_tokens: 0,
+                    captured_at: Instant::now(),
+                },
+            ),
+        ]),
+    );
+    let model = ModelId("tiny".into());
+    let ctx = SelectionContext::new(&model, None).with_load_snapshot(&snapshot);
+    let fused = FusedScorePolicy::new(vec![(Arc::new(LoadBasedPolicy::new()), None)])
+        .expect("load-based is fusable");
+
+    assert_eq!(
+        fused.select(&ws, &ctx).expect("must route").id,
+        ws[1].id,
+        "fused score must pass the request snapshot to the load-based term"
+    );
+}
+
+#[test]
+fn score_policy_forwards_the_request_snapshot_to_load_based() {
+    let ws = vec![worker("w0"), worker("w1")];
+    let _after_snapshot: Vec<_> = (0..10).map(|_| ws[1].load_guard()).collect();
+    let snapshot = EngineLoadSnapshot::from_workers(
+        31,
+        HashMap::from([
+            (
+                ws[0].url.clone(),
+                EngineWorkerLoad {
+                    num_running_reqs: 50,
+                    num_waiting_reqs: 0,
+                    num_tokens: 0,
+                    max_total_num_tokens: 0,
+                    captured_at: Instant::now(),
+                },
+            ),
+            (
+                ws[1].url.clone(),
+                EngineWorkerLoad {
+                    num_running_reqs: 1,
+                    num_waiting_reqs: 0,
+                    num_tokens: 0,
+                    max_total_num_tokens: 0,
+                    captured_at: Instant::now(),
+                },
+            ),
+        ]),
+    );
+    let model = ModelId("tiny".into());
+    let ctx = SelectionContext::new(&model, None).with_load_snapshot(&snapshot);
+    let score = ScorePolicy::new(Arc::new(LoadBasedPolicy::new()));
+
+    assert_eq!(
+        score.select(&ws, &ctx).expect("must route").id,
+        ws[1].id,
+        "ScorePolicy must preserve the load-based snapshot contract"
+    );
+}

experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs

@@ -41,6 +41,8 @@ async fn two_independent_subscribers_converge_to_same_tree_state() {
         topic: String::new(),
         block_size,
         dp_size: 1,
+        load_port_base: None,
+        load_topic: None,
         is_bigram: false,
     };
 
@@ -53,12 +55,7 @@ async fn two_independent_subscribers_converge_to_same_tree_state() {
     router_a.add_worker(worker_url, Some(cfg.clone())).await;
     router_b.add_worker(worker_url, Some(cfg.clone())).await;
 
-    // SUB-side handshake settle. Publishing before the subscribers
-    // finish their initial connect loses messages in PUB/SUB semantics;
-    // the polling loop below would then never converge.
-    tokio::time::sleep(Duration::from_millis(200)).await;
-
-    // 3. Publish a deterministic, multi-block event chain.
+    // 3. Build a deterministic, multi-block event chain.
     let tokens: Vec<u32> = (0..16).collect();
     let hashes = compute_block_hashes(&tokens, block_size as usize);
     assert!(
@@ -68,21 +65,21 @@ async fn two_independent_subscribers_converge_to_same_tree_state() {
     );
     let event_bytes = encode_block_stored_event(&hashes, None, &tokens, block_size);
     let payload = encode_event_batch(0.0, vec![event_bytes], Some(0));
-    publisher
-        .send(build_multipart(1, payload))
-        .await
-        .expect("publish BlockStored");
-
-    // 4. Poll both trees until both report the FULL chain matched. The
-    //    SUB→mpsc→pump→tree pipeline is async; loopback delivery is
-    //    reliable but not instantaneous.
+    // 4. Republish until both subscribers observe the event. PUB/SUB has
+    // no readiness acknowledgement, so a one-shot send can race a new
+    // subscriber's handshake under a parallel test load.
     let target = hashes.len();
     let key = KvWorkerId {
         url: worker_url.into(),
         dp_rank: 0,
     };
     let start = std::time::Instant::now();
+    let mut sequence = 1i64;
     loop {
+        publisher
+            .send(build_multipart(sequence, payload.clone()))
+            .await
+            .expect("publish BlockStored");
         let ma = router_a.tree().match_prefix(None, &hashes);
         let mb = router_b.tree().match_prefix(None, &hashes);
         let converged = ma.matched_blocks == target
@@ -112,6 +109,7 @@ async fn two_independent_subscribers_converge_to_same_tree_state() {
                 ma.matched_blocks, ma.workers, mb.matched_blocks, mb.workers,
             );
         }
+        sequence += 1;
         tokio::time::sleep(Duration::from_millis(20)).await;
     }
 
@@ -174,6 +172,8 @@ async fn two_subscribers_merge_events_from_two_publishers() {
         topic: String::new(),
         block_size,
         dp_size: 1,
+        load_port_base: None,
+        load_topic: None,
         is_bigram: false,
     };
     let cfg_y = EventConfig {
@@ -182,6 +182,8 @@ async fn two_subscribers_merge_events_from_two_publishers() {
         topic: String::new(),
         block_size,
         dp_size: 1,
+        load_port_base: None,
+        load_topic: None,
         is_bigram: false,
     };
 
@@ -193,10 +195,6 @@ async fn two_subscribers_merge_events_from_two_publishers() {
     router_b.add_worker(worker_x, Some(cfg_x.clone())).await;
     router_b.add_worker(worker_y, Some(cfg_y.clone())).await;
 
-    // Four SUB→PUB handshakes need to settle before publishing; missed
-    // SUBSCRIBE frames lose messages forever in PUB/SUB semantics.
-    tokio::time::sleep(Duration::from_millis(200)).await;
-
     // Two non-overlapping token streams → two distinct hash chains. The
     // gap between them (0..16 vs 1000..1016) keeps `compute_block_hashes`
     // outputs disjoint so a cross-attribution bug can't be masked by
@@ -221,15 +219,6 @@ async fn two_subscribers_merge_events_from_two_publishers() {
         )],
         Some(0),
     );
-    pub_x
-        .send(build_multipart(1, payload_x))
-        .await
-        .expect("publish on pub_x");
-    pub_y
-        .send(build_multipart(1, payload_y))
-        .await
-        .expect("publish on pub_y");
-
     let key_x = KvWorkerId {
         url: worker_x.into(),
         dp_rank: 0,
@@ -242,7 +231,16 @@ async fn two_subscribers_merge_events_from_two_publishers() {
     let target_y = hashes_y.len();
 
     let start = std::time::Instant::now();
+    let mut sequence = 1i64;
     loop {
+        pub_x
+            .send(build_multipart(sequence, payload_x.clone()))
+            .await
+            .expect("publish on pub_x");
+        pub_y
+            .send(build_multipart(sequence, payload_y.clone()))
+            .await
+            .expect("publish on pub_y");
         let ax = router_a.tree().match_prefix(None, &hashes_x);
         let ay = router_a.tree().match_prefix(None, &hashes_y);
         let bx = router_b.tree().match_prefix(None, &hashes_x);
@@ -298,6 +296,7 @@ async fn two_subscribers_merge_events_from_two_publishers() {
                 by.workers,
             );
         }
+        sequence += 1;
         tokio::time::sleep(Duration::from_millis(20)).await;
     }