From 6741cf588414bca232c2a28be7ef60e379396c4a Mon Sep 17 00:00:00 2001 From: Alex McC <319643551+hello-alexmcc@users.noreply.github.com> Date: Mon, 7 Sep 2026 22:51:04 -0700 Subject: [PATCH 1/3] fix(pd): fail fast on a failed leg and admit rooms the decode can take MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The gateway mints one bootstrap room per disaggregated request and dispatches both legs with tokio::join!. SGLang-lineage engines arm a bootstrap deadline on the prefill the moment it lands (120s on TokenSpeed) that clears only once the decode scheduler has admitted the request, and decode admission is bounded by the engine's running window. A burst wider than that window leaves the prefill deadline racing a queue the gateway built: prefill times out rooms the decode has not reached, and the decode then pre-allocates them and waits out its own 300s transfer deadline for a peer that is gone. Joining both legs withheld the prefill error until that second timeout, held the pair's slot throughout, and short-circuited before the decode's map_err so a decode-side failure was never logged. Seen in CI run 34173426995 (job 101902848341): 14 of 64 requests failed and MMLU took 1392s instead of ~35s. Dispatch now returns the first failure and retires the leg still in flight off the request path, dropping its stream as soon as it lands — the abort path a client disconnect already uses, and the only thing that releases the stranded room before the engine's own deadline. The retired leg's outcome still reaches its circuit breaker and its failure is logged. Before a pair is dispatched, the decode worker's in-flight count is checked against the engine's reported running window; at the window the request waits for a slot, bounded by --pd-admission-wait-secs (default 30, well under the bootstrap deadline), and sheds with the existing overload 503 if none frees. An engine that reports no window is not gated. TokenSpeed spells the window max_num_seqs, so it was never extracted from GetServerInfo and its workers looked like non-reporters to both admission and fleet capacity; both spellings are read now, and the servicer reports the configured window on GetLoads as SGLang already does. Signed-off-by: Alex McC <319643551+hello-alexmcc@users.noreply.github.com> --- bindings/python/src/lib.rs | 5 + bindings/python/src/smg/router_args.py | 15 + bindings/python/tests/test_arg_parser.py | 1 + .../smg_grpc_servicer/tokenspeed/loads.py | 57 +++ .../smg_grpc_servicer/tokenspeed/servicer.py | 38 +- grpc_servicer/tests/test_tokenspeed_loads.py | 75 ++++ model_gateway/src/app_context.rs | 7 +- model_gateway/src/config/builder.rs | 5 + model_gateway/src/config/types.rs | 35 ++ model_gateway/src/main.rs | 25 ++ model_gateway/src/observability/metrics.rs | 10 + model_gateway/src/routers/common/mod.rs | 3 + model_gateway/src/routers/common/overload.rs | 23 +- .../src/routers/common/pd_admission.rs | 283 +++++++++++++ model_gateway/src/routers/grpc/client.rs | 3 + .../grpc/common/stages/request_execution.rs | 398 +++++++++++++++--- model_gateway/src/routers/grpc/pipeline.rs | 106 ++++- model_gateway/src/worker/overload.rs | 10 + model_gateway/src/worker/worker.rs | 49 ++- 19 files changed, 1050 insertions(+), 98 deletions(-) create mode 100644 grpc_servicer/smg_grpc_servicer/tokenspeed/loads.py create mode 100644 grpc_servicer/tests/test_tokenspeed_loads.py create mode 100644 model_gateway/src/routers/common/pd_admission.rs 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..92cc4e7e76 --- /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_get_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..45e04f2229 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,25 @@ pub fn shed_if_worker_overloaded(worker: &dyn Worker, model_id: &str) -> Option< )) } +/// Shed a disaggregated dispatch the decode leg cannot admit: it is already +/// running `window` requests and no slot 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. +pub(crate) fn shed_pd_admission(worker: &str, model_id: &str, window: usize) -> Response { + shed( + BRANCH_PD_ADMISSION_SHED, + STAGE_PD_ADMISSION, + worker, + format!( + "Decode worker '{worker}' for model '{model_id}' is running its full \ + admission window ({window}) and no slot freed in time" + ), + ) +} + /// 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..d0551098ce --- /dev/null +++ b/model_gateway/src/routers/common/pd_admission.rs @@ -0,0 +1,283 @@ +//! 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 slot 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. +//! +//! 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}, + 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 a decode slot. 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 re-reads the worker's in-flight count. +/// +/// Polling, not a notifier: the event we would signal is a PD load guard +/// dropping, which happens in the worker layer with no channel back to the +/// router, 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 slot and the wait deadline below, +/// and the sleep is asynchronous — a waiting request occupies no thread. +const SLOT_POLL_INTERVAL: Duration = Duration::from_millis(50); + +/// Seconds a PD dispatch waits for a decode slot 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)) +} + +/// What the gate decided for one dispatch. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Slot { + /// The window is unknown, or there was room straight away. + Free, + /// A slot freed during the wait. + Waited, + /// The wait deadline passed with the window still full. + Full, +} + +/// Wait until `in_flight()` drops below `window`, or `wait` elapses. +/// +/// `window == 0` means the engine reported no running window, which is not the +/// same as a window of zero: the gate abstains. The caller increments the +/// worker's in-flight count immediately after a `Free`/`Waited` verdict with no +/// await in between, so the only overshoot is between dispatchers that read the +/// same last free slot concurrently — bounded by the number of them, and gone +/// by the next request's read. +async fn wait_for_slot( + window: usize, + wait: Duration, + in_flight: impl Fn() -> usize, +) -> (Slot, usize) { + if window == 0 { + return (Slot::Free, 0); + } + let observed = in_flight(); + if observed < window { + return (Slot::Free, observed); + } + let deadline = Instant::now() + wait; + loop { + let now = Instant::now(); + if now >= deadline { + return (Slot::Full, in_flight()); + } + tokio::time::sleep(SLOT_POLL_INTERVAL.min(deadline - now)).await; + let observed = in_flight(); + if observed < window { + return (Slot::Waited, observed); + } + } +} + +/// Gate one disaggregated dispatch on the decode leg's admission window. +/// +/// `Some(response)` is the shed the caller must return instead of dispatching; +/// `None` means the decode can take the request now. +pub(crate) async fn admit_decode(decode: &dyn Worker, model_id: &str) -> Option { + // `None` (unreported, or a nonsense zero) leaves admission to the engine. + let window = usize::from(decode.max_running_requests()?); + let wait = admission_wait(); + + match wait_for_slot(window, wait, || decode.load()).await { + (Slot::Free, _) => None, + (Slot::Waited, in_flight) => { + Metrics::record_pd_admission_wait(); + debug!( + worker = decode.url(), + model_id, window, in_flight, "PD admission waited for a decode slot" + ); + None + } + (Slot::Full, in_flight) => { + Metrics::record_pd_admission_shed(); + debug!( + worker = decode.url(), + model_id, + window, + in_flight, + waited_secs = wait.as_secs(), + "PD admission shed: no decode slot freed" + ); + Some(overload::shed_pd_admission(decode.url(), model_id, window)) + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::{ + atomic::{AtomicUsize, Ordering as AtomicOrdering}, + Arc, + }; + + use axum::http::{header::RETRY_AFTER, StatusCode}; + 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(); + } + assert!(admit_decode(decode.as_ref(), "m").await.is_none()); + } + + /// Below the window there is nothing to decide, and no sleep to pay for. + #[tokio::test(start_paused = true)] + async fn in_flight_below_window_admits_without_waiting() { + let decode = decode_worker("grpc://127.0.0.1:9902", Some(4)); + decode.increment_load(); + decode.increment_load(); + decode.increment_load(); + + let started = Instant::now(); + assert!(admit_decode(decode.as_ref(), "m").await.is_none()); + assert_eq!( + started.elapsed(), + Duration::ZERO, + "no wait below the window" + ); + } + + /// A slot freeing mid-wait releases the request instead of shedding it. + #[tokio::test(start_paused = true)] + async fn a_freed_slot_admits_the_waiting_request() { + let in_flight = Arc::new(AtomicUsize::new(4)); + let counter = Arc::clone(&in_flight); + let free_a_slot = async move { + tokio::time::sleep(Duration::from_millis(500)).await; + counter.fetch_sub(1, AtomicOrdering::SeqCst); + }; + let counter = Arc::clone(&in_flight); + let gate = wait_for_slot(4, Duration::from_secs(30), move || { + counter.load(AtomicOrdering::SeqCst) + }); + + let ((), (slot, observed)) = tokio::join!(free_a_slot, gate); + + assert_eq!(slot, Slot::Waited); + assert_eq!(observed, 3); + } + + /// 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 (slot, _) = wait_for_slot(2, Duration::from_secs(30), || 2).await; + assert_eq!(slot, Slot::Full); + 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 (slot, _) = wait_for_slot(2, Duration::ZERO, || 2).await; + assert_eq!(slot, Slot::Full); + 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:9903", Some(2)); + decode.increment_load(); + decode.increment_load(); + + let response = admit_decode(decode.as_ref(), "m") + .await + .expect("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..6da9721dcb 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,72 @@ 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, + }, +} + /// 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 { @@ -84,6 +148,18 @@ pub(crate) async fn execute_plan( ) })?; + // Admission for a disaggregated pair: the decode leg's engine window is + // what clears the prefill's bootstrap deadline, so a request the decode + // cannot admit yet must wait here rather than in the engine's queue — + // where it would burn the prefill's deadline and strand both rooms. The + // gate abstains when the engine reports no window, and returns with the + // slot still free: the load guards below claim it with no await between. + if let Some(decode) = workers.decode_worker() { + if let Some(shed) = pd_admission::admit_decode(decode.as_ref(), &ctx.model_id).await { + return Err(shed); + } + } + let sub_requests = match &execution_plan { ExecutionPlan::Batch { requests, .. } => requests.len(), _ => 1, @@ -378,65 +454,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 +934,71 @@ 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"), + } + } + #[test] fn pd_leg_labels_reflect_each_legs_transport() { let prefill: Arc = Arc::new( diff --git a/model_gateway/src/routers/grpc/pipeline.rs b/model_gateway/src/routers/grpc/pipeline.rs index a96ff93c9f..fb2825f401 100644 --- a/model_gateway/src/routers/grpc/pipeline.rs +++ b/model_gateway/src/routers/grpc/pipeline.rs @@ -1443,17 +1443,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 +1525,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 +1567,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 +1974,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..6b57d933ef 100644 --- a/model_gateway/src/worker/worker.rs +++ b/model_gateway/src/worker/worker.rs @@ -403,19 +403,26 @@ 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 — SGLang and vLLM advertise + /// `max_running_requests`, TokenSpeed advertises `max_num_seqs` — 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. /// /// `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 + let labels = &self.metadata().spec.labels; + labels .get("max_running_requests") + .or_else(|| labels.get("max_num_seqs")) .and_then(|s| s.parse::().ok()) .filter(|n| *n > 0) } @@ -2031,6 +2038,32 @@ 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_the_canonical_label() { + use crate::worker::BasicWorkerBuilder; + 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(256)); + } + #[test] fn test_is_realtime_capable_true_from_label() { use crate::worker::BasicWorkerBuilder; From 72fc876a3723f9039bcb160a1cf43033cf5d62b2 Mon Sep 17 00:00:00 2001 From: Alex McC <319643551+hello-alexmcc@users.noreply.github.com> Date: Tue, 8 Sep 2026 03:06:02 -0700 Subject: [PATCH 2/3] fix(pd): make decode admission an atomic claim sized to the plan MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review of #2466 found the gate read the in-flight count instead of reserving against it, so two dispatches racing for the last room in a window of one both saw it free and both sent — the over-window burst the gate exists to prevent. Admission now claims its rooms on the worker atomically and hands back a guard that releases them; the claim rides in the dispatch's LoadGuards, so it is returned on an early error, a failed leg, a client disconnect, a retry, or completion. The claim counter lives beside the load counter in WorkerRuntime, so a same-URL replacement inherits the rooms the engine still holds. The claim is sized to the plan: a batched completion fans out one PD dispatch per sub-request, so an n=64 request could post 64 bootstrap rooms past a gate that had asked about one. A plan wider than the window sheds at once, since no wait can free more rooms than the window holds. The gate also moves above the encode-output take, so an EPD request waiting on the decode window no longer pins the encode jobs' SHM for the whole wait and discards finished encode work on a shed. Label precedence now matches the servicer: max_num_seqs is authoritative for the engines that use it and is resolved first on both sides, so a worker's discovery label and its GetLoads report cannot disagree. Signed-off-by: Alex McC <319643551+hello-alexmcc@users.noreply.github.com> --- .../smg_grpc_servicer/tokenspeed/loads.py | 2 +- model_gateway/src/routers/common/overload.rs | 20 +- .../src/routers/common/pd_admission.rs | 312 +++++++++++++----- .../grpc/common/stages/request_execution.rs | 76 +++-- model_gateway/src/routers/grpc/context.rs | 23 +- model_gateway/src/routers/grpc/pipeline.rs | 2 + model_gateway/src/worker/worker.rs | 130 +++++++- 7 files changed, 441 insertions(+), 124 deletions(-) diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/loads.py b/grpc_servicer/smg_grpc_servicer/tokenspeed/loads.py index 92cc4e7e76..ad803a72fa 100644 --- a/grpc_servicer/smg_grpc_servicer/tokenspeed/loads.py +++ b/grpc_servicer/smg_grpc_servicer/tokenspeed/loads.py @@ -1,7 +1,7 @@ """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_get_loads.py), +TokenSpeed installed (see grpc_servicer/tests/test_tokenspeed_loads.py), the same split the SGLang side uses. """ diff --git a/model_gateway/src/routers/common/overload.rs b/model_gateway/src/routers/common/overload.rs index 45e04f2229..9de740b4e8 100644 --- a/model_gateway/src/routers/common/overload.rs +++ b/model_gateway/src/routers/common/overload.rs @@ -89,21 +89,29 @@ pub fn shed_if_worker_overloaded(worker: &dyn Worker, model_id: &str) -> Option< )) } -/// Shed a disaggregated dispatch the decode leg cannot admit: it is already -/// running `window` requests and no slot freed inside the admission wait. +/// 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. -pub(crate) fn shed_pd_admission(worker: &str, model_id: &str, window: usize) -> Response { +/// 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}' is running its full \ - admission window ({window}) and no slot freed in time" + "Decode worker '{worker}' for model '{model_id}' could not admit {rooms} \ + request(s) within its running window of {window}" ), ) } diff --git a/model_gateway/src/routers/common/pd_admission.rs b/model_gateway/src/routers/common/pd_admission.rs index d0551098ce..1a44906b59 100644 --- a/model_gateway/src/routers/common/pd_admission.rs +++ b/model_gateway/src/routers/common/pd_admission.rs @@ -13,16 +13,24 @@ //! //! 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 slot rather than joining the engine's queue, and sheds if +//! 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}, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, + }, time::Duration, }; @@ -32,22 +40,21 @@ use tracing::debug; use crate::{observability::metrics::Metrics, routers::common::overload, worker::Worker}; -/// Default seconds a PD dispatch may wait for a decode slot. Well under the +/// 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 re-reads the worker's in-flight count. +/// How often the wait retries its claim. /// -/// Polling, not a notifier: the event we would signal is a PD load guard -/// dropping, which happens in the worker layer with no channel back to the -/// router, and a per-worker registry of notifiers would be process-wide +/// 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 slot and the wait deadline below, +/// 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 SLOT_POLL_INTERVAL: Duration = Duration::from_millis(50); +const CLAIM_RETRY_INTERVAL: Duration = Duration::from_millis(50); -/// Seconds a PD dispatch waits for a decode slot before shedding. +/// 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 @@ -64,93 +71,146 @@ fn admission_wait() -> Duration { Duration::from_secs(PD_ADMISSION_WAIT_SECS.load(Ordering::Relaxed)) } -/// What the gate decided for one dispatch. +/// 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 Slot { - /// The window is unknown, or there was room straight away. - Free, - /// A slot freed during the wait. +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. - Full, + Refused, } -/// Wait until `in_flight()` drops below `window`, or `wait` elapses. +/// Retry `try_claim` until it succeeds or `wait` elapses. /// -/// `window == 0` means the engine reported no running window, which is not the -/// same as a window of zero: the gate abstains. The caller increments the -/// worker's in-flight count immediately after a `Free`/`Waited` verdict with no -/// await in between, so the only overshoot is between dispatchers that read the -/// same last free slot concurrently — bounded by the number of them, and gone -/// by the next request's read. -async fn wait_for_slot( - window: usize, - wait: Duration, - in_flight: impl Fn() -> usize, -) -> (Slot, usize) { - if window == 0 { - return (Slot::Free, 0); - } - let observed = in_flight(); - if observed < window { - return (Slot::Free, observed); +/// 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 (Slot::Full, in_flight()); + return Claim::Refused; } - tokio::time::sleep(SLOT_POLL_INTERVAL.min(deadline - now)).await; - let observed = in_flight(); - if observed < window { - return (Slot::Waited, observed); + 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. +/// Gate one disaggregated dispatch on the decode leg's admission window, +/// claiming the `rooms` bootstrap rooms it is about to post. /// -/// `Some(response)` is the shed the caller must return instead of dispatching; -/// `None` means the decode can take the request now. -pub(crate) async fn admit_decode(decode: &dyn Worker, model_id: &str) -> Option { +/// `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 window = usize::from(decode.max_running_requests()?); - let wait = admission_wait(); + 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); - match wait_for_slot(window, wait, || decode.load()).await { - (Slot::Free, _) => None, - (Slot::Waited, in_flight) => { + // 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, in_flight, "PD admission waited for a decode slot" + model_id, window, rooms, "PD admission waited for decode rooms" ); - None + Ok(Some(PdAdmissionGuard { + worker: Arc::clone(decode), + rooms, + })) } - (Slot::Full, in_flight) => { + Claim::Refused => { Metrics::record_pd_admission_shed(); debug!( worker = decode.url(), model_id, window, - in_flight, + rooms, + admitted = decode.pd_admitted(), waited_secs = wait.as_secs(), - "PD admission shed: no decode slot freed" + "PD admission shed: no decode rooms freed" ); - Some(overload::shed_pd_admission(decode.url(), model_id, window)) + Err(overload::shed_pd_admission( + decode.url(), + model_id, + window, + rooms, + )) } } } #[cfg(test)] mod tests { - use std::sync::{ - atomic::{AtomicUsize, Ordering as AtomicOrdering}, - Arc, - }; + 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::*; @@ -186,52 +246,131 @@ mod tests { for _ in 0..1_000 { decode.increment_load(); } - assert!(admit_decode(decode.as_ref(), "m").await.is_none()); + 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 decide, and no sleep to pay for. + /// 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 in_flight_below_window_admits_without_waiting() { + async fn a_free_window_admits_without_waiting_and_releases_on_drop() { let decode = decode_worker("grpc://127.0.0.1:9902", Some(4)); - decode.increment_load(); - decode.increment_load(); - decode.increment_load(); - let started = Instant::now(); - assert!(admit_decode(decode.as_ref(), "m").await.is_none()); + 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 slot freeing mid-wait releases the request instead of shedding it. + /// A room freeing mid-wait releases the request instead of shedding it. #[tokio::test(start_paused = true)] - async fn a_freed_slot_admits_the_waiting_request() { - let in_flight = Arc::new(AtomicUsize::new(4)); - let counter = Arc::clone(&in_flight); - let free_a_slot = async move { + 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; - counter.fetch_sub(1, AtomicOrdering::SeqCst); + claimed.store(0, AtomicOrdering::SeqCst); }; - let counter = Arc::clone(&in_flight); - let gate = wait_for_slot(4, Duration::from_secs(30), move || { - counter.load(AtomicOrdering::SeqCst) + let gate = claim_within(Duration::from_secs(30), || { + claimed + .compare_exchange(0, 1, AtomicOrdering::SeqCst, AtomicOrdering::SeqCst) + .is_ok() }); - let ((), (slot, observed)) = tokio::join!(free_a_slot, gate); + let ((), claim) = tokio::join!(free_a_room, gate); - assert_eq!(slot, Slot::Waited); - assert_eq!(observed, 3); + 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 (slot, _) = wait_for_slot(2, Duration::from_secs(30), || 2).await; - assert_eq!(slot, Slot::Full); + 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 {:?}", @@ -243,8 +382,8 @@ mod tests { #[tokio::test(start_paused = true)] async fn a_zero_wait_sheds_without_sleeping() { let started = Instant::now(); - let (slot, _) = wait_for_slot(2, Duration::ZERO, || 2).await; - assert_eq!(slot, Slot::Full); + let claim = claim_within(Duration::ZERO, || false).await; + assert_eq!(claim, Claim::Refused); assert_eq!(started.elapsed(), Duration::ZERO); } @@ -253,13 +392,12 @@ mod tests { #[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:9903", Some(2)); - decode.increment_load(); - decode.increment_load(); + 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.as_ref(), "m") + let response = admit_decode(&decode, "m", 1) .await - .expect("a full window sheds"); + .expect_err("a full window sheds"); assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); assert_eq!( extract_error_code_from_response(&response), 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 6da9721dcb..5331c931c9 100644 --- a/model_gateway/src/routers/grpc/common/stages/request_execution.rs +++ b/model_gateway/src/routers/grpc/common/stages/request_execution.rs @@ -97,6 +97,17 @@ enum PdDispatchOutcome { }, } +/// 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 { @@ -120,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 @@ -148,26 +183,9 @@ pub(crate) async fn execute_plan( ) })?; - // Admission for a disaggregated pair: the decode leg's engine window is - // what clears the prefill's bootstrap deadline, so a request the decode - // cannot admit yet must wait here rather than in the engine's queue — - // where it would burn the prefill's deadline and strand both rooms. The - // gate abstains when the engine reports no window, and returns with the - // slot still free: the load guards below claim it with no await between. - if let Some(decode) = workers.decode_worker() { - if let Some(shed) = pd_admission::admit_decode(decode.as_ref(), &ctx.model_id).await { - return Err(shed); - } - } - - 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. @@ -999,6 +1017,24 @@ mod tests { } } + /// 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 fb2825f401..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 { diff --git a/model_gateway/src/worker/worker.rs b/model_gateway/src/worker/worker.rs index 6b57d933ef..fd3e10ea5e 100644 --- a/model_gateway/src/worker/worker.rs +++ b/model_gateway/src/worker/worker.rs @@ -380,6 +380,21 @@ 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. + fn try_admit_pd(&self, count: usize, window: usize) -> bool; + + /// 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; + /// Get the current routing-key load cardinality. fn routing_key_load(&self) -> usize; @@ -408,11 +423,14 @@ pub trait Worker: Send + Sync + fmt::Debug + 'static { /// worker hasn't reported a value or reports zero (zero is meaningless /// for capacity accounting). /// - /// Engines spell the same window two ways — SGLang and vLLM advertise - /// `max_running_requests`, TokenSpeed advertises `max_num_seqs` — 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. + /// 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 @@ -421,8 +439,8 @@ pub trait Worker: Send + Sync + fmt::Debug + 'static { fn max_running_requests(&self) -> Option { let labels = &self.metadata().spec.labels; labels - .get("max_running_requests") - .or_else(|| labels.get("max_num_seqs")) + .get("max_num_seqs") + .or_else(|| labels.get("max_running_requests")) .and_then(|s| s.parse::().ok()) .filter(|n| *n > 0) } @@ -1053,6 +1071,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, @@ -1070,6 +1095,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), @@ -1145,6 +1171,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 { @@ -1474,6 +1525,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() } @@ -2053,15 +2116,64 @@ mod tests { } #[test] - fn test_max_running_requests_prefers_the_canonical_label() { + 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(256)); + 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] From 732c4b06e10f7e6eaaa8d58d435a4d9872ae7f4a Mon Sep 17 00:00:00 2001 From: Alex McC <319643551+hello-alexmcc@users.noreply.github.com> Date: Tue, 8 Sep 2026 04:17:20 -0700 Subject: [PATCH 3/3] fix(pd): give the admission claim methods defaults so every Worker keeps building The Go bindings carry their own Worker implementation, and the three new claim methods had no defaults, so smg-golang stopped compiling and the wheel build and the multimodal check failed. The defaults admit unconditionally and track nothing: the gate only reaches a worker that reports a running window, and an implementation with a shared runtime overrides them with the real claim. Signed-off-by: Alex McC <319643551+hello-alexmcc@users.noreply.github.com> --- model_gateway/src/worker/worker.rs | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/model_gateway/src/worker/worker.rs b/model_gateway/src/worker/worker.rs index fd3e10ea5e..df82ea0811 100644 --- a/model_gateway/src/worker/worker.rs +++ b/model_gateway/src/worker/worker.rs @@ -387,13 +387,23 @@ pub trait Worker: Send + Sync + fmt::Debug + 'static { /// [`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. - fn try_admit_pd(&self, count: usize, window: usize) -> bool; + /// + /// 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); + fn release_pd(&self, _count: usize) {} /// Rooms currently claimed on this worker. - fn pd_admitted(&self) -> usize; + fn pd_admitted(&self) -> usize { + 0 + } /// Get the current routing-key load cardinality. fn routing_key_load(&self) -> usize;