diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index e47620b159..d249ee1a82 100755 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -528,6 +528,7 @@ struct Router { kv_connector_annotation: String, kv_engine_id_annotation: String, mm_per_request_image_limit: Option, + pd_admission_wait_secs: u64, } impl Router { @@ -848,6 +849,7 @@ impl Router { .worker_overload_protection(self.worker_overload_protection) .disable_load_monitoring(self.disable_load_monitoring) .load_monitor_interval_secs(self.load_monitor_interval) + .pd_admission_wait_secs(self.pd_admission_wait_secs) .max_concurrent_requests(self.max_concurrent_requests) .queue_size(self.queue_size) .queue_timeout_secs(self.queue_timeout_secs) @@ -1098,6 +1100,7 @@ impl Router { kv_connector_annotation = String::from("smg.ai/kv-connector"), kv_engine_id_annotation = String::from("smg.ai/kv-engine-id"), mm_per_request_image_limit = None, + pd_admission_wait_secs = 30, ))] #[expect(clippy::too_many_arguments)] #[expect( @@ -1251,6 +1254,7 @@ impl Router { kv_connector_annotation: String, kv_engine_id_annotation: String, mm_per_request_image_limit: Option, + pd_admission_wait_secs: u64, ) -> PyResult { let mut all_urls = worker_urls.clone(); @@ -1418,6 +1422,7 @@ impl Router { kv_connector_annotation, kv_engine_id_annotation, mm_per_request_image_limit, + pd_admission_wait_secs, }) } diff --git a/bindings/python/src/smg/router_args.py b/bindings/python/src/smg/router_args.py index d64fb7091a..98bd8afad9 100644 --- a/bindings/python/src/smg/router_args.py +++ b/bindings/python/src/smg/router_args.py @@ -251,6 +251,9 @@ class RouterArgs: kv_engine_id_annotation: str = "smg.ai/kv-engine-id" # Per-request image-count limit replacing model spec limits; None keeps spec limits mm_per_request_image_limit: int | None = None + # Seconds a PD dispatch waits for a slot in the decode engine's running + # window before shedding; 0 sheds immediately + pd_admission_wait_secs: int = 30 @staticmethod def add_cli_args( @@ -728,6 +731,18 @@ def add_cli_args( " retries" ), ) + routing_group.add_argument( + f"--{prefix}pd-admission-wait-secs", + type=int, + default=RouterArgs.pd_admission_wait_secs, + help=( + "Seconds a prefill/decode dispatch waits for a free slot in" + " the decode engine's running window before shedding with 503" + " worker_overload_protection_shed. Keep it well under the" + " engine's bootstrap deadline. 0 sheds immediately; engines" + " that report no running window are never gated" + ), + ) routing_group.add_argument( f"--{prefix}stream-body-stall-timeout-secs", type=int, diff --git a/bindings/python/tests/test_arg_parser.py b/bindings/python/tests/test_arg_parser.py index 3ba1491df9..3d7dea98fb 100644 --- a/bindings/python/tests/test_arg_parser.py +++ b/bindings/python/tests/test_arg_parser.py @@ -1440,6 +1440,7 @@ class TestRouterArgsFieldOrder: "kv_connector_annotation", "kv_engine_id_annotation", "mm_per_request_image_limit", + "pd_admission_wait_secs", ] def test_complete_field_sequence_is_frozen(self): diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/loads.py b/grpc_servicer/smg_grpc_servicer/tokenspeed/loads.py new file mode 100644 index 0000000000..ad803a72fa --- /dev/null +++ b/grpc_servicer/smg_grpc_servicer/tokenspeed/loads.py @@ -0,0 +1,57 @@ +"""Conversion of TokenSpeed scheduler load replies into the gateway's protobuf. + +Kept free of engine imports so the field mapping can be unit-tested without +TokenSpeed installed (see grpc_servicer/tests/test_tokenspeed_loads.py), +the same split the SGLang side uses. +""" + +from __future__ import annotations + +from typing import Any + +from smg_grpc_proto import tokenspeed_scheduler_pb2 + + +def running_window(server_args: Any) -> int: + """The scheduler's configured admission window, or 0 when unknown. + + ``GetLoadReqOutput`` carries no admission bound, so the window has to come + from the server args. TokenSpeed spells it ``max_num_seqs``; the protobuf + field is the one SGLang fills with the same number, so a frontend reads a + single name across both engines. It matters most on a disaggregated decode + worker: that window is exactly the bound its prefill peer's bootstrap + deadline is racing, and a frontend that cannot see it cannot pace its + dispatch. + """ + return int( + getattr(server_args, "max_num_seqs", 0) + or getattr(server_args, "max_running_requests", 0) + or 0 + ) + + +def convert_load_to_protobuf( + load_output: Any, + *, + page_size: int, + max_total_num_tokens: int, + max_running_requests: int, +) -> tokenspeed_scheduler_pb2.SchedulerLoad: + """Convert one rank's ``GetLoadReqOutput`` to a protobuf ``SchedulerLoad``. + + The reply counts ``num_reqs`` as running + waiting and reports KV usage in + pages, so running requests and used tokens are derived here. + """ + num_waiting_reqs = int(load_output.num_waiting_reqs) + num_total_reqs = int(load_output.num_reqs) + num_used_tokens = int(load_output.num_pages) * page_size + return tokenspeed_scheduler_pb2.SchedulerLoad( + dp_rank=int(load_output.dp_rank), + num_running_reqs=max(0, num_total_reqs - num_waiting_reqs), + num_waiting_reqs=num_waiting_reqs, + num_total_reqs=num_total_reqs, + num_used_tokens=num_used_tokens, + max_total_num_tokens=max_total_num_tokens, + token_usage=(num_used_tokens / max_total_num_tokens if max_total_num_tokens > 0 else 0.0), + max_running_requests=max_running_requests, + ) diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py b/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py index 991626c902..23194ee9d5 100644 --- a/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py +++ b/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py @@ -47,6 +47,7 @@ from smg_grpc_servicer.tokenizer_bundle import CHUNK_SIZE, build_tokenizer_zip from smg_grpc_servicer.tokenspeed.health_servicer import TokenSpeedHealthServicer from smg_grpc_servicer.tokenspeed.kv_events import resolve_kv_events_config +from smg_grpc_servicer.tokenspeed.loads import convert_load_to_protobuf, running_window if TYPE_CHECKING: # Type-only — keeps these out of the cold-path graph when the servicer is @@ -674,31 +675,20 @@ async def GetLoads( or getattr(self.async_llm.server_args, "max_total_num_tokens", 0) or 0 ) - - scheduler_loads: list[tokenspeed_scheduler_pb2.SchedulerLoad] = [] - total_running = 0 - total_waiting = 0 - token_usages: list[float] = [] - for lo in load_outputs: - num_running = max(0, int(lo.num_reqs) - int(lo.num_waiting_reqs)) - num_used_tokens = int(lo.num_pages) * page_size - token_usage = ( - num_used_tokens / max_total_num_tokens if max_total_num_tokens > 0 else 0.0 - ) - scheduler_loads.append( - tokenspeed_scheduler_pb2.SchedulerLoad( - dp_rank=int(lo.dp_rank), - num_running_reqs=num_running, - num_waiting_reqs=int(lo.num_waiting_reqs), - num_total_reqs=int(lo.num_reqs), - num_used_tokens=num_used_tokens, - max_total_num_tokens=max_total_num_tokens, - token_usage=token_usage, - ) + max_running_requests = running_window(self.async_llm.server_args) + + scheduler_loads = [ + convert_load_to_protobuf( + lo, + page_size=page_size, + max_total_num_tokens=max_total_num_tokens, + max_running_requests=max_running_requests, ) - total_running += num_running - total_waiting += int(lo.num_waiting_reqs) - token_usages.append(token_usage) + for lo in load_outputs + ] + token_usages = [load.token_usage for load in scheduler_loads] + total_running = sum(load.num_running_reqs for load in scheduler_loads) + total_waiting = sum(load.num_waiting_reqs for load in scheduler_loads) aggregate = tokenspeed_scheduler_pb2.AggregateMetrics( total_running_reqs=total_running, diff --git a/grpc_servicer/tests/test_tokenspeed_loads.py b/grpc_servicer/tests/test_tokenspeed_loads.py new file mode 100644 index 0000000000..8d70b32dd2 --- /dev/null +++ b/grpc_servicer/tests/test_tokenspeed_loads.py @@ -0,0 +1,75 @@ +"""GetLoads must report the scheduler's configured running window. + +``GetLoadReqOutput`` carries no admission bound, so the servicer left +``max_running_requests`` at zero on every TokenSpeed load report. A frontend +could not tell how many requests the worker will run at once — and on a +disaggregated decode worker that window is exactly the bound its prefill +peer's bootstrap deadline is racing. + +Run with: pytest grpc_servicer/tests/test_tokenspeed_loads.py +""" + +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest + +pytest.importorskip("smg_grpc_proto") + + +@pytest.fixture(scope="module") +def loads_mod(): + """Load loads.py by path: the tokenspeed package __init__ imports the engine.""" + path = Path(__file__).resolve().parent.parent / "smg_grpc_servicer" / "tokenspeed" / "loads.py" + spec = importlib.util.spec_from_file_location("tokenspeed_loads_under_test", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class TestRunningWindow: + def test_max_num_seqs_is_the_window(self, loads_mod): + assert loads_mod.running_window(SimpleNamespace(max_num_seqs=16)) == 16 + + def test_max_running_requests_spelling_is_accepted(self, loads_mod): + assert loads_mod.running_window(SimpleNamespace(max_running_requests=64)) == 64 + + def test_absent_window_reports_zero_rather_than_guessing(self, loads_mod): + assert loads_mod.running_window(SimpleNamespace()) == 0 + assert loads_mod.running_window(SimpleNamespace(max_num_seqs=None)) == 0 + + +class TestLoadConversion: + def test_window_and_derived_counters_reach_the_protobuf(self, loads_mod): + load_output = SimpleNamespace(dp_rank=1, num_reqs=5, num_waiting_reqs=2, num_pages=8) + + load = loads_mod.convert_load_to_protobuf( + load_output, + page_size=16, + max_total_num_tokens=1024, + max_running_requests=16, + ) + + assert load.max_running_requests == 16 + assert load.dp_rank == 1 + # num_reqs counts running + waiting; used tokens come from pages. + assert load.num_running_reqs == 3 + assert load.num_waiting_reqs == 2 + assert load.num_total_reqs == 5 + assert load.num_used_tokens == 128 + assert load.max_total_num_tokens == 1024 + assert load.token_usage == pytest.approx(0.125) + + def test_unknown_capacity_reports_zero_usage_instead_of_dividing(self, loads_mod): + load_output = SimpleNamespace(dp_rank=0, num_reqs=1, num_waiting_reqs=0, num_pages=4) + + load = loads_mod.convert_load_to_protobuf( + load_output, + page_size=1, + max_total_num_tokens=0, + max_running_requests=0, + ) + + assert load.token_usage == 0.0 + assert load.max_running_requests == 0 diff --git a/model_gateway/src/app_context.rs b/model_gateway/src/app_context.rs index 3a91d7cd0e..1c34e8308f 100644 --- a/model_gateway/src/app_context.rs +++ b/model_gateway/src/app_context.rs @@ -21,7 +21,9 @@ use crate::{ policies::PolicyRegistry, rate_limit::RateLimitManager, routers::{ - common::{openai_bridge::FormatRegistry, overload, realtime::RealtimeRegistry}, + common::{ + openai_bridge::FormatRegistry, overload, pd_admission, realtime::RealtimeRegistry, + }, gateway::Gateway, grpc::multimodal::MultimodalConfigRegistry, }, @@ -631,6 +633,9 @@ impl AppContextBuilder { // The overload shed advertises the poll interval as Retry-After — the // veto cannot clear between polls. overload::set_shed_retry_after_secs(config.load_monitor_interval_secs); + // PD dispatch waits here, not in the decode engine's queue, when the + // pair's running window is full. + pd_admission::set_pd_admission_wait_secs(config.pd_admission_wait_secs); // Wire the backend load-snapshot feed into every policy that consumes // it; the monitor polls every group by default, conditionally under // `--disable-load-monitoring`. diff --git a/model_gateway/src/config/builder.rs b/model_gateway/src/config/builder.rs index 6dd32c9cc7..8e02f07fd6 100644 --- a/model_gateway/src/config/builder.rs +++ b/model_gateway/src/config/builder.rs @@ -263,6 +263,11 @@ impl RouterConfigBuilder { self } + pub fn pd_admission_wait_secs(mut self, secs: u64) -> Self { + self.config.pd_admission_wait_secs = secs; + self + } + pub fn disable_load_monitoring(mut self, disabled: bool) -> Self { self.config.disable_load_monitoring = disabled; self diff --git a/model_gateway/src/config/types.rs b/model_gateway/src/config/types.rs index 98cabc7996..c1a35d3fe1 100755 --- a/model_gateway/src/config/types.rs +++ b/model_gateway/src/config/types.rs @@ -10,6 +10,7 @@ pub use smg_data_connector::{ use super::{validation::ConfigValidator, ConfigResult}; use crate::{ + routers::common::pd_admission::DEFAULT_PD_ADMISSION_WAIT_SECS, tenant::DEFAULT_TENANT_HEADER_NAME, worker::{ConnectionMode, RuntimeType}, }; @@ -91,6 +92,14 @@ pub struct RouterConfig { pub job_queue_concurrency: usize, #[serde(default = "default_load_monitor_interval_secs")] pub load_monitor_interval_secs: u64, + /// How long a disaggregated (PD) dispatch waits for a slot in the decode + /// engine's running window before shedding. Must stay well under the + /// engine's bootstrap deadline (120s on TokenSpeed): a request that waits + /// out this budget and then dispatches still has the whole deadline ahead + /// of it. `0` sheds immediately instead of waiting. Ignored for engines + /// that report no running window. + #[serde(default = "default_pd_admission_wait_secs")] + pub pd_admission_wait_secs: u64, /// Restore the conditional load-monitor poll gate: only poll worker groups /// when a load-aware routing policy, `engine_metrics`, or overload /// protection needs the data. Default `false` — the monitor polls every @@ -333,6 +342,10 @@ fn default_load_monitor_interval_secs() -> u64 { 10 } +fn default_pd_admission_wait_secs() -> u64 { + DEFAULT_PD_ADMISSION_WAIT_SECS +} + fn default_job_queue_capacity() -> usize { 1000 } @@ -1064,6 +1077,7 @@ impl Default for RouterConfig { job_queue_capacity: default_job_queue_capacity(), job_queue_concurrency: default_job_queue_concurrency(), load_monitor_interval_secs: 10, + pd_admission_wait_secs: default_pd_admission_wait_secs(), disable_load_monitoring: false, worker_overload_protection: false, worker_overload_waiting_requests: None, @@ -1387,6 +1401,27 @@ mod tests { assert_eq!(with.stream_body_stall_timeout_secs, 0); } + #[test] + fn test_pd_admission_wait_serde_default_and_roundtrip() { + // Config files predating the field deserialize to the 30s default. + let mut json: serde_json::Value = serde_json::to_value(RouterConfig::default()).unwrap(); + json.as_object_mut() + .unwrap() + .remove("pd_admission_wait_secs") + .unwrap(); + let without: RouterConfig = serde_json::from_value(json).unwrap(); + assert_eq!(without.pd_admission_wait_secs, 30); + + // The shed-immediately zero round-trips instead of reverting. + let config = RouterConfig::builder() + .regular_mode(vec![]) + .pd_admission_wait_secs(0) + .build_unchecked(); + let json = serde_json::to_string(&config).unwrap(); + let with: RouterConfig = serde_json::from_str(&json).unwrap(); + assert_eq!(with.pd_admission_wait_secs, 0); + } + #[test] fn alias_field_spellings_deserialize_and_serialize_canonically() { // Config files may use the intent-revealing spellings; alias in, diff --git a/model_gateway/src/main.rs b/model_gateway/src/main.rs index 2fbf091278..d77c20fb91 100644 --- a/model_gateway/src/main.rs +++ b/model_gateway/src/main.rs @@ -535,6 +535,16 @@ struct CliArgs { #[arg(long, default_value_t = false, help_heading = "Load Monitoring")] engine_metrics: bool, + /// Seconds a prefill/decode dispatch waits for a free slot in the decode + /// engine's running window (--max-num-seqs / --max-running-requests) + /// before shedding with 503 worker_overload_protection_shed. Keep it well + /// under the engine's bootstrap deadline (120s on TokenSpeed) so a + /// request that waits still dispatches with the deadline ahead of it. 0 + /// sheds immediately. Engines that report no running window are never + /// gated + #[arg(long, default_value_t = 30, help_heading = "Load Monitoring")] + pd_admission_wait_secs: u64, + /// TTL in seconds for event-driven cache-aware indexer entries: entries /// neither stored nor read by a query within this window are pruned. /// Bounds index growth when a backend stops emitting removal events. @@ -1814,6 +1824,7 @@ impl CliArgs { .job_queue_capacity(self.job_queue_capacity) .job_queue_concurrency(self.job_queue_concurrency) .load_monitor_interval_secs(self.load_monitor_interval) + .pd_admission_wait_secs(self.pd_admission_wait_secs) .disable_load_monitoring(self.disable_load_monitoring) .worker_overload_protection(self.worker_overload_protection) .worker_overload_waiting_requests(self.worker_overload_waiting_requests) @@ -2267,6 +2278,20 @@ mod tests { } } + /// The PD admission wait is a router-only setting and must flow into + /// `RouterConfig`, where the dispatch path latches it at startup. + #[test] + fn pd_admission_wait_flag_flows_into_router_config() { + let cli = cli_args_from(&["--pd-admission-wait-secs", "5"]); + let router_config = cli.to_router_config(vec![], vec![]).unwrap(); + assert_eq!(router_config.pd_admission_wait_secs, 5); + let server_config = cli.to_server_config(router_config).unwrap(); + assert_eq!(server_config.router_config.pd_admission_wait_secs, 5); + + let defaults = cli_args_from(&[]).to_router_config(vec![], vec![]).unwrap(); + assert_eq!(defaults.pd_admission_wait_secs, 30); + } + /// The streamed-body stall timeout is a router-only setting and must /// flow into `RouterConfig`. #[test] diff --git a/model_gateway/src/observability/metrics.rs b/model_gateway/src/observability/metrics.rs index 4f23e4b985..78fb7fd0c9 100644 --- a/model_gateway/src/observability/metrics.rs +++ b/model_gateway/src/observability/metrics.rs @@ -1126,6 +1126,16 @@ impl Metrics { counter!("smg_pd_kv_transfer_failures_total").increment(1); } + /// Record a PD dispatch that had to wait for a decode admission slot. + pub fn record_pd_admission_wait() { + counter!("smg_pd_admission_waits_total").increment(1); + } + + /// Record a PD dispatch shed because no decode slot freed in time. + pub fn record_pd_admission_shed() { + counter!("smg_pd_admission_sheds_total").increment(1); + } + // ======================================================================== // Layer 3: Worker metrics // ======================================================================== diff --git a/model_gateway/src/routers/common/mod.rs b/model_gateway/src/routers/common/mod.rs index 74232f3f5a..9d1f9b69ab 100644 --- a/model_gateway/src/routers/common/mod.rs +++ b/model_gateway/src/routers/common/mod.rs @@ -18,6 +18,8 @@ //! session registry) shared by the OpenAI and HTTP routers //! - [`overload`] — shed responses for the absolute worker-overload //! guard, shared by the HTTP and gRPC selection paths +//! - [`pd_admission`] — the disaggregated dispatch's admission gate: never +//! post more bootstrap rooms to a pair than its decode engine can admit //! - [`placement`] — policy-driven worker placement over the routing //! pools (single worker and prefill/decode pairs), the one sequence the //! HTTP and gRPC families used to each carry a copy of @@ -40,6 +42,7 @@ pub use smg_external_router::header_utils; pub(crate) mod kv_transfer; pub use smg_external_router::{mcp_utils, openai_bridge}; pub mod overload; +pub mod pd_admission; pub use smg_external_router::persistence_utils; pub(crate) mod placement; pub use smg_external_router::realtime; diff --git a/model_gateway/src/routers/common/overload.rs b/model_gateway/src/routers/common/overload.rs index a56e2c6ace..9de740b4e8 100644 --- a/model_gateway/src/routers/common/overload.rs +++ b/model_gateway/src/routers/common/overload.rs @@ -22,8 +22,8 @@ use crate::{ routers::{common::retry::mark_non_retryable, error}, worker::{ overload::{ - BRANCH_ALL_OVERLOADED_SHED, BRANCH_OVERLOADED_AT_DISPATCH, STAGE_DISPATCH, - STAGE_SELECTION, + BRANCH_ALL_OVERLOADED_SHED, BRANCH_OVERLOADED_AT_DISPATCH, BRANCH_PD_ADMISSION_SHED, + STAGE_DISPATCH, STAGE_PD_ADMISSION, STAGE_SELECTION, }, Worker, }, @@ -89,6 +89,33 @@ pub fn shed_if_worker_overloaded(worker: &dyn Worker, model_id: &str) -> Option< )) } +/// Shed a disaggregated dispatch the decode leg cannot admit: the `rooms` +/// bootstrap rooms it needs do not fit in the engine's running `window`, and +/// none freed inside the admission wait. +/// +/// Same client-visible answer as the two vetoes above — the request was never +/// sent, and the wait already outlived any backoff a retry would add — under +/// its own decision branch, because here the *pair* is full rather than a +/// threshold being crossed. The counts are in the message because the two +/// causes need different operator responses: a transient full window versus a +/// batched request that is permanently wider than the pair. +pub(crate) fn shed_pd_admission( + worker: &str, + model_id: &str, + window: usize, + rooms: usize, +) -> Response { + shed( + BRANCH_PD_ADMISSION_SHED, + STAGE_PD_ADMISSION, + worker, + format!( + "Decode worker '{worker}' for model '{model_id}' could not admit {rooms} \ + request(s) within its running window of {window}" + ), + ) +} + /// One decision line, one counter, one response — marked non-retryable: the /// veto clears at the poll interval, which no backoff window outlives, and a /// terminal shed is what keeps the counter per-request rather than per-attempt. diff --git a/model_gateway/src/routers/common/pd_admission.rs b/model_gateway/src/routers/common/pd_admission.rs new file mode 100644 index 0000000000..1a44906b59 --- /dev/null +++ b/model_gateway/src/routers/common/pd_admission.rs @@ -0,0 +1,421 @@ +//! Admission control for disaggregated (PD) dispatch, bounded by the decode +//! engine's running window. +//! +//! SGLang-lineage engines start a bootstrap deadline on the prefill leg the +//! moment a request lands, and it clears only once the decode scheduler has +//! *admitted* that request and answered with its KV manifest. Decode admission +//! is bounded by the engine's running window (`--max-num-seqs` / +//! `--max-running-requests`), so a burst wider than that window leaves the +//! prefill deadline racing a queue the gateway itself created: prefill times +//! out rooms the decode has not reached yet, and the decode then pre-allocates +//! those same rooms and waits out its own transfer deadline for a peer that is +//! already gone. +//! +//! The gate below is the gateway's half of the fix — never post more rooms to +//! a pair than its decode can take. A request that arrives with the window +//! full waits for a room rather than joining the engine's queue, and sheds if +//! none frees in time, with the same 503 selection already answers when every +//! worker is vetoed. +//! +//! Admission is a *claim*, not a look: it reserves its rooms on the worker +//! atomically ([`Worker::try_admit_pd`]) and hands back a guard that releases +//! them. Reading the in-flight count and then sending would let two dispatches +//! both take the last free room, which is the burst this exists to stop. +//! +//! Nothing here runs when the engine does not report a window: admission is +//! not the gateway's to decide then, and dispatch behaves exactly as it did +//! before this module existed. + +use std::{ + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, + }, + time::Duration, +}; + +use axum::response::Response; +use tokio::time::Instant; +use tracing::debug; + +use crate::{observability::metrics::Metrics, routers::common::overload, worker::Worker}; + +/// Default seconds a PD dispatch may wait for decode rooms. Well under the +/// engines' bootstrap deadline (120 s on TokenSpeed), so a request that does +/// wait still dispatches with the whole deadline ahead of it. +pub const DEFAULT_PD_ADMISSION_WAIT_SECS: u64 = 30; + +/// How often the wait retries its claim. +/// +/// Polling, not a notifier: the event we would signal is a PD admission guard +/// dropping, and a per-worker registry of notifiers would be process-wide +/// mutable state with its own eviction problem. 50 ms is far finer than both +/// the engine step that actually frees a room and the wait deadline below, +/// and the sleep is asynchronous — a waiting request occupies no thread. +const CLAIM_RETRY_INTERVAL: Duration = Duration::from_millis(50); + +/// Seconds a PD dispatch waits for decode rooms before shedding. +/// +/// Process-wide for the same reason as the overload shed's `Retry-After`: the +/// gate is a free function on a dispatch path every router reaches, and the +/// value is one operator knob rather than a per-request input. Latched once at +/// startup from `--pd-admission-wait-secs`. +static PD_ADMISSION_WAIT_SECS: AtomicU64 = AtomicU64::new(DEFAULT_PD_ADMISSION_WAIT_SECS); + +/// Latch the admission wait. Called once at startup from the router config. +pub fn set_pd_admission_wait_secs(secs: u64) { + PD_ADMISSION_WAIT_SECS.store(secs, Ordering::Relaxed); +} + +fn admission_wait() -> Duration { + Duration::from_secs(PD_ADMISSION_WAIT_SECS.load(Ordering::Relaxed)) +} + +/// A claim on `rooms` slots in the decode engine's running window. +/// +/// Rides in the dispatch's `LoadGuards`, so the rooms are released on every +/// path the dispatch can end on: an error before the legs go out, a failed +/// leg, a client disconnect mid-stream, a retry that re-selects, or normal +/// completion. +pub(crate) struct PdAdmissionGuard { + worker: Arc, + rooms: usize, +} + +impl Drop for PdAdmissionGuard { + fn drop(&mut self) { + self.worker.release_pd(self.rooms); + } +} + +impl std::fmt::Debug for PdAdmissionGuard { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PdAdmissionGuard") + .field("worker", &self.worker.url()) + .field("rooms", &self.rooms) + .finish() + } +} + +/// How a claim was obtained. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Claim { + /// The rooms were free on the first try. + Immediate, + /// The rooms freed during the wait. + Waited, + /// The wait deadline passed with the window still full. + Refused, +} + +/// Retry `try_claim` until it succeeds or `wait` elapses. +/// +/// The claim itself is what makes admission safe under concurrency, so this +/// never reads a count: every attempt is the same all-or-nothing reservation, +/// and the first one to succeed owns the rooms. +async fn claim_within(wait: Duration, try_claim: impl Fn() -> bool) -> Claim { + if try_claim() { + return Claim::Immediate; + } + let deadline = Instant::now() + wait; + loop { + let now = Instant::now(); + if now >= deadline { + return Claim::Refused; + } + tokio::time::sleep(CLAIM_RETRY_INTERVAL.min(deadline - now)).await; + if try_claim() { + return Claim::Waited; + } + } +} + +/// Gate one disaggregated dispatch on the decode leg's admission window, +/// claiming the `rooms` bootstrap rooms it is about to post. +/// +/// `Ok(Some(guard))` holds the claim; `Ok(None)` means the engine reports no +/// window and admission is not the gateway's to decide; `Err` is the shed the +/// caller must return instead of dispatching. +pub(crate) async fn admit_decode( + decode: &Arc, + model_id: &str, + rooms: usize, +) -> Result, Response> { + // `None` (unreported, or a nonsense zero) leaves admission to the engine. + let Some(window) = decode.max_running_requests().map(usize::from) else { + return Ok(None); + }; + // A plan always posts at least one room, whatever its sub-request count + // claims. + let rooms = rooms.max(1); + + // A batched plan can demand more rooms than the pair will ever run at + // once. No wait can free more than the window holds, so answer now + // instead of sleeping out the budget to reach the same shed. + if rooms > window { + Metrics::record_pd_admission_shed(); + debug!( + worker = decode.url(), + model_id, window, rooms, "PD admission shed: more rooms than the decode window holds" + ); + return Err(overload::shed_pd_admission( + decode.url(), + model_id, + window, + rooms, + )); + } + + let wait = admission_wait(); + let claim = claim_within(wait, || decode.try_admit_pd(rooms, window)).await; + match claim { + Claim::Immediate => Ok(Some(PdAdmissionGuard { + worker: Arc::clone(decode), + rooms, + })), + Claim::Waited => { + Metrics::record_pd_admission_wait(); + debug!( + worker = decode.url(), + model_id, window, rooms, "PD admission waited for decode rooms" + ); + Ok(Some(PdAdmissionGuard { + worker: Arc::clone(decode), + rooms, + })) + } + Claim::Refused => { + Metrics::record_pd_admission_shed(); + debug!( + worker = decode.url(), + model_id, + window, + rooms, + admitted = decode.pd_admitted(), + waited_secs = wait.as_secs(), + "PD admission shed: no decode rooms freed" + ); + Err(overload::shed_pd_admission( + decode.url(), + model_id, + window, + rooms, + )) + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; + + use axum::http::{header::RETRY_AFTER, StatusCode}; + use futures::future::join_all; + use openai_protocol::{model_card::ModelCard, worker::HealthCheckConfig}; + + use super::*; + use crate::{ + routers::{common::retry::is_retryable_response, error::extract_error_code_from_response}, + worker::{BasicWorkerBuilder, ConnectionMode, WorkerType}, + }; + + fn decode_worker(url: &str, window: Option) -> Arc { + let mut labels = std::collections::HashMap::new(); + if let Some(window) = window { + labels.insert("max_running_requests".to_string(), window.to_string()); + } + Arc::new( + BasicWorkerBuilder::new(url) + .model(ModelCard::new("m")) + .worker_type(WorkerType::Decode) + .connection_mode(ConnectionMode::Grpc) + .labels(labels) + .health_config(HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) + .build(), + ) + } + + /// An engine that reports no window keeps the pre-gate behavior: dispatch, + /// however deep the gateway's in-flight count already is. + #[tokio::test] + async fn unknown_window_never_waits_or_sheds() { + let decode = decode_worker("grpc://127.0.0.1:9901", None); + for _ in 0..1_000 { + decode.increment_load(); + } + let admission = admit_decode(&decode, "m", 1).await.expect("no shed"); + assert!(admission.is_none(), "an unreported window claims nothing"); + assert_eq!(decode.pd_admitted(), 0); + } + + /// Below the window there is nothing to wait for — but the room is still + /// claimed, and released with the guard. + #[tokio::test(start_paused = true)] + async fn a_free_window_admits_without_waiting_and_releases_on_drop() { + let decode = decode_worker("grpc://127.0.0.1:9902", Some(4)); + let started = Instant::now(); + let admission = admit_decode(&decode, "m", 1) + .await + .expect("no shed") + .expect("a claim"); + assert_eq!( + started.elapsed(), + Duration::ZERO, + "no wait below the window" + ); + assert_eq!(decode.pd_admitted(), 1, "admission claims its room"); + + drop(admission); + assert_eq!(decode.pd_admitted(), 0, "the claim releases with the guard"); + } + + /// The claim, not a read, is what bounds the window: N requests racing a + /// window of one must produce exactly one admission, and the loser sheds. + #[tokio::test(start_paused = true)] + async fn concurrent_requests_never_exceed_a_window_of_one() { + set_pd_admission_wait_secs(1); + let decode = decode_worker("grpc://127.0.0.1:9903", Some(1)); + + let live = AtomicUsize::new(0); + let peak = AtomicUsize::new(0); + let attempts = (0..8).map(|_| async { + let outcome = admit_decode(&decode, "m", 1).await; + if outcome.as_ref().is_ok_and(Option::is_some) { + let now = live.fetch_add(1, AtomicOrdering::SeqCst) + 1; + peak.fetch_max(now, AtomicOrdering::SeqCst); + // Hold the claim past every other attempt's deadline. + tokio::time::sleep(Duration::from_secs(5)).await; + live.fetch_sub(1, AtomicOrdering::SeqCst); + } + outcome + }); + let outcomes = join_all(attempts).await; + + let admitted = outcomes + .iter() + .filter(|outcome| outcome.as_ref().is_ok_and(Option::is_some)) + .count(); + assert_eq!(admitted, 1, "a window of one admits exactly one request"); + assert_eq!(peak.load(AtomicOrdering::SeqCst), 1, "never two at once"); + assert_eq!( + outcomes.iter().filter(|o| o.is_err()).count(), + 7, + "every request that could not claim a room is shed" + ); + set_pd_admission_wait_secs(DEFAULT_PD_ADMISSION_WAIT_SECS); + } + + /// A batched plan claims one room per sub-request, so its siblings cannot + /// slip past a gate that only ever checked for one. + #[tokio::test(start_paused = true)] + async fn a_batched_plan_claims_one_room_per_sub_request() { + set_pd_admission_wait_secs(1); + let decode = decode_worker("grpc://127.0.0.1:9904", Some(4)); + + let batch = admit_decode(&decode, "m", 4) + .await + .expect("no shed") + .expect("a claim"); + assert_eq!(decode.pd_admitted(), 4, "one room per sub-request"); + + // The window is now full: a single-room request must shed, not slip in. + assert!( + admit_decode(&decode, "m", 1).await.is_err(), + "a full window admits nothing more" + ); + + drop(batch); + assert!(admit_decode(&decode, "m", 1).await.is_ok()); + set_pd_admission_wait_secs(DEFAULT_PD_ADMISSION_WAIT_SECS); + } + + /// A plan wider than the window can never be admitted, so it sheds at + /// once rather than sleeping out the budget to reach the same answer. + #[tokio::test(start_paused = true)] + async fn a_plan_wider_than_the_window_sheds_without_waiting() { + let decode = decode_worker("grpc://127.0.0.1:9905", Some(4)); + let started = Instant::now(); + + let response = admit_decode(&decode, "m", 5).await.expect_err("shed"); + + assert_eq!(started.elapsed(), Duration::ZERO); + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(decode.pd_admitted(), 0, "a shed claims nothing"); + } + + /// A room freeing mid-wait releases the request instead of shedding it. + #[tokio::test(start_paused = true)] + async fn a_freed_room_admits_the_waiting_request() { + let claimed = AtomicUsize::new(1); + let free_a_room = async { + tokio::time::sleep(Duration::from_millis(500)).await; + claimed.store(0, AtomicOrdering::SeqCst); + }; + let gate = claim_within(Duration::from_secs(30), || { + claimed + .compare_exchange(0, 1, AtomicOrdering::SeqCst, AtomicOrdering::SeqCst) + .is_ok() + }); + + let ((), claim) = tokio::join!(free_a_room, gate); + + assert_eq!(claim, Claim::Waited); + } + + /// A window that never frees sheds at the deadline, not before it. + #[tokio::test(start_paused = true)] + async fn a_full_window_sheds_at_the_deadline() { + let started = Instant::now(); + let claim = claim_within(Duration::from_secs(30), || false).await; + assert_eq!(claim, Claim::Refused); + assert!( + started.elapsed() >= Duration::from_secs(30), + "the shed must wait out the whole admission window, waited {:?}", + started.elapsed() + ); + } + + /// A zero wait is the "shed immediately" setting: no sleep, no admission. + #[tokio::test(start_paused = true)] + async fn a_zero_wait_sheds_without_sleeping() { + let started = Instant::now(); + let claim = claim_within(Duration::ZERO, || false).await; + assert_eq!(claim, Claim::Refused); + assert_eq!(started.elapsed(), Duration::ZERO); + } + + /// The shed is the overload guard's 503, terminal for the retry layer and + /// carrying the client's pacing hint. + #[tokio::test(start_paused = true)] + async fn the_shed_is_the_overload_503_with_retry_after() { + set_pd_admission_wait_secs(1); + let decode = decode_worker("grpc://127.0.0.1:9906", Some(2)); + assert!(decode.try_admit_pd(2, 2), "fill the window"); + + let response = admit_decode(&decode, "m", 1) + .await + .expect_err("a full window sheds"); + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!( + extract_error_code_from_response(&response), + "worker_overload_protection_shed" + ); + assert!( + response + .headers() + .get(RETRY_AFTER) + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()) + .is_some_and(|secs| secs >= 1), + "Retry-After must carry whole seconds" + ); + assert!( + !is_retryable_response(&response), + "the wait already outlived any backoff a retry would add" + ); + set_pd_admission_wait_secs(DEFAULT_PD_ADMISSION_WAIT_SECS); + } +} diff --git a/model_gateway/src/routers/grpc/client.rs b/model_gateway/src/routers/grpc/client.rs index 2bd1221b57..6fba4da7f0 100644 --- a/model_gateway/src/routers/grpc/client.rs +++ b/model_gateway/src/routers/grpc/client.rs @@ -865,6 +865,9 @@ const TOKENSPEED_GRPC_KEYS: &[&str] = &[ "pp_size", "context_length", "max_total_tokens", + // TokenSpeed's spelling of the scheduler's running window; the fleet + // capacity accounting and the PD admission gate read either name. + "max_num_seqs", "max_running_requests", "load_balance_method", "is_embedding", diff --git a/model_gateway/src/routers/grpc/common/stages/request_execution.rs b/model_gateway/src/routers/grpc/common/stages/request_execution.rs index 2c8d840663..5331c931c9 100644 --- a/model_gateway/src/routers/grpc/common/stages/request_execution.rs +++ b/model_gateway/src/routers/grpc/common/stages/request_execution.rs @@ -1,6 +1,6 @@ //! Request execution: dispatch one attempt of the retained execution plan. -use std::time::Instant; +use std::{future::Future, pin::Pin, sync::Arc, time::Instant}; use axum::response::Response; use futures::future::{join_all, try_join_all}; @@ -10,9 +10,12 @@ use super::pd_protocol::{DpPlacement, PdDispatch, PdProtocol}; use crate::{ observability::metrics::{metrics_labels, Metrics}, routers::{ - common::kv_transfer::{ - connector_mode_for_worker, mooncake_decode_params, mooncake_prefill_params, - KvConnectorMode, NIXL_PREFILL_KV_PARAMS, + common::{ + kv_transfer::{ + connector_mode_for_worker, mooncake_decode_params, mooncake_prefill_params, + KvConnectorMode, NIXL_PREFILL_KV_PARAMS, + }, + pd_admission, }, error, grpc::{ @@ -28,11 +31,83 @@ use crate::{ utils::tonic_ext::{TonicResultExt, TonicStatusExt}, }, }, - worker::ConnectionModeExt, + worker::{ConnectionModeExt, Worker}, }; type StreamResult = Result; +/// One leg's owned dispatch. Owned (rather than borrowed from the client +/// selection) so the leg still in flight when its partner fails can be moved +/// off the request path — see [`retire_pd_leg`]. +type PdLegDispatch = Pin + Send>>; + +/// Which leg of a disaggregated dispatch a result belongs to. The error code +/// and message are part of the client contract, so they live per leg here +/// rather than being reconstructed at each call site. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum PdLeg { + Prefill, + Decode, +} + +impl PdLeg { + fn name(self) -> &'static str { + match self { + Self::Prefill => metrics_labels::WORKER_PREFILL, + Self::Decode => metrics_labels::WORKER_DECODE, + } + } + + fn error_code(self) -> &'static str { + match self { + Self::Prefill => "prefill_worker_failed_to_start", + Self::Decode => "decode_worker_failed_to_start", + } + } + + fn error_message(self) -> &'static str { + match self { + Self::Prefill => "Prefill worker failed to start", + Self::Decode => "Decode worker failed to start", + } + } + + fn partner(self) -> Self { + match self { + Self::Prefill => Self::Decode, + Self::Decode => Self::Prefill, + } + } +} + +/// How the two parallel legs resolved. +#[expect( + clippy::large_enum_variant, + reason = "consumed by its caller on the next line; boxing would add an allocation to the success path to save one stack move" +)] +enum PdDispatchOutcome { + /// Both legs answered; either or both may carry an error. + Both(StreamResult, StreamResult), + /// One leg failed while its partner was still dispatching. The partner is + /// handed back so the caller can retire it instead of waiting it out. + FailedFirst { + leg: PdLeg, + error: tonic::Status, + partner: PdLegDispatch, + }, +} + +/// Backend requests one plan dispatches — one per batched prompt, else one. +/// +/// This is both the load-guard scale and the number of PD bootstrap rooms the +/// plan will post, which is why admission and the guards read the same count. +fn plan_sub_requests(plan: &ExecutionPlan) -> usize { + match plan { + ExecutionPlan::Batch { requests, .. } => requests.len(), + _ => 1, + } +} + /// Metric connection labels for the PD legs (a leg can be gRPC or ZMQ). fn pd_leg_labels(workers: &WorkerSelection) -> (&'static str, &'static str) { match workers { @@ -56,6 +131,30 @@ pub(crate) async fn execute_plan( ctx: &mut DispatchContext, execution_plan: ExecutionPlan, ) -> Result<(), Response> { + // One bootstrap room per backend request the plan will post: a batched + // completion fans out one PD dispatch per sub-request, so admission has + // to claim for all of them or the siblings walk past a gate that only + // ever asked about one. + let sub_requests = plan_sub_requests(&execution_plan); + + // Admission runs before this attempt claims anything else. The decode + // leg's engine window is what clears the prefill's bootstrap deadline, so + // a request the decode cannot admit yet waits here rather than in the + // engine's queue — where it would burn the prefill's deadline and strand + // both rooms. Waiting ahead of the encode take below also keeps the + // encode jobs' SHM unclaimed for the duration: under the burst this gate + // targets, holding it would queue a second scarce resource behind the + // decode window and throw the finished encode work away on a shed. The + // gate abstains when the engine reports no window. + let admission = match ctx + .workers + .as_ref() + .and_then(WorkerSelection::decode_worker) + { + Some(decode) => pd_admission::admit_decode(decode, &ctx.model_id, sub_requests).await?, + None => None, + }; + // `None` for non-EPD, text-only EPD, or an EPD retry (the first dispatch // consumed it). Taking it transfers the encode jobs' SHM Drop guards // here: dispatch consumes them, while an early error before dispatch @@ -84,14 +183,9 @@ pub(crate) async fn execute_plan( ) })?; - let sub_requests = match &execution_plan { - ExecutionPlan::Batch { requests, .. } => requests.len(), - _ => 1, - }; - ctx.load_guards = Some(LoadGuards::scaled( - workers, - ctx.sticky_key.as_deref(), - sub_requests, + ctx.load_guards = Some(LoadGuards::admitted( + admission, + LoadGuards::scaled(workers, ctx.sticky_key.as_deref(), sub_requests), )); // Extract dispatch metadata for the tracing span and PD metric labels. @@ -378,65 +472,211 @@ async fn execute_parallel_pd( // it on this path), so prefill duration cannot be measured — only TTFT, // recorded at the first decode token in streaming. prefill_start anchors it. let prefill_start = Instant::now(); - let (prefill_result, decode_result): (StreamResult, StreamResult) = tokio::join!( - prefill_client.generate(prefill_request), - decode_client.generate(decode_request) - ); + let (prefill_label, decode_label) = pd_leg_labels(workers); + // Each leg owns its client handle (a cheap channel clone, the same one + // batched dispatch takes per sub-request) so the leg still in flight when + // its partner fails can be moved off the request path. + let mut prefill_client = prefill_client.clone(); + let mut decode_client = decode_client.clone(); + let prefill_dispatch: PdLegDispatch = + Box::pin(async move { prefill_client.generate(prefill_request).await }); + let decode_dispatch: PdLegDispatch = + Box::pin(async move { decode_client.generate(decode_request).await }); + + match dispatch_pd_legs(prefill_dispatch, decode_dispatch).await { + PdDispatchOutcome::Both(prefill_result, decode_result) => { + // Record circuit breaker outcomes (client errors don't count as failures) + workers.record_prefill_decode_outcomes( + prefill_result.cb_status_code(), + decode_result.cb_status_code(), + ); - // Record circuit breaker outcomes (client errors don't count as failures) - workers.record_prefill_decode_outcomes( - prefill_result.cb_status_code(), - decode_result.cb_status_code(), - ); + // Both legs are translated before either is propagated: a + // decode-side failure is logged and counted even when the prefill + // error is the one the client sees. + let prefill = + prefill_result.map_err(|e| pd_leg_error(PdLeg::Prefill, prefill_label, &e)); + let decode = decode_result.map_err(|e| pd_leg_error(PdLeg::Decode, decode_label, &e)); + let (prefill_stream, decode_stream) = match (prefill, decode) { + (Ok(prefill), Ok(decode)) => (prefill, decode), + // The surviving leg's stream drops here, which aborts its + // room on the worker. + (Err(response), _) | (Ok(_), Err(response)) => return Err(response), + }; + + // A client disconnect drops both leg streams, which normally fires an + // immediate abort to each worker. The decode leg must not be aborted + // while it is still receiving the KV handoff from prefill — tearing + // the request down mid-transfer can crash or leak on the engine — so + // its abort is deferred until the first decode response (the proof + // the handoff completed). The prefill leg keeps the immediate abort: + // if prefill is still running there is nothing to hand off yet, and + // stopping it promptly frees capacity. + let decode_stream = decode_stream.defer_abort_until_first_item(); + + Ok(ExecutionResult::PrefillDecode { + prefill: prefill_stream, + decode: Box::new(decode_stream), + pd_timing: PdTiming { + prefill_start, + runtime, + }, + }) + } + PdDispatchOutcome::FailedFirst { + leg, + error, + partner, + } => { + let status = error.http_status().as_u16(); + // Only the leg that answered is recorded: the abandoned one has + // said nothing about its worker yet, and `retire_pd_leg` records + // it when it finally does. + let (label, partner_label, partner_worker) = match leg { + PdLeg::Prefill => { + workers.record_outcome_prefill(status); + (prefill_label, decode_label, workers.decode_worker()) + } + PdLeg::Decode => { + workers.record_outcome_decode(status); + (decode_label, prefill_label, workers.prefill_worker()) + } + }; + if let Some(worker) = partner_worker { + retire_pd_leg(leg.partner(), partner_label, Arc::clone(worker), partner); + } + Err(pd_leg_error(leg, label, &error)) + } + } +} - let (prefill_label, decode_label) = pd_leg_labels(workers); +/// Dispatch both legs together and answer with the first failure. +/// +/// The legs rendezvous on one bootstrap room, so a leg that cannot start +/// leaves its partner's room unreachable — and the engine holding that room +/// for the whole of its own deadline. Waiting for the second verdict before +/// reporting the first failure is what turned a 120 s prefill bootstrap +/// timeout into a 300 s decode transfer timeout for the client, with the +/// pair's slot held throughout. A leg that has already answered is never +/// discarded, so the success path is byte-for-byte what a join produced. +async fn dispatch_pd_legs( + mut prefill: PdLegDispatch, + mut decode: PdLegDispatch, +) -> PdDispatchOutcome { + /// One turn of the dispatch loop. The `select!` handlers may only + /// classify — moving a still-pending leg out happens after the macro's + /// borrows end. + enum LegStep { + /// The prefill leg answered; a failure here was beaten to the finish + /// by its partner, so both verdicts are in. + Prefill(StreamResult), + Decode(StreamResult), + /// This leg failed while its partner was still dispatching. + Alone(PdLeg, tonic::Status), + } - // Handle prefill result - let prefill_stream = prefill_result.map_err(|e| { - Metrics::record_worker_error( - metrics_labels::WORKER_PREFILL, - prefill_label, - metrics_labels::ERROR_BACKEND, - ); - error!(function = "execute_parallel_pd", error = %e, "Prefill worker failed to start"); - e.to_http_error( - "prefill_worker_failed_to_start", - format!("Prefill worker failed to start: {}", e.message()), - ) - })?; + // At most one of these is `Some` at the top of the loop: the second + // answer returns rather than being stored, which is also what keeps both + // `select!` branches from being disabled at once. + let mut prefill_result: Option = None; + let mut decode_result: Option = None; + + loop { + let step = tokio::select! { + result = &mut prefill, if prefill_result.is_none() => match result { + Err(error) if decode_result.is_none() => LegStep::Alone(PdLeg::Prefill, error), + result => LegStep::Prefill(result), + }, + result = &mut decode, if decode_result.is_none() => match result { + Err(error) if prefill_result.is_none() => LegStep::Alone(PdLeg::Decode, error), + result => LegStep::Decode(result), + }, + }; - // Handle decode result - let decode_stream = decode_result.map_err(|e| { - Metrics::record_worker_error( - metrics_labels::WORKER_DECODE, - decode_label, - metrics_labels::ERROR_BACKEND, - ); - error!(function = "execute_parallel_pd", error = %e, "Decode worker failed to start"); - e.to_http_error( - "decode_worker_failed_to_start", - format!("Decode worker failed to start: {}", e.message()), - ) - })?; + match step { + LegStep::Alone(leg, error) => { + let partner = match leg { + PdLeg::Prefill => decode, + PdLeg::Decode => prefill, + }; + return PdDispatchOutcome::FailedFirst { + leg, + error, + partner, + }; + } + LegStep::Prefill(result) => match decode_result.take() { + Some(decode) => return PdDispatchOutcome::Both(result, decode), + None => prefill_result = Some(result), + }, + LegStep::Decode(result) => match prefill_result.take() { + Some(prefill) => return PdDispatchOutcome::Both(prefill, result), + None => decode_result = Some(result), + }, + } + } +} - // A client disconnect drops both leg streams, which normally fires an - // immediate abort to each worker. The decode leg must not be aborted - // while it is still receiving the KV handoff from prefill — tearing - // the request down mid-transfer can crash or leak on the engine — so - // its abort is deferred until the first decode response (the proof - // the handoff completed). The prefill leg keeps the immediate abort: - // if prefill is still running there is nothing to hand off yet, and - // stopping it promptly frees capacity. - let decode_stream = decode_stream.defer_abort_until_first_item(); +/// Log, count and translate one leg's dispatch failure into the client answer. +fn pd_leg_error(leg: PdLeg, connection: &'static str, error: &tonic::Status) -> Response { + Metrics::record_worker_error(leg.name(), connection, metrics_labels::ERROR_BACKEND); + match leg { + PdLeg::Prefill => { + error!(function = "execute_parallel_pd", error = %error, "Prefill worker failed to start"); + } + PdLeg::Decode => { + error!(function = "execute_parallel_pd", error = %error, "Decode worker failed to start"); + } + } + error.to_http_error( + leg.error_code(), + format!("{}: {}", leg.error_message(), error.message()), + ) +} - Ok(ExecutionResult::PrefillDecode { - prefill: prefill_stream, - decode: Box::new(decode_stream), - pd_timing: PdTiming { - prefill_start, - runtime, - }, - }) +/// Finish off the leg that was still dispatching when its partner failed. +/// +/// The request is already answered; what is left is the engine-side room, +/// which the abandoned leg will allocate the moment its dispatch lands and +/// then hold for its own deadline. Dropping the leg's stream as soon as it +/// arrives is the abort path the router already uses on a client disconnect, +/// and it is the only thing that releases that room early. The abort is *not* +/// deferred here: deferral exists to protect an in-flight KV handoff, and the +/// peer that would have written it is the leg that just failed. +#[expect( + clippy::disallowed_methods, + reason = "the abandoned PD leg is retired off the request path; its dispatch is bounded by the engine's own deadline" +)] +fn retire_pd_leg( + leg: PdLeg, + connection: &'static str, + worker: Arc, + dispatch: PdLegDispatch, +) { + tokio::spawn(async move { + let result = dispatch.await; + worker.record_outcome(result.cb_status_code()); + match result { + Ok(stream) => { + debug!( + leg = leg.name(), + worker = worker.url(), + "Aborting the room of the PD leg abandoned after its partner failed" + ); + drop(stream); + } + Err(error) => { + Metrics::record_worker_error(leg.name(), connection, metrics_labels::ERROR_BACKEND); + error!( + function = "retire_pd_leg", + leg = leg.name(), + worker = worker.url(), + error = %error, + "PD leg failed after its partner had already failed" + ); + } + } + }); } /// Execute vLLM PD: send to prefill with max_tokens=1 first, wait for completion, @@ -712,13 +952,89 @@ async fn execute_sequential_pd( #[cfg(test)] mod tests { - use std::sync::Arc; + use std::{sync::Arc, time::Duration}; use smg_grpc_client::{tokenspeed_proto as ts, vllm_proto as vllm}; use super::*; use crate::worker::{BasicWorkerBuilder, ConnectionMode, RuntimeType, Worker, WorkerType}; + /// A leg that never answers — the decode leg stuck behind the engine's + /// transfer deadline while the prefill leg has already given up. + fn never_answers() -> PdLegDispatch { + Box::pin(std::future::pending()) + } + + fn fails_now(message: &'static str) -> PdLegDispatch { + Box::pin(async move { Err(tonic::Status::deadline_exceeded(message)) }) + } + + fn answers_after(delay: Duration) -> PdLegDispatch { + Box::pin(async move { + tokio::time::sleep(delay).await; + Err(tonic::Status::unavailable("late leg")) + }) + } + + /// A prefill that times out must be reported without waiting for the + /// decode leg, and the decode's dispatch must come back so its room can + /// be released instead of being left to the engine's own deadline. + #[tokio::test(start_paused = true)] + async fn a_failed_prefill_answers_without_waiting_for_the_decode_leg() { + let started = tokio::time::Instant::now(); + let outcome = dispatch_pd_legs(fails_now("bootstrap timeout"), never_answers()).await; + + assert_eq!( + started.elapsed(), + Duration::ZERO, + "the prefill error must not wait on the decode leg" + ); + match outcome { + PdDispatchOutcome::FailedFirst { leg, error, .. } => { + assert_eq!(leg, PdLeg::Prefill); + assert_eq!(error.message(), "bootstrap timeout"); + } + PdDispatchOutcome::Both(..) => panic!("expected a fail-fast outcome"), + } + } + + /// The mirror case: a decode leg that cannot start is answered without + /// waiting out the prefill's bootstrap deadline. + #[tokio::test(start_paused = true)] + async fn a_failed_decode_answers_without_waiting_for_the_prefill_leg() { + let outcome = dispatch_pd_legs( + answers_after(Duration::from_secs(120)), + fails_now("decode refused"), + ) + .await; + + match outcome { + PdDispatchOutcome::FailedFirst { leg, error, .. } => { + assert_eq!(leg, PdLeg::Decode); + assert_eq!(error.message(), "decode refused"); + } + PdDispatchOutcome::Both(..) => panic!("expected a fail-fast outcome"), + } + } + + /// Admission and the load guards read the same count, so a batched plan + /// claims one room per prompt instead of walking past a gate that only + /// asked about one. + #[test] + fn plan_sub_requests_counts_every_batched_prompt() { + let single = ExecutionPlan::PrefillDecode(ProtoGenerateRequest::Vllm(Box::default())); + assert_eq!(plan_sub_requests(&single), 1); + + let batch = ExecutionPlan::Batch { + kind: ExecutionPlanKind::PrefillDecode, + shared_request_id: "cmpl-1".to_string(), + requests: (0..4) + .map(|_| ProtoGenerateRequest::Vllm(Box::default())) + .collect(), + }; + assert_eq!(plan_sub_requests(&batch), 4); + } + #[test] fn pd_leg_labels_reflect_each_legs_transport() { let prefill: Arc = Arc::new( diff --git a/model_gateway/src/routers/grpc/context.rs b/model_gateway/src/routers/grpc/context.rs index ee8524fc67..755045ddd4 100644 --- a/model_gateway/src/routers/grpc/context.rs +++ b/model_gateway/src/routers/grpc/context.rs @@ -41,7 +41,7 @@ use super::{ use crate::{ middleware::TenantRequestMeta, policies::CacheNamespace, - routers::error::internal_error, + routers::{common::pd_admission::PdAdmissionGuard, error::internal_error}, worker::{ConnectionMode, RuntimeType, Worker, WorkerLoadGuard, WorkerRegistry}, }; @@ -626,6 +626,14 @@ pub(crate) enum LoadGuards { Batch { _guards: Vec, }, + /// A disaggregated dispatch whose bootstrap rooms the PD admission gate + /// claimed before it was allowed to send. The claim rides with the guards + /// so it is released on every path the dispatch can end on — an early + /// error, a failed leg, a client disconnect, a retry, or completion. + Admitted { + _admission: PdAdmissionGuard, + _guards: Box, + }, } impl LoadGuards { @@ -643,6 +651,19 @@ impl LoadGuards { } } + /// Bind an admission claim to the dispatch's guards, so the rooms it + /// reserved outlive nothing else. `None` (the engine reports no window) + /// returns the guards untouched. + pub fn admitted(admission: Option, guards: Self) -> Self { + match admission { + Some(admission) => Self::Admitted { + _admission: admission, + _guards: Box::new(guards), + }, + None => guards, + } + } + /// One guard set per concurrent sub-request. pub fn scaled(selection: &WorkerSelection, routing_key: Option<&str>, count: usize) -> Self { if count <= 1 { diff --git a/model_gateway/src/routers/grpc/pipeline.rs b/model_gateway/src/routers/grpc/pipeline.rs index a96ff93c9f..1edb5e9182 100644 --- a/model_gateway/src/routers/grpc/pipeline.rs +++ b/model_gateway/src/routers/grpc/pipeline.rs @@ -602,6 +602,8 @@ impl RequestPipeline { ); // The failed attempt's worker load must not stay elevated through // the backoff window (a fresh context dropped them here before). + // This releases its PD admission claim too, so the retry is not + // queued behind its own predecessor's bootstrap rooms. dctx.load_guards = None; let Some(config) = retry_config else { @@ -1443,17 +1445,23 @@ mod request_release_tests { /// probe reaches zero strong references or a deadline passes, recording /// the outcome in `released`. An ungated stub (no probe) answers /// immediately -- used for the PD prefill leg. `fail_first` makes the - /// first generate call return UNAVAILABLE, for retry-replay tests; every - /// call's input token ids and engine request id are recorded. + /// first generate call return UNAVAILABLE, for retry-replay tests; + /// `fail_always` fails every call, and `answer_after` stalls the generate + /// RPC, standing in for an engine that answers only at its own deadline. + /// Every call's input token ids and engine request id are recorded, as is + /// every aborted request id. #[derive(Clone, Default)] struct GatedScheduler { probe: Option>, gate_rpc: bool, released: Arc, fail_first: bool, + fail_always: bool, + answer_after: Option, calls: Arc, seen_input_ids: Arc>>>, seen_request_ids: Arc>>, + aborted_request_ids: Arc>>, } impl GatedScheduler { @@ -1519,9 +1527,14 @@ mod request_release_tests { .lock() .unwrap_or_else(PoisonError::into_inner) .push(request.request_id.clone()); - if self.fail_first && self.calls.fetch_add(1, Ordering::SeqCst) == 0 { + if self.fail_always + || (self.fail_first && self.calls.fetch_add(1, Ordering::SeqCst) == 0) + { return Err(Status::unavailable("release-test induced failure")); } + if let Some(delay) = self.answer_after { + tokio::time::sleep(delay).await; + } if self.gate_rpc { if let Some(probe) = &self.probe { Self::await_probe(probe, &self.released).await; @@ -1556,8 +1569,12 @@ mod request_release_tests { async fn abort( &self, - _request: TonicRequest, + request: TonicRequest, ) -> Result, Status> { + self.aborted_request_ids + .lock() + .unwrap_or_else(PoisonError::into_inner) + .push(request.into_inner().request_id); Ok(TonicResponse::new(ts::AbortResponse { success: true, message: String::new(), @@ -1959,6 +1976,89 @@ mod request_release_tests { assert_ne!(ids[0], ids[1], "each attempt gets a fresh engine id"); } } + + /// A prefill leg that cannot start must answer the client immediately + /// rather than waiting out the decode leg's own deadline, and the decode + /// room it stranded must be aborted rather than left for the engine to + /// time out. The decode stub here answers only after `DECODE_DELAY`, + /// standing in for an engine whose transfer deadline is what the client + /// used to wait on. + #[tokio::test] + async fn a_failed_prefill_answers_now_and_aborts_the_decode_room() { + const DECODE_DELAY: Duration = Duration::from_secs(2); + + let prefill_port = spawn_stub(GatedScheduler { + fail_always: true, + ..Default::default() + }) + .await; + let aborted = Arc::new(Mutex::new(Vec::new())); + let decode_request_ids = Arc::new(Mutex::new(Vec::new())); + let decode_port = spawn_stub(GatedScheduler { + answer_after: Some(DECODE_DELAY), + aborted_request_ids: Arc::clone(&aborted), + seen_request_ids: Arc::clone(&decode_request_ids), + ..Default::default() + }) + .await; + + let worker_registry = Arc::new(WorkerRegistry::new()); + register_worker(&worker_registry, prefill_port, WorkerType::Prefill); + register_worker(&worker_registry, decode_port, WorkerType::Decode); + let pipeline = completion_pipeline(&worker_registry, Mode::PrefillDecode); + let components = components(worker_registry).await; + + let started = Instant::now(); + let response = pipeline + .execute_completion( + completion_request(false), + None, + MODEL.to_string(), + components, + None, + None, + None, + ) + .await; + let answered_in = started.elapsed(); + + assert_eq!(response.status(), http::StatusCode::SERVICE_UNAVAILABLE); + assert_eq!( + error::extract_error_code_from_response(&response), + "prefill_worker_failed_to_start" + ); + assert!( + answered_in < DECODE_DELAY, + "the prefill failure must not wait for the decode leg, took {answered_in:?}" + ); + + // The decode leg is retired off the request path: once its dispatch + // lands, its stream drops and that is what sends the abort. + let deadline = Instant::now() + DECODE_DELAY + Duration::from_secs(8); + loop { + if !aborted + .lock() + .unwrap_or_else(PoisonError::into_inner) + .is_empty() + { + break; + } + assert!( + Instant::now() < deadline, + "the stranded decode room was never aborted" + ); + tokio::time::sleep(Duration::from_millis(20)).await; + } + let aborted = aborted.lock().unwrap_or_else(PoisonError::into_inner); + let dispatched = decode_request_ids + .lock() + .unwrap_or_else(PoisonError::into_inner); + assert_eq!( + aborted.as_slice(), + dispatched.as_slice(), + "the abort must name the decode leg's own request id" + ); + } } #[cfg(test)] diff --git a/model_gateway/src/worker/overload.rs b/model_gateway/src/worker/overload.rs index 707ae65652..88c22f5245 100644 --- a/model_gateway/src/worker/overload.rs +++ b/model_gateway/src/worker/overload.rs @@ -22,12 +22,22 @@ pub const BRANCH_ALL_OVERLOADED_SHED: &str = "all_overloaded_shed"; /// here every other worker may well be idle. pub const BRANCH_OVERLOADED_AT_DISPATCH: &str = "overloaded_at_dispatch"; +/// Decision-log branch: the decode leg of a disaggregated pair was already +/// running its full engine window and no slot freed inside the admission wait. +/// The pair is not overloaded by the threshold predicate above — the gateway +/// simply has nowhere to put another bootstrap room. +pub const BRANCH_PD_ADMISSION_SHED: &str = "pd_admission_shed"; + /// Shed detected while assembling the candidate pool. pub const STAGE_SELECTION: &str = "selection"; /// Shed detected by the dispatch-time re-check of the chosen worker. pub const STAGE_DISPATCH: &str = "dispatch"; +/// Shed detected by the PD admission gate, which runs between selection and +/// the paired dispatch. +pub const STAGE_PD_ADMISSION: &str = "pd_admission"; + /// Absolute thresholds above which a worker is vetoed from routing. /// /// Both fields unset disables the feature entirely: the flag is never written, diff --git a/model_gateway/src/worker/worker.rs b/model_gateway/src/worker/worker.rs index b63be148f3..df82ea0811 100644 --- a/model_gateway/src/worker/worker.rs +++ b/model_gateway/src/worker/worker.rs @@ -380,6 +380,31 @@ pub trait Worker: Send + Sync + fmt::Debug + 'static { /// Decrement the load counter fn decrement_load(&self); + /// Claim `count` PD bootstrap rooms while the worker's claimed total + /// stays within `window`; a refusal claims nothing. + /// + /// The PD admission gate reserves through this rather than comparing + /// [`Worker::load`] to the window: a read-then-send lets two dispatches + /// both take the last free slot, which is exactly the over-window burst + /// the gate exists to prevent. + /// + /// The defaults admit unconditionally and track nothing. They exist for + /// implementations that carry no shared runtime (the language bindings, + /// test doubles): the gate only reaches a worker that reports a running + /// window, and any implementation that does report one must override + /// these three with a real claim, or the window is not enforced. + fn try_admit_pd(&self, _count: usize, _window: usize) -> bool { + true + } + + /// Release `count` rooms claimed by [`Worker::try_admit_pd`]. + fn release_pd(&self, _count: usize) {} + + /// Rooms currently claimed on this worker. + fn pd_admitted(&self) -> usize { + 0 + } + /// Get the current routing-key load cardinality. fn routing_key_load(&self) -> usize; @@ -403,19 +428,29 @@ pub trait Worker: Send + Sync + fmt::Debug + 'static { /// Worker-reported in-flight capacity, if available. /// - /// Reads the `max_running_requests` label populated by the metadata - /// discovery pipeline (Step 4 of the worker lifecycle). Returns - /// `None` when the worker hasn't reported a value or reports zero - /// (zero is meaningless for capacity accounting). + /// Reads the running-window label populated by the metadata discovery + /// pipeline (Step 4 of the worker lifecycle). Returns `None` when the + /// worker hasn't reported a value or reports zero (zero is meaningless + /// for capacity accounting). + /// + /// Engines spell the same window two ways — TokenSpeed and vLLM + /// advertise `max_num_seqs`, SGLang advertises `max_running_requests` — + /// and both are the count of requests the scheduler will run at once, so + /// both are read here rather than leaving TokenSpeed workers looking like + /// non-reporters. `max_num_seqs` wins when a worker reports both: it is + /// the authoritative name for the engines that use it, and it is the + /// order `smg_grpc_servicer.tokenspeed.loads::running_window` reports in, + /// so the label and the `GetLoads` report cannot disagree. /// /// `WorkerCapacity` uses this to derive total fleet capacity when /// every worker reports; falls back to a configured per-worker - /// estimate otherwise. + /// estimate otherwise. The PD admission gate uses it as the decode + /// leg's admission bound. fn max_running_requests(&self) -> Option { - self.metadata() - .spec - .labels - .get("max_running_requests") + let labels = &self.metadata().spec.labels; + labels + .get("max_num_seqs") + .or_else(|| labels.get("max_running_requests")) .and_then(|s| s.parse::().ok()) .filter(|n| *n > 0) } @@ -1046,6 +1081,13 @@ pub struct WorkerRuntime { consecutive_successes: AtomicUsize, total_pending_probes: AtomicUsize, load_counter: AtomicUsize, + /// Bootstrap rooms the PD admission gate has claimed on this worker and + /// not yet released. Separate from `load_counter` because admission must + /// *claim* against the engine's running window rather than read it: two + /// dispatches that both saw the last free slot would both send, which is + /// the over-window burst the gate exists to prevent. Lives here so a + /// same-URL replacement inherits the rooms the engine still holds. + pd_admitted: AtomicUsize, processed_counter: AtomicUsize, worker_routing_key_load: WorkerRoutingKeyLoad, revision: AtomicU64, @@ -1063,6 +1105,7 @@ impl WorkerRuntime { consecutive_successes: AtomicUsize::new(0), total_pending_probes: AtomicUsize::new(0), load_counter: AtomicUsize::new(0), + pd_admitted: AtomicUsize::new(0), processed_counter: AtomicUsize::new(0), worker_routing_key_load: WorkerRoutingKeyLoad::new(url), revision: AtomicU64::new(0), @@ -1138,6 +1181,31 @@ impl WorkerRuntime { .is_ok() } + // ── PD admission claims ───────────────────────────────────────── + + pub fn pd_admitted(&self) -> usize { + self.pd_admitted.load(Ordering::Relaxed) + } + + /// Claim `count` rooms, but only while the claimed total stays within + /// `window`. All-or-nothing: a refusal claims nothing. + pub fn try_admit_pd(&self, count: usize, window: usize) -> bool { + self.pd_admitted + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |admitted| { + admitted.checked_add(count).filter(|total| *total <= window) + }) + .is_ok() + } + + /// Release `count` claimed rooms, saturating at zero. + pub fn release_pd(&self, count: usize) { + let _ = self + .pd_admitted + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |admitted| { + Some(admitted.saturating_sub(count)) + }); + } + // ── Routing-key load ──────────────────────────────────────────── pub fn routing_key_load(&self) -> usize { @@ -1467,6 +1535,18 @@ impl Worker for BasicWorker { self.update_running_requests_metrics(); } + fn try_admit_pd(&self, count: usize, window: usize) -> bool { + self.runtime.load().try_admit_pd(count, window) + } + + fn release_pd(&self, count: usize) { + self.runtime.load().release_pd(count); + } + + fn pd_admitted(&self) -> usize { + self.runtime.load().pd_admitted() + } + fn routing_key_load(&self) -> usize { self.runtime.load().routing_key_load() } @@ -2031,6 +2111,81 @@ mod tests { assert_eq!(worker.max_running_requests(), None); } + #[test] + fn test_max_running_requests_reads_the_tokenspeed_spelling() { + use crate::worker::BasicWorkerBuilder; + // TokenSpeed advertises the same scheduler window as `max_num_seqs`; + // without this a TokenSpeed worker looks like a non-reporter to both + // fleet capacity and PD admission. + let mut labels = std::collections::HashMap::new(); + labels.insert("max_num_seqs".to_string(), "16".to_string()); + let worker = BasicWorkerBuilder::new("grpc://w:30000") + .labels(labels) + .build(); + assert_eq!(worker.max_running_requests(), Some(16)); + } + + #[test] + fn test_max_running_requests_prefers_max_num_seqs_when_both_are_reported() { + use crate::worker::BasicWorkerBuilder; + // The servicer's `running_window` resolves in the same order, so a + // worker reporting both spellings cannot have its discovery label + // disagree with its GetLoads report. + let mut labels = std::collections::HashMap::new(); + labels.insert("max_running_requests".to_string(), "256".to_string()); + labels.insert("max_num_seqs".to_string(), "16".to_string()); + let worker = BasicWorkerBuilder::new("grpc://w:30000") + .labels(labels) + .build(); + assert_eq!(worker.max_running_requests(), Some(16)); + } + + #[test] + fn test_pd_admission_claims_are_atomic_and_bounded_by_the_window() { + use crate::worker::BasicWorkerBuilder; + let worker = BasicWorkerBuilder::new("grpc://w:30000").build(); + + assert!(worker.try_admit_pd(2, 3)); + assert_eq!(worker.pd_admitted(), 2); + // All-or-nothing: a refusal must leave the claim untouched. + assert!(!worker.try_admit_pd(2, 3)); + assert_eq!(worker.pd_admitted(), 2); + assert!(worker.try_admit_pd(1, 3)); + assert_eq!(worker.pd_admitted(), 3); + + worker.release_pd(3); + assert_eq!(worker.pd_admitted(), 0); + // Releases saturate rather than wrapping to usize::MAX. + worker.release_pd(1); + assert_eq!(worker.pd_admitted(), 0); + } + + #[test] + fn test_pd_admission_never_overshoots_the_window_under_real_parallelism() { + use std::{sync::Arc, thread}; + + use crate::worker::BasicWorkerBuilder; + // The whole point of claiming rather than reading: threads racing for + // the last rooms must not all win. + let worker: Arc = Arc::new(BasicWorkerBuilder::new("grpc://w:30000").build()); + let window = 8; + let threads: Vec<_> = (0..32) + .map(|_| { + let worker = Arc::clone(&worker); + thread::spawn(move || worker.try_admit_pd(1, window)) + }) + .collect(); + + let admitted = threads + .into_iter() + .map(|handle| handle.join().expect("claim thread")) + .filter(|claimed| *claimed) + .count(); + + assert_eq!(admitted, window, "exactly the window's worth may claim"); + assert_eq!(worker.pd_admitted(), window); + } + #[test] fn test_is_realtime_capable_true_from_label() { use crate::worker::BasicWorkerBuilder;