sgl-router: experimental Rust HTTP router for SGLang worker pools (#25851)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Kangyan-Zhou
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
aae04b1241
commit
6e8fe176be
@@ -0,0 +1,334 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
//! PD-disagg bootstrap-room injection + dual-dispatch — end-to-end
|
||||
//! at the HTTP layer using MockWorkers.
|
||||
//!
|
||||
//! Asserts the router-side contract for SGLang disagg-prefill HTTP mode:
|
||||
//!
|
||||
//! * Every PD-mode `/v1/chat/completions` request fans out to BOTH a
|
||||
//! prefill and a decode worker (the prefill is `tokio::spawn`'d in
|
||||
//! the background; the decode is awaited for the client response).
|
||||
//! * Both bodies carry the SAME flat top-level fields:
|
||||
//! - `bootstrap_host` = the chosen prefill worker's host
|
||||
//! - `bootstrap_port` = the chosen prefill worker's bootstrap port
|
||||
//! - `bootstrap_room` = a random u64 in `[0, i64::MAX]` (63-bit)
|
||||
//! * Plain-mode requests do NOT carry any `bootstrap_*` field — the
|
||||
//! injection step is gated on `worker.mode() == Prefill`.
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::http::{Request, StatusCode};
|
||||
use bytes::Bytes;
|
||||
use serde_json::{json, Value};
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, Config, DiscoveryBackend, DiscoveryConfig, ModelConfig, ObservabilityConfig,
|
||||
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
use sgl_router::policies::factory::build_registry_with_defaults;
|
||||
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::WorkerRegistry;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tower::ServiceExt;
|
||||
|
||||
fn config() -> Config {
|
||||
Config {
|
||||
server: ServerConfig {
|
||||
host: "0".into(),
|
||||
port: 0,
|
||||
},
|
||||
observability: ObservabilityConfig::default(),
|
||||
models: vec![ModelConfig {
|
||||
id: "tiny".into(),
|
||||
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
|
||||
policy: PolicyKind::RoundRobin,
|
||||
circuit_breaker: None,
|
||||
cache_aware: None,
|
||||
}],
|
||||
discovery: DiscoveryConfig {
|
||||
backend: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
},
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn build_ctx(specs: Vec<WorkerSpec>) -> Arc<AppContext> {
|
||||
let cfg = config();
|
||||
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
|
||||
let registry = Arc::new(WorkerRegistry::default());
|
||||
for s in specs {
|
||||
let _ = registry.add(s);
|
||||
}
|
||||
let policies = Arc::new(build_registry_with_defaults(&cfg).unwrap());
|
||||
let proxy = Arc::new(Proxy::new(Duration::from_secs(5)).unwrap());
|
||||
Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies))
|
||||
}
|
||||
|
||||
fn chat_request() -> Request<Body> {
|
||||
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()
|
||||
}
|
||||
|
||||
/// Pattern-B dispatch: prefill is `tokio::spawn`'d as a detached task
|
||||
/// so the client response can return as soon as decode is reachable —
|
||||
/// the prefill body is captured *eventually* but may not be present
|
||||
/// when the handler returns. Poll with a short bound rather than
|
||||
/// sleeping a fixed duration.
|
||||
async fn await_captured_body(
|
||||
mock: &crate::common::mock_worker::MockWorker,
|
||||
timeout: Duration,
|
||||
label: &str,
|
||||
) -> Bytes {
|
||||
let start = std::time::Instant::now();
|
||||
loop {
|
||||
// Release the `std::sync::Mutex` guard before the sleep.await
|
||||
// (clippy: await_holding_lock).
|
||||
let captured = mock.captured.lock().unwrap().last_body.clone();
|
||||
if let Some(b) = captured {
|
||||
return b;
|
||||
}
|
||||
if start.elapsed() > timeout {
|
||||
panic!("{label}: no request body captured within {timeout:?}");
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(5)).await;
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_body(b: &Bytes) -> Value {
|
||||
serde_json::from_slice(b).expect("body must be valid JSON")
|
||||
}
|
||||
|
||||
/// Helper: extract bootstrap_host as &str.
|
||||
fn bootstrap_host(v: &Value) -> Option<&str> {
|
||||
v.get("bootstrap_host").and_then(|x| x.as_str())
|
||||
}
|
||||
/// Helper: extract bootstrap_port as u16.
|
||||
fn bootstrap_port(v: &Value) -> Option<u16> {
|
||||
v.get("bootstrap_port")
|
||||
.and_then(|x| x.as_u64())
|
||||
.map(|p| p as u16)
|
||||
}
|
||||
/// Helper: extract bootstrap_room as u64.
|
||||
fn bootstrap_room(v: &Value) -> Option<u64> {
|
||||
v.get("bootstrap_room").and_then(|x| x.as_u64())
|
||||
}
|
||||
|
||||
/// PD-mode chat fans out to BOTH prefill and decode with identical
|
||||
/// bootstrap fields injected into both bodies.
|
||||
#[tokio::test]
|
||||
async fn pd_mode_chat_injects_bootstrap_fields_into_both_bodies() {
|
||||
let prefill = crate::common::mock_worker::MockWorker::start(vec![]).await;
|
||||
let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
|
||||
let ctx = build_ctx(vec![
|
||||
WorkerSpec {
|
||||
id: WorkerId("p1".into()),
|
||||
url: prefill.url.clone(),
|
||||
mode: WorkerMode::Prefill,
|
||||
model_ids: vec![ModelId("tiny".into())],
|
||||
bootstrap_port: Some(8997),
|
||||
},
|
||||
WorkerSpec {
|
||||
id: WorkerId("d1".into()),
|
||||
url: decode.url.clone(),
|
||||
mode: WorkerMode::Decode,
|
||||
model_ids: vec![ModelId("tiny".into())],
|
||||
bootstrap_port: None,
|
||||
},
|
||||
]);
|
||||
let app = build_router(ctx);
|
||||
|
||||
let res = app.oneshot(chat_request()).await.unwrap();
|
||||
assert_eq!(res.status(), StatusCode::OK, "decode side should 200");
|
||||
|
||||
let prefill_body = await_captured_body(&prefill, Duration::from_secs(2), "prefill").await;
|
||||
let decode_body = await_captured_body(&decode, Duration::from_secs(2), "decode").await;
|
||||
let pj = parse_body(&prefill_body);
|
||||
let dj = parse_body(&decode_body);
|
||||
|
||||
// Same bootstrap_room on both sides (one room minted per request).
|
||||
let p_room = bootstrap_room(&pj).expect("prefill body missing bootstrap_room");
|
||||
let d_room = bootstrap_room(&dj).expect("decode body missing bootstrap_room");
|
||||
assert_eq!(
|
||||
p_room, d_room,
|
||||
"prefill and decode must share the same bootstrap_room"
|
||||
);
|
||||
|
||||
// Room must be in [0, i64::MAX]: the SGLang prefill stores it as
|
||||
// i64 internally, so values with the top bit set wrap negative.
|
||||
assert!(
|
||||
p_room <= i64::MAX as u64,
|
||||
"bootstrap_room {p_room} exceeds 63-bit range; SGLang would mis-store as negative i64",
|
||||
);
|
||||
|
||||
// bootstrap_host on both sides == prefill worker's hostname
|
||||
// (MockWorker binds to 127.0.0.1).
|
||||
assert_eq!(bootstrap_host(&pj), Some("127.0.0.1"));
|
||||
assert_eq!(bootstrap_host(&dj), Some("127.0.0.1"));
|
||||
|
||||
// bootstrap_port on both sides == prefill's configured bootstrap_port.
|
||||
assert_eq!(bootstrap_port(&pj), Some(8997));
|
||||
assert_eq!(bootstrap_port(&dj), Some(8997));
|
||||
}
|
||||
|
||||
/// Plain-mode (non-PD) requests do NOT carry any `bootstrap_*` field.
|
||||
/// The injection step is gated on `worker.mode() == Prefill`; plain
|
||||
/// workers serve the chat route directly without disagg bootstrapping.
|
||||
#[tokio::test]
|
||||
async fn plain_mode_chat_does_not_inject_bootstrap_fields() {
|
||||
let plain = crate::common::mock_worker::MockWorker::start(vec![]).await;
|
||||
let ctx = build_ctx(vec![WorkerSpec {
|
||||
id: WorkerId("w1".into()),
|
||||
url: plain.url.clone(),
|
||||
mode: WorkerMode::Plain,
|
||||
model_ids: vec![ModelId("tiny".into())],
|
||||
bootstrap_port: None,
|
||||
}]);
|
||||
let app = build_router(ctx);
|
||||
|
||||
let res = app.oneshot(chat_request()).await.unwrap();
|
||||
assert_eq!(res.status(), StatusCode::OK);
|
||||
|
||||
let body = await_captured_body(&plain, Duration::from_secs(2), "plain").await;
|
||||
let v = parse_body(&body);
|
||||
assert!(
|
||||
v.get("bootstrap_room").is_none(),
|
||||
"plain-mode request must not carry bootstrap_room; got {v}"
|
||||
);
|
||||
assert!(
|
||||
v.get("bootstrap_host").is_none(),
|
||||
"plain-mode request must not carry bootstrap_host; got {v}"
|
||||
);
|
||||
assert!(
|
||||
v.get("bootstrap_port").is_none(),
|
||||
"plain-mode request must not carry bootstrap_port; got {v}"
|
||||
);
|
||||
}
|
||||
|
||||
/// PD-mode with multiple prefill workers + different `bootstrap_port`
|
||||
/// values: the bootstrap_port injected MUST match the actually-chosen
|
||||
/// prefill (not e.g. the first registered or a global config value).
|
||||
#[tokio::test]
|
||||
async fn pd_mode_bootstrap_port_matches_chosen_prefill_worker() {
|
||||
let prefill_a = crate::common::mock_worker::MockWorker::start(vec![]).await;
|
||||
let prefill_b = crate::common::mock_worker::MockWorker::start(vec![]).await;
|
||||
let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
|
||||
let ctx = build_ctx(vec![
|
||||
WorkerSpec {
|
||||
id: WorkerId("pA".into()),
|
||||
url: prefill_a.url.clone(),
|
||||
mode: WorkerMode::Prefill,
|
||||
model_ids: vec![ModelId("tiny".into())],
|
||||
bootstrap_port: Some(11111),
|
||||
},
|
||||
WorkerSpec {
|
||||
id: WorkerId("pB".into()),
|
||||
url: prefill_b.url.clone(),
|
||||
mode: WorkerMode::Prefill,
|
||||
model_ids: vec![ModelId("tiny".into())],
|
||||
bootstrap_port: Some(22222),
|
||||
},
|
||||
WorkerSpec {
|
||||
id: WorkerId("d1".into()),
|
||||
url: decode.url.clone(),
|
||||
mode: WorkerMode::Decode,
|
||||
model_ids: vec![ModelId("tiny".into())],
|
||||
bootstrap_port: None,
|
||||
},
|
||||
]);
|
||||
let app = build_router(ctx);
|
||||
|
||||
// Fire enough requests to ensure round-robin hits both prefill workers.
|
||||
for _ in 0..6 {
|
||||
let res = app.clone().oneshot(chat_request()).await.unwrap();
|
||||
assert_eq!(res.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
// Wait until both prefill workers have captured at least one body.
|
||||
let body_a = await_captured_body(&prefill_a, Duration::from_secs(2), "prefill_a").await;
|
||||
let body_b = await_captured_body(&prefill_b, Duration::from_secs(2), "prefill_b").await;
|
||||
let va = parse_body(&body_a);
|
||||
let vb = parse_body(&body_b);
|
||||
// Each prefill must see its OWN bootstrap_port — never the other's.
|
||||
assert_eq!(
|
||||
bootstrap_port(&va),
|
||||
Some(11111),
|
||||
"prefill_a body should carry its own bootstrap_port"
|
||||
);
|
||||
assert_eq!(
|
||||
bootstrap_port(&vb),
|
||||
Some(22222),
|
||||
"prefill_b body should carry its own bootstrap_port"
|
||||
);
|
||||
}
|
||||
|
||||
/// Pin Pattern B's "prefill failure is invisible to the client"
|
||||
/// contract: when the spawned prefill task gets a 5xx (or any other
|
||||
/// upstream error), the decode response still reaches the client
|
||||
/// unmodified. The router intentionally does not wire fail-fast here —
|
||||
/// the decode side will eventually hang on `bootstrap_room` and time
|
||||
/// out, but the chat handler itself doesn't propagate the prefill
|
||||
/// error. Matches llm-d / aibrix behaviour.
|
||||
#[tokio::test]
|
||||
async fn pd_mode_prefill_5xx_does_not_poison_decode_response() {
|
||||
let prefill = crate::common::mock_worker::MockWorker::start_returning_error(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
json!({"error": "simulated prefill failure"}),
|
||||
)
|
||||
.await;
|
||||
let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
|
||||
let ctx = build_ctx(vec![
|
||||
WorkerSpec {
|
||||
id: WorkerId("p1".into()),
|
||||
url: prefill.url.clone(),
|
||||
mode: WorkerMode::Prefill,
|
||||
model_ids: vec![ModelId("tiny".into())],
|
||||
bootstrap_port: Some(8997),
|
||||
},
|
||||
WorkerSpec {
|
||||
id: WorkerId("d1".into()),
|
||||
url: decode.url.clone(),
|
||||
mode: WorkerMode::Decode,
|
||||
model_ids: vec![ModelId("tiny".into())],
|
||||
bootstrap_port: None,
|
||||
},
|
||||
]);
|
||||
let app = build_router(ctx);
|
||||
|
||||
// Client must see decode's 200 — the failing prefill is invisible.
|
||||
let res = app.oneshot(chat_request()).await.unwrap();
|
||||
assert_eq!(
|
||||
res.status(),
|
||||
StatusCode::OK,
|
||||
"decode response should reach the client even when prefill returned 5xx",
|
||||
);
|
||||
|
||||
// Decode received its body (proves dual dispatch fired despite
|
||||
// the prefill failure).
|
||||
let decode_body = await_captured_body(&decode, Duration::from_secs(2), "decode").await;
|
||||
let v = parse_body(&decode_body);
|
||||
assert_eq!(bootstrap_port(&v), Some(8997));
|
||||
|
||||
// Prefill also received its body — it just returned 5xx. The
|
||||
// bootstrap fields are present so the engine WOULD have honoured
|
||||
// the bootstrap_room if the mock had succeeded.
|
||||
let prefill_body = await_captured_body(&prefill, Duration::from_secs(2), "prefill").await;
|
||||
let pv = parse_body(&prefill_body);
|
||||
assert_eq!(bootstrap_port(&pv), Some(8997));
|
||||
}
|
||||
Reference in New Issue
Block a user