Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions docs/concepts/reliability/priority-scheduling.md
Original file line number Diff line number Diff line change
Expand Up @@ -63,11 +63,16 @@ This is why `system` and `interactive` ship with reservations by default while `

## Per-class queues

When no slot is immediately available, a request does not fail right away — it joins a **per-class FIFO queue**. Each class has its own queue with its own depth limit and its own wait timeout:
When no slot is immediately available, a request does not fail right away. It joins its class lane. The configured class sizes add up to one work-conserving partition-wide queue budget, while each class retains its own wait timeout:

- If the queue is already at its configured depth, the request is rejected immediately (**429**).
- If the shared partition queue budget is exhausted, the request is rejected immediately (**429**).
- If the request waits longer than the class's timeout, it is rejected (**408**).

With fair sharing enabled, each class lane is further divided into FIFO tenant
subqueues. One tenant's occupancy ceiling is shared across all four classes in
the partition. This keeps a noisy tenant from filling the shared queue before
another tenant can become present for weighted selection.

A client that disconnects *while queued* is not currently detected — its place is held until that timeout fires, because the cancel signal isn't yet wired to client disconnect at this stage. (The **499** code exists for this case but isn't emitted today.)

Higher classes have shorter queues and shorter timeouts (fail fast, the latency matters); lower classes have deeper queues and longer timeouts (wait patiently, throughput matters). The dispatcher drains queues in **strict priority order** — `system` first, then `interactive`, `default`, and `bulk`, fully draining a higher tier before serving a lower one. A sustained higher-priority flood *will* hold off a lower tier; the [starvation guard](#starvation-promotion) below is what keeps that from lasting forever.
Expand Down
28 changes: 26 additions & 2 deletions docs/reference/priority-scheduler.md
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ The scheduler surfaces admission and preemption outcomes as HTTP status codes. E
| Status | Condition | `X-SMG-Error-Code` | Extra headers |
|--------|-----------|--------------------|---------------|
| **503** Service Unavailable | **Preempted** — admitted, then cancelled before its first byte to make room for a higher-priority request | `scheduler_preempted` | `X-SMG-Preempted: true`, `Retry-After: 1` |
| **429** Too Many Requests | **Queue full** — the request's per-class queue is at its configured depth | `scheduler_queue_full` | — |
| **429** Too Many Requests | **Queue full**: the partition-wide queue budget is exhausted, or fair sharing is enabled and this tenant reached its queue ceiling | `scheduler_queue_full` | None |
| **408** Request Timeout | **Queue timeout** — the request waited longer than its class's `queue_timeout` | `scheduler_queue_timeout` | — |
| **499** Client Closed Request | **Client gone** — the client disconnected before admission completed (nginx convention; never actually read) | `scheduler_client_cancelled` | — |

Expand Down Expand Up @@ -113,6 +113,7 @@ tenant_policies:
fair_share:
default_weight: 1
default_output_tokens: 256
max_queued_requests_per_tenant: 64
trust_output_token_estimate_header: false
trust_request_model_header: false
tenant_weights:
Expand All @@ -130,7 +131,7 @@ Each entry under `classes` accepts the following fields. All are per-class.
|-------|------|---------|
| `reserved_floor` | integer (slots) | Minimum slots reserved for this class — the value the effective reservation never drops below. A higher class's *unused* reservation is held back from lower classes; a class's own reservation never reduces its own headroom. |
| `reserved_per_slot` | float (0.0–1.0+) | Share of live capacity reserved for this class, on top of the floor: `effective = max(reserved_floor, ceil(reserved_per_slot × capacity))`, recomputed as capacity changes so the reservation tracks the fleet. `0.0` (the default) means purely absolute (just the floor). Must be finite and ≥ 0. At startup, if the floors + shares exceed capacity the scheduler fails safe to legacy admission; at runtime a capacity dip is absorbed by clamping the lowest classes first. |
| `queue_size` | integer | Per-class queue depth limit. A request that arrives when the queue is full is rejected with **429**. |
| `queue_size` | integer | Soft per-class share of the work-conserving partition queue. The four values are summed into one partition-wide occupancy budget, so an active class may borrow an idle class's share. A request arriving when the shared budget is full is rejected with **429**. |
| `queue_timeout_secs` | integer (seconds) | How long a queued request waits before it is rejected with **408**. Must be `> 0`. |
| `starvation_threshold_secs` | integer (seconds) | Head-of-queue age past which the dispatcher promotes a waiter out of normal priority order (and lets it use a reserved-but-unused slot) to avoid starvation. Must be `> 0`. |
| `can_preempt` | boolean | Whether admissions in this class may preempt a lower-class in-flight request that has not yet emitted its first byte. |
Expand All @@ -154,6 +155,7 @@ At startup the scheduler validates:

- `queue_timeout_secs > 0` for every class (else startup fails for that class).
- `starvation_threshold_secs > 0` for every class.
- `fair_share.max_queued_requests_per_tenant > 0` when fair sharing is configured.
- The sum of all `reserved` values must not exceed the live backend capacity. On a capacity *shrink* that would otherwise break this invariant, the scheduler scales reservations down proportionally rather than locking itself out.

Any validation failure triggers the [fail-safe fallback to legacy admission](#enabling-the-scheduler).
Expand All @@ -172,6 +174,7 @@ Priority class selection remains the outer policy.
|-------|---------|---------|
| `default_weight` | `1.0` | Weight assigned to a resolved tenant absent from `tenant_weights`. Must be finite and greater than zero. |
| `default_output_tokens` | `256` | Provisional output-token charge used when no trusted estimate is available. Must be greater than zero. |
| `max_queued_requests_per_tenant` | `64` | Maximum queued requests from one canonical tenant inside one scheduler partition, shared across all four priority classes. Must be greater than zero. |
| `trust_output_token_estimate_header` | `false` | Honor `x-smg-output-token-estimate`. Enable only behind a proxy that strips client copies and injects a validated estimate. |
| `trust_request_model_header` | `false` | Honor `x-smg-request-model` for per-model profile selection. Model profiles require this setting. Enable only behind a proxy that strips client copies and injects the parsed request model. |
| `tenant_weights` | `{}` | Relative weights keyed by canonical tenant key. Every value must be finite and greater than zero. |
Expand Down Expand Up @@ -200,6 +203,13 @@ fair_share:
other_weight: 20
```

Fair-share queues retain one FIFO subqueue per real tenant. The per-tenant
ceiling is acquired before the partition-wide queue budget, so one noisy
tenant cannot consume every waiting position and prevent another tenant from
becoming eligible for weighted selection. Cancellation, timeout, and dequeue
release both occupancy claims. Priority class remains the outer dispatch
policy.

The percentages apply while the corresponding buckets are simultaneously
backlogged for that model. Idle shares are borrowed, so a free model slot is
never held empty. Requests for models without a configured profile continue to
Expand All @@ -211,6 +221,16 @@ client disconnects, the backend fails, or terminal usage is missing or
truncated, the provisional charge remains. This prevents cancellation from
evading fair-share accounting.

When fair sharing runs with enforced `engine_feedback` adaptive admission, the
adaptive controller publishes a total fleet capacity ceiling to the scheduler.
The scheduler's own in-flight counters remain authoritative for used capacity,
and only telemetry revisions wake the capacity coordinator. A local adaptive
429 safety recheck cancels the provisional fair-share reservation because no
backend work began.
This coupling requires explicit admission partitions; the unpartitioned
scheduler retains its worker-capacity ceiling because one global adaptive cap
cannot safely combine pressure from non-fungible model pools.

Fairness credit accrues only during active contention. When a new or returning
tenant becomes active, its virtual finish is rebased to the current active-set
virtual time, so idle tenants do not bank unlimited catch-up credit. The queue
Expand Down Expand Up @@ -293,12 +313,16 @@ The scheduler exposes these Prometheus metrics (see the [Metrics Reference](metr
| `smg_scheduler_utilization` | Gauge | — | Total in-flight divided by backend capacity. |
| `smg_scheduler_class_capacity_pressure` | Gauge | `class` | Normalized 0.0–1.0 pressure (worse of queue and slot pressure). |
| `smg_fair_share_charged_output_tokens_total` | Counter | `tenant` | Actual or conservative fallback output tokens charged across the process-local ledger. |
| `smg_fair_share_model_charged_output_tokens_total` | Counter | `model`, `bucket`, `tenant` | Charged output tokens within configured model profiles. `bucket` is bounded to `explicit` or `other`. |
| `smg_fair_share_virtual_finish` | Gauge | `model`, `tenant` | Active-set normalized tenant virtual finish. Flat mode uses `model="global"`. |
| `smg_fair_share_other_bucket_virtual_finish` | Gauge | `model` | Outer virtual finish of a model profile's aggregate `other` bucket. |
| `smg_fair_share_reserved_output_tokens` | Gauge | `tenant` | Provisional output-token charges held by active requests. |
| `smg_fair_share_queue_wait_seconds` | Histogram | `model`, `tenant`, `class` | Fair-share queue wait by model profile, tenant, and priority class. |
| `smg_fair_share_fallback_total` | Counter | `reason` | Settlement or estimate paths that used a configured fallback. |
| `smg_fair_share_unknown_tenant_total` | Counter | `tenant` | Requests whose canonical tenant has no explicit configured weight. |
| `smg_fair_share_queue_rejections_total` | Counter | `reason` | Fair-share queue rejections. `reason` is bounded to `tenant_limit`, `global_limit`, or `missing_tenant`. |
| `smg_scheduler_partition_static_capacity` | Gauge | `partition` | Worker-derived partition safety ceiling before adaptive admission. |
| `smg_scheduler_partition_adaptive_capacity` | Gauge | `partition` | Effective engine-feedback scheduler capacity after the adaptive ceiling. |

---

Expand Down
100 changes: 96 additions & 4 deletions model_gateway/src/middleware/scheduler/admission.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,8 @@ use tracing::trace;
use super::{
fair_share::FairShareProfile, metrics as sched_metrics, state::SchedulerState, AdmitOutcome,
Class, GlobalFairShare, RejectionReason, SchedulerError, SchedulerGuardBody,
HEADER_X_SMG_PREEMPTED, OUTPUT_TOKEN_ESTIMATE_HEADER, PRIORITY_HEADER,
HEADER_X_SMG_LOCAL_ADAPTIVE_REJECTED, HEADER_X_SMG_PREEMPTED, OUTPUT_TOKEN_ESTIMATE_HEADER,
PRIORITY_HEADER,
};
use crate::{
middleware::{
Expand All @@ -49,6 +50,20 @@ fn next_registry_id() -> RequestId {
RequestId(format!("sched-{n}"))
}

fn consume_local_adaptive_rejection_marker(
response: &mut Response,
permit: &mut super::SchedulerPermit,
) -> bool {
let rejected = response
.headers_mut()
.remove(HEADER_X_SMG_LOCAL_ADAPTIVE_REJECTED)
.is_some();
if rejected {
permit.cancel_fair_share_reservation();
}
rejected
}

/// The class a request resolves to, plus what it asked for.
struct ResolvedPriority {
/// Post-clamp class the request is admitted under.
Expand Down Expand Up @@ -187,13 +202,17 @@ pub async fn priority_admission_middleware(
)
.await
{
AdmitOutcome::Admitted(permit) => {
AdmitOutcome::Admitted(mut permit) => {
pending_guard.resolve();
Metrics::record_http_admission_admitted();
let active_guard = AdmissionActiveGuard::new();
// Hand the handler the cancel token (for preemption select!).
req.extensions_mut().insert(permit.cancel_token());
let response = next.run(req).await;
let mut response = next.run(req).await;
// The local adaptive controller rejected before backend work
// began. Return the scheduler slot normally, but erase the
// provisional token charge so retries do not accumulate debt.
consume_local_adaptive_rejection_marker(&mut response, &mut permit);
// Best-effort: the handler's PreemptionGuard tags a *pre-response*
// preemption (a 503 carrying this header), which we count as
// `preempted`. A preemption that fires after the handler produced
Expand Down Expand Up @@ -253,13 +272,18 @@ pub async fn priority_admission_middleware(
mod tests {
use std::collections::HashMap;

use axum::http::StatusCode;

use super::*;
use crate::middleware::scheduler::FairShareConfig;
use crate::middleware::scheduler::{
ClassConfig, FairShareConfig, PriorityScheduler, PrioritySchedulerYaml, SchedulerSettings,
};

fn ledger(trust_header: bool) -> GlobalFairShare {
GlobalFairShare::from_config(&FairShareConfig {
default_weight: 1.0,
default_output_tokens: 256,
max_queued_requests_per_tenant: 64,
trust_output_token_estimate_header: trust_header,
trust_request_model_header: false,
tenant_weights: HashMap::new(),
Expand All @@ -280,4 +304,72 @@ mod tests {
headers.insert(OUTPUT_TOKEN_ESTIMATE_HEADER, "128".parse().unwrap());
assert_eq!(output_token_estimate(&headers, &ledger(true)), 128);
}

#[tokio::test]
async fn local_adaptive_rejection_cancels_provisional_fair_share_charge() {
let mut classes = HashMap::new();
for class in Class::ALL {
let mut config = ClassConfig::default_for(class);
config.reserved_floor = 0;
config.reserved_per_slot = 0.0;
config.queue_size = 8;
classes.insert(class, config);
}
let yaml = PrioritySchedulerYaml {
classes,
fair_share: Some(FairShareConfig {
default_weight: 1.0,
default_output_tokens: 256,
max_queued_requests_per_tenant: 8,
trust_output_token_estimate_header: false,
trust_request_model_header: false,
tenant_weights: HashMap::from([("header:alice".to_string(), 1.0)]),
model_profiles: HashMap::new(),
}),
..Default::default()
};
let settings =
SchedulerSettings::from_cli_and_yaml(true, Class::Default, 32, Some(&yaml)).unwrap();
let ledger = Arc::new(GlobalFairShare::from_settings(&settings).unwrap());
let scheduler =
PriorityScheduler::new_with_fair_share(&settings, 1, Some(Arc::clone(&ledger)))
.unwrap();
let tenant = TenantKey::new("header:alice");
let scope_id = scheduler.fair_share_scope_for_test().unwrap();
let AdmitOutcome::Admitted(mut permit) = scheduler
.admit_for_tenant(
Class::Default,
RequestId("adaptive-rejected".into()),
CancellationToken::new(),
tenant.clone(),
256,
)
.await
else {
panic!("fair-share request should be provisionally admitted");
};
assert_eq!(ledger.snapshot(scope_id, &tenant).1, 256);

let mut response = Response::builder()
.status(StatusCode::TOO_MANY_REQUESTS)
.header(HEADER_X_SMG_LOCAL_ADAPTIVE_REJECTED, "true")
.body(Body::empty())
.unwrap();
assert!(consume_local_adaptive_rejection_marker(
&mut response,
&mut permit
));
assert!(response
.headers()
.get(HEADER_X_SMG_LOCAL_ADAPTIVE_REJECTED)
.is_none());
assert_eq!(ledger.snapshot(scope_id, &tenant), (0, 0, 0, [0; 4]));

drop(permit);
assert_eq!(
ledger.snapshot(scope_id, &tenant),
(0, 0, 0, [0; 4]),
"permit drop remains idempotent after cancellation"
);
}
}
1 change: 1 addition & 0 deletions model_gateway/src/middleware/scheduler/body.rs
Original file line number Diff line number Diff line change
Expand Up @@ -238,6 +238,7 @@ mod tests {
fair_share: Some(FairShareConfig {
default_weight: 1.0,
default_output_tokens,
max_queued_requests_per_tenant: 64,
trust_output_token_estimate_header: false,
trust_request_model_header: false,
tenant_weights: HashMap::from([("header:alice".to_string(), 1.0)]),
Expand Down
Loading
Loading