diff --git a/experimental/sgl-router/POLICY_DESIGN.md b/experimental/sgl-router/POLICY_DESIGN.md index d0077d2cb..58cf904fb 100644 --- a/experimental/sgl-router/POLICY_DESIGN.md +++ b/experimental/sgl-router/POLICY_DESIGN.md @@ -222,24 +222,35 @@ no HTTP body, bucket resolver, state handles, snapshots, or backend configuratio Admission evaluates acceptance. It does not rank engines, choose replacements, change buckets, or mutate affinity. -| Check | Acceptance rule | -| --- | --- | -| `AllowAll` | Add no acceptance constraint | -| `CapacityAdmission` | Projected running requests and KV tokens fit reported capacity | -| `PendingPrefillAdmission` | Waiting uncached tokens plus incoming uncached work fit the budget | -| `InFlightLimitAdmission` | Router-local in-flight requests are below the limit | -| `QueueLimitAdmission` | Engine-reported waiting requests are below the limit | -| `AllOfAdmission` | Every attached check allows the request | +An admission policy is a set of per-engine caps, `AdmissionLimits`. Each cap +is optional; an unset cap is not checked, and the default admits everything. +A cap admits while the engine's current metric is below it. Request size is +not part of admission: buckets already select by input length and context +capacity, and admission only observes load without reserving it. -`EngineAdmission::check(engine, request, load)` checks one engine and returns -`Allow`, `Reject(reason)`, or an error for invalid inputs. Policies attach the -checker directly as `Arc`. There is no placement setting, -filtering wrapper, or before/after API; each policy decides where checking -belongs in its selection algorithm. The `load` argument is an -`Option<&EngineReportedWorkerLoad>` retained by the policy for this engine, including -request counts, token usage, capacity, and the report timestamp. `None` means -no usable observation, never zero load; each check defines its missing-data -behavior. Other required state handles belong to the checker. +| Limit | Engine metric | +| --- | --- | +| `max_running_requests` | Reported running requests | +| `max_waiting_requests` | Reported waiting requests | +| `max_kv_tokens` | Reported total KV tokens | +| `max_pending_prefill_tokens` | Reported waiting uncached tokens | +| `max_inflight_requests` | Router-local in-flight requests | + +```json +{"max_running_requests": 64, "max_kv_tokens": 1048576, "max_inflight_requests": 64} +``` + +Limits are absolute caps; they do not default to capacities reported by the +engine. Unknown fields are rejected during deserialization. + +The policy reads the selected engine's `EngineMetrics` from the load snapshot +it already captured for selection plus the live in-flight counter and calls +`EngineAdmission::check(engine, metrics)`, which returns `Allow`, +`Reject(limit name)`, or an error. Reported +metrics are `None` without a fresh, complete report, never zero, and such +limits fail open; the in-flight count is always known. Policies attach the +checker as `Arc` and decide where checking belongs in +their selection algorithm; there is no placement setting or filtering wrapper. Power-of-two first selects an engine, then calls admission exactly once on that engine. A rejection returns `AdmissionRejected` to the bucket loop; it does not @@ -257,26 +268,20 @@ lacks a fresh, complete native report with valid capacity, both are compared by router-local active requests instead. Basic reports from older publishers are still passed to admission when fresh, but do not supply native pressure metrics. -Prepare the signals needed by admission before checking. A pending-prefill check -uses per-engine uncached work when a prefix is known, and full input otherwise. -Decode capacity uses the expected peak sequence length when available, including -on a cache hit. Power-of-two retains the selected engine's load record from -selection and passes it to admission without another snapshot. A single candidate -still has its load read for admission, even though selection needs no comparison. -Neither the bucket nor HTTP handler supplies observations. Concrete load-aware -acceptance rules and additional cache-specific admission signals remain follow-up -work. Synchronous checks do not fetch telemetry over the network themselves. +Power-of-two retains the selected engine's load record from selection and +passes it to admission without another snapshot. A single candidate still has +its load read for admission, even though selection needs no comparison. +Neither the bucket nor HTTP handler supplies observations. Synchronous checks +do not fetch telemetry over the network themselves. -`AllowAll` is the default for new explicit policy attachments. It leaves health, -role, membership, and policy preferences in force. Migrated configurations must -retain their existing capacity and configured budget checks; see compatibility -below. Each check defines its missing-data behavior. Unknown load is not zero; -the existing capacity and pending-prefill checks allow requests without a fresh, -complete native report. +`AdmissionLimits::default()` is the default for new explicit policy attachments. +It leaves health, role, membership, and policy preferences in force. Migrated +configurations must retain their existing capacity and configured budget checks; +see compatibility below. -The cache policy's `worker_queue_limit` is a **soft preference**, not -`QueueLimitAdmission`. Saturation handling can reconsider a queued engine, but -cannot bypass attached hard admission. +The cache policy's `worker_queue_limit` is a **soft preference**; +`max_waiting_requests` is a hard rejection. Saturation handling can reconsider +a queued engine, but cannot bypass attached hard admission. Admission checks observe capacity; they do not reserve it. Concurrent requests may pass against the same observation. Strict reservations would require a @@ -432,12 +437,12 @@ buckets: worker_ids: [P1, P2] policy: type: cache_aware - admission: {type: capacity} + admission: {max_running_requests: 64, max_kv_tokens: 1048576} decode: worker_ids: [D1, D2] policy: type: power_of_two - admission: {type: capacity} + admission: {max_running_requests: 64, max_kv_tokens: 1048576} - id: long-context rank: 20 @@ -449,12 +454,12 @@ buckets: worker_ids: [P3, P4] policy: type: cache_aware - admission: {type: capacity} + admission: {max_running_requests: 64, max_kv_tokens: 1048576} decode: worker_ids: [D3, D4] policy: type: power_of_two - admission: {type: capacity} + admission: {max_running_requests: 64, max_kv_tokens: 1048576} ``` A request with 4k input tokens and a 16k expected peak cannot fit the short @@ -492,8 +497,8 @@ do not accept and ignore them. queue limit, and saturation floor. - Preserve session and sticky headers, idle timeouts, eviction cadence, and the four sticky fallback choices. Global modes need a bucket-first migration design. -- Translate `--filter overloaded` and `--max-in-flight` into - `InFlightLimitAdmission`, composed with other checks through `AllOfAdmission`. +- Map `--filter overloaded` and `--max-in-flight` to `max_inflight_requests`; + the existing router-local counter remains the source. - Preserve configured capacity, pending-prefill, and in-flight checks, including their missing-report behavior. Power-of-two applies admission to its selected engine; other policies explicitly place checks in their selection logic. @@ -558,7 +563,8 @@ Implemented here: - `EngineGroup::pick` owns live candidate filtering, policy invocation, and exact candidate validation, without cross-bucket fallback. - `Policy::pick`, within-group fallback interface, per-engine `EngineAdmission::check`, - and `AllowAll`. Power-of-two samples two distinct engines, compares stage pressure, + and `AdmissionLimits` over running, waiting, KV, pending-prefill and in-flight + metrics. Power-of-two samples two distinct engines, compares stage pressure, and checks its selected engine with no replacement on rejection. - Policy-owned load dependency and local observations. Power-of-two passes the selected engine's load record directly to admission, without another snapshot. @@ -586,12 +592,12 @@ 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 order: concrete admission (#40271), then bucket SLO ordering -in a separate PR, followed by remaining policies and production configuration. +Follow-up work includes bucket SLO ordering, remaining selection policies, +and production configuration. Not yet implemented in the reorg path: -- Other concrete policies and capacity/in-flight admission checks. +- 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. diff --git a/experimental/sgl-router/src/policies_reorg/admission.rs b/experimental/sgl-router/src/policies_reorg/admission.rs index f63896b38..854688a1d 100644 --- a/experimental/sgl-router/src/policies_reorg/admission.rs +++ b/experimental/sgl-router/src/policies_reorg/admission.rs @@ -1,12 +1,17 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors // SPDX-License-Identifier: Apache-2.0 +//! Per-engine admission: a set of caps compared against the selected engine's +//! current measurements. Request size is a bucket concern, not an admission one. + use std::fmt::Debug; -use crate::state::load_monitor::engine_reported_load::EngineReportedWorkerLoad; +use serde::{Deserialize, Serialize}; + +use crate::state::load_monitor::engine_reported_load::EngineReportedLoadSnapshot; use crate::workers::Worker; -use super::{PickError, PickRequest}; +use super::PickError; #[derive(Debug, Clone, PartialEq, Eq)] pub enum Decision { @@ -14,29 +19,84 @@ pub enum Decision { Reject(String), } -/// Checks one engine using the load observation retained by selection. -/// `None` means no usable load observation, never zero load. Each check defines -/// its missing-data behavior and owns any other state handles it needs. -/// Each policy decides when to check an engine and how to handle rejection. -pub trait EngineAdmission: Send + Sync + Debug { - fn check( - &self, - engine: &Worker, - request: &PickRequest<'_>, - load: Option<&EngineReportedWorkerLoad>, - ) -> Result; +/// One engine's measurements at pick time. Reported values are `None` without a +/// fresh, complete report, never zero. In-flight requests are counted by this +/// router and always known. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct EngineMetrics { + pub running_requests: Option, + pub waiting_requests: Option, + pub kv_tokens: Option, + pub pending_prefill_tokens: Option, + pub inflight_requests: u64, } -#[derive(Debug)] -pub struct AllowAll; +impl EngineMetrics { + /// Read `engine` from the load snapshot selection already captured. + pub fn observe(engine: &Worker, load: &EngineReportedLoadSnapshot) -> Self { + let basic = load.fresh_load_for_url(&engine.url); + let native = load.fresh_native_cache_load_for_url(&engine.url); + Self { + running_requests: basic.map(|load| load.num_running_reqs), + waiting_requests: basic.map(|load| load.num_waiting_reqs), + kv_tokens: native.map(|load| load.num_total_tokens), + pending_prefill_tokens: native.map(|load| load.num_waiting_uncached_tokens), + inflight_requests: engine.router_inflight_load() as u64, + } + } +} -impl EngineAdmission for AllowAll { - fn check( - &self, - _: &Worker, - _: &PickRequest<'_>, - _: Option<&EngineReportedWorkerLoad>, - ) -> Result { +/// Checks one selected engine. Policies decide when to check and how to handle +/// rejection; admission never selects replacements or reserves capacity. +pub trait EngineAdmission: Send + Sync + Debug { + fn check(&self, engine: &Worker, metrics: &EngineMetrics) -> Result; +} + +/// Per-engine caps; an unset limit is not checked. A limit admits while the +/// metric is below it. Unknown engine metrics fail open. The default allows +/// everything. Checks observe load; they do not reserve capacity. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct AdmissionLimits { + pub max_running_requests: Option, + pub max_waiting_requests: Option, + pub max_kv_tokens: Option, + pub max_pending_prefill_tokens: Option, + pub max_inflight_requests: Option, +} + +impl EngineAdmission for AdmissionLimits { + fn check(&self, _: &Worker, engine: &EngineMetrics) -> Result { + let limits = [ + ( + "max_running_requests", + self.max_running_requests, + engine.running_requests, + ), + ( + "max_waiting_requests", + self.max_waiting_requests, + engine.waiting_requests, + ), + ("max_kv_tokens", self.max_kv_tokens, engine.kv_tokens), + ( + "max_pending_prefill_tokens", + self.max_pending_prefill_tokens, + engine.pending_prefill_tokens, + ), + ( + "max_inflight_requests", + self.max_inflight_requests, + Some(engine.inflight_requests), + ), + ]; + for (name, max, current) in limits { + if let (Some(max), Some(current)) = (max, current) { + if current >= max { + return Ok(Decision::Reject(name.into())); + } + } + } Ok(Decision::Allow) } } diff --git a/experimental/sgl-router/src/policies_reorg/cache_aware.rs b/experimental/sgl-router/src/policies_reorg/cache_aware.rs index 199317be1..a5b7b2a96 100644 --- a/experimental/sgl-router/src/policies_reorg/cache_aware.rs +++ b/experimental/sgl-router/src/policies_reorg/cache_aware.rs @@ -26,7 +26,7 @@ use crate::state::load_monitor::engine_reported_load::{ }; use crate::workers::Worker; -use super::admission::{AllowAll, Decision, EngineAdmission}; +use super::admission::{AdmissionLimits, Decision, EngineAdmission, EngineMetrics}; use super::power_of_two::PowerOfTwoPolicy; use super::{Pick, PickError, PickRequest, Policy, Rejection, Stage}; @@ -173,7 +173,7 @@ impl CacheAwarePolicy { fallback: Arc::new(PowerOfTwoPolicy::new(Arc::clone(&engine_load))), engine_load, config, - admission: Arc::new(AllowAll), + admission: Arc::new(AdmissionLimits::default()), }) } @@ -236,11 +236,10 @@ impl CacheAwarePolicy { fn check( &self, engine: &Worker, - request: &PickRequest<'_>, load: &EngineReportedLoadSnapshot, ) -> Result, PickError> { - let load = load.fresh_load_for_url(&engine.url); - Ok(match self.admission.check(engine, request, load)? { + let metrics = EngineMetrics::observe(engine, load); + Ok(match self.admission.check(engine, &metrics)? { Decision::Allow => None, Decision::Reject(reason) => Some(Rejection { engine: engine.id.clone(), @@ -252,13 +251,12 @@ impl CacheAwarePolicy { fn admit<'e>( &self, candidates: &[Candidate<'e>], - request: &PickRequest<'_>, load: &EngineReportedLoadSnapshot, rejections: &mut Vec, ) -> Result>, PickError> { let mut admitted = Vec::new(); for &candidate in candidates { - match self.check(candidate.engine, request, load)? { + match self.check(candidate.engine, load)? { None => admitted.push(candidate), Some(rejection) => rejections.push(rejection), } @@ -316,7 +314,6 @@ impl CacheAwarePolicy { &self, candidates: &[Candidate<'_>], engines: &[Arc], - request: &PickRequest<'_>, load: &EngineReportedLoadSnapshot, ) -> Result, PickError> { let limit = self.config.worker_queue_limit; @@ -333,7 +330,7 @@ impl CacheAwarePolicy { evaluated.extend(&gated); } let mut rejections = Vec::new(); - let admitted = self.admit(&evaluated, request, load, &mut rejections)?; + let admitted = self.admit(&evaluated, load, &mut rejections)?; if let Some(&least) = admitted.iter().min_by_key(|c| c.uncached_tokens) { let loads = FreshLoadLookup::new(Some(load), evaluated.iter().map(|c| c.engine)); let guarded = self.config.pressure_guard @@ -377,7 +374,7 @@ impl CacheAwarePolicy { if pinned { let loads = FreshLoadLookup::new(Some(load), gated.iter().map(|c| c.engine)); let owner = self - .admit(&gated, request, load, &mut rejections)? + .admit(&gated, load, &mut rejections)? .into_iter() .min_by(|left, right| { loads @@ -430,7 +427,7 @@ impl Policy for CacheAwarePolicy { // Capture load after remote I/O; selection and admission share it. let load = self.engine_load.capture_snapshot(Instant::now()); let candidates = self.candidates(engines, request, signal.as_deref(), &load); - if let Some(pick) = self.resolve(&candidates, engines, request, &load)? { + if let Some(pick) = self.resolve(&candidates, engines, &load)? { return Ok(pick); } // Miss: fall back within the unqueued tier when one exists. @@ -448,7 +445,7 @@ impl Policy for CacheAwarePolicy { if !pool.iter().any(|e| Arc::ptr_eq(e, &pick.engine)) { return Err(PickError::OutsideCandidates(pick.engine.id.clone())); } - if let Some(rejection) = self.check(&pick.engine, request, &load)? { + if let Some(rejection) = self.check(&pick.engine, &load)? { return Err(PickError::AdmissionRejected(rejection)); } pick.reason = "no_cache_candidate"; diff --git a/experimental/sgl-router/src/policies_reorg/power_of_two.rs b/experimental/sgl-router/src/policies_reorg/power_of_two.rs index 4d6e7b06a..03d5bb38a 100644 --- a/experimental/sgl-router/src/policies_reorg/power_of_two.rs +++ b/experimental/sgl-router/src/policies_reorg/power_of_two.rs @@ -11,7 +11,7 @@ use crate::policies::admission::{compare_decode_pressure, compare_prefill_pressu use crate::state::load_monitor::engine_reported_load::EngineReportedLoadTable; use crate::workers::Worker; -use super::admission::{AllowAll, Decision, EngineAdmission}; +use super::admission::{AdmissionLimits, Decision, EngineAdmission, EngineMetrics}; use super::{Pick, PickError, PickRequest, Policy, Rejection, Stage}; /// Samples two distinct engines and selects the one with lower stage pressure. @@ -27,7 +27,7 @@ impl PowerOfTwoPolicy { pub fn new(engine_load: Arc) -> Self { Self { engine_load, - admission: Arc::new(AllowAll), + admission: Arc::new(AdmissionLimits::default()), } } } @@ -63,8 +63,8 @@ impl Policy for PowerOfTwoPolicy { Arc::clone(if pressure.is_gt() { right } else { left }) } }; - let engine_load = load.fresh_load_for_url(&engine.url); - if let Decision::Reject(reason) = self.admission.check(&engine, request, engine_load)? { + let metrics = EngineMetrics::observe(&engine, &load); + if let Decision::Reject(reason) = self.admission.check(&engine, &metrics)? { return Err(PickError::AdmissionRejected(Rejection { engine: engine.id.clone(), reason, diff --git a/experimental/sgl-router/src/policies_reorg/session_aware.rs b/experimental/sgl-router/src/policies_reorg/session_aware.rs index f56273e0d..bb0b1e4b3 100644 --- a/experimental/sgl-router/src/policies_reorg/session_aware.rs +++ b/experimental/sgl-router/src/policies_reorg/session_aware.rs @@ -15,7 +15,7 @@ use crate::state::load_monitor::engine_reported_load::{ use crate::state::AffinityStore; use crate::workers::Worker; -use super::admission::{AllowAll, Decision, EngineAdmission}; +use super::admission::{AdmissionLimits, Decision, EngineAdmission, EngineMetrics}; use super::power_of_two::PowerOfTwoPolicy; use super::{Pick, PickError, PickRequest, Policy, Rejection}; @@ -36,7 +36,7 @@ impl SessionAwarePolicy { store, fallback: PowerOfTwoPolicy::new(Arc::clone(&engine_load)), engine_load, - admission: Arc::new(AllowAll), + admission: Arc::new(AdmissionLimits::default()), } } @@ -55,16 +55,9 @@ impl SessionAwarePolicy { )) } - fn check( - &self, - engine: &Worker, - request: &PickRequest<'_>, - load: &EngineReportedLoadSnapshot, - ) -> Result<(), PickError> { - match self - .admission - .check(engine, request, load.fresh_load_for_url(&engine.url))? - { + fn check(&self, engine: &Worker, load: &EngineReportedLoadSnapshot) -> Result<(), PickError> { + let metrics = EngineMetrics::observe(engine, load); + match self.admission.check(engine, &metrics)? { Decision::Allow => Ok(()), Decision::Reject(reason) => Err(PickError::AdmissionRejected(Rejection { engine: engine.id.clone(), @@ -87,7 +80,7 @@ impl Policy for SessionAwarePolicy { let key = Self::assignment_key(request); if let Some(bound) = key.as_ref().and_then(|key| self.store.bound(key, engines)) { let load = self.engine_load.capture_snapshot(Instant::now()); - self.check(bound, request, &load)?; + self.check(bound, &load)?; return Ok(Pick { engine: Arc::clone(bound), reason: "session_primary", @@ -98,7 +91,7 @@ impl Policy for SessionAwarePolicy { // checks its chosen engine before creating or replacing a binding. let mut pick = self.pick_fallback(engines, request).await?; let load = self.engine_load.capture_snapshot(Instant::now()); - self.check(&pick.engine, request, &load)?; + self.check(&pick.engine, &load)?; let Some(key) = key else { pick.reason = "no_session"; return Ok(pick); @@ -108,7 +101,7 @@ impl Policy for SessionAwarePolicy { if !Arc::ptr_eq(effective, &pick.engine) { // A racing first assignment wins. Check it once, without // rewriting a rejected binding or retrying another engine. - self.check(effective, request, &load)?; + self.check(effective, &load)?; pick.reason = "session_primary"; } else { pick.reason = "assigned"; diff --git a/experimental/sgl-router/tests/component/main.rs b/experimental/sgl-router/tests/component/main.rs index d6acf13a7..62063c817 100644 --- a/experimental/sgl-router/tests/component/main.rs +++ b/experimental/sgl-router/tests/component/main.rs @@ -11,6 +11,7 @@ mod discovery; mod health; mod policies; mod policies_reorg; +mod policies_reorg_admission; mod policies_reorg_cache_aware; mod policies_reorg_load; mod policies_reorg_power_of_two; diff --git a/experimental/sgl-router/tests/component/policies_reorg.rs b/experimental/sgl-router/tests/component/policies_reorg.rs index fa49c61ec..9b0f298ca 100644 --- a/experimental/sgl-router/tests/component/policies_reorg.rs +++ b/experimental/sgl-router/tests/component/policies_reorg.rs @@ -8,9 +8,10 @@ use sgl_router::buckets_reorg::{ Bucket, BucketGroups, BucketRequest, BucketResolver, EngineGroup, TokenLimits, }; use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec}; -use sgl_router::policies_reorg::admission::{AllowAll, Decision, EngineAdmission}; +use sgl_router::policies_reorg::admission::{ + AdmissionLimits, Decision, EngineAdmission, EngineMetrics, +}; use sgl_router::policies_reorg::{Pick, PickError, PickRequest, Policy, Rejection, Stage}; -use sgl_router::state::load_monitor::engine_reported_load::EngineReportedWorkerLoad; use sgl_router::workers::{Worker, WorkerRegistry}; #[derive(Debug)] @@ -25,7 +26,7 @@ struct TestPolicy { impl Default for TestPolicy { fn default() -> Self { Self { - admission: Arc::new(AllowAll), + admission: Arc::new(AdmissionLimits::default()), result: None, miss: false, invalid: false, @@ -52,7 +53,9 @@ impl Policy for TestPolicy { return Err(PickError::NoCandidates); } let engine = self.result.clone().unwrap_or_else(|| engines[0].clone()); - if let Decision::Reject(reason) = self.admission.check(&engine, request, None)? { + if let Decision::Reject(reason) = + self.admission.check(&engine, &EngineMetrics::default())? + { return Err(PickError::AdmissionRejected(Rejection { engine: engine.id.clone(), reason, @@ -70,12 +73,7 @@ impl Policy for TestPolicy { struct Reject(&'static str); impl EngineAdmission for Reject { - fn check( - &self, - engine: &Worker, - _: &PickRequest<'_>, - _: Option<&EngineReportedWorkerLoad>, - ) -> Result { + fn check(&self, engine: &Worker, _: &EngineMetrics) -> Result { Ok(if engine.id.0 == self.0 { Decision::Reject("full".into()) } else { @@ -437,12 +435,7 @@ async fn power_of_two_checks_selected_engine_and_propagates_rejection_without_fa } impl EngineAdmission for Check { - fn check( - &self, - engine: &Worker, - _: &PickRequest<'_>, - _: Option<&EngineReportedWorkerLoad>, - ) -> Result { + fn check(&self, engine: &Worker, _: &EngineMetrics) -> Result { self.calls.lock().unwrap().push(engine.id.clone()); if self.invalid { Err(PickError::InvalidSignal("admission input".into())) diff --git a/experimental/sgl-router/tests/component/policies_reorg_admission.rs b/experimental/sgl-router/tests/component/policies_reorg_admission.rs new file mode 100644 index 000000000..bf9cf2d83 --- /dev/null +++ b/experimental/sgl-router/tests/component/policies_reorg_admission.rs @@ -0,0 +1,180 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; +use std::time::Instant; + +use serde_json::json; +use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec}; +use sgl_router::policies_reorg::admission::{ + AdmissionLimits, Decision, EngineAdmission, EngineMetrics, +}; +use sgl_router::policies_reorg::power_of_two::PowerOfTwoPolicy; +use sgl_router::policies_reorg::{PickError, PickRequest, Policy, Stage}; +use sgl_router::state::load_monitor::engine_reported_load::{ + EngineReportedLoadTable, LoadStat, NativeCacheRankLoad, +}; +use sgl_router::workers::Worker; + +fn engine() -> Arc { + Arc::new(Worker::new(WorkerSpec { + id: WorkerId("w".into()), + url: "http://w".into(), + mode: Stage::Plain, + model_ids: vec![ModelId("m".into())], + bootstrap_port: None, + })) +} + +fn report(running: u64, waiting: u64, kv_tokens: u64, pending: u64) -> LoadStat { + LoadStat { + num_running_reqs: running, + num_waiting_reqs: waiting, + num_tokens: kv_tokens, + max_total_num_tokens: 1000, + native_cache: Some(NativeCacheRankLoad { + num_waiting_uncached_tokens: pending, + num_total_tokens: kv_tokens, + max_running_requests: 100, + total_prefill_uncached_tokens: 0, + total_prefill_busy_us: 0, + }), + } +} + +fn reject(name: &str) -> Decision { + Decision::Reject(name.into()) +} + +#[test] +fn limits_are_flat_optional_fields() { + let limits: AdmissionLimits = + serde_json::from_value(json!({"max_inflight_requests": 4, "max_kv_tokens": 100})).unwrap(); + assert_eq!( + limits, + AdmissionLimits { + max_inflight_requests: Some(4), + max_kv_tokens: Some(100), + ..Default::default() + } + ); + assert_eq!( + serde_json::from_value::(json!({})).unwrap(), + AdmissionLimits::default() + ); + assert!(serde_json::from_value::(json!({"max_in_flight": 4})).is_err()); +} + +#[test] +fn each_limit_caps_its_metric_and_fails_open_when_unknown() { + let engine = engine(); + let known = EngineMetrics { + running_requests: Some(3), + waiting_requests: Some(1), + kv_tokens: Some(80), + pending_prefill_tokens: Some(90), + inflight_requests: 2, + }; + let unknown = EngineMetrics { + inflight_requests: 2, + ..Default::default() + }; + let limit = |field: &str, max| { + let mut limits = AdmissionLimits::default(); + *match field { + "max_running_requests" => &mut limits.max_running_requests, + "max_waiting_requests" => &mut limits.max_waiting_requests, + "max_kv_tokens" => &mut limits.max_kv_tokens, + "max_pending_prefill_tokens" => &mut limits.max_pending_prefill_tokens, + "max_inflight_requests" => &mut limits.max_inflight_requests, + _ => unreachable!(), + } = Some(max); + limits + }; + assert_eq!( + AdmissionLimits::default().check(&engine, &known).unwrap(), + Decision::Allow + ); + for (field, below, at) in [ + ("max_running_requests", 4, 3), + ("max_waiting_requests", 2, 1), + ("max_kv_tokens", 81, 80), + ("max_pending_prefill_tokens", 91, 90), + ("max_inflight_requests", 3, 2), + ] { + let (below, at) = (limit(field, below), limit(field, at)); + assert_eq!(below.check(&engine, &known).unwrap(), Decision::Allow); + assert_eq!(at.check(&engine, &known).unwrap(), reject(field)); + // Without a report only the router-local in-flight count applies. + let expected = if field == "max_inflight_requests" { + reject(field) + } else { + Decision::Allow + }; + assert_eq!(at.check(&engine, &unknown).unwrap(), expected); + } +} + +#[tokio::test] +async fn metrics_come_from_the_selection_snapshot_and_live_inflight_count() { + let engine = engine(); + let table = EngineReportedLoadTable::new(); + table.set(&engine.url, 0, report(1, 2, 80, 5), Instant::now()); + let guard = engine.load_guard(); + assert_eq!( + EngineMetrics::observe(&engine, &table.capture_snapshot(Instant::now())), + EngineMetrics { + running_requests: Some(1), + waiting_requests: Some(2), + kv_tokens: Some(80), + pending_prefill_tokens: Some(5), + inflight_requests: 1, + } + ); + drop(guard); + // Basic reports carry request counts but no KV or pending-prefill tokens. + let basic = LoadStat { + native_cache: None, + ..report(1, 2, 80, 5) + }; + table.set(&engine.url, 0, basic, Instant::now()); + let metrics = EngineMetrics::observe(&engine, &table.capture_snapshot(Instant::now())); + assert_eq!( + ( + metrics.running_requests, + metrics.kv_tokens, + metrics.pending_prefill_tokens + ), + (Some(1), None, None) + ); + + let mut policy = PowerOfTwoPolicy::new(table.clone()); + policy.admission = Arc::new(AdmissionLimits { + max_running_requests: Some(2), + max_kv_tokens: Some(100), + ..Default::default() + }); + let model = ModelId("m".into()); + let request = PickRequest::new(&model, Stage::Plain, 10); + let engines = [engine]; + for (running, kv_tokens, rejection) in [ + (1, 100, Some("max_kv_tokens")), + (1, 99, None), + (2, 0, Some("max_running_requests")), + ] { + table.set( + &engines[0].url, + 0, + report(running, 0, kv_tokens, 0), + Instant::now(), + ); + let result = policy.pick(&engines, &request).await; + match rejection { + None => assert!(result.is_ok()), + Some(reason) => assert!(matches!( + result, + Err(PickError::AdmissionRejected(rejected)) if rejected.reason == reason + )), + } + } +} diff --git a/experimental/sgl-router/tests/component/policies_reorg_cache_aware.rs b/experimental/sgl-router/tests/component/policies_reorg_cache_aware.rs index a002100dd..ade06911d 100644 --- a/experimental/sgl-router/tests/component/policies_reorg_cache_aware.rs +++ b/experimental/sgl-router/tests/component/policies_reorg_cache_aware.rs @@ -10,7 +10,7 @@ use sgl_router::buckets_reorg::{Bucket, BucketGroups, BucketResolver, EngineGrou use sgl_router::config::AffinityConfig; use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec}; use sgl_router::policies::prefix_provider::RadixTreePrefixProvider; -use sgl_router::policies_reorg::admission::{Decision, EngineAdmission}; +use sgl_router::policies_reorg::admission::{Decision, EngineAdmission, EngineMetrics}; use sgl_router::policies_reorg::cache_aware::{CacheAwarePolicy, CacheSource, PrefixMemo}; use sgl_router::policies_reorg::power_of_two::PowerOfTwoPolicy; use sgl_router::policies_reorg::{PickError, PickRequest, Policy, Stage}; @@ -18,7 +18,7 @@ use sgl_router::state::kv_events::{ compute_block_hashes, compute_block_hashes_bigram, BlockSizeOracle, HashTree, KvWorkerId, }; use sgl_router::state::load_monitor::engine_reported_load::{ - EngineReportedLoadTable, EngineReportedWorkerLoad, LoadStat, NativeCacheRankLoad, + EngineReportedLoadTable, LoadStat, NativeCacheRankLoad, }; use sgl_router::workers::Worker; @@ -116,16 +116,11 @@ impl Reject { } impl EngineAdmission for Reject { - fn check( - &self, - engine: &Worker, - _: &PickRequest<'_>, - load: Option<&EngineReportedWorkerLoad>, - ) -> Result { + fn check(&self, engine: &Worker, metrics: &EngineMetrics) -> Result { self.calls .lock() .unwrap() - .push((engine.id.0.clone(), load.map(|load| load.num_waiting_reqs))); + .push((engine.id.0.clone(), metrics.waiting_requests)); Ok(if engine.id.0 == self.id { Decision::Reject("full".into()) } else { diff --git a/experimental/sgl-router/tests/component/policies_reorg_load.rs b/experimental/sgl-router/tests/component/policies_reorg_load.rs index e73cbb313..b4cdfab01 100644 --- a/experimental/sgl-router/tests/component/policies_reorg_load.rs +++ b/experimental/sgl-router/tests/component/policies_reorg_load.rs @@ -5,12 +5,10 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec}; -use sgl_router::policies_reorg::admission::{Decision, EngineAdmission}; +use sgl_router::policies_reorg::admission::{Decision, EngineAdmission, EngineMetrics}; use sgl_router::policies_reorg::power_of_two::PowerOfTwoPolicy; use sgl_router::policies_reorg::{PickError, PickRequest, Policy, Stage}; -use sgl_router::state::load_monitor::engine_reported_load::{ - EngineReportedLoadTable, EngineReportedWorkerLoad, LoadStat, -}; +use sgl_router::state::load_monitor::engine_reported_load::{EngineReportedLoadTable, LoadStat}; use sgl_router::workers::Worker; const URL: &str = "http://engine"; @@ -43,21 +41,16 @@ fn engine() -> Arc { #[derive(Debug)] struct ObserveAdmission { table: Arc, - observations: Mutex>>, + observations: Mutex>, } impl EngineAdmission for ObserveAdmission { - fn check( - &self, - engine: &Worker, - _: &PickRequest<'_>, - load: Option<&EngineReportedWorkerLoad>, - ) -> Result { + fn check(&self, engine: &Worker, metrics: &EngineMetrics) -> Result { assert_eq!(engine.url, URL); // A new report arriving after selection must not change the observation // supplied to admission. The next pick should read the new report. report(&self.table, 0, 99, Instant::now()); - self.observations.lock().unwrap().push(load.cloned()); + self.observations.lock().unwrap().push(*metrics); Ok(Decision::Allow) } } @@ -95,17 +88,16 @@ async fn selected_load_reaches_admission_and_next_pick_reads_fresh_state() { } let observations = admission.observations.lock().unwrap(); assert_eq!(observations.len(), 2); + // Basic reports carry request counts but no KV or pending-prefill tokens. assert_eq!( observations[0], - Some(EngineReportedWorkerLoad { - num_running_reqs: 4, - num_waiting_reqs: 4, - num_tokens: 60, - max_total_num_tokens: 200, - captured_at: first_at, - }) + EngineMetrics { + running_requests: Some(4), + waiting_requests: Some(4), + ..EngineMetrics::default() + } ); - assert_eq!(observations[1].as_ref().unwrap().num_running_reqs, 102); + assert_eq!(observations[1].running_requests, Some(102)); } #[tokio::test] @@ -133,6 +125,10 @@ async fn missing_stale_and_incomplete_reports_reach_admission_as_unknown() { .pick(&[engine()], &PickRequest::new(&model, Stage::Plain, 10)) .await .unwrap(); - assert_eq!(*admission.observations.lock().unwrap(), [None], "{case}"); + assert_eq!( + *admission.observations.lock().unwrap(), + [EngineMetrics::default()], + "{case}" + ); } } diff --git a/experimental/sgl-router/tests/component/policies_reorg_session_aware.rs b/experimental/sgl-router/tests/component/policies_reorg_session_aware.rs index 0cae2dabf..00e9f1142 100644 --- a/experimental/sgl-router/tests/component/policies_reorg_session_aware.rs +++ b/experimental/sgl-router/tests/component/policies_reorg_session_aware.rs @@ -6,12 +6,10 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec}; -use sgl_router::policies_reorg::admission::{Decision, EngineAdmission}; +use sgl_router::policies_reorg::admission::{Decision, EngineAdmission, EngineMetrics}; use sgl_router::policies_reorg::session_aware::SessionAwarePolicy; use sgl_router::policies_reorg::{PickError, PickRequest, Policy, Stage}; -use sgl_router::state::load_monitor::engine_reported_load::{ - EngineReportedLoadTable, EngineReportedWorkerLoad, LoadStat, -}; +use sgl_router::state::load_monitor::engine_reported_load::{EngineReportedLoadTable, LoadStat}; use sgl_router::state::load_monitor::router_inflight_load::MockClock; use sgl_router::state::AffinityStore; use sgl_router::workers::Worker; @@ -52,16 +50,11 @@ struct Admission { } impl EngineAdmission for Admission { - fn check( - &self, - engine: &Worker, - _: &PickRequest<'_>, - load: Option<&EngineReportedWorkerLoad>, - ) -> Result { + fn check(&self, engine: &Worker, metrics: &EngineMetrics) -> Result { self.calls .lock() .unwrap() - .push((engine.id.0.clone(), load.map(|load| load.num_waiting_reqs))); + .push((engine.id.0.clone(), metrics.waiting_requests)); if self.invalid.load(Ordering::Relaxed) { return Err(PickError::InvalidSignal("invalid admission input".into())); } @@ -339,6 +332,7 @@ async fn shared_store_refreshes_active_sessions_and_expires_idle_ones() { #[derive(Debug)] struct RacingAdmission { + model: ModelId, competitor: SessionAwarePolicy, winner: Arc, reject_winner: bool, @@ -346,19 +340,14 @@ struct RacingAdmission { } impl EngineAdmission for RacingAdmission { - fn check( - &self, - engine: &Worker, - request: &PickRequest<'_>, - _: Option<&EngineReportedWorkerLoad>, - ) -> Result { + fn check(&self, engine: &Worker, _: &EngineMetrics) -> Result { self.calls.lock().unwrap().push(engine.id.0.clone()); if engine.id.0 == "a" { // Complete a competing first request after this request selected a, // but before it can commit. The effective binding is now b. futures::executor::block_on( self.competitor - .pick(std::slice::from_ref(&self.winner), request), + .pick(std::slice::from_ref(&self.winner), &request(&self.model)), )?; } Ok(if self.reject_winner && engine.id == self.winner.id { @@ -375,6 +364,7 @@ async fn concurrent_assignment_winner_is_checked_and_preserved_on_rejection() { let (mut policy, store) = policy(); let engines = [engine("a", 0), engine("b", 9)]; let admission = Arc::new(RacingAdmission { + model: ModelId("m".into()), competitor: SessionAwarePolicy::new(store.clone(), EngineReportedLoadTable::new()), winner: engines[1].clone(), reject_winner, diff --git a/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs b/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs index b00ea76ad..e71d66e38 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs @@ -6,10 +6,11 @@ use crate::common::mock_worker::MockWorker; use futures::future::BoxFuture; use sgl_router::buckets_reorg::{Bucket, BucketGroups, BucketResolver, EngineGroup}; use sgl_router::policies::PolicyRegistry; -use sgl_router::policies_reorg::admission::{AllowAll, Decision, EngineAdmission}; +use sgl_router::policies_reorg::admission::{ + AdmissionLimits, Decision, EngineAdmission, EngineMetrics, +}; use sgl_router::policies_reorg::{Pick, PickError, PickRequest, Policy, Rejection, Stage}; use sgl_router::server::app_context::ChatRouting; -use sgl_router::state::load_monitor::engine_reported_load::EngineReportedWorkerLoad; use std::sync::Mutex; mod session_aware; @@ -27,7 +28,7 @@ struct FirstPolicy { impl Default for FirstPolicy { fn default() -> Self { Self { - admission: Arc::new(AllowAll), + admission: Arc::new(AdmissionLimits::default()), miss: false, invalid: false, calls: Mutex::new(Vec::new()), @@ -58,7 +59,9 @@ impl Policy for FirstPolicy { return Err(PickError::NoCandidates); } let engine = engines[0].clone(); - if let Decision::Reject(reason) = self.admission.check(&engine, request, None)? { + if let Decision::Reject(reason) = + self.admission.check(&engine, &EngineMetrics::default())? + { return Err(PickError::AdmissionRejected(Rejection { engine: engine.id.clone(), reason, @@ -76,12 +79,7 @@ impl Policy for FirstPolicy { struct RejectAll; impl EngineAdmission for RejectAll { - fn check( - &self, - _: &Worker, - _: &PickRequest<'_>, - _: Option<&EngineReportedWorkerLoad>, - ) -> Result { + fn check(&self, _: &Worker, _: &EngineMetrics) -> Result { Ok(Decision::Reject("full".into())) } } @@ -147,6 +145,70 @@ fn body(content: &str) -> serde_json::Value { serde_json::json!({"model": "tiny", "messages": [{"role": "user", "content": content}]}) } +#[tokio::test] +async fn configured_limits_reject_before_dispatch_and_admit_after_load_drops() { + use sgl_router::policies_reorg::power_of_two::PowerOfTwoPolicy; + use sgl_router::state::load_monitor::engine_reported_load::{LoadStat, NativeCacheRankLoad}; + use std::time::Instant; + + let worker = MockWorker::start(vec![]).await; + let mut ctx = Arc::try_unwrap(context(&[("w", Stage::Plain, &worker)], vec![])) + .unwrap_or_else(|_| panic!("context is not shared yet")); + let mut policy = PowerOfTwoPolicy::new(ctx.engine_reported_load.clone()); + policy.admission = Arc::new(AdmissionLimits { + max_running_requests: Some(1), + max_kv_tokens: Some(100), + ..Default::default() + }); + ctx.chat_routing = ChatRouting::Reorg( + [( + ModelId("tiny".into()), + BucketResolver::new(vec![Bucket::new( + "default", + BucketGroups::Plain(EngineGroup { + worker_ids: Some([WorkerId("w".into())].into()), + policy: Arc::new(policy), + }), + )]), + )] + .into(), + ); + let ctx = Arc::new(ctx); + let app = build_router(ctx.clone()); + for (running, kv_tokens, expected) in [ + (1, 0, StatusCode::SERVICE_UNAVAILABLE), + (0, 100, StatusCode::SERVICE_UNAVAILABLE), + (0, 0, StatusCode::OK), + ] { + ctx.engine_reported_load.set( + &worker.url, + 0, + LoadStat { + num_running_reqs: running, + num_waiting_reqs: 0, + num_tokens: 0, + max_total_num_tokens: 1000, + native_cache: Some(NativeCacheRankLoad { + num_waiting_uncached_tokens: 0, + num_total_tokens: kv_tokens, + max_running_requests: 10, + total_prefill_uncached_tokens: 0, + total_prefill_busy_us: 0, + }), + }, + Instant::now(), + ); + let response = app.clone().oneshot(request(body("hi"))).await.unwrap(); + assert_eq!(response.status(), expected); + response.into_body().collect().await.unwrap(); + assert_eq!( + worker.captured.lock().unwrap().last_body.is_some(), + expected == StatusCode::OK + ); + assert_eq!(ctx.router_inflight_load.inflight_count(), 0); + } +} + #[tokio::test] async fn length_selects_plain_bucket_before_engine_selection() { let short_worker = MockWorker::start(vec![]).await;