[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,
|
Admission evaluates acceptance. It does not rank engines, choose replacements,
|
||||||
change buckets, or mutate affinity.
|
change buckets, or mutate affinity.
|
||||||
|
|
||||||
| Check | Acceptance rule |
|
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.
|
||||||
| `AllowAll` | Add no acceptance constraint |
|
A cap admits while the engine's current metric is below it. Request size is
|
||||||
| `CapacityAdmission` | Projected running requests and KV tokens fit reported capacity |
|
not part of admission: buckets already select by input length and context
|
||||||
| `PendingPrefillAdmission` | Waiting uncached tokens plus incoming uncached work fit the budget |
|
capacity, and admission only observes load without reserving it.
|
||||||
| `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 |
|
|
||||||
|
|
||||||
`EngineAdmission::check(engine, request, load)` checks one engine and returns
|
| Limit | Engine metric |
|
||||||
`Allow`, `Reject(reason)`, or an error for invalid inputs. Policies attach the
|
| --- | --- |
|
||||||
checker directly as `Arc<dyn EngineAdmission>`. There is no placement setting,
|
| `max_running_requests` | Reported running requests |
|
||||||
filtering wrapper, or before/after API; each policy decides where checking
|
| `max_waiting_requests` | Reported waiting requests |
|
||||||
belongs in its selection algorithm. The `load` argument is an
|
| `max_kv_tokens` | Reported total KV tokens |
|
||||||
`Option<&EngineReportedWorkerLoad>` retained by the policy for this engine, including
|
| `max_pending_prefill_tokens` | Reported waiting uncached tokens |
|
||||||
request counts, token usage, capacity, and the report timestamp. `None` means
|
| `max_inflight_requests` | Router-local in-flight requests |
|
||||||
no usable observation, never zero load; each check defines its missing-data
|
|
||||||
behavior. Other required state handles belong to the checker.
|
```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
|
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
|
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
|
router-local active requests instead. Basic reports from older publishers are
|
||||||
still passed to admission when fresh, but do not supply native pressure metrics.
|
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
|
Power-of-two retains the selected engine's load record from selection and
|
||||||
uses per-engine uncached work when a prefix is known, and full input otherwise.
|
passes it to admission without another snapshot. A single candidate still has
|
||||||
Decode capacity uses the expected peak sequence length when available, including
|
its load read for admission, even though selection needs no comparison.
|
||||||
on a cache hit. Power-of-two retains the selected engine's load record from
|
Neither the bucket nor HTTP handler supplies observations. Synchronous checks
|
||||||
selection and passes it to admission without another snapshot. A single candidate
|
do not fetch telemetry over the network themselves.
|
||||||
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.
|
|
||||||
|
|
||||||
`AllowAll` is the default for new explicit policy attachments. It leaves health,
|
`AdmissionLimits::default()` is the default for new explicit policy attachments.
|
||||||
role, membership, and policy preferences in force. Migrated configurations must
|
It leaves health, role, membership, and policy preferences in force. Migrated
|
||||||
retain their existing capacity and configured budget checks; see compatibility
|
configurations must retain their existing capacity and configured budget checks;
|
||||||
below. Each check defines its missing-data behavior. Unknown load is not zero;
|
see compatibility below.
|
||||||
the existing capacity and pending-prefill checks allow requests without a fresh,
|
|
||||||
complete native report.
|
|
||||||
|
|
||||||
The cache policy's `worker_queue_limit` is a **soft preference**, not
|
The cache policy's `worker_queue_limit` is a **soft preference**;
|
||||||
`QueueLimitAdmission`. Saturation handling can reconsider a queued engine, but
|
`max_waiting_requests` is a hard rejection. Saturation handling can reconsider
|
||||||
cannot bypass attached hard admission.
|
a queued engine, but cannot bypass attached hard admission.
|
||||||
|
|
||||||
Admission checks observe capacity; they do not reserve it. Concurrent requests
|
Admission checks observe capacity; they do not reserve it. Concurrent requests
|
||||||
may pass against the same observation. Strict reservations would require a
|
may pass against the same observation. Strict reservations would require a
|
||||||
@@ -432,12 +437,12 @@ buckets:
|
|||||||
worker_ids: [P1, P2]
|
worker_ids: [P1, P2]
|
||||||
policy:
|
policy:
|
||||||
type: cache_aware
|
type: cache_aware
|
||||||
admission: {type: capacity}
|
admission: {max_running_requests: 64, max_kv_tokens: 1048576}
|
||||||
decode:
|
decode:
|
||||||
worker_ids: [D1, D2]
|
worker_ids: [D1, D2]
|
||||||
policy:
|
policy:
|
||||||
type: power_of_two
|
type: power_of_two
|
||||||
admission: {type: capacity}
|
admission: {max_running_requests: 64, max_kv_tokens: 1048576}
|
||||||
|
|
||||||
- id: long-context
|
- id: long-context
|
||||||
rank: 20
|
rank: 20
|
||||||
@@ -449,12 +454,12 @@ buckets:
|
|||||||
worker_ids: [P3, P4]
|
worker_ids: [P3, P4]
|
||||||
policy:
|
policy:
|
||||||
type: cache_aware
|
type: cache_aware
|
||||||
admission: {type: capacity}
|
admission: {max_running_requests: 64, max_kv_tokens: 1048576}
|
||||||
decode:
|
decode:
|
||||||
worker_ids: [D3, D4]
|
worker_ids: [D3, D4]
|
||||||
policy:
|
policy:
|
||||||
type: power_of_two
|
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
|
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.
|
queue limit, and saturation floor.
|
||||||
- Preserve session and sticky headers, idle timeouts, eviction cadence, and the
|
- Preserve session and sticky headers, idle timeouts, eviction cadence, and the
|
||||||
four sticky fallback choices. Global modes need a bucket-first migration design.
|
four sticky fallback choices. Global modes need a bucket-first migration design.
|
||||||
- Translate `--filter overloaded` and `--max-in-flight` into
|
- Map `--filter overloaded` and `--max-in-flight` to `max_inflight_requests`;
|
||||||
`InFlightLimitAdmission`, composed with other checks through `AllOfAdmission`.
|
the existing router-local counter remains the source.
|
||||||
- Preserve configured capacity, pending-prefill, and in-flight checks, including
|
- Preserve configured capacity, pending-prefill, and in-flight checks, including
|
||||||
their missing-report behavior. Power-of-two applies admission to its selected
|
their missing-report behavior. Power-of-two applies admission to its selected
|
||||||
engine; other policies explicitly place checks in their selection logic.
|
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
|
- `EngineGroup::pick` owns live candidate filtering, policy invocation, and
|
||||||
exact candidate validation, without cross-bucket fallback.
|
exact candidate validation, without cross-bucket fallback.
|
||||||
- `Policy::pick`, within-group fallback interface, per-engine `EngineAdmission::check`,
|
- `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.
|
and checks its selected engine with no replacement on rejection.
|
||||||
- Policy-owned load dependency and local observations. Power-of-two passes the
|
- Policy-owned load dependency and local observations. Power-of-two passes the
|
||||||
selected engine's load record directly to admission, without another snapshot.
|
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
|
The caller owns expiry and sweeper lifecycle. A binding may remain after a
|
||||||
later PD group fails, because it records placement rather than dispatch.
|
later PD group fails, because it records placement rather than dispatch.
|
||||||
|
|
||||||
Follow-up order: concrete admission (#40271), then bucket SLO ordering
|
Follow-up work includes bucket SLO ordering, remaining selection policies,
|
||||||
in a separate PR, followed by remaining policies and production configuration.
|
and production configuration.
|
||||||
|
|
||||||
Not yet implemented in the reorg path:
|
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.
|
- SLO estimates, targets, and bucket preference ordering.
|
||||||
- CLI/configuration parsing, validation, and model-specific construction.
|
- CLI/configuration parsing, validation, and model-specific construction.
|
||||||
The YAML above is illustrative; reorg resolvers are installed in code.
|
The YAML above is illustrative; reorg resolvers are installed in code.
|
||||||
|
|||||||
@@ -1,12 +1,17 @@
|
|||||||
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||||
// SPDX-License-Identifier: Apache-2.0
|
// 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 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 crate::workers::Worker;
|
||||||
|
|
||||||
use super::{PickError, PickRequest};
|
use super::PickError;
|
||||||
|
|
||||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
pub enum Decision {
|
pub enum Decision {
|
||||||
@@ -14,29 +19,84 @@ pub enum Decision {
|
|||||||
Reject(String),
|
Reject(String),
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Checks one engine using the load observation retained by selection.
|
/// One engine's measurements at pick time. Reported values are `None` without a
|
||||||
/// `None` means no usable load observation, never zero load. Each check defines
|
/// fresh, complete report, never zero. In-flight requests are counted by this
|
||||||
/// its missing-data behavior and owns any other state handles it needs.
|
/// router and always known.
|
||||||
/// Each policy decides when to check an engine and how to handle rejection.
|
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||||
pub trait EngineAdmission: Send + Sync + Debug {
|
pub struct EngineMetrics {
|
||||||
fn check(
|
pub running_requests: Option<u64>,
|
||||||
&self,
|
pub waiting_requests: Option<u64>,
|
||||||
engine: &Worker,
|
pub kv_tokens: Option<u64>,
|
||||||
request: &PickRequest<'_>,
|
pub pending_prefill_tokens: Option<u64>,
|
||||||
load: Option<&EngineReportedWorkerLoad>,
|
pub inflight_requests: u64,
|
||||||
) -> Result<Decision, PickError>;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug)]
|
impl EngineMetrics {
|
||||||
pub struct AllowAll;
|
/// 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 {
|
/// Checks one selected engine. Policies decide when to check and how to handle
|
||||||
fn check(
|
/// rejection; admission never selects replacements or reserves capacity.
|
||||||
&self,
|
pub trait EngineAdmission: Send + Sync + Debug {
|
||||||
_: &Worker,
|
fn check(&self, engine: &Worker, metrics: &EngineMetrics) -> Result<Decision, PickError>;
|
||||||
_: &PickRequest<'_>,
|
}
|
||||||
_: Option<&EngineReportedWorkerLoad>,
|
|
||||||
) -> 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)
|
Ok(Decision::Allow)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -26,7 +26,7 @@ use crate::state::load_monitor::engine_reported_load::{
|
|||||||
};
|
};
|
||||||
use crate::workers::Worker;
|
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::power_of_two::PowerOfTwoPolicy;
|
||||||
use super::{Pick, PickError, PickRequest, Policy, Rejection, Stage};
|
use super::{Pick, PickError, PickRequest, Policy, Rejection, Stage};
|
||||||
|
|
||||||
@@ -173,7 +173,7 @@ impl CacheAwarePolicy {
|
|||||||
fallback: Arc::new(PowerOfTwoPolicy::new(Arc::clone(&engine_load))),
|
fallback: Arc::new(PowerOfTwoPolicy::new(Arc::clone(&engine_load))),
|
||||||
engine_load,
|
engine_load,
|
||||||
config,
|
config,
|
||||||
admission: Arc::new(AllowAll),
|
admission: Arc::new(AdmissionLimits::default()),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -236,11 +236,10 @@ impl CacheAwarePolicy {
|
|||||||
fn check(
|
fn check(
|
||||||
&self,
|
&self,
|
||||||
engine: &Worker,
|
engine: &Worker,
|
||||||
request: &PickRequest<'_>,
|
|
||||||
load: &EngineReportedLoadSnapshot,
|
load: &EngineReportedLoadSnapshot,
|
||||||
) -> Result<Option<Rejection>, PickError> {
|
) -> Result<Option<Rejection>, PickError> {
|
||||||
let load = load.fresh_load_for_url(&engine.url);
|
let metrics = EngineMetrics::observe(engine, load);
|
||||||
Ok(match self.admission.check(engine, request, load)? {
|
Ok(match self.admission.check(engine, &metrics)? {
|
||||||
Decision::Allow => None,
|
Decision::Allow => None,
|
||||||
Decision::Reject(reason) => Some(Rejection {
|
Decision::Reject(reason) => Some(Rejection {
|
||||||
engine: engine.id.clone(),
|
engine: engine.id.clone(),
|
||||||
@@ -252,13 +251,12 @@ impl CacheAwarePolicy {
|
|||||||
fn admit<'e>(
|
fn admit<'e>(
|
||||||
&self,
|
&self,
|
||||||
candidates: &[Candidate<'e>],
|
candidates: &[Candidate<'e>],
|
||||||
request: &PickRequest<'_>,
|
|
||||||
load: &EngineReportedLoadSnapshot,
|
load: &EngineReportedLoadSnapshot,
|
||||||
rejections: &mut Vec<Rejection>,
|
rejections: &mut Vec<Rejection>,
|
||||||
) -> Result<Vec<Candidate<'e>>, PickError> {
|
) -> Result<Vec<Candidate<'e>>, PickError> {
|
||||||
let mut admitted = Vec::new();
|
let mut admitted = Vec::new();
|
||||||
for &candidate in candidates {
|
for &candidate in candidates {
|
||||||
match self.check(candidate.engine, request, load)? {
|
match self.check(candidate.engine, load)? {
|
||||||
None => admitted.push(candidate),
|
None => admitted.push(candidate),
|
||||||
Some(rejection) => rejections.push(rejection),
|
Some(rejection) => rejections.push(rejection),
|
||||||
}
|
}
|
||||||
@@ -316,7 +314,6 @@ impl CacheAwarePolicy {
|
|||||||
&self,
|
&self,
|
||||||
candidates: &[Candidate<'_>],
|
candidates: &[Candidate<'_>],
|
||||||
engines: &[Arc<Worker>],
|
engines: &[Arc<Worker>],
|
||||||
request: &PickRequest<'_>,
|
|
||||||
load: &EngineReportedLoadSnapshot,
|
load: &EngineReportedLoadSnapshot,
|
||||||
) -> Result<Option<Pick>, PickError> {
|
) -> Result<Option<Pick>, PickError> {
|
||||||
let limit = self.config.worker_queue_limit;
|
let limit = self.config.worker_queue_limit;
|
||||||
@@ -333,7 +330,7 @@ impl CacheAwarePolicy {
|
|||||||
evaluated.extend(&gated);
|
evaluated.extend(&gated);
|
||||||
}
|
}
|
||||||
let mut rejections = Vec::new();
|
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) {
|
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 loads = FreshLoadLookup::new(Some(load), evaluated.iter().map(|c| c.engine));
|
||||||
let guarded = self.config.pressure_guard
|
let guarded = self.config.pressure_guard
|
||||||
@@ -377,7 +374,7 @@ impl CacheAwarePolicy {
|
|||||||
if pinned {
|
if pinned {
|
||||||
let loads = FreshLoadLookup::new(Some(load), gated.iter().map(|c| c.engine));
|
let loads = FreshLoadLookup::new(Some(load), gated.iter().map(|c| c.engine));
|
||||||
let owner = self
|
let owner = self
|
||||||
.admit(&gated, request, load, &mut rejections)?
|
.admit(&gated, load, &mut rejections)?
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.min_by(|left, right| {
|
.min_by(|left, right| {
|
||||||
loads
|
loads
|
||||||
@@ -430,7 +427,7 @@ impl Policy for CacheAwarePolicy {
|
|||||||
// Capture load after remote I/O; selection and admission share it.
|
// Capture load after remote I/O; selection and admission share it.
|
||||||
let load = self.engine_load.capture_snapshot(Instant::now());
|
let load = self.engine_load.capture_snapshot(Instant::now());
|
||||||
let candidates = self.candidates(engines, request, signal.as_deref(), &load);
|
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);
|
return Ok(pick);
|
||||||
}
|
}
|
||||||
// Miss: fall back within the unqueued tier when one exists.
|
// 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)) {
|
if !pool.iter().any(|e| Arc::ptr_eq(e, &pick.engine)) {
|
||||||
return Err(PickError::OutsideCandidates(pick.engine.id.clone()));
|
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));
|
return Err(PickError::AdmissionRejected(rejection));
|
||||||
}
|
}
|
||||||
pick.reason = "no_cache_candidate";
|
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::state::load_monitor::engine_reported_load::EngineReportedLoadTable;
|
||||||
use crate::workers::Worker;
|
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};
|
use super::{Pick, PickError, PickRequest, Policy, Rejection, Stage};
|
||||||
|
|
||||||
/// Samples two distinct engines and selects the one with lower stage pressure.
|
/// 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 {
|
pub fn new(engine_load: Arc<EngineReportedLoadTable>) -> Self {
|
||||||
Self {
|
Self {
|
||||||
engine_load,
|
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 })
|
Arc::clone(if pressure.is_gt() { right } else { left })
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
let engine_load = load.fresh_load_for_url(&engine.url);
|
let metrics = EngineMetrics::observe(&engine, &load);
|
||||||
if let Decision::Reject(reason) = self.admission.check(&engine, request, engine_load)? {
|
if let Decision::Reject(reason) = self.admission.check(&engine, &metrics)? {
|
||||||
return Err(PickError::AdmissionRejected(Rejection {
|
return Err(PickError::AdmissionRejected(Rejection {
|
||||||
engine: engine.id.clone(),
|
engine: engine.id.clone(),
|
||||||
reason,
|
reason,
|
||||||
|
|||||||
@@ -15,7 +15,7 @@ use crate::state::load_monitor::engine_reported_load::{
|
|||||||
use crate::state::AffinityStore;
|
use crate::state::AffinityStore;
|
||||||
use crate::workers::Worker;
|
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::power_of_two::PowerOfTwoPolicy;
|
||||||
use super::{Pick, PickError, PickRequest, Policy, Rejection};
|
use super::{Pick, PickError, PickRequest, Policy, Rejection};
|
||||||
|
|
||||||
@@ -36,7 +36,7 @@ impl SessionAwarePolicy {
|
|||||||
store,
|
store,
|
||||||
fallback: PowerOfTwoPolicy::new(Arc::clone(&engine_load)),
|
fallback: PowerOfTwoPolicy::new(Arc::clone(&engine_load)),
|
||||||
engine_load,
|
engine_load,
|
||||||
admission: Arc::new(AllowAll),
|
admission: Arc::new(AdmissionLimits::default()),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -55,16 +55,9 @@ impl SessionAwarePolicy {
|
|||||||
))
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
fn check(
|
fn check(&self, engine: &Worker, load: &EngineReportedLoadSnapshot) -> Result<(), PickError> {
|
||||||
&self,
|
let metrics = EngineMetrics::observe(engine, load);
|
||||||
engine: &Worker,
|
match self.admission.check(engine, &metrics)? {
|
||||||
request: &PickRequest<'_>,
|
|
||||||
load: &EngineReportedLoadSnapshot,
|
|
||||||
) -> Result<(), PickError> {
|
|
||||||
match self
|
|
||||||
.admission
|
|
||||||
.check(engine, request, load.fresh_load_for_url(&engine.url))?
|
|
||||||
{
|
|
||||||
Decision::Allow => Ok(()),
|
Decision::Allow => Ok(()),
|
||||||
Decision::Reject(reason) => Err(PickError::AdmissionRejected(Rejection {
|
Decision::Reject(reason) => Err(PickError::AdmissionRejected(Rejection {
|
||||||
engine: engine.id.clone(),
|
engine: engine.id.clone(),
|
||||||
@@ -87,7 +80,7 @@ impl Policy for SessionAwarePolicy {
|
|||||||
let key = Self::assignment_key(request);
|
let key = Self::assignment_key(request);
|
||||||
if let Some(bound) = key.as_ref().and_then(|key| self.store.bound(key, engines)) {
|
if let Some(bound) = key.as_ref().and_then(|key| self.store.bound(key, engines)) {
|
||||||
let load = self.engine_load.capture_snapshot(Instant::now());
|
let load = self.engine_load.capture_snapshot(Instant::now());
|
||||||
self.check(bound, request, &load)?;
|
self.check(bound, &load)?;
|
||||||
return Ok(Pick {
|
return Ok(Pick {
|
||||||
engine: Arc::clone(bound),
|
engine: Arc::clone(bound),
|
||||||
reason: "session_primary",
|
reason: "session_primary",
|
||||||
@@ -98,7 +91,7 @@ impl Policy for SessionAwarePolicy {
|
|||||||
// checks its chosen engine before creating or replacing a binding.
|
// checks its chosen engine before creating or replacing a binding.
|
||||||
let mut pick = self.pick_fallback(engines, request).await?;
|
let mut pick = self.pick_fallback(engines, request).await?;
|
||||||
let load = self.engine_load.capture_snapshot(Instant::now());
|
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 {
|
let Some(key) = key else {
|
||||||
pick.reason = "no_session";
|
pick.reason = "no_session";
|
||||||
return Ok(pick);
|
return Ok(pick);
|
||||||
@@ -108,7 +101,7 @@ impl Policy for SessionAwarePolicy {
|
|||||||
if !Arc::ptr_eq(effective, &pick.engine) {
|
if !Arc::ptr_eq(effective, &pick.engine) {
|
||||||
// A racing first assignment wins. Check it once, without
|
// A racing first assignment wins. Check it once, without
|
||||||
// rewriting a rejected binding or retrying another engine.
|
// rewriting a rejected binding or retrying another engine.
|
||||||
self.check(effective, request, &load)?;
|
self.check(effective, &load)?;
|
||||||
pick.reason = "session_primary";
|
pick.reason = "session_primary";
|
||||||
} else {
|
} else {
|
||||||
pick.reason = "assigned";
|
pick.reason = "assigned";
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ mod discovery;
|
|||||||
mod health;
|
mod health;
|
||||||
mod policies;
|
mod policies;
|
||||||
mod policies_reorg;
|
mod policies_reorg;
|
||||||
|
mod policies_reorg_admission;
|
||||||
mod policies_reorg_cache_aware;
|
mod policies_reorg_cache_aware;
|
||||||
mod policies_reorg_load;
|
mod policies_reorg_load;
|
||||||
mod policies_reorg_power_of_two;
|
mod policies_reorg_power_of_two;
|
||||||
|
|||||||
@@ -8,9 +8,10 @@ use sgl_router::buckets_reorg::{
|
|||||||
Bucket, BucketGroups, BucketRequest, BucketResolver, EngineGroup, TokenLimits,
|
Bucket, BucketGroups, BucketRequest, BucketResolver, EngineGroup, TokenLimits,
|
||||||
};
|
};
|
||||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec};
|
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::policies_reorg::{Pick, PickError, PickRequest, Policy, Rejection, Stage};
|
||||||
use sgl_router::state::load_monitor::engine_reported_load::EngineReportedWorkerLoad;
|
|
||||||
use sgl_router::workers::{Worker, WorkerRegistry};
|
use sgl_router::workers::{Worker, WorkerRegistry};
|
||||||
|
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
@@ -25,7 +26,7 @@ struct TestPolicy {
|
|||||||
impl Default for TestPolicy {
|
impl Default for TestPolicy {
|
||||||
fn default() -> Self {
|
fn default() -> Self {
|
||||||
Self {
|
Self {
|
||||||
admission: Arc::new(AllowAll),
|
admission: Arc::new(AdmissionLimits::default()),
|
||||||
result: None,
|
result: None,
|
||||||
miss: false,
|
miss: false,
|
||||||
invalid: false,
|
invalid: false,
|
||||||
@@ -52,7 +53,9 @@ impl Policy for TestPolicy {
|
|||||||
return Err(PickError::NoCandidates);
|
return Err(PickError::NoCandidates);
|
||||||
}
|
}
|
||||||
let engine = self.result.clone().unwrap_or_else(|| engines[0].clone());
|
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 {
|
return Err(PickError::AdmissionRejected(Rejection {
|
||||||
engine: engine.id.clone(),
|
engine: engine.id.clone(),
|
||||||
reason,
|
reason,
|
||||||
@@ -70,12 +73,7 @@ impl Policy for TestPolicy {
|
|||||||
struct Reject(&'static str);
|
struct Reject(&'static str);
|
||||||
|
|
||||||
impl EngineAdmission for Reject {
|
impl EngineAdmission for Reject {
|
||||||
fn check(
|
fn check(&self, engine: &Worker, _: &EngineMetrics) -> Result<Decision, PickError> {
|
||||||
&self,
|
|
||||||
engine: &Worker,
|
|
||||||
_: &PickRequest<'_>,
|
|
||||||
_: Option<&EngineReportedWorkerLoad>,
|
|
||||||
) -> Result<Decision, PickError> {
|
|
||||||
Ok(if engine.id.0 == self.0 {
|
Ok(if engine.id.0 == self.0 {
|
||||||
Decision::Reject("full".into())
|
Decision::Reject("full".into())
|
||||||
} else {
|
} else {
|
||||||
@@ -437,12 +435,7 @@ async fn power_of_two_checks_selected_engine_and_propagates_rejection_without_fa
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl EngineAdmission for Check {
|
impl EngineAdmission for Check {
|
||||||
fn check(
|
fn check(&self, engine: &Worker, _: &EngineMetrics) -> Result<Decision, PickError> {
|
||||||
&self,
|
|
||||||
engine: &Worker,
|
|
||||||
_: &PickRequest<'_>,
|
|
||||||
_: Option<&EngineReportedWorkerLoad>,
|
|
||||||
) -> Result<Decision, PickError> {
|
|
||||||
self.calls.lock().unwrap().push(engine.id.clone());
|
self.calls.lock().unwrap().push(engine.id.clone());
|
||||||
if self.invalid {
|
if self.invalid {
|
||||||
Err(PickError::InvalidSignal("admission input".into()))
|
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::config::AffinityConfig;
|
||||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec};
|
use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec};
|
||||||
use sgl_router::policies::prefix_provider::RadixTreePrefixProvider;
|
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::cache_aware::{CacheAwarePolicy, CacheSource, PrefixMemo};
|
||||||
use sgl_router::policies_reorg::power_of_two::PowerOfTwoPolicy;
|
use sgl_router::policies_reorg::power_of_two::PowerOfTwoPolicy;
|
||||||
use sgl_router::policies_reorg::{PickError, PickRequest, Policy, Stage};
|
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,
|
compute_block_hashes, compute_block_hashes_bigram, BlockSizeOracle, HashTree, KvWorkerId,
|
||||||
};
|
};
|
||||||
use sgl_router::state::load_monitor::engine_reported_load::{
|
use sgl_router::state::load_monitor::engine_reported_load::{
|
||||||
EngineReportedLoadTable, EngineReportedWorkerLoad, LoadStat, NativeCacheRankLoad,
|
EngineReportedLoadTable, LoadStat, NativeCacheRankLoad,
|
||||||
};
|
};
|
||||||
use sgl_router::workers::Worker;
|
use sgl_router::workers::Worker;
|
||||||
|
|
||||||
@@ -116,16 +116,11 @@ impl Reject {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl EngineAdmission for Reject {
|
impl EngineAdmission for Reject {
|
||||||
fn check(
|
fn check(&self, engine: &Worker, metrics: &EngineMetrics) -> Result<Decision, PickError> {
|
||||||
&self,
|
|
||||||
engine: &Worker,
|
|
||||||
_: &PickRequest<'_>,
|
|
||||||
load: Option<&EngineReportedWorkerLoad>,
|
|
||||||
) -> Result<Decision, PickError> {
|
|
||||||
self.calls
|
self.calls
|
||||||
.lock()
|
.lock()
|
||||||
.unwrap()
|
.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 {
|
Ok(if engine.id.0 == self.id {
|
||||||
Decision::Reject("full".into())
|
Decision::Reject("full".into())
|
||||||
} else {
|
} else {
|
||||||
|
|||||||
@@ -5,12 +5,10 @@ use std::sync::{Arc, Mutex};
|
|||||||
use std::time::{Duration, Instant};
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec};
|
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::power_of_two::PowerOfTwoPolicy;
|
||||||
use sgl_router::policies_reorg::{PickError, PickRequest, Policy, Stage};
|
use sgl_router::policies_reorg::{PickError, PickRequest, Policy, Stage};
|
||||||
use sgl_router::state::load_monitor::engine_reported_load::{
|
use sgl_router::state::load_monitor::engine_reported_load::{EngineReportedLoadTable, LoadStat};
|
||||||
EngineReportedLoadTable, EngineReportedWorkerLoad, LoadStat,
|
|
||||||
};
|
|
||||||
use sgl_router::workers::Worker;
|
use sgl_router::workers::Worker;
|
||||||
|
|
||||||
const URL: &str = "http://engine";
|
const URL: &str = "http://engine";
|
||||||
@@ -43,21 +41,16 @@ fn engine() -> Arc<Worker> {
|
|||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
struct ObserveAdmission {
|
struct ObserveAdmission {
|
||||||
table: Arc<EngineReportedLoadTable>,
|
table: Arc<EngineReportedLoadTable>,
|
||||||
observations: Mutex<Vec<Option<EngineReportedWorkerLoad>>>,
|
observations: Mutex<Vec<EngineMetrics>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl EngineAdmission for ObserveAdmission {
|
impl EngineAdmission for ObserveAdmission {
|
||||||
fn check(
|
fn check(&self, engine: &Worker, metrics: &EngineMetrics) -> Result<Decision, PickError> {
|
||||||
&self,
|
|
||||||
engine: &Worker,
|
|
||||||
_: &PickRequest<'_>,
|
|
||||||
load: Option<&EngineReportedWorkerLoad>,
|
|
||||||
) -> Result<Decision, PickError> {
|
|
||||||
assert_eq!(engine.url, URL);
|
assert_eq!(engine.url, URL);
|
||||||
// A new report arriving after selection must not change the observation
|
// A new report arriving after selection must not change the observation
|
||||||
// supplied to admission. The next pick should read the new report.
|
// supplied to admission. The next pick should read the new report.
|
||||||
report(&self.table, 0, 99, Instant::now());
|
report(&self.table, 0, 99, Instant::now());
|
||||||
self.observations.lock().unwrap().push(load.cloned());
|
self.observations.lock().unwrap().push(*metrics);
|
||||||
Ok(Decision::Allow)
|
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();
|
let observations = admission.observations.lock().unwrap();
|
||||||
assert_eq!(observations.len(), 2);
|
assert_eq!(observations.len(), 2);
|
||||||
|
// Basic reports carry request counts but no KV or pending-prefill tokens.
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
observations[0],
|
observations[0],
|
||||||
Some(EngineReportedWorkerLoad {
|
EngineMetrics {
|
||||||
num_running_reqs: 4,
|
running_requests: Some(4),
|
||||||
num_waiting_reqs: 4,
|
waiting_requests: Some(4),
|
||||||
num_tokens: 60,
|
..EngineMetrics::default()
|
||||||
max_total_num_tokens: 200,
|
}
|
||||||
captured_at: first_at,
|
|
||||||
})
|
|
||||||
);
|
);
|
||||||
assert_eq!(observations[1].as_ref().unwrap().num_running_reqs, 102);
|
assert_eq!(observations[1].running_requests, Some(102));
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[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))
|
.pick(&[engine()], &PickRequest::new(&model, Stage::Plain, 10))
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.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 std::time::{Duration, Instant};
|
||||||
|
|
||||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec};
|
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::session_aware::SessionAwarePolicy;
|
||||||
use sgl_router::policies_reorg::{PickError, PickRequest, Policy, Stage};
|
use sgl_router::policies_reorg::{PickError, PickRequest, Policy, Stage};
|
||||||
use sgl_router::state::load_monitor::engine_reported_load::{
|
use sgl_router::state::load_monitor::engine_reported_load::{EngineReportedLoadTable, LoadStat};
|
||||||
EngineReportedLoadTable, EngineReportedWorkerLoad, LoadStat,
|
|
||||||
};
|
|
||||||
use sgl_router::state::load_monitor::router_inflight_load::MockClock;
|
use sgl_router::state::load_monitor::router_inflight_load::MockClock;
|
||||||
use sgl_router::state::AffinityStore;
|
use sgl_router::state::AffinityStore;
|
||||||
use sgl_router::workers::Worker;
|
use sgl_router::workers::Worker;
|
||||||
@@ -52,16 +50,11 @@ struct Admission {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl EngineAdmission for Admission {
|
impl EngineAdmission for Admission {
|
||||||
fn check(
|
fn check(&self, engine: &Worker, metrics: &EngineMetrics) -> Result<Decision, PickError> {
|
||||||
&self,
|
|
||||||
engine: &Worker,
|
|
||||||
_: &PickRequest<'_>,
|
|
||||||
load: Option<&EngineReportedWorkerLoad>,
|
|
||||||
) -> Result<Decision, PickError> {
|
|
||||||
self.calls
|
self.calls
|
||||||
.lock()
|
.lock()
|
||||||
.unwrap()
|
.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) {
|
if self.invalid.load(Ordering::Relaxed) {
|
||||||
return Err(PickError::InvalidSignal("invalid admission input".into()));
|
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)]
|
#[derive(Debug)]
|
||||||
struct RacingAdmission {
|
struct RacingAdmission {
|
||||||
|
model: ModelId,
|
||||||
competitor: SessionAwarePolicy,
|
competitor: SessionAwarePolicy,
|
||||||
winner: Arc<Worker>,
|
winner: Arc<Worker>,
|
||||||
reject_winner: bool,
|
reject_winner: bool,
|
||||||
@@ -346,19 +340,14 @@ struct RacingAdmission {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl EngineAdmission for RacingAdmission {
|
impl EngineAdmission for RacingAdmission {
|
||||||
fn check(
|
fn check(&self, engine: &Worker, _: &EngineMetrics) -> Result<Decision, PickError> {
|
||||||
&self,
|
|
||||||
engine: &Worker,
|
|
||||||
request: &PickRequest<'_>,
|
|
||||||
_: Option<&EngineReportedWorkerLoad>,
|
|
||||||
) -> Result<Decision, PickError> {
|
|
||||||
self.calls.lock().unwrap().push(engine.id.0.clone());
|
self.calls.lock().unwrap().push(engine.id.0.clone());
|
||||||
if engine.id.0 == "a" {
|
if engine.id.0 == "a" {
|
||||||
// Complete a competing first request after this request selected a,
|
// Complete a competing first request after this request selected a,
|
||||||
// but before it can commit. The effective binding is now b.
|
// but before it can commit. The effective binding is now b.
|
||||||
futures::executor::block_on(
|
futures::executor::block_on(
|
||||||
self.competitor
|
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 {
|
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 (mut policy, store) = policy();
|
||||||
let engines = [engine("a", 0), engine("b", 9)];
|
let engines = [engine("a", 0), engine("b", 9)];
|
||||||
let admission = Arc::new(RacingAdmission {
|
let admission = Arc::new(RacingAdmission {
|
||||||
|
model: ModelId("m".into()),
|
||||||
competitor: SessionAwarePolicy::new(store.clone(), EngineReportedLoadTable::new()),
|
competitor: SessionAwarePolicy::new(store.clone(), EngineReportedLoadTable::new()),
|
||||||
winner: engines[1].clone(),
|
winner: engines[1].clone(),
|
||||||
reject_winner,
|
reject_winner,
|
||||||
|
|||||||
@@ -6,10 +6,11 @@ use crate::common::mock_worker::MockWorker;
|
|||||||
use futures::future::BoxFuture;
|
use futures::future::BoxFuture;
|
||||||
use sgl_router::buckets_reorg::{Bucket, BucketGroups, BucketResolver, EngineGroup};
|
use sgl_router::buckets_reorg::{Bucket, BucketGroups, BucketResolver, EngineGroup};
|
||||||
use sgl_router::policies::PolicyRegistry;
|
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::policies_reorg::{Pick, PickError, PickRequest, Policy, Rejection, Stage};
|
||||||
use sgl_router::server::app_context::ChatRouting;
|
use sgl_router::server::app_context::ChatRouting;
|
||||||
use sgl_router::state::load_monitor::engine_reported_load::EngineReportedWorkerLoad;
|
|
||||||
use std::sync::Mutex;
|
use std::sync::Mutex;
|
||||||
|
|
||||||
mod session_aware;
|
mod session_aware;
|
||||||
@@ -27,7 +28,7 @@ struct FirstPolicy {
|
|||||||
impl Default for FirstPolicy {
|
impl Default for FirstPolicy {
|
||||||
fn default() -> Self {
|
fn default() -> Self {
|
||||||
Self {
|
Self {
|
||||||
admission: Arc::new(AllowAll),
|
admission: Arc::new(AdmissionLimits::default()),
|
||||||
miss: false,
|
miss: false,
|
||||||
invalid: false,
|
invalid: false,
|
||||||
calls: Mutex::new(Vec::new()),
|
calls: Mutex::new(Vec::new()),
|
||||||
@@ -58,7 +59,9 @@ impl Policy for FirstPolicy {
|
|||||||
return Err(PickError::NoCandidates);
|
return Err(PickError::NoCandidates);
|
||||||
}
|
}
|
||||||
let engine = engines[0].clone();
|
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 {
|
return Err(PickError::AdmissionRejected(Rejection {
|
||||||
engine: engine.id.clone(),
|
engine: engine.id.clone(),
|
||||||
reason,
|
reason,
|
||||||
@@ -76,12 +79,7 @@ impl Policy for FirstPolicy {
|
|||||||
struct RejectAll;
|
struct RejectAll;
|
||||||
|
|
||||||
impl EngineAdmission for RejectAll {
|
impl EngineAdmission for RejectAll {
|
||||||
fn check(
|
fn check(&self, _: &Worker, _: &EngineMetrics) -> Result<Decision, PickError> {
|
||||||
&self,
|
|
||||||
_: &Worker,
|
|
||||||
_: &PickRequest<'_>,
|
|
||||||
_: Option<&EngineReportedWorkerLoad>,
|
|
||||||
) -> Result<Decision, PickError> {
|
|
||||||
Ok(Decision::Reject("full".into()))
|
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}]})
|
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]
|
#[tokio::test]
|
||||||
async fn length_selects_plain_bucket_before_engine_selection() {
|
async fn length_selects_plain_bucket_before_engine_selection() {
|
||||||
let short_worker = MockWorker::start(vec![]).await;
|
let short_worker = MockWorker::start(vec![]).await;
|
||||||
|
|||||||
Reference in New Issue
Block a user