minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Change routing policy API to be async to support more policies (#17048)

Browse Source
This commit is contained in:
fzyzcjy
2026-01-17 17:12:51 +08:00
committed by GitHub
parent d2c863878c
commit 9c2530642c
16 changed files with 350 additions and 217 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/benches/manual_policy_benchmark.rs
+57 -37
View File
@@ -1,10 +1,11 @@
use std::{sync::Arc, thread};
use std::sync::Arc;
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use smg::{
core::{BasicWorkerBuilder, Worker, WorkerType},
policies::{LoadBalancingPolicy, ManualPolicy, SelectWorkerInfo},
};
use tokio::runtime::Runtime;
// ============================================================================
// Test Helpers
@@ -22,19 +23,24 @@ fn create_workers(count: usize) -> Vec<Arc<dyn Worker>> {
.collect()
}
fn select_with_key(policy: &ManualPolicy, workers: &[Arc<dyn Worker>], key: &str) -> Option<usize> {
fn select_with_key(
rt: &Runtime,
policy: &ManualPolicy,
workers: &[Arc<dyn Worker>],
key: &str,
) -> Option<usize> {
let mut headers = http::HeaderMap::new();
headers.insert("x-smg-routing-key", key.parse().unwrap());
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
policy.select_worker(workers, &info)
rt.block_on(policy.select_worker(workers, &info))
}
fn warmup_keys(policy: &ManualPolicy, workers: &[Arc<dyn Worker>], keys: &[String]) {
fn warmup_keys(rt: &Runtime, policy: &ManualPolicy, workers: &[Arc<dyn Worker>], keys: &[String]) {
for key in keys {
select_with_key(policy, workers, key);
select_with_key(rt, policy, workers, key);
}
}
@@ -47,13 +53,14 @@ fn gen_keys(count: usize, prefix: &str) -> Vec<String> {
// ============================================================================
fn bench_fast_path_hit(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let mut group = c.benchmark_group("manual_policy/fast_path");
for worker_count in [4, 16, 64, 256] {
let policy = ManualPolicy::new();
let workers = create_workers(worker_count);
let keys = gen_keys(1000, "user-");
warmup_keys(&policy, &workers, &keys);
warmup_keys(&rt, &policy, &workers, &keys);
group.throughput(Throughput::Elements(1));
group.bench_with_input(
@@ -62,7 +69,7 @@ fn bench_fast_path_hit(c: &mut Criterion) {
|b, _| {
let mut idx = 0;
b.iter(|| {
let result = select_with_key(&policy, &workers, &keys[idx % keys.len()]);
let result = select_with_key(&rt, &policy, &workers, &keys[idx % keys.len()]);
idx += 1;
black_box(result)
});
@@ -73,6 +80,7 @@ fn bench_fast_path_hit(c: &mut Criterion) {
}
fn bench_slow_path_vacant(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let mut group = c.benchmark_group("manual_policy/slow_path_vacant");
for worker_count in [4, 16, 64, 256] {
@@ -87,7 +95,7 @@ fn bench_slow_path_vacant(c: &mut Criterion) {
let mut idx = 0;
b.iter(|| {
let key = format!("new-user-{}", idx);
let result = select_with_key(&policy, &workers, &key);
let result = select_with_key(&rt, &policy, &workers, &key);
idx += 1;
black_box(result)
});
@@ -98,6 +106,7 @@ fn bench_slow_path_vacant(c: &mut Criterion) {
}
fn bench_no_routing_key(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let mut group = c.benchmark_group("manual_policy/no_routing_key");
for worker_count in [4, 16, 64, 256] {
@@ -110,7 +119,7 @@ fn bench_no_routing_key(c: &mut Criterion) {
&worker_count,
|b, _| {
let info = SelectWorkerInfo::default();
b.iter(|| black_box(policy.select_worker(&workers, &info)));
b.iter(|| black_box(rt.block_on(policy.select_worker(&workers, &info))));
},
);
}
@@ -118,6 +127,7 @@ fn bench_no_routing_key(c: &mut Criterion) {
}
fn bench_failover(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let mut group = c.benchmark_group("manual_policy/failover");
group.sample_size(50);
@@ -130,12 +140,12 @@ fn bench_failover(c: &mut Criterion) {
|| {
let policy = ManualPolicy::new();
let workers = create_workers(count);
let idx = select_with_key(&policy, &workers, "failover-test").unwrap();
let idx = select_with_key(&rt, &policy, &workers, "failover-test").unwrap();
workers[idx].set_healthy(false);
(policy, workers)
},
|(policy, workers)| {
black_box(select_with_key(&policy, &workers, "failover-test"))
black_box(select_with_key(&rt, &policy, &workers, "failover-test"))
},
);
},
@@ -145,6 +155,12 @@ fn bench_failover(c: &mut Criterion) {
}
fn bench_concurrent(c: &mut Criterion) {
let rt = Arc::new(
tokio::runtime::Builder::new_multi_thread()
.worker_threads(4)
.build()
.unwrap(),
);
let mut group = c.benchmark_group("manual_policy/concurrent");
group.sample_size(50);
@@ -157,32 +173,35 @@ fn bench_concurrent(c: &mut Criterion) {
let policy = Arc::new(ManualPolicy::new());
let workers: Arc<Vec<Arc<dyn Worker>>> = Arc::new(create_workers(16));
let handles: Vec<_> = (0..threads)
.map(|t| {
let policy = Arc::clone(&policy);
let workers = Arc::clone(&workers);
thread::spawn(move || {
for i in 0..500 {
let key = if i % 5 == 0 {
format!("thread{}_user{}", t, i)
} else {
format!("shared_user{}", i % 50)
};
let mut headers = http::HeaderMap::new();
headers.insert("x-smg-routing-key", key.parse().unwrap());
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let _ = black_box(policy.select_worker(&workers, &info));
}
rt.block_on(async {
let handles: Vec<_> = (0..threads)
.map(|t| {
let policy = Arc::clone(&policy);
let workers = Arc::clone(&workers);
tokio::spawn(async move {
for i in 0..500 {
let key = if i % 5 == 0 {
format!("thread{}_user{}", t, i)
} else {
format!("shared_user{}", i % 50)
};
let mut headers = http::HeaderMap::new();
headers.insert("x-smg-routing-key", key.parse().unwrap());
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let _ =
black_box(policy.select_worker(&workers, &info).await);
}
})
})
})
.collect();
.collect();
for h in handles {
h.join().unwrap();
}
for h in handles {
h.await.unwrap();
}
});
});
},
);
@@ -191,19 +210,20 @@ fn bench_concurrent(c: &mut Criterion) {
}
fn bench_cache_size_impact(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let mut group = c.benchmark_group("manual_policy/cache_size");
for cache_size in [100, 1000, 10000, 100000] {
let policy = ManualPolicy::new();
let workers = create_workers(16);
let keys = gen_keys(cache_size, "user-");
warmup_keys(&policy, &workers, &keys);
warmup_keys(&rt, &policy, &workers, &keys);
group.throughput(Throughput::Elements(1));
group.bench_with_input(BenchmarkId::new("keys", cache_size), &cache_size, |b, _| {
let mut idx = 0;
b.iter(|| {
let result = select_with_key(&policy, &workers, &keys[idx % keys.len()]);
let result = select_with_key(&rt, &policy, &workers, &keys[idx % keys.len()]);
idx += 1;
black_box(result)
});