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:
co-authored by
Claude Opus 4.7
parent
aae04b1241
commit
6e8fe176be
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user