[sgl-router] refactor - generalized admission policy definitions (#40271)
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Fable 5.1
parent
70b5b03e78
commit
a9871012ac
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user