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,75 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
//! Policy-selection throughput microbench.
|
||||
//!
|
||||
//! Mirrors `sgl-model-gateway/benches/manual_policy_benchmark.rs` —
|
||||
//! measures how fast the routing layer returns a worker for a given
|
||||
//! request context, across the policies sgl-router actually ships
|
||||
//! (round-robin, random, power-of-two-choices). The cache-aware-zmq
|
||||
//! policy lives in `tree_lookup.rs`; this file targets the non-tree
|
||||
//! policies' steady-state hot path.
|
||||
|
||||
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
use sgl_router::policies::power_of_two::PowerOfTwoChoicesPolicy;
|
||||
use sgl_router::policies::random::RandomPolicy;
|
||||
use sgl_router::policies::round_robin::RoundRobinPolicy;
|
||||
use sgl_router::policies::{Policy, SelectionContext};
|
||||
use sgl_router::workers::{Worker, WorkerRegistry};
|
||||
use std::sync::Arc;
|
||||
|
||||
fn workers(n: usize, model: &str) -> Vec<Arc<Worker>> {
|
||||
let registry = WorkerRegistry::default();
|
||||
for i in 0..n {
|
||||
registry
|
||||
.add(WorkerSpec {
|
||||
id: WorkerId(format!("w{i}")),
|
||||
url: format!("http://w{i}:30000"),
|
||||
mode: WorkerMode::Plain,
|
||||
model_ids: vec![ModelId(model.into())],
|
||||
bootstrap_port: None,
|
||||
})
|
||||
.expect("test workers are unmixed");
|
||||
}
|
||||
registry.workers_for(&ModelId(model.into()))
|
||||
}
|
||||
|
||||
fn bench_policy(c: &mut Criterion, name: &str, policy: Arc<dyn Policy>) {
|
||||
let mut group = c.benchmark_group(format!("policy_select::{name}"));
|
||||
for &n in &[4usize, 16, 64, 256] {
|
||||
let workers = workers(n, "tiny");
|
||||
let model = ModelId("tiny".into());
|
||||
// Same body across iterations — measures the policy's per-call
|
||||
// cost rather than body-parsing overhead.
|
||||
let body = serde_json::to_vec(&serde_json::json!({
|
||||
"model": "tiny",
|
||||
"messages": [{"role": "user", "content": "hello world"}],
|
||||
}))
|
||||
.unwrap();
|
||||
group.throughput(Throughput::Elements(1));
|
||||
group.bench_with_input(BenchmarkId::from_parameter(n), &n, |b, _| {
|
||||
b.iter(|| {
|
||||
let ctx = SelectionContext::new(&model, Some(&body));
|
||||
let chosen = policy.select(black_box(&workers), &ctx);
|
||||
black_box(chosen);
|
||||
});
|
||||
});
|
||||
}
|
||||
group.finish();
|
||||
}
|
||||
|
||||
fn bench_round_robin(c: &mut Criterion) {
|
||||
bench_policy(c, "round_robin", Arc::new(RoundRobinPolicy::new()));
|
||||
}
|
||||
|
||||
fn bench_random(c: &mut Criterion) {
|
||||
bench_policy(c, "random", Arc::new(RandomPolicy::new()));
|
||||
}
|
||||
|
||||
fn bench_power_of_two(c: &mut Criterion) {
|
||||
bench_policy(c, "power_of_two", Arc::new(PowerOfTwoChoicesPolicy::new()));
|
||||
}
|
||||
|
||||
criterion_group!(benches, bench_round_robin, bench_random, bench_power_of_two);
|
||||
criterion_main!(benches);
|
||||
@@ -0,0 +1,89 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
//! Cache-aware tree-lookup microbench.
|
||||
//!
|
||||
//! Mirrors the shape of `sgl-model-gateway/benches/radix_tree_benchmark.rs`
|
||||
//! (specifically the `TokenTree` / `PositionalIndexer` paths — which serve
|
||||
//! the same role as sgl-router's `HashTree`). The bench measures:
|
||||
//!
|
||||
//! * `insert` — populate one worker's prefix.
|
||||
//! * `match_prefix` — score an incoming request against the tree.
|
||||
//!
|
||||
//! Output is `criterion`'s default (target/criterion/...). To run:
|
||||
//!
|
||||
//! cargo bench --bench tree_lookup
|
||||
//! cargo bench --bench tree_lookup -- --sample-size 30 # faster
|
||||
//!
|
||||
//! See `BENCHMARKS.md` for the SMG↔sgl-router comparison table.
|
||||
|
||||
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
|
||||
use rand::rngs::StdRng;
|
||||
use rand::{Rng, SeedableRng};
|
||||
use sgl_router::policies::kv_events::tree::{HashTree, KvWorkerId};
|
||||
|
||||
fn build_tree(num_workers: usize, blocks_per_worker: usize, seed: u64) -> HashTree {
|
||||
let tree = HashTree::new();
|
||||
let mut rng = StdRng::seed_from_u64(seed);
|
||||
for w in 0..num_workers {
|
||||
let worker = KvWorkerId::new(format!("http://w{w}:30000"), 0);
|
||||
// Each worker holds a distinct (random) prefix so the trees fan
|
||||
// out — this is the realistic case for cache-aware routing.
|
||||
let hashes: Vec<i64> = (0..blocks_per_worker).map(|_| rng.gen::<i64>()).collect();
|
||||
tree.insert(&worker, None, &hashes);
|
||||
}
|
||||
tree
|
||||
}
|
||||
|
||||
fn bench_insert(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("hashtree_insert");
|
||||
for &n_blocks in &[8usize, 32, 128, 512] {
|
||||
group.throughput(Throughput::Elements(n_blocks as u64));
|
||||
group.bench_with_input(BenchmarkId::from_parameter(n_blocks), &n_blocks, |b, &n| {
|
||||
let mut rng = StdRng::seed_from_u64(0xC0FFEE);
|
||||
let hashes: Vec<i64> = (0..n).map(|_| rng.gen::<i64>()).collect();
|
||||
b.iter_batched(
|
||||
HashTree::new,
|
||||
|tree| {
|
||||
let worker = KvWorkerId::new("http://w:30000".to_string(), 0);
|
||||
tree.insert(&worker, None, black_box(&hashes));
|
||||
tree
|
||||
},
|
||||
criterion::BatchSize::SmallInput,
|
||||
);
|
||||
});
|
||||
}
|
||||
group.finish();
|
||||
}
|
||||
|
||||
fn bench_match_prefix(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("hashtree_match_prefix");
|
||||
// (workers, blocks_per_worker, query_len) cases that span the
|
||||
// realistic operating window: small fleet w/ moderate prefixes,
|
||||
// medium fleet w/ long prefixes, and a stress case.
|
||||
let cases = [
|
||||
(4usize, 32usize, 8usize),
|
||||
(16, 64, 32),
|
||||
(64, 128, 64),
|
||||
(128, 256, 128),
|
||||
];
|
||||
for (workers, bpw, query_len) in cases {
|
||||
let label = format!("w{workers}_bpw{bpw}_q{query_len}");
|
||||
group.throughput(Throughput::Elements(query_len as u64));
|
||||
let tree = build_tree(workers, bpw, 0xDEADBEEF);
|
||||
// Pull one real worker's prefix so the query has a non-trivial
|
||||
// partial match — closer to the production hot path.
|
||||
let mut rng = StdRng::seed_from_u64(0x12345);
|
||||
let probe: Vec<i64> = (0..query_len).map(|_| rng.gen::<i64>()).collect();
|
||||
group.bench_function(label, |b| {
|
||||
b.iter(|| {
|
||||
let m = tree.match_prefix(None, black_box(&probe));
|
||||
black_box(m.matched_blocks)
|
||||
});
|
||||
});
|
||||
}
|
||||
group.finish();
|
||||
}
|
||||
|
||||
criterion_group!(benches, bench_insert, bench_match_prefix);
|
||||
criterion_main!(benches);
|
||||
Reference in New Issue
Block a user