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

[Router] Add bucket-aware policy domains and native cache indexing (#38108)

Browse Source 
Signed-off-by: Vincent Gao <vincentbo@linux.alibaba.com>
Co-authored-by: inkcherry <mingzhi.liu@amd.com>
Co-authored-by: yangbodong22011 <13137470+yangbodong22011@users.noreply.github.com>
This commit is contained in:
Vincent Gao
2026-09-06 19:47:51 +08:00
committed by GitHub
co-authored by inkcherry yangbodong22011
parent a176ba2f7b
commit 5bebe7a033
76 changed files with 5842 additions and 3639 deletions
Download Patch File Download Diff File Expand all files Collapse all files
experimental/sgl-router/tests/proxy/bucket_routing.rs
+794
View File
@@ -0,0 +1,794 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! HTTP contract for static P/D buckets.
//!
//! Buckets narrow the candidate domain before policy selection. Prefill SLO
//! profiles may override rank, while decode uses `input_tokens + max_tokens`.
use axum::body::Body;
use axum::http::{Request, StatusCode};
use sgl_kv_indexer::{PrefixIndex, PrefixIndexError, PrefixMatch, PrefixOutcome};
use sgl_router::config::{
ActiveLoadConfig, AffinityConfig, BucketConfig, BucketSpec, BucketStage, CacheAwareConfig,
CachePrefixProvider, Config, DiscoveryBackend, KvIndexerEndpointConfig, ModelConfig,
ObservabilityConfig, PolicyKind, ProxyConfig, ServerConfig, SessionAffinityMode,
SloBucketPolicy, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::engine_load::{LoadStat, NativeCacheRankLoad};
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::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tower::ServiceExt;
fn bucket(id: &str, stage: BucketStage, rank: u32, worker_id: &str) -> BucketSpec {
BucketSpec {
id: id.into(),
stage,
rank,
worker_ids: vec![worker_id.into()],
min_extend_tokens: None,
max_extend_tokens: None,
min_sequence_tokens: None,
max_sequence_tokens: None,
max_context_tokens: Some(16_384),
ttft_p95_at_capacity_ms: None,
tps_p05_at_capacity: None,
max_pending_prefill_tokens: None,
}
}
fn build_app_context(
specs: Vec<WorkerSpec>,
bucket_config: BucketConfig,
policy: PolicyKind,
affinity: Option<AffinityConfig>,
) -> AppContext {
let config = Config {
server: ServerConfig {
host: "0".into(),
port: 0,
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy,
decode_policy: Default::default(),
bucket_config: Some(bucket_config),
circuit_breaker: None,
cache_aware: None,
sticky: None,
affinity,
fused: None,
eligibility: None,
},
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
};
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&config).unwrap());
let registry = Arc::new(WorkerRegistry::default());
for spec in specs {
let _ = registry.add(spec);
}
let policies = Arc::new(build_registry_with_defaults(&config).unwrap());
let proxy = Arc::new(Proxy::new(Duration::from_secs(5)).unwrap());
AppContext::new(config, tokenizers, proxy, registry, policies)
}
fn build_ctx(
specs: Vec<WorkerSpec>,
bucket_config: BucketConfig,
policy: PolicyKind,
affinity: Option<AffinityConfig>,
) -> Arc<AppContext> {
Arc::new(build_app_context(specs, bucket_config, policy, affinity))
}
struct FakePrefixIndex {
address: Option<String>,
calls: AtomicUsize,
}
impl FakePrefixIndex {
fn matched(address: String) -> Arc<Self> {
Arc::new(Self {
address: Some(address),
calls: AtomicUsize::new(0),
})
}
fn no_signal() -> Arc<Self> {
Arc::new(Self {
address: None,
calls: AtomicUsize::new(0),
})
}
}
#[tonic::async_trait]
impl PrefixIndex for FakePrefixIndex {
async fn match_prefix(&self, hashes: Vec<i64>) -> Result<PrefixOutcome, PrefixIndexError> {
self.calls.fetch_add(1, Ordering::Relaxed);
let Some(address) = &self.address else {
return Ok(PrefixOutcome::Empty);
};
let matched_prefix_blocks =
u32::try_from(hashes.len().saturating_sub(1)).unwrap_or(u32::MAX);
Ok(PrefixOutcome::Matched {
matches: vec![PrefixMatch {
address: address.clone(),
matched_prefix_blocks,
worker_id: "fake-index-worker".into(),
}],
best_prefix_blocks: matched_prefix_blocks,
})
}
}
struct TwoPrefixIndex {
best_address: String,
lower_ranked_address: String,
}
impl TwoPrefixIndex {
fn new(best_address: String, lower_ranked_address: String) -> Arc<Self> {
Arc::new(Self {
best_address,
lower_ranked_address,
})
}
}
#[tonic::async_trait]
impl PrefixIndex for TwoPrefixIndex {
async fn match_prefix(&self, hashes: Vec<i64>) -> Result<PrefixOutcome, PrefixIndexError> {
let best_prefix_blocks = u32::try_from(hashes.len().saturating_sub(1)).unwrap_or(u32::MAX);
let lower_ranked_prefix_blocks = (best_prefix_blocks / 2).max(1);
Ok(PrefixOutcome::Matched {
matches: vec![
PrefixMatch {
address: self.best_address.clone(),
matched_prefix_blocks: best_prefix_blocks,
worker_id: "best-index-worker".into(),
},
PrefixMatch {
address: self.lower_ranked_address.clone(),
matched_prefix_blocks: lower_ranked_prefix_blocks,
worker_id: "lower-index-worker".into(),
},
],
best_prefix_blocks,
})
}
}
fn build_cache_ctx(
specs: Vec<WorkerSpec>,
bucket_config: BucketConfig,
prefix_index: Arc<dyn PrefixIndex>,
) -> Arc<AppContext> {
build_cache_ctx_with_affinity(
specs,
bucket_config,
prefix_index,
AffinityConfig::default(),
)
}
fn build_cache_ctx_with_affinity(
specs: Vec<WorkerSpec>,
bucket_config: BucketConfig,
prefix_index: Arc<dyn PrefixIndex>,
affinity: AffinityConfig,
) -> Arc<AppContext> {
let mut context =
build_app_context(specs, bucket_config, PolicyKind::CacheAware, Some(affinity));
context.config.model.cache_aware = Some(CacheAwareConfig {
prefix_provider: CachePrefixProvider::Indexer,
kv_indexer_endpoint: Some(KvIndexerEndpointConfig {
url: "http://fake-indexer".into(),
query_timeout_ms: 100,
query_max_inflight: 32,
}),
});
context.prefix_index = Some(prefix_index);
context.block_size_oracle.try_set(1).unwrap();
Arc::new(context)
}
fn worker_spec(id: &str, url: String, mode: WorkerMode) -> WorkerSpec {
WorkerSpec {
id: WorkerId(id.into()),
url,
mode,
model_ids: vec![ModelId("tiny".into())],
bootstrap_port: (mode == WorkerMode::Prefill).then_some(8997),
}
}
fn set_native_load(
ctx: &AppContext,
worker_url: &str,
num_total_tokens: u64,
max_total_num_tokens: u64,
) {
ctx.engine_load.set(
worker_url,
0,
LoadStat {
num_running_reqs: 0,
num_waiting_reqs: 0,
num_tokens: num_total_tokens,
max_total_num_tokens,
native_cache: Some(NativeCacheRankLoad {
num_waiting_uncached_tokens: 0,
num_total_tokens,
max_running_requests: 64,
total_prefill_uncached_tokens: 1,
total_prefill_busy_us: 1,
}),
},
Instant::now(),
);
}
fn chat_request(ttft_slo_ms: Option<u64>, max_tokens: Option<u64>) -> Request<Body> {
chat_request_with_content("bucket routing", ttft_slo_ms, max_tokens, None)
}
fn chat_request_with_content(
content: &str,
ttft_slo_ms: Option<u64>,
max_tokens: Option<u64>,
session_id: Option<&str>,
) -> Request<Body> {
let mut builder = Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header("content-type", "application/json");
if let Some(ttft_slo_ms) = ttft_slo_ms {
builder = builder.header("x-sgl-ttft-slo-ms", ttft_slo_ms.to_string());
}
if let Some(session_id) = session_id {
builder = builder.header("x-session-id", session_id);
}
builder
.body(Body::from(
serde_json::to_vec(&serde_json::json!({
"model": "tiny",
"messages": [{"role": "user", "content": content}],
"max_tokens": max_tokens,
}))
.unwrap(),
))
.unwrap()
}
async fn wait_for_prefill(mock: &crate::common::mock_worker::MockWorker) {
tokio::time::timeout(Duration::from_secs(2), async {
loop {
if mock.captured.lock().unwrap().last_body.is_some() {
return;
}
tokio::time::sleep(Duration::from_millis(5)).await;
}
})
.await
.expect("selected prefill worker must receive the detached request");
}
async fn wait_for_prefill_body_containing(
mock: &crate::common::mock_worker::MockWorker,
expected: &str,
) -> Vec<u8> {
tokio::time::timeout(Duration::from_secs(2), async {
loop {
let captured = mock.captured.lock().unwrap().last_body.clone();
if let Some(body) = captured {
if String::from_utf8_lossy(&body).contains(expected) {
return body.to_vec();
}
}
tokio::time::sleep(Duration::from_millis(5)).await;
}
})
.await
.expect("selected prefill worker must receive the expected request body")
}
#[tokio::test]
async fn prefill_slo_first_uses_eligible_ttft_bucket_before_lower_rank_bucket() {
let cheap = crate::common::mock_worker::MockWorker::start(vec![]).await;
let fast = crate::common::mock_worker::MockWorker::start(vec![]).await;
let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
let mut cheap_bucket = bucket("p-cheap", BucketStage::Prefill, 10, "p-cheap");
cheap_bucket.ttft_p95_at_capacity_ms = Some(400);
let mut fast_bucket = bucket("p-fast", BucketStage::Prefill, 20, "p-fast");
fast_bucket.ttft_p95_at_capacity_ms = Some(100);
let bucket_config = BucketConfig {
buckets: vec![
cheap_bucket,
fast_bucket,
bucket("d-catch-all", BucketStage::Decode, 30, "d"),
],
ttft_slo_policy: SloBucketPolicy::SloFirst,
tps_slo_policy: SloBucketPolicy::Disabled,
};
let ctx = build_ctx(
vec![
worker_spec("p-cheap", cheap.url.clone(), WorkerMode::Prefill),
worker_spec("p-fast", fast.url.clone(), WorkerMode::Prefill),
worker_spec("d", decode.url.clone(), WorkerMode::Decode),
],
bucket_config,
PolicyKind::PowerOfTwo,
None,
);
let response = build_router(ctx)
.oneshot(chat_request(Some(200), Some(16)))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
wait_for_prefill(&fast).await;
assert!(
cheap.captured.lock().unwrap().last_body.is_none(),
"lower-rank but TTFT-ineligible P Bucket must not be dispatched first"
);
}
#[tokio::test]
async fn prefill_tries_later_compatible_bucket_before_capacity_fallback() {
let full = crate::common::mock_worker::MockWorker::start(vec![]).await;
let available = crate::common::mock_worker::MockWorker::start(vec![]).await;
let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
let bucket_config = BucketConfig {
buckets: vec![
bucket("p-full", BucketStage::Prefill, 10, "p-full"),
bucket("p-available", BucketStage::Prefill, 20, "p-available"),
bucket("d", BucketStage::Decode, 30, "d"),
],
ttft_slo_policy: SloBucketPolicy::Disabled,
tps_slo_policy: SloBucketPolicy::Disabled,
};
let ctx = build_ctx(
vec![
worker_spec("p-full", full.url.clone(), WorkerMode::Prefill),
worker_spec("p-available", available.url.clone(), WorkerMode::Prefill),
worker_spec("d", decode.url.clone(), WorkerMode::Decode),
],
bucket_config,
PolicyKind::PowerOfTwo,
None,
);
set_native_load(&ctx, &full.url, 100, 100);
set_native_load(&ctx, &available.url, 0, 10_000);
let response = build_router(ctx)
.oneshot(chat_request(None, Some(16)))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
wait_for_prefill(&available).await;
assert!(
full.captured.lock().unwrap().last_body.is_none(),
"capacity fallback must wait until all compatible prefill buckets are exhausted"
);
}
#[tokio::test]
async fn decode_bucket_uses_input_plus_requested_output_budget() {
let prefill = crate::common::mock_worker::MockWorker::start(vec![]).await;
let short_decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
let long_decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
let mut short_bucket = bucket("d-short", BucketStage::Decode, 20, "d-short");
short_bucket.max_sequence_tokens = Some(1_024);
let mut long_bucket = bucket("d-long", BucketStage::Decode, 30, "d-long");
long_bucket.min_sequence_tokens = Some(1_025);
let bucket_config = BucketConfig {
buckets: vec![
bucket("p", BucketStage::Prefill, 10, "p"),
short_bucket,
long_bucket,
],
ttft_slo_policy: SloBucketPolicy::Disabled,
tps_slo_policy: SloBucketPolicy::Disabled,
};
let ctx = build_ctx(
vec![
worker_spec("p", prefill.url.clone(), WorkerMode::Prefill),
worker_spec("d-short", short_decode.url.clone(), WorkerMode::Decode),
worker_spec("d-long", long_decode.url.clone(), WorkerMode::Decode),
],
bucket_config,
PolicyKind::PowerOfTwo,
None,
);
let response = build_router(ctx)
.oneshot(chat_request(None, Some(2_000)))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get("x-sgl-decode-url")
.and_then(|value| value.to_str().ok()),
Some(long_decode.url.as_str()),
"peak sequence length must exclude the short Decode Bucket"
);
assert!(
long_decode.captured.lock().unwrap().last_body.is_some(),
"the selected long Decode worker is awaited before the response"
);
assert!(
short_decode.captured.lock().unwrap().last_body.is_none(),
"the incompatible short Decode Bucket must not receive the request"
);
}
#[tokio::test]
async fn decode_tries_later_compatible_bucket_before_capacity_fallback() {
let prefill = crate::common::mock_worker::MockWorker::start(vec![]).await;
let full = crate::common::mock_worker::MockWorker::start(vec![]).await;
let available = crate::common::mock_worker::MockWorker::start(vec![]).await;
let bucket_config = BucketConfig {
buckets: vec![
bucket("p", BucketStage::Prefill, 10, "p"),
bucket("d-full", BucketStage::Decode, 20, "d-full"),
bucket("d-available", BucketStage::Decode, 30, "d-available"),
],
ttft_slo_policy: SloBucketPolicy::Disabled,
tps_slo_policy: SloBucketPolicy::Disabled,
};
let ctx = build_ctx(
vec![
worker_spec("p", prefill.url.clone(), WorkerMode::Prefill),
worker_spec("d-full", full.url.clone(), WorkerMode::Decode),
worker_spec("d-available", available.url.clone(), WorkerMode::Decode),
],
bucket_config,
PolicyKind::PowerOfTwo,
None,
);
set_native_load(&ctx, &full.url, 100, 100);
set_native_load(&ctx, &available.url, 0, 10_000);
let response = build_router(ctx)
.oneshot(chat_request(None, Some(16)))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get("x-sgl-decode-url")
.and_then(|value| value.to_str().ok()),
Some(available.url.as_str()),
"capacity fallback must wait until all compatible decode buckets are exhausted"
);
}
#[tokio::test]
async fn prefill_only_bucket_configuration_keeps_global_decode_routing() {
let prefill = crate::common::mock_worker::MockWorker::start(vec![]).await;
let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
let bucket_config = BucketConfig {
buckets: vec![bucket("p", BucketStage::Prefill, 10, "p")],
ttft_slo_policy: SloBucketPolicy::Disabled,
tps_slo_policy: SloBucketPolicy::Disabled,
};
let ctx = build_ctx(
vec![
worker_spec("p", prefill.url.clone(), WorkerMode::Prefill),
worker_spec("d", decode.url.clone(), WorkerMode::Decode),
],
bucket_config,
PolicyKind::PowerOfTwo,
None,
);
let response = build_router(ctx)
.oneshot(chat_request(None, Some(16)))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get("x-sgl-decode-url")
.and_then(|value| value.to_str().ok()),
Some(decode.url.as_str()),
"a Prefill-only Bucket rollout must retain the Step 1 global Decode domain"
);
}
#[tokio::test]
async fn global_rebind_session_affinity_can_keep_a_cross_length_bucket_primary() {
let short = crate::common::mock_worker::MockWorker::start(vec![]).await;
let long = crate::common::mock_worker::MockWorker::start(vec![]).await;
let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
let mut short_bucket = bucket("p-short", BucketStage::Prefill, 10, "p-short");
short_bucket.max_extend_tokens = Some(256);
short_bucket.max_context_tokens = Some(16_384);
short_bucket.ttft_p95_at_capacity_ms = Some(80);
let mut long_bucket = bucket("p-long", BucketStage::Prefill, 20, "p-long");
long_bucket.min_extend_tokens = Some(257);
long_bucket.max_context_tokens = Some(16_384);
long_bucket.ttft_p95_at_capacity_ms = Some(300);
let bucket_config = BucketConfig {
buckets: vec![
short_bucket,
long_bucket,
bucket("d-catch-all", BucketStage::Decode, 30, "d"),
],
ttft_slo_policy: SloBucketPolicy::SloFirst,
tps_slo_policy: SloBucketPolicy::Disabled,
};
let ctx = build_ctx(
vec![
worker_spec("p-short", short.url.clone(), WorkerMode::Prefill),
worker_spec("p-long", long.url.clone(), WorkerMode::Prefill),
worker_spec("d", decode.url.clone(), WorkerMode::Decode),
],
bucket_config,
PolicyKind::SessionAware,
Some(AffinityConfig {
session_affinity_mode: SessionAffinityMode::GlobalRebind,
..Default::default()
}),
);
let app = build_router(ctx);
let first = app
.clone()
.oneshot(chat_request_with_content(
"short",
Some(120),
Some(8),
Some("s-1"),
))
.await
.unwrap();
assert_eq!(first.status(), StatusCode::OK);
wait_for_prefill(&short).await;
let long_content = "length ".repeat(128);
let second = app
.oneshot(chat_request_with_content(
&long_content,
Some(120),
Some(8),
Some("s-1"),
))
.await
.unwrap();
assert_eq!(second.status(), StatusCode::OK);
let short_body = wait_for_prefill_body_containing(&short, &long_content).await;
assert!(
String::from_utf8_lossy(&short_body).contains(&long_content),
"the second, long request must retain the existing cross-Bucket session primary"
);
assert!(
long.captured.lock().unwrap().last_body.is_none(),
"target length Bucket is skipped only because the primary's own Hard TTFT profile is eligible"
);
}
#[tokio::test]
async fn global_preserve_establishes_then_reuses_a_new_assignment() {
let prefill = crate::common::mock_worker::MockWorker::start(vec![]).await;
let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
let bucket_config = BucketConfig {
buckets: vec![
bucket("p", BucketStage::Prefill, 10, "p"),
bucket("d", BucketStage::Decode, 20, "d"),
],
ttft_slo_policy: SloBucketPolicy::Disabled,
tps_slo_policy: SloBucketPolicy::Disabled,
};
let ctx = build_ctx(
vec![
worker_spec("p", prefill.url.clone(), WorkerMode::Prefill),
worker_spec("d", decode.url.clone(), WorkerMode::Decode),
],
bucket_config,
PolicyKind::SessionAware,
Some(AffinityConfig {
session_affinity_mode: SessionAffinityMode::GlobalPreserve,
..Default::default()
}),
);
let app = build_router(Arc::clone(&ctx));
for content in ["first global request", "second global request"] {
let response = app
.clone()
.oneshot(chat_request_with_content(
content,
None,
Some(8),
Some("global-session"),
))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
}
let metrics = ctx.metrics.render();
assert!(
metrics.contains(
r#"sgl_router_policy_decisions_total{policy="session_aware",reason="assigned"} 1"#
),
"the first global-preserve request must establish an assignment: {metrics}"
);
assert!(
metrics.contains(
r#"sgl_router_policy_decisions_total{policy="session_aware",reason="session_primary"} 1"#
),
"the second global-preserve request must reuse the assignment: {metrics}"
);
}
#[tokio::test]
async fn cache_winner_uses_target_uncached_work_before_prompt_length_bucket() {
let short = crate::common::mock_worker::MockWorker::start(vec![]).await;
let long = crate::common::mock_worker::MockWorker::start(vec![]).await;
let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
let mut short_bucket = bucket("p-short", BucketStage::Prefill, 10, "p-short");
short_bucket.max_extend_tokens = Some(8);
let mut long_bucket = bucket("p-long", BucketStage::Prefill, 20, "p-long");
long_bucket.min_extend_tokens = Some(9);
let bucket_config = BucketConfig {
buckets: vec![
short_bucket,
long_bucket,
bucket("d-catch-all", BucketStage::Decode, 30, "d"),
],
ttft_slo_policy: SloBucketPolicy::Disabled,
tps_slo_policy: SloBucketPolicy::Disabled,
};
let index = FakePrefixIndex::matched(short.url.clone());
let prefix_index: Arc<dyn PrefixIndex> = index.clone();
let ctx = build_cache_ctx(
vec![
worker_spec("p-short", short.url.clone(), WorkerMode::Prefill),
worker_spec("p-long", long.url.clone(), WorkerMode::Prefill),
worker_spec("d", decode.url.clone(), WorkerMode::Decode),
],
bucket_config,
prefix_index,
);
let content = "cached-prefix ".repeat(128);
let response = build_router(ctx)
.oneshot(chat_request_with_content(&content, None, Some(8), None))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
wait_for_prefill(&short).await;
assert!(
long.captured.lock().unwrap().last_body.is_none(),
"a cache winner with small target-specific uncached work must not be replaced by the full-length Bucket"
);
assert_eq!(
index.calls.load(Ordering::Relaxed),
1,
"the async Indexer query must run once at ingress, not once per Bucket"
);
}
#[tokio::test]
async fn cache_candidate_bucket_binding_happens_before_candidate_limit() {
let best = crate::common::mock_worker::MockWorker::start(vec![]).await;
let lower_ranked = crate::common::mock_worker::MockWorker::start(vec![]).await;
let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
let mut best_bucket = bucket("p-best", BucketStage::Prefill, 10, "p-best");
best_bucket.min_extend_tokens = Some(32);
let mut lower_ranked_bucket = bucket("p-lower", BucketStage::Prefill, 20, "p-lower");
lower_ranked_bucket.min_extend_tokens = Some(32);
let bucket_config = BucketConfig {
buckets: vec![
best_bucket,
lower_ranked_bucket,
bucket("d-catch-all", BucketStage::Decode, 30, "d"),
],
ttft_slo_policy: SloBucketPolicy::Disabled,
tps_slo_policy: SloBucketPolicy::Disabled,
};
let index: Arc<dyn PrefixIndex> =
TwoPrefixIndex::new(best.url.clone(), lower_ranked.url.clone());
let ctx = build_cache_ctx_with_affinity(
vec![
worker_spec("p-best", best.url.clone(), WorkerMode::Prefill),
worker_spec("p-lower", lower_ranked.url.clone(), WorkerMode::Prefill),
worker_spec("d", decode.url.clone(), WorkerMode::Decode),
],
bucket_config,
index,
AffinityConfig {
cache_candidate_min_workers: 1,
cache_candidate_ratio: 0.0,
cache_candidate_max_workers: 1,
..AffinityConfig::default()
},
);
let content = "cached bucket candidate ".repeat(256);
let response = build_router(Arc::clone(&ctx))
.oneshot(chat_request_with_content(&content, None, Some(8), None))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
wait_for_prefill(&lower_ranked).await;
assert!(
best.captured.lock().unwrap().last_body.is_none(),
"the top Indexer hit is Bucket-incompatible and must not consume K=1"
);
assert!(
ctx.metrics.render().contains(
r#"sgl_router_policy_decisions_total{policy="cache_aware",reason="cache_candidate"} 1"#
),
"the compatible lower-ranked cache holder must remain a cache candidate"
);
}
#[tokio::test]
async fn cache_no_signal_restarts_normal_prompt_length_bucket_fallback() {
let short = crate::common::mock_worker::MockWorker::start(vec![]).await;
let long = crate::common::mock_worker::MockWorker::start(vec![]).await;
let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
let mut short_bucket = bucket("p-short", BucketStage::Prefill, 10, "p-short");
short_bucket.max_extend_tokens = Some(8);
let mut long_bucket = bucket("p-long", BucketStage::Prefill, 20, "p-long");
long_bucket.min_extend_tokens = Some(9);
let bucket_config = BucketConfig {
buckets: vec![
short_bucket,
long_bucket,
bucket("d-catch-all", BucketStage::Decode, 30, "d"),
],
ttft_slo_policy: SloBucketPolicy::Disabled,
tps_slo_policy: SloBucketPolicy::Disabled,
};
let index = FakePrefixIndex::no_signal();
let prefix_index: Arc<dyn PrefixIndex> = index.clone();
let ctx = build_cache_ctx(
vec![
worker_spec("p-short", short.url.clone(), WorkerMode::Prefill),
worker_spec("p-long", long.url.clone(), WorkerMode::Prefill),
worker_spec("d", decode.url.clone(), WorkerMode::Decode),
],
bucket_config,
prefix_index,
);
let content = "uncached-prompt ".repeat(128);
let response = build_router(ctx)
.oneshot(chat_request_with_content(&content, None, Some(8), None))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
wait_for_prefill(&long).await;
assert!(
short.captured.lock().unwrap().last_body.is_none(),
"without a cache winner the request must restart the normal full-input Bucket path"
);
assert_eq!(index.calls.load(Ordering::Relaxed), 1);
}
experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs
+4 -49
View File
@@ -12,20 +12,11 @@
//! doesn't render tool schemas, so its ids would diverge from the engine).
//! * A request with multimodal (array) content → `input_ids` omitted (a text
//! tokenizer can't represent image content).
//!
//! The model id contains `deepseek-v4` so the tokenizer registry auto-attaches
//! the built-in V4 chat encoder — the engine-equivalent path — without a
//! template fixture.
use axum::body::Body;
use axum::http::{Request, StatusCode};
use serde_json::{json, Value};
use sgl_router::config::{
ActiveLoadConfig, CacheAwareConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig,
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::engine_load::EngineLoadTable;
use sgl_router::policies::factory::build_registry;
use sgl_router::policies::kv_events::{BlockSizeOracle, HashTree};
use sgl_router::proxy::Proxy;
@@ -37,36 +28,9 @@ use std::sync::Arc;
use std::time::Duration;
use tower::ServiceExt;
use crate::common::cache_aware_fixture::{config, MODEL};
use crate::common::mock_worker::MockWorker;
const MODEL: &str = "deepseek-v4-tiny";
fn config() -> Config {
Config {
server: ServerConfig {
host: "0".into(),
port: 0,
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
id: MODEL.into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::CacheAwareZmq,
circuit_breaker: None,
cache_aware: Some(CacheAwareConfig::default()),
sticky: None,
affinity: None,
fused: None,
eligibility: None,
},
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
}
}
fn build_ctx(url: String) -> Arc<AppContext> {
let cfg = config();
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
@@ -82,18 +46,9 @@ fn build_ctx(url: String) -> Arc<AppContext> {
model_ids: vec![ModelId(MODEL.into())],
bootstrap_port: None,
});
// Use the real loaded tokenizers (not the empty-registry test default) so
// the cache-aware policy can tokenize at ingress.
let policies = Arc::new(
build_registry(
&cfg,
Arc::new(HashTree::new()),
Arc::clone(&tokenizers),
BlockSizeOracle::new(),
EngineLoadTable::new(),
)
.unwrap(),
);
// Use the configured tokenizer so the chat path can emit input_ids.
let policies =
Arc::new(build_registry(&cfg, Arc::new(HashTree::new()), BlockSizeOracle::new()).unwrap());
let proxy = Arc::new(Proxy::new(Duration::from_secs(5)).unwrap());
Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies))
}
experimental/sgl-router/tests/proxy/chat_routing.rs
+2
View File
@@ -34,6 +34,8 @@ fn config_for(_worker_url: &str) -> Config {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
decode_policy: Default::default(),
bucket_config: None,
circuit_breaker: None,
cache_aware: None,
sticky: None,
experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs
+5 -3
View File
@@ -13,7 +13,7 @@ use sgl_router::config::{
pub const MODEL: &str = "deepseek-v4-tiny";
/// A single-model `cache_aware_zmq` router. Discovery is a placeholder because
/// A single-model native `cache_aware` router. Discovery is a placeholder because
/// every caller installs its own `WorkerRegistry`.
pub fn config() -> Config {
Config {
@@ -25,11 +25,13 @@ pub fn config() -> Config {
model: ModelConfig {
id: MODEL.into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::CacheAwareZmq,
policy: PolicyKind::CacheAware,
decode_policy: Default::default(),
bucket_config: None,
circuit_breaker: None,
cache_aware: Some(CacheAwareConfig::default()),
sticky: None,
affinity: None,
sticky: None,
fused: None,
eligibility: None,
},
experimental/sgl-router/tests/proxy/external_indexer_routing.rs
+18 -12
View File
@@ -16,8 +16,8 @@ use sgl_kv_indexer::pb::{
use sgl_kv_indexer::{
server_builder, GrpcPrefixIndex, InMemoryKvIndexerBackend, KvIndexerService, PrefixIndexConfig,
};
use sgl_router::config::{AffinityConfig, CachePrefixProvider, PolicyKind};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::engine_load::EngineLoadTable;
use sgl_router::policies::factory::build_registry;
use sgl_router::policies::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree};
use sgl_router::policies::request_tokens_for;
@@ -36,7 +36,20 @@ use crate::common::mock_worker::MockWorker;
async fn external_indexer_routes_to_the_cached_worker() {
let cached = MockWorker::start(vec![]).await;
let uncached = MockWorker::start(vec![]).await;
let cfg = config();
let mut cfg = config();
cfg.model.policy = PolicyKind::CacheAware;
cfg.model
.cache_aware
.as_mut()
.expect("fixture includes cache-aware configuration")
.prefix_provider = CachePrefixProvider::Indexer;
cfg.model.affinity = Some(AffinityConfig {
cache_affinity_min_matched_tokens: Some(0),
cache_candidate_min_workers: 1,
cache_candidate_ratio: 1.0,
cache_candidate_max_workers: 1,
..Default::default()
});
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
let body = json!({
"model": MODEL,
@@ -68,6 +81,7 @@ async fn external_indexer_routes_to_the_cached_worker() {
hashes: hashes.clone(),
component_masks: Vec::new(),
block_sizes: Vec::new(),
parent_block_hash: None,
}],
worker_address: cached.url.clone(),
cache_spec: None,
@@ -89,16 +103,8 @@ async fn external_indexer_routes_to_the_cached_worker() {
}
let oracle = BlockSizeOracle::new();
oracle.try_set(1).unwrap();
let policies = Arc::new(
build_registry(
&cfg,
Arc::new(HashTree::new()),
Arc::clone(&tokenizers),
Arc::clone(&oracle),
EngineLoadTable::new(),
)
.unwrap(),
);
let policies =
Arc::new(build_registry(&cfg, Arc::new(HashTree::new()), Arc::clone(&oracle)).unwrap());
let mut ctx = AppContext::new(
cfg,
tokenizers,
experimental/sgl-router/tests/proxy/failover.rs
+2
View File
@@ -35,6 +35,8 @@ async fn failover_when_one_worker_dies() {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
decode_policy: Default::default(),
bucket_config: None,
circuit_breaker: Some(CircuitBreakerConfig {
threshold: std::num::NonZeroU32::new(1).unwrap(), // open after first failure
cool_down_secs: 30,
experimental/sgl-router/tests/proxy/graceful_shutdown.rs
+25 -23
View File
@@ -14,7 +14,7 @@
//! 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 futures::future::join_all;
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
@@ -44,6 +44,8 @@ fn build_ctx_with_worker(worker_url: &str) -> Arc<AppContext> {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
decode_policy: Default::default(),
bucket_config: None,
circuit_breaker: None,
cache_aware: None,
sticky: None,
@@ -130,12 +132,11 @@ async fn shutdown_drains_100_inflight_streaming_chat_completions() {
}))
.unwrap();
let mut handles = Vec::with_capacity(N);
for i in 0..N {
let responses = join_all((0..N).map(|i| {
let c = client.clone();
let u = url.clone();
let b = body.clone();
handles.push(tokio::spawn(async move {
async move {
let resp = c
.post(&u)
.header("content-type", "application/json")
@@ -146,33 +147,34 @@ async fn shutdown_drains_100_inflight_streaming_chat_completions() {
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)
}));
}
Ok::<_, String>((i, resp))
}
}))
.await;
let responses: Vec<_> = responses
.into_iter()
.collect::<Result<_, _>>()
.expect("every client received response headers before shutdown");
// 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.
// 4. Each response header confirms that its request is in flight. Trigger
// shutdown only after the full cohort connects, then verify that Axum
// drains all 100 existing streams.
let started = Instant::now();
shutdown_tx.send(()).unwrap();
// 6. Every in-flight request must complete with a `[DONE]` terminator
// 5. 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
for result in join_all(responses.into_iter().map(|(i, response)| async move {
response
.bytes()
.await
.expect("client task panicked")
.expect("client completed");
.map_err(|e| format!("client {i} body: {e}"))
}))
.await
{
let result = result.expect("client body completed");
bytes_total += result.len();
let body_str = String::from_utf8_lossy(&result);
if body_str.contains("data: [DONE]") {
experimental/sgl-router/tests/proxy/header_forwarding.rs
+2
View File
@@ -31,6 +31,8 @@ async fn forwards_whitelisted_headers_strips_others() {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
decode_policy: Default::default(),
bucket_config: None,
circuit_breaker: None,
cache_aware: None,
sticky: None,
experimental/sgl-router/tests/proxy/main.rs
+2
View File
@@ -10,6 +10,7 @@
mod common;
mod bucket_routing;
mod cache_aware_input_ids;
mod chat_routing;
mod external_indexer_routing;
@@ -18,6 +19,7 @@ mod graceful_shutdown;
mod header_forwarding;
mod pd_bootstrap_injection;
mod pd_pool_isolation;
mod radix_tree_routing;
mod roundrobin_input_ids;
mod shared_prefill_admission;
mod sticky_input_ids;
experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs
+41 -2
View File
@@ -32,7 +32,7 @@ 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 std::time::{Duration, Instant};
use tower::ServiceExt;
fn config() -> Config {
@@ -46,6 +46,8 @@ fn config() -> Config {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
decode_policy: Default::default(),
bucket_config: None,
circuit_breaker: None,
cache_aware: None,
sticky: None,
@@ -98,7 +100,7 @@ async fn await_captured_body(
timeout: Duration,
label: &str,
) -> Bytes {
let start = std::time::Instant::now();
let start = Instant::now();
loop {
// Release the `std::sync::Mutex` guard before the sleep.await
// (clippy: await_holding_lock).
@@ -189,6 +191,43 @@ async fn pd_mode_chat_injects_bootstrap_fields_into_both_bodies() {
assert_eq!(bootstrap_port(&dj), Some(8997));
}
#[tokio::test]
async fn round_robin_pd_prefill_does_not_track_dispatch_timestamps() {
let prefill =
crate::common::mock_worker::MockWorker::start_hanging(Duration::from_millis(200)).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 prefill_worker = ctx
.registry
.workers_for(&ModelId("tiny".into()))
.into_iter()
.find(|worker| worker.id.0 == "p1")
.expect("prefill worker is registered");
let cutoff = Instant::now() - Duration::from_secs(1);
let request = tokio::spawn(build_router(Arc::clone(&ctx)).oneshot(chat_request()));
await_captured_body(&prefill, Duration::from_secs(2), "prefill").await;
assert_eq!(prefill_worker.active_load(), 1);
assert_eq!(prefill_worker.slots_acquired_since(cutoff), 0);
assert_eq!(request.await.unwrap().unwrap().status(), StatusCode::OK);
}
/// 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.
experimental/sgl-router/tests/proxy/pd_pool_isolation.rs
+13 -24
View File
@@ -45,6 +45,8 @@ fn config() -> Config {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
decode_policy: Default::default(),
bucket_config: None,
circuit_breaker: None,
cache_aware: None,
sticky: None,
@@ -205,31 +207,19 @@ async fn pd_mode_chat_dispatch_fans_to_both_prefill_and_decode() {
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).
/// PD-mode chat request carries an `x-sgl-decode-url` header for the final
/// Decode decision. Step 1 defaults to Decode P2; the header remains an
/// observability contract regardless of which Decode policy produced it.
#[tokio::test]
async fn pd_mode_chat_dispatch_sets_decode_affinity_header() {
async fn pd_mode_chat_dispatch_sets_final_decode_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).
// MockWorker URLs all bind to `127.0.0.1`; this test deliberately does
// not assert a host relation. It pins only the HTTP wiring: the final D
// selected by the role-local policy is reflected on the P request.
let ctx = build_ctx(vec![
WorkerSpec {
id: WorkerId("p1".into()),
@@ -268,9 +258,8 @@ async fn pd_mode_chat_dispatch_sets_decode_affinity_header() {
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.
// Every request that hit a prefill mock MUST carry the final-decode
// header. The value MUST be one of the two registered Decode URLs.
let decode_urls: HashSet<String> = [decode_a.url.clone(), decode_b.url.clone()]
.into_iter()
.collect();
@@ -346,9 +335,9 @@ async fn pd_mode_prefill_only_returns_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
/// can observe final Decode selection 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`.
/// behavior asserted by `pd_mode_chat_dispatch_sets_final_decode_header`.
#[tokio::test]
async fn pd_mode_chat_response_carries_decode_affinity_header() {
use std::collections::HashSet;
experimental/sgl-router/tests/proxy/radix_tree_routing.rs
+94
View File
@@ -0,0 +1,94 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use std::sync::Arc;
use std::time::Duration;
use axum::body::Body;
use axum::http::{Request, StatusCode};
use serde_json::json;
use sgl_router::config::{AffinityConfig, CachePrefixProvider, PolicyKind};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::factory::build_registry;
use sgl_router::policies::kv_events::{
compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId,
};
use sgl_router::policies::prefix_provider::RadixTreePrefixProvider;
use sgl_router::policies::request_tokens_for;
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 tower::ServiceExt;
use crate::common::cache_aware_fixture::{config, MODEL};
use crate::common::mock_worker::MockWorker;
#[tokio::test]
async fn radix_tree_routes_cache_aware_request_to_cached_worker() {
let cached = MockWorker::start(vec![]).await;
let uncached = MockWorker::start(vec![]).await;
let mut cfg = config();
cfg.model.policy = PolicyKind::CacheAware;
cfg.model.cache_aware.as_mut().unwrap().prefix_provider = CachePrefixProvider::RadixTree;
cfg.model.affinity = Some(AffinityConfig {
cache_affinity_min_matched_tokens: Some(0),
cache_candidate_min_workers: 1,
cache_candidate_ratio: 1.0,
cache_candidate_max_workers: 1,
..Default::default()
});
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
let body = json!({
"model": MODEL,
"messages": [{"role": "user", "content": "local radix cache hit"}],
});
let tokens = request_tokens_for(&tokenizers, &ModelId(MODEL.into()), &body)
.expect("test prompt tokenizes");
let hashes = compute_block_hashes(&tokens.ids, 1);
assert!(!hashes.is_empty());
let tree = Arc::new(HashTree::new());
tree.insert(&KvWorkerId::new(cached.url.clone(), 0), None, &hashes);
let registry = Arc::new(WorkerRegistry::default());
for url in [&cached.url, &uncached.url] {
registry
.add(WorkerSpec {
id: WorkerId(url.clone()),
url: url.clone(),
mode: WorkerMode::Plain,
model_ids: vec![ModelId(MODEL.into())],
bootstrap_port: None,
})
.unwrap();
}
let oracle = BlockSizeOracle::new();
oracle.try_set(1).unwrap();
let policies = Arc::new(build_registry(&cfg, Arc::clone(&tree), Arc::clone(&oracle)).unwrap());
let mut ctx = AppContext::new(
cfg,
tokenizers,
Arc::new(Proxy::new(Duration::from_secs(5)).unwrap()),
registry,
policies,
);
ctx.radix_tree_prefix_provider = Some(RadixTreePrefixProvider::new(tree, Arc::clone(&oracle)));
ctx.block_size_oracle = oracle;
let response = build_router(Arc::new(ctx))
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header("content-type", "application/json")
.body(Body::from(serde_json::to_vec(&body).unwrap()))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
assert!(cached.captured.lock().unwrap().last_body.is_some());
assert!(uncached.captured.lock().unwrap().last_body.is_none());
}
experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs
+2
View File
@@ -42,6 +42,8 @@ fn config() -> Config {
id: MODEL.into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
decode_policy: Default::default(),
bucket_config: None,
circuit_breaker: None,
cache_aware: None,
sticky: None,
experimental/sgl-router/tests/proxy/shared_prefill_admission.rs
+38 -8
View File
@@ -11,7 +11,7 @@ use sgl_router::config::{
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::engine_load::LoadStat;
use sgl_router::policies::engine_load::{LoadStat, NativeCacheRankLoad};
use sgl_router::policies::{
CacheCandidate, CacheCandidateProposal, Policy, PolicyRegistry, PrefillProposal, ProposalKind,
SelectionContext, SelectionProposal,
@@ -128,6 +128,7 @@ impl Policy for CacheCandidatesPolicy {
max_pending_prefill_tokens: None,
}],
cache_switch_margin_tokens: 0,
..Default::default()
}))
}
@@ -147,6 +148,8 @@ fn config(policy: PolicyKind) -> Config {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy,
decode_policy: Default::default(),
bucket_config: None,
circuit_breaker: None,
cache_aware: None,
sticky: None,
@@ -286,17 +289,30 @@ async fn chat_commits_the_admitted_prefill_backup() {
})
})
.await;
let now = Instant::now();
let native_load = |total_prefill_uncached_tokens, total_prefill_busy_us| LoadStat {
num_running_reqs: 1,
num_waiting_reqs: 0,
num_tokens: 100,
max_total_num_tokens: 100,
native_cache: Some(NativeCacheRankLoad {
num_waiting_uncached_tokens: 0,
num_total_tokens: 100,
max_running_requests: 16,
total_prefill_uncached_tokens,
total_prefill_busy_us,
}),
};
fixture.ctx.engine_load.set(
&fixture.workers[0].url,
0,
LoadStat {
num_running_reqs: 1,
num_waiting_reqs: 0,
num_tokens: 100,
max_total_num_tokens: 100,
},
Instant::now(),
native_load(1, 1),
now - Duration::from_secs(1),
);
fixture
.ctx
.engine_load
.set(&fixture.workers[0].url, 0, native_load(2, 2), now);
assert_eq!(send_chat(&fixture.ctx).await, StatusCode::OK);
assert!(fixture.backends[0]
@@ -339,6 +355,13 @@ async fn capacity_exhaustion_does_not_return_503() {
num_waiting_reqs: 0,
num_tokens: 100,
max_total_num_tokens: 100,
native_cache: Some(NativeCacheRankLoad {
num_waiting_uncached_tokens: 0,
num_total_tokens: 100,
max_running_requests: 16,
total_prefill_uncached_tokens: 1,
total_prefill_busy_us: 1,
}),
},
Instant::now(),
);
@@ -415,6 +438,13 @@ async fn chat_records_cache_candidates_exhausted() {
num_waiting_reqs: 0,
num_tokens: 100,
max_total_num_tokens: 100,
native_cache: Some(NativeCacheRankLoad {
num_waiting_uncached_tokens: 0,
num_total_tokens: 100,
max_running_requests: 16,
total_prefill_uncached_tokens: 1,
total_prefill_busy_us: 1,
}),
},
Instant::now(),
);
experimental/sgl-router/tests/proxy/sticky_input_ids.rs
+2
View File
@@ -54,6 +54,8 @@ fn config() -> Config {
id: MODEL.into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::Sticky,
decode_policy: Default::default(),
bucket_config: None,
circuit_breaker: None,
cache_aware: None,
// Push eviction far out so the background sweeper never fires
experimental/sgl-router/tests/proxy/sticky_routing.rs
+2
View File
@@ -42,6 +42,8 @@ fn build_sticky_ctx(header_name: &str, worker_urls: &[String]) -> Arc<AppContext
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::Sticky,
decode_policy: Default::default(),
bucket_config: None,
circuit_breaker: None,
cache_aware: None,
sticky: Some(StickyConfig {
experimental/sgl-router/tests/proxy/timeout.rs
+2
View File
@@ -38,6 +38,8 @@ fn config(_worker_url: &str) -> Config {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
decode_policy: Default::default(),
bucket_config: None,
circuit_breaker: None,
cache_aware: None,
sticky: None,