[sgl-router] refactor - generalized admission policy definitions (#40271)

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Kan Wu
2026-09-21 16:26:30 +08:00
committed by GitHub
co-authored by Claude Fable 5.1
parent 70b5b03e78
commit a9871012ac
12 changed files with 445 additions and 172 deletions
+51 -45
View File
@@ -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<dyn EngineAdmission>`. 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<dyn EngineAdmission>` 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.
@@ -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<Decision, PickError>;
/// 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<u64>,
pub waiting_requests: Option<u64>,
pub kv_tokens: Option<u64>,
pub pending_prefill_tokens: Option<u64>,
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<Decision, PickError> {
/// 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<Decision, PickError>;
}
/// 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<u64>,
pub max_waiting_requests: Option<u64>,
pub max_kv_tokens: Option<u64>,
pub max_pending_prefill_tokens: Option<u64>,
pub max_inflight_requests: Option<u64>,
}
impl EngineAdmission for AdmissionLimits {
fn check(&self, _: &Worker, engine: &EngineMetrics) -> Result<Decision, PickError> {
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)
}
}
@@ -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<Option<Rejection>, 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<Rejection>,
) -> Result<Vec<Candidate<'e>>, 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<Worker>],
request: &PickRequest<'_>,
load: &EngineReportedLoadSnapshot,
) -> Result<Option<Pick>, 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";
@@ -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<EngineReportedLoadTable>) -> 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,
@@ -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";
@@ -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;
@@ -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<Decision, PickError> {
fn check(&self, engine: &Worker, _: &EngineMetrics) -> Result<Decision, PickError> {
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<Decision, PickError> {
fn check(&self, engine: &Worker, _: &EngineMetrics) -> Result<Decision, PickError> {
self.calls.lock().unwrap().push(engine.id.clone());
if self.invalid {
Err(PickError::InvalidSignal("admission input".into()))
@@ -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<Worker> {
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::<AdmissionLimits>(json!({})).unwrap(),
AdmissionLimits::default()
);
assert!(serde_json::from_value::<AdmissionLimits>(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
)),
}
}
}
@@ -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<Decision, PickError> {
fn check(&self, engine: &Worker, metrics: &EngineMetrics) -> Result<Decision, PickError> {
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 {
@@ -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<Worker> {
#[derive(Debug)]
struct ObserveAdmission {
table: Arc<EngineReportedLoadTable>,
observations: Mutex<Vec<Option<EngineReportedWorkerLoad>>>,
observations: Mutex<Vec<EngineMetrics>>,
}
impl EngineAdmission for ObserveAdmission {
fn check(
&self,
engine: &Worker,
_: &PickRequest<'_>,
load: Option<&EngineReportedWorkerLoad>,
) -> Result<Decision, PickError> {
fn check(&self, engine: &Worker, metrics: &EngineMetrics) -> Result<Decision, PickError> {
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}"
);
}
}
@@ -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<Decision, PickError> {
fn check(&self, engine: &Worker, metrics: &EngineMetrics) -> Result<Decision, PickError> {
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<Worker>,
reject_winner: bool,
@@ -346,19 +340,14 @@ struct RacingAdmission {
}
impl EngineAdmission for RacingAdmission {
fn check(
&self,
engine: &Worker,
request: &PickRequest<'_>,
_: Option<&EngineReportedWorkerLoad>,
) -> Result<Decision, PickError> {
fn check(&self, engine: &Worker, _: &EngineMetrics) -> Result<Decision, PickError> {
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,
@@ -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<Decision, PickError> {
fn check(&self, _: &Worker, _: &EngineMetrics) -> Result<Decision, PickError> {
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;