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

[sgl-router] refactor - SLO ordering for bucket selection (#40292)

Browse Source
This commit is contained in:
Kan Wu
2026-09-22 11:16:10 +08:00
committed by GitHub
parent 27f796ca6c
commit 59a723ef1e
8 changed files with 422 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
experimental/sgl-router/POLICY_DESIGN.md
+38 -13
View File
@@ -7,8 +7,9 @@ listed at the end.
## Principles
1. **Order compatible buckets by token length first.** `BucketResolver` returns
all matching buckets, smallest capacity first, without inspecting workers or policies.
1. **Filter buckets by token length, then order preferences.** `BucketResolver`
returns all compatible buckets, applying optional SLO preferences before
capacity/rank/ID ordering, without inspecting workers or policies.
2. **The bucket owns plain versus PD engine selection.** `Bucket::pick_engines`
calls its one plain group or both prefill and decode groups, returning a
complete selection. Both PD engines come from that same bucket.
@@ -79,7 +80,7 @@ tier executor, or separate selection framework is required.
## 2. Responsibilities and request flow
`BucketResolver::resolve(input_tokens, expected_peak_tokens)` returns an ordered
`BucketResolver::resolve(input_tokens, expected_peak_tokens, ttft_ms, tokens_per_second)` returns an ordered
list of compatible bucket references (possibly empty), or an invalid-signal error.
It does not receive a stage or a load view, resolve live engines, or invoke policies.
The handler iterates this list until a bucket supplies the complete engine selection.
@@ -153,9 +154,10 @@ the existing body-size estimate when tokenization is unavailable.
2. Keep buckets whose inclusive input-token range contains the input length.
3. Check the bucket context capacity against input plus requested output when
known, or against input length when the output budget is unknown.
4. Sort by ascending input capacity (the lesser of the input upper bound and
context capacity). Unbounded capacities sort last. Break ties by ascending
bucket rank, then ID, and return the entire ordered list.
4. Apply enabled SLO preferences to complete buckets. Count unmet preferences
equally, then sort by ascending input capacity (the lesser of the input upper
bound and context capacity), rank, and ID. Unbounded capacities sort last.
Return the entire ordered list, retaining nonpreferred buckets for fallback.
5. The handler calls each bucket's `pick_engines` until one supplies its complete
selection. A failed PD attempt never contributes an engine to a later pair.
@@ -172,7 +174,29 @@ bucket, both groups share this one request-length decision. Policy fallback on
a cache/affinity miss stays within that group's candidates. There is no second
pass with relaxed admission and no post-policy substitution.
SLO ordering, global session modes, and sticky policies are follow-ups. Their
### Optional SLO ordering
`Bucket` has optional `ttft_ms` and `tokens_per_second` estimates. The resolver
has independent `ttft_slo` and `tps_slo` preferences: `Disabled` (default),
`SloFirst` (matching first), and `BestEffort` (nonmatching first). The handler
parses `x-sgl-ttft-slo-ms` and `x-sgl-tps-slo` only when their preference is enabled;
invalid enabled headers return 400 before dispatch. Disabled headers are ignored.
SLO targets belong to bucket resolution, not the engine policy's `PickRequest`.
Absent targets are neutral. A bucket matches TTFT when its positive estimate is
at most the target, and throughput when its finite positive estimate is at least
the target. Missing or invalid estimates do not match a supplied target. Enabled
TTFT targets must be positive; throughput targets must be finite and positive.
Each unmet preference adds one ordering penalty. With both preferences set to
`SloFirst`, a bucket matching both comes before one matching either, followed by
buckets matching neither. Capacity/rank/ID breaks ties within these tiers. The
same logic applies to plain and PD buckets, and a PD bucket always supplies both
engines. TTFT and throughput preferences never independently resolve P/D groups.
Length constraints are applied first and admission rejection still advances to
the next complete bucket, including a bucket outside the preferred SLO tier.
Global session modes and sticky policies are follow-ups. Their
integration must preserve bucket-first selection and the same-bucket PD rule.
Cross-bucket affinity probing is not part of this interface. Session/routing
keys still pass through `PickRequest` for policies operating inside the selected
@@ -490,8 +514,8 @@ do not accept and ignore them.
- Keep `load_based` as the CLI name for `LeastLoadPolicy`.
- Preserve explicit engine membership and context constraints. Bucket-level
ranges and ordering replace independent per-stage selection.
Restore existing SLO behavior in the separate SLO PR before serving switchover;
legacy routing continues to support SLOs during this skeleton-only phase.
SLO preferences order whole buckets; deployment configuration remains explicit
until the production configuration factory and serving switchover are ready.
- Preserve cache-provider selection, endpoint validation, query timeout and
concurrency limits, and unavailable-backend fallback.
- Preserve cache thresholds and tuning: the 1,024-token default minimum hit,
@@ -558,7 +582,8 @@ This PR adds the side-by-side interfaces in `src/buckets_reorg.rs` and
Implemented here:
- `BucketResolver::resolve` returns all length-compatible buckets in capacity/rank/ID order.
- `BucketResolver::resolve` returns all length-compatible buckets in optional
SLO-preference tiers, with capacity/rank/ID order within each tier.
- `Bucket::pick_engines` owns plain/PD orchestration and stage-specific policy
requests; `BucketRequest` carries prepared facts and `BucketPick` retains picks.
- `Bucket` owns input limits, context capacity, rank, and plain-or-PD groups.
@@ -576,6 +601,8 @@ Implemented here:
dispatches only after one complete selection. Exhaustion retains admission reasons.
- `AppContext::chat_routing` configures legacy versus reorg routing on the same
endpoint and carries the reorg model-resolver map.
- Optional bucket TTFT/throughput estimates and preferences, with enabled-header
parsing and whole-bucket fallback in the shared chat route.
- `CacheAwarePolicy` reads local radix-tree or remote indexer prefixes, intersects
exact worker URLs with the current group, applies hit thresholds and candidate
bounds, and preserves the soft queue gate, saturation pin and pressure guard.
@@ -594,13 +621,11 @@ Implemented here:
The caller owns expiry and sweeper lifecycle. A binding may remain after a
later PD group fails, because it records placement rather than dispatch.
Follow-up work includes bucket SLO ordering, remaining selection policies,
and production configuration.
Follow-up work includes remaining selection policies and production configuration.
Not yet implemented in the reorg path:
- Other concrete selection policies.
- SLO estimates, targets, and bucket preference ordering.
- CLI/configuration parsing, validation, and model-specific construction.
The YAML above is illustrative; reorg resolvers are installed in code.
- Global session modes and sticky routing-key affinity.
experimental/sgl-router/src/buckets_reorg.rs
+63 -7
View File
@@ -1,11 +1,11 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Order buckets by request length; each bucket owns the groups used to pick engines.
//! Order length-compatible buckets by optional SLO preferences, then capacity and rank.
//!
//! ```text
//! BucketResolver (one model's buckets)
//! -> Bucket (token limits, context capacity, rank)
//! -> Bucket (token limits, context capacity, rank, SLO estimates)
//! -> Plain: one EngineGroup
//! -> PD: prefill + decode EngineGroups
//! -> each EngineGroup: worker membership + its own Policy
@@ -18,7 +18,7 @@
//! The handler tries buckets in order, advancing on missing candidates or admission
//! rejection. Both P/D picks must succeed in the same bucket before dispatch.
//! [`WorkerRegistry`] owns live workers; groups reference their IDs. Policies own
//! their load/KV/affinity dependencies and share observations within each attempt.
//! their load/KV/affinity dependencies and pass selected observations to admission.
use std::collections::HashSet;
use std::sync::Arc;
@@ -120,12 +120,15 @@ pub struct BucketPick {
#[derive(Debug)]
pub struct Bucket {
pub id: String,
/// Break ties between equally sized buckets; lower ranks win.
/// Break ties within an SLO tier between equally sized buckets; lower ranks win.
pub rank: u32,
/// Inclusive input-token range used to choose the bucket.
pub limits: TokenLimits,
/// Full sequence capacity, checked against the expected peak when known.
pub max_context_tokens: Option<u64>,
/// Optional service estimates used only for bucket ordering.
pub ttft_ms: Option<u64>,
pub tokens_per_second: Option<f64>,
pub groups: BucketGroups,
}
@@ -136,6 +139,8 @@ impl Bucket {
rank: 0,
limits: TokenLimits::default(),
max_context_tokens: None,
ttft_ms: None,
tokens_per_second: None,
groups,
}
}
@@ -225,10 +230,30 @@ impl Bucket {
}
}
/// Soft preference; nonpreferred buckets remain available for fallback.
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub enum SloPreference {
#[default]
Disabled,
SloFirst,
BestEffort,
}
impl SloPreference {
fn penalty(self, matches: Option<bool>) -> u8 {
match (self, matches) {
(Self::SloFirst, Some(false)) | (Self::BestEffort, Some(true)) => 1,
_ => 0,
}
}
}
/// Model-specific bucket configuration. Selection does not inspect engine state.
#[derive(Debug, Default)]
pub struct BucketResolver {
pub buckets: Vec<Bucket>,
pub ttft_slo: SloPreference,
pub tps_slo: SloPreference,
}
impl BucketResolver {
@@ -237,27 +262,58 @@ impl BucketResolver {
for bucket in &buckets {
bucket.validate()?;
}
Ok(Self { buckets })
Ok(Self {
buckets,
..Self::default()
})
}
/// Return all length-compatible buckets, ordered by input capacity, rank, and ID.
/// Return all length-compatible buckets, ordered by unmet SLO preferences,
/// then input capacity, rank, and ID. Both preferences have equal weight.
/// The caller tries their groups in order until a complete engine selection succeeds.
pub fn resolve(
&self,
input_tokens: u64,
expected_peak_tokens: Option<u64>,
ttft_ms: Option<u64>,
tokens_per_second: Option<f64>,
) -> Result<Vec<&Bucket>, PickError> {
if expected_peak_tokens.is_some_and(|tokens| tokens < input_tokens) {
return Err(PickError::InvalidSignal(
"expected peak tokens are below input length".into(),
));
}
if self.ttft_slo != SloPreference::Disabled && ttft_ms == Some(0) {
return Err(PickError::InvalidSignal(
"requested TTFT must be positive".into(),
));
}
if self.tps_slo != SloPreference::Disabled
&& tokens_per_second.is_some_and(|tps| !tps.is_finite() || tps <= 0.0)
{
return Err(PickError::InvalidSignal(
"requested tokens per second must be finite and positive".into(),
));
}
let mut buckets: Vec<_> = self
.buckets
.iter()
.filter(|bucket| bucket.fits(input_tokens, expected_peak_tokens))
.collect();
buckets.sort_by_key(|bucket| (bucket.input_capacity(), bucket.rank, &bucket.id));
buckets.sort_by_key(|bucket| {
let ttft_matches = ttft_ms.map(|target| {
bucket
.ttft_ms
.is_some_and(|estimate| estimate > 0 && estimate <= target)
});
let tps_matches = tokens_per_second.map(|target| {
bucket.tokens_per_second.is_some_and(|estimate| {
estimate.is_finite() && estimate > 0.0 && estimate >= target
})
});
let penalty = self.ttft_slo.penalty(ttft_matches) + self.tps_slo.penalty(tps_matches);
(penalty, bucket.input_capacity(), bucket.rank, &bucket.id)
});
Ok(buckets)
}
}
experimental/sgl-router/src/server/routes/chat/reorg.rs
+21 -3
View File
@@ -2,9 +2,12 @@
// SPDX-License-Identifier: Apache-2.0
use super::forward::{forward_chat_request, SelectedWorkers};
use super::nonempty_header;
use super::preparation::{parse_routing_fields, PreparedChatRequest};
use crate::buckets_reorg::{BucketRequest, BucketResolver};
use super::{
nonempty_header, parse_optional_positive_f64_header, parse_optional_positive_u64_header,
X_SGL_TPS_SLO, X_SGL_TTFT_SLO_MS,
};
use crate::buckets_reorg::{BucketRequest, BucketResolver, SloPreference};
use crate::discovery::ModelId;
use crate::policies_reorg::{PickError, Stage};
use crate::server::app_context::AppContext;
@@ -45,8 +48,23 @@ pub(super) async fn chat_completions(
})
.transpose()?;
let ttft_ms = if resolver.ttft_slo != SloPreference::Disabled {
parse_optional_positive_u64_header(&headers, &X_SGL_TTFT_SLO_MS, "TTFT SLO")?
} else {
None
};
let tokens_per_second = if resolver.tps_slo != SloPreference::Disabled {
parse_optional_positive_f64_header(&headers, &X_SGL_TPS_SLO, "TPS SLO")?
} else {
None
};
let buckets = resolver
.resolve(input_tokens, expected_peak_tokens)
.resolve(
input_tokens,
expected_peak_tokens,
ttft_ms,
tokens_per_second,
)
.map_err(|error| selection_error(error, &request.model, None))?;
if buckets.is_empty() {
return Err(selection_error(
experimental/sgl-router/tests/component/main.rs
+1
View File
@@ -16,5 +16,6 @@ mod policies_reorg_cache_aware;
mod policies_reorg_load;
mod policies_reorg_power_of_two;
mod policies_reorg_session_aware;
mod policies_reorg_slo;
mod tokenizer;
mod workers;
experimental/sgl-router/tests/component/policies_reorg.rs
+29 -14
View File
@@ -186,17 +186,20 @@ fn resolve_orders_all_length_fits_by_capacity_rank_and_id() {
.unwrap();
assert_eq!(
resolver
.resolve(10, None)
.resolve(10, None, None, None)
.unwrap()
.iter()
.map(|bucket| bucket.id.as_str())
.collect::<Vec<_>>(),
["a", "z", "later", "catch-all"]
);
assert_eq!(resolver.resolve(11, None).unwrap()[0].id, "min");
assert_eq!(resolver.resolve(15, None).unwrap()[0].id, "min");
assert_eq!(resolver.resolve(20, None).unwrap()[0].id, "a");
assert_eq!(resolver.resolve(21, None).unwrap()[0].id, "catch-all");
assert_eq!(resolver.resolve(11, None, None, None).unwrap()[0].id, "min");
assert_eq!(resolver.resolve(15, None, None, None).unwrap()[0].id, "min");
assert_eq!(resolver.resolve(20, None, None, None).unwrap()[0].id, "a");
assert_eq!(
resolver.resolve(21, None, None, None).unwrap()[0].id,
"catch-all"
);
}
#[test]
@@ -207,17 +210,29 @@ fn context_capacity_checks_peak_when_known_and_input_otherwise() {
let mut long = bucket("long", None, policy);
long.max_context_tokens = Some(30);
let resolver = BucketResolver::new(vec![long, short]).unwrap();
assert_eq!(resolver.resolve(10, None).unwrap()[0].id, "short");
assert_eq!(resolver.resolve(10, Some(20)).unwrap()[0].id, "short");
assert_eq!(resolver.resolve(10, Some(21)).unwrap()[0].id, "long");
assert!(resolver.resolve(10, Some(31)).unwrap().is_empty());
assert!(resolver.resolve(31, None).unwrap().is_empty());
assert_eq!(
resolver.resolve(10, None, None, None).unwrap()[0].id,
"short"
);
assert_eq!(
resolver.resolve(10, Some(20), None, None).unwrap()[0].id,
"short"
);
assert_eq!(
resolver.resolve(10, Some(21), None, None).unwrap()[0].id,
"long"
);
assert!(resolver
.resolve(10, Some(31), None, None)
.unwrap()
.is_empty());
assert!(resolver.resolve(31, None, None, None).unwrap().is_empty());
assert!(matches!(
resolver.resolve(10, Some(9)),
resolver.resolve(10, Some(9), None, None),
Err(PickError::InvalidSignal(_))
));
assert!(BucketResolver::default()
.resolve(1, None)
.resolve(1, None, None, None)
.unwrap()
.is_empty());
}
@@ -238,7 +253,7 @@ async fn selected_pd_bucket_owns_both_memberships_and_policies() {
},
)])
.unwrap();
let bucket = resolver.resolve(10, Some(20)).unwrap()[0];
let bucket = resolver.resolve(10, Some(20), None, None).unwrap()[0];
let request = BucketRequest {
prefix: None,
model: &model,
@@ -270,7 +285,7 @@ async fn resolver_includes_empty_groups_without_invoking_policies() {
};
let resolver =
BucketResolver::new(vec![empty, bucket("available", Some(20), policy.clone())]).unwrap();
let buckets = resolver.resolve(10, None).unwrap();
let buckets = resolver.resolve(10, None, None, None).unwrap();
assert_eq!(
buckets
.iter()
experimental/sgl-router/tests/component/policies_reorg_slo.rs
+135
View File
@@ -0,0 +1,135 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use std::sync::Arc;
use sgl_router::buckets_reorg::{Bucket, BucketGroups, BucketResolver, EngineGroup, SloPreference};
use sgl_router::policies_reorg::power_of_two::PowerOfTwoPolicy;
use sgl_router::policies_reorg::PickError;
use sgl_router::state::load_monitor::engine_reported_load::EngineReportedLoadTable;
fn bucket(id: &str, max: u64, rank: u32, ttft: Option<u64>, tps: Option<f64>) -> Bucket {
let mut bucket = Bucket::new(
id,
BucketGroups::Plain(EngineGroup::new(Arc::new(PowerOfTwoPolicy::new(
EngineReportedLoadTable::new(),
)))),
);
bucket.limits.max = Some(max);
bucket.rank = rank;
bucket.ttft_ms = ttft;
bucket.tokens_per_second = tps;
bucket
}
fn resolver() -> BucketResolver {
BucketResolver::new(vec![
bucket("both", 100, 9, Some(50), Some(100.0)),
bucket("ttft", 10, 1, Some(50), Some(10.0)),
bucket("tps", 10, 2, Some(100), Some(100.0)),
bucket("neither", 8, 0, Some(100), Some(10.0)),
bucket("missing", 10, 3, None, None),
bucket("too-short", 1, 0, Some(1), Some(1000.0)),
])
.unwrap()
}
fn ids(resolver: &BucketResolver, ttft: Option<u64>, tps: Option<f64>) -> Vec<&str> {
resolver
.resolve(5, Some(5), ttft, tps)
.unwrap()
.into_iter()
.map(|b| b.id.as_str())
.collect()
}
#[test]
fn slo_tiers_keep_length_constraints_and_capacity_rank_order() {
let mut resolver = resolver();
assert_eq!(
ids(&resolver, Some(50), Some(100.0)),
["neither", "ttft", "tps", "missing", "both"]
);
resolver.ttft_slo = SloPreference::SloFirst;
resolver.tps_slo = SloPreference::SloFirst;
assert_eq!(
ids(&resolver, Some(50), Some(100.0)),
["both", "ttft", "tps", "neither", "missing"]
);
resolver.ttft_slo = SloPreference::BestEffort;
resolver.tps_slo = SloPreference::BestEffort;
assert_eq!(
ids(&resolver, Some(50), Some(100.0)),
["neither", "missing", "ttft", "tps", "both"]
);
resolver.ttft_slo = SloPreference::SloFirst;
assert_eq!(
ids(&resolver, Some(50), Some(100.0)),
["ttft", "neither", "missing", "both", "tps"]
);
}
#[test]
fn absent_targets_are_neutral_and_each_preference_can_be_disabled() {
let mut resolver = resolver();
resolver.ttft_slo = SloPreference::SloFirst;
resolver.tps_slo = SloPreference::BestEffort;
assert_eq!(
ids(&resolver, None, None),
["neither", "ttft", "tps", "missing", "both"]
);
assert_eq!(
ids(&resolver, Some(50), None),
["ttft", "both", "neither", "tps", "missing"]
);
resolver.ttft_slo = SloPreference::Disabled;
resolver.tps_slo = SloPreference::SloFirst;
assert_eq!(
ids(&resolver, Some(1), Some(100.0)),
["tps", "both", "neither", "ttft", "missing"]
);
}
#[test]
fn invalid_enabled_targets_fail_and_invalid_estimates_do_not_match() {
let mut resolver = resolver();
resolver.ttft_slo = SloPreference::SloFirst;
resolver.tps_slo = SloPreference::SloFirst;
assert!(matches!(
resolver.resolve(5, None, Some(0), None),
Err(PickError::InvalidSignal(_))
));
for tps in [0.0, -1.0, f64::NAN, f64::INFINITY] {
assert!(matches!(
resolver.resolve(5, None, None, Some(tps)),
Err(PickError::InvalidSignal(_))
));
}
resolver.buckets[0].ttft_ms = Some(0);
resolver.buckets[0].tokens_per_second = Some(f64::INFINITY);
assert_eq!(
ids(&resolver, Some(50), Some(100.0)),
["ttft", "tps", "neither", "missing", "both"]
);
resolver.ttft_slo = SloPreference::Disabled;
resolver.tps_slo = SloPreference::Disabled;
assert!(resolver.resolve(5, None, Some(0), Some(f64::NAN)).is_ok());
}
#[test]
fn slo_preferences_never_relax_peak_capacity_or_input_range() {
let mut resolver = resolver();
resolver.ttft_slo = SloPreference::SloFirst;
resolver.tps_slo = SloPreference::SloFirst;
resolver.buckets[0].max_context_tokens = Some(6);
resolver.buckets[1].limits.min = Some(6);
let found = resolver.resolve(5, Some(7), Some(50), Some(100.0)).unwrap();
assert_eq!(
found.iter().map(|b| b.id.as_str()).collect::<Vec<_>>(),
["tps", "neither", "missing"]
);
assert!(matches!(
resolver.resolve(5, Some(4), None, None),
Err(PickError::InvalidSignal(_))
));
}
experimental/sgl-router/tests/proxy/chat_routing/reorg.rs
+1
View File
@@ -14,6 +14,7 @@ use sgl_router::server::app_context::ChatRouting;
use std::sync::Mutex;
mod session_aware;
mod slo;
type PickCall = (String, Stage, u64, Option<u64>);
experimental/sgl-router/tests/proxy/chat_routing/reorg/slo.rs
+134
View File
@@ -0,0 +1,134 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use super::*;
#[tokio::test]
async fn slo_headers_order_complete_pd_buckets_and_validate_before_dispatch() {
use sgl_router::buckets_reorg::SloPreference;
let fast_p = MockWorker::start(vec![]).await;
let slow_p = MockWorker::start(vec![]).await;
let fast_d = MockWorker::start(vec![]).await;
let slow_d = MockWorker::start(vec![]).await;
let mut ctx = Arc::try_unwrap(context(
&[
("fast-p", Stage::Prefill, &fast_p),
("slow-p", Stage::Prefill, &slow_p),
("fast-d", Stage::Decode, &fast_d),
("slow-d", Stage::Decode, &slow_d),
],
vec![],
))
.unwrap_or_else(|_| panic!("context is not shared yet"));
let mut prefill = Bucket::new(
"a-prefill",
BucketGroups::Pd {
prefill: group("fast-p", Arc::new(FirstPolicy::default())),
decode: group("slow-d", Arc::new(FirstPolicy::default())),
},
);
prefill.ttft_ms = Some(50);
prefill.tokens_per_second = Some(10.0);
let mut decode = Bucket::new(
"b-decode",
BucketGroups::Pd {
prefill: group("slow-p", Arc::new(FirstPolicy::default())),
decode: group("fast-d", Arc::new(FirstPolicy::default())),
},
);
decode.ttft_ms = Some(100);
decode.tokens_per_second = Some(100.0);
let mut resolver = BucketResolver::new(vec![decode, prefill]).unwrap();
resolver.ttft_slo = SloPreference::SloFirst;
resolver.tps_slo = SloPreference::SloFirst;
ctx.chat_routing = ChatRouting::Reorg([(ModelId("tiny".into()), resolver)].into());
let app = build_router(Arc::new(ctx));
for (ttft, tps) in [("0", "100"), ("50", "NaN"), ("50", "0"), ("abc", "100")] {
let mut req = request(body("hi"));
req.headers_mut()
.insert("x-sgl-ttft-slo-ms", ttft.parse().unwrap());
req.headers_mut()
.insert("x-sgl-tps-slo", tps.parse().unwrap());
assert_eq!(
app.clone().oneshot(req).await.unwrap().status(),
StatusCode::BAD_REQUEST
);
}
for worker in [&fast_p, &slow_p, &fast_d, &slow_d] {
assert!(worker.captured.lock().unwrap().last_body.is_none());
}
// With both targets unmet once, the existing bucket order breaks the tie.
// With TTFT unconstrained, throughput chooses the other complete bucket.
for (ttft, expected_d, expected_p, unused_d, unused_p) in [
("50", &slow_d, &fast_p, &fast_d, &slow_p),
("100", &fast_d, &slow_p, &slow_d, &fast_p),
] {
for worker in [&fast_p, &slow_p, &fast_d, &slow_d] {
worker.captured.lock().unwrap().last_body = None;
}
let mut req = request(body("hi"));
req.headers_mut()
.insert("x-sgl-ttft-slo-ms", ttft.parse().unwrap());
req.headers_mut()
.insert("x-sgl-tps-slo", "100".parse().unwrap());
let response = app.clone().oneshot(req).await.unwrap();
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(response.headers()["x-sgl-decode-url"], expected_d.url);
response.into_body().collect().await.unwrap();
tokio::time::timeout(Duration::from_secs(2), async {
while expected_p.captured.lock().unwrap().last_body.is_none() {
tokio::time::sleep(Duration::from_millis(5)).await;
}
})
.await
.unwrap();
assert!(unused_p.captured.lock().unwrap().last_body.is_none());
assert!(unused_d.captured.lock().unwrap().last_body.is_none());
}
}
#[tokio::test]
async fn preferred_bucket_rejection_falls_back_and_disabled_headers_are_ignored() {
use sgl_router::buckets_reorg::SloPreference;
for enabled in [false, true] {
let rejected = MockWorker::start(vec![]).await;
let accepted = MockWorker::start(vec![]).await;
let mut ctx = Arc::try_unwrap(context(
&[
("rejected", Stage::Plain, &rejected),
("accepted", Stage::Plain, &accepted),
],
vec![],
))
.unwrap_or_else(|_| panic!("context is not shared yet"));
let mut fast = Bucket::new(
"a-fast",
BucketGroups::Plain(group("rejected", rejecting_policy())),
);
fast.ttft_ms = Some(10);
let mut slow = Bucket::new(
"b-slow",
BucketGroups::Plain(group("accepted", Arc::new(FirstPolicy::default()))),
);
slow.ttft_ms = Some(100);
let mut resolver = BucketResolver::new(vec![slow, fast]).unwrap();
if enabled {
resolver.ttft_slo = SloPreference::SloFirst;
}
ctx.chat_routing = ChatRouting::Reorg([(ModelId("tiny".into()), resolver)].into());
let app = build_router(Arc::new(ctx));
let mut req = request(body("hi"));
req.headers_mut().insert(
"x-sgl-ttft-slo-ms",
if enabled { "50" } else { "invalid" }.parse().unwrap(),
);
// TPS is disabled even when TTFT is enabled.
req.headers_mut()
.insert("x-sgl-tps-slo", "NaN".parse().unwrap());
let response = app.oneshot(req).await.unwrap();
assert_eq!(response.status(), StatusCode::OK);
response.into_body().collect().await.unwrap();
assert!(rejected.captured.lock().unwrap().last_body.is_none());
assert!(accepted.captured.lock().unwrap().last_body.is_some());
}
}