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
2026-05-25 15:34:05 +08:00
committed by GitHub
co-authored by Claude Opus 4.7
parent aae04b1241
commit 6e8fe176be
131 changed files with 28623 additions and 55 deletions
@@ -0,0 +1,4 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
mod static_urls;
@@ -0,0 +1,169 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use sgl_router::config::StaticUrlsDiscoveryConfig;
use sgl_router::discovery::{DiscoveryEvent, WorkerMode};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::mpsc;
#[tokio::test]
async fn emits_one_added_per_url_with_plain_seed() {
let cfg = StaticUrlsDiscoveryConfig {
urls: vec!["http://x:30000".into(), "http://y:30000".into()],
};
let (tx, mut rx) = mpsc::channel(16);
let _h = sgl_router::discovery::static_urls::spawn(cfg, tx)
.await
.unwrap();
let mut seen = std::collections::HashSet::new();
for _ in 0..2 {
let event = tokio::time::timeout(Duration::from_secs(2), rx.recv())
.await
.unwrap()
.unwrap();
match event {
DiscoveryEvent::Added(spec) => {
// mode / model_ids / bootstrap_port are seeded as Plain/empty/None;
// the worker manager fills them from /server_info post-discovery.
assert_eq!(spec.mode, WorkerMode::Plain);
assert!(spec.model_ids.is_empty());
assert_eq!(spec.bootstrap_port, None);
// The URL doubles as the worker id — strings already have to
// be unique (rejected at config-load otherwise).
assert_eq!(spec.id.0, spec.url);
seen.insert(spec.url);
}
other => panic!("unexpected event: {other:?}"),
}
}
assert_eq!(
seen,
["http://x:30000".to_string(), "http://y:30000".to_string()].into(),
);
}
/// Single-URL list — the common dev deployment shape. The producer
/// emits exactly one event and then parks until the receiver is
/// dropped. Earlier versions exited as soon as fan-out completed,
/// which tripped `server::supervisor::supervise_critical_tasks` →
/// `mark_unready` → `/readyz` 503; the lib-side
/// `stays_alive_after_fanout_until_receiver_dropped` pins that
/// invariant in isolation, while this test pins the same contract
/// through the public `spawn` entry point used by the binary.
#[tokio::test]
async fn emits_one_event_and_parks_until_receiver_dropped() {
let cfg = StaticUrlsDiscoveryConfig {
urls: vec!["http://x:30000".into()],
};
let (tx, mut rx) = mpsc::channel(16);
let h = sgl_router::discovery::static_urls::spawn(cfg, tx)
.await
.unwrap();
let event = rx.recv().await.unwrap();
assert!(matches!(event, DiscoveryEvent::Added(_)));
assert!(rx.try_recv().is_err(), "exactly one event expected");
// Drop the receiver → producer's `tx.closed()` resolves → task
// exits cleanly.
drop(rx);
tokio::time::timeout(Duration::from_secs(2), h)
.await
.expect("static_urls task should exit after receiver is dropped")
.expect("join handle should not panic");
}
/// Spin up a fake worker that advertises
/// `disaggregation_mode = "prefill"` + `disaggregation_bootstrap_port`,
/// pipe it through `spawn_discovery` (StaticUrls backend) into
/// `manager::run_with_config`, and assert the worker lands in the
/// registry with `WorkerMode::Prefill` + the disclosed port.
///
/// This is the load-bearing end-to-end assertion for the refactor's
/// central claim — "prefill, decode, and plain workers can all appear
/// in the same `urls` list and end up classified correctly" — exercised
/// against the full discovery → introspect → registry pipeline rather
/// than just the in-isolation `register_one` unit test.
#[tokio::test]
async fn static_urls_pd_role_resolved_end_to_end() {
use axum::{routing::get, Json, Router};
use serde_json::json;
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, DiscoveryConfig, ObservabilityConfig,
ProxyConfig, ServerConfig,
};
use sgl_router::discovery::{spawn_discovery, WorkerId};
use sgl_router::workers::{manager, WorkerRegistry};
use tokio::net::TcpListener;
use tokio::sync::oneshot;
// Fake worker advertising a prefill role + bootstrap port.
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
let url = format!("http://127.0.0.1:{port}");
let app = Router::new().route(
"/server_info",
get(|| async {
Json(json!({
"served_model_name": "tiny",
"disaggregation_mode": "prefill",
"disaggregation_bootstrap_port": 8998,
}))
}),
);
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
tokio::spawn(async move {
let _ = axum::serve(listener, app)
.with_graceful_shutdown(async move {
let _ = shutdown_rx.await;
})
.await;
});
let cfg = Config {
server: ServerConfig {
host: "127.0.0.1".into(),
port: 0,
},
observability: ObservabilityConfig::default(),
models: vec![],
discovery: DiscoveryConfig {
backend: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec![url.clone()],
}),
},
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
};
let registry = Arc::new(WorkerRegistry::default());
let (event_rx, _disc) = spawn_discovery(&cfg).await.unwrap();
let _mgr = tokio::spawn(manager::run_with_config(
event_rx,
registry.clone(),
Some(Arc::new(cfg)),
None,
None,
));
let id = WorkerId(url);
let resolved = tokio::time::timeout(Duration::from_secs(2), async {
loop {
if let Some(w) = registry.get(&id) {
if w.mode() == WorkerMode::Prefill && w.bootstrap_port() == Some(8998) {
return true;
}
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await;
assert!(
resolved.is_ok(),
"expected mode=Prefill bootstrap_port=Some(8998); got {:?}",
registry.get(&id).map(|w| (w.mode(), w.bootstrap_port()))
);
let _ = shutdown_tx.send(());
}
@@ -0,0 +1,163 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use sgl_router::health::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig};
use std::time::Duration;
fn cb() -> CircuitBreaker {
CircuitBreaker::with_config(CircuitBreakerConfig {
threshold: std::num::NonZeroU32::new(3).unwrap(),
cool_down: Duration::from_millis(100),
})
}
#[test]
fn starts_closed_and_allows() {
let b = cb();
assert!(b.allow());
}
#[test]
fn three_failures_open_the_breaker() {
let b = cb();
b.record_failure();
b.record_failure();
assert!(b.allow(), "still closed before threshold");
b.record_failure();
assert!(!b.allow(), "open after threshold reached");
}
#[test]
fn intermittent_success_resets_failure_count() {
let b = cb();
b.record_failure();
b.record_failure();
b.record_success(); // resets
b.record_failure();
b.record_failure();
assert!(b.allow(), "should still be closed (2 failures since reset)");
}
#[tokio::test(start_paused = true)]
async fn open_breaker_recovers_via_half_open() {
let b = cb();
b.record_failure();
b.record_failure();
b.record_failure();
assert!(!b.allow());
// Wait past cool_down.
tokio::time::advance(Duration::from_millis(150)).await;
// Half-open: allow one probe.
assert!(b.allow(), "half-open allows the probe");
// While half-open, further allow() calls should reject (only one probe in flight).
assert!(!b.allow(), "half-open rejects second probe");
// Probe succeeded.
b.record_success();
assert!(b.allow(), "closed after successful probe");
assert!(b.allow(), "stays closed");
}
#[tokio::test(start_paused = true)]
async fn half_open_failure_reopens() {
let b = cb();
b.record_failure();
b.record_failure();
b.record_failure();
tokio::time::advance(Duration::from_millis(150)).await;
assert!(b.allow(), "half-open admit");
b.record_failure();
// Back to Open.
assert!(!b.allow(), "back to open");
}
#[tokio::test(start_paused = true)]
async fn would_allow_is_non_mutating_past_cool_down() {
// `would_allow()` answers "would `allow()` return true right now?" without
// claiming a probe slot. Enumeration / filtering paths (e.g.
// `WorkerRegistry::healthy_workers_for`) call it to inspect breakers
// without disturbing state.
let b = cb();
b.record_failure();
b.record_failure();
b.record_failure();
assert!(!b.allow(), "open after threshold");
tokio::time::advance(Duration::from_millis(150)).await;
// Repeated would_allow() returns true and leaves state untouched.
assert!(b.would_allow());
assert!(b.would_allow());
assert!(b.would_allow());
// The first allow() claims the half-open probe.
assert!(b.allow(), "allow() admits the probe");
// The probe is in flight — subsequent allow() (and would_allow()) reject.
assert!(!b.allow(), "only one probe in flight");
assert!(!b.would_allow(), "would_allow() agrees: no slot available");
}
#[tokio::test(start_paused = true)]
async fn enumeration_then_dispatch_preserves_probe() {
// Regression for the bug where `healthy_workers_for` filtered with
// mutating `allow()`. Once would_allow() is the filter, an enumeration
// pass over many workers must not steal the probe slot from the one
// worker that actually gets dispatched to.
let b = cb();
b.record_failure();
b.record_failure();
b.record_failure();
tokio::time::advance(Duration::from_millis(150)).await;
// Imagine 3 workers; enumeration filters each with would_allow().
for _ in 0..3 {
assert!(b.would_allow(), "filter sees the worker as available");
}
// Now the policy picks ONE worker and dispatch claims the probe.
assert!(b.allow(), "dispatch on the picked worker succeeds");
}
#[test]
fn would_allow_in_closed_state_is_true_and_non_mutating() {
let b = cb();
for _ in 0..5 {
assert!(b.would_allow());
}
// And allow() should still work afterwards.
assert!(b.allow());
}
#[tokio::test(start_paused = true)]
async fn open_breaker_recovery_is_not_delayed_by_continued_failures() {
// Regression: previously, record_failure on an already-Open breaker
// refreshed opened_at, so a failure storm pinned the breaker open
// forever. Now the cool_down is measured from first-open.
let b = CircuitBreaker::with_config(CircuitBreakerConfig {
threshold: std::num::NonZeroU32::new(3).unwrap(),
cool_down: Duration::from_millis(100),
});
// Open it.
b.record_failure();
b.record_failure();
b.record_failure();
assert!(!b.allow(), "breaker should be open");
// Advance halfway through cool_down, then record more failures.
tokio::time::advance(Duration::from_millis(50)).await;
b.record_failure();
b.record_failure();
b.record_failure();
// Advance just past the original cool_down.
tokio::time::advance(Duration::from_millis(60)).await;
// We're past the original cool_down → HalfOpen.
assert!(
b.allow(),
"breaker should be half-open after cool_down from first-open"
);
}
@@ -0,0 +1,4 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
mod circuit_breaker;
@@ -0,0 +1,14 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Component-scope integration tests.
//!
//! Each submodule exercises a single library component (policy, registry,
//! discovery, health, tokenizer) via the crate's public API. None of these
//! tests spin up the full HTTP router; for those see `tests/proxy/`.
mod discovery;
mod health;
mod policies;
mod tokenizer;
mod workers;
@@ -0,0 +1,182 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! E2E test for the cache-aware-zmq policy.
//!
//! Drives a real ZMQ PUB socket → `KvEventIndex` subscriber pipeline →
//! `HashTree` → `CacheAwareZmqPolicy::select`. Verifies that an event
//! published by one worker's PUB causes subsequent selection to route
//! to that worker (cache-aware affinity).
//!
//! API constraint: the subscriber registry builds endpoints as
//! `tcp://{host}:{port_base + dp_rank}` where `port_base` is in the
//! per-worker `EventConfig`. Both mock workers below share
//! `127.0.0.1` as host, so both subscribe to the same PUB socket and
//! both end up indexed in the tree. The tiebreak (lowest active_load)
//! picks the worker we want; same shape as the SMG version of this
//! test.
use std::sync::Arc;
use std::time::Duration;
use zeromq::SocketSend;
use sgl_router::config::CacheAwareConfig;
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::kv_events::{compute_block_hashes, discovery::EventConfig, KvEventIndex};
use sgl_router::policies::{Policy, SelectionContext};
use sgl_router::tokenizer::TokenizerRegistry;
use sgl_router::workers::Worker;
use super::zmq_helpers::{
build_multipart, encode_block_stored_event, encode_event_batch, make_pub_bound,
};
fn build_worker(url: &str, model: &str) -> Arc<Worker> {
Arc::new(Worker::new(WorkerSpec {
id: WorkerId(url.into()),
url: url.into(),
mode: WorkerMode::Plain,
model_ids: vec![ModelId(model.into())],
bootstrap_port: None,
}))
}
/// E2E: real PUB socket publishes a `BlockStored` for worker A's
/// hash chain. The `CacheAwareZmqPolicy`'s shared `KvEventIndex`
/// receives it, applies it to the tree, and the next `select` call
/// picks worker A.
///
/// Both workers share `127.0.0.1` as host so both subscribers connect
/// to the same PUB and both get indexed under their KvWorkerIds — the
/// same shape as the SMG e2e test. We tie-break on min-load: worker B
/// is bumped above worker A so the matched-worker pick prefers A.
#[tokio::test]
async fn zmq_indexer_routes_to_publishing_worker_e2e() {
let model_id = ModelId("tiny".into());
// 1. Tokenizer registry — use the in-tree tiny fixture.
let cfg = sgl_router::config::Config {
server: sgl_router::config::ServerConfig {
host: "0".into(),
port: 0,
},
observability: Default::default(),
models: vec![sgl_router::config::ModelConfig {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: sgl_router::config::PolicyKind::CacheAwareZmq,
circuit_breaker: None,
cache_aware: None,
}],
discovery: sgl_router::config::DiscoveryConfig {
backend: sgl_router::config::DiscoveryBackend::StaticUrls(
sgl_router::config::StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
},
),
},
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
};
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
// 2. Bind a real PUB socket on an OS-assigned port.
let (mut pub_a, port) = make_pub_bound().await;
// 3. Compute the hash chain for the routing prompt.
let text = "hello world hello world hello world";
let tok = tokenizers.get("tiny").unwrap();
let token_ids = sgl_router::tokenizer::adapter::encode(&tok, text).unwrap();
let block_size = 4u32;
let hashes = compute_block_hashes(&token_ids, block_size as usize);
assert!(!hashes.is_empty(), "tiny tokenizer must yield ≥1 block");
// 4. Build the KvEventIndex + policy. The policy holds an
// Arc<HashTree> that the index also owns; events the index
// receives mutate the same tree the policy reads.
let kv_index = KvEventIndex::new();
// Mirror what `KvEventIndex::add_worker` would do in production: seed
// the oracle with the worker-reported page_size before any cache
// lookup happens. The integration path calls `add_worker` further
// down, but here we want the policy to know `block_size` immediately.
let block_size_oracle = kv_index.block_size_oracle();
block_size_oracle.try_set(block_size).unwrap();
let policy = CacheAwareZmqPolicy::new(
CacheAwareConfig {
cache_threshold: 0.0,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
},
kv_index.tree(),
Arc::clone(&tokenizers),
block_size_oracle,
);
// 5. Register two workers. They share `127.0.0.1` so both
// subscribers connect to the same PUB; preresolved EventConfig
// points at the bound port.
let url_a = "http://127.0.0.1:30000";
let url_b = "http://127.0.0.1:30001";
let preresolved = EventConfig {
host: "127.0.0.1".to_string(),
port_base: port,
topic: String::new(),
block_size,
dp_size: 1,
};
kv_index.add_worker(url_a, Some(preresolved.clone())).await;
kv_index.add_worker(url_b, Some(preresolved)).await;
// SUB sockets take a moment to handshake. The polling loop below
// soaks up any extra latency; this is just a publish-before-SUB
// guard.
tokio::time::sleep(Duration::from_millis(150)).await;
// 6. Publish a BlockStored event for the routing prompt's chain.
let event_bytes = encode_block_stored_event(&hashes, None, &token_ids, block_size);
let payload = encode_event_batch(0.0, vec![event_bytes], Some(0));
pub_a
.send(build_multipart(1, payload))
.await
.expect("send block-stored event");
// 7. Bump worker B's load so the tie-break picks A among matched
// workers. The bump stays below balance_abs_threshold so the
// imbalance fast-path does not skip cache-aware selection.
// Bind the guards to a Vec held for the rest of the test scope
// so the counter stays > 0 through the polling loop.
let w_a = build_worker(url_a, "tiny");
let w_b = build_worker(url_b, "tiny");
let _b_load: Vec<_> = (0..3).map(|_| w_b.load_guard()).collect();
let workers = vec![Arc::clone(&w_a), Arc::clone(&w_b)];
// 8. Drive select until the event has been applied. The pipeline is
// asynchronous (publish → SUB recv → mpsc → pump → tree); a
// polling loop is less flaky than a fixed sleep.
let body = serde_json::to_vec(&serde_json::json!({"prompt": text})).unwrap();
let ctx = SelectionContext::new(&model_id, Some(&body));
let start = std::time::Instant::now();
let mut chose_a = false;
while start.elapsed() < Duration::from_secs(3) {
if let Some(w) = policy.select(&workers, &ctx) {
if w.url == url_a {
chose_a = true;
break;
}
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
assert!(
chose_a,
"policy did not route to publishing worker A within timeout",
);
// 9. Shutdown cleanly.
let r = tokio::time::timeout(Duration::from_secs(2), kv_index.shutdown()).await;
assert!(r.is_ok(), "kv_index shutdown should not hang");
}
@@ -0,0 +1,75 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Cross-implementation parity test for the KV-event block-hash algorithm.
//!
//! The Rust implementation at `src/policies/kv_events/hash.rs` must produce
//! the same i64 block hashes as SGLang's `radix_cache::RadixKey.hash_page`
//! followed by `hash_str_to_int64`. Hard-coded `cross_language_golden_*`
//! values inside `hash.rs` are correct but brittle: if either side's
//! algorithm changes, the comments don't get regenerated and the tests
//! pass with stale expectations.
//!
//! This test consumes a fixture produced by
//! `tests/scripts/generate_kv_events_hash_parity.py`, which replicates the
//! SGLang algorithm verbatim (see the script's docstring for authority
//! pointers). CI regenerates the fixture (see
//! `.github/workflows/pr-test-sgl-router.yml`) and diffs against the
//! committed file; this test asserts the Rust implementation matches
//! whatever fixture is checked in.
use serde::Deserialize;
use sgl_router::policies::kv_events::compute_block_hashes;
use std::path::PathBuf;
#[derive(Debug, Deserialize)]
struct ParityCase {
name: String,
tokens: Vec<u32>,
block_size: usize,
expected_i64_hashes: Vec<i64>,
}
fn fixture_path() -> PathBuf {
PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("tests")
.join("fixtures")
.join("kv_events_hash_parity.json")
}
fn load_cases() -> Vec<ParityCase> {
let path = fixture_path();
let bytes = std::fs::read(&path)
.unwrap_or_else(|e| panic!("read parity fixture {}: {e}", path.display()));
serde_json::from_slice(&bytes)
.unwrap_or_else(|e| panic!("decode parity fixture {}: {e}", path.display()))
}
#[test]
fn fixture_is_non_empty() {
let cases = load_cases();
assert!(
!cases.is_empty(),
"kv_events_hash_parity.json is empty — run \
tests/scripts/generate_kv_events_hash_parity.py",
);
}
/// Drives every case in the fixture through `compute_block_hashes` and
/// asserts equality with the Python-derived expectation.
#[test]
fn rust_block_hashes_match_python_radix_cache() {
for case in load_cases() {
// block_size of 0 is rejected by `compute_block_hashes` with a
// panic; the Python generator also rejects it. The fixture
// doesn't include a 0 case, so unwrap is safe.
let block_size = std::num::NonZeroUsize::new(case.block_size)
.unwrap_or_else(|| panic!("case {} has block_size=0 which is invalid", case.name));
let got = compute_block_hashes(&case.tokens, block_size.get());
assert_eq!(
got, case.expected_i64_hashes,
"case {}: tokens={:?} block_size={} — Rust produced {:?}, fixture says {:?}",
case.name, case.tokens, case.block_size, got, case.expected_i64_hashes,
);
}
}
@@ -0,0 +1,167 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Concurrent-mutation stress test for `HashTree`.
//!
//! The 19 inline tests in `policies::kv_events::tree` are all
//! single-threaded. Under production load, multiple worker subscribers
//! drive `insert` / `remove` / `clear_worker` against the same tree from
//! tokio worker threads while the chat handler simultaneously calls
//! `match_prefix` from many concurrent requests.
//!
//! The tree is documented as taking a write-lock for mutations and a
//! read-lock for `match_prefix`; this test exercises that contract under
//! heavy contention to catch:
//!
//! * Deadlocks between the reverse index and the arena's RwLock.
//! * Logical races where a removed worker still appears in the reverse
//! index (or vice versa).
//! * Panics from a node arena being mutated mid-read.
//!
//! After the storm settles, the tree must be self-consistent: every
//! worker that was fully cleared must be absent from every node's worker
//! set, and `node_count()` must converge to zero.
use std::sync::Arc;
use std::thread;
use sgl_router::policies::kv_events::{HashTree, KvWorkerId};
fn worker(i: usize) -> KvWorkerId {
KvWorkerId {
url: format!("http://w{i}:30000"),
dp_rank: 0,
}
}
/// 8 mutator threads × 200 ops + 4 reader threads × 500 match queries.
/// Each mutator inserts a chain, queries it, then clears the worker; the
/// invariant is that after every thread joins, the tree is empty (every
/// worker was cleared) and no thread panicked.
#[test]
fn tree_survives_concurrent_inserts_removes_and_matches() {
let tree = Arc::new(HashTree::new());
let mut handles = Vec::new();
for tid in 0..8 {
let tree = tree.clone();
handles.push(thread::spawn(move || {
let w = worker(tid);
for round in 0..200_u64 {
// Each round uses a fresh chain so different mutators
// don't trample each other's nodes — we want contention
// on the lock, not contention on the keys (those are
// covered by the single-threaded reinsert/remove tests).
let chain: Vec<i64> = (0..4)
.map(|i| ((tid as i64) << 32) | ((round as i64) << 8) | i as i64)
.collect();
tree.insert(&w, None, &chain);
let m = tree.match_prefix(None, &chain);
assert!(
m.matched_blocks <= chain.len(),
"match must never exceed query length",
);
// Half the rounds use remove(&chain); the rest use
// clear_worker — both must leave a consistent tree.
if round % 2 == 0 {
tree.remove(&w, &chain);
} else {
tree.clear_worker(&w);
}
}
// Final blanket clear in case the last iteration used `remove`
// on only part of the chain.
tree.clear_worker(&w);
}));
}
for tid in 0..4 {
let tree = tree.clone();
handles.push(thread::spawn(move || {
for round in 0..500_u64 {
let probe: Vec<i64> = (0..3)
.map(|i| ((tid as i64) << 40) | ((round as i64) << 8) | i as i64)
.collect();
// Readers must never block-walk and must never panic.
let _ = tree.match_prefix(None, &probe);
}
}));
}
for h in handles {
h.join()
.expect("worker thread panicked under concurrent load");
}
assert_eq!(
tree.node_count(),
0,
"tree must be empty after every worker was cleared; \
residual nodes indicate a missed clear_worker path",
);
// The arena and the reverse index must agree: zero non-root nodes
// means zero `by_hash` entries. A bug that prunes the arena but not
// the reverse index would leak memory and corrupt future inserts;
// this assertion turns that into an immediate test failure.
assert_eq!(
tree.reverse_index_size(),
0,
"by_hash reverse index must be empty when no non-root nodes remain",
);
}
/// A mutator races `clear_worker` against a reader that is mid-`match_prefix`
/// on a deep chain. The reader must never see a partially-mutated tree
/// (no panic, no double-counted workers in the result set).
#[test]
fn match_prefix_is_consistent_with_concurrent_clear() {
let tree = Arc::new(HashTree::new());
let w = worker(0);
let chain: Vec<i64> = (0..32).map(|i| 1_000 + i).collect();
// Pre-populate so the reader has something to walk.
tree.insert(&w, None, &chain);
let stop = Arc::new(std::sync::atomic::AtomicBool::new(false));
let mutator = {
let tree = tree.clone();
let stop = stop.clone();
let w = w.clone();
let chain = chain.clone();
thread::spawn(move || {
let mut round = 0u64;
while !stop.load(std::sync::atomic::Ordering::Relaxed) {
if round.is_multiple_of(2) {
tree.clear_worker(&w);
} else {
tree.insert(&w, None, &chain);
}
round += 1;
}
})
};
for _ in 0..2_000 {
let m = tree.match_prefix(None, &chain);
// Either the worker was present (matched_blocks == chain.len(),
// workers set contains w) or it was cleared mid-walk (matched_blocks
// == 0 OR matched_blocks > 0 with empty workers if the chain is
// partially present). Whichever — the result must be internally
// consistent.
if m.matched_blocks == chain.len() {
assert!(
m.workers.contains(&w),
"full match must include worker; got {:?}",
m.workers,
);
}
}
stop.store(true, std::sync::atomic::Ordering::Relaxed);
mutator.join().unwrap();
}
@@ -0,0 +1,305 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Two independent `KvEventIndex` instances subscribed to the same PUB
//! socket — the in-process surrogate for "two router replicas watching
//! the same SGLang worker's KV publisher."
//!
//! Why this matters: sgl-router v1 explicitly omits multi-replica state
//! sync (deferred to v2 in the slim-design spec). Independent ZMQ
//! subscription is the **only** mechanism by which two routers arrive at
//! a consistent cache-aware view today. If a future change accidentally
//! degraded that property — e.g. a worker that only allows one subscriber,
//! a switch from PUB/SUB to PUSH/PULL, or a teardown bug that drops
//! events to one of N subscribers — this test fails loudly.
//!
//! Property pinned: after publishing N `BlockStored` events, both trees
//! report the same `match_prefix(matched_blocks, workers)` for the
//! published key, and an unpublished key remains absent from both.
use std::sync::Arc;
use std::time::Duration;
use zeromq::SocketSend;
use sgl_router::policies::kv_events::discovery::EventConfig;
use sgl_router::policies::kv_events::{compute_block_hashes, KvEventIndex, KvWorkerId};
use super::zmq_helpers::{
build_multipart, encode_block_stored_event, encode_event_batch, make_pub_bound,
};
#[tokio::test]
async fn two_independent_subscribers_converge_to_same_tree_state() {
// 1. One PUB socket — the worker. Both router surrogates connect to it.
let (mut publisher, port) = make_pub_bound().await;
let worker_url = "http://127.0.0.1:30000";
let block_size = 4u32;
let cfg = EventConfig {
host: "127.0.0.1".into(),
port_base: port,
topic: String::new(),
block_size,
dp_size: 1,
};
// 2. Two independent router-process surrogates, each with its own
// `KvEventIndex` (own tree, own subscriber, own pump task). Both
// call `add_worker` with the same preresolved `EventConfig` — the
// same shape production wires through `WorkerManager`.
let router_a = KvEventIndex::new();
let router_b = KvEventIndex::new();
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.
let tokens: Vec<u32> = (0..16).collect();
let hashes = compute_block_hashes(&tokens, block_size as usize);
assert!(
hashes.len() >= 3,
"test needs ≥3 blocks; got {}",
hashes.len()
);
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.
let target = hashes.len();
let key = KvWorkerId {
url: worker_url.into(),
dp_rank: 0,
};
let start = std::time::Instant::now();
loop {
let ma = router_a.tree().match_prefix(None, &hashes);
let mb = router_b.tree().match_prefix(None, &hashes);
let converged = ma.matched_blocks == target
&& mb.matched_blocks == target
&& ma.workers.contains(&key)
&& mb.workers.contains(&key);
if converged {
// Both trees agree on count AND on the worker that holds the
// prefix. This is what the cache-aware-zmq policy reads to
// pick a worker; both routers picking the same key here
// means they would route the same prompt to the same worker.
assert_eq!(
ma.matched_blocks, mb.matched_blocks,
"subscribers disagreed on matched_blocks",
);
assert_eq!(
ma.workers, mb.workers,
"subscribers disagreed on worker set",
);
break;
}
if start.elapsed() > Duration::from_secs(3) {
panic!(
"subscribers did not converge within 3s: \
router_a={{matched={}, workers={:?}}}, \
router_b={{matched={}, workers={:?}}}, target={target}",
ma.matched_blocks, ma.workers, mb.matched_blocks, mb.workers,
);
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
// 5. Negative leg: a key that was never published must not appear in
// either tree. Guards against a future bug where one subscriber
// accidentally inherits another's state (shared static, etc.).
let unseen: Vec<i64> = vec![999_999_999_001, 999_999_999_002, 999_999_999_003];
let na = router_a.tree().match_prefix(None, &unseen);
let nb = router_b.tree().match_prefix(None, &unseen);
assert_eq!(na.matched_blocks, 0, "router_a leaked unpublished key");
assert_eq!(nb.matched_blocks, 0, "router_b leaked unpublished key");
// 6. Both shutdowns must complete cleanly — no hang from the second
// subscriber holding a reference to a shared resource. The first
// drains under a generous ceiling (worker thread joins, mpsc
// receiver drop); the second has nothing left to wait on and
// must complete promptly. A slow second shutdown indicates the
// two subscribers were sharing a resource that serialized them.
let r = tokio::time::timeout(Duration::from_secs(2), Arc::clone(&router_a).shutdown()).await;
assert!(r.is_ok(), "router_a shutdown hung");
let t = std::time::Instant::now();
let r = tokio::time::timeout(Duration::from_secs(2), Arc::clone(&router_b).shutdown()).await;
assert!(r.is_ok(), "router_b shutdown hung");
let elapsed = t.elapsed();
assert!(
elapsed < Duration::from_millis(100),
"router_b shutdown after router_a drained took {elapsed:?}; \
expected <100ms (no shared-resource contention)",
);
}
/// Two PUB sockets (two workers) + two `KvEventIndex` instances (two
/// routers), each subscribed to **both** publishers. This is the real
/// v1 HA shape: each router replica fans out subscriptions across the
/// worker pool and merges every publisher's `BlockStored` stream into
/// its own tree. The companion 1-PUB test above only verifies broadcast
/// fan-out; this test verifies the per-worker attribution stays correct
/// when events arrive from multiple sources concurrently.
///
/// Property pinned: after publishing prefix `X` on `pub_x` and prefix
/// `Y` on `pub_y`, both trees report
/// * `match_prefix(X) = {full, workers={worker_x}}`
/// * `match_prefix(Y) = {full, workers={worker_y}}`
/// with no cross-attribution (worker_x must NOT appear in match(Y)).
/// A regression that wires both subscribers to the same internal
/// channel — or that mis-keys events by their arrival socket rather
/// than their announced worker URL — would surface here as cross-
/// contamination of the worker sets.
#[tokio::test]
async fn two_subscribers_merge_events_from_two_publishers() {
let (mut pub_x, port_x) = make_pub_bound().await;
let (mut pub_y, port_y) = make_pub_bound().await;
let worker_x = "http://127.0.0.1:30001";
let worker_y = "http://127.0.0.1:30002";
let block_size = 4u32;
let cfg_x = EventConfig {
host: "127.0.0.1".into(),
port_base: port_x,
topic: String::new(),
block_size,
dp_size: 1,
};
let cfg_y = EventConfig {
host: "127.0.0.1".into(),
port_base: port_y,
topic: String::new(),
block_size,
dp_size: 1,
};
// Both routers subscribe to BOTH workers — the production fan-out.
let router_a = KvEventIndex::new();
let router_b = KvEventIndex::new();
router_a.add_worker(worker_x, Some(cfg_x.clone())).await;
router_a.add_worker(worker_y, Some(cfg_y.clone())).await;
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
// hash collision.
let tokens_x: Vec<u32> = (0..16).collect();
let tokens_y: Vec<u32> = (1000..1016).collect();
let hashes_x = compute_block_hashes(&tokens_x, block_size as usize);
let hashes_y = compute_block_hashes(&tokens_y, block_size as usize);
assert!(hashes_x.len() >= 3 && hashes_y.len() >= 3);
let payload_x = encode_event_batch(
0.0,
vec![encode_block_stored_event(
&hashes_x, None, &tokens_x, block_size,
)],
Some(0),
);
let payload_y = encode_event_batch(
0.0,
vec![encode_block_stored_event(
&hashes_y, None, &tokens_y, block_size,
)],
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,
};
let key_y = KvWorkerId {
url: worker_y.into(),
dp_rank: 0,
};
let target_x = hashes_x.len();
let target_y = hashes_y.len();
let start = std::time::Instant::now();
loop {
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);
let by = router_b.tree().match_prefix(None, &hashes_y);
let converged = ax.matched_blocks == target_x
&& ay.matched_blocks == target_y
&& bx.matched_blocks == target_x
&& by.matched_blocks == target_y
&& ax.workers.contains(&key_x)
&& ay.workers.contains(&key_y)
&& bx.workers.contains(&key_x)
&& by.workers.contains(&key_y);
if converged {
// Negative attribution: prefix X must not be attributed to
// worker_y in either tree, and vice versa. A regression that
// keyed events by arriving socket rather than announced
// worker URL would set BOTH worker keys on each prefix.
assert!(
!ax.workers.contains(&key_y),
"router_a cross-attributed worker_y to prefix X: {:?}",
ax.workers,
);
assert!(
!ay.workers.contains(&key_x),
"router_a cross-attributed worker_x to prefix Y: {:?}",
ay.workers,
);
assert!(
!bx.workers.contains(&key_y),
"router_b cross-attributed worker_y to prefix X: {:?}",
bx.workers,
);
assert!(
!by.workers.contains(&key_x),
"router_b cross-attributed worker_x to prefix Y: {:?}",
by.workers,
);
break;
}
if start.elapsed() > Duration::from_secs(3) {
panic!(
"trees did not converge within 3s:\n \
router_a: X={{matched={}, workers={:?}}}, Y={{matched={}, workers={:?}}}\n \
router_b: X={{matched={}, workers={:?}}}, Y={{matched={}, workers={:?}}}\n \
targets: X={target_x}, Y={target_y}",
ax.matched_blocks,
ax.workers,
ay.matched_blocks,
ay.workers,
bx.matched_blocks,
bx.workers,
by.matched_blocks,
by.workers,
);
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
let r = tokio::time::timeout(Duration::from_secs(2), Arc::clone(&router_a).shutdown()).await;
assert!(r.is_ok(), "router_a shutdown hung");
let r = tokio::time::timeout(Duration::from_secs(2), Arc::clone(&router_b).shutdown()).await;
assert!(r.is_ok(), "router_b shutdown hung");
}
@@ -0,0 +1,11 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
mod zmq_helpers;
mod cache_aware_zmq;
mod kv_events_hash_parity;
mod kv_events_tree_concurrent;
mod kv_events_two_subscribers;
mod power_of_two;
mod round_robin;
@@ -0,0 +1,73 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::power_of_two::PowerOfTwoChoicesPolicy;
use sgl_router::policies::{Policy, SelectionContext};
use sgl_router::workers::Worker;
use std::sync::atomic::Ordering;
use std::sync::Arc;
fn worker(id: &str) -> Arc<Worker> {
Arc::new(Worker::new(WorkerSpec {
id: WorkerId(id.into()),
url: format!("http://{id}"),
mode: WorkerMode::Plain,
model_ids: vec![ModelId("m".into())],
bootstrap_port: None,
}))
}
#[test]
fn selects_lower_load() {
let a = worker("a");
let b = worker("b");
a.active_requests.store(10, Ordering::Relaxed);
b.active_requests.store(2, Ordering::Relaxed);
let p = PowerOfTwoChoicesPolicy::new();
let ws = vec![a.clone(), b.clone()];
let model_id = ModelId("m".into());
let ctx = SelectionContext::new(&model_id, None);
let chosen = p.select(&ws, &ctx).unwrap();
assert_eq!(chosen.id.0, "b");
}
#[test]
fn distribution_skews_to_lower_load() {
// With 3 workers and one heavily loaded, the loaded one should win
// significantly less than 1/3 of selections.
let workers = vec![worker("a"), worker("b"), worker("c")];
workers[2].active_requests.store(100, Ordering::Relaxed); // c is loaded
let p = PowerOfTwoChoicesPolicy::new();
let model_id = ModelId("m".into());
let ctx = SelectionContext::new(&model_id, None);
let mut counts = std::collections::HashMap::new();
for _ in 0..1000 {
let w = p.select(&workers, &ctx).unwrap();
*counts.entry(w.id.0.clone()).or_insert(0) += 1;
}
let c_picks = *counts.get("c").unwrap_or(&0);
assert!(
c_picks < 200,
"loaded worker should be picked < 20% of the time, got {c_picks}"
);
}
#[test]
fn empty_returns_none() {
let p = PowerOfTwoChoicesPolicy::new();
let ws: Vec<Arc<Worker>> = vec![];
let model_id = ModelId("m".into());
let ctx = SelectionContext::new(&model_id, None);
assert!(p.select(&ws, &ctx).is_none());
}
#[test]
fn single_worker_returns_it() {
let p = PowerOfTwoChoicesPolicy::new();
let ws = vec![worker("only")];
let model_id = ModelId("m".into());
let ctx = SelectionContext::new(&model_id, None);
assert_eq!(p.select(&ws, &ctx).unwrap().id.0, "only");
}
@@ -0,0 +1,56 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::round_robin::RoundRobinPolicy;
use sgl_router::policies::{Policy, SelectionContext};
use sgl_router::workers::Worker;
use std::sync::Arc;
fn worker(id: &str) -> Arc<Worker> {
Arc::new(Worker::new(WorkerSpec {
id: WorkerId(id.into()),
url: format!("http://{id}"),
mode: WorkerMode::Plain,
model_ids: vec![ModelId("m".into())],
bootstrap_port: None,
}))
}
#[test]
fn cycles_through_workers() {
let p = RoundRobinPolicy::new();
let ws = vec![worker("a"), worker("b"), worker("c")];
let model_id = ModelId("m".into());
let ctx = SelectionContext::new(&model_id, None);
let picks: Vec<_> = (0..6)
.filter_map(|_| p.select(&ws, &ctx))
.map(|w| w.id.0.clone())
.collect();
assert_eq!(picks, vec!["a", "b", "c", "a", "b", "c"]);
}
#[test]
fn empty_pool_returns_none() {
let p = RoundRobinPolicy::new();
let ws: Vec<Arc<Worker>> = vec![];
let model_id = ModelId("m".into());
let ctx = SelectionContext::new(&model_id, None);
assert!(p.select(&ws, &ctx).is_none());
}
#[test]
fn distribution_across_100_calls() {
let p = RoundRobinPolicy::new();
let ws = vec![worker("a"), worker("b"), worker("c")];
let model_id = ModelId("m".into());
let ctx = SelectionContext::new(&model_id, None);
let mut counts = std::collections::HashMap::new();
for _ in 0..99 {
let w = p.select(&ws, &ctx).unwrap();
*counts.entry(w.id.0.clone()).or_insert(0) += 1;
}
assert_eq!(counts["a"], 33);
assert_eq!(counts["b"], 33);
assert_eq!(counts["c"], 33);
}
@@ -0,0 +1,88 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Shared ZMQ wire-format helpers for the `policies::kv_events` component
//! tests. Encodes events in the same msgspec layout SGLang emits, builds
//! the two-frame `[seq, payload]` ZMQ message a real publisher sends, and
//! binds a loopback PUB socket on an OS-assigned port.
#![allow(dead_code)]
use bytes::Bytes;
use rmp::encode as mp;
use zeromq::{Endpoint, PubSocket, Socket, ZmqMessage};
/// Bind a PUB socket to an OS-assigned 127.0.0.1 port. Returns
/// `(socket, port)`.
pub async fn make_pub_bound() -> (PubSocket, u16) {
let mut sock = PubSocket::new();
let endpoint = sock
.bind("tcp://127.0.0.1:0")
.await
.expect("bind PUB socket");
let port = match endpoint {
Endpoint::Tcp(_, p) => p,
other => panic!("unexpected endpoint: {other:?}"),
};
(sock, port)
}
/// Encode a single `BlockStored` event in the wire format msgspec
/// emits. Layout: `["BlockStored", block_hashes, parent, token_ids,
/// block_size, lora_id, medium]`.
pub fn encode_block_stored_event(
block_hashes: &[i64],
parent: Option<i64>,
token_ids: &[u32],
block_size: u32,
) -> Vec<u8> {
let mut buf = Vec::new();
mp::write_array_len(&mut buf, 7).unwrap();
mp::write_str(&mut buf, "BlockStored").unwrap();
mp::write_array_len(&mut buf, block_hashes.len() as u32).unwrap();
for v in block_hashes {
mp::write_sint(&mut buf, *v).unwrap();
}
match parent {
Some(v) => {
mp::write_sint(&mut buf, v).unwrap();
}
None => mp::write_nil(&mut buf).unwrap(),
}
mp::write_array_len(&mut buf, token_ids.len() as u32).unwrap();
for v in token_ids {
mp::write_uint(&mut buf, *v as u64).unwrap();
}
mp::write_uint(&mut buf, block_size as u64).unwrap();
mp::write_nil(&mut buf).unwrap(); // lora_id
mp::write_str(&mut buf, "GPU").unwrap();
buf
}
/// Wrap one or more pre-encoded events into a KVEventBatch with
/// timestamp + optional dp-rank.
pub fn encode_event_batch(ts: f64, events: Vec<Vec<u8>>, attn_dp_rank: Option<u32>) -> Vec<u8> {
let mut buf = Vec::new();
mp::write_array_len(&mut buf, 3).unwrap();
mp::write_f64(&mut buf, ts).unwrap();
mp::write_array_len(&mut buf, events.len() as u32).unwrap();
for ev in events {
buf.extend_from_slice(&ev);
}
match attn_dp_rank {
Some(v) => {
mp::write_uint(&mut buf, v as u64).unwrap();
}
None => mp::write_nil(&mut buf).unwrap(),
}
buf
}
/// Build the two-frame ZMQ message a real KV publisher sends:
/// `[seq (big-endian i64), payload]`.
pub fn build_multipart(seq: i64, payload: Vec<u8>) -> ZmqMessage {
let mut msg = ZmqMessage::from(Bytes::new());
msg.push_back(Bytes::copy_from_slice(&seq.to_be_bytes()));
msg.push_back(Bytes::from(payload));
msg
}
@@ -0,0 +1,4 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
mod parity;
@@ -0,0 +1,150 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Bit-parity check: dynamo-tokenizers must produce the same token_ids as
//! SGLang's reference (transformers.AutoTokenizer) for every (model, shape)
//! fixture. Any drift is a regression.
//!
//! ## Running
//!
//! `cargo test --release --test component tokenizer::parity` runs the test.
//!
//! Each fixture cell needs the model's `tokenizer.json` on disk; the test
//! looks in the local HuggingFace cache (`HF_HOME` or `~/.cache/huggingface`).
//! Cells whose snapshot isn't cached are skipped (with a warning); cells
//! whose snapshot IS cached are asserted bit-identical.
//!
//! Locally, when no fixtures can be checked (fresh cache) the test emits a
//! warning and passes — useful for contributors without the model snapshots.
//! In CI (`SGLANG_IS_IN_CI=true`) the same condition is a hard failure: a
//! parity matrix that validates nothing is worse than no test at all, since
//! it gives a false sense of coverage. The e2e HTTP tokenize test remains
//! the authoritative live-model parity gate, but this matrix must actually
//! run against cached snapshots when present in CI.
//!
//! ## Regenerating fixtures
//!
//! Run `tests/scripts/generate_parity_fixtures.py` after changing a prompt
//! shape or adding a model, then commit the new JSON.
use serde::Deserialize;
use std::path::PathBuf;
#[derive(Deserialize)]
struct Fixture {
model_id: String,
shape: String,
prompt_text: String,
expected_token_ids: Vec<u32>,
#[allow(dead_code)]
skip_special_tokens: bool,
}
fn fixture_root() -> PathBuf {
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/tokenizer_parity")
}
/// Resolve a model's tokenizer.json file from the local HF cache.
///
/// Strategy:
/// 1. Check HF_HOME env var, or default to ~/.cache/huggingface
/// 2. Look for models--<safe-name>/snapshots/<hash>/tokenizer.json
/// 3. Return None if not found — the test cell is skipped.
fn resolve_tokenizer_path(model_id: &str) -> Option<PathBuf> {
let hf_home = std::env::var("HF_HOME")
.ok()
.map(PathBuf::from)
.or_else(|| dirs::home_dir().map(|h| h.join(".cache/huggingface")))?;
let safe = model_id.replace('/', "--");
let candidate = hf_home.join("hub").join(format!("models--{safe}"));
if !candidate.exists() {
return None;
}
let snapshots = candidate.join("snapshots");
let snap = std::fs::read_dir(&snapshots).ok()?.next()?.ok()?.path();
let tj = snap.join("tokenizer.json");
tj.exists().then_some(tj)
}
/// Parity matrix: dynamo-tokenizers vs. transformers.AutoTokenizer.
///
/// Skips cells whose tokenizer.json isn't in the local HF cache. See
/// module-level docs.
#[test]
fn parity_matrix() {
let mut checked = 0;
let mut skipped = vec![];
for model_dir in std::fs::read_dir(fixture_root()).unwrap() {
let model_dir = model_dir.unwrap().path();
if !model_dir.is_dir() {
continue;
}
for shape_file in std::fs::read_dir(&model_dir).unwrap() {
let p = shape_file.unwrap().path();
if p.extension().and_then(|s| s.to_str()) != Some("json") {
continue;
}
let raw = std::fs::read_to_string(&p).unwrap();
let f: Fixture =
serde_json::from_str(&raw).unwrap_or_else(|e| panic!("parse {}: {e}", p.display()));
let Some(tp) = resolve_tokenizer_path(&f.model_id) else {
skipped.push((f.model_id.clone(), f.shape.clone()));
continue;
};
let tok = sgl_router::tokenizer::adapter::load(tp.to_str().unwrap()).unwrap();
let ids = sgl_router::tokenizer::adapter::encode(&tok, &f.prompt_text).unwrap();
assert_eq!(
ids, f.expected_token_ids,
"DRIFT on {}/{}",
f.model_id, f.shape
);
checked += 1;
}
}
let expected = std::fs::read_dir(fixture_root())
.unwrap()
.filter_map(|e| e.ok())
.filter(|e| e.path().is_dir())
.map(|e| {
std::fs::read_dir(e.path())
.unwrap()
.filter_map(|f| f.ok())
.filter(|f| f.path().extension().and_then(|s| s.to_str()) == Some("json"))
.count()
})
.sum::<usize>();
assert_eq!(
checked + skipped.len(),
expected,
"expected {expected} fixtures, found {}",
checked + skipped.len()
);
if checked == 0 {
let families: Vec<String> = skipped
.iter()
.map(|(m, _)| m.clone())
.collect::<std::collections::BTreeSet<_>>()
.into_iter()
.collect();
let msg = format!(
"parity_matrix: no fixtures could be checked — HF cache empty? skipped {} cells \
across model families: [{}]. The e2e HTTP tokenize test remains the \
authoritative live-model parity gate.",
skipped.len(),
families.join(", "),
);
if std::env::var("SGLANG_IS_IN_CI").as_deref() == Ok("true") {
panic!(
"{msg}\n\nThis is a hard failure in CI: a parity test that validates zero \
cells provides no coverage. Either pre-populate the HF cache for these \
model families on the runner, or remove the parity test."
);
}
eprintln!("{msg}");
} else {
eprintln!(
"parity: {checked} cells passed, {} skipped (no HF snapshot)",
skipped.len()
);
}
}
@@ -0,0 +1,148 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Concurrent-state invariants for the worker/registry/breaker layer.
//!
//! These tests stress the lock-free / single-Mutex paths that production
//! traffic exercises in parallel: many requests calling `breaker.allow()`,
//! many discovery events racing with workers_for() reads, and LoadGuard
//! lifecycles under panics.
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::health::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig};
use sgl_router::workers::{Worker, WorkerRegistry};
/// HalfOpen state must admit at most one probe at a time even under high
/// concurrency. N threads race `allow()` when the breaker is HalfOpen; the
/// invariant is that exactly one observes `true` (the probe holder); the
/// rest see `false` because `probe_in_flight` is already set.
#[tokio::test(start_paused = true)]
async fn breaker_half_open_admits_only_one_probe_concurrently() {
let cb = Arc::new(CircuitBreaker::with_config(CircuitBreakerConfig {
threshold: std::num::NonZeroU32::new(1).unwrap(),
cool_down: Duration::from_millis(50),
}));
// Trip into Open.
cb.record_failure();
assert!(!cb.allow(), "must be Open immediately after a failure");
// Advance the paused clock past cool_down so the next `allow()` will
// attempt the Open → HalfOpen transition.
tokio::time::advance(Duration::from_millis(60)).await;
let admitted = Arc::new(AtomicUsize::new(0));
let mut handles = Vec::new();
for _ in 0..32 {
let cb = cb.clone();
let admitted = admitted.clone();
handles.push(tokio::spawn(async move {
if cb.allow() {
admitted.fetch_add(1, Ordering::Relaxed);
}
}));
}
for h in handles {
h.await.unwrap();
}
assert_eq!(
admitted.load(Ordering::Relaxed),
1,
"exactly one probe must be admitted in HalfOpen",
);
}
/// Concurrent `add_with_cb` (upsert) and `remove` from many threads on the
/// same WorkerId must not panic, must not deadlock, and must leave a
/// consistent index — `workers_for(model)` may return 0 or 1 worker, but
/// must never resolve to a worker that has been removed.
#[test]
fn registry_concurrent_add_remove_keeps_indexes_consistent() {
let r = Arc::new(WorkerRegistry::default());
let model = ModelId("m".into());
let mut handles = Vec::new();
for i in 0..8 {
let r = r.clone();
let model = model.clone();
handles.push(std::thread::spawn(move || {
for _ in 0..200 {
let _ = r.add(WorkerSpec {
id: WorkerId(format!("w{i}")),
url: format!("http://w{i}:30000"),
mode: WorkerMode::Plain,
model_ids: vec![model.clone()],
bootstrap_port: None,
});
let snapshot = r.workers_for(&model);
for w in &snapshot {
// Cross-index invariant: an entry surfaced via
// `by_model[m]` must come from a Worker whose own
// `model_ids` includes `m`. An earlier version of
// this assertion checked `w.id.0.starts_with('w')`,
// which is a tautology — every id is `w0..w7` by
// construction — and a regression where `by_model`
// pointed at the wrong Worker (e.g., a stale entry
// left after an upsert that should have cleared its
// by_model membership for the dropped model) would
// pass silently. We can't `re-get by_id and ptr_eq`
// because a concurrent remove can drop the by_id
// entry between the two reads — `Arc` keeps the
// Worker alive on our side but the index map is
// gone. The model-membership claim, however, is a
// property of the Arc itself and stays stable.
assert!(
w.model_ids.contains(&model),
"cross-index drift: by_model[{model:?}] surfaced \
{:?} whose own model_ids = {:?}",
w.id,
w.model_ids,
);
}
r.remove(&WorkerId(format!("w{i}")));
}
}));
}
for h in handles {
h.join().unwrap();
}
// After every thread finishes, every removed worker must really be gone.
assert!(
r.workers_for(&model).is_empty(),
"registry must be empty after all threads finished their add/remove cycles",
);
}
/// `LoadGuard` must decrement the counter during a panic-unwind, not just
/// on a normal scope exit. Rust's RAII contract via `Drop` covers this,
/// but a future refactor (e.g. adding a manual decrement on a non-panic
/// path) could silently regress it. This test pins the invariant.
#[test]
fn load_guard_decrements_on_panic_unwind() {
let w = Arc::new(Worker::new(WorkerSpec {
id: WorkerId("w".into()),
url: "http://x:30000".into(),
mode: WorkerMode::Plain,
model_ids: vec![ModelId("m".into())],
bootstrap_port: None,
}));
assert_eq!(w.active_load(), 0);
let w_inner = w.clone();
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || {
let _g = w_inner.load_guard();
assert_eq!(w_inner.active_load(), 1);
panic!("synthetic panic to exercise Drop on unwind");
}));
assert!(result.is_err(), "the closure must have panicked");
assert_eq!(
w.active_load(),
0,
"LoadGuard's Drop must decrement even when the holder panics",
);
}
@@ -0,0 +1,511 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use axum::{routing::get, Json, Router};
use serde_json::{json, Value};
use sgl_router::discovery::{DiscoveryEvent, ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::workers::{manager, WorkerRegistry};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::net::TcpListener;
use tokio::sync::{mpsc, oneshot};
/// Spin up a tiny fake worker that returns `body` on `GET /server_info`.
/// Returns the worker base URL and a shutdown channel.
async fn spawn_fake_worker(body: Value) -> (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();
async move { 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)
}
fn spec_for(id: &str, url: &str, mode: WorkerMode) -> WorkerSpec {
// model_ids are intentionally empty: the manager resolves them via
// /server_info introspection. Pre-populating here would lie about
// what discovery backends actually emit.
WorkerSpec {
id: WorkerId(id.into()),
url: url.into(),
mode,
model_ids: Vec::new(),
bootstrap_port: None,
}
}
#[tokio::test]
async fn manager_processes_added_then_removed() {
let (url_a, _s_a) = spawn_fake_worker(json!({"served_model_name": "m"})).await;
let (url_b, _s_b) = spawn_fake_worker(json!({"served_model_name": "m"})).await;
let (tx, rx) = mpsc::channel(16);
let registry = Arc::new(WorkerRegistry::default());
let h = tokio::spawn(manager::run(rx, registry.clone()));
tx.send(DiscoveryEvent::Added(spec_for(
"w1",
&url_a,
WorkerMode::Plain,
)))
.await
.unwrap();
tx.send(DiscoveryEvent::Added(spec_for(
"w2",
&url_b,
WorkerMode::Plain,
)))
.await
.unwrap();
// Give the manager time to drain.
tokio::time::sleep(Duration::from_millis(200)).await;
assert_eq!(registry.workers_for(&ModelId("m".into())).len(), 2);
tx.send(DiscoveryEvent::Removed {
id: WorkerId("w1".into()),
})
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(registry.workers_for(&ModelId("m".into())).len(), 1);
drop(tx);
h.await.unwrap();
}
#[tokio::test]
async fn manager_handles_mode_changed() {
let (url, _s) = spawn_fake_worker(json!({"served_model_name": "m"})).await;
let (tx, rx) = mpsc::channel(16);
let registry = Arc::new(WorkerRegistry::default());
let h = tokio::spawn(manager::run(rx, registry.clone()));
tx.send(DiscoveryEvent::Added(spec_for(
"w1",
&url,
WorkerMode::Prefill,
)))
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(200)).await;
assert_eq!(
registry
.workers_for_mode(&ModelId("m".into()), WorkerMode::Prefill)
.len(),
1
);
tx.send(DiscoveryEvent::ModeChanged {
id: WorkerId("w1".into()),
mode: WorkerMode::Decode,
})
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(
registry
.workers_for_mode(&ModelId("m".into()), WorkerMode::Prefill)
.len(),
0
);
assert_eq!(
registry
.workers_for_mode(&ModelId("m".into()), WorkerMode::Decode)
.len(),
1
);
drop(tx);
h.await.unwrap();
}
#[tokio::test]
async fn mode_changed_preserves_active_requests_and_breaker() {
let (url, _s) = spawn_fake_worker(json!({"served_model_name": "m"})).await;
let (tx, rx) = mpsc::channel(16);
let registry = Arc::new(WorkerRegistry::default());
let h = tokio::spawn(manager::run(rx, registry.clone()));
tx.send(DiscoveryEvent::Added(spec_for(
"w1",
&url,
WorkerMode::Prefill,
)))
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(200)).await;
// Grab a handle, bump active_requests, and open the breaker.
let w = registry.get(&WorkerId("w1".into())).unwrap();
w.active_requests.fetch_add(5, Ordering::Relaxed);
// Default threshold is 3 — record 10 failures to guarantee Open state.
for _ in 0..10 {
w.breaker.record_failure();
}
let breaker_open_before = !w.breaker.allow();
assert!(
breaker_open_before,
"breaker should be open after 10 failures"
);
// Flip mode via ModeChanged.
tx.send(DiscoveryEvent::ModeChanged {
id: WorkerId("w1".into()),
mode: WorkerMode::Decode,
})
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
// Re-fetch the Worker handle from the registry.
let w_after = registry.get(&WorkerId("w1".into())).unwrap();
assert_eq!(
w_after.mode(),
WorkerMode::Decode,
"mode should have flipped to Decode"
);
assert_eq!(
w_after.active_requests.load(Ordering::Relaxed),
5,
"active_requests should be preserved across mode change"
);
assert!(
!w_after.breaker.allow(),
"breaker open state should be preserved across mode change"
);
// Critical: the Arc identity must be the same — mutation in place.
assert!(
Arc::ptr_eq(&w, &w_after),
"Worker handle should be the SAME Arc, not a fresh replacement"
);
drop(tx);
h.await.unwrap();
}
/// An out-of-order `ModeChanged` for a worker the registry does not know
/// about (e.g. a buggy discovery backend reordered `Removed` and
/// `ModeChanged`) must not panic, must not silently log INFO claiming the
/// mode flip happened, and must leave the registry untouched.
#[tokio::test]
async fn manager_handles_orphan_mode_changed_without_panic() {
let (tx, rx) = mpsc::channel(16);
let registry = Arc::new(WorkerRegistry::default());
let h = tokio::spawn(manager::run(rx, registry.clone()));
tx.send(DiscoveryEvent::ModeChanged {
id: WorkerId("ghost".into()),
mode: WorkerMode::Decode,
})
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
assert!(
registry.get(&WorkerId("ghost".into())).is_none(),
"an orphan ModeChanged must not create a phantom worker",
);
assert_eq!(
registry.workers_for(&ModelId("m".into())).len(),
0,
"registry must be empty after an orphan event",
);
drop(tx);
h.await.unwrap();
}
/// A `Removed` for an unknown id is a no-op — registry stays empty, manager
/// keeps running.
#[tokio::test]
async fn manager_handles_orphan_removed_without_panic() {
let (tx, rx) = mpsc::channel(16);
let registry = Arc::new(WorkerRegistry::default());
let h = tokio::spawn(manager::run(rx, registry.clone()));
tx.send(DiscoveryEvent::Removed {
id: WorkerId("ghost".into()),
})
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
assert!(registry.is_empty());
drop(tx);
h.await.unwrap();
}
/// Duplicate `Added` for the same id is an upsert — the registry ends up
/// with exactly one worker. The model resolved by /server_info wins on
/// re-add (a different worker may advertise a different served model).
#[tokio::test]
async fn manager_handles_duplicate_added_as_upsert() {
let (url_first, _s_first) = spawn_fake_worker(json!({"served_model_name": "m1"})).await;
let (url_second, _s_second) = spawn_fake_worker(json!({"served_model_name": "m1"})).await;
let (tx, rx) = mpsc::channel(16);
let registry = Arc::new(WorkerRegistry::default());
let h = tokio::spawn(manager::run(rx, registry.clone()));
tx.send(DiscoveryEvent::Added(spec_for(
"w1",
&url_first,
WorkerMode::Plain,
)))
.await
.unwrap();
tx.send(DiscoveryEvent::Added(spec_for(
"w1",
&url_second,
WorkerMode::Plain,
)))
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(300)).await;
assert_eq!(
registry.workers_for(&ModelId("m1".into())).len(),
1,
"w1 still serves m1 after the second Added",
);
drop(tx);
h.await.unwrap();
}
/// Spawn a fake worker whose `/server_info` returns `body` only after
/// sleeping for `delay`. Returns the worker URL and a shutdown channel.
async fn spawn_slow_worker(body: Value, delay: Duration) -> (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();
async move {
tokio::time::sleep(delay).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.
async fn spawn_counting_worker(body: Value) -> (String, Arc<AtomicUsize>, oneshot::Sender<()>) {
let body = Arc::new(body);
let counter = Arc::new(AtomicUsize::new(0));
let counter_clone = counter.clone();
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 counter = counter_clone.clone();
async move {
counter.fetch_add(1, Ordering::SeqCst);
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}"), 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.
#[tokio::test]
async fn added_events_run_in_parallel() {
let delay = Duration::from_millis(200);
let n = 5;
let mut workers = Vec::new();
for _ in 0..n {
workers.push(spawn_slow_worker(json!({"served_model_name": "m"}), delay).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}"),
url,
WorkerMode::Plain,
)))
.await
.unwrap();
}
let registered = tokio::time::timeout(Duration::from_secs(5), async {
loop {
if registry.workers_for(&ModelId("m".into())).len() == n {
return true;
}
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();
}
/// A `Removed` issued while the matching `Added` is still mid-fetch
/// must await the in-flight registration handle before removing.
/// Without that ordering the removal runs first (registry has nothing
/// to remove), then the Added's deferred registry write leaks the
/// worker.
#[tokio::test]
async fn removed_awaits_pending_added() {
let (url, _s) = spawn_slow_worker(
json!({"served_model_name": "m"}),
Duration::from_millis(300),
)
.await;
let (tx, rx) = mpsc::channel(16);
let registry = Arc::new(WorkerRegistry::default());
let h = tokio::spawn(manager::run(rx, registry.clone()));
tx.send(DiscoveryEvent::Added(spec_for(
"w-slow",
&url,
WorkerMode::Plain,
)))
.await
.unwrap();
tx.send(DiscoveryEvent::Removed {
id: WorkerId("w-slow".into()),
})
.await
.unwrap();
// Wait long enough for the Added's /server_info to complete (300ms),
// then assert the worker is gone. If Removed ran before Added's
// registry write, the post-fetch write would leak the entry.
tokio::time::sleep(Duration::from_millis(600)).await;
assert!(
registry.get(&WorkerId("w-slow".into())).is_none(),
"Removed must await the in-flight Added; otherwise the deferred \
registry write leaks the worker"
);
drop(tx);
h.await.unwrap();
}
/// The manager must make exactly ONE `/server_info` request per worker.
/// Before this fix the worker manager fetched `served_model_name` and
/// `KvEventIndex::add_worker` fetched the `kv_events` block
/// independently — 2N round-trips for N workers.
#[tokio::test]
async fn manager_emits_single_server_info_fetch_per_worker() {
use sgl_router::policies::kv_events::KvEventIndex;
let body = json!({
"served_model_name": "m",
"kv_events": {
"publisher": "zmq",
"endpoint_host": "127.0.0.1",
"endpoint_port_base": 60100,
"topic": "",
"block_size": 64,
"dp_size": 1,
}
});
let (url, counter, _s) = spawn_counting_worker(body).await;
let (tx, rx) = mpsc::channel(16);
let registry = Arc::new(WorkerRegistry::default());
let kv_index = KvEventIndex::new();
let h = tokio::spawn(manager::run_with_config(
rx,
registry.clone(),
None,
Some(kv_index.clone()),
None,
));
tx.send(DiscoveryEvent::Added(spec_for(
"w1",
&url,
WorkerMode::Plain,
)))
.await
.unwrap();
// Wait for both the registry and kv-events index to reflect the worker.
let ready = tokio::time::timeout(Duration::from_secs(2), async {
loop {
if registry.get(&WorkerId("w1".into())).is_some() && kv_index.known_worker_count() == 1
{
return true;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await;
assert!(
ready.is_ok(),
"manager did not finish onboarding the worker"
);
let hits = counter.load(Ordering::SeqCst);
assert_eq!(
hits, 1,
"manager must fetch /server_info exactly once per worker (got {hits})"
);
drop(tx);
h.await.unwrap();
kv_index.shutdown().await;
}
@@ -0,0 +1,5 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
mod concurrent_state;
mod manager;
@@ -0,0 +1,266 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
# SPDX-License-Identifier: Apache-2.0
"""Content-based cross-router routing test for cache-aware-zmq.
Two routers + two SGLang workers + one shared model. Each router runs an
independent ``cache_aware_zmq`` policy whose ``KvEventIndex`` subscribes
to **both** workers' KV publishers.
The test warms each worker with a DIFFERENT prefix DIRECTLY (bypassing
both routers), then sends those prefixes through each router and
asserts that routing follows the prefix CONTENT: ``PREFIX_X`` lands on
the worker holding X, ``PREFIX_Y`` lands on the worker holding Y, on
both routers.
# Why content-based, not convergence
An earlier version of this test asserted that both routers converged on
the *same dominant worker* after a one-prefix warmup. That property
sounds like it pins the ZMQ-fan-out contract, but it doesn't: when the
KV-event path is broken (subscribers never opened, e.g. a worker's
``/server_info`` lacks the ``kv_events`` block), ``cache_aware_zmq``
silently degrades to **min-load** — which, with sequential requests
holding ``active_load`` at zero, picks the same worker deterministically
on every call within a router. Both routers' min-load picks happened to
agree often enough (about half the time, modulo HashSet seed) to make
the convergence assertion pass even when no event ever flowed.
Content-based routing is uniquely sensitive to the KV-event path. Two
disjoint prefixes warmed on two different workers can only be routed
correctly if the router knows *which worker holds which content* — the
only mechanism that supplies that information is the ``BlockStored``
event stream. Under min-load fallback, both prefixes route to the same
default worker on each router, so the ``PREFIX_Y → worker_y`` assertion
fails regardless of which worker min-load defaults to.
"""
from __future__ import annotations
import re
import time
import httpx
import pytest
from infra.gateway import Gateway
from infra.model_pool import PASSTHROUGH_CHAT_TEMPLATE_PATH, spawn_worker
from infra.model_specs import get_model_spec
# Disjoint prefixes — share no common opening text, so block 0 hashes
# differ from the first block onward and each worker's HashTree
# contribution is uniquely identifying.
#
# Length matters: each prefix must span ≥2 SGLang blocks at the default
# block_size of 64 tokens so the worker actually emits BlockStored
# events. Below that, the publisher stays quiet and we'd be testing
# min-load by accident — the exact failure mode this test exists to
# rule out.
_PREFIX_X_BODY = (
"Apricot bouquet cinnamon dewdrop elderflower fennel garlic "
"hibiscus indigo jasmine kumquat lavender mint nutmeg oregano "
"paprika quince rosemary saffron tarragon. "
)
PREFIX_X = (_PREFIX_X_BODY * 8).strip()
_PREFIX_Y_BODY = (
"Zephyr yellow xylophone wombat vortex umbrella thistle saffron "
"quartz peppermint orchid nightshade marigold lemongrass kale "
"juniper iris hyacinth gardenia foxglove. "
)
PREFIX_Y = (_PREFIX_Y_BODY * 8).strip()
_REQ_TOTAL_RE = re.compile(
r"^sgl_router_requests_total\{([^}]*)\}\s+(\d+(?:\.\d+)?)\s*$"
)
_LABEL_RE = re.compile(r'(\w+)="([^"]*)"')
def _success_counts_by_worker(router_url: str) -> dict[str, int]:
"""Scrape ``/metrics`` and return ``{worker_url: success_count}``."""
r = httpx.get(f"{router_url}/metrics", timeout=5.0)
r.raise_for_status()
counts: dict[str, int] = {}
for line in r.text.splitlines():
m = _REQ_TOTAL_RE.match(line)
if not m:
continue
labels = dict(_LABEL_RE.findall(m.group(1)))
if labels.get("outcome") != "success":
continue
worker = labels.get("worker_url")
if not worker:
continue
try:
counts[worker] = counts.get(worker, 0) + int(float(m.group(2)))
except ValueError:
continue
return counts
def _send_chat(url: str, model_id: str, prompt: str) -> int:
"""POST one chat completion; return the HTTP status."""
r = httpx.post(
f"{url}/v1/chat/completions",
json={
"model": model_id,
"messages": [{"role": "user", "content": prompt}],
"max_tokens": 4,
"stream": False,
},
timeout=60.0,
)
return r.status_code
def _direct_warm(worker_url: str, model_id: str, prefix: str) -> None:
"""Send one ``/v1/chat/completions`` request with ``prefix`` DIRECTLY to a worker.
The KV-event publisher emits ``BlockStored`` as the request's
prompt blocks commit to that worker's cache; routers subscribed to
the publisher receive the event and add ``(block_hash → worker)``
entries to their ``HashTree``. The test then exercises those
entries by routing through the router.
Direct-warming (rather than going through a router) is the load-
bearing detail: routing through a router would itself choose which
worker to populate, so the two workers' HashTree state would no
longer be uniquely identifying.
Token alignment with the router — ``cache_aware_zmq`` hashes
``messages[*].content`` RAW (``cache_aware_zmq.rs::extract_prompt_text``)
using ``add_special_tokens=false``. By default SGLang's chat
endpoint would wrap ``prefix`` in the model's chat template before
tokenizing — adding role tags, end-of-turn markers, and a
generation prompt — and the resulting block hashes would never
match what the router computes from raw content.
The test launches each worker with ``--chat-template
<PASSTHROUGH_CHAT_TEMPLATE_PATH>``: a Jinja template that emits
only ``messages[*].content`` (the same shape the router extracts),
and which combines with Transformers' ``apply_chat_template(
tokenize=True, add_special_tokens=False)`` to produce the same
token stream the router will compute. So warm and route hash the
same blocks via the same endpoint.
"""
r = httpx.post(
f"{worker_url}/v1/chat/completions",
json={
"model": model_id,
"messages": [{"role": "user", "content": prefix}],
"max_tokens": 4,
"stream": False,
},
timeout=60.0,
)
assert (
r.status_code == 200
), f"direct warm to {worker_url} failed: HTTP {r.status_code} {r.text!r}"
def _route_through(router_url: str, model_id: str, prompt: str) -> str:
"""Send one request through ``router_url``; return which worker handled it.
Computed by diffing the per-worker success-counter on ``/metrics``
around the call. Asserts exactly one worker absorbed the request
(no partial counts, no cancellation race).
"""
before = _success_counts_by_worker(router_url)
code = _send_chat(router_url, model_id, prompt)
assert code == 200, f"request to {router_url} failed: HTTP {code}"
after = _success_counts_by_worker(router_url)
deltas = {w: after.get(w, 0) - before.get(w, 0) for w in set(after) | set(before)}
winners = [w for w, d in deltas.items() if d > 0]
assert (
len(winners) == 1
), f"expected exactly one worker delta on {router_url}, got {deltas}"
return winners[0]
@pytest.mark.real_gpu
@pytest.mark.slow
def test_two_routers_route_by_prefix_content(
router_binary, # noqa: ARG001 — fixture forces release-binary presence
gpu_allocator,
):
"""Each router must route by prefix CONTENT, agreeing across routers.
With each worker direct-warmed by a different disjoint prefix, the
only way a router can route ``PREFIX_X → worker_x`` AND
``PREFIX_Y → worker_y`` is by consulting a HashTree populated from
the BlockStored events the workers emit. Min-load fallback (the
failure mode when no SUB socket opened) is content-blind and would
route both prefixes to whichever worker its tiebreaker prefers.
"""
spec = get_model_spec("qwen3-0.6b")
gpus = gpu_allocator.acquire(2)
# Passthrough chat template — see _direct_warm for the rationale. Both
# workers must run with the same template; otherwise their KV blocks
# would hash template-wrapped tokens while the router hashes raw
# content, and every lookup would miss the tree.
worker_chat_template_args = ["--chat-template", PASSTHROUGH_CHAT_TEMPLATE_PATH]
try:
with (
spawn_worker(
"qwen3-0.6b",
gpu_ids=[gpus[0]],
enable_kv_events=True,
extra_args=worker_chat_template_args,
) as worker_x,
spawn_worker(
"qwen3-0.6b",
gpu_ids=[gpus[1]],
enable_kv_events=True,
extra_args=worker_chat_template_args,
) as worker_y,
Gateway() as router_a,
Gateway() as router_b,
):
worker_urls = [worker_x.url, worker_y.url]
for gw in (router_a, router_b):
gw.start_regular(
model_id=spec["model"],
tokenizer_path=spec["model"],
worker_urls=worker_urls,
policy="cache_aware_zmq",
timeout=120.0,
)
# 1. Direct-warm each worker with its own prefix. Must happen
# AFTER both routers have started — ZMQ PUB/SUB doesn't
# replay messages emitted before SUB attaches, so any
# BlockStored event predating subscription is lost and
# the HashTree never sees it.
_direct_warm(worker_x.url, spec["model"], PREFIX_X)
_direct_warm(worker_y.url, spec["model"], PREFIX_Y)
# 2. Drain the SUB mpsc + pump-apply path. Sub-second under
# loopback ZMQ; 2 s leaves comfortable headroom.
time.sleep(2.0)
# 3. Content-routing assertion (×4): each prefix must land
# on the worker that holds it, on either router.
#
# The four assertions below are independently strong:
# min-load fallback routes both prefixes on a given
# router to a single default worker, so for ANY broken-
# fan-out scenario at least one of the four fails.
for router, label in ((router_a, "A"), (router_b, "B")):
landed = _route_through(router.base_url, spec["model"], PREFIX_X)
assert landed == worker_x.url, (
f"router {label}: PREFIX_X must route to worker_x "
f"({worker_x.url}); landed on {landed}. "
f"Likely cause: HashTree is empty — KV-event "
f"subscriber never opened, or BlockStored events "
f"never reached the pump."
)
landed = _route_through(router.base_url, spec["model"], PREFIX_Y)
assert landed == worker_y.url, (
f"router {label}: PREFIX_Y must route to worker_y "
f"({worker_y.url}); landed on {landed}. "
f"Likely cause: HashTree is empty — KV-event "
f"subscriber never opened, or BlockStored events "
f"never reached the pump."
)
finally:
gpu_allocator.release(gpus)
@@ -0,0 +1,98 @@
"""Basic chat-completions correctness — ported from SMG's
``e2e_test/chat_completions/test_validation.py``, narrowed to the
subset that exercises sgl-router (not SMG's per-message validators).
The shape:
- single-worker regular-mode router
- non-streaming + streaming chat completion
- assistant message non-empty, role correct, finish_reason set
These are the smoke tests that run first; if they pass, the heavier
multi-worker acceptance tests are worth running.
"""
from __future__ import annotations
import httpx
import pytest
from infra.gateway import Gateway
from infra.model_pool import spawn_worker
from infra.model_specs import get_model_spec
@pytest.mark.real_gpu
def test_chat_non_streaming_returns_assistant_message(
router_binary, # noqa: ARG001
gpu_allocator,
):
gpu = gpu_allocator.acquire(1)
try:
with spawn_worker("qwen3-0.6b", gpu_ids=gpu) as worker:
spec = get_model_spec("qwen3-0.6b")
with Gateway() as gw:
gw.start_regular(
model_id=spec["model"],
tokenizer_path=spec["model"],
worker_urls=[worker.url],
timeout=120.0,
)
resp = httpx.post(
f"{gw.base_url}/v1/chat/completions",
json={
"model": spec["model"],
"messages": [{"role": "user", "content": "Say hi."}],
"max_tokens": 16,
"stream": False,
},
timeout=60.0,
)
assert resp.status_code == 200, resp.text
body = resp.json()
choice = body["choices"][0]
assert choice["message"]["role"] == "assistant"
assert choice["message"][
"content"
], f"empty assistant content: {choice!r}"
assert choice.get("finish_reason"), choice
finally:
gpu_allocator.release(gpu)
@pytest.mark.real_gpu
def test_chat_streaming_emits_sse_chunks_with_done(
router_binary, # noqa: ARG001
gpu_allocator,
):
gpu = gpu_allocator.acquire(1)
try:
with spawn_worker("qwen3-0.6b", gpu_ids=gpu) as worker:
spec = get_model_spec("qwen3-0.6b")
with Gateway() as gw:
gw.start_regular(
model_id=spec["model"],
tokenizer_path=spec["model"],
worker_urls=[worker.url],
timeout=120.0,
)
chunks: list[str] = []
with httpx.stream(
"POST",
f"{gw.base_url}/v1/chat/completions",
json={
"model": spec["model"],
"messages": [{"role": "user", "content": "Say hi."}],
"max_tokens": 16,
"stream": True,
},
timeout=60.0,
) as resp:
assert resp.status_code == 200, resp.read().decode()
for line in resp.iter_lines():
if line.startswith("data:"):
chunks.append(line.strip())
assert len(chunks) >= 2, f"expected >=2 SSE chunks, got: {chunks}"
assert any(
"[DONE]" in c for c in chunks
), f"no [DONE] terminator in stream: {chunks}"
finally:
gpu_allocator.release(gpu)
@@ -0,0 +1,327 @@
"""Pytest fixtures for ``experimental/sgl-router/tests/e2e/``.
Two flavors of fixtures coexist here:
1. **Session-scoped smoke fixtures** (``sglang_server`` + ``router``) —
launch ONE SGLang worker + ONE router on fixed ports for the whole
test session. Used by the lightweight ``test_chat_smoke.py`` /
``test_tokenize_smoke.py`` files. These are the cheap "did the
binary start at all" sanity tests.
2. **Per-test multi-worker fixtures** (``router_binary`` +
``gpu_allocator``) — just enough infra for the acceptance tests in
``chat_completions/`` to bring up their own multi-worker
topologies. Backed by the ``infra.gateway.Gateway`` and
``infra.model_pool.spawn_worker`` helpers.
Both sets share the same release binary; ``SGL_ROUTER_BINARY`` env var
overrides the path for both.
"""
from __future__ import annotations
import logging
import os
import signal
import subprocess
import sys
import tempfile
import threading
import time
from collections.abc import Iterator
from pathlib import Path
import httpx
import pytest
logger = logging.getLogger(__name__)
# Make `from infra import gateway, model_pool, model_specs` resolve from
# tests under tests/e2e/ without requiring a sibling `__init__.py` chain.
# Mirrors SMG's e2e_test/conftest.py sys.path setup.
_E2E_DIR = Path(__file__).resolve().parent
if str(_E2E_DIR) not in sys.path:
sys.path.insert(0, str(_E2E_DIR))
MODEL = "Qwen/Qwen3-0.6B"
SGLANG_PORT = 30000
ROUTER_PORT = 8090
# Path to the release binary. This file lives at
# `experimental/sgl-router/tests/e2e/conftest.py`, so:
# parent = tests/e2e/
# parent.parent = tests/
# parent.parent.parent = experimental/sgl-router/ ← cargo workspace root
# A previous version used `parent.parent / "target"`, which pointed at
# `experimental/sgl-router/tests/target/` and silently broke every
# fixture that tries to launch the router binary (CI's
# `cargo build --release` lands the artifact at
# `experimental/sgl-router/target/release/sgl-router`, not under
# `tests/`).
_SGL_ROUTER_ROOT = Path(__file__).parent.parent.parent
_BINARY = (
Path(os.environ.get("CARGO_TARGET_DIR", str(_SGL_ROUTER_ROOT / "target")))
/ "release"
/ "sgl-router"
)
def _wait_http(url: str, timeout: int = 120) -> None:
"""Poll *url* until it returns 2xx or raises RuntimeError on timeout."""
deadline = time.time() + timeout
last_exc: Exception | None = None
while time.time() < deadline:
try:
resp = httpx.get(url, timeout=5)
if resp.status_code < 300:
return
except Exception as exc: # noqa: BLE001
last_exc = exc
time.sleep(5)
raise RuntimeError(
f"Timed out waiting for {url} after {timeout}s (last error: {last_exc})"
)
@pytest.fixture(scope="session")
def sglang_server():
"""Launch a real SGLang server on port 30000 and wait until healthy."""
# Stream the server's stdout/stderr to a file rather than capturing
# to subprocess.PIPE. The launch_server startup log is verbose (model
# download, JIT warmup, NCCL init); once a PIPE'd output fills its
# ~64 KB OS buffer with nothing reading it, the SGLang process
# blocks on stdout write and never reaches "Server started" — the
# health probe then times out at 300 s and we have no visibility
# into *why*. A real log file fixes both (no buffer pressure, and
# the file is dumped on failure for triage).
log_path = Path(tempfile.gettempdir()) / f"sglang-server-{SGLANG_PORT}.log"
log_handle = open(log_path, "w", buffering=1) # line-buffered
proc = subprocess.Popen(
[
"python3",
"-m",
"sglang.launch_server",
"--model-path",
MODEL,
"--port",
str(SGLANG_PORT),
"--tp",
"1",
],
stdout=log_handle,
stderr=subprocess.STDOUT,
)
try:
_wait_http(f"http://localhost:{SGLANG_PORT}/health", timeout=300)
except Exception:
# Dump the server log so the operator can see why startup failed
# (model download error, port conflict, OOM, JIT crash, etc.).
proc.send_signal(signal.SIGTERM)
try:
proc.wait(timeout=30)
except subprocess.TimeoutExpired:
proc.kill()
proc.wait()
log_handle.flush()
log_handle.close()
try:
tail = log_path.read_text(errors="replace").splitlines()[-200:]
except OSError:
tail = ["(server log unreadable)"]
logger.error(
"sglang_server fixture failed; last 200 log lines from %s:\n%s",
log_path,
"\n".join(tail),
)
raise
yield f"http://localhost:{SGLANG_PORT}"
proc.send_signal(signal.SIGTERM)
try:
proc.wait(timeout=30)
except subprocess.TimeoutExpired:
proc.kill()
proc.wait()
log_handle.flush()
log_handle.close()
def _find_tokenizer_path(model: str) -> str:
"""Locate the tokenizer.json for *model* from the local HF Hub cache.
Falls back to the model string itself (a valid HF Hub repo identifier
that dynamo-tokenizers can resolve at runtime) when the cache is absent.
"""
try:
from huggingface_hub import try_to_load_from_cache # type: ignore[import]
path = try_to_load_from_cache(model, "tokenizer.json")
if path and Path(path).is_file():
return str(path)
except Exception: # noqa: BLE001
pass
# Let dynamo-tokenizers resolve the repo identifier directly.
return model
def build_smoke_router_config(
*,
host: str,
port: int,
model: str,
tokenizer_path: str,
sglang_url: str,
) -> str:
"""Build the TOML the smoke `router` fixture writes to disk.
Returns ``main_config_text`` carrying ``[server]``, ``[[models]]``,
and ``[discovery] backend = "static_urls"`` with the worker URL
inline. The Rust ``Config`` struct requires a ``[discovery]``
section (``DiscoveryConfig`` has no ``#[serde(default)]``) and has
no top-level ``workers`` field. The previous ``static_file``
backend was replaced by ``static_urls`` (which holds the URL list
inline rather than via a side-car file).
"""
return f"""\
[server]
host = "{host}"
port = {port}
[[models]]
id = "{model}"
tokenizer_path = "{tokenizer_path}"
[discovery]
backend = "static_urls"
[discovery.static_urls]
urls = ["{sglang_url}"]
"""
@pytest.fixture(scope="session")
def router(sglang_server): # noqa: ARG001 (sglang_server must start first)
"""Launch sgl-router on port 8090 pointed at the SGLang worker."""
tok_path = _find_tokenizer_path(MODEL)
cfg_handle = tempfile.NamedTemporaryFile(mode="w", suffix=".toml", delete=False)
cfg_path = Path(cfg_handle.name)
main_text = build_smoke_router_config(
host="0.0.0.0",
port=ROUTER_PORT,
model=MODEL,
tokenizer_path=tok_path,
sglang_url=f"http://localhost:{SGLANG_PORT}",
)
cfg_handle.write(main_text)
cfg_handle.close()
try:
proc = subprocess.Popen(
[str(_BINARY), "--config", str(cfg_path)],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
)
try:
_wait_http(f"http://localhost:{ROUTER_PORT}/readyz", timeout=60)
except Exception:
proc.send_signal(signal.SIGTERM)
proc.wait(timeout=30)
raise
yield f"http://localhost:{ROUTER_PORT}"
proc.send_signal(signal.SIGTERM)
try:
proc.wait(timeout=30)
except subprocess.TimeoutExpired:
proc.kill()
proc.wait()
finally:
cfg_path.unlink(missing_ok=True)
# ---------------------------------------------------------------------------
# Per-test multi-worker acceptance fixtures
# ---------------------------------------------------------------------------
def _detect_gpu_count() -> int:
"""Count visible GPUs via ``nvidia-smi``. Returns 0 when no NVIDIA GPU
is available (CI on CPU-only runners, dev laptops, etc.).
"""
try:
out = subprocess.check_output(
["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"],
stderr=subprocess.DEVNULL,
timeout=5.0,
)
except (FileNotFoundError, subprocess.SubprocessError):
return 0
return len([ln for ln in out.decode().splitlines() if ln.strip()])
class GPUAllocator:
"""Single-process GPU index allocator. Test-scoped; not safe for
cross-process use (pytest-xdist) — each worker would race over the
full GPU set. Acceptance tests run serially, so this is fine.
"""
def __init__(self, total: int):
self.total = total
self._free: list[int] = list(range(total))
self._lock = threading.Lock()
def acquire(self, n: int = 1) -> list[int]:
with self._lock:
if n > len(self._free):
raise pytest.skip.Exception(
f"requested {n} GPUs, only {len(self._free)}/{self.total} free"
)
picked = self._free[:n]
self._free = self._free[n:]
return picked
def release(self, ids: list[int]) -> None:
with self._lock:
self._free.extend(ids)
self._free.sort()
@pytest.fixture(scope="session")
def router_binary() -> Path:
"""Locate the release ``sgl-router`` binary or skip the session.
Used by the multi-worker acceptance tests (which spawn their own
Gateway per test instead of using the session-scoped ``router``
fixture).
"""
env_path = os.environ.get("SGL_ROUTER_BINARY")
candidates: list[Path] = []
if env_path:
candidates.append(Path(env_path))
candidates.append(_BINARY)
for c in candidates:
if c.exists():
return c
pytest.skip(
"sgl-router release binary not found at any of: "
+ ", ".join(str(c) for c in candidates)
+ ". Build with `cargo build --release` in experimental/sgl-router/."
)
@pytest.fixture(scope="session")
def gpu_allocator() -> Iterator[GPUAllocator]:
"""Session-scoped GPU index allocator. Skips the entire session when
no GPUs are visible — acceptance tests under chat_completions/ are
real-GPU.
"""
n = _detect_gpu_count()
if n == 0:
pytest.skip(
"no NVIDIA GPUs visible to nvidia-smi; acceptance tests are GPU-only"
)
yield GPUAllocator(n)
@@ -0,0 +1,404 @@
"""Minimal sgl-router Gateway class — adapted from SMG's e2e_test/infra/gateway.py.
Differences from SMG:
- SMG drives a Python launcher (`python3 -m sglang_router.launch_router`)
with worker URLs on the CLI.
- sgl-router uses a Rust binary (`experimental/sgl-router/target/release/sgl-router`)
with a TOML config file. Worker discovery is config-file-based; this
Gateway writes a TOML to a tempfile and execs the binary with
`--config <tempfile>`.
Supported lifecycles:
- Regular mode: one model, N worker URLs, single policy.
- PD mode: one model, prefill_workers + decode_workers (lists of URLs),
discovery emits separate `WorkerMode::Prefill` / `WorkerMode::Decode`
entries. The router resolves PD pool isolation at request time.
Use as a context manager:
with Gateway() as gw:
gw.start_regular(model_path="...", worker_urls=[...])
resp = httpx.post(f"{gw.base_url}/v1/chat/completions", json=...)
or pytest fixture style (see e2e_test/conftest.py).
"""
from __future__ import annotations
import logging
import os
import signal
import socket
import subprocess
import tempfile
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
import httpx
logger = logging.getLogger(__name__)
# Repo-relative path to the release binary. Set ``SGL_ROUTER_BINARY`` to
# override (e.g. a debug build, or a non-default ``CARGO_TARGET_DIR``).
# This file is at `experimental/sgl-router/tests/e2e/infra/gateway.py`,
# so four `.parent` hops to reach the sgl-router workspace root
# (infra → e2e → tests → sgl-router). Cargo lands the binary at
# `experimental/sgl-router/target/release/sgl-router`. A previous
# version used three hops and pointed at `tests/target/`, which
# would have broken any test that actually launches the router via
# this helper.
DEFAULT_BINARY = (
Path(__file__).resolve().parent.parent.parent.parent
/ "target"
/ "release"
/ "sgl-router"
)
def _get_open_port() -> int:
"""Reserve an ephemeral TCP port in [20000, 55535].
The router itself doesn't have the ``port + 10000`` gRPC-derivation
constraint that SGLang's launch_server does, but we cap the range
anyway so the e2e helpers behave consistently across components.
"""
for _ in range(50):
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
port = s.getsockname()[1]
if 20000 <= port <= 55535:
return port
raise RuntimeError(
"could not allocate an ephemeral port in [20000, 55535] after 50 tries"
)
def _resolve_tokenizer_path(tokenizer_path: str) -> str:
"""Resolve a HuggingFace repo ID to a local ``tokenizer.json`` path.
sgl-router's tokenizer loader treats the input as a filesystem path and
inspects its extension; a bare HF id like ``Qwen/Qwen3-0.6B`` looks
like a file with extension ``.6B`` and is rejected. When the HF Hub
cache already has the tokenizer, point the loader at the on-disk
``tokenizer.json`` directly. Pass paths/URLs through unchanged.
"""
p = Path(tokenizer_path)
if p.exists():
return str(p)
try:
from huggingface_hub import try_to_load_from_cache # type: ignore[import]
cached = try_to_load_from_cache(tokenizer_path, "tokenizer.json")
if cached and Path(cached).is_file():
return str(cached)
except Exception: # noqa: BLE001
pass
return tokenizer_path
@dataclass
class WorkerInfo:
"""Worker visible to the gateway via ``/v1/models``-style introspection.
Mirrors SMG's WorkerInfo shape so test code reads the same. sgl-router
does not currently surface a `/v1/workers` admin API — this is a
placeholder for a future admin surface; current tests scrape
`/metrics` for per-worker observability instead.
"""
id: str
url: str
model: str | None = None
status: str = "unknown"
metadata: dict[str, Any] = field(default_factory=dict)
class Gateway:
"""Lifecycle-managed sgl-router instance for e2e tests.
Not thread-safe; assume one Gateway per test (or per fixture scope).
"""
def __init__(
self,
host: str = "127.0.0.1",
port: int | None = None,
binary: Path | None = None,
proxy_request_timeout_secs: int | None = None,
stale_request_timeout_secs: int | None = None,
):
self.host = host
self.port = port or _get_open_port()
self.base_url = f"http://{self.host}:{self.port}"
# Resolve binary from env override, explicit arg, or repo default.
env_binary = os.environ.get("SGL_ROUTER_BINARY")
if binary is not None:
self.binary = Path(binary)
elif env_binary:
self.binary = Path(env_binary)
else:
self.binary = DEFAULT_BINARY
# Test-side overrides for the router's tunables. Both default to
# `None`, in which case the router uses its production defaults
# (60 s proxy timeout, 300 s stale-request timeout). Tests set
# these short so per-request failures and stale-request expiry
# surface within the test's wall-time budget.
self.proxy_request_timeout_secs = proxy_request_timeout_secs
self.stale_request_timeout_secs = stale_request_timeout_secs
self.process: subprocess.Popen | None = None
self._config_path: Path | None = None
self._started: bool = False
# Track child workers we spawned so __exit__ can tear them down.
self._owned_workers: list[subprocess.Popen] = []
# ----- context manager -------------------------------------------------
def __enter__(self) -> "Gateway":
return self
def __exit__(self, *exc) -> None:
self.shutdown()
# ----- start ----------------------------------------------------------
def start_regular(
self,
*,
model_id: str,
tokenizer_path: str,
worker_urls: list[str],
policy: str = "round_robin",
extra_models: list[dict] | None = None,
timeout: float = 60.0,
) -> None:
"""Start the router in regular (non-PD) mode.
Args:
model_id: The model identifier the router will dispatch under.
tokenizer_path: Path or HF ID for the tokenizer the router uses
for cache-aware tokenization.
worker_urls: URLs of already-running ``sglang.launch_server``
instances. The router uses ``static_urls`` discovery;
each worker's mode (plain) and any disaggregation
metadata are learned from ``/server_info``.
policy: Policy kind — ``round_robin``, ``random``, ``power_of_two``,
or ``cache_aware_zmq``.
timeout: How long to wait for ``/readyz`` before giving up.
"""
self._launch(
self._build_config(
model_id=model_id,
tokenizer_path=tokenizer_path,
urls=list(worker_urls),
policy=policy,
extra_models=extra_models or [],
),
timeout=timeout,
)
def start_pd(
self,
*,
model_id: str,
tokenizer_path: str,
prefill_urls: list[str],
decode_urls: list[str],
policy: str = "round_robin",
timeout: float = 60.0,
) -> None:
"""Start the router in PD-disaggregated mode.
All prefill + decode URLs go into one ``static_urls`` list. The
router seeds each worker as ``WorkerMode::Plain`` and the
manager's ``/server_info`` introspect step overrides mode +
``bootstrap_port`` from the worker's self-disclosure. Workers
must have been launched with ``--disaggregation-mode`` and
``--disaggregation-bootstrap-port`` for the PD role to be
picked up (see ``model_pool.spawn_worker``); modern SGLang is
assumed.
"""
self._launch(
self._build_config(
model_id=model_id,
tokenizer_path=tokenizer_path,
urls=list(prefill_urls) + list(decode_urls),
policy=policy,
extra_models=[],
),
timeout=timeout,
)
# ----- shutdown --------------------------------------------------------
def shutdown(self) -> None:
"""SIGTERM the router; SIGKILL after 30s. Idempotent."""
if self.process is not None and self.process.poll() is None:
try:
self.process.send_signal(signal.SIGTERM)
try:
self.process.wait(timeout=30)
except subprocess.TimeoutExpired:
self.process.kill()
self.process.wait()
except ProcessLookupError:
pass
self.process = None
if self._config_path and self._config_path.exists():
self._config_path.unlink(missing_ok=True)
self._config_path = None
self._started = False
# Tear down any owned upstream workers.
for w in self._owned_workers:
if w.poll() is None:
try:
w.send_signal(signal.SIGTERM)
try:
w.wait(timeout=30)
except subprocess.TimeoutExpired:
w.kill()
w.wait()
except ProcessLookupError:
pass
self._owned_workers.clear()
# ----- HTTP introspection helpers -------------------------------------
def healthy(self, timeout: float = 5.0) -> bool:
try:
resp = httpx.get(f"{self.base_url}/healthz", timeout=timeout)
return resp.status_code == 200
except (httpx.RequestError, httpx.TimeoutException):
return False
def ready(self, timeout: float = 5.0) -> bool:
try:
resp = httpx.get(f"{self.base_url}/readyz", timeout=timeout)
return resp.status_code == 200
except (httpx.RequestError, httpx.TimeoutException):
return False
def metrics_text(self, timeout: float = 5.0) -> str | None:
try:
resp = httpx.get(f"{self.base_url}/metrics", timeout=timeout)
if resp.status_code == 200:
return resp.text
return None
except (httpx.RequestError, httpx.TimeoutException):
return None
# ----- internals ------------------------------------------------------
def _build_config(
self,
*,
model_id: str,
tokenizer_path: str,
urls: list[str],
policy: str,
extra_models: list[dict],
) -> str:
resolved_tokenizer = _resolve_tokenizer_path(tokenizer_path)
extra_model_toml = ""
for em in extra_models:
extra_model_toml += (
f'\n[[models]]\nid = "{em["id"]}"\n'
f'tokenizer_path = "{_resolve_tokenizer_path(em["tokenizer_path"])}"\n'
f'policy = "{em.get("policy", policy)}"\n'
)
# Optional tunables — only emit the [proxy] and [active_load]
# sections if a test has overridden them, so production defaults
# apply otherwise.
proxy_section = ""
if self.proxy_request_timeout_secs is not None:
proxy_section = (
f"\n[proxy]\nrequest_timeout_secs = {self.proxy_request_timeout_secs}\n"
)
active_load_section = ""
if self.stale_request_timeout_secs is not None:
active_load_section = (
f"\n[active_load]\nstale_request_timeout_secs = "
f"{self.stale_request_timeout_secs}\n"
)
urls_toml = ", ".join(f'"{u}"' for u in urls)
return f"""\
[server]
host = "{self.host}"
port = {self.port}
[[models]]
id = "{model_id}"
tokenizer_path = "{resolved_tokenizer}"
policy = "{policy}"
{extra_model_toml}
[discovery]
backend = "static_urls"
[discovery.static_urls]
urls = [{urls_toml}]
{proxy_section}{active_load_section}"""
def _launch(self, config_text: str, *, timeout: float) -> None:
if not self.binary.exists():
raise RuntimeError(
f"sgl-router binary not found at {self.binary}. "
"Build it first: `cd experimental/sgl-router && cargo build --release` "
"or set SGL_ROUTER_BINARY to the binary path."
)
# Write the main config.
fd, path = tempfile.mkstemp(suffix=".toml", prefix="sgl-router-")
os.close(fd)
self._config_path = Path(path)
self._config_path.write_text(config_text, encoding="utf-8")
logger.info("sgl-router config: %s", self._config_path)
logger.debug("sgl-router config text:\n%s", config_text)
self.process = subprocess.Popen(
[str(self.binary), "--config", str(self._config_path)],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
start_new_session=True,
)
try:
self._wait_ready(timeout=timeout)
except Exception:
self.shutdown()
raise
self._started = True
def _wait_ready(self, *, timeout: float) -> None:
deadline = time.time() + timeout
last_exc: Exception | None = None
while time.time() < deadline:
if self.process is not None and self.process.poll() is not None:
# Process exited early — surface stdout/stderr.
out = b""
try:
if self.process.stdout is not None:
out = self.process.stdout.read() or b""
except Exception: # noqa: BLE001
pass
raise RuntimeError(
f"sgl-router exited during startup with code "
f"{self.process.returncode}. output:\n{out.decode(errors='replace')}",
)
try:
resp = httpx.get(f"{self.base_url}/readyz", timeout=2.0)
if resp.status_code == 200:
return
except (httpx.RequestError, httpx.TimeoutException) as exc:
last_exc = exc
time.sleep(0.5)
raise TimeoutError(
f"sgl-router did not become ready at {self.base_url} within {timeout}s "
f"(last error: {last_exc})"
)
@@ -0,0 +1,228 @@
"""Minimal SGLang worker spawner for sgl-router e2e tests.
Adapted from SMG's e2e_test/infra/model_pool.py — the 1200-line original
manages a pool of long-lived workers across many tests; here we only
need a thin wrapper around ``sglang.launch_server`` that:
- allocates GPU(s) for the worker (via ``CUDA_VISIBLE_DEVICES``),
- spawns ``python3 -m sglang.launch_server`` with the right args,
- waits for ``/health`` to come up,
- optionally injects ``--kv-events-config`` so the worker exposes
the ``kv_events`` block on ``/server_info``.
A test owns a ``ModelInstance`` for its duration; teardown shuts the
worker down. No cross-test pooling — the acceptance tests are slow
enough already (model load dominates) that pooling complexity wasn't
worth porting.
"""
from __future__ import annotations
import json
import logging
import os
import signal
import socket
import subprocess
import time
from dataclasses import dataclass, field
from pathlib import Path
import httpx
from .model_specs import get_model_spec
logger = logging.getLogger(__name__)
# Passthrough Jinja chat template that emits ONLY `messages[*].content`
# joined with `\n` — matching the router's cache_aware_zmq prompt
# extraction. A worker launched with
# ``--chat-template <PASSTHROUGH_CHAT_TEMPLATE_PATH>`` tokenizes the
# raw content string, so its KV-block hashes align with what the
# router computes from the same chat-completions request. Test-only.
PASSTHROUGH_CHAT_TEMPLATE_PATH = str(
Path(__file__).parent / "passthrough_chat_template.jinja"
)
def _get_open_port() -> int:
"""Allocate an ephemeral TCP port in the range [20000, 55535].
SGLang derives its internal gRPC port as ``http_port + 10000``; if the
kernel hands us an ephemeral port above 55535, that derivation overflows
65535 and ``ServerArgs.__post_init__`` rejects it. Retrying a bounded
number of times keeps us safely below the ceiling without hand-rolling
a port registry.
"""
for _ in range(50):
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
port = s.getsockname()[1]
if 20000 <= port <= 55535:
return port
raise RuntimeError(
"could not allocate an ephemeral port in [20000, 55535] after 50 tries; "
"SGLang derives its internal gRPC port as http_port + 10000 and "
"rejects values above 65535"
)
@dataclass
class ModelInstance:
"""A running ``sglang.launch_server`` process.
Use as a context manager:
with spawn_worker("qwen3-0.6b", gpu_ids=[0]) as inst:
httpx.post(f"{inst.url}/generate", ...)
"""
url: str
port: int
process: subprocess.Popen
model_id: str
gpu_ids: list[int] = field(default_factory=list)
kv_events_endpoint: str | None = None
def __enter__(self) -> "ModelInstance":
return self
def __exit__(self, *exc) -> None:
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
def spawn_worker(
model_id: str,
*,
gpu_ids: list[int],
port: int | None = None,
enable_kv_events: bool = False,
kv_events_port: int | None = None,
disagg_mode: str | None = None,
bootstrap_port: int | None = None,
extra_args: list[str] | None = None,
timeout: float = 600.0,
) -> ModelInstance:
"""Spawn a single ``sglang.launch_server`` and wait for ``/health``.
Args:
model_id: Key into :data:`model_specs.MODEL_SPECS`.
gpu_ids: Concrete GPU indices to bind via ``CUDA_VISIBLE_DEVICES``.
port: HTTP port; auto-assigned if None.
enable_kv_events: If True, inject ``--kv-events-config`` with a
ZMQ publisher so the router's introspection picks up the
kv_events block from ``/server_info`` (Patch 1).
kv_events_port: ZMQ publisher port. Auto-assigned if None and
``enable_kv_events`` is True.
disagg_mode: "prefill" or "decode" for PD-disagg launches; passed
through as ``--disaggregation-mode``.
bootstrap_port: PD-disagg bootstrap port (prefill side only).
extra_args: Additional CLI args appended verbatim.
timeout: Health-check timeout. Cold-start on a fresh GPU can be
slow; default is 10 minutes.
"""
spec = get_model_spec(model_id)
port = port or _get_open_port()
base_url = f"http://127.0.0.1:{port}"
cmd = [
"python3",
"-m",
"sglang.launch_server",
"--model-path",
spec["model"],
"--port",
str(port),
"--host",
"127.0.0.1",
"--tp",
str(spec.get("tp", 1)),
]
cmd.extend(spec.get("worker_args", []) or [])
kv_events_endpoint: str | None = None
if enable_kv_events:
kv_port = kv_events_port or _get_open_port()
kv_events_endpoint = f"tcp://*:{kv_port}"
kv_cfg = {
"publisher": "zmq",
"endpoint": kv_events_endpoint,
"topic": "kv",
}
cmd.extend(["--kv-events-config", json.dumps(kv_cfg)])
if disagg_mode is not None:
cmd.extend(["--disaggregation-mode", disagg_mode])
if bootstrap_port is not None:
cmd.extend(["--disaggregation-bootstrap-port", str(bootstrap_port)])
if extra_args:
cmd.extend(extra_args)
env = os.environ.copy()
env["CUDA_VISIBLE_DEVICES"] = ",".join(str(g) for g in gpu_ids)
logger.info(
"spawning sglang worker: model=%s port=%d gpus=%s disagg=%s",
model_id,
port,
gpu_ids,
disagg_mode,
)
proc = subprocess.Popen(
cmd,
env=env,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
start_new_session=True,
)
inst = ModelInstance(
url=base_url,
port=port,
process=proc,
model_id=model_id,
gpu_ids=list(gpu_ids),
kv_events_endpoint=kv_events_endpoint,
)
# Wait for /health. Cold-start on H200 with weights uncached can take
# ~5 minutes; CI configurations should pre-warm.
deadline = time.time() + timeout
while time.time() < deadline:
if proc.poll() is not None:
out = b""
try:
if proc.stdout is not None:
out = proc.stdout.read() or b""
except Exception: # noqa: BLE001
pass
raise RuntimeError(
f"sglang worker exited during startup with code {proc.returncode}; "
f"cmd: {' '.join(cmd)}\noutput:\n{out.decode(errors='replace')}",
)
try:
resp = httpx.get(f"{base_url}/health", timeout=2.0)
if resp.status_code == 200:
logger.info("sglang worker ready at %s", base_url)
return inst
except (httpx.RequestError, httpx.TimeoutException):
pass
time.sleep(2.0)
inst.shutdown()
raise TimeoutError(
f"sglang worker did not become healthy at {base_url} within {timeout}s",
)
@@ -0,0 +1,77 @@
"""Model specifications for sgl-router e2e tests.
Adapted from SMG's e2e_test/infra/model_specs.py. The same dict-of-dicts
shape (so test code reads the same) but the entries are narrower —
sgl-router tests today target small/medium models only; the larger
function-calling / reasoning models from SMG are out of scope.
Each entry:
- model: HuggingFace path or local path (env-resolved)
- memory_gb: estimated single-GPU footprint
- tp: tensor-parallel size (= GPUs needed)
- features: feature tags for filtering
- worker_args: optional extra `sglang.launch_server` flags
"""
from __future__ import annotations
import os
# Local-cache root for CI / cluster nodes that pre-download HF weights.
# Mirrors the SMG `ROUTER_LOCAL_MODEL_PATH` env var.
ROUTER_LOCAL_MODEL_PATH = os.environ.get("ROUTER_LOCAL_MODEL_PATH", "")
def _resolve_model_path(hf_path: str) -> str:
"""Prefer a local copy of the model when one exists under
``ROUTER_LOCAL_MODEL_PATH``; otherwise fall back to the HuggingFace ID.
"""
if ROUTER_LOCAL_MODEL_PATH:
local_path = os.path.join(ROUTER_LOCAL_MODEL_PATH, hf_path)
if os.path.exists(local_path):
return local_path
return hf_path
MODEL_SPECS: dict[str, dict] = {
# Fast-start tiny model for convergence / decode-affinity / stale-request
# tests. Single GPU, ~2 GB weights, sub-30s start on a warm cache.
"qwen3-0.6b": {
"model": _resolve_model_path("Qwen/Qwen3-0.6B"),
"memory_gb": 4,
"tp": 1,
"features": ["chat", "streaming"],
},
# Standard small chat model — matches SMG's `llama-1b` entry.
"llama-1b": {
"model": _resolve_model_path("meta-llama/Llama-3.2-1B-Instruct"),
"memory_gb": 4,
"tp": 1,
"features": ["chat", "streaming"],
},
# Primary 8B chat model — matches SMG's `llama-8b`.
"llama-8b": {
"model": _resolve_model_path("meta-llama/Llama-3.1-8B-Instruct"),
"memory_gb": 16,
"tp": 1,
"features": ["chat", "streaming"],
},
}
def get_model_spec(model_id: str) -> dict:
"""Return the spec dict for ``model_id``; KeyError if absent."""
if model_id not in MODEL_SPECS:
raise KeyError(
f"Unknown model: {model_id}. Available: {list(MODEL_SPECS.keys())}"
)
return MODEL_SPECS[model_id]
def get_models_with_feature(feature: str) -> list[str]:
"""Filter model IDs by feature tag (e.g. ``streaming``, ``chat``)."""
return [
model_id
for model_id, spec in MODEL_SPECS.items()
if feature in spec.get("features", [])
]
@@ -0,0 +1,13 @@
{#-
Passthrough chat template for cache-aware-zmq e2e tests.
Emits ONLY `messages[*].content` joined with `\n` — no role markers,
no special tokens, no generation prompt. This is the SAME shape the
router's cache_aware_zmq policy produces in `extract_prompt_text`,
so a worker launched with `--chat-template <this file>` tokenizes the
same string the router will tokenize for routing — making block
hashes align across worker KV cache and router HashTree.
Use only for tests; not appropriate for any real chat workload.
-#}
{{- messages | map(attribute='content') | join('\n') -}}
@@ -0,0 +1,6 @@
FROM python:3.12-slim
WORKDIR /app
RUN pip install --no-cache-dir fastapi uvicorn
COPY fake_worker.py .
EXPOSE 30000
CMD ["python", "fake_worker.py"]
@@ -0,0 +1,39 @@
# syntax=docker/dockerfile:1.6
# Build sgl-router binary for k8s integration E2E.
# Context root: repo root (one level above experimental/sgl-router/).
# Matches rust-toolchain.toml's pinned channel, avoiding an in-build rustup channel-sync.
FROM rust:1.90-bookworm AS builder
# Pin to the exact toolchain pre-installed in the base image so rustup
# doesn't try to sync the channel manifest when it sees rust-toolchain.toml's
# `channel = "1.90"`.
ENV RUSTUP_TOOLCHAIN=1.90.0
# libssl-dev + pkg-config ship with rust:1.90-bookworm already; no apt-get needed.
WORKDIR /build
# Copy just the sgl-router crate (context is the repo root)
COPY experimental/sgl-router /build/experimental/sgl-router
RUN --mount=type=cache,target=/usr/local/cargo/registry \
--mount=type=cache,target=/usr/local/cargo/git \
--mount=type=cache,target=/build/experimental/sgl-router/target \
cd /build/experimental/sgl-router \
&& cargo build --release --bin sgl-router \
&& cp target/release/sgl-router /usr/local/bin/sgl-router
FROM debian:bookworm-slim
RUN apt-get update && apt-get install -y ca-certificates \
&& rm -rf /var/lib/apt/lists/*
COPY --from=builder /usr/local/bin/sgl-router /usr/local/bin/sgl-router
# Tiny tokenizer fixture used by the E2E config
COPY experimental/sgl-router/tests/fixtures/tiny_tokenizer.json /etc/tokenizer/tiny.json
EXPOSE 8090
ENTRYPOINT ["sgl-router"]
@@ -0,0 +1,251 @@
"""Pytest configuration for sgl-router K8s integration tests.
These tests require:
- A kind cluster named 'sgl-router-kind'
- The sgl-router:e2e and sgl-router-fake-worker:e2e images loaded into kind
- kubectl configured to use the kind-sgl-router-kind context
Setup: ./tests/e2e/k8s_integration/setup.sh
Teardown: ./tests/e2e/k8s_integration/setup.sh teardown
"""
from __future__ import annotations
import logging
import socket
import subprocess
import time
import httpx
import pytest
logger = logging.getLogger(__name__)
NAMESPACE = "sgl-router-test"
CLUSTER_NAME = "sgl-router-kind"
KUBECTL_CONTEXT = f"kind-{CLUSTER_NAME}"
# sgl-router discovery reconciliation: if the watcher misses an event the
# reconciler fires within ~60s. Tests that exercise removal wait up to 90s.
RECONCILIATION_WAIT_SECS = 90
# Errors safe to retry while polling (transport-level only — HTTP 4xx/5xx
# are intentionally NOT included so real regressions surface immediately).
_TRANSIENT_ERRORS = (
httpx.TransportError,
httpx.TimeoutException,
ConnectionError,
OSError,
)
def pytest_configure(config):
config.addinivalue_line(
"markers",
"slow: marks tests that wait for multiple reconciliation cycles "
"(deselect with '-m \"not slow\"')",
)
def _kubectl(
*args: str,
check: bool = True,
capture: bool = True,
) -> subprocess.CompletedProcess:
cmd = ["kubectl", "--context", KUBECTL_CONTEXT, *args]
logger.debug("Running: %s", " ".join(cmd))
return subprocess.run(cmd, capture_output=capture, text=True, check=check)
def _apply_from_stdin(yaml_content: str) -> subprocess.CompletedProcess:
return subprocess.run(
["kubectl", "--context", KUBECTL_CONTEXT, "apply", "-f", "-"],
input=yaml_content,
capture_output=True,
text=True,
check=True,
)
def _wait_for_deployment_ready(
name: str,
namespace: str = NAMESPACE,
timeout: int = 180,
) -> None:
_kubectl(
"rollout",
"status",
f"deployment/{name}",
"-n",
namespace,
f"--timeout={timeout}s",
)
def _wait_for_pod_ready(
name: str,
namespace: str = NAMESPACE,
timeout: int = 120,
) -> None:
_kubectl(
"wait",
"--for=condition=Ready",
f"pod/{name}",
"-n",
namespace,
f"--timeout={timeout}s",
)
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
while time.time() < deadline:
if proc.poll() is not None:
stderr = proc.stderr.read().decode() if proc.stderr else ""
raise RuntimeError(f"port-forward process exited early: {stderr}")
try:
with socket.create_connection(("127.0.0.1", port), timeout=1):
return
except OSError:
time.sleep(0.5)
raise TimeoutError(f"Port {port} not ready after {timeout}s")
def _port_forward_start(
namespace: str,
service: str,
local_port: int,
remote_port: int,
) -> subprocess.Popen:
"""Start kubectl port-forward and wait until the port is reachable."""
cmd = [
"kubectl",
"--context",
KUBECTL_CONTEXT,
"port-forward",
f"svc/{service}",
f"{local_port}:{remote_port}",
"-n",
namespace,
]
logger.info("Starting port-forward: %s", " ".join(cmd))
proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
_wait_for_port(local_port, proc)
return proc
def _cleanup_port_forward(name: str, pf: subprocess.Popen) -> None:
try:
pf.terminate()
pf.wait(timeout=10)
except subprocess.TimeoutExpired:
logger.warning(
"Port-forward %s did not exit on SIGTERM after 10s; killing", name
)
pf.kill()
try:
pf.wait(timeout=5)
except subprocess.TimeoutExpired:
logger.warning("Port-forward %s still running after SIGKILL", name)
except Exception as exc:
logger.warning("Error cleaning up %s port-forward: %s", name, exc)
rc = pf.returncode
stderr = pf.stderr.read().decode() if pf.stderr else ""
if rc != -15:
suffix = f": {stderr.strip()}" if stderr.strip() else ""
logger.warning("Port-forward %s exited rc=%s%s", name, rc, suffix)
else:
logger.debug("Port-forward %s exited cleanly (rc=%s)", name, rc)
def _poll_until(
predicate,
description: str,
timeout: int,
interval: float = 5,
) -> bool:
"""Poll predicate until True, or raise TimeoutError.
Only transient network errors are retried; HTTP status errors and
programming errors propagate immediately.
"""
deadline = time.time() + timeout
last_error = None
attempts = 0
while time.time() < deadline:
try:
attempts += 1
if predicate():
logger.info(
"Condition met: %s (after %d attempts)", description, attempts
)
return True
except _TRANSIENT_ERRORS as exc:
last_error = exc
logger.debug("Transient error on attempt %d: %s", attempts, exc)
time.sleep(interval)
msg = f"Timeout waiting for: {description} (after {timeout}s, {attempts} attempts)"
if last_error:
msg += f" — last error: {last_error}"
raise TimeoutError(msg)
def _get_router_url(router_base: str) -> str:
return router_base
def _router_is_healthy(router_base: str) -> bool:
try:
r = httpx.get(f"{router_base}/healthz", timeout=3.0)
return r.status_code == 200
except Exception:
return False
@pytest.fixture(scope="session")
def k8s_cluster():
"""Assert the kind cluster exists and kubectl context is reachable."""
result = subprocess.run(
["kind", "get", "clusters"],
capture_output=True,
text=True,
check=True,
)
if CLUSTER_NAME not in result.stdout.splitlines():
pytest.skip(
f"kind cluster '{CLUSTER_NAME}' not found — run "
f"./tests/e2e/k8s_integration/setup.sh first"
)
_kubectl("cluster-info")
return True
@pytest.fixture(scope="function")
def router_port_forward(k8s_cluster):
"""Per-test port-forward to sgl-router service.
Function-scoped because some tests (notably
test_lifecycle.TestRouterRestart) force-delete the router pod;
a session-scoped port-forward would be bound to the deleted pod's
network namespace and stay dead for all subsequent tests in the
suite. Per-test setup costs ~1-2s.
"""
_wait_for_deployment_ready("sgl-router")
pf = _port_forward_start(NAMESPACE, "sgl-router", 8090, 8090)
try:
_poll_until(
lambda: _router_is_healthy("http://127.0.0.1:8090"),
"sgl-router /healthz returns 200",
timeout=30,
interval=1,
)
yield "http://127.0.0.1:8090"
finally:
_cleanup_port_forward("sgl-router", pf)
@pytest.fixture(scope="function")
def router_url(router_port_forward):
return router_port_forward
@@ -0,0 +1,73 @@
"""Minimal fake SGLang worker for kind E2E integration testing.
Responds to:
GET /health -> {"status": "ok"}
GET /server_info -> {"served_model_name": MODEL_ID}
GET /v1/models -> list with a single MODEL_ID model entry
POST /v1/chat/completions -> echoes the last user message back
"""
from __future__ import annotations
import os
import uvicorn
from fastapi import FastAPI, Request
app = FastAPI()
MODEL_ID = os.environ.get("MODEL_ID", "tiny")
@app.get("/health")
async def health():
return {"status": "ok"}
@app.get("/server_info")
async def server_info():
# The sgl-router worker manager fetches this on every Added event and
# uses `served_model_name` to populate the registry's model index.
return {"served_model_name": MODEL_ID}
@app.get("/v1/models")
async def models():
return {
"object": "list",
"data": [
{
"id": MODEL_ID,
"object": "model",
"created": 0,
"owned_by": "sglang",
}
],
}
@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
payload = await request.json()
messages = payload.get("messages", [])
last_content = messages[-1]["content"] if messages else ""
return {
"id": "chatcmpl-mock",
"object": "chat.completion",
"model": payload.get("model", MODEL_ID),
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": f"echo: {last_content}",
},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
}
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=30000)
@@ -0,0 +1,4 @@
apiVersion: v1
kind: Namespace
metadata:
name: sgl-router-test
@@ -0,0 +1,33 @@
# Cluster-wide RBAC for the cross-namespace discovery test.
# Distinct ServiceAccount/ClusterRole names to avoid collision with
# the namespace-scoped Role in rbac.yaml used by the default router.
apiVersion: v1
kind: ServiceAccount
metadata:
name: sgl-router-cluster
namespace: sgl-router-test
---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRole
metadata:
name: sgl-router-cluster
rules:
- apiGroups: ["discovery.k8s.io"]
resources: ["endpointslices"]
verbs: ["get", "list", "watch"]
- apiGroups: [""]
resources: ["services", "pods"]
verbs: ["get", "list", "watch"]
---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRoleBinding
metadata:
name: sgl-router-cluster
roleRef:
apiGroup: rbac.authorization.k8s.io
kind: ClusterRole
name: sgl-router-cluster
subjects:
- kind: ServiceAccount
name: sgl-router-cluster
namespace: sgl-router-test
@@ -0,0 +1,34 @@
apiVersion: v1
kind: ServiceAccount
metadata:
name: sgl-router
namespace: sgl-router-test
---
apiVersion: rbac.authorization.k8s.io/v1
kind: Role
metadata:
name: sgl-router
namespace: sgl-router-test
rules:
# EndpointSlice watch (k8s discovery backend)
- apiGroups: ["discovery.k8s.io"]
resources: ["endpointslices"]
verbs: ["get", "list", "watch"]
# Service list/watch (needed to resolve EndpointSlice owner)
- apiGroups: [""]
resources: ["services", "pods"]
verbs: ["get", "list", "watch"]
---
apiVersion: rbac.authorization.k8s.io/v1
kind: RoleBinding
metadata:
name: sgl-router
namespace: sgl-router-test
roleRef:
apiGroup: rbac.authorization.k8s.io
kind: Role
name: sgl-router
subjects:
- kind: ServiceAccount
name: sgl-router
namespace: sgl-router-test
@@ -0,0 +1,60 @@
# sgl-router deployment with ClusterRole for cross-namespace discovery test.
# Watches workers in ALL namespaces via cluster-scoped EndpointSlice access.
apiVersion: apps/v1
kind: Deployment
metadata:
name: sgl-router-cluster
namespace: sgl-router-test
spec:
replicas: 1
selector:
matchLabels:
app: sgl-router-cluster
template:
metadata:
labels:
app: sgl-router-cluster
spec:
serviceAccountName: sgl-router-cluster
containers:
- name: router
image: sgl-router:e2e
imagePullPolicy: Never
args:
- "--config"
- "/etc/config/router-cluster.toml"
ports:
- containerPort: 8091
name: http
readinessProbe:
httpGet:
path: /readyz
port: 8091
initialDelaySeconds: 3
periodSeconds: 3
livenessProbe:
httpGet:
path: /healthz
port: 8091
initialDelaySeconds: 5
periodSeconds: 10
volumeMounts:
- name: config
mountPath: /etc/config
volumes:
- name: config
configMap:
name: sgl-router-cluster-config
---
apiVersion: v1
kind: Service
metadata:
name: sgl-router-cluster
namespace: sgl-router-test
spec:
selector:
app: sgl-router-cluster
ports:
- name: http
port: 8091
targetPort: 8091
@@ -0,0 +1,58 @@
apiVersion: apps/v1
kind: Deployment
metadata:
name: sgl-router
namespace: sgl-router-test
spec:
replicas: 1
selector:
matchLabels:
app: sgl-router
template:
metadata:
labels:
app: sgl-router
spec:
serviceAccountName: sgl-router
containers:
- name: router
image: sgl-router:e2e
imagePullPolicy: Never
args:
- "--config"
- "/etc/config/router.toml"
ports:
- containerPort: 8090
name: http
readinessProbe:
httpGet:
path: /readyz
port: 8090
initialDelaySeconds: 3
periodSeconds: 3
livenessProbe:
httpGet:
path: /healthz
port: 8090
initialDelaySeconds: 5
periodSeconds: 10
volumeMounts:
- name: config
mountPath: /etc/config
volumes:
- name: config
configMap:
name: sgl-router-config
---
apiVersion: v1
kind: Service
metadata:
name: sgl-router
namespace: sgl-router-test
spec:
selector:
app: sgl-router
ports:
- name: http
port: 8090
targetPort: 8090
@@ -0,0 +1,2 @@
httpx==0.27.2
pytest==8.3.3
+191
View File
@@ -0,0 +1,191 @@
#!/usr/bin/env bash
# Bootstrap a kind cluster for sgl-router K8s integration E2E tests.
#
# Prerequisites: Docker, kind, kubectl
#
# Usage:
# ./tests/e2e/k8s_integration/setup.sh # full setup
# ./tests/e2e/k8s_integration/setup.sh teardown # delete the cluster
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../../.." && pwd)" # repo root (above experimental/)
SGL_ROUTER_DIR="${REPO_ROOT}/experimental/sgl-router"
CLUSTER_NAME="${CLUSTER:-sgl-router-kind}"
NAMESPACE="${NAMESPACE:-sgl-router-test}"
CONTEXT="kind-${CLUSTER_NAME}"
MANIFESTS_DIR="${SCRIPT_DIR}/manifests"
log() { echo "==> $*"; }
teardown() {
log "Tearing down cluster '${CLUSTER_NAME}'..."
if kind get clusters 2>/dev/null | grep -q "^${CLUSTER_NAME}$"; then
kind delete cluster --name "${CLUSTER_NAME}"
else
log "Cluster '${CLUSTER_NAME}' not found, nothing to tear down."
fi
log "Done."
}
if [[ "${1:-}" == "teardown" ]]; then
teardown
exit 0
fi
# ---------------------------------------------------------------------------
# Step 1: Create kind cluster (idempotent)
# ---------------------------------------------------------------------------
if kind get clusters 2>/dev/null | grep -q "^${CLUSTER_NAME}$"; then
log "Kind cluster '${CLUSTER_NAME}' already exists — reusing."
else
log "Creating kind cluster '${CLUSTER_NAME}'..."
kind create cluster --name "${CLUSTER_NAME}" --wait 60s
fi
kubectl config use-context "${CONTEXT}"
# ---------------------------------------------------------------------------
# Step 2: Build Docker images (unless SKIP_DOCKER_BUILD=1)
# ---------------------------------------------------------------------------
if [[ "${SKIP_DOCKER_BUILD:-}" == "1" ]]; then
log "SKIP_DOCKER_BUILD=1 — skipping docker build; expecting images to exist locally."
for img in sgl-router:e2e sgl-router-fake-worker:e2e; do
if ! docker image inspect "${img}" >/dev/null 2>&1; then
log "ERROR: ${img} not found locally; cannot continue without building."
exit 1
fi
done
else
log "Building sgl-router:e2e from ${REPO_ROOT} ..."
docker build \
-f "${SCRIPT_DIR}/Dockerfile.router" \
-t sgl-router:e2e \
"${REPO_ROOT}"
log "Building sgl-router-fake-worker:e2e ..."
docker build \
-f "${SCRIPT_DIR}/Dockerfile.fake_worker" \
-t sgl-router-fake-worker:e2e \
"${SCRIPT_DIR}"
fi
# ---------------------------------------------------------------------------
# Step 3: Load images into kind
# ---------------------------------------------------------------------------
log "Loading images into kind cluster '${CLUSTER_NAME}'..."
kind load docker-image sgl-router:e2e --name "${CLUSTER_NAME}"
kind load docker-image sgl-router-fake-worker:e2e --name "${CLUSTER_NAME}"
# ---------------------------------------------------------------------------
# Step 4: Apply namespace and RBAC
# ---------------------------------------------------------------------------
log "Applying namespace and RBAC..."
kubectl --context "${CONTEXT}" apply -f "${MANIFESTS_DIR}/namespace.yaml"
kubectl --context "${CONTEXT}" apply -f "${MANIFESTS_DIR}/rbac.yaml"
# ---------------------------------------------------------------------------
# Step 5: Deploy 3 fake-worker replicas behind a Service
# The Service causes K8s to auto-create an EndpointSlice, which
# the sgl-router K8s discovery backend watches.
# ---------------------------------------------------------------------------
log "Deploying fake-worker Deployment + Service (3 replicas, app=sglang)..."
kubectl --context "${CONTEXT}" -n "${NAMESPACE}" apply -f - <<EOF
apiVersion: apps/v1
kind: Deployment
metadata:
name: fake-worker
namespace: ${NAMESPACE}
labels:
app: sglang
spec:
replicas: 3
selector:
matchLabels:
app: sglang
template:
metadata:
labels:
app: sglang
spec:
containers:
- name: worker
image: sgl-router-fake-worker:e2e
imagePullPolicy: Never
ports:
- containerPort: 30000
readinessProbe:
httpGet:
path: /health
port: 30000
initialDelaySeconds: 2
periodSeconds: 3
---
apiVersion: v1
kind: Service
metadata:
name: fake-worker
namespace: ${NAMESPACE}
labels:
app: sglang
spec:
selector:
app: sglang
ports:
- port: 30000
targetPort: 30000
EOF
log "Waiting for fake-worker rollout..."
kubectl --context "${CONTEXT}" -n "${NAMESPACE}" rollout status deployment/fake-worker --timeout=120s
# ---------------------------------------------------------------------------
# Step 6: Create sgl-router ConfigMap with k8s discovery pointing at the
# namespace where fake-worker pods live.
# ---------------------------------------------------------------------------
log "Creating sgl-router-config ConfigMap..."
ROUTER_CONFIG="[server]
host = \"0.0.0.0\"
port = 8090
[[models]]
id = \"tiny\"
tokenizer_path = \"/etc/tokenizer/tiny.json\"
policy = \"round_robin\"
# Aggressive breaker so a terminating pod's connection-refused
# immediately excludes it from the next request's candidate set —
# the reconciliation tests scale workers rapidly and depend on
# fast worker eviction to absorb the churn.
circuit_breaker = { threshold = 1, cool_down_secs = 5 }
[discovery]
backend = \"k8s\"
[discovery.k8s]
namespace = \"${NAMESPACE}\"
label_selector = \"app=sglang\""
kubectl --context "${CONTEXT}" -n "${NAMESPACE}" create configmap sgl-router-config \
--from-literal=router.toml="${ROUTER_CONFIG}" \
--dry-run=client -o yaml \
| kubectl --context "${CONTEXT}" apply -f -
# ---------------------------------------------------------------------------
# Step 7: Deploy sgl-router
# ---------------------------------------------------------------------------
log "Deploying sgl-router..."
kubectl --context "${CONTEXT}" apply -f "${MANIFESTS_DIR}/router.yaml"
log "Waiting for sgl-router rollout..."
kubectl --context "${CONTEXT}" -n "${NAMESPACE}" rollout status deployment/sgl-router --timeout=300s
# ---------------------------------------------------------------------------
# Done
# ---------------------------------------------------------------------------
log ""
log "Setup complete! Run the integration tests with:"
log " pytest tests/e2e/k8s_integration/ -v -s"
log ""
log "To tear down:"
log " ./tests/e2e/k8s_integration/setup.sh teardown"
@@ -0,0 +1,265 @@
"""Cross-namespace service discovery integration test.
Validates that a sgl-router instance with cluster-wide RBAC and no namespace
filter in its k8s discovery config watches EndpointSlices in all namespaces.
Workers deployed in a second namespace (sgl-router-test-extra) must be
discovered alongside those in the primary namespace.
This test deploys a separate router Deployment (sgl-router-cluster) with a
ClusterRole that grants EndpointSlice access across all namespaces.
Run with:
pytest tests/e2e/k8s_integration/test_cross_namespace.py -v -s
"""
from __future__ import annotations
import json
import logging
import subprocess
from pathlib import Path
import httpx
import pytest
from conftest import (
KUBECTL_CONTEXT,
NAMESPACE,
_apply_from_stdin,
_cleanup_port_forward,
_kubectl,
_poll_until,
_port_forward_start,
_wait_for_deployment_ready,
logger,
)
MANIFESTS_DIR = Path(__file__).parent / "manifests"
EXTRA_NAMESPACE = "sgl-router-test-extra"
CLUSTER_ROUTER_PORT = 8093
def _deploy_fake_worker_in_ns(name: str, namespace: str) -> None:
"""Deploy a fake-worker pod with imagePullPolicy=Never in the given namespace."""
pod_manifest = {
"apiVersion": "v1",
"kind": "Pod",
"metadata": {
"name": name,
"namespace": namespace,
"labels": {"app": "sglang", "cross-ns-test": "true"},
},
"spec": {
"containers": [
{
"name": "worker",
"image": "sgl-router-fake-worker:e2e",
"imagePullPolicy": "Never",
"ports": [{"containerPort": 30000}],
"readinessProbe": {
"httpGet": {"path": "/health", "port": 30000},
"initialDelaySeconds": 2,
"periodSeconds": 3,
},
}
]
},
}
proc = subprocess.run(
["kubectl", "--context", KUBECTL_CONTEXT, "apply", "-f", "-"],
input=json.dumps(pod_manifest),
capture_output=True,
text=True,
check=False,
)
if proc.returncode != 0:
raise RuntimeError(
f"Failed to deploy pod {name} in namespace {namespace} "
f"(rc={proc.returncode}): {proc.stderr.strip()!r}"
)
logger.info("Deployed worker %s in namespace %s", name, namespace)
def _safe_delete_pod(name: str, namespace: str) -> None:
try:
_kubectl(
"delete",
"pod",
name,
"-n",
namespace,
"--ignore-not-found",
"--force",
"--grace-period=0",
)
except Exception as exc:
logger.warning("Cleanup failed for pod %s in ns %s: %s", name, namespace, exc)
def _ensure_namespace(name: str) -> None:
manifest = {"apiVersion": "v1", "kind": "Namespace", "metadata": {"name": name}}
_apply_from_stdin(json.dumps(manifest))
def _ensure_service_in_ns(namespace: str, selector: str = "app=sglang") -> None:
"""Create a Service so K8s auto-creates an EndpointSlice for cross-ns workers.
Service `metadata.labels` propagates to the auto-created EndpointSlice's
labels — and the cluster-scoped router filters slices server-side by
`app=sglang,cross-ns-test=true`. Without those labels on the Service,
its EndpointSlice gets filtered out and the cross-ns worker is invisible.
"""
svc_manifest = {
"apiVersion": "v1",
"kind": "Service",
"metadata": {
"name": "fake-worker",
"namespace": namespace,
"labels": {"app": "sglang", "cross-ns-test": "true"},
},
"spec": {
"selector": {"app": "sglang", "cross-ns-test": "true"},
"ports": [{"port": 30000, "targetPort": 30000}],
},
}
_apply_from_stdin(json.dumps(svc_manifest))
def _can_route(router_url: str) -> bool:
try:
r = httpx.post(
f"{router_url}/v1/chat/completions",
json={
"model": "tiny",
"messages": [{"role": "user", "content": "cross-ns"}],
},
timeout=8.0,
)
return r.status_code == 200
except Exception:
return False
@pytest.fixture(scope="module")
def cluster_scoped_router(k8s_cluster):
"""Deploy the cluster-scoped RBAC + router, plus a second namespace."""
rbac_manifest = MANIFESTS_DIR / "rbac-cluster-scoped.yaml"
router_manifest = MANIFESTS_DIR / "router-cluster-scoped.yaml"
_kubectl("apply", "-f", str(rbac_manifest))
_ensure_namespace(EXTRA_NAMESPACE)
_ensure_service_in_ns(EXTRA_NAMESPACE)
# ConfigMap for the cluster-scoped router: empty namespace = watch all
cluster_config = """[server]
host = "0.0.0.0"
port = 8091
[[models]]
id = "tiny"
tokenizer_path = "/etc/tokenizer/tiny.json"
policy = "round_robin"
[discovery]
backend = "k8s"
[discovery.k8s]
namespace = ""
label_selector = "app=sglang,cross-ns-test=true"
"""
_kubectl(
"create",
"configmap",
"sgl-router-cluster-config",
f"--from-literal=router-cluster.toml={cluster_config}",
"-n",
NAMESPACE,
"--dry-run=client",
"-o",
"yaml",
check=True,
)
# pipe through apply
proc = _kubectl(
"create",
"configmap",
"sgl-router-cluster-config",
f"--from-literal=router-cluster.toml={cluster_config}",
"-n",
NAMESPACE,
"--dry-run=client",
"-o",
"yaml",
)
_apply_from_stdin(proc.stdout)
_kubectl("apply", "-f", str(router_manifest))
# The cluster-scoped router's /readyz blocks on registry-not-empty, so
# without at least one matching worker the rollout-status check below
# would hang for 180s. Deploy a "bootstrap" worker in EXTRA_NAMESPACE
# with the label_selector match (app=sglang,cross-ns-test=true) so the
# router's k8s discovery picks it up before the readiness probe runs.
# The test body adds a SECOND worker later to verify dynamic discovery.
bootstrap_worker = "cross-ns-worker-bootstrap"
_deploy_fake_worker_in_ns(bootstrap_worker, EXTRA_NAMESPACE)
pf = None
try:
_wait_for_deployment_ready("sgl-router-cluster")
pf = _port_forward_start(
NAMESPACE, "sgl-router-cluster", CLUSTER_ROUTER_PORT, 8091
)
yield f"http://127.0.0.1:{CLUSTER_ROUTER_PORT}"
finally:
if pf is not None:
_cleanup_port_forward("cluster_router", pf)
_safe_delete_pod(bootstrap_worker, EXTRA_NAMESPACE)
_kubectl(
"delete", "-f", str(router_manifest), "--ignore-not-found", check=False
)
_kubectl("delete", "-f", str(rbac_manifest), "--ignore-not-found", check=False)
_kubectl(
"delete",
"namespace",
EXTRA_NAMESPACE,
"--ignore-not-found",
"--wait=true",
"--timeout=60s",
check=False,
)
class TestClusterWideDiscovery:
"""Router with ClusterRole and no namespace filter sees workers in every namespace."""
def test_router_routes_to_worker_in_extra_namespace(self, cluster_scoped_router):
"""Deploy one fake-worker pod in the extra namespace behind a Service;
the cluster-scoped router must discover it (via its EndpointSlice) and
successfully route a chat completion to it."""
router_url = cluster_scoped_router
worker_name = "cross-ns-worker-extra"
try:
_deploy_fake_worker_in_ns(worker_name, EXTRA_NAMESPACE)
_poll_until(
lambda: _can_route(router_url),
"cluster-scoped router routes to worker in extra namespace",
timeout=60,
interval=3,
)
r = httpx.post(
f"{router_url}/v1/chat/completions",
json={
"model": "tiny",
"messages": [
{"role": "user", "content": "cross-namespace routing"}
],
},
timeout=15.0,
)
assert r.status_code == 200, f"expected 200, got {r.status_code}: {r.text}"
assert "echo:" in r.json()["choices"][0]["message"]["content"]
finally:
_safe_delete_pod(worker_name, EXTRA_NAMESPACE)
@@ -0,0 +1,84 @@
"""E2E: sgl-router K8s discovery — basic routing.
Verifies that sgl-router, configured with the k8s EndpointSlice backend,
discovers the 3 fake-worker replicas deployed by setup.sh and successfully
routes chat-completion requests to them.
"""
from __future__ import annotations
import httpx
import pytest
from conftest import NAMESPACE, _kubectl, _poll_until, logger
def _scale_fake_worker(replicas: int) -> None:
_kubectl(
"scale",
"deployment/fake-worker",
f"--replicas={replicas}",
"-n",
NAMESPACE,
)
def test_router_routes_chat_to_a_worker(router_url):
"""A /v1/chat/completions request through the router returns 200 with the
fake-worker echo payload, proving end-to-end routing works."""
r = httpx.post(
f"{router_url}/v1/chat/completions",
json={
"model": "tiny",
"messages": [{"role": "user", "content": "hello"}],
"stream": False,
},
timeout=15.0,
)
assert r.status_code == 200, f"expected 200, got {r.status_code}: {r.text}"
body = r.json()
assert "echo:" in body["choices"][0]["message"]["content"]
def test_router_lists_model(router_url):
"""GET /v1/models returns the 'tiny' model entry from the router config."""
r = httpx.get(f"{router_url}/v1/models", timeout=10.0)
assert r.status_code == 200, f"expected 200, got {r.status_code}: {r.text}"
body = r.json()
ids = [m["id"] for m in body["data"]]
assert "tiny" in ids, f"expected 'tiny' in model list, got {ids}"
def test_router_discovers_multiple_workers(router_url):
"""Scale down from 3 to 1 and back to 3 replicas; router must continue
routing successfully after each transition (EndpointSlice watch reflects
the change)."""
# First confirm baseline routing
r = httpx.post(
f"{router_url}/v1/chat/completions",
json={
"model": "tiny",
"messages": [{"role": "user", "content": "scale-test"}],
},
timeout=15.0,
)
assert r.status_code == 200
# Scale down to 1 — router should still route after reconverging
_scale_fake_worker(1)
_poll_until(
lambda: httpx.post(
f"{router_url}/v1/chat/completions",
json={
"model": "tiny",
"messages": [{"role": "user", "content": "post-scale-down"}],
},
timeout=10.0,
).status_code
== 200,
"router routes after scale-down to 1",
timeout=60,
interval=3,
)
# Restore to 3
_scale_fake_worker(3)
@@ -0,0 +1,150 @@
"""Worker lifecycle integration tests.
Covers:
1. Scaling replicas up — new EndpointSlice entries are discovered.
2. Scaling replicas down — removed endpoints are deregistered.
3. Router restart — after the router pod is killed, the Deployment restarts
it and it re-lists the existing EndpointSlice entries without duplicates.
These tests DO NOT use a /workers admin API (sgl-router does not expose
one). They verify behaviour through /v1/chat/completions responses and
by driving the deployment scale.
"""
from __future__ import annotations
import logging
import httpx
import pytest
from conftest import (
NAMESPACE,
_cleanup_port_forward,
_kubectl,
_poll_until,
_port_forward_start,
_wait_for_deployment_ready,
logger,
)
ROUTER_RESTART_PORT = 8092
def _scale(deployment: str, replicas: int) -> None:
_kubectl(
"scale", f"deployment/{deployment}", f"--replicas={replicas}", "-n", NAMESPACE
)
def _can_route(router_url: str) -> bool:
try:
r = httpx.post(
f"{router_url}/v1/chat/completions",
json={"model": "tiny", "messages": [{"role": "user", "content": "ping"}]},
timeout=8.0,
)
return r.status_code == 200
except Exception:
return False
class TestScaleUp:
"""Scaling fake-worker replicas up must not break routing."""
def test_router_routes_after_scale_up(self, router_url):
"""Restore 3 replicas (in case a prior test left 1), verify routing."""
_scale("fake-worker", 3)
_poll_until(
lambda: _can_route(router_url),
"router routes after scale-up to 3",
timeout=60,
interval=3,
)
class TestScaleDown:
"""Scaling to 0 then back up must restore routing."""
def test_router_recovers_after_scale_to_zero_and_back(self, router_url):
try:
_scale("fake-worker", 0)
# After scale-to-0 the router may return 503 (no healthy workers)
# That is expected behaviour — assert it transitions back on scale-up.
_scale("fake-worker", 2)
_poll_until(
lambda: _can_route(router_url),
"router routes again after scale-up from 0",
timeout=90,
interval=3,
)
finally:
_scale("fake-worker", 3)
class TestRouterRestart:
"""Killing the router pod forces a Deployment restart; the new pod must
re-discover workers via the EndpointSlice watch without duplicates."""
def test_router_rediscovers_workers_after_restart(self, k8s_cluster):
# Use a dedicated port to avoid clashing with the session fixture
pf_holder: list = [None]
try:
_wait_for_deployment_ready("sgl-router")
pf_holder[0] = _port_forward_start(
NAMESPACE, "sgl-router", ROUTER_RESTART_PORT, 8090
)
restart_url = f"http://127.0.0.1:{ROUTER_RESTART_PORT}"
# Baseline: routing works pre-restart
_poll_until(
lambda: _can_route(restart_url),
"baseline routing works pre-restart",
timeout=30,
interval=2,
)
# Kill the router pod — the Deployment ReplicaSet will restart it
res = _kubectl(
"get",
"pod",
"-n",
NAMESPACE,
"-l",
"app=sgl-router",
"-o",
"jsonpath={.items[0].metadata.name}",
check=False,
)
old_pod = res.stdout.strip()
if old_pod:
_kubectl(
"delete",
"pod",
old_pod,
"-n",
NAMESPACE,
"--force",
"--grace-period=0",
)
# Tear down the old port-forward before waiting for the new pod
if pf_holder[0] is not None:
_cleanup_port_forward("router-restart-pre-kill", pf_holder[0])
pf_holder[0] = None
_wait_for_deployment_ready("sgl-router")
pf_holder[0] = _port_forward_start(
NAMESPACE, "sgl-router", ROUTER_RESTART_PORT, 8090
)
# After restart, routing must come back (EndpointSlice re-watch)
_poll_until(
lambda: _can_route(restart_url),
"routing restored after router restart",
timeout=60,
interval=3,
)
finally:
if pf_holder[0] is not None:
_cleanup_port_forward("router-restart", pf_holder[0])
@@ -0,0 +1,163 @@
"""K8s discovery reconciliation integration tests.
Tests verify that:
1. The K8s EndpointSlice watcher correctly discovers new workers as Services
and backing Deployments are updated.
2. Workers are removed from the router's registry after the backing EndpointSlice
entries disappear (pod deleted / deployment scaled to 0).
3. After a simulated watch-connection interruption (router restarted), the
registry converges back to the correct worker set.
Note: sgl-router does not currently expose a Prometheus /metrics endpoint,
so the SMG-style metric assertions are not used here. Disconnect/reconnect
coverage is provided by test_lifecycle.TestRouterRestart.
"""
from __future__ import annotations
import logging
import time
import httpx
import pytest
from conftest import (
NAMESPACE,
RECONCILIATION_WAIT_SECS,
_kubectl,
_poll_until,
logger,
)
def _scale_fake_worker(replicas: int) -> None:
_kubectl(
"scale", "deployment/fake-worker", f"--replicas={replicas}", "-n", NAMESPACE
)
def _can_route(router_url: str) -> bool:
try:
r = httpx.post(
f"{router_url}/v1/chat/completions",
json={
"model": "tiny",
"messages": [{"role": "user", "content": "reconcile"}],
},
timeout=8.0,
)
return r.status_code == 200
except Exception:
return False
class TestWatcherDiscovery:
"""The EndpointSlice watcher discovers new endpoints on Deployment scale-up."""
def test_watcher_discovers_new_endpoints_on_scale_up(self, router_url):
"""Scale from 1 to 3 replicas; router must continue routing successfully."""
_scale_fake_worker(1)
# Wait for scale-down to propagate and routing to stabilise
_poll_until(
lambda: _can_route(router_url),
"router routes with 1 replica",
timeout=60,
interval=3,
)
_scale_fake_worker(3)
_poll_until(
lambda: _can_route(router_url),
"router routes with 3 replicas (after scale-up)",
timeout=60,
interval=3,
)
class TestStaleEndpointRemoval:
"""When fake-worker replicas drop, the router must stop routing to the
removed endpoints.
Because sgl-router has no /workers admin API, we verify removal
indirectly: scale to 0, assert the router returns non-200 (or at least
that scaling back to 2 restores routing), then restore.
"""
def test_routing_restores_after_scale_down_and_back_up(self, router_url):
"""Scale to 0 (no workers → expect non-200), then restore to 2.
After restore the router must route again within the reconciliation window.
"""
try:
_scale_fake_worker(0)
# Expect routing to fail eventually (503 or connection error)
deadline = time.time() + RECONCILIATION_WAIT_SECS
routing_failed = False
while time.time() < deadline:
try:
r = httpx.post(
f"{router_url}/v1/chat/completions",
json={
"model": "tiny",
"messages": [{"role": "user", "content": "no-workers"}],
},
timeout=5.0,
)
if r.status_code != 200:
routing_failed = True
break
except Exception:
routing_failed = True
break
time.sleep(3)
# If after RECONCILIATION_WAIT_SECS the router is still routing,
# that means old endpoints are cached — not necessarily wrong for
# a watcher that hasn't ticked yet, but log a warning.
if not routing_failed:
logger.warning(
"Router still returning 200 after scale-to-0; "
"EndpointSlice event may be delayed — continuing test."
)
# Restore workers and verify routing comes back
_scale_fake_worker(2)
_poll_until(
lambda: _can_route(router_url),
"routing restored after scale back up to 2",
timeout=RECONCILIATION_WAIT_SECS,
interval=3,
)
finally:
_scale_fake_worker(3)
class TestReconciliationConsistency:
"""Routing remains stable over multiple reconciliation windows with steady
worker state — no spurious deregistrations or duplicate registrations."""
@pytest.mark.slow
def test_routing_stable_over_multiple_reconciliation_cycles(self, router_url):
"""Deploy 3 workers, sample routing success over ~150s (2 reconciliation
cycles + margin), assert no interruptions."""
_scale_fake_worker(3)
_poll_until(
lambda: _can_route(router_url),
"baseline routing with 3 workers",
timeout=30,
interval=2,
)
# Sample every 15s for 150s
wait_secs = RECONCILIATION_WAIT_SECS + 60
end_time = time.time() + wait_secs
failures = []
while time.time() < end_time:
ok = _can_route(router_url)
if not ok:
failures.append(time.time())
time.sleep(15)
assert not failures, (
f"Routing failed at {len(failures)} sample(s) during stability window; "
f"timestamps: {failures}"
)
@@ -0,0 +1,26 @@
# Pytest configuration for sgl-router tests/e2e/.
# Lives next to conftest.py so `pytest experimental/sgl-router/tests/e2e/`
# picks it up automatically.
[tool.pytest.ini_options]
minversion = "8.0"
# Default discovery: smoke tests (top-level test_*.py) and the
# multi-worker chat_completions suite. k8s_integration is intentionally
# not in the default set — it requires a kind/k8s cluster and is
# invoked explicitly.
testpaths = [
".",
"chat_completions",
]
python_files = ["test_*.py"]
python_classes = ["Test*"]
python_functions = ["test_*"]
markers = [
"real_gpu: requires at least one NVIDIA GPU (skipped on CPU-only hosts)",
"pd_mode: requires the router started in PD-disaggregation mode",
"slow: takes >30s (model load, multi-request convergence checks)",
]
log_cli = true
log_cli_level = "INFO"
log_cli_format = "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
log_cli_date_format = "%H:%M:%S"
@@ -0,0 +1,15 @@
httpx==0.27.2
pytest==8.3.3
pytest-asyncio==0.24.0
# huggingface_hub is intentionally NOT pinned here. SGLang's
# `scripts/ci/cuda/ci_install_dependency.sh` already installs a
# version compatible with the rest of its transitive deps
# (transformers / diffusers / kernels, which require
# huggingface_hub >= 1.5 / >= 0.34 / >= 1.3 respectively). An earlier
# pin of `huggingface_hub==0.26.2` here got installed AFTER the SGLang
# deps and downgraded huggingface_hub past `is_offline_mode`'s top-
# level export, which broke `from sglang.srt.server_args import …`
# at module import time and turned every smoke test into a 5-minute
# `/health` timeout with no actionable signal.
# The e2e suite only uses huggingface_hub's `try_to_load_from_cache`,
# which is available in every release SGLang would install.
@@ -0,0 +1,64 @@
"""
Smoke tests for /v1/models and /v1/chat/completions (streaming + non-streaming).
"""
from __future__ import annotations
import httpx
import pytest
MODEL = "Qwen/Qwen3-0.6B"
def test_models(router: str) -> None:
"""GET /v1/models must list the configured model."""
resp = httpx.get(f"{router}/v1/models", timeout=30)
assert resp.status_code == 200, resp.text
data = resp.json()
ids = [m["id"] for m in data.get("data", [])]
assert any(
MODEL in mid for mid in ids
), f"Model {MODEL!r} not found in /v1/models response: {ids}"
def test_chat_non_streaming(router: str) -> None:
"""POST /v1/chat/completions (stream=False) returns an assistant message."""
payload = {
"model": MODEL,
"messages": [{"role": "user", "content": "Say hi."}],
"max_tokens": 10,
"stream": False,
}
resp = httpx.post(f"{router}/v1/chat/completions", json=payload, timeout=60)
assert resp.status_code == 200, resp.text
body = resp.json()
choice = body["choices"][0]
assert choice["message"]["role"] == "assistant"
assert choice["message"]["content"], "Expected non-empty assistant content"
def test_chat_streaming(router: str) -> None:
"""POST /v1/chat/completions (stream=True) returns >=2 SSE chunks incl. [DONE]."""
payload = {
"model": MODEL,
"messages": [{"role": "user", "content": "Say hi."}],
"max_tokens": 10,
"stream": True,
}
chunks: list[str] = []
with httpx.stream(
"POST",
f"{router}/v1/chat/completions",
json=payload,
timeout=60,
) as resp:
assert resp.status_code == 200, resp.read().decode()
for line in resp.iter_lines():
line = line.strip()
if line.startswith("data:"):
chunks.append(line)
assert len(chunks) >= 2, f"Expected >=2 SSE chunks, got {len(chunks)}: {chunks}"
assert any(
"[DONE]" in c for c in chunks
), f"No [DONE] chunk found in SSE stream: {chunks}"
@@ -0,0 +1,37 @@
"""
Smoke test for /v1/tokenize and /v1/detokenize round-trip.
"""
from __future__ import annotations
import httpx
MODEL = "Qwen/Qwen3-0.6B"
TEXT = "Hello, world!"
def test_tokenize_round_trip(router: str) -> None:
"""POST /v1/tokenize then /v1/detokenize must recover the original text."""
# Tokenize
tok_resp = httpx.post(
f"{router}/v1/tokenize",
json={"model": MODEL, "prompt": TEXT},
timeout=30,
)
assert tok_resp.status_code == 200, tok_resp.text
tokens = tok_resp.json()["tokens"]
assert (
isinstance(tokens, list) and len(tokens) > 0
), f"Expected non-empty token list, got: {tokens}"
# Detokenize
detok_resp = httpx.post(
f"{router}/v1/detokenize",
json={"model": MODEL, "tokens": tokens},
timeout=30,
)
assert detok_resp.status_code == 200, detok_resp.text
recovered = detok_resp.json()["text"]
assert (
TEXT in recovered or recovered in TEXT
), f"Round-trip mismatch: original={TEXT!r}, recovered={recovered!r}"
@@ -0,0 +1,232 @@
[
{
"name": "single_full_block",
"tokens": [
1,
2,
3,
4
],
"block_size": 4,
"expected_i64_hashes": [
-3488128144981237669
]
},
{
"name": "partial_last_block",
"tokens": [
1,
2,
3,
4,
5
],
"block_size": 4,
"expected_i64_hashes": [
-3488128144981237669,
-3787494577174227566
]
},
{
"name": "multi_block",
"tokens": [
10,
20,
30,
40,
50,
60,
70,
80
],
"block_size": 2,
"expected_i64_hashes": [
978178666101069530,
-895308556211281782,
-8033692805846017938,
835415944263129316
]
},
{
"name": "empty_tokens",
"tokens": [],
"block_size": 4,
"expected_i64_hashes": []
},
{
"name": "block_size_one",
"tokens": [
7,
8,
9
],
"block_size": 1,
"expected_i64_hashes": [
-1702009526849766914,
903318264012984157,
-8265893088400908305
]
},
{
"name": "odd_boundary",
"tokens": [
100,
200,
300,
400,
500,
600,
700
],
"block_size": 3,
"expected_i64_hashes": [
-7293070039731858224,
-5869816562584529365,
7513319606423624955
]
},
{
"name": "long_sequence",
"tokens": [
1,
2,
3,
4,
5,
6,
7,
8,
9,
10,
11,
12,
13,
14,
15,
16,
17,
18,
19,
20,
21,
22,
23,
24,
25,
26,
27,
28,
29,
30,
31,
32,
33,
34,
35,
36,
37,
38,
39,
40,
41,
42,
43,
44,
45,
46,
47,
48,
49,
50,
51,
52,
53,
54,
55,
56,
57,
58,
59,
60,
61,
62,
63,
64,
65,
66,
67,
68,
69,
70,
71,
72,
73,
74,
75,
76,
77,
78,
79,
80,
81,
82,
83,
84,
85,
86,
87,
88,
89,
90,
91,
92,
93,
94,
95,
96,
97,
98,
99,
100,
101,
102,
103,
104,
105,
106,
107,
108,
109,
110,
111,
112,
113,
114,
115,
116,
117,
118,
119,
120,
121,
122,
123,
124,
125,
126,
127,
128
],
"block_size": 16,
"expected_i64_hashes": [
8635429971592222890,
1256577331724852459,
5689809685380680247,
3927976462491479733,
5639345789331840936,
-4601255381563393033,
3368460852864325515,
1233425155141659070
]
}
]
File diff suppressed because one or more lines are too long
@@ -0,0 +1,849 @@
{
"model_id": "deepseek-ai/DeepSeek-V3",
"shape": "long",
"prompt_text": "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. ",
"expected_token_ids": [
83240,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
101339,
55848,
39208,
10434,
57037,
14,
67956,
109387,
51320,
14,
10012,
696,
312,
4667,
5158,
14408,
121876,
329,
3992,
3404,
7314,
492,
1231,
95691,
127631,
86798,
67,
16,
223
],
"skip_special_tokens": false
}
@@ -0,0 +1,83 @@
{
"model_id": "deepseek-ai/DeepSeek-V3",
"shape": "multi_turn_with_tools",
"prompt_text": "<|im_start|>system\nYou have tools.<|im_end|>\n<|im_start|>user\nWeather in Paris?<|im_end|>\n<|im_start|>assistant\n<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Paris\"}}\n</tool_call><|im_end|>\n",
"expected_token_ids": [
30,
94,
328,
37864,
94,
32,
27824,
201,
3476,
611,
6704,
32334,
94,
328,
42616,
94,
1018,
30,
94,
328,
37864,
94,
32,
5265,
201,
58565,
295,
11111,
33,
30,
94,
328,
42616,
94,
1018,
30,
94,
328,
37864,
94,
32,
624,
15059,
201,
30,
72461,
112042,
1018,
24313,
2852,
3362,
582,
1133,
65,
50219,
1760,
582,
83772,
3362,
28612,
37399,
3362,
582,
51119,
4,
30316,
1718,
72461,
112042,
5451,
94,
328,
42616,
94,
1018
],
"skip_special_tokens": false
}
@@ -0,0 +1,12 @@
{
"model_id": "deepseek-ai/DeepSeek-V3",
"shape": "short",
"prompt_text": "Hello, world!",
"expected_token_ids": [
19923,
14,
2058,
3
],
"skip_special_tokens": false
}
@@ -0,0 +1,63 @@
{
"model_id": "deepseek-ai/DeepSeek-V3",
"shape": "special_token_heavy",
"prompt_text": "<|im_start|>system\nYou are helpful.<|im_end|>\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\nHello<|im_end|>\n<|endoftext|>",
"expected_token_ids": [
30,
94,
328,
37864,
94,
32,
27824,
201,
3476,
477,
11502,
32334,
94,
328,
42616,
94,
1018,
30,
94,
328,
37864,
94,
32,
5265,
201,
23166,
30,
94,
328,
42616,
94,
1018,
30,
94,
328,
37864,
94,
32,
624,
15059,
201,
19923,
30,
94,
328,
42616,
94,
1018,
30,
94,
523,
2154,
2067,
94,
32
],
"skip_special_tokens": false
}
@@ -0,0 +1,669 @@
{
"model_id": "openai/gpt-oss-20b",
"shape": "long",
"prompt_text": "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. ",
"expected_token_ids": [
61495,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
86529,
38714,
25840,
2353,
36204,
11,
54472,
91785,
45688,
11,
10412,
621,
160226,
14725,
173578,
4518,
110546,
859,
79682,
78404,
151394,
13,
220
],
"skip_special_tokens": false
}
@@ -0,0 +1,80 @@
{
"model_id": "openai/gpt-oss-20b",
"shape": "multi_turn_with_tools",
"prompt_text": "<|im_start|>system\nYou have tools.<|im_end|>\n<|im_start|>user\nWeather in Paris?<|im_end|>\n<|im_start|>assistant\n<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Paris\"}}\n</tool_call><|im_end|>\n",
"expected_token_ids": [
27,
91,
321,
10949,
91,
29,
17360,
198,
3575,
679,
8437,
30502,
91,
321,
13707,
91,
523,
27,
91,
321,
10949,
91,
29,
1428,
198,
29602,
306,
12650,
190440,
91,
321,
13707,
91,
523,
27,
91,
321,
10949,
91,
29,
173781,
198,
27,
17952,
25158,
523,
10848,
897,
1243,
392,
522,
170154,
672,
392,
34317,
1243,
10494,
17500,
1243,
392,
72782,
18583,
739,
808,
17952,
25158,
3784,
91,
321,
13707,
91,
523
],
"skip_special_tokens": false
}
@@ -0,0 +1,12 @@
{
"model_id": "openai/gpt-oss-20b",
"shape": "short",
"prompt_text": "Hello, world!",
"expected_token_ids": [
13225,
11,
2375,
0
],
"skip_special_tokens": false
}
@@ -0,0 +1,56 @@
{
"model_id": "openai/gpt-oss-20b",
"shape": "special_token_heavy",
"prompt_text": "<|im_start|>system\nYou are helpful.<|im_end|>\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\nHello<|im_end|>\n<|endoftext|>",
"expected_token_ids": [
27,
91,
321,
10949,
91,
29,
17360,
198,
3575,
553,
10297,
30502,
91,
321,
13707,
91,
523,
27,
91,
321,
10949,
91,
29,
1428,
198,
12194,
27,
91,
321,
13707,
91,
523,
27,
91,
321,
10949,
91,
29,
173781,
198,
13225,
27,
91,
321,
13707,
91,
523,
199999
],
"skip_special_tokens": false
}
@@ -0,0 +1,669 @@
{
"model_id": "Qwen/Qwen3-30B-A3B",
"shape": "long",
"prompt_text": "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. ",
"expected_token_ids": [
32783,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
46931,
26342,
23655,
2444,
27212,
11,
35140,
57924,
30060,
11,
10923,
653,
79122,
18965,
86404,
8621,
72204,
1842,
57296,
58917,
85927,
13,
220
],
"skip_special_tokens": false
}
@@ -0,0 +1,50 @@
{
"model_id": "Qwen/Qwen3-30B-A3B",
"shape": "multi_turn_with_tools",
"prompt_text": "<|im_start|>system\nYou have tools.<|im_end|>\n<|im_start|>user\nWeather in Paris?<|im_end|>\n<|im_start|>assistant\n<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Paris\"}}\n</tool_call><|im_end|>\n",
"expected_token_ids": [
151644,
8948,
198,
2610,
614,
7375,
13,
151645,
198,
151644,
872,
198,
28981,
304,
12095,
30,
151645,
198,
151644,
77091,
198,
151657,
198,
4913,
606,
788,
330,
455,
69364,
497,
330,
16370,
788,
5212,
8926,
788,
330,
59604,
95642,
151658,
151645,
198
],
"skip_special_tokens": false
}
@@ -0,0 +1,12 @@
{
"model_id": "Qwen/Qwen3-30B-A3B",
"shape": "short",
"prompt_text": "Hello, world!",
"expected_token_ids": [
9707,
11,
1879,
0
],
"skip_special_tokens": false
}
@@ -0,0 +1,30 @@
{
"model_id": "Qwen/Qwen3-30B-A3B",
"shape": "special_token_heavy",
"prompt_text": "<|im_start|>system\nYou are helpful.<|im_end|>\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\nHello<|im_end|>\n<|endoftext|>",
"expected_token_ids": [
151644,
8948,
198,
2610,
525,
10950,
13,
151645,
198,
151644,
872,
198,
13048,
151645,
198,
151644,
77091,
198,
9707,
151645,
198,
151643
],
"skip_special_tokens": false
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,437 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Minimal axum mock of an SGLang HTTP worker for routing tests.
use axum::body::Body;
use axum::extract::State;
use axum::http::{HeaderMap, HeaderName, HeaderValue, StatusCode};
use axum::response::{IntoResponse, Response};
use axum::routing::{get, post};
use axum::Json;
use bytes::Bytes;
use serde_json::Value;
use std::collections::{HashMap, HashSet};
use std::net::SocketAddr;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::oneshot;
/// Headers captured from the most recent inbound request.
#[derive(Default)]
pub struct CapturedHeaders {
pub seen: HashSet<String>, // names (kept for backwards compat)
pub headers: HashMap<String, String>, // name -> value (last write wins)
pub last_body: Option<Bytes>,
}
#[derive(Clone)]
#[allow(dead_code)] // Only used by some test files; mock_worker is shared.
pub struct MockWorkerState {
pub captured: Arc<Mutex<CapturedHeaders>>,
pub stream_chunks: Arc<Vec<&'static str>>,
}
/// A running mock SGLang worker. Shuts down on Drop via the oneshot sender.
pub struct MockWorker {
pub url: String,
// Used in header_forwarding_test; not every test file reads captured headers.
#[allow(dead_code)]
pub captured: Arc<Mutex<CapturedHeaders>>,
_shutdown: oneshot::Sender<()>,
}
impl MockWorker {
/// Bind to a random port on 127.0.0.1 and start serving.
///
/// `stream_chunks` are the raw SSE bytes returned when a streaming
/// chat-completion request arrives.
#[allow(dead_code)] // Only used by some test files.
pub async fn start(stream_chunks: Vec<&'static str>) -> Self {
let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
let state = MockWorkerState {
captured: captured.clone(),
stream_chunks: Arc::new(stream_chunks),
};
// /server_info advertises served_model_name="tiny" so the
// worker-manager introspect step resolves model_ids for the
// "tiny" model the tests register a tokenizer + policy under.
let app = axum::Router::new()
.route("/v1/chat/completions", post(chat))
.route("/server_info", get(serve_tiny_server_info))
.with_state(state);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr: SocketAddr = listener.local_addr().unwrap();
let url = format!("http://{addr}");
let (tx, rx) = oneshot::channel::<()>();
tokio::spawn(async move {
axum::serve(listener, app)
.with_graceful_shutdown(async {
let _ = rx.await;
})
.await
.unwrap();
});
Self {
url,
captured,
_shutdown: tx,
}
}
/// Bind to a random port and start a worker that accepts the request,
/// sleeps for `delay`, then returns `200 OK` with an empty JSON object.
/// Used to test router behaviour when the upstream wedges after accepting
/// the TCP connection but before sending response headers.
#[allow(dead_code)]
pub async fn start_hanging(delay: Duration) -> Self {
let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
#[derive(Clone)]
struct HangState {
captured: Arc<Mutex<CapturedHeaders>>,
delay: Duration,
}
async fn hang_handler(
State(s): State<HangState>,
headers: HeaderMap,
body: Bytes,
) -> Response<Body> {
{
let mut g = s.captured.lock().unwrap();
g.last_body = Some(body.clone());
for (k, v) in headers.iter() {
g.seen.insert(k.as_str().to_string());
if let Ok(val) = v.to_str() {
g.headers.insert(k.as_str().to_string(), val.to_string());
}
}
}
tokio::time::sleep(s.delay).await;
let mut r = Response::new(Body::from("{}"));
*r.status_mut() = StatusCode::OK;
r.headers_mut().insert(
HeaderName::from_static("content-type"),
HeaderValue::from_static("application/json"),
);
r
}
let state = HangState {
captured: captured.clone(),
delay,
};
let app = axum::Router::new()
.route("/v1/chat/completions", post(hang_handler))
.route("/server_info", get(serve_tiny_server_info))
.with_state(state);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr: SocketAddr = listener.local_addr().unwrap();
let url = format!("http://{addr}");
let (tx, rx) = oneshot::channel::<()>();
tokio::spawn(async move {
axum::serve(listener, app)
.with_graceful_shutdown(async {
let _ = rx.await;
})
.await
.unwrap();
});
Self {
url,
captured,
_shutdown: tx,
}
}
/// Bind to a random port and start a worker that streams `chunks` with a
/// fixed `delay` between each chunk. Used to test that load guards survive
/// the full body lifetime for streaming responses.
#[allow(dead_code)]
pub async fn start_slow_stream(chunks: Vec<&'static str>, delay: Duration) -> Self {
let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
#[derive(Clone)]
struct SlowState {
captured: Arc<Mutex<CapturedHeaders>>,
chunks: Arc<Vec<&'static str>>,
delay: Duration,
}
async fn slow_chat(
State(s): State<SlowState>,
headers: HeaderMap,
body: Bytes,
) -> Response<Body> {
{
let mut g = s.captured.lock().unwrap();
g.last_body = Some(body.clone());
for (k, v) in headers.iter() {
g.seen.insert(k.as_str().to_string());
if let Ok(val) = v.to_str() {
g.headers.insert(k.as_str().to_string(), val.to_string());
}
}
}
let chunks = s.chunks.clone();
let delay = s.delay;
// Stream chunks via a channel, sleeping between each send.
let (tx, rx) = tokio::sync::mpsc::channel::<Result<Bytes, std::io::Error>>(4);
tokio::spawn(async move {
for chunk in chunks.iter() {
tokio::time::sleep(delay).await;
if tx.send(Ok(Bytes::from(*chunk))).await.is_err() {
break;
}
}
});
let body = Body::from_stream(tokio_stream::wrappers::ReceiverStream::new(rx));
let mut r = Response::new(body);
*r.status_mut() = StatusCode::OK;
r.headers_mut().insert(
HeaderName::from_static("content-type"),
"text/event-stream".parse().unwrap(),
);
r
}
let state = SlowState {
captured: captured.clone(),
chunks: Arc::new(chunks),
delay,
};
let app = axum::Router::new()
.route("/v1/chat/completions", post(slow_chat))
.route("/server_info", get(serve_tiny_server_info))
.with_state(state);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr: SocketAddr = listener.local_addr().unwrap();
let url = format!("http://{addr}");
let (tx, rx) = oneshot::channel::<()>();
tokio::spawn(async move {
axum::serve(listener, app)
.with_graceful_shutdown(async {
let _ = rx.await;
})
.await
.unwrap();
});
Self {
url,
captured,
_shutdown: tx,
}
}
/// Bind to a raw TCP listener and start a worker that writes a status
/// line + headers with a large declared `Content-Length`, then writes
/// only `partial_body_bytes` of body before closing the connection.
///
/// Used to test router behaviour when the upstream replies with a status
/// but drops the connection mid-body. We can't build this with axum
/// directly (it owns the response lifecycle); raw TCP gives us frame-level
/// control to short-write the body and close.
///
/// NOTE: unlike the axum-based variants, this helper does NOT serve
/// `/server_info` (one-shot raw-TCP accept, no path routing). Callers
/// that wire this through `spawn_discovery` will see introspect fail
/// with empty `model_ids`. All current callers inject the worker via
/// `registry.add()` directly, which bypasses introspect.
#[allow(dead_code)]
pub async fn start_returning_partial_body(
status: StatusCode,
partial_body_bytes: &'static [u8],
) -> Self {
let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr: SocketAddr = listener.local_addr().unwrap();
let url = format!("http://{addr}");
let (tx, mut rx) = oneshot::channel::<()>();
tokio::spawn(async move {
// Accept one connection (or exit on shutdown).
tokio::select! {
_ = &mut rx => (),
accept = listener.accept() => {
let (mut sock, _) = match accept {
Ok(v) => v,
Err(_) => return,
};
// Drain the request bytes until we see end-of-headers
// (`\r\n\r\n`). We deliberately do NOT fully consume the
// request body — the router has already sent it before
// awaiting our response, and we want to write the
// truncated response promptly.
let mut buf = [0u8; 4096];
let mut acc: Vec<u8> = Vec::new();
while !acc.windows(4).any(|w| w == b"\r\n\r\n") {
let n = match sock.read(&mut buf).await {
Ok(0) | Err(_) => return,
Ok(n) => n,
};
acc.extend_from_slice(&buf[..n]);
if acc.len() > 64 * 1024 {
// Defensive: don't loop forever if the request
// never produces a header terminator.
break;
}
}
// Write a response with a Content-Length larger than the
// bytes we will actually write, then drop the socket
// before the body completes.
let declared_len = partial_body_bytes.len() + 1024;
let head = format!(
"HTTP/1.1 {status_u16} {phrase}\r\n\
content-type: application/json\r\n\
content-length: {declared_len}\r\n\
connection: close\r\n\
\r\n",
status_u16 = status.as_u16(),
phrase = status.canonical_reason().unwrap_or("OK"),
);
if sock.write_all(head.as_bytes()).await.is_err() {
return;
}
if sock.write_all(partial_body_bytes).await.is_err() {
return;
}
// Flush, then drop — the client should see content-length
// mismatch as a transport-level body read failure.
let _ = sock.flush().await;
drop(sock);
}
}
});
Self {
url,
captured,
_shutdown: tx,
}
}
/// Bind to a random port and start a worker that ALWAYS returns the given
/// HTTP status code and JSON body with `Content-Type: application/json`.
/// Used to test router behaviour when the upstream returns an error.
#[allow(dead_code)]
pub async fn start_returning_error(status: StatusCode, body: Value) -> Self {
let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
let body_arc = Arc::new(body.to_string());
#[derive(Clone)]
struct ErrorState {
captured: Arc<Mutex<CapturedHeaders>>,
body_str: Arc<String>,
status: StatusCode,
}
async fn error_handler(
State(s): State<ErrorState>,
headers: HeaderMap,
body: Bytes,
) -> Response<Body> {
{
let mut g = s.captured.lock().unwrap();
g.last_body = Some(body);
for (k, v) in headers.iter() {
g.seen.insert(k.as_str().to_string());
if let Ok(val) = v.to_str() {
g.headers.insert(k.as_str().to_string(), val.to_string());
}
}
}
let mut r = Response::new(Body::from(s.body_str.as_ref().clone()));
*r.status_mut() = s.status;
r.headers_mut().insert(
HeaderName::from_static("content-type"),
HeaderValue::from_static("application/json"),
);
r
}
let state = ErrorState {
captured: captured.clone(),
body_str: body_arc,
status,
};
let app = axum::Router::new()
.route("/v1/chat/completions", post(error_handler))
.route("/server_info", get(serve_tiny_server_info))
.with_state(state);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr: SocketAddr = listener.local_addr().unwrap();
let url = format!("http://{addr}");
let (tx, rx) = oneshot::channel::<()>();
tokio::spawn(async move {
axum::serve(listener, app)
.with_graceful_shutdown(async {
let _ = rx.await;
})
.await
.unwrap();
});
Self {
url,
captured,
_shutdown: tx,
}
}
}
/// Stateless `/server_info` handler shared by every axum-based
/// `MockWorker::start_*` variant. Advertising `served_model_name="tiny"`
/// lets the worker manager's introspect step resolve `model_ids` for any
/// variant that flows through `spawn_discovery`, instead of burning 3 ×
/// `SERVER_INFO_TIMEOUT` of retries before registering with empty
/// `model_ids`. Adding it unconditionally is cheaper than tracking which
/// variants do or don't get introspected.
#[allow(dead_code)] // shared across all axum variants
async fn serve_tiny_server_info() -> Json<Value> {
Json(serde_json::json!({"served_model_name": "tiny"}))
}
#[allow(dead_code)] // Used by `MockWorker::start`, only some test files need it.
async fn chat(State(s): State<MockWorkerState>, headers: HeaderMap, body: Bytes) -> Response<Body> {
{
let mut g = s.captured.lock().unwrap();
g.last_body = Some(body.clone());
for (k, v) in headers.iter() {
g.seen.insert(k.as_str().to_string());
if let Ok(val) = v.to_str() {
g.headers.insert(k.as_str().to_string(), val.to_string());
}
}
}
let v: Value = serde_json::from_slice(&body).unwrap_or(Value::Null);
let streaming = v.get("stream").and_then(|x| x.as_bool()).unwrap_or(false);
if streaming {
let chunks: Vec<_> = s
.stream_chunks
.iter()
.map(|c| Ok::<_, std::io::Error>(Bytes::from(*c)))
.collect();
let body = Body::from_stream(futures::stream::iter(chunks));
let mut r = Response::new(body);
*r.status_mut() = StatusCode::OK;
r.headers_mut().insert(
HeaderName::from_static("content-type"),
"text/event-stream".parse().unwrap(),
);
return r;
}
let resp = serde_json::json!({
"id": "chatcmpl-test",
"object": "chat.completion",
"model": v["model"].as_str().unwrap_or("unknown"),
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop"
}]
});
Json(resp).into_response()
}
@@ -0,0 +1,7 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Shared test harness re-exports.
pub mod mock_worker;
pub mod streaming;
@@ -0,0 +1,61 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! SSE parsing and body-collection helpers for integration tests.
use bytes::Bytes;
/// Parse an SSE stream's `data: …` payloads (one per event).
#[allow(dead_code)]
pub fn parse_sse_data(raw: &[u8]) -> Vec<String> {
let s = std::str::from_utf8(raw).unwrap_or("");
s.lines()
.filter_map(|l| l.strip_prefix("data: "))
.map(|l| l.to_string())
.collect()
}
/// Collect an axum Body to bytes in tests.
#[allow(dead_code)]
pub async fn collect_body(body: axum::body::Body) -> Bytes {
use http_body_util::BodyExt;
body.collect().await.unwrap().to_bytes()
}
#[cfg(test)]
mod tests {
use super::*;
/// Ported from SMG tests/api/streaming_tests.rs::test_sse_format_parsing.
/// Verifies that parse_sse_data:
/// 1. Extracts standard `data: …` lines.
/// 2. Silently ignores SSE `event: …` type fields (not data lines).
/// 3. Silently ignores SSE `: …` comment lines.
/// 4. Correctly parses `[DONE]` sentinel.
///
/// These edge-cases matter because SGLang workers may emit `event: message`
/// fields in their SSE frames. A parser that accidentally leaks those into
/// the payload list would cause clients to fail on JSON-parse.
#[test]
fn parse_sse_data_extracts_data_lines_only() {
// Basic: three data lines including the [DONE] sentinel.
let basic =
b"data: {\"text\":\"Hello\"}\n\ndata: {\"text\":\" world\"}\n\ndata: [DONE]\n\n";
let events = parse_sse_data(basic);
assert_eq!(events.len(), 3, "expected 3 data events, got: {events:?}");
assert_eq!(events[0], "{\"text\":\"Hello\"}");
assert_eq!(events[1], "{\"text\":\" world\"}");
assert_eq!(events[2], "[DONE]");
// Mixed: event: type field + comment line — neither must appear in output.
let mixed = b"event: message\ndata: {\"test\":true}\n\n: comment line\ndata: [DONE]\n\n";
let events = parse_sse_data(mixed);
assert_eq!(
events.len(),
2,
"event: and : comment lines must be ignored; got: {events:?}"
);
assert_eq!(events[0], "{\"test\":true}");
assert_eq!(events[1], "[DONE]");
}
}
@@ -0,0 +1,143 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use axum::body::Body;
use axum::http::Request;
use sgl_router::config::*;
use sgl_router::discovery::{spawn_discovery, ModelId};
use sgl_router::policies::factory::build_registry_with_defaults as build_policy_registry;
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::manager;
use sgl_router::workers::WorkerRegistry;
use std::sync::Arc;
use std::time::Duration;
use tower::ServiceExt;
#[tokio::test]
async fn failover_when_one_worker_dies() {
// Three mock workers. Each advertises served_model_name = "tiny" on
// /server_info, so the worker manager's introspect step resolves the
// registry's model_ids without us having to hand-declare them here.
let w1 = crate::common::mock_worker::MockWorker::start(vec![]).await;
let w2 = crate::common::mock_worker::MockWorker::start(vec![]).await;
let w3 = crate::common::mock_worker::MockWorker::start(vec![]).await;
let cfg = Config {
server: ServerConfig {
host: "0".into(),
port: 0,
},
observability: Default::default(),
models: vec![ModelConfig {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
circuit_breaker: Some(CircuitBreakerConfig {
threshold: std::num::NonZeroU32::new(1).unwrap(), // open after first failure
cool_down_secs: 30,
}),
cache_aware: None,
}],
discovery: DiscoveryConfig {
backend: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec![w1.url.clone(), w2.url.clone(), w3.url.clone()],
}),
},
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
};
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
let registry = Arc::new(WorkerRegistry::default());
let policies = Arc::new(build_policy_registry(&cfg).unwrap());
let (event_rx, _disc) = spawn_discovery(&cfg).await.unwrap();
let _mgr = tokio::spawn(manager::run_with_config(
event_rx,
registry.clone(),
Some(Arc::new(cfg.clone())),
None,
None,
));
// Poll for the registry to converge — `register_one` introspect is
// a per-task spawn (manager.rs:127), so order of registration is
// non-deterministic under load. Cap the wait so a real hang surfaces
// instead of becoming a flake.
let converged = tokio::time::timeout(Duration::from_secs(5), async {
loop {
if registry.workers_for(&ModelId("tiny".into())).len() == 3 {
return;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await;
assert!(
converged.is_ok(),
"registry should contain all 3 workers after discovery; have {}",
registry.workers_for(&ModelId("tiny".into())).len()
);
let proxy = Arc::new(Proxy::new(Duration::from_secs(5)).unwrap());
let ctx = Arc::new(AppContext::new(
cfg,
tokenizers,
proxy,
registry.clone(),
policies,
));
ctx.mark_ready();
let app = build_router(ctx);
// Kill w2 by dropping its handle, then poll until its socket
// actually refuses connections. Without this, the first request
// routed to w2 can race against the listener's graceful shutdown
// and succeed, masking the failover assertion below.
let w2_url = w2.url.clone();
drop(w2);
let host_port = w2_url.trim_start_matches("http://");
let down = tokio::time::timeout(Duration::from_secs(2), async {
loop {
if tokio::net::TcpStream::connect(host_port).await.is_err() {
return;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await;
assert!(down.is_ok(), "w2 socket never went down");
// Send 6 requests; round-robin would route 2 to w2 → connection refused →
// breaker opens (threshold=1); subsequent round-robin picks rotate among
// the 2 healthy workers (#1 and #3) because healthy_workers_for filters out w2.
let mut errs = 0usize;
let mut oks = 0usize;
for i in 0..6 {
let body = serde_json::to_vec(&serde_json::json!({
"model": "tiny",
"messages": [{"role": "user", "content": format!("hi {i}")}],
}))
.unwrap();
let req = Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap();
let res = app.clone().oneshot(req).await.unwrap();
if res.status().is_success() {
oks += 1;
} else {
errs += 1;
}
}
// We expect exactly 1 error — the first call routed to w2 fails and opens
// its breaker; subsequent round-robin picks rotate among the 2 healthy
// workers since registry.healthy_workers_for filters out the open breaker.
assert_eq!(errs, 1, "exactly the first w2 pick should error");
assert_eq!(oks, 5, "remaining 5 picks should succeed via filtered RR");
}
@@ -0,0 +1,230 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Pins the contract that `axum::serve(...).with_graceful_shutdown(...)` —
//! exactly as wired in `src/main.rs` — drains every in-flight streaming
//! request through the **real** `build_router(ctx)` stack before the
//! server future resolves. A k8s SIGTERM must not truncate streaming
//! completions.
//!
//! Why route the test through the real router (chat handler + proxy +
//! SSE pump) rather than a synthetic `Router::new().route(...)`: a
//! truncation regression could live in `forward_streaming_to`'s
//! `bytes_stream_to_body` completion hook, in `chat::chat_completions`'
//! guards, or in the SSE pump's `tx.send().await` race — all of which
//! would be silently skipped by a synthetic-handler test.
use bytes::Bytes;
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, Instant};
use tokio::net::TcpListener;
use tokio::sync::oneshot;
const TEST_TIMEOUT: Duration = Duration::from_secs(15);
fn build_ctx_with_worker(worker_url: &str) -> Arc<AppContext> {
let cfg = Config {
server: ServerConfig {
host: "127.0.0.1".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(),
};
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
let registry = Arc::new(WorkerRegistry::default());
registry
.add(WorkerSpec {
id: WorkerId("w1".into()),
url: worker_url.to_string(),
mode: WorkerMode::Plain,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: None,
})
.expect("test worker accepted");
let policies = Arc::new(build_registry_with_defaults(&cfg).unwrap());
let proxy = Arc::new(Proxy::new(TEST_TIMEOUT).unwrap());
let ctx = AppContext::new(cfg, tokenizers, proxy, registry, policies);
ctx.mark_ready();
Arc::new(ctx)
}
/// Streaming chat-completions body the worker hands back chunk-by-chunk.
/// One ~60 ms delay per chunk × 8 chunks ≈ ~480 ms per request, long
/// enough that we can race in ~100 concurrent clients and trigger
/// shutdown while every stream is still mid-flight.
const SLOW_CHUNKS: &[&str] = &[
"data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n",
"data: {\"choices\":[{\"delta\":{\"content\":\"b\"}}]}\n\n",
"data: {\"choices\":[{\"delta\":{\"content\":\"c\"}}]}\n\n",
"data: {\"choices\":[{\"delta\":{\"content\":\"d\"}}]}\n\n",
"data: {\"choices\":[{\"delta\":{\"content\":\"e\"}}]}\n\n",
"data: {\"choices\":[{\"delta\":{\"content\":\"f\"}}]}\n\n",
"data: {\"choices\":[{\"delta\":{\"content\":\"g\"}}]}\n\n",
"data: [DONE]\n\n",
];
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn shutdown_drains_100_inflight_streaming_chat_completions() {
// 1. Spin up a slow streaming worker.
let worker = crate::common::mock_worker::MockWorker::start_slow_stream(
SLOW_CHUNKS.to_vec(),
Duration::from_millis(60),
)
.await;
let ctx = build_ctx_with_worker(&worker.url);
// 2. Serve the REAL `build_router(ctx)` on a random port with the
// `with_graceful_shutdown` wiring main.rs uses.
let app = build_router(ctx);
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let url = format!("http://{addr}/v1/chat/completions");
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
let server = tokio::spawn(async move {
axum::serve(listener, app)
.with_graceful_shutdown(async move {
let _ = shutdown_rx.await;
})
.await
.expect("axum::serve cleanly resolves on shutdown");
});
// 3. Fire 100 concurrent streaming clients.
const N: usize = 100;
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(10))
.build()
.unwrap();
let body = serde_json::to_vec(&serde_json::json!({
"model": "tiny",
"messages": [{"role": "user", "content": "hi"}],
"stream": true,
}))
.unwrap();
let mut handles = Vec::with_capacity(N);
for i in 0..N {
let c = client.clone();
let u = url.clone();
let b = body.clone();
handles.push(tokio::spawn(async move {
let resp = c
.post(&u)
.header("content-type", "application/json")
.body(b)
.send()
.await
.map_err(|e| format!("client {i} send: {e}"))?;
if !resp.status().is_success() {
return Err(format!("client {i} non-2xx: {}", resp.status()));
}
let bytes: Bytes = resp
.bytes()
.await
.map_err(|e| format!("client {i} body: {e}"))?;
Ok::<Bytes, String>(bytes)
}));
}
// 4. Let every request grab a connection and start receiving data.
// 100 ms is past the first chunk delay (60 ms) for every stream
// but well before the last chunk fires.
tokio::time::sleep(Duration::from_millis(100)).await;
// 5. Trigger shutdown. axum stops accepting new connections but
// MUST drain the 100 already-attached streams.
let started = Instant::now();
shutdown_tx.send(()).unwrap();
// 6. Every in-flight request must complete with a `[DONE]` terminator
// — proving the stream was NOT truncated by shutdown.
let mut bytes_total: usize = 0;
let mut done_count: usize = 0;
for h in handles {
let result = h
.await
.expect("client task panicked")
.expect("client completed");
bytes_total += result.len();
let body_str = String::from_utf8_lossy(&result);
if body_str.contains("data: [DONE]") {
done_count += 1;
}
}
// Server task must exit cleanly once all 100 in-flight requests drained.
server.await.expect("server task joins after shutdown");
let elapsed = started.elapsed();
assert_eq!(
done_count, N,
"all {N} streams must terminate with `data: [DONE]` during graceful shutdown (got {done_count})"
);
assert!(
bytes_total > 0,
"expected non-zero body bytes across {N} clients"
);
// Drain MUST have taken at least ~400 ms (7 remaining chunks * 60ms).
// A shorter wait implies the streams were truncated.
assert!(
elapsed >= Duration::from_millis(300),
"graceful shutdown returned too fast ({elapsed:?}) — likely truncated streams"
);
}
#[tokio::test]
async fn shutdown_with_no_inflight_returns_promptly() {
// Complement of the load test: when nothing is in flight, the
// shutdown future resolves quickly. Catches a regression where the
// server might hang waiting on an idle connection pool.
let worker = crate::common::mock_worker::MockWorker::start(vec![]).await;
let ctx = build_ctx_with_worker(&worker.url);
let app = build_router(ctx);
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
let server = tokio::spawn(async move {
axum::serve(listener, app)
.with_graceful_shutdown(async move {
let _ = shutdown_rx.await;
})
.await
.unwrap();
});
let started = Instant::now();
shutdown_tx.send(()).unwrap();
tokio::time::timeout(Duration::from_secs(2), server)
.await
.expect("server resolves within 2s when idle")
.expect("server task joined cleanly");
let elapsed = started.elapsed();
assert!(
elapsed < Duration::from_secs(1),
"idle shutdown took too long: {elapsed:?}"
);
}
@@ -0,0 +1,125 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use axum::body::Body;
use axum::http::Request;
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 as build_policy_registry;
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;
#[tokio::test]
async fn forwards_whitelisted_headers_strips_others() {
let worker = crate::common::mock_worker::MockWorker::start(vec![]).await;
let cfg = 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(),
};
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
let registry = Arc::new(WorkerRegistry::default());
let _ = registry.add(WorkerSpec {
id: WorkerId("w1".into()),
url: worker.url.clone(),
mode: WorkerMode::Plain,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: None,
});
let policies = Arc::new(build_policy_registry(&cfg).unwrap());
let proxy = Arc::new(Proxy::new(Duration::from_secs(5)).unwrap());
let app = build_router(Arc::new(AppContext::new(
cfg, tokenizers, proxy, registry, policies,
)));
let body = serde_json::to_vec(&serde_json::json!({
"model":"tiny","messages":[{"role":"user","content":"hi"}]
}))
.unwrap();
// Use a spoofed content-length that differs from the real body length so we
// can distinguish "inbound value forwarded" from "reqwest auto-computed it".
let spoofed_content_length = "99999";
let req = Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header("content-type", "application/json")
.header("authorization", "Bearer test")
.header("x-request-id", "abc-123")
.header("x-sgl-route-key", "k1")
.header("cookie", "should-not-forward=true")
.header("host", "example.com")
.header("content-length", spoofed_content_length)
.header("transfer-encoding", "chunked")
.body(Body::from(body))
.unwrap();
app.oneshot(req).await.unwrap();
let seen = worker.captured.lock().unwrap();
// Whitelisted headers are forwarded with their inbound VALUES intact —
// a regression that mangles, uppercases, or drops the value (e.g.,
// forwarding the name but not the value) must fail this assertion.
assert_eq!(
seen.headers.get("authorization").map(String::as_str),
Some("Bearer test"),
"authorization must be forwarded with its inbound value verbatim",
);
assert_eq!(
seen.headers.get("x-request-id").map(String::as_str),
Some("abc-123"),
"x-request-id must be forwarded with its inbound value verbatim",
);
assert_eq!(
seen.headers.get("x-sgl-route-key").map(String::as_str),
Some("k1"),
"x-sgl-route-key must be forwarded with its inbound value verbatim",
);
// Cookie must be stripped.
assert!(!seen.seen.contains("cookie"));
// transfer-encoding is hop-by-hop and must not be forwarded (reqwest does not
// re-add it for a regular body, so absence check is reliable here).
assert!(
!seen.seen.contains("transfer-encoding"),
"transfer-encoding is hop-by-hop and must be stripped"
);
// content-length: the inbound spoofed value must not reach the upstream.
// reqwest may auto-compute its own content-length for the outbound body,
// so we assert value-inequality rather than absence.
assert_ne!(
seen.headers.get("content-length").map(|s| s.as_str()),
Some(spoofed_content_length),
"router must not forward the inbound content-length value to upstream"
);
// Host: the inbound value must not reach the upstream.
let captured_host: Option<&String> = seen.headers.get("host");
assert_ne!(
captured_host,
Some(&"example.com".to_string()),
"router must not forward the inbound Host header to upstream"
);
}
@@ -0,0 +1,19 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Full HTTP proxy integration tests.
//!
//! Each submodule spins up the router via `build_router(AppContext)` and
//! drives real requests through a `common::mock_worker::MockWorker`
//! backend. For component-scope tests that don't need the router, see
//! `tests/component/`.
mod common;
mod chat_routing;
mod failover;
mod graceful_shutdown;
mod header_forwarding;
mod pd_bootstrap_injection;
mod pd_pool_isolation;
mod timeout;
@@ -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));
}
@@ -0,0 +1,425 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! PD pool isolation — end-to-end at the HTTP layer using MockWorker.
//!
//! Drives the chat handler with:
//!
//! * A model whose registered workers are all `WorkerMode::Decode`. The
//! handler dispatches **prefill** traffic (chat-completions is the
//! prefill phase of a PD request), so it must return 503 with
//! `no_prefill_workers_available`.
//! * A model with no workers at all → 503 `no_healthy_workers`
//! (existing code path; pinned here so a future PD wiring change
//! doesn't silently swap codes).
//! * A PD-disagg model with both pools healthy → request flows to the
//! prefill worker (smoke; the decode worker MUST NOT be selected for
//! the chat route).
use axum::body::Body;
use axum::http::{Request, StatusCode};
use http_body_util::BodyExt;
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()
}
/// Gap closer #1: PD mode with only decode workers → 503 with
/// `no_prefill_workers_available`. The chat route is a prefill
/// dispatch, so a decode-only pool means partial failure.
#[tokio::test]
async fn pd_mode_decode_only_returns_no_prefill_workers_available() {
let worker = crate::common::mock_worker::MockWorker::start(vec![]).await;
let ctx = build_ctx(vec![WorkerSpec {
id: WorkerId("d1".into()),
url: worker.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::SERVICE_UNAVAILABLE);
assert_eq!(
res.headers().get("x-router-error-code").unwrap(),
"no_prefill_workers_available",
);
let body = res.into_body().collect().await.unwrap().to_bytes();
let body_str = String::from_utf8_lossy(&body);
assert!(
body_str.contains("\"code\":\"no_prefill_workers_available\""),
"body: {body_str}"
);
}
/// Pin the existing-code-path branch: no workers at all → 503 with
/// `no_healthy_workers`. Ensures the new PD code path didn't swap the
/// code for the "model has zero workers" case.
#[tokio::test]
async fn no_workers_returns_no_healthy_workers() {
let ctx = build_ctx(vec![]);
let app = build_router(ctx);
let res = app.oneshot(chat_request()).await.unwrap();
assert_eq!(res.status(), StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(
res.headers().get("x-router-error-code").unwrap(),
"no_healthy_workers",
);
}
/// PD-disagg deployment with both pools healthy → chat dispatch fans
/// out to BOTH the prefill and the decode worker (Pattern B: prefill
/// in a detached task, decode awaited for the client response). Both
/// receive the same bootstrap-injected body so the SGLang engine can
/// match KV transfers via `bootstrap_room`. Pool *isolation* — the
/// guarantee that the policy's prefill candidate set excludes decode
/// workers — is exercised at the resolver layer
/// (`policies::registry::tests::pd_resolution_returns_distinct_pools`).
/// Here we only assert the HTTP-layer wiring of the dual dispatch.
#[tokio::test]
async fn pd_mode_chat_dispatch_fans_to_both_prefill_and_decode() {
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);
// Fire a single request; both prefill (spawn-and-forget) and
// decode (awaited) must receive a body with the injected
// bootstrap fields. The decode body is what the client sees on
// the response.
let res = app.oneshot(chat_request()).await.unwrap();
assert_eq!(
res.status(),
StatusCode::OK,
"decode response status should reach the client",
);
// Decode receives its body synchronously (we awaited it), so it's
// guaranteed captured by the time the response returned. Scope
// the lock guard to this block so it doesn't span the `.await`
// below (clippy: await_holding_lock).
{
let decode_seen = decode.captured.lock().unwrap();
assert!(
decode_seen.last_body.is_some(),
"decode worker must receive the bootstrap-injected request body in PD mode",
);
}
// Prefill is detached; poll briefly until its capture lands. The
// prefill task races the HTTP response back to the client. The
// local binding releases the `std::sync::Mutex` guard before the
// `.await` — holding a sync mutex across an await would let one
// task pin the lock while another tries to acquire it.
let prefill_body = tokio::time::timeout(Duration::from_secs(2), async {
loop {
let captured = prefill.captured.lock().unwrap().last_body.clone();
if let Some(b) = captured {
return b;
}
tokio::time::sleep(Duration::from_millis(5)).await;
}
})
.await
.expect("prefill MUST eventually receive its body via the detached task");
assert!(!prefill_body.is_empty());
}
/// Task C: PD-mode chat request carries an `x-sgl-decode-url` header
/// pointing at the host-affinity decode peer. With two prefill workers
/// on different hosts and a decode worker on each, the affinity helper
/// MUST pick the decode peer co-located with the chosen prefill.
///
/// Round-robin will select prefill workers deterministically (alphabetic
/// dashmap order is not guaranteed; the test fires several requests so
/// at least one lands on each prefill, and asserts the per-host pairing
/// holds across all of them).
#[tokio::test]
async fn pd_mode_chat_dispatch_sets_decode_affinity_header() {
use std::collections::HashSet;
let prefill_a = crate::common::mock_worker::MockWorker::start(vec![]).await;
let prefill_b = crate::common::mock_worker::MockWorker::start(vec![]).await;
let decode_a = crate::common::mock_worker::MockWorker::start(vec![]).await;
let decode_b = crate::common::mock_worker::MockWorker::start(vec![]).await;
// MockWorker URLs always bind to `127.0.0.1`, so every worker
// shares the same host string and the affinity helper's
// same-host branch is moot here — the helper still returns a
// decode peer via the load-tiebreak fallback. The unit tests in
// `policies::registry::tests::decoder_picks_same_host_when_available`
// carry the real burden of pinning the host-affinity rules; this
// integration test only asserts the wiring is in place (the
// `x-sgl-decode-url` header IS set on PD requests, and the
// value is one of the registered decode worker URLs).
let ctx = build_ctx(vec![
WorkerSpec {
id: WorkerId("p1".into()),
url: prefill_a.url.clone(),
mode: WorkerMode::Prefill,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: None,
},
WorkerSpec {
id: WorkerId("p2".into()),
url: prefill_b.url.clone(),
mode: WorkerMode::Prefill,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: None,
},
WorkerSpec {
id: WorkerId("d1".into()),
url: decode_a.url.clone(),
mode: WorkerMode::Decode,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: None,
},
WorkerSpec {
id: WorkerId("d2".into()),
url: decode_b.url.clone(),
mode: WorkerMode::Decode,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: None,
},
]);
let app = build_router(ctx);
// Fire 4 requests; both prefill workers see traffic via round-robin.
for _ in 0..4 {
let res = app.clone().oneshot(chat_request()).await.unwrap();
assert_eq!(res.status(), StatusCode::OK);
}
// Every request that hit a prefill mock MUST carry the decode-hint
// header. The header value MUST be one of the two registered
// decode worker URLs.
let decode_urls: HashSet<String> = [decode_a.url.clone(), decode_b.url.clone()]
.into_iter()
.collect();
for (label, p) in [("prefill_a", &prefill_a), ("prefill_b", &prefill_b)] {
let g = p.captured.lock().unwrap();
if g.last_body.is_none() {
// This prefill didn't receive a request — round-robin's
// dashmap iteration is non-deterministic, so one side may
// skip in a 4-request fire. Continue.
continue;
}
let hdr = g.headers.get("x-sgl-decode-url").unwrap_or_else(|| {
panic!(
"{label} did not receive an x-sgl-decode-url header. headers: {:?}",
g.headers
)
});
assert!(
decode_urls.contains(hdr),
"{label} got decode hint {hdr}, expected one of {decode_urls:?}",
);
}
}
/// Task C: plain-mode (non-PD) request does NOT carry the
/// `x-sgl-decode-url` header. Pin: the affinity step is gated on
/// `worker.mode() == Prefill` so plain workers are not asked to
/// bootstrap nonexistent decode peers.
#[tokio::test]
async fn plain_mode_chat_dispatch_omits_decode_affinity_header() {
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 g = plain.captured.lock().unwrap();
assert!(
!g.headers.contains_key("x-sgl-decode-url"),
"plain-mode worker must not receive a decode-affinity header. headers: {:?}",
g.headers,
);
}
/// Task C: PD-mode prefill request with NO decode workers → 503
/// `no_decode_workers_available`. Pin: failure mode is loud and
/// distinct from the existing `no_prefill_workers_available` path.
#[tokio::test]
async fn pd_mode_prefill_only_returns_no_decode_workers_available() {
let prefill = 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: None,
}]);
let app = build_router(ctx);
let res = app.oneshot(chat_request()).await.unwrap();
assert_eq!(res.status(), StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(
res.headers().get("x-router-error-code").unwrap(),
"no_decode_workers_available",
);
}
/// PD-mode chat response carries `x-sgl-decode-url` so external tests
/// can observe decode affinity end-to-end (without sniffing the proxy
/// hop into the upstream prefill worker). Mirrors the request-side
/// behavior asserted by `pd_mode_chat_dispatch_sets_decode_affinity_header`.
#[tokio::test]
async fn pd_mode_chat_response_carries_decode_affinity_header() {
use std::collections::HashSet;
let prefill = crate::common::mock_worker::MockWorker::start(vec![]).await;
let decode_a = crate::common::mock_worker::MockWorker::start(vec![]).await;
let decode_b = 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: None,
},
WorkerSpec {
id: WorkerId("d1".into()),
url: decode_a.url.clone(),
mode: WorkerMode::Decode,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: None,
},
WorkerSpec {
id: WorkerId("d2".into()),
url: decode_b.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);
let decode_urls: HashSet<String> = [decode_a.url.clone(), decode_b.url.clone()]
.into_iter()
.collect();
let hdr = res
.headers()
.get("x-sgl-decode-url")
.unwrap_or_else(|| {
panic!(
"PD-mode chat response did not carry x-sgl-decode-url; headers: {:?}",
res.headers(),
)
})
.to_str()
.unwrap()
.to_owned();
assert!(
decode_urls.contains(&hdr),
"response carried decode hint {hdr}, expected one of {decode_urls:?}",
);
}
/// Plain-mode chat response does NOT carry `x-sgl-decode-url`. Pin: the
/// response-side mirror is gated on PD-mode dispatch.
#[tokio::test]
async fn plain_mode_chat_response_omits_decode_affinity_header() {
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);
assert!(
!res.headers().contains_key("x-sgl-decode-url"),
"plain-mode chat response must not carry x-sgl-decode-url; headers: {:?}",
res.headers(),
);
}
@@ -0,0 +1,115 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Tests that the router does not wedge indefinitely when an upstream
//! worker accepts the TCP connection but never sends response headers.
//!
//! Without a configured `.timeout(...)` on the reqwest client, a stalled
//! backend hangs the axum handler future forever and the test harness
//! would just timeout. We assert here that the router returns a fast,
//! clean 502 (`upstream_timeout`) instead.
use axum::body::Body;
use axum::http::{Request, StatusCode};
use http_body_util::BodyExt;
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 as build_policy_registry;
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(_worker_url: &str) -> 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(),
}
}
#[tokio::test]
async fn non_streaming_request_times_out_when_worker_hangs() {
// Worker accepts and then sleeps for 5s; router timeout is 200ms.
let worker =
crate::common::mock_worker::MockWorker::start_hanging(Duration::from_secs(5)).await;
let cfg = config(&worker.url);
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
let registry = Arc::new(WorkerRegistry::default());
let _ = registry.add(WorkerSpec {
id: WorkerId("w1".into()),
url: worker.url.clone(),
mode: WorkerMode::Plain,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: None,
});
let policies = Arc::new(build_policy_registry(&cfg).unwrap());
let proxy = Arc::new(Proxy::new(Duration::from_millis(200)).unwrap());
let ctx = Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies));
let app = build_router(ctx);
let req = 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"}],
"stream": false
}))
.unwrap(),
))
.unwrap();
let started = std::time::Instant::now();
// Outer guard so a regression doesn't wedge CI forever.
let res = tokio::time::timeout(Duration::from_secs(2), app.oneshot(req))
.await
.expect("router must return within 2s when proxy timeout is 200ms")
.unwrap();
let elapsed = started.elapsed();
assert!(
elapsed < Duration::from_secs(1),
"router must short-circuit on upstream timeout; elapsed {elapsed:?}"
);
assert_eq!(res.status(), StatusCode::BAD_GATEWAY);
assert_eq!(
res.headers().get("x-router-error-code").unwrap(),
"upstream_timeout"
);
let bytes = res.into_body().collect().await.unwrap().to_bytes();
let body_str = String::from_utf8_lossy(&bytes);
assert!(
body_str.contains("\"code\":\"upstream_timeout\""),
"body: {body_str}"
);
// No leak of worker URL or reqwest source chain to the client.
assert!(
!body_str.contains(&worker.url),
"worker URL must not leak in client-visible body: {body_str}"
);
}
@@ -0,0 +1,237 @@
"""
Generator + validator for KV-event block-hash parity fixtures.
Two modes:
python3 experimental/sgl-router/tests/scripts/generate_kv_events_hash_parity.py
Regenerate the committed JSON fixture from the locally-replicated
algorithm. Run this when changing block-hash logic or adding new
shape coverage. CI's drift-check step runs this in --check mode.
python3 experimental/sgl-router/tests/scripts/generate_kv_events_hash_parity.py --validate-against-sglang
Import the real `sglang.srt.mem_cache.radix_cache.RadixKey.hash_page`
and assert it agrees with the locally-replicated algorithm on every
fixture case. This is the only place the replica and the real
SGLang implementation are checked against each other. Run it
nightly (or whenever sglang is available on the Python path).
# Authority
Source-of-truth implementation:
- `python/sglang/srt/mem_cache/radix_cache.py::RadixKey.hash_page`
- `python/sglang/srt/mem_cache/utils.py::hash_str_to_int64`
`hash_page_chain` below replicates that algorithm verbatim (no `import
sglang`) so the script runs without the heavy SGLang dependency tree and
can be audited at a glance. The algorithm is intentionally tiny:
sha256(prior_digest_bytes ++ token_LE_u32 ++ token_LE_u32 ++ ...)
truncate to i64 = signed(first 16 hex chars)
If SGLang ever changes the algorithm, update both the SGLang side AND
this script in the same commit; the Rust port in
`src/policies/kv_events/hash.rs` will then need the corresponding
update. The nightly `--validate-against-sglang` job is the safety net
that catches an SGLang-side change the human forgot to mirror here.
# Output format
A JSON array of cases. Each case is:
{
"name": "<descriptive label>",
"tokens": [<u32>, ...],
"block_size": <usize>,
"expected_i64_hashes": [<i64>, ...]
}
"""
from __future__ import annotations
import argparse
import hashlib
import json
import pathlib
import sys
def hash_page_chain(tokens: list[int], block_size: int) -> list[int]:
"""Compute the i64-truncated block hashes for `tokens` using SGLang's
`RadixKey.hash_page` algorithm + `hash_str_to_int64`.
Returns one i64 per full or partial block. A partial last block (when
`len(tokens) % block_size != 0`) chains against the previous block's
full 32-byte SHA256 digest, matching SGLang's behaviour.
"""
if block_size == 0:
raise ValueError("block_size must be positive")
out: list[int] = []
prior_digest: bytes | None = None
n = len(tokens)
if n == 0:
return out
# Walk every page boundary, including a trailing partial page.
start = 0
while start < n:
end = min(start + block_size, n)
hasher = hashlib.sha256()
if prior_digest is not None:
hasher.update(prior_digest)
for t in tokens[start:end]:
hasher.update(t.to_bytes(4, byteorder="little", signed=False))
digest = hasher.digest()
prior_digest = digest
# hash_str_to_int64: first 16 hex chars (top 64 bits) -> signed i64.
hex_digest = digest.hex()
uint64_val = int(hex_digest[:16], 16)
if uint64_val >= 2**63:
i64 = uint64_val - 2**64
else:
i64 = uint64_val
out.append(i64)
start = end
return out
# Cases mirror the three existing `cross_language_golden_*` tests plus
# additional shape coverage that exercises (a) zero-token edge, (b)
# block_size = 1, (c) very long sequences, (d) odd boundaries.
CASES: list[dict] = [
{
"name": "single_full_block",
"tokens": [1, 2, 3, 4],
"block_size": 4,
},
{
"name": "partial_last_block",
"tokens": [1, 2, 3, 4, 5],
"block_size": 4,
},
{
"name": "multi_block",
"tokens": [10, 20, 30, 40, 50, 60, 70, 80],
"block_size": 2,
},
{
"name": "empty_tokens",
"tokens": [],
"block_size": 4,
},
{
"name": "block_size_one",
"tokens": [7, 8, 9],
"block_size": 1,
},
{
"name": "odd_boundary",
"tokens": [100, 200, 300, 400, 500, 600, 700],
"block_size": 3,
},
{
"name": "long_sequence",
# 128 tokens at block_size 16 → 8 blocks exactly.
"tokens": list(range(1, 129)),
"block_size": 16,
},
]
def _materialize_cases() -> list[dict]:
return [
{
"name": c["name"],
"tokens": c["tokens"],
"block_size": c["block_size"],
"expected_i64_hashes": hash_page_chain(c["tokens"], c["block_size"]),
}
for c in CASES
]
def _validate_against_sglang() -> int:
"""Import the real SGLang `RadixKey.hash_page` and compare its output
case-by-case against the locally-replicated `hash_page_chain`. Exits
non-zero (and prints a diff-friendly summary) on any mismatch.
Returns 0 on success. This is the parity safety net for nightly CI.
"""
try:
from sglang.srt.mem_cache.radix_cache import RadixKey
except ImportError as e:
print(
f"--validate-against-sglang: cannot import sglang ({e}). "
"Install sglang into the Python path before running this mode.",
file=sys.stderr,
)
return 2
failures: list[str] = []
for c in CASES:
local = hash_page_chain(c["tokens"], c["block_size"])
if c["block_size"] == 0 or not c["tokens"]:
# `RadixKey.hash_page` requires a non-empty page; the local
# replica handles edge cases (empty input → empty list)
# which the SGLang oracle would refuse. Skip these cases
# under validation — the replica owns the boundary semantics.
continue
sglang_hashes: list[int] = []
prior_hex: str | None = None
for start in range(0, len(c["tokens"]), c["block_size"]):
page = c["tokens"][start : start + c["block_size"]]
key = RadixKey(token_ids=page, extra_key=None)
hex_digest = key.hash_page(prior_hex)
# SGLang's hash_page returns the hex digest; truncate to i64
# the same way `hash_str_to_int64` does.
uint64_val = int(hex_digest[:16], 16)
i64 = uint64_val - (1 << 64) if uint64_val >= (1 << 63) else uint64_val
sglang_hashes.append(i64)
prior_hex = hex_digest
if sglang_hashes != local:
failures.append(f"case {c['name']}: local={local} sglang={sglang_hashes}")
if failures:
print(
"--validate-against-sglang: replica/SGLang DRIFT detected:",
file=sys.stderr,
)
for f in failures:
print(f" {f}", file=sys.stderr)
return 1
print(f"--validate-against-sglang: OK ({len(CASES)} cases agreed)")
return 0
def _write_fixture(cases_out: list[dict]) -> pathlib.Path:
out_path = (
pathlib.Path(__file__).resolve().parent.parent
/ "fixtures"
/ "kv_events_hash_parity.json"
)
out_path.parent.mkdir(parents=True, exist_ok=True)
with out_path.open("w") as f:
json.dump(cases_out, f, indent=2, sort_keys=False)
f.write("\n")
return out_path
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--validate-against-sglang",
action="store_true",
help="Compare the local replica to the imported SGLang implementation "
"and exit non-zero on drift. Requires sglang on the Python path.",
)
args = parser.parse_args()
if args.validate_against_sglang:
return _validate_against_sglang()
cases_out = _materialize_cases()
out_path = _write_fixture(cases_out)
print(f"wrote {len(cases_out)} cases to {out_path}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
@@ -0,0 +1,115 @@
"""
One-shot generator for tokenizer parity fixtures.
Run manually when adding a model or changing a prompt shape:
python3 -m venv /tmp/parity-fixture-venv
/tmp/parity-fixture-venv/bin/pip install transformers
/tmp/parity-fixture-venv/bin/python experimental/sgl-router/tests/scripts/generate_parity_fixtures.py
CI does NOT run this — it consumes the committed JSON.
Model substitutions (gated models → public siblings of same family):
- Qwen/Qwen3-30B-A3B (gated) → Qwen/Qwen3-0.6B (same Qwen3 family, public)
- deepseek-ai/DeepSeek-V3.2-Exp (gated) → deepseek-ai/DeepSeek-V3 (older public sibling)
- openai/gpt-oss-20b → openai/gpt-oss-20b (public, used as-is)
The acceptance criterion is "3 production model families × 4 shapes".
Using a smaller model from the same family satisfies the tokenizer parity
requirement because they share the same tokenizer.json vocabulary and merges.
"""
import json
import pathlib
import sys
try:
from transformers import AutoTokenizer
except ImportError:
sys.exit("pip install transformers first")
ROOT = pathlib.Path(__file__).resolve().parents[1] / "fixtures" / "tokenizer_parity"
# Primary model ids (may be gated). Fallbacks used automatically if 401/403.
MODELS = [
# (primary_hf_id, fallback_hf_id, slug)
("Qwen/Qwen3-30B-A3B", "Qwen/Qwen3-0.6B", "qwen3-30b"),
("deepseek-ai/DeepSeek-V3.2-Exp", "deepseek-ai/DeepSeek-V3", "deepseek-v3p2"),
("openai/gpt-oss-20b", None, "gpt-oss-20b"),
]
LOREM = (
"Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod "
"tempor incididunt ut labore et dolore magna aliqua. " * 30
)
SHAPES = {
"short": "Hello, world!",
"long": LOREM,
"special_token_heavy": (
"<|im_start|>system\nYou are helpful.<|im_end|>\n"
"<|im_start|>user\nHi<|im_end|>\n"
"<|im_start|>assistant\nHello<|im_end|>\n<|endoftext|>"
),
"multi_turn_with_tools": (
"<|im_start|>system\nYou have tools.<|im_end|>\n"
"<|im_start|>user\nWeather in Paris?<|im_end|>\n"
"<|im_start|>assistant\n<tool_call>\n"
'{"name": "get_weather", "arguments": {"city": "Paris"}}\n'
"</tool_call><|im_end|>\n"
),
}
def load_tokenizer_with_fallback(primary, fallback, slug):
"""Try primary model id; fall back to sibling on any load failure.
Failure modes handled:
- 401/403/gated: access denied on HuggingFace
- ValueError/KeyError: model type too new for installed transformers
- AttributeError: broken config chain in transformers compatibility layer
- OSError/requests errors: network / hub issues
"""
for hf_id in filter(None, [primary, fallback]):
try:
print(f" Trying {hf_id}...", flush=True)
tok = AutoTokenizer.from_pretrained(hf_id, trust_remote_code=True)
print(f" Loaded {hf_id}", flush=True)
return hf_id, tok
except (ValueError, KeyError, AttributeError, OSError) as e:
msg = str(e)
print(
f" {hf_id}: load failed ({type(e).__name__}: {msg[:120]}), trying fallback...",
flush=True,
)
if fallback is None:
raise
continue
raise RuntimeError(
f"No accessible tokenizer for slug={slug} " f"(tried: {primary}, {fallback})"
)
def main():
total = 0
for primary, fallback, slug in MODELS:
out = ROOT / slug
out.mkdir(parents=True, exist_ok=True)
print(f"\nLoading tokenizer for slug={slug}:", flush=True)
actual_hf_id, tok = load_tokenizer_with_fallback(primary, fallback, slug)
for shape, text in SHAPES.items():
ids = tok.encode(text, add_special_tokens=False)
fixture = {
"model_id": actual_hf_id,
"shape": shape,
"prompt_text": text,
"expected_token_ids": ids,
"skip_special_tokens": False,
}
(out / f"{shape}.json").write_text(json.dumps(fixture, indent=2))
print(f" {slug}/{shape}: {len(ids)} tokens", flush=True)
total += 1
print(f"\nDone: {total} fixtures written to {ROOT}", flush=True)
if __name__ == "__main__":
main()