diff --git a/experimental/sgl-router/POLICY_DESIGN.md b/experimental/sgl-router/POLICY_DESIGN.md index f49d383b2..bd6fba52b 100644 --- a/experimental/sgl-router/POLICY_DESIGN.md +++ b/experimental/sgl-router/POLICY_DESIGN.md @@ -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. diff --git a/experimental/sgl-router/src/buckets_reorg.rs b/experimental/sgl-router/src/buckets_reorg.rs index ed7caa9dd..9d9a1b151 100644 --- a/experimental/sgl-router/src/buckets_reorg.rs +++ b/experimental/sgl-router/src/buckets_reorg.rs @@ -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, + /// Optional service estimates used only for bucket ordering. + pub ttft_ms: Option, + pub tokens_per_second: Option, 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) -> 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, + 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, + ttft_ms: Option, + tokens_per_second: Option, ) -> Result, 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) } } diff --git a/experimental/sgl-router/src/server/routes/chat/reorg.rs b/experimental/sgl-router/src/server/routes/chat/reorg.rs index 0822ebf07..6911071c2 100644 --- a/experimental/sgl-router/src/server/routes/chat/reorg.rs +++ b/experimental/sgl-router/src/server/routes/chat/reorg.rs @@ -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( diff --git a/experimental/sgl-router/tests/component/main.rs b/experimental/sgl-router/tests/component/main.rs index 62063c817..3caf5059e 100644 --- a/experimental/sgl-router/tests/component/main.rs +++ b/experimental/sgl-router/tests/component/main.rs @@ -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; diff --git a/experimental/sgl-router/tests/component/policies_reorg.rs b/experimental/sgl-router/tests/component/policies_reorg.rs index 9b0f298ca..60ac33f03 100644 --- a/experimental/sgl-router/tests/component/policies_reorg.rs +++ b/experimental/sgl-router/tests/component/policies_reorg.rs @@ -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::>(), ["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() diff --git a/experimental/sgl-router/tests/component/policies_reorg_slo.rs b/experimental/sgl-router/tests/component/policies_reorg_slo.rs new file mode 100644 index 000000000..45d70ff52 --- /dev/null +++ b/experimental/sgl-router/tests/component/policies_reorg_slo.rs @@ -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, tps: Option) -> 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, tps: Option) -> 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::>(), + ["tps", "neither", "missing"] + ); + assert!(matches!( + resolver.resolve(5, Some(4), None, None), + Err(PickError::InvalidSignal(_)) + )); +} diff --git a/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs b/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs index 6ffdb65d7..e1c4711be 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs @@ -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); diff --git a/experimental/sgl-router/tests/proxy/chat_routing/reorg/slo.rs b/experimental/sgl-router/tests/proxy/chat_routing/reorg/slo.rs new file mode 100644 index 000000000..6332c506b --- /dev/null +++ b/experimental/sgl-router/tests/proxy/chat_routing/reorg/slo.rs @@ -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()); + } +}