minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Vincent Gao
2026-09-05 11:53:26 +08:00
committed by GitHub
co-authored by Kangyan Zhou Claude Opus 4.8 Shangming Cai
parent d50e9a9756
commit ecd97de1fc
55 changed files with 6834 additions and 348 deletions
Download Patch File Download Diff File Expand all files Collapse all files
experimental/sgl-router/tests/component/discovery/static_urls.rs
+1
View File
@@ -133,6 +133,7 @@ async fn static_urls_pd_role_resolved_end_to_end() {
circuit_breaker: None,
cache_aware: None,
sticky: None,
affinity: None,
fused: None,
eligibility: None,
},
experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs
+5
View File
@@ -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
+87 -2
View File
@@ -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
+26 -27
View File
@@ -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;
}
experimental/sgl-router/tests/component/workers/manager.rs
+96 -24
View File
@@ -7,9 +7,9 @@ use sgl_router::discovery::{DiscoveryEvent, ModelId, WorkerId, WorkerMode, Worke
use sgl_router::workers::{manager, WorkerRegistry};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use std::time::Duration;
use tokio::net::TcpListener;
use tokio::sync::{mpsc, oneshot};
use tokio::sync::{mpsc, oneshot, Barrier};
/// Spin up a tiny fake worker that returns `body` on `GET /server_info`.
/// Returns the worker base URL and a shutdown channel.
@@ -48,6 +48,19 @@ fn spec_for(id: &str, url: &str, mode: WorkerMode) -> WorkerSpec {
}
}
async fn wait_until(condition: impl Fn() -> bool, description: &str) {
tokio::time::timeout(Duration::from_secs(2), async {
loop {
if condition() {
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.unwrap_or_else(|_| panic!("timed out waiting for {description}"));
}
#[tokio::test]
async fn manager_processes_added_then_removed() {
let (url_a, _s_a) = spawn_fake_worker(json!({"served_model_name": "m"})).await;
@@ -72,8 +85,11 @@ async fn manager_processes_added_then_removed() {
.await
.unwrap();
// Give the manager time to drain.
tokio::time::sleep(Duration::from_millis(200)).await;
wait_until(
|| registry.workers_for(&ModelId("m".into())).len() == 2,
"both workers to register",
)
.await;
assert_eq!(registry.workers_for(&ModelId("m".into())).len(), 2);
tx.send(DiscoveryEvent::Removed {
@@ -81,7 +97,11 @@ async fn manager_processes_added_then_removed() {
})
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
wait_until(
|| registry.workers_for(&ModelId("m".into())).len() == 1,
"removed worker to leave the registry",
)
.await;
assert_eq!(registry.workers_for(&ModelId("m".into())).len(), 1);
drop(tx);
@@ -103,7 +123,16 @@ async fn manager_handles_mode_changed() {
)))
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(200)).await;
wait_until(
|| {
registry
.workers_for_mode(&ModelId("m".into()), WorkerMode::Prefill)
.len()
== 1
},
"prefill worker to register",
)
.await;
assert_eq!(
registry
.workers_for_mode(&ModelId("m".into()), WorkerMode::Prefill)
@@ -117,7 +146,19 @@ async fn manager_handles_mode_changed() {
})
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
wait_until(
|| {
registry
.workers_for_mode(&ModelId("m".into()), WorkerMode::Prefill)
.is_empty()
&& registry
.workers_for_mode(&ModelId("m".into()), WorkerMode::Decode)
.len()
== 1
},
"worker mode to change to decode",
)
.await;
assert_eq!(
registry
.workers_for_mode(&ModelId("m".into()), WorkerMode::Prefill)
@@ -280,7 +321,15 @@ async fn manager_handles_duplicate_added_as_upsert() {
)))
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(300)).await;
wait_until(
|| {
registry
.get(&WorkerId("w1".into()))
.is_some_and(|worker| worker.url == url_second)
},
"replacement Added event to update the worker",
)
.await;
assert_eq!(
registry.workers_for(&ModelId("m1".into())).len(),
@@ -319,6 +368,32 @@ async fn spawn_slow_worker(body: Value, delay: Duration) -> (String, oneshot::Se
(format!("http://127.0.0.1:{port}"), tx)
}
async fn spawn_gated_worker(body: Value, gate: Arc<Barrier>) -> (String, oneshot::Sender<()>) {
let body = Arc::new(body);
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
let app = Router::new().route(
"/server_info",
get(move || {
let body = body.clone();
let gate = Arc::clone(&gate);
async move {
gate.wait().await;
Json((*body).clone())
}
}),
);
let (tx, rx) = oneshot::channel::<()>();
tokio::spawn(async move {
let _ = axum::serve(listener, app)
.with_graceful_shutdown(async move {
let _ = rx.await;
})
.await;
});
(format!("http://127.0.0.1:{port}"), tx)
}
/// Spawn a fake worker that counts each `GET /server_info` hit in the
/// returned `AtomicUsize`. Used to assert the manager makes exactly
/// one round-trip per worker.
@@ -350,25 +425,22 @@ async fn spawn_counting_worker(body: Value) -> (String, Arc<AtomicUsize>, onesho
(format!("http://127.0.0.1:{port}"), counter, tx)
}
/// Registration must run in parallel across multiple `Added` events.
/// Each fake worker delays its `/server_info` by 200ms; with sequential
/// processing the manager would take ≥1000ms for 5 workers. We allow
/// up to 600ms (3x the per-fetch delay) as a generous bound that still
/// rejects the sequential implementation.
/// Registration must start every `/server_info` fetch before any response is
/// released. A sequential manager stalls at the first worker's barrier.
#[tokio::test]
async fn added_events_run_in_parallel() {
let delay = Duration::from_millis(200);
let n = 5;
let gate = Arc::new(Barrier::new(n + 1));
let mut workers = Vec::new();
for _ in 0..n {
workers.push(spawn_slow_worker(json!({"served_model_name": "m"}), delay).await);
workers
.push(spawn_gated_worker(json!({"served_model_name": "m"}), Arc::clone(&gate)).await);
}
let (tx, rx) = mpsc::channel(16);
let registry = Arc::new(WorkerRegistry::default());
let h = tokio::spawn(manager::run(rx, registry.clone()));
let start = Instant::now();
for (i, (url, _s)) in workers.iter().enumerate() {
tx.send(DiscoveryEvent::Added(spec_for(
&format!("w{i}"),
@@ -378,22 +450,22 @@ async fn added_events_run_in_parallel() {
.await
.unwrap();
}
let registered = tokio::time::timeout(Duration::from_secs(5), async {
assert!(
tokio::time::timeout(Duration::from_secs(2), gate.wait())
.await
.is_ok(),
"manager did not start all {n} /server_info requests concurrently"
);
let registered = tokio::time::timeout(Duration::from_secs(2), async {
loop {
if registry.workers_for(&ModelId("m".into())).len() == n {
return true;
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await;
let elapsed = start.elapsed();
assert!(registered.is_ok(), "manager failed to register {n} workers");
assert!(
elapsed < Duration::from_millis(600),
"registration of {n} workers took {elapsed:?}; sequential per-worker /server_info \
fetches would take ≥1000ms — parallel spawn is required"
);
drop(tx);
h.await.unwrap();