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();
experimental/sgl-router/tests/e2e/infra/model_pool.py
+38 -10
View File
@@ -34,6 +34,18 @@ from .model_specs import get_model_spec
logger = logging.getLogger(__name__)
def _wait_for_process_group_exit(pgid: int, timeout: float) -> bool:
deadline = time.monotonic() + timeout
while True:
try:
os.killpg(pgid, 0)
except ProcessLookupError:
return True
if time.monotonic() >= deadline:
return False
time.sleep(0.1)
def _get_open_port() -> int:
"""Allocate an ephemeral TCP port in the range [20000, 55535].
@@ -72,6 +84,7 @@ class ModelInstance:
model_id: str
gpu_ids: list[int] = field(default_factory=list)
kv_events_endpoint: str | None = None
_shutdown_started: bool = field(default=False, init=False, repr=False)
def __enter__(self) -> "ModelInstance":
return self
@@ -80,16 +93,31 @@ class ModelInstance:
self.shutdown()
def shutdown(self) -> None:
if self.process is not None and self.process.poll() is None:
try:
self.process.send_signal(signal.SIGTERM)
try:
self.process.wait(timeout=60)
except subprocess.TimeoutExpired:
self.process.kill()
self.process.wait()
except ProcessLookupError:
pass
if self.process is None or self._shutdown_started:
return
self._shutdown_started = True
pgid = self.process.pid
try:
os.killpg(pgid, signal.SIGTERM)
except ProcessLookupError:
return
try:
self.process.wait(timeout=60)
except subprocess.TimeoutExpired:
pass
if _wait_for_process_group_exit(pgid, timeout=30):
return
try:
os.killpg(pgid, signal.SIGKILL)
except ProcessLookupError:
return
self.process.wait()
if not _wait_for_process_group_exit(pgid, timeout=5):
raise RuntimeError(f"worker process group {pgid} did not exit")
def spawn_worker(
experimental/sgl-router/tests/e2e/infra/test_model_pool.py
+50
View File
@@ -0,0 +1,50 @@
import signal
from infra import model_pool
class _Process:
pid = 1234
def __init__(self):
self.wait_timeouts = []
def poll(self):
return None
def send_signal(self, sig):
raise AssertionError(f"signaled only the parent process: {sig}")
def wait(self, timeout=None):
self.wait_timeouts.append(timeout)
return 0
def test_shutdown_waits_for_the_worker_process_group(monkeypatch):
process = _Process()
signals = []
probes = iter([True, True, False])
def killpg(pgid, sig):
signals.append((pgid, sig))
if sig == 0 and not next(probes):
raise ProcessLookupError
monkeypatch.setattr(model_pool.os, "killpg", killpg)
monkeypatch.setattr(model_pool.time, "sleep", lambda _: None)
instance = model_pool.ModelInstance(
url="http://127.0.0.1:30000",
port=30000,
process=process,
model_id="qwen3-0.6b",
)
instance.shutdown()
assert signals == [
(process.pid, signal.SIGTERM),
(process.pid, 0),
(process.pid, 0),
(process.pid, 0),
]
assert process.wait_timeouts == [60]
experimental/sgl-router/tests/e2e/k8s_integration/conftest.py
+55
View File
@@ -11,6 +11,7 @@ Teardown: ./tests/e2e/k8s_integration/setup.sh teardown
from __future__ import annotations
import json
import logging
import socket
import subprocess
@@ -97,6 +98,60 @@ def _wait_for_pod_ready(
)
def _wait_for_replacement_pod_ready(
old_pod: str,
selector: str,
namespace: str = NAMESPACE,
timeout: int = 120,
interval: float = 0.5,
) -> str:
deadline = time.time() + timeout
last_observed = "no pods"
while time.time() < deadline:
result = _kubectl(
"get",
"pods",
"-n",
namespace,
"-l",
selector,
"-o",
"json",
check=False,
)
if getattr(result, "returncode", 0) == 0:
pods = json.loads(result.stdout or "{}").get("items", [])
names = [pod.get("metadata", {}).get("name", "") for pod in pods]
last_observed = ", ".join(filter(None, names)) or "no pods"
if old_pod not in names:
for pod in sorted(
pods, key=lambda item: item.get("metadata", {}).get("name", "")
):
metadata = pod.get("metadata", {})
status = pod.get("status", {})
ready = any(
condition.get("type") == "Ready"
and condition.get("status") == "True"
for condition in status.get("conditions", [])
)
if (
metadata.get("name") != old_pod
and not metadata.get("deletionTimestamp")
and status.get("phase") == "Running"
and ready
):
return metadata["name"]
time.sleep(interval)
raise TimeoutError(
f"No ready replacement for pod {old_pod!r} with selector {selector!r} "
f"after {timeout}s; last observed: {last_observed}"
)
def _wait_for_port(port: int, proc: subprocess.Popen, timeout: int = 15) -> None:
"""Poll until a TCP connection to localhost:port succeeds."""
deadline = time.time() + timeout
experimental/sgl-router/tests/e2e/k8s_integration/test_conftest_helpers.py
+54
View File
@@ -0,0 +1,54 @@
import json
from types import SimpleNamespace
import conftest as k8s_conftest
def _pod(name: str, phase: str, ready: bool) -> dict:
return {
"metadata": {"name": name},
"status": {
"phase": phase,
"conditions": [
{
"type": "Ready",
"status": "True" if ready else "False",
}
],
},
}
def test_wait_for_replacement_pod_ignores_old_and_pending_pods(monkeypatch):
old_pod = "sgl-router-old"
new_pod = "sgl-router-new"
responses = iter(
[
[_pod(old_pod, "Running", True)],
[
_pod(old_pod, "Running", True),
_pod(new_pod, "Running", True),
],
[_pod(new_pod, "Pending", False)],
[_pod(new_pod, "Running", True)],
]
)
calls = []
def fake_kubectl(*args, **kwargs):
calls.append((args, kwargs))
return SimpleNamespace(stdout=json.dumps({"items": next(responses)}))
monkeypatch.setattr(k8s_conftest, "_kubectl", fake_kubectl)
monkeypatch.setattr(k8s_conftest.time, "sleep", lambda _: None)
replacement = k8s_conftest._wait_for_replacement_pod_ready(
old_pod,
"app=sgl-router",
timeout=5,
interval=0,
)
assert replacement == new_pod
assert len(calls) == 4
assert all("-o" in args and "json" in args for args, _ in calls)
experimental/sgl-router/tests/e2e/k8s_integration/test_lifecycle.py
+5 -5
View File
@@ -13,10 +13,7 @@ by driving the deployment scale.
from __future__ import annotations
import logging
import httpx
import pytest
from conftest import (
NAMESPACE,
_cleanup_port_forward,
@@ -24,7 +21,7 @@ from conftest import (
_poll_until,
_port_forward_start,
_wait_for_deployment_ready,
logger,
_wait_for_replacement_pod_ready,
)
ROUTER_RESTART_PORT = 8092
@@ -132,7 +129,10 @@ class TestRouterRestart:
_cleanup_port_forward("router-restart-pre-kill", pf_holder[0])
pf_holder[0] = None
_wait_for_deployment_ready("sgl-router")
if old_pod:
_wait_for_replacement_pod_ready(old_pod, "app=sgl-router")
else:
_wait_for_deployment_ready("sgl-router")
pf_holder[0] = _port_forward_start(
NAMESPACE, "sgl-router", ROUTER_RESTART_PORT, 8090
experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs
+3
View File
@@ -25,6 +25,7 @@ use sgl_router::config::{
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::engine_load::EngineLoadTable;
use sgl_router::policies::factory::build_registry;
use sgl_router::policies::kv_events::{BlockSizeOracle, HashTree};
use sgl_router::proxy::Proxy;
@@ -54,6 +55,7 @@ fn config() -> Config {
circuit_breaker: None,
cache_aware: Some(CacheAwareConfig::default()),
sticky: None,
affinity: None,
fused: None,
eligibility: None,
},
@@ -88,6 +90,7 @@ fn build_ctx(url: String) -> Arc<AppContext> {
Arc::new(HashTree::new()),
Arc::clone(&tokenizers),
BlockSizeOracle::new(),
EngineLoadTable::new(),
)
.unwrap(),
);
experimental/sgl-router/tests/proxy/chat_routing.rs
+1
View File
@@ -37,6 +37,7 @@ fn config_for(_worker_url: &str) -> Config {
circuit_breaker: None,
cache_aware: None,
sticky: None,
affinity: None,
fused: None,
eligibility: None,
},
experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs
+1
View File
@@ -29,6 +29,7 @@ pub fn config() -> Config {
circuit_breaker: None,
cache_aware: Some(CacheAwareConfig::default()),
sticky: None,
affinity: None,
fused: None,
eligibility: None,
},
experimental/sgl-router/tests/proxy/external_indexer_routing.rs
+2
View File
@@ -17,6 +17,7 @@ use sgl_kv_indexer::{
server_builder, GrpcPrefixIndex, InMemoryKvIndexerBackend, KvIndexerService, PrefixIndexConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::engine_load::EngineLoadTable;
use sgl_router::policies::factory::build_registry;
use sgl_router::policies::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree};
use sgl_router::policies::request_tokens_for;
@@ -94,6 +95,7 @@ async fn external_indexer_routes_to_the_cached_worker() {
Arc::new(HashTree::new()),
Arc::clone(&tokenizers),
Arc::clone(&oracle),
EngineLoadTable::new(),
)
.unwrap(),
);
experimental/sgl-router/tests/proxy/failover.rs
+1
View File
@@ -41,6 +41,7 @@ async fn failover_when_one_worker_dies() {
}),
cache_aware: None,
sticky: None,
affinity: None,
fused: None,
eligibility: None,
},
experimental/sgl-router/tests/proxy/graceful_shutdown.rs
+1
View File
@@ -47,6 +47,7 @@ fn build_ctx_with_worker(worker_url: &str) -> Arc<AppContext> {
circuit_breaker: None,
cache_aware: None,
sticky: None,
affinity: None,
fused: None,
eligibility: None,
},
experimental/sgl-router/tests/proxy/header_forwarding.rs
+1
View File
@@ -34,6 +34,7 @@ async fn forwards_whitelisted_headers_strips_others() {
circuit_breaker: None,
cache_aware: None,
sticky: None,
affinity: None,
fused: None,
eligibility: None,
},
experimental/sgl-router/tests/proxy/main.rs
+1
View File
@@ -19,6 +19,7 @@ mod header_forwarding;
mod pd_bootstrap_injection;
mod pd_pool_isolation;
mod roundrobin_input_ids;
mod shared_prefill_admission;
mod sticky_input_ids;
mod sticky_routing;
mod timeout;
experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs
+1
View File
@@ -49,6 +49,7 @@ fn config() -> Config {
circuit_breaker: None,
cache_aware: None,
sticky: None,
affinity: None,
fused: None,
eligibility: None,
},
experimental/sgl-router/tests/proxy/pd_pool_isolation.rs
+1
View File
@@ -48,6 +48,7 @@ fn config() -> Config {
circuit_breaker: None,
cache_aware: None,
sticky: None,
affinity: None,
fused: None,
eligibility: None,
},
experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs
+1
View File
@@ -45,6 +45,7 @@ fn config() -> Config {
circuit_breaker: None,
cache_aware: None,
sticky: None,
affinity: None,
fused: None,
eligibility: None,
},
experimental/sgl-router/tests/proxy/shared_prefill_admission.rs
+427
View File
@@ -0,0 +1,427 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use axum::body::Body;
use axum::http::{Request, StatusCode};
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::engine_load::LoadStat;
use sgl_router::policies::{
CacheCandidate, CacheCandidateProposal, Policy, PolicyRegistry, PrefillProposal, ProposalKind,
SelectionContext, SelectionProposal,
};
use sgl_router::proxy::Proxy;
use sgl_router::server::app::build_router;
use sgl_router::server::app_context::AppContext;
use sgl_router::tokenizer::TokenizerRegistry;
use sgl_router::workers::{Worker, WorkerRegistry};
use tower::ServiceExt;
use crate::common::mock_worker::MockWorker;
#[derive(Debug)]
struct AdmissionProbePolicy {
primary: Arc<Worker>,
backup: Arc<Worker>,
committed: Arc<Mutex<Option<String>>>,
}
impl Policy for AdmissionProbePolicy {
fn select(&self, _: &[Arc<Worker>], _: &SelectionContext<'_>) -> Option<Arc<Worker>> {
panic!("chat routing must resolve the prefill proposal before selection")
}
fn propose(&self, _: &[Arc<Worker>], _: &SelectionContext<'_>) -> Option<SelectionProposal> {
Some(
SelectionProposal::with_backup(Arc::clone(&self.primary), Arc::clone(&self.backup))
.with_kind(ProposalKind::SessionAffinity),
)
}
fn uses_shared_prefill_admission(&self) -> bool {
true
}
fn commit_prefill_selection(
&self,
_: &SelectionContext<'_>,
_: ProposalKind,
selected: &Arc<Worker>,
) {
*self.committed.lock().unwrap() = Some(selected.id.0.clone());
}
}
#[derive(Debug)]
struct EmptyPolicy;
impl Policy for EmptyPolicy {
fn select(&self, _: &[Arc<Worker>], _: &SelectionContext<'_>) -> Option<Arc<Worker>> {
None
}
}
#[derive(Debug)]
struct InvalidPairPolicy {
outsider: Arc<Worker>,
}
impl Policy for InvalidPairPolicy {
fn select(&self, _: &[Arc<Worker>], _: &SelectionContext<'_>) -> Option<Arc<Worker>> {
panic!("chat routing must use the invalid prefill proposal")
}
fn propose(&self, _: &[Arc<Worker>], _: &SelectionContext<'_>) -> Option<SelectionProposal> {
Some(SelectionProposal::primary(Arc::clone(&self.outsider)))
}
fn uses_shared_prefill_admission(&self) -> bool {
true
}
}
#[derive(Debug)]
struct CacheCandidatesPolicy {
worker: Arc<Worker>,
}
#[derive(Debug)]
struct SnapshotProbePolicy {
worker: Arc<Worker>,
needs_snapshot: bool,
observed_snapshot: Arc<Mutex<Option<bool>>>,
}
impl Policy for SnapshotProbePolicy {
fn select(&self, _: &[Arc<Worker>], ctx: &SelectionContext<'_>) -> Option<Arc<Worker>> {
*self.observed_snapshot.lock().unwrap() = Some(ctx.load_snapshot().is_some());
Some(Arc::clone(&self.worker))
}
fn needs_load_snapshot(&self) -> bool {
self.needs_snapshot
}
}
impl Policy for CacheCandidatesPolicy {
fn select(&self, _: &[Arc<Worker>], _: &SelectionContext<'_>) -> Option<Arc<Worker>> {
panic!("chat routing must use the cache-candidate proposal")
}
fn propose_prefill(
&self,
_: &[Arc<Worker>],
_: &SelectionContext<'_>,
) -> Option<PrefillProposal> {
Some(PrefillProposal::CacheCandidates(CacheCandidateProposal {
candidates: vec![CacheCandidate {
worker: Arc::clone(&self.worker),
matched_prefix_tokens: 1,
uncached_tokens: 1,
candidate_range_id: "global".into(),
max_pending_prefill_tokens: None,
}],
cache_switch_margin_tokens: 0,
}))
}
fn needs_load_snapshot(&self) -> bool {
true
}
}
fn config(policy: PolicyKind) -> Config {
Config {
server: ServerConfig {
host: "0".into(),
port: 0,
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy,
circuit_breaker: None,
cache_aware: None,
sticky: None,
affinity: None,
fused: None,
eligibility: None,
},
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
}
}
struct TestFixture {
ctx: Arc<AppContext>,
backends: Vec<MockWorker>,
workers: Vec<Arc<Worker>>,
}
async fn fixture(
policy_kind: PolicyKind,
build_policy: impl FnOnce(&[Arc<Worker>]) -> Arc<dyn Policy>,
) -> TestFixture {
let backends = vec![
MockWorker::start(vec![]).await,
MockWorker::start(vec![]).await,
];
let cfg = config(policy_kind);
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
let registry = Arc::new(WorkerRegistry::default());
for (index, backend) in backends.iter().enumerate() {
registry
.add(WorkerSpec {
id: WorkerId(if index == 0 { "primary" } else { "backup" }.into()),
url: backend.url.clone(),
mode: WorkerMode::Plain,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: None,
})
.unwrap();
}
let registered = registry.workers_for(&ModelId("tiny".into()));
let workers = ["primary", "backup"]
.into_iter()
.map(|id| {
registered
.iter()
.find(|worker| worker.id.0 == id)
.cloned()
.unwrap()
})
.collect::<Vec<_>>();
let policies = Arc::new(PolicyRegistry::default());
policies.insert(ModelId("tiny".into()), build_policy(&workers));
let ctx = Arc::new(AppContext::new(
cfg,
tokenizers,
Arc::new(Proxy::new(Duration::from_secs(5)).unwrap()),
registry,
policies,
));
TestFixture {
ctx,
backends,
workers,
}
}
async fn send_chat(ctx: &Arc<AppContext>) -> StatusCode {
build_router(Arc::clone(ctx))
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header("content-type", "application/json")
.body(Body::from(
serde_json::to_vec(&serde_json::json!({
"model": "tiny",
"messages": [{"role": "user", "content": "hi"}],
}))
.unwrap(),
))
.unwrap(),
)
.await
.unwrap()
.status()
}
#[tokio::test]
async fn chat_attaches_load_snapshot_only_when_the_policy_needs_it() {
for needs_snapshot in [false, true] {
let observed_snapshot = Arc::new(Mutex::new(None));
let fixture = fixture(PolicyKind::RoundRobin, |workers| {
Arc::new(SnapshotProbePolicy {
worker: Arc::clone(&workers[0]),
needs_snapshot,
observed_snapshot: Arc::clone(&observed_snapshot),
})
})
.await;
assert_eq!(send_chat(&fixture.ctx).await, StatusCode::OK);
assert_eq!(*observed_snapshot.lock().unwrap(), Some(needs_snapshot));
}
}
fn assert_failure_metric(ctx: &AppContext, policy: &str, expected_reason: &str) {
let metrics = ctx.metrics.render();
assert!(metrics.contains(&format!(
"sgl_router_policy_selection_failures_total{{policy=\"{policy}\",reason=\"{expected_reason}\"}} 1"
)));
for other in [
"prefill_admission_exhausted",
"cache_candidates_exhausted",
"proposal_empty",
] {
if other != expected_reason {
assert!(
!metrics.contains(&format!("reason=\"{other}\"")),
"{metrics}"
);
}
}
}
#[tokio::test]
async fn chat_commits_the_admitted_prefill_backup() {
let committed = Arc::new(Mutex::new(None));
let fixture = fixture(PolicyKind::SessionAware, |workers| {
Arc::new(AdmissionProbePolicy {
primary: Arc::clone(&workers[0]),
backup: Arc::clone(&workers[1]),
committed: Arc::clone(&committed),
})
})
.await;
fixture.ctx.engine_load.set(
&fixture.workers[0].url,
0,
LoadStat {
num_running_reqs: 1,
num_waiting_reqs: 0,
num_tokens: 100,
max_total_num_tokens: 100,
},
Instant::now(),
);
assert_eq!(send_chat(&fixture.ctx).await, StatusCode::OK);
assert!(fixture.backends[0]
.captured
.lock()
.unwrap()
.last_body
.is_none());
assert!(fixture.backends[1]
.captured
.lock()
.unwrap()
.last_body
.is_some());
assert_eq!(committed.lock().unwrap().as_deref(), Some("backup"));
assert!(!fixture
.ctx
.metrics
.render()
.contains("sgl_router_policy_selection_failures_total{"));
}
#[tokio::test]
async fn capacity_exhaustion_does_not_return_503() {
let committed = Arc::new(Mutex::new(None));
let fixture = fixture(PolicyKind::SessionAware, |workers| {
Arc::new(AdmissionProbePolicy {
primary: Arc::clone(&workers[0]),
backup: Arc::clone(&workers[1]),
committed: Arc::clone(&committed),
})
})
.await;
for worker in &fixture.workers {
fixture.ctx.engine_load.set(
&worker.url,
0,
LoadStat {
num_running_reqs: 0,
num_waiting_reqs: 0,
num_tokens: 100,
max_total_num_tokens: 100,
},
Instant::now(),
);
}
assert_eq!(send_chat(&fixture.ctx).await, StatusCode::OK);
assert_eq!(
fixture
.backends
.iter()
.filter(|backend| backend.captured.lock().unwrap().last_body.is_some())
.count(),
1,
"the request must be dispatched to exactly one legal backend"
);
assert!(matches!(
committed.lock().unwrap().as_deref(),
Some("primary" | "backup")
));
assert!(!fixture
.ctx
.metrics
.render()
.contains("sgl_router_policy_selection_failures_total{"));
}
#[tokio::test]
async fn chat_records_proposal_empty() {
let fixture = fixture(PolicyKind::SessionAware, |_| Arc::new(EmptyPolicy)).await;
assert_eq!(
send_chat(&fixture.ctx).await,
StatusCode::SERVICE_UNAVAILABLE
);
assert_failure_metric(&fixture.ctx, "session_aware", "proposal_empty");
}
#[tokio::test]
async fn chat_records_prefill_admission_exhausted_for_out_of_range_primary() {
let outsider = Arc::new(Worker::new(WorkerSpec {
id: WorkerId("outsider".into()),
url: "http://outsider:30000".into(),
mode: WorkerMode::Plain,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: None,
}));
let fixture = fixture(PolicyKind::SessionAware, |_| {
Arc::new(InvalidPairPolicy {
outsider: Arc::clone(&outsider),
})
})
.await;
assert_eq!(
send_chat(&fixture.ctx).await,
StatusCode::SERVICE_UNAVAILABLE
);
assert_failure_metric(&fixture.ctx, "session_aware", "prefill_admission_exhausted");
}
#[tokio::test]
async fn chat_records_cache_candidates_exhausted() {
let fixture = fixture(PolicyKind::CacheAware, |workers| {
Arc::new(CacheCandidatesPolicy {
worker: Arc::clone(&workers[0]),
})
})
.await;
fixture.ctx.engine_load.set(
&fixture.workers[0].url,
0,
LoadStat {
num_running_reqs: 1,
num_waiting_reqs: 0,
num_tokens: 100,
max_total_num_tokens: 100,
},
Instant::now(),
);
assert_eq!(
send_chat(&fixture.ctx).await,
StatusCode::SERVICE_UNAVAILABLE
);
assert_failure_metric(&fixture.ctx, "cache_aware", "cache_candidates_exhausted");
}
experimental/sgl-router/tests/proxy/sticky_input_ids.rs
+1
View File
@@ -64,6 +64,7 @@ fn config() -> Config {
idle_secs: 3600,
eviction_interval_secs: 3600,
}),
affinity: None,
fused: None,
eligibility: None,
},
experimental/sgl-router/tests/proxy/sticky_routing.rs
+1
View File
@@ -50,6 +50,7 @@ fn build_sticky_ctx(header_name: &str, worker_urls: &[String]) -> Arc<AppContext
idle_secs: 3600,
eviction_interval_secs: 3600,
}),
affinity: None,
fused: None,
eligibility: None,
},
experimental/sgl-router/tests/proxy/timeout.rs
+1
View File
@@ -41,6 +41,7 @@ fn config(_worker_url: &str) -> Config {
circuit_breaker: None,
cache_aware: None,
sticky: None,
affinity: None,
fused: None,
eligibility: None,
},