From f97255aaae9bfec31c375e50b384d0a78b3d5319 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Tue, 4 Aug 2026 00:20:34 +0000 Subject: [PATCH 01/88] feat(router): per-worker circuit breaker in both data planes Closes #82. Failover already retries a failed dispatch on another worker, but its memory is a per-request `tried` set that is discarded when the request returns. A worker that is broken for inference yet healthy to the platform -- it accepts the connection, answers /health, stays ACTIVE in discovery -- is therefore re-picked by the very next request, and every request after that pays the failover cost again. Kubernetes cannot see this worker and neither can discovery; the router is the only component that knows, and until now it forgot immediately. Adds a three-state breaker (closed / open / half-open) to both routers, with identical semantics and identical flags: after N consecutive faults a worker leaves the candidate list for a cooldown, then exactly one probe is admitted; success restores it, failure doubles the cooldown up to a cap. Backing off matters because the usual cause -- a worker wedged on a bad KV handoff -- does not clear on the first retry, and a fixed cooldown becomes a probe every N seconds forever. Scoped deliberately: * Only pre-first-byte failures count. Once bytes have streamed the failure is already non-retryable, and counting it would open the breaker on ordinary client disconnects. * 4xx never counts. A malformed request returns 400 from every worker it reaches, so feeding it in would trip the entire healthy fleet on one bad client. 429 is excluded too: it means "full right now", which the policy's load accounting already routes around. * If every candidate is open the router dispatches anyway rather than returning 503 -- a request served by a probably-bad worker beats turning a partial outage into a total one. * PD filters its two pools independently: a wedged prefill and a wedged decode are different events, and one must not evict the other role's healthy workers. * WorkerStatus is never written. That field belongs to discovery; this is the router's private opinion, exported instead as infera_router_worker_breaker_state / _trips_total. A worker tripping repeatedly while discovery still calls it ACTIVE is the alert worth having. --breaker-failure-threshold=0 disables the whole thing. Tests assert the cost, not the status code: failover made every request in these cases succeed both before and after the fix, so a status-only test proves nothing. Against a dead worker and a healthy one, the dead worker was dispatched to 10 times out of 10 before and 3 times after (Python), and 5 of 10 before and 3 after under round-robin (Rust). Both suites were re-run against the pre-fix behaviour to confirm they fail. The PD tests exist because lint, not the suite, caught two of the five record sites referring to a local name that does not exist in those scopes -- a NameError that would have fired only while handling a decode outage. Nothing had ever executed those generators. Signed-off-by: Zhang, Jiejing --- infera/router/auto.py | 8 +- infera/router/base.py | 7 + infera/router/breaker.py | 238 +++++++++++ infera/router/disagg.py | 10 + infera/router/mixed.py | 15 +- infera/server/__main__.py | 10 + infera/server/args.py | 24 ++ infera/server/launch_rust.py | 6 + infera/server/metrics.py | 18 + manual/features/routing_and_transport.md | 37 ++ manual/reference/cli.md | 3 + rust/router/src/breaker.rs | 509 +++++++++++++++++++++++ rust/router/src/config.rs | 13 + rust/router/src/disagg.rs | 88 +++- rust/router/src/handlers.rs | 22 +- rust/router/src/lib.rs | 1 + rust/router/src/main.rs | 8 +- rust/router/src/proxy.rs | 22 +- rust/router/tests/functional.rs | 83 ++++ tests/unit/router/test_breaker.py | 228 ++++++++++ tests/unit/router/test_disagg_breaker.py | 115 +++++ tests/unit/router/test_failover.py | 75 ++++ 22 files changed, 1515 insertions(+), 25 deletions(-) create mode 100644 infera/router/breaker.py create mode 100644 rust/router/src/breaker.rs create mode 100644 tests/unit/router/test_breaker.py create mode 100644 tests/unit/router/test_disagg_breaker.py diff --git a/infera/router/auto.py b/infera/router/auto.py index 918884b9..0f840c10 100644 --- a/infera/router/auto.py +++ b/infera/router/auto.py @@ -38,11 +38,17 @@ def __init__(self, *args, **kwargs) -> None: self.policy, nats_client=self.nats_client, request_max_retries=self.request_max_retries, + # One breaker shared by both sub-routers: otherwise each would build + # its own default and the configured thresholds would never reach + # them, since AutoRouter is what the server actually constructs. + breaker=self.breaker, ) # Pass the NATS request client to the PD router too, so disaggregated # (prefill/decode) dispatch uses the per-instance NATS transport when # configured (it falls back to HTTP only when nats_client is None). - self._disagg = DisaggRouter(self.pool, self.policy, nats_client=self.nats_client) + self._disagg = DisaggRouter( + self.pool, self.policy, nats_client=self.nats_client, breaker=self.breaker + ) async def aclose(self) -> None: await self._mixed.aclose() diff --git a/infera/router/base.py b/infera/router/base.py index ad2f91d9..5553dbe3 100644 --- a/infera/router/base.py +++ b/infera/router/base.py @@ -10,6 +10,7 @@ from fastapi import Response from infera.common.worker_pool import WorkerPool +from infera.router.breaker import CircuitBreaker from infera.router.policy.base import Policy @@ -28,6 +29,7 @@ def __init__( policy: Policy, nats_client=None, request_max_retries: int = 1, + breaker: CircuitBreaker | None = None, ) -> None: self.pool = pool self.policy = policy @@ -40,6 +42,11 @@ def __init__( # disables retries (single attempt). Mid-stream failures are never # retried (output already partially sent). self.request_max_retries = max(0, request_max_retries) + # Per-worker failure memory across requests. Failover alone forgets a + # bad worker the moment the request ends, so the next one re-picks it. + # Subclasses that select their own target consult this when filtering + # candidates; DirectRouter does not select and leaves it unused. + self.breaker = breaker if breaker is not None else CircuitBreaker() @abstractmethod async def dispatch( diff --git a/infera/router/breaker.py b/infera/router/breaker.py new file mode 100644 index 00000000..91764ff2 --- /dev/null +++ b/infera/router/breaker.py @@ -0,0 +1,238 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Per-worker circuit breaker for routers that select their own target. + +Failover alone is not enough. It retries a failed dispatch on another worker, +but the memory of that failure lives in a per-request ``tried`` set that is +discarded when the request returns -- so the next request scores the same broken +worker as if nothing happened, picks it again if cache locality says so, and +pays the failover cost again. A worker that is healthy to the platform and +broken for inference therefore taxes *every* request, indefinitely. + +That worker is not hypothetical: it accepts the connection, answers ``/health``, +stays ``ACTIVE`` in discovery, and fails before the first byte. Kubernetes +cannot see it and neither can discovery. The router is the only component that +knows, and without this it forgets immediately. + +Scope, deliberately narrow: + +* **Only pre-first-byte failures trip it.** A failure after bytes have been + streamed is already non-retryable by design -- the worker demonstrably served + part of the request, and treating that as a health signal would open the + breaker on ordinary client disconnects. +* **It never touches ``WorkerStatus``.** That field is owned by discovery; this + is the router's private opinion, applied when filtering candidates. +* **Only routers that select use it.** ``direct.py`` has no failover because the + gateway owns selection there, so a breaker would be wrong. + +States are the usual three. ``closed`` routes normally. After +``failure_threshold`` consecutive failures the breaker goes ``open`` and the +worker is excluded for ``cooldown`` seconds. It then becomes ``half_open`` and +admits exactly one probe: success closes it and clears the count, failure +reopens it with the cooldown doubled, up to ``max_cooldown``. Backing off +matters because the common cause -- a worker wedged on a bad KV handoff -- does +not resolve on the first retry, and a fixed cooldown turns into a probe every +``cooldown`` seconds forever. +""" + +from __future__ import annotations + +import logging +import time +from dataclasses import dataclass, field +from enum import Enum + +from infera.server import metrics + +logger = logging.getLogger(__name__) + +_STATE_VALUE = {"closed": 0, "half_open": 1, "open": 2} + + +def _observe(worker_id: str, state) -> None: + """Mirror a state change to Prometheus. Never allowed to fail a request -- + an unregistered collector or a duplicate registry must not take out the + data plane.""" + try: + metrics.worker_breaker_state.labels(worker_id=worker_id).set(_STATE_VALUE[state.value]) + except Exception: # pragma: no cover + pass + + +def is_worker_fault(status: int) -> bool: + """True if an HTTP status is evidence about the *worker*, not the request. + + Failover retries on any pre-first-byte error, including 4xx -- that is + correct, since a 400 costs nothing to re-ask. Feeding 4xx to the breaker is + not: a malformed request returns 400 from every worker it touches, so the + breaker would trip the entire healthy fleet on one bad client. 429 is + excluded for a different reason -- it means "full right now", which the + policy's load accounting already routes around, and a 5s cooldown with + doubling is far too heavy a response to transient backpressure. + """ + return status >= 500 or status == 0 + + +class BreakerState(str, Enum): + CLOSED = "closed" + OPEN = "open" + HALF_OPEN = "half_open" + + +@dataclass +class _Entry: + consecutive_failures: int = 0 + state: BreakerState = BreakerState.CLOSED + # Wall time after which an open breaker becomes half-open. + opens_until: float = 0.0 + # Cooldown applied on the *next* trip; doubles each time a probe fails. + next_cooldown: float = 0.0 + # True while a half-open probe is in flight, so only one is admitted. + probe_in_flight: bool = False + trips: int = 0 + + +@dataclass +class CircuitBreaker: + """Tracks per-worker health as seen by dispatch outcomes. + + Not thread-safe by design: routers drive it from a single asyncio loop, and + a lock here would sit on the hot path of every request for no benefit. + """ + + failure_threshold: int = 3 + cooldown: float = 5.0 + max_cooldown: float = 60.0 + #: Injectable clock, so tests do not sleep. + now: object = field(default=time.monotonic) + _entries: dict[str, _Entry] = field(default_factory=dict, init=False) + + # --- queries ------------------------------------------------------------ + + def _entry(self, worker_id: str) -> _Entry: + e = self._entries.get(worker_id) + if e is None: + e = _Entry(next_cooldown=self.cooldown) + self._entries[worker_id] = e + return e + + @property + def enabled(self) -> bool: + """A threshold of 0 or less turns the breaker off entirely, so an + operator can fall back to plain failover without a code change.""" + return self.failure_threshold > 0 + + def allows(self, worker_id: str) -> bool: + """True if this worker may be dispatched to right now. + + Transitions open -> half_open as a side effect when the cooldown has + elapsed, because the alternative is a separate timer whose only job is + to flip a flag that this function already has to check. + """ + if not self.enabled: + return True + e = self._entries.get(worker_id) + if e is None or e.state is BreakerState.CLOSED: + return True + if e.state is BreakerState.OPEN: + if self.now() < e.opens_until: + return False + e.state = BreakerState.HALF_OPEN + e.probe_in_flight = False + _observe(worker_id, e.state) + logger.info("breaker: worker %s half-open, admitting one probe", worker_id) + # half_open: admit exactly one probe. + if e.probe_in_flight: + return False + e.probe_in_flight = True + return True + + def filter(self, workers): + """Drop workers whose breaker is open. Returns a list. + + If every candidate is open, returns them all rather than nothing: a + request served by a probably-bad worker beats a 503 when there is no + alternative, and refusing to route would turn a partial outage into a + total one. + """ + allowed = [w for w in workers if self.allows(self._id_of(w))] + if allowed: + return allowed + if workers: + logger.warning( + "breaker: all %d candidate(s) open; routing anyway rather than failing", + len(workers), + ) + return list(workers) + + @staticmethod + def _id_of(w) -> str: + # Accepts a WorkerInfo or anything exposing .worker_id. + return getattr(w, "worker_id", None) or str(w) + + def state_of(self, worker_id: str) -> BreakerState: + e = self._entries.get(worker_id) + return e.state if e else BreakerState.CLOSED + + # --- outcomes ----------------------------------------------------------- + + def record_success(self, worker_id: str) -> None: + e = self._entries.get(worker_id) + if e is None: + return + if e.state is not BreakerState.CLOSED: + logger.info("breaker: worker %s recovered, closing", worker_id) + e.consecutive_failures = 0 + e.state = BreakerState.CLOSED + e.probe_in_flight = False + e.next_cooldown = self.cooldown + _observe(worker_id, e.state) + + def record_failure(self, worker_id: str) -> None: + """Record a pre-first-byte dispatch failure.""" + if not self.enabled: + return + e = self._entry(worker_id) + e.consecutive_failures += 1 + was_probe = e.state is BreakerState.HALF_OPEN + e.probe_in_flight = False + + if was_probe: + # A failed probe reopens immediately and backs off further, without + # waiting for the threshold again -- we already know it is bad. + e.next_cooldown = min(e.next_cooldown * 2, self.max_cooldown) + self._open(worker_id, e) + return + if e.consecutive_failures >= self.failure_threshold: + self._open(worker_id, e) + + def _open(self, worker_id: str, e: _Entry) -> None: + e.state = BreakerState.OPEN + e.opens_until = self.now() + e.next_cooldown + e.trips += 1 + _observe(worker_id, e.state) + try: + metrics.worker_breaker_trips_total.labels(worker_id=worker_id).inc() + except Exception: # pragma: no cover - metrics must never break routing + pass + logger.warning( + "breaker: worker %s open for %.1fs after %d consecutive failure(s)", + worker_id, + e.next_cooldown, + e.consecutive_failures, + ) + + # --- introspection for metrics / tests ----------------------------------- + + def snapshot(self) -> dict[str, dict]: + return { + wid: { + "state": e.state.value, + "consecutive_failures": e.consecutive_failures, + "trips": e.trips, + } + for wid, e in self._entries.items() + } diff --git a/infera/router/disagg.py b/infera/router/disagg.py index a340e744..1b3ef9da 100644 --- a/infera/router/disagg.py +++ b/infera/router/disagg.py @@ -119,6 +119,11 @@ async def dispatch( model = body.get("model") prefills = self.pool.list_active(model=model, mode=DisaggMode.PREFILL) decodes = self.pool.list_active(model=model, mode=DisaggMode.DECODE) + # Independently per role: a wedged prefill and a wedged decode are + # different events against different pools, and one open breaker + # must not remove the other role's healthy workers. + prefills = self.breaker.filter(prefills) + decodes = self.breaker.filter(decodes) if not prefills or not decodes: obs["outcome"] = "503" metrics.pd_bootstrap_failures_total.labels(reason="no_pd_workers").inc() @@ -341,6 +346,7 @@ async def _post(url, leg, worker_id, leg_body, leg_headers): except httpx.HTTPError as exc: obs["outcome"] = "502" metrics.pd_bootstrap_failures_total.labels(reason="worker_unreachable").inc() + self.breaker.record_failure(p.worker_id) return _sanitized_error("PD request failed", exc, status_code=502) if p_resp.status_code >= 400: @@ -547,6 +553,7 @@ async def _dispatch_serial( p_failed = True obs["outcome"] = "502" metrics.pd_bootstrap_failures_total.labels(reason="prefill_unreachable").inc() + self.breaker.record_failure(p.worker_id) return _sanitized_error("prefill leg failed", exc, status_code=502) if p_resp.status_code >= 400: @@ -632,6 +639,7 @@ async def _dispatch_serial( except httpx.HTTPError as exc: obs["outcome"] = "502" metrics.pd_bootstrap_failures_total.labels(reason="decode_unreachable").inc() + self.breaker.record_failure(d.worker_id) return _sanitized_error("decode leg failed", exc, status_code=502) try: @@ -677,6 +685,7 @@ async def _stream_decode_only( exc or "", ) metrics.pd_bootstrap_failures_total.labels(reason="decode_unreachable").inc() + self.breaker.record_failure(d_target.worker.worker_id) err = json.dumps({"error": "decode unreachable"}) yield f"data: {err}\n\n".encode() return @@ -831,6 +840,7 @@ async def _stream_dual( exc or "", ) metrics.pd_bootstrap_failures_total.labels(reason="decode_unreachable").inc() + self.breaker.record_failure(d_target.worker.worker_id) # json.dumps: exc text may contain chars that break SSE. err = json.dumps({"error": "decode unreachable"}) yield f"data: {err}\n\n".encode() diff --git a/infera/router/mixed.py b/infera/router/mixed.py index 6f1836a1..5c3e66d2 100644 --- a/infera/router/mixed.py +++ b/infera/router/mixed.py @@ -16,6 +16,7 @@ from infera.common.nats_request import TYPE_DATA, TYPE_DONE, TYPE_ERROR from infera.common.worker_pool import DisaggMode from infera.router.base import BaseRouter +from infera.router.breaker import is_worker_fault from infera.router.cache_control import parse_cache_hints from infera.router.dp_routing import dp_rank_header from infera.router.engine_priority import inject_engine_priority @@ -82,13 +83,25 @@ async def dispatch( for w in self.pool.list_active(model=model, mode=DisaggMode.MIXED) if w.worker_id not in tried ] + # Drop workers the breaker has open. Falls back to the unfiltered + # list when every candidate is open -- a request served by a + # probably-bad worker beats turning a partial outage into a 503. + candidates = self.breaker.filter(candidates) if not candidates: break target, blocks = self.policy.pick(candidates, body) tried.add(target.worker.worker_id) try: - return await self._attempt(target, blocks, body, hints, path, stream, obs) + resp = await self._attempt(target, blocks, body, hints, path, stream, obs) + self.breaker.record_success(target.worker.worker_id) + return resp except _Retry as r: + # Pre-first-byte only: _Retry is never raised once bytes have + # been streamed, so a mid-stream failure cannot trip this. + # 4xx is retried but not held against the worker -- see + # is_worker_fault(). + if is_worker_fault(getattr(r.response, "status_code", 0)): + self.breaker.record_failure(target.worker.worker_id) last_error = r.response logger.info( "failover: worker %s failed before first byte; %d worker(s) tried", diff --git a/infera/server/__main__.py b/infera/server/__main__.py index f9400561..5465fa00 100644 --- a/infera/server/__main__.py +++ b/infera/server/__main__.py @@ -21,6 +21,7 @@ from infera.kv.subscriber import KvEventSubscriberPool from infera.kv.writer import KvIndexWriter from infera.router.auto import AutoRouter +from infera.router.breaker import CircuitBreaker from infera.router.direct import DirectRouter from infera.router.policy.factory import build_policy from infera.server.app import init_app @@ -223,6 +224,13 @@ def on_worker_removed(worker_id: str) -> None: logger.info("request transport: nats (per-instance subjects)") # --- Router + FastAPI app --- + # Per-worker failure memory, shared by every router this process builds. + # DirectRouter never selects, so it holds one but does not consult it. + breaker = CircuitBreaker( + failure_threshold=args.breaker_failure_threshold, + cooldown=args.breaker_cooldown_s, + max_cooldown=args.breaker_max_cooldown_s, + ) # router-mode=direct trusts an upstream GAIE EPP's per-request worker # selection (x-worker-instance-id header); auto selects in-process. if args.router_mode == "direct": @@ -231,6 +239,7 @@ def on_worker_removed(worker_id: str) -> None: policy, nats_client=nats_request_client, request_max_retries=args.request_max_retries, + breaker=breaker, ) logger.info("router-mode=direct (honouring GAIE EPP x-worker-instance-id)") else: @@ -239,6 +248,7 @@ def on_worker_removed(worker_id: str) -> None: policy, nats_client=nats_request_client, request_max_retries=args.request_max_retries, + breaker=breaker, ) app = init_app( registry, diff --git a/infera/server/args.py b/infera/server/args.py index 59b7266a..9297c5d4 100644 --- a/infera/server/args.py +++ b/infera/server/args.py @@ -205,4 +205,28 @@ def parse_server_args(argv: list[str] | None = None) -> argparse.Namespace: "backlog). Mid-stream failures are never retried. Default 1; 0 disables. " "Overrides $INFERA_REQUEST_MAX_RETRIES.", ) + parser.add_argument( + "--breaker-failure-threshold", + type=int, + default=int(os.environ.get("INFERA_BREAKER_FAILURE_THRESHOLD", "3") or 3), + help="Consecutive pre-first-byte worker faults (5xx / unreachable; 4xx " + "and 429 excluded) before the router takes a worker out of rotation. " + "Failover alone forgets between requests, so a worker that is ACTIVE in " + "discovery but broken for inference is otherwise re-picked forever. " + "0 disables the breaker. Overrides $INFERA_BREAKER_FAILURE_THRESHOLD.", + ) + parser.add_argument( + "--breaker-cooldown-s", + type=float, + default=float(os.environ.get("INFERA_BREAKER_COOLDOWN_S", "5") or 5), + help="Seconds a tripped worker is excluded before one probe request is " + "admitted. Overrides $INFERA_BREAKER_COOLDOWN_S.", + ) + parser.add_argument( + "--breaker-max-cooldown-s", + type=float, + default=float(os.environ.get("INFERA_BREAKER_MAX_COOLDOWN_S", "60") or 60), + help="Ceiling for the cooldown, which doubles on each failed probe. " + "Overrides $INFERA_BREAKER_MAX_COOLDOWN_S.", + ) return parser.parse_args(argv) diff --git a/infera/server/launch_rust.py b/infera/server/launch_rust.py index fbc23bce..16ea9a1d 100644 --- a/infera/server/launch_rust.py +++ b/infera/server/launch_rust.py @@ -80,6 +80,12 @@ def exec_rust(args: argparse.Namespace) -> None: args.request_transport, "--request-max-retries", str(args.request_max_retries), + "--breaker-failure-threshold", + str(args.breaker_failure_threshold), + "--breaker-cooldown-s", + str(args.breaker_cooldown_s), + "--breaker-max-cooldown-s", + str(args.breaker_max_cooldown_s), ] # kv-aware needs the tokenizer + overlap weights, or it degrades to # load-only routing (no cache locality). Resolve HF ids to a local path diff --git a/infera/server/metrics.py b/infera/server/metrics.py index ebd8af9b..b130bcd5 100644 --- a/infera/server/metrics.py +++ b/infera/server/metrics.py @@ -133,6 +133,24 @@ ) +worker_breaker_state = Gauge( + "infera_router_worker_breaker_state", + "Router-side circuit breaker per worker: 0=closed, 1=half_open, 2=open. " + "Non-zero means the router is routing around a worker that discovery still " + "reports ACTIVE — the gap this metric exists to make visible.", + labelnames=("worker_id",), + registry=REGISTRY, +) + +worker_breaker_trips_total = Counter( + "infera_router_worker_breaker_trips_total", + "Times a worker's breaker has opened. A worker tripping repeatedly while " + "staying ACTIVE is broken for inference but healthy to the platform.", + labelnames=("worker_id",), + registry=REGISTRY, +) + + # ---------------------------------------------------------------------- # KV-aware policy internals # ---------------------------------------------------------------------- diff --git a/manual/features/routing_and_transport.md b/manual/features/routing_and_transport.md index 7ec103d9..ef249024 100644 --- a/manual/features/routing_and_transport.md +++ b/manual/features/routing_and_transport.md @@ -113,6 +113,43 @@ idle-timeout-before-first-token, or a 429 admission reject. It never retries mid-stream (once tokens flow, a failure surfaces to the client). Raise it for more resilience; set `0` to fail fast. +### Circuit breaker + +Failover on its own has no memory. Its `tried` set lives for one request, so a +worker that is broken for inference but healthy to discovery — it accepts the +connection, answers `/health`, stays `ACTIVE` — gets re-picked by the *next* +request, and every request after that pays the failover cost again. + +The breaker is that missing memory. After `--breaker-failure-threshold` +consecutive faults a worker is dropped from the candidate list for +`--breaker-cooldown-s`, then one probe request is admitted: if it succeeds the +worker is restored, if it fails the cooldown doubles, up to +`--breaker-max-cooldown-s`. + +| Flag | Env | Default | Meaning | +|---|---|---|---| +| `--breaker-failure-threshold` | `INFERA_BREAKER_FAILURE_THRESHOLD` | `3` | consecutive faults before removal; `0` disables | +| `--breaker-cooldown-s` | `INFERA_BREAKER_COOLDOWN_S` | `5` | exclusion window before a probe | +| `--breaker-max-cooldown-s` | `INFERA_BREAKER_MAX_COOLDOWN_S` | `60` | cap on the doubling backoff | + +Two exclusions are deliberate. **4xx never counts** — a malformed request returns +400 from every worker it reaches, so counting it would trip the entire healthy +fleet on one bad client. **429 never counts** either: it means "full right now", +which the policy's load accounting already routes around, and a doubling cooldown +is far too heavy a response to transient backpressure. + +If *every* candidate is open the router dispatches anyway rather than returning +503 — a request served by a probably-bad worker beats turning a partial outage +into a total one. + +The breaker never writes `WorkerStatus`; that field belongs to discovery. This is +the router's private opinion, and it is visible as +`infera_router_worker_breaker_state` (0 closed / 1 half-open / 2 open) and +`infera_router_worker_breaker_trips_total`. A worker tripping repeatedly while +discovery still reports it `ACTIVE` is the signal worth alerting on. + +Both the Python and Rust routers implement this identically, with the same flags. + ## KV-event transport Powers [KV-aware routing](kv_aware_routing.md). `--kv-event-transport`: diff --git a/manual/reference/cli.md b/manual/reference/cli.md index ed4afd56..184171c5 100644 --- a/manual/reference/cli.md +++ b/manual/reference/cli.md @@ -22,6 +22,9 @@ the **same** on the server and every worker. See | `--kv-prefill-overlap-weight` | (unset) | KV-aware PD: prefill-side weight (typical `20.0`); overrides the global | | `--kv-decode-overlap-weight` | (unset) | KV-aware PD: decode-side weight (typical `2.0`); overrides the global | | `--request-max-retries` | `1` | retry on an alternate worker on pre-response failure (never mid-stream); `0` disables | +| `--breaker-failure-threshold` | `3` | consecutive worker faults (5xx / unreachable) before a worker leaves rotation; `0` disables the breaker | +| `--breaker-cooldown-s` | `5` | how long a tripped worker is excluded before one probe request is admitted | +| `--breaker-max-cooldown-s` | `60` | ceiling for that cooldown, which doubles on each failed probe | | `--discovery-backend` | `kubernetes` | `kubernetes` \| `etcd` | | `--etcd-endpoint` | — | required for `--discovery-backend etcd` | | `--etcd-prefix` | `/infera/workers/` | etcd key prefix the fleet registers under | diff --git a/rust/router/src/breaker.rs b/rust/router/src/breaker.rs new file mode 100644 index 00000000..92d49d17 --- /dev/null +++ b/rust/router/src/breaker.rs @@ -0,0 +1,509 @@ +/////////////////////////////////////////////////////////////////////////////// +// Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +// +// SPDX-License-Identifier: MIT +/////////////////////////////////////////////////////////////////////////////// +//! Per-worker circuit breaker. Mirrors `infera/router/breaker.py` — same three +//! states, same thresholds, same all-open fallback — so the two data planes +//! behave identically under a wedged worker. +//! +//! Failover on its own is not enough. It retries a failed dispatch elsewhere, +//! but the memory of that failure lives in a per-request `tried` set that is +//! dropped when the request returns, so the next request scores the same broken +//! worker as if nothing had happened. A worker that answers `/health`, stays +//! ACTIVE in etcd, and fails before the first byte therefore taxes *every* +//! request, indefinitely. +//! +//! Unlike the Python side this is shared across tokio worker threads, so the +//! map lives behind a `Mutex`. The critical sections are a hash lookup and a +//! few integer writes; contention is not a concern at any plausible request +//! rate, and a lock-free design here would buy nothing for the complexity. + +use std::collections::HashMap; +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +/// Statuses that are evidence about the *worker*, not the request. +/// +/// Failover retries on any pre-first-byte error including 4xx, which is +/// correct — re-asking costs nothing. Feeding 4xx to the breaker is not: a +/// malformed request returns 400 from every worker it touches, so one bad +/// client would trip the whole healthy fleet. 429 is excluded for a different +/// reason: it means "full right now", which load accounting already routes +/// around, and a doubling cooldown is far too heavy for transient backpressure. +pub fn is_worker_fault(status: u16) -> bool { + status >= 500 || status == 0 +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BreakerState { + Closed, + Open, + HalfOpen, +} + +impl BreakerState { + pub fn as_str(self) -> &'static str { + match self { + BreakerState::Closed => "closed", + BreakerState::Open => "open", + BreakerState::HalfOpen => "half_open", + } + } +} + +#[derive(Debug)] +struct Entry { + consecutive_failures: u32, + state: BreakerState, + /// Instant after which an open breaker becomes half-open. + opens_until: Instant, + /// Cooldown applied on the *next* trip; doubles each time a probe fails. + next_cooldown: Duration, + /// Set while a half-open probe is in flight, so only one is admitted. + probe_in_flight: bool, + trips: u64, +} + +pub struct CircuitBreaker { + failure_threshold: u32, + cooldown: Duration, + max_cooldown: Duration, + entries: Mutex>, +} + +impl CircuitBreaker { + pub fn new(failure_threshold: u32, cooldown: Duration, max_cooldown: Duration) -> Self { + Self { + failure_threshold, + cooldown, + max_cooldown, + entries: Mutex::new(HashMap::new()), + } + } + + /// Whether this worker may be dispatched to right now. + /// + /// Transitions Open -> HalfOpen as a side effect once the cooldown has + /// elapsed, because the alternative is a background timer whose only job is + /// to flip a flag this function already has to read. Call it once per + /// candidate per request: in HalfOpen it *consumes* the single probe slot. + pub fn allows(&self, worker_id: &str) -> bool { + self.allows_at(worker_id, Instant::now()) + } + + /// A threshold of 0 turns the breaker off entirely, so an operator can fall + /// back to plain failover without a code change. + fn enabled(&self) -> bool { + self.failure_threshold > 0 + } + + fn allows_at(&self, worker_id: &str, now: Instant) -> bool { + if !self.enabled() { + return true; + } + let mut map = self.entries.lock().expect("breaker mutex poisoned"); + let Some(e) = map.get_mut(worker_id) else { + return true; + }; + match e.state { + BreakerState::Closed => return true, + BreakerState::Open => { + if now < e.opens_until { + return false; + } + e.state = BreakerState::HalfOpen; + e.probe_in_flight = false; + tracing::info!(worker = worker_id, "breaker half-open, admitting one probe"); + } + BreakerState::HalfOpen => {} + } + if e.probe_in_flight { + return false; + } + e.probe_in_flight = true; + true + } + + /// Drop workers whose breaker is open. + /// + /// If *every* candidate is open the full list is returned instead of an + /// empty one: a request served by a probably-bad worker beats a guaranteed + /// 503, and refusing to route would turn a partial outage into a total one. + pub fn filter(&self, workers: &[W], id_of: impl Fn(&W) -> &str) -> Vec { + let allowed: Vec = workers + .iter() + .filter(|w| self.allows(id_of(w))) + .cloned() + .collect(); + if !allowed.is_empty() { + return allowed; + } + if !workers.is_empty() { + tracing::warn!( + candidates = workers.len(), + "breaker: all candidates open; routing anyway rather than failing" + ); + } + workers.to_vec() + } + + pub fn record_success(&self, worker_id: &str) { + let mut map = self.entries.lock().expect("breaker mutex poisoned"); + if let Some(e) = map.get_mut(worker_id) { + if e.state != BreakerState::Closed { + tracing::info!(worker = worker_id, "breaker: worker recovered, closing"); + } + e.consecutive_failures = 0; + e.state = BreakerState::Closed; + e.probe_in_flight = false; + e.next_cooldown = self.cooldown; + } + } + + /// Record a pre-first-byte dispatch failure. Callers must gate this on + /// [`is_worker_fault`] when the failure carries an HTTP status. + pub fn record_failure(&self, worker_id: &str) { + self.record_failure_at(worker_id, Instant::now()); + } + + fn record_failure_at(&self, worker_id: &str, now: Instant) { + if !self.enabled() { + return; + } + let mut map = self.entries.lock().expect("breaker mutex poisoned"); + let e = map.entry(worker_id.to_string()).or_insert_with(|| Entry { + consecutive_failures: 0, + state: BreakerState::Closed, + opens_until: now, + next_cooldown: self.cooldown, + probe_in_flight: false, + trips: 0, + }); + e.consecutive_failures += 1; + let was_probe = e.state == BreakerState::HalfOpen; + e.probe_in_flight = false; + + if was_probe { + // A failed probe reopens immediately and backs off further, without + // waiting out the threshold again — we already know it is bad. The + // common cause (a worker wedged on a bad KV handoff) does not clear + // on the first retry, and a fixed cooldown would probe it at a + // constant rate forever. + e.next_cooldown = (e.next_cooldown * 2).min(self.max_cooldown); + } else if e.consecutive_failures < self.failure_threshold { + return; + } + e.state = BreakerState::Open; + e.opens_until = now + e.next_cooldown; + e.trips += 1; + tracing::warn!( + worker = worker_id, + cooldown_s = e.next_cooldown.as_secs_f64(), + failures = e.consecutive_failures, + "breaker: worker open" + ); + } + + pub fn state_of(&self, worker_id: &str) -> BreakerState { + self.entries + .lock() + .expect("breaker mutex poisoned") + .get(worker_id) + .map(|e| e.state) + .unwrap_or(BreakerState::Closed) + } + + /// `(worker_id, state, trips)` for metrics export. + pub fn snapshot(&self) -> Vec<(String, BreakerState, u64)> { + let map = self.entries.lock().expect("breaker mutex poisoned"); + let mut out: Vec<_> = map + .iter() + .map(|(k, e)| (k.clone(), e.state, e.trips)) + .collect(); + out.sort_by(|a, b| a.0.cmp(&b.0)); + out + } +} + +impl Default for CircuitBreaker { + fn default() -> Self { + Self::new(3, Duration::from_secs(5), Duration::from_secs(60)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Time is driven forward explicitly rather than slept, so the cooldown + /// behaviour is tested at full speed and deterministically. + fn cb() -> CircuitBreaker { + CircuitBreaker::new(3, Duration::from_secs(5), Duration::from_secs(20)) + } + + #[derive(Clone)] + struct W(&'static str); + + #[test] + fn threshold_zero_disables_it() { + // Without the guard, `failures >= 0` would trip on the first failure — + // the exact opposite of what --breaker-failure-threshold=0 promises. + let b = CircuitBreaker::new(0, Duration::from_secs(5), Duration::from_secs(20)); + for _ in 0..20 { + b.record_failure("w1"); + } + assert!(b.allows("w1")); + assert_eq!(b.state_of("w1"), BreakerState::Closed); + let ws = vec![W("a"), W("b")]; + assert_eq!(b.filter(&ws, |w| w.0).len(), 2); + } + + #[test] + fn unknown_worker_is_allowed() { + let b = cb(); + assert!(b.allows("w1")); + assert_eq!(b.state_of("w1"), BreakerState::Closed); + } + + #[test] + fn failures_below_threshold_do_not_open() { + let b = cb(); + b.record_failure("w1"); + b.record_failure("w1"); + assert!(b.allows("w1")); + assert_eq!(b.state_of("w1"), BreakerState::Closed); + } + + #[test] + fn opens_at_threshold_and_excludes() { + let b = cb(); + for _ in 0..3 { + b.record_failure("w1"); + } + assert_eq!(b.state_of("w1"), BreakerState::Open); + assert!(!b.allows("w1")); + } + + #[test] + fn success_resets_the_count() { + let b = cb(); + b.record_failure("w1"); + b.record_failure("w1"); + b.record_success("w1"); + b.record_failure("w1"); + b.record_failure("w1"); + assert_eq!(b.state_of("w1"), BreakerState::Closed); + } + + #[test] + fn half_open_admits_exactly_one_probe() { + let b = cb(); + let t0 = Instant::now(); + for _ in 0..3 { + b.record_failure_at("w1", t0); + } + let t1 = t0 + Duration::from_millis(5_100); + assert!(b.allows_at("w1", t1), "cooldown elapsed -> one probe"); + assert_eq!(b.state_of("w1"), BreakerState::HalfOpen); + assert!( + !b.allows_at("w1", t1), + "a second concurrent request must not also probe" + ); + } + + #[test] + fn successful_probe_closes() { + let b = cb(); + let t0 = Instant::now(); + for _ in 0..3 { + b.record_failure_at("w1", t0); + } + let t1 = t0 + Duration::from_millis(5_100); + b.allows_at("w1", t1); + b.record_success("w1"); + assert_eq!(b.state_of("w1"), BreakerState::Closed); + assert!(b.allows_at("w1", t1)); + } + + #[test] + fn failed_probe_reopens_with_doubled_cooldown() { + let b = cb(); + let t0 = Instant::now(); + for _ in 0..3 { + b.record_failure_at("w1", t0); + } + let t1 = t0 + Duration::from_millis(5_100); + b.allows_at("w1", t1); + b.record_failure_at("w1", t1); // probe fails -> reopen, 5s -> 10s + assert_eq!(b.state_of("w1"), BreakerState::Open); + + let t2 = t1 + Duration::from_millis(5_100); // old cooldown would be up + assert!(!b.allows_at("w1", t2), "backoff must have doubled"); + let t3 = t1 + Duration::from_millis(10_100); + assert!(b.allows_at("w1", t3)); + } + + #[test] + fn cooldown_is_capped() { + let b = cb(); + let mut t = Instant::now(); + for _ in 0..3 { + b.record_failure_at("w1", t); + } + for _ in 0..6 { + t += Duration::from_secs(1000); + b.allows_at("w1", t); + b.record_failure_at("w1", t); + } + assert!( + b.allows_at("w1", t + Duration::from_millis(20_100)), + "cooldown must not grow without bound" + ); + } + + #[test] + fn filter_drops_open_workers() { + let b = cb(); + for _ in 0..3 { + b.record_failure("bad"); + } + let ws = vec![W("good"), W("bad")]; + let got = b.filter(&ws, |w| w.0); + assert_eq!(got.len(), 1); + assert_eq!(got[0].0, "good"); + } + + #[test] + fn filter_returns_all_when_every_worker_is_open() { + let b = cb(); + for id in ["a", "b"] { + for _ in 0..3 { + b.record_failure(id); + } + } + let ws = vec![W("a"), W("b")]; + assert_eq!(b.filter(&ws, |w| w.0).len(), 2); + } + + #[test] + fn filter_of_empty_is_empty() { + let b = cb(); + let ws: Vec = vec![]; + assert!(b.filter(&ws, |w| w.0).is_empty()); + } + + #[test] + fn workers_are_independent() { + let b = cb(); + for _ in 0..3 { + b.record_failure("bad"); + } + assert!(b.allows("good")); + } + + #[test] + fn success_on_unknown_worker_is_harmless() { + let b = cb(); + b.record_success("never-seen"); + assert!(b.allows("never-seen")); + } + + #[test] + fn snapshot_reports_trips() { + let b = cb(); + let t0 = Instant::now(); + for _ in 0..3 { + b.record_failure_at("w1", t0); + } + let t1 = t0 + Duration::from_millis(5_100); + b.allows_at("w1", t1); + b.record_failure_at("w1", t1); + let snap = b.snapshot(); + assert_eq!(snap.len(), 1); + assert_eq!(snap[0].1, BreakerState::Open); + assert_eq!(snap[0].2, 2, "initial trip plus the failed probe"); + } + + #[test] + fn client_errors_are_not_worker_faults() { + for s in [400u16, 404, 422, 429] { + assert!(!is_worker_fault(s), "{s} must not trip the breaker"); + } + for s in [0u16, 500, 502, 503, 504] { + assert!(is_worker_fault(s), "{s} must trip the breaker"); + } + } + + #[test] + fn a_bad_client_cannot_trip_the_fleet() { + let b = cb(); + let ws = vec![W("a"), W("b"), W("c")]; + for _ in 0..10 { + for w in b.filter(&ws, |w| w.0) { + if is_worker_fault(400) { + b.record_failure(w.0); + } + } + } + for w in &ws { + assert_eq!(b.state_of(w.0), BreakerState::Closed); + } + } + + #[test] + fn the_regression_this_exists_for() { + // A worker that fails every dispatch must stop being selected. Before + // this type existed `tried` was per-request, so `bad` was offered on + // all ten requests. + let b = cb(); + let ws = vec![W("good"), W("bad")]; + let mut offered_bad = 0; + for _ in 0..10 { + let cands = b.filter(&ws, |w| w.0); + if cands.iter().any(|w| w.0 == "bad") { + offered_bad += 1; + b.record_failure("bad"); + } + b.record_success("good"); + } + assert_eq!(offered_bad, 3, "bad worker must stop being offered"); + } + + #[test] + fn concurrent_probes_admit_only_one() { + // The Python breaker is single-loop; this one is shared across tokio + // threads, so the half-open slot has to be safe under real contention. + use std::sync::Arc; + let b = Arc::new(cb()); + for _ in 0..3 { + b.record_failure("w1"); + } + // Force half-open by driving the clock through the private hook. + let t = Instant::now() + Duration::from_secs(6); + assert!(b.allows_at("w1", t)); + b.record_failure_at("w1", t); // back to open, then reopen at 10s + let t2 = t + Duration::from_secs(11); + + let admitted = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let hs: Vec<_> = (0..16) + .map(|_| { + let b = b.clone(); + let admitted = admitted.clone(); + std::thread::spawn(move || { + if b.allows_at("w1", t2) { + admitted.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + } + }) + }) + .collect(); + for h in hs { + h.join().unwrap(); + } + assert_eq!( + admitted.load(std::sync::atomic::Ordering::SeqCst), + 1, + "exactly one of 16 racing threads may probe" + ); + } +} diff --git a/rust/router/src/config.rs b/rust/router/src/config.rs index 123ab838..2e25f01b 100644 --- a/rust/router/src/config.rs +++ b/rust/router/src/config.rs @@ -27,6 +27,19 @@ pub struct Config { #[arg(long, default_value_t = 1)] pub request_max_retries: usize, + /// Consecutive pre-first-byte worker faults before a worker is taken out + /// of rotation. Failover alone forgets between requests; this remembers. + #[arg(long, default_value_t = 3)] + pub breaker_failure_threshold: u32, + + /// Seconds a tripped worker is excluded before one probe is admitted. + #[arg(long, default_value_t = 5.0)] + pub breaker_cooldown_s: f64, + + /// Ceiling for the cooldown, which doubles on each failed probe. + #[arg(long, default_value_t = 60.0)] + pub breaker_max_cooldown_s: f64, + /// `round-robin` or `kv-aware` (DP-attention cache-locality routing). #[arg(long, default_value = "round-robin")] pub router_policy: String, diff --git a/rust/router/src/disagg.rs b/rust/router/src/disagg.rs index fcf7b9d2..7f141c4b 100644 --- a/rust/router/src/disagg.rs +++ b/rust/router/src/disagg.rs @@ -12,6 +12,7 @@ //! KVPoll until a ~300s timeout. A detached `tokio::spawn` gives us exactly //! that: it outlives the client connection. +use std::sync::Arc; use std::time::Duration; use axum::body::{Body, Bytes}; @@ -19,6 +20,7 @@ use axum::http::{header, StatusCode}; use axum::response::Response; use serde_json::{Map, Value}; +use crate::breaker::{is_worker_fault, CircuitBreaker}; use crate::dp; use crate::handlers::AppState; use crate::policy::{ActiveGuard, Role}; @@ -41,16 +43,21 @@ pub async fn dispatch( ) -> Response { // role_hint lets a cost-aware policy weight P (cache-heavy: a hit skips a // whole prefill pass) differently from D (route by load). - let p_pick = state.policy.pick( - snap.list_active(model, DisaggMode::Prefill), - request, - Role::Prefill, - ); - let d_pick = state.policy.pick( - snap.list_active(model, DisaggMode::Decode), - request, - Role::Decode, - ); + // Each pool is filtered against the breaker independently: a wedged prefill + // and a wedged decode are different events against different pools, and one + // open breaker must not remove the other role's healthy workers. + let p_avail = state + .breaker + .filter(snap.list_active(model, DisaggMode::Prefill), |w| { + w.worker_id.as_str() + }); + let d_avail = state + .breaker + .filter(snap.list_active(model, DisaggMode::Decode), |w| { + w.worker_id.as_str() + }); + let p_pick = state.policy.pick(&p_avail, request, Role::Prefill); + let d_pick = state.policy.pick(&d_avail, request, Role::Decode); let p = p_pick.target; let d = d_pick.target; // One guard for both legs; dropped when the decode body finishes streaming @@ -121,7 +128,14 @@ async fn stream_dual( d_body: Map, guard: ActiveGuard, ) -> Response { - spawn_prefill_drain(state.http.clone(), p_url, p_body, p.dp_rank); + spawn_prefill_drain( + state.http.clone(), + state.breaker.clone(), + p.worker.worker_id.clone(), + p_url, + p_body, + p.dp_rank, + ); match open_decode(state, d, &d_url, &d_body).await { Ok(resp) => Response::builder() @@ -167,13 +181,26 @@ async fn unary_dual( st.as_u16() ); } + if is_worker_fault(st.as_u16()) { + state.breaker.record_failure(&p.worker.worker_id); + } else { + state.breaker.record_success(&p.worker.worker_id); + } + } + Err(e) => { + tracing::warn!("prefill {} failed: {e}", p_url); + state.breaker.record_failure(&p.worker.worker_id); } - Err(e) => tracing::warn!("prefill {} failed: {e}", p_url), } match d_res { Ok(resp) => { let st = resp.status(); + if is_worker_fault(st.as_u16()) { + state.breaker.record_failure(&d.worker.worker_id); + } else { + state.breaker.record_success(&d.worker.worker_id); + } let ct = content_type(&resp); match resp.bytes().await { Ok(bytes) => Response::builder() @@ -187,10 +214,13 @@ async fn unary_dual( ), } } - Err(e) => json_error( - StatusCode::BAD_GATEWAY, - &format!("decode {} unreachable: {e}", d.worker.worker_id), - ), + Err(e) => { + state.breaker.record_failure(&d.worker.worker_id); + json_error( + StatusCode::BAD_GATEWAY, + &format!("decode {} unreachable: {e}", d.worker.worker_id), + ) + } } } @@ -198,6 +228,8 @@ async fn unary_dual( /// Never awaited by the request path, so a client disconnect can't cancel it. fn spawn_prefill_drain( http: reqwest::Client, + breaker: Arc, + worker_id: String, url: String, body: Map, dp_rank: Option, @@ -217,8 +249,19 @@ fn spawn_prefill_drain( st.as_u16() ); } + // This leg is detached, so its outcome never reaches the client + // — but a prefill that 5xx's still leaves decode hanging on + // KVPoll, which is exactly the failure worth remembering. + if is_worker_fault(st.as_u16()) { + breaker.record_failure(&worker_id); + } else { + breaker.record_success(&worker_id); + } + } + Err(e) => { + tracing::warn!("prefill {url} failed: {e} (decode may hang on KVPoll)"); + breaker.record_failure(&worker_id); } - Err(e) => tracing::warn!("prefill {url} failed: {e} (decode may hang on KVPoll)"), } }); } @@ -237,6 +280,9 @@ async fn open_decode( Ok(resp) => { let st = resp.status(); if st.is_client_error() || st.is_server_error() { + if is_worker_fault(st.as_u16()) { + state.breaker.record_failure(&d.worker.worker_id); + } let txt = resp.text().await.unwrap_or_default(); return Err(format!( "decode {} error {}: {}", @@ -245,6 +291,7 @@ async fn open_decode( &txt[..txt.len().min(300)] )); } + state.breaker.record_success(&d.worker.worker_id); return Ok(resp); } Err(e) if attempt < DECODE_OPEN_RETRIES => { @@ -255,7 +302,12 @@ async fn open_decode( tokio::time::sleep(backoff).await; backoff = (backoff * 2).min(Duration::from_millis(500)); } - Err(e) => return Err(format!("decode {} unreachable: {e}", d.worker.worker_id)), + Err(e) => { + // Exhausted the in-request retries: this worker is not merely + // slow to accept a connection. + state.breaker.record_failure(&d.worker.worker_id); + return Err(format!("decode {} unreachable: {e}", d.worker.worker_id)); + } } } unreachable!("loop returns on the final attempt") diff --git a/rust/router/src/handlers.rs b/rust/router/src/handlers.rs index d731267f..1e8f2b7d 100644 --- a/rust/router/src/handlers.rs +++ b/rust/router/src/handlers.rs @@ -15,6 +15,7 @@ use axum::routing::{get, post}; use axum::{Json, Router}; use serde_json::json; +use crate::breaker::CircuitBreaker; use crate::policy::Policy; use crate::pool::SharedPool; use crate::proxy; @@ -26,6 +27,9 @@ pub struct AppState { pub http: reqwest::Client, pub started: Instant, pub retries: usize, + /// Per-worker failure memory. Shared across threads and across requests — + /// that persistence across requests is the whole point (see breaker.rs). + pub breaker: Arc, } pub fn app(state: AppState) -> Router { @@ -75,11 +79,25 @@ async fn models(State(st): State) -> impl IntoResponse { async fn metrics(State(st): State) -> impl IntoResponse { let snap = st.pool.load(); - format!( + let mut out = format!( "# infera-router (rust)\n\ infera_router_active_workers {}\n\ infera_router_uptime_seconds {}\n", snap.active_count(), st.started.elapsed().as_secs() - ) + ); + // Non-zero state means the router is routing around a worker that + // discovery still reports ACTIVE — the gap this metric exists to show. + for (worker_id, state, trips) in st.breaker.snapshot() { + let v = match state { + crate::breaker::BreakerState::Closed => 0, + crate::breaker::BreakerState::HalfOpen => 1, + crate::breaker::BreakerState::Open => 2, + }; + out.push_str(&format!( + "infera_router_worker_breaker_state{{worker_id=\"{worker_id}\"}} {v}\n\ + infera_router_worker_breaker_trips_total{{worker_id=\"{worker_id}\"}} {trips}\n" + )); + } + out } diff --git a/rust/router/src/lib.rs b/rust/router/src/lib.rs index 2555db29..0983afca 100644 --- a/rust/router/src/lib.rs +++ b/rust/router/src/lib.rs @@ -15,6 +15,7 @@ //! Modules are `pub` so the binary and the `tests/` suite share one API. pub mod block_hasher; +pub mod breaker; pub mod cache_control; pub mod config; pub mod disagg; diff --git a/rust/router/src/main.rs b/rust/router/src/main.rs index 3b728175..0ede6af1 100644 --- a/rust/router/src/main.rs +++ b/rust/router/src/main.rs @@ -7,12 +7,13 @@ //! crate (see `lib.rs`); this just wires config → discovery → server. use std::sync::Arc; -use std::time::Instant; +use std::time::{Duration, Instant}; use arc_swap::ArcSwap; use tracing_subscriber::EnvFilter; use infera_router::block_hasher::BlockHasher; +use infera_router::breaker; use infera_router::config::Config; use infera_router::handlers::{app, AppState}; use infera_router::kv_event::KvEventClient; @@ -69,6 +70,11 @@ async fn main() -> anyhow::Result<()> { http: proxy::build_upstream_client()?, started: Instant::now(), retries: cfg.request_max_retries, + breaker: Arc::new(breaker::CircuitBreaker::new( + cfg.breaker_failure_threshold, + Duration::from_secs_f64(cfg.breaker_cooldown_s), + Duration::from_secs_f64(cfg.breaker_max_cooldown_s), + )), }; let addr = format!("{}:{}", cfg.host, cfg.port); diff --git a/rust/router/src/proxy.rs b/rust/router/src/proxy.rs index 7bab99df..05a67675 100644 --- a/rust/router/src/proxy.rs +++ b/rust/router/src/proxy.rs @@ -18,6 +18,7 @@ use axum::response::Response; use futures::Stream; use serde_json::Value; +use crate::breaker::is_worker_fault; use crate::dp; use crate::handlers::AppState; use crate::policy::{ActiveGuard, Role}; @@ -108,6 +109,10 @@ async fn mixed_dispatch( if avail.is_empty() { break; } + // Drop workers the breaker has open. Falls back to the unfiltered list + // when every candidate is open — a request served by a probably-bad + // worker beats turning a partial outage into a 503. + let avail = state.breaker.filter(&avail, |w| w.worker_id.as_str()); let pick = state.policy.pick(&avail, request, Role::Mixed); tried.insert(pick.target.worker.worker_id.clone()); // Load guard: started here, dropped when this attempt fails (fail-over) @@ -116,9 +121,22 @@ async fn mixed_dispatch( state.policy.clone(), vec![(pick.target.route_key(), pick.blocks.clone())], ); + let wid = pick.target.worker.worker_id.clone(); match attempt(state, &pick.target, &raw, stream, path, guard).await { - Ok(resp) => return resp, - Err(err_resp) => last_err = Some(err_resp), + Ok(resp) => { + state.breaker.record_success(&wid); + return resp; + } + Err(err_resp) => { + // `attempt` only returns Err before any byte reached the client, + // so a mid-stream failure can never trip the breaker. 4xx is + // failed over but not held against the worker — see + // is_worker_fault(). + if is_worker_fault(err_resp.status().as_u16()) { + state.breaker.record_failure(&wid); + } + last_err = Some(err_resp); + } } } last_err.unwrap_or_else(|| json_error(StatusCode::SERVICE_UNAVAILABLE, "all workers failed")) diff --git a/rust/router/tests/functional.rs b/rust/router/tests/functional.rs index 86b0fac9..8cdebda9 100644 --- a/rust/router/tests/functional.rs +++ b/rust/router/tests/functional.rs @@ -21,6 +21,7 @@ use axum::{Json, Router}; use serde_json::{json, Value}; use infera_router::block_hasher::BlockHasher; +use infera_router::breaker::CircuitBreaker; use infera_router::handlers::{app, AppState}; use infera_router::kv_event::KvEventClient; use infera_router::policy::{KvEventAwarePolicy, RoundRobin}; @@ -106,6 +107,7 @@ fn make_state(workers: Vec>, retries: usize) -> AppState { http: proxy::build_upstream_client().unwrap(), started: Instant::now(), retries, + breaker: Arc::new(CircuitBreaker::default()), } } @@ -199,6 +201,7 @@ fn make_kv_state(workers: Vec>, retries: usize) -> AppState { http: proxy::build_upstream_client().unwrap(), started: Instant::now(), retries, + breaker: Arc::new(CircuitBreaker::default()), } } @@ -287,6 +290,86 @@ async fn mixed_failover_to_healthy_worker() { assert_eq!(ok.hit_count(), 1); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn breaker_stops_reselecting_a_dead_worker() { + // The regression behind issue #82, end to end through the real router + // rather than against the breaker in isolation. + // + // Failover already made every one of these ten requests succeed, so a test + // that only checked status codes passed before the fix and after it. What + // was broken is the *cost*: `tried` is per-request, so RoundRobin kept + // offering the dead worker its turn -- 5 of 10 requests paid a wasted + // upstream round trip. The assertion that matters is bad.hit_count(), which + // is 5 without the breaker and 3 with it. + let (url_bad, bad) = spawn_mock(500, false, json!(null)).await; + let (url_ok, ok) = spawn_mock(200, false, json!({"ok": true})).await; + let state = make_state( + vec![ + worker( + json!({"worker_id": "bad", "url": url_bad, "model_name": "m", "disagg_mode": "mixed"}), + ), + worker( + json!({"worker_id": "ok", "url": url_ok, "model_name": "m", "disagg_mode": "mixed"}), + ), + ], + 1, + ); + let router = spawn_router(state).await; + + for _ in 0..10 { + let resp = client() + .post(format!("{router}/v1/chat/completions")) + .json(&json!({"model": "m"})) + .send() + .await + .unwrap(); + assert_eq!( + resp.status(), + 200, + "failover must still serve every request" + ); + } + + // Default threshold is 3. RoundRobin offers `bad` on every other request, + // so it takes 3 of its turns to trip; after that it is out of rotation and + // the 5s cooldown does not elapse within the test. + assert_eq!( + bad.hit_count(), + 3, + "dead worker must stop being re-picked after the threshold (was 5 before the fix)" + ); + assert_eq!(ok.hit_count(), 10, "healthy worker still serves everything"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn breaker_ignores_client_errors() { + // A 400 comes from the request, not the worker: it would be returned by + // every worker in the fleet, so counting it would circuit-break all of + // them. Ten bad requests must leave the worker in rotation. + let (url, w) = spawn_mock(400, false, json!(null)).await; + let state = make_state( + vec![worker( + json!({"worker_id": "w1", "url": url, "model_name": "m", "disagg_mode": "mixed"}), + )], + 0, + ); + let router = spawn_router(state).await; + + for _ in 0..10 { + let _ = client() + .post(format!("{router}/v1/chat/completions")) + .json(&json!({"model": "m"})) + .send() + .await + .unwrap(); + } + assert_eq!( + w.hit_count(), + 10, + "4xx must not take a healthy worker out of rotation" + ); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn mixed_streaming_relays_sse() { let (url, _mock) = spawn_mock(200, true, json!(null)).await; diff --git a/tests/unit/router/test_breaker.py b/tests/unit/router/test_breaker.py new file mode 100644 index 00000000..e8a87af8 --- /dev/null +++ b/tests/unit/router/test_breaker.py @@ -0,0 +1,228 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Circuit breaker: a failing worker must stop being re-picked. + +The bug this guards against is not "failover is broken" -- failover works. It is +that failover's memory is per-request, so the *next* request scores a wedged +worker as if nothing happened. Every test here is written against that: the +assertions are about what happens on the second and third request, not the +first. + +Time is injected rather than slept, so the cooldown behaviour is tested at full +speed and deterministically. +""" + +from __future__ import annotations + +import pytest + +from infera.router.breaker import BreakerState, CircuitBreaker, is_worker_fault + + +class FakeClock: + def __init__(self) -> None: + self.t = 1000.0 + + def __call__(self) -> float: + return self.t + + def advance(self, dt: float) -> None: + self.t += dt + + +class W: + """Minimal stand-in for WorkerInfo -- the breaker only reads worker_id.""" + + def __init__(self, wid: str) -> None: + self.worker_id = wid + + +@pytest.fixture +def clock(): + return FakeClock() + + +@pytest.fixture +def cb(clock): + return CircuitBreaker(failure_threshold=3, cooldown=5.0, max_cooldown=20.0, now=clock) + + +def test_unknown_worker_is_allowed(cb): + """A worker never seen must route; the breaker is not an allowlist.""" + assert cb.allows("w1") is True + assert cb.state_of("w1") is BreakerState.CLOSED + + +def test_failures_below_threshold_do_not_open(cb): + for _ in range(2): + cb.record_failure("w1") + assert cb.allows("w1") is True + assert cb.state_of("w1") is BreakerState.CLOSED + + +def test_opens_at_threshold_and_excludes(cb): + """The actual bug: after N failures the worker must stop being offered.""" + for _ in range(3): + cb.record_failure("w1") + assert cb.state_of("w1") is BreakerState.OPEN + assert cb.allows("w1") is False + + +def test_success_resets_the_count(cb): + """Intermittent failures must not accumulate into a trip.""" + cb.record_failure("w1") + cb.record_failure("w1") + cb.record_success("w1") + cb.record_failure("w1") + cb.record_failure("w1") + assert cb.state_of("w1") is BreakerState.CLOSED + assert cb.allows("w1") is True + + +def test_half_open_admits_exactly_one_probe(cb, clock): + for _ in range(3): + cb.record_failure("w1") + clock.advance(5.1) + assert cb.allows("w1") is True, "cooldown elapsed -> one probe admitted" + assert cb.state_of("w1") is BreakerState.HALF_OPEN + assert cb.allows("w1") is False, "a second concurrent request must not also probe" + + +def test_successful_probe_closes(cb, clock): + for _ in range(3): + cb.record_failure("w1") + clock.advance(5.1) + cb.allows("w1") + cb.record_success("w1") + assert cb.state_of("w1") is BreakerState.CLOSED + assert cb.allows("w1") is True + + +def test_failed_probe_reopens_with_doubled_cooldown(cb, clock): + """A wedged worker does not recover on the first retry. A fixed cooldown + would probe it forever at a constant rate; this asserts the backoff.""" + for _ in range(3): + cb.record_failure("w1") + clock.advance(5.1) + cb.allows("w1") + cb.record_failure("w1") # probe fails -> reopen, cooldown 5 -> 10 + assert cb.state_of("w1") is BreakerState.OPEN + + clock.advance(5.1) # old cooldown would have elapsed + assert cb.allows("w1") is False, "backoff must have doubled" + clock.advance(5.0) # now past 10s + assert cb.allows("w1") is True + + +def test_cooldown_is_capped(cb, clock): + for _ in range(3): + cb.record_failure("w1") + for _ in range(6): # drive the doubling past max_cooldown + clock.advance(1000.0) + cb.allows("w1") + cb.record_failure("w1") + clock.advance(20.1) # max_cooldown = 20 + assert cb.allows("w1") is True, "cooldown must not grow without bound" + + +def test_filter_drops_open_workers(cb): + ws = [W("good"), W("bad")] + for _ in range(3): + cb.record_failure("bad") + assert [w.worker_id for w in cb.filter(ws)] == ["good"] + + +def test_filter_returns_all_when_every_worker_is_open(cb): + """Refusing to route would turn a partial outage into a total one. A + request served by a probably-bad worker beats a guaranteed 503.""" + ws = [W("a"), W("b")] + for wid in ("a", "b"): + for _ in range(3): + cb.record_failure(wid) + assert {w.worker_id for w in cb.filter(ws)} == {"a", "b"} + + +def test_filter_of_empty_is_empty(cb): + assert cb.filter([]) == [] + + +def test_workers_are_independent(cb): + for _ in range(3): + cb.record_failure("bad") + assert cb.allows("good") is True + assert cb.state_of("good") is BreakerState.CLOSED + + +def test_success_on_unknown_worker_is_harmless(cb): + cb.record_success("never-seen") + assert cb.allows("never-seen") is True + + +def test_snapshot_reports_trips(cb, clock): + for _ in range(3): + cb.record_failure("w1") + clock.advance(5.1) + cb.allows("w1") + cb.record_failure("w1") + snap = cb.snapshot()["w1"] + assert snap["state"] == "open" + assert snap["trips"] == 2, "initial trip plus the failed probe" + + +@pytest.mark.parametrize("status", [500, 502, 503, 504, 0]) +def test_server_errors_are_worker_faults(status): + assert is_worker_fault(status) is True + + +@pytest.mark.parametrize("status", [400, 404, 422, 429]) +def test_client_errors_are_not_worker_faults(status): + """A malformed request 400s on every worker it reaches. Counting that as a + health signal would trip the breaker across an entirely healthy fleet. + 429 is excluded separately: it means "full now", which load accounting + already routes around, and a doubling cooldown is far too heavy for it.""" + assert is_worker_fault(status) is False + + +def test_a_bad_client_cannot_trip_the_fleet(cb): + """Ten malformed requests against three healthy workers must leave all + three closed.""" + ws = [W("a"), W("b"), W("c")] + for _ in range(10): + for w in cb.filter(ws): + if is_worker_fault(400): + cb.record_failure(w.worker_id) + assert all(cb.state_of(w.worker_id) is BreakerState.CLOSED for w in ws) + + +def test_the_regression_this_exists_for(cb): + """End to end in breaker terms: a worker that fails every dispatch stops + being selected, instead of being re-picked on every subsequent request. + + Before this class existed, `tried` was per-request, so the loop below would + have offered `bad` on all ten requests. + """ + ws = [W("good"), W("bad")] + offered_bad = 0 + for _ in range(10): + candidates = cb.filter(ws) + if any(w.worker_id == "bad" for w in candidates): + offered_bad += 1 + cb.record_failure("bad") + cb.record_success("good") + assert offered_bad == 3, f"bad worker offered {offered_bad} times, expected 3 (the threshold)" + + +def test_threshold_zero_disables_it(): + """The documented off switch. Without this, threshold=0 would satisfy + `failures >= threshold` on the very first failure and trip immediately -- + the exact opposite of what --breaker-failure-threshold=0 promises.""" + off = CircuitBreaker(failure_threshold=0) + for _ in range(20): + off.record_failure("w1") + assert off.allows("w1") is True + assert off.state_of("w1") is BreakerState.CLOSED + ws = [W("a"), W("b")] + assert len(off.filter(ws)) == 2 diff --git a/tests/unit/router/test_disagg_breaker.py b/tests/unit/router/test_disagg_breaker.py new file mode 100644 index 00000000..0856714a --- /dev/null +++ b/tests/unit/router/test_disagg_breaker.py @@ -0,0 +1,115 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""The breaker's PD call sites, exercised on the paths that actually reach them. + +These are the streaming generators in ``disagg.py``. They only run when a +decode leg is unreachable -- which is exactly the condition the breaker exists +for, and exactly what no existing test drove. Two of the five record sites were +written against a local name (``d``) that does not exist in these scopes; the +result would have been a ``NameError`` raised *while handling* a decode outage, +turning a clean SSE error into a 500. Lint caught it, but nothing executed it. +""" + +from __future__ import annotations + +import httpx +import pytest + +from infera.common.worker_pool import EngineType, WorkerInfo +from infera.router.disagg import DisaggRouter +from infera.router.policy.target import RouteTarget + + +class _FakePolicy: + def pick(self, candidates, body): + return RouteTarget(candidates[0]), [] + + def on_request_started(self, route_key, blocks): + pass + + def on_request_finished(self, route_key, blocks): + pass + + +class _FakePool: + def __init__(self, workers): + self._workers = workers + + def list_active(self, model=None, mode=None): + return list(self._workers) + + +def _w(wid): + return WorkerInfo( + worker_id=wid, + url=f"http://{wid}", + model_name="m", + engine=EngineType.SGLANG, + request_transport="http", + ) + + +def _router(): + r = DisaggRouter(_FakePool([_w("p1"), _w("d1")]), _FakePolicy()) + # Every send fails at the transport layer: the decode leg is unreachable. + r._client = httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: (_ for _ in ()).throw(httpx.ConnectError("refused", request=request)) + ) + ) + # The pre-flight retry loop sleeps between attempts; not worth the wall time. + r._DECODE_OPEN_MAX_RETRIES = 0 + return r + + +async def _drain(agen) -> bytes: + out = b"" + async for chunk in agen: + out += chunk if isinstance(chunk, bytes) else chunk.encode() + return out + + +@pytest.mark.asyncio +async def test_decode_only_stream_records_failure_on_unreachable(): + r = _router() + d_target = RouteTarget(_w("d1")) + body = await _drain( + r._stream_decode_only(d_target, [], "http://d1/v1/chat/completions", {"model": "m"}) + ) + assert b"decode unreachable" in body, "client must get a clean SSE error, not a traceback" + assert r.breaker.state_of("d1").value == "closed", "one failure is below the threshold" + for _ in range(2): + await _drain( + r._stream_decode_only(d_target, [], "http://d1/v1/chat/completions", {"model": "m"}) + ) + assert r.breaker.state_of("d1").value == "open", "three unreachable decodes must trip it" + await r.aclose() + + +@pytest.mark.asyncio +async def test_dual_stream_records_failure_on_unreachable_decode(): + r = _router() + p_target = RouteTarget(_w("p1")) + d_target = RouteTarget(_w("d1")) + for _ in range(3): + body = await _drain( + r._stream_dual( + p_target, + [], + d_target, + [], + "http://p1/v1/chat/completions", + "http://d1/v1/chat/completions", + {"model": "m"}, + {"model": "m"}, + ) + ) + assert b"decode unreachable" in body + + assert r.breaker.state_of("d1").value == "open" + # The prefill leg is a separate pool: a wedged decode must not evict it. + assert r.breaker.state_of("p1").value == "closed" + await r.aclose() diff --git a/tests/unit/router/test_failover.py b/tests/unit/router/test_failover.py index 27ebe9d9..4596895a 100644 --- a/tests/unit/router/test_failover.py +++ b/tests/unit/router/test_failover.py @@ -208,3 +208,78 @@ def handler(request: httpx.Request) -> httpx.Response: assert resp.status_code == 200 assert json.loads(bytes(resp.body))["id"] == "ok-http" await r.aclose() + + +# --- circuit breaker: failure memory ACROSS requests (issue #82) --------------- + + +@pytest.mark.asyncio +async def test_breaker_stops_reselecting_a_dead_worker(): + """The regression behind issue #82, through the real MixedRouter. + + Failover already made all ten of these requests succeed, so a test that + only checked status codes passed both before and after the fix. What was + broken is the cost: ``tried`` is per-request, so a worker that is broken + for inference but healthy to discovery was re-picked on *every* request and + every one of them paid a wasted round trip. ``nats.streamed`` is the + assertion that matters. + """ + scripts = { + "w1": [(TYPE_ERROR, 502, b"wedged")], + "w2": [(TYPE_DATA, None, b"ok"), (TYPE_DONE, 200, b"")], + } + nats = _FakeNats(scripts) + r = _router([_w("w1"), _w("w2")], nats, retries=1) + for _ in range(10): + resp = await r.dispatch({"model": "m"}, stream=True) + assert await _drain_stream(resp) == b"ok", "failover must still serve every request" + + # _FakePolicy always picks candidates[0], so w1 is offered every request + # until the breaker (threshold 3) takes it out; its 5s cooldown does not + # elapse during the test. + assert nats.streamed.count("w1") == 3, ( + f"w1 dispatched {nats.streamed.count('w1')} times; expected 3 " + "(it was 10 before the breaker existed)" + ) + assert nats.streamed.count("w2") == 10 + await r.aclose() + + +@pytest.mark.asyncio +async def test_breaker_ignores_client_errors(): + """A 400 is the request's fault, and every worker would return it. Counting + it would circuit-break an entirely healthy fleet on one bad client.""" + nats = _FakeNats({"w1": [(TYPE_ERROR, 400, b"bad request")]}) + r = _router([_w("w1")], nats, retries=0) + for _ in range(10): + await r.dispatch({"model": "m"}, stream=True) + assert nats.streamed.count("w1") == 10, "4xx must not take a healthy worker out of rotation" + await r.aclose() + + +@pytest.mark.asyncio +async def test_breaker_recovers_after_cooldown(): + """A worker that comes back must be picked up again, not stay excluded.""" + scripts = { + "w1": [(TYPE_ERROR, 502, b"wedged")], + "w2": [(TYPE_DATA, None, b"ok"), (TYPE_DONE, 200, b"")], + } + nats = _FakeNats(scripts) + r = _router([_w("w1"), _w("w2")], nats, retries=1) + + clock = [1000.0] + r.breaker.now = lambda: clock[0] + + for _ in range(3): + await _drain_stream(await r.dispatch({"model": "m"}, stream=True)) + assert nats.streamed.count("w1") == 3 # tripped + + await _drain_stream(await r.dispatch({"model": "m"}, stream=True)) + assert nats.streamed.count("w1") == 3, "still open during the cooldown" + + scripts["w1"] = [(TYPE_DATA, None, b"back"), (TYPE_DONE, 200, b"")] + clock[0] += 5.1 + body = await _drain_stream(await r.dispatch({"model": "m"}, stream=True)) + assert body == b"back", "the half-open probe must reach the recovered worker" + assert r.breaker.state_of("w1").value == "closed" + await r.aclose() From 546663b9816d5d71f63f530ffec669b6f0b7d121 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Tue, 4 Aug 2026 03:37:04 +0000 Subject: [PATCH 02/88] feat(tools): GPU-free fake worker for testing everything above the engine Discovery, routing, failover, the circuit breaker, drain and any future autoscaling loop are all testable without a model. What blocked that was that there was no way to *be* a worker without loading one: the e2e harness starts real engines in containers, so the cheapest fleet anyone could assemble cost a GPU and a multi-minute weight load per member. `infera-fake-worker` registers through the real registration clients with a real EngineConfig, so the payload is built by build_worker_payload and parsed back by worker_info_from_json -- the same functions every engine and every discovery backend use. It cannot drift from the contract without failing its own tests. Only what happens after a request arrives is faked. Three flags exist specifically to make expensive problems cheap: * --startup-delay-s simulates the 5-15 minute weight load, holding /health at 503. That delay is the main reason naive autoscaling overshoots, since an unready replica still counts in the fleet while consuming 0% of the metric. * --max-concurrency gives requests somewhere to queue, which is what makes num_requests_waiting non-zero. Without a queue, the metric the whole industry scales on is identically zero and a scaling test would pass without testing anything. * --fail-first drives the circuit breaker through open -> half-open -> closed without killing a process. /metrics uses engine-native names selected by --engine, so a scaling rule written against fakes transfers to a real fleet unchanged. The SGLang names are second-hand and flagged as unverified in the README. Deliberately not simulated: KV transfer. A PD fake takes part in pool membership and routing and nothing else, which the README says plainly so nobody concludes Mooncake works because a fake fleet answered. Verified end to end, not just unit tested: three fakes registered into a real etcd, a real infera server discovered all three, and unary, SSE and round-robin distribution (4/4 across two workers) all behaved. Building that exposed a bug in an earlier draft -- register() writes the record but the caller must run heartbeat_loop(), so without it the 30s lease expired and workers silently vanished from the pool. Fixed, and written up in the README since it presents as a discovery bug and is not one. Signed-off-by: Zhang, Jiejing --- infera/tools/fakeworker/README.md | 75 +++++ infera/tools/fakeworker/__init__.py | 6 + infera/tools/fakeworker/__main__.py | 9 + infera/tools/fakeworker/server.py | 391 +++++++++++++++++++++++++++ pyproject.toml | 1 + tests/unit/tools/test_fake_worker.py | 221 +++++++++++++++ 6 files changed, 703 insertions(+) create mode 100644 infera/tools/fakeworker/README.md create mode 100644 infera/tools/fakeworker/__init__.py create mode 100644 infera/tools/fakeworker/__main__.py create mode 100644 infera/tools/fakeworker/server.py create mode 100644 tests/unit/tools/test_fake_worker.py diff --git a/infera/tools/fakeworker/README.md b/infera/tools/fakeworker/README.md new file mode 100644 index 00000000..fc78e306 --- /dev/null +++ b/infera/tools/fakeworker/README.md @@ -0,0 +1,75 @@ +# Fake worker + +A worker that joins the fleet and serves tokens, with no GPU and no weights. + +Everything above the engine — discovery, routing, failover, the circuit breaker, +drain, autoscaling — is testable without a model. What blocked that was simply +that there was no way to *be* a worker without loading one: `tests/e2e/harness` +starts real engines in containers, so the cheapest fleet anyone could build cost +a GPU and a multi-minute weight load per member. + +```bash +infera-fake-worker --model-name my-model --port 9101 \ + --discovery-backend etcd --etcd-endpoint http://127.0.0.1:2379 +``` + +## Why it can be trusted + +It registers through the **real** registration clients with a real +`EngineConfig`, so the payload is built by `build_worker_payload` — the same +function every engine uses — and parsed back by `worker_info_from_json`, the +same function every discovery backend uses. If that contract changes, this +changes with it or the test suite fails. Only what happens *after* a request +arrives is faked. + +## The knobs that matter + +Most flags are obvious. These three exist because they make otherwise expensive +problems reproducible on a laptop: + +| Flag | Why | +|---|---| +| `--startup-delay-s` | Simulates the 5–15 minute weight load. `/health` stays 503 until it elapses. This is the single biggest reason naive autoscaling overshoots — an unready replica still counts in the fleet but consumes 0% of the metric, so the loop keeps asking for more. Reproducing it here costs nothing; reproducing it on real hardware costs a GPU-hour per iteration. | +| `--max-concurrency` | Gives requests somewhere to queue. `num_requests_waiting` is the metric the entire industry autoscales on, and without a queue it is identically zero — so a scaling test against a fake fleet would pass without testing anything. | +| `--fail-first N` | Refuse the first N requests with 503, then recover. Drives the router's circuit breaker through open → half-open → closed without killing a process. | + +`/metrics` exposes engine-native names (`vllm:num_requests_waiting`, +`sglang:num_queue_reqs`, …) chosen by `--engine`, so a scaling rule written +against fakes transfers to a real fleet unchanged. + +> The vLLM names are from its published metrics documentation. **The SGLang +> names are second-hand from a research pass and have not been checked against a +> live SGLang.** Verify before depending on them. + +## Limits — read these before drawing conclusions + +**No KV transfer is simulated.** A `--disagg-mode prefill` / `decode` fake takes +part in pool membership and routing, and that is all. No KV moves between the +legs. Useful for testing that P and D pools exist, are discovered, and are +routed to independently; **useless for testing Mooncake, bootstrap handshakes, +or anything about the transfer itself.** Do not conclude that PD "works" because +a fake fleet answered. + +**`--kv` synthesizes a tokenizer canary from the model name.** All fakes for one +model agree, which is what they need to do — `CanaryVerifier` silently drops a +worker whose canary differs from the first-registered one, so disagreeing fakes +would produce a fleet that is half the expected size for no visible reason. But +that synthetic canary will never match a **real** worker's, so do not mix fakes +and real workers under one model name. Without `--kv` there is no canary at all +and the fakes join anything. + +**It answers instantly by construction.** `--ttft-ms` and `--itl-ms` are a +latency model, not a performance model: TTFT scales linearly with prompt length +and nothing contends for memory. Do not use it to predict real throughput. + +## Gotchas found while building this + +- **`--advertise-host` must be routable from the router.** It defaults to + `$POD_IP`. `0.0.0.0` registers a URL no peer can reach. +- **The server requires `--router-tokenizer-path` even for `round-robin`**, and + resolves it eagerly. Any existing directory satisfies it, which is enough to + bring a router up against fakes. +- **Registration alone is not enough** — `heartbeat_loop()` has to be running or + the etcd lease (30 s) expires and the worker vanishes from the pool about half + a minute after it appears. This looks exactly like a discovery bug and is not + one. The real worker entrypoint starts it too; so does this. diff --git a/infera/tools/fakeworker/__init__.py b/infera/tools/fakeworker/__init__.py new file mode 100644 index 00000000..7b323fac --- /dev/null +++ b/infera/tools/fakeworker/__init__.py @@ -0,0 +1,6 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""GPU-free worker for testing everything above the engine.""" diff --git a/infera/tools/fakeworker/__main__.py b/infera/tools/fakeworker/__main__.py new file mode 100644 index 00000000..e0171656 --- /dev/null +++ b/infera/tools/fakeworker/__main__.py @@ -0,0 +1,9 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +from infera.tools.fakeworker.server import main + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/infera/tools/fakeworker/server.py b/infera/tools/fakeworker/server.py new file mode 100644 index 00000000..bfee6998 --- /dev/null +++ b/infera/tools/fakeworker/server.py @@ -0,0 +1,391 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""A worker that registers into the fleet and serves tokens, without a GPU. + +Everything above the engine — discovery, routing, failover, the circuit breaker, +drain, and any autoscaling loop — is testable without weights. What blocked that +until now was simply that there was no way to *be* a worker without loading a +model: ``tests/e2e/harness`` starts real engines in containers, so the cheapest +fleet anyone could build cost a GPU and a multi-minute weight load per member. + +This is that missing piece. It registers through the **real** registration +clients with a real :class:`EngineConfig`, so the fake cannot drift from the +contract a genuine worker satisfies — if the payload changes, this changes with +it or fails loudly. What it fakes is only what happens after a request arrives. + +The parts worth having are the ones that make the *hard* problems reproducible: + +* ``--startup-delay-s`` simulates the 5-15 minute weight load. That delay is the + single biggest reason naive autoscaling overshoots — an unready replica counts + as consuming 0% of the metric, so the loop keeps asking for more. Reproducing + it costs nothing here and a GPU-hour on real hardware. +* ``--max-concurrency`` gives requests somewhere to queue, which is what makes + ``num_requests_waiting`` mean anything. Without a queue the metric everyone + scales on is identically zero. +* SIGTERM deregisters *before* draining, mirroring the real worker, so the + window where a terminating pod is still receiving traffic is observable. + +Not simulated, deliberately: real KV transfer. A PD fake accepts the bootstrap +fields and answers, but no KV moves between prefill and decode. That makes it +useful for testing pool membership, routing and scaling, and useless for testing +the transfer itself. Do not use it to conclude anything about Mooncake. +""" + +from __future__ import annotations + +import argparse +import asyncio +import hashlib +import json +import logging +import os +import signal +import time +from dataclasses import dataclass, field + +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse + +from infera.common.worker_pool import DisaggMode, EngineType, KvRegistrationMetadata +from infera.engine.base import EngineConfig + +logger = logging.getLogger("infera.fakeworker") + +# Roughly the shape of an English word, so token counts and byte counts stay in +# a believable ratio for anything measuring throughput. +_FILLER = "token " + + +@dataclass +class Behaviour: + """What this worker pretends about its own performance.""" + + ttft_ms: float = 50.0 + itl_ms: float = 10.0 + max_concurrency: int = 8 + #: Requests beyond max_concurrency wait here; this is what makes the + #: queue-depth metric — the one the whole industry scales on — non-zero. + max_kv_blocks: int = 1024 + blocks_per_request: int = 8 + fail_rate: float = 0.0 + #: Serve 5xx for the first N requests, then recover. For exercising the + #: circuit breaker's half-open probe without killing the process. + fail_first: int = 0 + + +@dataclass +class State: + running: int = 0 + waiting: int = 0 + served: int = 0 + failed: int = 0 + ready: bool = False + started_at: float = field(default_factory=time.monotonic) + _sem: asyncio.Semaphore | None = None + draining: bool = False + + +def deterministic_canary(model_name: str) -> list[int]: + """A stand-in for the real tokenizer canary. + + Real workers tokenize a fixed probe string and register the ids, so that a + fleet running mismatched tokenizers under one model name is rejected at + registration rather than at inference time. A fake has no tokenizer, but it + still has to agree with *other fakes* for the same model, or the second one + to register is silently dropped from the pool. + + Deriving it from the model name gives exactly that: all fakes for a model + agree, and different models differ. It will not match a real worker's canary + — see the README; do not mix fakes and real workers under one model name. + """ + h = hashlib.sha256(f"infera-fake-canary::{model_name}".encode()).digest() + return [int.from_bytes(h[i : i + 2], "big") for i in range(0, 16, 2)] + + +def build_app(cfg: EngineConfig, behaviour: Behaviour, state: State) -> FastAPI: + app = FastAPI(title="infera fake worker") + + def _engine_metric(name: str) -> str: + # Engine-native metric names, so a scaling rule written against a fake + # fleet transfers to a real one unchanged. vLLM's names are from its + # published metrics doc; SGLang's are second-hand (see README) and worth + # checking against a live engine before relying on them. + prefix = "vllm" if cfg.engine == EngineType.VLLM else "sglang" + sglang = { + "num_requests_waiting": "num_queue_reqs", + "num_requests_running": "num_running_reqs", + "gpu_cache_usage_perc": "token_usage", + } + return f"{prefix}:{sglang[name] if prefix == 'sglang' else name}" + + async def _admit() -> bool: + """Returns False if this request should be rejected outright.""" + if state.draining: + return False + if behaviour.fail_first and state.served + state.failed < behaviour.fail_first: + state.failed += 1 + return False + return True + + async def _generate(prompt_tokens: int, max_tokens: int): + """Occupy a concurrency slot for a believable amount of time.""" + assert state._sem is not None + state.waiting += 1 + async with state._sem: + state.waiting -= 1 + state.running += 1 + try: + # TTFT scales with prompt length the way a real prefill does, + # so prefill-heavy and decode-heavy load look different to + # anything measuring them. + await asyncio.sleep(behaviour.ttft_ms / 1000.0 * max(1.0, prompt_tokens / 512)) + for _ in range(max_tokens): + yield _FILLER + await asyncio.sleep(behaviour.itl_ms / 1000.0) + state.served += 1 + finally: + state.running -= 1 + + def _parse(body: dict) -> tuple[int, int]: + text = json.dumps(body.get("messages") or body.get("prompt") or "") + prompt_tokens = max(1, len(text) // 4) + return prompt_tokens, int(body.get("max_tokens") or 16) + + @app.post("/v1/chat/completions") + @app.post("/v1/completions") + async def completions(request: Request): + body = await request.json() + if not await _admit(): + return JSONResponse({"error": "fake worker refusing"}, status_code=503) + prompt_tokens, max_tokens = _parse(body) + + if body.get("stream"): + + async def sse(): + async for chunk in _generate(prompt_tokens, max_tokens): + payload = {"choices": [{"delta": {"content": chunk}}]} + yield f"data: {json.dumps(payload)}\n\n".encode() + yield b"data: [DONE]\n\n" + + return StreamingResponse(sse(), media_type="text/event-stream") + + out = "".join([c async for c in _generate(prompt_tokens, max_tokens)]) + return JSONResponse( + { + "id": f"fake-{state.served}", + "model": cfg.model_name, + "choices": [{"message": {"role": "assistant", "content": out}}], + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": max_tokens, + "total_tokens": prompt_tokens + max_tokens, + }, + } + ) + + @app.get("/v1/models") + async def models(): + return {"object": "list", "data": [{"id": cfg.model_name, "object": "model"}]} + + @app.get("/health") + async def health(): + # Unready until the simulated weight load finishes. This is the whole + # point of --startup-delay-s: a replica that exists but cannot serve is + # exactly what makes autoscalers overshoot. + if not state.ready: + return JSONResponse({"status": "loading"}, status_code=503) + return {"status": "ok", "running": state.running, "waiting": state.waiting} + + @app.get("/metrics") + async def metrics(): + used = min(1.0, (state.running * behaviour.blocks_per_request) / behaviour.max_kv_blocks) + lines = [ + f"{_engine_metric('num_requests_running')} {state.running}", + f"{_engine_metric('num_requests_waiting')} {state.waiting}", + f"{_engine_metric('gpu_cache_usage_perc')} {used:.4f}", + f"infera_fake_worker_served_total {state.served}", + f"infera_fake_worker_refused_total {state.failed}", + f"infera_fake_worker_ready {1 if state.ready else 0}", + f"infera_fake_worker_draining {1 if state.draining else 0}", + ] + return PlainTextResponse("\n".join(lines) + "\n") + + return app + + +def build_config(args) -> EngineConfig: + kv = None + if args.kv: + canary = deterministic_canary(args.model_name) + kv = KvRegistrationMetadata( + engine_block_size=args.kv_block_size, + index_block_size=args.kv_block_size, + tokenizer=args.model_name, + tokenizer_digest=hashlib.sha256(args.model_name.encode()).hexdigest(), + tokenizer_canary=canary, + supports_events=False, # no ZMQ publisher; router falls back to load-only + ) + return EngineConfig( + model_name=args.model_name, + host=args.advertise_host or args.host, + port=args.port, + engine=EngineType(args.engine), + disagg_mode=DisaggMode(args.disagg_mode), + disagg_meta={"protocol": args.pd_protocol} if args.disagg_mode != "mixed" else {}, + kv=kv, + kv_block_size=args.kv_block_size if args.kv else None, + dp_rank=args.dp_rank, + dp_size=args.dp_size, + request_transport="http", + ) + + +def parse_args(argv=None): + p = argparse.ArgumentParser( + prog="infera-fake-worker", + description="Register into an Infera fleet and serve tokens without a GPU.", + ) + p.add_argument("--model-name", required=True, help="must match what the router routes for") + p.add_argument("--host", default="0.0.0.0") + p.add_argument( + "--advertise-host", + default=os.environ.get("POD_IP"), + help="address peers use to reach this worker; defaults to $POD_IP. " + "0.0.0.0 is never routable from another pod.", + ) + p.add_argument("--port", type=int, default=8080) + p.add_argument("--engine", default="sglang", choices=[e.value for e in EngineType]) + p.add_argument("--disagg-mode", default="mixed", choices=[m.value for m in DisaggMode]) + p.add_argument("--pd-protocol", default="sglang_bootstrap") + p.add_argument("--dp-rank", type=int, default=None) + p.add_argument("--dp-size", type=int, default=None) + + p.add_argument( + "--discovery-backend", + default=os.environ.get("INFERA_DISCOVERY_BACKEND", "kubernetes"), + choices=["kubernetes", "etcd"], + ) + p.add_argument("--etcd-endpoint", default=os.environ.get("INFERA_ETCD_ENDPOINT")) + p.add_argument("--etcd-prefix", default="/infera/workers/") + + p.add_argument( + "--kv", + action="store_true", + help="register a KV metadata block with a synthetic tokenizer canary. " + "All fakes for a model agree; a fake and a REAL worker will not -- the " + "second to register is silently dropped. See the README.", + ) + p.add_argument("--kv-block-size", type=int, default=64) + + p.add_argument("--ttft-ms", type=float, default=50.0) + p.add_argument("--itl-ms", type=float, default=10.0) + p.add_argument("--max-concurrency", type=int, default=8) + p.add_argument("--max-kv-blocks", type=int, default=1024) + p.add_argument( + "--startup-delay-s", + type=float, + default=0.0, + help="seconds of simulated weight loading before /health goes green. " + "Set this to your real cold start to reproduce autoscaler overshoot.", + ) + p.add_argument( + "--fail-first", + type=int, + default=0, + help="refuse the first N requests with 503, then recover -- exercises " + "the router's circuit breaker and its half-open probe.", + ) + p.add_argument("--drain-timeout", type=float, default=30.0) + return p.parse_args(argv) + + +async def _serve(args) -> None: + import uvicorn + + cfg = build_config(args) + behaviour = Behaviour( + ttft_ms=args.ttft_ms, + itl_ms=args.itl_ms, + max_concurrency=args.max_concurrency, + max_kv_blocks=args.max_kv_blocks, + fail_first=args.fail_first, + ) + state = State() + state._sem = asyncio.Semaphore(args.max_concurrency) + app = build_app(cfg, behaviour, state) + + server = uvicorn.Server( + uvicorn.Config(app, host=args.host, port=args.port, log_level="warning") + ) + serve_task = asyncio.create_task(server.serve()) + + if args.startup_delay_s > 0: + logger.info("simulating weight load for %.0fs", args.startup_delay_s) + await asyncio.sleep(args.startup_delay_s) + state.ready = True + + # Register only once ready, exactly like a real worker: a worker that is in + # the pool but cannot serve is a routing black hole. + if args.discovery_backend == "etcd": + from infera.common.registration import RegistrationClient + + reg = RegistrationClient(args.etcd_endpoint, prefix=args.etcd_prefix) + else: + from infera.common.registration_k8s import K8sRegistrationClient + + reg = K8sRegistrationClient() + worker_id = await reg.register(cfg) + # register() only writes the record; keeping it alive is the caller's job in + # both backends, exactly as in the real worker entrypoint. Without this the + # etcd lease (30s) expires and the worker silently vanishes from the pool + # about half a minute after it appears -- which looks like a discovery bug + # and is not one. + hb_task = asyncio.create_task(reg.heartbeat_loop(), name="fake-worker-heartbeat") + logger.info( + "registered %s model=%s mode=%s via %s", + worker_id, + cfg.model_name, + cfg.disagg_mode.value, + args.discovery_backend, + ) + + stop = asyncio.Event() + + async def _shutdown() -> None: + # Deregister BEFORE draining, matching the real worker: the router must + # stop sending new work before we stop accepting it, or the gap shows up + # to clients as failures rather than as a clean drain. + state.draining = True + hb_task.cancel() + try: + await reg.deregister() + except Exception as exc: # noqa: BLE001 - shutdown must not raise + logger.warning("deregister failed: %s", exc) + deadline = time.monotonic() + args.drain_timeout + while state.running and time.monotonic() < deadline: + await asyncio.sleep(0.1) + if state.running: + logger.warning("drain timeout with %d request(s) still in flight", state.running) + server.should_exit = True + stop.set() + + loop = asyncio.get_running_loop() + for sig in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(sig, lambda: asyncio.create_task(_shutdown())) + + await stop.wait() + await serve_task + + +def main(argv=None) -> int: + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + args = parse_args(argv) + asyncio.run(_serve(args)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/pyproject.toml b/pyproject.toml index e71ed3ed..75c1868b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,7 @@ infera-kvd-probe = "infera.kvd.bench.probe:main" infera-kvd-l3-bench = "infera.kvd.bench.l3_bench:main" # node + PD preflight suite (gpu / network / storage / firmware / host probes) infera-preflight = "infera.tools.preflight.cli:main" +infera-fake-worker = "infera.tools.fakeworker.server:main" [build-system] requires = ["setuptools>=69", "setuptools_scm[toml]>=8", "wheel"] diff --git a/tests/unit/tools/test_fake_worker.py b/tests/unit/tools/test_fake_worker.py new file mode 100644 index 00000000..e1df1d8e --- /dev/null +++ b/tests/unit/tools/test_fake_worker.py @@ -0,0 +1,221 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""The fake worker has to be trustworthy, or every test built on it is too. + +Two properties matter more than the rest. It must register through the *real* +contract, so a fleet of fakes exercises the same discovery path a real fleet +does -- that is checked by building an actual ``EngineConfig`` and running it +through ``build_worker_payload``, the same function every engine uses. And its +queue must be real, because ``num_requests_waiting`` is the metric the entire +industry autoscales on, and a fake that always reports zero would make every +scaling test vacuously pass. +""" + +from __future__ import annotations + +import asyncio + +import httpx +import pytest + +from infera.common.discovery import worker_info_from_json +from infera.common.registration import build_worker_payload +from infera.common.worker_pool import DisaggMode, EngineType +from infera.tools.fakeworker.server import ( + Behaviour, + State, + build_app, + build_config, + deterministic_canary, + parse_args, +) + + +def _args(*extra): + return parse_args(["--model-name", "m", *extra]) + + +def _stack(**behaviour_kw): + args = _args() + cfg = build_config(args) + state = State(ready=True) + state._sem = asyncio.Semaphore(behaviour_kw.get("max_concurrency", 8)) + b = Behaviour(**behaviour_kw) + return cfg, b, state, build_app(cfg, b, state) + + +def _client(app): + return httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://fake") + + +# --- the registration contract ------------------------------------------------ + + +def test_registers_through_the_real_payload_builder(): + """If this breaks, the fake has drifted from what a real worker registers -- + which is the one failure that would silently invalidate everything else.""" + cfg = build_config(_args("--engine", "vllm", "--disagg-mode", "prefill")) + payload = build_worker_payload(cfg) + assert payload["model_name"] == "m" + assert payload["engine"] == EngineType.VLLM + assert payload["disagg_mode"] == DisaggMode.PREFILL + # And discovery must be able to parse it back: worker_info_from_json is the + # single function every backend (etcd, kubernetes) uses on the wire record, + # so a round-trip through it is the real contract, not an approximation. + info = worker_info_from_json(payload) + assert info.worker_id == payload["worker_id"] + assert info.disagg_mode is DisaggMode.PREFILL + + +def test_kv_block_is_absent_unless_asked_for(): + """Without --kv there is no canary, so fakes can join any fleet. With it, + canary verification applies and mixing with real workers breaks.""" + assert build_config(_args()).kv is None + assert build_config(_args("--kv")).kv is not None + + +def test_fakes_for_one_model_agree_on_the_canary(): + """Disagreeing fakes would be silently dropped from the pool by + CanaryVerifier -- a fleet that looks half its intended size for no visible + reason.""" + assert deterministic_canary("llama") == deterministic_canary("llama") + assert deterministic_canary("llama") != deterministic_canary("qwen") + + +# --- the serving surface ------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_unary_completion_shape(): + _, _, _, app = _stack(ttft_ms=0, itl_ms=0) + async with _client(app) as c: + r = await c.post("/v1/chat/completions", json={"model": "m", "max_tokens": 5}) + assert r.status_code == 200 + body = r.json() + assert body["choices"][0]["message"]["content"] + assert body["usage"]["completion_tokens"] == 5 + + +@pytest.mark.asyncio +async def test_streaming_emits_sse_and_terminates(): + """The router's failover and circuit breaker are both first-byte-sensitive, + so a fake that cannot stream cannot exercise either.""" + _, _, _, app = _stack(ttft_ms=0, itl_ms=0) + async with _client(app) as c: + r = await c.post( + "/v1/chat/completions", json={"model": "m", "max_tokens": 3, "stream": True} + ) + body = (await r.aread()).decode() + assert r.headers["content-type"].startswith("text/event-stream") + assert body.count("data: ") == 4 # 3 chunks + [DONE] + assert body.endswith("data: [DONE]\n\n") + + +@pytest.mark.asyncio +async def test_health_is_503_until_ready(): + """--startup-delay-s exists to reproduce autoscaler overshoot, which only + happens because an unready replica still counts in the fleet.""" + cfg, b, state, app = _stack() + state.ready = False + async with _client(app) as c: + assert (await c.get("/health")).status_code == 503 + state.ready = True + assert (await c.get("/health")).status_code == 200 + + +# --- the metrics an autoscaler would read ------------------------------------- + + +@pytest.mark.asyncio +async def test_queue_depth_is_real(): + """The whole point. Drive more concurrent requests than the worker admits + and the waiting count must actually rise -- a fake that always reports 0 + would make every scaling test pass without testing anything.""" + _, _, state, app = _stack(max_concurrency=2, ttft_ms=200, itl_ms=0) + async with _client(app) as c: + tasks = [ + asyncio.create_task(c.post("/v1/completions", json={"model": "m", "max_tokens": 1})) + for _ in range(6) + ] + await asyncio.sleep(0.05) + peak_waiting = state.waiting + peak_running = state.running + await asyncio.gather(*tasks) + + assert peak_running == 2, f"admitted {peak_running}, expected the concurrency cap" + assert peak_waiting == 4, f"queued {peak_waiting}, expected the other 4" + assert state.waiting == 0 and state.running == 0, "must settle back to idle" + + +@pytest.mark.asyncio +async def test_metrics_use_engine_native_names(): + """A scaling rule written against fakes should transfer to a real fleet + unchanged, so the names have to be the engine's, not ours.""" + for engine, expected in ( + ("vllm", "vllm:num_requests_waiting"), + ("sglang", "sglang:num_queue_reqs"), + ): + args = _args("--engine", engine) + cfg = build_config(args) + state = State(ready=True) + state._sem = asyncio.Semaphore(1) + app = build_app(cfg, Behaviour(), state) + async with _client(app) as c: + text = (await c.get("/metrics")).text + assert expected in text, f"{engine}: missing {expected}\n{text}" + + +@pytest.mark.asyncio +async def test_kv_usage_tracks_inflight(): + _, _, state, app = _stack(max_concurrency=4, ttft_ms=200, itl_ms=0) + async with _client(app) as c: + idle = (await c.get("/metrics")).text + tasks = [ + asyncio.create_task(c.post("/v1/completions", json={"model": "m", "max_tokens": 1})) + for _ in range(4) + ] + await asyncio.sleep(0.05) + busy = (await c.get("/metrics")).text + await asyncio.gather(*tasks) + + def usage(t): + return float( + [ln for ln in t.splitlines() if "cache_usage" in ln or "token_usage" in ln][0].split()[ + -1 + ] + ) + + assert usage(idle) == 0.0 + assert usage(busy) > 0.0 + + +# --- failure injection, for the circuit breaker ------------------------------- + + +@pytest.mark.asyncio +async def test_fail_first_then_recovers(): + """Mirrors the breaker's half-open probe: a worker that is broken, stays + broken for a while, then comes back.""" + _, _, state, app = _stack(ttft_ms=0, itl_ms=0, fail_first=3) + async with _client(app) as c: + codes = [ + (await c.post("/v1/completions", json={"model": "m", "max_tokens": 1})).status_code + for _ in range(5) + ] + assert codes[:3] == [503, 503, 503] + assert codes[3:] == [200, 200] + + +@pytest.mark.asyncio +async def test_draining_refuses_new_work(): + """SIGTERM deregisters before draining; until the router notices, arriving + requests must be refused rather than accepted and then cut.""" + _, _, state, app = _stack(ttft_ms=0, itl_ms=0) + state.draining = True + async with _client(app) as c: + r = await c.post("/v1/completions", json={"model": "m", "max_tokens": 1}) + assert r.status_code == 503 + assert "infera_fake_worker_draining 1" in (await c.get("/metrics")).text From e6152b7d36c09bd8d72b465ea5e96a886bb259b6 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Tue, 4 Aug 2026 03:49:01 +0000 Subject: [PATCH 03/88] feat(tools): fake worker supports PD and DP-attention routing Follows the GPU-free fake worker. Pool membership was already there; what was missing was everything that makes PD and DP *routing* assertable. The fake now advertises the same disagg_meta a real worker does -- {"protocol": ..., "params": {"bootstrap_addr": ...}}, prefill only -- with --pd-protocol constrained to the router's own protocol registry, so a typo dies at argparse instead of surfacing as a protocol error on the first request. The earlier draft had both wrong (underscore instead of hyphen, no params at all) and PD simply did not dispatch. GET /debug/routing reports what the *router* decided: per-rank request counts keyed on the X-Data-Parallel-Rank header it sent, and the handoff fields it injected. None of that is observable with a real engine -- a malformed handoff does not raise, it hangs on KVPoll until a ~300s timeout, and the failure surfaces nowhere near the router that caused it. Verified end to end against a real router and real etcd, no GPU: * PD dual-dispatch: both legs receive the same bootstrap_room, with bootstrap_host/port pointing at the prefill worker's advertised endpoint. * DP attention, rank-multiplexed: 8 requests over a dp_size=4 prefill fanned out 2/2/2/2 across ranks via the header, decode received disagg_prefill_dp_rank, and room % dp_size == dp_rank held on 12 of 12 requests -- the invariant SGLang's follow_bootstrap_room balancer enforces with a KVTransferError. Two things found by running it. A bind failure still registered the worker: uvicorn logs "address already in use" and gives up, but the process keeps running, so a port collision produced a worker that was in the pool and served nothing. Registration now waits for the socket and exits 3 on collision. And DP_RANK_HEADER is extracted from dp_routing so the header name has one definition rather than a literal the fake could silently stop matching after a rename; Rust already had this. Also documents the DP shape distinction that makes the path look dead: is_rank_multiplexed() is dp_size > 1 AND dp_rank is None, so a worker that registers its own rank is an endpoint and correctly opts out of router-side DP routing. Signed-off-by: Zhang, Jiejing --- infera/router/dp_routing.py | 8 ++- infera/tools/fakeworker/README.md | 48 +++++++++++++++ infera/tools/fakeworker/server.py | 97 ++++++++++++++++++++++++++++++- 3 files changed, 150 insertions(+), 3 deletions(-) diff --git a/infera/router/dp_routing.py b/infera/router/dp_routing.py index ffd905f2..6ee0397b 100644 --- a/infera/router/dp_routing.py +++ b/infera/router/dp_routing.py @@ -15,6 +15,12 @@ from infera.common.worker_pool import EngineType from infera.router.policy.target import RouteTarget +#: Both engines honour this (SGLang ``DataParallelController``, vLLM +#: ``_get_data_parallel_rank``, case-insensitively). Named rather than inlined +#: so anything asserting on it -- tests, the fake worker -- breaks loudly on a +#: rename instead of silently never matching. Mirrors Rust's DP_RANK_HEADER. +DP_RANK_HEADER = "X-Data-Parallel-Rank" + def dp_rank_header(target: RouteTarget) -> dict[str, str] | None: """``X-Data-Parallel-Rank`` header pinning the request to a DP rank. Both @@ -22,7 +28,7 @@ def dp_rank_header(target: RouteTarget) -> dict[str, str] | None: case-insensitive) honour it; no-op when the target carries no rank.""" if target.dp_rank is None: return None - return {"X-Data-Parallel-Rank": str(target.dp_rank)} + return {DP_RANK_HEADER: str(target.dp_rank)} def inject_disagg_prefill_dp_rank( diff --git a/infera/tools/fakeworker/README.md b/infera/tools/fakeworker/README.md index fc78e306..85d9d1b0 100644 --- a/infera/tools/fakeworker/README.md +++ b/infera/tools/fakeworker/README.md @@ -41,6 +41,48 @@ against fakes transfers to a real fleet unchanged. > names are second-hand from a research pass and have not been checked against a > live SGLang.** Verify before depending on them. +## PD and DP attention + +Both work, and both are verified end to end against a real router with no GPU. + +**PD.** `--disagg-mode prefill|decode` advertises the same `disagg_meta` a real +worker does — `{"protocol": ..., "params": {"bootstrap_addr": ...}}`, prefill +only — with `--pd-protocol` constrained to the router's own protocol registry so +a typo dies at argparse rather than at the first request. + +**DP attention.** Two deployment shapes exist and they behave differently, which +is worth knowing before concluding anything is broken: + +| Shape | Registration | What the router does | +|---|---|---| +| per-rank endpoints | `--dp-rank R --dp-size N` | Nothing. The address already selects the rank, so no header is pinned and no room alignment happens. `dp_rank` on the `RouteTarget` stays `None`. | +| rank-multiplexed | `--dp-size N`, **no** `--dp-rank` | `expand_targets` fans one worker into N targets; the router pins `X-Data-Parallel-Rank`, aligns `bootstrap_room % dp_size == dp_rank`, and injects `disagg_prefill_dp_rank` for the decode leg. | + +`is_rank_multiplexed()` is `dp_size > 1 and dp_rank is None` — so registering a +rank makes the worker an endpoint and opts *out* of router-side DP routing. That +is correct, not a bug, but the first time you see it the DP path looks dead. + +### `GET /debug/routing` + +The reason PD and DP are testable at all. It reports what the *router* decided — +per-rank request counts keyed on the header it sent, and the handoff fields it +injected on the last request: + +```json +{"dp_rank": null, "dp_size": 4, + "requests_by_dp_rank": {"0": 2, "1": 2, "2": 2, "3": 2}, + "last_handoff": {"bootstrap_host": "127.0.0.1", "bootstrap_port": 18430, + "bootstrap_room": 7442254485466660987, + "disagg_prefill_dp_rank": 3}} +``` + +None of that is observable with a real engine: a malformed handoff does not +raise, it hangs on KVPoll until a ~300 s timeout, and the failure surfaces +nowhere near the router that caused it. Here you can assert on it directly — +e.g. that `bootstrap_room % dp_size == disagg_prefill_dp_rank` holds on every +request, which is the invariant SGLang's `follow_bootstrap_room` balancer +enforces with a `KVTransferError`. + ## Limits — read these before drawing conclusions **No KV transfer is simulated.** A `--disagg-mode prefill` / `decode` fake takes @@ -69,6 +111,12 @@ and nothing contends for memory. Do not use it to predict real throughput. - **The server requires `--router-tokenizer-path` even for `round-robin`**, and resolves it eagerly. Any existing directory satisfies it, which is enough to bring a router up against fakes. +- **A bind failure used to still register.** uvicorn logs `address already in + use` and gives up, but the process keeps running — so a port collision + produced a worker that was in the pool and served nothing. Now the socket must + be bound before registration, and a collision exits 3. Worth remembering + because the symptom was a router error about the *worker* returning garbage, + which pointed nowhere near the real cause. - **Registration alone is not enough** — `heartbeat_loop()` has to be running or the etcd lease (30 s) expires and the worker vanishes from the pool about half a minute after it appears. This looks exactly like a discovery bug and is not diff --git a/infera/tools/fakeworker/server.py b/infera/tools/fakeworker/server.py index bfee6998..36c68e21 100644 --- a/infera/tools/fakeworker/server.py +++ b/infera/tools/fakeworker/server.py @@ -51,6 +51,8 @@ from infera.common.worker_pool import DisaggMode, EngineType, KvRegistrationMetadata from infera.engine.base import EngineConfig +from infera.router.disagg_protocols import _PROTOCOLS +from infera.router.dp_routing import DP_RANK_HEADER logger = logging.getLogger("infera.fakeworker") @@ -87,6 +89,16 @@ class State: _sem: asyncio.Semaphore | None = None draining: bool = False + #: Per-DP-rank request counts, keyed by the X-Data-Parallel-Rank header the + #: router sent. With a real engine you cannot easily see which rank the + #: router *intended* -- the engine just serves. Counting them here is what + #: makes DP-attention routing assertable at all. + by_dp_rank: dict[str, int] = field(default_factory=dict) + #: The PD handoff fields the router injected on the last request. Whether + #: the router shaped the body correctly is invisible from outside; a real + #: engine either works or hangs on KVPoll with no explanation. + last_handoff: dict = field(default_factory=dict) + def deterministic_canary(model_name: str) -> list[int]: """A stand-in for the real tokenizer canary. @@ -158,6 +170,25 @@ def _parse(body: dict) -> tuple[int, int]: @app.post("/v1/completions") async def completions(request: Request): body = await request.json() + + # Record what the router decided *before* deciding whether to serve, so + # a refused request still shows up in the routing evidence. + rank = request.headers.get(DP_RANK_HEADER) or "-" + state.by_dp_rank[rank] = state.by_dp_rank.get(rank, 0) + 1 + handoff = { + k: body[k] + for k in ( + "bootstrap_host", + "bootstrap_port", + "bootstrap_room", + "disagg_prefill_dp_rank", + "kv_transfer_params", + ) + if k in body + } + if handoff: + state.last_handoff = handoff + if not await _admit(): return JSONResponse({"error": "fake worker refusing"}, status_code=503) prompt_tokens, max_tokens = _parse(body) @@ -211,11 +242,47 @@ async def metrics(): f"infera_fake_worker_ready {1 if state.ready else 0}", f"infera_fake_worker_draining {1 if state.draining else 0}", ] + for rank, n in sorted(state.by_dp_rank.items()): + lines.append(f'infera_fake_worker_requests_by_dp_rank{{dp_rank="{rank}"}} {n}') return PlainTextResponse("\n".join(lines) + "\n") + @app.get("/debug/routing") + async def routing(): + """What the router actually decided, which is otherwise unobservable. + + A real engine given a malformed PD handoff does not complain -- it hangs + on KVPoll until a ~300s timeout, and the failure surfaces nowhere near + the router that caused it. This turns that into an assertion. + """ + return { + "worker_id": f"{cfg.host}:{cfg.port}", + "disagg_mode": cfg.disagg_mode.value, + "dp_rank": cfg.dp_rank, + "dp_size": cfg.dp_size, + "requests_by_dp_rank": state.by_dp_rank, + "last_handoff": state.last_handoff, + } + return app +def _disagg_meta(args) -> dict: + """Mirror what a real worker advertises. + + Only PREFILL carries a bootstrap endpoint; DECODE tags the protocol so the + router can fail fast on a cross-protocol pairing, and has nothing else to + say. Getting this wrong does not fail loudly at registration -- it fails at + the first PD request, as a protocol error that reads like a router bug. + """ + if args.disagg_mode == "mixed": + return {} + params: dict = {} + if args.disagg_mode == "prefill": + host = args.advertise_host or args.host + params["bootstrap_addr"] = f"{host}:{args.bootstrap_port}" + return {"protocol": args.pd_protocol, "params": params} + + def build_config(args) -> EngineConfig: kv = None if args.kv: @@ -234,7 +301,7 @@ def build_config(args) -> EngineConfig: port=args.port, engine=EngineType(args.engine), disagg_mode=DisaggMode(args.disagg_mode), - disagg_meta={"protocol": args.pd_protocol} if args.disagg_mode != "mixed" else {}, + disagg_meta=_disagg_meta(args), kv=kv, kv_block_size=args.kv_block_size if args.kv else None, dp_rank=args.dp_rank, @@ -259,7 +326,20 @@ def parse_args(argv=None): p.add_argument("--port", type=int, default=8080) p.add_argument("--engine", default="sglang", choices=[e.value for e in EngineType]) p.add_argument("--disagg-mode", default="mixed", choices=[m.value for m in DisaggMode]) - p.add_argument("--pd-protocol", default="sglang_bootstrap") + p.add_argument( + "--pd-protocol", + default="sglang-bootstrap", + choices=sorted(_PROTOCOLS), + help="must match the router's registry; a decode worker advertising a " + "different one is rejected as a protocol mismatch", + ) + p.add_argument( + "--bootstrap-port", + type=int, + default=8998, + help="advertised in disagg_meta by a prefill worker. Nothing listens on " + "it -- no KV is transferred (see README).", + ) p.add_argument("--dp-rank", type=int, default=None) p.add_argument("--dp-size", type=int, default=None) @@ -322,6 +402,19 @@ async def _serve(args) -> None: ) serve_task = asyncio.create_task(server.serve()) + # Do not register until the socket is actually bound. uvicorn logs a bind + # failure and gives up, but the process keeps running -- so without this + # check a port collision produces a worker that is in the pool and serves + # nothing. That is a routing black hole, and it is exactly the failure this + # tool exists to help find rather than create. + for _ in range(100): + if server.started or serve_task.done(): + break + await asyncio.sleep(0.05) + if not server.started: + serve_task.cancel() + raise SystemExit(f"failed to bind {args.host}:{args.port} -- not registering") + if args.startup_delay_s > 0: logger.info("simulating weight load for %.0fs", args.startup_delay_s) await asyncio.sleep(args.startup_delay_s) From 8af53f6a80ec0cef5881f708783237258f48e211 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Tue, 4 Aug 2026 05:44:34 +0000 Subject: [PATCH 04/88] =?UTF-8?q?feat:=20graceful=20scale-down=20=E2=80=94?= =?UTF-8?q?=20drain=20in-flight=20generations=20before=20stopping?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Scale-down cut live requests. Three separate gaps, each independently sufficient to lose a generation mid-stream. **A condemned Pod stayed a routing candidate.** discovery_k8s removed a worker on DELETE, on a cleared annotation, or on phase != Running -- never on deletionTimestamp. A terminating Pod keeps phase Running until its containers exit, and the operator injects a preStop sleep before SIGTERM, so for that entire delay the worker was condemned, healthy in the router's eyes, and still being assigned new work. The sleep meant to make shutdown graceful was instead buying more requests guaranteed to be cut. Reading deletionTimestamp turns that window into drain time. **HTTP transport never drained.** On NATS infera owns the request path and stop(drain=True) waits for in-flight work. On HTTP the router talks straight to the engine, so infera never sees the request and cannot count it -- shutdown went deregister() then engine.stop(), severing everything active. It now asks the engine, which does know: poll its /metrics for running+queued until zero or --drain-timeout. Queued requests count, because a request the engine accepted but has not started is work a client is waiting on. An unreadable metric returns immediately with a loud warning rather than waiting out the timeout: a rolling update that stalls on a parse failure is worse than one that cuts a request, and a silent full-timeout wait is indistinguishable from a genuinely busy worker. Engine metric names live in common/engine_metrics so drain, the fake worker and any future autoscaler share one mapping; ATOM is deliberately absent rather than guessed, since a wrong name reads as an idle engine and an idle engine is what makes a drain cut requests. **WorkerStatus.DRAINING was never set.** It has been in the enum and filtered from list_active since the beginning with nothing writing it. Both registration clients gain announce_draining(), and shutdown is reordered to announce -> drain -> deregister. Note this adds little on the Kubernetes path now that deletionTimestamp is honoured; the value is on etcd, where nothing else observes that the process is going away, and in observability -- a worker that vanishes looks like one that crashed, while /v1/workers showing "active, draining" says a rollout is progressing. The status field is omitted entirely when ACTIVE so the record stays identical to what older workers wrote. Verified end to end against a real router and real etcd, no GPU, using the fake worker: SIGTERM mid-generation, 30 requests through the drain window, 0 failures, 0 reached the draining worker, /v1/workers reported one active and one draining throughout, and the in-flight 250-chunk stream completed with [DONE]. Router removal latency measured at 15 ms. Signed-off-by: Zhang, Jiejing --- infera/common/discovery_k8s.py | 26 ++- infera/common/engine_metrics.py | 89 ++++++++++ infera/common/registration.py | 40 ++++- infera/common/registration_k8s.py | 26 +++ infera/engine/drain.py | 113 ++++++++++++ infera/engine/sglang/__main__.py | 22 ++- infera/engine/vllm/__main__.py | 22 ++- infera/tools/fakeworker/server.py | 17 +- .../common/test_discovery_k8s_terminating.py | 114 ++++++++++++ tests/unit/common/test_draining_status.py | 121 +++++++++++++ tests/unit/engine/test_drain.py | 162 ++++++++++++++++++ 11 files changed, 741 insertions(+), 11 deletions(-) create mode 100644 infera/common/engine_metrics.py create mode 100644 infera/engine/drain.py create mode 100644 tests/unit/common/test_discovery_k8s_terminating.py create mode 100644 tests/unit/common/test_draining_status.py create mode 100644 tests/unit/engine/test_drain.py diff --git a/infera/common/discovery_k8s.py b/infera/common/discovery_k8s.py index 5c3bfe8d..31613802 100644 --- a/infera/common/discovery_k8s.py +++ b/infera/common/discovery_k8s.py @@ -211,6 +211,26 @@ def _pod_running(pod: dict) -> bool: phase = ((pod.get("status") or {}).get("phase")) or "" return phase == "Running" + @staticmethod + def _pod_terminating(pod: dict) -> bool: + """True once the API server has stamped the Pod for deletion. + + A terminating Pod keeps ``phase: Running`` until its containers exit, so + a liveness check alone cannot see it. That gap is not academic: the + operator injects a ``preStop sleep`` before SIGTERM, and for the whole + of that delay the Pod is condemned, still Running, and — without this + check — still a routing candidate. The router would keep assigning new + work right up to the moment the process is killed, so the sleep that + exists to make shutdown graceful instead buys more requests that are + guaranteed to be cut. + + Reading ``deletionTimestamp`` closes that: the worker leaves the pool + the instant deletion is requested, in-flight work finishes on the + connections it already has, and the preStop delay becomes what it was + meant to be — drain time. + """ + return bool((pod.get("metadata") or {}).get("deletionTimestamp")) + def _handle_pod(self, pod: dict, *, deleted: bool) -> None: meta = pod.get("metadata") or {} pod_name = meta.get("name") or "" @@ -219,8 +239,10 @@ def _handle_pod(self, pod: dict, *, deleted: bool) -> None: annotations = meta.get("annotations") or {} raw = annotations.get(WORKER_INFO_ANNOTATION) - # Removal: explicit DELETE, pod no longer Running, or annotation gone. - if deleted or raw is None or not self._pod_running(pod): + # Removal: explicit DELETE, deletion requested, pod no longer Running, + # or annotation gone. Terminating is checked separately from Running + # because a condemned Pod stays Running until its containers exit. + if deleted or raw is None or self._pod_terminating(pod) or not self._pod_running(pod): worker_id = self._pod_to_worker.pop(pod_name, None) if worker_id is not None: self._remove(worker_id) diff --git a/infera/common/engine_metrics.py b/infera/common/engine_metrics.py new file mode 100644 index 00000000..53aca1d8 --- /dev/null +++ b/infera/common/engine_metrics.py @@ -0,0 +1,89 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""What each engine calls the metrics we need, in one place. + +Every engine exposes the same three facts — requests running, requests queued, +KV cache in use — under a different name, and the names drift between releases. +Anything that reads them (graceful drain, an autoscaler, the fake worker) needs +the same mapping, so it lives here rather than being spelled out at each call +site where one of them would quietly rot. + +Provenance, because it is uneven and matters: + +* **vLLM** — from its published metrics documentation. +* **SGLang** — second-hand, not verified against a running engine. Treat a + lookup failure as "unknown", never as "zero"; the difference decides whether a + drain waits or gives up. +* **ATOM** — unknown. Deliberately absent rather than guessed: a wrong name + reads as an idle engine, and an idle engine is exactly the answer that makes a + drain cut live requests. +""" + +from __future__ import annotations + +import re + +from infera.common.worker_pool import EngineType + +#: metric key -> per-engine exposition name. A missing engine means "we do not +#: know", which callers must distinguish from "the value is zero". +_NAMES: dict[str, dict[EngineType, str]] = { + "requests_running": { + EngineType.VLLM: "vllm:num_requests_running", + EngineType.SGLANG: "sglang:num_running_reqs", + }, + "requests_waiting": { + EngineType.VLLM: "vllm:num_requests_waiting", + EngineType.SGLANG: "sglang:num_queue_reqs", + }, + "kv_cache_usage": { + EngineType.VLLM: "vllm:gpu_cache_usage_perc", + EngineType.SGLANG: "sglang:token_usage", + }, +} + + +def metric_name(key: str, engine: EngineType) -> str | None: + """Exposition name for ``key`` on ``engine``, or None if not known.""" + return _NAMES[key].get(engine) + + +def parse_metric(text: str, name: str) -> float | None: + """Read a single unlabelled gauge out of Prometheus text exposition. + + Returns None when the series is absent, which is not the same as 0.0 — a + caller draining in-flight work must not read "metric missing" as "idle". + """ + m = re.search(rf"^{re.escape(name)}(?:\{{[^}}]*\}})?\s+([0-9.eE+-]+)\s*$", text, re.MULTILINE) + if m is None: + return None + try: + return float(m.group(1)) + except ValueError: + return None + + +def inflight_from_metrics(text: str, engine: EngineType) -> float | None: + """Requests the engine is running plus those it has queued. + + Queued requests count: a request the engine has accepted but not started is + still work the client is waiting on, and killing the process loses it just + as surely as one mid-generation. + + None means the engine's in-flight count could not be determined. + """ + total = 0.0 + seen = False + for key in ("requests_running", "requests_waiting"): + name = metric_name(key, engine) + if name is None: + continue + value = parse_metric(text, name) + if value is None: + continue + total += value + seen = True + return total if seen else None diff --git a/infera/common/registration.py b/infera/common/registration.py index 26059c05..bbb5a377 100644 --- a/infera/common/registration.py +++ b/infera/common/registration.py @@ -12,6 +12,7 @@ import httpx from infera.common.discovery import DEFAULT_PREFIX, _b64, _normalize_endpoint +from infera.common.worker_pool import WorkerStatus from infera.engine.base import EngineConfig logger = logging.getLogger(__name__) @@ -19,12 +20,17 @@ _DEFAULT_LEASE_TTL = 30 # seconds -def build_worker_payload(config: EngineConfig) -> dict: +def build_worker_payload(config: EngineConfig, *, status: WorkerStatus | None = None) -> dict: """Build the worker registration record shared by every backend. The same dict is PUT to etcd (RegistrationClient) or stored in the worker Pod annotation (K8sRegistrationClient), so the server-side parse (discovery.worker_info_from_json) is transport-agnostic. + + ``status`` is omitted for a healthy worker, which keeps the record identical + to what older workers wrote and lets the parser's ACTIVE default stand. It + is set only to announce DRAINING, so the field's presence means something + happened rather than being ambient. """ worker_id = f"{config.host}:{config.port}" payload: dict = { @@ -40,6 +46,8 @@ def build_worker_payload(config: EngineConfig) -> dict: "dp_size": config.dp_size, "request_transport": getattr(config, "request_transport", "http"), } + if status is not None and status is not WorkerStatus.ACTIVE: + payload["status"] = status.value if config.kv is not None: payload["kv"] = config.kv.to_dict() return payload @@ -98,6 +106,36 @@ async def register(self, config: EngineConfig) -> str: ) return worker_id + async def announce_draining(self) -> bool: + """Rewrite the record as DRAINING, keeping the lease alive. + + ``list_active`` filters DRAINING out, so this stops new work being + routed here without deleting the record — which is the difference that + matters during a shutdown. A worker that simply vanishes is + indistinguishable from one that crashed; one that is visibly draining + tells an operator (and ``/v1/workers``) that a rolling update is + proceeding normally and roughly how far along it is. + + On the Kubernetes backend this is largely redundant with the registry's + ``deletionTimestamp`` check, which removes a condemned Pod without the + worker having to say anything. It is not redundant on etcd, where + nothing else observes that the process is going away. + """ + if self._lease_id is None or self._key is None or self._config is None: + return False + try: + value = json.dumps(build_worker_payload(self._config, status=WorkerStatus.DRAINING)) + r = await self._http.post( + "/v3/kv/put", + json={"key": _b64(self._key), "value": _b64(value), "lease": self._lease_id}, + ) + r.raise_for_status() + logger.info("worker %s announced DRAINING", self._worker_id) + return True + except Exception as exc: # noqa: BLE001 - shutdown must continue + logger.warning("could not announce DRAINING: %s", exc) + return False + async def deregister(self) -> None: if self._lease_id is not None: try: diff --git a/infera/common/registration_k8s.py b/infera/common/registration_k8s.py index 747abf83..04777fc4 100644 --- a/infera/common/registration_k8s.py +++ b/infera/common/registration_k8s.py @@ -26,6 +26,7 @@ from infera.common.discovery_k8s import WORKER_INFO_ANNOTATION from infera.common.k8s_client import in_cluster_namespace, make_client from infera.common.registration import build_worker_payload +from infera.common.worker_pool import WorkerStatus from infera.engine.base import EngineConfig logger = logging.getLogger(__name__) @@ -80,6 +81,31 @@ async def register(self, config: EngineConfig) -> str: ) return worker_id + async def announce_draining(self) -> bool: + """Rewrite the annotation as DRAINING instead of clearing it. + + ``list_active`` filters DRAINING out, so new work stops being routed + here while the record — and therefore the worker's visibility in + ``/v1/workers`` — survives the drain. A worker that vanishes looks the + same as one that crashed. + + Largely belt-and-braces on this backend: ``KubernetesRegistry`` already + drops a Pod the moment it carries a ``deletionTimestamp``, so a + kubectl-initiated shutdown has stopped receiving work before the worker + is even signalled. This covers the paths that do not go through Pod + deletion at all. + """ + if self._config is None: + return False + try: + payload = build_worker_payload(self._config, status=WorkerStatus.DRAINING) + await self._patch_annotation(json.dumps(payload)) + logger.info("worker %s announced DRAINING", self._worker_id) + return True + except Exception as exc: # noqa: BLE001 - shutdown must continue + logger.warning("could not announce DRAINING: %s", exc) + return False + async def deregister(self) -> None: # Best-effort: clear the annotation so a terminating-but-lingering Pod # stops being routed before its DELETE event lands. diff --git a/infera/engine/drain.py b/infera/engine/drain.py new file mode 100644 index 00000000..032cc165 --- /dev/null +++ b/infera/engine/drain.py @@ -0,0 +1,113 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Let in-flight generations finish before the engine is stopped. + +On the NATS transport infera owns the request path, so it knows exactly what is +in flight and ``NatsRequestServer.stop(drain=True)`` waits for it. On HTTP the +router talks straight to the engine's own server: infera never sees the request, +cannot count it, and so — until this — did not wait for it. Shutdown went +``deregister()`` then ``engine.stop()``, cutting every active generation. + +The way out is to ask the engine, which does know. It publishes its running and +queued request counts on ``/metrics``; poll until both reach zero or the timeout +expires. That is a poll rather than a signal, so it is bounded by +``poll_interval`` rather than exact — acceptable, because the alternative is not +draining at all. + +Two behaviours are deliberate: + +* **Deregister first, then drain.** Ordering is the whole point. Draining while + still registered just means more work arrives; the router has to stop choosing + this worker before waiting for the work it already has. +* **An unreadable metric does not block shutdown.** If the engine's in-flight + count cannot be determined — an unknown engine, a renamed series, a dead HTTP + server — this logs loudly and returns rather than hanging until the timeout. + A rolling update that stalls on a parse failure is a worse outcome than one + that cuts a request, and a silent full-timeout wait would look identical to a + genuinely busy worker. +""" + +from __future__ import annotations + +import asyncio +import logging +import time + +import httpx + +from infera.common.engine_metrics import inflight_from_metrics +from infera.common.worker_pool import EngineType + +logger = logging.getLogger(__name__) + +_POLL_INTERVAL_S = 0.5 + + +async def drain_engine_inflight( + *, + host: str, + port: int, + engine: EngineType, + timeout: float, + poll_interval: float = _POLL_INTERVAL_S, +) -> bool: + """Wait until the engine reports no in-flight work, bounded by ``timeout``. + + Returns True if it drained, False if it timed out or could not be measured. + Never raises: this runs on the shutdown path, where an exception would skip + the engine teardown that follows. + """ + if timeout <= 0: + return False + + # The engine binds the advertised port, but 0.0.0.0 is not a destination. + probe_host = "127.0.0.1" if host in ("0.0.0.0", "", "::") else host + url = f"http://{probe_host}:{port}/metrics" + deadline = time.monotonic() + timeout + peak = 0.0 + + try: + async with httpx.AsyncClient(timeout=5.0) as client: + while True: + try: + resp = await client.get(url) + inflight = ( + inflight_from_metrics(resp.text, engine) + if resp.status_code == 200 + else None + ) + except httpx.HTTPError as exc: + logger.info("drain: engine metrics unreachable (%s); not waiting", exc) + return False + + if inflight is None: + logger.warning( + "drain: cannot read in-flight count for %s from %s -- shutting down " + "WITHOUT draining. In-flight generations will be cut.", + engine.value, + url, + ) + return False + + peak = max(peak, inflight) + if inflight <= 0: + if peak > 0: + logger.info("drain: engine idle, %.0f request(s) completed", peak) + return True + + if time.monotonic() >= deadline: + logger.warning( + "drain: timeout after %.0fs with %.0f request(s) still in flight; " + "they will be cut", + timeout, + inflight, + ) + return False + + await asyncio.sleep(min(poll_interval, max(0.0, deadline - time.monotonic()))) + except Exception as exc: # noqa: BLE001 - shutdown must continue regardless + logger.warning("drain: aborted (%s: %s); not waiting", type(exc).__name__, exc) + return False diff --git a/infera/engine/sglang/__main__.py b/infera/engine/sglang/__main__.py index 75f6ae69..1873434a 100644 --- a/infera/engine/sglang/__main__.py +++ b/infera/engine/sglang/__main__.py @@ -30,6 +30,7 @@ ) from infera.common.registration import RegistrationClient from infera.engine.base import watch_engine_death +from infera.engine.drain import drain_engine_inflight from infera.engine.sglang.args import SglangWorkerArgs, parse_sglang_args from infera.engine.sglang.kv_wiring import ( SglangKvWiring, @@ -414,9 +415,28 @@ async def _run_after_start(args: SglangWorkerArgs, engine: SglangEngine, config) except asyncio.CancelledError: pass - await reg_client.deregister() + # Announce first, drain second, deregister last. Announcing DRAINING takes + # this worker out of routing (list_active filters it) while leaving the + # record in place, so for the whole drain it is visibly draining rather than + # simply gone -- which is what distinguishes an orderly rollout from a crash + # in /v1/workers. Deregistering first would work too, but it throws away + # that signal at exactly the moment someone is watching for it. + await reg_client.announce_draining() if nats_req_server is not None: await nats_req_server.stop(drain=True, drain_timeout=args.drain_timeout) + else: + # HTTP transport: the router talks straight to the engine, so infera + # never saw these requests and has to ask the engine what is still in + # flight. The DRAINING announcement above already stopped new work being + # routed here; this waits for the work already accepted. + await drain_engine_inflight( + host=config.host, + port=config.port, + engine=config.engine, + timeout=args.drain_timeout, + ) + + await reg_client.deregister() if kv_relay is not None: await kv_relay.stop() diff --git a/infera/engine/vllm/__main__.py b/infera/engine/vllm/__main__.py index b8f33c48..7179ee23 100644 --- a/infera/engine/vllm/__main__.py +++ b/infera/engine/vllm/__main__.py @@ -24,6 +24,7 @@ from infera.common.registration import RegistrationClient from infera.common.worker_pool import DisaggMode, KvRegistrationMetadata from infera.engine.base import watch_engine_death +from infera.engine.drain import drain_engine_inflight from infera.engine.vllm.args import VllmWorkerArgs, parse_vllm_args from infera.engine.vllm.worker import VllmEngine @@ -360,9 +361,28 @@ async def main() -> None: except asyncio.CancelledError: pass - await reg_client.deregister() + # Announce first, drain second, deregister last. Announcing DRAINING takes + # this worker out of routing (list_active filters it) while leaving the + # record in place, so for the whole drain it is visibly draining rather than + # simply gone -- which is what distinguishes an orderly rollout from a crash + # in /v1/workers. Deregistering first would work too, but it throws away + # that signal at exactly the moment someone is watching for it. + await reg_client.announce_draining() if nats_req_server is not None: await nats_req_server.stop(drain=True, drain_timeout=args.drain_timeout) + else: + # HTTP transport: the router talks straight to the engine, so infera + # never saw these requests and has to ask the engine what is still in + # flight. The DRAINING announcement above already stopped new work being + # routed here; this waits for the work already accepted. + await drain_engine_inflight( + host=config.host, + port=config.port, + engine=config.engine, + timeout=args.drain_timeout, + ) + + await reg_client.deregister() if kv_relay is not None: await kv_relay.stop() await engine.stop() diff --git a/infera/tools/fakeworker/server.py b/infera/tools/fakeworker/server.py index 36c68e21..77f5c98d 100644 --- a/infera/tools/fakeworker/server.py +++ b/infera/tools/fakeworker/server.py @@ -448,20 +448,25 @@ async def _serve(args) -> None: stop = asyncio.Event() async def _shutdown() -> None: - # Deregister BEFORE draining, matching the real worker: the router must - # stop sending new work before we stop accepting it, or the gap shows up - # to clients as failures rather than as a clean drain. + # Announce, drain, then deregister -- the same order as the real worker + # entrypoints. Stopping new work has to come first or the drain just + # races arrivals; keeping the record until the end is what makes the + # worker visibly draining rather than simply gone. state.draining = True - hb_task.cancel() try: - await reg.deregister() + await reg.announce_draining() except Exception as exc: # noqa: BLE001 - shutdown must not raise - logger.warning("deregister failed: %s", exc) + logger.warning("announce_draining failed: %s", exc) deadline = time.monotonic() + args.drain_timeout while state.running and time.monotonic() < deadline: await asyncio.sleep(0.1) if state.running: logger.warning("drain timeout with %d request(s) still in flight", state.running) + hb_task.cancel() + try: + await reg.deregister() + except Exception as exc: # noqa: BLE001 - shutdown must not raise + logger.warning("deregister failed: %s", exc) server.should_exit = True stop.set() diff --git a/tests/unit/common/test_discovery_k8s_terminating.py b/tests/unit/common/test_discovery_k8s_terminating.py new file mode 100644 index 00000000..8e3b05df --- /dev/null +++ b/tests/unit/common/test_discovery_k8s_terminating.py @@ -0,0 +1,114 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""A condemned Pod must leave the pool before its process is killed. + +Kubernetes keeps ``phase: Running`` on a terminating Pod until its containers +exit, so liveness alone cannot tell a healthy worker from one that is seconds +from SIGTERM. The operator makes that window long on purpose -- it injects a +``preStop sleep`` so in-flight work has time to finish -- which means that +without a ``deletionTimestamp`` check the router spends the entire drain window +assigning new requests to a worker that is guaranteed to be killed. + +These tests pin the removal rules rather than the implementation: what matters +is which observable Pod states take a worker out of rotation. +""" + +from __future__ import annotations + +import json + +from infera.common.discovery_k8s import WORKER_INFO_ANNOTATION, KubernetesRegistry + + +def _payload(worker_id: str = "10.0.0.1:8080") -> str: + host, port = worker_id.split(":") + return json.dumps( + { + "worker_id": worker_id, + "url": f"http://{worker_id}", + "model_name": "m", + "engine": "sglang", + "disagg_mode": "mixed", + "disagg_meta": {}, + "kv_events_endpoint": None, + "kv_block_size": None, + "dp_rank": None, + "dp_size": None, + "request_transport": "http", + } + ) + + +def _pod(name="w-0", *, annotated=True, phase="Running", terminating=False): + meta: dict = {"name": name} + if annotated: + meta["annotations"] = {WORKER_INFO_ANNOTATION: _payload()} + if terminating: + meta["deletionTimestamp"] = "2026-08-04T03:00:00Z" + return {"metadata": meta, "status": {"phase": phase}} + + +def _registry(): + removed: list[str] = [] + reg = KubernetesRegistry( + "app=infera", + namespace="infera", + on_worker_removed=removed.append, + ) + return reg, removed + + +def _ids(reg): + return [w.worker_id for w in reg.pool.list_active()] + + +def test_running_pod_registers(): + reg, _ = _registry() + reg._handle_pod(_pod(), deleted=False) + assert _ids(reg) == ["10.0.0.1:8080"] + + +def test_terminating_pod_is_removed_while_still_running(): + """The case this exists for. `phase` is still Running -- only the deletion + timestamp distinguishes a healthy worker from one inside its preStop delay, + and every request routed to it in that window is work that gets cut.""" + reg, removed = _registry() + reg._handle_pod(_pod(), deleted=False) + assert _ids(reg) == ["10.0.0.1:8080"] + + reg._handle_pod(_pod(terminating=True, phase="Running"), deleted=False) + assert _ids(reg) == [], "a condemned Pod must not stay a routing candidate" + assert removed == ["10.0.0.1:8080"] + + +def test_terminating_pod_never_enters_the_pool(): + """A relist during a rolling update can surface an already-terminating Pod + the registry has never seen. It must not be admitted.""" + reg, _ = _registry() + reg._handle_pod(_pod(terminating=True), deleted=False) + assert _ids(reg) == [] + + +def test_removal_is_idempotent(): + """Watch events are re-delivered after a 410/relist, so the same + terminating Pod arrives more than once.""" + reg, removed = _registry() + reg._handle_pod(_pod(), deleted=False) + for _ in range(3): + reg._handle_pod(_pod(terminating=True), deleted=False) + assert removed == ["10.0.0.1:8080"], "must not fire the removal callback repeatedly" + + +def test_other_removal_rules_still_hold(): + for label, kwargs, deleted in ( + ("explicit DELETE", {}, True), + ("annotation cleared", {"annotated": False}, False), + ("no longer Running", {"phase": "Failed"}, False), + ): + reg, _ = _registry() + reg._handle_pod(_pod(), deleted=False) + reg._handle_pod(_pod(**kwargs), deleted=deleted) + assert _ids(reg) == [], f"{label} must still deregister" diff --git a/tests/unit/common/test_draining_status.py b/tests/unit/common/test_draining_status.py new file mode 100644 index 00000000..852b6cdd --- /dev/null +++ b/tests/unit/common/test_draining_status.py @@ -0,0 +1,121 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Announcing DRAINING, and what it is actually for. + +``WorkerStatus.DRAINING`` has been in the enum and filtered out of +``list_active`` since the beginning, and until now nothing ever set it. The +value it adds over simply deleting the record is not routing -- both stop new +work -- it is that the worker stays *visible* while it drains. A worker that +vanishes looks identical to one that crashed; one that reports DRAINING tells an +operator a rollout is proceeding and roughly how far along it is. +""" + +from __future__ import annotations + +import json + +import pytest + +from infera.common.discovery import worker_info_from_json +from infera.common.registration import build_worker_payload +from infera.common.worker_pool import EngineType, WorkerPool, WorkerStatus +from infera.engine.base import EngineConfig + + +def _cfg(): + return EngineConfig(model_name="m", host="10.0.0.1", port=8080, engine=EngineType.SGLANG) + + +def test_healthy_payload_omits_status_entirely(): + """Keeps the record byte-identical to what older workers wrote, so the + parser's ACTIVE default stands and the field's presence means something.""" + assert "status" not in build_worker_payload(_cfg()) + assert "status" not in build_worker_payload(_cfg(), status=WorkerStatus.ACTIVE) + + +def test_draining_payload_carries_the_status(): + payload = build_worker_payload(_cfg(), status=WorkerStatus.DRAINING) + assert payload["status"] == "draining" + + +def test_round_trip_through_discovery(): + """The wire record has to survive the same parse every backend uses.""" + payload = build_worker_payload(_cfg(), status=WorkerStatus.DRAINING) + info = worker_info_from_json(json.loads(json.dumps(payload))) + assert info.status is WorkerStatus.DRAINING + + +def test_draining_worker_is_excluded_but_still_visible(): + """The whole point: out of rotation, still in the fleet listing.""" + pool = WorkerPool() + pool.add(worker_info_from_json(build_worker_payload(_cfg()))) + assert [w.worker_id for w in pool.list_active()] == ["10.0.0.1:8080"] + + pool.add(worker_info_from_json(build_worker_payload(_cfg(), status=WorkerStatus.DRAINING))) + assert pool.list_active() == [], "a draining worker must not be routed to" + assert pool.get("10.0.0.1:8080") is not None, "but it must still be observable" + + +# --- the etcd client ---------------------------------------------------------- + + +class _FakeHttp: + def __init__(self): + self.puts: list[dict] = [] + self.fail = False + + async def post(self, path, json=None): # noqa: A002 - mirrors httpx + if self.fail: + raise RuntimeError("etcd unreachable") + self.puts.append({"path": path, "json": json}) + + class R: + @staticmethod + def raise_for_status(): + pass + + return R() + + +@pytest.mark.asyncio +async def test_etcd_announce_writes_draining_on_the_same_lease(): + from infera.common.registration import RegistrationClient + + c = RegistrationClient("http://etcd:2379") + c._http = _FakeHttp() + c._lease_id, c._key, c._worker_id, c._config = 42, "/infera/workers/w", "w", _cfg() + + assert await c.announce_draining() is True + (put,) = c._http.puts + assert put["path"] == "/v3/kv/put" + assert put["json"]["lease"] == 42, "must keep the lease, not orphan the key" + + import base64 + + value = json.loads(base64.b64decode(put["json"]["value"])) + assert value["status"] == "draining" + + +@pytest.mark.asyncio +async def test_announce_never_raises_on_the_shutdown_path(): + """It runs immediately before the drain; raising here would skip it.""" + from infera.common.registration import RegistrationClient + + c = RegistrationClient("http://etcd:2379") + c._http = _FakeHttp() + c._http.fail = True + c._lease_id, c._key, c._worker_id, c._config = 1, "/k", "w", _cfg() + assert await c.announce_draining() is False + + +@pytest.mark.asyncio +async def test_announce_before_register_is_a_no_op(): + from infera.common.registration import RegistrationClient + + c = RegistrationClient("http://etcd:2379") + c._http = _FakeHttp() + assert await c.announce_draining() is False + assert c._http.puts == [] diff --git a/tests/unit/engine/test_drain.py b/tests/unit/engine/test_drain.py new file mode 100644 index 00000000..ca443f78 --- /dev/null +++ b/tests/unit/engine/test_drain.py @@ -0,0 +1,162 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Graceful drain on the HTTP transport. + +On NATS infera owns the request path and knows what is in flight. On HTTP the +router talks straight to the engine, so infera has to ask the engine — which +means every failure mode is a *measurement* failure, and the interesting cases +are all about what happens when the number cannot be trusted. + +The rule these tests pin down: never treat "unknown" as "idle". An unreadable +metric that defaults to zero would make the drain pass instantly and cut live +generations, and it would do so silently. +""" + +from __future__ import annotations + +import asyncio + +import httpx +import pytest + +from infera.common.engine_metrics import inflight_from_metrics, metric_name, parse_metric +from infera.common.worker_pool import EngineType +from infera.engine import drain as drain_mod +from infera.engine.drain import drain_engine_inflight + +# --- reading the engine's numbers --------------------------------------------- + + +def test_parses_plain_and_labelled_gauges(): + assert parse_metric("vllm:num_requests_running 3.0\n", "vllm:num_requests_running") == 3.0 + assert parse_metric('x:g{a="b"} 7\n', "x:g") == 7.0 + + +def test_missing_series_is_none_not_zero(): + """The distinction the whole drain rests on.""" + assert parse_metric("something_else 1\n", "vllm:num_requests_running") is None + + +def test_inflight_counts_running_plus_waiting(): + """A queued request is work a client is waiting on; killing the process + loses it just as surely as one mid-generation.""" + text = "vllm:num_requests_running 2\nvllm:num_requests_waiting 5\n" + assert inflight_from_metrics(text, EngineType.VLLM) == 7.0 + + +def test_unknown_engine_yields_none(): + """ATOM has no mapping. Guessing one would report an idle engine, and an + idle engine is exactly the answer that makes a drain cut live requests.""" + assert metric_name("requests_running", EngineType.ATOM) is None + assert inflight_from_metrics("vllm:num_requests_running 4\n", EngineType.ATOM) is None + + +def test_partial_metrics_still_count(): + text = "sglang:num_running_reqs 1\n" + assert inflight_from_metrics(text, EngineType.SGLANG) == 1.0 + + +# --- the drain loop ----------------------------------------------------------- + + +def _patch_client(monkeypatch, handler): + """Route drain's httpx client at a mock transport.""" + real = httpx.AsyncClient + + def factory(*a, **kw): + kw.pop("timeout", None) + return real(transport=httpx.MockTransport(handler), timeout=5.0) + + monkeypatch.setattr(drain_mod.httpx, "AsyncClient", factory) + + +@pytest.mark.asyncio +async def test_returns_when_engine_goes_idle(monkeypatch): + counts = iter([3, 2, 0]) + + def handler(request): + return httpx.Response(200, text=f"vllm:num_requests_running {next(counts)}\n") + + _patch_client(monkeypatch, handler) + drained = await drain_engine_inflight( + host="1.2.3.4", port=8000, engine=EngineType.VLLM, timeout=5, poll_interval=0.01 + ) + assert drained is True + + +@pytest.mark.asyncio +async def test_times_out_while_still_busy(monkeypatch): + def handler(request): + return httpx.Response(200, text="vllm:num_requests_running 4\n") + + _patch_client(monkeypatch, handler) + drained = await drain_engine_inflight( + host="1.2.3.4", port=8000, engine=EngineType.VLLM, timeout=0.2, poll_interval=0.01 + ) + assert drained is False, "a busy engine must not report a clean drain" + + +@pytest.mark.asyncio +async def test_unreadable_metric_does_not_hang(monkeypatch, caplog): + """A rolling update that stalls on a parse failure is worse than one that + cuts a request -- and a silent full-timeout wait is indistinguishable from + a genuinely busy worker.""" + + def handler(request): + return httpx.Response(200, text="totally_different_metric 1\n") + + _patch_client(monkeypatch, handler) + started = asyncio.get_running_loop().time() + drained = await drain_engine_inflight( + host="1.2.3.4", port=8000, engine=EngineType.VLLM, timeout=30, poll_interval=0.01 + ) + elapsed = asyncio.get_running_loop().time() - started + assert drained is False + assert elapsed < 1.0, f"returned after {elapsed:.1f}s; must not wait out the timeout" + assert "WITHOUT draining" in caplog.text + + +@pytest.mark.asyncio +async def test_unreachable_engine_does_not_hang(monkeypatch): + def handler(request): + raise httpx.ConnectError("refused", request=request) + + _patch_client(monkeypatch, handler) + drained = await drain_engine_inflight( + host="1.2.3.4", port=8000, engine=EngineType.VLLM, timeout=30, poll_interval=0.01 + ) + assert drained is False + + +@pytest.mark.asyncio +async def test_zero_timeout_is_a_no_op(monkeypatch): + calls = [] + + def handler(request): + calls.append(1) + return httpx.Response(200, text="vllm:num_requests_running 0\n") + + _patch_client(monkeypatch, handler) + assert ( + await drain_engine_inflight(host="1.2.3.4", port=8000, engine=EngineType.VLLM, timeout=0) + is False + ) + assert calls == [], "--drain-timeout 0 must not even probe" + + +@pytest.mark.asyncio +async def test_never_raises_on_the_shutdown_path(monkeypatch): + """This runs immediately before engine.stop(); an exception here would skip + the teardown that follows.""" + + def factory(*a, **kw): + raise RuntimeError("boom") + + monkeypatch.setattr(drain_mod.httpx, "AsyncClient", factory) + assert ( + await drain_engine_inflight(host="1.2.3.4", port=8000, engine=EngineType.VLLM, timeout=1) + is False + ) From 05e51110931661e972336ce14aba463d31e0aa94 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Tue, 4 Aug 2026 07:27:18 +0000 Subject: [PATCH 05/88] fix(drain): corrections found by running a real SGLang engine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Validated graceful drain against real SGLang 0.5.15 serving Qwen3-8B on MI355X, three instances, one GPU each, registering through the real worker entrypoint. Six concurrent 4000-token generations were in flight (num_running_reqs = 6.0) when the container got SIGTERM; all six returned 200 with full-length output, 22s later. The log shows the intended sequence: announced DRAINING -> "engine idle for 6s, 6 request(s) completed" -> deregistered -> engine stopping. Getting there exposed four things unit tests could not. **sglang serves /metrics only with --enable-metrics; otherwise it 404s.** Drain reads the in-flight count from there, so on a default deployment it would have found nothing and shut down without waiting — correct behaviour on my side, feature silently absent in practice. The worker now injects the flag, since infera depends on that endpoint for a correct shutdown. **The gauges are labelled per rank.** Real output is sglang:num_running_reqs{engine_type=...,tp_rank="0",...}, one series per rank, and the parser took only the first match — so with tp_size > 1 an idle rank could hide a busy one and the drain would exit immediately. It now sums every label set, which is safe here because the sum is zero exactly when every rank is zero. **The gauges lag.** num_running_reqs stayed at 12 for 5-15s after the last HTTP response completed; SGLang refreshes them on its own schedule. Stale-high is harmless, but stale-low at the start is not: a request accepted moments before SIGTERM may not be counted yet, so a single zero reading cannot distinguish idle from not-yet-counted. Drain now requires zero to persist for a settle window (6s, chosen from the measured lag) before believing it. **The metric names were right.** sglang:num_running_reqs, sglang:num_queue_reqs and sglang:token_usage all exist as guessed — they had been carrying a "second-hand, unverified" caveat since they were added. That caveat is now discharged for SGLang. Reading the real exposition also surfaced the PD handoff queues (num_prefill_bootstrap/inflight_queue_reqs, num_decode_prealloc/transfer_queue_reqs), which are now counted as in-flight: a prefill worker can show no running and no queued requests while KV transfers are still outstanding, and killing it there strands the decode workers waiting on that KV. Signed-off-by: Zhang, Jiejing --- infera/common/engine_metrics.py | 54 ++++++++++++++++++++++++++------- infera/engine/drain.py | 32 +++++++++++++++++-- infera/engine/sglang/worker.py | 8 +++++ tests/unit/engine/test_drain.py | 38 +++++++++++++++++++++-- 4 files changed, 116 insertions(+), 16 deletions(-) diff --git a/infera/common/engine_metrics.py b/infera/common/engine_metrics.py index 53aca1d8..3949affc 100644 --- a/infera/common/engine_metrics.py +++ b/infera/common/engine_metrics.py @@ -45,6 +45,21 @@ }, } +#: Extra per-engine gauges that also represent unfinished work, counted only +#: when draining. These are the PD handoff queues: a prefill worker can show no +#: running and no queued requests while KV transfers are still outstanding, and +#: killing it there strands the decode workers waiting on that KV -- the failure +#: every PD system in the field documents and none of them prevents. +#: Verified present on SGLang 0.5.15 (`--enable-metrics`). +_DRAIN_EXTRA: dict[EngineType, tuple[str, ...]] = { + EngineType.SGLANG: ( + "sglang:num_prefill_bootstrap_queue_reqs", + "sglang:num_prefill_inflight_queue_reqs", + "sglang:num_decode_prealloc_queue_reqs", + "sglang:num_decode_transfer_queue_reqs", + ), +} + def metric_name(key: str, engine: EngineType) -> str | None: """Exposition name for ``key`` on ``engine``, or None if not known.""" @@ -52,18 +67,28 @@ def metric_name(key: str, engine: EngineType) -> str | None: def parse_metric(text: str, name: str) -> float | None: - """Read a single unlabelled gauge out of Prometheus text exposition. + """Sum every label set of a gauge in Prometheus text exposition. - Returns None when the series is absent, which is not the same as 0.0 — a + Engines label these per rank -- SGLang emits + ``sglang:num_running_reqs{tp_rank="0",...}`` and one series per rank -- so + reading only the first match would let a busy rank hide behind an idle one. + Summing is safe for the question a drain asks, because the sum is zero + exactly when every rank is zero. + + Returns None when the series is absent, which is not the same as 0.0: a caller draining in-flight work must not read "metric missing" as "idle". """ - m = re.search(rf"^{re.escape(name)}(?:\{{[^}}]*\}})?\s+([0-9.eE+-]+)\s*$", text, re.MULTILINE) - if m is None: - return None - try: - return float(m.group(1)) - except ValueError: - return None + total = 0.0 + found = False + for m in re.finditer( + rf"^{re.escape(name)}(?:\{{[^}}]*\}})?\s+([0-9.eE+-]+)\s*$", text, re.MULTILINE + ): + try: + total += float(m.group(1)) + except ValueError: + continue + found = True + return total if found else None def inflight_from_metrics(text: str, engine: EngineType) -> float | None: @@ -71,7 +96,8 @@ def inflight_from_metrics(text: str, engine: EngineType) -> float | None: Queued requests count: a request the engine has accepted but not started is still work the client is waiting on, and killing the process loses it just - as surely as one mid-generation. + as surely as one mid-generation. So do the PD handoff queues, where the + request may be finished locally while its KV is still in transit. None means the engine's in-flight count could not be determined. """ @@ -86,4 +112,10 @@ def inflight_from_metrics(text: str, engine: EngineType) -> float | None: continue total += value seen = True - return total if seen else None + if not seen: + return None + # Absent PD queues are genuinely zero here rather than unknown: the engine + # published a metrics page and simply is not running disaggregated. + for name in _DRAIN_EXTRA.get(engine, ()): + total += parse_metric(text, name) or 0.0 + return total diff --git a/infera/engine/drain.py b/infera/engine/drain.py index 032cc165..e6cdc13b 100644 --- a/infera/engine/drain.py +++ b/infera/engine/drain.py @@ -45,6 +45,18 @@ _POLL_INTERVAL_S = 0.5 +#: How long the engine must report zero before we believe it. +#: +#: These gauges are refreshed on the engine's own schedule, not per request. +#: Measured on SGLang 0.5.15: ``num_running_reqs`` stayed at 12 for 5-15s after +#: the last HTTP response completed. The lag is safe in the direction that +#: matters (stale-high just makes the drain wait), but it is dangerous at the +#: start: a request accepted moments before SIGTERM may not be in the gauge yet, +#: so a single zero reading can mean "idle" or "not counted yet". Requiring the +#: zero to persist past one refresh cycle tells those apart. Costs a few seconds +#: on every shutdown; cheap against cutting a live generation. +_SETTLE_S = 6.0 + async def drain_engine_inflight( *, @@ -53,6 +65,7 @@ async def drain_engine_inflight( engine: EngineType, timeout: float, poll_interval: float = _POLL_INTERVAL_S, + settle: float = _SETTLE_S, ) -> bool: """Wait until the engine reports no in-flight work, bounded by ``timeout``. @@ -68,6 +81,7 @@ async def drain_engine_inflight( url = f"http://{probe_host}:{port}/metrics" deadline = time.monotonic() + timeout peak = 0.0 + zero_since: float | None = None try: async with httpx.AsyncClient(timeout=5.0) as client: @@ -93,10 +107,22 @@ async def drain_engine_inflight( return False peak = max(peak, inflight) + now = time.monotonic() if inflight <= 0: - if peak > 0: - logger.info("drain: engine idle, %.0f request(s) completed", peak) - return True + if zero_since is None: + zero_since = now + elif now - zero_since >= settle: + if peak > 0: + logger.info( + "drain: engine idle for %.0fs, %.0f request(s) completed", + settle, + peak, + ) + return True + else: + # A late gauge refresh revealed work we had not seen; the + # settle window has to start over. + zero_since = None if time.monotonic() >= deadline: logger.warning( diff --git a/infera/engine/sglang/worker.py b/infera/engine/sglang/worker.py index 2fb3f533..613b73dd 100644 --- a/infera/engine/sglang/worker.py +++ b/infera/engine/sglang/worker.py @@ -80,6 +80,14 @@ def __init__( async def start(self) -> EngineConfig: argv = list(self.sglang_argv) + # sglang serves /metrics only with --enable-metrics; without it the + # endpoint 404s. Graceful shutdown reads the in-flight request count + # from there, so leaving it off silently downgrades every scale-down + # and rolling update to "kill in-flight generations". Cheap enough to + # always enable, and the caller can still have passed it explicitly. + if not any(a == "--enable-metrics" for a in argv): + argv.append("--enable-metrics") + if self.enable_kv_events: dp_size = int(getattr(self.server_args, "dp_size", 1) or 1) self._kv_events_port = free_tcp_port_block(dp_size) if dp_size > 1 else free_tcp_port() diff --git a/tests/unit/engine/test_drain.py b/tests/unit/engine/test_drain.py index ca443f78..86754ae1 100644 --- a/tests/unit/engine/test_drain.py +++ b/tests/unit/engine/test_drain.py @@ -75,18 +75,52 @@ def factory(*a, **kw): @pytest.mark.asyncio async def test_returns_when_engine_goes_idle(monkeypatch): - counts = iter([3, 2, 0]) + counts = iter([3, 2] + [0] * 100) def handler(request): return httpx.Response(200, text=f"vllm:num_requests_running {next(counts)}\n") _patch_client(monkeypatch, handler) drained = await drain_engine_inflight( - host="1.2.3.4", port=8000, engine=EngineType.VLLM, timeout=5, poll_interval=0.01 + host="1.2.3.4", + port=8000, + engine=EngineType.VLLM, + timeout=5, + poll_interval=0.01, + settle=0.05, ) assert drained is True +@pytest.mark.asyncio +async def test_a_single_zero_reading_is_not_enough(monkeypatch): + """Measured on SGLang: the gauge lags the work by 5-15s, so one zero can + mean "idle" or "not counted yet". A late non-zero must restart the window + rather than being ignored.""" + counts = iter([0, 0, 4] + [0] * 100) + seen: list[float] = [] + + def handler(request): + v = next(counts) + seen.append(v) + return httpx.Response(200, text=f"vllm:num_requests_running {v}\n") + + _patch_client(monkeypatch, handler) + drained = await drain_engine_inflight( + host="1.2.3.4", + port=8000, + engine=EngineType.VLLM, + timeout=5, + poll_interval=0.01, + settle=0.05, + ) + assert drained is True + assert 4 in seen, "the late non-zero reading must have been observed" + # It must have kept polling well past the point where the first two zeros + # would have satisfied a naive implementation. + assert len(seen) > 3 + + @pytest.mark.asyncio async def test_times_out_while_still_busy(monkeypatch): def handler(request): From 0b697f2e72bd331a4f81070e89cab40dfdf62f49 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Tue, 4 Aug 2026 08:58:51 +0000 Subject: [PATCH 06/88] docs(scaling): how to add and remove workers, with measured numbers Adds manual/features/scaling.md. Every figure on the page came from a run on this hardware; nothing is projected. The page is organised around the asymmetry that actually shapes the problem: scale-up is bounded by model load (measured 140 s for an 8B on one MI355X) while scale-down is bounded by the longest in-flight generation, with the router dropping the worker in 15 ms. That gap is why "react to load by starting a worker" does not work for bursts shorter than a cold start, and the page says so rather than leaving it implied. Also corrects a real bug found while validating vLLM. The mapping had vllm:gpu_cache_usage_perc; the running engine exposes vllm:kv_cache_usage_perc. Reading the wrong name returns "no KV in use" silently, so both spellings are now listed and callers take whichever is present. This is exactly the metric-name drift the SGLang entry was already flagged for. Both engines' names are now verified against running engines rather than documentation, and the caveats say so. Measurements on the page: SGLang 0.5.15 and vLLM 0.1.dev19253, Qwen3-8B, one MI355X per instance, HTTP transport, etcd discovery, real router. Drain under load -- six concurrent 4000-token generations in flight at SIGTERM, 6/6 completed with 200 and full-length output on both engines. Scale up then down -- two instances under continuous traffic, a third added and one removed, 260 requests and 0 failures including the windows around each transition. What was not measured is called out in a warning block: multi-node, TP > 1, PD under scaling, and scale-down during an active KV transfer. The PD handoff queues are counted in the drain but that path has not been exercised on hardware. The autoscaling section states the position plainly -- infera ships no autoscaler, and an external one cannot drive an InferaDeployment today because the operator reconciles replicas from the CR on every pass. Signed-off-by: Zhang, Jiejing --- infera/common/engine_metrics.py | 39 ++++- manual/features/routing_and_transport.md | 4 + manual/features/scaling.md | 181 +++++++++++++++++++++++ manual/sphinx/_toc.yml.in | 2 + 4 files changed, 218 insertions(+), 8 deletions(-) create mode 100644 manual/features/scaling.md diff --git a/infera/common/engine_metrics.py b/infera/common/engine_metrics.py index 3949affc..2e5a9837 100644 --- a/infera/common/engine_metrics.py +++ b/infera/common/engine_metrics.py @@ -13,10 +13,13 @@ Provenance, because it is uneven and matters: -* **vLLM** — from its published metrics documentation. -* **SGLang** — second-hand, not verified against a running engine. Treat a - lookup failure as "unknown", never as "zero"; the difference decides whether a - drain waits or gives up. +* **vLLM** — verified against a running engine (vLLM 0.1.dev19253, Qwen3-8B on + MI355X). Note ``kv_cache_usage_perc`` was ``gpu_cache_usage_perc`` in older + builds; the alias list below covers both. +* **SGLang** — verified against a running engine (SGLang 0.5.15, Qwen3-8B on + MI355X). Note sglang serves ``/metrics`` only with ``--enable-metrics``; the + worker entrypoint injects it. Treat a lookup failure as "unknown", never as + "zero"; the difference decides whether a drain waits or gives up. * **ATOM** — unknown. Deliberately absent rather than guessed: a wrong name reads as an idle engine, and an idle engine is exactly the answer that makes a drain cut live requests. @@ -28,8 +31,16 @@ from infera.common.worker_pool import EngineType -#: metric key -> per-engine exposition name. A missing engine means "we do not -#: know", which callers must distinguish from "the value is zero". +#: metric key -> per-engine exposition name(s). A missing engine means "we do +#: not know", which callers must distinguish from "the value is zero". Several +#: names per entry means the engine renamed the series between releases and both +#: spellings are in the wild. +_ALIASES: dict[str, dict[EngineType, tuple[str, ...]]] = { + "kv_cache_usage": { + EngineType.VLLM: ("vllm:kv_cache_usage_perc", "vllm:gpu_cache_usage_perc"), + }, +} + _NAMES: dict[str, dict[EngineType, str]] = { "requests_running": { EngineType.VLLM: "vllm:num_requests_running", @@ -39,8 +50,11 @@ EngineType.VLLM: "vllm:num_requests_waiting", EngineType.SGLANG: "sglang:num_queue_reqs", }, + # vLLM renamed this: older builds expose gpu_cache_usage_perc, current ones + # kv_cache_usage_perc. Both are listed and callers sum whichever is present, + # because pinning one silently returns "no KV in use" on the other. "kv_cache_usage": { - EngineType.VLLM: "vllm:gpu_cache_usage_perc", + EngineType.VLLM: "vllm:kv_cache_usage_perc", EngineType.SGLANG: "sglang:token_usage", }, } @@ -62,10 +76,19 @@ def metric_name(key: str, engine: EngineType) -> str | None: - """Exposition name for ``key`` on ``engine``, or None if not known.""" + """Primary exposition name for ``key`` on ``engine``, or None if unknown.""" return _NAMES[key].get(engine) +def metric_names(key: str, engine: EngineType) -> tuple[str, ...]: + """Every spelling of ``key`` on ``engine``, newest first.""" + alias = _ALIASES.get(key, {}).get(engine) + if alias: + return alias + name = _NAMES[key].get(engine) + return (name,) if name else () + + def parse_metric(text: str, name: str) -> float | None: """Sum every label set of a gauge in Prometheus text exposition. diff --git a/manual/features/routing_and_transport.md b/manual/features/routing_and_transport.md index ef249024..a1437c92 100644 --- a/manual/features/routing_and_transport.md +++ b/manual/features/routing_and_transport.md @@ -150,6 +150,10 @@ discovery still reports it `ACTIVE` is the signal worth alerting on. Both the Python and Rust routers implement this identically, with the same flags. +The breaker is the router's view of a worker that is failing. For the orderly +case — a worker being removed on purpose — see [Scaling a fleet](scaling.md), +which covers draining in-flight generations before shutdown. + ## KV-event transport Powers [KV-aware routing](kv_aware_routing.md). `--kv-event-transport`: diff --git a/manual/features/scaling.md b/manual/features/scaling.md new file mode 100644 index 00000000..0ec09775 --- /dev/null +++ b/manual/features/scaling.md @@ -0,0 +1,181 @@ +# Scaling a fleet + +Adding and removing workers while traffic is flowing. Every number on this page +was measured on the hardware described in [Measurements](#measurements) — none +of it is projected. + +## How it works + +There is no scaling controller. Workers **self-register** into discovery when +they are ready and **deregister** when they shut down, and the router routes to +whatever is registered at that instant. Scaling is therefore just starting and +stopping worker processes; nothing has to be told about it. + +``` +worker ready ──► register (etcd lease / Pod annotation) ──► router's watch fires + ──► receives traffic +SIGTERM ──► announce DRAINING ──► drain in-flight ──► deregister ──► exit +``` + +That shape is why scale-up and scale-down have very different costs. Scale-up is +bounded by **model load**, which is minutes. Scale-down is bounded by the +**longest in-flight generation**, which is seconds — and the router stops +choosing the worker in milliseconds, long before it stops serving. + +## Scaling up + +Start another worker with the same `--model-name` and the same discovery +settings. It joins when it is ready, and not before: registration happens after +the engine has loaded weights, so a worker in the pool is always a worker that +can serve. + +```bash +infera-worker ... --port 20002 --etcd-endpoint http://etcd:2379 +``` + +On Kubernetes, raise `replicas` on the worker service in the `InferaDeployment`. + +**Budget minutes, not seconds.** Measured cold start for an 8B model on one +MI355X was **140 s** from `docker run` to appearing in `/v1/workers`, almost all +of it weight loading. Anything that reacts to load by starting a worker has to +tolerate that delay — a rule that scales up when a queue is deep will still be +scaling up long after the queue drained. + +The corollary matters more than it looks: for a burst shorter than the cold +start, **adding workers cannot help**. Either keep headroom, or shift traffic +between roles that are already running (see +[PD disaggregation](pd_disaggregation.md)). + +## Scaling down + +Send `SIGTERM`. Do not `SIGKILL`, and do not simply delete the Pod without a +grace period. + +The worker then, in this order: + +1. **Announces `DRAINING`.** The router filters draining workers out of routing + immediately, so no new work arrives. The record stays, so the worker remains + visible in `/v1/workers` — a worker that vanishes looks exactly like one that + crashed. +2. **Drains.** On the NATS transport infera tracks in-flight requests directly. + On HTTP the router talks straight to the engine, so infera asks the engine + instead, polling its `/metrics` until running, queued, and PD-handoff queues + all reach zero. Bounded by `--drain-timeout` (default 30 s). +3. **Deregisters**, then stops the engine. + +Requests already in flight run to completion. Requests that arrive during the +drain go to other workers. + +```{note} +`--drain-timeout` is a **ceiling, not a delay** — a worker with nothing in flight +exits in about six seconds regardless. Set it above your p99 generation time. +Anything still running when it expires is cut, with a warning naming the count. +``` + +### Why in-flight work is visible at all + +The engine's own gauges are the only source of truth on the HTTP path, and they +have three properties worth knowing: + +- **SGLang serves `/metrics` only with `--enable-metrics`.** Without it the + endpoint 404s and the drain has nothing to read. The worker entrypoint injects + the flag, so this is handled — but a hand-rolled deployment that bypasses it + will silently lose the drain. +- **The gauges lag.** Measured on SGLang: `num_running_reqs` stayed at 12 for + 5–15 s after the last response completed. The drain therefore requires the + count to read zero continuously for a settle window before believing it, + which also protects against a request accepted moments before `SIGTERM` that + has not been counted yet. +- **PD handoff queues count as in-flight.** A prefill worker can show no running + and no queued requests while KV transfers are still outstanding. Stopping it + there strands the decode workers waiting on that KV, so + `num_prefill_bootstrap_queue_reqs`, `num_prefill_inflight_queue_reqs`, + `num_decode_prealloc_queue_reqs` and `num_decode_transfer_queue_reqs` are + included in the count. + +If the in-flight count cannot be read at all — an unknown engine, a renamed +series, a dead HTTP server — the worker logs a warning naming the metric it +looked for and shuts down **without** draining rather than blocking. A rolling +update that stalls on a parse failure is worse than one that cuts a request, and +a silent full-timeout wait would be indistinguishable from a genuinely busy +worker. + +### On Kubernetes + +Two things beyond `SIGTERM`: + +- The registry drops a Pod as soon as it carries a `deletionTimestamp`, without + waiting for the container to exit. A terminating Pod keeps `phase: Running`, + so without this it would stay a routing candidate for the whole `preStop` + delay — turning a hook meant to make shutdown graceful into extra seconds of + accepting work that is about to be killed. +- `terminationGracePeriodSeconds` must exceed `preStop` + `--drain-timeout`, or + the kubelet `SIGKILL`s mid-drain. The operator sets 120 s with a 15 s preStop; + raise it if your generations are longer. + +## PD and DP + +Prefill and decode register into separate pools and are selected per request, so +they scale **independently** — add prefill for longer inputs, decode for more +concurrent users. Two constraints: + +- **Neither pool can go to zero.** PD dispatch fails closed when either side is + empty. `minReplicas: 0` on either is an outage, not an idle saving. +- **A DP worker's shape decides who picks the rank.** A worker registering + `dp_size > 1` with **no** `dp_rank` is rank-multiplexed: the router fans it + out into one target per rank and pins `X-Data-Parallel-Rank`. A worker that + registers its own `dp_rank` is a plain endpoint and opts out — its address + already selects the rank. Both are valid; only the first involves the router. + +## Measurements + +SGLang 0.5.15 and vLLM 0.1.dev19253, Qwen3-8B, one MI355X per instance, HTTP +transport, etcd discovery, real router. + +| | | +|---|---| +| Cold start (`docker run` → in `/v1/workers`) | **140 s** | +| Scale-down (`SIGTERM` → out of `/v1/workers`) | **30 s** | +| Router reaction to a worker leaving | **15 ms** | +| Drain settle window | 6 s | + +Two runs, both with traffic flowing throughout: + +**Drain under load.** Six concurrent 4000-token generations in flight at +`SIGTERM`. Both engines: **6/6 completed with HTTP 200** and full-length output +(15–19 k characters). SGLang 22 s, vLLM 19 s from signal to last response. + +**Scale up then down.** Two instances, continuous traffic, a third added and +then one removed. **260 requests, 0 failures**, including in the 5-second +windows around each transition. The removed instance's log shows the intended +sequence: + +``` +worker 127.0.0.1:20001 announced DRAINING +drain: engine idle for 6s, 1 request(s) completed +deregistered worker 127.0.0.1:20001 (lease revoked) +``` + +```{warning} +**Not measured:** multi-node workers, TP > 1, PD under scaling, and scale-down +during an active KV transfer. The PD handoff queues are counted in the drain, +but that path has not been exercised on hardware. +``` + +## Autoscaling + +Infera does not ship an autoscaler, and an external one cannot currently drive +an `InferaDeployment`: the operator reconciles `replicas` from the CR on every +pass, so a `HorizontalPodAutoscaler` writing to the child Deployment is reverted +within seconds. Scaling today is a deliberate act — `kubectl scale` on the CR, +or starting and stopping workers. + +The mechanics an autoscaler would need are in place: workers join and leave +cleanly under load, and the signals worth scaling on +(`vllm:num_requests_waiting`, `sglang:num_queue_reqs`, KV cache utilisation) are +exposed by the engines and read by the drain path already. + +The unsolved part is not the plumbing. It is that a **140-second cold start sits +inside a control loop that ticks every 15 seconds**, and that nothing in +Kubernetes lets a scaler choose *which* replica to remove — so the one holding +the warmest KV cache is as likely to go as any other. diff --git a/manual/sphinx/_toc.yml.in b/manual/sphinx/_toc.yml.in index cb76e25c..84cde5c5 100644 --- a/manual/sphinx/_toc.yml.in +++ b/manual/sphinx/_toc.yml.in @@ -45,6 +45,8 @@ subtrees: title: KV-aware routing - file: features/routing_and_transport.md title: Routing and transport + - file: features/scaling.md + title: Scaling a fleet - caption: Serving entries: From 088af554d444ca8ad0f1d8bd1717f75921ba28b9 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Tue, 4 Aug 2026 09:14:23 +0000 Subject: [PATCH 07/88] fix(router): fail over on a unary 5xx, so the breaker actually sees it A worker 5xx before any byte reached the client is precisely what failover exists for, and the non-streaming paths returned it verbatim instead. Both of them: direct HTTP, and NATS unary. Only transport errors and non-JSON bodies raised _Retry. Two consequences, the second worse than the first. A unary request never failed over, so a single wedged worker turned 1/N of all traffic into errors even with healthy workers idle. And because the circuit breaker records failures from the same _Retry path, it never saw those failures either -- the breaker shipped in the previous commit was inert on the most common configuration there is, non-streaming over HTTP. The streaming path and the Rust router both already retried this; this brings the remaining two into line. 4xx still passes straight through: the request is bad, not the worker, so every worker would answer the same and retrying only triples the latency of an error the client needs. Measured against a real SGLang Qwen3-8B worker plus one worker that always 503s, 20 unary requests: 20/20 succeeded and the bad worker was tried 3 times before the breaker opened it. The same setup before this change lost 6 of 20 and left the breaker with no samples at all. Signed-off-by: Zhang, Jiejing --- infera/router/mixed.py | 14 +++++++ tests/unit/router/test_failover.py | 60 ++++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+) diff --git a/infera/router/mixed.py b/infera/router/mixed.py index 5c3e66d2..a90ecbee 100644 --- a/infera/router/mixed.py +++ b/infera/router/mixed.py @@ -255,6 +255,10 @@ async def _attempt_unary( ) ) from None obs["outcome"] = "ok" if status < 400 else f"{status // 100}xx" + # Same rule as the HTTP path below: a 5xx before any data is a + # worker fault and retryable; a 4xx belongs to the request. + if is_worker_fault(status): + raise _Retry(JSONResponse(content=payload_json, status_code=status)) return JSONResponse(content=payload_json, status_code=status) # Direct HTTP forward. @@ -288,6 +292,16 @@ async def _attempt_unary( ) ) from None obs["outcome"] = "ok" if resp.status_code < 400 else f"{resp.status_code // 100}xx" + # A 5xx here is the worker failing before a single byte reached the + # client, which is exactly the case failover exists for -- and until now + # this path returned it verbatim instead, so a unary request over HTTP + # never failed over and never fed the circuit breaker. The streaming + # path and the Rust router both retry it; this brings the third one into + # line. 4xx still passes straight through: the request itself is bad and + # every worker would say the same, so retrying only triples the latency + # of an error the client needs to see. + if is_worker_fault(resp.status_code): + raise _Retry(JSONResponse(content=payload_json, status_code=resp.status_code)) return JSONResponse(content=payload_json, status_code=resp.status_code) async def _normalized_stream( diff --git a/tests/unit/router/test_failover.py b/tests/unit/router/test_failover.py index 4596895a..5027c047 100644 --- a/tests/unit/router/test_failover.py +++ b/tests/unit/router/test_failover.py @@ -283,3 +283,63 @@ async def test_breaker_recovers_after_cooldown(): assert body == b"back", "the half-open probe must reach the recovered worker" assert r.breaker.state_of("w1").value == "closed" await r.aclose() + + +# --- unary 5xx must fail over too (the gap the breaker fell through) --------- + + +@pytest.mark.asyncio +async def test_unary_http_fails_over_on_5xx(): + """A worker 500 before any byte reached the client is exactly what failover + is for. This path used to return it verbatim, so a non-streaming request + over HTTP -- the default in every k8s example -- never failed over and never + reached the circuit breaker.""" + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.host == "w1": + return httpx.Response(500, json={"error": "boom"}) + return httpx.Response(200, json={"id": "ok"}) + + r = _router([_w("w1", transport="http"), _w("w2", transport="http")], nats=None, retries=1) + r._client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + resp = await r.dispatch({"model": "m"}, stream=False) + assert resp.status_code == 200 + assert json.loads(bytes(resp.body))["id"] == "ok" + await r.aclose() + + +@pytest.mark.asyncio +async def test_unary_http_does_not_fail_over_on_4xx(): + """The request is bad, not the worker. Every worker would answer the same, + so retrying only triples the latency of an error the client must see.""" + hits: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + hits.append(request.url.host) + return httpx.Response(400, json={"error": "bad request"}) + + r = _router([_w("w1", transport="http"), _w("w2", transport="http")], nats=None, retries=1) + r._client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + resp = await r.dispatch({"model": "m"}, stream=False) + assert resp.status_code == 400 + assert hits == ["w1"], f"4xx must not be retried, but hit {hits}" + await r.aclose() + + +@pytest.mark.asyncio +async def test_unary_5xx_trips_the_breaker(): + """The consequence that made this worth fixing: without failover the + breaker never saw a unary failure, so a wedged worker was re-picked + forever on the most common configuration.""" + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.host == "bad": + return httpx.Response(503, json={"error": "wedged"}) + return httpx.Response(200, json={"id": "ok"}) + + r = _router([_w("bad", transport="http"), _w("good", transport="http")], nats=None, retries=1) + r._client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + for _ in range(3): + assert (await r.dispatch({"model": "m"}, stream=False)).status_code == 200 + assert r.breaker.state_of("bad").value == "open" + await r.aclose() From ab05f4f3b90e3bfc8b74fd0358fa5523cab65866 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Tue, 4 Aug 2026 09:21:45 +0000 Subject: [PATCH 08/88] =?UTF-8?q?docs(scaling):=20multi-node=20=E2=80=94?= =?UTF-8?q?=20the=20two=20things=20that=20actually=20differ?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an "Across machines" section. The mechanism does not change when workers live on different hosts, because discovery is already the coordination point; what changes is two pieces of configuration, and both fail quietly enough to be worth naming. --advertise-host must be the node's routable address. Leaving it at the single-node default registers a URL that resolves to the wrong machine everywhere else, and the symptom is a router that lists the worker and cannot reach it -- which reads as a broken worker rather than a misconfiguration. Discovery has the mirror-image problem: an etcd advertising a loopback client URL works perfectly on its own node and is invisible from the others. Both are one curl to check, so the section gives the curls. Measured on chi2800 / chi2866, one MI355X each, workers advertising their own IPs, etcd and router on the first node: both registered with distinct addresses, 12 requests split 7/7 across the machines, and a SIGTERM to the remote worker drained cleanly -- three in-flight 3000-token generations all completed (13.7-14.4k characters), it left the fleet after 30s, and 100 requests flowing through the router across the whole transition saw 0 failures. Scoped honestly in a warning block: this is workers on separate machines, not one worker spanning machines (numberOfNodes > 1 / LeaderWorkerSet) and not PD over RDMA between nodes. Neither was exercised. Also notes that rdma/hca is not an allocatable resource on this cluster, so a PD deployment here would need host networking rather than a device plugin. Signed-off-by: Zhang, Jiejing --- manual/features/scaling.md | 42 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/manual/features/scaling.md b/manual/features/scaling.md index 0ec09775..b569fcd3 100644 --- a/manual/features/scaling.md +++ b/manual/features/scaling.md @@ -127,6 +127,48 @@ concurrent users. Two constraints: registers its own `dp_rank` is a plain endpoint and opts out — its address already selects the rank. Both are valid; only the first involves the router. +## Across machines + +Nothing about scaling changes when workers live on different hosts — discovery +is already the coordination point, so a worker on another machine joins the same +way. Two things do change, and both are configuration rather than mechanism: + +**`--advertise-host` must be the node's routable address.** It is the URL peers +dial, and the single-node habit of leaving it at `127.0.0.1` registers an +address that resolves to the wrong machine everywhere else. The failure is +quiet in the worst way: the router *lists* the worker and cannot reach it, so it +looks like a broken worker rather than a misconfiguration. On Kubernetes, take +it from the downward API (`POD_IP`). + +**Discovery must be reachable from every node.** An etcd bound only to loopback, +or advertising a loopback client URL, works perfectly on the node running it and +is invisible from the others. + +Checking both before deploying costs nothing: + +```bash +# from each worker node +curl -s -o /dev/null -w '%{http_code}\n' -X POST http://:2379/v3/kv/range -d '{"key":"Lw=="}' +# from the router node, once a worker has registered +curl -s http://:8000/v1/workers | jq -r '.workers[].url' # must be dialable +``` + +Measured on two nodes (chi2800 / chi2866, one MI355X each, workers advertising +their own IPs, etcd and router on the first node): both workers registered with +distinct addresses, 12 requests distributed 7/7 across the machines, and a +`SIGTERM` to the **remote** worker drained cleanly — its three in-flight +3000-token generations all completed (13.7–14.4 k characters), it left the fleet +after 30 s, and 100 requests flowing through the router during the whole +transition saw **0 failures**. + +```{warning} +This covers workers on separate machines. It does **not** cover a single worker +*spanning* machines (`numberOfNodes > 1`, LeaderWorkerSet) or PD over RDMA +between nodes — neither has been exercised here. Note also that on this cluster +`rdma/hca` is not advertised as an allocatable resource, so a PD deployment +would need host networking and direct device access rather than a device plugin. +``` + ## Measurements SGLang 0.5.15 and vLLM 0.1.dev19253, Qwen3-8B, one MI355X per instance, HTTP From 84f3c85ef7c31106eb2fc3447d926d436cf5fb54 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Tue, 4 Aug 2026 15:34:41 +0000 Subject: [PATCH 09/88] docs(scaling): separate "stops receiving" from "record disappears" The scale-down figure conflated two timings that differ by two orders of magnitude, and the smaller one is the operationally important one. A worker stops being routed to within a second of SIGTERM -- that is the DRAINING announcement plus the router's watch, and it answers the question that matters: is traffic still being sent to something about to die. How long the process then lives is set by the longest generation it was already serving, and is unrelated. Measured with a fake worker holding a 40-second generation: under a second to stop receiving new requests, 38 s until the record disappeared, with the generation completing in between (400/400 chunks). The earlier "30 s" came from watching /v1/workers, which lists draining workers too -- deliberately, so a rollout is visible while it happens. Counting rows there measures process lifetime, not routing. The page now says so and points at the status field instead, and the multi-node paragraph carries the same clarification. Signed-off-by: Zhang, Jiejing --- manual/features/scaling.md | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/manual/features/scaling.md b/manual/features/scaling.md index b569fcd3..1ae31cae 100644 --- a/manual/features/scaling.md +++ b/manual/features/scaling.md @@ -66,6 +66,19 @@ The worker then, in this order: Requests already in flight run to completion. Requests that arrive during the drain go to other workers. +**Two different timings, easily conflated.** A worker stops *receiving* new +requests within a second of `SIGTERM` — that is the `DRAINING` announcement plus +the router's watch, and it is the number that decides whether traffic is still +being sent somewhere that is about to die. How long the *process* then lives is +a separate and much larger number, set by the longest generation it was already +serving. Measured: under a second to stop receiving, 38 s until the record +disappeared, while a 40-second generation ran to completion in between. + +Watching `/v1/workers` measures the second one, not the first: it lists every +worker including draining ones, precisely so a rollout is visible while it +happens. To see the transition, read the `status` field rather than counting +rows. + ```{note} `--drain-timeout` is a **ceiling, not a delay** — a worker with nothing in flight exits in about six seconds regardless. Set it above your p99 generation time. @@ -157,9 +170,10 @@ Measured on two nodes (chi2800 / chi2866, one MI355X each, workers advertising their own IPs, etcd and router on the first node): both workers registered with distinct addresses, 12 requests distributed 7/7 across the machines, and a `SIGTERM` to the **remote** worker drained cleanly — its three in-flight -3000-token generations all completed (13.7–14.4 k characters), it left the fleet -after 30 s, and 100 requests flowing through the router during the whole -transition saw **0 failures**. +3000-token generations all completed (13.7–14.4 k characters), its record +disappeared after 30 s, and 100 requests flowing through the router during the +whole transition saw **0 failures**. (As above, the record surviving 30 s is the +generations finishing, not 30 s of continuing to receive work.) ```{warning} This covers workers on separate machines. It does **not** cover a single worker @@ -177,8 +191,9 @@ transport, etcd discovery, real router. | | | |---|---| | Cold start (`docker run` → in `/v1/workers`) | **140 s** | -| Scale-down (`SIGTERM` → out of `/v1/workers`) | **30 s** | -| Router reaction to a worker leaving | **15 ms** | +| Scale-down: `SIGTERM` → stops receiving new requests | **< 1 s** | +| Scale-down: `SIGTERM` → record gone from `/v1/workers` | **30–38 s** | +| Router reaction to a worker's record being deleted | **15 ms** | | Drain settle window | 6 s | Two runs, both with traffic flowing throughout: From 6364a01aba2344112bbed7aed288ad92e06ebc0a Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Tue, 4 Aug 2026 16:18:41 +0000 Subject: [PATCH 10/88] fix(operator): derive terminationGracePeriodSeconds from --drain-timeout The grace period was a fixed 120s carrying a comment that it "must exceed preStop + the worker --drain-timeout". Nothing parsed that flag -- it lives in free-form args -- so the invariant was documented and unenforced, and raising --drain-timeout for long generations, which is the only reason anyone raises it, pushed shutdown past the grace and turned the drain back into a kill. Measured on the live cluster before this change: a service declaring --drain-timeout 300 was rendered with terminationGracePeriodSeconds 120. The worst-case budget needs 365s; the kubelet would SIGKILL at 120, roughly 105s into the drain. The grace is now preStop + --drain-timeout + 50s of teardown headroom, floored at the previous 120s and still only ever raised, never lowered. The headroom is not padding: engine.stop() alone waits up to 30s for the engine's process group before escalating to SIGKILL, and deregistration plus KV-plane teardown follow the drain. The flag is read from ServiceSpec.Args and from the container's own command/args, because extraPodSpec templates are passed through verbatim and a deployment that tuned the drain is likely to have written it there. Adds the first Go tests in deploy/operator, covering the arithmetic, both flag spellings, fractional and malformed values, the floor, and the extraPodSpec case. Verified with go build, go vet, gofmt and go test in a golang:1.25 container, since Go is not installed here. Also documents the Kubernetes shutdown sequence and the per-stage worst-case budget in manual/features/scaling.md -- the recipes deploy with kubernetes discovery and HTTP transport, which is a different path from the etcd deployments the drain was first measured on. Signed-off-by: Zhang, Jiejing --- .../operator/internal/controller/builders.go | 64 ++++++++++-- .../internal/controller/builders_test.go | 99 +++++++++++++++++++ manual/features/scaling.md | 63 ++++++++++-- 3 files changed, 209 insertions(+), 17 deletions(-) create mode 100644 deploy/operator/internal/controller/builders_test.go diff --git a/deploy/operator/internal/controller/builders.go b/deploy/operator/internal/controller/builders.go index d88384a3..cf474f7f 100644 --- a/deploy/operator/internal/controller/builders.go +++ b/deploy/operator/internal/controller/builders.go @@ -7,6 +7,9 @@ package controller import ( "fmt" + "math" + "strconv" + "strings" appsv1 "k8s.io/api/apps/v1" corev1 "k8s.io/api/core/v1" @@ -30,10 +33,51 @@ const ( lwsKind = "LeaderWorkerSet" // Graceful rolling-upgrade tuning for GPU worker pods. - workerPreStopDrainSeconds = 15 // preStop sleep: let the router drop us before SIGTERM - workerTerminationGraceSeconds int64 = 120 // must exceed preStop + the worker --drain-timeout + workerPreStopDrainSeconds = 15 // preStop sleep: let the router drop us before SIGTERM + workerDefaultDrainTimeoutSeconds = 30 // matches the worker's --drain-timeout default + // Teardown after the drain finishes: deregistering, stopping the KV plane, + // and engine.stop(), which SIGTERMs the engine's process group and waits up + // to 30s before escalating to SIGKILL. + workerTeardownHeadroomSeconds = 50 + // Floor, so short drain timeouts still leave room for a slow engine exit. + workerTerminationGraceSeconds int64 = 120 ) +// graceSecondsFor sizes terminationGracePeriodSeconds so the kubelet cannot +// SIGKILL a worker in the middle of shutting down. +// +// The budget is preStop + the worker's --drain-timeout + teardown. That last +// term is not small: engine.stop() alone waits up to 30s for the engine's +// process group before escalating. Leaving the grace at a fixed 120s was fine +// for the default 30s drain, but --drain-timeout lives in free-form args that +// nothing here parsed -- so raising it for long generations (the exact reason +// anyone raises it) silently pushed shutdown past the grace and turned a +// graceful drain back into a kill. +func graceSecondsFor(args []string) int64 { + drain := workerDefaultDrainTimeoutSeconds + for i, a := range args { + v := "" + if a == "--drain-timeout" && i+1 < len(args) { + v = args[i+1] + } else if strings.HasPrefix(a, "--drain-timeout=") { + v = strings.TrimPrefix(a, "--drain-timeout=") + } + if v == "" { + continue + } + // The worker takes a float; round up so a fractional value never + // shortens the budget. + if f, err := strconv.ParseFloat(v, 64); err == nil && f > 0 { + drain = int(math.Ceil(f)) + } + } + need := int64(workerPreStopDrainSeconds + drain + workerTeardownHeadroomSeconds) + if need < workerTerminationGraceSeconds { + return workerTerminationGraceSeconds + } + return need +} + // labelsFor returns the selector/identity labels for a service's workload. func labelsFor(idepName, svcName string) map[string]string { return map[string]string{ @@ -216,11 +260,16 @@ var mainContainerNames = map[string]struct{}{"main": {}, "infera": {}} // generations, plus a /health readiness probe for single-node workers (skipped // for multi-node LWS groups whose follower ranks > 0 do not serve /health). // Existing values are preserved; the grace is only raised, never lowered. -func injectWorkerRolloutDefaults(spec *corev1.PodSpec, idx int, port int32, addReadiness bool) { +func injectWorkerRolloutDefaults(spec *corev1.PodSpec, idx int, port int32, addReadiness bool, args []string) { if idx < 0 || idx >= len(spec.Containers) { return } c := &spec.Containers[idx] + // The flag can arrive two ways: via ServiceSpec.Args on the rendered path, + // or written straight into the container by an extraPodSpec template, which + // is passed through verbatim. Reading only the first would miss exactly the + // deployments most likely to have tuned it. + drainArgs := append(append(append([]string{}, args...), c.Command...), c.Args...) if addReadiness && c.ReadinessProbe == nil { // SGLang's /health runs a tiny prefill self-check that often takes // >1s, so a 1s probe timeout (the k8s default) flaps the pod between @@ -246,8 +295,9 @@ func injectWorkerRolloutDefaults(spec *corev1.PodSpec, idx int, port int32, addR }, } } - if spec.TerminationGracePeriodSeconds == nil || *spec.TerminationGracePeriodSeconds < workerTerminationGraceSeconds { - grace := workerTerminationGraceSeconds + if want := graceSecondsFor(drainArgs); spec.TerminationGracePeriodSeconds == nil || + *spec.TerminationGracePeriodSeconds < want { + grace := want spec.TerminationGracePeriodSeconds = &grace } } @@ -296,7 +346,7 @@ func podTemplateFromExtra(idep *inferav1alpha1.InferaDeployment, svcName string, // Graceful rolling-upgrade defaults for worker pods rendered by an external // template: inject readiness/preStop/grace the template omitted. if svc.ComponentType == inferav1alpha1.ComponentTypeWorker { - injectWorkerRolloutDefaults(&spec, idx, port, svc.NumberOfNodes <= 1 && !svc.SkipReadinessProbe) + injectWorkerRolloutDefaults(&spec, idx, port, svc.NumberOfNodes <= 1 && !svc.SkipReadinessProbe, svc.Args) } return corev1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{Labels: podLabelsFor(idep.Name, svcName, svc)}, @@ -355,7 +405,7 @@ func podTemplate(idep *inferav1alpha1.InferaDeployment, svcName string, svc infe // readiness is skipped for multi-node LWS groups (follower ranks have no // /health). The server (CPU-only) keeps the default fast shutdown. if svc.ComponentType == inferav1alpha1.ComponentTypeWorker { - injectWorkerRolloutDefaults(&podSpec, 0, port, svc.NumberOfNodes <= 1 && !svc.SkipReadinessProbe) + injectWorkerRolloutDefaults(&podSpec, 0, port, svc.NumberOfNodes <= 1 && !svc.SkipReadinessProbe, svc.Args) } tmpl := corev1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{Labels: podLabelsFor(idep.Name, svcName, svc)}, diff --git a/deploy/operator/internal/controller/builders_test.go b/deploy/operator/internal/controller/builders_test.go new file mode 100644 index 00000000..160448d8 --- /dev/null +++ b/deploy/operator/internal/controller/builders_test.go @@ -0,0 +1,99 @@ +/* +Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + +SPDX-License-Identifier: MIT +*/ + +package controller + +import ( + "testing" + + corev1 "k8s.io/api/core/v1" +) + +// The grace period is the only thing standing between a graceful drain and a +// SIGKILL halfway through one. It has to cover preStop, the worker's own +// --drain-timeout, and the teardown that follows -- of which engine.stop() +// alone can take 30s waiting on the engine's process group. +// +// The failure this guards against is quiet: raising --drain-timeout is exactly +// what an operator does when generations are long, and until the grace was +// derived from it that made shutdown *less* graceful, not more. +func TestGraceSecondsFor(t *testing.T) { + cases := []struct { + name string + args []string + want int64 + }{ + {"no args uses the floor", nil, 120}, + {"default drain stays at the floor", []string{"--drain-timeout", "30"}, 120}, + { + "a long drain raises the grace above the floor", + []string{"--model-path", "/m", "--drain-timeout", "120"}, + 185, // 15 preStop + 120 drain + 50 teardown + }, + {"equals form is parsed too", []string{"--drain-timeout=120"}, 185}, + { + "fractional values round up rather than shortening the budget", + []string{"--drain-timeout", "60.5"}, + 126, // 15 + 61 + 50 + }, + {"a short drain does not lower the floor", []string{"--drain-timeout", "1"}, 120}, + {"garbage falls back to the default", []string{"--drain-timeout", "abc"}, 120}, + {"a trailing flag with no value is ignored", []string{"--drain-timeout"}, 120}, + {"non-positive is ignored", []string{"--drain-timeout", "0"}, 120}, + {"the last occurrence wins", []string{"--drain-timeout", "5", "--drain-timeout", "200"}, 265}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + if got := graceSecondsFor(c.args); got != c.want { + t.Fatalf("graceSecondsFor(%v) = %d, want %d", c.args, got, c.want) + } + }) + } +} + +// The budget must actually hold, not merely be larger than the old constant. +func TestGraceCoversTheWholeShutdown(t *testing.T) { + for _, drain := range []int{30, 60, 120, 300} { + args := []string{"--drain-timeout", itoa(drain)} + grace := graceSecondsFor(args) + need := int64(workerPreStopDrainSeconds + drain + workerTeardownHeadroomSeconds) + if grace < need { + t.Fatalf("drain=%d: grace %d < required %d -- kubelet would SIGKILL mid-drain", + drain, grace, need) + } + } +} + +func itoa(i int) string { + if i == 0 { + return "0" + } + var b []byte + for i > 0 { + b = append([]byte{byte('0' + i%10)}, b...) + i /= 10 + } + return string(b) +} + +// extraPodSpec templates are passed through verbatim, so --drain-timeout may +// live on the container rather than in ServiceSpec.Args. Reading only the +// latter would miss precisely the deployments that tuned it. +func TestGraceReadsDrainTimeoutFromTheContainerToo(t *testing.T) { + spec := &corev1.PodSpec{Containers: []corev1.Container{{ + Name: "main", + Command: []string{"python3", "-m", "infera.engine.sglang"}, + Args: []string{"--model-path", "/m", "--drain-timeout", "240"}, + }}} + injectWorkerRolloutDefaults(spec, 0, 8080, false, nil) + if spec.TerminationGracePeriodSeconds == nil { + t.Fatal("grace not set") + } + want := int64(workerPreStopDrainSeconds + 240 + workerTeardownHeadroomSeconds) + if *spec.TerminationGracePeriodSeconds != want { + t.Fatalf("grace = %d, want %d", *spec.TerminationGracePeriodSeconds, want) + } +} diff --git a/manual/features/scaling.md b/manual/features/scaling.md index 1ae31cae..d8381185 100644 --- a/manual/features/scaling.md +++ b/manual/features/scaling.md @@ -115,16 +115,59 @@ worker. ### On Kubernetes -Two things beyond `SIGTERM`: - -- The registry drops a Pod as soon as it carries a `deletionTimestamp`, without - waiting for the container to exit. A terminating Pod keeps `phase: Running`, - so without this it would stay a routing candidate for the whole `preStop` - delay — turning a hook meant to make shutdown graceful into extra seconds of - accepting work that is about to be killed. -- `terminationGracePeriodSeconds` must exceed `preStop` + `--drain-timeout`, or - the kubelet `SIGKILL`s mid-drain. The operator sets 120 s with a 15 s preStop; - raise it if your generations are longer. +The recipes deploy with `discoveryBackend: kubernetes` and +`--request-transport http`, so the shutdown path differs from a bare +etcd deployment in two ways — and gains one stage. + +**Discovery is a Pod annotation, not an etcd lease.** Registering writes +`infera.amd.com/worker-info` on the worker's own Pod; deregistering clears it. +The registry additionally drops a Pod the moment it carries a +`deletionTimestamp`, without waiting for the container to exit. That matters +because a terminating Pod keeps `phase: Running` — without the check it would +stay a routing candidate for the whole `preStop` delay, turning a hook meant to +make shutdown graceful into extra seconds of accepting work about to be killed. + +**There is a `preStop` delay before `SIGTERM`.** The operator injects +`sleep 15`, so the full sequence is: + +``` +deletion requested ──► deletionTimestamp set ──► registry drops the worker + ──► preStop sleep 15 (still serving what it has) + ──► SIGTERM ──► DRAINING ──► drain ──► deregister ──► engine.stop() + ──► [kubelet SIGKILL at terminationGracePeriodSeconds] +``` + +### Worst case, and the budget + +Every stage is individually bounded: + +| Stage | Bound | Set by | +|---|---|---| +| `preStop` | 15 s | operator | +| announce `DRAINING` | 10 s | registration HTTP client timeout | +| drain | `--drain-timeout` (default 30 s) | flag | +| deregister | 10 s | registration HTTP client timeout | +| `engine.stop()` | 30 s | `SIGTERM` to the engine's process group, then `SIGKILL` | +| **total** | **≈95 s at defaults** | | + +`terminationGracePeriodSeconds` has to cover that whole sum, because the kubelet +`SIGKILL`s the moment it expires — mid-drain if that is where things are. The +operator now **derives** it as `preStop + --drain-timeout + 50 s` of teardown +headroom, with a 120 s floor, reading the flag from `ServiceSpec.Args` or from +the container directly when an `extraPodSpec` template supplies it. + +```{warning} +This used to be a fixed 120 s with a comment saying it "must exceed preStop + +the worker `--drain-timeout`" — and nothing parsed that flag, so the invariant +was documented and unenforced. Raising `--drain-timeout` for long generations +(the only reason anyone raises it) pushed shutdown past the grace and turned the +drain back into a kill. Measured on a live cluster before the change: a worker +declaring `--drain-timeout 300` still received `terminationGracePeriodSeconds: +120`, i.e. 365 s of budget granted 120. +``` + +If you set the grace period yourself it is respected as long as it is **larger** +than the derived value; it is only ever raised, never lowered. ## PD and DP From 50aa4492df9fc979d28109a8cb76713d3dbf14ec Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Tue, 4 Aug 2026 16:38:40 +0000 Subject: [PATCH 11/88] docs(scaling): measure the Kubernetes drain path on a live cluster The Kubernetes path was the one with the least verification -- the deletionTimestamp check had unit tests and nothing else, while the recipes deploy on exactly that path (kubernetes discovery, HTTP transport). Exercised end to end on the live k3s cluster with no GPU: fake workers registering by Pod annotation, the real infera server watching them with --discovery-backend kubernetes, and a Pod deleted while holding an in-flight generation. kubectl delete pod took the worker out of routing in 93 ms, against the 15 000 ms preStop delay the operator injects. That gap is the entire reason for reading deletionTimestamp: the alternatives -- the DELETE event, or phase leaving Running -- only fire once the container has exited, so without it the router would have kept assigning work for the whole preStop window and then had it killed. The 300-chunk generation already in flight completed in full, and the replacement Pod had registered before the drain finished. Signed-off-by: Zhang, Jiejing --- manual/features/scaling.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/manual/features/scaling.md b/manual/features/scaling.md index d8381185..900ad1e7 100644 --- a/manual/features/scaling.md +++ b/manual/features/scaling.md @@ -166,6 +166,14 @@ declaring `--drain-timeout 300` still received `terminationGracePeriodSeconds: 120`, i.e. 365 s of budget granted 120. ``` +Measured on a live k3s cluster, workers registering by Pod annotation and a real +infera server watching them: `kubectl delete pod` took the worker out of routing +in **93 ms**, against a 15 000 ms `preStop` delay — the whole point of reading +`deletionTimestamp`, since the alternatives (the `DELETE` event, or `phase` +leaving `Running`) only fire after the container has already exited. A +300-chunk generation in flight on that Pod completed in full while it drained, +and its replacement had registered before the drain finished. + If you set the grace period yourself it is respected as long as it is **larger** than the derived value; it is only ever raised, never lowered. From 5c7fe8ce9fe6e729f527e96bec74610d8a06b6d7 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Tue, 4 Aug 2026 16:53:26 +0000 Subject: [PATCH 12/88] docs(scaling): measure the drain with a real engine on the recipe path The Kubernetes path had been measured only with fake workers. This runs it as the recipes actually deploy: an InferaDeployment reconciled by the operator, real SGLang serving Qwen3-8B on two MI355X, Kubernetes discovery by Pod annotation, HTTP transport, with the repo mounted over the image's installed infera so the drain under test is the one running. kubectl delete pod took the worker out of routing in 87 ms against the 15 000 ms preStop delay, and the four concurrent 2500-token generations it was serving all completed with 200 and full-length output (6.5-13.3 kB). The replacement Pod had registered before the drain finished. Also documents something that cost an hour to find and will cost the next person the same: spec.services..resources is silently ignored when extraPodSpec is set, because the template is passed through verbatim. A worker that relies on it schedules and starts and then dies with "No accelerator available", which reads as a driver or device-plugin problem rather than a manifest one. Signed-off-by: Zhang, Jiejing --- manual/features/scaling.md | 30 +++++++++++++++++++++++------- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/manual/features/scaling.md b/manual/features/scaling.md index 900ad1e7..193f537b 100644 --- a/manual/features/scaling.md +++ b/manual/features/scaling.md @@ -166,13 +166,29 @@ declaring `--drain-timeout 300` still received `terminationGracePeriodSeconds: 120`, i.e. 365 s of budget granted 120. ``` -Measured on a live k3s cluster, workers registering by Pod annotation and a real -infera server watching them: `kubectl delete pod` took the worker out of routing -in **93 ms**, against a 15 000 ms `preStop` delay — the whole point of reading -`deletionTimestamp`, since the alternatives (the `DELETE` event, or `phase` -leaving `Running`) only fire after the container has already exited. A -300-chunk generation in flight on that Pod completed in full while it drained, -and its replacement had registered before the drain finished. +Measured on a live k3s cluster: `kubectl delete pod` took the worker out of +routing in **87–93 ms**, against the 15 000 ms `preStop` delay. That gap is the +whole point of reading `deletionTimestamp` — the alternatives (the `DELETE` +event, or `phase` leaving `Running`) only fire once the container has already +exited, so without it the router would keep assigning work for the entire +`preStop` window and then have it killed. + +Two runs, both with a Pod deleted while holding in-flight work: + +- **Real SGLang Qwen3-8B** deployed by the operator (`InferaDeployment`, two + workers, one MI355X each, Kubernetes discovery, HTTP transport): four + concurrent 2500-token generations in flight, **4/4 completed with HTTP 200** + and full-length output (6.5–13.3 kB), replacement Pod registered before the + drain finished. +- **Fake workers**, same path without a GPU: a 300-chunk generation completed + in full across the drain. + +```{note} +`spec.services..resources` is **ignored when `extraPodSpec` is set** — the +template is passed through verbatim, so the GPU request has to live on your own +container. A worker that omits it schedules, starts, and then fails with "No +accelerator available". +``` If you set the grace period yourself it is respected as long as it is **larger** than the derived value; it is only ever raised, never lowered. From 5cac72293a53ac8bbbd30a3f1bdf80f458705337 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Tue, 4 Aug 2026 18:48:05 +0000 Subject: [PATCH 13/88] style: strip trailing whitespace in the GLM-5.2 1P1D example Not from this branch. examples/sglang_1p1d_glm5.2/README.md arrived on main in 7609bb15 (#84) with a trailing-whitespace line, which the pre-commit hook rejects -- so the lint job fails on the merge result for every open PR, including this one, while each branch passes on its own. Fixed here only because it blocks this PR's CI. One line, whitespace only, no content change. Signed-off-by: Zhang, Jiejing --- examples/sglang_1p1d_glm5.2/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/sglang_1p1d_glm5.2/README.md b/examples/sglang_1p1d_glm5.2/README.md index 620d4687..f438b4ef 100644 --- a/examples/sglang_1p1d_glm5.2/README.md +++ b/examples/sglang_1p1d_glm5.2/README.md @@ -37,7 +37,7 @@ on two different fabrics, is in [`results/`](results/README.md). node-P node-D ├─ etcd ◀── discovery ──▶ ├─ infera router :8100 ◀── clients -├─ kvd daemon +├─ kvd daemon └─ prefill leg :30000 ══ KV over RDMA ══▶ decode leg :30001 TP8, DPA off, TP8, DPA on (dp8), kvd L2/L3 MTP (EAGLE) From 7289dc14fe9bcc20393127785a2490a7c80f6cc9 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Tue, 4 Aug 2026 19:12:34 +0000 Subject: [PATCH 14/88] fix(router): name the empty pool when half a PD deployment is left Scaling either PD pool to zero fails closed, which is right. The diagnostic was not: AutoRouter falls through to the mixed router when only one pool is populated, and that router answers "no active mixed worker for model=..." -- pointing the reader at something they never deployed while the surviving decode (or prefill) pool sits right there. Found by testing PD scale-down, which the docs had listed as unverified. Now, when exactly one PD pool has workers and there are no mixed workers to absorb the traffic, the 503 says which pool is empty and how many workers the other one has, and the same goes to the log. A mixed worker alongside half a PD fleet still routes normally -- that is the rolling upgrade case and must not trip this. Also records the PD scaling measurement in the docs: a 1P1D fake fleet grown to 2P2D and shrunk back under continuous traffic, 200 requests and 0 failures, both pools scaling independently. The warning block now distinguishes what was measured with fakes from what still needs a real engine, since no KV moved in that run. Signed-off-by: Zhang, Jiejing --- infera/router/auto.py | 40 +++++++++++++++++-- manual/features/scaling.md | 16 ++++++-- tests/unit/router/test_failover.py | 62 +++++++++++++++++++++++++++++- 3 files changed, 110 insertions(+), 8 deletions(-) diff --git a/infera/router/auto.py b/infera/router/auto.py index 0f840c10..0db4a494 100644 --- a/infera/router/auto.py +++ b/infera/router/auto.py @@ -5,20 +5,28 @@ ############################################################################### from __future__ import annotations +import logging + from fastapi import Response +from fastapi.responses import JSONResponse from infera.common.worker_pool import DisaggMode from infera.router.base import BaseRouter from infera.router.disagg import DisaggRouter from infera.router.mixed import MixedRouter +logger = logging.getLogger(__name__) + class AutoRouter(BaseRouter): """Per-request router selector. - Selection policy (v0.1: PD-preferred with mixed fallback): - - If the model has BOTH prefill and decode workers → DisaggRouter - - Otherwise (only mixed, partial PD, or empty) → MixedRouter + Selection policy (PD-preferred with mixed fallback): + - BOTH prefill and decode workers → DisaggRouter + - Exactly one PD pool, and no mixed workers → 503 naming the + empty pool (half a PD deployment cannot serve, and saying "no mixed + worker" would describe something the operator never deployed) + - Otherwise (mixed workers present, or nothing at all) → MixedRouter (MixedRouter itself returns 503 if no mixed worker is available) This supports mixed deployments (some models PD, others mixed) and rolling @@ -66,4 +74,30 @@ async def dispatch( has_d = self.pool.list_active(model=model, mode=DisaggMode.DECODE) if has_p and has_d: return await self._disagg.dispatch(body, stream=stream, path=path) + # Exactly one PD pool populated: the deployment is disaggregated but + # half of it is gone. Falling through to the mixed router would be + # correct-but-useless -- there are no mixed workers either, so it + # answers "no active mixed worker", which sends the reader looking for + # something they never deployed while a decode (or prefill) pool sits + # right there. Scaling either side to zero is the usual cause. + if bool(has_p) != bool(has_d) and not self.pool.list_active( + model=model, mode=DisaggMode.MIXED + ): + present, missing = ("prefill", "decode") if has_p else ("decode", "prefill") + logger.warning( + "model=%r has %d %s worker(s) but no %s worker: PD dispatch needs both", + model, + len(has_p or has_d), + present, + missing, + ) + return JSONResponse( + content={ + "error": ( + f"model={model!r} has {len(has_p or has_d)} {present} worker(s) " + f"but no {missing} worker; PD dispatch requires both pools" + ) + }, + status_code=503, + ) return await self._mixed.dispatch(body, stream=stream, path=path) diff --git a/manual/features/scaling.md b/manual/features/scaling.md index 193f537b..f6966453 100644 --- a/manual/features/scaling.md +++ b/manual/features/scaling.md @@ -200,7 +200,9 @@ they scale **independently** — add prefill for longer inputs, decode for more concurrent users. Two constraints: - **Neither pool can go to zero.** PD dispatch fails closed when either side is - empty. `minReplicas: 0` on either is an outage, not an idle saving. + empty — `minReplicas: 0` on either is an outage, not an idle saving. The 503 + names the empty pool (`has 1 decode worker(s) but no prefill worker`), so the + cause is visible without reading the fleet. - **A DP worker's shape decides who picks the rank.** A worker registering `dp_size > 1` with **no** `dp_rank` is rank-multiplexed: the router fans it out into one target per rank and pins `X-Data-Parallel-Rank`. A worker that @@ -280,10 +282,16 @@ drain: engine idle for 6s, 1 request(s) completed deregistered worker 127.0.0.1:20001 (lease revoked) ``` +**PD scaling, measured.** A 1P1D fake fleet grown to 2P2D and shrunk back under +continuous traffic: **200 requests, 0 failures**, both pools scaling +independently and the drained workers finishing their in-flight work. Taking the +last prefill away then returns 503 naming the empty pool. + ```{warning} -**Not measured:** multi-node workers, TP > 1, PD under scaling, and scale-down -during an active KV transfer. The PD handoff queues are counted in the drain, -but that path has not been exercised on hardware. +**Not measured:** multi-node workers, TP > 1, PD scaling with a *real* engine +(the run above used fake workers, so no KV moved), and scale-down during an +active KV transfer. The PD handoff queues are counted in the drain, but that +path has not been exercised on hardware. ``` ## Autoscaling diff --git a/tests/unit/router/test_failover.py b/tests/unit/router/test_failover.py index 5027c047..65a2b99f 100644 --- a/tests/unit/router/test_failover.py +++ b/tests/unit/router/test_failover.py @@ -14,7 +14,7 @@ import pytest from infera.common.nats_request import TYPE_DATA, TYPE_DONE, TYPE_ERROR -from infera.common.worker_pool import EngineType, WorkerInfo +from infera.common.worker_pool import DisaggMode, EngineType, WorkerInfo from infera.router.mixed import MixedRouter from infera.router.policy.target import RouteTarget @@ -343,3 +343,63 @@ def handler(request: httpx.Request) -> httpx.Response: assert (await r.dispatch({"model": "m"}, stream=False)).status_code == 200 assert r.breaker.state_of("bad").value == "open" await r.aclose() + + +# --- half a PD deployment must say so -------------------------------------- + + +class _ModePool: + """Pool that can answer per-disagg-mode, unlike _FakePool.""" + + def __init__(self, workers): + self._w = workers + + def list_active(self, model=None, mode=None): + return [w for w in self._w if mode is None or w.disagg_mode == mode] + + +def _pd_worker(wid, mode): + return WorkerInfo( + worker_id=wid, + url=f"http://{wid}", + model_name="m", + engine=EngineType.SGLANG, + disagg_mode=mode, + request_transport="http", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "present,missing", + [(DisaggMode.PREFILL, "decode"), (DisaggMode.DECODE, "prefill")], +) +async def test_half_a_pd_deployment_names_the_empty_pool(present, missing): + """Scaling either PD pool to zero fails closed -- correctly -- but used to + report "no active mixed worker", which points at something the operator + never deployed while the surviving pool sits right there.""" + from infera.router.auto import AutoRouter + + r = AutoRouter(_ModePool([_pd_worker("w1", present)]), _FakePolicy()) + resp = await r.dispatch({"model": "m"}, stream=False) + assert resp.status_code == 503 + body = json.loads(bytes(resp.body))["error"] + assert missing in body and "PD dispatch requires both pools" in body, body + assert "mixed" not in body, f"must not blame mixed workers: {body}" + await r.aclose() + + +@pytest.mark.asyncio +async def test_a_mixed_worker_still_absorbs_a_half_pd_fleet(): + """A mixed worker alongside half a PD pool can serve, so the 503 must not + fire -- this is the rolling-upgrade case.""" + from infera.router.auto import AutoRouter + + pool = _ModePool([_pd_worker("p", DisaggMode.PREFILL), _pd_worker("m1", DisaggMode.MIXED)]) + r = AutoRouter(pool, _FakePolicy(), nats_client=None, request_max_retries=0) + r._mixed._client = httpx.AsyncClient( + transport=httpx.MockTransport(lambda req: httpx.Response(200, json={"id": "ok"})) + ) + resp = await r.dispatch({"model": "m"}, stream=False) + assert resp.status_code == 200 + await r.aclose() From 81ff397f48f5fe7190cae5c59ab4188d8f813719 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Wed, 5 Aug 2026 04:25:52 +0000 Subject: [PATCH 15/88] feat(tools): fake worker supports the NATS request transport The fake could only be reached over direct HTTP, so the NATS transport -- the path with the better drain semantics, and the one a colleague correctly flagged as important -- was untestable without a real engine. --request-transport nats starts the real NatsRequestServer, which proxies to this process's own HTTP surface exactly as it proxies to a real engine's. The transport under test is the production one; only what answers at the far end is fake. Shutdown likewise goes through the real stop(drain=True), so the drain being exercised is the shipped one. Measured against a real NATS 2.10 broker and a real router: two workers registered with request_transport=nats, unary and SSE both flowed through the broker (12 in / 10 out messages, 3 connections), and SIGTERM during a 300-chunk generation drained cleanly -- the log reads "draining 1 in-flight NATS request(s)", the generation completed 300/300, and the worker deregistered 21.3s later, which is just the remaining generation time with no polling overhead. The README says plainly that the fake's *HTTP* drain is not representative: this process serves its own requests, so it knows its in-flight count exactly, while a real worker on HTTP has to poll the engine's lagging /metrics behind a settle window. Comparing the fake's two transports therefore measures nothing -- both are exact. That difference only appears with a real engine, and is already measured there. Signed-off-by: Zhang, Jiejing --- infera/tools/fakeworker/README.md | 32 +++++++++++++++++++++++++++++++ infera/tools/fakeworker/server.py | 28 ++++++++++++++++++++++++++- 2 files changed, 59 insertions(+), 1 deletion(-) diff --git a/infera/tools/fakeworker/README.md b/infera/tools/fakeworker/README.md index 85d9d1b0..58bfc51f 100644 --- a/infera/tools/fakeworker/README.md +++ b/infera/tools/fakeworker/README.md @@ -83,6 +83,38 @@ e.g. that `bootstrap_room % dp_size == disagg_prefill_dp_rank` holds on every request, which is the invariant SGLang's `follow_bootstrap_room` balancer enforces with a `KVTransferError`. +## NATS transport + +`--request-transport nats` routes requests through a broker instead of having +the router dial this worker. It uses the **real** `NatsRequestServer`, which +proxies to this process's own HTTP surface exactly as it proxies to a real +engine's — so the transport under test is the production one, not a stand-in. + +```bash +infera-fake-worker --model-name m --port 9101 \ + --request-transport nats --nats-server nats://127.0.0.1:4222 \ + --discovery-backend etcd --etcd-endpoint http://127.0.0.1:2379 +``` + +Shutdown then goes through the real NATS drain — unsubscribe first, then wait on +the in-flight set infera actually holds: + +``` +worker 127.0.0.1:19951 announced DRAINING +draining 1 in-flight NATS request(s), up to 60s +deregistered worker 127.0.0.1:19951 +``` + +```{note} +**The fake's HTTP drain is not representative of a real worker's.** This process +serves the requests itself, so it knows its own in-flight count exactly. A real +worker on HTTP transport does not: the router talks straight to the engine, so +infera has to poll the engine's `/metrics` and wait out a settle window because +those gauges lag. Comparing the fake's HTTP drain against its NATS drain +therefore measures nothing — both are exact. The difference only shows up with a +real engine. +``` + ## Limits — read these before drawing conclusions **No KV transfer is simulated.** A `--disagg-mode prefill` / `decode` fake takes diff --git a/infera/tools/fakeworker/server.py b/infera/tools/fakeworker/server.py index 77f5c98d..d896ff7d 100644 --- a/infera/tools/fakeworker/server.py +++ b/infera/tools/fakeworker/server.py @@ -306,7 +306,7 @@ def build_config(args) -> EngineConfig: kv_block_size=args.kv_block_size if args.kv else None, dp_rank=args.dp_rank, dp_size=args.dp_size, - request_transport="http", + request_transport=args.request_transport, ) @@ -378,6 +378,16 @@ def parse_args(argv=None): help="refuse the first N requests with 503, then recover -- exercises " "the router's circuit breaker and its half-open probe.", ) + p.add_argument( + "--request-transport", + default="http", + choices=["http", "nats"], + help="nats routes requests through a broker instead of the router " + "dialling this worker directly. Uses the real NatsRequestServer, which " + "proxies to this process's own HTTP surface exactly as it proxies to a " + "real engine -- so the transport under test is the production one.", + ) + p.add_argument("--nats-server", default=os.environ.get("NATS_SERVER")) p.add_argument("--drain-timeout", type=float, default=30.0) return p.parse_args(argv) @@ -415,6 +425,16 @@ async def _serve(args) -> None: serve_task.cancel() raise SystemExit(f"failed to bind {args.host}:{args.port} -- not registering") + nats_req_server = None + if args.request_transport == "nats": + from infera.common.nats_request import NatsRequestServer + + nats_req_server = NatsRequestServer( + f"{cfg.host}:{cfg.port}", args.port, url=args.nats_server + ) + await nats_req_server.start() + logger.info("nats request consumer started for %s:%s", cfg.host, cfg.port) + if args.startup_delay_s > 0: logger.info("simulating weight load for %.0fs", args.startup_delay_s) await asyncio.sleep(args.startup_delay_s) @@ -457,6 +477,12 @@ async def _shutdown() -> None: await reg.announce_draining() except Exception as exc: # noqa: BLE001 - shutdown must not raise logger.warning("announce_draining failed: %s", exc) + if nats_req_server is not None: + # The real drain: unsubscribe first so nothing new arrives, then + # wait on the in-flight set infera actually holds. No polling and no + # settle window -- unlike HTTP, where the count has to be inferred + # from the engine's lagging gauges. + await nats_req_server.stop(drain=True, drain_timeout=args.drain_timeout) deadline = time.monotonic() + args.drain_timeout while state.running and time.monotonic() < deadline: await asyncio.sleep(0.1) From fc954a57b194354846d735f54df2aa98ebd1942d Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Wed, 5 Aug 2026 05:03:46 +0000 Subject: [PATCH 16/88] docs(scaling): the transport decides how good the drain can be Adds a section measuring what NATS buys over HTTP for scale-down, after a review comment made the point that dropping NATS would be a loss. The measurements support that, and for a more specific reason than implementation quality: the two transports differ in *where the information lives*. On NATS infera owns the request path and holds the in-flight set, so the drain is exact. On HTTP the router dials the engine directly and never sees the request, so infera has to infer the count from the engine's gauges -- which lag, hence the settle window. Measured with the same fake worker and the same generation: * NATS, one in-flight generation: "draining 1 in-flight NATS request(s)" -- it knows the count -- 300/300 chunks completed, deregistered 21.3s later, which is just the remaining generation time. * NATS, nothing in flight: announce to deregister in 3 ms. * HTTP with a real engine, nothing in flight: at least the 6 s settle window, because one zero reading cannot be told from a stale gauge. Also notes what this costs and what else it buys -- a broker, and request cancellation the HTTP path lacks -- and that the Rust router does not implement the NATS transport, so the Rust data plane and the NATS drain are currently an either/or. That tradeoff is worth being explicit about rather than discovered later. Signed-off-by: Zhang, Jiejing --- manual/features/scaling.md | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/manual/features/scaling.md b/manual/features/scaling.md index f6966453..24d1b7ff 100644 --- a/manual/features/scaling.md +++ b/manual/features/scaling.md @@ -85,6 +85,39 @@ exits in about six seconds regardless. Set it above your p99 generation time. Anything still running when it expires is cut, with a warning naming the count. ``` +### The transport decides how well this works + +Draining is only as good as the router's view of what is in flight, and that +differs by transport — not by implementation quality, but by where the +information lives. + +| | who knows what is in flight | drain | +|---|---|---| +| **NATS** (`--request-transport nats`) | infera — it owns the request path and holds the in-flight set | exact, no polling | +| **HTTP** (default in the recipes) | only the engine — the router dials it directly and never sees the request | poll the engine's `/metrics`, behind a settle window | + +Measured, same fake worker, same generation: + +- **NATS, one in-flight generation**: the log reads `draining 1 in-flight NATS + request(s)` — it knows the count — the 300-chunk generation completed in full, + and the worker deregistered **21.3 s** later, which is just the remaining + generation time with no overhead. +- **NATS, nothing in flight**: announce → deregister in **3 ms**. +- **HTTP with a real engine, nothing in flight**: at least the **6 s** settle + window, because a single zero reading cannot be told apart from a gauge that + has not refreshed yet. + +So NATS costs a broker and buys a drain that is exact rather than inferred. It +also buys request cancellation the HTTP path does not have — a timeout or client +disconnect publishes to `infera.cancel.` and the worker tears down the +engine connection, instead of leaving it generating. + +```{note} +The Rust router does not implement the NATS transport (`lib.rs`: "Configs +outside this set (NATS transport, ...) are served by the Python backend"), so +the Rust data plane and the NATS drain are currently an either/or. +``` + ### Why in-flight work is visible at all The engine's own gauges are the only source of truth on the HTTP path, and they From b89b1fd5951c8f6bc5ee1e9aad51b272539b7926 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Wed, 5 Aug 2026 05:18:36 +0000 Subject: [PATCH 17/88] fix(kv): pull a newly registered snapshot target immediately Under the ZMQ event transport the router rebuilds each worker's cache view by pulling its /v1/kv-snapshot, on a 30 s SnapshotReconciler loop. register_target only added the target to the dict; the loop was sitting in wait_for(kick, timeout=interval_s), so a worker that joined while the reconciler was running went unpulled for up to a full interval. The one code path that does kick, trigger_gap_recovery, has no caller in production -- only a test. So nothing shortened that wait. For a genuinely new worker this is harmless: its cache is empty, so an empty routing view is accurate. It is not harmless on a router restart or a rolling upgrade, which is where it actually bites. Every existing worker arrives through this same path with a cache that is warm, and until its snapshot lands kv-aware routing scores all of them as holding nothing -- so the first 30 s after a router comes back routes as if the fleet had no cache at all. The NATS event transport does not have this gap: NatsKvEventClient bootstraps from a JetStream KV bucket watchall, which pushes all initial values on subscribe. Found while measuring what NATS buys over ZMQ, after a review comment argued the NATS path matters. Two tests: a target registered while the loop is running is pulled within 0.3 s (fails on the old code, waiting out the interval), and re-registering a known target does not pull again -- registration is re-asserted routinely, since the Kubernetes backend rewrites its Pod annotation every 30 s and etcd redelivers on relist, so kicking on each would turn a self-heal into a stampede proportional to fleet size. Signed-off-by: Zhang, Jiejing --- infera/kv/snapshot.py | 18 +++++- tests/unit/kv/test_snapshot.py | 115 ++++++++++++++++++++++++++++++++- 2 files changed, 130 insertions(+), 3 deletions(-) diff --git a/infera/kv/snapshot.py b/infera/kv/snapshot.py index 16301c74..5d0c71d7 100644 --- a/infera/kv/snapshot.py +++ b/infera/kv/snapshot.py @@ -243,9 +243,25 @@ def register_target( model: str, compat_key: str, ) -> None: - """Tell the reconciler to periodically pull this (publisher, tree).""" + """Tell the reconciler to periodically pull this (publisher, tree). + + Also pulls it now rather than on the next tick. The loop waits out the + full interval between sweeps, so without this a worker that joins while + the reconciler is running is invisible to kv-aware routing for up to + ``interval_s`` -- 30 s in production. + + For a genuinely new worker that would be harmless, since its cache is + empty and an empty view is accurate. It is not harmless on a router + restart or a rolling upgrade: every existing worker arrives through this + same path with a warm cache, and until its snapshot lands the policy + scores them all as holding nothing. + """ key = (publisher_id, endpoint, model, compat_key) + new = key not in self._targets self._targets[key] = None + if new: + self._urgent.add(key) + self._kick.set() def unregister_target( self, diff --git a/tests/unit/kv/test_snapshot.py b/tests/unit/kv/test_snapshot.py index 3eef1970..9413f509 100644 --- a/tests/unit/kv/test_snapshot.py +++ b/tests/unit/kv/test_snapshot.py @@ -289,12 +289,14 @@ def test_producer_empty_snapshot_for_unknown_stream() -> None: # ---------------------------------------------------------------------- -async def _make_reconciler(pull_fn) -> tuple[KVIndex, KvIndexWriter, SnapshotReconciler]: +async def _make_reconciler( + pull_fn, interval_s: float = 10_000 +) -> tuple[KVIndex, KvIndexWriter, SnapshotReconciler]: index = KVIndex() queue: asyncio.Queue = asyncio.Queue(maxsize=100) writer = KvIndexWriter(index=index, queue=queue) await writer.start() - rec = SnapshotReconciler(index=index, writer=writer, pull_fn=pull_fn, interval_s=10_000) + rec = SnapshotReconciler(index=index, writer=writer, pull_fn=pull_fn, interval_s=interval_s) return index, writer, rec @@ -685,3 +687,112 @@ def test_index_drop_tree() -> None: matches_b = index.find_matches(model="m", compat_key="ckB", chain=[chain[0]], candidates=["w1"]) assert matches_a["w1"] == OverlapBlocks() assert matches_b["w1"].device == 1 + + +async def test_target_registered_while_running_is_not_left_for_a_whole_interval() -> None: + """A worker that joins after the reconciler is already running must have its + snapshot pulled promptly, not on the next periodic tick. + + This is the router-restart / rolling-upgrade case, and it is where the delay + actually costs something. A *newly started* worker has an empty cache, so an + empty routing view of it is correct. But when the router restarts, every + existing worker arrives through the same path with a cache that is genuinely + warm -- and until its snapshot lands, kv-aware routing scores all of them as + holding nothing. + + The loop sits in `wait_for(self._kick.wait(), timeout=interval_s)`, so a + registration that does not set the kick waits out the full interval (30 s in + production). + """ + chain = hash_token_blocks(list(range(4)), block_size=4) + snap = Snapshot( + publisher_id="w-late", + publisher_type="worker", + model_name="m", + compat_key="ck", + index_block_size=4, + batch_id=0, + blocks=( + SnapshotBlock( + sequence_hash=chain[0].sequence_hash, + parent_sequence_hash=None, + block_hash=chain[0].block_hash, + tiers=("device",), + ), + ), + ) + pulled: list[str] = [] + + async def pull_fn(publisher_id, *args, **kwargs): + pulled.append(publisher_id) + return snap + + # A long interval stands in for production's 30 s: if registration relies on + # the periodic tick, this test waits for it and fails on the assertion. + index, writer, rec = await _make_reconciler(pull_fn, interval_s=5.0) + try: + await rec.start() + await asyncio.sleep(0.05) + pulled.clear() + + rec.register_target(publisher_id="w-late", endpoint="ignored", model="m", compat_key="ck") + await asyncio.sleep(0.3) + assert "w-late" in pulled, ( + "a worker registered while the reconciler was running was not pulled " + "within 0.3s; it is waiting out the periodic interval" + ) + finally: + await rec.stop() + await writer.stop() + + +async def test_reregistering_a_known_target_does_not_re_pull() -> None: + """Registration is re-asserted routinely -- the Kubernetes backend rewrites + its Pod annotation every 30 s, and etcd watch redelivers on relist. Kicking + on every one of those would turn a self-heal into a snapshot stampede + proportional to fleet size. + """ + chain = hash_token_blocks(list(range(4)), block_size=4) + snap = Snapshot( + publisher_id="w1", + publisher_type="worker", + model_name="m", + compat_key="ck", + index_block_size=4, + batch_id=0, + blocks=( + SnapshotBlock( + sequence_hash=chain[0].sequence_hash, + parent_sequence_hash=None, + block_hash=chain[0].block_hash, + tiers=("device",), + ), + ), + ) + pulls = 0 + + async def pull_fn(*args, **kwargs): + nonlocal pulls + pulls += 1 + return snap + + index, writer, rec = await _make_reconciler(pull_fn, interval_s=5.0) + try: + await rec.start() + await asyncio.sleep(0.05) + kw = dict(publisher_id="w1", endpoint="ignored", model="m", compat_key="ck") + rec.register_target(**kw) + await asyncio.sleep(0.2) + after_first = pulls + assert after_first > 0, "the first registration must pull" + + for _ in range(5): + rec.register_target(**kw) + await asyncio.sleep(0.2) + assert pulls == after_first, ( + f"re-registering pulled {pulls - after_first} more time(s); " + "only a new target should kick" + ) + finally: + await rec.stop() + await writer.stop() From f2d7f7a8fcc5a5351b7bfe596d9d288d68b03365 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Wed, 5 Aug 2026 05:32:28 +0000 Subject: [PATCH 18/88] docs(scaling): measure the NATS JetStream admission throttle Documents what INFERA_NATS_REQ_MAX_PENDING actually does, after verifying it end to end against a real NATS 2.10 broker with JetStream. It matters for scaling because it covers the window scaling cannot: a burst shorter than a 140 s cold start cannot be answered by adding workers, so the only choice is between queueing behind a saturated worker and steering away from one. The section leads with how to observe it, because the obvious signal is the wrong one and cost me two failed measurements. A refusal raises the same retryable failure as any other pre-first-byte error, so the request fails over to a freer worker and the client sees 200 -- a 429 only reaches the client when every worker is over the limit and retries are exhausted. Looking for 429s finds nothing and reads as "the throttle does not work". The distribution is the real signal. With one deliberately saturated worker (concurrency 1, 300 ms/token) and one fast one at limit 3, twenty requests sent under backlog went +0 / +20 where round-robin would have been +10 / +10. The saturated consumer read num_ack_pending = 10 at the time; the ack is deliberately after the request is fully proxied so the gauge reflects in-flight work rather than delivery. Also notes the limit of the mechanism: the check is per dispatch, so a simultaneous burst is admitted in full -- every admission check runs before any of them has built backlog. Signed-off-by: Zhang, Jiejing --- manual/features/scaling.md | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/manual/features/scaling.md b/manual/features/scaling.md index 24d1b7ff..23dd9b01 100644 --- a/manual/features/scaling.md +++ b/manual/features/scaling.md @@ -112,6 +112,38 @@ also buys request cancellation the HTTP path does not have — a timeout or clie disconnect publishes to `infera.cancel.` and the worker tears down the engine connection, instead of leaving it generating. +### Admission control + +Setting `INFERA_NATS_REQ_MAX_PENDING` (or `--nats-req-max-pending`) above zero +on **both** the server and the workers makes the request path JetStream-backed: +a WorkQueue stream with one durable consumer per worker. The router reads that +consumer's backlog before dispatching and refuses a worker over the limit. + +This matters for scaling because it covers the window scaling cannot: a burst +shorter than a 140 s cold start cannot be answered by adding workers, so the +choice is between queueing behind a saturated worker and steering away from it. + +**Look at the distribution, not the status codes.** A refusal raises the same +retryable failure as any other pre-first-byte error, so the request fails over +to a freer worker and the client sees `200`. Only when every worker is over the +limit and retries are exhausted does a `429` reach the client. Measured with one +deliberately saturated worker (concurrency 1) and one fast one, limit 3: + +| | saturated worker | fast worker | +|---|---|---| +| 20 requests under backlog | **+0** | **+20** | +| round-robin without the throttle | +10 | +10 | + +The worker's consumer showed `num_ack_pending = 10` against a limit of 3 at the +time — the ack happens after the request is fully proxied, precisely so the +backlog gauge reflects genuinely in-flight work rather than mere delivery. + +```{note} +The check is per dispatch, so it steers *new* requests. Requests already +dispatched are unaffected, and a simultaneous burst is all admitted — every +admission check runs before any of them has built backlog. +``` + ```{note} The Rust router does not implement the NATS transport (`lib.rs`: "Configs outside this set (NATS transport, ...) are served by the Python backend"), so From cf6b4488bae33121981b29fa605fc77fbf4bd9d2 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Wed, 5 Aug 2026 05:48:23 +0000 Subject: [PATCH 19/88] docs(readme): say how scaling works, with the measured numbers The README described discovery ("workers register into etcd and heartbeat") but said nothing about what happens when the fleet changes size -- which is the question anyone evaluating this for Kubernetes asks first. Adds a Scaling paragraph to the Kubernetes section and a key-features bullet. The design is worth stating plainly because it is unusual: there is no scaling controller. Workers self-register when ready and deregister when they drain, and the router routes to whatever is registered at that instant, so scaling is just starting and stopping workers. Leads with the asymmetry rather than burying it, because it constrains everything built on top: scale up is 140 s (weight loading), scale down stops receiving in under a second, and adding plus removing instances under continuous traffic cost 0 failed requests out of 260. The consequence -- a burst shorter than a cold start cannot be answered by adding workers -- is the thing a reader most needs to know before designing around this, so it is said outright. Also states that infera ships no autoscaler, rather than leaving it to be inferred, and links the manual page for what is and is not in place for one. Links point at the published docs site, not repo-relative paths, which 404 there. Signed-off-by: Zhang, Jiejing --- README.md | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 51a97b55..081baa9f 100644 --- a/README.md +++ b/README.md @@ -49,7 +49,8 @@ Around them: - **Multi-engine** — run vLLM, SGLang, or ATOM behind one common serving interface. - **OpenAI- and Anthropic-compatible API** — `/v1/chat/completions`, `/v1/completions`, and `/v1/messages` (Anthropic Messages, translated in-process). -- **Self-registering fleet** — workers register into etcd and heartbeat, so the router works from a live view and never routes to a worker that is gone; run any number of stateless server replicas. +- **Self-registering fleet** — workers register into etcd (or their own Pod annotation on Kubernetes) and heartbeat, so the router works from a live view and never routes to a worker that is gone; run any number of stateless server replicas. +- **Scale without dropping requests** — a worker joins when it is ready and leaves by draining: it announces `DRAINING`, finishes the generations it already accepted, and only then deregisters. Measured on MI355X: a worker stops receiving new work **under a second** after `SIGTERM` while its in-flight 4000-token generations all complete, and adding or removing instances under continuous traffic costs **zero failed requests**. See [Scaling a fleet](https://rocm.docs.amd.com/projects/infera/en/latest/features/scaling.html). - **Kubernetes-native** — an operator reconciles an `InferaDeployment` CRD (aggregated / PD / multi-node), with an optional Gateway API (GAIE) endpoint picker. ## Architecture @@ -146,6 +147,24 @@ kubectl apply -f examples/k8s-deployments/single-node-aggregated.yaml Ready-to-fill deployment templates (single-node, prefill/decode, multi-node TP, GAIE) and their placeholders are in [`examples/k8s-deployments/`](examples/k8s-deployments/README.md). +**Scaling.** There is no scaling controller — workers self-register when ready and deregister when +they drain, and the router routes to whatever is registered at that instant. Scaling is therefore +`kubectl scale` on the CR, or starting and stopping workers; nothing has to be told about it. + +The two directions cost very different things, and it shapes everything built on top: + +| | measured | +|---|---| +| scale up: `docker run` → serving | **140 s**, almost all of it weight loading | +| scale down: `SIGTERM` → stops receiving | **< 1 s** | +| scale down: in-flight generations | run to completion, bounded by `--drain-timeout` | +| adding + removing under traffic | **260 requests, 0 failures** | + +Because a cold start is minutes and the control loop is seconds, **a burst shorter than a cold +start cannot be answered by adding workers** — keep headroom, or steer traffic to instances that +are already running. Infera does not ship an autoscaler; [Scaling a +fleet](https://rocm.docs.amd.com/projects/infera/en/latest/features/scaling.html) documents what is in place for one, and what is not. + ## Engine images Prebuilt images are published to the `rocm/infera` repository on From 6d5de0c038ad3e71b974ba5a5355245f93d9bb09 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Wed, 5 Aug 2026 06:25:42 +0000 Subject: [PATCH 20/88] feat(operator): standard Kubernetes /scale surface via a scaling adapter Nothing could drive scaling from outside. An InferaDeployment exposes only status; the child Deployment and LeaderWorkerSet expose /scale natively, but the reconciler assigns their whole .Spec every pass and Owns() them, so an external write is reverted almost immediately -- measured, a kubectl scale to 3 went back to 1 in under 3 seconds, not on the 15s resync. Every HPA, KEDA ScaledObject and custom planner loses that race. An InferaDeployment also cannot carry /scale itself, and that is structural rather than unfinished: spec.services is a map with user-chosen keys, while the scale subresource requires specReplicasPath to be a static dot-notation JSONPath. There is no way to name the replicas of an arbitrary map entry. Dynamo hit the identical wall with spec.components[] and solved it the same way. Adds InferaScalingAdapter -- one object per scalable service, carrying specpath=.spec.replicas, statuspath=.status.replicas, selectorpath=.status.selector. kubectl scale, HPA, KEDA and a planner all work through it with no per-tool support in the operator. Single writer by construction: while an adapter has spec.replicas set, the InferaDeployment reconciler reads it instead of the CR's own count. The adapter controller never touches the workload -- it owns only the half /scale reads back. An adapter without spec.replicas is deliberately inert, so an autoscaler can be attached and observed before it is trusted. Also fixes status.replicas on both the adapter and ServiceStatus to report the observed count rather than echoing the desired one. An autoscaler computes desired = ceil(current * metric/target); with current equal to what it just asked for it cannot tell a scale-up has not landed, and keeps multiplying through a 140s model load. Verified on the live k3s cluster: the API now serves inferascalingadapters/scale, kubectl scale works, GET /scale returns a proper autoscaling/v1 Scale, and an HPA targeting it reports AbleToScale=True / SucceededGetScale. With status.selector empty the HPA says InvalidSelector and refuses -- populating it clears that and the only remaining error is a missing metrics source, i.e. the scale interface itself is fully accepted. Go build, vet, gofmt and tests clean in a golang:1.25 container. The four pre-existing gofmt offenders on main are untouched. Signed-off-by: Zhang, Jiejing --- .../v1alpha1/inferascalingadapter_types.go | 115 +++++++++ .../api/v1alpha1/zz_generated.deepcopy.go | 161 ++++++++++--- deploy/operator/cmd/main.go | 8 + .../infera.amd.com_inferascalingadapters.yaml | 207 ++++++++++++++++ deploy/operator/config/rbac/role.yaml | 45 ++-- .../infera.amd.com_inferascalingadapters.yaml | 207 ++++++++++++++++ .../helm/infera-operator/templates/rbac.yaml | 13 ++ .../operator/internal/controller/builders.go | 23 +- .../internal/controller/builders_test.go | 30 +++ .../controller/inferadeployment_controller.go | 62 ++++- .../inferascalingadapter_controller.go | 221 ++++++++++++++++++ manual/features/scaling.md | 78 +++++-- 12 files changed, 1096 insertions(+), 74 deletions(-) create mode 100644 deploy/operator/api/v1alpha1/inferascalingadapter_types.go create mode 100644 deploy/operator/config/crd/bases/infera.amd.com_inferascalingadapters.yaml create mode 100644 deploy/operator/helm/infera-operator/crds/infera.amd.com_inferascalingadapters.yaml create mode 100644 deploy/operator/internal/controller/inferascalingadapter_controller.go diff --git a/deploy/operator/api/v1alpha1/inferascalingadapter_types.go b/deploy/operator/api/v1alpha1/inferascalingadapter_types.go new file mode 100644 index 00000000..4e9ab7a7 --- /dev/null +++ b/deploy/operator/api/v1alpha1/inferascalingadapter_types.go @@ -0,0 +1,115 @@ +/* +Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + +SPDX-License-Identifier: MIT +*/ + +package v1alpha1 + +import ( + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" +) + +// InferaScalingAdapterSpec points at one service inside an InferaDeployment and +// owns its replica count. +type InferaScalingAdapterSpec struct { + // DeploymentRef is the InferaDeployment to scale, in this namespace. + // +kubebuilder:validation:MinLength=1 + DeploymentRef string `json:"deploymentRef"` + + // ServiceName is the key in that deployment's `spec.services` map. + // +kubebuilder:validation:MinLength=1 + ServiceName string `json:"serviceName"` + + // Replicas is the desired count, and the field `/scale` writes. + // + // Left unset the adapter is inert: the InferaDeployment's own + // `spec.services[].replicas` still applies. That makes adding an + // adapter a safe no-op until something actually scales, so an autoscaler + // can be attached and observed before it is trusted. + // +optional + // +kubebuilder:validation:Minimum=0 + Replicas *int32 `json:"replicas,omitempty"` +} + +// InferaScalingAdapterStatus carries what `/scale` reads back. +type InferaScalingAdapterStatus struct { + // Replicas is the count observed on the workload -- not the desired count + // echoed back. + // + // The distinction is the whole reason this field exists. HorizontalPodAutoscaler + // computes `desired = ceil(current * metric/target)`; if `current` is really + // the desired value it never lags reality, so during a multi-minute model load + // the autoscaler cannot tell that a scale-up has not landed yet and keeps + // multiplying. + // +optional + Replicas int32 `json:"replicas"` + + // ReadyReplicas is how many of those are actually serving. + // +optional + ReadyReplicas int32 `json:"readyReplicas,omitempty"` + + // Selector is a serialized label selector matching the scaled pods. + // HorizontalPodAutoscaler requires this to be a *string*, not a structured + // selector, and refuses to scale a resource whose scale subresource does not + // provide one. + // +optional + Selector string `json:"selector,omitempty"` + + // ObservedGeneration is the adapter generation this status reflects. + // +optional + ObservedGeneration int64 `json:"observedGeneration,omitempty"` + + // Conditions carries `Ready` (the target resolves and is being driven) and + // `Degraded` (it does not). + // +optional + Conditions []metav1.Condition `json:"conditions,omitempty"` +} + +// InferaScalingAdapter gives one service of an InferaDeployment a standard +// Kubernetes `/scale` subresource. +// +// An InferaDeployment cannot carry `/scale` itself, and this is a property of +// its shape rather than a missing feature: `spec.services` is a map with +// user-chosen keys, while the scale subresource requires `specReplicasPath` to +// be a *static* dot-notation JSONPath under `.spec`. There is no way to write +// "the replicas of an arbitrary map entry". +// +// So scaling gets its own object, one per scalable service. That makes +// `kubectl scale`, HorizontalPodAutoscaler, KEDA and a custom planner all work +// through the same standard interface, with no per-tool support in this +// operator. +// +// While an adapter exists with `spec.replicas` set, it is the single writer of +// that service's replica count: the InferaDeployment reconciler reads the +// adapter instead of the CR's own `replicas`, so the two cannot fight. Delete +// the adapter, or clear `spec.replicas`, and the CR is back in charge. +// +// +kubebuilder:object:root=true +// +kubebuilder:subresource:status +// +kubebuilder:subresource:scale:specpath=.spec.replicas,statuspath=.status.replicas,selectorpath=.status.selector +// +kubebuilder:resource:shortName=isa +// +kubebuilder:printcolumn:name="Target",type=string,JSONPath=`.spec.deploymentRef` +// +kubebuilder:printcolumn:name="Service",type=string,JSONPath=`.spec.serviceName` +// +kubebuilder:printcolumn:name="Desired",type=integer,JSONPath=`.spec.replicas` +// +kubebuilder:printcolumn:name="Current",type=integer,JSONPath=`.status.replicas` +// +kubebuilder:printcolumn:name="Ready",type=integer,JSONPath=`.status.readyReplicas` +// +kubebuilder:printcolumn:name="Age",type=date,JSONPath=`.metadata.creationTimestamp` +type InferaScalingAdapter struct { + metav1.TypeMeta `json:",inline"` + metav1.ObjectMeta `json:"metadata,omitempty"` + + Spec InferaScalingAdapterSpec `json:"spec,omitempty"` + Status InferaScalingAdapterStatus `json:"status,omitempty"` +} + +// +kubebuilder:object:root=true +type InferaScalingAdapterList struct { + metav1.TypeMeta `json:",inline"` + metav1.ListMeta `json:"metadata,omitempty"` + Items []InferaScalingAdapter `json:"items"` +} + +func init() { + SchemeBuilder.Register(&InferaScalingAdapter{}, &InferaScalingAdapterList{}) +} diff --git a/deploy/operator/api/v1alpha1/zz_generated.deepcopy.go b/deploy/operator/api/v1alpha1/zz_generated.deepcopy.go index 6de23549..cc4f5032 100644 --- a/deploy/operator/api/v1alpha1/zz_generated.deepcopy.go +++ b/deploy/operator/api/v1alpha1/zz_generated.deepcopy.go @@ -45,36 +45,6 @@ func (in *GAIEStatus) DeepCopy() *GAIEStatus { return out } -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *NATSSpec) DeepCopyInto(out *NATSSpec) { - *out = *in -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new NATSSpec. -func (in *NATSSpec) DeepCopy() *NATSSpec { - if in == nil { - return nil - } - out := new(NATSSpec) - in.DeepCopyInto(out) - return out -} - -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *Resources) DeepCopyInto(out *Resources) { - *out = *in -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new Resources. -func (in *Resources) DeepCopy() *Resources { - if in == nil { - return nil - } - out := new(Resources) - in.DeepCopyInto(out) - return out -} - // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *InferaDeployment) DeepCopyInto(out *InferaDeployment) { *out = *in @@ -207,6 +177,137 @@ func (in *InferaDeploymentStatus) DeepCopy() *InferaDeploymentStatus { return out } +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *InferaScalingAdapter) DeepCopyInto(out *InferaScalingAdapter) { + *out = *in + out.TypeMeta = in.TypeMeta + in.ObjectMeta.DeepCopyInto(&out.ObjectMeta) + in.Spec.DeepCopyInto(&out.Spec) + in.Status.DeepCopyInto(&out.Status) +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new InferaScalingAdapter. +func (in *InferaScalingAdapter) DeepCopy() *InferaScalingAdapter { + if in == nil { + return nil + } + out := new(InferaScalingAdapter) + in.DeepCopyInto(out) + return out +} + +// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object. +func (in *InferaScalingAdapter) DeepCopyObject() runtime.Object { + if c := in.DeepCopy(); c != nil { + return c + } + return nil +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *InferaScalingAdapterList) DeepCopyInto(out *InferaScalingAdapterList) { + *out = *in + out.TypeMeta = in.TypeMeta + in.ListMeta.DeepCopyInto(&out.ListMeta) + if in.Items != nil { + in, out := &in.Items, &out.Items + *out = make([]InferaScalingAdapter, len(*in)) + for i := range *in { + (*in)[i].DeepCopyInto(&(*out)[i]) + } + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new InferaScalingAdapterList. +func (in *InferaScalingAdapterList) DeepCopy() *InferaScalingAdapterList { + if in == nil { + return nil + } + out := new(InferaScalingAdapterList) + in.DeepCopyInto(out) + return out +} + +// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object. +func (in *InferaScalingAdapterList) DeepCopyObject() runtime.Object { + if c := in.DeepCopy(); c != nil { + return c + } + return nil +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *InferaScalingAdapterSpec) DeepCopyInto(out *InferaScalingAdapterSpec) { + *out = *in + if in.Replicas != nil { + in, out := &in.Replicas, &out.Replicas + *out = new(int32) + **out = **in + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new InferaScalingAdapterSpec. +func (in *InferaScalingAdapterSpec) DeepCopy() *InferaScalingAdapterSpec { + if in == nil { + return nil + } + out := new(InferaScalingAdapterSpec) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *InferaScalingAdapterStatus) DeepCopyInto(out *InferaScalingAdapterStatus) { + *out = *in + if in.Conditions != nil { + in, out := &in.Conditions, &out.Conditions + *out = make([]metav1.Condition, len(*in)) + for i := range *in { + (*in)[i].DeepCopyInto(&(*out)[i]) + } + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new InferaScalingAdapterStatus. +func (in *InferaScalingAdapterStatus) DeepCopy() *InferaScalingAdapterStatus { + if in == nil { + return nil + } + out := new(InferaScalingAdapterStatus) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *NATSSpec) DeepCopyInto(out *NATSSpec) { + *out = *in +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new NATSSpec. +func (in *NATSSpec) DeepCopy() *NATSSpec { + if in == nil { + return nil + } + out := new(NATSSpec) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *Resources) DeepCopyInto(out *Resources) { + *out = *in +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new Resources. +func (in *Resources) DeepCopy() *Resources { + if in == nil { + return nil + } + out := new(Resources) + in.DeepCopyInto(out) + return out +} + // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *ServiceSpec) DeepCopyInto(out *ServiceSpec) { *out = *in diff --git a/deploy/operator/cmd/main.go b/deploy/operator/cmd/main.go index b0da1976..867e05c2 100644 --- a/deploy/operator/cmd/main.go +++ b/deploy/operator/cmd/main.go @@ -55,6 +55,14 @@ func main() { os.Exit(1) } + if err := (&controller.InferaScalingAdapterReconciler{ + Client: mgr.GetClient(), + Scheme: mgr.GetScheme(), + }).SetupWithManager(mgr); err != nil { + setupLog.Error(err, "unable to create controller", "controller", "InferaScalingAdapter") + os.Exit(1) + } + if err := (&controller.InferaDeploymentReconciler{ Client: mgr.GetClient(), Scheme: mgr.GetScheme(), diff --git a/deploy/operator/config/crd/bases/infera.amd.com_inferascalingadapters.yaml b/deploy/operator/config/crd/bases/infera.amd.com_inferascalingadapters.yaml new file mode 100644 index 00000000..8b722246 --- /dev/null +++ b/deploy/operator/config/crd/bases/infera.amd.com_inferascalingadapters.yaml @@ -0,0 +1,207 @@ +--- +apiVersion: apiextensions.k8s.io/v1 +kind: CustomResourceDefinition +metadata: + annotations: + controller-gen.kubebuilder.io/version: v0.17.1 + name: inferascalingadapters.infera.amd.com +spec: + group: infera.amd.com + names: + kind: InferaScalingAdapter + listKind: InferaScalingAdapterList + plural: inferascalingadapters + shortNames: + - isa + singular: inferascalingadapter + scope: Namespaced + versions: + - additionalPrinterColumns: + - jsonPath: .spec.deploymentRef + name: Target + type: string + - jsonPath: .spec.serviceName + name: Service + type: string + - jsonPath: .spec.replicas + name: Desired + type: integer + - jsonPath: .status.replicas + name: Current + type: integer + - jsonPath: .status.readyReplicas + name: Ready + type: integer + - jsonPath: .metadata.creationTimestamp + name: Age + type: date + name: v1alpha1 + schema: + openAPIV3Schema: + description: |- + InferaScalingAdapter gives one service of an InferaDeployment a standard + Kubernetes `/scale` subresource. + + An InferaDeployment cannot carry `/scale` itself, and this is a property of + its shape rather than a missing feature: `spec.services` is a map with + user-chosen keys, while the scale subresource requires `specReplicasPath` to + be a *static* dot-notation JSONPath under `.spec`. There is no way to write + "the replicas of an arbitrary map entry". + + So scaling gets its own object, one per scalable service. That makes + `kubectl scale`, HorizontalPodAutoscaler, KEDA and a custom planner all work + through the same standard interface, with no per-tool support in this + operator. + + While an adapter exists with `spec.replicas` set, it is the single writer of + that service's replica count: the InferaDeployment reconciler reads the + adapter instead of the CR's own `replicas`, so the two cannot fight. Delete + the adapter, or clear `spec.replicas`, and the CR is back in charge. + properties: + apiVersion: + description: |- + APIVersion defines the versioned schema of this representation of an object. + Servers should convert recognized schemas to the latest internal value, and + may reject unrecognized values. + More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#resources + type: string + kind: + description: |- + Kind is a string value representing the REST resource this object represents. + Servers may infer this from the endpoint the client submits requests to. + Cannot be updated. + In CamelCase. + More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#types-kinds + type: string + metadata: + type: object + spec: + description: |- + InferaScalingAdapterSpec points at one service inside an InferaDeployment and + owns its replica count. + properties: + deploymentRef: + description: DeploymentRef is the InferaDeployment to scale, in this + namespace. + minLength: 1 + type: string + replicas: + description: |- + Replicas is the desired count, and the field `/scale` writes. + + Left unset the adapter is inert: the InferaDeployment's own + `spec.services[].replicas` still applies. That makes adding an + adapter a safe no-op until something actually scales, so an autoscaler + can be attached and observed before it is trusted. + format: int32 + minimum: 0 + type: integer + serviceName: + description: ServiceName is the key in that deployment's `spec.services` + map. + minLength: 1 + type: string + required: + - deploymentRef + - serviceName + type: object + status: + description: InferaScalingAdapterStatus carries what `/scale` reads back. + properties: + conditions: + description: |- + Conditions carries `Ready` (the target resolves and is being driven) and + `Degraded` (it does not). + items: + description: Condition contains details for one aspect of the current + state of this API Resource. + properties: + lastTransitionTime: + description: |- + lastTransitionTime is the last time the condition transitioned from one status to another. + This should be when the underlying condition changed. If that is not known, then using the time when the API field changed is acceptable. + format: date-time + type: string + message: + description: |- + message is a human readable message indicating details about the transition. + This may be an empty string. + maxLength: 32768 + type: string + observedGeneration: + description: |- + observedGeneration represents the .metadata.generation that the condition was set based upon. + For instance, if .metadata.generation is currently 12, but the .status.conditions[x].observedGeneration is 9, the condition is out of date + with respect to the current state of the instance. + format: int64 + minimum: 0 + type: integer + reason: + description: |- + reason contains a programmatic identifier indicating the reason for the condition's last transition. + Producers of specific condition types may define expected values and meanings for this field, + and whether the values are considered a guaranteed API. + The value should be a CamelCase string. + This field may not be empty. + maxLength: 1024 + minLength: 1 + pattern: ^[A-Za-z]([A-Za-z0-9_,:]*[A-Za-z0-9_])?$ + type: string + status: + description: status of the condition, one of True, False, Unknown. + enum: + - "True" + - "False" + - Unknown + type: string + type: + description: type of condition in CamelCase or in foo.example.com/CamelCase. + maxLength: 316 + pattern: ^([a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*/)?(([A-Za-z0-9][-A-Za-z0-9_.]*)?[A-Za-z0-9])$ + type: string + required: + - lastTransitionTime + - message + - reason + - status + - type + type: object + type: array + observedGeneration: + description: ObservedGeneration is the adapter generation this status + reflects. + format: int64 + type: integer + readyReplicas: + description: ReadyReplicas is how many of those are actually serving. + format: int32 + type: integer + replicas: + description: |- + Replicas is the count observed on the workload -- not the desired count + echoed back. + + The distinction is the whole reason this field exists. HorizontalPodAutoscaler + computes `desired = ceil(current * metric/target)`; if `current` is really + the desired value it never lags reality, so during a multi-minute model load + the autoscaler cannot tell that a scale-up has not landed yet and keeps + multiplying. + format: int32 + type: integer + selector: + description: |- + Selector is a serialized label selector matching the scaled pods. + HorizontalPodAutoscaler requires this to be a *string*, not a structured + selector, and refuses to scale a resource whose scale subresource does not + provide one. + type: string + type: object + type: object + served: true + storage: true + subresources: + scale: + labelSelectorPath: .status.selector + specReplicasPath: .spec.replicas + statusReplicasPath: .status.replicas + status: {} diff --git a/deploy/operator/config/rbac/role.yaml b/deploy/operator/config/rbac/role.yaml index 662cdf59..d5c44e2f 100644 --- a/deploy/operator/config/rbac/role.yaml +++ b/deploy/operator/config/rbac/role.yaml @@ -52,9 +52,10 @@ rules: - update - watch - apiGroups: - - inference.networking.k8s.io + - infera.amd.com resources: - - inferencepools + - inferadeployments + - inferascalingadapters verbs: - create - delete @@ -64,22 +65,25 @@ rules: - update - watch - apiGroups: - - leaderworkerset.x-k8s.io + - infera.amd.com resources: - - leaderworkersets + - inferadeployments/finalizers + verbs: + - update +- apiGroups: + - infera.amd.com + resources: + - inferadeployments/status + - inferascalingadapters/scale + - inferascalingadapters/status verbs: - - create - - delete - get - - list - patch - update - - watch - apiGroups: - - rbac.authorization.k8s.io + - inference.networking.k8s.io resources: - - rolebindings - - roles + - inferencepools verbs: - create - delete @@ -89,9 +93,9 @@ rules: - update - watch - apiGroups: - - infera.amd.com + - leaderworkerset.x-k8s.io resources: - - inferadeployments + - leaderworkersets verbs: - create - delete @@ -101,16 +105,15 @@ rules: - update - watch - apiGroups: - - infera.amd.com - resources: - - inferadeployments/finalizers - verbs: - - update -- apiGroups: - - infera.amd.com + - rbac.authorization.k8s.io resources: - - inferadeployments/status + - rolebindings + - roles verbs: + - create + - delete - get + - list - patch - update + - watch diff --git a/deploy/operator/helm/infera-operator/crds/infera.amd.com_inferascalingadapters.yaml b/deploy/operator/helm/infera-operator/crds/infera.amd.com_inferascalingadapters.yaml new file mode 100644 index 00000000..8b722246 --- /dev/null +++ b/deploy/operator/helm/infera-operator/crds/infera.amd.com_inferascalingadapters.yaml @@ -0,0 +1,207 @@ +--- +apiVersion: apiextensions.k8s.io/v1 +kind: CustomResourceDefinition +metadata: + annotations: + controller-gen.kubebuilder.io/version: v0.17.1 + name: inferascalingadapters.infera.amd.com +spec: + group: infera.amd.com + names: + kind: InferaScalingAdapter + listKind: InferaScalingAdapterList + plural: inferascalingadapters + shortNames: + - isa + singular: inferascalingadapter + scope: Namespaced + versions: + - additionalPrinterColumns: + - jsonPath: .spec.deploymentRef + name: Target + type: string + - jsonPath: .spec.serviceName + name: Service + type: string + - jsonPath: .spec.replicas + name: Desired + type: integer + - jsonPath: .status.replicas + name: Current + type: integer + - jsonPath: .status.readyReplicas + name: Ready + type: integer + - jsonPath: .metadata.creationTimestamp + name: Age + type: date + name: v1alpha1 + schema: + openAPIV3Schema: + description: |- + InferaScalingAdapter gives one service of an InferaDeployment a standard + Kubernetes `/scale` subresource. + + An InferaDeployment cannot carry `/scale` itself, and this is a property of + its shape rather than a missing feature: `spec.services` is a map with + user-chosen keys, while the scale subresource requires `specReplicasPath` to + be a *static* dot-notation JSONPath under `.spec`. There is no way to write + "the replicas of an arbitrary map entry". + + So scaling gets its own object, one per scalable service. That makes + `kubectl scale`, HorizontalPodAutoscaler, KEDA and a custom planner all work + through the same standard interface, with no per-tool support in this + operator. + + While an adapter exists with `spec.replicas` set, it is the single writer of + that service's replica count: the InferaDeployment reconciler reads the + adapter instead of the CR's own `replicas`, so the two cannot fight. Delete + the adapter, or clear `spec.replicas`, and the CR is back in charge. + properties: + apiVersion: + description: |- + APIVersion defines the versioned schema of this representation of an object. + Servers should convert recognized schemas to the latest internal value, and + may reject unrecognized values. + More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#resources + type: string + kind: + description: |- + Kind is a string value representing the REST resource this object represents. + Servers may infer this from the endpoint the client submits requests to. + Cannot be updated. + In CamelCase. + More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#types-kinds + type: string + metadata: + type: object + spec: + description: |- + InferaScalingAdapterSpec points at one service inside an InferaDeployment and + owns its replica count. + properties: + deploymentRef: + description: DeploymentRef is the InferaDeployment to scale, in this + namespace. + minLength: 1 + type: string + replicas: + description: |- + Replicas is the desired count, and the field `/scale` writes. + + Left unset the adapter is inert: the InferaDeployment's own + `spec.services[].replicas` still applies. That makes adding an + adapter a safe no-op until something actually scales, so an autoscaler + can be attached and observed before it is trusted. + format: int32 + minimum: 0 + type: integer + serviceName: + description: ServiceName is the key in that deployment's `spec.services` + map. + minLength: 1 + type: string + required: + - deploymentRef + - serviceName + type: object + status: + description: InferaScalingAdapterStatus carries what `/scale` reads back. + properties: + conditions: + description: |- + Conditions carries `Ready` (the target resolves and is being driven) and + `Degraded` (it does not). + items: + description: Condition contains details for one aspect of the current + state of this API Resource. + properties: + lastTransitionTime: + description: |- + lastTransitionTime is the last time the condition transitioned from one status to another. + This should be when the underlying condition changed. If that is not known, then using the time when the API field changed is acceptable. + format: date-time + type: string + message: + description: |- + message is a human readable message indicating details about the transition. + This may be an empty string. + maxLength: 32768 + type: string + observedGeneration: + description: |- + observedGeneration represents the .metadata.generation that the condition was set based upon. + For instance, if .metadata.generation is currently 12, but the .status.conditions[x].observedGeneration is 9, the condition is out of date + with respect to the current state of the instance. + format: int64 + minimum: 0 + type: integer + reason: + description: |- + reason contains a programmatic identifier indicating the reason for the condition's last transition. + Producers of specific condition types may define expected values and meanings for this field, + and whether the values are considered a guaranteed API. + The value should be a CamelCase string. + This field may not be empty. + maxLength: 1024 + minLength: 1 + pattern: ^[A-Za-z]([A-Za-z0-9_,:]*[A-Za-z0-9_])?$ + type: string + status: + description: status of the condition, one of True, False, Unknown. + enum: + - "True" + - "False" + - Unknown + type: string + type: + description: type of condition in CamelCase or in foo.example.com/CamelCase. + maxLength: 316 + pattern: ^([a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*/)?(([A-Za-z0-9][-A-Za-z0-9_.]*)?[A-Za-z0-9])$ + type: string + required: + - lastTransitionTime + - message + - reason + - status + - type + type: object + type: array + observedGeneration: + description: ObservedGeneration is the adapter generation this status + reflects. + format: int64 + type: integer + readyReplicas: + description: ReadyReplicas is how many of those are actually serving. + format: int32 + type: integer + replicas: + description: |- + Replicas is the count observed on the workload -- not the desired count + echoed back. + + The distinction is the whole reason this field exists. HorizontalPodAutoscaler + computes `desired = ceil(current * metric/target)`; if `current` is really + the desired value it never lags reality, so during a multi-minute model load + the autoscaler cannot tell that a scale-up has not landed yet and keeps + multiplying. + format: int32 + type: integer + selector: + description: |- + Selector is a serialized label selector matching the scaled pods. + HorizontalPodAutoscaler requires this to be a *string*, not a structured + selector, and refuses to scale a resource whose scale subresource does not + provide one. + type: string + type: object + type: object + served: true + storage: true + subresources: + scale: + labelSelectorPath: .status.selector + specReplicasPath: .spec.replicas + statusReplicasPath: .status.replicas + status: {} diff --git a/deploy/operator/helm/infera-operator/templates/rbac.yaml b/deploy/operator/helm/infera-operator/templates/rbac.yaml index 54d8f55e..38fd46f9 100644 --- a/deploy/operator/helm/infera-operator/templates/rbac.yaml +++ b/deploy/operator/helm/infera-operator/templates/rbac.yaml @@ -43,6 +43,19 @@ rules: - apiGroups: ["infera.amd.com"] resources: ["inferadeployments/status"] verbs: ["get", "patch", "update"] +# Scaling adapters: the standard /scale surface an HPA, KEDA or a custom planner +# drives. The operator reads spec.replicas and owns status; /scale is listed +# because that is the endpoint external scalers write, and it is a distinct +# subresource from status for RBAC purposes. +- apiGroups: ["infera.amd.com"] + resources: ["inferascalingadapters"] + verbs: ["create", "delete", "get", "list", "patch", "update", "watch"] +- apiGroups: ["infera.amd.com"] + resources: ["inferascalingadapters/status"] + verbs: ["get", "patch", "update"] +- apiGroups: ["infera.amd.com"] + resources: ["inferascalingadapters/scale"] + verbs: ["get", "patch", "update"] # Leader election lease + event recording. - apiGroups: ["coordination.k8s.io"] resources: ["leases"] diff --git a/deploy/operator/internal/controller/builders.go b/deploy/operator/internal/controller/builders.go index cf474f7f..f5b6bb8d 100644 --- a/deploy/operator/internal/controller/builders.go +++ b/deploy/operator/internal/controller/builders.go @@ -152,6 +152,21 @@ func replicasOf(svc inferav1alpha1.ServiceSpec) int32 { return 1 } +// effectiveReplicas is the count to write onto the workload: the scaling +// adapter's if one owns this service, otherwise the CR's own. +// +// Routing the adapter through here rather than letting it write the workload +// directly keeps a single writer. The reconciler already assigns the whole +// child `.Spec` on every pass, so a second writer would simply be reverted -- +// which is exactly what happens today to anyone pointing an HPA at the child +// Deployment. +func effectiveReplicas(svc inferav1alpha1.ServiceSpec, svcName string, overrides map[string]int32) int32 { + if n, ok := overrides[svcName]; ok { + return n + } + return replicasOf(svc) +} + // containerCommand builds the infera entrypoint + operator-injected flags, // then appends the user's free-form Args (model-path, tokenizer, tp-size, ...). func containerCommand(idep *inferav1alpha1.InferaDeployment, svc inferav1alpha1.ServiceSpec) []string { @@ -415,8 +430,8 @@ func podTemplate(idep *inferav1alpha1.InferaDeployment, svcName string, svc infe return tmpl } -func buildDeployment(idep *inferav1alpha1.InferaDeployment, svcName string, svc inferav1alpha1.ServiceSpec) *appsv1.Deployment { - reps := replicasOf(svc) +func buildDeployment(idep *inferav1alpha1.InferaDeployment, svcName string, svc inferav1alpha1.ServiceSpec, overrides map[string]int32) *appsv1.Deployment { + reps := effectiveReplicas(svc, svcName, overrides) lbls := labelsFor(idep.Name, svcName) // Worker services use surge-free RollingUpdate (maxSurge=0, maxUnavailable=1): // the default RollingUpdate brings up a surge pod first, which on a @@ -455,8 +470,8 @@ func buildDeployment(idep *inferav1alpha1.InferaDeployment, svcName string, svc // buildLeaderWorkerSet returns an unstructured LeaderWorkerSet so the operator // does not take a compile-time dependency on the LWS Go module (keeps Infera // self-contained; the LWS CRD must be installed in the cluster). -func buildLeaderWorkerSet(idep *inferav1alpha1.InferaDeployment, svcName string, svc inferav1alpha1.ServiceSpec) *unstructured.Unstructured { - reps := replicasOf(svc) +func buildLeaderWorkerSet(idep *inferav1alpha1.InferaDeployment, svcName string, svc inferav1alpha1.ServiceSpec, overrides map[string]int32) *unstructured.Unstructured { + reps := effectiveReplicas(svc, svcName, overrides) lbls := labelsFor(idep.Name, svcName) tmpl := podTemplate(idep, svcName, svc) // Convert the typed PodTemplateSpec to a map for embedding. diff --git a/deploy/operator/internal/controller/builders_test.go b/deploy/operator/internal/controller/builders_test.go index 160448d8..c32f6c02 100644 --- a/deploy/operator/internal/controller/builders_test.go +++ b/deploy/operator/internal/controller/builders_test.go @@ -10,6 +10,8 @@ import ( "testing" corev1 "k8s.io/api/core/v1" + + inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" ) // The grace period is the only thing standing between a graceful drain and a @@ -97,3 +99,31 @@ func TestGraceReadsDrainTimeoutFromTheContainerToo(t *testing.T) { t.Fatalf("grace = %d, want %d", *spec.TerminationGracePeriodSeconds, want) } } + +// An adapter owns its service's replica count; everything else keeps using the +// CR's. Getting this wrong in either direction is bad: ignoring the adapter +// makes `/scale` a no-op, and applying it too broadly makes a single autoscaler +// silently resize pools nobody pointed it at. +func TestEffectiveReplicas(t *testing.T) { + three := int32(3) + svc := inferav1alpha1.ServiceSpec{Replicas: &three} + + if got := effectiveReplicas(svc, "worker", nil); got != 3 { + t.Fatalf("no adapters: got %d, want the CR's 3", got) + } + if got := effectiveReplicas(svc, "worker", map[string]int32{"worker": 7}); got != 7 { + t.Fatalf("adapter present: got %d, want 7", got) + } + if got := effectiveReplicas(svc, "worker", map[string]int32{"prefill": 7}); got != 3 { + t.Fatalf("adapter for another service: got %d, want the CR's 3", got) + } + // Zero is a legitimate target, not "unset" -- an autoscaler scaling a pool + // to zero must not silently fall back to the CR's count. + if got := effectiveReplicas(svc, "worker", map[string]int32{"worker": 0}); got != 0 { + t.Fatalf("adapter asking for 0: got %d, want 0", got) + } + // The CR default when it says nothing either. + if got := effectiveReplicas(inferav1alpha1.ServiceSpec{}, "worker", nil); got != 1 { + t.Fatalf("nothing set anywhere: got %d, want the default 1", got) + } +} diff --git a/deploy/operator/internal/controller/inferadeployment_controller.go b/deploy/operator/internal/controller/inferadeployment_controller.go index 07c032b3..148bdd3a 100644 --- a/deploy/operator/internal/controller/inferadeployment_controller.go +++ b/deploy/operator/internal/controller/inferadeployment_controller.go @@ -13,6 +13,8 @@ import ( appsv1 "k8s.io/api/apps/v1" corev1 "k8s.io/api/core/v1" rbacv1 "k8s.io/api/rbac/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + "k8s.io/apimachinery/pkg/api/meta" "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" "k8s.io/apimachinery/pkg/runtime" ctrl "sigs.k8s.io/controller-runtime" @@ -86,17 +88,27 @@ func (r *InferaDeploymentReconciler) Reconcile(ctx context.Context, req ctrl.Req } // 2. Each service -> Deployment (single-node) or LeaderWorkerSet (multi-node). + // + // Scaling adapters are resolved first: where one owns a service, its + // replica count wins over the CR's. Reading them here rather than letting + // the adapter write the workload keeps a single writer -- this reconciler + // assigns the whole child `.Spec` every pass, so a second writer would just + // be reverted. + overrides, err := r.replicaOverrides(ctx, idep) + if err != nil { + return ctrl.Result{}, err + } status := map[string]inferav1alpha1.ServiceStatus{} for _, name := range sortedKeys(idep.Spec.Services) { svc := idep.Spec.Services[name] if svc.NumberOfNodes > 1 { - lws := buildLeaderWorkerSet(idep, name, svc) + lws := buildLeaderWorkerSet(idep, name, svc, overrides) if err := r.applyUnstructured(ctx, idep, lws); err != nil { return ctrl.Result{}, err } status[name] = r.lwsStatus(ctx, idep, name, svc) } else { - dep := buildDeployment(idep, name, svc) + dep := buildDeployment(idep, name, svc, overrides) if err := r.applyObject(ctx, idep, dep); err != nil { return ctrl.Result{}, err } @@ -144,6 +156,39 @@ func (r *InferaDeploymentReconciler) Reconcile(ctx context.Context, req ctrl.Req return ctrl.Result{RequeueAfter: 15 * time.Second}, nil } +// replicaOverrides maps service name -> replica count for services owned by a +// scaling adapter. +// +// An adapter with no `spec.replicas` is deliberately absent from the map rather +// than contributing a zero: creating an adapter must be a no-op until something +// actually scales, so an autoscaler can be attached and watched before it is +// trusted. A missing CRD is likewise not an error -- the adapter is optional, +// and an operator that refused to reconcile without it would break every +// existing deployment on upgrade. +func (r *InferaDeploymentReconciler) replicaOverrides( + ctx context.Context, idep *inferav1alpha1.InferaDeployment, +) (map[string]int32, error) { + list := &inferav1alpha1.InferaScalingAdapterList{} + if err := r.List(ctx, list, client.InNamespace(idep.Namespace)); err != nil { + if meta.IsNoMatchError(err) || apierrors.IsNotFound(err) { + return nil, nil + } + return nil, err + } + out := map[string]int32{} + for i := range list.Items { + a := &list.Items[i] + if a.Spec.DeploymentRef != idep.Name || a.Spec.Replicas == nil { + continue + } + if _, ok := idep.Spec.Services[a.Spec.ServiceName]; !ok { + continue // dangling adapter; its own controller reports Degraded + } + out[a.Spec.ServiceName] = *a.Spec.Replicas + } + return out, nil +} + // applyObject create-or-updates a typed object, setting the owner reference. func (r *InferaDeploymentReconciler) applyObject(ctx context.Context, idep *inferav1alpha1.InferaDeployment, desired client.Object) error { // Build a fresh empty object of the same kind keyed by name/namespace. @@ -172,19 +217,28 @@ func (r *InferaDeploymentReconciler) applyUnstructured(ctx context.Context, idep } func (r *InferaDeploymentReconciler) deploymentStatus(ctx context.Context, idep *inferav1alpha1.InferaDeployment, name string, svc inferav1alpha1.ServiceSpec) inferav1alpha1.ServiceStatus { - st := inferav1alpha1.ServiceStatus{Kind: "Deployment", Replicas: replicasOf(svc)} + // Replicas is what the workload reports, not what the spec asked for. + // Echoing the desired value back makes status useless for exactly the + // reader that needs it most -- an autoscaler computing + // `desired = ceil(current * metric/target)` cannot tell a scale-up has not + // landed if `current` is the number it just asked for. + st := inferav1alpha1.ServiceStatus{Kind: "Deployment"} dep := &appsv1.Deployment{} if err := r.Get(ctx, client.ObjectKey{Name: idep.Name + "-" + name, Namespace: idep.Namespace}, dep); err == nil { + st.Replicas = dep.Status.Replicas st.ReadyReplicas = dep.Status.ReadyReplicas } return st } func (r *InferaDeploymentReconciler) lwsStatus(ctx context.Context, idep *inferav1alpha1.InferaDeployment, name string, svc inferav1alpha1.ServiceSpec) inferav1alpha1.ServiceStatus { - st := inferav1alpha1.ServiceStatus{Kind: "LeaderWorkerSet", Replicas: replicasOf(svc)} + st := inferav1alpha1.ServiceStatus{Kind: "LeaderWorkerSet"} u := &unstructured.Unstructured{} u.SetGroupVersionKind(lwsGVK()) if err := r.Get(ctx, client.ObjectKey{Name: idep.Name + "-" + name, Namespace: idep.Namespace}, u); err == nil { + if v, ok, _ := unstructured.NestedInt64(u.Object, "status", "replicas"); ok { + st.Replicas = int32(v) + } if v, ok, _ := unstructured.NestedInt64(u.Object, "status", "readyReplicas"); ok { st.ReadyReplicas = int32(v) } diff --git a/deploy/operator/internal/controller/inferascalingadapter_controller.go b/deploy/operator/internal/controller/inferascalingadapter_controller.go new file mode 100644 index 00000000..ea35eb7c --- /dev/null +++ b/deploy/operator/internal/controller/inferascalingadapter_controller.go @@ -0,0 +1,221 @@ +/* +Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + +SPDX-License-Identifier: MIT +*/ + +package controller + +import ( + "context" + "fmt" + "time" + + appsv1 "k8s.io/api/apps/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/labels" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/types" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/handler" + "sigs.k8s.io/controller-runtime/pkg/log" + "sigs.k8s.io/controller-runtime/pkg/reconcile" + + inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" +) + +// InferaScalingAdapterReconciler keeps an adapter's status in step with the +// workload it scales. +// +// It deliberately does not write the workload. The InferaDeployment reconciler +// reads adapters when it builds children, so there is exactly one writer of any +// child `.Spec`. Two writers would not merely race -- that reconciler assigns +// the whole spec on every pass, so the loser is reverted within seconds, which +// is precisely the failure an HPA pointed at the child Deployment hits today. +// +// What this controller owns is the half `/scale` reads back: `status.replicas` +// from the live workload, and `status.selector`, without which HorizontalPod- +// Autoscaler refuses to scale the resource at all. +type InferaScalingAdapterReconciler struct { + client.Client + Scheme *runtime.Scheme +} + +// +kubebuilder:rbac:groups=infera.amd.com,resources=inferascalingadapters,verbs=get;list;watch;create;update;patch;delete +// +kubebuilder:rbac:groups=infera.amd.com,resources=inferascalingadapters/status,verbs=get;update;patch +// +kubebuilder:rbac:groups=infera.amd.com,resources=inferascalingadapters/scale,verbs=get;update;patch + +func (r *InferaScalingAdapterReconciler) Reconcile( + ctx context.Context, req ctrl.Request, +) (ctrl.Result, error) { + lg := log.FromContext(ctx) + + adapter := &inferav1alpha1.InferaScalingAdapter{} + if err := r.Get(ctx, req.NamespacedName, adapter); err != nil { + return ctrl.Result{}, client.IgnoreNotFound(err) + } + if !adapter.DeletionTimestamp.IsZero() { + return ctrl.Result{}, nil + } + + st := inferav1alpha1.InferaScalingAdapterStatus{ + ObservedGeneration: adapter.Generation, + } + + idep := &inferav1alpha1.InferaDeployment{} + err := r.Get(ctx, types.NamespacedName{ + Name: adapter.Spec.DeploymentRef, Namespace: adapter.Namespace, + }, idep) + switch { + case apierrors.IsNotFound(err): + return r.degraded(ctx, adapter, st, "TargetNotFound", + fmt.Sprintf("no InferaDeployment %q in this namespace", adapter.Spec.DeploymentRef)) + case err != nil: + return ctrl.Result{}, err + } + + svc, ok := idep.Spec.Services[adapter.Spec.ServiceName] + if !ok { + return r.degraded(ctx, adapter, st, "ServiceNotFound", + fmt.Sprintf("InferaDeployment %q has no service %q", + adapter.Spec.DeploymentRef, adapter.Spec.ServiceName)) + } + + // The selector must be a serialized string, not a structured selector -- + // that is what the scale subresource contract requires, and HPA rejects a + // target without one. + st.Selector = labels.SelectorFromSet( + labelsFor(idep.Name, adapter.Spec.ServiceName)).String() + + name := idep.Name + "-" + adapter.Spec.ServiceName + key := types.NamespacedName{Name: name, Namespace: idep.Namespace} + if svc.NumberOfNodes > 1 { + u := &unstructured.Unstructured{} + u.SetGroupVersionKind(lwsGVK()) + if err := r.Get(ctx, key, u); err == nil { + if v, ok, _ := unstructured.NestedInt64(u.Object, "status", "replicas"); ok { + st.Replicas = int32(v) + } + if v, ok, _ := unstructured.NestedInt64(u.Object, "status", "readyReplicas"); ok { + st.ReadyReplicas = int32(v) + } + } + } else { + dep := &appsv1.Deployment{} + if err := r.Get(ctx, key, dep); err == nil { + st.Replicas = dep.Status.Replicas + st.ReadyReplicas = dep.Status.ReadyReplicas + } + } + + msg := "adapter is inert: spec.replicas unset, the InferaDeployment's own replicas apply" + if adapter.Spec.Replicas != nil { + msg = fmt.Sprintf("driving %s/%s to %d replica(s)", + adapter.Spec.DeploymentRef, adapter.Spec.ServiceName, *adapter.Spec.Replicas) + } + setCondition(&st.Conditions, adapter.Generation, "Ready", metav1.ConditionTrue, "Resolved", msg) + setCondition(&st.Conditions, adapter.Generation, "Degraded", metav1.ConditionFalse, "Resolved", msg) + + lg.V(1).Info("scaling adapter reconciled", "target", adapter.Spec.DeploymentRef, + "service", adapter.Spec.ServiceName, "observed", st.Replicas) + return r.writeStatus(ctx, adapter, st) +} + +func (r *InferaScalingAdapterReconciler) degraded( + ctx context.Context, a *inferav1alpha1.InferaScalingAdapter, + st inferav1alpha1.InferaScalingAdapterStatus, reason, msg string, +) (ctrl.Result, error) { + setCondition(&st.Conditions, a.Generation, "Ready", metav1.ConditionFalse, reason, msg) + setCondition(&st.Conditions, a.Generation, "Degraded", metav1.ConditionTrue, reason, msg) + res, err := r.writeStatus(ctx, a, st) + if err != nil { + return res, err + } + // A dangling adapter usually means the target has not been created yet, so + // retry rather than waiting for an event on an object that does not exist. + return ctrl.Result{RequeueAfter: 30 * time.Second}, nil +} + +func (r *InferaScalingAdapterReconciler) writeStatus( + ctx context.Context, a *inferav1alpha1.InferaScalingAdapter, + st inferav1alpha1.InferaScalingAdapterStatus, +) (ctrl.Result, error) { + if equalStatus(a.Status, st) { + return ctrl.Result{}, nil + } + a.Status = st + return ctrl.Result{}, r.Status().Update(ctx, a) +} + +func equalStatus(a, b inferav1alpha1.InferaScalingAdapterStatus) bool { + if a.Replicas != b.Replicas || a.ReadyReplicas != b.ReadyReplicas || + a.Selector != b.Selector || a.ObservedGeneration != b.ObservedGeneration || + len(a.Conditions) != len(b.Conditions) { + return false + } + for i := range a.Conditions { + if a.Conditions[i].Type != b.Conditions[i].Type || + a.Conditions[i].Status != b.Conditions[i].Status || + a.Conditions[i].Reason != b.Conditions[i].Reason || + a.Conditions[i].Message != b.Conditions[i].Message { + return false + } + } + return true +} + +func setCondition( + conds *[]metav1.Condition, gen int64, typ string, + status metav1.ConditionStatus, reason, msg string, +) { + for i := range *conds { + if (*conds)[i].Type == typ { + c := &(*conds)[i] + if c.Status != status { + c.LastTransitionTime = metav1.Now() + } + c.Status, c.Reason, c.Message, c.ObservedGeneration = status, reason, msg, gen + return + } + } + *conds = append(*conds, metav1.Condition{ + Type: typ, Status: status, Reason: reason, Message: msg, + ObservedGeneration: gen, LastTransitionTime: metav1.Now(), + }) +} + +func (r *InferaScalingAdapterReconciler) SetupWithManager(mgr ctrl.Manager) error { + // Watching the InferaDeployment matters as much as the adapter itself: a + // scale write only changes `spec.replicas` here, and the workload does not + // move until the other reconciler runs. Without this the adapter's status + // would lag by a resync period after every scale. + return ctrl.NewControllerManagedBy(mgr). + For(&inferav1alpha1.InferaScalingAdapter{}). + Watches( + &inferav1alpha1.InferaDeployment{}, + handler.EnqueueRequestsFromMapFunc(r.adaptersForDeployment), + ). + Complete(r) +} + +func (r *InferaScalingAdapterReconciler) adaptersForDeployment( + ctx context.Context, obj client.Object, +) []reconcile.Request { + list := &inferav1alpha1.InferaScalingAdapterList{} + if err := r.List(ctx, list, client.InNamespace(obj.GetNamespace())); err != nil { + return nil + } + var out []reconcile.Request + for i := range list.Items { + if list.Items[i].Spec.DeploymentRef != obj.GetName() { + continue + } + out = append(out, reconcile.Request{NamespacedName: types.NamespacedName{ + Name: list.Items[i].Name, Namespace: list.Items[i].Namespace, + }}) + } + return out +} diff --git a/manual/features/scaling.md b/manual/features/scaling.md index 23dd9b01..7fe30027 100644 --- a/manual/features/scaling.md +++ b/manual/features/scaling.md @@ -361,18 +361,66 @@ path has not been exercised on hardware. ## Autoscaling -Infera does not ship an autoscaler, and an external one cannot currently drive -an `InferaDeployment`: the operator reconciles `replicas` from the CR on every -pass, so a `HorizontalPodAutoscaler` writing to the child Deployment is reverted -within seconds. Scaling today is a deliberate act — `kubectl scale` on the CR, -or starting and stopping workers. - -The mechanics an autoscaler would need are in place: workers join and leave -cleanly under load, and the signals worth scaling on -(`vllm:num_requests_waiting`, `sglang:num_queue_reqs`, KV cache utilisation) are -exposed by the engines and read by the drain path already. - -The unsolved part is not the plumbing. It is that a **140-second cold start sits -inside a control loop that ticks every 15 seconds**, and that nothing in -Kubernetes lets a scaler choose *which* replica to remove — so the one holding -the warmest KV cache is as likely to go as any other. +Infera ships no autoscaler. It does ship the standard interface one drives. + +### The `/scale` subresource + +An `InferaDeployment` cannot carry `/scale` itself, and that is a property of +its shape rather than a missing feature: `spec.services` is a map with +user-chosen keys, while the scale subresource requires `specReplicasPath` to be +a *static* dot-notation JSONPath. There is no way to name "the replicas of an +arbitrary map entry". + +So scaling gets its own object — one `InferaScalingAdapter` per scalable +service: + +```yaml +apiVersion: infera.amd.com/v1alpha1 +kind: InferaScalingAdapter +metadata: {name: qwen-decode, namespace: infera} +spec: + deploymentRef: qwen # the InferaDeployment + serviceName: decode # a key in its spec.services + replicas: 2 +``` + +```bash +kubectl scale inferascalingadapter/qwen-decode --replicas=5 +``` + +That is the whole integration. `kubectl scale`, `HorizontalPodAutoscaler`, KEDA +and a custom planner all work through it with no per-tool support in the +operator — a `scaleTargetRef` of kind `InferaScalingAdapter` is all an HPA +needs. + +**One writer, always.** While an adapter has `spec.replicas` set it owns that +service's count: the InferaDeployment reconciler reads the adapter instead of +the CR's own `replicas`. Clearing `spec.replicas`, or deleting the adapter, hands +control back. Creating an adapter without `spec.replicas` is deliberately inert, +so an autoscaler can be attached and watched before it is trusted. + +This matters because the reconciler assigns the whole child `.Spec` on every +pass. Anything else writing that Deployment loses — which is exactly what +happens to an HPA pointed straight at it: **measured, a `kubectl scale` to 3 was +reverted to 1 in under 3 seconds.** Going through the adapter is not a +convention, it is the only thing that survives. + +`status.replicas` reports what the workload *observes*, not the desired count +echoed back. An autoscaler computes `desired = ceil(current × metric/target)`; +if `current` were the number it just asked for, it could not tell a scale-up had +not landed and would keep multiplying through a 140-second model load. + +### What is still missing + +The plumbing is not the hard part. Two things are: + +- **A 140-second cold start sits inside a control loop that ticks every 15 + seconds.** A burst shorter than the cold start cannot be answered by adding + workers at all. +- **Nothing in Kubernetes lets a scaler choose *which* replica to remove**, so + the one holding the warmest KV cache is as likely to go as any other. Upstream + has declined to fix this (k8s#123541, closed as not planned). + +The signals worth scaling on (`vllm:num_requests_waiting`, +`sglang:num_queue_reqs`, KV utilisation) are exposed by the engines and already +read by the drain path, but nothing polls them continuously yet. From cb5f913e37b7db027e1126b895413b9ae91c55b7 Mon Sep 17 00:00:00 2001 From: liyingli Date: Wed, 5 Aug 2026 08:45:39 +0000 Subject: [PATCH 21/88] fix(kvd): resolve a mount's devices through sysfs when lsblk cannot lsblk answers by device name, and two shapes we run in have no name it can open. On LVM, findmnt hands back /dev/mapper/-, a symlink udev creates; the forward walk reports the leaf as -, and the --inverse walk that would have supplied the transport is handed /dev/-, which is not a path that exists. md never showed this -- md0 and /dev/md0 agree -- and md is what the walk was written against, as its docstring says. Inside a container it fails a step earlier: /dev is a private devtmpfs holding only what --device named. privileged does not rescue LVM either, because the host devtmpfs it exposes carries /dev/dm-N but not the /dev/mapper/ symlinks, those being udev's work and containers not running udev. Either way the probe answers "unknown device" and the region takes the conservative buffered branch: 3.70 GB/s against the 14.56 GB/s O_DIRECT measured on the same gfx942 mount, a 4x cut paid for a naming detail, and one that until now had to be worked around with an explicit --io-mode direct. Walk sysfs when lsblk returns nothing or returns devices it could not put a transport on. sysfs is indexed by the major:minor that stat() already gives us rather than by name, so it needs no device node, no udev and no subprocess, and the block layer is not namespaced, so an unprivileged container sees the full topology. lsblk stays the fast path -- the transport table was calibrated against its TRAN column -- and a complete answer from it is never second-guessed. Verified on an 8-NVMe LVM: auto now resolves all eight members and picks DIRECT both on the host and from an unprivileged container, where it picked buffered in both before. Co-authored-by: Cursor Signed-off-by: liyingli --- infera/kvd/storage_classify.py | 190 +++++++++++++++- tests/unit/kvd/test_storage_classify.py | 277 ++++++++++++++++++++++++ 2 files changed, 456 insertions(+), 11 deletions(-) diff --git a/infera/kvd/storage_classify.py b/infera/kvd/storage_classify.py index 73e93c7d..df10d13d 100644 --- a/infera/kvd/storage_classify.py +++ b/infera/kvd/storage_classify.py @@ -46,16 +46,24 @@ ``max(2, cpu_count() // n_shards)`` so 8-shard configs on small boxes never spin up more threads than the box can productively schedule. -The probe is a chain of cheap subprocess calls — ``findmnt`` and -``lsblk`` — both shipped by util-linux on every modern distro. If -either binary is missing (containers, exotic systems) we fall back to -the conservative "buffered" + workers=4 default and log a WARN. - -The chain handles md-raid, LVM, dm-crypt, and bind mounts transparently -by walking ``lsblk -no NAME,TRAN,ROTA`` recursively. For mixed-device -arrays (e.g. NVMe + SATA in the same md0) the worst-case device wins -— a single SATA member in an mdraid pulls the whole array into the -"buffered" bucket. +The probe starts from ``findmnt`` to name the mount, then resolves that +mount to its physical devices two ways. ``lsblk -no NAME,TRAN,ROTA`` is +tried first, since the transport table above was calibrated against its +TRAN column. It answers by device *name*, though, which fails on an LVM +mount (the name findmnt hands back is a udev symlink) and inside a +container (``/dev`` holds only what ``--device`` put there). So whenever +lsblk comes back empty, or with devices it could not name a transport +for, we walk sysfs instead: ``major:minor`` from ``stat()`` into +``/sys/dev/block``, then down ``slaves/`` to the physical disks. sysfs +needs no device node, no udev and no subprocess, and the block layer is +not namespaced, so an unprivileged container sees the whole topology. + +Between them the two cover md-raid, LVM, dm-crypt, partitions and bind +mounts. If both come up empty we fall back to the conservative +"buffered" + workers=4 default and log a WARN. For mixed-device arrays +(e.g. NVMe + SATA in the same md0) the worst-case device wins — a single +SATA member in an mdraid pulls the whole array into the "buffered" +bucket. Public API: @@ -323,6 +331,151 @@ def _lsblk_inverse_parent(devpath: str) -> tuple[str, bool | None]: return "", None +# ---------------------------------------------------------------------- +# sysfs device walk — the fallback for everything lsblk cannot answer +# ---------------------------------------------------------------------- +# +# lsblk resolves a device by *name*, which is where it loses two cases we +# actually run in: +# +# - **LVM.** ``findmnt`` reports ``/dev/mapper/-``, a symlink udev +# creates; the kernel's own node is ``/dev/dm-N``. The forward walk then +# reports the leaf as ``-``, and ``/dev/-`` is not a path +# that exists, so the ``--inverse`` walk that would have supplied the +# transport dead-ends. md is unaffected — ``md0`` and ``/dev/md0`` agree — +# which is why this went unnoticed. +# - **Containers.** ``/dev`` is a private, near-empty devtmpfs; Docker only +# populates what ``--device`` names. lsblk cannot open the node at all. +# ``privileged: true`` does not rescue LVM either: the host devtmpfs it +# exposes carries ``/dev/dm-N`` but not the ``/dev/mapper/`` symlinks, +# because those are udev's work and containers do not run udev. +# +# sysfs sidesteps both. It is indexed by ``major:minor`` rather than by name, +# and those numbers come straight from ``stat()`` — no node, no udev, no +# subprocess. It is also global: the block layer is not namespaced, so a +# container sees the full topology even unprivileged. + +# Every sysfs lookup below hangs off this one root so the tests can aim the +# whole walk at a fixture tree instead of the machine they happen to run on. +_SYSFS_ROOT = "/sys" + + +def _sysfs_read(path: str) -> str: + """Read a sysfs attribute; "" if it is missing or unreadable.""" + try: + with open(path) as f: + return f.read().strip() + except OSError: + return "" + + +def _sysfs_whole_disk(name: str) -> str: + """Resolve a partition to the disk that carries it. + + Partitions have neither ``queue/`` nor ``device/`` of their own — both + live on the whole-disk node one level up (``.../block/sda/sda1`` → ``sda``). + Non-partitions are returned unchanged. + """ + real = os.path.realpath(os.path.join(f"{_SYSFS_ROOT}/class/block", name)) + if os.path.exists(os.path.join(real, "partition")): + return os.path.basename(os.path.dirname(real)) + return name + + +def _sysfs_transport(name: str) -> str: + """Best-effort transport for a leaf device, mirroring lsblk's TRAN. + + Returns "" when the bus cannot be identified, which the callers already + treat as "unknown → buffered". Since this whole path only runs after + lsblk has failed to answer, an unrecognised bus is no worse than the + status quo — every transport we *do* resolve is a strict improvement. + """ + disk = _sysfs_whole_disk(name) + devlink = os.path.join(f"{_SYSFS_ROOT}/class/block", disk, "device") + if not os.path.exists(devlink): + # No backing device: dm/md nodes with no slaves, loop, zram, ramdisk. + return "" + devpath = os.path.realpath(devlink) + subsystem = os.path.basename(os.path.realpath(os.path.join(devlink, "subsystem"))) + if subsystem in ("nvme", "virtio", "mmc"): + return subsystem + if subsystem != "scsi": + return "" + # SCSI multiplexes every serial bus, so the subsystem alone says nothing. + # util-linux keys off the host's proc_name; we follow it, then fall back + # to the shape of the device path for the buses that do not set one. + m = re.search(r"/host(\d+)/", devpath + "/") + proc = _sysfs_read(f"{_SYSFS_ROOT}/class/scsi_host/host{m.group(1)}/proc_name") if m else "" + if proc.startswith("iscsi"): + return "iscsi" + if os.path.exists(os.path.join(devpath, "sas_address")): + return "sas" + if "/usb" in devpath: + return "usb" + if proc in ("ahci", "ata_piix") or proc.startswith(("sata_", "pata_")) or "/ata" in devpath: + return "sata" + fc_host = f"{_SYSFS_ROOT}/class/fc_host" + if os.path.isdir(fc_host) and os.listdir(fc_host): + return "fc" + return "" + + +def _sysfs_leaves(name: str, seen: set[str] | None = None) -> list[str]: + """Descend ``slaves/`` until devices that have none — the physical disks. + + ``slaves/`` is how the block layer records "this virtual device is built + on those"; dm stacks (LVM, dm-crypt) and md arrays both populate it, and + nesting them just deepens the recursion. ``seen`` guards against a device + appearing under two parents (an LVM RAID mirrors each leg through its own + rimage) rather than against a real cycle, which the kernel forbids. + """ + seen = seen if seen is not None else set() + if name in seen: + return [] + seen.add(name) + slaves = os.path.join(f"{_SYSFS_ROOT}/class/block", name, "slaves") + try: + children = sorted(os.listdir(slaves)) if os.path.isdir(slaves) else [] + except OSError: + children = [] + if not children: + return [name] + out: list[str] = [] + for child in children: + out += _sysfs_leaves(child, seen) + return out + + +def _sysfs_rotational(name: str) -> bool: + """``queue/rotational`` for a leaf, read off its whole disk.""" + disk = _sysfs_whole_disk(name) + return _sysfs_read(f"{_SYSFS_ROOT}/class/block/{disk}/queue/rotational") == "1" + + +def _sysfs_for_path(path: Path) -> list[DeviceInfo]: + """Resolve ``path`` to its physical devices through sysfs alone. + + Returns [] when sysfs cannot answer — no /sys mounted, or a filesystem + with no block device behind it (NFS and friends report a synthetic + ``st_dev`` that has no ``/sys/dev/block`` entry). Never raises. + """ + try: + st = os.stat(path) + devno = f"{os.major(st.st_dev)}:{os.minor(st.st_dev)}" + link = f"{_SYSFS_ROOT}/dev/block/{devno}" + if not os.path.exists(link): + return [] + top = os.path.basename(os.path.realpath(link)) + # dict.fromkeys dedupes while keeping the walk order stable. + return [ + DeviceInfo(dev=d, transport=_sysfs_transport(d), rotational=_sysfs_rotational(d)) + for d in dict.fromkeys(_sysfs_leaves(top)) + ] + except OSError as exc: + logger.debug("storage_classify: sysfs walk of %s failed: %s", path, exc) + return [] + + def _nfs_mount_opts(source: str, path: Path) -> str: """Return the raw mount-options string for the NFS mount covering ``path``. Empty string if /proc/mounts isn't readable or no NFS @@ -429,9 +582,24 @@ def classify_storage(path: Path) -> StorageInfo: if fstype in _NO_DIRECT_FSTYPES: return info devices = _lsblk_for_source(source) + # lsblk is kept as the fast path because it is what the transport table + # above was calibrated against, but it answers by name and so cannot see + # through an LVM mount or into a container's empty /dev. Take the sysfs + # answer whenever lsblk returned nothing, or returned devices it could not + # put a transport on — both of those end at "conservative buffered", and + # on NVMe that is a ~4x throughput cut taken for a naming detail. + if not devices or any(not d.transport for d in devices): + from_sysfs = _sysfs_for_path(path) + if from_sysfs and any(d.transport for d in from_sysfs): + logger.debug( + "storage_classify: lsblk could not classify source=%r; sysfs resolved %d device(s)", + source, + len(from_sysfs), + ) + devices = from_sysfs if not devices: info.warnings.append( - f"lsblk returned no devices for source={source!r}; defaulting to buffered" + f"neither lsblk nor sysfs found devices for source={source!r}; defaulting to buffered" ) info.devices = devices return info diff --git a/tests/unit/kvd/test_storage_classify.py b/tests/unit/kvd/test_storage_classify.py index ee2e1529..bc57c26a 100644 --- a/tests/unit/kvd/test_storage_classify.py +++ b/tests/unit/kvd/test_storage_classify.py @@ -17,6 +17,7 @@ from __future__ import annotations +import os import subprocess from pathlib import Path @@ -750,3 +751,279 @@ def test_format_workers_decision_nfs_includes_nconnect(fake_run, fake_proc_mount out = format_workers_decision(info, workers, rationale) assert "nconnect" in out assert "8" in out + + +# ---------------------------------------------------------------------- +# sysfs device walk — the path that covers what lsblk cannot resolve. +# +# Every case here builds a real directory tree shaped like sysfs and points +# the module at it, so the walk is exercised end to end (symlinks, slaves/ +# recursion, partition→disk resolution) without depending on whatever disks +# the test machine happens to have. +# ---------------------------------------------------------------------- + + +@pytest.fixture +def fake_sysfs(monkeypatch, tmp_path): + """Build a sysfs fixture tree and aim the module's walk at it. + + ``devices`` maps each block device name to its shape: + + rotational "0" / "1" ``queue/rotational``; whole disks only, + since that is where the kernel puts it + slaves [names] the devices this one is built on + subsystem "nvme"/... the bus its backing device sits on + devpath "a/b/c" where under ``/sys/devices`` that lands — + this is what carries ``/hostN/`` and ``/usb`` + sas_address True the attribute that marks a SAS disk + parent_disk "sda" makes this device a partition of that disk + + ``scsi_hosts`` maps a host name to its ``proc_name``, the attribute + util-linux reads to tell iSCSI and SATA apart. + + Returns the path to probe: ``tmp_path`` itself, registered in the fake + ``/sys/dev/block`` under its own real ``st_dev`` so ``stat()`` lands on + the fixture topology on any machine. + """ + + def build(devices: dict, target: str, scsi_hosts: dict | None = None) -> Path: + root = tmp_path / "sys" + blk = root / "class" / "block" + blk.mkdir(parents=True) + (root / "dev" / "block").mkdir(parents=True) + devices_dir = root / "devices" + + def node_dir(name: str, spec: dict) -> Path: + parent = spec.get("parent_disk") + if not parent: + d = blk / name + d.mkdir(parents=True, exist_ok=True) + return d + # A partition lives inside its disk's directory, and /sys/class + # only links to it — which is exactly what the walk relies on to + # find the disk that carries the queue/ and device/ attributes. + d = devices_dir / "block" / parent / name + d.mkdir(parents=True, exist_ok=True) + (d / "partition").write_text("1\n") + if not (blk / name).exists(): + (blk / name).symlink_to(d) + return d + + # Whole disks first: a partition's entry has to be able to point into + # its parent's directory. + ordered = sorted(devices.items(), key=lambda kv: bool(kv[1].get("parent_disk"))) + for name, spec in ordered: + d = node_dir(name, spec) + if "rotational" in spec: + (d / "queue").mkdir(exist_ok=True) + (d / "queue" / "rotational").write_text(spec["rotational"] + "\n") + for child in spec.get("slaves", []): + (d / "slaves").mkdir(exist_ok=True) + (d / "slaves" / child).mkdir(exist_ok=True) + subsystem = spec.get("subsystem") + if subsystem: + phys = devices_dir / spec.get("devpath", f"pci0000:00/{name}") + phys.mkdir(parents=True, exist_ok=True) + bus = root / "bus" / subsystem + bus.mkdir(parents=True, exist_ok=True) + if not (phys / "subsystem").exists(): + (phys / "subsystem").symlink_to(bus) + if spec.get("sas_address"): + (phys / "sas_address").write_text("0x5000c500a1b2c3d4\n") + if not (d / "device").exists(): + (d / "device").symlink_to(phys) + + for host, proc_name in (scsi_hosts or {}).items(): + hd = root / "class" / "scsi_host" / host + hd.mkdir(parents=True, exist_ok=True) + (hd / "proc_name").write_text(proc_name + "\n") + + probe = tmp_path / "probe" + probe.mkdir(exist_ok=True) + st = probe.stat() + devno = f"{os.major(st.st_dev)}:{os.minor(st.st_dev)}" + (root / "dev" / "block" / devno).symlink_to(blk / target) + + monkeypatch.setattr(storage_classify, "_SYSFS_ROOT", str(root)) + return probe + + return build + + +def test_sysfs_walks_lvm_stack_down_to_nvme_members(fake_sysfs): + """The case lsblk cannot do: the logical volume's name is not a path, + so only a major:minor lookup reaches the members.""" + probe = fake_sysfs( + { + "dm-8": {"slaves": ["dm-0", "dm-1"], "rotational": "0"}, + "dm-0": {"slaves": ["nvme0n1"], "rotational": "0"}, + "dm-1": {"slaves": ["nvme1n1"], "rotational": "0"}, + "nvme0n1": {"subsystem": "nvme", "rotational": "0"}, + "nvme1n1": {"subsystem": "nvme", "rotational": "0"}, + }, + target="dm-8", + ) + devices = storage_classify._sysfs_for_path(probe) + assert [d.dev for d in devices] == ["nvme0n1", "nvme1n1"] + assert all(d.transport == "nvme" and not d.rotational for d in devices) + + +def test_sysfs_walks_md_stack(fake_sysfs): + probe = fake_sysfs( + { + "md0": {"slaves": ["sda", "sdb"], "rotational": "0"}, + "sda": {"subsystem": "scsi", "devpath": "pci0000:00/ata1/host1/sda", "rotational": "0"}, + "sdb": {"subsystem": "scsi", "devpath": "pci0000:00/ata2/host2/sdb", "rotational": "0"}, + }, + target="md0", + scsi_hosts={"host1": "ahci", "host2": "ahci"}, + ) + devices = storage_classify._sysfs_for_path(probe) + assert [d.dev for d in devices] == ["sda", "sdb"] + assert all(d.transport == "sata" for d in devices) + + +def test_sysfs_dedupes_a_leaf_reached_through_two_parents(fake_sysfs): + """An LVM RAID mirrors each leg through its own rimage, so the same + physical disk can be reachable more than once.""" + probe = fake_sysfs( + { + "dm-9": {"slaves": ["dm-2", "dm-3"], "rotational": "0"}, + "dm-2": {"slaves": ["nvme0n1"], "rotational": "0"}, + "dm-3": {"slaves": ["nvme0n1"], "rotational": "0"}, + "nvme0n1": {"subsystem": "nvme", "rotational": "0"}, + }, + target="dm-9", + ) + assert [d.dev for d in storage_classify._sysfs_for_path(probe)] == ["nvme0n1"] + + +def test_sysfs_resolves_a_partition_to_its_parent_disk(fake_sysfs): + """Partitions carry neither queue/ nor device/ — both belong to the disk.""" + probe = fake_sysfs( + { + "sda": { + "subsystem": "scsi", + "devpath": "platform/host0/session1/target0:0:0/0:0:0:1", + "rotational": "1", + }, + "sda1": {"parent_disk": "sda"}, + }, + target="sda1", + scsi_hosts={"host0": "iscsi_tcp"}, + ) + devices = storage_classify._sysfs_for_path(probe) + assert [(d.dev, d.transport, d.rotational) for d in devices] == [("sda1", "iscsi", True)] + + +def test_sysfs_reads_sas_from_the_device_attribute(fake_sysfs): + probe = fake_sysfs( + { + "sdc": { + "subsystem": "scsi", + "devpath": "pci0000:00/host3/sdc", + "sas_address": True, + "rotational": "0", + } + }, + target="sdc", + scsi_hosts={"host3": "mpt3sas"}, + ) + assert storage_classify._sysfs_for_path(probe)[0].transport == "sas" + + +def test_sysfs_returns_nothing_when_the_devno_is_not_registered(fake_sysfs, tmp_path): + """A filesystem with no block device behind it — NFS reports a synthetic + st_dev that has no /sys/dev/block entry.""" + fake_sysfs({"nvme0n1": {"subsystem": "nvme", "rotational": "0"}}, target="nvme0n1") + monkey_root = Path(str(tmp_path / "sys" / "dev" / "block")) + for entry in monkey_root.iterdir(): + entry.unlink() + assert storage_classify._sysfs_for_path(tmp_path / "probe") == [] + + +def test_sysfs_leaves_transport_blank_for_an_unrecognised_bus(fake_sysfs): + """Unknown is still a useful answer — it keeps the conservative default + rather than guessing O_DIRECT onto something that would hate it.""" + probe = fake_sysfs( + {"xvda": {"subsystem": "xen", "rotational": "0"}}, + target="xvda", + ) + devices = storage_classify._sysfs_for_path(probe) + assert [(d.dev, d.transport) for d in devices] == [("xvda", "")] + + +# --- and the same thing end to end, through pick_io_mode --------------- + + +def _lvm_on_nvme(fake_sysfs): + return fake_sysfs( + { + "dm-8": {"slaves": ["nvme0n1"], "rotational": "0"}, + "nvme0n1": {"subsystem": "nvme", "rotational": "0"}, + }, + target="dm-8", + ) + + +def test_lsblk_failing_outright_is_recovered_by_sysfs(fake_run, fake_sysfs): + """Inside a container /dev holds only what --device put there, so lsblk + cannot open the node at all. Before the sysfs fallback this mount was + classified 'unknown device' and ran buffered on NVMe.""" + probe = _lvm_on_nvme(fake_sysfs) + fake_run( + { + "findmnt": "/dev/mapper/nvme_vg-nvme_lv xfs\n", + "lsblk": ("lsblk: not a block device\n", 32), + } + ) + o_direct, rationale = pick_io_mode(probe) + assert o_direct is True + assert "nvme" in rationale + + +def test_lsblk_answering_without_a_transport_is_recovered_by_sysfs(fake_run, fake_sysfs): + """On the host lsblk does resolve the LV, but the --inverse walk that + would supply the transport is handed /dev/-, which is not a + path that exists. The row comes back bare.""" + probe = _lvm_on_nvme(fake_sysfs) + fake_run( + { + "findmnt": "/dev/mapper/nvme_vg-nvme_lv xfs\n", + "lsblk": "nvme_vg-nvme_lv 0\n", # blank TRAN, and no inverse answer + } + ) + info = classify_storage(probe) + assert [d.dev for d in info.devices] == ["nvme0n1"] + assert pick_io_mode(probe)[0] is True + + +def test_a_complete_lsblk_answer_is_not_second_guessed(fake_run, fake_sysfs): + """sysfs is a fallback, not an override: when lsblk names a transport it + wins, even where the fixture topology would say something else.""" + probe = _lvm_on_nvme(fake_sysfs) # sysfs here would say nvme → O_DIRECT + fake_run( + { + "findmnt": "/dev/sdb1 ext4\n", + "lsblk": "sdb sata 0\n", + } + ) + o_direct, rationale = pick_io_mode(probe) + assert o_direct is False + assert "sata" in rationale + + +def test_both_probes_failing_still_lands_on_conservative_buffered(fake_run, monkeypatch, tmp_path): + """No lsblk answer and no sysfs either — the posture that protects an + HDD or a SAN from having O_DIRECT guessed onto it.""" + fake_run( + { + "findmnt": "/dev/mapper/vg-lv xfs\n", + "lsblk": ("", 32), + } + ) + monkeypatch.setattr(storage_classify, "_SYSFS_ROOT", str(tmp_path / "no-sysfs-here")) + info = classify_storage(tmp_path) + assert info.devices == [] + assert any("sysfs" in w for w in info.warnings) + assert pick_io_mode(tmp_path) == (False, "unknown device, conservative buffered") From 8a9289eb77db9f857696631174b73bd1208fcbc9 Mon Sep 17 00:00:00 2001 From: liyingli Date: Wed, 5 Aug 2026 10:29:14 +0000 Subject: [PATCH 22/88] fix(kvd): scope the sysfs FC check to the device being classified The check read machine-wide state: any fc_host anywhere in the box made every otherwise-unclassified SCSI disk answer "fc", so a local RAID volume behind a megaraid controller came back labelled a SAN. The decision never changed -- fc and unknown are both buffered -- but the rationale did, and a startup line calling a local disk a SAN is the same kind of misdirection the io_mode reporting fix went after. Ask instead whether *this* device's SCSI host is the FC HBA, which is what util-linux does: /sys/class/fc_host/host for the H out of the device's own path. Dropping the check would have worked too, since unknown is buffered either way, but then a real SAN loses its label and picks up the "unknown transport" WARN that pick_io_mode logs. Reported in review. Verified unchanged against the two transports this host can show: /dev/sda still resolves iscsi and /dev/nvme0n1 still resolves nvme, both matching lsblk. Co-authored-by: Cursor Signed-off-by: liyingli --- infera/kvd/storage_classify.py | 20 ++++++++----- tests/unit/kvd/test_storage_classify.py | 40 +++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 7 deletions(-) diff --git a/infera/kvd/storage_classify.py b/infera/kvd/storage_classify.py index df10d13d..4ce04092 100644 --- a/infera/kvd/storage_classify.py +++ b/infera/kvd/storage_classify.py @@ -389,6 +389,10 @@ def _sysfs_transport(name: str) -> str: treat as "unknown → buffered". Since this whole path only runs after lsblk has failed to answer, an unrecognised bus is no worse than the status quo — every transport we *do* resolve is a strict improvement. + + Every check below is a per-device one, and answers the same way lsblk + would. There is deliberately no guessing beyond that: a bus we cannot + name from this device's own sysfs entries stays "", and "" is buffered. """ disk = _sysfs_whole_disk(name) devlink = os.path.join(f"{_SYSFS_ROOT}/class/block", disk, "device") @@ -402,10 +406,15 @@ def _sysfs_transport(name: str) -> str: if subsystem != "scsi": return "" # SCSI multiplexes every serial bus, so the subsystem alone says nothing. - # util-linux keys off the host's proc_name; we follow it, then fall back - # to the shape of the device path for the buses that do not set one. - m = re.search(r"/host(\d+)/", devpath + "/") - proc = _sysfs_read(f"{_SYSFS_ROOT}/class/scsi_host/host{m.group(1)}/proc_name") if m else "" + # util-linux keys off the SCSI host this device hangs from; we follow it, + # then fall back to the shape of the device's own path for the buses that + # set no proc_name. Every one of these is scoped to this device — the + # host number comes out of its path, and the path is its own. + host = re.search(r"/host(\d+)/", devpath + "/") + hostn = host.group(1) if host else "" + proc = _sysfs_read(f"{_SYSFS_ROOT}/class/scsi_host/host{hostn}/proc_name") if hostn else "" + if hostn and os.path.isdir(f"{_SYSFS_ROOT}/class/fc_host/host{hostn}"): + return "fc" if proc.startswith("iscsi"): return "iscsi" if os.path.exists(os.path.join(devpath, "sas_address")): @@ -414,9 +423,6 @@ def _sysfs_transport(name: str) -> str: return "usb" if proc in ("ahci", "ata_piix") or proc.startswith(("sata_", "pata_")) or "/ata" in devpath: return "sata" - fc_host = f"{_SYSFS_ROOT}/class/fc_host" - if os.path.isdir(fc_host) and os.listdir(fc_host): - return "fc" return "" diff --git a/tests/unit/kvd/test_storage_classify.py b/tests/unit/kvd/test_storage_classify.py index bc57c26a..45bd4e02 100644 --- a/tests/unit/kvd/test_storage_classify.py +++ b/tests/unit/kvd/test_storage_classify.py @@ -932,6 +932,46 @@ def test_sysfs_reads_sas_from_the_device_attribute(fake_sysfs): assert storage_classify._sysfs_for_path(probe)[0].transport == "sas" +def test_sysfs_does_not_name_a_bus_from_machine_wide_state(fake_sysfs, tmp_path): + """A RAID card presenting a logical volume supplies no per-device signal, + so it has to stay unknown. Answering from something machine-wide instead — + "this box has an fc_host, so call it fc" — labels a local disk a SAN. + """ + probe = fake_sysfs( + { + "sdf": { + "subsystem": "scsi", + "devpath": "pci0000:00/host2/target2:2:0/2:2:0:0", + "rotational": "0", + } + }, + target="sdf", + scsi_hosts={"host2": "megaraid_sas"}, + ) + # An FC HBA elsewhere in the box, on a host this disk has nothing to do with. + (tmp_path / "sys" / "class" / "fc_host" / "host9").mkdir(parents=True) + assert storage_classify._sysfs_for_path(probe)[0].transport == "" + + +def test_sysfs_names_fc_when_the_hba_is_this_disk_s_own_host(fake_sysfs, tmp_path): + """The other half of the pair: a real SAN is worth naming, because + 'SAN fc → buffered' reads a great deal better in the startup log than an + unknown transport, which also logs a WARN.""" + probe = fake_sysfs( + { + "sdg": { + "subsystem": "scsi", + "devpath": "pci0000:00/host4/rport-4:0-0/target4:0:0/4:0:0:0", + "rotational": "0", + } + }, + target="sdg", + scsi_hosts={"host4": "lpfc"}, + ) + (tmp_path / "sys" / "class" / "fc_host" / "host4").mkdir(parents=True) + assert storage_classify._sysfs_for_path(probe)[0].transport == "fc" + + def test_sysfs_returns_nothing_when_the_devno_is_not_registered(fake_sysfs, tmp_path): """A filesystem with no block device behind it — NFS reports a synthetic st_dev that has no /sys/dev/block entry.""" From b0a1799b0c7003e310ca16b4b044122e2e3bdd30 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Fri, 7 Aug 2026 13:09:10 +0000 Subject: [PATCH 23/88] build(tools): check gofmt in pre-commit `gofmt -l` prints the files it would rewrite and still exits 0, so a bare `entry: gofmt -l` enforces nothing. This tree had accumulated Go files main had never formatted, noticed only when something unrelated ran gofmt over the directory; the two label maps here are that backlog. Checks and never rewrites: a hook that formats mid-commit leaves the staged and working copies disagreeing, so the diff you reviewed is not the one you commit. Co-authored-by: Cursor Signed-off-by: leiwei12 --- .pre-commit-config.yaml | 8 ++++ deploy/operator/internal/controller/gaie.go | 2 +- deploy/operator/internal/controller/nats.go | 4 +- scripts/check-gofmt.sh | 45 +++++++++++++++++++++ 4 files changed, 56 insertions(+), 3 deletions(-) create mode 100755 scripts/check-gofmt.sh diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5abd08ba..d3750dcd 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -39,6 +39,14 @@ repos: files: ^rust/.*\.rs$ pass_filenames: false + # Go formatting gate (needs a local toolchain). Only the operator is Go. + # `go vet` and the tests are heavier, so they run in CI. + - id: gofmt + name: gofmt (go) + entry: scripts/check-gofmt.sh + language: script + files: ^deploy/operator/.*\.go$ + # Refuse a commit whose git email is a machine-generated local hostname # (git's fallback when user.email is unset), so an internal build-host # name can't leak into permanent public history. Skipped in CI. diff --git a/deploy/operator/internal/controller/gaie.go b/deploy/operator/internal/controller/gaie.go index 6cf3c2bd..adfe54db 100644 --- a/deploy/operator/internal/controller/gaie.go +++ b/deploy/operator/internal/controller/gaie.go @@ -237,7 +237,7 @@ func buildInferencePool(idep *inferav1alpha1.InferaDeployment) *unstructured.Uns "selector": map[string]any{ "matchLabels": map[string]any{ "infera.amd.com/deployment": idep.Name, - gaieFrontendLabel: "true", + gaieFrontendLabel: "true", }, }, "endpointPickerRef": map[string]any{ diff --git a/deploy/operator/internal/controller/nats.go b/deploy/operator/internal/controller/nats.go index 545f4f20..d234cfe4 100644 --- a/deploy/operator/internal/controller/nats.go +++ b/deploy/operator/internal/controller/nats.go @@ -19,8 +19,8 @@ import ( func natsLabels(idepName string) map[string]string { return map[string]string{ "app.kubernetes.io/managed-by": "infera-operator", - "infera.amd.com/deployment": idepName, - "infera.amd.com/component": "nats", + "infera.amd.com/deployment": idepName, + "infera.amd.com/component": "nats", } } diff --git a/scripts/check-gofmt.sh b/scripts/check-gofmt.sh new file mode 100755 index 00000000..e4387b20 --- /dev/null +++ b/scripts/check-gofmt.sh @@ -0,0 +1,45 @@ +#!/usr/bin/env bash +# pre-commit hook: refuse Go sources that gofmt would rewrite. +# +# Worth a script rather than an inline entry, because `gofmt -l` prints the +# offending files and *still exits 0*. A bare `entry: gofmt -l` therefore passes +# on every commit and enforces nothing — which is how this tree accumulated Go +# files main had never formatted, discovered only when something unrelated ran +# gofmt over the directory. +# +# Formatting is checked, never applied: a hook that rewrites files mid-commit +# leaves the staged and working copies disagreeing, and the diff you reviewed is +# not the diff you commit. +set -euo pipefail + +if ! command -v gofmt >/dev/null 2>&1; then + cat >&2 <<'MSG' +gofmt not found, but this commit touches Go sources. + +Install the Go toolchain (https://go.dev/dl/), or skip this one check with: + + SKIP=gofmt git commit ... +MSG + exit 1 +fi + +# pre-commit passes the staged files matching `files:`; nothing to do otherwise. +[ "$#" -eq 0 ] && exit 0 + +unformatted="$(gofmt -l "$@")" +[ -z "$unformatted" ] && exit 0 + +{ + echo + echo "Refusing the commit: gofmt would rewrite these files." + echo + printf ' %s\n' $unformatted + echo + echo "Format them in place, then re-stage:" + echo + # Unquoted on purpose: word-splitting rejoins the newline-separated list into + # a single space-separated command line. + echo " gofmt -w" $unformatted + echo +} >&2 +exit 1 From 73ef1e661e40cba6c734c9d4eb477ffc6d303f26 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Fri, 7 Aug 2026 13:09:58 +0000 Subject: [PATCH 24/88] fix(k8s): show a draining worker instead of dropping its record A condemned Pod left routing and vanished from /v1/workers in the same instant, indistinguishable from one that had already gone. The window it spends finishing in-flight generations is exactly when someone is watching, and it read as if the worker had died. It now leaves routing just as promptly but is marked DRAINING and keeps its record until the Pod actually goes. Removal is announced once -- on the mark, not again on the delete -- so a caller that already reacted is not told twice. Keeping the record makes a lost delete event permanent, where removal on sight used to bound it, so a re-list now drops what a full snapshot does not mention. That is routine rather than defensive: the re-list exists because etcd compaction expires the watch every few minutes, and a Pod deleted inside a reconnect window produces no event anyone sees. A phantom is out of routing by its status, but it would be reported forever and hold its model's tokenizer canary pinned against a genuinely different worker later. Co-authored-by: Cursor Signed-off-by: leiwei12 --- infera/common/discovery_k8s.py | 88 ++++++++++++++++-- .../common/test_discovery_k8s_terminating.py | 90 +++++++++++++++++++ 2 files changed, 169 insertions(+), 9 deletions(-) diff --git a/infera/common/discovery_k8s.py b/infera/common/discovery_k8s.py index 31613802..1b993a01 100644 --- a/infera/common/discovery_k8s.py +++ b/infera/common/discovery_k8s.py @@ -34,6 +34,7 @@ CanaryVerifier, WorkerInfo, WorkerPool, + WorkerStatus, ) logger = logging.getLogger(__name__) @@ -94,8 +95,13 @@ async def _relist(self) -> str | None: resp = await self._http.get(f"/api/v1/namespaces/{self._namespace}/pods", params=params) resp.raise_for_status() body = resp.json() + seen: set[str] = set() for pod in body.get("items", []) or []: + name = ((pod.get("metadata") or {}).get("name")) or "" + if name: + seen.add(name) self._handle_pod(pod, deleted=False) + self._reconcile_absent(seen) rv = (body.get("metadata") or {}).get("resourceVersion") logger.info( "k8s (re)list: %d worker(s) for selector %r in ns %s (rv=%s)", @@ -239,15 +245,23 @@ def _handle_pod(self, pod: dict, *, deleted: bool) -> None: annotations = meta.get("annotations") or {} raw = annotations.get(WORKER_INFO_ANNOTATION) - # Removal: explicit DELETE, deletion requested, pod no longer Running, - # or annotation gone. Terminating is checked separately from Running - # because a condemned Pod stays Running until its containers exit. - if deleted or raw is None or self._pod_terminating(pod) or not self._pod_running(pod): + # Gone for good: explicit DELETE, annotation cleared (the worker + # deregistered, so its drain is over), or no longer Running. + if deleted or raw is None or not self._pod_running(pod): worker_id = self._pod_to_worker.pop(pod_name, None) if worker_id is not None: self._remove(worker_id) return + # Condemned but still serving. Checked separately from Running because a + # terminating Pod stays Running until its containers exit; it leaves + # routing now and the record survives until it actually goes. + if self._pod_terminating(pod): + worker_id = self._pod_to_worker.get(pod_name) + if worker_id is not None: + self._mark_draining(worker_id) + return + try: info = worker_info_from_json(json.loads(raw)) except Exception as exc: @@ -295,17 +309,73 @@ def _handle_pod(self, pod: dict, *, deleted: bool) -> None: except Exception: logger.exception("on_worker_added callback failed") + def _reconcile_absent(self, seen: set[str]) -> None: + """Drop workers whose Pod is missing from a full list. + + A list is a complete snapshot, so anything still tracked that it does + not mention has been deleted -- and the event saying so was lost. That + is a routine occurrence rather than an edge case: the re-list exists + because etcd compaction expires the watch's resourceVersion every few + minutes, and any Pod deleted inside a reconnect window produces no + event anyone observes. + + It matters more now that a draining worker keeps its record. Removal + used to happen the instant a Pod was condemned, which bounded how long + a stale entry could survive; waiting for a later event instead makes a + missed one permanent. A phantom is filtered out of routing by its + DRAINING status, so nothing is dispatched to it -- but it is reported + by `/v1/workers` forever, and it keeps its model's tokenizer canary + pinned, which would reject a genuinely different worker later. + """ + for pod_name in [p for p in self._pod_to_worker if p not in seen]: + worker_id = self._pod_to_worker.pop(pod_name) + logger.info("k8s: pod %s absent from list; dropping worker %s", pod_name, worker_id) + self._remove(worker_id) + + def _mark_draining(self, worker_id: str) -> None: + """Take a condemned worker out of routing without dropping its record. + + ``list_active`` filters DRAINING, so this stops new work reaching the + worker just as removal would -- but ``list_all`` still shows it, and + that difference is the whole point. Deleting the record makes a Pod + that is finishing its in-flight generations look exactly like one that + crashed, so ``/v1/workers`` cannot distinguish an orderly rollout from + a fleet losing workers, at precisely the moment someone is watching. + + The record does not linger: the worker clears its own annotation when + the drain completes, which lands here as "annotation gone" and removes + it for real. + + Callbacks fire here rather than at that later removal because routing + is what they act on -- the KV subscriber and the policy's block + accounting must stop treating a departing worker as a target now, not + when its Pod object finally disappears. ``_remove`` therefore skips + them if it already announced, so each worker is announced exactly once. + """ + existing = self._pool.get(worker_id) + if existing is None or existing.status is WorkerStatus.DRAINING: + return # never registered, or already announced + existing.status = WorkerStatus.DRAINING + logger.info("k8s: worker %s draining (pod terminating, record kept)", worker_id) + self._notify_removed(worker_id) + def _remove(self, worker_id: str) -> None: existing = self._pool.get(worker_id) if existing is None: return + announced = existing.status is WorkerStatus.DRAINING self._pool.remove(worker_id) remaining = [w for w in self._pool.list_all() if w.model_name == existing.model_name] if not remaining: self._canary.forget(existing.model_name) logger.info("k8s: worker %s removed (pod deleted / not ready)", worker_id) - if self._on_removed is not None: - try: - self._on_removed(worker_id) - except Exception: - logger.exception("on_worker_removed callback failed") + if not announced: + self._notify_removed(worker_id) + + def _notify_removed(self, worker_id: str) -> None: + if self._on_removed is None: + return + try: + self._on_removed(worker_id) + except Exception: + logger.exception("on_worker_removed callback failed") diff --git a/tests/unit/common/test_discovery_k8s_terminating.py b/tests/unit/common/test_discovery_k8s_terminating.py index 8e3b05df..eeb91676 100644 --- a/tests/unit/common/test_discovery_k8s_terminating.py +++ b/tests/unit/common/test_discovery_k8s_terminating.py @@ -21,6 +21,7 @@ import json from infera.common.discovery_k8s import WORKER_INFO_ANNOTATION, KubernetesRegistry +from infera.common.worker_pool import WorkerStatus def _payload(worker_id: str = "10.0.0.1:8080") -> str: @@ -112,3 +113,92 @@ def test_other_removal_rules_still_hold(): reg._handle_pod(_pod(), deleted=False) reg._handle_pod(_pod(**kwargs), deleted=deleted) assert _ids(reg) == [], f"{label} must still deregister" + + +def _all(reg): + return {w.worker_id: w.status for w in reg.pool.list_all()} + + +def test_a_draining_worker_stays_visible(): + """Out of routing, still on the books. + + Dropping the record entirely makes a worker finishing its in-flight + generations look exactly like one that crashed, so `/v1/workers` cannot + tell an orderly rollout from a fleet losing workers -- at exactly the + moment someone is watching one happen. + """ + reg, _ = _registry() + reg._handle_pod(_pod(), deleted=False) + reg._handle_pod(_pod(terminating=True), deleted=False) + + assert _ids(reg) == [], "still must not be a routing candidate" + assert _all(reg) == {"10.0.0.1:8080": WorkerStatus.DRAINING} + + +def test_the_record_goes_when_the_drain_finishes(): + """The worker clears its own annotation once drained, which lands here as + 'annotation gone'. Without that the draining record would be immortal.""" + reg, removed = _registry() + reg._handle_pod(_pod(), deleted=False) + reg._handle_pod(_pod(terminating=True), deleted=False) + assert _all(reg), "precondition: the record survives the terminating event" + + reg._handle_pod(_pod(terminating=True, annotated=False), deleted=False) + assert _all(reg) == {}, "a drained worker must leave the pool" + assert removed == ["10.0.0.1:8080"], "announced once, not once per stage" + + +def test_announced_once_across_draining_then_delete(): + """The callbacks stop a KV subscriber and clear block accounting. Firing + them twice for one worker is not free, and firing them late (only at the + DELETE) would leave the router accounting for a worker it no longer routes + to for the whole drain.""" + reg, removed = _registry() + reg._handle_pod(_pod(), deleted=False) + reg._handle_pod(_pod(terminating=True), deleted=False) + assert removed == ["10.0.0.1:8080"], "must announce as soon as it leaves routing" + + reg._handle_pod(_pod(terminating=True), deleted=True) + assert _all(reg) == {} + assert removed == ["10.0.0.1:8080"], "the DELETE must not re-announce" + + +def test_a_worker_that_never_drained_still_announces_on_delete(): + """The drain path is not the only way out: a crash or an evicted Pod goes + straight to removal, and that still has to reach the callbacks.""" + reg, removed = _registry() + reg._handle_pod(_pod(), deleted=False) + reg._handle_pod(_pod(), deleted=True) + assert _all(reg) == {} + assert removed == ["10.0.0.1:8080"] + + +def test_a_list_reconciles_away_a_worker_whose_delete_was_missed(): + """A list is a complete snapshot, so a tracked Pod it does not mention is + gone. + + This is not hypothetical. Keeping a draining record alive means its removal + now depends on a later event, and the watch drops out routinely -- the + re-list exists precisely because etcd compaction expires the + resourceVersion every few minutes. A Pod deleted inside that window + produces no event anyone sees, so without reconciling against the list the + record is immortal: `/v1/workers` reports a worker that does not exist, and + the model's canary is never forgotten because a phantom still holds it. + """ + reg, removed = _registry() + reg._handle_pod(_pod(), deleted=False) + reg._handle_pod(_pod(terminating=True), deleted=False) + assert _all(reg), "precondition: the draining record is being kept" + + # The Pod is gone; a fresh list simply does not contain it. + reg._reconcile_absent(seen=set()) + + assert _all(reg) == {}, "a tracked Pod missing from a full list must be dropped" + assert removed == ["10.0.0.1:8080"], "already announced when it started draining" + + +def test_a_list_keeps_workers_it_still_sees(): + reg, _ = _registry() + reg._handle_pod(_pod(name="w-0"), deleted=False) + reg._reconcile_absent(seen={"w-0"}) + assert _ids(reg) == ["10.0.0.1:8080"], "a Pod present in the list must survive" From 8067f6c1bc265ed1f2d8c7f25dd3d77d01dc097f Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Fri, 7 Aug 2026 13:10:25 +0000 Subject: [PATCH 25/88] fix(nats): don't strand queued work or leave the consumer behind Under the throttle a request is pulled from a WorkQueue stream, so it can be accepted and queued while still invisible to `_inflight`, which only tracks what was delivered here. Draining closed the door first, and `unsubscribe()` discards what is left -- so a request the router had already handed us was never served and never refused. The client just waits out its idle timeout, 900s by default, for a reply nobody will send. The backlog is now handed over before the door closes. The durable consumer went with it. `unsubscribe()` only tears down the local subscription; a durable outlives it by definition, and its name derives from the worker id, which a rebuilt Pod never reuses because the IP changes. Left behind, every rollout stranded one more consumer holding WorkQueue quota that nothing would ever read -- measured against a live broker, three rollouts left three. Both waits now share one deadline. Charging each the full --drain-timeout would let a worker outlive the grace period it was sized for and be SIGKILLed mid-generation, which is the outcome draining exists to avoid. Co-authored-by: Cursor Signed-off-by: leiwei12 --- infera/common/nats_request.py | 73 +++++++++++++++- tests/unit/common/test_nats_drain.py | 123 ++++++++++++++++++++++++++- 2 files changed, 191 insertions(+), 5 deletions(-) diff --git a/infera/common/nats_request.py b/infera/common/nats_request.py index 9e7cad40..0894d951 100644 --- a/infera/common/nats_request.py +++ b/infera/common/nats_request.py @@ -71,6 +71,10 @@ # stream's per-worker consumer pending count is a live backlog gauge. REQUEST_STREAM = "INFERA_REQUESTS" +# How often a draining worker re-checks its JetStream backlog. Short, because +# the messages are already accepted and every poll is one round-trip. +_QUEUED_POLL_INTERVAL_S = 0.2 + # Reply framing headers. HDR_TYPE = "rs-type" HDR_STATUS = "rs-status" @@ -440,6 +444,7 @@ def __init__( self._url = url self._nc = None self._sub = None + self._js = None # set only under the throttle (JetStream-backed path) self._cancel_sub = None self._http: httpx.AsyncClient | None = None # flag (entry-point CLI) > env > built-in default. @@ -470,6 +475,7 @@ async def start(self) -> None: from nats.js.api import AckPolicy, ConsumerConfig js = self._nc.jetstream() + self._js = js await _ensure_request_stream(js) self._sub = await js.subscribe( subject, @@ -502,6 +508,18 @@ async def stop(self, *, drain: bool = False, drain_timeout: float = 0.0) -> None generations finish for up to ``drain_timeout`` seconds before cancelling any leftovers, so a worker being rolled does not sever active streams. With ``drain=False`` (default) in-flight tasks are cancelled at once.""" + deadline = time.monotonic() + drain_timeout if (drain and drain_timeout > 0) else None + + # 0. Under the throttle, requests land in a WorkQueue stream and are + # pulled from it -- so one can be accepted, queued, and invisible to + # `_inflight`, which only tracks what has been *delivered* here. + # `unsubscribe()` discards remaining messages ("remaining messages will + # be discarded", nats-py), so closing the door first would strand work + # the router already handed us: the client waits out the full idle + # timeout (900s by default) for a reply nobody will ever send. + if deadline is not None: + await self._await_queued(deadline) + # 1. Stop accepting NEW requests immediately so nothing new lands while # we drain (unsubscribe the request subject first). if self._sub is not None: @@ -511,15 +529,16 @@ async def stop(self, *, drain: bool = False, drain_timeout: float = 0.0) -> None pass self._sub = None # 2. Optionally let in-flight requests finish (bounded by drain_timeout). - if drain and drain_timeout > 0: + if deadline is not None: inflight = [t for t in self._inflight.values() if not t.done()] if inflight: + remaining = max(0.0, deadline - time.monotonic()) logger.info( "draining %d in-flight NATS request(s), up to %.0fs", len(inflight), - drain_timeout, + remaining, ) - _done, pending = await asyncio.wait(inflight, timeout=drain_timeout) + _done, pending = await asyncio.wait(inflight, timeout=remaining) if pending: logger.warning( "drain timeout; cancelling %d unfinished request(s)", len(pending) @@ -529,7 +548,18 @@ async def stop(self, *, drain: bool = False, drain_timeout: float = 0.0) -> None if not task.done(): task.cancel() self._inflight.clear() - # 4. Drop the cancel listener and the connection. + # 4. Delete this worker's durable consumer. `unsubscribe()` only tears + # down the local subscription -- a durable survives on the server by + # definition, and its name is derived from worker_id, which a rebuilt + # Pod never reuses (the IP changes). Left behind, every rollout adds an + # orphan holding WorkQueue quota that nothing will ever consume. + if self._js is not None: + try: + await self._js.delete_consumer(REQUEST_STREAM, request_durable(self._worker_id)) + except Exception as exc: # noqa: BLE001 - shutdown must continue + logger.debug("could not delete request consumer: %s", exc) + self._js = None + # 5. Drop the cancel listener and the connection. if self._cancel_sub is not None: try: await self._cancel_sub.unsubscribe() @@ -546,6 +576,41 @@ async def stop(self, *, drain: bool = False, drain_timeout: float = 0.0) -> None pass self._nc = None + async def _await_queued(self, deadline: float) -> None: + """Let JetStream hand over everything already queued for this worker. + + Only ``num_pending`` (accepted, not yet delivered) is waited on: + ``num_ack_pending`` is work already delivered, which is exactly what + ``_inflight`` tracks and step 2 waits for. Counting both would double + the wait for the same requests. + + Never raises, and gives up rather than hanging when the consumer cannot + be read -- a shutdown that stalls on a broker hiccup is worse than one + that drops a queued request, and the caller's deadline is shared with + the in-flight wait that follows. + """ + if self._js is None: + return + durable = request_durable(self._worker_id) + while True: + try: + info = await self._js.consumer_info(REQUEST_STREAM, durable) + except Exception as exc: # noqa: BLE001 - shutdown must continue + logger.debug("drain: cannot read consumer backlog (%s); not waiting", exc) + return + queued = int(getattr(info, "num_pending", 0) or 0) + if queued <= 0: + return + if time.monotonic() >= deadline: + logger.warning( + "drain timeout with %d request(s) still queued in JetStream; " + "they will not be served", + queued, + ) + return + logger.info("drain: waiting for %d queued request(s) to be delivered", queued) + await asyncio.sleep(min(_QUEUED_POLL_INTERVAL_S, max(0.0, deadline - time.monotonic()))) + async def _reply( self, inbox: str, rtype: str, data: bytes = b"", status: int | None = None ) -> None: diff --git a/tests/unit/common/test_nats_drain.py b/tests/unit/common/test_nats_drain.py index 39203a70..b075ce7f 100644 --- a/tests/unit/common/test_nats_drain.py +++ b/tests/unit/common/test_nats_drain.py @@ -10,10 +10,11 @@ from __future__ import annotations import asyncio +from types import SimpleNamespace import pytest -from infera.common.nats_request import NatsRequestServer +from infera.common.nats_request import REQUEST_STREAM, NatsRequestServer, request_durable @pytest.mark.asyncio @@ -40,3 +41,123 @@ async def test_no_drain_cancels_in_flight_immediately(): await asyncio.gather(t, return_exceptions=True) assert t.cancelled() + + +# --- JetStream (admission throttle) path ------------------------------------- +# +# Under the throttle the router publishes into a WorkQueue stream and this +# worker pulls from it, so a request can be accepted, queued, and invisible to +# `_inflight` -- which only tracks what has already been delivered here. + + +class _FakeJs: + """JetStream stand-in: a scripted num_pending sequence + a call log.""" + + def __init__(self, pending: list[int] | None = None, *, fail: bool = False) -> None: + self._pending = list(pending or []) + self._fail = fail + self.deleted: list[tuple[str, str]] = [] + self.info_calls = 0 + + async def consumer_info(self, stream, consumer, timeout=None): + self.info_calls += 1 + if self._fail: + raise RuntimeError("broker hiccup") + value = self._pending.pop(0) if self._pending else 0 + return SimpleNamespace(num_pending=value, num_ack_pending=0) + + async def delete_consumer(self, stream, consumer): + self.deleted.append((stream, consumer)) + return True + + +class _FakeSub: + """Records when the subject was unsubscribed, in units of backlog polls.""" + + def __init__(self, js: _FakeJs) -> None: + self._js = js + self.unsubscribed_after_polls: int | None = None + + async def unsubscribe(self): + self.unsubscribed_after_polls = self._js.info_calls + + +@pytest.mark.asyncio +async def test_the_door_closes_only_after_the_backlog_clears(): + """`unsubscribe()` discards whatever is left in the stream, so it must not + run until the backlog has been handed over -- otherwise a request the + router already accepted is silently dropped and the client waits out the + full idle timeout for a reply nobody will send. + + Pinned by *when* the unsubscribe happened, not just that the backlog was + polled: polling and then closing the door anyway would pass a weaker check. + """ + js = _FakeJs([2, 1, 0]) + sub = _FakeSub(js) + srv = NatsRequestServer("w:1", 30000) + srv._js = js + srv._sub = sub + + await srv.stop(drain=True, drain_timeout=5) + + assert sub.unsubscribed_after_polls is not None, "unsubscribe never ran" + assert sub.unsubscribed_after_polls >= 3, ( + "unsubscribed while the stream still held work: only " + f"{sub.unsubscribed_after_polls} poll(s) had happened, backlog clears on the 3rd" + ) + + +@pytest.mark.asyncio +async def test_a_backlog_that_never_clears_is_bounded_by_the_deadline(): + """A stuck backlog must not hold the process past its drain budget -- the + kubelet's SIGKILL does not wait, and everything after this still has to + run.""" + js = _FakeJs([5] * 1000) + srv = NatsRequestServer("w:1", 30000) + srv._js = js + + start = asyncio.get_running_loop().time() + await srv.stop(drain=True, drain_timeout=0.5) + elapsed = asyncio.get_running_loop().time() - start + + assert elapsed < 3.0, f"stop() overran its 0.5s budget by far: {elapsed:.1f}s" + + +@pytest.mark.asyncio +async def test_an_unreadable_backlog_does_not_block_shutdown(): + """A stalled rollout is worse than a dropped queued request, and a broker + that cannot answer looks identical to a genuinely busy one.""" + js = _FakeJs(fail=True) + srv = NatsRequestServer("w:1", 30000) + srv._js = js + + await asyncio.wait_for(srv.stop(drain=True, drain_timeout=30), timeout=2.0) + + assert js.info_calls == 1, "should give up after the first failed read" + + +@pytest.mark.asyncio +async def test_stop_deletes_the_durable_consumer(): + """The durable outlives the subscription by definition, and its name comes + from worker_id -- which a rebuilt Pod never reuses. Left behind, every + rollout adds an orphan holding WorkQueue quota nothing will consume.""" + js = _FakeJs([0]) + srv = NatsRequestServer("10.0.0.1:30000", 30000) + srv._js = js + + await srv.stop(drain=True, drain_timeout=1) + + assert len(js.deleted) == 1, f"expected one delete_consumer call, got {js.deleted}" + stream, consumer = js.deleted[0] + assert stream == REQUEST_STREAM + assert consumer == request_durable("10.0.0.1:30000") + + +@pytest.mark.asyncio +async def test_core_nats_path_touches_no_jetstream(): + """Without the throttle there is no stream and no durable; the shutdown + path must not assume otherwise.""" + srv = NatsRequestServer("w:1", 30000) + assert srv._js is None + + await srv.stop(drain=True, drain_timeout=1) # must not raise From cfbf7c0a349e5ade7f0aa0fa7dd7ace586ab82e6 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Fri, 7 Aug 2026 13:18:27 +0000 Subject: [PATCH 26/88] fix(operator): read the drain timeout from the environment too Sizing the grace period off --drain-timeout closed one door and left the next one open. The worker takes $INFERA_DRAIN_TIMEOUT as that flag's default, so setting the variable raises the drain exactly as effectively -- and the grace stayed at the 120s floor, so a worker draining for its full timeout was SIGKILLed partway through, which is the outcome the sizing exists to prevent. Environment is read first and the flag overrides it, matching argparse rather than taking the larger of the two: safe, but it would leave a pod lingering minutes past what its config says. A valueFrom reference resolves in the kubelet and is unknowable here, so the budget falls back rather than reading the empty value as zero. Container env is read alongside ServiceSpec.Env for the same reason the flag already is -- an extraPodSpec template passes through verbatim, and those are the deployments most likely to have tuned the drain at all. Co-authored-by: Cursor Signed-off-by: leiwei12 --- .../operator/internal/controller/builders.go | 68 +++++++++++++----- .../internal/controller/builders_test.go | 71 ++++++++++++++++++- 2 files changed, 120 insertions(+), 19 deletions(-) diff --git a/deploy/operator/internal/controller/builders.go b/deploy/operator/internal/controller/builders.go index f5b6bb8d..df4c1a79 100644 --- a/deploy/operator/internal/controller/builders.go +++ b/deploy/operator/internal/controller/builders.go @@ -33,7 +33,7 @@ const ( lwsKind = "LeaderWorkerSet" // Graceful rolling-upgrade tuning for GPU worker pods. - workerPreStopDrainSeconds = 15 // preStop sleep: let the router drop us before SIGTERM + workerPreStopDrainSeconds = 15 // preStop sleep: let the router drop us before SIGTERM workerDefaultDrainTimeoutSeconds = 30 // matches the worker's --drain-timeout default // Teardown after the drain finishes: deregistering, stopping the KV plane, // and engine.stop(), which SIGTERMs the engine's process group and waits up @@ -41,8 +41,22 @@ const ( workerTeardownHeadroomSeconds = 50 // Floor, so short drain timeouts still leave room for a slow engine exit. workerTerminationGraceSeconds int64 = 120 + + // The worker reads this as the default for --drain-timeout, so it sets the + // drain just as effectively as the flag does. + drainTimeoutEnvVar = "INFERA_DRAIN_TIMEOUT" ) +// drainSeconds parses a worker --drain-timeout value. The worker takes a +// float; round up so a fractional value never shortens the budget. +func drainSeconds(v string) (int, bool) { + f, err := strconv.ParseFloat(v, 64) + if err != nil || f <= 0 { + return 0, false + } + return int(math.Ceil(f)), true +} + // graceSecondsFor sizes terminationGracePeriodSeconds so the kubelet cannot // SIGKILL a worker in the middle of shutting down. // @@ -53,8 +67,27 @@ const ( // nothing here parsed -- so raising it for long generations (the exact reason // anyone raises it) silently pushed shutdown past the grace and turned a // graceful drain back into a kill. -func graceSecondsFor(args []string) int64 { +// +// The drain can be set two ways and both have to be read. The worker takes +// $INFERA_DRAIN_TIMEOUT as the flag's *default*, so an env var raises the drain +// exactly as effectively as the flag does -- and parsing only the flag left the +// same silent overrun through a different door. +func graceSecondsFor(args []string, env []corev1.EnvVar) int64 { drain := workerDefaultDrainTimeoutSeconds + // Environment first so an explicit flag overrides it, matching argparse: + // the variable supplies the default, the flag replaces it. + for _, e := range env { + if e.Name != drainTimeoutEnvVar { + continue + } + // A valueFrom reference is resolved by the kubelet, not here, so its + // value is unknowable at build time and the budget falls back to the + // flag or the default. Worth knowing if a drain is ever cut short + // despite a ConfigMap saying otherwise. + if d, ok := drainSeconds(e.Value); ok { + drain = d + } + } for i, a := range args { v := "" if a == "--drain-timeout" && i+1 < len(args) { @@ -65,10 +98,8 @@ func graceSecondsFor(args []string) int64 { if v == "" { continue } - // The worker takes a float; round up so a fractional value never - // shortens the budget. - if f, err := strconv.ParseFloat(v, 64); err == nil && f > 0 { - drain = int(math.Ceil(f)) + if d, ok := drainSeconds(v); ok { + drain = d } } need := int64(workerPreStopDrainSeconds + drain + workerTeardownHeadroomSeconds) @@ -82,8 +113,8 @@ func graceSecondsFor(args []string) int64 { func labelsFor(idepName, svcName string) map[string]string { return map[string]string{ "app.kubernetes.io/managed-by": "infera-operator", - "infera.amd.com/deployment": idepName, - "infera.amd.com/service": svcName, + "infera.amd.com/deployment": idepName, + "infera.amd.com/service": svcName, } } @@ -275,16 +306,21 @@ var mainContainerNames = map[string]struct{}{"main": {}, "infera": {}} // generations, plus a /health readiness probe for single-node workers (skipped // for multi-node LWS groups whose follower ranks > 0 do not serve /health). // Existing values are preserved; the grace is only raised, never lowered. -func injectWorkerRolloutDefaults(spec *corev1.PodSpec, idx int, port int32, addReadiness bool, args []string) { +func injectWorkerRolloutDefaults( + spec *corev1.PodSpec, idx int, port int32, addReadiness bool, + args []string, env []corev1.EnvVar, +) { if idx < 0 || idx >= len(spec.Containers) { return } c := &spec.Containers[idx] - // The flag can arrive two ways: via ServiceSpec.Args on the rendered path, - // or written straight into the container by an extraPodSpec template, which - // is passed through verbatim. Reading only the first would miss exactly the - // deployments most likely to have tuned it. + // The drain can arrive several ways: via ServiceSpec.Args/Env on the + // rendered path, or written straight into the container by an extraPodSpec + // template, which is passed through verbatim. Reading only the first would + // miss exactly the deployments most likely to have tuned it. The container's + // own values go last so they win, being what the process actually sees. drainArgs := append(append(append([]string{}, args...), c.Command...), c.Args...) + drainEnv := append(append([]corev1.EnvVar{}, env...), c.Env...) if addReadiness && c.ReadinessProbe == nil { // SGLang's /health runs a tiny prefill self-check that often takes // >1s, so a 1s probe timeout (the k8s default) flaps the pod between @@ -310,7 +346,7 @@ func injectWorkerRolloutDefaults(spec *corev1.PodSpec, idx int, port int32, addR }, } } - if want := graceSecondsFor(drainArgs); spec.TerminationGracePeriodSeconds == nil || + if want := graceSecondsFor(drainArgs, drainEnv); spec.TerminationGracePeriodSeconds == nil || *spec.TerminationGracePeriodSeconds < want { grace := want spec.TerminationGracePeriodSeconds = &grace @@ -361,7 +397,7 @@ func podTemplateFromExtra(idep *inferav1alpha1.InferaDeployment, svcName string, // Graceful rolling-upgrade defaults for worker pods rendered by an external // template: inject readiness/preStop/grace the template omitted. if svc.ComponentType == inferav1alpha1.ComponentTypeWorker { - injectWorkerRolloutDefaults(&spec, idx, port, svc.NumberOfNodes <= 1 && !svc.SkipReadinessProbe, svc.Args) + injectWorkerRolloutDefaults(&spec, idx, port, svc.NumberOfNodes <= 1 && !svc.SkipReadinessProbe, svc.Args, svc.Env) } return corev1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{Labels: podLabelsFor(idep.Name, svcName, svc)}, @@ -420,7 +456,7 @@ func podTemplate(idep *inferav1alpha1.InferaDeployment, svcName string, svc infe // readiness is skipped for multi-node LWS groups (follower ranks have no // /health). The server (CPU-only) keeps the default fast shutdown. if svc.ComponentType == inferav1alpha1.ComponentTypeWorker { - injectWorkerRolloutDefaults(&podSpec, 0, port, svc.NumberOfNodes <= 1 && !svc.SkipReadinessProbe, svc.Args) + injectWorkerRolloutDefaults(&podSpec, 0, port, svc.NumberOfNodes <= 1 && !svc.SkipReadinessProbe, svc.Args, svc.Env) } tmpl := corev1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{Labels: podLabelsFor(idep.Name, svcName, svc)}, diff --git a/deploy/operator/internal/controller/builders_test.go b/deploy/operator/internal/controller/builders_test.go index c32f6c02..f767b2da 100644 --- a/deploy/operator/internal/controller/builders_test.go +++ b/deploy/operator/internal/controller/builders_test.go @@ -49,7 +49,7 @@ func TestGraceSecondsFor(t *testing.T) { } for _, c := range cases { t.Run(c.name, func(t *testing.T) { - if got := graceSecondsFor(c.args); got != c.want { + if got := graceSecondsFor(c.args, nil); got != c.want { t.Fatalf("graceSecondsFor(%v) = %d, want %d", c.args, got, c.want) } }) @@ -60,7 +60,7 @@ func TestGraceSecondsFor(t *testing.T) { func TestGraceCoversTheWholeShutdown(t *testing.T) { for _, drain := range []int{30, 60, 120, 300} { args := []string{"--drain-timeout", itoa(drain)} - grace := graceSecondsFor(args) + grace := graceSecondsFor(args, nil) need := int64(workerPreStopDrainSeconds + drain + workerTeardownHeadroomSeconds) if grace < need { t.Fatalf("drain=%d: grace %d < required %d -- kubelet would SIGKILL mid-drain", @@ -90,7 +90,7 @@ func TestGraceReadsDrainTimeoutFromTheContainerToo(t *testing.T) { Command: []string{"python3", "-m", "infera.engine.sglang"}, Args: []string{"--model-path", "/m", "--drain-timeout", "240"}, }}} - injectWorkerRolloutDefaults(spec, 0, 8080, false, nil) + injectWorkerRolloutDefaults(spec, 0, 8080, false, nil, nil) if spec.TerminationGracePeriodSeconds == nil { t.Fatal("grace not set") } @@ -127,3 +127,68 @@ func TestEffectiveReplicas(t *testing.T) { t.Fatalf("nothing set anywhere: got %d, want the default 1", got) } } + +// The worker takes $INFERA_DRAIN_TIMEOUT as the default for --drain-timeout, so +// setting it raises the drain exactly as the flag does. Sizing the grace from +// the flag alone left the same silent overrun through a different door: the +// worker would drain for its full timeout and be SIGKILLed partway through. +func TestGraceReadsDrainTimeoutFromTheEnvironment(t *testing.T) { + env := []corev1.EnvVar{ + {Name: "HF_HOME", Value: "/models"}, + {Name: drainTimeoutEnvVar, Value: "300"}, + } + want := int64(workerPreStopDrainSeconds + 300 + workerTeardownHeadroomSeconds) + if got := graceSecondsFor(nil, env); got != want { + t.Fatalf("env-set drain: grace = %d, want %d", got, want) + } +} + +// argparse reads the variable as the flag's default, so an explicit flag wins. +// Sizing the budget off the larger of the two would be safe but wrong, and +// wrong here means a pod that lingers minutes longer than its config says. +func TestGraceFlagOverridesTheEnvironment(t *testing.T) { + env := []corev1.EnvVar{{Name: drainTimeoutEnvVar, Value: "300"}} + args := []string{"--drain-timeout", "60"} + want := int64(workerPreStopDrainSeconds + 60 + workerTeardownHeadroomSeconds) + if got := graceSecondsFor(args, env); got != want { + t.Fatalf("flag with env set: grace = %d, want the flag's %d", got, want) + } +} + +func TestGraceIgnoresUnreadableEnvValues(t *testing.T) { + // valueFrom resolves in the kubelet; nothing is readable here, so the + // budget has to fall back rather than treat the empty value as zero. + from := []corev1.EnvVar{{ + Name: drainTimeoutEnvVar, + ValueFrom: &corev1.EnvVarSource{ + ConfigMapKeyRef: &corev1.ConfigMapKeySelector{Key: "drain"}, + }, + }} + if got := graceSecondsFor(nil, from); got != workerTerminationGraceSeconds { + t.Fatalf("valueFrom: grace = %d, want the floor %d", got, workerTerminationGraceSeconds) + } + for _, v := range []string{"", "abc", "0", "-5"} { + env := []corev1.EnvVar{{Name: drainTimeoutEnvVar, Value: v}} + if got := graceSecondsFor(nil, env); got != workerTerminationGraceSeconds { + t.Fatalf("env %q: grace = %d, want the floor %d", v, got, workerTerminationGraceSeconds) + } + } +} + +// An extraPodSpec template is passed through verbatim, so the variable may sit +// on the container rather than in ServiceSpec.Env -- the same asymmetry the +// flag has, and the deployments most likely to have tuned the drain. +func TestGraceReadsDrainEnvFromTheContainerToo(t *testing.T) { + spec := &corev1.PodSpec{Containers: []corev1.Container{{ + Name: "main", + Env: []corev1.EnvVar{{Name: drainTimeoutEnvVar, Value: "240"}}, + }}} + injectWorkerRolloutDefaults(spec, 0, 8080, false, nil, nil) + if spec.TerminationGracePeriodSeconds == nil { + t.Fatal("grace not set") + } + want := int64(workerPreStopDrainSeconds + 240 + workerTeardownHeadroomSeconds) + if *spec.TerminationGracePeriodSeconds != want { + t.Fatalf("grace = %d, want %d", *spec.TerminationGracePeriodSeconds, want) + } +} From 24073a59a21a6aaecbbf4121a617d555552e3d22 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Wed, 5 Aug 2026 07:23:53 +0000 Subject: [PATCH 27/88] build(operator): regenerate what the release actually ships Packaging the operator was four manual steps with nothing enforcing them, and the two that matter most are silent when skipped. The chart carries its own copy of the CRD (helm/infera-operator/crds/), a mirror of config/crd/bases that `make manifests` never touched -- so editing api/v1alpha1 and running the generator left the chart on the old schema, and `helm install` would create a CRD that rejects the very fields the manager already understands. Nothing fails at package time; it fails on someone else's cluster. `docker build` has the same shape. The Dockerfile compiles the source tree as-is, so a missed `make generate` bakes a stale zz_generated.deepcopy.go into the image. Both are now dependencies rather than steps a README asks you to remember: docker-build: generate manifests helm-package: manifests manifests -> sync-chart-crd sync-chart-crd is its own target rather than a line inside manifests, so a chart copy that drifted without the types changing can be repaired without a regeneration. The two paths it copies between move into CRD_DIR and CHART_CRD_DIR, which also retires the third hard-coded spelling of config/crd/bases, in `install`. .gitignore drops the *.tgz that `helm package` leaves in the working tree. It was showing up untracked beside the sources, one `git add .` away from being committed. Also retargets the manager image and the chart at docker.io/rocm -- the image as an operator- tag of the shared rocm/infera repo, matching the engine images (sglang-v0.1.1, server-v0.1.1) -- and bumps the chart to 0.1.3. The zz_generated.deepcopy.go and config/rbac/role.yaml churn is controller-gen reordering only: both files are identical to their committed versions once sorted. Verified: `make -n docker-build` emits generate -> manifests -> sync-chart-crd -> docker build in that order, and `make helm-package` produces infera-operator-0.1.3.tgz carrying crds/infera.amd.com_inferadeployments.yaml byte-identical to config/crd/bases. Co-authored-by: Cursor Signed-off-by: leiwei12 --- deploy/operator/.gitignore | 2 + deploy/operator/Makefile | 39 +++++++++--- .../api/v1alpha1/zz_generated.deepcopy.go | 60 +++++++++---------- deploy/operator/config/rbac/role.yaml | 42 ++++++------- .../operator/helm/infera-operator/Chart.yaml | 4 +- .../operator/helm/infera-operator/values.yaml | 10 ++-- 6 files changed, 90 insertions(+), 67 deletions(-) diff --git a/deploy/operator/.gitignore b/deploy/operator/.gitignore index 59ae33d9..ec8dae45 100644 --- a/deploy/operator/.gitignore +++ b/deploy/operator/.gitignore @@ -2,3 +2,5 @@ bin/ testbin/ *.test cover.out +# Packaged Helm chart produced by `make helm-package`. +*.tgz diff --git a/deploy/operator/Makefile b/deploy/operator/Makefile index 56937a92..663bee1a 100644 --- a/deploy/operator/Makefile +++ b/deploy/operator/Makefile @@ -1,24 +1,39 @@ # infera-operator — minimal build/codegen targets. # Operator manager image + chart registry. Override on the CLI, e.g. # make docker-build docker-push IMG=docker.io/inferaimage/infera:0.1.0 -IMG ?= docker.io/inferaimage/infera-operator-manager:0.1.1 +# The manager ships as a tag of the shared rocm/infera repo, matching the +# engine images (sglang-v0.1.1, server-v0.1.1, ...). +IMG ?= docker.io/rocm/infera:operator-v0.1.3 # OCI registry namespace the packaged Helm chart is pushed to (chart repo name # comes from Chart.yaml: infera-operator). -CHART_REGISTRY ?= oci://docker.io/inferaimage +CHART_REGISTRY ?= oci://docker.io/rocm CHART_DIR ?= helm/infera-operator +CRD_DIR ?= config/crd/bases +# The chart ships its own copy of the CRD, so `helm install` can create it +# without a separate kubectl apply. It is a mirror of CRD_DIR, never edited by +# hand -- `manifests` refreshes both so a chart cannot lag behind the types. +CHART_CRD_DIR ?= $(CHART_DIR)/crds CONTROLLER_GEN_VERSION ?= v0.17.1 CONTROLLER_GEN = go run sigs.k8s.io/controller-tools/cmd/controller-gen@$(CONTROLLER_GEN_VERSION) -.PHONY: generate manifests build vet test docker-build docker-push helm-package helm-push install deploy sample +.PHONY: generate manifests sync-chart-crd build vet test docker-build docker-push helm-package helm-push install deploy sample ## Generate DeepCopy methods. generate: $(CONTROLLER_GEN) object paths=./api/... -## Generate CRD + RBAC manifests. +## Generate CRD + RBAC manifests, then mirror the CRD into the chart. manifests: - $(CONTROLLER_GEN) crd paths=./api/... output:crd:dir=config/crd/bases + $(CONTROLLER_GEN) crd paths=./api/... output:crd:dir=$(CRD_DIR) $(CONTROLLER_GEN) rbac:roleName=infera-operator-role paths=./internal/... output:rbac:dir=config/rbac + $(MAKE) sync-chart-crd + +## Copy the generated CRD into the chart. Splitting this out keeps it runnable +## on its own, for the case where the chart copy drifted without the types +## changing. +sync-chart-crd: + mkdir -p $(CHART_CRD_DIR) + cp $(CRD_DIR)/*.yaml $(CHART_CRD_DIR)/ build: generate go build -o bin/manager ./cmd @@ -29,15 +44,21 @@ vet: test: go test ./... -docker-build: +## Build the manager image. Depends on the generators because the Dockerfile +## compiles from the source tree as-is: a stale zz_generated.deepcopy.go would +## be baked into the image, and a stale CRD would ship a chart that rejects the +## fields the manager already understands. +docker-build: generate manifests docker build -t $(IMG) . ## Push the operator manager image (requires `docker login`). docker-push: docker push $(IMG) -## Package the Helm chart into a local .tgz. -helm-package: +## Package the Helm chart into a local .tgz. Depends on manifests so a chart +## can never be cut against a stale CRD -- the failure that produces is quiet: +## the chart installs, and the operator rejects fields the CRD has not heard of. +helm-package: manifests helm package $(CHART_DIR) ## Push the packaged chart as an OCI artifact (requires `helm registry login`). @@ -46,7 +67,7 @@ helm-push: helm-package ## Apply the CRD to the current kube context. install: manifests - kubectl apply -f config/crd/bases + kubectl apply -f $(CRD_DIR) ## Apply a sample deployment. sample: diff --git a/deploy/operator/api/v1alpha1/zz_generated.deepcopy.go b/deploy/operator/api/v1alpha1/zz_generated.deepcopy.go index 6de23549..769bfb75 100644 --- a/deploy/operator/api/v1alpha1/zz_generated.deepcopy.go +++ b/deploy/operator/api/v1alpha1/zz_generated.deepcopy.go @@ -45,36 +45,6 @@ func (in *GAIEStatus) DeepCopy() *GAIEStatus { return out } -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *NATSSpec) DeepCopyInto(out *NATSSpec) { - *out = *in -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new NATSSpec. -func (in *NATSSpec) DeepCopy() *NATSSpec { - if in == nil { - return nil - } - out := new(NATSSpec) - in.DeepCopyInto(out) - return out -} - -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *Resources) DeepCopyInto(out *Resources) { - *out = *in -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new Resources. -func (in *Resources) DeepCopy() *Resources { - if in == nil { - return nil - } - out := new(Resources) - in.DeepCopyInto(out) - return out -} - // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *InferaDeployment) DeepCopyInto(out *InferaDeployment) { *out = *in @@ -207,6 +177,36 @@ func (in *InferaDeploymentStatus) DeepCopy() *InferaDeploymentStatus { return out } +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *NATSSpec) DeepCopyInto(out *NATSSpec) { + *out = *in +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new NATSSpec. +func (in *NATSSpec) DeepCopy() *NATSSpec { + if in == nil { + return nil + } + out := new(NATSSpec) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *Resources) DeepCopyInto(out *Resources) { + *out = *in +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new Resources. +func (in *Resources) DeepCopy() *Resources { + if in == nil { + return nil + } + out := new(Resources) + in.DeepCopyInto(out) + return out +} + // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *ServiceSpec) DeepCopyInto(out *ServiceSpec) { *out = *in diff --git a/deploy/operator/config/rbac/role.yaml b/deploy/operator/config/rbac/role.yaml index 662cdf59..4fcb9e46 100644 --- a/deploy/operator/config/rbac/role.yaml +++ b/deploy/operator/config/rbac/role.yaml @@ -52,9 +52,9 @@ rules: - update - watch - apiGroups: - - inference.networking.k8s.io + - infera.amd.com resources: - - inferencepools + - inferadeployments verbs: - create - delete @@ -64,22 +64,23 @@ rules: - update - watch - apiGroups: - - leaderworkerset.x-k8s.io + - infera.amd.com resources: - - leaderworkersets + - inferadeployments/finalizers + verbs: + - update +- apiGroups: + - infera.amd.com + resources: + - inferadeployments/status verbs: - - create - - delete - get - - list - patch - update - - watch - apiGroups: - - rbac.authorization.k8s.io + - inference.networking.k8s.io resources: - - rolebindings - - roles + - inferencepools verbs: - create - delete @@ -89,9 +90,9 @@ rules: - update - watch - apiGroups: - - infera.amd.com + - leaderworkerset.x-k8s.io resources: - - inferadeployments + - leaderworkersets verbs: - create - delete @@ -101,16 +102,15 @@ rules: - update - watch - apiGroups: - - infera.amd.com - resources: - - inferadeployments/finalizers - verbs: - - update -- apiGroups: - - infera.amd.com + - rbac.authorization.k8s.io resources: - - inferadeployments/status + - rolebindings + - roles verbs: + - create + - delete - get + - list - patch - update + - watch diff --git a/deploy/operator/helm/infera-operator/Chart.yaml b/deploy/operator/helm/infera-operator/Chart.yaml index e229d165..74dbc370 100644 --- a/deploy/operator/helm/infera-operator/Chart.yaml +++ b/deploy/operator/helm/infera-operator/Chart.yaml @@ -7,5 +7,5 @@ apiVersion: v2 name: infera-operator description: Infera Kubernetes operator. Reconciles InferaDeployment (infera.amd.com/v1alpha1) into Deployments / LeaderWorkerSets for aggregated and disaggregated (prefill/decode) LLM serving on AMD MI300X. Self-contained — no external inference-operator dependency. type: application -version: 0.1.2 -appVersion: "0.1.2" +version: 0.1.3 +appVersion: "0.1.3" diff --git a/deploy/operator/helm/infera-operator/values.yaml b/deploy/operator/helm/infera-operator/values.yaml index bb65a2f3..110dfb6e 100644 --- a/deploy/operator/helm/infera-operator/values.yaml +++ b/deploy/operator/helm/infera-operator/values.yaml @@ -8,11 +8,11 @@ replicaCount: 1 image: - # Operator manager image. NOTE: distinct repo name from the chart - # (infera-operator) to avoid an OCI tag collision when both are pushed to - # the same namespace. - repository: docker.io/inferaimage/infera-operator-manager - tag: "0.1.2" + # Operator manager image. It lives as an `operator-` tag of the shared + # rocm/infera repo; the chart itself is a separate OCI repo + # (rocm/infera-operator), so their tags cannot collide. + repository: docker.io/rocm/infera + tag: "operator-v0.1.3" pullPolicy: IfNotPresent imagePullSecrets: [] From 978654101154190d7184c1f6e74245afb48848c4 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Fri, 7 Aug 2026 13:20:29 +0000 Subject: [PATCH 28/88] revert(operator): drop the scaling adapter, keep one writer The adapter existed to give autoscalers a standard /scale surface, but an InferaDeployment cannot carry one and that follows from its shape rather than being an omission: spec.services is a map with user-chosen keys, while the scale subresource needs a *static* specReplicasPath and a CRD may declare only one. A single path could name one service, leaving every other pool unscalable. What it cost meanwhile was the obvious thing working. Where an adapter owned a service its count won outright, so editing replicas in the CR -- the way the field reads and the way everyone reaches for first -- silently did nothing. Trading that away bought a surface no autoscaler could drive. So the CR is the single writer again. This does not block a future /scale; it removes a half-built path to it. Anything scaling today already edits the CR, and the drain on scale-down is unchanged. The reconciler now also watches LeaderWorkerSet, so a multi-node service's status reaches the CR when it changes rather than on the next resync. The watch is registered only when the CRD is served: controller-runtime builds an informer per watched type at startup and fails the manager outright for a kind the API server does not know, and LWS is optional here. Co-authored-by: Cursor Signed-off-by: leiwei12 --- .../v1alpha1/inferascalingadapter_types.go | 115 ------- .../api/v1alpha1/zz_generated.deepcopy.go | 161 ++------- deploy/operator/cmd/main.go | 10 +- .../infera.amd.com_inferascalingadapters.yaml | 207 ------------ deploy/operator/config/rbac/role.yaml | 45 ++- .../infera.amd.com_inferascalingadapters.yaml | 207 ------------ .../helm/infera-operator/templates/rbac.yaml | 13 - .../operator/internal/controller/builders.go | 62 ++-- .../internal/controller/builders_test.go | 30 -- .../controller/inferadeployment_controller.go | 61 +--- .../inferascalingadapter_controller.go | 221 ------------ .../internal/controller/scale_paths_test.go | 315 ++++++++++++++++++ .../internal/controller/watches_test.go | 51 +++ manual/features/scaling.md | 85 +++-- 14 files changed, 510 insertions(+), 1073 deletions(-) delete mode 100644 deploy/operator/api/v1alpha1/inferascalingadapter_types.go delete mode 100644 deploy/operator/config/crd/bases/infera.amd.com_inferascalingadapters.yaml delete mode 100644 deploy/operator/helm/infera-operator/crds/infera.amd.com_inferascalingadapters.yaml delete mode 100644 deploy/operator/internal/controller/inferascalingadapter_controller.go create mode 100644 deploy/operator/internal/controller/scale_paths_test.go create mode 100644 deploy/operator/internal/controller/watches_test.go diff --git a/deploy/operator/api/v1alpha1/inferascalingadapter_types.go b/deploy/operator/api/v1alpha1/inferascalingadapter_types.go deleted file mode 100644 index 4e9ab7a7..00000000 --- a/deploy/operator/api/v1alpha1/inferascalingadapter_types.go +++ /dev/null @@ -1,115 +0,0 @@ -/* -Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. - -SPDX-License-Identifier: MIT -*/ - -package v1alpha1 - -import ( - metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" -) - -// InferaScalingAdapterSpec points at one service inside an InferaDeployment and -// owns its replica count. -type InferaScalingAdapterSpec struct { - // DeploymentRef is the InferaDeployment to scale, in this namespace. - // +kubebuilder:validation:MinLength=1 - DeploymentRef string `json:"deploymentRef"` - - // ServiceName is the key in that deployment's `spec.services` map. - // +kubebuilder:validation:MinLength=1 - ServiceName string `json:"serviceName"` - - // Replicas is the desired count, and the field `/scale` writes. - // - // Left unset the adapter is inert: the InferaDeployment's own - // `spec.services[].replicas` still applies. That makes adding an - // adapter a safe no-op until something actually scales, so an autoscaler - // can be attached and observed before it is trusted. - // +optional - // +kubebuilder:validation:Minimum=0 - Replicas *int32 `json:"replicas,omitempty"` -} - -// InferaScalingAdapterStatus carries what `/scale` reads back. -type InferaScalingAdapterStatus struct { - // Replicas is the count observed on the workload -- not the desired count - // echoed back. - // - // The distinction is the whole reason this field exists. HorizontalPodAutoscaler - // computes `desired = ceil(current * metric/target)`; if `current` is really - // the desired value it never lags reality, so during a multi-minute model load - // the autoscaler cannot tell that a scale-up has not landed yet and keeps - // multiplying. - // +optional - Replicas int32 `json:"replicas"` - - // ReadyReplicas is how many of those are actually serving. - // +optional - ReadyReplicas int32 `json:"readyReplicas,omitempty"` - - // Selector is a serialized label selector matching the scaled pods. - // HorizontalPodAutoscaler requires this to be a *string*, not a structured - // selector, and refuses to scale a resource whose scale subresource does not - // provide one. - // +optional - Selector string `json:"selector,omitempty"` - - // ObservedGeneration is the adapter generation this status reflects. - // +optional - ObservedGeneration int64 `json:"observedGeneration,omitempty"` - - // Conditions carries `Ready` (the target resolves and is being driven) and - // `Degraded` (it does not). - // +optional - Conditions []metav1.Condition `json:"conditions,omitempty"` -} - -// InferaScalingAdapter gives one service of an InferaDeployment a standard -// Kubernetes `/scale` subresource. -// -// An InferaDeployment cannot carry `/scale` itself, and this is a property of -// its shape rather than a missing feature: `spec.services` is a map with -// user-chosen keys, while the scale subresource requires `specReplicasPath` to -// be a *static* dot-notation JSONPath under `.spec`. There is no way to write -// "the replicas of an arbitrary map entry". -// -// So scaling gets its own object, one per scalable service. That makes -// `kubectl scale`, HorizontalPodAutoscaler, KEDA and a custom planner all work -// through the same standard interface, with no per-tool support in this -// operator. -// -// While an adapter exists with `spec.replicas` set, it is the single writer of -// that service's replica count: the InferaDeployment reconciler reads the -// adapter instead of the CR's own `replicas`, so the two cannot fight. Delete -// the adapter, or clear `spec.replicas`, and the CR is back in charge. -// -// +kubebuilder:object:root=true -// +kubebuilder:subresource:status -// +kubebuilder:subresource:scale:specpath=.spec.replicas,statuspath=.status.replicas,selectorpath=.status.selector -// +kubebuilder:resource:shortName=isa -// +kubebuilder:printcolumn:name="Target",type=string,JSONPath=`.spec.deploymentRef` -// +kubebuilder:printcolumn:name="Service",type=string,JSONPath=`.spec.serviceName` -// +kubebuilder:printcolumn:name="Desired",type=integer,JSONPath=`.spec.replicas` -// +kubebuilder:printcolumn:name="Current",type=integer,JSONPath=`.status.replicas` -// +kubebuilder:printcolumn:name="Ready",type=integer,JSONPath=`.status.readyReplicas` -// +kubebuilder:printcolumn:name="Age",type=date,JSONPath=`.metadata.creationTimestamp` -type InferaScalingAdapter struct { - metav1.TypeMeta `json:",inline"` - metav1.ObjectMeta `json:"metadata,omitempty"` - - Spec InferaScalingAdapterSpec `json:"spec,omitempty"` - Status InferaScalingAdapterStatus `json:"status,omitempty"` -} - -// +kubebuilder:object:root=true -type InferaScalingAdapterList struct { - metav1.TypeMeta `json:",inline"` - metav1.ListMeta `json:"metadata,omitempty"` - Items []InferaScalingAdapter `json:"items"` -} - -func init() { - SchemeBuilder.Register(&InferaScalingAdapter{}, &InferaScalingAdapterList{}) -} diff --git a/deploy/operator/api/v1alpha1/zz_generated.deepcopy.go b/deploy/operator/api/v1alpha1/zz_generated.deepcopy.go index cc4f5032..6de23549 100644 --- a/deploy/operator/api/v1alpha1/zz_generated.deepcopy.go +++ b/deploy/operator/api/v1alpha1/zz_generated.deepcopy.go @@ -45,6 +45,36 @@ func (in *GAIEStatus) DeepCopy() *GAIEStatus { return out } +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *NATSSpec) DeepCopyInto(out *NATSSpec) { + *out = *in +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new NATSSpec. +func (in *NATSSpec) DeepCopy() *NATSSpec { + if in == nil { + return nil + } + out := new(NATSSpec) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *Resources) DeepCopyInto(out *Resources) { + *out = *in +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new Resources. +func (in *Resources) DeepCopy() *Resources { + if in == nil { + return nil + } + out := new(Resources) + in.DeepCopyInto(out) + return out +} + // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *InferaDeployment) DeepCopyInto(out *InferaDeployment) { *out = *in @@ -177,137 +207,6 @@ func (in *InferaDeploymentStatus) DeepCopy() *InferaDeploymentStatus { return out } -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *InferaScalingAdapter) DeepCopyInto(out *InferaScalingAdapter) { - *out = *in - out.TypeMeta = in.TypeMeta - in.ObjectMeta.DeepCopyInto(&out.ObjectMeta) - in.Spec.DeepCopyInto(&out.Spec) - in.Status.DeepCopyInto(&out.Status) -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new InferaScalingAdapter. -func (in *InferaScalingAdapter) DeepCopy() *InferaScalingAdapter { - if in == nil { - return nil - } - out := new(InferaScalingAdapter) - in.DeepCopyInto(out) - return out -} - -// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object. -func (in *InferaScalingAdapter) DeepCopyObject() runtime.Object { - if c := in.DeepCopy(); c != nil { - return c - } - return nil -} - -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *InferaScalingAdapterList) DeepCopyInto(out *InferaScalingAdapterList) { - *out = *in - out.TypeMeta = in.TypeMeta - in.ListMeta.DeepCopyInto(&out.ListMeta) - if in.Items != nil { - in, out := &in.Items, &out.Items - *out = make([]InferaScalingAdapter, len(*in)) - for i := range *in { - (*in)[i].DeepCopyInto(&(*out)[i]) - } - } -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new InferaScalingAdapterList. -func (in *InferaScalingAdapterList) DeepCopy() *InferaScalingAdapterList { - if in == nil { - return nil - } - out := new(InferaScalingAdapterList) - in.DeepCopyInto(out) - return out -} - -// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object. -func (in *InferaScalingAdapterList) DeepCopyObject() runtime.Object { - if c := in.DeepCopy(); c != nil { - return c - } - return nil -} - -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *InferaScalingAdapterSpec) DeepCopyInto(out *InferaScalingAdapterSpec) { - *out = *in - if in.Replicas != nil { - in, out := &in.Replicas, &out.Replicas - *out = new(int32) - **out = **in - } -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new InferaScalingAdapterSpec. -func (in *InferaScalingAdapterSpec) DeepCopy() *InferaScalingAdapterSpec { - if in == nil { - return nil - } - out := new(InferaScalingAdapterSpec) - in.DeepCopyInto(out) - return out -} - -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *InferaScalingAdapterStatus) DeepCopyInto(out *InferaScalingAdapterStatus) { - *out = *in - if in.Conditions != nil { - in, out := &in.Conditions, &out.Conditions - *out = make([]metav1.Condition, len(*in)) - for i := range *in { - (*in)[i].DeepCopyInto(&(*out)[i]) - } - } -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new InferaScalingAdapterStatus. -func (in *InferaScalingAdapterStatus) DeepCopy() *InferaScalingAdapterStatus { - if in == nil { - return nil - } - out := new(InferaScalingAdapterStatus) - in.DeepCopyInto(out) - return out -} - -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *NATSSpec) DeepCopyInto(out *NATSSpec) { - *out = *in -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new NATSSpec. -func (in *NATSSpec) DeepCopy() *NATSSpec { - if in == nil { - return nil - } - out := new(NATSSpec) - in.DeepCopyInto(out) - return out -} - -// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *Resources) DeepCopyInto(out *Resources) { - *out = *in -} - -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new Resources. -func (in *Resources) DeepCopy() *Resources { - if in == nil { - return nil - } - out := new(Resources) - in.DeepCopyInto(out) - return out -} - // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *ServiceSpec) DeepCopyInto(out *ServiceSpec) { *out = *in diff --git a/deploy/operator/cmd/main.go b/deploy/operator/cmd/main.go index 867e05c2..f370c031 100644 --- a/deploy/operator/cmd/main.go +++ b/deploy/operator/cmd/main.go @@ -9,9 +9,9 @@ import ( "flag" "os" - clientgoscheme "k8s.io/client-go/kubernetes/scheme" "k8s.io/apimachinery/pkg/runtime" utilruntime "k8s.io/apimachinery/pkg/util/runtime" + clientgoscheme "k8s.io/client-go/kubernetes/scheme" ctrl "sigs.k8s.io/controller-runtime" "sigs.k8s.io/controller-runtime/pkg/healthz" "sigs.k8s.io/controller-runtime/pkg/log/zap" @@ -55,14 +55,6 @@ func main() { os.Exit(1) } - if err := (&controller.InferaScalingAdapterReconciler{ - Client: mgr.GetClient(), - Scheme: mgr.GetScheme(), - }).SetupWithManager(mgr); err != nil { - setupLog.Error(err, "unable to create controller", "controller", "InferaScalingAdapter") - os.Exit(1) - } - if err := (&controller.InferaDeploymentReconciler{ Client: mgr.GetClient(), Scheme: mgr.GetScheme(), diff --git a/deploy/operator/config/crd/bases/infera.amd.com_inferascalingadapters.yaml b/deploy/operator/config/crd/bases/infera.amd.com_inferascalingadapters.yaml deleted file mode 100644 index 8b722246..00000000 --- a/deploy/operator/config/crd/bases/infera.amd.com_inferascalingadapters.yaml +++ /dev/null @@ -1,207 +0,0 @@ ---- -apiVersion: apiextensions.k8s.io/v1 -kind: CustomResourceDefinition -metadata: - annotations: - controller-gen.kubebuilder.io/version: v0.17.1 - name: inferascalingadapters.infera.amd.com -spec: - group: infera.amd.com - names: - kind: InferaScalingAdapter - listKind: InferaScalingAdapterList - plural: inferascalingadapters - shortNames: - - isa - singular: inferascalingadapter - scope: Namespaced - versions: - - additionalPrinterColumns: - - jsonPath: .spec.deploymentRef - name: Target - type: string - - jsonPath: .spec.serviceName - name: Service - type: string - - jsonPath: .spec.replicas - name: Desired - type: integer - - jsonPath: .status.replicas - name: Current - type: integer - - jsonPath: .status.readyReplicas - name: Ready - type: integer - - jsonPath: .metadata.creationTimestamp - name: Age - type: date - name: v1alpha1 - schema: - openAPIV3Schema: - description: |- - InferaScalingAdapter gives one service of an InferaDeployment a standard - Kubernetes `/scale` subresource. - - An InferaDeployment cannot carry `/scale` itself, and this is a property of - its shape rather than a missing feature: `spec.services` is a map with - user-chosen keys, while the scale subresource requires `specReplicasPath` to - be a *static* dot-notation JSONPath under `.spec`. There is no way to write - "the replicas of an arbitrary map entry". - - So scaling gets its own object, one per scalable service. That makes - `kubectl scale`, HorizontalPodAutoscaler, KEDA and a custom planner all work - through the same standard interface, with no per-tool support in this - operator. - - While an adapter exists with `spec.replicas` set, it is the single writer of - that service's replica count: the InferaDeployment reconciler reads the - adapter instead of the CR's own `replicas`, so the two cannot fight. Delete - the adapter, or clear `spec.replicas`, and the CR is back in charge. - properties: - apiVersion: - description: |- - APIVersion defines the versioned schema of this representation of an object. - Servers should convert recognized schemas to the latest internal value, and - may reject unrecognized values. - More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#resources - type: string - kind: - description: |- - Kind is a string value representing the REST resource this object represents. - Servers may infer this from the endpoint the client submits requests to. - Cannot be updated. - In CamelCase. - More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#types-kinds - type: string - metadata: - type: object - spec: - description: |- - InferaScalingAdapterSpec points at one service inside an InferaDeployment and - owns its replica count. - properties: - deploymentRef: - description: DeploymentRef is the InferaDeployment to scale, in this - namespace. - minLength: 1 - type: string - replicas: - description: |- - Replicas is the desired count, and the field `/scale` writes. - - Left unset the adapter is inert: the InferaDeployment's own - `spec.services[].replicas` still applies. That makes adding an - adapter a safe no-op until something actually scales, so an autoscaler - can be attached and observed before it is trusted. - format: int32 - minimum: 0 - type: integer - serviceName: - description: ServiceName is the key in that deployment's `spec.services` - map. - minLength: 1 - type: string - required: - - deploymentRef - - serviceName - type: object - status: - description: InferaScalingAdapterStatus carries what `/scale` reads back. - properties: - conditions: - description: |- - Conditions carries `Ready` (the target resolves and is being driven) and - `Degraded` (it does not). - items: - description: Condition contains details for one aspect of the current - state of this API Resource. - properties: - lastTransitionTime: - description: |- - lastTransitionTime is the last time the condition transitioned from one status to another. - This should be when the underlying condition changed. If that is not known, then using the time when the API field changed is acceptable. - format: date-time - type: string - message: - description: |- - message is a human readable message indicating details about the transition. - This may be an empty string. - maxLength: 32768 - type: string - observedGeneration: - description: |- - observedGeneration represents the .metadata.generation that the condition was set based upon. - For instance, if .metadata.generation is currently 12, but the .status.conditions[x].observedGeneration is 9, the condition is out of date - with respect to the current state of the instance. - format: int64 - minimum: 0 - type: integer - reason: - description: |- - reason contains a programmatic identifier indicating the reason for the condition's last transition. - Producers of specific condition types may define expected values and meanings for this field, - and whether the values are considered a guaranteed API. - The value should be a CamelCase string. - This field may not be empty. - maxLength: 1024 - minLength: 1 - pattern: ^[A-Za-z]([A-Za-z0-9_,:]*[A-Za-z0-9_])?$ - type: string - status: - description: status of the condition, one of True, False, Unknown. - enum: - - "True" - - "False" - - Unknown - type: string - type: - description: type of condition in CamelCase or in foo.example.com/CamelCase. - maxLength: 316 - pattern: ^([a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*/)?(([A-Za-z0-9][-A-Za-z0-9_.]*)?[A-Za-z0-9])$ - type: string - required: - - lastTransitionTime - - message - - reason - - status - - type - type: object - type: array - observedGeneration: - description: ObservedGeneration is the adapter generation this status - reflects. - format: int64 - type: integer - readyReplicas: - description: ReadyReplicas is how many of those are actually serving. - format: int32 - type: integer - replicas: - description: |- - Replicas is the count observed on the workload -- not the desired count - echoed back. - - The distinction is the whole reason this field exists. HorizontalPodAutoscaler - computes `desired = ceil(current * metric/target)`; if `current` is really - the desired value it never lags reality, so during a multi-minute model load - the autoscaler cannot tell that a scale-up has not landed yet and keeps - multiplying. - format: int32 - type: integer - selector: - description: |- - Selector is a serialized label selector matching the scaled pods. - HorizontalPodAutoscaler requires this to be a *string*, not a structured - selector, and refuses to scale a resource whose scale subresource does not - provide one. - type: string - type: object - type: object - served: true - storage: true - subresources: - scale: - labelSelectorPath: .status.selector - specReplicasPath: .spec.replicas - statusReplicasPath: .status.replicas - status: {} diff --git a/deploy/operator/config/rbac/role.yaml b/deploy/operator/config/rbac/role.yaml index d5c44e2f..662cdf59 100644 --- a/deploy/operator/config/rbac/role.yaml +++ b/deploy/operator/config/rbac/role.yaml @@ -52,10 +52,9 @@ rules: - update - watch - apiGroups: - - infera.amd.com + - inference.networking.k8s.io resources: - - inferadeployments - - inferascalingadapters + - inferencepools verbs: - create - delete @@ -65,25 +64,22 @@ rules: - update - watch - apiGroups: - - infera.amd.com - resources: - - inferadeployments/finalizers - verbs: - - update -- apiGroups: - - infera.amd.com + - leaderworkerset.x-k8s.io resources: - - inferadeployments/status - - inferascalingadapters/scale - - inferascalingadapters/status + - leaderworkersets verbs: + - create + - delete - get + - list - patch - update + - watch - apiGroups: - - inference.networking.k8s.io + - rbac.authorization.k8s.io resources: - - inferencepools + - rolebindings + - roles verbs: - create - delete @@ -93,9 +89,9 @@ rules: - update - watch - apiGroups: - - leaderworkerset.x-k8s.io + - infera.amd.com resources: - - leaderworkersets + - inferadeployments verbs: - create - delete @@ -105,15 +101,16 @@ rules: - update - watch - apiGroups: - - rbac.authorization.k8s.io + - infera.amd.com resources: - - rolebindings - - roles + - inferadeployments/finalizers + verbs: + - update +- apiGroups: + - infera.amd.com + resources: + - inferadeployments/status verbs: - - create - - delete - get - - list - patch - update - - watch diff --git a/deploy/operator/helm/infera-operator/crds/infera.amd.com_inferascalingadapters.yaml b/deploy/operator/helm/infera-operator/crds/infera.amd.com_inferascalingadapters.yaml deleted file mode 100644 index 8b722246..00000000 --- a/deploy/operator/helm/infera-operator/crds/infera.amd.com_inferascalingadapters.yaml +++ /dev/null @@ -1,207 +0,0 @@ ---- -apiVersion: apiextensions.k8s.io/v1 -kind: CustomResourceDefinition -metadata: - annotations: - controller-gen.kubebuilder.io/version: v0.17.1 - name: inferascalingadapters.infera.amd.com -spec: - group: infera.amd.com - names: - kind: InferaScalingAdapter - listKind: InferaScalingAdapterList - plural: inferascalingadapters - shortNames: - - isa - singular: inferascalingadapter - scope: Namespaced - versions: - - additionalPrinterColumns: - - jsonPath: .spec.deploymentRef - name: Target - type: string - - jsonPath: .spec.serviceName - name: Service - type: string - - jsonPath: .spec.replicas - name: Desired - type: integer - - jsonPath: .status.replicas - name: Current - type: integer - - jsonPath: .status.readyReplicas - name: Ready - type: integer - - jsonPath: .metadata.creationTimestamp - name: Age - type: date - name: v1alpha1 - schema: - openAPIV3Schema: - description: |- - InferaScalingAdapter gives one service of an InferaDeployment a standard - Kubernetes `/scale` subresource. - - An InferaDeployment cannot carry `/scale` itself, and this is a property of - its shape rather than a missing feature: `spec.services` is a map with - user-chosen keys, while the scale subresource requires `specReplicasPath` to - be a *static* dot-notation JSONPath under `.spec`. There is no way to write - "the replicas of an arbitrary map entry". - - So scaling gets its own object, one per scalable service. That makes - `kubectl scale`, HorizontalPodAutoscaler, KEDA and a custom planner all work - through the same standard interface, with no per-tool support in this - operator. - - While an adapter exists with `spec.replicas` set, it is the single writer of - that service's replica count: the InferaDeployment reconciler reads the - adapter instead of the CR's own `replicas`, so the two cannot fight. Delete - the adapter, or clear `spec.replicas`, and the CR is back in charge. - properties: - apiVersion: - description: |- - APIVersion defines the versioned schema of this representation of an object. - Servers should convert recognized schemas to the latest internal value, and - may reject unrecognized values. - More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#resources - type: string - kind: - description: |- - Kind is a string value representing the REST resource this object represents. - Servers may infer this from the endpoint the client submits requests to. - Cannot be updated. - In CamelCase. - More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#types-kinds - type: string - metadata: - type: object - spec: - description: |- - InferaScalingAdapterSpec points at one service inside an InferaDeployment and - owns its replica count. - properties: - deploymentRef: - description: DeploymentRef is the InferaDeployment to scale, in this - namespace. - minLength: 1 - type: string - replicas: - description: |- - Replicas is the desired count, and the field `/scale` writes. - - Left unset the adapter is inert: the InferaDeployment's own - `spec.services[].replicas` still applies. That makes adding an - adapter a safe no-op until something actually scales, so an autoscaler - can be attached and observed before it is trusted. - format: int32 - minimum: 0 - type: integer - serviceName: - description: ServiceName is the key in that deployment's `spec.services` - map. - minLength: 1 - type: string - required: - - deploymentRef - - serviceName - type: object - status: - description: InferaScalingAdapterStatus carries what `/scale` reads back. - properties: - conditions: - description: |- - Conditions carries `Ready` (the target resolves and is being driven) and - `Degraded` (it does not). - items: - description: Condition contains details for one aspect of the current - state of this API Resource. - properties: - lastTransitionTime: - description: |- - lastTransitionTime is the last time the condition transitioned from one status to another. - This should be when the underlying condition changed. If that is not known, then using the time when the API field changed is acceptable. - format: date-time - type: string - message: - description: |- - message is a human readable message indicating details about the transition. - This may be an empty string. - maxLength: 32768 - type: string - observedGeneration: - description: |- - observedGeneration represents the .metadata.generation that the condition was set based upon. - For instance, if .metadata.generation is currently 12, but the .status.conditions[x].observedGeneration is 9, the condition is out of date - with respect to the current state of the instance. - format: int64 - minimum: 0 - type: integer - reason: - description: |- - reason contains a programmatic identifier indicating the reason for the condition's last transition. - Producers of specific condition types may define expected values and meanings for this field, - and whether the values are considered a guaranteed API. - The value should be a CamelCase string. - This field may not be empty. - maxLength: 1024 - minLength: 1 - pattern: ^[A-Za-z]([A-Za-z0-9_,:]*[A-Za-z0-9_])?$ - type: string - status: - description: status of the condition, one of True, False, Unknown. - enum: - - "True" - - "False" - - Unknown - type: string - type: - description: type of condition in CamelCase or in foo.example.com/CamelCase. - maxLength: 316 - pattern: ^([a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*/)?(([A-Za-z0-9][-A-Za-z0-9_.]*)?[A-Za-z0-9])$ - type: string - required: - - lastTransitionTime - - message - - reason - - status - - type - type: object - type: array - observedGeneration: - description: ObservedGeneration is the adapter generation this status - reflects. - format: int64 - type: integer - readyReplicas: - description: ReadyReplicas is how many of those are actually serving. - format: int32 - type: integer - replicas: - description: |- - Replicas is the count observed on the workload -- not the desired count - echoed back. - - The distinction is the whole reason this field exists. HorizontalPodAutoscaler - computes `desired = ceil(current * metric/target)`; if `current` is really - the desired value it never lags reality, so during a multi-minute model load - the autoscaler cannot tell that a scale-up has not landed yet and keeps - multiplying. - format: int32 - type: integer - selector: - description: |- - Selector is a serialized label selector matching the scaled pods. - HorizontalPodAutoscaler requires this to be a *string*, not a structured - selector, and refuses to scale a resource whose scale subresource does not - provide one. - type: string - type: object - type: object - served: true - storage: true - subresources: - scale: - labelSelectorPath: .status.selector - specReplicasPath: .spec.replicas - statusReplicasPath: .status.replicas - status: {} diff --git a/deploy/operator/helm/infera-operator/templates/rbac.yaml b/deploy/operator/helm/infera-operator/templates/rbac.yaml index 38fd46f9..54d8f55e 100644 --- a/deploy/operator/helm/infera-operator/templates/rbac.yaml +++ b/deploy/operator/helm/infera-operator/templates/rbac.yaml @@ -43,19 +43,6 @@ rules: - apiGroups: ["infera.amd.com"] resources: ["inferadeployments/status"] verbs: ["get", "patch", "update"] -# Scaling adapters: the standard /scale surface an HPA, KEDA or a custom planner -# drives. The operator reads spec.replicas and owns status; /scale is listed -# because that is the endpoint external scalers write, and it is a distinct -# subresource from status for RBAC purposes. -- apiGroups: ["infera.amd.com"] - resources: ["inferascalingadapters"] - verbs: ["create", "delete", "get", "list", "patch", "update", "watch"] -- apiGroups: ["infera.amd.com"] - resources: ["inferascalingadapters/status"] - verbs: ["get", "patch", "update"] -- apiGroups: ["infera.amd.com"] - resources: ["inferascalingadapters/scale"] - verbs: ["get", "patch", "update"] # Leader election lease + event recording. - apiGroups: ["coordination.k8s.io"] resources: ["leases"] diff --git a/deploy/operator/internal/controller/builders.go b/deploy/operator/internal/controller/builders.go index df4c1a79..0d36f320 100644 --- a/deploy/operator/internal/controller/builders.go +++ b/deploy/operator/internal/controller/builders.go @@ -14,9 +14,11 @@ import ( appsv1 "k8s.io/api/apps/v1" corev1 "k8s.io/api/core/v1" rbacv1 "k8s.io/api/rbac/v1" + "k8s.io/apimachinery/pkg/api/meta" "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/runtime/schema" "k8s.io/apimachinery/pkg/util/intstr" inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" @@ -109,15 +111,48 @@ func graceSecondsFor(args []string, env []corev1.EnvVar) int64 { return need } +// Identity labels on every workload this operator builds. They are the only +// link back from a Deployment/LeaderWorkerSet to the CR and service that +// produced it, so the watch handlers that map a workload event to the objects +// interested in it read these rather than re-deriving the name. +const ( + labelKeyDeployment = "infera.amd.com/deployment" + labelKeyService = "infera.amd.com/service" +) + // labelsFor returns the selector/identity labels for a service's workload. func labelsFor(idepName, svcName string) map[string]string { return map[string]string{ "app.kubernetes.io/managed-by": "infera-operator", - "infera.amd.com/deployment": idepName, - "infera.amd.com/service": svcName, + labelKeyDeployment: idepName, + labelKeyService: svcName, } } +// lwsInstalled reports whether the LeaderWorkerSet CRD is served by the API. +// +// It gates registering a watch on LWS: controller-runtime builds an informer +// for every watched type at startup, and one for a kind the API server does not +// serve fails the manager outright. LWS is an optional dependency here -- only +// multi-node services use it -- so a single-node cluster without the CRD must +// still be able to run the operator. +// +// The check runs once, at setup. Installing the CRD afterwards therefore needs +// an operator restart to pick up the watch; until then multi-node status still +// refreshes on the reconciler's periodic resync, just not immediately. +func lwsInstalled(mapper meta.RESTMapper) bool { + _, err := mapper.RESTMapping( + schema.GroupKind{Group: lwsGVK().Group, Kind: lwsGVK().Kind}, lwsGVK().Version) + return err == nil +} + +// lwsObject returns an empty LeaderWorkerSet for use as a watch target. +func lwsObject() *unstructured.Unstructured { + u := &unstructured.Unstructured{} + u.SetGroupVersionKind(lwsGVK()) + return u +} + // podLabelsFor returns the operator's selector labels merged with any // caller-supplied ServiceSpec.PodLabels (e.g. an external orchestrator's // workload-id label used by its pod syncer). Operator selector labels always @@ -183,21 +218,6 @@ func replicasOf(svc inferav1alpha1.ServiceSpec) int32 { return 1 } -// effectiveReplicas is the count to write onto the workload: the scaling -// adapter's if one owns this service, otherwise the CR's own. -// -// Routing the adapter through here rather than letting it write the workload -// directly keeps a single writer. The reconciler already assigns the whole -// child `.Spec` on every pass, so a second writer would simply be reverted -- -// which is exactly what happens today to anyone pointing an HPA at the child -// Deployment. -func effectiveReplicas(svc inferav1alpha1.ServiceSpec, svcName string, overrides map[string]int32) int32 { - if n, ok := overrides[svcName]; ok { - return n - } - return replicasOf(svc) -} - // containerCommand builds the infera entrypoint + operator-injected flags, // then appends the user's free-form Args (model-path, tokenizer, tp-size, ...). func containerCommand(idep *inferav1alpha1.InferaDeployment, svc inferav1alpha1.ServiceSpec) []string { @@ -466,8 +486,8 @@ func podTemplate(idep *inferav1alpha1.InferaDeployment, svcName string, svc infe return tmpl } -func buildDeployment(idep *inferav1alpha1.InferaDeployment, svcName string, svc inferav1alpha1.ServiceSpec, overrides map[string]int32) *appsv1.Deployment { - reps := effectiveReplicas(svc, svcName, overrides) +func buildDeployment(idep *inferav1alpha1.InferaDeployment, svcName string, svc inferav1alpha1.ServiceSpec) *appsv1.Deployment { + reps := replicasOf(svc) lbls := labelsFor(idep.Name, svcName) // Worker services use surge-free RollingUpdate (maxSurge=0, maxUnavailable=1): // the default RollingUpdate brings up a surge pod first, which on a @@ -506,8 +526,8 @@ func buildDeployment(idep *inferav1alpha1.InferaDeployment, svcName string, svc // buildLeaderWorkerSet returns an unstructured LeaderWorkerSet so the operator // does not take a compile-time dependency on the LWS Go module (keeps Infera // self-contained; the LWS CRD must be installed in the cluster). -func buildLeaderWorkerSet(idep *inferav1alpha1.InferaDeployment, svcName string, svc inferav1alpha1.ServiceSpec, overrides map[string]int32) *unstructured.Unstructured { - reps := effectiveReplicas(svc, svcName, overrides) +func buildLeaderWorkerSet(idep *inferav1alpha1.InferaDeployment, svcName string, svc inferav1alpha1.ServiceSpec) *unstructured.Unstructured { + reps := replicasOf(svc) lbls := labelsFor(idep.Name, svcName) tmpl := podTemplate(idep, svcName, svc) // Convert the typed PodTemplateSpec to a map for embedding. diff --git a/deploy/operator/internal/controller/builders_test.go b/deploy/operator/internal/controller/builders_test.go index f767b2da..6db09a4f 100644 --- a/deploy/operator/internal/controller/builders_test.go +++ b/deploy/operator/internal/controller/builders_test.go @@ -10,8 +10,6 @@ import ( "testing" corev1 "k8s.io/api/core/v1" - - inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" ) // The grace period is the only thing standing between a graceful drain and a @@ -100,34 +98,6 @@ func TestGraceReadsDrainTimeoutFromTheContainerToo(t *testing.T) { } } -// An adapter owns its service's replica count; everything else keeps using the -// CR's. Getting this wrong in either direction is bad: ignoring the adapter -// makes `/scale` a no-op, and applying it too broadly makes a single autoscaler -// silently resize pools nobody pointed it at. -func TestEffectiveReplicas(t *testing.T) { - three := int32(3) - svc := inferav1alpha1.ServiceSpec{Replicas: &three} - - if got := effectiveReplicas(svc, "worker", nil); got != 3 { - t.Fatalf("no adapters: got %d, want the CR's 3", got) - } - if got := effectiveReplicas(svc, "worker", map[string]int32{"worker": 7}); got != 7 { - t.Fatalf("adapter present: got %d, want 7", got) - } - if got := effectiveReplicas(svc, "worker", map[string]int32{"prefill": 7}); got != 3 { - t.Fatalf("adapter for another service: got %d, want the CR's 3", got) - } - // Zero is a legitimate target, not "unset" -- an autoscaler scaling a pool - // to zero must not silently fall back to the CR's count. - if got := effectiveReplicas(svc, "worker", map[string]int32{"worker": 0}); got != 0 { - t.Fatalf("adapter asking for 0: got %d, want 0", got) - } - // The CR default when it says nothing either. - if got := effectiveReplicas(inferav1alpha1.ServiceSpec{}, "worker", nil); got != 1 { - t.Fatalf("nothing set anywhere: got %d, want the default 1", got) - } -} - // The worker takes $INFERA_DRAIN_TIMEOUT as the default for --drain-timeout, so // setting it raises the drain exactly as the flag does. Sizing the grace from // the flag alone left the same silent overrun through a different door: the diff --git a/deploy/operator/internal/controller/inferadeployment_controller.go b/deploy/operator/internal/controller/inferadeployment_controller.go index 148bdd3a..80878764 100644 --- a/deploy/operator/internal/controller/inferadeployment_controller.go +++ b/deploy/operator/internal/controller/inferadeployment_controller.go @@ -13,8 +13,6 @@ import ( appsv1 "k8s.io/api/apps/v1" corev1 "k8s.io/api/core/v1" rbacv1 "k8s.io/api/rbac/v1" - apierrors "k8s.io/apimachinery/pkg/api/errors" - "k8s.io/apimachinery/pkg/api/meta" "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" "k8s.io/apimachinery/pkg/runtime" ctrl "sigs.k8s.io/controller-runtime" @@ -88,27 +86,17 @@ func (r *InferaDeploymentReconciler) Reconcile(ctx context.Context, req ctrl.Req } // 2. Each service -> Deployment (single-node) or LeaderWorkerSet (multi-node). - // - // Scaling adapters are resolved first: where one owns a service, its - // replica count wins over the CR's. Reading them here rather than letting - // the adapter write the workload keeps a single writer -- this reconciler - // assigns the whole child `.Spec` every pass, so a second writer would just - // be reverted. - overrides, err := r.replicaOverrides(ctx, idep) - if err != nil { - return ctrl.Result{}, err - } status := map[string]inferav1alpha1.ServiceStatus{} for _, name := range sortedKeys(idep.Spec.Services) { svc := idep.Spec.Services[name] if svc.NumberOfNodes > 1 { - lws := buildLeaderWorkerSet(idep, name, svc, overrides) + lws := buildLeaderWorkerSet(idep, name, svc) if err := r.applyUnstructured(ctx, idep, lws); err != nil { return ctrl.Result{}, err } status[name] = r.lwsStatus(ctx, idep, name, svc) } else { - dep := buildDeployment(idep, name, svc, overrides) + dep := buildDeployment(idep, name, svc) if err := r.applyObject(ctx, idep, dep); err != nil { return ctrl.Result{}, err } @@ -156,39 +144,6 @@ func (r *InferaDeploymentReconciler) Reconcile(ctx context.Context, req ctrl.Req return ctrl.Result{RequeueAfter: 15 * time.Second}, nil } -// replicaOverrides maps service name -> replica count for services owned by a -// scaling adapter. -// -// An adapter with no `spec.replicas` is deliberately absent from the map rather -// than contributing a zero: creating an adapter must be a no-op until something -// actually scales, so an autoscaler can be attached and watched before it is -// trusted. A missing CRD is likewise not an error -- the adapter is optional, -// and an operator that refused to reconcile without it would break every -// existing deployment on upgrade. -func (r *InferaDeploymentReconciler) replicaOverrides( - ctx context.Context, idep *inferav1alpha1.InferaDeployment, -) (map[string]int32, error) { - list := &inferav1alpha1.InferaScalingAdapterList{} - if err := r.List(ctx, list, client.InNamespace(idep.Namespace)); err != nil { - if meta.IsNoMatchError(err) || apierrors.IsNotFound(err) { - return nil, nil - } - return nil, err - } - out := map[string]int32{} - for i := range list.Items { - a := &list.Items[i] - if a.Spec.DeploymentRef != idep.Name || a.Spec.Replicas == nil { - continue - } - if _, ok := idep.Spec.Services[a.Spec.ServiceName]; !ok { - continue // dangling adapter; its own controller reports Degraded - } - out[a.Spec.ServiceName] = *a.Spec.Replicas - } - return out, nil -} - // applyObject create-or-updates a typed object, setting the owner reference. func (r *InferaDeploymentReconciler) applyObject(ctx context.Context, idep *inferav1alpha1.InferaDeployment, desired client.Object) error { // Build a fresh empty object of the same kind keyed by name/namespace. @@ -308,10 +263,16 @@ func sortedKeys(m map[string]inferav1alpha1.ServiceSpec) []string { // SetupWithManager registers the controller. func (r *InferaDeploymentReconciler) SetupWithManager(mgr ctrl.Manager) error { - return ctrl.NewControllerManagedBy(mgr). + b := ctrl.NewControllerManagedBy(mgr). For(&inferav1alpha1.InferaDeployment{}). Owns(&appsv1.Deployment{}). Owns(&appsv1.StatefulSet{}). - Owns(&corev1.Service{}). - Complete(r) + Owns(&corev1.Service{}) + // Multi-node services are LeaderWorkerSets, so their status only reaches + // InferaDeployment.status on a resync unless we watch them. Guarded because + // the CRD is optional -- see lwsInstalled. + if lwsInstalled(mgr.GetRESTMapper()) { + b = b.Owns(lwsObject()) + } + return b.Complete(r) } diff --git a/deploy/operator/internal/controller/inferascalingadapter_controller.go b/deploy/operator/internal/controller/inferascalingadapter_controller.go deleted file mode 100644 index ea35eb7c..00000000 --- a/deploy/operator/internal/controller/inferascalingadapter_controller.go +++ /dev/null @@ -1,221 +0,0 @@ -/* -Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. - -SPDX-License-Identifier: MIT -*/ - -package controller - -import ( - "context" - "fmt" - "time" - - appsv1 "k8s.io/api/apps/v1" - apierrors "k8s.io/apimachinery/pkg/api/errors" - metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" - "k8s.io/apimachinery/pkg/labels" - "k8s.io/apimachinery/pkg/runtime" - "k8s.io/apimachinery/pkg/types" - ctrl "sigs.k8s.io/controller-runtime" - "sigs.k8s.io/controller-runtime/pkg/client" - "sigs.k8s.io/controller-runtime/pkg/handler" - "sigs.k8s.io/controller-runtime/pkg/log" - "sigs.k8s.io/controller-runtime/pkg/reconcile" - - inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" -) - -// InferaScalingAdapterReconciler keeps an adapter's status in step with the -// workload it scales. -// -// It deliberately does not write the workload. The InferaDeployment reconciler -// reads adapters when it builds children, so there is exactly one writer of any -// child `.Spec`. Two writers would not merely race -- that reconciler assigns -// the whole spec on every pass, so the loser is reverted within seconds, which -// is precisely the failure an HPA pointed at the child Deployment hits today. -// -// What this controller owns is the half `/scale` reads back: `status.replicas` -// from the live workload, and `status.selector`, without which HorizontalPod- -// Autoscaler refuses to scale the resource at all. -type InferaScalingAdapterReconciler struct { - client.Client - Scheme *runtime.Scheme -} - -// +kubebuilder:rbac:groups=infera.amd.com,resources=inferascalingadapters,verbs=get;list;watch;create;update;patch;delete -// +kubebuilder:rbac:groups=infera.amd.com,resources=inferascalingadapters/status,verbs=get;update;patch -// +kubebuilder:rbac:groups=infera.amd.com,resources=inferascalingadapters/scale,verbs=get;update;patch - -func (r *InferaScalingAdapterReconciler) Reconcile( - ctx context.Context, req ctrl.Request, -) (ctrl.Result, error) { - lg := log.FromContext(ctx) - - adapter := &inferav1alpha1.InferaScalingAdapter{} - if err := r.Get(ctx, req.NamespacedName, adapter); err != nil { - return ctrl.Result{}, client.IgnoreNotFound(err) - } - if !adapter.DeletionTimestamp.IsZero() { - return ctrl.Result{}, nil - } - - st := inferav1alpha1.InferaScalingAdapterStatus{ - ObservedGeneration: adapter.Generation, - } - - idep := &inferav1alpha1.InferaDeployment{} - err := r.Get(ctx, types.NamespacedName{ - Name: adapter.Spec.DeploymentRef, Namespace: adapter.Namespace, - }, idep) - switch { - case apierrors.IsNotFound(err): - return r.degraded(ctx, adapter, st, "TargetNotFound", - fmt.Sprintf("no InferaDeployment %q in this namespace", adapter.Spec.DeploymentRef)) - case err != nil: - return ctrl.Result{}, err - } - - svc, ok := idep.Spec.Services[adapter.Spec.ServiceName] - if !ok { - return r.degraded(ctx, adapter, st, "ServiceNotFound", - fmt.Sprintf("InferaDeployment %q has no service %q", - adapter.Spec.DeploymentRef, adapter.Spec.ServiceName)) - } - - // The selector must be a serialized string, not a structured selector -- - // that is what the scale subresource contract requires, and HPA rejects a - // target without one. - st.Selector = labels.SelectorFromSet( - labelsFor(idep.Name, adapter.Spec.ServiceName)).String() - - name := idep.Name + "-" + adapter.Spec.ServiceName - key := types.NamespacedName{Name: name, Namespace: idep.Namespace} - if svc.NumberOfNodes > 1 { - u := &unstructured.Unstructured{} - u.SetGroupVersionKind(lwsGVK()) - if err := r.Get(ctx, key, u); err == nil { - if v, ok, _ := unstructured.NestedInt64(u.Object, "status", "replicas"); ok { - st.Replicas = int32(v) - } - if v, ok, _ := unstructured.NestedInt64(u.Object, "status", "readyReplicas"); ok { - st.ReadyReplicas = int32(v) - } - } - } else { - dep := &appsv1.Deployment{} - if err := r.Get(ctx, key, dep); err == nil { - st.Replicas = dep.Status.Replicas - st.ReadyReplicas = dep.Status.ReadyReplicas - } - } - - msg := "adapter is inert: spec.replicas unset, the InferaDeployment's own replicas apply" - if adapter.Spec.Replicas != nil { - msg = fmt.Sprintf("driving %s/%s to %d replica(s)", - adapter.Spec.DeploymentRef, adapter.Spec.ServiceName, *adapter.Spec.Replicas) - } - setCondition(&st.Conditions, adapter.Generation, "Ready", metav1.ConditionTrue, "Resolved", msg) - setCondition(&st.Conditions, adapter.Generation, "Degraded", metav1.ConditionFalse, "Resolved", msg) - - lg.V(1).Info("scaling adapter reconciled", "target", adapter.Spec.DeploymentRef, - "service", adapter.Spec.ServiceName, "observed", st.Replicas) - return r.writeStatus(ctx, adapter, st) -} - -func (r *InferaScalingAdapterReconciler) degraded( - ctx context.Context, a *inferav1alpha1.InferaScalingAdapter, - st inferav1alpha1.InferaScalingAdapterStatus, reason, msg string, -) (ctrl.Result, error) { - setCondition(&st.Conditions, a.Generation, "Ready", metav1.ConditionFalse, reason, msg) - setCondition(&st.Conditions, a.Generation, "Degraded", metav1.ConditionTrue, reason, msg) - res, err := r.writeStatus(ctx, a, st) - if err != nil { - return res, err - } - // A dangling adapter usually means the target has not been created yet, so - // retry rather than waiting for an event on an object that does not exist. - return ctrl.Result{RequeueAfter: 30 * time.Second}, nil -} - -func (r *InferaScalingAdapterReconciler) writeStatus( - ctx context.Context, a *inferav1alpha1.InferaScalingAdapter, - st inferav1alpha1.InferaScalingAdapterStatus, -) (ctrl.Result, error) { - if equalStatus(a.Status, st) { - return ctrl.Result{}, nil - } - a.Status = st - return ctrl.Result{}, r.Status().Update(ctx, a) -} - -func equalStatus(a, b inferav1alpha1.InferaScalingAdapterStatus) bool { - if a.Replicas != b.Replicas || a.ReadyReplicas != b.ReadyReplicas || - a.Selector != b.Selector || a.ObservedGeneration != b.ObservedGeneration || - len(a.Conditions) != len(b.Conditions) { - return false - } - for i := range a.Conditions { - if a.Conditions[i].Type != b.Conditions[i].Type || - a.Conditions[i].Status != b.Conditions[i].Status || - a.Conditions[i].Reason != b.Conditions[i].Reason || - a.Conditions[i].Message != b.Conditions[i].Message { - return false - } - } - return true -} - -func setCondition( - conds *[]metav1.Condition, gen int64, typ string, - status metav1.ConditionStatus, reason, msg string, -) { - for i := range *conds { - if (*conds)[i].Type == typ { - c := &(*conds)[i] - if c.Status != status { - c.LastTransitionTime = metav1.Now() - } - c.Status, c.Reason, c.Message, c.ObservedGeneration = status, reason, msg, gen - return - } - } - *conds = append(*conds, metav1.Condition{ - Type: typ, Status: status, Reason: reason, Message: msg, - ObservedGeneration: gen, LastTransitionTime: metav1.Now(), - }) -} - -func (r *InferaScalingAdapterReconciler) SetupWithManager(mgr ctrl.Manager) error { - // Watching the InferaDeployment matters as much as the adapter itself: a - // scale write only changes `spec.replicas` here, and the workload does not - // move until the other reconciler runs. Without this the adapter's status - // would lag by a resync period after every scale. - return ctrl.NewControllerManagedBy(mgr). - For(&inferav1alpha1.InferaScalingAdapter{}). - Watches( - &inferav1alpha1.InferaDeployment{}, - handler.EnqueueRequestsFromMapFunc(r.adaptersForDeployment), - ). - Complete(r) -} - -func (r *InferaScalingAdapterReconciler) adaptersForDeployment( - ctx context.Context, obj client.Object, -) []reconcile.Request { - list := &inferav1alpha1.InferaScalingAdapterList{} - if err := r.List(ctx, list, client.InNamespace(obj.GetNamespace())); err != nil { - return nil - } - var out []reconcile.Request - for i := range list.Items { - if list.Items[i].Spec.DeploymentRef != obj.GetName() { - continue - } - out = append(out, reconcile.Request{NamespacedName: types.NamespacedName{ - Name: list.Items[i].Name, Namespace: list.Items[i].Namespace, - }}) - } - return out -} diff --git a/deploy/operator/internal/controller/scale_paths_test.go b/deploy/operator/internal/controller/scale_paths_test.go new file mode 100644 index 00000000..c3a6555d --- /dev/null +++ b/deploy/operator/internal/controller/scale_paths_test.go @@ -0,0 +1,315 @@ +/* +Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + +SPDX-License-Identifier: MIT +*/ + +package controller + +import ( + "context" + "testing" + + appsv1 "k8s.io/api/apps/v1" + rbacv1 "k8s.io/api/rbac/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/types" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + + inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" +) + +// Which of the three ways to write a replica count actually reaches the pods. +// These are easy to conflate and they behave differently on purpose, so the +// distinction is pinned here rather than left to a README. + +func scaleScheme(t *testing.T) *runtime.Scheme { + t.Helper() + s := testScheme(t) + if err := rbacv1.AddToScheme(s); err != nil { + t.Fatalf("add rbac scheme: %v", err) + } + return s +} + +func idepWith(replicas int32) *inferav1alpha1.InferaDeployment { + r := replicas + return &inferav1alpha1.InferaDeployment{ + ObjectMeta: metav1.ObjectMeta{Name: "qwen", Namespace: "ns"}, + Spec: inferav1alpha1.InferaDeploymentSpec{ + Image: "infera:test", + Services: map[string]inferav1alpha1.ServiceSpec{ + "decode": { + ComponentType: inferav1alpha1.ComponentTypeWorker, + Replicas: &r, + NumberOfNodes: 1, + }, + }, + }, + } +} + +func reconcileOnce(t *testing.T, cl client.Client, s *runtime.Scheme) { + t.Helper() + r := &InferaDeploymentReconciler{Client: cl, Scheme: s} + _, err := r.Reconcile(context.Background(), ctrl.Request{ + NamespacedName: types.NamespacedName{Name: "qwen", Namespace: "ns"}, + }) + if err != nil { + t.Fatalf("reconcile: %v", err) + } +} + +func childReplicas(t *testing.T, cl client.Client) int32 { + t.Helper() + dep := &appsv1.Deployment{} + key := types.NamespacedName{Name: "qwen-decode", Namespace: "ns"} + if err := cl.Get(context.Background(), key, dep); err != nil { + t.Fatalf("get child Deployment: %v", err) + } + if dep.Spec.Replicas == nil { + t.Fatal("child Deployment has no replicas set") + } + return *dep.Spec.Replicas +} + +// Editing the CR is the normal path and it works: the CR is the desired state, +// so a change to it is what reconciliation exists to propagate. +func TestEditingTheCRScales(t *testing.T) { + s := scaleScheme(t) + idep := idepWith(2) + cl := fake.NewClientBuilder().WithScheme(s).WithObjects(idep). + WithStatusSubresource(idep).Build() + + reconcileOnce(t, cl, s) + if got := childReplicas(t, cl); got != 2 { + t.Fatalf("initial: child has %d replicas, want 2", got) + } + + // The user edits the CR. + live := &inferav1alpha1.InferaDeployment{} + key := types.NamespacedName{Name: "qwen", Namespace: "ns"} + if err := cl.Get(context.Background(), key, live); err != nil { + t.Fatalf("get idep: %v", err) + } + five := int32(5) + svc := live.Spec.Services["decode"] + svc.Replicas = &five + live.Spec.Services["decode"] = svc + if err := cl.Update(context.Background(), live); err != nil { + t.Fatalf("update idep: %v", err) + } + + reconcileOnce(t, cl, s) + if got := childReplicas(t, cl); got != 5 { + t.Fatalf("after editing the CR: child has %d replicas, want 5", got) + } +} + +// Editing the *child* is the path that does not survive, and that is the +// intended behaviour of any operator: the child is derived state, so the next +// pass restores it from the CR. This is why an HPA has to be pointed at a +// scaling adapter and never at the generated Deployment. +func TestEditingTheChildIsReverted(t *testing.T) { + s := scaleScheme(t) + idep := idepWith(2) + cl := fake.NewClientBuilder().WithScheme(s).WithObjects(idep). + WithStatusSubresource(idep).Build() + + reconcileOnce(t, cl, s) + + // Something scales the generated Deployment directly. + dep := &appsv1.Deployment{} + key := types.NamespacedName{Name: "qwen-decode", Namespace: "ns"} + if err := cl.Get(context.Background(), key, dep); err != nil { + t.Fatalf("get child: %v", err) + } + three := int32(3) + dep.Spec.Replicas = &three + if err := cl.Update(context.Background(), dep); err != nil { + t.Fatalf("update child: %v", err) + } + + reconcileOnce(t, cl, s) + if got := childReplicas(t, cl); got != 2 { + t.Fatalf("child edit survived reconciliation: %d, want it reverted to the CR's 2", got) + } +} + +// Scaling down through the CR has to land on pods that drain, not pods that +// get cut. The two features are built separately -- the replica count comes +// from the CR, the graceful shutdown from what the operator injects into the +// pod template -- so this checks they meet: a Deployment produced by a normal +// reconcile carries the preStop delay and a grace period long enough to cover +// the whole shutdown. +// +// Without preStop the router keeps assigning work for the entire termination; +// without the grace covering preStop + drain + teardown the kubelet SIGKILLs +// mid-drain. Either one silently turns a graceful scale-down back into a kill. +func TestScalingDownThroughTheCRLandsOnDrainablePods(t *testing.T) { + s := scaleScheme(t) + idep := idepWith(3) + // A long drain, the case where a fixed grace period used to fall short. + svc := idep.Spec.Services["decode"] + svc.Args = []string{"--drain-timeout", "300"} + idep.Spec.Services["decode"] = svc + + cl := fake.NewClientBuilder().WithScheme(s).WithObjects(idep). + WithStatusSubresource(idep).Build() + reconcileOnce(t, cl, s) + + dep := &appsv1.Deployment{} + key := types.NamespacedName{Name: "qwen-decode", Namespace: "ns"} + if err := cl.Get(context.Background(), key, dep); err != nil { + t.Fatalf("get child: %v", err) + } + spec := dep.Spec.Template.Spec + + if len(spec.Containers) == 0 { + t.Fatal("no containers in the pod template") + } + lc := spec.Containers[0].Lifecycle + if lc == nil || lc.PreStop == nil || lc.PreStop.Exec == nil { + t.Fatal("no preStop hook: the router would keep routing to a condemned pod") + } + + if spec.TerminationGracePeriodSeconds == nil { + t.Fatal("no terminationGracePeriodSeconds: the kubelet default is 30s, far under a 300s drain") + } + want := int64(workerPreStopDrainSeconds + 300 + workerTeardownHeadroomSeconds) + if got := *spec.TerminationGracePeriodSeconds; got != want { + t.Fatalf("grace = %ds, want %ds (preStop %d + drain 300 + teardown %d)", + got, want, workerPreStopDrainSeconds, workerTeardownHeadroomSeconds) + } + + // And the scale-down itself still works on that same object. + live := &inferav1alpha1.InferaDeployment{} + if err := cl.Get(context.Background(), types.NamespacedName{Name: "qwen", Namespace: "ns"}, live); err != nil { + t.Fatalf("get idep: %v", err) + } + one := int32(1) + svc = live.Spec.Services["decode"] + svc.Replicas = &one + live.Spec.Services["decode"] = svc + if err := cl.Update(context.Background(), live); err != nil { + t.Fatalf("update idep: %v", err) + } + reconcileOnce(t, cl, s) + if got := childReplicas(t, cl); got != 1 { + t.Fatalf("scale down through the CR: child has %d replicas, want 1", got) + } +} + +// Multi-node workers are torn down a whole group at a time, so the same +// guarantees have to hold on the LWS path -- where the pod template travels +// through a different builder. +func TestMultiNodePodsAlsoDrain(t *testing.T) { + s := scaleScheme(t) + s.AddKnownTypeWithName(lwsGVK(), &unstructured.Unstructured{}) + + idep := idepWith(2) + svc := idep.Spec.Services["decode"] + svc.NumberOfNodes = 3 + svc.Args = []string{"--drain-timeout", "180"} + idep.Spec.Services["decode"] = svc + + cl := fake.NewClientBuilder().WithScheme(s).WithObjects(idep). + WithStatusSubresource(idep).Build() + reconcileOnce(t, cl, s) + + u := &unstructured.Unstructured{} + u.SetGroupVersionKind(lwsGVK()) + if err := cl.Get(context.Background(), + types.NamespacedName{Name: "qwen-decode", Namespace: "ns"}, u); err != nil { + t.Fatalf("get child LWS: %v", err) + } + + grace, found, err := unstructured.NestedInt64(u.Object, + "spec", "leaderWorkerTemplate", "workerTemplate", "spec", "terminationGracePeriodSeconds") + if err != nil || !found { + t.Fatalf("LWS pod template has no terminationGracePeriodSeconds (found=%v, err=%v)", found, err) + } + want := int64(workerPreStopDrainSeconds + 180 + workerTeardownHeadroomSeconds) + if grace != want { + t.Fatalf("LWS grace = %ds, want %ds", grace, want) + } + + containers, found, err := unstructured.NestedSlice(u.Object, + "spec", "leaderWorkerTemplate", "workerTemplate", "spec", "containers") + if err != nil || !found || len(containers) == 0 { + t.Fatalf("LWS pod template has no containers (found=%v, err=%v)", found, err) + } + c, _ := containers[0].(map[string]any) + if _, ok := c["lifecycle"]; !ok { + t.Fatal("LWS container has no lifecycle/preStop: a condemned group keeps receiving work") + } +} + +// A LeaderWorkerSet carries a real scale subresource of its own, so an HPA can +// be pointed straight at the generated LWS and the write will succeed. It still +// does not work, for the same reason it does not work on the generated +// Deployment: reconciliation assigns the whole child spec every pass, replicas +// included. The scale write lands, and the next reconcile overwrites it. +// +// This is worth pinning because the LWS case looks different from the outside +// -- `kubectl get lws` shows a scale subresource, HPA reports success, nothing +// errors -- and the only symptom is a replica count that keeps snapping back. +func TestEditingTheChildLWSIsAlsoReverted(t *testing.T) { + s := scaleScheme(t) + s.AddKnownTypeWithName(lwsGVK(), &unstructured.Unstructured{}) + s.AddKnownTypeWithName(lwsGVK().GroupVersion().WithKind(lwsKind+"List"), + &unstructured.UnstructuredList{}) + + idep := idepWith(2) + svc := idep.Spec.Services["decode"] + svc.NumberOfNodes = 3 // multi-node -> LeaderWorkerSet instead of Deployment + idep.Spec.Services["decode"] = svc + + cl := fake.NewClientBuilder().WithScheme(s).WithObjects(idep). + WithStatusSubresource(idep).Build() + reconcileOnce(t, cl, s) + + get := func() *unstructured.Unstructured { + u := &unstructured.Unstructured{} + u.SetGroupVersionKind(lwsGVK()) + key := types.NamespacedName{Name: "qwen-decode", Namespace: "ns"} + if err := cl.Get(context.Background(), key, u); err != nil { + t.Fatalf("get child LWS: %v", err) + } + return u + } + replicas := func(u *unstructured.Unstructured) int64 { + v, found, err := unstructured.NestedInt64(u.Object, "spec", "replicas") + if err != nil || !found { + t.Fatalf("LWS has no spec.replicas (found=%v, err=%v)", found, err) + } + return v + } + + lws := get() + if got := replicas(lws); got != 2 { + t.Fatalf("initial: LWS has %d groups, want 2", got) + } + + // An HPA scales the LWS directly -- exactly what its scale subresource + // invites, and exactly what does not survive. + if err := unstructured.SetNestedField(lws.Object, int64(6), "spec", "replicas"); err != nil { + t.Fatalf("set replicas: %v", err) + } + if err := cl.Update(context.Background(), lws); err != nil { + t.Fatalf("update child LWS: %v", err) + } + if got := replicas(get()); got != 6 { + t.Fatalf("precondition: the scale write itself must land, got %d", got) + } + + reconcileOnce(t, cl, s) + + if got := replicas(get()); got != 2 { + t.Fatalf("LWS scale survived reconciliation: %d groups, want it reverted to 2", got) + } +} diff --git a/deploy/operator/internal/controller/watches_test.go b/deploy/operator/internal/controller/watches_test.go new file mode 100644 index 00000000..4725d912 --- /dev/null +++ b/deploy/operator/internal/controller/watches_test.go @@ -0,0 +1,51 @@ +/* +Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + +SPDX-License-Identifier: MIT +*/ + +package controller + +import ( + "testing" + + appsv1 "k8s.io/api/apps/v1" + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/meta" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/runtime/schema" + + inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" +) + +func testScheme(t *testing.T) *runtime.Scheme { + t.Helper() + s := runtime.NewScheme() + if err := inferav1alpha1.AddToScheme(s); err != nil { + t.Fatalf("add infera scheme: %v", err) + } + if err := appsv1.AddToScheme(s); err != nil { + t.Fatalf("add apps scheme: %v", err) + } + if err := corev1.AddToScheme(s); err != nil { + t.Fatalf("add core scheme: %v", err) + } + return s +} + +// The LWS watch is registered only when the CRD is served. controller-runtime +// builds an informer per watched type at startup and one for an unserved kind +// fails the manager, so a single-node cluster without LWS installed must not +// have the operator refuse to start. +func TestLwsInstalled(t *testing.T) { + empty := meta.NewDefaultRESTMapper(nil) + if lwsInstalled(empty) { + t.Fatal("no LWS CRD: reported installed, the manager would fail to start") + } + + withLWS := meta.NewDefaultRESTMapper([]schema.GroupVersion{lwsGVK().GroupVersion()}) + withLWS.Add(lwsGVK(), meta.RESTScopeNamespace) + if !lwsInstalled(withLWS) { + t.Fatal("LWS CRD present: reported missing, multi-node status would lag a resync") + } +} diff --git a/manual/features/scaling.md b/manual/features/scaling.md index 7fe30027..3849a2ac 100644 --- a/manual/features/scaling.md +++ b/manual/features/scaling.md @@ -186,12 +186,17 @@ etcd deployment in two ways — and gains one stage. **Discovery is a Pod annotation, not an etcd lease.** Registering writes `infera.amd.com/worker-info` on the worker's own Pod; deregistering clears it. -The registry additionally drops a Pod the moment it carries a +The registry additionally marks a Pod `DRAINING` the moment it carries a `deletionTimestamp`, without waiting for the container to exit. That matters because a terminating Pod keeps `phase: Running` — without the check it would stay a routing candidate for the whole `preStop` delay, turning a hook meant to make shutdown graceful into extra seconds of accepting work about to be killed. +The mark, rather than an outright removal, is what keeps the two timings above +distinguishable on this backend too: the worker leaves routing immediately and +its record stays until it clears its own annotation at the end of the drain, so +`/v1/workers` shows a rollout in progress instead of a worker that vanished. + **There is a `preStop` delay before `SIGTERM`.** The operator injects `sleep 15`, so the full sequence is: @@ -359,60 +364,50 @@ active KV transfer. The PD handoff queues are counted in the drain, but that path has not been exercised on hardware. ``` -## Autoscaling - -Infera ships no autoscaler. It does ship the standard interface one drives. +## Scaling a deployment -### The `/scale` subresource - -An `InferaDeployment` cannot carry `/scale` itself, and that is a property of -its shape rather than a missing feature: `spec.services` is a map with -user-chosen keys, while the scale subresource requires `specReplicasPath` to be -a *static* dot-notation JSONPath. There is no way to name "the replicas of an -arbitrary map entry". - -So scaling gets its own object — one `InferaScalingAdapter` per scalable -service: - -```yaml -apiVersion: infera.amd.com/v1alpha1 -kind: InferaScalingAdapter -metadata: {name: qwen-decode, namespace: infera} -spec: - deploymentRef: qwen # the InferaDeployment - serviceName: decode # a key in its spec.services - replicas: 2 -``` +Edit the service's `replicas` in the `InferaDeployment`. That is the only +supported way in, and it is the only write that survives: ```bash -kubectl scale inferascalingadapter/qwen-decode --replicas=5 +kubectl patch inferadeployment qwen --type=merge \ + -p '{"spec":{"services":{"decode":{"replicas":5}}}}' ``` -That is the whole integration. `kubectl scale`, `HorizontalPodAutoscaler`, KEDA -and a custom planner all work through it with no per-tool support in the -operator — a `scaleTargetRef` of kind `InferaScalingAdapter` is all an HPA -needs. +For a multi-node service the count is **groups**, not pods: `replicas: 5` with +`numberOfNodes: 3` is fifteen pods and five servable instances, since only +node-rank 0 of each group registers. -**One writer, always.** While an adapter has `spec.replicas` set it owns that -service's count: the InferaDeployment reconciler reads the adapter instead of -the CR's own `replicas`. Clearing `spec.replicas`, or deleting the adapter, hands -control back. Creating an adapter without `spec.replicas` is deliberately inert, -so an autoscaler can be attached and watched before it is trusted. +Pods removed by a scale-down drain first — the operator injects the `preStop` +delay and a grace period sized from `--drain-timeout`, so the sequence is the +same one `kubectl delete pod` follows. -This matters because the reconciler assigns the whole child `.Spec` on every -pass. Anything else writing that Deployment loses — which is exactly what -happens to an HPA pointed straight at it: **measured, a `kubectl scale` to 3 was -reverted to 1 in under 3 seconds.** Going through the adapter is not a -convention, it is the only thing that survives. +```{warning} +Do **not** scale the generated `Deployment` or `LeaderWorkerSet` directly. Both +carry a real `/scale` subresource, so the write succeeds and nothing reports an +error — and then the next reconcile reverts it, because this reconciler assigns +the whole child `.Spec` on every pass. Measured: a `kubectl scale` to 3 went +back to 1 in under 3 seconds. The only symptom is a replica count that keeps +snapping back. +``` + +## Autoscaling + +Infera ships no autoscaler, and there is currently no `/scale` surface for an +external one to drive. -`status.replicas` reports what the workload *observes*, not the desired count -echoed back. An autoscaler computes `desired = ceil(current × metric/target)`; -if `current` were the number it just asked for, it could not tell a scale-up had -not landed and would keep multiplying through a 140-second model load. +An `InferaDeployment` cannot carry `/scale` itself, and that is a property of +its shape rather than an omission: `spec.services` is a map with user-chosen +keys, while the scale subresource requires `specReplicasPath` to be a *static* +dot-notation JSONPath, and a CRD may declare only one. A single path could name +one service — hardcoding `decode`, say — which leaves every other pool, and in +a PD deployment specifically the prefill pool, with no handle at all. -### What is still missing +Pointing an autoscaler at the generated workload does not work either, for the +reason in the warning above: those objects are derived state and are rewritten +every pass. -The plumbing is not the hard part. Two things are: +Two harder problems sit behind the plumbing anyway: - **A 140-second cold start sits inside a control loop that ticks every 15 seconds.** A burst shorter than the cold start cannot be answered by adding From df01c7cc7a4acd5503e6424b2d0e3797efa35c9e Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Wed, 5 Aug 2026 07:31:52 +0000 Subject: [PATCH 29/88] build(operator): restart the version line at 0.1.0 688e608 retargeted the manager image and the chart at docker.io/rocm. Opening that line at 0.1.0 rather than carrying 0.1.3 across keeps the published numbering readable: the first release under the new name is the first number. Resets all three copies of the version together -- Chart.yaml (version and appVersion), values.yaml (image.tag), and the Makefile's IMG default. They are three spellings of one number, and moving fewer than all three is the quiet failure: the chart installs and then pulls a tag nobody pushed. Verified: `helm template` renders docker.io/rocm/infera:operator-v0.1.0 alongside app.kubernetes.io/version 0.1.0, and `make helm-package` produces infera-operator-0.1.0.tgz. Co-authored-by: Cursor Signed-off-by: leiwei12 --- deploy/operator/Makefile | 2 +- deploy/operator/helm/infera-operator/Chart.yaml | 4 ++-- deploy/operator/helm/infera-operator/values.yaml | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/deploy/operator/Makefile b/deploy/operator/Makefile index 663bee1a..c779242c 100644 --- a/deploy/operator/Makefile +++ b/deploy/operator/Makefile @@ -3,7 +3,7 @@ # make docker-build docker-push IMG=docker.io/inferaimage/infera:0.1.0 # The manager ships as a tag of the shared rocm/infera repo, matching the # engine images (sglang-v0.1.1, server-v0.1.1, ...). -IMG ?= docker.io/rocm/infera:operator-v0.1.3 +IMG ?= docker.io/rocm/infera:operator-v0.1.0 # OCI registry namespace the packaged Helm chart is pushed to (chart repo name # comes from Chart.yaml: infera-operator). CHART_REGISTRY ?= oci://docker.io/rocm diff --git a/deploy/operator/helm/infera-operator/Chart.yaml b/deploy/operator/helm/infera-operator/Chart.yaml index 74dbc370..730b7978 100644 --- a/deploy/operator/helm/infera-operator/Chart.yaml +++ b/deploy/operator/helm/infera-operator/Chart.yaml @@ -7,5 +7,5 @@ apiVersion: v2 name: infera-operator description: Infera Kubernetes operator. Reconciles InferaDeployment (infera.amd.com/v1alpha1) into Deployments / LeaderWorkerSets for aggregated and disaggregated (prefill/decode) LLM serving on AMD MI300X. Self-contained — no external inference-operator dependency. type: application -version: 0.1.3 -appVersion: "0.1.3" +version: 0.1.0 +appVersion: "0.1.0" diff --git a/deploy/operator/helm/infera-operator/values.yaml b/deploy/operator/helm/infera-operator/values.yaml index 110dfb6e..55b9530e 100644 --- a/deploy/operator/helm/infera-operator/values.yaml +++ b/deploy/operator/helm/infera-operator/values.yaml @@ -12,7 +12,7 @@ image: # rocm/infera repo; the chart itself is a separate OCI repo # (rocm/infera-operator), so their tags cannot collide. repository: docker.io/rocm/infera - tag: "operator-v0.1.3" + tag: "operator-v0.1.0" pullPolicy: IfNotPresent imagePullSecrets: [] From 1b4d33ff003a1b0ed4ea234b3091b66e10d82c12 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Sat, 8 Aug 2026 11:21:43 +0000 Subject: [PATCH 30/88] fix(router): stop the breaker wedging a worker that has recovered Claiming and releasing the half-open probe slot were never paired. `filter` claims one for every candidate it lets through, and the policy then dispatches to exactly one of them, so the rest are never told how they did -- the slot stays held, the worker stays half-open, and it never routes again until the router restarts. Measured through the real dispatch loop: a worker that recovered saw none of the next 50 requests over 50 minutes. The claim is now bounded by probe_timeout instead of waiting for an outcome that may never come. That also covers the other ways a slot leaked: a 4xx, a client disconnect, or a probe that simply never returns -- which is the behaviour of the wedged worker this whole class exists to route around. A 4xx now frees the slot without scoring the worker. It was counted as recovery, which resets the failure count and closes an open breaker: a worker alternating 500s and 400s could never reach three *consecutive* failures, and one bad client could reopen a breaker holding back a genuinely broken worker. PD recorded failures in five places and successes in none, so on that path "three in a row" had quietly become "three ever" -- a healthy worker failing once a day tripped on day three -- and a successful probe never closed the breaker. Both entry points funnel through one dispatch, so the outcome is recorded there rather than at a dozen returns, which is how it came to be missing in the first place. Anything but a clean response is neutral for both roles: one of the two may be at fault and this layer cannot tell which. Entries can now be forgotten. Worker ids are addresses no rebuilt Pod reuses, so each rollout stranded another entry and another pair of Prometheus series; wiring it to discovery is the next commit. Co-authored-by: Cursor Signed-off-by: leiwei12 --- infera/router/breaker.py | 78 +++++++++++++++++++---- infera/router/disagg.py | 37 +++++++++++ infera/router/mixed.py | 12 +++- tests/unit/router/test_breaker.py | 50 +++++++++++++++ tests/unit/router/test_disagg_breaker.py | 81 +++++++++++++++++++++++- 5 files changed, 244 insertions(+), 14 deletions(-) diff --git a/infera/router/breaker.py b/infera/router/breaker.py index 91764ff2..fadf97c5 100644 --- a/infera/router/breaker.py +++ b/infera/router/breaker.py @@ -31,11 +31,15 @@ States are the usual three. ``closed`` routes normally. After ``failure_threshold`` consecutive failures the breaker goes ``open`` and the worker is excluded for ``cooldown`` seconds. It then becomes ``half_open`` and -admits exactly one probe: success closes it and clears the count, failure +admits one probe at a time: success closes it and clears the count, failure reopens it with the cooldown doubled, up to ``max_cooldown``. Backing off matters because the common cause -- a worker wedged on a bad KV handoff -- does not resolve on the first retry, and a fixed cooldown turns into a probe every ``cooldown`` seconds forever. + +"One at a time" is bounded by ``probe_timeout`` rather than by waiting for an +outcome, because the outcome may never arrive: the slot is claimed while +filtering candidates, and only the one the policy dispatches to reports back. """ from __future__ import annotations @@ -90,8 +94,9 @@ class _Entry: opens_until: float = 0.0 # Cooldown applied on the *next* trip; doubles each time a probe fails. next_cooldown: float = 0.0 - # True while a half-open probe is in flight, so only one is admitted. - probe_in_flight: bool = False + # When the outstanding half-open probe was admitted, so only one is in + # flight at a time. None means the slot is free. + probe_started_at: float | None = None trips: int = 0 @@ -106,6 +111,16 @@ class CircuitBreaker: failure_threshold: int = 3 cooldown: float = 5.0 max_cooldown: float = 60.0 + #: How long a claimed probe slot is honoured before it is reclaimed. + #: + #: Claiming and releasing the slot are not paired: ``filter`` claims one for + #: every candidate it lets through, and the policy then dispatches exactly + #: one of them, so the rest are never told how they did. A 4xx, a client + #: disconnect or a request that never returns leaves the slot held too. + #: Bounding the claim keeps any of those from wedging a healthy worker out + #: of rotation permanently. The cost of reclaiming too early is one extra + #: probe; the cost of never reclaiming is a worker lost until restart. + probe_timeout: float = 60.0 #: Injectable clock, so tests do not sleep. now: object = field(default=time.monotonic) _entries: dict[str, _Entry] = field(default_factory=dict, init=False) @@ -137,17 +152,26 @@ def allows(self, worker_id: str) -> bool: e = self._entries.get(worker_id) if e is None or e.state is BreakerState.CLOSED: return True + now = self.now() if e.state is BreakerState.OPEN: - if self.now() < e.opens_until: + if now < e.opens_until: return False e.state = BreakerState.HALF_OPEN - e.probe_in_flight = False + e.probe_started_at = None _observe(worker_id, e.state) logger.info("breaker: worker %s half-open, admitting one probe", worker_id) - # half_open: admit exactly one probe. - if e.probe_in_flight: - return False - e.probe_in_flight = True + # half_open: one probe at a time, and only for as long as a probe could + # plausibly still be running -- see probe_timeout for why the claim has + # to expire rather than wait for an outcome that may never come. + if e.probe_started_at is not None: + if now - e.probe_started_at < self.probe_timeout: + return False + logger.info( + "breaker: worker %s probe slot unclaimed after %.0fs; admitting another", + worker_id, + self.probe_timeout, + ) + e.probe_started_at = now return True def filter(self, workers): @@ -187,10 +211,42 @@ def record_success(self, worker_id: str) -> None: logger.info("breaker: worker %s recovered, closing", worker_id) e.consecutive_failures = 0 e.state = BreakerState.CLOSED - e.probe_in_flight = False + e.probe_started_at = None e.next_cooldown = self.cooldown _observe(worker_id, e.state) + def record_neutral(self, worker_id: str) -> None: + """Release the probe slot without scoring the worker either way. + + For an outcome that says nothing about worker health -- a 4xx, which + every worker would answer identically, or a 429, which is backpressure + the policy already routes around. Counting it as recovery is as wrong + as counting it as failure: it would reset the failure count and close + an open breaker, so a worker alternating 500s and 400s would never + accumulate the consecutive failures needed to trip. But the slot such a + request consumed must still come back, or one bad client can wedge a + recovering worker out of rotation. + """ + e = self._entries.get(worker_id) + if e is None: + return + e.probe_started_at = None + + def forget(self, worker_id: str) -> None: + """Drop everything remembered about a worker that has left the fleet. + + Worker ids are addresses and a rebuilt Pod never reuses one, so without + this every rollout strands another entry -- and another pair of + Prometheus series, since both are labelled by worker id. + """ + if self._entries.pop(worker_id, None) is None: + return + for collector in (metrics.worker_breaker_state, metrics.worker_breaker_trips_total): + try: + collector.remove(worker_id) + except Exception: # noqa: BLE001 - never registered, or already gone + pass + def record_failure(self, worker_id: str) -> None: """Record a pre-first-byte dispatch failure.""" if not self.enabled: @@ -198,7 +254,7 @@ def record_failure(self, worker_id: str) -> None: e = self._entry(worker_id) e.consecutive_failures += 1 was_probe = e.state is BreakerState.HALF_OPEN - e.probe_in_flight = False + e.probe_started_at = None if was_probe: # A failed probe reopens immediately and backs off further, without diff --git a/infera/router/disagg.py b/infera/router/disagg.py index 1b3ef9da..ad66a5b0 100644 --- a/infera/router/disagg.py +++ b/infera/router/disagg.py @@ -181,6 +181,43 @@ async def _run_pd( body: dict, stream: bool, path: str, + ) -> Response: + """Record the breaker outcome for both roles around the dual dispatch. + + Every exit of :meth:`_dispatch_pd` passes through here -- both the + policy-driven and gateway-driven entry points reach it -- which is the + only reason this is one place rather than a dozen. Failures are already + recorded inside, at the specific worker that caused them; what was + missing is the other half. Without it ``consecutive_failures`` never + resets, so "three failures in a row" quietly becomes "three failures + ever", and a worker that has probed successfully never closes. + + Anything other than a clean response is neutral rather than a success + for either role: one of the two may well be at fault, and this layer + cannot tell which. Neutral is idempotent against the failure already + recorded below, and still frees the probe slot. + """ + resp = await self._dispatch_pd( + obs, p_target, d_target, p_blocks, d_blocks, body, stream, path + ) + ok = getattr(resp, "status_code", 200) < 400 + for worker_id in (p_target.worker.worker_id, d_target.worker.worker_id): + if ok: + self.breaker.record_success(worker_id) + else: + self.breaker.record_neutral(worker_id) + return resp + + async def _dispatch_pd( + self, + obs, + p_target: RouteTarget, + d_target: RouteTarget, + p_blocks: list[int], + d_blocks: list[int], + body: dict, + stream: bool, + path: str, ) -> Response: """Run the PD dual-dispatch for already-selected prefill/decode targets: resolve the protocol, forge the bootstrap room/request id, and hand off diff --git a/infera/router/mixed.py b/infera/router/mixed.py index a90ecbee..9f2fe4fa 100644 --- a/infera/router/mixed.py +++ b/infera/router/mixed.py @@ -93,7 +93,15 @@ async def dispatch( tried.add(target.worker.worker_id) try: resp = await self._attempt(target, blocks, body, hints, path, stream, obs) - self.breaker.record_success(target.worker.worker_id) + # Only a clean response is evidence the worker is healthy. + # A 4xx says the request was bad -- every worker would + # answer the same -- so scoring it as recovery would reset + # the failure count and reopen a breaker that should stay + # shut. It still has to free the probe slot it took. + if getattr(resp, "status_code", 200) < 400: + self.breaker.record_success(target.worker.worker_id) + else: + self.breaker.record_neutral(target.worker.worker_id) return resp except _Retry as r: # Pre-first-byte only: _Retry is never raised once bytes have @@ -102,6 +110,8 @@ async def dispatch( # is_worker_fault(). if is_worker_fault(getattr(r.response, "status_code", 0)): self.breaker.record_failure(target.worker.worker_id) + else: + self.breaker.record_neutral(target.worker.worker_id) last_error = r.response logger.info( "failover: worker %s failed before first byte; %d worker(s) tried", diff --git a/tests/unit/router/test_breaker.py b/tests/unit/router/test_breaker.py index e8a87af8..3beb1fa2 100644 --- a/tests/unit/router/test_breaker.py +++ b/tests/unit/router/test_breaker.py @@ -226,3 +226,53 @@ def test_threshold_zero_disables_it(): assert off.state_of("w1") is BreakerState.CLOSED ws = [W("a"), W("b")] assert len(off.filter(ws)) == 2 + + +def test_a_probe_slot_taken_but_never_dispatched_is_reclaimed(cb, clock): + """The wedge: filter() claims the probe slot for *every* candidate it lets + through, while the policy dispatches exactly one of them. + + So a recovering worker routinely has its slot taken by a request that then + went elsewhere, and nothing records an outcome for it. Without a time bound + on the claim, that worker sits in half_open with the slot held forever -- + permanently out of rotation while perfectly healthy, and only a router + restart brings it back. + """ + good, bad = W("good"), W("bad") + for _ in range(3): + cb.record_failure("bad") + clock.advance(5.1) + + # A request arrives, both are offered, the policy picks `good`. + assert bad in cb.filter([good, bad]), "cooldown elapsed -> bad is due a probe" + cb.record_success("good") + + # `bad` never learns how its probe went, because it never got one. + clock.advance(cb.probe_timeout + 1.0) + assert bad in cb.filter([good, bad]), "an unused probe claim must not be permanent" + + +def test_a_neutral_outcome_frees_the_slot_without_scoring_it(cb, clock): + """4xx is evidence about the request, not the worker, so it must neither + trip the breaker nor count as a recovery -- but the probe slot it consumed + still has to come back, or the worker is wedged by one bad client.""" + for _ in range(3): + cb.record_failure("w1") + clock.advance(5.1) + assert cb.allows("w1") is True, "cooldown elapsed -> one probe admitted" + + cb.record_neutral("w1") + assert cb.state_of("w1") is BreakerState.HALF_OPEN, "a 4xx is not a recovery" + assert cb.allows("w1") is True, "but the slot is free for a real probe" + + +def test_forgetting_a_worker_drops_its_entry(cb): + """Worker ids are addresses, and a rebuilt pod never reuses one. Entries + kept for workers discovery has dropped grow without bound, and each also + pins a Prometheus series.""" + for _ in range(3): + cb.record_failure("gone") + assert "gone" in cb.snapshot() + cb.forget("gone") + assert "gone" not in cb.snapshot() + assert cb.state_of("gone") is BreakerState.CLOSED diff --git a/tests/unit/router/test_disagg_breaker.py b/tests/unit/router/test_disagg_breaker.py index 0856714a..0f0b76c0 100644 --- a/tests/unit/router/test_disagg_breaker.py +++ b/tests/unit/router/test_disagg_breaker.py @@ -18,13 +18,13 @@ import httpx import pytest -from infera.common.worker_pool import EngineType, WorkerInfo +from infera.common.worker_pool import DisaggMode, EngineType, WorkerInfo from infera.router.disagg import DisaggRouter from infera.router.policy.target import RouteTarget class _FakePolicy: - def pick(self, candidates, body): + def pick(self, candidates, body, role_hint=None): return RouteTarget(candidates[0]), [] def on_request_started(self, route_key, blocks): @@ -113,3 +113,80 @@ async def test_dual_stream_records_failure_on_unreachable_decode(): # The prefill leg is a separate pool: a wedged decode must not evict it. assert r.breaker.state_of("p1").value == "closed" await r.aclose() + + +class _RolePool: + """Unlike _FakePool, hands back the pool the caller actually asked for, so + a dispatch gets a real prefill/decode pair rather than the same worker twice.""" + + def __init__(self, prefill, decode): + self._by_mode = {DisaggMode.PREFILL: [prefill], DisaggMode.DECODE: [decode]} + + def list_active(self, model=None, mode=None): + return list(self._by_mode.get(mode, [])) + + +def _pd_worker(wid, mode): + meta = {"protocol": "sglang-bootstrap"} + if mode is DisaggMode.PREFILL: + meta["params"] = {"bootstrap_addr": f"{wid}:9000"} + return WorkerInfo( + worker_id=wid, + url=f"http://{wid}", + model_name="m", + engine=EngineType.SGLANG, + request_transport="http", + disagg_mode=mode, + disagg_meta=meta, + ) + + +def _ok_router(): + """A PD router whose every leg answers 200.""" + r = DisaggRouter( + _RolePool(_pd_worker("p1", DisaggMode.PREFILL), _pd_worker("d1", DisaggMode.DECODE)), + _FakePolicy(), + ) + r._client = httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: httpx.Response( + 200, json={"id": "x", "choices": [{"message": {"content": "hi"}}]} + ) + ) + ) + return r + + +@pytest.mark.asyncio +async def test_a_served_request_clears_the_failure_count(): + """Without a success recorded anywhere, "three consecutive failures" decays + into "three failures ever": the counter only climbs, so a healthy worker + that fails once a day trips on day three.""" + r = _ok_router() + for _ in range(2): + r.breaker.record_failure("d1") + assert r.breaker.snapshot()["d1"]["consecutive_failures"] == 2 + + await r.dispatch({"model": "m"}, stream=False) + + assert r.breaker.snapshot()["d1"]["consecutive_failures"] == 0 + assert r.breaker.state_of("d1").value == "closed" + await r.aclose() + + +@pytest.mark.asyncio +async def test_a_tripped_pd_worker_recovers_after_a_good_probe(): + """The probe is dispatched and succeeds; if nothing records that, the + worker stays half-open with its probe slot held and never routes again.""" + r = _ok_router() + for _ in range(3): + r.breaker.record_failure("d1") + assert r.breaker.state_of("d1").value == "open" + + # Let the cooldown lapse so the next dispatch is the half-open probe. + r.breaker._entries["d1"].opens_until = 0.0 + + await r.dispatch({"model": "m"}, stream=False) + + assert r.breaker.state_of("d1").value == "closed", "a good probe must close it" + await r.aclose() From 21de36b60112d0240b1a1cad4becf2b32553c8bf Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Wed, 5 Aug 2026 08:08:51 +0000 Subject: [PATCH 31/88] build(operator): publish to staging, as the engine images already do The release workflow builds engine images into docker.io/inferaimage/infera and a separate process promotes them to docker.io/rocm/infera after review; release.yml defaults to the staging repo for exactly that reason. The operator's Makefile defaulted to the promotion target instead, so the one address a bare `make docker-push` would reach was the one nobody working here has rights to write -- and if they did, it would land in the public repo without passing the review the two-repo split exists to enforce. IMG and CHART_REGISTRY now both name inferaimage. What the artefacts *say* is unchanged and deliberately so: values.yaml still embeds rocm/infera:operator-v0.1.0 and the docs still install from oci://docker.io/rocm/infera-operator, because a chart describes where it will live, not where it is staged. Embedding the staging address would survive the promotion and point every user at it. Testing a staged chart end to end therefore needs one override: helm install infera-operator oci://docker.io/inferaimage/infera-operator \ --version 0.1.0 -n infera-system --create-namespace \ --set image.repository=docker.io/inferaimage/infera Note for whoever runs the promotion: this ships two artefacts, and the second is a Helm chart rather than a container image, so an image-only promotion would leave it behind. inferaimage/infera:operator-v0.1.0 container image inferaimage/infera-operator:0.1.0 Helm chart (OCI artefact) Co-authored-by: Cursor Signed-off-by: leiwei12 --- deploy/operator/Makefile | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/deploy/operator/Makefile b/deploy/operator/Makefile index c779242c..ff571c88 100644 --- a/deploy/operator/Makefile +++ b/deploy/operator/Makefile @@ -1,12 +1,19 @@ # infera-operator — minimal build/codegen targets. -# Operator manager image + chart registry. Override on the CLI, e.g. -# make docker-build docker-push IMG=docker.io/inferaimage/infera:0.1.0 -# The manager ships as a tag of the shared rocm/infera repo, matching the -# engine images (sglang-v0.1.1, server-v0.1.1, ...). -IMG ?= docker.io/rocm/infera:operator-v0.1.0 +# The manager is a tag of the shared rocm/infera repo, alongside the engine +# images, so one repo holds everything that runs. The chart is pushed to +# rocm/infera-operator instead -- a separate OCI repo, so image tags and chart +# versions cannot collide in the same namespace. +# +# docker pull rocm/infera:operator-v0.1.0 +# helm pull oci://docker.io/rocm/infera-operator --version 0.1.0 +# +# Builds publish to the staging repo, as the engine images do; promotion to +# rocm/infera is a separate step. Override on the CLI, e.g. +# make docker-build docker-push IMG=docker.io//infera:operator-v0.1.0 +IMG ?= docker.io/inferaimage/infera:operator-v0.1.0 # OCI registry namespace the packaged Helm chart is pushed to (chart repo name -# comes from Chart.yaml: infera-operator). -CHART_REGISTRY ?= oci://docker.io/rocm +# comes from Chart.yaml: infera-operator). Staging, like IMG. +CHART_REGISTRY ?= oci://docker.io/inferaimage CHART_DIR ?= helm/infera-operator CRD_DIR ?= config/crd/bases # The chart ships its own copy of the CRD, so `helm install` can create it From 1c49ed7aebb8ee01915913babd7209f91f905106 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Sat, 8 Aug 2026 11:28:46 +0000 Subject: [PATCH 32/88] fix(server): forget a departed worker, and say when a drain metric is missing Breaker entries are now dropped when discovery drops the worker. A worker id is an address and a rebuilt Pod never reuses one, so every rollout stranded another entry and another pair of Prometheus series labelled by that id -- the label cardinality being the part that actually hurts. The breaker is built before the registry starts because the removal callback closes over it, and a worker leaving during startup would otherwise hit an unbound name. A drain reading a metrics page with only some of its series present now says so. Counting what is readable stays the behaviour -- refusing a partial page would let one renamed series stop the drain waiting at all, which is worse -- but the absent series contributes zero, so queued work can be cut while the drain reports itself done. Silence made that indistinguishable from an idle engine, which is the exact failure this module was written to prevent, and these names do drift: the vLLM KV gauge was renamed under this module. Also corrects the shutdown ordering in drain.py's header, which described deregistering before draining. The implementation announces DRAINING, drains, then deregisters -- deliberately, since a record that vanishes at the start of a drain looks exactly like a worker that crashed. Getting that backwards in the one paragraph explaining why the order matters is worth fixing. Co-authored-by: Cursor Signed-off-by: leiwei12 --- infera/common/engine_metrics.py | 18 ++++++++++++++++++ infera/engine/drain.py | 9 ++++++--- infera/server/__main__.py | 24 +++++++++++++++++------- tests/unit/engine/test_drain.py | 21 +++++++++++++++++++++ 4 files changed, 62 insertions(+), 10 deletions(-) diff --git a/infera/common/engine_metrics.py b/infera/common/engine_metrics.py index 2e5a9837..858145b8 100644 --- a/infera/common/engine_metrics.py +++ b/infera/common/engine_metrics.py @@ -27,10 +27,13 @@ from __future__ import annotations +import logging import re from infera.common.worker_pool import EngineType +logger = logging.getLogger(__name__) + #: metric key -> per-engine exposition name(s). A missing engine means "we do #: not know", which callers must distinguish from "the value is zero". Several #: names per entry means the engine renamed the series between releases and both @@ -126,17 +129,32 @@ def inflight_from_metrics(text: str, engine: EngineType) -> float | None: """ total = 0.0 seen = False + missing: list[str] = [] for key in ("requests_running", "requests_waiting"): name = metric_name(key, engine) if name is None: continue value = parse_metric(text, name) if value is None: + missing.append(name) continue total += value seen = True if not seen: return None + if missing: + # Partial readings are still worth acting on -- refusing them outright + # would mean one renamed series stops the drain waiting for anything at + # all, which is the worse failure. But it must not pass silently: the + # absent series contributes zero, so a drain can finish while the work + # it describes is still outstanding. These names do drift; the vLLM KV + # gauge was renamed under exactly this module. + logger.warning( + "drain: %s not found on the %s metrics page; its work counts as zero, " + "so in-flight requests may be cut. Check the exposition names.", + ", ".join(missing), + engine.value, + ) # Absent PD queues are genuinely zero here rather than unknown: the engine # published a metrics page and simply is not running disaggregated. for name in _DRAIN_EXTRA.get(engine, ()): diff --git a/infera/engine/drain.py b/infera/engine/drain.py index e6cdc13b..7d7c2901 100644 --- a/infera/engine/drain.py +++ b/infera/engine/drain.py @@ -19,9 +19,12 @@ Two behaviours are deliberate: -* **Deregister first, then drain.** Ordering is the whole point. Draining while - still registered just means more work arrives; the router has to stop choosing - this worker before waiting for the work it already has. +* **Announce first, then drain, then deregister.** Ordering is the whole point. + Draining while still a routing candidate just means more work arrives, so the + worker announces DRAINING -- which takes it out of the candidate list -- before + waiting for the work it already has. Deregistering comes last, because a record + that disappears at the start of the drain is indistinguishable from a worker + that crashed. * **An unreadable metric does not block shutdown.** If the engine's in-flight count cannot be determined — an unknown engine, a renamed series, a dead HTTP server — this logs loudly and returns rather than hanging until the timeout. diff --git a/infera/server/__main__.py b/infera/server/__main__.py index 5465fa00..499d346f 100644 --- a/infera/server/__main__.py +++ b/infera/server/__main__.py @@ -93,6 +93,18 @@ async def main(args) -> None: # older than what's already in the index. writer.set_reconciler(reconciler) + # Per-worker failure memory, shared by every router this process builds. + # DirectRouter never selects, so it holds one but does not consult it. + # + # Built here rather than beside the routers because on_worker_removed + # closes over it and the registry starts first, so a worker that leaves + # during startup would otherwise hit an unbound name. + breaker = CircuitBreaker( + failure_threshold=args.breaker_failure_threshold, + cooldown=args.breaker_cooldown_s, + max_cooldown=args.breaker_max_cooldown_s, + ) + # Per-worker snapshot of fields we need at removal time. The Registry # has already evicted the WorkerInfo from its pool by the time # on_worker_removed fires, so we stash the kv block at registration. @@ -153,6 +165,11 @@ def on_worker_removed(worker_id: str) -> None: except Exception: logger.exception("policy.on_worker_removed failed for %s", worker_id) + # A worker id is an address, and a rebuilt Pod never reuses one, so an + # entry left here outlives the fleet member it describes -- one more + # per rollout, each pinning a Prometheus series labelled by that id. + breaker.forget(worker_id) + # Phase 1 reconciler/subscriber cleanup. snap = kv_snapshots.pop(worker_id, None) if snap is None: @@ -224,13 +241,6 @@ def on_worker_removed(worker_id: str) -> None: logger.info("request transport: nats (per-instance subjects)") # --- Router + FastAPI app --- - # Per-worker failure memory, shared by every router this process builds. - # DirectRouter never selects, so it holds one but does not consult it. - breaker = CircuitBreaker( - failure_threshold=args.breaker_failure_threshold, - cooldown=args.breaker_cooldown_s, - max_cooldown=args.breaker_max_cooldown_s, - ) # router-mode=direct trusts an upstream GAIE EPP's per-request worker # selection (x-worker-instance-id header); auto selects in-process. if args.router_mode == "direct": diff --git a/tests/unit/engine/test_drain.py b/tests/unit/engine/test_drain.py index 86754ae1..0acb5823 100644 --- a/tests/unit/engine/test_drain.py +++ b/tests/unit/engine/test_drain.py @@ -18,6 +18,7 @@ from __future__ import annotations import asyncio +import logging import httpx import pytest @@ -59,6 +60,26 @@ def test_partial_metrics_still_count(): assert inflight_from_metrics(text, EngineType.SGLANG) == 1.0 +def test_a_missing_series_says_so(caplog): + """Counting only what is readable is the right call -- refusing a partial + page would let one renamed series stop the drain waiting at all. But the + absent series contributes zero, so queued work can be cut while the drain + reports itself finished. That has to be audible: these names do drift, and + the vLLM KV gauge was renamed under this very module. + """ + text = "sglang:num_running_reqs 1\n" # num_queue_reqs absent + with caplog.at_level(logging.WARNING, logger="infera.common.engine_metrics"): + assert inflight_from_metrics(text, EngineType.SGLANG) == 1.0 + assert "sglang:num_queue_reqs" in caplog.text + caplog.clear() + + # A complete page must stay quiet, or the warning is noise on every poll. + both = "sglang:num_running_reqs 1\nsglang:num_queue_reqs 2\n" + with caplog.at_level(logging.WARNING, logger="infera.common.engine_metrics"): + assert inflight_from_metrics(both, EngineType.SGLANG) == 3.0 + assert caplog.text == "" + + # --- the drain loop ----------------------------------------------------------- From bb81c3381893112aae9554804c602e49638c6d68 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Wed, 5 Aug 2026 08:09:10 +0000 Subject: [PATCH 33/88] docs(recipes): install the operator from the published chart Every recipe told the reader to install from deploy/operator/helm/infera-operator, a path that only resolves inside a clone of this repo. That made a git checkout a prerequisite for a step whose only purpose is to get the CRD onto a cluster, and it silently installed whatever the working tree happened to contain rather than a released version. The chart is published as an OCI artefact, so the six recipes now name it: helm install infera-operator oci://docker.io/rocm/infera-operator --version 0.1.0 \ -n infera-system --create-namespace This also matches how the same pages already refer to the engine images -- rocm/infera:sglang-v0.1.1 and friends, never the staging repo. The two `helm upgrade --install` call sites keep that form; they are idempotent on purpose, since the operator may already be installed. Co-authored-by: Cursor Signed-off-by: leiwei12 --- examples/recipes/glm5.2/README.md | 3 ++- examples/recipes/kimi-k3-optimized/README.md | 2 +- examples/recipes/kimi-k3/README.md | 3 ++- manual/recipes/glm5.2.md | 3 ++- manual/recipes/kimi-k3-optimized.md | 2 +- manual/recipes/kimi-k3.md | 3 ++- 6 files changed, 10 insertions(+), 6 deletions(-) diff --git a/examples/recipes/glm5.2/README.md b/examples/recipes/glm5.2/README.md index 84e5a639..b0811317 100644 --- a/examples/recipes/glm5.2/README.md +++ b/examples/recipes/glm5.2/README.md @@ -52,7 +52,8 @@ pinned host memory charged to the sidecar's limit. kubectl get nodes -o custom-columns=NODE:.metadata.name,GPU:.status.allocatable.'amd\.com/gpu' # the infera operator (provides the InferaDeployment CRD) -helm install infera-operator deploy/operator/helm/infera-operator -n infera-system --create-namespace +helm install infera-operator oci://docker.io/rocm/infera-operator --version 0.1.0 \ + -n infera-system --create-namespace kubectl -n infera-system rollout status deploy/infera-operator ``` diff --git a/examples/recipes/kimi-k3-optimized/README.md b/examples/recipes/kimi-k3-optimized/README.md index 81dd207d..350e225f 100644 --- a/examples/recipes/kimi-k3-optimized/README.md +++ b/examples/recipes/kimi-k3-optimized/README.md @@ -178,7 +178,7 @@ kubectl create namespace infera --dry-run=client -o yaml | kubectl apply -f - # the operator (provides the InferaDeployment CRD). Skip if already installed — # `helm install` fails on name reuse; use `helm upgrade --install` to be idempotent. -helm upgrade --install infera-operator deploy/operator/helm/infera-operator \ +helm upgrade --install infera-operator oci://docker.io/rocm/infera-operator --version 0.1.0 \ -n infera-system --create-namespace ``` diff --git a/examples/recipes/kimi-k3/README.md b/examples/recipes/kimi-k3/README.md index d79ea5ea..284329db 100644 --- a/examples/recipes/kimi-k3/README.md +++ b/examples/recipes/kimi-k3/README.md @@ -41,7 +41,8 @@ NVMe**, not NFS. Loading over NFS took about an hour here; local disk is minutes kubectl get nodes -o custom-columns=NODE:.metadata.name,GPU:.status.allocatable.'amd\.com/gpu' # the infera operator (provides the InferaDeployment CRD) -helm install infera-operator deploy/operator/helm/infera-operator -n infera-system --create-namespace +helm install infera-operator oci://docker.io/rocm/infera-operator --version 0.1.0 \ + -n infera-system --create-namespace kubectl -n infera-system rollout status deploy/infera-operator ``` diff --git a/manual/recipes/glm5.2.md b/manual/recipes/glm5.2.md index 973265d1..f170f0cd 100644 --- a/manual/recipes/glm5.2.md +++ b/manual/recipes/glm5.2.md @@ -82,7 +82,8 @@ kubectl -n infera get pods -w kubectl get nodes -o custom-columns=NODE:.metadata.name,GPU:.status.allocatable.'amd\.com/gpu' # the operator (provides the InferaDeployment CRD) -helm install infera-operator deploy/operator/helm/infera-operator -n infera-system --create-namespace +helm install infera-operator oci://docker.io/rocm/infera-operator --version 0.1.0 \ + -n infera-system --create-namespace ``` The weights are expected in the `model-cache` PVC. On k3s, install with `--data-dir` diff --git a/manual/recipes/kimi-k3-optimized.md b/manual/recipes/kimi-k3-optimized.md index 151db114..d06cba65 100644 --- a/manual/recipes/kimi-k3-optimized.md +++ b/manual/recipes/kimi-k3-optimized.md @@ -249,7 +249,7 @@ kubectl create namespace infera --dry-run=client -o yaml | kubectl apply -f - # on k3s, helm needs KUBECONFIG spelled out — kubectl finds it implicitly, helm does not export KUBECONFIG=/etc/rancher/k3s/k3s.yaml -helm upgrade --install infera-operator deploy/operator/helm/infera-operator \ +helm upgrade --install infera-operator oci://docker.io/rocm/infera-operator --version 0.1.0 \ -n infera-system --create-namespace hf download moonshotai/Kimi-K3 --local-dir /Kimi-K3 diff --git a/manual/recipes/kimi-k3.md b/manual/recipes/kimi-k3.md index 3bc73ca3..517c9944 100644 --- a/manual/recipes/kimi-k3.md +++ b/manual/recipes/kimi-k3.md @@ -95,7 +95,8 @@ kubectl -n infera get pods -w kubectl get nodes -o custom-columns=NODE:.metadata.name,GPU:.status.allocatable.'amd\.com/gpu' # the operator (provides the InferaDeployment CRD) -helm install infera-operator deploy/operator/helm/infera-operator -n infera-system --create-namespace +helm install infera-operator oci://docker.io/rocm/infera-operator --version 0.1.0 \ + -n infera-system --create-namespace ``` The weights are expected in the `model-cache` PVC. On k3s, install with `--data-dir` From da178560fda16959d90ad801fd273904ad3dce73 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Sat, 8 Aug 2026 11:41:40 +0000 Subject: [PATCH 34/88] fix(operator): report ready only when the requested replicas are there Three defects, all on the paths this branch touched. ServiceStatus.Replicas now reports what the workload has rather than what was asked for, which is right for a reader watching a scale-up land but is the wrong side of the readiness comparison: `ready < observed` only sees Pods that exist and are not ready. A replica that was never created -- unschedulable, out of quota, no GPU -- is missing from both numbers, so they agree and a 3-replica service reports ready on 2. `.status.state` is a printcolumn and the obvious readiness gate for anything orchestrating on top, so this decides whether traffic gets sent. The roll-up compares against the spec now; a service scaled to zero, or one the spec no longer mentions, does not hold the deployment back. applyUnstructured replaced the whole .spec, but these objects are only partly ours: three fields are written here and the API server defaults nine more from the LWS CRD. Every pass stripped those, the API server restored them, and CreateOrUpdate saw a difference again -- a wasted write per resync, and a write loop now that the reconciler watches LeaderWorkerSet, since each write enqueues the reconcile that makes the next one. Fields the operator sets are merged in and the rest is left alone. Drain timeout precedence is resolved per source instead of globally. On the extraPodSpec path the template is passed through verbatim, so a --drain-timeout in ServiceSpec.Args is never rendered into the container and changes nothing -- yet it outranked the INFERA_DRAIN_TIMEOUT the container actually reads. A container draining for 600s was given a 120s grace and SIGKILLed partway through, which is the failure graceSecondsFor exists to prevent. What the process sees wins; within one source a flag still beats a variable, as argparse does. Co-authored-by: Cursor Signed-off-by: leiwei12 --- .../controller/apply_idempotence_test.go | 139 ++++++++++++++++++ .../operator/internal/controller/builders.go | 75 +++++++--- .../internal/controller/builders_test.go | 49 +++++- .../controller/inferadeployment_controller.go | 77 ++++++++-- .../internal/controller/scale_paths_test.go | 5 +- .../internal/controller/status_rollup_test.go | 89 +++++++++++ 6 files changed, 393 insertions(+), 41 deletions(-) create mode 100644 deploy/operator/internal/controller/apply_idempotence_test.go create mode 100644 deploy/operator/internal/controller/status_rollup_test.go diff --git a/deploy/operator/internal/controller/apply_idempotence_test.go b/deploy/operator/internal/controller/apply_idempotence_test.go new file mode 100644 index 00000000..dcafdb25 --- /dev/null +++ b/deploy/operator/internal/controller/apply_idempotence_test.go @@ -0,0 +1,139 @@ +/* +Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + +SPDX-License-Identifier: MIT +*/ + +package controller + +import ( + "context" + "testing" + + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/types" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + + inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" +) + +// The operator writes three fields of a LeaderWorkerSet; the API server fills +// in the rest from the CRD's defaults -- startupPolicy, rolloutStrategy, +// leaderWorkerTemplate.restartPolicy and more, nine of them on LWS v1. +// +// Replacing the whole .spec strips every one of those on each pass, the API +// server puts them back, and the next pass strips them again. That was a +// wasted write every resync; it becomes a hot loop now that the reconciler +// watches LeaderWorkerSet, because the write it just made enqueues the +// request that makes the next one. +// +// So: reconciling an object that is already in the desired state must not +// write to it. +func TestApplyingTheSameLwsTwiceDoesNotWriteAgain(t *testing.T) { + s := testScheme(t) + if err := inferav1alpha1.AddToScheme(s); err != nil { + t.Fatalf("scheme: %v", err) + } + c := fake.NewClientBuilder().WithScheme(s).Build() + r := &InferaDeploymentReconciler{Client: c, Scheme: s} + ctx := context.Background() + + idep := &inferav1alpha1.InferaDeployment{} + idep.Name = "qwen" + idep.Namespace = "default" + idep.UID = "uid-1" + + desired := func() *unstructured.Unstructured { + u := &unstructured.Unstructured{} + u.SetGroupVersionKind(lwsGVK()) + u.SetName("qwen-worker") + u.SetNamespace("default") + _ = unstructured.SetNestedField(u.Object, int64(2), "spec", "replicas") + _ = unstructured.SetNestedField(u.Object, int64(2), + "spec", "leaderWorkerTemplate", "size") + return u + } + + if err := r.applyUnstructured(ctx, idep, desired()); err != nil { + t.Fatalf("first apply: %v", err) + } + + // Stand in for the API server defaulting the fields the operator omits. + live := &unstructured.Unstructured{} + live.SetGroupVersionKind(lwsGVK()) + if err := c.Get(ctx, types.NamespacedName{Name: "qwen-worker", Namespace: "default"}, live); err != nil { + t.Fatalf("get after create: %v", err) + } + _ = unstructured.SetNestedField(live.Object, "LeaderCreated", "spec", "startupPolicy") + _ = unstructured.SetNestedField(live.Object, "RollingUpdate", "spec", "rolloutStrategy", "type") + _ = unstructured.SetNestedField(live.Object, "RecreateGroupOnPodRestart", + "spec", "leaderWorkerTemplate", "restartPolicy") + if err := c.Update(ctx, live); err != nil { + t.Fatalf("apply defaults: %v", err) + } + before := live.GetResourceVersion() + + if err := r.applyUnstructured(ctx, idep, desired()); err != nil { + t.Fatalf("second apply: %v", err) + } + + after := &unstructured.Unstructured{} + after.SetGroupVersionKind(lwsGVK()) + if err := c.Get(ctx, types.NamespacedName{Name: "qwen-worker", Namespace: "default"}, after); err != nil { + t.Fatalf("get after reconcile: %v", err) + } + + if got := after.GetResourceVersion(); got != before { + t.Errorf("reconcile rewrote an unchanged object (resourceVersion %s -> %s); "+ + "with the LWS watch registered this is a write loop", before, got) + } + for _, f := range [][]string{ + {"spec", "startupPolicy"}, + {"spec", "rolloutStrategy", "type"}, + {"spec", "leaderWorkerTemplate", "restartPolicy"}, + } { + if v, ok, _ := unstructured.NestedString(after.Object, f...); !ok || v == "" { + t.Errorf("%v was stripped; the API server will re-default it and the "+ + "next pass strips it again", f) + } + } +} + +// Merging must not turn into "never update": a genuine spec change still has +// to reach the child, or scaling through the CR would silently do nothing. +func TestApplyStillPushesAChangedField(t *testing.T) { + s := testScheme(t) + c := fake.NewClientBuilder().WithScheme(s).Build() + r := &InferaDeploymentReconciler{Client: c, Scheme: s} + ctx := context.Background() + + idep := &inferav1alpha1.InferaDeployment{} + idep.Name = "qwen" + idep.Namespace = "default" + idep.UID = "uid-1" + + build := func(replicas int64) *unstructured.Unstructured { + u := &unstructured.Unstructured{} + u.SetGroupVersionKind(lwsGVK()) + u.SetName("qwen-worker") + u.SetNamespace("default") + _ = unstructured.SetNestedField(u.Object, replicas, "spec", "replicas") + return u + } + + if err := r.applyUnstructured(ctx, idep, build(2)); err != nil { + t.Fatalf("create: %v", err) + } + if err := r.applyUnstructured(ctx, idep, build(5)); err != nil { + t.Fatalf("scale: %v", err) + } + + got := &unstructured.Unstructured{} + got.SetGroupVersionKind(lwsGVK()) + if err := c.Get(ctx, types.NamespacedName{Name: "qwen-worker", Namespace: "default"}, got); err != nil { + t.Fatalf("get: %v", err) + } + if v, _, _ := unstructured.NestedInt64(got.Object, "spec", "replicas"); v != 5 { + t.Fatalf("replicas = %d, want 5 -- scaling through the CR did not land", v) + } +} diff --git a/deploy/operator/internal/controller/builders.go b/deploy/operator/internal/controller/builders.go index 0d36f320..55523799 100644 --- a/deploy/operator/internal/controller/builders.go +++ b/deploy/operator/internal/controller/builders.go @@ -74,11 +74,42 @@ func drainSeconds(v string) (int, bool) { // $INFERA_DRAIN_TIMEOUT as the flag's *default*, so an env var raises the drain // exactly as effectively as the flag does -- and parsing only the flag left the // same silent overrun through a different door. -func graceSecondsFor(args []string, env []corev1.EnvVar) int64 { +// +// Sources are given in increasing priority, and precedence is resolved per +// source rather than globally. It has to be: on the extraPodSpec path the +// template is passed through verbatim, so a --drain-timeout in ServiceSpec.Args +// is never rendered into the container and does not affect the drain at all. +// Letting that inert flag outrank the variable the container really reads sizes +// the budget for a drain that never happens, while the real one runs long and +// is killed partway through -- the exact failure this function exists to stop. +func graceSecondsFor(sources ...drainSource) int64 { drain := workerDefaultDrainTimeoutSeconds - // Environment first so an explicit flag overrides it, matching argparse: - // the variable supplies the default, the flag replaces it. - for _, e := range env { + for _, s := range sources { + if d, ok := s.drainTimeout(); ok { + drain = d + } + } + need := int64(workerPreStopDrainSeconds + drain + workerTeardownHeadroomSeconds) + if need < workerTerminationGraceSeconds { + return workerTerminationGraceSeconds + } + return need +} + +// drainSource is one place a drain timeout can be configured: a set of args and +// env vars that travel together, either both from ServiceSpec or both from the +// container itself. +type drainSource struct { + args []string + env []corev1.EnvVar +} + +// drainTimeout resolves this source alone, reporting whether it set anything. +// Env first so an explicit flag overrides it, matching argparse: the variable +// supplies the default, the flag replaces it. +func (s drainSource) drainTimeout() (int, bool) { + out, found := 0, false + for _, e := range s.env { if e.Name != drainTimeoutEnvVar { continue } @@ -87,13 +118,13 @@ func graceSecondsFor(args []string, env []corev1.EnvVar) int64 { // flag or the default. Worth knowing if a drain is ever cut short // despite a ConfigMap saying otherwise. if d, ok := drainSeconds(e.Value); ok { - drain = d + out, found = d, true } } - for i, a := range args { + for i, a := range s.args { v := "" - if a == "--drain-timeout" && i+1 < len(args) { - v = args[i+1] + if a == "--drain-timeout" && i+1 < len(s.args) { + v = s.args[i+1] } else if strings.HasPrefix(a, "--drain-timeout=") { v = strings.TrimPrefix(a, "--drain-timeout=") } @@ -101,14 +132,10 @@ func graceSecondsFor(args []string, env []corev1.EnvVar) int64 { continue } if d, ok := drainSeconds(v); ok { - drain = d + out, found = d, true } } - need := int64(workerPreStopDrainSeconds + drain + workerTeardownHeadroomSeconds) - if need < workerTerminationGraceSeconds { - return workerTerminationGraceSeconds - } - return need + return out, found } // Identity labels on every workload this operator builds. They are the only @@ -334,13 +361,17 @@ func injectWorkerRolloutDefaults( return } c := &spec.Containers[idx] - // The drain can arrive several ways: via ServiceSpec.Args/Env on the - // rendered path, or written straight into the container by an extraPodSpec - // template, which is passed through verbatim. Reading only the first would - // miss exactly the deployments most likely to have tuned it. The container's - // own values go last so they win, being what the process actually sees. - drainArgs := append(append(append([]string{}, args...), c.Command...), c.Args...) - drainEnv := append(append([]corev1.EnvVar{}, env...), c.Env...) + // The drain can arrive two ways: via ServiceSpec.Args/Env on the rendered + // path, or written straight into the container by an extraPodSpec template, + // which is passed through verbatim. Reading only the first would miss + // exactly the deployments most likely to have tuned it. They stay separate + // sources, listed in increasing priority, because the container's is what + // the process actually reads -- see graceSecondsFor. + fromService := drainSource{args: args, env: env} + fromContainer := drainSource{ + args: append(append([]string{}, c.Command...), c.Args...), + env: c.Env, + } if addReadiness && c.ReadinessProbe == nil { // SGLang's /health runs a tiny prefill self-check that often takes // >1s, so a 1s probe timeout (the k8s default) flaps the pod between @@ -366,7 +397,7 @@ func injectWorkerRolloutDefaults( }, } } - if want := graceSecondsFor(drainArgs, drainEnv); spec.TerminationGracePeriodSeconds == nil || + if want := graceSecondsFor(fromService, fromContainer); spec.TerminationGracePeriodSeconds == nil || *spec.TerminationGracePeriodSeconds < want { grace := want spec.TerminationGracePeriodSeconds = &grace diff --git a/deploy/operator/internal/controller/builders_test.go b/deploy/operator/internal/controller/builders_test.go index 6db09a4f..80718add 100644 --- a/deploy/operator/internal/controller/builders_test.go +++ b/deploy/operator/internal/controller/builders_test.go @@ -47,7 +47,7 @@ func TestGraceSecondsFor(t *testing.T) { } for _, c := range cases { t.Run(c.name, func(t *testing.T) { - if got := graceSecondsFor(c.args, nil); got != c.want { + if got := graceSecondsFor(drainSource{args: c.args}); got != c.want { t.Fatalf("graceSecondsFor(%v) = %d, want %d", c.args, got, c.want) } }) @@ -58,7 +58,7 @@ func TestGraceSecondsFor(t *testing.T) { func TestGraceCoversTheWholeShutdown(t *testing.T) { for _, drain := range []int{30, 60, 120, 300} { args := []string{"--drain-timeout", itoa(drain)} - grace := graceSecondsFor(args, nil) + grace := graceSecondsFor(drainSource{args: args}) need := int64(workerPreStopDrainSeconds + drain + workerTeardownHeadroomSeconds) if grace < need { t.Fatalf("drain=%d: grace %d < required %d -- kubelet would SIGKILL mid-drain", @@ -108,7 +108,7 @@ func TestGraceReadsDrainTimeoutFromTheEnvironment(t *testing.T) { {Name: drainTimeoutEnvVar, Value: "300"}, } want := int64(workerPreStopDrainSeconds + 300 + workerTeardownHeadroomSeconds) - if got := graceSecondsFor(nil, env); got != want { + if got := graceSecondsFor(drainSource{env: env}); got != want { t.Fatalf("env-set drain: grace = %d, want %d", got, want) } } @@ -120,7 +120,7 @@ func TestGraceFlagOverridesTheEnvironment(t *testing.T) { env := []corev1.EnvVar{{Name: drainTimeoutEnvVar, Value: "300"}} args := []string{"--drain-timeout", "60"} want := int64(workerPreStopDrainSeconds + 60 + workerTeardownHeadroomSeconds) - if got := graceSecondsFor(args, env); got != want { + if got := graceSecondsFor(drainSource{args: args, env: env}); got != want { t.Fatalf("flag with env set: grace = %d, want the flag's %d", got, want) } } @@ -134,12 +134,12 @@ func TestGraceIgnoresUnreadableEnvValues(t *testing.T) { ConfigMapKeyRef: &corev1.ConfigMapKeySelector{Key: "drain"}, }, }} - if got := graceSecondsFor(nil, from); got != workerTerminationGraceSeconds { + if got := graceSecondsFor(drainSource{env: from}); got != workerTerminationGraceSeconds { t.Fatalf("valueFrom: grace = %d, want the floor %d", got, workerTerminationGraceSeconds) } for _, v := range []string{"", "abc", "0", "-5"} { env := []corev1.EnvVar{{Name: drainTimeoutEnvVar, Value: v}} - if got := graceSecondsFor(nil, env); got != workerTerminationGraceSeconds { + if got := graceSecondsFor(drainSource{env: env}); got != workerTerminationGraceSeconds { t.Fatalf("env %q: grace = %d, want the floor %d", v, got, workerTerminationGraceSeconds) } } @@ -162,3 +162,40 @@ func TestGraceReadsDrainEnvFromTheContainerToo(t *testing.T) { t.Fatalf("grace = %d, want %d", *spec.TerminationGracePeriodSeconds, want) } } + +// On the extraPodSpec path the template is passed through verbatim, so +// ServiceSpec.Args is never rendered into the container -- a --drain-timeout +// sitting there is inert. It must not outrank the variable the container will +// actually read, or the budget is sized for a drain that never happens while +// the real one runs long and gets SIGKILLed partway through. Precedence is by +// source: what the process sees wins, and only within a source does a flag +// beat a variable. +func TestAnInertServiceSpecFlagDoesNotOutrankTheContainer(t *testing.T) { + spec := &corev1.PodSpec{Containers: []corev1.Container{{ + Name: "main", + Env: []corev1.EnvVar{{Name: drainTimeoutEnvVar, Value: "600"}}, + }}} + injectWorkerRolloutDefaults(spec, 0, 8080, false, []string{"--drain-timeout", "30"}, nil) + + want := int64(workerPreStopDrainSeconds + 600 + workerTeardownHeadroomSeconds) + if got := *spec.TerminationGracePeriodSeconds; got != want { + t.Fatalf("grace = %d, want %d -- the container drains for 600s, so %d "+ + "leaves the kubelet killing it partway through", got, want, got) + } +} + +// The same precedence, the other way round: a flag the container really runs +// with beats a variable from ServiceSpec. +func TestTheContainerFlagBeatsAServiceSpecVariable(t *testing.T) { + spec := &corev1.PodSpec{Containers: []corev1.Container{{ + Name: "main", + Args: []string{"--drain-timeout=300"}, + }}} + env := []corev1.EnvVar{{Name: drainTimeoutEnvVar, Value: "45"}} + injectWorkerRolloutDefaults(spec, 0, 8080, false, nil, env) + + want := int64(workerPreStopDrainSeconds + 300 + workerTeardownHeadroomSeconds) + if got := *spec.TerminationGracePeriodSeconds; got != want { + t.Fatalf("grace = %d, want %d", got, want) + } +} diff --git a/deploy/operator/internal/controller/inferadeployment_controller.go b/deploy/operator/internal/controller/inferadeployment_controller.go index 80878764..ec5200fa 100644 --- a/deploy/operator/internal/controller/inferadeployment_controller.go +++ b/deploy/operator/internal/controller/inferadeployment_controller.go @@ -136,7 +136,7 @@ func (r *InferaDeploymentReconciler) Reconcile(ctx context.Context, req ctrl.Req } else { idep.Status.GAIE = nil } - idep.Status.State = rollupState(status) + idep.Status.State = rollupState(status, idep.Spec.Services) if err := r.Status().Update(ctx, idep); err != nil { lg.Error(err, "status update failed") return ctrl.Result{RequeueAfter: 5 * time.Second}, nil @@ -164,13 +164,49 @@ func (r *InferaDeploymentReconciler) applyUnstructured(ctx context.Context, idep existing.SetNamespace(desired.GetNamespace()) _, err := controllerutil.CreateOrUpdate(ctx, r.Client, existing, func() error { spec, _, _ := unstructured.NestedMap(desired.Object, "spec") - _ = unstructured.SetNestedMap(existing.Object, spec, "spec") + current, _, _ := unstructured.NestedMap(existing.Object, "spec") + if current == nil { + current = map[string]any{} + } + _ = unstructured.SetNestedMap(existing.Object, mergeSpec(current, spec), "spec") existing.SetLabels(desired.GetLabels()) return controllerutil.SetControllerReference(idep, existing, r.Scheme) }) return err } +// mergeSpec overlays the fields the operator sets onto what is already there, +// leaving anything it does not mention alone. +// +// Replacing .spec wholesale would be simpler, but these objects are only +// partly ours: buildLeaderWorkerSet writes three fields and the API server +// defaults the other nine from the CRD. Overwriting the whole map strips those +// defaults on every pass, the API server restores them, and the next pass +// strips them again -- so CreateOrUpdate sees a difference every single time +// and issues a write. Harmless-but-wasteful once per resync; a write loop now +// that the reconciler watches LeaderWorkerSet, since each write enqueues the +// reconcile that produces the next one. +// +// Nested maps merge; anything else replaces. Lists are owned outright -- a +// container list merged element-wise would be neither what was asked for nor +// what was there. +func mergeSpec(into, from map[string]any) map[string]any { + for k, v := range from { + sub, isMap := v.(map[string]any) + if !isMap { + into[k] = v + continue + } + existing, ok := into[k].(map[string]any) + if !ok { + into[k] = v + continue + } + into[k] = mergeSpec(existing, sub) + } + return into +} + func (r *InferaDeploymentReconciler) deploymentStatus(ctx context.Context, idep *inferav1alpha1.InferaDeployment, name string, svc inferav1alpha1.ServiceSpec) inferav1alpha1.ServiceStatus { // Replicas is what the workload reports, not what the spec asked for. // Echoing the desired value back makes status useless for exactly the @@ -236,20 +272,39 @@ func copySpec(existing, desired client.Object) { } } -func rollupState(svcs map[string]inferav1alpha1.ServiceStatus) inferav1alpha1.DeploymentState { +// rollupState answers whether every service has the capacity the spec asked +// for, so it compares against the spec rather than against the status. +// +// ServiceStatus.Replicas is what the workload reports, which is the right +// thing for a reader watching a scale-up land but the wrong side of this +// comparison: `ReadyReplicas < Replicas` only sees Pods that exist and are +// not ready. A replica that was never created -- unschedulable, out of quota, +// no GPU -- is absent from both numbers, so they agree and the deployment +// calls itself ready on a fraction of its fleet. +func rollupState( + svcs map[string]inferav1alpha1.ServiceStatus, + specs map[string]inferav1alpha1.ServiceSpec, +) inferav1alpha1.DeploymentState { if len(svcs) == 0 { return inferav1alpha1.StatePending } - allReady := true - for _, s := range svcs { - if s.ReadyReplicas < s.Replicas || s.Replicas == 0 { - allReady = false + for name, s := range svcs { + spec, ok := specs[name] + if !ok { + // Reported but no longer in the spec: on its way out, and not a + // reason to hold the deployment back. + continue + } + want := replicasOf(spec) + if want == 0 { + // Deliberately scaled to zero; nothing to wait for. + continue + } + if s.ReadyReplicas < want { + return inferav1alpha1.StatePending } } - if allReady { - return inferav1alpha1.StateReady - } - return inferav1alpha1.StatePending + return inferav1alpha1.StateReady } func sortedKeys(m map[string]inferav1alpha1.ServiceSpec) []string { diff --git a/deploy/operator/internal/controller/scale_paths_test.go b/deploy/operator/internal/controller/scale_paths_test.go index c3a6555d..1d8a48d7 100644 --- a/deploy/operator/internal/controller/scale_paths_test.go +++ b/deploy/operator/internal/controller/scale_paths_test.go @@ -112,8 +112,9 @@ func TestEditingTheCRScales(t *testing.T) { // Editing the *child* is the path that does not survive, and that is the // intended behaviour of any operator: the child is derived state, so the next -// pass restores it from the CR. This is why an HPA has to be pointed at a -// scaling adapter and never at the generated Deployment. +// pass restores it from the CR. The write succeeds and nothing reports an +// error, which is why pointing an autoscaler at the generated Deployment looks +// like it works right up until the next reconcile. func TestEditingTheChildIsReverted(t *testing.T) { s := scaleScheme(t) idep := idepWith(2) diff --git a/deploy/operator/internal/controller/status_rollup_test.go b/deploy/operator/internal/controller/status_rollup_test.go new file mode 100644 index 00000000..826971c0 --- /dev/null +++ b/deploy/operator/internal/controller/status_rollup_test.go @@ -0,0 +1,89 @@ +/* +Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + +SPDX-License-Identifier: MIT +*/ + +package controller + +import ( + "testing" + + inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" +) + +// ServiceStatus.Replicas reports what the workload has, not what was asked +// for -- an autoscaler cannot tell a scale-up has not landed if `current` is +// the number it just requested. That makes it the wrong side of the readiness +// comparison: `ready < observed` only catches Pods that exist and are not +// ready, and says nothing about Pods that were never created at all. +// +// Which is the case that matters. A worker Pod that cannot be scheduled -- no +// GPU, quota exhausted, a node taint -- never reaches the ReplicaSet's +// status.replicas, so ready equals observed and the whole deployment reports +// itself ready on a fraction of its capacity. `.status.state` is a +// printcolumn and the natural readiness gate for anything orchestrating on +// top, so this is what decides whether traffic is sent. +func TestReadyNeedsTheReplicaCountThatWasAskedFor(t *testing.T) { + three := int32(3) + svcs := map[string]inferav1alpha1.ServiceSpec{ + "decode": {Replicas: &three}, + } + + cases := []struct { + name string + observed inferav1alpha1.ServiceStatus + want inferav1alpha1.DeploymentState + }{ + { + name: "every requested replica is up", + observed: inferav1alpha1.ServiceStatus{Replicas: 3, ReadyReplicas: 3}, + want: inferav1alpha1.StateReady, + }, + { + name: "a replica could not be scheduled, so it is not in the workload at all", + observed: inferav1alpha1.ServiceStatus{Replicas: 2, ReadyReplicas: 2}, + want: inferav1alpha1.StatePending, + }, + { + name: "all present, one still starting", + observed: inferav1alpha1.ServiceStatus{Replicas: 3, ReadyReplicas: 2}, + want: inferav1alpha1.StatePending, + }, + { + name: "nothing up yet", + observed: inferav1alpha1.ServiceStatus{Replicas: 0, ReadyReplicas: 0}, + want: inferav1alpha1.StatePending, + }, + } + + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + got := rollupState(map[string]inferav1alpha1.ServiceStatus{"decode": c.observed}, svcs) + if got != c.want { + t.Fatalf("state = %q, want %q (spec asked for %d, workload has %d/%d)", + got, c.want, three, c.observed.ReadyReplicas, c.observed.Replicas) + } + }) + } +} + +// A service the spec no longer mentions must not hold the deployment back, +// and one with no status yet must not read as satisfied. +func TestRollupHandlesServicesMissingFromEitherSide(t *testing.T) { + one := int32(1) + specs := map[string]inferav1alpha1.ServiceSpec{"decode": {Replicas: &one}} + + if got := rollupState(map[string]inferav1alpha1.ServiceStatus{}, specs); got != inferav1alpha1.StatePending { + t.Fatalf("no status reported yet: state = %q, want pending", got) + } + + // Status carries a service the spec dropped; the live one is satisfied. + svcs := map[string]inferav1alpha1.ServiceStatus{ + "decode": {Replicas: 1, ReadyReplicas: 1}, + "stale": {Replicas: 0, ReadyReplicas: 0}, + } + if got := rollupState(svcs, specs); got != inferav1alpha1.StateReady { + t.Fatalf("a service no longer in the spec blocked readiness: state = %q", got) + } +} From ea60d7c5d27e22e1a359dad8c6731e3178482700 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Wed, 5 Aug 2026 08:20:42 +0000 Subject: [PATCH 35/88] fix(operator): generate the RBAC the manager itself needs config/rbac/role.yaml is what external consumers -- Primus-SaFE among them -- build their RBAC from, and it was missing two grants the manager cannot start without: coordination.k8s.io/leases create,delete,get,list,patch,update,watch /events create,patch Both belong to the controller-runtime manager rather than the reconciler: it takes a lease whenever --leader-elect is set, which the chart sets by default, and it records events. Nothing in internal/ mentions either, so controller-gen had nothing to emit and the generated role has been short since the beginning. It went unnoticed because the deployed RBAC does not come from this file. The chart carries a hand-written ClusterRole that already listed both, so every install from the chart worked -- while anyone generating from config/rbac got a manager that fails to acquire its lease and exits. That failure reads like a broken operator, not like a missing RBAC rule. Two markers fix it at the source. Verified by expanding both roles into (apiGroup, resource, verb) triples and comparing the sets: 12 rules, 87 permissions, identical. Before this they were 10 and 78. A sorted line comparison would not have caught the original gap, or proved this closes it, since neither distinguishes a reordering from a verb moving between groups. Co-authored-by: Cursor Signed-off-by: leiwei12 --- deploy/operator/config/rbac/role.yaml | 19 +++++++++++++++++++ .../controller/inferadeployment_controller.go | 7 +++++++ 2 files changed, 26 insertions(+) diff --git a/deploy/operator/config/rbac/role.yaml b/deploy/operator/config/rbac/role.yaml index 4fcb9e46..160d6dcf 100644 --- a/deploy/operator/config/rbac/role.yaml +++ b/deploy/operator/config/rbac/role.yaml @@ -4,6 +4,13 @@ kind: ClusterRole metadata: name: infera-operator-role rules: +- apiGroups: + - "" + resources: + - events + verbs: + - create + - patch - apiGroups: - "" resources: @@ -39,6 +46,18 @@ rules: - patch - update - watch +- apiGroups: + - coordination.k8s.io + resources: + - leases + verbs: + - create + - delete + - get + - list + - patch + - update + - watch - apiGroups: - gateway.networking.k8s.io resources: diff --git a/deploy/operator/internal/controller/inferadeployment_controller.go b/deploy/operator/internal/controller/inferadeployment_controller.go index 07c032b3..167c3570 100644 --- a/deploy/operator/internal/controller/inferadeployment_controller.go +++ b/deploy/operator/internal/controller/inferadeployment_controller.go @@ -43,6 +43,13 @@ type InferaDeploymentReconciler struct { // (RBAC escalation-prevention: a grantor must hold what it grants) and so a // future operator path could read Pod status directly. // +kubebuilder:rbac:groups="",resources=pods,verbs=get;list;watch;patch +// The manager runtime needs the next two, not the reconciler: controller-runtime +// takes a lease when --leader-elect is set (the chart sets it by default) and +// records events. Nothing else here would emit them, so without these markers +// config/rbac/role.yaml describes a manager that cannot acquire its lease -- +// and it is the file external consumers build their RBAC from. +// +kubebuilder:rbac:groups=coordination.k8s.io,resources=leases,verbs=get;list;watch;create;update;patch;delete +// +kubebuilder:rbac:groups="",resources=events,verbs=create;patch func (r *InferaDeploymentReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) { lg := log.FromContext(ctx) From 66d52d710ce426ec3c84c210cf4430290bdbcf0f Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Sat, 8 Aug 2026 11:49:55 +0000 Subject: [PATCH 36/88] fix(router): bring the Rust breaker's probe slot in line with Python Same defect as the Python side, and two more the Rust paths add on their own. The half-open probe slot was claimed and never reliably released. `filter` claims one per candidate it lets through while the policy dispatches to one of them, so a recovering worker routinely had its slot taken by a request that went elsewhere and stayed half-open with nothing to release it. The claim is now bounded by a timeout rather than waiting for an outcome. That also covers the cancellation case, which has no Python equivalent: axum drops the handler future when a client disconnects, so neither record call runs -- and a probe to a worker wedged after accepting the connection, the exact condition this guards against, never returns at all. A 4xx reached `attempt`'s Err arm, which records nothing, leaking the slot deterministically rather than by chance. It now records a neutral outcome. The PD paths had the opposite bug: `else { record_success }` treated 400, 404 and 429 as evidence of health, resetting the failure count and closing an open breaker. A worker alternating 500s and 400s could never accumulate three consecutive failures, and a 429 from a worker the all-open fallback reached undid its backoff. Only 2xx counts as recovery now; open_decode's 4xx, which recorded neither outcome, is neutral. Entries are pruned against the live fleet on each discovery snapshot, beside the existing policy.sync_workers call. Worker ids are addresses no rebuilt Pod reuses, so every rollout stranded another entry and another /metrics series labelled by that id. cargo test 94 passed, fmt and clippy --all-targets -D warnings clean. The reclaim test was confirmed against the unbounded claim. Co-authored-by: Cursor Signed-off-by: leiwei12 --- rust/router/src/breaker.rs | 142 ++++++++++++++++++++++++++++++++--- rust/router/src/disagg.rs | 14 +++- rust/router/src/discovery.rs | 27 +++++-- rust/router/src/main.rs | 17 +++-- rust/router/src/proxy.rs | 5 ++ 5 files changed, 181 insertions(+), 24 deletions(-) diff --git a/rust/router/src/breaker.rs b/rust/router/src/breaker.rs index 92d49d17..185ea3a4 100644 --- a/rust/router/src/breaker.rs +++ b/rust/router/src/breaker.rs @@ -19,7 +19,7 @@ //! few integer writes; contention is not a concern at any plausible request //! rate, and a lock-free design here would buy nothing for the complexity. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Mutex; use std::time::{Duration, Instant}; @@ -60,15 +60,28 @@ struct Entry { opens_until: Instant, /// Cooldown applied on the *next* trip; doubles each time a probe fails. next_cooldown: Duration, - /// Set while a half-open probe is in flight, so only one is admitted. - probe_in_flight: bool, + /// When the outstanding half-open probe was admitted, so only one runs at + /// a time. `None` means the slot is free. + probe_started_at: Option, trips: u64, } +/// How long a claimed probe slot is honoured before it is reclaimed. +/// +/// Claiming and releasing are not paired: `filter` claims a slot for every +/// candidate it lets through, and the policy dispatches to exactly one of them, +/// so the rest are never told how they did. A 4xx records neither outcome +/// either, and a cancelled request -- the client hung up, or the worker never +/// answered, which is the very condition this guards against -- unwinds without +/// reaching any record call. Any of those would otherwise hold the slot +/// forever, leaving a recovered worker permanently out of rotation. +const PROBE_TIMEOUT: Duration = Duration::from_secs(60); + pub struct CircuitBreaker { failure_threshold: u32, cooldown: Duration, max_cooldown: Duration, + probe_timeout: Duration, entries: Mutex>, } @@ -78,6 +91,7 @@ impl CircuitBreaker { failure_threshold, cooldown, max_cooldown, + probe_timeout: PROBE_TIMEOUT, entries: Mutex::new(HashMap::new()), } } @@ -113,15 +127,24 @@ impl CircuitBreaker { return false; } e.state = BreakerState::HalfOpen; - e.probe_in_flight = false; + e.probe_started_at = None; tracing::info!(worker = worker_id, "breaker half-open, admitting one probe"); } BreakerState::HalfOpen => {} } - if e.probe_in_flight { - return false; + // One probe at a time, but only for as long as one could plausibly + // still be running: the claim expires rather than waiting for an + // outcome that may never arrive. See PROBE_TIMEOUT. + if let Some(started) = e.probe_started_at { + if now.duration_since(started) < self.probe_timeout { + return false; + } + tracing::info!( + worker = worker_id, + "breaker: probe slot unclaimed, admitting another" + ); } - e.probe_in_flight = true; + e.probe_started_at = Some(now); true } @@ -156,11 +179,40 @@ impl CircuitBreaker { } e.consecutive_failures = 0; e.state = BreakerState::Closed; - e.probe_in_flight = false; + e.probe_started_at = None; e.next_cooldown = self.cooldown; } } + /// Release the probe slot without scoring the worker either way. + /// + /// For an outcome that says nothing about worker health: a 4xx, which every + /// worker would answer identically, or a 429, which is backpressure the + /// policy already routes around. Counting either as recovery is as wrong as + /// counting it as failure -- it would reset the failure count and close an + /// open breaker, so a worker alternating 500s and 400s could never reach + /// the consecutive failures needed to trip, and one 429 from a worker the + /// all-open fallback reached would undo its backoff. The slot such a + /// request consumed still has to come back. + pub fn record_neutral(&self, worker_id: &str) { + let mut map = self.entries.lock().expect("breaker mutex poisoned"); + if let Some(e) = map.get_mut(worker_id) { + e.probe_started_at = None; + } + } + + /// Drop everything remembered about workers no longer in the fleet. + /// + /// Called with the full active set on each discovery snapshot, mirroring + /// `Policy::sync_workers`. Worker ids are addresses and a rebuilt Pod never + /// reuses one, so without this every rollout strands another entry -- and + /// another pair of Prometheus series, since /metrics exports one per entry + /// labelled by worker id. + pub fn retain_workers(&self, active: &HashSet) { + let mut map = self.entries.lock().expect("breaker mutex poisoned"); + map.retain(|id, _| active.contains(id)); + } + /// Record a pre-first-byte dispatch failure. Callers must gate this on /// [`is_worker_fault`] when the failure carries an HTTP status. pub fn record_failure(&self, worker_id: &str) { @@ -177,12 +229,12 @@ impl CircuitBreaker { state: BreakerState::Closed, opens_until: now, next_cooldown: self.cooldown, - probe_in_flight: false, + probe_started_at: None, trips: 0, }); e.consecutive_failures += 1; let was_probe = e.state == BreakerState::HalfOpen; - e.probe_in_flight = false; + e.probe_started_at = None; if was_probe { // A failed probe reopens immediately and backs off further, without @@ -506,4 +558,74 @@ mod tests { "exactly one of 16 racing threads may probe" ); } + + // filter() claims the probe slot for every candidate it lets through, but + // the policy dispatches to exactly one of them, so the others are never + // told how they did. Without a bound on the claim those workers sit in + // half-open holding a slot nothing will ever release -- healthy, and + // permanently unroutable until the process restarts. + #[test] + fn a_probe_slot_taken_but_never_dispatched_is_reclaimed() { + let b = CircuitBreaker::new(3, Duration::from_secs(5), Duration::from_secs(60)); + let t = Instant::now(); + for _ in 0..3 { + b.record_failure_at("w1", t); + } + let due = t + Duration::from_secs(6); + assert!(b.allows_at("w1", due), "cooldown elapsed -> a probe is due"); + + // The request went to another worker; nothing reports back for w1. + let later = due + PROBE_TIMEOUT + Duration::from_secs(1); + assert!( + b.allows_at("w1", later), + "an unused probe claim must not be permanent" + ); + } + + // 4xx says the request was bad, not the worker. Scoring it as recovery + // would reset the failure count and close an open breaker, so a worker + // alternating 500s and 400s could never reach three consecutive failures. + #[test] + fn a_neutral_outcome_frees_the_slot_without_scoring_it() { + let b = CircuitBreaker::new(3, Duration::from_secs(5), Duration::from_secs(60)); + let t = Instant::now(); + for _ in 0..3 { + b.record_failure_at("w1", t); + } + let due = t + Duration::from_secs(6); + assert!(b.allows_at("w1", due)); + + b.record_neutral("w1"); + assert_eq!( + b.state_of("w1"), + BreakerState::HalfOpen, + "a 4xx is not a recovery" + ); + assert!( + b.allows_at("w1", due), + "but the slot is free for a real probe" + ); + } + + // Worker ids are addresses and a rebuilt Pod never reuses one, so entries + // for departed workers accumulate for the process lifetime -- each also + // pinning a Prometheus series, which is the part that actually hurts. + #[test] + fn workers_gone_from_the_fleet_are_forgotten() { + let b = CircuitBreaker::new(3, Duration::from_secs(5), Duration::from_secs(60)); + let t = Instant::now(); + for _ in 0..3 { + b.record_failure_at("gone", t); + b.record_failure_at("stay", t); + } + assert_eq!(b.snapshot().len(), 2); + + let active: HashSet = ["stay".to_string()].into_iter().collect(); + b.retain_workers(&active); + + let snap = b.snapshot(); + assert_eq!(snap.len(), 1); + assert!(snap.iter().any(|(id, _, _)| id == "stay")); + assert_eq!(b.state_of("gone"), BreakerState::Closed); + } } diff --git a/rust/router/src/disagg.rs b/rust/router/src/disagg.rs index 7f141c4b..d2371532 100644 --- a/rust/router/src/disagg.rs +++ b/rust/router/src/disagg.rs @@ -183,8 +183,10 @@ async fn unary_dual( } if is_worker_fault(st.as_u16()) { state.breaker.record_failure(&p.worker.worker_id); - } else { + } else if st.is_success() { state.breaker.record_success(&p.worker.worker_id); + } else { + state.breaker.record_neutral(&p.worker.worker_id); } } Err(e) => { @@ -198,8 +200,10 @@ async fn unary_dual( let st = resp.status(); if is_worker_fault(st.as_u16()) { state.breaker.record_failure(&d.worker.worker_id); - } else { + } else if st.is_success() { state.breaker.record_success(&d.worker.worker_id); + } else { + state.breaker.record_neutral(&d.worker.worker_id); } let ct = content_type(&resp); match resp.bytes().await { @@ -254,8 +258,10 @@ fn spawn_prefill_drain( // KVPoll, which is exactly the failure worth remembering. if is_worker_fault(st.as_u16()) { breaker.record_failure(&worker_id); - } else { + } else if st.is_success() { breaker.record_success(&worker_id); + } else { + breaker.record_neutral(&worker_id); } } Err(e) => { @@ -282,6 +288,8 @@ async fn open_decode( if st.is_client_error() || st.is_server_error() { if is_worker_fault(st.as_u16()) { state.breaker.record_failure(&d.worker.worker_id); + } else { + state.breaker.record_neutral(&d.worker.worker_id); } let txt = resp.text().await.unwrap_or_default(); return Err(format!( diff --git a/rust/router/src/discovery.rs b/rust/router/src/discovery.rs index fa357568..6852e55c 100644 --- a/rust/router/src/discovery.rs +++ b/rust/router/src/discovery.rs @@ -17,10 +17,17 @@ use futures::StreamExt; use serde::Serialize; use serde_json::Value; +use crate::breaker::CircuitBreaker; use crate::policy::Policy; use crate::pool::{SharedPool, Snapshot, Worker}; -pub async fn run(base: String, prefix: String, pool: SharedPool, policy: Arc) { +pub async fn run( + base: String, + prefix: String, + pool: SharedPool, + policy: Arc, + breaker: Arc, +) { let prefix = if prefix.ends_with('/') { prefix } else { @@ -28,7 +35,7 @@ pub async fn run(base: String, prefix: String, pool: SharedPool, policy: Arc backoff = 1, Err(e) => { tracing::warn!("etcd discovery error: {e}; retry in {backoff}s"); @@ -64,6 +71,7 @@ async fn discover_once( prefix: &str, pool: &SharedPool, policy: &Arc, + breaker: &Arc, ) -> anyhow::Result<()> { let client = reqwest::Client::builder().build()?; let re = range_end(prefix); @@ -95,7 +103,7 @@ async fn discover_once( workers.len(), prefix ); - publish(pool, policy, &workers); + publish(pool, policy, breaker, &workers); // 2. watch for changes (long-lived NDJSON stream) let create = serde_json::json!({ @@ -121,7 +129,7 @@ async fn discover_once( } if let Ok(msg) = serde_json::from_slice::(line) { if apply_watch(prefix, &msg, &mut workers) { - publish(pool, policy, &workers); + publish(pool, policy, breaker, &workers); } } } @@ -129,11 +137,20 @@ async fn discover_once( Ok(()) } -fn publish(pool: &SharedPool, policy: &Arc, workers: &HashMap>) { +fn publish( + pool: &SharedPool, + policy: &Arc, + breaker: &Arc, + workers: &HashMap>, +) { let all: Vec> = workers.values().cloned().collect(); // Let cost-aware policies reconcile per-worker state (kv-event subscriptions, // load bookkeeping) against the new fleet before we swap the snapshot in. policy.sync_workers(&all); + // Same reconcile for the breaker: a worker id is an address no rebuilt Pod + // reuses, so entries for departed workers would otherwise accumulate for + // the process lifetime, each pinning a Prometheus series. + breaker.retain_workers(&workers.keys().cloned().collect()); pool.store(Arc::new(Snapshot::build(all))); } diff --git a/rust/router/src/main.rs b/rust/router/src/main.rs index 0ede6af1..35c123dd 100644 --- a/rust/router/src/main.rs +++ b/rust/router/src/main.rs @@ -56,12 +56,21 @@ async fn main() -> anyhow::Result<()> { // without locking, so reads scale across cores. let pool = Arc::new(ArcSwap::from_pointee(Snapshot::empty())); + // Built before discovery starts: the reconcile loop prunes its entries + // against the live fleet, the same way it reconciles the policy's. + let breaker = Arc::new(breaker::CircuitBreaker::new( + cfg.breaker_failure_threshold, + Duration::from_secs_f64(cfg.breaker_cooldown_s), + Duration::from_secs_f64(cfg.breaker_max_cooldown_s), + )); + { let pool = pool.clone(); let policy = policy.clone(); + let breaker = breaker.clone(); let base = cfg.etcd_base(); let prefix = cfg.etcd_prefix.clone(); - tokio::spawn(async move { discovery::run(base, prefix, pool, policy).await }); + tokio::spawn(async move { discovery::run(base, prefix, pool, policy, breaker).await }); } let state = AppState { @@ -70,11 +79,7 @@ async fn main() -> anyhow::Result<()> { http: proxy::build_upstream_client()?, started: Instant::now(), retries: cfg.request_max_retries, - breaker: Arc::new(breaker::CircuitBreaker::new( - cfg.breaker_failure_threshold, - Duration::from_secs_f64(cfg.breaker_cooldown_s), - Duration::from_secs_f64(cfg.breaker_max_cooldown_s), - )), + breaker, }; let addr = format!("{}:{}", cfg.host, cfg.port); diff --git a/rust/router/src/proxy.rs b/rust/router/src/proxy.rs index 05a67675..29924349 100644 --- a/rust/router/src/proxy.rs +++ b/rust/router/src/proxy.rs @@ -134,6 +134,11 @@ async fn mixed_dispatch( // is_worker_fault(). if is_worker_fault(err_resp.status().as_u16()) { state.breaker.record_failure(&wid); + } else { + // A 4xx is not held against the worker, but the probe slot + // it consumed has to come back or one bad client wedges a + // recovering worker out of rotation. + state.breaker.record_neutral(&wid); } last_err = Some(err_resp); } From 6f42a2db8d77a8a6391b2eb755cff8bab4029e10 Mon Sep 17 00:00:00 2001 From: xiaobochen-amd Date: Thu, 6 Aug 2026 15:27:32 +0000 Subject: [PATCH 37/88] fix(tests): a tier that cannot run must not report PASS _dispatch_slurm and run_e2e_disagg returned 0 when srun was absent, or when the host python3 could not import pytest/pytest-asyncio/httpx. The tier then propagated rc=0 all the way to RESULT: PASS having executed nothing. This is not hypothetical. On 2026-08-06 a runner fleet whose python3 lacked pytest turned all three e2e-disag legs green in 7 seconds each -- a required check passing without running a single test. The run was only caught because the mixed and engine tiers failed hard for an unrelated reason and took the whole run red with them. Both guards now fail and name the cause. The dependency message prints the python3 that was actually consulted plus its real ImportError: the old "missing host deps" line is true of every python3 on the box and does not say which one was asked, which is precisely what sent that triage the wrong way. The skip behaviour still exists for a dev box that genuinely has no SLURM, but it has to be asked for: INFERA_E2E_ALLOW_SKIP=1. Even then it is not silent -- skipped tiers are recorded and the last line reads "RESULT: PASS (SKIPPED: e2e disagg)" instead of a clean pass. Verified on five paths: no srun (disag), no srun (engine, the _dispatch_slurm guard), missing deps, missing deps with the opt-out, and a healthy host where neither guard fires and dispatch still reaches "mode=resv". No workflow or script parses the RESULT string, so changing the last line is safe. Co-authored-by: Cursor Signed-off-by: xiaobochen-amd --- tests/run_tests.sh | 49 +++++++++++++++++++++++++++++++++++++++------- 1 file changed, 42 insertions(+), 7 deletions(-) diff --git a/tests/run_tests.sh b/tests/run_tests.sh index 52da8d3d..50e118bf 100755 --- a/tests/run_tests.sh +++ b/tests/run_tests.sh @@ -195,6 +195,24 @@ QOS_WAIT="${INFERA_E2E_QOS_WAIT:-30}" # than a single-node srun, and giving up early only churns the pair-hold race. HOLD_WAIT="${INFERA_E2E_HOLD_WAIT:-60}" +# A tier that could not run is not a tier that passed: returning 0 here is how a +# runner whose python3 lacked pytest turned every e2e-disag leg green in 7s. Fail +# and name the cause; a dev box that really has no SLURM opts out explicitly. +# $1=label $2=what is wrong $3=how to fix it +_SKIPPED_TIERS="" +_skip_or_fail() { + local label="$1" why="$2" fix="$3" + if [ "${INFERA_E2E_ALLOW_SKIP:-}" = 1 ]; then + _SKIPPED_TIERS="${_SKIPPED_TIERS:+$_SKIPPED_TIERS, }$label" + echo "[$label] SKIPPED (INFERA_E2E_ALLOW_SKIP=1): $why" >&2 + return 0 + fi + echo "[$label] FATAL: $why" >&2 + echo "[$label] fix: $fix" >&2 + echo "[$label] (or INFERA_E2E_ALLOW_SKIP=1 to skip this tier instead of failing)" >&2 + return 1 +} + _have_slurm() { command -v srun >/dev/null 2>&1; } # The nodes reservation $1 covers, one per line ('' if it is gone/expired). # Spur ignores the NAME arg and dumps all reservations; match the exact block. @@ -358,8 +376,10 @@ _watch_job() { _dispatch_slurm() { local label="$1"; shift if ! _have_slurm; then - echo "[$label] WARNING: no SLURM (srun) — skipping" >&2 - return 0 + _skip_or_fail "$label" \ + "no SLURM: srun is not on PATH, so this tier cannot be dispatched to a GPU node" \ + "expose the SLURM client on this host, or run where docker + >=8 AMD GPUs are present" + return $? fi # srun's own client banners/errors (job id, "running on ", ...). local out="$SCRATCH/.dispatch-$label.out" @@ -603,11 +623,20 @@ run_e2e_disagg() { local engines=("$@") echo "===== e2e PD-disaggregated (cross-node, 2 nodes): ${engines[*]} =====" if ! _have_slurm; then - echo "[e2e disagg] WARNING: no SLURM (srun) — skipping PD-disaggregated tests" >&2 - return 0 + _skip_or_fail "e2e disagg" \ + "no SLURM: srun is not on PATH, so the PD-disaggregated tests cannot run" \ + "expose the SLURM client on this host" + return $? + fi + # Name the interpreter actually consulted and quote its ImportError: "missing + # host deps" is true of every python3 on the box, and sent the last triage wrong. + local deps_err + if ! deps_err=$(python3 -c "import pytest, pytest_asyncio, httpx" 2>&1); then + _skip_or_fail "e2e disagg" \ + "the disagg orchestrator runs pytest on THIS host, and $(command -v python3 || echo 'python3 (not on PATH)') cannot import its deps: ${deps_err##*$'\n'}" \ + "pip install pytest pytest-asyncio httpx" + return $? fi - python3 -c "import pytest, pytest_asyncio, httpx" >/dev/null 2>&1 \ - || { echo "[e2e disagg] WARNING: missing host deps (pytest/pytest-asyncio/httpx) — skipping" >&2; return 0; } if [ -n "$SHARED_LOG_DIR" ]; then exec > >(stdbuf -oL tee -a "$SHARED_LOG_DIR/dispatch-disag-$$.log") 2>&1 @@ -777,5 +806,11 @@ if [ -d "$E2E_LOG_DIR" ]; then ls -1 "$E2E_LOG_DIR"/*.log 2>/dev/null | sed 's|^| |' || true fi -[ "$rc" -eq 0 ] && echo "RESULT: PASS" || echo "RESULT: FAIL" +if [ "$rc" -ne 0 ]; then + echo "RESULT: FAIL" +elif [ -n "$_SKIPPED_TIERS" ]; then + echo "RESULT: PASS (SKIPPED: $_SKIPPED_TIERS)" +else + echo "RESULT: PASS" +fi exit "$rc" From 11397e1338ddef2a1ee8d4ac75199fdfe8420995 Mon Sep 17 00:00:00 2001 From: liyingli Date: Thu, 6 Aug 2026 08:52:19 +0000 Subject: [PATCH 38/88] docs(recipes): add the GLM-5.2-FP8 gfx942 PD+kvd Kubernetes recipe A `docker` + shell deployment of GLM-5.2-FP8 on 2 x MI300X was brought up and benchmarked, but nothing carried it into Kubernetes. Anyone repeating it had to re-derive the manifest from four launch scripts, and the parts that are easiest to get wrong there are exactly the parts that fail quietly -- a KV transfer that falls back to TCP, an L3 tier that is written and never read, a hicache page silently rejected for being larger than its tablespace slot. This is that deployment as one InferaDeployment: router + prefill (with a kvd sidecar) + decode, TP8/DP8 with DP-attention, MTP, fp8 KV, Mooncake RDMA between the legs, and KV offload to host RAM and node-local NVMe. Every engine, router and kvd flag is the docker recipe's, unchanged. The README's section 6 lists each place the substrate forced something different and why, so a reader can check the claim rather than trust it. Three of those are worth calling out because they are not mechanical: kvd runs as a NATIVE sidecar (initContainers + restartPolicy: Always, hence k8s 1.29+). The engine probes the kvd socket once with a 5 s timeout and refuses to start if nothing answers -- it does not retry -- so "kvd first" has to be a scheduling guarantee, not a convention. An ordinary container makes it a race. The kvd sidecar gets 136Gi, not the 64G its --max-bytes suggests. kvd holds two independent budgets: --max-bytes caps the inline store, while the shared arena is sized separately, defaults to the same value, and is mmap'd and mlocked whole at startup. Sizing the limit to one of them gets the sidecar OOM-killed mid-run, which reads as a kvd bug. The RDMA rail and GID index are placeholders, not the validated cluster's mlx5_0 and 3. leg.sh takes both as required variables with no fallback, and the GID index is documented as per-node -- so there are two of those, one per leg. Hard-coding them would have been less faithful to the recipe, not more. This recipe also does not use the overlay, which every other one does. GLM-5.2 on the v0.5.16 gfx942 base needs a rebuilt Mooncake engine.so and four SGLang source patches; the payload carries patches/vllm/ only and its patch loop is gated on `import vllm`, so on an SGLang base it does not run at all. Section 1 covers what breaks without each patch and how to read the markers out of a built image. Not yet run in its Kubernetes form. The validation table says so per row rather than in a footnote, and the two docker-side numbers still being re-measured are marked TODO instead of quoted. Signed-off-by: liyingli Co-authored-by: Cursor --- examples/recipes/README.md | 19 +- examples/recipes/glm5.2-fp8-gfx942/README.md | 284 +++++++++++++++++ .../disaggregated-kvd/deploy.yaml | 298 ++++++++++++++++++ 3 files changed, 597 insertions(+), 4 deletions(-) create mode 100644 examples/recipes/glm5.2-fp8-gfx942/README.md create mode 100644 examples/recipes/glm5.2-fp8-gfx942/disaggregated-kvd/deploy.yaml diff --git a/examples/recipes/README.md b/examples/recipes/README.md index 474f6e97..6d92900d 100644 --- a/examples/recipes/README.md +++ b/examples/recipes/README.md @@ -6,6 +6,7 @@ serve them. Pick a model, pick a combo, `kubectl apply`. | Model | Engine | Recipe | |---|---|---| | GLM-5.2-MXFP4 | SGLang | [`glm5.2/`](glm5.2/README.md) | +| GLM-5.2-FP8 (gfx942) | SGLang | [`glm5.2-fp8-gfx942/`](glm5.2-fp8-gfx942/README.md) | | Kimi-K3 | vLLM | [`kimi-k3/`](kimi-k3/README.md) | @@ -21,8 +22,10 @@ as the wire format, and the directory name as what it means. ## The four combos -Every recipe comes in the same four shapes. They compose two independent choices: -**how requests are split across GPUs**, and **whether KV survives past the GPU**. +Recipes come in the same four shapes, composing two independent choices: **how +requests are split across GPUs**, and **whether KV survives past the GPU**. A +recipe pinned to one validated configuration ships only the shapes it was run in — +`glm5.2-fp8-gfx942` is `disaggregated-kvd` alone. | Combo | Serving | KV cache | Use it when | |---|---|---|---| @@ -36,8 +39,8 @@ RDMA, and there is no TCP fallback. If you are on one box, use `aggregated`. ## How these manifests are built -Every recipe is **stock vendor image + overlay + (optional) sidecar**. The vendor -image is never forked: +Recipes are **stock vendor image + overlay + (optional) sidecar**. The vendor image +is never forked: ``` initContainer infera-overlay busybox carrying /payload ──cp──▶ emptyDir @@ -50,6 +53,14 @@ ROCm major, then execs the engine. So following an upstream vLLM or SGLang relea is an image-tag edit here — no rebuild of ours, and no repeat of the incident where forking the base for one model broke every other model. +**`glm5.2-fp8-gfx942` is the exception**, and its README §1 says why: GLM-5.2 on the +v0.5.16 gfx942 base needs a rebuilt Mooncake `engine.so` plus four SGLang source +patches, while the payload carries `deploy/docker/patches/vllm/` only and +`infera-exec` runs that loop just when `vllm` imports. A payload-mounted stock base +would come up green there and then die on the first cross-node KV transfer, corrupt +long prompts, or kill the prefill scheduler. Where the overlay can carry what a +deployment needs, it does; that recipe is what it looks like when it cannot. + Build the overlay before deploying: ```bash diff --git a/examples/recipes/glm5.2-fp8-gfx942/README.md b/examples/recipes/glm5.2-fp8-gfx942/README.md new file mode 100644 index 00000000..6041ed06 --- /dev/null +++ b/examples/recipes/glm5.2-fp8-gfx942/README.md @@ -0,0 +1,284 @@ +# GLM-5.2-FP8 on gfx942 (Kubernetes) + +Serve GLM-5.2-FP8 across two gfx942 nodes (MI300X / MI325X): SGLang prefill/decode +disaggregation over Mooncake RDMA, DP-attention, MTP speculative decoding, kv-aware +routing, and KV offload to host RAM + node-local NVMe through `infera-kvd`. + +Every flag here is lifted from a `docker` + shell deployment of the same topology, +validated on 2 × MI300X and referred to throughout as **the docker recipe**. Nothing +was retuned for Kubernetes: §6 lists every difference and why the substrate forced +it, and there are no others. + +| Combo | Serving | KV cache | Manifest | +|---|---|---|---| +| **disaggregated + kvd** | prefill and decode on separate nodes | + kvd L2 pinned host RAM, L3 on node NVMe | [`disaggregated-kvd/deploy.yaml`](disaggregated-kvd/deploy.yaml) | + +The docker recipe's `KVD=0` A/B baseline is the same deployment minus the offload +tier — §5 has the four-line edit rather than a second copy of the manifest. + +## 1. This recipe does not use the overlay + +Every other recipe here runs a **stock vendor image** with the overlay payload +mounted in. This one runs an **infera-built engine image** instead, because GLM-5.2 +on the v0.5.16 gfx942 base needs a rebuilt native library and four source patches, +none of which a mounted payload can supply. The overlay payload carries +`deploy/docker/patches/vllm/` only, and `infera-exec`'s patch loop is additionally +gated on `import vllm` — so on an SGLang base it does not merely miss them, it does +not run at all. + +**The Mooncake rebuild is the part no payload could ever carry.** The base bundles +Mooncake at upstream #2682, which installs a HIP IPC transport unconditionally and +prefers it over RDMA — so cross-node PD dies on the first request inside +`hipIpcOpenMemHandle`, which cannot open a peer node's handle. +`Dockerfile.sglang.gfx942` rebuilds `engine.so` in place with the transport gated; +the build step is self-verifying and fails if the gate did not compile in. + +**Four SGLang source patches**, all under `deploy/docker/patches/`: + +| Fix | Patch | What happens without it | +|---|---|---| +| DSA indexer row count | `sglang_dsa/patch_dsa_indexer_hip_dp_padded_rows.py` | top-k dies with `Expected lengths.size(0) == B` as soon as concurrency > 1 | +| mooncake early-send KV wait event | `sglang_disagg/patch_mooncake_early_send_wait_event.py` | every prefill chunk but the last is RDMA-read while the forward is still writing it; prompts longer than one chunk come back **partially wrong, with nothing in any log** | +| hicache staged write-back gate | `sglang_rocm/patch_hicache_rocm_staged_write_back.py` | the prefill scheduler dies (exit −3, `Tensor match failed … device=rocm:0`) on the first request that reuses a prefix | +| hicache host allocator | `sglang_rocm/patch_hicache_rocm_host_alloc.py` | preventive on gfx942 rather than a fix for a crash seen here: hicache hands host `data_ptr()`s to GPU kernels, which is only correct while `hipHostRegister` maps the pages at the host VA. It did in every measurement on this base; the patch moves to `hipHostMalloc`, where that identity is an API guarantee instead of an accident of the driver. On gfx950 the two differ and the first kvd write-back aborts with `Memory access fault by GPU node-N` | + +Build the image: + +```bash +docker build -f deploy/docker/Dockerfile.sglang.gfx942 -t infera:sglang-gfx942-glm52 . +``` + +Load it on **both** nodes (or push it to a registry the cluster can pull; the +manifest uses `imagePullPolicy: IfNotPresent`, so a locally-loaded image is used +as-is). Three of the four patches leave greppable markers, which is the cheap way to +tell a correctly built image from one where a patch silently no-op'd against a moved +anchor: + +```bash +docker run --rm --entrypoint bash infera:sglang-gfx942-glm52 -c ' +P=$(python3 -c "import sglang, os; print(os.path.dirname(sglang.__file__))") +for m in GLM52_ROCM_HOST_ALLOC GLM52_ROCM_STAGED_WRITE_BACK GLM52_P1V3; do + grep -rql "$m" "$P" && echo "ok $m" || echo "MISSING $m" +done' +``` + +The mooncake wait-event patch leaves no named marker — its replacement text is its +own idempotency check — so it, and the Mooncake rebuild, are the two to confirm from +the build log. + +## 2. Prerequisites + +**Hardware.** Two nodes, 8× gfx942 each, on a mutually routable RoCE fabric — the +KV handoff is RDMA with no TCP fallback. The prefill node carries three Pods and so +needs ~**670 GiB** of free host RAM: 512 for the engine (256 GB of hicache host tier +plus load buffers), 136 for kvd (§6), 16 for the router. Those are requests; the +node's `/dev/shm` is on top of them and is charged to no limit. The prefill node +also needs a **node-local NVMe** directory for L3. + +Check the fabric from inside the image on **both** nodes before deploying, because +zero visible ports is how RDMA fails here — `ibv_get_device_list()` returns nothing, +Mooncake falls back to TCP, and the deployment still comes up: + +```bash +docker run --rm --network host --device=/dev/infiniband --cap-add=IPC_LOCK \ + --entrypoint bash infera:sglang-gfx942-glm52 -c 'ibv_devinfo | grep -c PORT_ACTIVE' +``` + +[`sglang_1p1d_glm5.2/preflight_rdma.sh`](../../sglang_1p1d_glm5.2/preflight_rdma.sh) +wraps this plus a cross-node bandwidth and mooncake KV probe. It takes the image as +`IMAGE=` and tests the hosts, so it applies here unchanged even though it ships with +the MI355X example. + +That preflight report is also where the manifest's RDMA placeholders come from: + +- `` — the rail(s) Mooncake may use, comma-separated, from + `ibv_devices`. A rail that is physically down must **not** be listed. +- `` / `` — from `show_gids `, the index + whose type is `RoCE v2`. There are two because the index is **per node, not per + cluster**; two identical machines routinely expose different ones. They are + usually equal, but check both — the wrong index pins KV to an interface that + never carries it and the transfer simply times out. + +The docker recipe's cluster answered `mlx5_0` and `3`, which is what its validation +ran on. Those are its numbers, not defaults: `leg.sh` takes both as required +variables with no fallback, so a manifest that hard-coded them would be less +faithful to the recipe, not more. + +**Cluster.** Kubernetes **1.29+** — the kvd daemon is a native sidecar +(`initContainers` with `restartPolicy: Always`), which is what makes it reach a +healthy startupProbe *before* the engine starts. That ordering is load-bearing: the +engine probes the kvd socket once with a 5 s timeout and refuses to start if +nothing answers, without retrying. On an older cluster the sidecar has to be an +ordinary container and the ordering becomes a race the engine loses about as often +as it wins. + +```bash +# AMD GPU device plugin — nodes must advertise amd.com/gpu +kubectl get nodes -o custom-columns=NODE:.metadata.name,GPU:.status.allocatable.'amd\.com/gpu' + +# the infera operator (provides the InferaDeployment CRD) +helm install infera-operator deploy/operator/helm/infera-operator -n infera-system --create-namespace +kubectl -n infera-system rollout status deploy/infera-operator +``` + +**`memlock`.** The docker recipe passed `--ulimit memlock=-1` and Kubernetes has no per-Pod +equivalent, but the capability is what actually carries this: `CAP_IPC_LOCK` — granted +to the kvd sidecar explicitly and to the engines through `privileged` — makes the +kernel skip `RLIMIT_MEMLOCK` accounting for both `mlock(2)` and RDMA registration. If +you do hit the limit the failure is soft and quiet: kvd logs `mlock failed (errno=…)` +at INFO and runs its arena **unpinned** (still correct, slower DMA staging). Raising +the limit itself needs `default_ulimits` in the container runtime's config; there is +no Pod-spec field for it. + +**Weights.** `hostPath`, not a PVC, at the **same path on both nodes**. If the path +is a HuggingFace cache symlink, mount the directory the links resolve into as well +— otherwise the inner relative links dangle and `transformers` rejects the model +with `Should have a model_type key in its config.json`, four minutes into startup +and far from its cause. + +## 3. Deploy + +```bash +kubectl create namespace infera --dry-run=client -o yaml | kubectl apply -f - + +sed -e "s||node-a|" -e "s||node-b|" \ + -e "s||/mnt/models|" -e "s||/mnt/nvme/kvd-l3|" \ + -e "s||mlx5_0|" \ + -e "s||3|" -e "s||3|" \ + examples/recipes/glm5.2-fp8-gfx942/disaggregated-kvd/deploy.yaml | kubectl apply -f - +``` + +The RDMA values above are the docker recipe's; substitute what §2 reported on your +own fabric. Any placeholder left unsubstituted fails loudly, which is the point — +`kubectl` rejects `` outright, and a literal `` is +not a device `ibv_open_device` will accept. + +Cold start is 15–25 min: both legs load GLM-5.2 plus the MTP nextn layer, and the +log goes quiet while they do. That is why the workers use a `startupProbe` with a +90-minute budget (matching `INFERA_ENGINE_READY_TIMEOUT`) and no readiness probe. +Don't kill a slow load. + +```bash +kubectl -n infera get pods -w +kubectl -n infera logs -f -c main \ + -l infera.amd.com/deployment=glm52-fp8-pd-kvd,infera.amd.com/service=prefill +``` + +## 4. Smoke test + +```bash +kubectl -n infera port-forward svc/glm52-fp8-pd-kvd-server 8000:8000 & + +curl -s localhost:8000/v1/workers | jq # expect one prefill + one decode +curl -s localhost:8000/v1/chat/completions -H 'Content-Type: application/json' \ + -d '{"model":"/models/GLM-5.2-FP8", + "messages":[{"role":"user","content":"What is 127 * 31? Answer with the number only."}], + "max_tokens":128,"temperature":0, + "chat_template_kwargs":{"enable_thinking":false}}' | jq -r '.choices[0].message.content' +``` + +`3937`, plus `installTransport, type=rdma` in the decode leg's log, means the router +paired the legs and KV moves over RDMA rather than silently falling back to TCP: + +```bash +kubectl -n infera logs -c main \ + -l infera.amd.com/deployment=glm52-fp8-pd-kvd,infera.amd.com/service=decode \ + | grep -aE 'GID index|installTransport' +``` + +The manifest passes no `--served-model-name`, matching the docker recipe, so the +served name **is** the model path — `/models/GLM-5.2-FP8` in the request above. + +Benchmarking needs nothing further: the router is an ordinary OpenAI-compatible +endpoint and does not care how the engines were started, so point a load generator +at this same forwarded port. For a KV-reuse benchmark specifically, read §5 first — +prompts shorter than a hicache page never reach kvd at all. + +## 5. kvd + +```bash +POD=$(kubectl -n infera get pod -o name \ + -l infera.amd.com/deployment=glm52-fp8-pd-kvd,infera.amd.com/service=prefill | head -1) +kubectl -n infera exec $POD -c kvd -- \ + python3 -m infera.kvd.statctl --socket /tmp/infera-kvd/kvd.sock +``` + +`sets_total` climbing means the engine writes to kvd; `gets_total` / `hits_total` +climbing means it reads back. Writes alone prove only half the path. + +Send a prompt long enough to fill several hicache pages before reading these — one +shorter than a page produces no kvd traffic at all and proves nothing. + +Two counters that will mislead you: + +- **`misses_total` counts failed gets only.** SGLang never gets what `batch_exists` + did not confirm, and a prefetch abandoned before its query reaches the backend + leaves *every* counter untouched. `0 misses` is compatible with L3 having served + nothing. Read the scorer's `cached tokens by tier` instead. +- **`entries: 0` with a healthy-looking deployment** means kvd rejected the KV + layout. A value larger than the biggest tablespace pool is rejected, not split, + and `--tablespace-pools 1M,4M` is sized for GLM-5.2-FP8's 2.74 MiB KV pages and + 624 KiB indexer pages at `page_size 64`. Grep the kvd container for + `value_exceeds_largest_pool` (a single-pool tablespace says + `value_exceeds_slot_bytes` instead). + +**The `KVD=0` A/B baseline** is this manifest minus the offload tier: delete the +`kvd` entry under `prefill.extraPodSpec.initContainers`, drop +`--infera-kvd-socket /tmp/infera-kvd/kvd.sock --hicache-size 32` from the prefill +command, drop the `kvd-sock` / `kvd-l3` volumes, and drop the `kvd-sock` +`volumeMount` from the prefill `main` container — miss that last one and the Pod +is rejected for referencing a volume that no longer exists. Worth running only once the +deployment is above its pressure point — below it the 54 GB device pool per rank +answers everything and both arms are identical. That is what the docker recipe's +agentic trace showed; **TODO — kvd's share of the hits is being re-measured, so no +figure is quoted here yet.** + +## 6. What changed from the docker recipe, and why + +Every engine, router and kvd flag is identical. These are substrate translations: + +| `docker` form | Kubernetes form | Why | +|---|---|---| +| etcd container + `--etcd-endpoint` | `discoveryBackend: kubernetes` | the operator's own backend: workers self-register on their Pod annotation, so there is no etcd to run. Both publish the same `WorkerInfo` | +| `--advertise-host $PREFILL_IP` | `--advertise-host $(POD_IP)` | downward API. With `hostNetwork` this is the node IP, which is what the peer dials for the Mooncake bootstrap handshake. **Override it if your RoCE rail is on a different address than the node's primary IP** | +| `--network host` | `hostNetwork: true` | the RoCE rails are host interfaces; the pod network cannot reach them | +| `--ipc host --shm-size 128g` | `hostIPC: true` | `--ipc host` already makes `/dev/shm` the host's and docker ignores `--shm-size` alongside it, so this is the whole translation. The node's `/dev/shm` must be large enough for TP8 | +| `--device=/dev/infiniband` | `privileged: true` + `hostPath /dev/infiniband` | no unprivileged equivalent without an RDMA device plugin. Same pattern as the validated [`pd-1p1d-mooncake.yaml`](../../k8s-deployments/pd-1p1d-mooncake.yaml) | +| `--device=/dev/kfd --device=/dev/dri`, `HIP_VISIBLE_DEVICES=0..7` | `amd.com/gpu: 8` | the device plugin owns these. It exposes the allocated GPUs renumbered from 0, so the visible-device variables are dropped rather than translated — pinning them would silently mask GPUs on any node where the allocation is not 0–7 | +| `bash launch_kvd.sh` before `launch_prefill.sh` | native sidecar + `startupProbe` | the shell ordering becomes a scheduling guarantee instead of a convention | +| `-v $KVD_L3_DIR:$KVD_L3_DIR` | `hostPath` at `/kvd-l3` | still node-local NVMe. A PVC would work only with a node-local StorageClass, and anything shared classifies as buffered — 3.70 GB/s against 14.56 GB/s with `O_DIRECT`, measured on the LVM-over-7-NVMe xfs the docker recipe ran on | +| `--ulimit memlock=-1` | *nothing* | no Pod-spec equivalent, and mostly moot: `CAP_IPC_LOCK` already exempts both containers from `RLIMIT_MEMLOCK`. See §2 | +| `RDMA_IB_DEVICES` / `MC_GID_INDEX`, both `require_env` in `leg.sh` | ``, ``, `` | same contract — the recipe never had defaults for these — expressed as placeholders instead of required environment. Split in two because the GID index is per node, and `leg.sh` was invoked once per node while one manifest covers both | +| image ENTRYPOINT bypassed by `docker exec` | image ENTRYPOINT bypassed by `command:` | same net effect. The ENTRYPOINT only matches a host `libionic` ABI, and the image bakes the ABI-4 build (`INSTALL_LIBIONIC=1`), so skipping it costs nothing on Mellanox and nothing on ionic either unless the image was built with `INSTALL_LIBIONIC=0` | +| `KVD_IO_MODE=auto` default | `--io-mode direct` | what the docker recipe actually ran. `auto`'s classifier walks the mount to the block device and, from inside a container, ends at a `/dev/mapper` node it cannot see — so it takes the conservative branch and picks buffered even on NVMe | + +Resource limits are new; the docker recipe ran without any, which makes the kvd +sidecar's the one worth checking before you copy it. kvd holds **two** independent +budgets: `--max-bytes 64G` caps the inline store, and the shared arena is sized +separately — it defaults to `--max-bytes`, and it is `mmap`'d and `mlock`'d whole at +startup, so it is never reclaimable under pressure. 64 + 64 + headroom is the +136 GiB in the manifest; size it to one of them and the OOM killer takes the sidecar +mid-run, which reads as a kvd bug rather than a limit. `prefill` gets 512 GiB +because `--hicache-size 32` is **GB per DP rank**, so 8 ranks pin 256 GB of host +tier; `decode` runs no host tier and gets half. `cpu: 32` matches `num_threads` in +`--model-loader-extra-config`. + +## Validation status + +| What | Status | +|---|---| +| This configuration in its **`docker` form** | brought up and benchmarked on 2 × MI300X. **TODO — numbers withheld pending a re-run.** Request count, cache efficiency and the kvd counters are being re-measured on one consistent run; quote nothing here until they land | +| **This manifest** | **not run.** Derived from that deployment flag for flag, and every deviation is in §6, but the Kubernetes form has not been brought up | +| Native kvd sidecar ordering | not run. The mechanism is standard k8s 1.29+; the claim that it removes the startup race is reasoned, not measured | + +Two things worth re-reading before a first bring-up, because both fail quietly +rather than loudly: `--advertise-host` resolving to a node IP that is not on the +RoCE rail (§6), and an engine image missing any of the three source fixes (§1) — +which is why §1 ends with a command that reads the markers out of the built image +rather than trusting the build log. + +## Source + +[`examples/recipes/glm5.2-fp8-gfx942/`](.) in +[AMD-AGI/Infera](https://github.com/AMD-AGI/Infera) · [all recipes](../README.md) · +[the same shape on MI355X, in `docker` form](../../sglang_1p1d_glm5.2/README.md) diff --git a/examples/recipes/glm5.2-fp8-gfx942/disaggregated-kvd/deploy.yaml b/examples/recipes/glm5.2-fp8-gfx942/disaggregated-kvd/deploy.yaml new file mode 100644 index 00000000..6287e662 --- /dev/null +++ b/examples/recipes/glm5.2-fp8-gfx942/disaggregated-kvd/deploy.yaml @@ -0,0 +1,298 @@ +# GLM-5.2-FP8 on gfx942 — disaggregated-kvd — Kubernetes recipe +# +# image infera:sglang-gfx942-glm52 built from deploy/docker/Dockerfile.sglang.gfx942 +# engine infera.engine.sglang TP8 / DP8 with DP-attention, MTP, fp8 KV +# KV move Mooncake RDMA over RoCE prefill node -> decode node +# KV tiers GPU -> SGLang hicache host RAM -> infera-kvd L2 -> L3 on node NVMe +# +# Docs: examples/recipes/glm5.2-fp8-gfx942/README.md — §1 why this recipe builds +# an engine image instead of mounting the overlay, §6 every host->k8s difference. +# Engine, router and kvd flags are lifted unchanged from a docker + shell +# deployment of the same topology, called "the docker recipe" in the comments +# below; anything the substrate forced is commented where it appears. +# +# Fill in before applying: +# / kubernetes.io/hostname of two gfx942 nodes on a +# mutually routable RoCE fabric. The KV handoff is +# RDMA and there is no TCP fallback. +# node directory holding GLM-5.2-FP8/, present at +# this path on BOTH nodes; mounted read-only at +# /models. +# node-local NVMe directory on the prefill node for +# kvd's L3. Must not be a shared filesystem. +# the RoCE rail(s) Mooncake may use, comma-separated, +# from `ibv_devices`. The docker recipe's own +# RDMA_IB_DEVICES; a rail that is down must not be +# listed. +# the GID index on that rail whose type is RoCE v2, +# from `show_gids`. Two placeholders because the +# index is PER NODE — two identical machines +# routinely differ. Usually the same value; check +# both anyway, because a wrong one pins KV to an +# interface that never carries it and the transfer +# just times out. +apiVersion: infera.amd.com/v1alpha1 +kind: InferaDeployment +metadata: + name: glm52-fp8-pd-kvd + namespace: infera +spec: + backendFramework: sglang + # Replaces the etcd the docker recipe ran: workers self-register on their Pod + # annotation and the router watches them. Same WorkerInfo either way. + discoveryBackend: kubernetes + nats: + deploy: false + services: + # ---- router: infera.server, kv-aware ------------------------------------ + server: + componentType: server + port: 8000 + extraPodSpec: + # Next to prefill, as the docker recipe ran it. Not required — the router + # reaches both legs over the node network — and it does put these + # requests on the prefill node's budget. See README §2. + nodeSelector: {kubernetes.io/hostname: } + containers: + - name: main + image: infera:sglang-gfx942-glm52 + imagePullPolicy: IfNotPresent + command: ["python3","-m","infera.server", + "--host","0.0.0.0","--port","8000","--router-backend","python", + "--router-tokenizer-path","/models/GLM-5.2-FP8", + "--discovery-backend","kubernetes", + "--request-transport","http","--kv-event-transport","zmq", + "--router-policy","kv-aware", + "--kv-prefill-overlap-weight","20.0", + "--kv-decode-overlap-weight","2.0"] + env: + - {name: POD_NAME, valueFrom: {fieldRef: {fieldPath: metadata.name}}} + - {name: POD_NAMESPACE, valueFrom: {fieldRef: {fieldPath: metadata.namespace}}} + - {name: POD_IP, valueFrom: {fieldRef: {fieldPath: status.podIP}}} + # The operator injects a readiness probe for workers only, so without + # this the Service takes traffic before uvicorn has bound the port. + readinessProbe: + httpGet: {path: /health, port: 8000} + periodSeconds: 5 + resources: + requests: {cpu: "8", memory: 16Gi} + limits: {cpu: "8", memory: 16Gi} + volumeMounts: + - {name: model, mountPath: /models, readOnly: true} + volumes: + - {name: model, hostPath: {path: , type: Directory}} + # ---- prefill leg + the kvd daemon --------------------------------------- + prefill: + componentType: worker + role: prefill + replicas: 1 + port: 30001 + # Weight load is 15-25 min cold. The startupProbe below is the readiness + # signal instead, with a budget matching INFERA_ENGINE_READY_TIMEOUT. + skipReadinessProbe: true + extraPodSpec: + nodeSelector: {kubernetes.io/hostname: } + # `--network host`. Required, not tuning: the RoCE rails are host + # interfaces, and the address this leg advertises for the Mooncake + # bootstrap handshake has to be one decode can open a QP to. + hostNetwork: true + dnsPolicy: ClusterFirstWithHostNet + # `--ipc host`, which is also what supplies /dev/shm. The node's must be + # large enough for TP8, and it is charged to the node, not to any limit. + hostIPC: true + initContainers: + # ---- kvd: L2 in host RAM + L3 on node NVMe ---- + # NATIVE sidecar (restartPolicy: Always, needs k8s 1.29+): the engine + # probes this socket once with a 5 s timeout and refuses to start if + # nothing answers, so kvd has to pass its startupProbe before `main` + # runs. That is the "kvd first" rule the docker recipe kept by hand. + - name: kvd + restartPolicy: Always + image: infera:sglang-gfx942-glm52 + imagePullPolicy: IfNotPresent + # --io-mode: `auto`'s probe ends at a /dev/mapper node it cannot see + # from inside a container and falls back to buffered. On the docker + # recipe's NVMe: 3.70 GB/s buffered vs 14.56 GB/s with O_DIRECT. + # --tablespace-pools: one slot must hold one whole hicache page, and + # GLM-5.2-FP8 at page_size 64 writes both 2.74 MiB KV pages and + # 624 KiB indexer pages. Oversize values are REJECTED, not split, so + # one pool would leave L3 silently empty for one of the two. + command: ["python3","-m","infera.kvd", + "--socket","/tmp/infera-kvd/kvd.sock", + "--max-bytes","64G", + "--long-path","/kvd-l3","--long-bytes","512G", + "--io-mode","direct", + "--use-tablespace","--tablespace-pools","1M,4M"] + startupProbe: + exec: + command: ["python3","-m","infera.kvd.statctl","--socket","/tmp/infera-kvd/kvd.sock"] + periodSeconds: 5 + timeoutSeconds: 5 # the default 1 s is short for a cold interpreter + failureThreshold: 36 # 3 min: L3 recovery replays the tablespace journal + securityContext: + capabilities: {add: ["IPC_LOCK"]} # required to mlock the arena + resources: + # Two independent budgets live here, not one. --max-bytes caps the + # inline store; the shared arena is sized separately and defaults to + # the same 64 GiB, and it is mmap'd and mlocked whole at startup, so + # it is never reclaimable. 64 + 64 + headroom. + requests: {cpu: "8", memory: 136Gi} + limits: {cpu: "8", memory: 136Gi} + volumeMounts: + - {name: kvd-sock, mountPath: /tmp/infera-kvd} + - {name: kvd-l3, mountPath: /kvd-l3} + containers: + - name: main + image: infera:sglang-gfx942-glm52 + imagePullPolicy: IfNotPresent + # `--device=/dev/infiniband` has no unprivileged equivalent without an + # RDMA device plugin; privileged + the hostPath below is the pattern + # validated in examples/k8s-deployments/pd-1p1d-mooncake.yaml. + securityContext: {privileged: true, capabilities: {add: ["IPC_LOCK","SYS_PTRACE"]}} + # `command` replaces the ENTRYPOINT, so infera-inject-host-ionic does + # not run — harmless, the image already bakes the ABI (README §6). + # --json-model-override-args is not tuning: it stands in for two + # sglang_dsa patches this base cannot take. Leave IndexShare on and + # the decode leg deadlocks on the first request. + # --hicache-size is GB PER DP RANK, so 32 is 256 GB across 8 ranks — + # that is what sets the memory limit below, and it is deliberately + # under the 54 GB/rank device pool that kvd's L3 already backs. + command: ["python3","-m","infera.engine.sglang", + "--model-path","/models/GLM-5.2-FP8", + "--host","0.0.0.0","--port","30001","--advertise-host","$(POD_IP)", + "--discovery-backend","kubernetes", + "--request-transport","http", + "--enable-kv-events","--kv-events","on","--kv-event-transport","zmq", + "--tp-size","8","--dp-size","8","--enable-dp-attention", + "--trust-remote-code","--kv-cache-dtype","fp8_e4m3", + "--reasoning-parser","glm45","--tool-call-parser","glm47", + "--dsa-prefill-backend","tilelang","--dsa-decode-backend","tilelang", + "--mem-fraction-static","0.85","--max-running-requests","128", + "--chunked-prefill-size","131072","--watchdog-timeout","1200", + "--disable-custom-all-reduce","--enable-cache-report", + "--speculative-algorithm","EAGLE","--speculative-num-steps","3", + "--speculative-eagle-topk","1","--speculative-num-draft-tokens","4", + "--json-model-override-args",'{"index_share_for_mtp_iteration":false}', + "--infera-kvd-socket","/tmp/infera-kvd/kvd.sock","--hicache-size","32", + "--weight-loader-prefetch-checkpoints", + "--model-loader-extra-config",'{"enable_multithread_load": true, "num_threads": 32}', + "--disaggregation-mode","prefill","--disaggregation-bootstrap-port","8998", + "--disaggregation-transfer-backend","mooncake", + "--disaggregation-ib-device",""] + env: + # POD_IP first: --advertise-host and the two SGLANG/HOST entries + # below expand it. With hostNetwork it is the node IP, which is what + # decode dials for the Mooncake bootstrap handshake — override + # --advertise-host if your RoCE rail is on a different address. + - {name: POD_IP, valueFrom: {fieldRef: {fieldPath: status.podIP}}} + - {name: POD_NAME, valueFrom: {fieldRef: {fieldPath: metadata.name}}} + - {name: POD_NAMESPACE, valueFrom: {fieldRef: {fieldPath: metadata.namespace}}} + - {name: SGLANG_HOST_IP, value: "$(POD_IP)"} + - {name: HOST_IP, value: "$(POD_IP)"} + # Per node, not per cluster — see the header and README §2. + - {name: MC_GID_INDEX, value: ""} + - {name: SGLANG_DSA_TRITON_PREFILL, value: "1"} + - {name: SGLANG_USE_AITER, value: "1"} + - {name: SAFETENSORS_FAST_GPU, value: "1"} + - {name: HSA_NO_SCRATCH_RECLAIM, value: "1"} + - {name: SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT, value: "3600"} + - {name: INFERA_ENGINE_READY_TIMEOUT, value: "5400"} + startupProbe: + httpGet: {path: /health, port: 30001} + periodSeconds: 15 + timeoutSeconds: 10 + failureThreshold: 360 # 90 min, matching INFERA_ENGINE_READY_TIMEOUT + resources: + # memory > 256 GiB of hicache host tier (32 GB x 8 DP ranks) plus + # the weight-load buffers. cpu matches num_threads in + # --model-loader-extra-config. + requests: {cpu: "32", memory: 512Gi, amd.com/gpu: 8} + limits: {cpu: "32", memory: 512Gi, amd.com/gpu: 8} + volumeMounts: + - {name: model, mountPath: /models, readOnly: true} + - {name: kvd-sock, mountPath: /tmp/infera-kvd} + - {name: ib, mountPath: /dev/infiniband} + volumes: + - {name: model, hostPath: {path: , type: Directory}} + # engine <-> kvd UDS. The zero-copy arena is a memfd handed over this + # socket via SCM_RIGHTS, so sharing the socket path is the whole + # requirement — there is no shared mount or /dev/shm object to line up. + - {name: kvd-sock, emptyDir: {}} + # hostPath, not a PVC: L3 must be node-local NVMe. A shared volume + # classifies as buffered and the reload lands inside the TTFT budget + # instead of under it. The Pod is already pinned to one node anyway. + - {name: kvd-l3, hostPath: {path: , type: DirectoryOrCreate}} + - {name: ib, hostPath: {path: /dev/infiniband, type: Directory}} + # ---- decode leg --------------------------------------------------------- + decode: + componentType: worker + role: decode + replicas: 1 + port: 31501 + skipReadinessProbe: true + extraPodSpec: + nodeSelector: {kubernetes.io/hostname: } + hostNetwork: true + dnsPolicy: ClusterFirstWithHostNet + hostIPC: true + # No kvd here, deliberately. SGLang issues storage prefetch on its + # aggregated and prefill branches only, so a decode-side L3 would be + # written and never read; infera detects a decode leg and refuses to + # wire kvd even if handed the socket. + containers: + - name: main + image: infera:sglang-gfx942-glm52 + imagePullPolicy: IfNotPresent + securityContext: {privileged: true, capabilities: {add: ["IPC_LOCK","SYS_PTRACE"]}} + # KV events stay OFF on this leg: prefill-side prefix locality is the + # win, and enabling them here can make SGLang reject the speculative + # disagg flags. + command: ["python3","-m","infera.engine.sglang", + "--model-path","/models/GLM-5.2-FP8", + "--host","0.0.0.0","--port","31501","--advertise-host","$(POD_IP)", + "--discovery-backend","kubernetes", + "--request-transport","http", + "--no-enable-kv-events","--kv-events","off", + "--tp-size","8","--dp-size","8","--enable-dp-attention", + "--trust-remote-code","--kv-cache-dtype","fp8_e4m3", + "--reasoning-parser","glm45","--tool-call-parser","glm47", + "--dsa-prefill-backend","tilelang","--dsa-decode-backend","tilelang", + "--mem-fraction-static","0.85","--max-running-requests","128", + "--chunked-prefill-size","131072","--watchdog-timeout","1200", + "--disable-custom-all-reduce","--enable-cache-report", + "--speculative-algorithm","EAGLE","--speculative-num-steps","3", + "--speculative-eagle-topk","1","--speculative-num-draft-tokens","4", + "--json-model-override-args",'{"index_share_for_mtp_iteration":false}', + "--weight-loader-prefetch-checkpoints", + "--model-loader-extra-config",'{"enable_multithread_load": true, "num_threads": 32}', + "--disaggregation-mode","decode", + "--disaggregation-transfer-backend","mooncake", + "--disaggregation-ib-device",""] + env: + - {name: POD_IP, valueFrom: {fieldRef: {fieldPath: status.podIP}}} + - {name: POD_NAME, valueFrom: {fieldRef: {fieldPath: metadata.name}}} + - {name: POD_NAMESPACE, valueFrom: {fieldRef: {fieldPath: metadata.namespace}}} + - {name: SGLANG_HOST_IP, value: "$(POD_IP)"} + - {name: HOST_IP, value: "$(POD_IP)"} + - {name: MC_GID_INDEX, value: ""} + - {name: SGLANG_DSA_TRITON_PREFILL, value: "1"} + - {name: SGLANG_USE_AITER, value: "1"} + - {name: SAFETENSORS_FAST_GPU, value: "1"} + - {name: HSA_NO_SCRATCH_RECLAIM, value: "1"} + - {name: SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT, value: "3600"} + - {name: INFERA_ENGINE_READY_TIMEOUT, value: "5400"} + startupProbe: + httpGet: {path: /health, port: 31501} + periodSeconds: 15 + timeoutSeconds: 10 + failureThreshold: 360 + resources: + # Half the prefill leg's memory: this one runs no hicache host tier. + requests: {cpu: "32", memory: 256Gi, amd.com/gpu: 8} + limits: {cpu: "32", memory: 256Gi, amd.com/gpu: 8} + volumeMounts: + - {name: model, mountPath: /models, readOnly: true} + - {name: ib, mountPath: /dev/infiniband} + volumes: + - {name: model, hostPath: {path: , type: Directory}} + - {name: ib, hostPath: {path: /dev/infiniband, type: Directory}} From 6ffa7152b39c6ea5cfdc2de6e95512cccd754472 Mon Sep 17 00:00:00 2001 From: xiaobochen-amd Date: Thu, 6 Aug 2026 15:38:42 +0000 Subject: [PATCH 39/88] =?UTF-8?q?ci:=20a=20failed=20squeue=20is=20not=20an?= =?UTF-8?q?=20empty=20queue=20=E2=80=94=20reclaim=20has=20to=20keep=20tryi?= =?UTF-8?q?ng?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The reclaim step existed to survive a transient Spur controller error, and was defeated by exactly that. `ids=$(squeue ... 2>/dev/null)` leaves a failed query looking identical to an empty one, so the loop read the error as "nothing to reclaim", printed that line, and broke out on its first pass. On 2026-08-06 four jobs were sitting in the queue when their reclaim ran and survived it. They held 3 of the reservation's 4 nodes until their 1h50m limit expired; every PR behind them failed waiting for a node that was not coming. Two smaller holes helped: the loop only spanned ~25s, and `scancel ... || true` kept the step green whether or not anything was actually cancelled. The five inline copies are now one script, called with the job-name prefix and suffix each job already computes: - retry on a non-zero squeue, printing its stderr instead of discarding it - exit 0 only after a *successful* query comes back empty - on timeout (RECLAIM_TIMEOUT, default 120s) emit ::error:: and dump the queue - never exit 0 unconfirmed: a leaked job squats a reserved GPU node, which is otherwise only discovered later, as a reservation that looks idle and is not The match keeps the old prefix+suffix shape, so infera-ci-hold-* (a -N2 holder, two nodes) and the spill and wipe jobs are still caught; a length guard stops a suffix longer than the job name from matching through substr. Verified with a stubbed squeue/scancel: clean queue exits 0 immediately; a permanently failing squeue retries and exits 1 where the old code reported success; of three queued jobs only the one matching both ends is cancelled, and the run exits 0 once a follow-up query confirms it is gone. Co-authored-by: Cursor Signed-off-by: xiaobochen-amd --- .github/scripts/reclaim_slurm_jobs.sh | 47 +++++++++++++++++++++++++++ .github/workflows/ci.yml | 36 +++----------------- .github/workflows/release.yml | 20 ++---------- 3 files changed, 54 insertions(+), 49 deletions(-) create mode 100755 .github/scripts/reclaim_slurm_jobs.sh diff --git a/.github/scripts/reclaim_slurm_jobs.sh b/.github/scripts/reclaim_slurm_jobs.sh new file mode 100755 index 00000000..8f07181c --- /dev/null +++ b/.github/scripts/reclaim_slurm_jobs.sh @@ -0,0 +1,47 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +# Cancel the SLURM jobs one CI job dispatched, and keep at it until the queue +# confirms they are gone. +# reclaim_slurm_jobs.sh +# +# The five inline copies this replaces discarded squeue's stderr, which made a +# FAILED query indistinguishable from an EMPTY one: a transient controller error +# read as "nothing to reclaim" and broke the retry loop on its first pass -- that +# error being the only reason the loop existed. On 2026-08-06 four jobs sitting in +# the queue when reclaim ran survived it and held 3 of the reservation's 4 nodes +# until their time limit expired. +set -uo pipefail + +prefix="${1:?usage: $0 }" +suffix="${2:?usage: $0 }" +budget="${RECLAIM_TIMEOUT:-120}" +interval="${RECLAIM_INTERVAL:-5}" +me="$(id -un)" +deadline=$(( SECONDS + budget )) + +echo "reclaiming SLURM jobs named ${prefix}*${suffix}" + +while :; do + # Exit code, not emptiness, is what separates an unreachable controller from a + # clean queue; stderr is folded in so the CI log names the failure. + if ! queue=$(squeue -h -u "$me" -o '%i %j' 2>&1); then + echo "squeue failed, retrying (this is NOT an empty queue): $queue" + else + ids=$(printf '%s\n' "$queue" | awk -v p="$prefix" -v s="$suffix" ' + index($2, p) == 1 && length($2) >= length(s) && + substr($2, length($2) - length(s) + 1) == s { print $1 }') + [ -z "$ids" ] && { echo "confirmed: no ${prefix}*${suffix} jobs left"; exit 0; } + echo "cancelling: $ids" + scancel $ids 2>&1 || echo "scancel returned non-zero, retrying" + fi + if [ "$SECONDS" -ge "$deadline" ]; then + # A leaked job holds a reserved GPU node until its time limit, so this has to + # be findable in the log rather than inferred later from a reservation that + # looks idle and is not. + echo "::error::could not confirm reclaim of ${prefix}*${suffix} within ${budget}s; check for leaked SLURM jobs" + squeue -u "$me" -o '%.10i %.44j %.2t %.10M %R' 2>&1 || true + exit 1 + fi + sleep "$interval" +done diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9c820eb8..4cfb4f28 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -220,15 +220,7 @@ jobs: run: exec bash tests/run_tests.sh engine - name: reclaim this job's SLURM jobs (on cancel/failure) if: always() && (cancelled() || failure()) - run: | - suf="-${{ github.run_id }}-engine" - for i in 1 2 3 4 5; do - ids=$(squeue -h -u "$(id -un)" -o '%i %j' 2>/dev/null \ - | awk -v suf="$suf" '$2 ~ /^infera-ci-/ && substr($2, length($2)-length(suf)+1)==suf {print $1}') - [ -z "$ids" ] && { echo "no (more) SLURM jobs to reclaim"; break; } - echo "reclaiming SLURM job(s): $ids (try $i)"; scancel $ids 2>&1 || true - sleep 5 - done + run: .github/scripts/reclaim_slurm_jobs.sh infera-ci- "-${{ github.run_id }}-engine" # Full PD-mixed e2e (per engine, parallel). When it runs is e2e_gate's call: # every PR into main, plus a push that lands untested code on main. @@ -272,16 +264,7 @@ jobs: run: exec bash tests/run_tests.sh e2e ${{ matrix.engine }} mixed - name: reclaim this job's SLURM jobs (on cancel/failure) if: always() && (cancelled() || failure()) - run: | - # Retry: a single scancel can hit a transient Spur controller error. - suf="-${{ github.run_id }}-${{ matrix.engine }}" - for i in 1 2 3 4 5; do - ids=$(squeue -h -u "$(id -un)" -o '%i %j' 2>/dev/null \ - | awk -v suf="$suf" '$2 ~ /^infera-ci-/ && substr($2, length($2)-length(suf)+1)==suf {print $1}') - [ -z "$ids" ] && { echo "no (more) SLURM jobs to reclaim"; break; } - echo "reclaiming SLURM job(s): $ids (try $i)"; scancel $ids 2>&1 || true - sleep 5 - done + run: .github/scripts/reclaim_slurm_jobs.sh infera-ci- "-${{ github.run_id }}-${{ matrix.engine }}" e2e-disag: # Gates mirror e2e-mixed, `!cancelled()` included: `always()` would keep this @@ -328,18 +311,9 @@ jobs: run: exec bash tests/run_tests.sh e2e ${{ matrix.engine }} disag - name: reclaim this job's SLURM jobs (on cancel/failure) if: always() && (cancelled() || failure()) - run: | - # Catches the infera-ci-hold-* pair holder too: it is a -N2 --gres=gpu:8 - # batch job, so a leaked one keeps TWO reserved nodes out of the pool. - # Retry: a single scancel can hit a transient Spur controller error. - suf="-${{ github.run_id }}-${{ matrix.engine }}-disag" - for i in 1 2 3 4 5; do - ids=$(squeue -h -u "$(id -un)" -o '%i %j' 2>/dev/null \ - | awk -v suf="$suf" '$2 ~ /^infera-ci-/ && substr($2, length($2)-length(suf)+1)==suf {print $1}') - [ -z "$ids" ] && { echo "no (more) SLURM jobs to reclaim"; break; } - echo "reclaiming SLURM job(s): $ids (try $i)"; scancel $ids 2>&1 || true - sleep 5 - done + # Catches the infera-ci-hold-* pair holder too: it is a -N2 --gres=gpu:8 + # batch job, so a leaked one keeps TWO reserved nodes out of the pool. + run: .github/scripts/reclaim_slurm_jobs.sh infera-ci- "-${{ github.run_id }}-${{ matrix.engine }}-disag" unit-torch-cpu: needs: [lint, pre_check, changes] diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index c456dd15..79d1af61 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -166,15 +166,7 @@ jobs: - name: reclaim this job's SLURM job (on cancel/failure) if: always() && (cancelled() || failure()) - run: | - suf="-${{ github.run_id }}-${{ matrix.engine }}" - for i in 1 2 3 4 5; do - ids=$(squeue -h -u "$(id -un)" -o '%i %j' 2>/dev/null \ - | awk -v suf="$suf" '$2 ~ /^infera-build-/ && substr($2, length($2)-length(suf)+1)==suf {print $1}') - [ -z "$ids" ] && { echo "no (more) SLURM jobs to reclaim"; break; } - echo "reclaiming SLURM job(s): $ids (try $i)"; scancel $ids 2>&1 || true - sleep 5 - done + run: .github/scripts/reclaim_slurm_jobs.sh infera-build- "-${{ github.run_id }}-${{ matrix.engine }}" # The base-agnostic overlay payload (deploy/overlay/). Unlike the engine # images this one is not independent: it builds its Python trees inside the @@ -236,15 +228,7 @@ jobs: - name: reclaim this job's SLURM job (on cancel/failure) if: always() && (cancelled() || failure()) - run: | - suf="-${{ github.run_id }}-overlay" - for i in 1 2 3 4 5; do - ids=$(squeue -h -u "$(id -un)" -o '%i %j' 2>/dev/null \ - | awk -v suf="$suf" '$2 ~ /^infera-build-/ && substr($2, length($2)-length(suf)+1)==suf {print $1}') - [ -z "$ids" ] && { echo "no (more) SLURM jobs to reclaim"; break; } - echo "reclaiming SLURM job(s): $ids (try $i)"; scancel $ids 2>&1 || true - sleep 5 - done + run: .github/scripts/reclaim_slurm_jobs.sh infera-build- "-${{ github.run_id }}-overlay" # Build the /manual Sphinx site alongside the images. Always uploads the HTML # as a workflow artifact; on a tag (release) it also attaches a tarball to the From 0be5deed66b57d25726cfbf678fdbedad7f10a9d Mon Sep 17 00:00:00 2001 From: liyingli Date: Thu, 6 Aug 2026 09:22:44 +0000 Subject: [PATCH 40/88] fix(recipes): the gfx942 recipe's io-mode rationale describes a fixed bug Rebasing onto main brought in f9a6a12, which teaches kvd's storage classifier to fall back to a major:minor walk through sysfs when lsblk cannot name a device. That is the exact failure the recipe cites as its reason for pinning --io-mode direct: an LVM mount seen from inside a container had no name lsblk could open, so the probe answered "unknown device" and the region took the conservative buffered branch. sysfs is indexed by numbers rather than names and the block layer is not namespaced, so auto reaches direct on its own now -- including from an unprivileged container, which is what the kvd sidecar is. The flag stays. It is the value the docker recipe ran and therefore the validated one, and a misclassification is a silent 4x on write-back, so pinning it is worth the line. Only the justification changes: it was a workaround, it is now a pin, and a reader who checks the classifier against the old wording would find they disagree. Left alone: the KVD_L3_DIR row's claim that a shared filesystem classifies as buffered. Nothing in that path resolves to a block device at all, so the sysfs fallback does not reach it. Also left alone is the "LVM-over-7-NVMe" figure -- it traces to the bench that produced the 3.70 / 14.56 GB/s pair on this same gfx942 mount, and the 8-NVMe rig in f9a6a12's message is where that fix was verified, not this node. Signed-off-by: liyingli Co-authored-by: Cursor --- examples/recipes/glm5.2-fp8-gfx942/README.md | 2 +- .../glm5.2-fp8-gfx942/disaggregated-kvd/deploy.yaml | 8 +++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/examples/recipes/glm5.2-fp8-gfx942/README.md b/examples/recipes/glm5.2-fp8-gfx942/README.md index 6041ed06..5d433c68 100644 --- a/examples/recipes/glm5.2-fp8-gfx942/README.md +++ b/examples/recipes/glm5.2-fp8-gfx942/README.md @@ -250,7 +250,7 @@ Every engine, router and kvd flag is identical. These are substrate translations | `--ulimit memlock=-1` | *nothing* | no Pod-spec equivalent, and mostly moot: `CAP_IPC_LOCK` already exempts both containers from `RLIMIT_MEMLOCK`. See §2 | | `RDMA_IB_DEVICES` / `MC_GID_INDEX`, both `require_env` in `leg.sh` | ``, ``, `` | same contract — the recipe never had defaults for these — expressed as placeholders instead of required environment. Split in two because the GID index is per node, and `leg.sh` was invoked once per node while one manifest covers both | | image ENTRYPOINT bypassed by `docker exec` | image ENTRYPOINT bypassed by `command:` | same net effect. The ENTRYPOINT only matches a host `libionic` ABI, and the image bakes the ABI-4 build (`INSTALL_LIBIONIC=1`), so skipping it costs nothing on Mellanox and nothing on ionic either unless the image was built with `INSTALL_LIBIONIC=0` | -| `KVD_IO_MODE=auto` default | `--io-mode direct` | what the docker recipe actually ran. `auto`'s classifier walks the mount to the block device and, from inside a container, ends at a `/dev/mapper` node it cannot see — so it takes the conservative branch and picks buffered even on NVMe | +| `KVD_IO_MODE=auto` default | `--io-mode direct` | what the docker recipe ran, and at the time it had to: the classifier resolved a mount by device *name*, and an LVM mount inside a container has no name it can open, so it took the conservative branch and picked buffered even on NVMe. That is fixed — the classifier now falls back to a `major:minor` walk through sysfs, which needs no device node and is not namespaced, so `auto` would reach `direct` on its own here. Kept pinned regardless: it is the validated value, and a misclassification is a silent 4x | Resource limits are new; the docker recipe ran without any, which makes the kvd sidecar's the one worth checking before you copy it. kvd holds **two** independent diff --git a/examples/recipes/glm5.2-fp8-gfx942/disaggregated-kvd/deploy.yaml b/examples/recipes/glm5.2-fp8-gfx942/disaggregated-kvd/deploy.yaml index 6287e662..bd1b6a9a 100644 --- a/examples/recipes/glm5.2-fp8-gfx942/disaggregated-kvd/deploy.yaml +++ b/examples/recipes/glm5.2-fp8-gfx942/disaggregated-kvd/deploy.yaml @@ -110,9 +110,11 @@ spec: restartPolicy: Always image: infera:sglang-gfx942-glm52 imagePullPolicy: IfNotPresent - # --io-mode: `auto`'s probe ends at a /dev/mapper node it cannot see - # from inside a container and falls back to buffered. On the docker - # recipe's NVMe: 3.70 GB/s buffered vs 14.56 GB/s with O_DIRECT. + # --io-mode: `auto` used to lose an LVM mount from inside a container + # and fall back to buffered — 3.70 GB/s against 14.56 with O_DIRECT + # here. It resolves that through sysfs now, so this is a pin rather + # than a workaround: it is what the docker recipe ran, and a + # misclassification costs 4x while nothing warns. # --tablespace-pools: one slot must hold one whole hicache page, and # GLM-5.2-FP8 at page_size 64 writes both 2.74 MiB KV pages and # 624 KiB indexer pages. Oversize values are REJECTED, not split, so From 331ab01a7ca0e9db35d07adf4ee01cd56bf8fcbb Mon Sep 17 00:00:00 2001 From: xiaobochen-amd Date: Thu, 6 Aug 2026 15:54:33 +0000 Subject: [PATCH 41/88] fix(tests): an unreachable controller is not a deleted reservation _reservation_nodes already returned non-zero when scontrol failed -- pipefail saw to that -- but all four callers reached for it as `[ -z "$(...)" ]`, which keeps the output and throws the status away. An unreachable controller and a reservation that is genuinely gone both arrive as an empty string, and the callers acted on the more destructive reading of the two. On 2026-08-06 a runner without SPUR_CONTROLLER_ADDR made every scontrol call answer "failed to connect to spurctld". The dispatcher read that as the pool having disappeared, printed mode=resv-gone->open, and sent the tier to the open partition; _e2e_preflight said "reservation does not exist (gone or expired)", which is where triage went first and lost time. Spur keeps the two distinguishable, which is what makes this fixable: it ignores the NAME argument and dumps every reservation, so asking for a name that is not there still exits 0 and the awk simply matches nothing. Only an unreachable controller exits non-zero. Verified on the live cluster. Each caller now separates them: _reservation_free -3 for a failed query, still -1 for a missing one _dispatch_slurm -3 joins the "keep the reservation" branch; only a query that answered may drop it _candidate_nodes yields nothing so the caller keeps waiting, instead of falling through to the open partition's idle list and handing the PD pair unreserved nodes run_e2e_disagg unsets INFERA_E2E_RESERVATION only on a confirmed absence; a failed query warns and keeps it _e2e_preflight says which of the two actually happened _reservation_nodes captures scontrol's output before parsing it. The status was already correct, but callers now depend on it, and under pipefail any stage of that pipeline can set it -- one `grep` added later would report every query as failed and leave the run clinging to a reservation that really had been deleted. Verified against the live cluster: an existing reservation resolves 4 nodes and _reservation_free returns 4; a name that does not exist exits 0 empty and returns -1; with SPUR_CONTROLLER_ADDR unset the query exits non-zero and returns -3, and the dispatcher prints mode=resv where it used to print mode=resv-gone->open. A healthy host is unchanged. Co-authored-by: Cursor Signed-off-by: xiaobochen-amd --- tests/run_tests.sh | 55 ++++++++++++++++++++++++++++++++-------------- 1 file changed, 39 insertions(+), 16 deletions(-) diff --git a/tests/run_tests.sh b/tests/run_tests.sh index 50e118bf..9e98900b 100755 --- a/tests/run_tests.sh +++ b/tests/run_tests.sh @@ -214,10 +214,16 @@ _skip_or_fail() { } _have_slurm() { command -v srun >/dev/null 2>&1; } -# The nodes reservation $1 covers, one per line ('' if it is gone/expired). -# Spur ignores the NAME arg and dumps all reservations; match the exact block. +# The nodes reservation $1 covers, one per line. Non-zero means the QUERY failed; +# exit 0 with no output means the reservation genuinely is not there. Spur keeps +# the two separable: it ignores the NAME arg and dumps every reservation, so a +# missing name still exits 0 and the awk below simply matches nothing. +# Capture before parsing: under pipefail a later stage's status would otherwise +# masquerade as a failed query, and callers now act on that distinction. _reservation_nodes() { - scontrol show reservation "$1" 2>/dev/null | awk -v r="ReservationName=$1" ' + local out + out=$(scontrol show reservation "$1" 2>/dev/null) || return 1 + printf '%s\n' "$out" | awk -v r="ReservationName=$1" ' BEGIN{RS="";FS="\n"} $1==r { for(i=1;i<=NF;i++) if($i ~ /Nodes=/){ n=$i; sub(/.*Nodes=/,"",n); sub(/[[:space:]].*/,"",n); print n; exit } }' \ | tr ',' '\n' | sed '/^$/d' @@ -235,7 +241,11 @@ _node_free() { # partition's idle ones. _candidate_nodes() { local n nodes="" - [ -n "${INFERA_E2E_RESERVATION:-}" ] && nodes="$(_reservation_nodes "$INFERA_E2E_RESERVATION")" + if [ -n "${INFERA_E2E_RESERVATION:-}" ]; then + # Query failed: offer nothing and let the caller keep waiting. Falling + # through would hand the PD pair unreserved nodes off the open partition. + nodes=$(_reservation_nodes "$INFERA_E2E_RESERVATION") || return 0 + fi if [ -z "$nodes" ]; then sinfo -h -N -p "$SLURM_PART" -t idle -o '%n' 2>/dev/null | awk 'NF && !seen[$0]++' return @@ -328,11 +338,11 @@ _amd_gpu_count() { _local_eligible() { [ "$(_amd_gpu_count)" -ge 8 ] && command -v docker >/dev/null 2>&1; } # Spill helper (Spur has no srun --immediate): free count, -1 if the reservation -# is gone/expired, -2 if scontrol is unavailable. +# is gone/expired, -2 if scontrol is unavailable, -3 if the query itself failed. _reservation_free() { local rname="$1" nodes n free=0 command -v scontrol >/dev/null 2>&1 || { echo -2; return; } - nodes=$(_reservation_nodes "$rname") + nodes=$(_reservation_nodes "$rname") || { echo -3; return; } [ -n "$nodes" ] || { echo -1; return; } for n in $nodes; do _node_free "$n" && free=$((free + 1)) @@ -413,10 +423,12 @@ _dispatch_slurm() { rfree=$(_reservation_free "$INFERA_E2E_RESERVATION") smax="${INFERA_E2E_SPILL_MAX:-2}" if [ "$rfree" = "-1" ]; then - echo "[$label] WARNING: reservation '$INFERA_E2E_RESERVATION' not found — falling back to open partition '$SLURM_PART'" >&2 + echo "[$label] WARNING: reservation '$INFERA_E2E_RESERVATION' does not exist — falling back to open partition '$SLURM_PART'" >&2 mode="resv-gone->open" elif [ "$rfree" != "0" ]; then - # free>0, or -2 (no scontrol): use the reservation. + # free>0, or -2/-3 (cannot tell): keep the reservation. Only a query that + # answered may drop it -- reading a controller blink as "gone" is what + # sent a whole run to the open partition on 2026-08-06. resv=(--reservation="$INFERA_E2E_RESERVATION"); mode="resv" else inflight=$(_spill_inflight) @@ -643,10 +655,16 @@ run_e2e_disagg() { fi # An expired reservation is worse than none — every step's `srun --reservation` - # would fail. Drop it, as _dispatch_slurm does for the mixed tier. - if [ -n "${INFERA_E2E_RESERVATION:-}" ] && [ -z "$(_reservation_nodes "$INFERA_E2E_RESERVATION")" ]; then - echo "[e2e disagg] WARNING: reservation '$INFERA_E2E_RESERVATION' not found — falling back to open partition '$SLURM_PART'" >&2 - unset INFERA_E2E_RESERVATION + # would fail. Drop it, as _dispatch_slurm does for the mixed tier, but only on + # a query that answered: a failed one says nothing about the pool. + local resv_nodes + if [ -n "${INFERA_E2E_RESERVATION:-}" ]; then + if ! resv_nodes=$(_reservation_nodes "$INFERA_E2E_RESERVATION"); then + echo "[e2e disagg] WARNING: cannot reach the scheduler to check reservation '$INFERA_E2E_RESERVATION' — keeping it" >&2 + elif [ -z "$resv_nodes" ]; then + echo "[e2e disagg] WARNING: reservation '$INFERA_E2E_RESERVATION' does not exist — falling back to open partition '$SLURM_PART'" >&2 + unset INFERA_E2E_RESERVATION + fi fi local rc=0 e prc out="$SCRATCH/.e2e-disag.out" @@ -713,10 +731,15 @@ run_e2e_disagg() { # Report-only: both tiers can still run (degraded) without a reservation or with # a nearly full /home, and a hard exit here would cost a whole CI run to find out. _e2e_preflight() { - local avail - if [ -n "${INFERA_E2E_RESERVATION:-}" ] && command -v scontrol >/dev/null 2>&1 \ - && [ -z "$(_reservation_nodes "$INFERA_E2E_RESERVATION")" ]; then - echo "[e2e] ERROR: reservation '$INFERA_E2E_RESERVATION' does not exist (gone or expired)" >&2 + local avail resv_nodes + if [ -n "${INFERA_E2E_RESERVATION:-}" ] && command -v scontrol >/dev/null 2>&1; then + if ! resv_nodes=$(_reservation_nodes "$INFERA_E2E_RESERVATION"); then + # This line used to say "does not exist" for an unreachable controller too, + # which sent triage looking for a deleted reservation. + echo "[e2e] ERROR: cannot reach the scheduler to check reservation '$INFERA_E2E_RESERVATION' (scontrol failed)" >&2 + elif [ -z "$resv_nodes" ]; then + echo "[e2e] ERROR: reservation '$INFERA_E2E_RESERVATION' does not exist (gone or expired)" >&2 + fi fi avail=$(df -Pk /home 2>/dev/null | awk 'NR==2{print $4}') case "$avail" in From fea6d466d94e910e17cd68697fe5de727ce701fe Mon Sep 17 00:00:00 2001 From: liyingli Date: Thu, 6 Aug 2026 11:24:46 +0000 Subject: [PATCH 42/88] docs(examples): add the GLM-5.2-FP8 gfx942 PD bring-up recipe The Kubernetes recipe under examples/recipes/glm5.2-fp8-gfx942/ lifts every flag from a `docker` + shell deployment of GLM-5.2-FP8 on 2 x MI300X, and then had nothing to point at: the scripts that deployment was tuned with lived in a working directory. A reader could see the values but not how they were reached, and could not run the thing the manifest is a translation of. This is that deployment: env.sh plus five launch scripts for etcd, an optional kvd daemon, the two SGLang legs and the Infera router, then verify.sh and bench.sh. Every value in env.sh is the measured one rather than SGLang's default, and carries the measurement that chose it -- chunked prefill at 1,024/rank (-23.8% duration, -34.4% TTFT against 16,384/rank), MTP 5/1/6 (-8.5% duration at 4.64 acceptance on a 4.00 break-even), the rust router (same routing decisions, 27% faster end to end), one RDMA rail (striping measured 11.9% slower), dp-attention on both legs (pure TP8 prefill 25.9% slower). Together, -45% on the trace they were tuned against. Two things shaped the scripts more than the flags did. Every interesting failure in this stack is silent. A KV hand-off that drops the prefix returns fluent text about something else, not an HTTP error; kv-aware without a tokenizer routes on load and looks healthy; a disagg pair that disagrees on the MTP shape just stops speculating. So verify.sh asserts an arithmetic answer only reachable through an intact prefix, greps the router's own pick log for a non-zero block count, and reads spec_accept_length off the decode leg -- five checks, each aimed at one of those, exiting non-zero rather than printing. kvd ships here but defaults to KVD=0, which is the measured-best shape for this workload and not a hedge: the tier cost 12% and served zero reads, the 54 GB/rank device pool already answering ~100% of the reuse the trace had. README section 6 says when to turn it on and how to tell, and env.sh refuses a KVD value that is neither 0 nor 1 -- KVD=true would otherwise skip the daemon and produce a baseline run wearing the kvd name. bench.sh is deliberately the simple one: sglang.bench_serving on random prompts, which sizes the deployment and reproduces none of the figures above. The agentic multi-turn trace that produced them is being generalised into a standalone tool and is not part of this. Co-authored-by: Cursor Signed-off-by: liyingli --- examples/glm5.2_gfx942/.gitignore | 4 + examples/glm5.2_gfx942/README.md | 322 ++++++++++++++++++ examples/glm5.2_gfx942/bench.sh | 49 +++ examples/glm5.2_gfx942/build_image.sh | 19 ++ examples/glm5.2_gfx942/env.sh | 144 ++++++++ examples/glm5.2_gfx942/host_container.sh | 96 ++++++ .../glm5.2_gfx942/launch/launch_decode.sh | 70 ++++ examples/glm5.2_gfx942/launch/launch_etcd.sh | 22 ++ examples/glm5.2_gfx942/launch/launch_kvd.sh | 56 +++ .../glm5.2_gfx942/launch/launch_prefill.sh | 77 +++++ .../glm5.2_gfx942/launch/launch_router.sh | 58 ++++ examples/glm5.2_gfx942/preflight_rdma.sh | 35 ++ examples/glm5.2_gfx942/stop.sh | 29 ++ examples/glm5.2_gfx942/verify.sh | 200 +++++++++++ 14 files changed, 1181 insertions(+) create mode 100644 examples/glm5.2_gfx942/.gitignore create mode 100644 examples/glm5.2_gfx942/README.md create mode 100644 examples/glm5.2_gfx942/bench.sh create mode 100644 examples/glm5.2_gfx942/build_image.sh create mode 100644 examples/glm5.2_gfx942/env.sh create mode 100644 examples/glm5.2_gfx942/host_container.sh create mode 100644 examples/glm5.2_gfx942/launch/launch_decode.sh create mode 100644 examples/glm5.2_gfx942/launch/launch_etcd.sh create mode 100644 examples/glm5.2_gfx942/launch/launch_kvd.sh create mode 100644 examples/glm5.2_gfx942/launch/launch_prefill.sh create mode 100644 examples/glm5.2_gfx942/launch/launch_router.sh create mode 100644 examples/glm5.2_gfx942/preflight_rdma.sh create mode 100644 examples/glm5.2_gfx942/stop.sh create mode 100644 examples/glm5.2_gfx942/verify.sh diff --git a/examples/glm5.2_gfx942/.gitignore b/examples/glm5.2_gfx942/.gitignore new file mode 100644 index 00000000..0b5adb9f --- /dev/null +++ b/examples/glm5.2_gfx942/.gitignore @@ -0,0 +1,4 @@ +# Written by the launch scripts and bench.sh. Engine logs run to megabytes and +# trip the large-file hook long before anyone reviews the diff. +logs/ +results/ diff --git a/examples/glm5.2_gfx942/README.md b/examples/glm5.2_gfx942/README.md new file mode 100644 index 00000000..59292322 --- /dev/null +++ b/examples/glm5.2_gfx942/README.md @@ -0,0 +1,322 @@ +# GLM-5.2-FP8 SGLang PD on gfx942 — bring-up, verification, benchmark + +Runnable package for GLM-5.2-FP8 on two gfx942 (MI300X) nodes: SGLang +prefill/decode disaggregation over Mooncake RDMA, DP-attention, EAGLE speculative +decoding, and Infera kv-aware routing. It covers bringing the service up, proving +it is actually correct, and sizing it with a simple serving benchmark. + +Every script is driven by environment variables, so you should not need to edit +any of them; `env.sh` holds the defaults, and they are the tuned recipe below +rather than SGLang's defaults. The node names (`node-0`, `node-1`) are +placeholders — substitute your own. + +KV offload below the GPU cache (`kvd`) ships here too, but **off by default** — +on the workload this was tuned against it cost 12% and served zero reads. §6 +covers turning it on and how to tell whether your workload is one that wants it. + +The agentic multi-turn benchmark that produced the numbers below is out of scope +here, and is being generalised into a standalone tool. + +## Scripts + +Half of these run on the host and half inside the long-lived engine container; +the column says which, and mixing the two is the most common way to get stuck. + +| Script | Runs on | Purpose | +| --- | --- | --- | +| `build_image.sh` | host | Build the engine image from `deploy/docker/Dockerfile.sglang.gfx942`. | +| `preflight_rdma.sh` | host | RDMA preflight: container port visibility, plus an optional cross-node fabric check. | +| `host_container.sh` | host | Create / inspect / remove the long-lived engine container. | +| `launch/launch_etcd.sh` | host | Start etcd on the prefill node (PD's shared registry). | +| `launch/launch_kvd.sh` | container | Optional KV offload daemon; a no-op at the default `KVD=0` (§6). | +| `launch/launch_prefill.sh` | container | SGLang prefill leg, TP8 / DP8 + dp-attention. | +| `launch/launch_decode.sh` | container | SGLang decode leg, same shape. | +| `launch/launch_router.sh` | container | Infera kv-aware router. | +| `verify.sh` | container | Correctness checks; exits non-zero on any failure. | +| `bench.sh` | container | SGLang `bench_serving` on a random dataset, through the router. | +| `stop.sh` | container | Stop the router and engine processes on this node. | + +## Topology + +| Node | Role | +| :--- | :--- | +| `node-0` | etcd + router + prefill leg + verify/benchmark entrypoint | +| `node-1` | decode leg | + +Both legs run TP8 / DP8 with `--enable-dp-attention` on 8 GPUs, and move the KV +cache to each other over one RDMA rail with Mooncake. + +## The tuned recipe + +These are the `env.sh` defaults. They were chosen on an **agentic multi-turn +trace** (32 conversations / 225 turns, concurrency 16, ~68k-token median input), +one axis at a time against a locked baseline: + +| Setting | Value | What it bought | +| --- | --- | --- | +| `ROUTER_BACKEND` | `rust` | Same routing decisions as the python backend request for request, 27% faster end to end. | +| `CHUNK` | `8192` (1,024/rank) | The largest single lever: −23.8% duration, −34.4% TTFT against 16,384/rank. | +| `MTP_STEPS`/`TOPK`/`DRAFT_TOKENS` | `5`/`1`/`6` | −8.5% duration, −13.8% TPOT. Acceptance 4.64 against a 4.00 break-even. | +| `IB_DEVICE` | one rail | Striping KV over every NIC measured 11.9% *slower*; KV uses 4.5% of one 200 Gb/s port. | +| dp-attention | on both legs | Pure TP8 prefill measured 25.9% slower: concurrency beats per-request latency here. | +| `KVD` | `0` | The offload tier cost 12% and served zero reads on this trace (§6). | +| `MEM_FRAC` / `MAX_RUNNING` | `0.85` / `128` | Baseline values, unchanged by the sweep. | + +Together those took that trace from 764 s to 420 s (−45%) and output throughput +from 64.8 to 118.0 tok/s. One metric moved the wrong way: ITL p90 rose 13.7%, +because a deeper draft emits tokens in burstier groups. That is free for batch +work and worth weighing for interactive streaming. + +`bench.sh` runs a *different*, simpler workload and will not reproduce those +figures — see §5. + +## 1. Prerequisites + +### 1.1 Hardware / software + +```text +Hardware: 2 nodes, 8x gfx942 (MI300X) each, RoCE between them, Docker with GPU access +Model: GLM-5.2-FP8 (glm_moe_dsa, MLA + DSA indexer, 78 layers) +Image: built here from deploy/docker/Dockerfile.sglang.gfx942 +``` + +### 1.2 Model weights + +`MODEL` must be a **local directory** — the scripts bind-mount it read-only: + +```bash +export MODEL=/your/path/GLM-5.2-FP8 +``` + +A HuggingFace cache path works too; `host_container.sh` detects the snapshot +symlinks and mounts the blobs alongside so they still resolve in the container. + +### 1.3 Build the image + +On both nodes (or build once and push). The image carries the ROCm hicache fixes +and the `infera-router` binary the rust backend needs, so do not substitute a +stock SGLang image: + +```bash +bash build_image.sh +``` + +## 2. Adapt to your cluster + +Export these on **both** nodes before running anything: + +```bash +export PREFILL_IP=10.0.0.1 # node-0, on the data network +export DECODE_IP=10.0.0.2 # node-1, on the data network +export MODEL=/your/path/GLM-5.2-FP8 +export IB_DEVICE=mlx5_0 # the RDMA rail the two nodes share +export MC_GID_INDEX=3 # RoCE GID index on that device +``` + +If your nodes resolve by name, `PREFILL_NODE` / `DECODE_NODE` derive the IPs +instead. The addresses must be the ones the peers can reach on the data network, +not a management NIC. + +## 3. Verify the RDMA fabric + +Cross-node PD moves the KV cache over the fabric on every request, and a mismatch +between the container's RDMA provider and the host driver degrades it to TCP +*silently* — the pair still answers, just slower than one node would be. Run this +on each node's host shell before bringing anything up: + +```bash +bash preflight_rdma.sh +``` + +The reported count of active RDMA ports must match the node's, not be `0`. For +the cross-node netperf and Mooncake probes, set a shared `DUMP_PATH` and run one +task per node (see `infera/tools/preflight/README.md`). + +## 4. Bring-up + +```text +node-0 (host): host_container.sh -> launch/launch_etcd.sh +node-1 (host): host_container.sh +node-0 (container): [launch/launch_kvd.sh] -> launch/launch_prefill.sh -> launch/launch_router.sh +node-1 (container): launch/launch_decode.sh +node-0 (container): verify.sh -> bench.sh +``` + +The two legs discover each other through etcd, so the decode leg can start in +parallel with the prefill leg. The router only needs both to be registered by the +time it takes traffic. The bracketed step is a no-op unless you set `KVD=1`. + +### 4.1 Containers + +On both nodes' host shells: + +```bash +bash host_container.sh +``` + +It checks the image, the weight mounts and the in-container imports before +reporting success, so a failure here is cheap compared to finding the same +problem four minutes into engine startup. Then, on `node-0` only: + +```bash +bash launch/launch_etcd.sh +``` + +### 4.2 Engines and router + +Enter the container on each node (`docker exec -it infera-glm52-gfx942 bash`), +then on `node-0`: + +```bash +bash launch/launch_prefill.sh +bash launch/launch_router.sh +``` + +and on `node-1`: + +```bash +bash launch/launch_decode.sh +``` + +**Cold start takes 15–25 minutes** — weights, then CUDA-graph capture. Do not +kill a slow launch. Follow it with `tail -f logs/prefill.log`. + +### 4.3 Verify + +On `node-0`, inside the container: + +```bash +bash verify.sh +``` + +Every check targets a failure this stack produces *without* returning an error: + +1. **Workers** — both legs registered in etcd. +2. **Correctness** — a padded prompt with a known answer. A broken KV hand-off + does not return an HTTP error; the decode leg reads a corrupt prefix and + produces fluent text unrelated to the prompt, which only an answer check + catches. +3. **kv-aware steering** — the router logged a prefill pick with + `request_blocks > 0`. Without block hashes it routes on load alone and looks + perfectly healthy doing it. +4. **MTP** — the decode leg's `/metrics` carries `sglang:spec_accept_length`. + Speculative decoding is dropped silently if the two legs disagree on its shape. +5. **kvd** — skipped at the default `KVD=0`; with the tier on, a fresh + multi-page prompt must leave writes behind in the daemon's counters. +6. **RDMA hand-off** — Mooncake transport lines in the decode log. + +### 4.4 Stop + +Inside the container on both nodes: + +```bash +bash stop.sh +``` + +Then, from the host, `bash host_container.sh --rm` and +`docker rm -f infera-glm52-etcd`. Remove the engine processes before relaunching +or the next run OOMs against VRAM the old one still holds. + +## 5. Benchmark + +```bash +bash bench.sh # defaults from env.sh +ISL=8192 OSL=512 CONC=32 bash bench.sh +for C in 8 16 32 64; do CONC=$C bash bench.sh; done # concurrency sweep +``` + +Defaults are `ISL=4096 OSL=1024 CONC=16`, and `NUM_PROMPTS` follows `CONC` at four +waves (64 prompts by default), so raising the concurrency alone keeps the run +length roughly constant. Results land in `results/.json` and `.log`. + +Two things to read correctly: + +- **The cache-hit line will be ~0, and that is right.** `--dataset-name random` + generates prompts that share no prefix, so a kv-aware router has nothing to + reuse. This benchmark sizes raw serving throughput; measuring cache reuse needs + a workload with real shared prefixes. +- **It will not reproduce the numbers in "The tuned recipe".** Those came from + the agentic trace described there, whose inputs are ~17× longer and heavily + prefix-shared. Use this to check the deployment is healthy and to compare + concurrencies against each other, not against those figures. + +## 6. Optional: KV offload below the GPU cache (`kvd`) + +`KVD=1` runs an `infera-kvd` daemon beside the prefill engine and points SGLang's +hierarchical cache at it, so a prefix evicted from the GPU survives in pinned host +RAM (L2) and on node-local NVMe (L3) instead of being recomputed. Prefill leg +only: SGLang issues storage prefetch on its aggregated and prefill branches, never +on the decode branch, so a decode-side daemon would be write-only. + +**It is off by default because it measured slower here.** On the agentic trace +above, `KVD=1` ran 12.0% slower and served *zero* reads while writing 100.8 GB: +the 54 GB-per-rank device pool already answered ~100% of the reuse that trace had +to offer, so every byte the tier stored was pure write cost. That is a property of +the workload, not of the tier — it earns its keep when the reuse horizon is longer +than the GPU pool can hold. Look at the prefill leg's cache hit rate first: if it +is already near its ceiling, offload has nothing left to catch. + +To turn it on, set these before creating the container, on the prefill node: + +```bash +export KVD=1 +export KVD_L3_DIR=/mnt/nvme/kvd-l3 # node-local NVMe; NFS/weka classifies as buffered +export KVD_IO_MODE=direct # pin it when you know that path is local NVMe +``` + +`KVD_L3_DIR` is bind-mounted, so an already-running container has to be recreated +(`bash host_container.sh --rm && bash host_container.sh`). Then, inside it, the +daemon goes up **before** the engine — the engine probes the socket at startup and +refuses to run without an answer: + +```bash +bash launch/launch_kvd.sh +bash launch/launch_prefill.sh +``` + +`verify.sh` then asserts the tier actually stores pages, and `bench.sh` writes the +daemon's counters next to each result as `.kvd.json`. You can read them at any +time with: + +```bash +python3 -m infera.kvd.statctl --socket /tmp/infera-kvd/kvd.sock +``` + +`hits_total` against `sets_total` is the whole question: writes with no hits over a +full run is exactly the 12% regression above. The remaining knobs — `KVD_RAM_BYTES`, +`KVD_LONG_BYTES`, `KVD_TABLESPACE_POOLS`, `HICACHE_SIZE` — are documented in +`env.sh`. For the same tier under Kubernetes, see +[`examples/recipes/glm5.2-fp8-gfx942/`](../recipes/glm5.2-fp8-gfx942/README.md). + +## Notes & gotchas + +1. **`CHUNK` is an aggregate, not a per-rank value.** Under dp-attention SGLang + splits it `CHUNK / dp_size`, so the default `8192` runs 1,024 per rank. The + engine log says `adjusted from … to …`, which reads like the setting was + rejected — it was not, it was divided. This is the single most misread value + in the recipe, so read back what actually took effect: + + ```bash + curl -s $PREFILL_URL/get_server_info \ + | python3 -c 'import json,sys; d=json.load(sys.stdin); print(d["chunked_prefill_size"], "per rank x", d["dp_size"], "ranks")' + ``` +2. **Both legs must agree on the MTP shape.** SGLang rejects a disaggregated pair + whose speculative config differs, so change `MTP_STEPS` / `MTP_DRAFT_TOKENS` + in `env.sh` (which both read) rather than on one leg's command line. +3. **The rust router requires etcd discovery.** `infera.server` validates the + supported subset before it execs the binary and fails with a pointer to + `--router-backend python`. This example uses etcd, so it is inside that subset; + a Kubernetes deployment is not, which is why the k8s recipe runs the python + backend. +4. **kv-aware fails soft.** Without a tokenizer it warns once and routes on load + alone. `launch_router.sh` refuses to start on that warning and `verify.sh` + re-checks it, because a run scored against load balancing while labelled + kv-aware is worse than a launch that stops. +5. **Advertise the data-network IP.** `--advertise-host` is what the peer dials + for the Mooncake bootstrap handshake; a management-NIC address there fails at + hand-off time, not at startup. +6. **`Ctrl-C` on a `tail -f` does not stop an engine.** The launch scripts run + them under `nohup`; use `stop.sh`. +7. **`kvd` outlives a restart, on disk.** `stop.sh` kills the daemon after the + engines, so nothing is pulled from under a live one, but L3 is journalled and + is recovered on the next start. Delete `KVD_L3_DIR` to start cold. diff --git a/examples/glm5.2_gfx942/bench.sh b/examples/glm5.2_gfx942/bench.sh new file mode 100644 index 00000000..7a4138f0 --- /dev/null +++ b/examples/glm5.2_gfx942/bench.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +# Simple serving benchmark through the router, using SGLang's own bench_serving on +# a random-token dataset. Run INSIDE the engine container on the PREFILL node: +# bash bench.sh +# ISL=8192 OSL=512 CONC=32 bash bench.sh +# +# This sizes the deployment; it is not the agentic workload the tuned recipe was +# chosen on. Random prompts share no prefix, so the kv-aware router has nothing to +# reuse and the cache-hit line below reads ~0 by construction -- that is correct, +# not a fault. See README §5. +set -euo pipefail +HERE="$(cd "$(dirname "$0")" && pwd)" +source "$HERE/env.sh" + +TAG="${TAG:-random_isl${ISL}_osl${OSL}_c${CONC}_n${NUM_PROMPTS}}" +OUT="$RESULT_DIR/$TAG" + +curl -sf -m 10 "$ROUTER_URL/health" >/dev/null \ + || { echo "[bench] router is not answering at $ROUTER_URL -- run verify.sh first" >&2; exit 1; } + +echo "[bench] $ROUTER_URL ISL=$ISL OSL=$OSL CONC=$CONC prompts=$NUM_PROMPTS -> $OUT.{json,log}" + +# --warmup-requests 1 keeps CUDA-graph capture and the first-token path out of the +# measured window. It cannot pre-warm a prefix here because the prompts are random. +python3 -m sglang.bench_serving \ + --backend sglang-oai-chat \ + --host "$PREFILL_IP" --port "$ROUTER_PORT" \ + --model "$MODEL" --tokenizer "$MODEL" \ + --dataset-name random \ + --random-input-len "$ISL" --random-output-len "$OSL" \ + --random-range-ratio "${RANGE:-0.8}" \ + --num-prompts "$NUM_PROMPTS" --max-concurrency "$CONC" \ + --request-rate "${RATE:-inf}" --warmup-requests "${WARMUP:-1}" \ + --seed "${SEED:-42}" \ + --cache-report --output-details --output-file "$OUT.json" \ + ${EXTRA_BENCH_ARGS:-} 2>&1 | tee "$OUT.log" + +# The offload tier's counters are the only place its side of the run is recorded, +# and they are cumulative, so capture them next to the result rather than leaving +# them to be read later against a different total. See README §6. +if [[ "$KVD" == "1" ]]; then + python3 -m infera.kvd.statctl --socket "$KVD_SOCKET" > "$OUT.kvd.json" 2>/dev/null \ + && echo "[bench] kvd counters: $OUT.kvd.json" \ + || echo "[bench] kvd counters unavailable on $KVD_SOCKET" >&2 +fi + +echo "[bench] done: $OUT.json" diff --git a/examples/glm5.2_gfx942/build_image.sh b/examples/glm5.2_gfx942/build_image.sh new file mode 100644 index 00000000..0e7a9f62 --- /dev/null +++ b/examples/glm5.2_gfx942/build_image.sh @@ -0,0 +1,19 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +# Build the engine image. Run on the host of each node, or build once and push. +# +# The image comes straight from deploy/docker/Dockerfile.sglang.gfx942. It already +# carries the ROCm hicache fixes and the infera-router binary; do not layer a +# second Dockerfile or runtime patches on top of it. +set -euo pipefail +HERE="$(cd "$(dirname "$0")" && pwd)" +source "$HERE/env.sh" + +DOCKERFILE="${DOCKERFILE:-$REPO/deploy/docker/Dockerfile.sglang.gfx942}" +[[ -f "$DOCKERFILE" ]] || { echo "[build] missing Dockerfile: $DOCKERFILE" >&2; exit 1; } + +echo "[build] image=$IMAGE" +echo "[build] dockerfile=$DOCKERFILE" +echo "[build] context=$REPO" +docker build -f "$DOCKERFILE" -t "$IMAGE" "$REPO" "$@" diff --git a/examples/glm5.2_gfx942/env.sh b/examples/glm5.2_gfx942/env.sh new file mode 100644 index 00000000..86e9190f --- /dev/null +++ b/examples/glm5.2_gfx942/env.sh @@ -0,0 +1,144 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +# Shared configuration for GLM-5.2-FP8 on two gfx942 nodes: SGLang prefill/decode +# disaggregation over Mooncake RDMA with Infera kv-aware routing. Sourced by every +# script here, so you should not need to edit any of them. Set at least +# PREFILL_IP, DECODE_IP and MODEL for your cluster. +# +# The engine values below are the tuned recipe (README "The tuned recipe"), not +# SGLang defaults; each carries the measurement that chose it. +set -uo pipefail + +# --- topology --------------------------------------------------------------- +# One prefill node, which also hosts etcd and the router, and one decode node. +# If your nodes resolve by name, setting PREFILL_NODE / DECODE_NODE is enough. +export PREFILL_NODE="${PREFILL_NODE:-node-0}" +export DECODE_NODE="${DECODE_NODE:-node-1}" +export PREFILL_IP="${PREFILL_IP:-$(getent ahostsv4 "$PREFILL_NODE" 2>/dev/null | awk 'NR==1{print $1}')}" +export DECODE_IP="${DECODE_IP:-$(getent ahostsv4 "$DECODE_NODE" 2>/dev/null | awk 'NR==1{print $1}')}" +: "${PREFILL_IP:=127.0.0.1}" +: "${DECODE_IP:=$PREFILL_IP}" + +export ETCD_ENDPOINT="${ETCD_ENDPOINT:-${PREFILL_IP}:2379}" +export ROUTER_PORT="${ROUTER_PORT:-8000}" +export PREFILL_PORT="${PREFILL_PORT:-30001}" +export DECODE_PORT="${DECODE_PORT:-31501}" +export BOOTSTRAP_PORT="${BOOTSTRAP_PORT:-8998}" + +export ROUTER_URL="${ROUTER_URL:-http://${PREFILL_IP}:${ROUTER_PORT}}" +export PREFILL_URL="${PREFILL_URL:-http://${PREFILL_IP}:${PREFILL_PORT}}" +export DECODE_URL="${DECODE_URL:-http://${DECODE_IP}:${DECODE_PORT}}" + +# --- image / container ------------------------------------------------------ +# Built by build_image.sh straight from deploy/docker/Dockerfile.sglang.gfx942. +# Do not layer runtime SGLang patches on it. +export IMAGE="${IMAGE:-infera:sglang-gfx942-glm52}" +export CONTAINER="${CONTAINER:-infera-glm52-gfx942}" +export ETCD_CONTAINER="${ETCD_CONTAINER:-infera-glm52-etcd}" + +# --- paths ------------------------------------------------------------------ +# REPO is bind-mounted at the same path inside the container, so these scripts +# resolve identically on the host and in the container. +EXAMPLE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +export REPO="${REPO:-$(cd "$EXAMPLE_DIR/../.." && pwd)}" +export MODEL="${MODEL:-/your/path/GLM-5.2-FP8}" # local weights dir, mounted read-only +export LOG_DIR="${LOG_DIR:-${EXAMPLE_DIR}/logs}" +export RESULT_DIR="${RESULT_DIR:-${EXAMPLE_DIR}/results}" + +# --- fabric ----------------------------------------------------------------- +# One rail carries the KV transfer. Striping it over every NIC was measured 11.9% +# slower and cannot help: KV uses 4.5% of a single 200 Gb/s port on this workload. +export IB_DEVICE="${IB_DEVICE:-mlx5_0}" +export MC_GID_INDEX="${MC_GID_INDEX:-3}" # RoCE GID index on that device +export TP="${TP:-8}" +export DP="${DP:-8}" + +# --- engine ----------------------------------------------------------------- +export MEM_FRAC="${MEM_FRAC:-0.85}" +export MAX_RUNNING="${MAX_RUNNING:-128}" + +# Aggregate value, NOT per rank: dp-attention splits it CHUNK/DP, so 8192 is +# 1,024/rank. Largest single lever on this deployment -- 1,024/rank beat the +# 16,384/rank this recipe used before by 23.8% on duration and 34.4% on TTFT. +# 512/rank is past the knee, where scheduling more chunks costs more than the +# smaller chunk saves. +export CHUNK="${CHUNK:-8192}" + +# EAGLE draft depth. TPOT = decode step time / accept_length, and each extra draft +# step measured 6.53 ms here, so deepening pays only while acceptance rises faster +# than the step cost. 5/1/6 accepts 4.64 against a 4.00 break-even; 7/1/8 both +# misses its break-even and runs prefill out of activation memory. Both legs must +# agree -- SGLang rejects a disagg pair whose speculative config differs. +export MTP="${MTP:-1}" # 0 disables speculative decoding +export MTP_STEPS="${MTP_STEPS:-5}" +export MTP_TOPK="${MTP_TOPK:-1}" +export MTP_DRAFT_TOKENS="${MTP_DRAFT_TOKENS:-6}" + +# Prometheus on both engines, off by default in SGLang. Without it /metrics 404s, +# and that endpoint is the only network-reachable source of the MTP acceptance +# length verify.sh checks. --enable-metrics-for-all-schedulers matters under +# dp-attention: TP 0 is the only scheduler that reports otherwise, so 1 of 8 DP +# ranks would stand in for the fleet. +export ENGINE_METRICS="${ENGINE_METRICS:-1}" + +# --- router ----------------------------------------------------------------- +# rust execs the infera-router binary the image carries at /usr/local/bin. It +# makes the same routing decisions as the python backend request for request and +# measured 27% faster end to end, so it is the default. Its supported subset is +# what this example passes anyway (etcd discovery, http transport, kv-aware). +export ROUTER_BACKEND="${ROUTER_BACKEND:-rust}" +export ROUTER_POLICY="${ROUTER_POLICY:-kv-aware}" + +# --- kvd: KV offload below the GPU cache ------------------------------------ +# KVD=1 runs an infera-kvd daemon beside the prefill engine and points SGLang's +# hierarchical cache at it, so prefixes evicted from the GPU survive in host RAM +# (L2) and on local NVMe (L3). +# +# KVD=0 is the default because the tier measured SLOWER on the workload this +# recipe was tuned on: it cost 12% and served zero reads, the GPU pool alone +# already covering ~100% of the reuse that trace had to offer. Turn it on when +# your working set outgrows 54 GB/rank; see README §6 for how to tell. +# +# Prefill leg only. SGLang issues storage prefetch on its aggregated and prefill +# branches, never on the decode branch, so kvd on a decode leg is write-only -- +# infera refuses to wire it there and says so in the log. +export KVD="${KVD:-0}" +export KVD_SOCKET="${KVD_SOCKET:-/tmp/infera-kvd/kvd.sock}" +# Must be node-local NVMe. Anything shared (NFS, weka) classifies as buffered and +# the reload lands in the TTFT budget instead of under it. +export KVD_L3_DIR="${KVD_L3_DIR:-/your/path/kvd-l3}" +export KVD_RAM_BYTES="${KVD_RAM_BYTES:-64G}" # L2 arena, pinned host RAM +export KVD_LONG_BYTES="${KVD_LONG_BYTES:-512G}" # L3 budget under KVD_L3_DIR +# O_DIRECT vs buffered for L3. `auto` classifies the mount by walking sysfs from +# its major:minor, which is not namespaced and so works unprivileged in a +# container; it falls back to `buffered` whenever it cannot identify the device, +# an overlay path being the usual reason. Pin `direct` when you know KVD_L3_DIR +# is local NVMe: a misclassification is a silent 4x, 14.56 GB/s against 3.70 on +# the LVM-over-NVMe xfs this was measured on. launch_kvd.sh prints the verdict. +export KVD_IO_MODE="${KVD_IO_MODE:-auto}" +# One tablespace slot must hold one whole hicache page, and a page holds every +# layer: GLM-5.2-FP8 at page_size 64 writes 2.74 MiB per KV page and 624 KiB per +# DSA-indexer page. Two pools cover both without wasting a slot on the smaller. +# A value that outgrows its largest pool is REJECTED, not split, which leaves L3 +# silently empty -- verify.sh fails on that rather than letting you find it later. +export KVD_TABLESPACE_POOLS="${KVD_TABLESPACE_POOLS:-1M,4M}" +# SGLang's own host tier, in GB PER DP RANK (so x8 here). It sits between the GPU +# pool and kvd and stages L3 reads. Deliberately smaller than the 54 GB device +# pool per rank: matching it would pin ~870 GB of host RAM for an L2 that kvd's L3 +# already backs. +export HICACHE_SIZE="${HICACHE_SIZE:-32}" +# KVD=true is not "1": the daemon would be skipped, the engine would start with +# no tier, and the run would look like a kvd run. Refuse rather than guess. +case "$KVD" in 0|1) ;; *) echo "[env] KVD='$KVD' is not '0' or '1'" >&2; exit 1 ;; esac + +# --- bench sizing ----------------------------------------------------------- +# Deliberately small: bench.sh sizes the deployment, it is not a sweep. Four waves +# of 16 concurrent is a few minutes. NUM_PROMPTS follows CONC so that raising the +# concurrency alone keeps the same number of waves. See README §5. +export ISL="${ISL:-4096}" +export OSL="${OSL:-1024}" +export CONC="${CONC:-16}" +export NUM_PROMPTS="${NUM_PROMPTS:-$((CONC * 4))}" + +mkdir -p "$LOG_DIR" "$RESULT_DIR" 2>/dev/null || true diff --git a/examples/glm5.2_gfx942/host_container.sh b/examples/glm5.2_gfx942/host_container.sh new file mode 100644 index 00000000..f51ebce2 --- /dev/null +++ b/examples/glm5.2_gfx942/host_container.sh @@ -0,0 +1,96 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +# Create or remove the long-lived engine container. Run on the HOST shell of both +# nodes; the launch, verify and bench scripts then run inside it. +# bash host_container.sh # create +# bash host_container.sh --status +# bash host_container.sh --rm +set -euo pipefail +HERE="$(cd "$(dirname "$0")" && pwd)" +source "$HERE/env.sh" + +case "${1:-}" in + --status) + docker ps -a --filter "name=^/${CONTAINER}$" \ + --format 'name={{.Names}} state={{.State}} image={{.Image}} up={{.Status}}' + exit 0 ;; + --rm) + docker rm -f "$CONTAINER" >/dev/null 2>&1 && echo "[container] removed $CONTAINER" \ + || echo "[container] $CONTAINER not present" + exit 0 ;; + "") ;; + *) echo "usage: bash $0 [--status|--rm]" >&2; exit 2 ;; +esac + +if [[ -n "$(docker ps -q --filter "name=^/${CONTAINER}$")" ]]; then + echo "[container] $CONTAINER already running" + exit 0 +fi + +docker image inspect "$IMAGE" >/dev/null \ + || { echo "[container] image missing: $IMAGE; run ./build_image.sh or docker pull it" >&2; exit 1; } +[[ -d "$MODEL" ]] || { echo "[container] model not found: $MODEL" >&2; exit 1; } + +# A HuggingFace cache is a tree of symlinks -- every file in a snapshot dir is a +# relative link into a sibling blobs/, and MODEL is often itself a link into that +# snapshot. `-v $MODEL:$MODEL` makes docker resolve the outer link and mount the +# snapshot at MODEL's path, which leaves every inner link dangling one directory +# short of blobs/. The symptom is not a missing file but transformers refusing the +# model with "Should have a `model_type` key in its config.json". +# So when MODEL is a link, mount what the links actually point at, each at its own +# path, and let MODEL resolve inside the container the same way it does outside. +MODEL_ARGS=(-v "$MODEL:$MODEL:ro") +MODEL_REAL="$(readlink -f "$MODEL")" +if [[ "$MODEL_REAL" != "$MODEL" ]]; then + HF_REPO="$MODEL_REAL" + [[ "$HF_REPO" == */snapshots/* ]] && HF_REPO="${HF_REPO%/snapshots/*}" # keep blobs/ + MODEL_ARGS=(-v "$(dirname "$MODEL"):$(dirname "$MODEL"):ro" -v "$HF_REPO:$HF_REPO:ro") + echo "[container] $MODEL -> $MODEL_REAL; mounting $HF_REPO so its blobs resolve" +fi + +docker rm -f "$CONTAINER" >/dev/null 2>&1 || true + +IONIC_ARGS=() +if [[ -e "${HOST_LIBIONIC:-/usr/lib/x86_64-linux-gnu/libionic.so}" ]]; then + IONIC_ARGS=(-v "${HOST_LIBIONIC:-/usr/lib/x86_64-linux-gnu/libionic.so}:/host-libionic/libionic.so:ro") +fi + +# Nothing when KVD=0, which is the default. When on, kvd's L3 must land on the +# host's NVMe rather than the container's overlay, so bind-mount it. Harmless on +# the decode node, which runs no daemon; mounting it there anyway keeps one +# container command for both nodes. The socket stays inside the container -- +# only the daemon and the prefill engine, both in here, ever open it. +KVD_ARGS=() +if [[ "$KVD" == "1" ]]; then + mkdir -p "$KVD_L3_DIR" \ + || { echo "[container] cannot create KVD_L3_DIR: $KVD_L3_DIR; set it for your cluster" >&2; exit 1; } + KVD_ARGS=(-v "$KVD_L3_DIR:$KVD_L3_DIR") +fi + +docker run -d --name "$CONTAINER" \ + --network host --ipc host --shm-size 128g \ + --device=/dev/kfd --device=/dev/dri --device=/dev/infiniband \ + --group-add video --group-add render \ + --cap-add=IPC_LOCK --cap-add=SYS_PTRACE \ + --security-opt seccomp=unconfined \ + --ulimit memlock=-1 --ulimit stack=67108864 --ulimit nofile=1048576 \ + "${IONIC_ARGS[@]}" "${KVD_ARGS[@]}" \ + -v "$REPO:$REPO" \ + "${MODEL_ARGS[@]}" \ + -w "$HERE" \ + "$IMAGE" sleep infinity >/dev/null + +sleep 2 +docker exec "$CONTAINER" bash -lc \ + 'python3 -c "import infera, sglang, torch; print(\"infera ok\", \"sglang\", sglang.__version__, \"torch\", torch.__version__)"' \ + || { echo "[container] import check failed"; docker logs "$CONTAINER" | tail -40; exit 1; } + +# Read the config the way the engine will. A mount that leaves the weights' +# symlinks dangling otherwise surfaces minutes into engine startup as a +# transformers error about model_type, far from its cause. +docker exec -e M="$MODEL" "$CONTAINER" bash -lc \ + 'python3 -c "import json, os; print(\"model ok\", json.load(open(os.environ[\"M\"] + \"/config.json\"))[\"model_type\"])"' \ + || { echo "[container] cannot read $MODEL/config.json inside the container -- check the mounts above" >&2; exit 1; } + +echo "[container] $CONTAINER up on $(hostname); enter with: docker exec -it $CONTAINER bash" diff --git a/examples/glm5.2_gfx942/launch/launch_decode.sh b/examples/glm5.2_gfx942/launch/launch_decode.sh new file mode 100644 index 00000000..c43da278 --- /dev/null +++ b/examples/glm5.2_gfx942/launch/launch_decode.sh @@ -0,0 +1,70 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +# Launch the SGLang decode leg on the decode node. +# Run INSIDE the engine container: bash launch/launch_decode.sh +set -euo pipefail +HERE="$(cd "$(dirname "$0")/.." && pwd)" +source "$HERE/env.sh" + +LOG="${LOG:-$LOG_DIR/decode.log}" +HOST_IP="${HOST_IP:-$DECODE_IP}" +PORT="${PORT:-$DECODE_PORT}" + +# kv events stay off on this leg. Prefill-side prefix locality is what the router +# steers by, and decode events can make SGLang reject the speculative disagg flags. +KV_EVENT_ARGS=(--no-enable-kv-events --kv-events off) + +# There is no kvd block here on purpose, even at KVD=1: SGLang issues storage +# prefetch on its aggregated and prefill branches only, so an offload tier on this +# leg would be write-only. + +MTP_ARGS=() +if [[ "$MTP" != "0" ]]; then + MTP_ARGS=(--speculative-algorithm EAGLE --speculative-num-steps "$MTP_STEPS" + --speculative-eagle-topk "$MTP_TOPK" --speculative-num-draft-tokens "$MTP_DRAFT_TOKENS" + --json-model-override-args '{"index_share_for_mtp_iteration":false}') +fi + +# This is the leg that matters for MTP: speculative verification happens here, so +# `accept len` is in THIS log and sglang:spec_accept_length on THIS /metrics. +METRICS_ARGS=() +if [[ "$ENGINE_METRICS" == "1" ]]; then + METRICS_ARGS=(--enable-metrics --enable-metrics-for-all-schedulers) +fi + +export HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}" +export CUDA_VISIBLE_DEVICES="$HIP_VISIBLE_DEVICES" +export SGLANG_HOST_IP="$HOST_IP" HOST_IP +export SGLANG_DSA_TRITON_PREFILL=1 SAFETENSORS_FAST_GPU=1 +export HSA_NO_SCRATCH_RECLAIM=1 SGLANG_USE_AITER="${SGLANG_USE_AITER:-1}" +export MC_GID_INDEX +export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT="${SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT:-3600}" +export INFERA_ENGINE_READY_TIMEOUT="${INFERA_ENGINE_READY_TIMEOUT:-5400}" + +pkill -f "(infera.engine.sglang|sglang.launch_server) .*--port ${PORT}( |$)" 2>/dev/null || true +sleep 5 + +# Both legs are the two ends of one Mooncake transfer, so IB_DEVICE, the MTP shape +# and the KV dtype must match launch_prefill.sh. +nohup python3 -m infera.engine.sglang \ + --model-path "$MODEL" --host 0.0.0.0 --port "$PORT" --advertise-host "$HOST_IP" \ + --etcd-endpoint "$ETCD_ENDPOINT" --discovery-backend etcd \ + --request-transport http "${KV_EVENT_ARGS[@]}" \ + --tp-size "$TP" --dp-size "$DP" --enable-dp-attention \ + --trust-remote-code --kv-cache-dtype fp8_e4m3 \ + --reasoning-parser glm45 --tool-call-parser glm47 \ + --dsa-prefill-backend tilelang --dsa-decode-backend tilelang \ + --mem-fraction-static "$MEM_FRAC" --max-running-requests "$MAX_RUNNING" \ + --chunked-prefill-size "$CHUNK" --watchdog-timeout 1200 \ + --disable-custom-all-reduce --enable-cache-report \ + "${MTP_ARGS[@]}" "${METRICS_ARGS[@]}" ${EXTRA_ARGS:-} \ + --weight-loader-prefetch-checkpoints \ + --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 32}' \ + --disaggregation-mode decode \ + --disaggregation-transfer-backend mooncake --disaggregation-ib-device "$IB_DEVICE" \ + > "$LOG" 2>&1 & + +echo "[decode] loading on ${HOST_IP}:${PORT}, TP=$TP DP=$DP MTP=$MTP($MTP_STEPS/$MTP_TOPK/$MTP_DRAFT_TOKENS)" +echo "[decode] rail=$IB_DEVICE metrics=$ENGINE_METRICS; log=$LOG" +echo "[decode] cold start can take 15-25 min; follow with: tail -f $LOG" diff --git a/examples/glm5.2_gfx942/launch/launch_etcd.sh b/examples/glm5.2_gfx942/launch/launch_etcd.sh new file mode 100644 index 00000000..838ff5e8 --- /dev/null +++ b/examples/glm5.2_gfx942/launch/launch_etcd.sh @@ -0,0 +1,22 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +# Start etcd on the prefill node, the registry both legs and the router discover +# each other through. Run on the HOST shell -- intentionally not in the engine +# container, which stays focused on infera/sglang and carries no etcd binary. +set -euo pipefail +HERE="$(cd "$(dirname "$0")/.." && pwd)" +source "$HERE/env.sh" + +HOST_IP="${ETCD_HOST_IP:-$PREFILL_IP}" +ETCD_IMAGE="${ETCD_IMAGE:-quay.io/coreos/etcd:v3.5.14}" + +docker rm -f "$ETCD_CONTAINER" >/dev/null 2>&1 || true +docker run -d --name "$ETCD_CONTAINER" --net host "$ETCD_IMAGE" \ + etcd --advertise-client-urls "http://${HOST_IP}:2379" \ + --listen-client-urls "http://0.0.0.0:2379" >/dev/null + +sleep 3 +docker exec "$ETCD_CONTAINER" etcdctl endpoint health \ + || { docker logs "$ETCD_CONTAINER" 2>&1 | tail -40; exit 1; } +echo "[etcd] ready: ${HOST_IP}:2379" diff --git a/examples/glm5.2_gfx942/launch/launch_kvd.sh b/examples/glm5.2_gfx942/launch/launch_kvd.sh new file mode 100644 index 00000000..464a5320 --- /dev/null +++ b/examples/glm5.2_gfx942/launch/launch_kvd.sh @@ -0,0 +1,56 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +# Launch the infera-kvd daemon on the prefill node, BEFORE the prefill engine. +# Run INSIDE the engine container: bash launch/launch_kvd.sh +# +# No-op unless KVD=1, so the bring-up sequence in the README can call it +# unconditionally. The prefill engine's socket probe refuses to start when the +# daemon is not answering, hence the order: kvd, then prefill. +set -euo pipefail +HERE="$(cd "$(dirname "$0")/.." && pwd)" +source "$HERE/env.sh" + +if [[ "$KVD" != "1" ]]; then + echo "[kvd] KVD=0 -- no offload tier, GPU radix tree is the only cache" + exit 0 +fi + +LOG="${LOG:-$LOG_DIR/kvd.log}" + +mkdir -p "$(dirname "$KVD_SOCKET")" "$KVD_L3_DIR" \ + || { echo "[kvd] cannot create $(dirname "$KVD_SOCKET") or $KVD_L3_DIR" >&2; exit 1; } + +pkill -f "infera\.kvd .*--socket ${KVD_SOCKET}( |$)" 2>/dev/null || true +# A stale socket file from a killed daemon still passes `test -S`, and the +# engine's probe then fails with Connection refused instead of waiting. +rm -f "$KVD_SOCKET" +sleep 1 + +# --use-tablespace is load-bearing: without it --long-path builds the legacy +# file-per-block region instead of the container-file tablespace (bounded file +# count, O_DIRECT, journal-recovered index). +nohup python3 -m infera.kvd \ + --socket "$KVD_SOCKET" \ + --max-bytes "$KVD_RAM_BYTES" \ + --long-path "$KVD_L3_DIR" --long-bytes "$KVD_LONG_BYTES" \ + --io-mode "$KVD_IO_MODE" \ + --use-tablespace --tablespace-pools "$KVD_TABLESPACE_POOLS" \ + > "$LOG" 2>&1 & + +# Wait until the daemon answers an RPC, not merely until the socket file appears: +# the engine's own startup probe connects for real, so a weaker wait here would +# just move the failure into engine startup. +READY=0 +for _ in $(seq 60); do + if python3 -m infera.kvd.statctl --socket "$KVD_SOCKET" >/dev/null 2>&1; then READY=1; break; fi + sleep 1 +done +[[ "$READY" == "1" ]] \ + || { echo "[kvd] daemon not answering on $KVD_SOCKET" >&2; tail -40 "$LOG"; exit 1; } + +echo "[kvd] up: socket=$KVD_SOCKET ram=$KVD_RAM_BYTES l3=$KVD_L3_DIR ($KVD_LONG_BYTES); log=$LOG" +# The classifier picks O_DIRECT or buffered for L3, and a shared mount landing on +# buffered is the difference between reading L3 under the TTFT budget and inside +# it. Print the verdict rather than leaving it in the log. +grep -aE -m5 "io_mode|selfcheck" "$LOG" || true diff --git a/examples/glm5.2_gfx942/launch/launch_prefill.sh b/examples/glm5.2_gfx942/launch/launch_prefill.sh new file mode 100644 index 00000000..90457ad2 --- /dev/null +++ b/examples/glm5.2_gfx942/launch/launch_prefill.sh @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +# Launch the SGLang prefill leg on the prefill node. +# Run INSIDE the engine container: bash launch/launch_prefill.sh +set -euo pipefail +HERE="$(cd "$(dirname "$0")/.." && pwd)" +source "$HERE/env.sh" + +LOG="${LOG:-$LOG_DIR/prefill.log}" +HOST_IP="${HOST_IP:-$PREFILL_IP}" +PORT="${PORT:-$PREFILL_PORT}" + +# kv events feed the router's cache-affinity view, and this is the leg that owns +# the prefix cache the policy steers by. Without them kv-aware still answers but +# routes on load alone. +KV_EVENT_ARGS=(--enable-kv-events --kv-events on --kv-event-transport zmq) + +MTP_ARGS=() +if [[ "$MTP" != "0" ]]; then + MTP_ARGS=(--speculative-algorithm EAGLE --speculative-num-steps "$MTP_STEPS" + --speculative-eagle-topk "$MTP_TOPK" --speculative-num-draft-tokens "$MTP_DRAFT_TOKENS" + --json-model-override-args '{"index_share_for_mtp_iteration":false}') +fi + +METRICS_ARGS=() +if [[ "$ENGINE_METRICS" == "1" ]]; then + METRICS_ARGS=(--enable-metrics --enable-metrics-for-all-schedulers) +fi + +# --infera-kvd-socket is the whole seam for the offload tier: infera probes the +# daemon, refuses to start if it is down, and appends SGLang's +# --enable-hierarchical-cache --hicache-storage-backend dynamic plus the backend +# module path itself. Do not hand-write those. --hicache-size caps SGLang's own +# host tier, which is otherwise derived from the GPU pool and asks for hundreds +# of GB per rank. KVD=0 (the default) leaves the GPU radix tree as the only +# cache, which is the measured-best shape here -- see env.sh. +KVD_ARGS=() +if [[ "$KVD" == "1" ]]; then + [[ -S "$KVD_SOCKET" ]] || { echo "[prefill] no kvd at $KVD_SOCKET; run launch/launch_kvd.sh first" >&2; exit 1; } + KVD_ARGS=(--infera-kvd-socket "$KVD_SOCKET" --hicache-size "$HICACHE_SIZE") +fi + +export HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}" +export CUDA_VISIBLE_DEVICES="$HIP_VISIBLE_DEVICES" +export SGLANG_HOST_IP="$HOST_IP" HOST_IP +export SGLANG_DSA_TRITON_PREFILL=1 SAFETENSORS_FAST_GPU=1 +export HSA_NO_SCRATCH_RECLAIM=1 SGLANG_USE_AITER="${SGLANG_USE_AITER:-1}" +export MC_GID_INDEX +# Cold start is weights plus graph capture; the defaults time out under it. +export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT="${SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT:-3600}" +export INFERA_ENGINE_READY_TIMEOUT="${INFERA_ENGINE_READY_TIMEOUT:-5400}" + +pkill -f "(infera.engine.sglang|sglang.launch_server) .*--port ${PORT}( |$)" 2>/dev/null || true +sleep 5 + +nohup python3 -m infera.engine.sglang \ + --model-path "$MODEL" --host 0.0.0.0 --port "$PORT" --advertise-host "$HOST_IP" \ + --etcd-endpoint "$ETCD_ENDPOINT" --discovery-backend etcd \ + --request-transport http "${KV_EVENT_ARGS[@]}" \ + --tp-size "$TP" --dp-size "$DP" --enable-dp-attention \ + --trust-remote-code --kv-cache-dtype fp8_e4m3 \ + --reasoning-parser glm45 --tool-call-parser glm47 \ + --dsa-prefill-backend tilelang --dsa-decode-backend tilelang \ + --mem-fraction-static "$MEM_FRAC" --max-running-requests "$MAX_RUNNING" \ + --chunked-prefill-size "$CHUNK" --watchdog-timeout 1200 \ + --disable-custom-all-reduce --enable-cache-report \ + "${MTP_ARGS[@]}" "${KVD_ARGS[@]}" "${METRICS_ARGS[@]}" ${EXTRA_ARGS:-} \ + --weight-loader-prefetch-checkpoints \ + --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 32}' \ + --disaggregation-mode prefill --disaggregation-bootstrap-port "$BOOTSTRAP_PORT" \ + --disaggregation-transfer-backend mooncake --disaggregation-ib-device "$IB_DEVICE" \ + > "$LOG" 2>&1 & + +echo "[prefill] loading on ${HOST_IP}:${PORT}, TP=$TP DP=$DP MTP=$MTP($MTP_STEPS/$MTP_TOPK/$MTP_DRAFT_TOKENS)" +echo "[prefill] chunk=$CHUNK ($((CHUNK / DP))/rank) rail=$IB_DEVICE metrics=$ENGINE_METRICS kvd=$KVD; log=$LOG" +echo "[prefill] cold start can take 15-25 min; follow with: tail -f $LOG" diff --git a/examples/glm5.2_gfx942/launch/launch_router.sh b/examples/glm5.2_gfx942/launch/launch_router.sh new file mode 100644 index 00000000..28a4e51c --- /dev/null +++ b/examples/glm5.2_gfx942/launch/launch_router.sh @@ -0,0 +1,58 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +# Launch the Infera kv-aware router on the prefill node. +# Run INSIDE the engine container: bash launch/launch_router.sh +# +# The router is a separate process that rediscovers both legs from etcd, so it can +# be restarted on its own -- the engines never move. +set -euo pipefail +HERE="$(cd "$(dirname "$0")/.." && pwd)" +source "$HERE/env.sh" + +LOG="${LOG:-$LOG_DIR/router.log}" + +# `rust` execs the infera-router binary over this python process. infera.server +# checks the supported subset before the exec and fails with a pointer to +# --router-backend python; everything passed here is inside it. One thing worth +# naming: the rust data plane ignores --kv-event-transport and always subscribes +# over ZMQ, which is the transport configured here anyway. +if [[ "$ROUTER_BACKEND" == "rust" ]]; then + command -v infera-router >/dev/null 2>&1 || [[ -n "${INFERA_ROUTER_BIN:-}" ]] \ + || { echo "[router] ROUTER_BACKEND=rust but no infera-router on PATH." >&2 + echo " This image builds it to /usr/local/bin; set" >&2 + echo " INFERA_ROUTER_BIN if yours does not." >&2; exit 1; } +fi + +# After the exec the process is named infera-router and no longer matches the +# python pattern, so kill both. +pkill -f "infera.server .*--port ${ROUTER_PORT}" 2>/dev/null || true +pkill -f "infera-router.*--port ${ROUTER_PORT}" 2>/dev/null || true +sleep 1 + +nohup python3 -m infera.server \ + --host 0.0.0.0 --port "$ROUTER_PORT" --router-backend "$ROUTER_BACKEND" \ + --etcd-endpoint "$ETCD_ENDPOINT" --router-tokenizer-path "$MODEL" \ + --discovery-backend etcd --request-transport http --kv-event-transport zmq \ + --router-policy "$ROUTER_POLICY" \ + --kv-prefill-overlap-weight 20.0 --kv-decode-overlap-weight 2.0 \ + > "$LOG" 2>&1 & + +for _ in $(seq 60); do + sleep 1 + curl -sf "http://127.0.0.1:${ROUTER_PORT}/health" >/dev/null && break +done +curl -sf "http://127.0.0.1:${ROUTER_PORT}/health" >/dev/null \ + || { echo "[router] did not come up"; tail -40 "$LOG"; exit 1; } + +# kv-aware needs the tokenizer to compute block hashes. Without it the router does +# not fail -- it warns once and routes on load alone, which looks like a healthy +# router that has quietly stopped doing the thing you deployed it for. +if [[ "$ROUTER_POLICY" == "kv-aware" ]] \ + && grep -qa "degrades to pure load balancing" "$LOG" 2>/dev/null; then + echo "[router] kv-aware degraded to load balancing -- tokenizer not loaded from $MODEL" >&2 + grep -a "degrades to pure load balancing" "$LOG" >&2 + exit 1 +fi + +echo "[router] ${ROUTER_POLICY}/${ROUTER_BACKEND} up on ${PREFILL_IP}:${ROUTER_PORT}; log=$LOG" diff --git a/examples/glm5.2_gfx942/preflight_rdma.sh b/examples/glm5.2_gfx942/preflight_rdma.sh new file mode 100644 index 00000000..c380aee6 --- /dev/null +++ b/examples/glm5.2_gfx942/preflight_rdma.sh @@ -0,0 +1,35 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +# RDMA preflight. Run on the host of each node BEFORE the PD bring-up: cross-node +# PD moves the KV cache over the fabric on every request, and a container whose +# RDMA provider does not match the host driver falls back to TCP silently -- the +# pair still answers, just slower than a single node. +# +# With DUMP_PATH set and one task per node it also runs infera's cross-node +# netperf and Mooncake checks. +set -euo pipefail +HERE="$(cd "$(dirname "$0")" && pwd)" +source "$HERE/env.sh" + +echo "[preflight] active RDMA ports visible in $IMAGE (expect the node's port count, not 0):" +IONIC_ARGS=() +if [[ -e "${HOST_LIBIONIC:-/usr/lib/x86_64-linux-gnu/libionic.so}" ]]; then + IONIC_ARGS=(-v "${HOST_LIBIONIC:-/usr/lib/x86_64-linux-gnu/libionic.so}:/host-libionic/libionic.so:ro") +fi +docker run --rm --network host --device=/dev/infiniband --cap-add=IPC_LOCK \ + "${IONIC_ARGS[@]}" "$IMAGE" bash -lc "ibv_devinfo | grep -c PORT_ACTIVE" + +if [[ -z "${DUMP_PATH:-}" ]]; then + echo "[preflight] DUMP_PATH not set; skipping cross-node netperf/mooncake checks." + exit 0 +fi + +: "${SLURM_PROCID:?set SLURM_PROCID= or launch with srun}" +: "${SLURM_NNODES:?set SLURM_NNODES= or launch with srun}" +export SLURMD_NODENAME="${SLURMD_NODENAME:-$(hostname)}" + +docker run --rm --network host --device=/dev/infiniband --cap-add=IPC_LOCK \ + -e MC_GID_INDEX -e IB_DEVICE -e SLURM_PROCID -e SLURM_NNODES -e SLURMD_NODENAME \ + "${IONIC_ARGS[@]}" -v "$DUMP_PATH:$DUMP_PATH" "$IMAGE" \ + python -m infera.tools.preflight --dump-path "$DUMP_PATH" --netperf --mooncake diff --git a/examples/glm5.2_gfx942/stop.sh b/examples/glm5.2_gfx942/stop.sh new file mode 100644 index 00000000..8b31b52d --- /dev/null +++ b/examples/glm5.2_gfx942/stop.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +# Stop the router and engine processes. Run INSIDE the engine container on BOTH +# nodes. The container itself survives; remove it from the host with +# `bash host_container.sh --rm`, and etcd with `docker rm -f infera-glm52-etcd`. +set -euo pipefail + +pkill -f "infera.engine.sglang" 2>/dev/null || true +pkill -f "infera.server" 2>/dev/null || true +# ROUTER_BACKEND=rust execs this binary over the python process, so after the exec +# the "infera.server" pattern above no longer matches the router. +pkill -f "infera-router" 2>/dev/null || true +pkill -f "sglang.launch_server" 2>/dev/null || true + +# Give the engines time to release VRAM; a relaunch that races them OOMs. +for _ in $(seq 60); do + pgrep -f "infera.engine.sglang|sglang.launch_server" >/dev/null || break + sleep 2 +done +pkill -9 -f "infera.engine.sglang|sglang.launch_server" 2>/dev/null || true + +# kvd after the engines, so its cache backend does not vanish from under one that +# is still running. Only the RAM tier dies here: L3 is on disk with a journal and +# is recovered on the next start. Delete KVD_L3_DIR to discard it. +pkill -f "infera.kvd" 2>/dev/null || true + +echo "[stop] remaining engine procs: $(pgrep -cf 'infera.engine.sglang|sglang.launch_server' || true)" +rocm-smi --showmemuse 2>/dev/null | grep "VRAM%" || true diff --git a/examples/glm5.2_gfx942/verify.sh b/examples/glm5.2_gfx942/verify.sh new file mode 100644 index 00000000..57d5916b --- /dev/null +++ b/examples/glm5.2_gfx942/verify.sh @@ -0,0 +1,200 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +# Verify the deployment before spending a benchmark on it. Run INSIDE the engine +# container on the PREFILL node: bash verify.sh +# +# Five assertions and one readout, each aimed at a failure this stack produces +# WITHOUT returning an error: a leg that never registered, a KV hand-off that +# yields fluent nonsense, kv-aware silently degraded to load balancing, MTP +# silently dropped, and -- only when KVD=1 -- an offload tier that stores nothing. +# The last prints the decode leg's RDMA lines to read by eye. Exits non-zero if +# any assertion fails. +set -euo pipefail +HERE="$(cd "$(dirname "$0")" && pwd)" +source "$HERE/env.sh" + +SERVER="${SERVER:-$ROUTER_URL}" +FAILED=0 +pass() { echo " PASS: $*"; } +fail() { echo " FAIL: $*"; FAILED=1; } + +# ---------------------------------------------------------------- 1. workers +echo "== workers (expect 1 prefill + 1 decode) ==" +WORKERS="$(curl -sf -m 10 "$SERVER/v1/workers" || true)" +if [[ -z "$WORKERS" ]]; then + fail "$SERVER/v1/workers did not answer -- is the router up? ($LOG_DIR/router.log)" +else + echo "$WORKERS" | python3 -m json.tool 2>/dev/null || echo "$WORKERS" + # {"workers": [{"disagg_mode": "prefill", "status": "active", ...}, ...]} + COUNTS="$(echo "$WORKERS" | python3 -c ' +import json, sys +try: + items = json.load(sys.stdin)["workers"] + modes = [w.get("disagg_mode") for w in items if w.get("status") == "active"] +except Exception: + modes = [] +print(modes.count("prefill"), modes.count("decode")) +' 2>/dev/null || echo "0 0")" + read -r N_PREFILL N_DECODE <<<"$COUNTS" + if (( N_PREFILL >= 1 && N_DECODE >= 1 )); then + pass "$N_PREFILL prefill, $N_DECODE decode registered" + else + fail "$N_PREFILL prefill, $N_DECODE decode -- a leg has not registered in etcd yet" + fi +fi + +# ------------------------------------------------------------ 2. correctness +# A broken KV hand-off does not return an HTTP error: the decode leg reads a +# corrupt or empty prefix and produces fluent text that has nothing to do with the +# prompt. So assert an answer that is only reachable if the prefix survived the +# transfer, rather than just checking for HTTP 200. +# +# The filler pads the prompt past one 64-token router block: a short question +# hashes to zero blocks, which would make check 3 report "not steering" on a +# perfectly healthy router. The prompt is fixed rather than overridable because +# the check below asserts its answer. +echo; echo "== correctness (127 * 31 = 3937, temperature 0) ==" +FILLER="$(python3 -c 'print(("The quick brown fox jumps over the lazy dog while the engine warms its caches. " * 40).strip())')" +PROMPT="$FILLER +What is 127 * 31? Answer with the number only." +BODY="$(MODEL="$MODEL" PROMPT="$PROMPT" python3 -c ' +import json, os +print(json.dumps({ + "model": os.environ["MODEL"], + "messages": [{"role": "user", "content": os.environ["PROMPT"]}], + "max_tokens": 64, "temperature": 0, + "chat_template_kwargs": {"enable_thinking": False}, +}))')" + +REPLY="$(curl -sf -m 300 "$SERVER/v1/chat/completions" \ + -H 'Content-Type: application/json' -d "$BODY" || true)" +if [[ -z "$REPLY" ]]; then + fail "no completion -- check $LOG_DIR/{router,prefill}.log and the decode node's log" +else + # GLM routes text through content or reasoning_content depending on the parser; + # read both so the check does not depend on which one this build populates. + TEXT="$(echo "$REPLY" | python3 -c ' +import json, sys +msg = json.load(sys.stdin)["choices"][0]["message"] +print(" ".join(filter(None, (msg.get("content"), msg.get("reasoning_content"))))) +' 2>/dev/null || true)" + echo " reply: ${TEXT:0:200}" + if [[ -z "$TEXT" ]]; then + fail "empty completion body" + elif [[ "$TEXT" == *3937* ]]; then + pass "correct answer through the PD pair" + else + fail "wrong answer -- fluent output with a bad prefix means the KV hand-off is broken" + fi +fi + +# ------------------------------------------------------- 3. kv-aware steering +# kv-aware fails soft: with no block hashes the router still answers, still looks +# healthy, and routes on load alone. Prove it is hashing before trusting a number. +echo; echo "== router: $ROUTER_POLICY / $ROUTER_BACKEND ==" +if [[ "$ROUTER_POLICY" != "kv-aware" ]]; then + echo " (skipped: policy is $ROUTER_POLICY)" +elif [[ ! -f "$LOG_DIR/router.log" ]]; then + fail "$LOG_DIR/router.log is missing -- run this on the prefill node" +elif grep -qa "degrades to pure load balancing" "$LOG_DIR/router.log"; then + fail "kv-aware degraded to load balancing -- tokenizer not loaded from $MODEL" +else + # Both backends log request_blocks= per pick, but the rust one wraps every + # field in ANSI colour codes and capitalises the role, so match on the number + # rather than on role=prefill. Decode picks always report 0; a non-zero anywhere + # means the tokenizer and block hasher are live. + BLOCKS="$(sed 's/\x1b\[[0-9;]*m//g' "$LOG_DIR/router.log" \ + | grep -aoE 'request_blocks=[0-9]+' | cut -d= -f2 | sort -n | tail -1 || true)" + if [[ -z "$BLOCKS" ]]; then + fail "no pick logged with request_blocks -- the request above should have made one" + elif (( BLOCKS > 0 )); then + pass "$BLOCKS blocks hashed on the largest pick, so the block hasher is live" + else + fail "every pick reports request_blocks=0 on a prompt padded past one block -- routing on load, not cache overlap" + fi +fi + +# ------------------------------------------------------------------- 4. MTP +# Speculative decoding is silently dropped if the two legs disagree on its shape. +# The acceptance series only exists on the decode leg, and only with metrics on. +echo; echo "== MTP acceptance (decode /metrics) ==" +if [[ "$MTP" == "0" ]]; then + echo " (skipped: MTP=0)" +else + METRICS="$(curl -sf -m 10 "$DECODE_URL/metrics" 2>/dev/null || true)" + if [[ -z "$METRICS" ]]; then + fail "$DECODE_URL/metrics does not answer -- relaunch the decode leg with ENGINE_METRICS=1" + elif grep -q "^sglang:spec_accept_length" <<<"$METRICS"; then + # -m3 rather than `| head -3`: head closing the pipe early SIGPIPEs grep, and + # under pipefail that failure would take the script down here. + grep -m3 "^sglang:spec_accept_length" <<<"$METRICS" | sed 's/^/ /' + # The series existing is the assertion; the values are 0 until those ranks + # have decoded, which a freshly launched fleet legitimately has not. + pass "MTP $MTP_STEPS/$MTP_TOPK/$MTP_DRAFT_TOKENS accepted by the decode leg" + else + fail "/metrics carries no sglang:spec_accept_length -- both legs must agree on the MTP shape" + fi +fi + +# --------------------------------------------------------------------- 5. kvd +# Skipped unless the offload tier is on; KVD=0 is the default (see env.sh). When +# on, it fails silently in both directions: a hicache that never wired up writes +# nothing at all, and a page too large for its tablespace slot is REJECTED rather +# than split, which leaves L3 empty while every other check here stays green. +echo; echo "== kvd offload tier (KVD=$KVD) ==" +if [[ "$KVD" != "1" ]]; then + echo " (skipped: the GPU radix tree is the only cache)" +else + sets_total() { python3 -m infera.kvd.statctl --socket "$KVD_SOCKET" \ + | python3 -c 'import json, sys; print(json.load(sys.stdin)["sets_total"])'; } + BEFORE="$(sets_total 2>/dev/null || true)" + if [[ -z "$BEFORE" ]]; then + fail "no daemon answering on $KVD_SOCKET -- run launch/launch_kvd.sh, then relaunch prefill" + else + # hicache writes back a page at a time, so a prompt shorter than one page + # produces no kvd traffic and would prove nothing; this one fills several. + # The nonce matters as much: the same text on a second run is already stored, + # writes nothing, and would fail this check on a healthy tier. + KVD_PROMPT="run $(date +%s%N). $(python3 -c "print('The quick brown fox jumps over the lazy dog. ' * 120)") Reply with OK." + BODY="$(MODEL="$MODEL" PROMPT="$KVD_PROMPT" python3 -c ' +import json, os +print(json.dumps({ + "model": os.environ["MODEL"], + "messages": [{"role": "user", "content": os.environ["PROMPT"]}], + "max_tokens": 8, "temperature": 0, + "chat_template_kwargs": {"enable_thinking": False}, +}))')" + curl -sf -m 300 "$SERVER/v1/chat/completions" \ + -H 'Content-Type: application/json' -d "$BODY" >/dev/null || true + sleep 5 # write-back to storage is asynchronous w.r.t. the response + AFTER="$(sets_total 2>/dev/null || echo "$BEFORE")" + + if grep -aq "value_exceeds" "$LOG_DIR/kvd.log" 2>/dev/null; then + fail "kvd rejected pages that outgrew their slot -- raise KVD_TABLESPACE_POOLS" + grep -a "value_exceeds" "$LOG_DIR/kvd.log" | tail -2 | sed 's/^/ /' + elif (( AFTER > BEFORE )); then + pass "kvd took $((AFTER - BEFORE)) writes from the prefill engine" + else + fail "no kvd writes ($BEFORE -> $AFTER) -- prefill should log a hicache storage backend" + grep -aiE "hierarchical|hicache|kvd" "$LOG_DIR/prefill.log" 2>/dev/null | tail -5 | sed 's/^/ /' || true + fi + fi +fi + +# ------------------------------------------------------------ 6. RDMA hand-off +# Informational: decode.log lives on the other node unless LOG_DIR is shared. +echo; echo "== RDMA hand-off (decode log, rail=$IB_DEVICE) ==" +if [[ -f "$LOG_DIR/decode.log" ]]; then + grep -aE "GID index|installTransport|mooncake" "$LOG_DIR/decode.log" | tail -5 | sed 's/^/ /' \ + || echo " (no hand-off lines yet)" +else + echo " (decode.log is on $DECODE_NODE; check it there for 'mooncake' transport lines)" +fi + +echo +if (( FAILED )); then + echo "VERIFY: FAILED -- do not benchmark this deployment yet." + exit 1 +fi +echo "VERIFY: all checks passed." From 50d8dccd35f6b897305fdf1081940ca059046852 Mon Sep 17 00:00:00 2001 From: xiaobochen-amd Date: Thu, 6 Aug 2026 16:05:44 +0000 Subject: [PATCH 43/88] fix(tests): finish separating a failed SLURM query from an empty result Three sites left over from the reservation and reclaim fixes, all reading a query that did not answer as a result that came back empty. _cancel_dispatched. Two of them, and this one runs from the INT/TERM trap, so it fires exactly when the scheduler is most likely to be unhappy. Listing the run's jobs returned as though there were nothing to cancel; confirming the cancel read a failed squeue as "gone". It now says what happened and hands off to the workflow's reclaim step, and only an answered query may confirm. That confirmation leans on Spur behaviour worth writing down: a job that is gone still exits 0 with no output, so empty-and-successful is real proof. Stock SLURM errors on an invalid id instead, so there the loop never confirms and the new warning is a false alarm -- noted in the comment, and the workflow's reclaim step covers both. _spill_inflight returned 0 on a failed count, which reads as "nothing borrowed" and authorises another spill -- borrowing harder from the open partition at the one moment the scheduler is already struggling. It returns non-zero now and the caller queues on the reservation instead. _hold_pair collapsed every outcome into 1, so a refused sbatch, a hold that never started, and genuinely losing the pair to another engine all surfaced as "lost the node-hold race N times". That sent triage looking for contention when the scheduler was the problem -- it is why the review could not use that line as evidence of a full pool. 2 now means SLURM never placed the hold, 1 stays for a real race, and the caller words the two differently. The caller uses `_hold_pair ...; hold_rc=$?` rather than `if ! _hold_pair`, where $? is the negation's status and never the function's. Verified, and commented so it does not come back. Verified with a stubbed squeue/sbatch: a failing squeue makes _cancel_dispatched report and return non-zero instead of claiming success, and makes _spill_inflight return non-zero so no spill is authorised; a gone job (exit 0, empty) confirms the cancel; a permanently refused sbatch returns 2. A healthy host still reaches mode=resv. Co-authored-by: Cursor Signed-off-by: xiaobochen-amd --- tests/run_tests.sh | 47 ++++++++++++++++++++++++++++++++++------------ 1 file changed, 35 insertions(+), 12 deletions(-) diff --git a/tests/run_tests.sh b/tests/run_tests.sh index 9e98900b..c2bb29f5 100755 --- a/tests/run_tests.sh +++ b/tests/run_tests.sh @@ -100,16 +100,21 @@ _cleanup_scratch() { # has to scancel it: job id from $_CUR_DISPATCH_OUT, else from the job tag. _CUR_DISPATCH_OUT="" _cancel_dispatched() { - local jids="" i suf csv + local jids="" i suf csv left if [ -n "$_CUR_DISPATCH_OUT" ] && [ -f "$_CUR_DISPATCH_OUT" ]; then jids=$(grep -oE 'srun: job [0-9]+' "$_CUR_DISPATCH_OUT" 2>/dev/null \ | grep -oE '[0-9]+' | sort -u | tr '\n' ' ') fi if [ -z "$jids" ] && [ -n "${INFERA_E2E_JOB_TAG:-}" ]; then suf="-${INFERA_E2E_JOB_TAG}" - jids=$(squeue -h -u "$(id -un)" -o '%i %j' 2>/dev/null \ + # A failed lookup is not an empty queue. Say so and leave it to ci.yml's + # reclaim step rather than returning as if there were nothing to cancel. + if ! jids=$(squeue -h -u "$(id -un)" -o '%i %j' 2>/dev/null \ | awk -v suf="$suf" '$2 ~ /^infera-ci-/ && substr($2, length($2)-length(suf)+1)==suf {print $1}' \ - | tr '\n' ' ') + | tr '\n' ' '); then + echo "[cleanup] squeue failed — cannot list this run's jobs, leaving them to the workflow's reclaim step" >&2 + return 1 + fi fi [ -n "$jids" ] || return 0 echo "[cleanup] cancelling dispatched SLURM job(s): $jids" >&2 @@ -118,8 +123,13 @@ _cancel_dispatched() { for i in 1 2 3 4 5; do scancel $jids >/dev/null 2>&1 || true sleep 2 - [ -z "$(squeue -h -j "$csv" -o '%i' 2>/dev/null)" ] && return 0 + # Only a query that answered may confirm the cancel: on Spur a gone job still + # exits 0 with no output. (Stock SLURM errors on an invalid id, so there this + # never confirms and the warning below is a false alarm -- the workflow's + # reclaim step is the backstop either way.) + left=$(squeue -h -j "$csv" -o '%i' 2>/dev/null) && [ -z "$left" ] && return 0 done + echo "[cleanup] could not confirm the cancel of $jids" >&2 } # Nodes the running PD-disagg attempt placed containers on. A killed run skips # pytest's teardown, so without this a cancel leaves prefill+decode on the GPUs. @@ -291,6 +301,9 @@ _rival_holder() { # Hold both PD nodes' GPUs for the whole run: disagg's per-step sruns leave them # idle in between, so SLURM would hand one out and the fixed ports (etcd 2379, # router 8000, ...) collide. Our own no-gres steps co-schedule. Sets _HOLDER_JID. +# 0 = held, 1 = another holder won the pair, 2 = SLURM never placed the hold. +# The caller reports 1 and 2 differently: they used to read alike, so a refused +# sbatch was announced as a lost race and pointed triage away from the scheduler. _hold_pair() { local pair="$1" script="$SCRATCH/hold.sh" jid st rs waited qos=() i other # A real script file, not --wrap: on Spur --wrap always NODE_FAILs at -N2. @@ -323,7 +336,7 @@ _hold_pair() { scancel "$jid" >/dev/null 2>&1 echo "[e2e disagg] hold attempt $i on $pair not started (${st:-?}/${rs:-?}) — retrying" >&2 done - return 1 + return 2 } # One renderD* per GPU; PCI vendor 0x1002 == AMD. _amd_gpu_count() { @@ -350,8 +363,12 @@ _reservation_free() { echo "$free" } # Caps borrowed nodes at INFERA_E2E_SPILL_MAX; concurrent dispatchers can race it. +# Non-zero if the query failed, so the caller does not read that as "none in +# flight" and borrow past the cap exactly when the scheduler is already unwell. _spill_inflight() { - squeue -h -u "$(id -un)" -o '%j' 2>/dev/null | grep -c -- 'spill' || true + local out + out=$(squeue -h -u "$(id -un)" -o '%j' 2>/dev/null) || return 1 + printf '%s\n' "$out" | grep -c -- 'spill' || true } # Report why the dispatch is still queued (a waiting job prints NOTHING, so a CI @@ -431,8 +448,8 @@ _dispatch_slurm() { # sent a whole run to the open partition on 2026-08-06. resv=(--reservation="$INFERA_E2E_RESERVATION"); mode="resv" else - inflight=$(_spill_inflight) - if [ "$smax" -gt 0 ] && [ "$inflight" -lt "$smax" ]; then + # A failed count must not authorise a spill: queue on the reservation. + if inflight=$(_spill_inflight) && [ "$smax" -gt 0 ] && [ "$inflight" -lt "$smax" ]; then # spill marker sits before the run_id-engine suffix so ci.yml reclaim matches. jobname="infera-ci-${label}-spill${INFERA_E2E_JOB_TAG:+-$INFERA_E2E_JOB_TAG}" mode="spill($((inflight + 1))/$smax)" @@ -669,7 +686,7 @@ run_e2e_disagg() { local rc=0 e prc out="$SCRATCH/.e2e-disag.out" local max_attempts=3 attempt exclude n1 n2 nodes ok - local races max_races="${INFERA_E2E_HOLD_RACE_MAX:-10}" + local races max_races="${INFERA_E2E_HOLD_RACE_MAX:-10}" hold_rc for e in "${engines[@]}"; do echo "----- e2e disagg — tests/e2e/pd_disag/$e -----" attempt=0; ok=0; exclude=""; races=0 @@ -690,13 +707,19 @@ run_e2e_disagg() { # Losing the race is not a node fault, so the pair must NOT join $exclude: # with a small pool the engine would exclude every node and then starve on # an idle cluster. Bounded so a pathological loser fails loudly instead. - if ! _hold_pair "$n1,$n2"; then + # Not `if ! _hold_pair`: inside that, $? is the negation's, not the call's. + _hold_pair "$n1,$n2"; hold_rc=$? + if [ "$hold_rc" -ne 0 ]; then races=$((races + 1)) if [ "$races" -ge "$max_races" ]; then - echo "[e2e disagg] lost the node-hold race $races times — giving up on $e" >&2 + if [ "$hold_rc" -eq 2 ]; then + echo "[e2e disagg] SLURM never placed a node hold in $races attempts — giving up on $e" >&2 + else + echo "[e2e disagg] lost the node-hold race $races times — giving up on $e" >&2 + fi break fi - echo "[e2e disagg] could not hold $n1,$n2 (race $races/$max_races) — re-picking in 30s" >&2 + echo "[e2e disagg] could not hold $n1,$n2 (attempt $races/$max_races) — re-picking in 30s" >&2 attempt=$((attempt - 1)); sleep 30; continue fi races=0 From 2ddb8b695733f47d401643ea8dafc90a8436d30e Mon Sep 17 00:00:00 2001 From: liyingli Date: Thu, 6 Aug 2026 11:25:03 +0000 Subject: [PATCH 44/88] fix(recipes): carry the measured chunk and MTP shape into the gfx942 manifest The manifest landed with the docker recipe's values as they stood before its tuning sweep finished, and claimed they were identical. Two are not, and both are the sweep's largest levers. --chunked-prefill-size 131072 is 16,384 per rank under DP-attention, which splits the value dp_size ways. 8192 -- 1,024 per rank -- measured 23.8% lower duration and 34.4% lower TTFT on the trace both recipes are tuned against. The aggregate-vs-per-rank division is also why the number looks alarming either way, so the README now shows how to read back what actually took effect from /get_server_info rather than trusting the engine's "adjusted from ... to ..." line. EAGLE 3/1/4 -> 5/1/6, worth 8.5% duration and 13.8% TPOT. Each extra draft step costs 6.53 ms here, so the shape pays only while acceptance outruns it: 5 steps accept 4.64 against a 4.00 break-even, and 7/1/8 both misses its break-even and runs prefill out of activation memory. Applied to both legs -- SGLang rejects a disaggregated pair whose speculative configs differ. Section 5's kvd A/B no longer says the measurement is pending: KVD=1 ran 12.0% slower with gets_total = 0, 100.8 GB written to L3 and not one page read back, because the 54 GB/rank device pool had already taken ~100% of the reuse that trace offered. That is a statement about the workload, and the section now says so along with what to check before paying for the tier. Section 6 also stops claiming every router flag is identical; the manifest runs --router-backend python because the Rust binary supports only etcd discovery while the operator's is kubernetes, which is a substrate translation like the rest and now sits with them. The intro links the docker recipe it is a translation of, now that examples/glm5.2_gfx942/ gives it somewhere to point. Co-authored-by: Cursor Signed-off-by: liyingli --- examples/recipes/glm5.2-fp8-gfx942/README.md | 29 ++++++++++++++----- .../disaggregated-kvd/deploy.yaml | 12 ++++---- 2 files changed, 28 insertions(+), 13 deletions(-) diff --git a/examples/recipes/glm5.2-fp8-gfx942/README.md b/examples/recipes/glm5.2-fp8-gfx942/README.md index 5d433c68..d4d07bad 100644 --- a/examples/recipes/glm5.2-fp8-gfx942/README.md +++ b/examples/recipes/glm5.2-fp8-gfx942/README.md @@ -5,9 +5,10 @@ disaggregation over Mooncake RDMA, DP-attention, MTP speculative decoding, kv-aw routing, and KV offload to host RAM + node-local NVMe through `infera-kvd`. Every flag here is lifted from a `docker` + shell deployment of the same topology, -validated on 2 × MI300X and referred to throughout as **the docker recipe**. Nothing -was retuned for Kubernetes: §6 lists every difference and why the substrate forced -it, and there are no others. +validated on 2 × MI300X and referred to throughout as **the docker recipe**: +[`examples/glm5.2_gfx942/`](../../glm5.2_gfx942/README.md), which also carries the +bring-up, verification and benchmark scripts. Nothing was retuned for Kubernetes: +§6 lists every difference and why the substrate forced it, and there are no others. | Combo | Serving | KV cache | Manifest | |---|---|---|---| @@ -229,13 +230,27 @@ command, drop the `kvd-sock` / `kvd-l3` volumes, and drop the `kvd-sock` `volumeMount` from the prefill `main` container — miss that last one and the Pod is rejected for referencing a volume that no longer exists. Worth running only once the deployment is above its pressure point — below it the 54 GB device pool per rank -answers everything and both arms are identical. That is what the docker recipe's -agentic trace showed; **TODO — kvd's share of the hits is being re-measured, so no -figure is quoted here yet.** +answers everything and both arms are identical. + +**That is what the docker recipe's agentic trace measured, and the result was +negative:** on a 32-conversation / 225-turn trace at `CONC=16`, `KVD=1` ran +**12.0% slower** than `KVD=0` and served **`gets_total = 0`** — 100.8 GB written +to L3 and not one page read back. Nothing was wrong with the offload path; the +trace simply had no misses left for it, its scorer efficiency against the +achievable ideal already sitting at ~100% on the GPU pool alone. Run this +manifest for a workload whose working set outgrows 54 GB/rank, and run the +`KVD=0` variant above to confirm yours does before paying for the tier. ## 6. What changed from the docker recipe, and why -Every engine, router and kvd flag is identical. These are substrate translations: +Every engine and kvd flag is identical. The one router flag that differs is +`--router-backend`: the docker recipe defaults to `rust`, which this manifest +cannot use, because the Rust binary supports only `--discovery-backend etcd` +while the operator's backend is `kubernetes` (first row below). The two make +identical routing decisions request for request; `python` reaches them ~27% +slower end to end, which is the price of the substrate here. + +The rest are substrate translations: | `docker` form | Kubernetes form | Why | |---|---|---| diff --git a/examples/recipes/glm5.2-fp8-gfx942/disaggregated-kvd/deploy.yaml b/examples/recipes/glm5.2-fp8-gfx942/disaggregated-kvd/deploy.yaml index bd1b6a9a..2b2e643d 100644 --- a/examples/recipes/glm5.2-fp8-gfx942/disaggregated-kvd/deploy.yaml +++ b/examples/recipes/glm5.2-fp8-gfx942/disaggregated-kvd/deploy.yaml @@ -170,10 +170,10 @@ spec: "--reasoning-parser","glm45","--tool-call-parser","glm47", "--dsa-prefill-backend","tilelang","--dsa-decode-backend","tilelang", "--mem-fraction-static","0.85","--max-running-requests","128", - "--chunked-prefill-size","131072","--watchdog-timeout","1200", + "--chunked-prefill-size","8192","--watchdog-timeout","1200", "--disable-custom-all-reduce","--enable-cache-report", - "--speculative-algorithm","EAGLE","--speculative-num-steps","3", - "--speculative-eagle-topk","1","--speculative-num-draft-tokens","4", + "--speculative-algorithm","EAGLE","--speculative-num-steps","5", + "--speculative-eagle-topk","1","--speculative-num-draft-tokens","6", "--json-model-override-args",'{"index_share_for_mtp_iteration":false}', "--infera-kvd-socket","/tmp/infera-kvd/kvd.sock","--hicache-size","32", "--weight-loader-prefetch-checkpoints", @@ -260,10 +260,10 @@ spec: "--reasoning-parser","glm45","--tool-call-parser","glm47", "--dsa-prefill-backend","tilelang","--dsa-decode-backend","tilelang", "--mem-fraction-static","0.85","--max-running-requests","128", - "--chunked-prefill-size","131072","--watchdog-timeout","1200", + "--chunked-prefill-size","8192","--watchdog-timeout","1200", "--disable-custom-all-reduce","--enable-cache-report", - "--speculative-algorithm","EAGLE","--speculative-num-steps","3", - "--speculative-eagle-topk","1","--speculative-num-draft-tokens","4", + "--speculative-algorithm","EAGLE","--speculative-num-steps","5", + "--speculative-eagle-topk","1","--speculative-num-draft-tokens","6", "--json-model-override-args",'{"index_share_for_mtp_iteration":false}', "--weight-loader-prefetch-checkpoints", "--model-loader-extra-config",'{"enable_multithread_load": true, "num_threads": 32}', From d4b6b990bc48d30dd6d5e52b59ead6f4b7e9d06c Mon Sep 17 00:00:00 2001 From: xiaobochen-amd Date: Thu, 6 Aug 2026 16:38:10 +0000 Subject: [PATCH 45/88] fix(tests): say why a SLURM query failed, not just that it did Review feedback on #97: _cancel_dispatched still had squeue's stderr going to /dev/null, so the new message read "squeue failed" and stopped there. Knowing a query failed without knowing why means reproducing a scheduler blink to debug one, which is the thing these commits set out to remove. It was also inconsistent with the same series: reclaim_slurm_jobs.sh already keeps stderr, and the review that motivated all of this said to stop discarding it. Three more sites had the same gap, so all four are fixed together: _cancel_dispatched both the job listing and the post-scancel confirmation _reservation_nodes forwards scontrol's own words; callers could only say "cannot reach the scheduler", which is not actionable _spill_inflight same, for the borrowed-node count Merging stderr into the pipeline would not have worked for the listing: awk filters to job names, so the error text is exactly what gets dropped. Each of these captures the command's output first and parses the copy -- the shape _reservation_nodes already used, now applied consistently. Verified with a stubbed squeue/scontrol: a refused connection now appears in full at every one of the four sites, including the "could not confirm the cancel" line, which previously ended without a reason. A healthy host still reaches mode=resv. Co-authored-by: Cursor Signed-off-by: xiaobochen-amd --- tests/run_tests.sh | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/tests/run_tests.sh b/tests/run_tests.sh index c2bb29f5..81af2458 100755 --- a/tests/run_tests.sh +++ b/tests/run_tests.sh @@ -100,21 +100,23 @@ _cleanup_scratch() { # has to scancel it: job id from $_CUR_DISPATCH_OUT, else from the job tag. _CUR_DISPATCH_OUT="" _cancel_dispatched() { - local jids="" i suf csv left + local jids="" i suf csv left queue if [ -n "$_CUR_DISPATCH_OUT" ] && [ -f "$_CUR_DISPATCH_OUT" ]; then jids=$(grep -oE 'srun: job [0-9]+' "$_CUR_DISPATCH_OUT" 2>/dev/null \ | grep -oE '[0-9]+' | sort -u | tr '\n' ' ') fi if [ -z "$jids" ] && [ -n "${INFERA_E2E_JOB_TAG:-}" ]; then suf="-${INFERA_E2E_JOB_TAG}" - # A failed lookup is not an empty queue. Say so and leave it to ci.yml's - # reclaim step rather than returning as if there were nothing to cancel. - if ! jids=$(squeue -h -u "$(id -un)" -o '%i %j' 2>/dev/null \ - | awk -v suf="$suf" '$2 ~ /^infera-ci-/ && substr($2, length($2)-length(suf)+1)==suf {print $1}' \ - | tr '\n' ' '); then - echo "[cleanup] squeue failed — cannot list this run's jobs, leaving them to the workflow's reclaim step" >&2 + # A failed lookup is not an empty queue. Capture before parsing: awk would + # filter the error away, leaving a bare "squeue failed" that has to be + # reproduced to diagnose. Then hand off to ci.yml's reclaim step. + if ! queue=$(squeue -h -u "$(id -un)" -o '%i %j' 2>&1); then + echo "[cleanup] squeue failed, leaving this run's jobs to the workflow's reclaim step: $queue" >&2 return 1 fi + jids=$(printf '%s\n' "$queue" \ + | awk -v suf="$suf" '$2 ~ /^infera-ci-/ && substr($2, length($2)-length(suf)+1)==suf {print $1}' \ + | tr '\n' ' ') fi [ -n "$jids" ] || return 0 echo "[cleanup] cancelling dispatched SLURM job(s): $jids" >&2 @@ -127,9 +129,9 @@ _cancel_dispatched() { # exits 0 with no output. (Stock SLURM errors on an invalid id, so there this # never confirms and the warning below is a false alarm -- the workflow's # reclaim step is the backstop either way.) - left=$(squeue -h -j "$csv" -o '%i' 2>/dev/null) && [ -z "$left" ] && return 0 + left=$(squeue -h -j "$csv" -o '%i' 2>&1) && [ -z "$left" ] && return 0 done - echo "[cleanup] could not confirm the cancel of $jids" >&2 + echo "[cleanup] could not confirm the cancel of $jids: ${left:-no output}" >&2 } # Nodes the running PD-disagg attempt placed containers on. A killed run skips # pytest's teardown, so without this a cancel leaves prefill+decode on the GPUs. @@ -232,7 +234,9 @@ _have_slurm() { command -v srun >/dev/null 2>&1; } # masquerade as a failed query, and callers now act on that distinction. _reservation_nodes() { local out - out=$(scontrol show reservation "$1" 2>/dev/null) || return 1 + # Forward scontrol's own words: callers can only say "cannot reach the + # scheduler", which is not enough to act on. + out=$(scontrol show reservation "$1" 2>&1) || { printf '%s\n' "$out" >&2; return 1; } printf '%s\n' "$out" | awk -v r="ReservationName=$1" ' BEGIN{RS="";FS="\n"} $1==r { for(i=1;i<=NF;i++) if($i ~ /Nodes=/){ n=$i; sub(/.*Nodes=/,"",n); sub(/[[:space:]].*/,"",n); print n; exit } }' \ @@ -367,7 +371,7 @@ _reservation_free() { # flight" and borrow past the cap exactly when the scheduler is already unwell. _spill_inflight() { local out - out=$(squeue -h -u "$(id -un)" -o '%j' 2>/dev/null) || return 1 + out=$(squeue -h -u "$(id -un)" -o '%j' 2>&1) || { printf '%s\n' "$out" >&2; return 1; } printf '%s\n' "$out" | grep -c -- 'spill' || true } From c6b26b9b5803f673e25e49576e47f3f62ced80a0 Mon Sep 17 00:00:00 2001 From: liyingli Date: Fri, 7 Aug 2026 03:28:08 +0000 Subject: [PATCH 46/88] fix(examples): refuse an unresolved node address instead of defaulting to loopback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit env.sh fell back to 127.0.0.1 when a node name did not resolve, and DECODE_IP fell back to PREFILL_IP on top of that. Both are the wrong shape of default for a deployment whose two halves find each other only through the addresses they register. The full fallback hides itself on the node you are watching: etcd runs on the prefill node, so 127.0.0.1:2379 resolves there, the prefill leg registers, the router comes up, and nothing on that node looks wrong. The decode leg is the one that fails, on the other machine, and not until 20 minutes in -- registration happens after the weights load and the graphs capture. Setting only PREFILL_IP is the likelier trap and the worse one. DECODE_IP silently becomes the prefill node's address, the decode leg advertises it, and both legs register: the worker count verify.sh checks first is correct, and only a real request finds that nothing is listening at the address decode published. So require_ips(), called by the six scripts that dial one of these addresses and by no other. A blanket check in env.sh would have been shorter and wrong: all twelve scripts source it, but build_image.sh and host_container.sh need no address at all, and the README has the image build (§1.3) before the cluster adaptation (§2) that sets them. Failing there would trade this footgun for one in the first documented step. Co-authored-by: Cursor Signed-off-by: liyingli --- examples/glm5.2_gfx942/README.md | 11 +++++++++- examples/glm5.2_gfx942/bench.sh | 1 + examples/glm5.2_gfx942/env.sh | 22 +++++++++++++++++-- .../glm5.2_gfx942/launch/launch_decode.sh | 1 + examples/glm5.2_gfx942/launch/launch_etcd.sh | 1 + .../glm5.2_gfx942/launch/launch_prefill.sh | 1 + .../glm5.2_gfx942/launch/launch_router.sh | 1 + examples/glm5.2_gfx942/verify.sh | 1 + 8 files changed, 36 insertions(+), 3 deletions(-) diff --git a/examples/glm5.2_gfx942/README.md b/examples/glm5.2_gfx942/README.md index 59292322..3bd71620 100644 --- a/examples/glm5.2_gfx942/README.md +++ b/examples/glm5.2_gfx942/README.md @@ -117,6 +117,9 @@ If your nodes resolve by name, `PREFILL_NODE` / `DECODE_NODE` derive the IPs instead. The addresses must be the ones the peers can reach on the data network, not a management NIC. +Every script that dials one of them refuses to start until both resolve, rather +than defaulting the missing one — see gotcha 7 for what that default would cost. + ## 3. Verify the RDMA fabric Cross-node PD moves the KV cache over the fabric on every request, and a mismatch @@ -317,6 +320,12 @@ full run is exactly the 12% regression above. The remaining knobs — `KVD_RAM_B hand-off time, not at startup. 6. **`Ctrl-C` on a `tail -f` does not stop an engine.** The launch scripts run them under `nohup`; use `stop.sh`. -7. **`kvd` outlives a restart, on disk.** `stop.sh` kills the daemon after the +7. **A missing IP is refused, not defaulted.** Both legs find each other only + through the addresses they register in etcd, and a wrong one costs a full cold + start to discover: registration happens *after* the weights load, on the other + node. Setting only `PREFILL_IP` is the trap worth naming — the decode leg would + advertise the prefill node's address, both legs would register, and only a real + request would find the hole. `require_ips` in `env.sh` stops that at launch. +8. **`kvd` outlives a restart, on disk.** `stop.sh` kills the daemon after the engines, so nothing is pulled from under a live one, but L3 is journalled and is recovered on the next start. Delete `KVD_L3_DIR` to start cold. diff --git a/examples/glm5.2_gfx942/bench.sh b/examples/glm5.2_gfx942/bench.sh index 7a4138f0..30f8c35b 100644 --- a/examples/glm5.2_gfx942/bench.sh +++ b/examples/glm5.2_gfx942/bench.sh @@ -13,6 +13,7 @@ set -euo pipefail HERE="$(cd "$(dirname "$0")" && pwd)" source "$HERE/env.sh" +require_ips TAG="${TAG:-random_isl${ISL}_osl${OSL}_c${CONC}_n${NUM_PROMPTS}}" OUT="$RESULT_DIR/$TAG" diff --git a/examples/glm5.2_gfx942/env.sh b/examples/glm5.2_gfx942/env.sh index 86e9190f..62cb5c5b 100644 --- a/examples/glm5.2_gfx942/env.sh +++ b/examples/glm5.2_gfx942/env.sh @@ -17,8 +17,6 @@ export PREFILL_NODE="${PREFILL_NODE:-node-0}" export DECODE_NODE="${DECODE_NODE:-node-1}" export PREFILL_IP="${PREFILL_IP:-$(getent ahostsv4 "$PREFILL_NODE" 2>/dev/null | awk 'NR==1{print $1}')}" export DECODE_IP="${DECODE_IP:-$(getent ahostsv4 "$DECODE_NODE" 2>/dev/null | awk 'NR==1{print $1}')}" -: "${PREFILL_IP:=127.0.0.1}" -: "${DECODE_IP:=$PREFILL_IP}" export ETCD_ENDPOINT="${ETCD_ENDPOINT:-${PREFILL_IP}:2379}" export ROUTER_PORT="${ROUTER_PORT:-8000}" @@ -30,6 +28,26 @@ export ROUTER_URL="${ROUTER_URL:-http://${PREFILL_IP}:${ROUTER_PORT}}" export PREFILL_URL="${PREFILL_URL:-http://${PREFILL_IP}:${PREFILL_PORT}}" export DECODE_URL="${DECODE_URL:-http://${DECODE_IP}:${DECODE_PORT}}" +# Called by every script that dials one of the addresses above, and by no other -- +# build_image.sh and host_container.sh source this file and need no IP at all. +# +# The tempting default, loopback for a node that does not resolve, is the worst +# kind of wrong here: etcd runs on the prefill node, so THAT leg registers happily +# and its whole node looks healthy while the decode leg finds nothing listening on +# its own loopback -- 20 minutes later, since registration comes after the weights +# load. A half-set pair is worse still: DECODE_IP falling back to PREFILL_IP has +# the decode leg advertise the prefill node's address, and then BOTH legs register +# and only a real request finds the hole. +require_ips() { + local bad=0 + [[ -n "$PREFILL_IP" ]] || { echo "[env] PREFILL_IP unset and '$PREFILL_NODE' does not resolve" >&2; bad=1; } + [[ -n "$DECODE_IP" ]] || { echo "[env] DECODE_IP unset and '$DECODE_NODE' does not resolve" >&2; bad=1; } + (( bad == 0 )) || { + echo "[env] export both, or point PREFILL_NODE/DECODE_NODE at names that resolve" >&2 + exit 1 + } +} + # --- image / container ------------------------------------------------------ # Built by build_image.sh straight from deploy/docker/Dockerfile.sglang.gfx942. # Do not layer runtime SGLang patches on it. diff --git a/examples/glm5.2_gfx942/launch/launch_decode.sh b/examples/glm5.2_gfx942/launch/launch_decode.sh index c43da278..40c76b6a 100644 --- a/examples/glm5.2_gfx942/launch/launch_decode.sh +++ b/examples/glm5.2_gfx942/launch/launch_decode.sh @@ -6,6 +6,7 @@ set -euo pipefail HERE="$(cd "$(dirname "$0")/.." && pwd)" source "$HERE/env.sh" +require_ips LOG="${LOG:-$LOG_DIR/decode.log}" HOST_IP="${HOST_IP:-$DECODE_IP}" diff --git a/examples/glm5.2_gfx942/launch/launch_etcd.sh b/examples/glm5.2_gfx942/launch/launch_etcd.sh index 838ff5e8..7e1d6ee8 100644 --- a/examples/glm5.2_gfx942/launch/launch_etcd.sh +++ b/examples/glm5.2_gfx942/launch/launch_etcd.sh @@ -7,6 +7,7 @@ set -euo pipefail HERE="$(cd "$(dirname "$0")/.." && pwd)" source "$HERE/env.sh" +require_ips HOST_IP="${ETCD_HOST_IP:-$PREFILL_IP}" ETCD_IMAGE="${ETCD_IMAGE:-quay.io/coreos/etcd:v3.5.14}" diff --git a/examples/glm5.2_gfx942/launch/launch_prefill.sh b/examples/glm5.2_gfx942/launch/launch_prefill.sh index 90457ad2..19c0f457 100644 --- a/examples/glm5.2_gfx942/launch/launch_prefill.sh +++ b/examples/glm5.2_gfx942/launch/launch_prefill.sh @@ -6,6 +6,7 @@ set -euo pipefail HERE="$(cd "$(dirname "$0")/.." && pwd)" source "$HERE/env.sh" +require_ips LOG="${LOG:-$LOG_DIR/prefill.log}" HOST_IP="${HOST_IP:-$PREFILL_IP}" diff --git a/examples/glm5.2_gfx942/launch/launch_router.sh b/examples/glm5.2_gfx942/launch/launch_router.sh index 28a4e51c..808728c6 100644 --- a/examples/glm5.2_gfx942/launch/launch_router.sh +++ b/examples/glm5.2_gfx942/launch/launch_router.sh @@ -9,6 +9,7 @@ set -euo pipefail HERE="$(cd "$(dirname "$0")/.." && pwd)" source "$HERE/env.sh" +require_ips LOG="${LOG:-$LOG_DIR/router.log}" diff --git a/examples/glm5.2_gfx942/verify.sh b/examples/glm5.2_gfx942/verify.sh index 57d5916b..096fcd68 100644 --- a/examples/glm5.2_gfx942/verify.sh +++ b/examples/glm5.2_gfx942/verify.sh @@ -13,6 +13,7 @@ set -euo pipefail HERE="$(cd "$(dirname "$0")" && pwd)" source "$HERE/env.sh" +require_ips SERVER="${SERVER:-$ROUTER_URL}" FAILED=0 From fb7fa66a17ccc5a3ef1f7fc40349c9afab79a8cf Mon Sep 17 00:00:00 2001 From: llying-001 Date: Thu, 6 Aug 2026 19:07:19 +0800 Subject: [PATCH 47/88] =?UTF-8?q?fix(router):=20render=20chat=20templates?= =?UTF-8?q?=20the=20way=20transformers=20does,=20or=20kv-a=E2=80=A6=20(#92?= =?UTF-8?q?)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(router): render chat templates the way transformers does, or kv-aware is blind The router reproduces the engine's prompt in order to hash it: block hashes are chained over token ids and matched against the kv-events the engine publishes from its own tokens, so the two renders have to agree byte for byte or every lookup misses. Chaining means there is no partial credit -- one byte of divergence and every block after it is wrong. The Python router gets this for free by calling tokenizer.apply_chat_template(), the same function the engine uses; the same code cannot disagree with itself. Rust has no interpreter for it and re-runs the template through minijinja, so every semantic of jinja2-plus-transformers has to be reproduced by hand. Five were not. Two of them abort the render outright rather than skewing it, and a failed render is not loud: the prompt comes back empty, the request hashes to nothing, and the policy falls back to load-only routing with every health signal green. GLM-5.2's chat path calls str.strip/lstrip/rstrip/split and its tool-call path calls dict.items(); minijinja has none of Python's data model. On the same trace that cost 0.00% predicted hits against the Python router's 83.33%. The other three only move bytes, which a chained hash punishes just as hard. transformers replaces jinja2's tojson with json.dumps(ensure_ascii=False), explicitly so HTML characters survive, and Python's separators carry a space; minijinja's builtin escapes < > & and packs them. Object key order is dropped twice over, by serde_json on parse and again by minijinja, so a tool call's arguments arrived alphabetised -- hence preserve_order on both. And transformers renders with trim_blocks and lstrip_blocks on, which minijinja defaults off; that one hides on the chat path, whose template trims explicitly with {%- -%}, and surfaces as two stray newlines in the tool-call branch, which does not. The tool-call half of this survived three rounds of benchmarking because the agentic trace has no tool calls in it, so the branch never ran. Verified by rendering the same payloads through the environment that transformers.utils.chat_template_utils._compile_jinja_template builds and diffing byte for byte: the chat path was already identical at 153 characters, the tool-call path now matches at 467. Both expectations are pinned as tests, with hermetic ones underneath for the string methods, the dict views and tojson. The payload carries a nested object, a list, a '<' and a non-ASCII character, with keys deliberately out of alphabetical order, so any one of the five regressing breaks it. Kimi's weights are not on this host so its parity is unverified, but trim_blocks is what transformers uses for every model, so enabling it can only move closer. Co-authored-by: Cursor Signed-off-by: liyingli * fix(router): stop preserve_order turning a key removal into a reorder Enabling preserve_order for the chat template also redefines serde_json::Map::remove as swap_remove, which fills the hole with the last entry. The signature and the call site are unchanged, so this lands silently. One call site is affected today, dropping stream_options from the body forwarded to a vLLM prefill leg. It is harmless -- a JSON object's key order carries no meaning to the engine parsing it, the messages array is untouched, and the block hashes are computed before any of this rewriting -- but the surprise is worth removing rather than documenting. shift_remove is gated on the feature, so this also converts "someone drops preserve_order later" from a silent reordering into a build failure. The Cargo.toml note points the next caller at it. Reported in review of the preceding commit. Co-authored-by: Cursor Signed-off-by: liyingli * perf(router): write escapes into the buffer instead of a String per code unit escape_non_ascii built a throwaway String for every UTF-16 code unit through format! and then copied it in. write! puts the same bytes straight into the output buffer. Only reachable when a template asks for ensure_ascii=True, which the filter transformers installs defaults away from and GLM-5.2 passes False for, so this is not on any path we measure -- but the allocation had no reason to be there. tojson_matches_transformers_json_dumps already pins the escaped output. Reported by Copilot in review. Co-authored-by: Cursor Signed-off-by: liyingli --------- Signed-off-by: liyingli Co-authored-by: Cursor --- rust/Cargo.lock | 2 + rust/router/Cargo.toml | 13 +- rust/router/src/block_hasher.rs | 418 ++++++++++++++++++++++++++++++-- rust/router/src/protocol.rs | 5 +- 4 files changed, 416 insertions(+), 22 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 77aa2983..ac6d109f 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -1067,6 +1067,7 @@ version = "2.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cb3d648e68cea56d9858d535ee28f9538404e2dd8cb08ed0bd05dca379477f39" dependencies = [ + "indexmap", "memo-map", "serde", "serde_json", @@ -1593,6 +1594,7 @@ version = "1.0.150" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" dependencies = [ + "indexmap", "itoa", "memchr", "serde", diff --git a/rust/router/Cargo.toml b/rust/router/Cargo.toml index 42492b1c..36213c77 100644 --- a/rust/router/Cargo.toml +++ b/rust/router/Cargo.toml @@ -16,7 +16,11 @@ axum = "0.7" # don't need it, but https etcd / workers keep working. reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls"] } serde = { version = "1", features = ["derive"] } -serde_json = "1" +# preserve_order: without it serde_json sorts object keys on parse, so a tool +# call's arguments reach the chat template in a different order than the engine +# saw them and every block hash downstream of it misses. It also redefines +# Map::remove as a swap that perturbs order -- reach for shift_remove instead. +serde_json = { version = "1", features = ["preserve_order"] } arc-swap = "1" clap = { version = "4", features = ["derive", "env"] } tracing = "0.1" @@ -35,4 +39,9 @@ tokenizers = { version = "0.20", default-features = false, features = ["onig"] } # prompts natively. Oniguruma handles the `&&` class intersection + `(?!\S)` # lookahead in Kimi's pat_str that the `regex` crate can't. onig = "6" -minijinja = { version = "2", features = ["json", "loader", "loop_controls"] } +minijinja = { version = "2", features = [ + "json", + "loader", + "loop_controls", + "preserve_order", +] } diff --git a/rust/router/src/block_hasher.rs b/rust/router/src/block_hasher.rs index 1c22a052..3d8cb0da 100644 --- a/rust/router/src/block_hasher.rs +++ b/rust/router/src/block_hasher.rs @@ -17,15 +17,68 @@ //! cost function falls back to load-only routing — never a 500 — exactly like //! the Python side. +use std::fmt::Write as _; +use std::io; use std::path::Path; use minijinja::{context, Environment}; +use serde::Serialize; use serde_json::Value; use tokenizers::Tokenizer; use crate::hasher::hash_request; use crate::tiktoken::KimiTokenizer; +/// `json.dumps` default separators: `", "` between items, `": "` after a key. +/// serde_json packs both, and the engine's prompt has the spaces. +struct PyJsonFormatter; + +impl serde_json::ser::Formatter for PyJsonFormatter { + fn begin_array_value( + &mut self, + w: &mut W, + first: bool, + ) -> io::Result<()> { + if first { + Ok(()) + } else { + w.write_all(b", ") + } + } + fn begin_object_key( + &mut self, + w: &mut W, + first: bool, + ) -> io::Result<()> { + if first { + Ok(()) + } else { + w.write_all(b", ") + } + } + fn begin_object_value(&mut self, w: &mut W) -> io::Result<()> { + w.write_all(b": ") + } +} + +/// Python's `ensure_ascii=True`: every non-ASCII scalar becomes `\uXXXX`, and +/// anything above the BMP becomes a surrogate pair. +fn escape_non_ascii(s: &str) -> String { + let mut out = String::with_capacity(s.len()); + for c in s.chars() { + if c.is_ascii() { + out.push(c); + } else { + for unit in c.encode_utf16(&mut [0u16; 2]) { + // write! rather than push_str(&format!(..)): straight into the + // buffer instead of a String per code unit. Infallible for String. + let _ = write!(out, "\\u{unit:04x}"); + } + } + } + out +} + pub struct BlockHasher { tokenizer: Option, /// Kimi-style tiktoken tokenizer (no `tokenizer.json`). Takes precedence @@ -180,24 +233,137 @@ impl BlockHasher { fn apply_chat_template(&self, messages: &Value) -> Option { let template = self.chat_template.as_ref()?; let mut env = Environment::new(); - // HF chat templates use the Python dict method `msg.get('key')` (7x in - // Kimi's), which minijinja lacks natively -> the template errors -> empty - // render -> no cache locality for chat. Supply just `.get(key[,default])`. + // transformers renders with both of these on. minijinja defaults them off, + // which leaves a newline after every `{% %}` that is alone on its line -- + // invisible in the plain chat path, which trims explicitly with `{%- -%}`, + // and two stray newlines in GLM-5.2's tool-call branch, which does not. + env.set_trim_blocks(true); + env.set_lstrip_blocks(true); + // HF chat templates are written against Python's data model and call + // methods minijinja does not have. A missing one is not a loud failure: + // the template errors, the render comes back empty, the prompt hashes to + // nothing, and kv-aware quietly becomes load-only routing at full health. + // Kimi's needs `msg.get('key')` (7x); GLM-5.2's needs `content.strip()` + // and `content.split('')[-1]`. env.set_unknown_method_callback(|_state, value, method, args| { + use minijinja::value::ValueKind; use minijinja::{Error, ErrorKind, Value}; - if method == "get" { - let key = args.first().cloned().unwrap_or(Value::UNDEFINED); - let default = args.get(1).cloned().unwrap_or_else(|| Value::from(())); - return Ok(match value.get_item(&key) { - Ok(v) if !v.is_undefined() => v, - _ => default, - }); + let as_text = |v: &Value| -> Result { + v.as_str().map(str::to_owned).ok_or_else(|| { + Error::new( + ErrorKind::InvalidOperation, + format!("{method}() expects a string, got {}", v.kind()), + ) + }) + }; + match method { + "get" => { + let key = args.first().cloned().unwrap_or(Value::UNDEFINED); + let default = args.get(1).cloned().unwrap_or_else(|| Value::from(())); + Ok(match value.get_item(&key) { + Ok(v) if !v.is_undefined() => v, + _ => default, + }) + } + // Python strips whitespace with no argument and the given + // character SET (not substring) with one. + "strip" | "lstrip" | "rstrip" => { + let s = as_text(value)?; + let chars: Option> = match args.first() { + Some(a) if !a.is_undefined() && !a.is_none() => { + Some(as_text(a)?.chars().collect()) + } + _ => None, + }; + let matches = |c: char| match &chars { + Some(set) => set.contains(&c), + None => c.is_whitespace(), + }; + Ok(Value::from(match method { + "lstrip" => s.trim_start_matches(matches).to_owned(), + "rstrip" => s.trim_end_matches(matches).to_owned(), + _ => s.trim_matches(matches).to_owned(), + })) + } + // `s.split(sep)` keeps empty fields, so "ab".split("") is + // ["a", "b"] and "".split("") is [""] -- str::split agrees. + // Bare `s.split()` is the different, whitespace-collapsing form. + "split" => { + let s = as_text(value)?; + let parts: Vec = match args.first() { + Some(a) if !a.is_undefined() && !a.is_none() => { + let sep = as_text(a)?; + if sep.is_empty() { + return Err(Error::new( + ErrorKind::InvalidOperation, + "split() with an empty separator", + )); + } + s.split(sep.as_str()).map(str::to_owned).collect() + } + _ => s.split_whitespace().map(str::to_owned).collect(), + }; + Ok(Value::from(parts)) + } + // Python's dict views. minijinja iterates a map as its keys, so + // the other two are built from that. + "items" | "keys" | "values" => { + if value.kind() != ValueKind::Map { + return Err(Error::new( + ErrorKind::InvalidOperation, + format!("{method}() expects a mapping, got {}", value.kind()), + )); + } + let at = |k: &Value| value.get_item(k).unwrap_or(Value::UNDEFINED); + let keys = value.try_iter()?; + Ok(Value::from(match method { + "keys" => keys.collect::>(), + "values" => keys.map(|k| at(&k)).collect(), + _ => keys + .map(|k| { + let v = at(&k); + Value::from(vec![k, v]) + }) + .collect(), + })) + } + _ => Err(Error::new( + ErrorKind::UnknownMethod, + format!("object has no method {method}"), + )), } - Err(Error::new( - ErrorKind::UnknownMethod, - format!("object has no method {method}"), - )) }); + // transformers does not use jinja2's tojson either -- it installs + // `json.dumps(x, ensure_ascii=False, sort_keys=False)`, explicitly to stop + // HTML characters being escaped. minijinja's builtin escapes `<`, `>` and + // `&` and packs the separators, so both of them would move the token + // stream away from the engine's. + env.add_filter( + "tojson", + |v: minijinja::Value, kwargs: minijinja::value::Kwargs| { + use minijinja::{Error, ErrorKind}; + let ascii = kwargs.get::>("ensure_ascii")?.unwrap_or(false); + // Ignoring an argument silently would only shift the divergence + // somewhere harder to see; failing outright would blind the + // router completely, which is worse. + if let Err(e) = kwargs.assert_all_used() { + tracing::warn!(err = %e, "kv-aware: tojson() argument ignored"); + } + let mut buf = Vec::new(); + v.serialize(&mut serde_json::Serializer::with_formatter( + &mut buf, + PyJsonFormatter, + )) + .map_err(|e| Error::new(ErrorKind::InvalidOperation, format!("tojson: {e}")))?; + let out = String::from_utf8(buf) + .map_err(|e| Error::new(ErrorKind::InvalidOperation, format!("tojson: {e}")))?; + Ok(minijinja::Value::from(if ascii { + escape_non_ascii(&out) + } else { + out + })) + }, + ); // HF templates call raise_exception(msg) on malformed input. env.add_function( "raise_exception", @@ -315,18 +481,37 @@ mod tests { assert!(h.hash_for(&json!({"prompt": "text"}), 4).is_empty()); } + /// Model dir for a weight-backed test, from the environment. + /// + /// These used to carry a hardcoded absolute path from whichever machine the + /// test was written on. That is worse than no test: everywhere else the + /// path is missing, the test returns early, and the run reports a pass + /// having checked nothing. Naming the variable at least makes the skip + /// legible and lets CI opt in. + fn model_dir_from_env(var: &str) -> Option { + match std::env::var(var) { + Ok(p) if Path::new(&p).join("chat_template.jinja").exists() => Some(p), + Ok(p) => { + eprintln!("skip: {var}={p} has no chat_template.jinja"); + None + } + Err(_) => { + eprintln!("skip: set {var} to a model dir to run this test"); + None + } + } + } + // Kimi ships chat_template.jinja (not embedded) + tiktoken. This guards the // whole chat path: standalone-template fallback, the `.get()` method, and the // `{% break %}` loop control all have to work or a chat request hashes to // empty (load-only routing, no cache locality). Skips if weights absent. #[test] fn kimi_chat_request_renders_and_hashes() { - const KIMI_DIR: &str = "/mnt/vast/john/huggingface/amd-Kimi-K2.6-MXFP4"; - if !Path::new(KIMI_DIR).join("chat_template.jinja").exists() { - eprintln!("skip: {KIMI_DIR} not present"); + let Some(kimi_dir) = model_dir_from_env("INFERA_TEST_KIMI_DIR") else { return; - } - let h = BlockHasher::load(KIMI_DIR); + }; + let h = BlockHasher::load(&kimi_dir); assert!(h.is_enabled()); // A multi-turn chat (incl. a tool message → exercises break/.get) must // render to a non-trivial token stream -> at least one 16-token block. @@ -344,4 +529,199 @@ mod tests { "kimi chat template must render + tokenize (else 0 cache locality)" ); } + + /// GLM-5.2's template splits the thinking block out and strips the rest: + /// + /// {%- set content = content.split('')[-1] %} + /// {%- if content.strip() -%}{{ content.strip() }}{%- endif -%} + /// + /// minijinja has neither method natively. Without them the render errors, + /// the prompt yields no tokens, and kv-aware degrades to load-only routing + /// while every health signal stays green -- measured as 0.00% predicted + /// hits against the Python router's 83.33% on the same trace. No weights + /// needed: the template text is the whole subject. + #[test] + fn python_str_methods_render_glm_style_template() { + let h = BlockHasher { + tokenizer: None, + tiktoken: None, + chat_template: Some( + "{%- for m in messages -%}\ + {%- set content = m['content'] -%}\ + {%- set reasoning = content.split('')[0].split('')[-1] -%}\ + {%- set content = content.split('')[-1] -%}\ + {%- if content.strip() -%}[{{ content.strip() }}|{{ reasoning.strip() }}]\ + {%- endif -%}{%- endfor -%}" + .to_string(), + ), + bos_token: None, + eos_token: None, + }; + let messages = json!([{"content": " weighing it the answer "}]); + assert_eq!( + h.apply_chat_template(&messages).as_deref(), + Some("[the answer|weighing it]"), + "split() must keep empty fields and index from the end; strip() must \ + trim both ends -- anything else changes the token stream and every \ + block hash with it" + ); + } + + /// The pieces of Python's string semantics the templates actually lean on, + /// where minijinja's nearest builtin differs: `split(sep)` KEEPS empty + /// fields (bare `split()` does not), and `strip(chars)` takes a character + /// SET, not a suffix. + #[test] + fn str_methods_follow_python_semantics() { + let h = BlockHasher { + tokenizer: None, + tiktoken: None, + chat_template: Some( + "{%- set s = messages[0]['content'] -%}\ + {{ s.split(',') | length }}|{{ s.split() | length }}|\ + {{ s.strip(' x') }}|{{ s.lstrip(' ') }}|{{ s.rstrip(' ') }}" + .to_string(), + ), + bos_token: None, + eos_token: None, + }; + // "a,,b " -> split(',') = ["a","","b "] (3, empties kept) + // -> split() = ["a,,b"] (1, whitespace-collapsing) + // -> strip(" x") trims spaces and 'x' from both ends + let messages = json!([{"content": "a,,b "}]); + assert_eq!( + h.apply_chat_template(&messages).as_deref(), + Some("3|1|a,,b|a,,b |a,,b") + ); + } + + /// `transformers` swaps jinja2's tojson for `json.dumps(..., ensure_ascii= + /// False)`, so the engine's prompt has Python's `", "` / `": "` spacing and + /// literal `<`, `>`, `&`. minijinja's builtin does the opposite on both + /// counts, and object order has to survive serde_json and minijinja (each + /// sorts keys unless told otherwise) or the arguments come out alphabetised. + #[test] + fn tojson_matches_transformers_json_dumps() { + let h = BlockHasher { + tokenizer: None, + tiktoken: None, + chat_template: Some( + "{{ messages[0] | tojson }}|{{ messages[0] | tojson(ensure_ascii=True) }}" + .to_string(), + ), + bos_token: None, + eos_token: None, + }; + let messages = json!([{"z": "", "a": [1, 2], "u": "中文"}]); + assert_eq!( + h.apply_chat_template(&messages).as_deref(), + Some( + r#"{"z": "", "a": [1, 2], "u": "中文"}|{"z": "", "a": [1, 2], "u": "\u4e2d\u6587"}"# + ), + "must equal json.dumps(x, ensure_ascii=...) byte for byte" + ); + } + + /// minijinja has none of Python's dict views, and GLM-5.2 iterates + /// `arguments.items()` on every tool call, so their absence takes the whole + /// render down rather than just that branch. + #[test] + fn dict_views_follow_python_semantics() { + let h = BlockHasher { + tokenizer: None, + tiktoken: None, + chat_template: Some( + "{%- set d = messages[0] -%}\ + {%- for k, v in d.items() -%}{{ k }}={{ v }};{%- endfor -%}\ + |{{ d.keys() | join(',') }}|{{ d.values() | join(',') }}" + .to_string(), + ), + bos_token: None, + eos_token: None, + }; + let messages = json!([{"path": "/etc/hosts", "limit": 40}]); + assert_eq!( + h.apply_chat_template(&messages).as_deref(), + Some("path=/etc/hosts;limit=40;|path,limit|/etc/hosts,40"), + "items() must unpack as (key, value), agree with keys()/values(), and \ + keep the request's own key order rather than alphabetising it" + ); + } + + /// End-to-end on real GLM-5.2 weights: template render -> tokenize -> block + /// hashes. Guards the whole chat path the way the Kimi test does. The + /// hermetic test above is the one that actually pins the bug; this one + /// catches a template that changes under us. Opt in with + /// `INFERA_TEST_GLM_DIR=/path/to/GLM-5.2-*`. + #[test] + fn glm52_chat_request_renders_and_hashes() { + let Some(glm_dir) = model_dir_from_env("INFERA_TEST_GLM_DIR") else { + return; + }; + // The render failure this guards against is only ever a warn log, so + // without a subscriber the assert below says "blind" and not why. + let _ = tracing_subscriber::fmt() + .with_max_level(tracing::Level::WARN) + .try_init(); + let h = BlockHasher::load(&glm_dir); + assert!(h.is_enabled()); + let body = json!({"messages": [ + {"role": "user", "content": "read the config file ".repeat(40)}, + {"role": "assistant", "content": "lookinghere it is ".repeat(40)}, + {"role": "user", "content": "now summarize it ".repeat(40)}, + ]}); + assert!( + !h.hash_for(&body, 64).is_empty(), + "GLM-5.2 chat template must render + tokenize, else kv-aware is blind" + ); + } + + /// The tool-call branch reaches template code the plain chat turns never do + /// (`arguments.items()`, `tojson`, the tool role, and `{% %}` tags with no + /// explicit whitespace control), and the trace this router was tuned on has + /// no tool calls -- so the path stayed broken with every benchmark green. + /// + /// A hit needs the router's tokens to equal the engine's, so "it rendered" + /// is not the bar; this pins the exact string. Regenerate the expected value + /// by rendering `chat_template.jinja` through the environment built in + /// `transformers.utils.chat_template_utils._compile_jinja_template`. + #[test] + fn glm52_tool_call_render_matches_transformers() { + let Some(glm_dir) = model_dir_from_env("INFERA_TEST_GLM_DIR") else { + return; + }; + let h = BlockHasher::load(&glm_dir); + // Ordered so that alphabetising the keys shows up, and carrying values + // that separate Python's json.dumps from serde_json's: a nested object, + // a list, an HTML character and a non-ASCII one. + let messages = json!([ + {"role": "user", "content": "read the config file"}, + {"role": "assistant", "content": "", "tool_calls": [{ + "type": "function", + "function": {"name": "read_file", "arguments": { + "path": "/etc/hosts", + "opts": {"depth": 2, "glob": "<*.py>", "note": "中文"}, + "tags": ["a", "b"], + "limit": 40, + }}, + }]}, + {"role": "tool", "content": "127.0.0.1 localhost"}, + ]); + assert_eq!( + h.apply_chat_template(&messages).as_deref(), + Some(concat!( + "[gMASK]<|system|>Reasoning Effort: Max", + "<|user|>read the config file", + "<|assistant|>", + "read_file", + "path/etc/hosts", + r#"opts{"depth": 2, "glob": "<*.py>", "note": "中文"}"#, + r#"tags["a", "b"]"#, + "limit40", + "", + "<|observation|>127.0.0.1 localhost", + "<|assistant|>", + )) + ); + } } diff --git a/rust/router/src/protocol.rs b/rust/router/src/protocol.rs index 6396d591..c5136ec6 100644 --- a/rust/router/src/protocol.rs +++ b/rust/router/src/protocol.rs @@ -108,7 +108,10 @@ fn required_param<'a>(w: &'a Worker, key: &str) -> anyhow::Result<&'a str> { pub fn annotate_vllm_prefill(body: &mut Map, room: u64) { body.insert("max_tokens".into(), Value::from(1)); body.insert("stream".into(), Value::from(false)); - body.remove("stream_options"); + // Not remove(): preserve_order redefines it as a swap, which would reorder + // the forwarded body. shift_remove is gated on that feature, so dropping it + // breaks the build rather than silently changing what we send. + body.shift_remove("stream_options"); if body.contains_key("max_completion_tokens") { body.insert("max_completion_tokens".into(), Value::from(1)); } From 8b6f98290838dff4fd5539ad8371ada43e7cc50d Mon Sep 17 00:00:00 2001 From: xiaobochen-amd Date: Fri, 7 Aug 2026 01:59:47 +0000 Subject: [PATCH 48/88] ci: one event per state of the code, not two per commit push fired on every branch alongside pull_request, so a push to a branch with an open PR started two runs for the same commit. To keep the work from doubling, the jobs were split by hand between them -- unit, rust and torch-cpu on the push event, GPU and DCO on the PR -- and three of them carried a same-repo/fork test to enforce the split. That split has costs the split itself cannot pay back. The PR's own check list shows "skipped" for unit, rust and torch-cpu: their real result lives in a push run the PR does not link to. Both runs report check runs named `unit` on one commit, which GitHub's own docs call ambiguous and advise against, and the PR run's skip resolves in seconds while the push run's real test takes a minute -- long enough for auto-merge to act on the wrong one. And `changes`, `e2e_gate` and `lint` simply ran twice. Scope push to main and let the PR event carry everything it gates. The fork conditions go with it, since there is nothing left to deduplicate. A branch with no PR now runs nothing; a draft PR gives the same feedback and puts it on the PR. e2e_gate loses the tree comparison. It existed to avoid re-running e2e on main when the merge produced the same content the PR had tested; the team's call is that main re-tests unconditionally, because two PRs can each pass alone and fail together and only the merged result shows it. What is left is the event, plus a draft check so iterating on a draft does not cost a GPU run per push -- ready_for_review is in `types` so marking it ready picks the tiers back up. pre_check goes too. skip-duplicate-actions matched on tree alone, ignoring the event, and it was there for the double-run this commit removes. Left in, it would have read the post-merge run as a duplicate of the PR that produced it and skipped the very re-test main exists to provide. It was also how `engine` could silently not run: a successful branch-push run (where engine never runs, e2e_gate being false off main) made the later PR look like a duplicate. Its `actions: read` permission goes with it. One pre-existing bug had to be fixed here rather than left: `changes` has no diff base on workflow_dispatch, so it fell back to HEAD~1 and would skip the whole run off a docs-only last commit. That was survivable while a branch push started CI. It is not now that the button is the only way to force a run, so a manual run is treated as a code change. Co-authored-by: Cursor Signed-off-by: xiaobochen-amd --- .github/workflows/ci.yml | 134 +++++++++++++++------------------------ 1 file changed, 51 insertions(+), 83 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4cfb4f28..116cfa07 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -1,10 +1,19 @@ name: CI on: + # One event per state of the code: the PR is the pre-merge gate, main is the + # post-merge check. Firing push on every branch ran both for the same commit, + # so the jobs were split by hand -- unit on push, GPU on the PR -- which left + # the PR's own check list reading "skipped" for tests that ran in a push run + # it does not link to, and two check runs named `unit` on one commit. + # A branch with no PR now runs nothing; open a draft to get CI. push: - branches: ["**"] # e2e only for code reaching main untested (see e2e_gate) + branches: [main] pull_request: - branches: [main] # e2e runs here, pre-merge — the usual path + branches: [main] + # ready_for_review so marking a draft ready re-runs and picks up the GPU + # tiers that drafts skip (see e2e_gate). + types: [opened, synchronize, reopened, ready_for_review] workflow_dispatch: # manual "Run workflow" button (Actions tab) inputs: run_e2e_mixed: @@ -22,7 +31,6 @@ concurrency: permissions: contents: read - actions: read jobs: # Classify the change: does it touch package code, or only docs/examples? A @@ -39,6 +47,16 @@ jobs: fetch-depth: 0 - id: f run: | + # Someone pressed "Run workflow": run everything, whatever the last + # commit happened to touch. There is no diff base on this event, so the + # logic below would fall back to HEAD~1 and skip the whole run off a + # docs-only commit -- and since a branch push no longer starts CI, this + # button is the only way to force one. + if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then + echo "manual run — treating as a code change" + echo "code=true" >> "$GITHUB_OUTPUT" + exit 0 + fi if [ "${{ github.event_name }}" = "pull_request" ]; then range="${{ github.event.pull_request.base.sha }}...${{ github.event.pull_request.head.sha }}" else @@ -63,65 +81,35 @@ jobs: echo "=> code=$code" echo "code=$code" >> "$GITHUB_OUTPUT" - # Should the e2e tiers run for THIS event? Every PR into main is tested - # pre-merge, so landing that same code must not test it twice — but "same" has - # to mean the CONTENT, not the PR: if main moved on while the PR sat open, what - # lands is a combination no e2e ever saw. Compare git TREES, which are exactly - # the content, and are equal iff the merge changed nothing versus the PR head. + # Should the GPU tiers run for THIS event? Every PR into main, and every merge + # into main. Re-running on main is deliberate duplication: two PRs can each + # pass alone and break together, and only the merged result shows that. + # Drafts are the exception — iterating on one must not cost a GPU run per + # push, so they get lint and unit only until they are marked ready. e2e_gate: runs-on: ubuntu-latest - permissions: - contents: read - pull-requests: read outputs: run: ${{ steps.decide.outputs.run }} steps: - id: decide env: - GH_TOKEN: ${{ github.token }} - REPO: ${{ github.repository }} - SHA: ${{ github.sha }} + EVENT: ${{ github.event_name }} + REF: ${{ github.ref }} + DRAFT: ${{ github.event.pull_request.draft }} run: | - tree_of() { gh api "repos/$REPO/commits/$1" --jq .commit.tree.sha 2>/dev/null; } - run=false - if [ "${{ github.event_name }}" = "pull_request" ]; then - run=true - echo "pull_request into main — e2e runs pre-merge" - elif [ "${{ github.event_name }}" = "push" ] && [ "${{ github.ref }}" = "refs/heads/main" ]; then - # Squash and merge commits both report the PR they came from. - head=$(gh api "repos/$REPO/commits/$SHA/pulls" --jq '.[0].head.sha // empty' 2>/dev/null) - landed=$(tree_of "$SHA") - tested=""; [ -n "$head" ] && tested=$(tree_of "$head") - if [ -z "$head" ]; then - run=true; echo "no PR behind this commit — its code was never e2e'd" - elif [ -z "$landed" ] || [ -z "$tested" ]; then - # Never infer "already tested" from a failed lookup: re-testing costs - # GPU minutes, shipping untested code costs more. - run=true; echo "could not read both trees — running e2e to be safe" - elif [ "$landed" = "$tested" ]; then - echo "tree $landed is what PR head $head already e2e'd — skipping" - else - run=true - echo "tree $landed != PR head $head's $tested — main moved under the PR" - fi + if [ "$EVENT" = pull_request ] && [ "$DRAFT" = true ]; then + echo "draft pull request — lint and unit only until it is marked ready" + elif [ "$EVENT" = pull_request ]; then + run=true; echo "pull request into main — GPU tiers run pre-merge" + elif [ "$EVENT" = push ] && [ "$REF" = refs/heads/main ]; then + run=true; echo "merged into main — GPU tiers run again on the result" else - echo "not a PR into main, and not a push to main — no e2e" + echo "neither a pull request into main nor a push to main — no GPU tiers" fi echo "=> run_e2e=$run" echo "run=$run" >> "$GITHUB_OUTPUT" - pre_check: - runs-on: ubuntu-latest - outputs: - should_skip: ${{ steps.skip.outputs.should_skip }} - steps: - - id: skip - uses: fkirc/skip-duplicate-actions@f75f66ce1886f00957d99748a42c724f4330bdcf # v5 - with: - skip_after_successful_duplicate: "true" - concurrent_skipping: "never" - # Sign-off gate for the GPU tiers below — `needs:` cannot reach a job in # another workflow, so dco.yml is called here as one. Skipped off a PR (there # is nothing to check), which the tiers below read as "did not fail". @@ -130,10 +118,8 @@ jobs: uses: ./.github/workflows/dco.yml lint: - needs: [pre_check, changes] - if: >- - needs.pre_check.outputs.should_skip != 'true' && - needs.changes.outputs.code == 'true' + needs: [changes] + if: needs.changes.outputs.code == 'true' runs-on: ubuntu-latest steps: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 @@ -145,14 +131,9 @@ jobs: run: pre-commit run --all-files --show-diff-on-failure unit: - needs: [lint, pre_check, changes] - # Skip on docs-only changes, and on same-repo PRs (already covered by the - # branch push event); still run on push and on fork PRs. - if: >- - needs.pre_check.outputs.should_skip != 'true' && - needs.changes.outputs.code == 'true' && - (github.event_name != 'pull_request' || - github.event.pull_request.head.repo.full_name != github.repository) + needs: [lint, changes] + # Skip on docs-only changes. + if: needs.changes.outputs.code == 'true' runs-on: ubuntu-latest steps: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 @@ -168,14 +149,9 @@ jobs: # Rust router: clippy + unit/integration tests. GH-hosted ubuntu-latest # ships a stable toolchain (with clippy), so no toolchain setup is needed. rust: - needs: [lint, pre_check, changes] - # Skip on docs-only changes, and on same-repo PRs (already covered by the - # branch push event); still run on push and on fork PRs. - if: >- - needs.pre_check.outputs.should_skip != 'true' && - needs.changes.outputs.code == 'true' && - (github.event_name != 'pull_request' || - github.event.pull_request.head.repo.full_name != github.repository) + needs: [lint, changes] + # Skip on docs-only changes. + if: needs.changes.outputs.code == 'true' runs-on: ubuntu-latest defaults: run: @@ -190,13 +166,11 @@ jobs: run: cargo test # Engine GPU tests. Same schedule as the e2e tiers (see e2e_gate): every PR - # into main, and any push that lands untested code on main — never on a plain - # branch push, which is where these GPU minutes used to go. + # into main once it is out of draft, and every merge into main. engine: - needs: [lint, pre_check, changes, dco, e2e_gate] + needs: [lint, changes, dco, e2e_gate] if: >- !cancelled() && - needs.pre_check.outputs.should_skip != 'true' && needs.changes.outputs.code == 'true' && needs.lint.result != 'failure' && needs.lint.result != 'cancelled' && needs.dco.result != 'failure' && needs.dco.result != 'cancelled' && @@ -222,12 +196,11 @@ jobs: if: always() && (cancelled() || failure()) run: .github/scripts/reclaim_slurm_jobs.sh infera-ci- "-${{ github.run_id }}-engine" - # Full PD-mixed e2e (per engine, parallel). When it runs is e2e_gate's call: - # every PR into main, plus a push that lands untested code on main. + # Full PD-mixed e2e (per engine, parallel). When it runs is e2e_gate's call. e2e-mixed: # Skipped for docs-only changes. Gated behind lint and dco: run only if neither # failed. `!cancelled()` + result checks (instead of a plain success dependency) - # is needed so e2e still runs when lint is *skipped* as a duplicate (pre_check) + # is needed so e2e still runs when lint is *skipped* for a docs-only change # or dco is skipped off a PR, but is held back when either fails. NOT # `always()`: on cancel the server re-evaluates job-level `if`, and `always()` # evaluates true, so the job is never cancelled — it keeps (or even starts) @@ -316,14 +289,9 @@ jobs: run: .github/scripts/reclaim_slurm_jobs.sh infera-ci- "-${{ github.run_id }}-${{ matrix.engine }}-disag" unit-torch-cpu: - needs: [lint, pre_check, changes] - # Skip on docs-only changes, and on same-repo PRs (already covered by the - # branch push event); still run on push and on fork PRs. - if: >- - needs.pre_check.outputs.should_skip != 'true' && - needs.changes.outputs.code == 'true' && - (github.event_name != 'pull_request' || - github.event.pull_request.head.repo.full_name != github.repository) + needs: [lint, changes] + # Skip on docs-only changes. + if: needs.changes.outputs.code == 'true' runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 From e1a8892b7b5947830b913480342afc7697b9b368 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Fri, 7 Aug 2026 02:07:03 +0000 Subject: [PATCH 49/88] fix(router): index only the attention KV-cache groups kv-aware routing reported cache_hits=0 on every decision against Kimi-K3 on two mixed workers, and pinned all traffic to one of them. Not a degraded hit rate -- half the fleet idle. vLLM emits one BlockStored PER KV-CACHE GROUP. On a hybrid model only the attention group pairs one hash with one block; Kimi-K3 has 24 MLA layers and 69 KDA layers, and the KDA groups run prefix caching in "align" mode, which nulls all but one block per step. Those events go out with token_ids spanning the whole chunk and a single surviving hash, and nothing on the wire says which block that hash covers. Captured off a live worker (block_size 768, --max-num-batched-tokens 4096, so chunked prefill splits at 3840): kv_cache_spec_kind group_idx token_ids block_hashes mamba 0 3840 1 mamba 1 3840 1 mamba 2 3840 1 mla_attention 3 3840 5 Indexing all four is worse than indexing none of them. vLLM mixes no group id into the block hash, so at equal block sizes a Mamba hash COLLIDES with an attention hash and overwrites its entry in map[engine_hash] -> router_hash; the next chunk then resolves its parent to the wrong node and every block after it is hashed off a poisoned chain. Which of the four wins depended on intra-batch emission order, which vLLM does not promise. So filter by kv_cache_spec_kind, the way Dynamo's is_main_attention() does -- which is why Dynamo survives hybrid models and we did not. The two fields the filter needs, group_idx and kv_cache_spec_kind, were on the wire all along; our msgspec struct simply did not declare them, so they were discarded before anything could look at them. Same change in the Rust decoder. mla_attention is NOT optional in the whitelist. Kimi-K3's attention layers are MLA, so a filter written as == "full_attention" discards 100% of that model's usable events -- the same empty view, reached from the other side. The filter fails OPEN when the field is absent (SGLang, older vLLM builds), because those engines were indexed before this change and a closed default would switch kv-aware routing off for them silently. On a length disagreement the event is NOT dropped whole. Its view entries are correct -- chained from a resolved parent over a contiguous token span -- and dropping them measurably costs hits: on Qwen3.5-0.8B (same event shapes), 18 of 32 requests hit a full prefix with the whole-event drop versus 22 of 32 without the filter at all. What is not correct is the hash-to-block pairing, so the view is indexed and the map is not. A later event then misses its parent and is dropped, which under-reports hits instead of mis-reporting them. Signed-off-by: Zhang, Jiejing --- infera/router/kv_event/client.py | 65 ++++- infera/router/kv_event/events.py | 20 ++ rust/router/src/kv_event.rs | 236 +++++++++++++++++- .../router/test_kv_block_size_mismatch.py | 39 ++- .../unit/router/test_kv_event_group_filter.py | 184 ++++++++++++++ 5 files changed, 524 insertions(+), 20 deletions(-) create mode 100644 tests/unit/router/test_kv_event_group_filter.py diff --git a/infera/router/kv_event/client.py b/infera/router/kv_event/client.py index 200d2401..9537e591 100644 --- a/infera/router/kv_event/client.py +++ b/infera/router/kv_event/client.py @@ -91,6 +91,8 @@ class WorkerSubscription: # Latched: the block-size mismatch is per-subscription, not per-event, so # logging it on every event would bury the fact under thousands of copies. block_size_mismatch_logged: bool = False + length_mismatch_logged: bool = False + non_indexable_group_logged: bool = False def view_for(self, rank: int | None) -> set[int]: return self.views.setdefault(rank or 0, set()) @@ -242,7 +244,40 @@ def _handle_event(self, sub: WorkerSubscription, ev: object, rank: int | None = sub.view_for(rank).clear() sub.map_for(rank).clear() + # vLLM emits one BlockStored per KV-CACHE GROUP. Only the attention groups + # pair one hash with one block; see `events.BlockStored`. Mirrors Dynamo's + # `is_main_attention()` filter, which is why Dynamo survives hybrid models. + # + # `mla_attention` is NOT optional here. Kimi-K3's attention layers are MLA, + # so `get_kv_cache_spec_kind` reports MLA_ATTENTION and a filter written as + # `== "full_attention"` would discard 100% of that model's usable events -- + # the same empty view this filter exists to prevent, arrived at from the + # other direction. + _INDEXABLE_SPEC_KINDS = frozenset({"full_attention", "mla_attention", "sink_full_attention"}) + + def _is_indexable_group(self, sub: WorkerSubscription, ev: object) -> bool: + kind = getattr(ev, "kv_cache_spec_kind", None) + if kind is None: + # SGLang, and vLLM builds before the field existed. Fail OPEN: those + # engines were being indexed before this filter and a closed default + # would silently switch kv-aware routing off for them. + return True + if kind in self._INDEXABLE_SPEC_KINDS: + return True + if not sub.non_indexable_group_logged: + sub.non_indexable_group_logged = True + logger.info( + "kv events from %s include non-attention groups (first seen: %s, " + "group_idx=%s); indexing only the attention group(s)", + sub.worker_id, + kind, + getattr(ev, "group_idx", None), + ) + return False + def _on_block_stored(self, sub: WorkerSubscription, ev: object, rank: int | None) -> None: + if not self._is_indexable_group(sub, ev): + return view, m = sub.view_for(rank), sub.map_for(rank) if ev.parent_block_hash is None: parent = ROUTER_SEED @@ -277,19 +312,39 @@ def _on_block_stored(self, sub: WorkerSubscription, ev: object, rank: int | None # then used to index block_hashes, so any disagreement between them was an # IndexError — surfacing as "list index out of range", a message that points # nowhere near the block size. - n = min(len(tokens) // bs, len(ev.block_hashes)) - if n * bs != len(tokens) or n != len(ev.block_hashes): + # Index every block the token span covers, but only trust a block hash + # when the two lengths agree. + # + # Measured on a hybrid model (Qwen3.5-0.8B, same event shapes as + # Kimi-K3): dropping a length-disagreeing event whole costs more than it + # saves -- 18/32 requests hit a full prefix versus 22/32 when the + # leading blocks are indexed. The view entries ARE correct; they are + # chained from a resolved parent over a contiguous token span. What is + # not correct is the hash-to-block pairing. + # + # So index the view and skip the map. `map[engine_hash] -> router_hash` + # is what a later event resolves its parent against, and on a sparse + # event the surviving hash need not describe the leading chunk -- vLLM + # gives no offset. Writing it binds an engine hash to the wrong block + # and poisons the chain from there on; withholding it makes a later + # event miss its parent and be dropped, which under-reports hits instead + # of mis-reporting them. + n = len(tokens) // bs + aligned = n == len(ev.block_hashes) + if not aligned and not sub.length_mismatch_logged: + sub.length_mismatch_logged = True logger.warning( "kv event from %s covers %d tokens and %d block hashes, which do not " - "agree at block_size=%d; indexing %d block(s) and dropping the rest", + "agree at block_size=%d; indexing the blocks but not their hashes, so " + "later events chaining off them will be dropped", sub.worker_id, len(tokens), len(ev.block_hashes), bs, - n, ) for i in range(n): chunk = tokens[i * bs : (i + 1) * bs] parent = hash_chunk(parent, chunk) view.add(parent) - m[ev.block_hashes[i]] = parent + if aligned: + m[ev.block_hashes[i]] = parent diff --git a/infera/router/kv_event/events.py b/infera/router/kv_event/events.py index 3ae26374..ed797fb2 100644 --- a/infera/router/kv_event/events.py +++ b/infera/router/kv_event/events.py @@ -42,6 +42,26 @@ class BlockStored(_VllmKVCacheEvent): block_size: int lora_id: int | None medium: str | None = None + # vLLM emits one event PER KV-CACHE GROUP, and only the attention groups + # carry a usable hash-per-block. On a hybrid model (Kimi-K3: 3 KDA/Mamba + # groups + 1 MLA group) the Mamba groups run prefix caching in "align" + # mode, where all but one block per step is a null block that is skipped + # when the hash list is built -- while ``token_ids`` still spans the whole + # range. Measured on Kimi-K3: the Mamba groups report 3840 tokens against + # ONE hash at block_size=768, the MLA group 3840 against five. + # + # There is no field saying which chunk the surviving hash covers, so a + # Mamba event cannot be indexed at all. Worse, vLLM's block hash does not + # mix in the group id, so with equal block sizes a Mamba hash COLLIDES with + # an attention hash and overwrites its entry in the engine-hash -> router- + # hash map, breaking the parent chain for every later block. Filtering on + # these two fields is what keeps the one usable stream intact; see + # ``client._on_block_stored``. + # + # Both are absent on SGLang and on vLLM builds predating them, so they + # default to None and the filter must fail open. Upstream: vllm#44451. + group_idx: int | None = None + kv_cache_spec_kind: str | None = None class BlockRemoved(_VllmKVCacheEvent): diff --git a/rust/router/src/kv_event.rs b/rust/router/src/kv_event.rs index cc6de242..09c663f5 100644 --- a/rust/router/src/kv_event.rs +++ b/rust/router/src/kv_event.rs @@ -58,6 +58,9 @@ enum Event { block_hashes: Vec, parent_block_hash: Option, token_ids: Vec, + /// vLLM's `kv_cache_spec_kind`. `None` on SGLang and on vLLM builds + /// predating the field; see `is_indexable_spec_kind`. + spec_kind: Option, }, Removed { block_hashes: Vec, @@ -313,7 +316,30 @@ fn apply_events( block_hashes, parent_block_hash, token_ids, + spec_kind, } => { + if !is_indexable_spec_kind(spec_kind.as_deref()) { + continue; + } + // Index every block the token span covers, but only trust a + // block hash when the two lengths agree. + // + // Measured on a hybrid model (Qwen3.5-0.8B, the same shapes + // Kimi-K3 emits): dropping a length-disagreeing event whole + // costs more than it saves -- 18/32 requests hit a full prefix + // versus 22/32 when the leading blocks are indexed. The view + // entries ARE correct: chained from a resolved parent over a + // contiguous span. The hash-to-block PAIRING is what is not. + // + // So fill the view and skip the map. `map` is what a later + // event resolves its parent against, and on a sparse event the + // surviving hash need not describe the leading chunk -- vLLM + // gives no offset. Writing it binds an engine hash to the wrong + // block and poisons the chain from there; withholding it makes + // a later event miss its parent and be dropped, which + // under-reports hits instead of mis-reporting them. + let n = token_ids.len() / bs; + let aligned = n == block_hashes.len(); let mut parent = match parent_block_hash { None => ROUTER_SEED, Some(ph) => match map.get(ph) { @@ -321,13 +347,12 @@ fn apply_events( None => continue, // chain broken: missing parent, drop }, }; - let n = token_ids.len() / bs; for i in 0..n { let chunk = &token_ids[i * bs..(i + 1) * bs]; parent = hash_chunk(parent, chunk); view.insert(parent); - if let Some(wh) = block_hashes.get(i) { - map.insert(*wh, parent); + if aligned { + map.insert(block_hashes[i], parent); } } } @@ -377,6 +402,8 @@ fn parse_event(ev: &rmpv::Value) -> Option { match tag { // [tag, block_hashes, parent_block_hash, token_ids, block_size, lora_id, medium?] "BlockStored" => Some(Event::Stored { + // SGLang's array form has no group fields; None => fail open. + spec_kind: None, block_hashes: a.get(1).map(as_u64_vec).unwrap_or_default(), parent_block_hash: a.get(2).and_then(as_u64_any), token_ids: a.get(3).map(as_u32_vec).unwrap_or_default(), @@ -390,9 +417,44 @@ fn parse_event(ev: &rmpv::Value) -> Option { } } +/// Which KV-cache groups carry a usable hash-per-block. +/// +/// vLLM emits one `BlockStored` per group. On a hybrid model -- Kimi-K3 is 3 +/// KDA/Mamba groups plus 1 MLA group -- the Mamba groups run prefix caching in +/// "align" mode, where all but one block per step is a null block skipped when +/// the hash list is built, while `token_ids` still spans the whole range. +/// Measured: Mamba groups report 3840 tokens against ONE hash at +/// block_size=768; the MLA group reports 3840 against five. +/// +/// Those events cannot be indexed -- nothing says which chunk the surviving +/// hash covers. They also actively corrupt the view: vLLM's block hash does not +/// mix in the group id, so at equal block sizes a Mamba hash collides with an +/// attention hash and overwrites its entry in the engine-hash -> router-hash +/// map, breaking the parent chain for every block after it. +/// +/// `mla_attention` is mandatory in this set. Kimi-K3's attention layers are +/// MLA, so a filter written as `== "full_attention"` would drop 100% of its +/// usable events -- the same empty view, reached from the other side. +/// +/// `None` means SGLang or a vLLM build predating the field: fail OPEN, since +/// those streams were indexed before this filter existed. Upstream: vllm#44451. +fn is_indexable_spec_kind(kind: Option<&str>) -> bool { + match kind { + None => true, + Some(k) => matches!( + k, + "full_attention" | "mla_attention" | "sink_full_attention" + ), + } +} + /// vLLM tagged-MAP event: `{"type": , : , ...}` (msgspec -/// `tag=True` map; the tag key is "type"). Fields are matched by NAME, so vLLM's -/// extra fields (lora_name, extra_keys, group_idx, kv_cache_spec_*) are ignored. +/// `tag=True` map; the tag key is "type"). Fields are matched by NAME. +/// +/// `kv_cache_spec_kind` IS read: vLLM emits one event per KV-cache group and +/// only the attention groups pair one hash with one block. Ignoring it, as this +/// did, meant a hybrid model's Mamba groups were indexed too -- see +/// `is_indexable_spec_kind`. fn parse_event_map(ev: &rmpv::Value) -> Option { let map = ev.as_map()?; let get = |k: &str| { @@ -405,6 +467,9 @@ fn parse_event_map(ev: &rmpv::Value) -> Option { block_hashes: get("block_hashes").map(as_u64_vec).unwrap_or_default(), parent_block_hash: get("parent_block_hash").and_then(as_u64_any), token_ids: get("token_ids").map(as_u32_vec).unwrap_or_default(), + spec_kind: get("kv_cache_spec_kind") + .and_then(|v| v.as_str()) + .map(str::to_owned), }), "BlockRemoved" => Some(Event::Removed { block_hashes: get("block_hashes").map(as_u64_vec).unwrap_or_default(), @@ -714,6 +779,164 @@ mod tests { c.shutdown(); } + /// vLLM emits one BlockStored per KV-cache GROUP. Only the attention groups + /// pair one hash with one block; a hybrid model's Mamba groups report the + /// whole token span against a single hash (measured on Kimi-K3: 3840 tokens, + /// 1 hash, block_size 768, against the MLA group's 3840/5). + #[test] + fn non_attention_groups_are_not_indexed() { + let c = KvEventClient::new(); + c.on_worker_added(&worker("w", Some("tcp://127.0.0.1:5601"), 4, None)); + + for kind in ["mamba", "sliding_window", "encoder_only_attention"] { + apply_events( + &c.state, + "w", + 0, + &[Event::Stored { + block_hashes: vec![111, 222], + parent_block_hash: None, + spec_kind: Some(kind.to_string()), + token_ids: vec![1, 2, 3, 4, 5, 6, 7, 8], + }], + ); + } + let q = crate::hasher::hash_request(&[1, 2, 3, 4, 5, 6, 7, 8], 4); + assert_eq!(c.prefix_hits("w", None, &q), 0); + c.shutdown(); + } + + /// `mla_attention` is mandatory in the allowed set: Kimi-K3's attention + /// layers are MLA, so a filter written as `== "full_attention"` would drop + /// 100% of that model's usable events. + #[test] + fn attention_groups_are_indexed() { + for kind in ["full_attention", "mla_attention", "sink_full_attention"] { + let c = KvEventClient::new(); + c.on_worker_added(&worker("w", Some("tcp://127.0.0.1:5602"), 4, None)); + apply_events( + &c.state, + "w", + 0, + &[Event::Stored { + block_hashes: vec![111, 222], + parent_block_hash: None, + spec_kind: Some(kind.to_string()), + token_ids: vec![1, 2, 3, 4, 5, 6, 7, 8], + }], + ); + let q = crate::hasher::hash_request(&[1, 2, 3, 4, 5, 6, 7, 8], 4); + assert_eq!(c.prefix_hits("w", None, &q), 2, "kind={kind}"); + c.shutdown(); + } + } + + /// A sparse event -- more token blocks than hashes -- still contributes its + /// blocks to the view; only the hashes are withheld, because nothing says + /// which chunk the surviving one describes. Measured: dropping the event + /// whole cut full-prefix hits from 22/32 to 18/32 on a hybrid model. + #[test] + fn sparse_event_indexes_blocks_but_not_hashes() { + let c = KvEventClient::new(); + c.on_worker_added(&worker("w", Some("tcp://127.0.0.1:5603"), 4, None)); + apply_events( + &c.state, + "w", + 0, + &[Event::Stored { + block_hashes: vec![111], // one hash for two blocks of tokens + parent_block_hash: None, + spec_kind: Some("mla_attention".to_string()), + token_ids: vec![1, 2, 3, 4, 5, 6, 7, 8], + }], + ); + let q = crate::hasher::hash_request(&[1, 2, 3, 4, 5, 6, 7, 8], 4); + assert_eq!(c.prefix_hits("w", None, &q), 2, "blocks are visible"); + + // The hash was not mapped, so a child naming it as parent is dropped + // rather than chained off a block it may not describe. + apply_events( + &c.state, + "w", + 0, + &[Event::Stored { + block_hashes: vec![222], + parent_block_hash: Some(111), + spec_kind: Some("mla_attention".to_string()), + token_ids: vec![9, 10, 11, 12], + }], + ); + let q2 = crate::hasher::hash_request(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], 4); + assert_eq!( + c.prefix_hits("w", None, &q2), + 2, + "child dropped, not mis-chained" + ); + c.shutdown(); + } + + /// The property the filter exists for. Both groups get the SAME block hash + /// (vLLM mixes no group id in), so a Mamba event arriving AFTER the + /// attention event overwrites `map[hash] -> router_hash` with a hash over + /// its own tokens; the next chunk then chains off the wrong node and the + /// view holds a block no query can reproduce. + #[test] + fn a_later_mamba_event_cannot_poison_the_attention_chain() { + let c = KvEventClient::new(); + c.on_worker_added(&worker("w", Some("tcp://127.0.0.1:5604"), 4, None)); + + let stored = |hashes: Vec, parent, kind: &str, toks: Vec| Event::Stored { + block_hashes: hashes, + parent_block_hash: parent, + spec_kind: Some(kind.to_string()), + token_ids: toks, + }; + apply_events( + &c.state, + "w", + 0, + &[stored( + vec![111, 222], + None, + "mla_attention", + vec![1, 2, 3, 4, 5, 6, 7, 8], + )], + ); + // same hashes, different tokens, arriving second + apply_events( + &c.state, + "w", + 0, + &[stored( + vec![111, 222], + None, + "mamba", + vec![90, 91, 92, 93, 94, 95, 96, 97], + )], + ); + // follow-on chunk, parented on the attention group's second block + apply_events( + &c.state, + "w", + 0, + &[stored( + vec![333], + Some(222), + "mla_attention", + vec![9, 10, 11, 12], + )], + ); + + let q = crate::hasher::hash_request(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], 4); + assert_eq!(q.len(), 3); + assert_eq!( + c.prefix_hits("w", None, &q), + 3, + "a later non-attention event overwrote the attention group's hash map" + ); + c.shutdown(); + } + // Drive the view-maintenance chain directly (no sockets) — this is the core // correctness property: a BlockStored feeds token_ids through the SAME chain // as the query side, so prefix_hits matches hash_request. @@ -730,6 +953,7 @@ mod tests { &[Event::Stored { block_hashes: vec![111, 222], parent_block_hash: None, + spec_kind: None, token_ids: vec![1, 2, 3, 4, 5, 6, 7, 8], }], ); @@ -756,6 +980,7 @@ mod tests { &[Event::Stored { block_hashes: vec![111, 222], parent_block_hash: None, + spec_kind: None, token_ids: vec![1, 2, 3, 4, 5, 6, 7, 8], }], ); @@ -787,6 +1012,7 @@ mod tests { &[Event::Stored { block_hashes: vec![1], parent_block_hash: None, + spec_kind: None, token_ids: vec![1, 2, 3, 4], }], ); diff --git a/tests/unit/router/test_kv_block_size_mismatch.py b/tests/unit/router/test_kv_block_size_mismatch.py index ae59c225..535b30c0 100644 --- a/tests/unit/router/test_kv_block_size_mismatch.py +++ b/tests/unit/router/test_kv_block_size_mismatch.py @@ -82,18 +82,30 @@ def test_worker_without_block_size_is_not_subscribed(caplog): assert client._subs == {} -async def test_mismatched_event_does_not_raise_and_indexes_what_it_can(): +async def test_mismatched_event_indexes_blocks_but_not_hashes(): """The real shape of the bug: 768 tokens, 1 block hash, subscriber at 1. - Before the fix this raised IndexError on i=1 and killed the subscriber. + The first fix stopped this raising IndexError. The second dropped the event + whole, on the reasoning that a partial index binds an engine hash to a block + it may not describe. That reasoning is right about the HASH and wrong about + the BLOCK, and measuring it showed the cost: on a hybrid model emitting + these shapes, dropping cut full-prefix hits from 22/32 to 18/32. + + The view entries are correct -- each is chained from a resolved parent over + a contiguous token span, which is exactly what a query reproduces. Only the + hash-to-block pairing is unknown, and only `map` uses it. + + So: index the blocks, withhold the hashes. A later event naming one of those + hashes as its parent then misses and is dropped, which under-reports hits + rather than chaining off the wrong node. """ client = KvEventClient() sub = await _subscribed(client, 1) client._handle_event(sub, _stored(tokens=768, hashes=1, block_size=768), rank=0) - assert len(sub.map_for(0)) == 1, "only the hashes actually supplied may be indexed" - assert len(sub.view_for(0)) == 1 + assert len(sub.view_for(0)) == 768, "the blocks the span covers are indexable" + assert len(sub.map_for(0)) == 0, "their hashes are not" async def test_block_size_disagreement_is_reported_once(caplog): @@ -123,16 +135,23 @@ async def test_agreeing_event_is_unaffected(): @pytest.mark.parametrize( - "tokens,hashes,expect", + "tokens,hashes,view,mapped", [ - (8, 1, 1), # more tokens than hashes cover — the observed failure - (4, 3, 1), # fewer tokens than hashes — the mirror case - (0, 0, 0), # empty event + (8, 2, 2, 2), # agreeing: 2 hashes x block_size 4 == 8 tokens + (8, 1, 2, 0), # more tokens than hashes cover — the observed failure + (4, 3, 1, 0), # fewer tokens than hashes — the mirror case + (0, 0, 0, 0), # empty event ], ) -async def test_bound_is_the_minimum_of_both(tokens, hashes, expect): +async def test_view_follows_the_tokens_and_the_map_follows_agreement(tokens, hashes, view, mapped): + """The view is bounded by the token span, the map by the lengths agreeing. + + Both disagreement directions withhold the hashes: the direction of the error + does not tell you which block the surviving hash belongs to. + """ client = KvEventClient() sub = await _subscribed(client, 4) client._handle_event(sub, _stored(tokens=tokens, hashes=hashes, block_size=4), rank=0) - assert len(sub.map_for(0)) == expect + assert len(sub.view_for(0)) == view + assert len(sub.map_for(0)) == mapped diff --git a/tests/unit/router/test_kv_event_group_filter.py b/tests/unit/router/test_kv_event_group_filter.py new file mode 100644 index 00000000..b2873174 --- /dev/null +++ b/tests/unit/router/test_kv_event_group_filter.py @@ -0,0 +1,184 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""vLLM emits one BlockStored per KV-cache group; only attention groups are indexable. + +Measured on Kimi-K3 (3 KDA/Mamba groups + 1 MLA group, block_size 768, one +prefill chunk): + + kv_cache_spec_kind group_idx token_ids block_hashes + mamba 0 3840 1 + mamba 1 3840 1 + mamba 2 3840 1 + mla_attention 3 3840 5 + +The Mamba groups run prefix caching in "align" mode, where all but one block per +step is a null block that is skipped when the hash list is built, while +``token_ids`` still spans the whole range. There is no field saying which chunk +the surviving hash covers, so the event cannot be indexed. + +Indexing it anyway is not merely useless. vLLM's block hash does not mix in the +group id, so at equal block sizes a Mamba hash COLLIDES with an attention hash +and overwrites its entry in the engine-hash -> router-hash map; every later +attention event naming that hash as its parent then chains off the wrong node. +That is how 3 of 4 event streams silently destroyed the one usable stream, and +the symptom was ``cache_hits=0`` on every routing decision with all traffic +pinned to a single worker. + +Upstream: vllm#44451 (open). NVIDIA Dynamo survives hybrid models only because +its `is_main_attention()` filter drops these groups before decoding. +""" + +from __future__ import annotations + +import logging + +import pytest + +from infera.common.worker_pool import EngineType, WorkerInfo +from infera.router.kv_event.client import KvEventClient +from infera.router.kv_event.events import BlockStored +from infera.router.kv_event.hasher import hash_request + +BS = 4 + + +def _worker(): + return WorkerInfo( + worker_id="w1:30000", + url="http://w1:30000", + model_name="m", + engine=EngineType.VLLM, + kv_events_endpoint="tcp://w1:5555", + kv_block_size=BS, + ) + + +def _stored(*, blocks, spec_kind, group_idx=0, parent=None, first_hash=0, first_token=0): + """An event whose lengths AGREE, so only the group filter can reject it. + + ``first_token`` matters: a follow-on chunk carries the NEXT tokens, not the + same ones again, and hashing the wrong span makes the chain diverge from + what a query reproduces. + """ + return BlockStored( + block_hashes=[bytes([first_hash + i]) for i in range(blocks)], + parent_block_hash=parent, + token_ids=list(range(first_token, first_token + blocks * BS)), + block_size=BS, + lora_id=None, + group_idx=group_idx, + kv_cache_spec_kind=spec_kind, + ) + + +async def _subscribed(client): + client.on_worker_added(_worker()) + sub = client._subs["w1:30000"] + for t in sub.tasks: + t.cancel() + return sub + + +@pytest.mark.parametrize("kind", ["full_attention", "mla_attention", "sink_full_attention"]) +async def test_attention_groups_are_indexed(kind): + """``mla_attention`` is not optional: Kimi-K3's attention layers are MLA, so + a filter written as ``== "full_attention"`` drops 100% of its usable events + — the same empty view the filter exists to prevent, reached from the other + side.""" + client = KvEventClient() + sub = await _subscribed(client) + + client._handle_event(sub, _stored(blocks=2, spec_kind=kind), rank=0) + + assert len(sub.map_for(0)) == 2 + assert len(sub.view_for(0)) == 2 + + +@pytest.mark.parametrize("kind", ["mamba", "sliding_window", "encoder_only_attention"]) +async def test_non_attention_groups_are_dropped(kind): + client = KvEventClient() + sub = await _subscribed(client) + + client._handle_event(sub, _stored(blocks=2, spec_kind=kind), rank=0) + + assert len(sub.map_for(0)) == 0 + assert len(sub.view_for(0)) == 0 + + +async def test_absent_spec_kind_fails_open(): + """SGLang, and vLLM builds predating the field, send no kind at all. Those + streams were indexed before this filter existed, so a closed default would + silently switch kv-aware routing off for them — the same class of regression + this filter is fixing.""" + client = KvEventClient() + sub = await _subscribed(client) + + client._handle_event(sub, _stored(blocks=2, spec_kind=None), rank=0) + + assert len(sub.map_for(0)) == 2 + + +async def test_a_later_mamba_event_cannot_poison_the_attention_chain(): + """The property the filter exists for, asserted through the query path. + + vLLM mixes no group id into the block hash, so at equal block sizes the + Mamba and attention groups hand out the SAME hash for the same position. + Order decides the damage: when the Mamba event arrives AFTER the attention + event, it overwrites ``map[hash] -> router_hash`` with a hash computed over + its own token span, and the next chunk -- which names that hash as its + parent -- chains off the wrong node. The view then holds a block hash that + no query can ever reproduce. + + Asserted the way routing actually asks: hash the full prompt and count how + many of its blocks are in the view. Checking ``len(map)`` would not catch + this, because the clobber replaces an entry rather than adding one. + """ + client = KvEventClient() + sub = await _subscribed(client) + + # attention group: blocks 0-1 of the prompt, tokens 0..7 + client._handle_event(sub, _stored(blocks=2, spec_kind="mla_attention", group_idx=3), rank=0) + # mamba group: same hashes, DIFFERENT tokens, arriving second + mamba = _stored(blocks=2, spec_kind="mamba", group_idx=0, first_token=100) + client._handle_event(sub, mamba, rank=0) + + # the next chunk: block 2, tokens 8..11, parented on block 1 + client._handle_event( + sub, + _stored( + blocks=1, + spec_kind="mla_attention", + group_idx=3, + parent=bytes([1]), + first_hash=2, + first_token=2 * BS, + ), + rank=0, + ) + + # A query over the whole 3-block prompt must find all three. + want = hash_request(list(range(3 * BS)), BS) + assert len(want) == 3 + hits = sum(1 for h in want if h in sub.view_for(0)) + assert hits == 3, ( + f"{hits}/3 blocks visible: a later non-attention event overwrote the " + "attention group's hash map and the chain continued from the wrong node" + ) + + +async def test_dropping_is_reported_once(caplog): + """Silence here is how the original defect survived: a router that quietly + indexes nothing looks exactly like one with a cold cache.""" + client = KvEventClient() + sub = await _subscribed(client) + + with caplog.at_level(logging.INFO): + for _ in range(3): + client._handle_event(sub, _stored(blocks=2, spec_kind="mamba"), rank=0) + + hits = [r for r in caplog.records if "non-attention groups" in r.getMessage()] + assert len(hits) == 1, "latched: three events must not produce three copies" + assert "mamba" in hits[0].getMessage() From 3ebad45543265ba123875dd45733dd345cf414e4 Mon Sep 17 00:00:00 2001 From: liyingli Date: Fri, 7 Aug 2026 03:28:16 +0000 Subject: [PATCH 50/88] docs(recipes): the validation-status row contradicted the numbers it gates The row said the docker-form figures were being re-measured and that nothing should be quoted until they landed. They landed, and the section 5 rewrite earlier on this branch quotes them -- including the kvd counters the row names explicitly. A reader hitting both has no way to tell which one is current. State what is actually true instead: the figures come from the sweep, and the kvd counters from its single KVD=1 run. The other two rows are unchanged and still say "not run", because the Kubernetes form still has not been. Co-authored-by: Cursor Signed-off-by: liyingli --- examples/recipes/glm5.2-fp8-gfx942/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/recipes/glm5.2-fp8-gfx942/README.md b/examples/recipes/glm5.2-fp8-gfx942/README.md index d4d07bad..c7969a0f 100644 --- a/examples/recipes/glm5.2-fp8-gfx942/README.md +++ b/examples/recipes/glm5.2-fp8-gfx942/README.md @@ -282,7 +282,7 @@ tier; `decode` runs no host tier and gets half. `cpu: 32` matches `num_threads` | What | Status | |---|---| -| This configuration in its **`docker` form** | brought up and benchmarked on 2 × MI300X. **TODO — numbers withheld pending a re-run.** Request count, cache efficiency and the kvd counters are being re-measured on one consistent run; quote nothing here until they land | +| This configuration in its **`docker` form** | brought up and benchmarked on 2 × MI300X. The figures quoted in §5 and in the docker recipe come from that sweep, one axis at a time against a locked baseline; the kvd counters in §5 are from the single `KVD=1` run of it | | **This manifest** | **not run.** Derived from that deployment flag for flag, and every deviation is in §6, but the Kubernetes form has not been brought up | | Native kvd sidecar ordering | not run. The mechanism is standard k8s 1.29+; the claim that it removes the startup race is reasoned, not measured | From 6aab24d92c66755435cab0ae74b0c8d38756383e Mon Sep 17 00:00:00 2001 From: llying-001 Date: Thu, 6 Aug 2026 19:07:59 +0800 Subject: [PATCH 51/88] =?UTF-8?q?fix(kvd,sglang):=20return=20sglang's=20St?= =?UTF-8?q?orageMetrics=20from=20the=20adapter's=20ge=E2=80=A6=20(#93)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(kvd,sglang): return sglang's StorageMetrics from the adapter's get_stats SGLang's hicache observability polls storage_backend.get_stats() once per scheduler step and hands the result straight to StorageMetricsCollector.log_storage_metrics, which does `assert isinstance(storage_metrics, StorageMetrics)`. The adapter returned a dict of daemon counters, so the prefill scheduler died on a bare AssertionError -- exit -3, then SIGQUIT -- about two minutes after startup: scheduler.py _get_new_batch_prefill_raw unified_radix_cache.py check_hicache_events metrics_collector.py assert isinstance(..., StorageMetrics) It needs both switches: a dynamic hicache storage backend AND --enable-metrics, which is what sets enable_storage_metrics. That is why earlier kvd runs never saw it -- the recipe passed no metrics flag, the collector was never built and get_stats() was never called. It is not a version regression: the assert and the per-step poll are identical in v0.5.14, v0.5.15, v0.5.15.post1 and v0.5.16, so both engine images (mi30x on v0.5.16, mi35x on v0.5.15.post1) carry it. get_stats() now returns StorageMetrics, drain-on-read like the built-in backends so every sample is observed exactly once: one sample per batch io, taken in the v2 funnel (_batch_io_v2) and on v1 batch_get / batch_set, carrying pages and GB/s -- per-page bytes reuse _get_pool_buffer_info(). The accumulators are deque(maxlen=4096), not lists: with metrics off nobody drains them and a list would grow for the length of the run. The daemon counters are not lost, they move to daemon_counters(). They are the only way to confirm kvd is really being read and written, and sglang has nowhere to put them. SGLang is the sole get_stats() consumer in the tree; infera/kv/api.py's same-named function is a FastAPI route handler, not a caller. On the import: `except ImportError` used to install a local look-alike dataclass unconditionally. That is right when sglang is absent (unit tests) and wrong when sglang is present but has moved the class, since the assert is against sglang's own type -- a look-alike is the original crash again. The two cases are now distinguished; the second logs critical at import and makes get_stats() return None, the one other value log_storage_metrics accepts. Four tests pin the contract: return type, drain-on-read, bounded accumulators, and None when the type cannot be imported. Four existing assertions that read daemon counters now go through daemon_counters(). Exercised on 2 x MI300X, GLM-5.2-FP8 PD with kvd and --enable-metrics both on: 225/225 requests, no scheduler death. Co-authored-by: Cursor Signed-off-by: liyingli * fix(kvd,sglang): lock the hicache metric deques, drop samples we cannot size Two findings from review of the get_stats() contract fix. Locking. SGLang records into the metric deques from its prefetch and backup worker threads -- cache_controller spawns both -- and drains them from the scheduler thread via check_hicache_events. deque.append is atomic under the GIL, but the drain is an extend-then-clear pair, so a sample landing between its two halves was lost. One lock now covers recording and draining. Sizing. _batch_io_v2 counted a pool's pages even when _get_pool_buffer_info() returned None for it, i.e. for a pool shape we do not recognize and therefore have no page stride for. The byte count then came out zero and the sample published 0 GB/s, which reads as a stalled transfer rather than as absent data. Keeping only the recognized pools' bytes is not better: `elapsed` covers every pool in the batch, so that understates the rate. The whole sample is dropped instead when any transferred pool cannot be sized, and _record_io() rejects a zero byte count on the v1 paths for the same reason. Three tests: _record_io drops an unsized batch; a v2 transfer against a pool with no buffer attribute leaves no sample; a mixed batch -- one sizable pool, one not -- records for the sizable pool alone and then records nothing once the unsized one joins it. The last fails against the previous commit. The get_stats() contract tests stay ungated on _STORAGE_METRICS_MISPLACED on purpose. That flag is only true where sglang is installed and has moved StorageMetrics, and a red suite is exactly the signal wanted there -- gating the assertions would let the import rot while CI stayed green. The degraded path has its own test, which sets the flag explicitly and so runs everywhere. Co-authored-by: Cursor Signed-off-by: liyingli --------- Signed-off-by: liyingli Co-authored-by: Cursor --- infera/engine/sglang/kvd_adapter.py | 156 +++++++++++++- .../engine/sglang/test_sglang_kvd_adapter.py | 199 +++++++++++++++++- 2 files changed, 344 insertions(+), 11 deletions(-) diff --git a/infera/engine/sglang/kvd_adapter.py b/infera/engine/sglang/kvd_adapter.py index fbad0603..03cdb5fc 100644 --- a/infera/engine/sglang/kvd_adapter.py +++ b/infera/engine/sglang/kvd_adapter.py @@ -52,6 +52,9 @@ import logging import os import threading +import time +from collections import deque +from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any @@ -98,6 +101,33 @@ class HiCacheStorageConfig: # type: ignore[no-redef] _SGLANG_AVAILABLE = False +# `get_stats()` is polled every scheduler step by SGLang's hicache +# observability and its return value goes straight into +# `StorageMetricsCollector.log_storage_metrics`, which asserts the exact type. +# Returning anything else (we used to return a dict of daemon counters) kills +# the scheduler with a bare AssertionError. Absent sglang, the local stand-in +# keeps this module importable for unit tests. +try: + from sglang.srt.observability.metrics_collector import StorageMetrics + + _STORAGE_METRICS_FROM_SGLANG = True +except ImportError: # pragma: no cover — exercised only without sglang installed + _STORAGE_METRICS_FROM_SGLANG = False + + @dataclass + class StorageMetrics: # type: ignore[no-redef] + prefetch_pgs: list[int] = field(default_factory=list) + backup_pgs: list[int] = field(default_factory=list) + prefetch_bandwidth: list[float] = field(default_factory=list) + backup_bandwidth: list[float] = field(default_factory=list) + + +# sglang is here but its StorageMetrics is not where we look for it — a future +# release moved or renamed it. The stand-in above satisfies nothing in that +# case: the assert is against sglang's own class, so handing one back is the +# original crash again. Report no metrics instead. +_STORAGE_METRICS_MISPLACED = _SGLANG_AVAILABLE and not _STORAGE_METRICS_FROM_SGLANG + def _torch(): """Lazy torch handle. Raises a clean error if called when torch @@ -110,6 +140,20 @@ def _torch(): logger = logging.getLogger(__name__) +if _STORAGE_METRICS_MISPLACED: # pragma: no cover — needs a future sglang + logger.critical( + "infera: sglang is installed but " + "sglang.srt.observability.metrics_collector.StorageMetrics is not " + "importable. hicache storage metrics (prefetch/backup pages and " + "bandwidth) will be empty until the import in kvd_adapter.py is " + "pointed at the new location." + ) + +# Per-series cap on undrained hicache metric samples. At one sample per batch +# io this is a few seconds of traffic; the deques only fill up when SGLang is +# not polling get_stats(), i.e. when metrics are off and nobody reads them. +_METRICS_SAMPLE_CAP = 4096 + _DEFAULT_SOCKET_PATH = "/var/run/infera-kvd.sock" @@ -308,6 +352,21 @@ def __init__( self._hipfile_read_warned: set[Any] = set() self._hipfile_write_warned: set[Any] = set() + # ----- hicache storage metrics ----- + # One sample per batch io, drained by get_stats(). Bounded rather than + # gated on a flag: SGLang only polls get_stats() when metrics are on, + # and unbounded lists would grow for the entire run when they are off. + # + # The lock is load-bearing, not defensive: SGLang records from its + # prefetch and backup worker threads and drains from the scheduler + # thread, and the drain is a read-then-clear that would otherwise lose + # whatever lands between its two halves. + self._metrics_lock = threading.Lock() + self._prefetch_pgs: deque[int] = deque(maxlen=_METRICS_SAMPLE_CAP) + self._backup_pgs: deque[int] = deque(maxlen=_METRICS_SAMPLE_CAP) + self._prefetch_bandwidth: deque[float] = deque(maxlen=_METRICS_SAMPLE_CAP) + self._backup_bandwidth: deque[float] = deque(maxlen=_METRICS_SAMPLE_CAP) + self._start_background_loop() self._connect_or_raise() @@ -355,7 +414,16 @@ def batch_get( connection IS the throughput path; we'll add it when measured.""" if target_locations is None or len(target_locations) != len(keys): raise ValueError("InferaKvdBackend.batch_get requires aligned target_locations") - return [self.get(k, t) for k, t in zip(keys, target_locations, strict=True)] + started = time.perf_counter() + out = [self.get(k, t) for k, t in zip(keys, target_locations, strict=True)] + hits = [t for t in out if t is not None] + self._record_io( + pages=len(hits), + nbytes=sum(int(t.numel()) * int(t.element_size()) for t in hits), + elapsed=time.perf_counter() - started, + write=False, + ) + return out def set( self, @@ -406,10 +474,22 @@ def batch_set( PoolTransfer can pipeline.""" if values is None or len(values) != len(keys): raise ValueError("InferaKvdBackend.batch_set requires aligned values") + started = time.perf_counter() success = True + written = 0 + nbytes = 0 for k, v in zip(keys, values, strict=True): - if not self.set(k, v): + if self.set(k, v): + written += 1 + nbytes += int(v.numel()) * int(v.element_size()) + else: success = False + self._record_io( + pages=written, + nbytes=nbytes, + elapsed=time.perf_counter() - started, + write=True, + ) return success # ------------------------------------------------------------------ @@ -548,6 +628,10 @@ def _batch_io_v2(self, transfers: Any, *, write: bool, extra_info: Any) -> dict: if write else self._retention_default ) + started = time.perf_counter() + pages_done = 0 + bytes_done = 0 + unsized_pages = False for transfer in transfers: name = transfer.name keys = list(transfer.keys or []) @@ -589,6 +673,25 @@ def _batch_io_v2(self, transfers: Any, *, write: bool, extra_info: Any) -> dict: else: per_key.append(self._v2_read_page(name, key, host_pool, page_offset)) results[name] = per_key + pages = sum(per_key) + info = self._get_pool_buffer_info(host_pool) + if info is None: + # Pool shape we don't recognize, so we have no page stride and + # no way to size the transfer. `elapsed` covers every pool in + # the batch, so counting this one's pages with zero bytes — or + # dropping only its bytes — understates the rate. Drop the + # whole sample; registration already logged the unknown shape. + unsized_pages = unsized_pages or pages > 0 + continue + pages_done += pages + bytes_done += pages * info[2] + if not unsized_pages: + self._record_io( + pages=pages_done, + nbytes=bytes_done, + elapsed=time.perf_counter() - started, + write=write, + ) return results def _pool_storage_key(self, pool_name: Any, key: str) -> str: @@ -949,10 +1052,51 @@ def clear(self) -> None: on the same kvd are untouched.""" self._run_async(self._client.clear(model=self._model, compat_key=self._compat_key)) - def get_stats(self) -> dict | None: - """Daemon-side counters. SGLang's hicache observability picks - this up; we surface the same numbers to Prometheus on the server - via `/v1/kv-stats`.""" + def _record_io(self, *, pages: int, nbytes: int, elapsed: float, write: bool) -> None: + """Append one batch-io sample. Bandwidth is GB/s over the whole + batch, matching what SGLang's built-in backends report. A batch we + cannot size is dropped rather than published as 0 GB/s — an absent + sample reads as "no data", a zero reads as "the transfer stalled".""" + if pages <= 0 or nbytes <= 0 or elapsed <= 0: + return + gb_per_s = (nbytes / elapsed) / 1e9 + with self._metrics_lock: + if write: + self._backup_pgs.append(pages) + self._backup_bandwidth.append(gb_per_s) + else: + self._prefetch_pgs.append(pages) + self._prefetch_bandwidth.append(gb_per_s) + + def get_stats(self) -> StorageMetrics | None: + """Drain the batch-io samples for SGLang's hicache observability. + + Polled every scheduler step. The return type is load-bearing: + `StorageMetricsCollector.log_storage_metrics` asserts on it, so a + wrong type takes the scheduler down rather than degrading metrics. + `None` is the one other value it accepts, and the only safe answer + when we cannot construct sglang's own class. Daemon-side counters + live in `daemon_counters()` — they are a different thing and SGLang + has nowhere to put them. + """ + if _STORAGE_METRICS_MISPLACED: # pragma: no cover — needs a future sglang + return None + metrics = StorageMetrics() + with self._metrics_lock: + metrics.prefetch_pgs.extend(self._prefetch_pgs) + metrics.backup_pgs.extend(self._backup_pgs) + metrics.prefetch_bandwidth.extend(self._prefetch_bandwidth) + metrics.backup_bandwidth.extend(self._backup_bandwidth) + self._prefetch_pgs.clear() + self._backup_pgs.clear() + self._prefetch_bandwidth.clear() + self._backup_bandwidth.clear() + return metrics + + def daemon_counters(self) -> dict | None: + """Daemon-side totals (entries, bytes, hits/misses/evictions). + Used to confirm kvd is actually being read and written, which the + per-step `get_stats()` histograms cannot show on their own.""" try: stats = self._run_async(self._client.stats()) except (KvdConnectionError, KvdProtocolError): diff --git a/tests/engine/sglang/test_sglang_kvd_adapter.py b/tests/engine/sglang/test_sglang_kvd_adapter.py index 47a07b08..8a433e65 100644 --- a/tests/engine/sglang/test_sglang_kvd_adapter.py +++ b/tests/engine/sglang/test_sglang_kvd_adapter.py @@ -217,7 +217,7 @@ async def test_adapter_init_connects(kvd_daemon, patched_adapter_codecs): ) try: # Round-trip a stats call as a liveness probe. - stats = await asyncio.to_thread(backend.get_stats) + stats = await asyncio.to_thread(backend.daemon_counters) assert stats is not None assert stats["entries"] == 0 finally: @@ -349,7 +349,7 @@ async def test_adapter_get_size_mismatch_returns_none(kvd_daemon, patched_adapte @pytest.mark.asyncio -async def test_adapter_get_stats_after_traffic(kvd_daemon, patched_adapter_codecs): +async def test_adapter_daemon_counters_after_traffic(kvd_daemon, patched_adapter_codecs): socket = kvd_daemon backend = await asyncio.to_thread(InferaKvdBackend, _make_config(), socket_path=socket) try: @@ -358,7 +358,7 @@ async def test_adapter_get_stats_after_traffic(kvd_daemon, patched_adapter_codec await asyncio.to_thread(backend.get, "a", target) # hit await asyncio.to_thread(backend.get, "absent", _FakeTensor(b"\x00")) # miss - stats = await asyncio.to_thread(backend.get_stats) + stats = await asyncio.to_thread(backend.daemon_counters) assert stats is not None # Counts include the auxiliary stats call's own GETs etc; just # sanity-check ranges. @@ -369,6 +369,106 @@ async def test_adapter_get_stats_after_traffic(kvd_daemon, patched_adapter_codec await asyncio.to_thread(backend.close) +# ---------------------------------------------------------------------- +# hicache storage metrics +# +# SGLang polls get_stats() every scheduler step and feeds the result +# straight into StorageMetricsCollector.log_storage_metrics, which does +# `assert isinstance(storage_metrics, StorageMetrics)`. Returning the +# daemon-counter dict here used to take the prefill scheduler down with a +# bare AssertionError the moment KVD and --enable-metrics were both on. +# ---------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_stats_returns_storage_metrics_type(kvd_daemon, patched_adapter_codecs): + """The type is the contract: anything else crashes the scheduler.""" + backend = await asyncio.to_thread(InferaKvdBackend, _make_config(), socket_path=kvd_daemon) + try: + metrics = await asyncio.to_thread(backend.get_stats) + assert isinstance(metrics, kvd_adapter.StorageMetrics) + # Every field the collector iterates must exist and be iterable. + assert metrics.prefetch_pgs == [] + assert metrics.backup_pgs == [] + assert metrics.prefetch_bandwidth == [] + assert metrics.backup_bandwidth == [] + finally: + await asyncio.to_thread(backend.close) + + +@pytest.mark.asyncio +async def test_get_stats_records_then_drains_batch_io(kvd_daemon, patched_adapter_codecs): + """Batch ops leave one sample per call; get_stats drains them so the + histograms observe each sample exactly once.""" + backend = await asyncio.to_thread(InferaKvdBackend, _make_config(), socket_path=kvd_daemon) + try: + keys = ["m1", "m2"] + values = [_FakeTensor(b"aaaa"), _FakeTensor(b"bbbb")] + await asyncio.to_thread(backend.batch_set, keys, values) + await asyncio.to_thread(backend.batch_get, keys, [_FakeTensor(b"\x00" * 4) for _ in keys]) + + metrics = await asyncio.to_thread(backend.get_stats) + assert metrics.backup_pgs == [2], "one sample per batch_set, 2 pages written" + assert metrics.prefetch_pgs == [2], "one sample per batch_get, 2 pages read" + assert len(metrics.backup_bandwidth) == 1 + assert len(metrics.prefetch_bandwidth) == 1 + assert all(b > 0 for b in metrics.backup_bandwidth + metrics.prefetch_bandwidth) + + # Drained: a second poll with no traffic in between is empty. + again = await asyncio.to_thread(backend.get_stats) + assert again.prefetch_pgs == [] + assert again.backup_pgs == [] + finally: + await asyncio.to_thread(backend.close) + + +@pytest.mark.asyncio +async def test_get_stats_returns_none_when_sglang_moved_the_type( + kvd_daemon, patched_adapter_codecs, monkeypatch +): + """If a future sglang moves StorageMetrics, our local stand-in would + fail the same isinstance assert. None is the only other value + log_storage_metrics accepts, so metrics go empty instead of fatal.""" + monkeypatch.setattr(kvd_adapter, "_STORAGE_METRICS_MISPLACED", True) + backend = await asyncio.to_thread(InferaKvdBackend, _make_config(), socket_path=kvd_daemon) + try: + assert await asyncio.to_thread(backend.get_stats) is None + finally: + await asyncio.to_thread(backend.close) + + +@pytest.mark.asyncio +async def test_record_io_drops_a_batch_it_cannot_size(kvd_daemon, patched_adapter_codecs): + """A transfer whose byte count we could not derive must not surface as + a 0 GB/s sample: absent reads as "no data", zero reads as "stalled".""" + backend = await asyncio.to_thread(InferaKvdBackend, _make_config(), socket_path=kvd_daemon) + try: + backend._record_io(pages=4, nbytes=0, elapsed=0.001, write=False) + backend._record_io(pages=4, nbytes=0, elapsed=0.001, write=True) + + metrics = await asyncio.to_thread(backend.get_stats) + assert metrics.prefetch_pgs == [] + assert metrics.backup_pgs == [] + assert metrics.prefetch_bandwidth == [] + assert metrics.backup_bandwidth == [] + finally: + await asyncio.to_thread(backend.close) + + +@pytest.mark.asyncio +async def test_get_stats_samples_are_bounded(kvd_daemon, patched_adapter_codecs): + """With metrics off nobody drains, so the accumulators must not grow + without bound for the length of a run.""" + backend = await asyncio.to_thread(InferaKvdBackend, _make_config(), socket_path=kvd_daemon) + try: + cap = kvd_adapter._METRICS_SAMPLE_CAP + for _ in range(cap + 10): + backend._record_io(pages=1, nbytes=1024, elapsed=0.001, write=False) + assert len(backend._prefetch_pgs) == cap + finally: + await asyncio.to_thread(backend.close) + + # ---------------------------------------------------------------------- # Torch-required: real tensor round-trip # ---------------------------------------------------------------------- @@ -705,6 +805,23 @@ def set_from_flat_data_page(self, index, data_page: _FakeTensor) -> None: self._storage[i : i + self.page_size] = bytes(data_page.payload) +class _SizedFakeHostPool(_FakeHostPool): + """A pool the adapter *can* size, i.e. one that carries the MLA-shaped + `kv_buffer` that `_get_pool_buffer_info` duck-types on. Plain + `_FakeHostPool` has no buffer attribute at all, so the two together + cover both arms of the metric sampling.""" + + def __init__(self, n_slots: int, page_size: int = 2) -> None: + super().__init__(n_slots, page_size) + self.kv_buffer = SimpleNamespace( + data_ptr=lambda: 4096, + numel=lambda: n_slots, + element_size=lambda: 1, + ) + self.kv_cache_dim = 1 + self.dtype = SimpleNamespace(itemsize=1) + + class _FakeIndices: """Stand-in for a torch.LongTensor of host indices. Supports `numel()` and integer indexing returning an obj with `.item()`.""" @@ -831,7 +948,7 @@ async def test_batch_set_v2_roundtrip_kv_and_indexer( assert set_res[PoolName.INDEXER] == [True, True] # Daemon should report 4 stored entries (2 pages × 2 pools). - stats = await asyncio.to_thread(backend.get_stats) + stats = await asyncio.to_thread(backend.daemon_counters) assert stats["entries"] == 4 # Wipe and read back through v2 → byte-perfect restore. @@ -846,12 +963,84 @@ async def test_batch_set_v2_roundtrip_kv_and_indexer( assert idx_pool._storage[s] == 0xA0 + s, f"idx slot {s} mismatch" # kvd hit counter should advance by 4 (we read 2×2 pages). - stats_after = await asyncio.to_thread(backend.get_stats) + stats_after = await asyncio.to_thread(backend.daemon_counters) assert stats_after["hits_total"] >= 4 finally: await asyncio.to_thread(backend.close) +@pytest.mark.asyncio +async def test_batch_io_v2_records_no_sample_for_a_pool_it_cannot_size( + kvd_daemon, patched_adapter_codecs, fake_sglang_v2_types +): + """`_FakeHostPool` carries neither buffer attribute, so the adapter has + no page stride for it. The transfer still has to work, and it has to + leave no bandwidth sample rather than one reading 0 GB/s.""" + PoolName = fake_sglang_v2_types.PoolName + PoolHitPolicy = fake_sglang_v2_types.PoolHitPolicy + + backend = await asyncio.to_thread( + kvd_adapter.InferaKvdBackend, + _make_config(), + socket_path=kvd_daemon, + client_id="v2-unsized-metrics", + ) + try: + kv_pool = _FakeHostPool(n_slots=8, page_size=2) + assert backend._get_pool_buffer_info(kv_pool) is None, "pool must be unrecognized" + backend.registered_pools = {PoolName.KV: kv_pool} + + keys = ["unsized_alpha", "unsized_beta"] + kv_t = _make_transfer(PoolName, PoolHitPolicy, PoolName.KV, keys, [0, 1, 2, 3]) + assert await asyncio.to_thread(backend.batch_set_v2, [kv_t]) == {PoolName.KV: [True, True]} + + metrics = await asyncio.to_thread(backend.get_stats) + assert metrics.backup_pgs == [] + assert metrics.backup_bandwidth == [] + finally: + await asyncio.to_thread(backend.close) + + +@pytest.mark.asyncio +async def test_batch_io_v2_drops_the_sample_when_one_pool_is_unsized( + kvd_daemon, patched_adapter_codecs, fake_sglang_v2_types +): + """A batch spans several pools but is timed as a whole. If one pool's + bytes are unknowable, keeping the others' would divide real bytes by + the whole batch's time and understate the rate, so the sample goes.""" + PoolName = fake_sglang_v2_types.PoolName + PoolHitPolicy = fake_sglang_v2_types.PoolHitPolicy + + backend = await asyncio.to_thread( + kvd_adapter.InferaKvdBackend, + _make_config(), + socket_path=kvd_daemon, + client_id="v2-mixed-metrics", + ) + try: + sized = _SizedFakeHostPool(n_slots=8, page_size=2) + unsized = _FakeHostPool(n_slots=8, page_size=2) + assert backend._get_pool_buffer_info(sized) is not None + assert backend._get_pool_buffer_info(unsized) is None + backend.registered_pools = {PoolName.KV: sized, PoolName.INDEXER: unsized} + + # Sized pool alone: one sample, and it is not a zero. + kv_t = _make_transfer(PoolName, PoolHitPolicy, PoolName.KV, ["mixed_a"], [0, 1]) + await asyncio.to_thread(backend.batch_set_v2, [kv_t]) + metrics = await asyncio.to_thread(backend.get_stats) + assert metrics.backup_pgs == [1] + assert metrics.backup_bandwidth and metrics.backup_bandwidth[0] > 0 + + # Same batch plus the unsized pool: nothing recorded at all. + idx_t = _make_transfer(PoolName, PoolHitPolicy, PoolName.INDEXER, ["mixed_b"], [0, 1]) + await asyncio.to_thread(backend.batch_set_v2, [kv_t, idx_t]) + after = await asyncio.to_thread(backend.get_stats) + assert after.backup_pgs == [] + assert after.backup_bandwidth == [] + finally: + await asyncio.to_thread(backend.close) + + @pytest.mark.asyncio async def test_batch_exists_v2_truncates_on_missing_kv( kvd_daemon, patched_adapter_codecs, fake_sglang_v2_types From 221701cdfbcb96c63bd1ee30e925a3696663ab27 Mon Sep 17 00:00:00 2001 From: xiaobochen-amd Date: Fri, 7 Aug 2026 02:00:13 +0000 Subject: [PATCH 52/88] ci: key the concurrency group on the PR, not the branch name github.head_ref is the source branch name with no repository attached, so two pull requests from different forks that both use a branch called `main` or `fix-ci` -- neither an unusual name -- land in one concurrency group and cancel each other. The symptom is somebody else's push killing your run, with nothing in either run pointing at the other. The PR number is unique across forks. Pushes to main have no pull_request in the payload and fall through to the ref, which is what they used before. Co-authored-by: Cursor Signed-off-by: xiaobochen-amd --- .github/workflows/ci.yml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 116cfa07..c46ed61e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -26,7 +26,11 @@ on: default: false concurrency: - group: ${{ github.workflow }}-${{ github.head_ref || github.ref }} + # PR number, not branch name: head_ref carries no repository, so two forks + # that both call a branch `main` or `fix-ci` would share a group and cancel + # each other's runs -- which reads as "my CI vanished" and is near impossible + # to trace back. Falls back to the ref for pushes to main. + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} cancel-in-progress: true permissions: From 421a5f4bf43ba77856d1eb5a85a25dea0b59ed62 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Fri, 7 Aug 2026 02:07:03 +0000 Subject: [PATCH 53/88] test(router): replay a real Kimi-K3 event stream, not just the filter's rejects The per-field tests prove the group filter rejects what it should. They do not prove a request can HIT afterwards, and that is the property that matters -- production still reported cache_hits=0 with the filter in place, and nothing in the suite would have noticed. This replays what a live worker actually emits, captured off the wire: four events per prefill chunk at block_size 768, three Mamba groups reporting the whole 3840-token span against a single hash and one MLA group reporting it against five, with the Mamba groups repeating the attention group's hashes because vLLM mixes no group id into them. Two chunks, not one, because a single chunk never exercises the parent map: chunk two names chunk one's last block as its parent, and resolving that lookup against a colliding Mamba hash is the failure mode the filter exists to prevent. Then it asks the question routing asks -- hash the prompt, count its blocks in the view -- and requires all ten. Also asserts a divergent tail hits the shared prefix and stops, since a filter that over-matched would route a new prompt onto another prefix's KV. Signed-off-by: Zhang, Jiejing --- .../router/test_kv_event_kimi_k3_shapes.py | 140 ++++++++++++++++++ 1 file changed, 140 insertions(+) create mode 100644 tests/unit/router/test_kv_event_kimi_k3_shapes.py diff --git a/tests/unit/router/test_kv_event_kimi_k3_shapes.py b/tests/unit/router/test_kv_event_kimi_k3_shapes.py new file mode 100644 index 00000000..76dfbc02 --- /dev/null +++ b/tests/unit/router/test_kv_event_kimi_k3_shapes.py @@ -0,0 +1,140 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""End-to-end replay of the event shapes a real Kimi-K3 worker emits. + +The per-field unit tests in ``test_kv_event_group_filter`` prove the filter +rejects what it should. They do not prove a request can HIT afterwards, which is +the property that actually matters and the one that stayed broken in production +after the filter went in: ``cache_hits=0`` on every routing decision. + +Captured from a live worker (block_size 768, ``--max-num-batched-tokens 4096``, +so chunked prefill splits at ``4096 // 768 * 768 = 3840``): + + kv_cache_spec_kind group_idx token_ids block_hashes + mamba 0 3840 1 + mamba 1 3840 1 + mamba 2 3840 1 + mla_attention 3 3840 5 + +Four groups per chunk, all at the same block size, and vLLM mixes no group id +into the hash -- so the Mamba groups' hashes collide with the attention group's. +""" + +from __future__ import annotations + +from infera.common.worker_pool import EngineType, WorkerInfo +from infera.router.kv_event.client import KvEventClient +from infera.router.kv_event.events import BlockStored +from infera.router.kv_event.hasher import hash_request + +BS = 768 +CHUNK_BLOCKS = 5 # 3840 tokens / 768 +CHUNK = BS * CHUNK_BLOCKS + + +def _worker(): + return WorkerInfo( + worker_id="w1:30000", + url="http://w1:30000", + model_name="m", + engine=EngineType.VLLM, + kv_events_endpoint="tcp://w1:5555", + kv_block_size=BS, + ) + + +async def _subscribed(client): + client.on_worker_added(_worker()) + sub = client._subs["w1:30000"] + for t in sub.tasks: + t.cancel() + return sub + + +def _chunk_events(tokens, first_block, parent): + """One prefill chunk as four events, in the order vLLM emits them. + + The Mamba groups repeat the attention group's hashes because the engine + derives them from the same token span with no group id mixed in. Each + reports the whole span against a single hash -- align mode keeps exactly one + real block per step and nulls the rest. + """ + hashes = [bytes([first_block + i]) for i in range(len(tokens) // BS)] + common = dict( + parent_block_hash=parent, + token_ids=list(tokens), + block_size=BS, + lora_id=None, + ) + return [ + BlockStored(block_hashes=[hashes[-1]], group_idx=g, kv_cache_spec_kind="mamba", **common) + for g in (0, 1, 2) + ] + [ + BlockStored(block_hashes=hashes, group_idx=3, kv_cache_spec_kind="mla_attention", **common) + ] + + +async def test_a_two_chunk_prefill_is_fully_hittable(): + """The property production needs: replay a prompt's events, then ask for the + same prompt and get every block back. + + Two chunks, because a single one never exercises the parent map: chunk two + names chunk one's last block as its parent, and that lookup is exactly what + a colliding Mamba hash used to break. + """ + client = KvEventClient() + sub = await _subscribed(client) + + prompt = list(range(2 * CHUNK)) + for ev in _chunk_events(prompt[:CHUNK], first_block=0, parent=None): + client._handle_event(sub, ev, rank=0) + for ev in _chunk_events( + prompt[CHUNK:], first_block=CHUNK_BLOCKS, parent=bytes([CHUNK_BLOCKS - 1]) + ): + client._handle_event(sub, ev, rank=0) + + want = hash_request(prompt, BS) + assert len(want) == 2 * CHUNK_BLOCKS + hits = sum(1 for h in want if h in sub.view_for(0)) + assert hits == len(want), ( + f"{hits}/{len(want)} blocks visible after replaying a two-chunk prefill; " + "a request for this prompt would route as if the cache were cold" + ) + + +async def test_a_partial_prefix_hits_its_prefix_only(): + """A shorter prompt sharing the first chunk must hit exactly that chunk -- + not more (which would be a false hit onto another prefix) and not less.""" + client = KvEventClient() + sub = await _subscribed(client) + + prompt = list(range(2 * CHUNK)) + for ev in _chunk_events(prompt[:CHUNK], first_block=0, parent=None): + client._handle_event(sub, ev, rank=0) + for ev in _chunk_events( + prompt[CHUNK:], first_block=CHUNK_BLOCKS, parent=bytes([CHUNK_BLOCKS - 1]) + ): + client._handle_event(sub, ev, rank=0) + + shared = hash_request(prompt[:CHUNK], BS) + assert sum(1 for h in shared if h in sub.view_for(0)) == CHUNK_BLOCKS + + divergent = hash_request(list(range(CHUNK)) + list(range(9_000, 9_000 + CHUNK)), BS) + hits = sum(1 for h in divergent if h in sub.view_for(0)) + assert hits == CHUNK_BLOCKS, f"{hits} hits: a divergent tail must not match" + + +async def test_the_view_holds_only_the_attention_group(): + """Four groups arrive per chunk; only one is indexable, so the view must be + one chunk's worth of blocks, not four.""" + client = KvEventClient() + sub = await _subscribed(client) + + for ev in _chunk_events(list(range(CHUNK)), first_block=0, parent=None): + client._handle_event(sub, ev, rank=0) + + assert len(sub.view_for(0)) == CHUNK_BLOCKS + assert len(sub.map_for(0)) == CHUNK_BLOCKS From de65a9de4265c664f151ac52f5ea73d7923becda Mon Sep 17 00:00:00 2001 From: yihou Date: Thu, 6 Aug 2026 05:51:11 +0000 Subject: [PATCH 54/88] chore(examples): drop the GLM-5.2 results/ directory The three files under examples/sglang_1p1d_glm5.2/results/ carry customer-benchmark numbers and cluster detail that should not ship in a public example kit. Remove the directory and drop the six links that pointed into it (five in README.md, one in engine/bench.sh), keeping the surrounding claims intact without the published numbers. Signed-off-by: yihou --- examples/sglang_1p1d_glm5.2/README.md | 14 +- examples/sglang_1p1d_glm5.2/engine/bench.sh | 3 +- examples/sglang_1p1d_glm5.2/results/README.md | 88 -------- .../results/customer_agentx_caseA_conc8.md | 212 ------------------ .../results/infera_agenticbench_conc8.md | 145 ------------ 5 files changed, 7 insertions(+), 455 deletions(-) delete mode 100644 examples/sglang_1p1d_glm5.2/results/README.md delete mode 100644 examples/sglang_1p1d_glm5.2/results/customer_agentx_caseA_conc8.md delete mode 100644 examples/sglang_1p1d_glm5.2/results/infera_agenticbench_conc8.md diff --git a/examples/sglang_1p1d_glm5.2/README.md b/examples/sglang_1p1d_glm5.2/README.md index f438b4ef..4cd4d3c3 100644 --- a/examples/sglang_1p1d_glm5.2/README.md +++ b/examples/sglang_1p1d_glm5.2/README.md @@ -14,8 +14,8 @@ bash cluster/cluster.peermem.sh up # (or cluster.dmabuf.sh) — bring bash cluster/cluster.peermem.sh smoke # prove it works ``` -Measured performance for this exact shape, under two independent agentic benchmarks -on two different fabrics, is in [`results/`](results/README.md). +This exact shape has been measured under two independent agentic benchmarks on two +different fabrics; those numbers are not published with this kit. ## Contents @@ -29,7 +29,6 @@ on two different fabrics, is in [`results/`](results/README.md). | `engine/bench.sh` | reference throughput sweep using SGLang's own `bench_serving` | | `engine/down.sh` | tear down and wait for VRAM to actually free | | `preflight_rdma.sh` | RDMA preflight: registration-mode probe + cross-node fabric measurement | -| [`results/`](results/README.md) | measured agentic-benchmark numbers at concurrency 8 | ## Topology @@ -216,14 +215,13 @@ flags are load-bearing in ways that are not obvious from the flag name: passes it). Its column is nonetheless meaningless on this dataset: `--dataset-name random` builds every prompt independently, so there is **no shared prefix by construction** and any nonzero value is residue from the previous round. Prefix reuse - is an agentic-workload property — see [`results/`](results/README.md). + is an agentic-workload property. `--num-prompts` is recomputed per concurrency (`10 × C`), so each arm of a sweep gets enough requests to reach steady state. -**This kit ships no agentic benchmark client**, by design. `results/` documents what -the agentic numbers look like and how to point the customer's harness at this -deployment. +**This kit ships no agentic benchmark client**, by design. Point the customer's own +harness at the router endpoint. ## 6. Tear down @@ -327,7 +325,7 @@ Stated plainly rather than implied. | what | status | |---|---| -| the deployment **shape** this kit encodes (1P1D + mooncake + DPA + MTP + kvd + kv-aware) | **validated end-to-end on two clusters**, both fabric types, with the agentic results in [`results/`](results/README.md) | +| the deployment **shape** this kit encodes (1P1D + mooncake + DPA + MTP + kvd + kv-aware) | **validated end-to-end on two clusters**, both fabric types, under a real agentic workload | | the tuned values (GMU, chunk, ctx, EAGLE settings, DSA env, router weights) | **validated** — each is carried over from a run that completed cleanly | | the three traps in Notes 1–3 | **first-hand**, each found by a run that failed or silently mis-measured | | **these scripts as written** | **validated** — `preflight_rdma.sh mode` → `up` → `smoke` → `bench` → `down` on a 2-node MI355X mode-B cluster, with no edits outside `cluster/cluster.dmabuf.sh`. Long context checked separately (needle, to 238K tokens) and under a real agentic workload at concurrency 8 | diff --git a/examples/sglang_1p1d_glm5.2/engine/bench.sh b/examples/sglang_1p1d_glm5.2/engine/bench.sh index a73be1cf..1d3c1d1a 100755 --- a/examples/sglang_1p1d_glm5.2/engine/bench.sh +++ b/examples/sglang_1p1d_glm5.2/engine/bench.sh @@ -3,8 +3,7 @@ # SPDX-License-Identifier: MIT # what: a REFERENCE throughput sweep against the router, using sglang's own bench_serving. # why : it ships inside the engine image, so there is nothing extra to install, and it is -# enough to confirm the deployment performs sanely. It is NOT the agentic benchmark — -# see results/ for what the agentic numbers look like and where those harnesses live. +# enough to confirm the deployment performs sanely. It is NOT the agentic benchmark. # how : bash cluster/.sh bench [conc ...] set -euo pipefail DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"; source "$DIR/../common.sh" diff --git a/examples/sglang_1p1d_glm5.2/results/README.md b/examples/sglang_1p1d_glm5.2/results/README.md deleted file mode 100644 index 7928c4b4..00000000 --- a/examples/sglang_1p1d_glm5.2/results/README.md +++ /dev/null @@ -1,88 +0,0 @@ -# Measured results — GLM-5.2 1P1D at concurrency 8 - -What this deployment actually did, under **two independent agentic benchmarks**, on -**two clusters with different RDMA fabrics**. Every number here comes from a -completed run against the deployment shape this kit ships; nothing is projected. - -| file | what | -|---|---| -| [`infera_agenticbench_conc8.md`](infera_agenticbench_conc8.md) | infera's own agentic benchmark — closed-loop session driver, ~67-minute window, both clusters | -| [`customer_agentx_caseA_conc8.md`](customer_agentx_caseA_conc8.md) | the customer's AgentX Case-A — frozen-trace replay, 900 s window, both clusters | - -## Read this before comparing any two numbers here - -Three axes vary across these runs, and only one of them is the deployment. - -**1. The two benchmarks apply load differently, so their concurrency numbers are -not the same operating point.** - -Both are **closed-loop** — a request is not issued until the previous one on that -session/worker has returned, so a slower server reduces offered load in either. What -differs is what "8" counts: - -| | infera AgenticBench | customer AgentX Case-A | -|---|---|---| -| what "conc 8" means | 8 *initial sessions*, population **grows** to a cap of 32; in-flight is emergent and lands wherever the server allows (mean 12.1) | 8 *fixed workers*; in-flight is bounded by 8 and lands below it (mean 5.1) | -| requests | built live by the driver | frozen in 200 session files, byte-identical every replay | -| window | ramp 400 s + sustain 3,600 s | 900 s | - -So the two ran at very different occupancies on the same server, and a TTFT gap -between them is expected — it is **not** evidence about the deployment. - -The two agree where the load model does not matter: **ITL/TPOT p50 ≈ 14–16 ms** and -**prefix cache ≈ 88 %**, measured independently by both drivers on both clusters. - -**2. The two clusters differ in the fabric, and in two configuration values that -follow from it.** See [Cross-cluster](#cross-cluster-why-the-single-rail-cluster-looks-slower) -below — the short version is that the cross-cluster rows are **context, not a -measurement**, because more than one variable moves. - -**3. The workload shape is the same in both benches, and reproduces.** Both are -built from the same Case-A profile — ISL p50/p90/p99 ≈ 74K/155K/235K, OSL p50/p90 ≈ -320/3,300, ~89 % prefix reuse — and the realised distributions land within ~7 % of -each other on every axis. Workload shape is therefore never the explanation for a -latency difference between them. The one axis where the customer harness's corpus -does *not* meet the spec is the prefix-reuse rate, which it hits per turn but not -overall; see -[the construction analysis](customer_agentx_caseA_conc8.md#the-harnesss-cache-hit-construction-does-not-meet-the-workload-spec). - -## Cross-cluster: why the single-rail cluster looks slower - -Running the same workload on both clusters, the single-rail one is **~1.6× slower on -TTFT p50** while matching on TPOT, cache hit rate and success rate. - -| | multi-rail cluster | single-rail cluster | -|---|---|---| -| RDMA | 8 rails, peer-mem loaded (mode A) | 1 ODP rail, no peer-mem, dma-buf (mode B) | -| aggregate KV bandwidth (preflight's own metric) | ~3,200 Gb/s | ~200 Gb/s | -| TTFT p50 (infera bench, sustain) | 1,365 ms | 2,239 ms | -| TTFT p90 | 4,903 ms | 6,389 ms | -| TPOT p50 | 14.8 ms | 16.1 ms | -| cache hit | 88.8 % | 88.7 % | - -**No single-variable experiment separates these causes.** Four candidates, with the -evidence behind each stated honestly: - -| # | candidate | evidence | what would settle it | -|---|---|---|---| -| 1 | **Fabric.** 1 rail vs 8, and dma-buf vs peer-mem registration. | First-hand node facts on both clusters. **No controlled experiment.** | Not separable — the fabric is the cluster. | -| 2 | **`--chunked-prefill-size` differs**: 65,536 global on the single-rail runs, 16,384 on the multi-rail ones — 4× the per-forward prefill work. | First-hand from both runs' resolved args. The two source recipes genuinely disagreed on this value and the disagreement was recorded rather than resolved. | Re-run one cluster at the other's chunk. Cheapest of the four. | -| 3 | **`--mem-fraction-static` differs**: 0.70 vs 0.80 on prefill → a smaller KV pool (−19 % measured when this was changed within one cluster). | First-hand. Forced, not chosen: 0.80 does not boot the single-rail cluster's prefill leg. | Only separable if 0.80 can be made to boot there. | -| 4 | **MTP and decode-side radix cache are mutually exclusive upstream**, so `decode_prefix_len` is always 0 and **every turn re-transfers the entire prompt KV**. A prefill-side cache hit saves *compute*, not *bytes*. | First-hand: SGLang raises on `--disaggregation-decode-enable-radix-cache` together with `--speculative-algorithm`. | This does not act alone — it **amplifies** #1, by putting a full-prompt KV transfer on every turn's critical path. | - -Candidate 4 is the reason to expect #1 to matter *on this workload specifically*: at -~86K mean input tokens and ~89 % prefix reuse, the deployment re-sends the whole -prompt's KV every turn regardless of how well the cache is working. A workload with -short prompts would be far less fabric-sensitive. - -**None of the four is confirmed.** The honest summary is that the single-rail cluster -is slower on TTFT by a factor that is consistent with its fabric, and that two -configuration deltas ride along with the fabric and are not controlled for. - -## Provenance - -Numbers are recomputed from raw per-request records where those exist, not copied -from a summary line. The customer-bench ladders come from aiperf's -`profile_export.jsonl`; the infera-bench ladders from the driver's own -`summary.json` sustain-phase block, with the ramp window excluded exactly as that -driver defines it. diff --git a/examples/sglang_1p1d_glm5.2/results/customer_agentx_caseA_conc8.md b/examples/sglang_1p1d_glm5.2/results/customer_agentx_caseA_conc8.md deleted file mode 100644 index 297e1e35..00000000 --- a/examples/sglang_1p1d_glm5.2/results/customer_agentx_caseA_conc8.md +++ /dev/null @@ -1,212 +0,0 @@ -# Customer benchmark — AgentX Case-A at concurrency 8 - -The customer-supplied agentic benchmark from -[ROCm/MAD PR #173](https://github.com/ROCm/MAD/pull/173) (`scripts/AgentX_CaseA/`), -replayed **unmodified** against exactly the deployment shape this kit ships. - -**This kit ships no runner for it.** The harness is the customer's, lives upstream, -and is configured entirely through environment variables — see -[Pointing it at this deployment](#pointing-it-at-this-deployment) at the bottom. - -## What the benchmark is - -A **deterministic, spec-constructed replay trace** plus a generic driver: - -- `gen_caseA_conformance.py` synthesises 200 sessions / 1,778 requests from the - Case-A parameters at a fixed seed, so the corpus is byte-identical every run. -- The corpus records **demand only** — per-request input/output token counts, turn - structure, think-times, KV prefix reuse. Nothing engine- or topology-specific. -- `replay_caseA.sh` replays it via **aiperf** (a SemiAnalysis fork, scenario - `inferencex-agentx-mvp`, `--custom-dataset-type weka_trace`) against any - OpenAI-compatible endpoint. - -It is **closed-loop**, like the infera bench: `--concurrency N` runs N workers, and a -worker does not issue its next request until the previous one has returned. Offered -load is therefore bounded by the server, not pushed at it — which the measurement -confirms, in-flight sitting at mean 5.13 against a cap of 8. What differs from the -infera bench is not the loop but *what N counts*: N is a fixed worker population here, -where the infera bench starts 8 sessions and lets the population grow. - -The frozen trace is what makes this benchmark useful: replaying it against a -different deployment changes exactly one variable. - -## Results at c8 - -Both clusters, the same frozen trace, the same deployment shape: - -| | multi-rail cluster | single-rail cluster | -|---|---|---| -| profiling requests | 231 | 231 | -| window | 901 s | 907 s | -| **errors / cancelled / context overflows** | **0 / 0 / 0** | **2 / 0 / 0** | -| in-flight max / mean | 8 / 5.13 | 8 / 4.98 | -| request rate | 0.256 req/s | 0.255 req/s | -| output token rate | 268 tok/s | 265 tok/s | -| **TTFT p50** | **5,146 ms** | 6,698 ms | -| TTFT p90 | 19,780 ms | 23,871 ms | -| TTFT p99 | 31,014 ms | 33,972 ms | -| **E2E p50** | **12,556 ms** | 13,874 ms | -| E2E p90 | 44,522 ms | 42,501 ms | -| **ITL p50** | **13.81 ms** | 13.26 ms | -| ITL p90 | 27.43 ms | 18.52 ms | -| ISL p50 / mean | 70,624 / 80,989 | 69,911 / 80,811 | -| OSL p50 / mean | 223 / — | 230 / 1,050 | -| **server-reported cache hit** (per-request p50) | **88.1 %** | **88.1 %** | - -Ladders are recomputed from aiperf's raw per-request records -(`profile_export.jsonl`), not copied from a summary line. - -**The two clusters agree closely at this concurrency** — TTFT p50 within 1.3×, ITL -within 4 %. That is a narrower gap than the infera bench shows between the same two -clusters, which is consistent with this driver's fixed worker population holding -in-flight at ~5 on both, where the infera bench's growing session population let the -faster cluster run at higher occupancy. - -### The harness's cache-hit construction does not meet the workload spec - -The Case-A spec asks for a **~88–89 % prefix-cache hit rate**. The corpus generator -aims at that number per turn and reaches it *as a per-turn median* — but the corpus it -produces does not deliver it as a workload property. Measured over all 200 sessions / -1,778 requests of the shipped corpus: - -| how the corpus's own `hash_ids` are counted | hit rate | -|---|---| -| per-turn ratio, median | 88.0 % | -| per-turn ratio, median, excluding turn 0 | 88.2 % | -| per-turn ratio, **mean** | 66.0 % | -| **token-weighted over the whole corpus** | **59.7 %** | - -Three things in `gen_caseA_conformance.py` drive the gap, and each is worth reporting -upstream: - -1. **Cold start is inside the target.** Turn 0 is all-miss by construction and is - 11.2 % of requests. The generator applies its 0.88–0.90 draw only from turn 1, so - the corpus can hit the target per-turn and still miss it overall. -2. **No token weighting.** The target is drawn per *turn*, uniformly, but turns differ - by two orders of magnitude in input size (ISL is sampled lognormal per turn, - independently of the reuse draw). Cheap turns and 240K-token turns count equally - toward a rate whose cost is entirely token-proportional. -3. **The reuse accounting does not survive a shrinking turn.** Reuse is - `min(len(prefix_blocks), int(total_blocks × uniform(0.88, 0.90)))` against a - `total_blocks` resampled independently each turn, and the surviving prefix is then - *overwritten* with this turn's `hash_ids` (`prefix_blocks = hash_ids`). Because - consecutive ISL draws are independent, **51 % of turns are shorter than their - predecessor**, and every one of those truncates the accumulated prefix permanently. - A real agentic session's context grows; this one random-walks. - -The consequence for reading the results below: the harness's own summary line reports -**`Theoretical Prefix Cache Hit` ≈ 50.8 %**, and that figure is *not* a deployment -measurement — it is computed from the trace file's `hash_ids` and never asks the -server, so it reads the same on both clusters regardless of how the cache performs. -Its closeness to the corpus's own 59.7 % token-weighted rate is the point: both are -properties of the corpus. - -The server-side number aiperf does record is `usage_prompt_cache_read_tokens`. Taken -as a per-request ratio it is **88.1 % p50 / 89.4 % p90 on both clusters** — but that -median is over the 177 of 231 records that carry the field, and the 54 that do not are -exactly the `turn_index == 0` requests, i.e. the cold ones. Counted with those in at -0 % it is 72.4 %; token-weighted over all requests, 50.3 %. **All three are real -numbers about different populations**, and the deployment comparison in this document -uses the first because it is the one that holds the population fixed across clusters. - -A second, unrelated defect is worth reporting with the above: `replay_caseA.sh` -silently loses results when `OUT` is set outside the script's own directory, because -the container mounts only that directory — the sweep then prints `FAILED` for a run -that actually succeeded. - -## What this benchmark establishes that the infera one cannot - -**1. Per-turn attribution — the prefix cache is worth 2.5× on TTFT.** aiperf tags -every record with `turn_index`, so first-turn (cold) and later-turn (cached) requests -can be split at matched input size: - -| | n | ISL p50 | TTFT p50 | -|---|---|---|---| -| first turn (cold) | 54 | 69,907 | **8,981 ms** | -| turn ≥ 1 (cached) | 177 | 70,998 | **3,568 ms** | - -2.5× faster for a 1.6 % *larger* prompt. That is the prefix cache priced directly. -The infera bench has no turn index and cannot produce this. - -**2. Independent confirmation of the cache rate.** 88.1 % measured from the server's -own usage field, by a third-party driver, on a third-party trace, on **both** -clusters — against the infera bench's 88.9 %. Same population caveat as above: this is -the median over the warm requests. - -**3. A structurally clean input distribution.** `max(in + out) = 258,303` by -construction, below the 262,144 context, so context overflow is impossible. The -infera bench samples input and output independently with no joint clamp, which is -where its 0.5 % error rate comes from. - -**4. TTFT-by-input-size, and what changes with concurrency.** Comparing c8 against a -c16 run on the same deployment, the curve changes *shape*: - -| | 0–50K | 50–100K | 100–160K | 160–220K | 220–300K | spread | -|---|---|---|---|---|---|---| -| **c8** TTFT p50 | 3,177 | 6,161 | 11,305 | 20,674 | 22,639 | **7.1×** | -| **c16** TTFT p50 | 14,425 | 18,594 | 22,723 | 33,162 | 40,192 | **2.8×** | - -*(single-rail cluster; the multi-rail cluster shows the same transition, 10.0× → 2.3×)* - -At c8 the curve is **prefill-shaped** — monotone and super-linear, the deployment is -computing. At c16 it **flattens** and the smallest bucket already costs 14 s: a 40K -request cannot take that long to prefill on a leg that serves 240K in 40 s, so that -time is queueing rather than compute. - -TTFT here is entirely server-side: `http_req_sending` p50 is 0.2 ms. - -## Why its TTFT is higher than the infera bench's - -At c8 the customer bench reads TTFT p50 5,146 ms against the infera bench's 1,365 ms -on the same cluster, **against the same server processes in the same hour**. - -Part of it is the load model — but not all of it, and the residual is honestly -unexplained. At c8 the customer bench's mean in-flight is **5.1, less than half** the -infera bench's 12.1, and its TTFT is still 3.8× worse. Lower offered load with worse -latency is not explained by queueing. - -Candidates, none of which the available data discriminates between: - -| candidate | what would settle it | -|---|---| -| the 900 s window is too short for the prefill radix tree to reach the steady state the 3,600 s run measures | run the customer bench at c8 for 3,600 s; compare its first 900 s against its last | -| the scenario injects a unique marker into every trajectory's first turn, deliberately defeating cross-trajectory prefix sharing the other driver allows | run with the scenario's cache-bust disabled | -| the driver may not honour the trace's `think_time`, so turns arrive far denser than the session profile specifies | compare measured inter-arrival per conversation against the trace's think-time field | -| measurement definition — first *token* vs first *chunk* | inspect the first streamed chunk of a known request under both drivers | - -**Do not pick one of these without running the experiment.** - -## The right posture: run both - -The two benchmarks answer different questions and should not be collapsed into one -number. - -- The **infera bench** discovers capacity: its session population grows, so it finds - the concurrency the workload naturally sustains, runs a long steady-state window, - and captures router and KV internals. The customer bench must be *told* a worker - count, so it measures the point you chose. -- The **customer bench** compares deployments: the frozen trace means replaying it - against a different topology changes exactly one variable. It also models real - per-turn structure and applies no client timeout. - -They agree where they should — ITL ≈ 14 ms, cache ≈ 88 % — and diverge exactly where -the load model differs. - -## Pointing it at this deployment - -Get the harness from upstream ([ROCm/MAD PR #173](https://github.com/ROCm/MAD/pull/173), -`scripts/AgentX_CaseA/`) and configure it by environment only — no code change is -needed, and none should be made. - -| var | set it to | -|---|---| -| `URL` | this deployment's router: `http://:8100` | -| `SERVED` | the served-model-name this kit launches with (`glm5.2-mxfp4`), **not** the harness default | -| `TOK` | a tokenizer path the harness's container can see | -| `CONCS` | `8` for the numbers above | -| `DUR` | `900` — the scenario enforces this as its minimum and rejects the script's own default | -| `OUT` | a path **inside the harness's own directory** — see the defect noted above | - -Two prerequisites the deployment side must satisfy, both of which this kit already -does: `--enable-cache-report` on the engine (or every cache-hit column reads 0), and a -context length that covers the trace's 258,303-token maximum. diff --git a/examples/sglang_1p1d_glm5.2/results/infera_agenticbench_conc8.md b/examples/sglang_1p1d_glm5.2/results/infera_agenticbench_conc8.md deleted file mode 100644 index db189757..00000000 --- a/examples/sglang_1p1d_glm5.2/results/infera_agenticbench_conc8.md +++ /dev/null @@ -1,145 +0,0 @@ -# infera AgenticBench — Case-A request shape at concurrency 8 - -infera's own agentic benchmark, run against exactly the deployment shape this kit -ships. **Closed-loop**: each session issues one request, waits for the response, -sleeps its inter-turn delay, then issues the next. Offered load is therefore set by -the live-session population, not by a QPS target. - -## Workload - -The Case-A agentic profile at reduced load. Long inputs, a heavy output tail, -realistic think-time, and a large shared prefix: - -| axis | value | -|---|---| -| input tokens | p50 74,000 · p90 155,000 · p99 235,000 (clamped at 260,000) | -| output tokens | p50 320 · p90 3,300 · p99 17,000 | -| turns per session | p50 3 · p90 20 · p99 103 | -| inter-turn delay | p50 4 s · p90 31 s · p99 240 s | -| target prefix-cache hit | 0.89 | -| initial sessions / max sessions / max in-flight | 8 / 32 / 24 | -| window | ramp 400 s (excluded) + sustain 3,600 s | - -All figures below are the **sustain phase only**; the ramp is a warm-up exclusion -window, sized so the synchronised start cohort has died off and the shared prefix is -resident before measurement begins. - -## Results - -Three runs. All three ran the workload byte-identical; they differ in the cluster -and in the deployment shape. - -| | **multi-rail cluster**
prefill DPA off + kv-aware | **single-rail cluster**
prefill DPA off + kv-aware | **single-rail cluster**
prefill DPA on + round-robin | -|---|---|---|---| -| requests sent / completed | 2,884 / 2,850 | 2,907 / 2,861 | 2,323 / 2,289 | -| success rate | 0.988 | 0.984 | 0.985 | -| QPS (sustain) | 0.74 | **0.75** | 0.60 | -| **TTFT p50** | **1,365 ms** | 2,239 ms | 3,504 ms | -| **TTFT p90** | **4,903 ms** | 6,389 ms | 7,079 ms | -| TTFT p99 | 9,066 ms | 10,606 ms | 16,602 ms | -| **TPOT p50** | **14.8 ms** | 16.1 ms | 16.5 ms | -| TPOT p90 | 17.7 ms | 21.1 ms | 23.3 ms | -| cache hit (actual / ideal) | 88.9 % / 89.0 % | 88.7 % / 89.0 % | 88.2 % / 89.0 % | -| cache efficiency | 100.0 % | 99.6 % | 99.1 % | -| MTP acceptance length | 2.02 (per-request) | 2.79 (engine) | 2.72 (engine) | -| engine faults | 0 | 0 | 0 | -| prefill `--mem-fraction-static` | 0.80 | 0.70 | 0.70 | -| `--chunked-prefill-size` (global) | 16,384 | 65,536 | 65,536 | - -### Against the workload's own SLA block - -| bar | target | multi-rail | single-rail (DPA off) | single-rail (DPA on) | -|---|---|---|---|---| -| success rate | ≥ 0.97 | 0.988 **PASS** | 0.984 **PASS** | 0.985 **PASS** | -| TTFT p90 | < 30,000 ms | 4,903 **PASS** (6.1×) | 6,389 **PASS** (4.7×) | 7,079 **PASS** (4.2×) | -| E2E p50 | < 4,500 ms | 7,400 ms **FAIL** | not recorded | not recorded | -| in-flight not pinned at cap | — | max 22 / 24 **PASS** | not pinned **PASS** | brushed 24 on 0.8 % of ticks **PASS** | - -The E2E miss is not a regression. `e2e_p50_ms: 4500` is a **latency-floor** spec that -is met only at concurrency 1; at ~12 mean in-flight it is the wrong bar, and it fails -identically in every loaded run on this stack. It is reported rather than quietly -dropped. - -## What these runs establish - -**1. Prefill scales with input size, with no pathology.** On the multi-rail run, -TTFT p50 by input-size bucket: - -| input | 0–50K | 50–100K | 100–160K | 160–220K | 220–300K | -|---|---|---|---|---|---| -| TTFT p50 | 623 ms | 1,036 ms | 1,815 ms | 3,411 ms | 5,863 ms | - -Monotone, 9.4× across a 4.5× size span, **no stall bucket**. The deployment is -computing, not queueing, at this load. - -**2. Decode is not the bottleneck.** TPOT p99/p50 = 1.45 on the multi-rail run — an -exceptionally tight ladder. A contended decode leg shows a fat upper tail. - -**3. The prefix cache works as specified.** 88.2–88.9 % actual against an ideal of -89.0 %, i.e. 99–100 % cache efficiency and 0.2–0.9 % eviction. The workload nests -every request inside the same prefix, so this verifies cache *accounting* under load; -it does not exercise eviction pressure. - -**4. MTP is healthy, not degenerate.** Acceptance length 2.02–2.79. **A steady 4.00 -would be bad news**, not a better result — it means the draft model is predicting a -repetition loop perfectly. - -**5. The load cap never bound**, so the workload set the offered load rather than -backpressure, and the measurement window is valid. - -## The two things this data cannot tell you - -**The DPA/routing comparison is two-variable.** The two single-rail arms differ in -*both* prefill DP-attention *and* router policy, because that is how they were -specified. Arm B (DPA off + kv-aware) is faster on every latency percentile and -carries 25 % more throughput, but that gap is the **combined** effect and this data -cannot split it. The 2×2 is missing its other two cells. - -One prior controlled result bears on the DPA half alone: at **concurrency 1**, in a -single-variable comparison, prefill DP-attention cost 1.65–1.93× on TTFT. The -direction is consistent with the gap above; the load there is 24× lower and nothing -licenses transferring the magnitude. - -**Routing policy and a memory knob are coupled, and it is not documented anywhere -else.** The DPA-on + round-robin arm **would not boot** at `--mem-fraction-static -0.80`; it aborted 60 s in with `HSA_STATUS_ERROR_OUT_OF_RESOURCES` while token usage -read 0.05 — an empty KV pool, so activation memory, not KV exhaustion. The mechanism -is the spreading itself: under round-robin 4–5 DP ranks prefill concurrently, each -holding its own chunk's activations, where kv-aware concentrates on 1–2. At 0.70 it -ran the full 4,007 s with zero faults, at a cost of 19 % of the KV pool — which, at a -peak token usage of ~0.05, was never the binding resource. - -**If you switch this kit to `ROUTER_POLICY=round-robin`, lower `GMU_PREFILL`.** That -is why the shipped default is 0.70 rather than the higher value a kv-aware-only -deployment could sustain. - -## One structural finding worth knowing before you tune - -**kv-aware routing performed no cache steering in these runs, on either leg.** Not a -misconfiguration — a consequence of the deployment: - -- **Prefill**: with DP-attention off, `dp_size=1`, so the router has exactly one - target to choose between. -- **Decode**: MTP forces `ChunkCache` instead of a radix tree, because SGLang raises - on `--disaggregation-decode-enable-radix-cache` together with - `--speculative-algorithm`. The router's KV view of that worker is therefore - permanently empty and the cost function's overlap term cancels, leaving pure - least-loaded routing. Verified on the wire: the decode leg's kv-event socket - emitted 0 messages in 15 s while prefill's emitted 30. - -The consequence is the one in the [cross-cluster analysis](README.md#cross-cluster-why-the-single-rail-cluster-looks-slower): -`decode_prefix_len` is always 0, so **every turn re-transfers the entire prompt KV**. -A prefill-side cache hit saves compute, not bytes. - -This does not make kv-aware pointless — it steers prefill whenever DP-attention is on -there, which the round-robin arm demonstrated by contrast (round-robin spread picks -±4 % across all 8 ranks; kv-aware concentrated on 2 of 8 and left the other six at one -batch each for an entire run, because the concentration is self-reinforcing: no -traffic → empty cache view → never the cheapest candidate → no traffic). - -## Running it yourself - -This kit deliberately ships **no agentic bench client** — only the service self-check -(`smoke`) and a reference sweep with SGLang's own `bench_serving` (`bench`). The -agentic harness above is a separate internal tool. What this kit gives you is the -deployment those numbers were measured against. From 871ca002173dd3e5527a3bb123c66265d5d1b281 Mon Sep 17 00:00:00 2001 From: xiaobochen-amd Date: Fri, 7 Aug 2026 02:00:34 +0000 Subject: [PATCH 55/88] ci: run the DCO check once per pull request, not twice dco.yml triggers itself on any pull request, and ci.yml called it again as a job so the GPU tiers could gate on the sign-off through `needs:`. Every PR therefore carried two DCO check runs under two names, `dco` and `dco / dco`. The gate it bought is redundant. lint already holds those tiers back, so a PR with a missing sign-off is stopped well before a GPU node is touched; the second DCO run only added a check name to keep straight in branch protection. Drop the call. dco.yml's own trigger is unchanged, so every PR into any branch is still checked -- once. Verified against the repository rulesets first: no required check is named `dco / dco`, so removing it cannot leave a PR waiting on a check that will never report. Co-authored-by: Cursor Signed-off-by: xiaobochen-amd --- .github/workflows/ci.yml | 29 +++++++++-------------------- .github/workflows/dco.yml | 6 ++++-- 2 files changed, 13 insertions(+), 22 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c46ed61e..b652010f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -114,13 +114,6 @@ jobs: echo "=> run_e2e=$run" echo "run=$run" >> "$GITHUB_OUTPUT" - # Sign-off gate for the GPU tiers below — `needs:` cannot reach a job in - # another workflow, so dco.yml is called here as one. Skipped off a PR (there - # is nothing to check), which the tiers below read as "did not fail". - dco: - if: github.event_name == 'pull_request' - uses: ./.github/workflows/dco.yml - lint: needs: [changes] if: needs.changes.outputs.code == 'true' @@ -172,12 +165,11 @@ jobs: # Engine GPU tests. Same schedule as the e2e tiers (see e2e_gate): every PR # into main once it is out of draft, and every merge into main. engine: - needs: [lint, changes, dco, e2e_gate] + needs: [lint, changes, e2e_gate] if: >- !cancelled() && needs.changes.outputs.code == 'true' && needs.lint.result != 'failure' && needs.lint.result != 'cancelled' && - needs.dco.result != 'failure' && needs.dco.result != 'cancelled' && (needs.e2e_gate.outputs.run == 'true' || github.event_name == 'workflow_dispatch') runs-on: [self-hosted, crusoe] timeout-minutes: 60 @@ -202,19 +194,17 @@ jobs: # Full PD-mixed e2e (per engine, parallel). When it runs is e2e_gate's call. e2e-mixed: - # Skipped for docs-only changes. Gated behind lint and dco: run only if neither - # failed. `!cancelled()` + result checks (instead of a plain success dependency) - # is needed so e2e still runs when lint is *skipped* for a docs-only change - # or dco is skipped off a PR, but is held back when either fails. NOT - # `always()`: on cancel the server re-evaluates job-level `if`, and `always()` - # evaluates true, so the job is never cancelled — it keeps (or even starts) - # burning GPU nodes after "Cancel workflow". - needs: [lint, changes, dco, e2e_gate] + # Skipped for docs-only changes, and held back if lint failed. Checking + # lint's *result* rather than depending on its success keeps `!cancelled()` + # meaningful; `always()` would not work here, because on cancel the server + # re-evaluates job-level `if` and `always()` is true, so the job would never + # be cancelled — it would keep (or even start) burning GPU nodes after + # "Cancel workflow". + needs: [lint, changes, e2e_gate] if: >- !cancelled() && needs.changes.outputs.code == 'true' && needs.lint.result != 'failure' && needs.lint.result != 'cancelled' && - needs.dco.result != 'failure' && needs.dco.result != 'cancelled' && (needs.e2e_gate.outputs.run == 'true' || (github.event_name == 'workflow_dispatch' && inputs.run_e2e_mixed)) strategy: @@ -246,12 +236,11 @@ jobs: e2e-disag: # Gates mirror e2e-mixed, `!cancelled()` included: `always()` would keep this # holding a two-node pair after "Cancel workflow". - needs: [lint, changes, dco, e2e_gate] + needs: [lint, changes, e2e_gate] if: >- !cancelled() && needs.changes.outputs.code == 'true' && needs.lint.result != 'failure' && needs.lint.result != 'cancelled' && - needs.dco.result != 'failure' && needs.dco.result != 'cancelled' && (needs.e2e_gate.outputs.run == 'true' || (github.event_name == 'workflow_dispatch' && inputs.run_e2e_disag)) strategy: diff --git a/.github/workflows/dco.yml b/.github/workflows/dco.yml index ded8c652..a8e0b863 100644 --- a/.github/workflows/dco.yml +++ b/.github/workflows/dco.yml @@ -5,8 +5,10 @@ name: DCO # See CONTRIBUTING.md > Developer Certificate of Origin. on: - # Standalone on a PR into any branch; ci.yml additionally calls this one as a - # job, which is what lets its GPU tiers gate on the sign-off via `needs:`. + # One check, on a PR into any branch. ci.yml used to call this as a job too so + # its GPU tiers could gate on the sign-off, which ran it twice per PR under two + # different check names; lint already holds those tiers back, so the second + # copy bought nothing. workflow_call stays for any future caller. pull_request: branches: ["**"] workflow_call: From 14b866af27910d6bc5214c39f3cad237f24589f4 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Fri, 7 Aug 2026 02:07:03 +0000 Subject: [PATCH 56/88] test(router): drive the hybrid pipeline over real ZMQ, in both event orders Three drafts of this test passed against the unfixed code before it caught anything, and each failure mode is worth keeping in mind when reading it. Draft 1 constructed BlockStored objects in Python. That cannot catch the defect at all: group_idx and kv_cache_spec_kind were on the wire the whole time and were dropped because the struct did not declare them, so a test that builds the struct itself starts from the fixed state. This one publishes msgpack dicts over a real socket and lets the decoder do its job. Draft 2 emitted the four per-chunk events in the order the capture happened to show, Mamba first and attention last. In that order the attention group's five correct hashes overwrite whatever the Mamba groups wrote, and the view comes out right even with no filter at all. vLLM promises no order; the test now runs both and the fix is what makes the outcome independent of it. Draft 3 asserted the view's SIZE and which worker won. Both survive a poisoned chain: ten entries still appear, five of them hashed from a clobbered parent, and a worker with five real blocks still beats an empty one. It now counts how many of the prompt's blocks actually match, which is the number routing acts on. Against the pre-fix implementation the mamba-last case reports 5/10 blocks visible on the worker that just served the prompt. Fixed: 10/10, both orders. Worth stating plainly: on the shapes captured from this cluster, the Mamba events are always sparse, so the length check alone rejects them and the group filter is not what makes this test pass. The filter earns its place on the orders and group kinds the length check cannot see -- a non-attention group whose lengths happen to agree, which sliding-window groups produce under VLLM_PREFIX_CACHE_RETENTION_INTERVAL. Signed-off-by: Zhang, Jiejing --- tests/unit/router/test_kv_event_e2e_hybrid.py | 252 ++++++++++++++++++ 1 file changed, 252 insertions(+) create mode 100644 tests/unit/router/test_kv_event_e2e_hybrid.py diff --git a/tests/unit/router/test_kv_event_e2e_hybrid.py b/tests/unit/router/test_kv_event_e2e_hybrid.py new file mode 100644 index 00000000..bc9b26e4 --- /dev/null +++ b/tests/unit/router/test_kv_event_e2e_hybrid.py @@ -0,0 +1,252 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Two workers, real ZMQ, a hybrid model's event stream — does routing follow the cache? + +``test_kv_event_e2e.py`` covers the same pipeline for SGLang. This is the vLLM +hybrid case, which fails differently and was found in production rather than by +a test: kv-aware routing reported ``cache_hits=0`` on every decision and pinned +all traffic to one worker, leaving the other node's GPUs idle. That is not a +degraded hit rate, it is half the fleet. + +The shapes below are what a live Kimi-K3 worker publishes, captured off the +wire (block_size 768, ``--max-num-batched-tokens 4096`` so chunked prefill +splits at ``4096 // 768 * 768 = 3840``): + + kv_cache_spec_kind group_idx token_ids block_hashes + mamba 0 3840 1 + mamba 1 3840 1 + mamba 2 3840 1 + mla_attention 3 3840 5 + +Three of the four groups are unusable: prefix caching runs in "align" mode on +the KDA layers, so all but one block per step is a null block skipped when the +hash list is built, while ``token_ids`` still spans everything. Nothing says +which chunk the surviving hash covers. + +They are also actively destructive. vLLM mixes no group id into the block hash, +so at equal block sizes a Mamba hash COLLIDES with an attention hash and +overwrites its entry in the engine-hash -> router-hash map; the next chunk then +resolves its parent to the wrong node. Three streams silently destroyed the +fourth, which is why the fix is a filter and not a lenient decoder. + +Uses real ZMQ rather than calling ``_handle_event`` directly, so the msgspec +schema is exercised: the two fields the filter depends on, ``group_idx`` and +``kv_cache_spec_kind``, were on the wire all along and were being discarded +because the struct did not declare them. A test that constructs events in +Python cannot catch that. +""" + +from __future__ import annotations + +import asyncio +import socket +from typing import Any + +import msgspec +import pytest +import zmq + +from infera.common.worker_pool import ( + DisaggMode, + EngineType, + WorkerInfo, + WorkerStatus, +) +from infera.router.kv_event.client import KvEventClient +from infera.router.policy.kv_event_aware import KvEventAwarePolicy +from infera.router.policy.target import RouteTarget + +_TOPIC = b"kv-events" +BS = 768 +CHUNK_BLOCKS = 5 +CHUNK = BS * CHUNK_BLOCKS + + +def _free_port() -> int: + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + s.close() + return port + + +def _worker(worker_id: str, endpoint: str) -> WorkerInfo: + return WorkerInfo( + worker_id=worker_id, + url=f"http://{worker_id}", + model_name="test/m", + engine=EngineType.VLLM, + status=WorkerStatus.ACTIVE, + disagg_mode=DisaggMode.MIXED, + kv_events_endpoint=endpoint, + kv_block_size=BS, + ) + + +class _IdentityHasher: + """Treat ``body["token_ids"]`` as already tokenized, so request/worker + alignment is controlled by the test rather than by a tokenizer.""" + + def hash_for(self, body: dict, *, block_size: int, engine=None) -> list[int]: + from infera.router.kv_event.hasher import hash_request + + return hash_request(body.get("token_ids", []), block_size) + + +def _chunk_payload( + tokens: list[int], first_block: int, parent: bytes | None, *, mamba_last: bool = False +) -> bytes: + """One prefill chunk, encoded exactly as vLLM puts it on the wire. + + Built as plain dicts, not our own structs, so the test does not inherit the + schema it is meant to be checking. + + ``mamba_last`` exists because ORDER DECIDED CORRECTNESS before the fix, and + an earlier draft of this test missed the bug entirely by picking the lucky + order. With the attention event last, its five correct hashes overwrite + whatever the Mamba events wrote and the view comes out right by accident. + Put the Mamba events last and they clobber the attention group's map, so + the next chunk resolves its parent to the wrong node. + + vLLM does not promise either order. A router whose cache view depends on it + is broken whichever way the engine happens to emit today. + """ + hashes = [bytes([first_block + i]) for i in range(len(tokens) // BS)] + base = { + "type": "BlockStored", + "parent_block_hash": parent, + "token_ids": tokens, + "block_size": BS, + "lora_id": None, + } + mamba = [ + {**base, "block_hashes": [hashes[-1]], "group_idx": g, "kv_cache_spec_kind": "mamba"} + for g in (0, 1, 2) + ] + attn = {**base, "block_hashes": hashes, "group_idx": 3, "kv_cache_spec_kind": "mla_attention"} + events = [attn, *mamba] if mamba_last else [*mamba, attn] + return msgspec.msgpack.encode([0.0, events, None]) + + +async def _publish_until( + pub: Any, payloads: list[bytes], predicate, *, deadline_s: float = 5.0 +) -> bool: + """PUB/SUB drops messages sent before the subscriber attaches, so resend + until the effect is visible or the deadline passes.""" + loop = asyncio.get_running_loop() + deadline = loop.time() + deadline_s + while loop.time() < deadline: + for p in payloads: + pub.send_multipart([_TOPIC, p]) + await asyncio.sleep(0.05) + if predicate(): + return True + return False + + +@pytest.mark.parametrize("mamba_last", [False, True], ids=["attn-last", "mamba-last"]) +@pytest.mark.asyncio +async def test_routing_follows_the_cache_across_two_workers(mamba_last): + """The production scenario, end to end, under BOTH intra-batch orders. + + Worker A serves a prompt; its events go out. A request for the SAME prompt + must then pick A over an idle B. + + Both orders are run because only one of them fails without the fix. With + the attention event last its correct hashes overwrite the Mamba groups' + and everything works by luck; with the Mamba events last they overwrite the + attention group's map and the second chunk chains off the wrong parent. The + fix makes the outcome independent of an order vLLM never promised. + """ + ctx = zmq.Context.instance() + port_a, port_b = _free_port(), _free_port() + pub_a = ctx.socket(zmq.PUB) + pub_a.bind(f"tcp://127.0.0.1:{port_a}") + pub_b = ctx.socket(zmq.PUB) + pub_b.bind(f"tcp://127.0.0.1:{port_b}") + + client = KvEventClient() + policy = KvEventAwarePolicy(client, _IdentityHasher()) + wa = _worker("a:30000", f"tcp://127.0.0.1:{port_a}") + wb = _worker("b:30000", f"tcp://127.0.0.1:{port_b}") + policy.on_worker_added(wa) + policy.on_worker_added(wb) + + prompt = list(range(2 * CHUNK)) + payloads = [ + _chunk_payload(prompt[:CHUNK], 0, None, mamba_last=mamba_last), + _chunk_payload( + prompt[CHUNK:], CHUNK_BLOCKS, bytes([CHUNK_BLOCKS - 1]), mamba_last=mamba_last + ), + ] + try: + ok = await _publish_until( + pub_a, + payloads, + lambda: len(client._subs["a:30000"].view_for(None)) >= 2 * CHUNK_BLOCKS, + ) + assert ok, ( + "worker A's view never filled: the attention group's events are not " + "being indexed at all" + ) + assert len(client._subs["b:30000"].view_for(None)) == 0, "B was never fed" + + # The assertion that matters. View SIZE is not it: a poisoned chain + # still produces ten entries, five of them hashed from the wrong parent, + # and A still beats an empty B -- so "the right worker won" passes while + # half the prompt is invisible. Count what actually matches. + from infera.router.kv_event.hasher import hash_request + + want = hash_request(prompt, BS) + view = client._subs["a:30000"].view_for(None) + matched = sum(1 for h in want if h in view) + assert matched == 2 * CHUNK_BLOCKS, ( + f"{matched}/{2 * CHUNK_BLOCKS} blocks of the prompt are visible on the " + "worker that just served it; the chain continued from a clobbered parent" + ) + + target, blocks = policy.pick([wa, wb], {"model": "test/m", "token_ids": prompt}) + assert isinstance(target, RouteTarget) + assert target.worker.worker_id == "a:30000" + assert len(blocks) == 2 * CHUNK_BLOCKS + finally: + await client.aclose() + pub_a.close(linger=0) + pub_b.close(linger=0) + + +@pytest.mark.asyncio +async def test_the_view_is_one_groups_worth_not_four(): + """Four events arrive per chunk and three are unusable. If they were all + indexed the view would be inflated with blocks hashed from a Mamba group's + span -- which is how the collision corrupts the map.""" + ctx = zmq.Context.instance() + port = _free_port() + pub = ctx.socket(zmq.PUB) + pub.bind(f"tcp://127.0.0.1:{port}") + + client = KvEventClient() + policy = KvEventAwarePolicy(client, _IdentityHasher()) + policy.on_worker_added(_worker("a:30000", f"tcp://127.0.0.1:{port}")) + try: + payload = _chunk_payload(list(range(CHUNK)), 0, None) + ok = await _publish_until( + pub, + [payload], + lambda: len(client._subs["a:30000"].view_for(None)) >= CHUNK_BLOCKS, + ) + assert ok, "the attention group's chunk never landed" + # Give the other three groups every chance to be indexed too. + for _ in range(5): + pub.send_multipart([_TOPIC, payload]) + await asyncio.sleep(0.05) + + sub = client._subs["a:30000"] + assert len(sub.view_for(None)) == CHUNK_BLOCKS + assert len(sub.map_for(None)) == CHUNK_BLOCKS + finally: + await client.aclose() + pub.close(linger=0) From d443cd5aa54e1708e6d725ff4eebe7ff108f3129 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Fri, 7 Aug 2026 00:48:31 -0700 Subject: [PATCH 57/88] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: leiwei12 Signed-off-by: Zhang, Jiejing --- deploy/operator/Makefile | 1 + 1 file changed, 1 insertion(+) diff --git a/deploy/operator/Makefile b/deploy/operator/Makefile index ff571c88..f2a5d276 100644 --- a/deploy/operator/Makefile +++ b/deploy/operator/Makefile @@ -40,6 +40,7 @@ manifests: ## changing. sync-chart-crd: mkdir -p $(CHART_CRD_DIR) + rm -f $(CHART_CRD_DIR)/*.yaml cp $(CRD_DIR)/*.yaml $(CHART_CRD_DIR)/ build: generate From 85b84704576efb040a7df8ae66683000b880d6f4 Mon Sep 17 00:00:00 2001 From: xiaobochen-amd Date: Fri, 7 Aug 2026 02:00:54 +0000 Subject: [PATCH 58/88] ci: pin the last three actions to a commit Eighteen of the twenty-one action references are pinned to a commit with the version in a trailing comment. Three were left on a moving tag: checkout@v4 and setup-python@v5 in unit-torch-cpu, upload-artifact@v4 in release.yml. A tag can be repointed by the action's owner, or by whoever takes over the account, and the workflow would run different code with nothing changed in this repository. The two in unit-torch-cpu take the same pins the rest of the file already uses, which also brings that job up from checkout v4 and setup-python v5 to the v7.0.0 and v6.3.0 every other job runs. Co-authored-by: Cursor Signed-off-by: xiaobochen-amd --- .github/workflows/ci.yml | 4 ++-- .github/workflows/release.yml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b652010f..10710ba0 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -287,8 +287,8 @@ jobs: if: needs.changes.outputs.code == 'true' runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 with: python-version: "3.10" - run: pip install -e ".[dev]" diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 79d1af61..35ccd8d1 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -321,7 +321,7 @@ jobs: run: pip install -r manual/sphinx/requirements.txt - name: Build manual (warnings = errors) run: make -C manual html - - uses: actions/upload-artifact@v4 + - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4 with: name: infera-manual-html path: manual/_build/html From bb22008b863944998a98454d451cf2f171468786 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Fri, 7 Aug 2026 02:07:19 +0000 Subject: [PATCH 59/88] test(router): pin routing outcomes that are decidable by hand The suite could answer "does a hit happen". It could not answer "does the RIGHT worker get it", because no test ever gave the policy a choice. That gap is exactly the shape of the production failure this branch started from: 32 of 32 requests to one worker while the second node sat idle, and nothing red. These cases have one correct answer each, arrived at by hand rather than by degree: a prompt one worker holds entirely (20/20 to the holder, and the same with the candidate list reversed, so order cannot be what decided it); two disjoint prefixes that must sort themselves 10/10 BY CONTENT -- round-robin also produces 10/10 overall while sending half of each prefix to the wrong worker, so the split is asserted per prefix; a strict subset losing to a longer match, down to a single block of difference; a cold fleet that must spread instead of piling up; and the load-versus-locality tradeoff at the default weight, where one cached block and one in-flight request are both worth 1, plus its release on request completion. Checked against two degenerate policies to confirm the cases are not vacuous: replacing the pick with targets[0] fails 8 of 10, and dropping the cache term so only load counts fails 5 of 10. Signed-off-by: Zhang, Jiejing --- .../router/test_kv_aware_routing_extremes.py | 224 ++++++++++++++++++ 1 file changed, 224 insertions(+) create mode 100644 tests/unit/router/test_kv_aware_routing_extremes.py diff --git a/tests/unit/router/test_kv_aware_routing_extremes.py b/tests/unit/router/test_kv_aware_routing_extremes.py new file mode 100644 index 00000000..9c951ec7 --- /dev/null +++ b/tests/unit/router/test_kv_aware_routing_extremes.py @@ -0,0 +1,224 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Routing outcomes that are decidable by hand, so a wrong one is obviously wrong. + +The pipeline tests answer "does a hit happen". These answer "does the RIGHT +worker get it", on inputs where the correct answer is not a matter of degree: +everything to one worker, an exact half-and-half split, a strict subset losing +to a superset. A policy that quietly degenerates to a fixed tiebreak passes a +hit-rate test on one worker and fails these. + +That degeneration is not hypothetical. In production kv-aware routing sent +32 of 32 requests to a single worker while the second node sat idle, and every +decision logged ``cache_hits=0``. Nothing in the suite would have caught it, +because no test asked where requests went when there was a choice. + +Cost, from ``KvEventAwarePolicy.pick``: + + cost(t) = w_overlap * (blocks_in_request - blocks_t_already_has) + + w_mm * images_t_lacks + + in_flight_blocks_on_t + +with ties broken by the lower in-flight count. At the default weight of 1.0, +one uncached block and one in-flight request are worth the same, which is what +makes the load-versus-locality cases below decidable rather than a judgement. +""" + +from __future__ import annotations + +import pytest + +from infera.common.worker_pool import ( + DisaggMode, + EngineType, + WorkerInfo, + WorkerStatus, +) +from infera.router.kv_event.client import KvEventClient +from infera.router.kv_event.events import BlockStored +from infera.router.policy.kv_event_aware import KvEventAwarePolicy + +BS = 4 + + +def _worker(wid: str) -> WorkerInfo: + return WorkerInfo( + worker_id=wid, + url=f"http://{wid}", + model_name="m", + engine=EngineType.VLLM, + status=WorkerStatus.ACTIVE, + disagg_mode=DisaggMode.MIXED, + kv_events_endpoint=f"tcp://{wid}:5555", + kv_block_size=BS, + ) + + +class _IdentityHasher: + def hash_for(self, body: dict, *, block_size: int, engine=None) -> list[int]: + from infera.router.kv_event.hasher import hash_request + + return hash_request(body.get("token_ids", []), block_size) + + +@pytest.fixture +async def rig(): + """Two workers behind one policy, with their event subscriptions stubbed out + so blocks can be placed directly and the outcome is not a race.""" + client = KvEventClient() + policy = KvEventAwarePolicy(client, _IdentityHasher()) + a, b = _worker("a:1"), _worker("b:1") + for w in (a, b): + policy.on_worker_added(w) + for t in client._subs[w.worker_id].tasks: + t.cancel() + return policy, client, a, b + + +def _store(client, wid: str, tokens: list[int], first_hash: int = 0) -> None: + """Give a worker the blocks for `tokens`, as one aligned event.""" + n = len(tokens) // BS + client._handle_event( + client._subs[wid], + BlockStored( + block_hashes=[bytes([first_hash + i]) for i in range(n)], + parent_block_hash=None, + token_ids=list(tokens), + block_size=BS, + lora_id=None, + group_idx=0, + kv_cache_spec_kind="full_attention", + ), + rank=None, + ) + + +def _pick(policy, workers, tokens): + target, _ = policy.pick(list(workers), {"model": "m", "token_ids": list(tokens)}) + return target.worker.worker_id + + +# --- everything to one worker ------------------------------------------------ + + +async def test_every_request_goes_to_the_only_worker_that_has_the_prefix(rig): + """One worker holds the whole prompt, the other holds nothing. There is no + tradeoff to weigh: 20/20 to the holder.""" + policy, client, a, b = rig + prompt = list(range(40)) # 10 blocks + _store(client, "a:1", prompt) + + picks = [_pick(policy, (a, b), prompt) for _ in range(20)] + assert picks.count("a:1") == 20, f"{picks.count('b:1')}/20 went to the cold worker" + + +async def test_the_holder_still_wins_when_it_is_the_second_candidate(rig): + """Candidate order must not decide this. A degenerate tiebreak looks correct + whenever the holder happens to be listed first.""" + policy, client, a, b = rig + prompt = list(range(40)) + _store(client, "b:1", prompt) + + assert _pick(policy, (a, b), prompt) == "b:1" + assert _pick(policy, (b, a), prompt) == "b:1" + + +# --- an exact half-and-half split ------------------------------------------- + + +async def test_two_disjoint_prefixes_split_exactly_down_the_middle(rig): + """Each worker owns one prefix. Requests must sort themselves 10/10 by + content, not alternate: a round-robin router also produces 10/10 overall, + but sends half of each prefix to the wrong worker.""" + policy, client, a, b = rig + p1 = list(range(40)) + p2 = list(range(1000, 1040)) + _store(client, "a:1", p1, first_hash=0) + _store(client, "b:1", p2, first_hash=100) + + by_prefix = {"p1": [], "p2": []} + for _ in range(10): + by_prefix["p1"].append(_pick(policy, (a, b), p1)) + by_prefix["p2"].append(_pick(policy, (a, b), p2)) + + assert by_prefix["p1"] == ["a:1"] * 10, "prefix 1 must always follow its holder" + assert by_prefix["p2"] == ["b:1"] * 10, "prefix 2 must always follow its holder" + + +# --- more of the prefix wins ------------------------------------------------ + + +@pytest.mark.parametrize( + "a_blocks,b_blocks,winner", + [ + (2, 8, "b:1"), # strict subset loses to the longer match + (8, 2, "a:1"), # and the same the other way round + (0, 1, "b:1"), # even one block beats nothing + (9, 10, "b:1"), # a single block of difference still decides + ], +) +async def test_the_longer_cached_prefix_wins(rig, a_blocks, b_blocks, winner): + policy, client, a, b = rig + prompt = list(range(40)) # 10 blocks + if a_blocks: + _store(client, "a:1", prompt[: a_blocks * BS], first_hash=0) + if b_blocks: + _store(client, "b:1", prompt[: b_blocks * BS], first_hash=100) + + assert _pick(policy, (a, b), prompt) == winner + + +# --- cold start must spread, not pile up ------------------------------------ + + +async def test_a_cold_fleet_spreads_instead_of_piling_onto_one_worker(rig): + """Neither worker has anything, so every candidate costs the same and the + in-flight count is the only signal left. This is the exact shape of the + production failure -- 32/32 to one worker -- and it is what the tiebreak + exists to prevent.""" + policy, _client, a, b = rig + counts = {"a:1": 0, "b:1": 0} + for i in range(20): + prompt = list(range(i * 100, i * 100 + 40)) # a fresh prefix each time + target, blocks = policy.pick([a, b], {"model": "m", "token_ids": prompt}) + wid = target.worker.worker_id + counts[wid] += 1 + policy.on_request_started(target.route_key, blocks) # stays in flight + + assert counts == {"a:1": 10, "b:1": 10}, f"cold fleet did not spread: {counts}" + + +async def test_load_outweighs_a_one_block_cache_edge(rig): + """One cached block and one in-flight request are both worth 1 at the + default weight, so a worker holding one extra block but carrying two more + in-flight requests must lose. Stated as a test because it is the tradeoff + someone will change the weight to alter.""" + policy, client, a, b = rig + prompt = list(range(40)) + _store(client, "a:1", prompt[: 1 * BS], first_hash=0) # A: 1 block, 9 misses + + # Put 2 blocks' worth of in-flight work on A. Cost(A) = 9 + 2 = 11 > Cost(B) = 10. + ta, blocks = policy.pick([a], {"model": "m", "token_ids": prompt}) + policy.on_request_started(ta.route_key, blocks[:2]) + + assert _pick(policy, (a, b), prompt) == "b:1" + + +async def test_finishing_a_request_returns_the_worker_to_contention(rig): + """In-flight cost must be released, or a worker that served a burst is + written off long after it went idle.""" + policy, client, a, b = rig + prompt = list(range(40)) + _store(client, "a:1", prompt[: 1 * BS], first_hash=0) + + ta, blocks = policy.pick([a], {"model": "m", "token_ids": prompt}) + policy.on_request_started(ta.route_key, blocks[:2]) + assert _pick(policy, (a, b), prompt) == "b:1" + + policy.on_request_finished(ta.route_key, blocks[:2]) + assert _pick(policy, (a, b), prompt) == "a:1", ( + "A holds a block and is idle again; it must win once its load is released" + ) From a4bf10ca8b3873eb265c7a259209d4fa0ce8054b Mon Sep 17 00:00:00 2001 From: xiaobochen-amd Date: Fri, 7 Aug 2026 02:42:09 +0000 Subject: [PATCH 60/88] fix(tests): the engine tier must not pass on zero collected tests Two ways run_engine reported PASS having tested nothing. `for f in $(find "$INFERA_TEST_SCOPE" ...)` iterates zero times when the scope matches nothing, and rc stays 0. find does report a missing path -- it exits 1 -- but nothing was in a position to see it: the loop reads only the output, and the `| sort` had already replaced find's status with its own. engine_tier hard-codes tests/engine/vllm and tests/engine/sglang, so one rename is all it takes for the tier to go green having run nothing. Capture the list before sorting and fail on either a find that could not read the scope or a scope with no test files -- the first case is not hypothetical padding: an unreadable subdirectory makes find list the files it could reach and still exit 1, which the old code would have run as if it were the whole suite. pytest's exit 5, "collected nothing", was written down as "whole file skipped -- not a failure". Three files carry a module-level importorskip: the two test_disagg_allow_tcp_args.py on vllm and sglang, and test_kvd_fp8_passthrough.py on torch. Each guard names a module that the image running that scope ships, so a 5 there means the image is broken -- precisely the moment the tier must go red. It also has to be labelled rather than merely counted, because pytest words a module-level skip as "1 skipped": left alone, the per-file summary would repeat that verbatim while the tier failed for a reason the log never states. (A fourth module-level guard, on test_kv_metadata_block_size.py, already failed correctly: importing infera.engine.vllm.__main__ raises something other than ImportError, which pytest reports as a collection error, exit 2.) Neither path has ever fired. No dispatch log under the shared CI log directory contains the old branch's "no tests ran" line, and all four complete engine runs report 21 files and 21 passed, the three guarded files among them. That is also what makes this safe: no file currently depends on exit 5 being forgiven. Co-authored-by: Cursor Signed-off-by: xiaobochen-amd --- tests/run_tests.sh | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/tests/run_tests.sh b/tests/run_tests.sh index 81af2458..3fec8b08 100755 --- a/tests/run_tests.sh +++ b/tests/run_tests.sh @@ -568,7 +568,15 @@ run_engine() { cd /workspace PYT="python3 -m pytest -p no:cacheprovider -o addopts= -q -rfE" rc=0 - for f in $(find "$INFERA_TEST_SCOPE" -name "test_*.py" | sort); do + # A scope that matches nothing iterates zero times and exits 0, so a rename + # or a typo in engine_tier would report PASS having tested nothing. Capture + # before sorting: through a pipe, find's own exit code would be sort's 0. + if ! files=$(find "$INFERA_TEST_SCOPE" -name "test_*.py") || [ -z "$files" ]; then + echo "[engine $INFERA_TEST_SCOPE] FATAL: scope unreadable or holds no test_*.py" >&2 + echo "[engine $INFERA_TEST_SCOPE] scope unreadable or holds no test_*.py" >> /scratch/failures.txt + exit 1 + fi + for f in $(printf "%s\n" "$files" | sort); do echo "----- pytest $f -----" # tee: stream live for CI, keep a copy for the classification below. $PYT "$f" 2>&1 | stdbuf -oL tee /scratch/.engine_f.out; code=${PIPESTATUS[0]} @@ -577,7 +585,11 @@ run_engine() { 139|134|137) line="CRASH(exit=$code)"; rc=1 echo "[engine $INFERA_TEST_SCOPE] CRASH(exit=$code) $f" >> /scratch/failures.txt ;; 0) line=$(printf "%s" "$out" | grep -E "passed|failed|skipped|no tests ran" | tail -1) ;; - 5) line="no tests ran (whole file skipped — not a failure)" ;; + # Nothing collected. Each guarded file importorskips a module its own + # image ships, so a 5 means the image is broken — exactly when this must + # go red. Say so plainly: pytest words it "1 skipped", which reads benign. + 5) line="FAIL: no tests collected (exit=5)"; rc=1 + echo "[engine $INFERA_TEST_SCOPE] $f (exit=5, no tests collected)" >> /scratch/failures.txt ;; *) line=$(printf "%s" "$out" | grep -E "passed|failed|error|skipped" | tail -1) [ -z "$line" ] && line="(exit=$code)"; rc=1 fails=$(printf "%s\n" "$out" | grep -aE "^(FAILED|ERROR) ") From fa4a7ef4d6235f4feafd7b980bb17300e70d513e Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Fri, 7 Aug 2026 23:43:25 +0000 Subject: [PATCH 61/88] ci(release): stop the pypi job racing the Release into existence Two independent bugs found while working out why the v0.2.4 release left kvd/server/pypi unpublished. 1. `pypi` fails on a tag push with "release not found" The job uploads the wheel with `gh release upload `, but the only job that CREATES the GitHub Release is `manual`, and `pypi` has no `needs:` on it -- so on a tag push it reaches the upload first and dies. That is what happened on v0.2.4 (run 31217988270): every build/verify step in the job passed and the release failed on the last line. Adding `needs: [build, overlay]` would fix the ordering but is wrong: a wheel depends on neither the engine images nor the overlay, and gating it there makes the Python release wait on the GPU queue -- exactly the coupling the job's own comment says it wants to avoid. So make `pypi` create the Release if it is absent, the same way `manual` already does. Both jobs can now reach `create` concurrently, so each re-checks with `release view` if its own create fails. Losing that race is fine; a genuine failure (bad permissions, wrong tag) still exits non-zero instead of resurfacing as a confusing "release not found" on the upload. 2. The overlay build claims to harvest an sglang native tree it never harvests `build_test_push.sh` resolved SGLANG_NATIVE_IMAGE, passed it as a --build-arg, and logged `overlay: harvest sglang=`. No FROM in Dockerfile.payload consumes it: the SGLang (CPython 3.10) Mooncake is COMPILED in the `mooncake310` stage, deliberately, so the HIP-transport gate is guaranteed present rather than inherited from whatever the engine image happened to ship (the reasoning is in that stage's comment). The log misdirects anyone debugging the overlay's sglang payload, and the dead arg also made the script name -- and therefore pull -- an infera:sglang- image the build has no use for. Drop the arg and say what actually happens. The same incorrect claim is corrected in the recipes README, which is user-facing. Signed-off-by: Zhang, Jiejing --- .github/scripts/build_test_push.sh | 14 ++++++++----- .github/workflows/release.yml | 33 +++++++++++++++++++++++++++--- deploy/overlay/Dockerfile.payload | 6 ++++-- examples/recipes/README.md | 10 +++++---- 4 files changed, 49 insertions(+), 14 deletions(-) diff --git a/.github/scripts/build_test_push.sh b/.github/scripts/build_test_push.sh index 301e22f8..1fcecdc2 100755 --- a/.github/scripts/build_test_push.sh +++ b/.github/scripts/build_test_push.sh @@ -72,8 +72,13 @@ cd "$(dirname "$0")/../.." # inside the STOCK vendor bases so they can be pruned down to what those bases # lack. So it needs four refs, and getting any of them wrong is silent: # -# NATIVE_IMAGE / SGLANG_NATIVE_IMAGE <- the engine images from THIS run, so a -# release never ships an overlay carrying a previous release's native code. +# NATIVE_IMAGE <- the vLLM engine image from THIS run, so a release +# never ships an overlay carrying a previous release's native code. Only +# the vLLM (CPython 3.12) tree is harvested: the SGLang (3.10) family's +# Mooncake is COMPILED in Dockerfile.payload's `mooncake310` stage, so +# that the HIP-transport gate is guaranteed present rather than inherited +# from whatever the engine image happened to ship. There is deliberately +# no SGLANG_NATIVE_IMAGE here. # VLLM_BASE_IMAGE / SGLANG_BASE_IMAGE <- read out of the engine Dockerfiles # rather than repeated here. The prune keeps exactly what the base lacks, # so a base ref that drifts from the one the engine image was built on @@ -95,15 +100,14 @@ if [ "$engine" = overlay ]; then else esuf="local" fi native_vllm="${NATIVE_IMAGE:-${IMAGE}:vllm-${esuf}}" - native_sglang="${SGLANG_NATIVE_IMAGE:-${IMAGE}:sglang-${esuf}}" build_args=( --build-arg "VLLM_BASE_IMAGE=${vllm_base}" --build-arg "SGLANG_BASE_IMAGE=${sglang_base}" --build-arg "NATIVE_IMAGE=${native_vllm}" - --build-arg "SGLANG_NATIVE_IMAGE=${native_sglang}" ) - echo "overlay: harvest vllm=${native_vllm} sglang=${native_sglang}" + echo "overlay: harvest vllm=${native_vllm}" + echo "overlay: compile sglang mooncake in-image (mooncake310 stage)" echo "overlay: deps on vllm=${vllm_base}" echo "overlay: deps on sglang=${sglang_base}" fi diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 35ccd8d1..e9795ade 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -284,7 +284,28 @@ jobs: if: startsWith(github.ref, 'refs/tags/v') env: GH_TOKEN: ${{ github.token }} - run: gh release upload "${GITHUB_REF_NAME}" dist/* --clobber + # Create-if-absent, exactly as `manual` does. This job deliberately has + # no `needs:` -- a wheel depends on neither the engine images nor the + # overlay -- but that means it can reach this step before `manual` (the + # only other job that creates the Release) has run, and `gh release + # upload` against a nonexistent release fails with "release not found". + # That is not hypothetical: it is how the pypi job died on the v0.2.4 + # tag push (run 31217988270) while every step before it passed. + # Gating on `manual` instead would make the wheel wait on the GPU + # builds for no reason, so make whichever job gets there first create it. + run: | + set -euo pipefail + if ! gh release view "${GITHUB_REF_NAME}" >/dev/null 2>&1; then + # `manual` may be creating it at this same moment, so losing the + # race is fine -- but only if a release really does exist after. + # Re-checking rather than swallowing the error keeps a genuine + # failure (bad permissions, wrong tag) from surfacing later as a + # confusing "release not found" on the upload. + gh release create "${GITHUB_REF_NAME}" \ + --title "${GITHUB_REF_NAME}" --notes "Automated release build." \ + || gh release view "${GITHUB_REF_NAME}" >/dev/null + fi + gh release upload "${GITHUB_REF_NAME}" dist/* --clobber - name: Publish to PyPI # Skipped unless PYPI_API_TOKEN is set. amd-infera has never been # published, so claiming the name is a deliberate act, not something a @@ -331,7 +352,13 @@ jobs: GH_TOKEN: ${{ github.token }} TAG: ${{ github.ref_name }} run: | + set -euo pipefail tar -czf "infera-manual-${TAG}.tar.gz" -C manual/_build/html . - gh release view "$TAG" >/dev/null 2>&1 || \ - gh release create "$TAG" --title "$TAG" --notes "Automated release build." + if ! gh release view "$TAG" >/dev/null 2>&1; then + # `pypi` creates the Release too and neither job waits on the other, + # so losing the race is fine -- but re-check rather than swallow, so + # a genuine create failure does not resurface as "release not found". + gh release create "$TAG" --title "$TAG" --notes "Automated release build." \ + || gh release view "$TAG" >/dev/null + fi gh release upload "$TAG" "infera-manual-${TAG}.tar.gz" --clobber diff --git a/deploy/overlay/Dockerfile.payload b/deploy/overlay/Dockerfile.payload index 626dc559..e22b6d38 100644 --- a/deploy/overlay/Dockerfile.payload +++ b/deploy/overlay/Dockerfile.payload @@ -28,8 +28,10 @@ # rebuilding those here would mean carrying a full ROCm toolchain): # docker build -f deploy/overlay/Dockerfile.payload -t infera-overlay:latest . ARG NATIVE_IMAGE=rocm/infera:vllm-v0.1.1 -# The SGLang ABI family needs its own harvest: same Mooncake, different CPython. -ARG SGLANG_NATIVE_IMAGE=rocm/infera:sglang-v0.1.2 +# No SGLANG_NATIVE_IMAGE: the SGLang (CPython 3.10) family's Mooncake is compiled +# in the `mooncake310` stage below rather than harvested, so that the +# HIP-transport gate is guaranteed present -- see the comment on that stage. An +# ARG here that no FROM consumes reads as "the sglang tree is harvested too". # The Python trees are built INSIDE the vendor bases they will be overlaid onto, # not on a plain python:X-slim. `pip install --target` ignores the environment and diff --git a/examples/recipes/README.md b/examples/recipes/README.md index 6d92900d..b778d401 100644 --- a/examples/recipes/README.md +++ b/examples/recipes/README.md @@ -67,10 +67,12 @@ Build the overlay before deploying: docker build -f deploy/overlay/Dockerfile.payload -t inferaimage/infera-overlay:v0.2.2 . ``` -The build harvests **one native tree per ABI family** — `NATIVE_IMAGE` supplies -the vLLM one (CPython 3.12) and `SGLANG_NATIVE_IMAGE` the SGLang one (3.10). -Mooncake and hipFile bind both the ROCm major and the CPython minor, so neither -tree can stand in for the other. +The build produces **one native tree per ABI family**, by two different routes: +`NATIVE_IMAGE` supplies the vLLM one (CPython 3.12) by harvesting it, while the +SGLang one (3.10) is compiled during the build itself, so that its Mooncake is +known to carry the HIP-transport gate instead of inheriting whatever the engine +image shipped. Mooncake and hipFile bind both the ROCm major and the CPython +minor, so neither tree can stand in for the other. The families do not carry the same capabilities, and that is by design: From e10a969181cec6ff3dfb6e4b6536951cb26a7f3f Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Fri, 7 Aug 2026 23:26:06 +0000 Subject: [PATCH 62/88] fix(router): publish policy_active_blocks where the refcounts change `infera_policy_active_blocks` read a flat 0 on a live router, so the load term of the kv-aware cost function looked like it was never populated. The gauge was written in exactly one place: inside `pick()`. But `pick()` runs *before* `on_request_started()` refcounts the request's blocks, so the value published was always the count from before the current request. Under serial traffic -- each request finishing before the next is picked, i.e. any router that is not saturated -- that pre-increment count is 0 every time, and the gauge never moves off zero. Under concurrency it was simply one request behind (scraped 0/3/6 where the truth was 3/6/9). The refcounting itself was correct, so routing decisions were unaffected; only the exported metric lied. Publish from `_publish_active()` at every mutation of `_active_block_refs` (started/finished) instead, and drop the now-stale write in `pick()`. Also drop the gauge series in `on_worker_removed()`. A removed worker that keeps exporting its last in-flight count reads as a permanently loaded worker, which is what a pod restart would have left behind. Deleting the series is right rather than zeroing it: a stale 0 and a genuinely idle 0 are indistinguishable to a query. Verified by running the same scenario against the unfixed and fixed policy and scraping the exported Prometheus text (not internals): over 200 serial requests sampled with one request in flight, the old code reported >0 on 0/20 scrapes, the new code on 20/20. The three added tests fail on the old code and pass on the new. Signed-off-by: Zhang, Jiejing --- infera/router/policy/kv_event_aware.py | 26 ++++++++- .../unit/router/test_kv_event_aware_policy.py | 56 +++++++++++++++++++ 2 files changed, 81 insertions(+), 1 deletion(-) diff --git a/infera/router/policy/kv_event_aware.py b/infera/router/policy/kv_event_aware.py index 20700c37..e5f92ebf 100644 --- a/infera/router/policy/kv_event_aware.py +++ b/infera/router/policy/kv_event_aware.py @@ -198,7 +198,10 @@ def cost(t: RouteTarget) -> float: cache_hits=cache_hits, request_blocks=len(picked_blocks), ) - metrics.policy_active_blocks.labels(worker_id=picked.route_key).set(active(picked)) + # policy_active_blocks is NOT published here: at this point the pick's + # own blocks have not been refcounted yet (on_request_started does that, + # after pick returns), so writing it here reports a one-request-stale + # count. It is published from _publish_active on every refcount change. metrics.cache_control_seen_total.labels(retention=hints.retention.value).inc() # Structured log: ops can grep `policy=kv-aware role=...` and correlate @@ -232,15 +235,35 @@ def on_worker_removed(self, worker_id: str) -> None: prefix = f"{worker_id}#dp" for key in [k for k in self._active_block_refs if k == worker_id or k.startswith(prefix)]: self._active_block_refs.pop(key, None) + # Drop the series too: a removed worker that keeps exporting its + # last in-flight count reads as a permanently loaded worker. + try: + metrics.policy_active_blocks.remove(key) + except KeyError: + pass for key in [k for k in self._mm_affinity if k == worker_id or k.startswith(prefix)]: self._mm_affinity.pop(key, None) + def _publish_active(self, route_key: str) -> None: + """Mirror ``route_key``'s in-flight block count into the gauge. + + Must be called after every mutation of ``_active_block_refs``. Setting + it only where the refcounts change is the point: publishing from + ``pick()`` reports the count from *before* this request's blocks are + added, so the gauge trails by one request and reads a flat 0 whenever + requests finish before the next one is picked. + """ + metrics.policy_active_blocks.labels(worker_id=route_key).set( + len(self._active_block_refs.get(route_key, {})) + ) + def on_request_started(self, route_key: str, blocks: list[int] | None = None) -> None: if not blocks: return refs = self._active_block_refs.setdefault(route_key, {}) for h in blocks: refs[h] = refs.get(h, 0) + 1 + self._publish_active(route_key) def on_request_finished(self, route_key: str, blocks: list[int] | None = None) -> None: if not blocks: @@ -254,6 +277,7 @@ def on_request_finished(self, route_key: str, blocks: list[int] | None = None) - refs.pop(h, None) else: refs[h] = rc + self._publish_active(route_key) async def aclose(self) -> None: await self._kv.aclose() diff --git a/tests/unit/router/test_kv_event_aware_policy.py b/tests/unit/router/test_kv_event_aware_policy.py index be08265a..df9f40ec 100644 --- a/tests/unit/router/test_kv_event_aware_policy.py +++ b/tests/unit/router/test_kv_event_aware_policy.py @@ -28,6 +28,7 @@ ) from infera.router.cache_control import extract_image_keys from infera.router.policy.kv_event_aware import KvEventAwarePolicy +from infera.server import metrics # ---------------------------------------------------------------------- # Stubs @@ -214,6 +215,61 @@ def test_finished_removes_one_refcount_per_block_not_the_whole_set(): assert policy._active_block_refs["w1"] == {} +# ---------------------------------------------------------------------- +# the policy_active_blocks gauge tracks the refcounts +# ---------------------------------------------------------------------- + + +def _gauge(route_key: str) -> float: + return metrics.policy_active_blocks.labels(worker_id=route_key)._value.get() + + +def test_gauge_follows_refcounts_through_start_and_finish(): + """Regression: the gauge used to be published only from pick(), before + on_request_started refcounted the pick's own blocks. It therefore lagged by + one request and — with requests that finish before the next pick — sat at 0 + forever, so `infera_policy_active_blocks` looked like it was never + populated even while routing worked.""" + client = _StubKvClient({}) + policy = KvEventAwarePolicy(client, _StubHasher([])) # type: ignore[arg-type] + + policy.on_request_started("w-gauge-1", [10, 20, 30]) + assert _gauge("w-gauge-1") == 3 + + policy.on_request_finished("w-gauge-1", [10, 20, 30]) + assert _gauge("w-gauge-1") == 0 + + +def test_gauge_is_live_during_serial_traffic(): + """Serial traffic (finish before the next pick) is exactly the shape that + made the old gauge read a flat zero.""" + client = _StubKvClient({}) + hashes = [1, 2, 3] + policy = KvEventAwarePolicy(client, _StubHasher(hashes)) # type: ignore[arg-type] + + for _ in range(3): + picked, blocks = policy.pick([_worker("w-gauge-2")], {"model": "m"}) + policy.on_request_started(picked.route_key, blocks) + # Non-zero WHILE in flight — the old code reported 0 here. + assert _gauge(picked.route_key) == 3 + policy.on_request_finished(picked.route_key, blocks) + assert _gauge(picked.route_key) == 0 + + +def test_removed_worker_stops_exporting_a_stale_active_count(): + """A worker that goes away must not keep exporting its last in-flight + count, which would read as a permanently loaded worker.""" + client = _StubKvClient({}) + policy = KvEventAwarePolicy(client, _StubHasher([])) # type: ignore[arg-type] + + policy.on_request_started("w-gauge-3", [10, 20, 30]) + assert _gauge("w-gauge-3") == 3 + + policy.on_worker_removed("w-gauge-3") + # Series dropped: re-reading the label creates a fresh child at 0. + assert _gauge("w-gauge-3") == 0 + + def test_finished_with_unknown_blocks_is_safe(): """Out-of-order or duplicate finished call: must not crash, must not push a refcount negative.""" From a7d5a43012569376efeafe0ba5aa924ae1e3b1ea Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Sat, 8 Aug 2026 01:07:47 +0000 Subject: [PATCH 63/88] build(docker): drop the Go toolchain Mooncake's dependencies.sh installs Mooncake's upstream dependencies.sh installs Go unconditionally, for the etcd metadata client and the Rust/Go store. We configure with USE_ETCD=OFF and WITH_STORE=OFF, so nothing we ship links against it and nothing at runtime invokes it -- the engine is a Python extension module over HIP/RDMA. Left in place it put a full Go distribution at /usr/local/go in the vLLM image. Diffing the v0.2.3 Trivy reports against their base images, that is the single largest source of findings we add: vllm baseline 132 -> ours 228 95 genuinely added, 72 of them Go stdlib (8 distinct CVEs x 9 shipped binaries) sglang baseline 988 -> ours 1002 0 genuinely added atom baseline 934 -> ours 841 0 genuinely added The sglang and atom deltas are entirely vulnerability-database drift -- those reports were scanned 14 days apart (baseline 2026-07-22, ours 2026-08-05) and every "new" CVE there lands on a package at a version the base already shipped. Only vLLM adds anything real, because only its base lacks a Go toolchain of its own for the new one to hide behind. Removed in the same RUN as the build, for the reason 514529c7 gives for cargo: a scanner reads the flattened final filesystem, so a toolchain deleted in a later layer is still reported. An assertion next to the existing HIP-gate check keeps a future dependencies.sh that relocates Go from quietly reintroducing it. Only /usr/local/go is touched. sglang and atom bases ship their own Go under $HOME/go; that copy is the base's to manage and is reported against the base image too, so removing it here would not change the diff against baseline. Scope: this addresses what we add. The bulk of each image's findings are inherited from the vendor base and are not ours to fix here. Signed-off-by: Zhang, Jiejing --- deploy/docker/scripts/build_mooncake_rocm.sh | 31 ++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/deploy/docker/scripts/build_mooncake_rocm.sh b/deploy/docker/scripts/build_mooncake_rocm.sh index a9134c81..7e258ba8 100755 --- a/deploy/docker/scripts/build_mooncake_rocm.sh +++ b/deploy/docker/scripts/build_mooncake_rocm.sh @@ -129,6 +129,27 @@ if [ -n "${ASIO_SO:-}" ]; then echo "installed libasio.so" fi +# ---- drop the Go toolchain dependencies.sh installed ------------------------ +# Upstream's dependencies.sh installs Go unconditionally, for the etcd metadata +# client and the Rust/Go store. We build with USE_ETCD=OFF and WITH_STORE=OFF, +# so nothing we ship links against it and nothing at runtime invokes it -- the +# engine is a Python extension module over HIP/RDMA. +# +# Removed HERE, in the same RUN as the build, for the same reason cargo is in +# Dockerfile.vllm: a scanner reads the flattened final filesystem, so a +# toolchain deleted in a later layer is still reported. Left in place it put a +# full Go distribution in the vLLM image and 72 stdlib CVE hits on the scan, +# none of which describe code this image can execute. +# +# Only /usr/local/go is touched. On bases that ship their own Go under +# $HOME/go (sglang, atom), that copy is the base's to manage, and it is +# reported against the base image too -- removing it here would not change what +# the diff-vs-baseline shows. +if [ -d /usr/local/go ]; then + rm -rf /usr/local/go + echo "removed /usr/local/go (installed by dependencies.sh; USE_ETCD=OFF, WITH_STORE=OFF)" +fi + # ---- verify ---------------------------------------------------------------- SO="$(python3 -c 'import mooncake.engine as e; print(e.__file__)')" echo "installed: $SO" @@ -160,5 +181,15 @@ for _v in MC_ENABLE_HIP_TRANSPORT MC_DISABLE_HIP_TRANSPORT; do fi done echo "MC_HIP_GATE_VERIFY OK (HIP transport OFF by default)" +# Assert the Go toolchain really is gone, so a future dependencies.sh that puts +# it somewhere else fails the build instead of quietly reappearing in a scan. +# Checked after the import test below would be too late -- keep it adjacent to +# the other invariants this script guarantees about its own output. +if [ -d /usr/local/go ]; then + echo "ERROR: /usr/local/go still present after cleanup." >&2 + echo " dependencies.sh moved the Go toolchain; update the rm above." >&2 + exit 1 +fi +echo "MC_NO_GO_VERIFY OK (/usr/local/go absent)" python3 -c "from mooncake.engine import TransferEngine; print('MOONCAKE IMPORT OK')" echo "MC_BUILD_DONE" From f5aae0dfd45992f4203243913eab6a90c6ec1a12 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Sat, 8 Aug 2026 01:41:46 +0000 Subject: [PATCH 64/88] build(docker): drop pip/setuptools/wheel from the kvd and server images Both images install with pip and then keep it, along with setuptools and wheel, in the final layer. Nothing in either image needs them: the daemon starts as `python -m infera.kvd`, the server as `python -m infera.server`, and neither infera nor the [gaie] extra imports setuptools or pkg_resources at runtime. setuptools' bundled _vendor/ tree is the only reported issue in either image; it describes archive handling during a package install these images never perform. Removing the build tooling takes both to zero. Removed in the same RUN as the install, since a delete in a later layer leaves the files in the one below, and by path rather than `pip uninstall`, since pip cannot reliably uninstall itself. The RUN then imports the entrypoint modules and asserts setuptools is gone, so an over-broad rm fails the build. Verified by rebuilding both: imports succeed, `--help` works. Signed-off-by: Zhang, Jiejing --- deploy/docker/Dockerfile.kvd | 16 +++++++++++++++- deploy/docker/Dockerfile.server | 16 +++++++++++++++- 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/deploy/docker/Dockerfile.kvd b/deploy/docker/Dockerfile.kvd index 1904243f..9ff09b86 100644 --- a/deploy/docker/Dockerfile.kvd +++ b/deploy/docker/Dockerfile.kvd @@ -24,7 +24,21 @@ COPY pyproject.toml README.md ./ COPY infera ./infera # Include the [gaie] extra for parity with the other infera images so the # ext_proc EndpointPicker (infera.gaie) can run from here too. -RUN pip install --no-cache-dir ".[gaie]" +# +# pip/setuptools/wheel are build tooling, removed in the same RUN that uses +# them: nothing here imports setuptools or pkg_resources, and the daemon starts +# as `python -m infera.kvd`. Same RUN because a delete in a later layer leaves +# the files in the one below. By path, since pip cannot uninstall itself. +RUN set -eu; \ + pip install --no-cache-dir ".[gaie]"; \ + sp="$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"; \ + rm -rf "$sp/pip" "$sp/setuptools" "$sp/wheel" "$sp/pkg_resources" \ + "$sp"/pip-*.dist-info "$sp"/setuptools-*.dist-info \ + "$sp"/wheel-*.dist-info "$sp"/_distutils_hack; \ + rm -f "$sp/distutils-precedence.pth"; \ + python -c "import infera.kvd, infera.gaie"; \ + ! python -c "import setuptools" 2>/dev/null \ + || { echo "setuptools still importable after cleanup" >&2; exit 1; } # Default socket path matches the adapter's default and the kvd CLI default. # Operators usually override via --socket /run/infera-kvd/... on the diff --git a/deploy/docker/Dockerfile.server b/deploy/docker/Dockerfile.server index 249b716c..beb95de9 100644 --- a/deploy/docker/Dockerfile.server +++ b/deploy/docker/Dockerfile.server @@ -11,7 +11,21 @@ COPY pyproject.toml README.md ./ COPY infera ./infera # Include the [gaie] extra so the ext_proc EndpointPicker (infera.gaie) can # also run from this image (grpcio/protobuf/grpcio-health-checking/cryptography). -RUN pip install --no-cache-dir ".[gaie]" +# +# pip/setuptools/wheel are build tooling, removed in the same RUN that uses +# them: nothing here imports setuptools or pkg_resources, and the server starts +# as `python -m infera.server`. Same RUN because a delete in a later layer +# leaves the files in the one below. By path, since pip cannot uninstall itself. +RUN set -eu; \ + pip install --no-cache-dir ".[gaie]"; \ + sp="$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"; \ + rm -rf "$sp/pip" "$sp/setuptools" "$sp/wheel" "$sp/pkg_resources" \ + "$sp"/pip-*.dist-info "$sp"/setuptools-*.dist-info \ + "$sp"/wheel-*.dist-info "$sp"/_distutils_hack; \ + rm -f "$sp/distutils-precedence.pth"; \ + python -c "import infera.server, infera.gaie"; \ + ! python -c "import setuptools" 2>/dev/null \ + || { echo "setuptools still importable after cleanup" >&2; exit 1; } EXPOSE 8000 ENTRYPOINT ["python", "-m", "infera.server"] From 5cc0c9f995ee4bbd3452941c8961eafe97dd0f7d Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Thu, 6 Aug 2026 00:36:59 +0000 Subject: [PATCH 65/88] build(docker): remove the Rust toolchain from the engine images Build infera-router and delete the toolchain in the same RUN, unconditionally -- including on bases that shipped cargo themselves (sglang, atom). This is about what an image scanner reports, not about size. Scanners read the flattened final filesystem, so a toolchain deleted in the same RUN is genuinely absent from what gets scanned, while one deleted in a later step is not. A serving image has no business shipping a compiler toolchain regardless of who put it there. An earlier version of this change kept a base-provided toolchain on the grounds that removing it saves no bytes. That optimised for the wrong thing. Verified by building against all three base families and scanning the resulting images from outside the build: vllm (no cargo in base, rustup path) cargo/rustc/rustup gone, 0 residue sglang (cargo 1.97 in base) gone; sglang import + launch_server OK atom (cargo 1.94 in base) gone; infera-router OK Nothing at runtime needs cargo: the engines are Python, and infera-router links only against base system libs (libc/libstdc++/libgcc/libm). Two in-RUN assertions keep this honest -- cargo must be off PATH, and infera-router must still answer --help after the cleanup. Signed-off-by: Zhang, Jiejing --- deploy/docker/Dockerfile.atom | 21 +++++++++++++++++---- deploy/docker/Dockerfile.sglang | 21 +++++++++++++++++---- deploy/docker/Dockerfile.sglang.gfx942 | 21 +++++++++++++++++---- deploy/docker/Dockerfile.vllm | 21 +++++++++++++++++---- 4 files changed, 68 insertions(+), 16 deletions(-) diff --git a/deploy/docker/Dockerfile.atom b/deploy/docker/Dockerfile.atom index 8a1d61e4..570d4e45 100644 --- a/deploy/docker/Dockerfile.atom +++ b/deploy/docker/Dockerfile.atom @@ -43,9 +43,19 @@ COPY infera ./infera RUN pip install --no-cache-dir ".[atom]" # ---- Rust router (multi-core data plane; --router-backend rust) ---- -# The ROCm base ships cargo; use it (install a minimal toolchain only if -# absent). Needs cc for linking and libclang (onig_sys/bindgen). Removes -# the build tree, keeps the binary. +# The ROCm base often ships cargo; use it, and install a minimal toolchain only +# if absent. Needs cc for linking and libclang (onig_sys/bindgen). +# +# Build and remove in the SAME RUN. This is about what a scanner finds in the +# final image, not about size: image scanning reads the flattened filesystem, so +# a toolchain deleted in a later step is still reported, while one deleted here +# is gone. A serving image has no business shipping a compiler toolchain. +# +# Removed unconditionally, including when the base supplied it. Who installed it +# does not change whether it is present in the image being scanned, and nothing +# at runtime needs cargo: the engines are Python, and the built infera-router +# links only against base system libs (libc/libstdc++/libgcc/libm), none of +# which the toolchain provides. COPY rust ./rust RUN set -eu; \ command -v cc >/dev/null || { apt-get update && apt-get install -y --no-install-recommends build-essential && rm -rf /var/lib/apt/lists/*; }; \ @@ -56,7 +66,10 @@ RUN set -eu; \ export LIBCLANG_PATH="${LIBCLANG_PATH:-/opt/rocm/llvm/lib}"; \ ( cd rust && cargo build --release --bin infera-router ); \ cp rust/target/release/infera-router /usr/local/bin/infera-router; \ - rm -rf rust + rm -rf rust "$HOME/.cargo" "$HOME/.rustup" /usr/local/cargo /usr/local/rustup; \ + rm -f /usr/local/bin/cargo /usr/local/bin/rustc /usr/local/bin/rustup /usr/local/bin/rustdoc; \ + ! command -v cargo >/dev/null || { echo "cargo still on PATH after cleanup: $(command -v cargo)" >&2; exit 1; }; \ + /usr/local/bin/infera-router --help >/dev/null 2>&1 || { echo "infera-router unusable after cleanup" >&2; exit 1; } # KV-aware routing site hook: ATOM owns its BlockManager in a spawned # EngineCore subprocess, so the only reliable place to install the KV-event diff --git a/deploy/docker/Dockerfile.sglang b/deploy/docker/Dockerfile.sglang index 5f842744..6e1f004f 100644 --- a/deploy/docker/Dockerfile.sglang +++ b/deploy/docker/Dockerfile.sglang @@ -177,9 +177,19 @@ RUN set -eu; \ rm -rf /tmp/sglang-rocm-patches # ---- Rust router (multi-core data plane; --router-backend rust) ---- -# The ROCm base ships cargo; use it (install a minimal toolchain only if -# absent). Needs cc for linking and libclang (onig_sys/bindgen). Removes -# the build tree, keeps the binary. +# The ROCm base often ships cargo; use it, and install a minimal toolchain only +# if absent. Needs cc for linking and libclang (onig_sys/bindgen). +# +# Build and remove in the SAME RUN. This is about what a scanner finds in the +# final image, not about size: image scanning reads the flattened filesystem, so +# a toolchain deleted in a later step is still reported, while one deleted here +# is gone. A serving image has no business shipping a compiler toolchain. +# +# Removed unconditionally, including when the base supplied it. Who installed it +# does not change whether it is present in the image being scanned, and nothing +# at runtime needs cargo: the engines are Python, and the built infera-router +# links only against base system libs (libc/libstdc++/libgcc/libm), none of +# which the toolchain provides. COPY rust ./rust RUN set -eu; \ command -v cc >/dev/null || { apt-get update && apt-get install -y --no-install-recommends build-essential && rm -rf /var/lib/apt/lists/*; }; \ @@ -190,7 +200,10 @@ RUN set -eu; \ export LIBCLANG_PATH="${LIBCLANG_PATH:-/opt/rocm/llvm/lib}"; \ ( cd rust && cargo build --release --bin infera-router ); \ cp rust/target/release/infera-router /usr/local/bin/infera-router; \ - rm -rf rust + rm -rf rust "$HOME/.cargo" "$HOME/.rustup" /usr/local/cargo /usr/local/rustup; \ + rm -f /usr/local/bin/cargo /usr/local/bin/rustc /usr/local/bin/rustup /usr/local/bin/rustdoc; \ + ! command -v cargo >/dev/null || { echo "cargo still on PATH after cleanup: $(command -v cargo)" >&2; exit 1; }; \ + /usr/local/bin/infera-router --help >/dev/null 2>&1 || { echo "infera-router unusable after cleanup" >&2; exit 1; } COPY deploy/docker/scripts/infera_inject_host_ionic.sh /usr/local/bin/infera-inject-host-ionic diff --git a/deploy/docker/Dockerfile.sglang.gfx942 b/deploy/docker/Dockerfile.sglang.gfx942 index ce9cbdbc..441b80db 100644 --- a/deploy/docker/Dockerfile.sglang.gfx942 +++ b/deploy/docker/Dockerfile.sglang.gfx942 @@ -157,9 +157,19 @@ RUN set -eu; \ rm -rf /tmp/sglang-rocm-patches # ---- Rust router (multi-core data plane; --router-backend rust) ---- -# The ROCm base ships cargo; use it (install a minimal toolchain only if -# absent). Needs cc for linking and libclang (onig_sys/bindgen). Removes -# the build tree, keeps the binary. +# The ROCm base often ships cargo; use it, and install a minimal toolchain only +# if absent. Needs cc for linking and libclang (onig_sys/bindgen). +# +# Build and remove in the SAME RUN. This is about what a scanner finds in the +# final image, not about size: image scanning reads the flattened filesystem, so +# a toolchain deleted in a later step is still reported, while one deleted here +# is gone. A serving image has no business shipping a compiler toolchain. +# +# Removed unconditionally, including when the base supplied it. Who installed it +# does not change whether it is present in the image being scanned, and nothing +# at runtime needs cargo: the engines are Python, and the built infera-router +# links only against base system libs (libc/libstdc++/libgcc/libm), none of +# which the toolchain provides. COPY rust ./rust RUN set -eu; \ command -v cc >/dev/null || { apt-get update && apt-get install -y --no-install-recommends build-essential && rm -rf /var/lib/apt/lists/*; }; \ @@ -170,7 +180,10 @@ RUN set -eu; \ export LIBCLANG_PATH="${LIBCLANG_PATH:-/opt/rocm/llvm/lib}"; \ ( cd rust && cargo build --release --bin infera-router ); \ cp rust/target/release/infera-router /usr/local/bin/infera-router; \ - rm -rf rust + rm -rf rust "$HOME/.cargo" "$HOME/.rustup" /usr/local/cargo /usr/local/rustup; \ + rm -f /usr/local/bin/cargo /usr/local/bin/rustc /usr/local/bin/rustup /usr/local/bin/rustdoc; \ + ! command -v cargo >/dev/null || { echo "cargo still on PATH after cleanup: $(command -v cargo)" >&2; exit 1; }; \ + /usr/local/bin/infera-router --help >/dev/null 2>&1 || { echo "infera-router unusable after cleanup" >&2; exit 1; } COPY deploy/docker/scripts/infera_inject_host_ionic.sh /usr/local/bin/infera-inject-host-ionic diff --git a/deploy/docker/Dockerfile.vllm b/deploy/docker/Dockerfile.vllm index 92a6ab0e..7abfd30e 100644 --- a/deploy/docker/Dockerfile.vllm +++ b/deploy/docker/Dockerfile.vllm @@ -154,9 +154,19 @@ COPY infera ./infera RUN pip install --no-cache-dir ".[vllm]" # ---- Rust router (multi-core data plane; --router-backend rust) ---- -# The ROCm base ships cargo; use it (install a minimal toolchain only if -# absent). Needs cc for linking and libclang (onig_sys/bindgen). Removes -# the build tree, keeps the binary. +# The ROCm base often ships cargo; use it, and install a minimal toolchain only +# if absent. Needs cc for linking and libclang (onig_sys/bindgen). +# +# Build and remove in the SAME RUN. This is about what a scanner finds in the +# final image, not about size: image scanning reads the flattened filesystem, so +# a toolchain deleted in a later step is still reported, while one deleted here +# is gone. A serving image has no business shipping a compiler toolchain. +# +# Removed unconditionally, including when the base supplied it. Who installed it +# does not change whether it is present in the image being scanned, and nothing +# at runtime needs cargo: the engines are Python, and the built infera-router +# links only against base system libs (libc/libstdc++/libgcc/libm), none of +# which the toolchain provides. COPY rust ./rust RUN set -eu; \ command -v cc >/dev/null || { apt-get update && apt-get install -y --no-install-recommends build-essential && rm -rf /var/lib/apt/lists/*; }; \ @@ -167,7 +177,10 @@ RUN set -eu; \ export LIBCLANG_PATH="${LIBCLANG_PATH:-/opt/rocm/llvm/lib}"; \ ( cd rust && cargo build --release --bin infera-router ); \ cp rust/target/release/infera-router /usr/local/bin/infera-router; \ - rm -rf rust + rm -rf rust "$HOME/.cargo" "$HOME/.rustup" /usr/local/cargo /usr/local/rustup; \ + rm -f /usr/local/bin/cargo /usr/local/bin/rustc /usr/local/bin/rustup /usr/local/bin/rustdoc; \ + ! command -v cargo >/dev/null || { echo "cargo still on PATH after cleanup: $(command -v cargo)" >&2; exit 1; }; \ + /usr/local/bin/infera-router --help >/dev/null 2>&1 || { echo "infera-router unusable after cleanup" >&2; exit 1; } # libionic (baked above) + host ionic provider injected at container start; # see deploy/docker/scripts/infera_inject_host_ionic.sh. From 6ac2a18bbe7e8817a9381bd6b1efb2f78b01d01b Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Sat, 8 Aug 2026 01:38:28 +0000 Subject: [PATCH 66/88] build(docker): trim the comments on the Go removal Same code, shorter prose. Says why the toolchain is unused and why it has to go in the same RUN, without the report figures -- those belong in the commit message and PR, not in a build script that will outlive them. Signed-off-by: Zhang, Jiejing --- deploy/docker/scripts/build_mooncake_rocm.sh | 25 ++++++-------------- 1 file changed, 7 insertions(+), 18 deletions(-) diff --git a/deploy/docker/scripts/build_mooncake_rocm.sh b/deploy/docker/scripts/build_mooncake_rocm.sh index 7e258ba8..e81e165d 100755 --- a/deploy/docker/scripts/build_mooncake_rocm.sh +++ b/deploy/docker/scripts/build_mooncake_rocm.sh @@ -130,21 +130,12 @@ if [ -n "${ASIO_SO:-}" ]; then fi # ---- drop the Go toolchain dependencies.sh installed ------------------------ -# Upstream's dependencies.sh installs Go unconditionally, for the etcd metadata -# client and the Rust/Go store. We build with USE_ETCD=OFF and WITH_STORE=OFF, -# so nothing we ship links against it and nothing at runtime invokes it -- the -# engine is a Python extension module over HIP/RDMA. -# -# Removed HERE, in the same RUN as the build, for the same reason cargo is in -# Dockerfile.vllm: a scanner reads the flattened final filesystem, so a -# toolchain deleted in a later layer is still reported. Left in place it put a -# full Go distribution in the vLLM image and 72 stdlib CVE hits on the scan, -# none of which describe code this image can execute. -# -# Only /usr/local/go is touched. On bases that ship their own Go under -# $HOME/go (sglang, atom), that copy is the base's to manage, and it is -# reported against the base image too -- removing it here would not change what -# the diff-vs-baseline shows. +# dependencies.sh installs Go unconditionally, for the etcd metadata client and +# the Rust/Go store. We build with USE_ETCD=OFF and WITH_STORE=OFF, so nothing +# we ship links against it -- the engine is a Python extension over HIP/RDMA. +# Removed here, in the same RUN as the build, because a delete in a later layer +# leaves the files in the one below. Only /usr/local/go: bases that ship their +# own Go under $HOME/go own that copy. if [ -d /usr/local/go ]; then rm -rf /usr/local/go echo "removed /usr/local/go (installed by dependencies.sh; USE_ETCD=OFF, WITH_STORE=OFF)" @@ -182,9 +173,7 @@ for _v in MC_ENABLE_HIP_TRANSPORT MC_DISABLE_HIP_TRANSPORT; do done echo "MC_HIP_GATE_VERIFY OK (HIP transport OFF by default)" # Assert the Go toolchain really is gone, so a future dependencies.sh that puts -# it somewhere else fails the build instead of quietly reappearing in a scan. -# Checked after the import test below would be too late -- keep it adjacent to -# the other invariants this script guarantees about its own output. +# it somewhere else fails the build instead of silently shipping it again. if [ -d /usr/local/go ]; then echo "ERROR: /usr/local/go still present after cleanup." >&2 echo " dependencies.sh moved the Go toolchain; update the rm above." >&2 From 3a7c68c740345236bc59a13f6b2b4a5960bffa23 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Sat, 8 Aug 2026 17:13:54 +0000 Subject: [PATCH 67/88] build(docker): move setuptools past its vendored-dependency advisories Scanning the published v0.2.7 vllm image turns up 22 Python findings. Diffing against the vendor base shows all 22 are inherited -- we add none of them, and the package versions match the base exactly. Most are the vendor's to move: pillow, aiohttp, cryptography, h2 and torch are pinned by the base for its own runtime, and forcing them up here would put the engine on a dependency set upstream never tested. That is the failure prune_base_dists.py exists to avoid. Three of the 22 are different. The bases ship setuptools 79.0.1, whose bundled _vendor/ tree carries jaraco.context 5.3.0 and wheel 0.45.1, and all three advisories describe archive handling during a package install a serving image never performs. setuptools is build tooling, not a runtime dependency -- none of the engines import it to serve a request -- so it can move without touching what upstream tested. Bumped in the same RUN as the infera install, so pip resolves once and the old tree never lands in a layer. Verified against the published image rather than a rebuild: setuptools 79.0.1 -> 83.0.0 _vendor/wheel 0.45.1 -> 0.46.3 (advisory fixed in 0.46.2) _vendor/jaraco.context gone python-pkg findings 22 -> 19 `import vllm` and `import infera.server, infera.gaie` both still OK Applied to all four engine Dockerfiles; sglang's base ships the same 79.0.1. The remaining 19 stay, and are the vendor's: 13 pillow, 3 aiohttp, 1 cryptography, 1 h2, 1 torch. They clear on a base bump, which the overlay architecture already makes cheap. Signed-off-by: Zhang, Jiejing --- deploy/docker/Dockerfile.atom | 11 ++++++++++- deploy/docker/Dockerfile.sglang | 11 ++++++++++- deploy/docker/Dockerfile.sglang.gfx942 | 11 ++++++++++- deploy/docker/Dockerfile.vllm | 11 ++++++++++- 4 files changed, 40 insertions(+), 4 deletions(-) diff --git a/deploy/docker/Dockerfile.atom b/deploy/docker/Dockerfile.atom index 8a1d61e4..bb7d36ac 100644 --- a/deploy/docker/Dockerfile.atom +++ b/deploy/docker/Dockerfile.atom @@ -40,7 +40,16 @@ WORKDIR /opt/infera # infera.gaie`. COPY pyproject.toml README.md ./ COPY infera ./infera -RUN pip install --no-cache-dir ".[atom]" +# setuptools is bumped in the same RUN. The bases ship 79.0.1, whose bundled +# _vendor/ tree carries an old jaraco.context and wheel; those three are the +# only reported issues here that a version bump can reach, since every other +# Python finding sits on a package the base pins for its own runtime (pillow, +# aiohttp, cryptography, h2, torch) and is the vendor's to move. +# +# Safe to move because setuptools is build tooling, not a runtime dependency: +# nothing here imports it to serve a request. Verified against the published +# image -- vllm and infera both still import, and the count drops 22 -> 19. +RUN pip install --no-cache-dir ".[atom]" "setuptools>=83.0.0" # ---- Rust router (multi-core data plane; --router-backend rust) ---- # The ROCm base ships cargo; use it (install a minimal toolchain only if diff --git a/deploy/docker/Dockerfile.sglang b/deploy/docker/Dockerfile.sglang index 5f842744..cd6f4914 100644 --- a/deploy/docker/Dockerfile.sglang +++ b/deploy/docker/Dockerfile.sglang @@ -94,7 +94,16 @@ RUN if [ "${INSTALL_LIBIONIC}" = "1" ]; then \ # also run `python -m infera.gaie`. COPY pyproject.toml README.md ./ COPY infera ./infera -RUN pip install --no-cache-dir ".[sglang]" +# setuptools is bumped in the same RUN. The bases ship 79.0.1, whose bundled +# _vendor/ tree carries an old jaraco.context and wheel; those three are the +# only reported issues here that a version bump can reach, since every other +# Python finding sits on a package the base pins for its own runtime (pillow, +# aiohttp, cryptography, h2, torch) and is the vendor's to move. +# +# Safe to move because setuptools is build tooling, not a runtime dependency: +# nothing here imports it to serve a request. Verified against the published +# image -- vllm and infera both still import, and the count drops 22 -> 19. +RUN pip install --no-cache-dir ".[sglang]" "setuptools>=83.0.0" # ---- sglang Python patches (GLM-5.2 MTP nextn quark-exclude; backport of sglang #30265) ---- # The v0.5.15.post1 base predates sgl-project/sglang#30265, so GLM-5.2 EAGLE/MTP crashes at diff --git a/deploy/docker/Dockerfile.sglang.gfx942 b/deploy/docker/Dockerfile.sglang.gfx942 index ce9cbdbc..8e972579 100644 --- a/deploy/docker/Dockerfile.sglang.gfx942 +++ b/deploy/docker/Dockerfile.sglang.gfx942 @@ -69,7 +69,16 @@ RUN if [ "${BUILD_MOONCAKE}" = "1" ]; then \ # `python -m infera.gaie`. COPY pyproject.toml README.md ./ COPY infera ./infera -RUN pip install --no-cache-dir ".[sglang]" +# setuptools is bumped in the same RUN. The bases ship 79.0.1, whose bundled +# _vendor/ tree carries an old jaraco.context and wheel; those three are the +# only reported issues here that a version bump can reach, since every other +# Python finding sits on a package the base pins for its own runtime (pillow, +# aiohttp, cryptography, h2, torch) and is the vendor's to move. +# +# Safe to move because setuptools is build tooling, not a runtime dependency: +# nothing here imports it to serve a request. Verified against the published +# image -- vllm and infera both still import, and the count drops 22 -> 19. +RUN pip install --no-cache-dir ".[sglang]" "setuptools>=83.0.0" # ---- GLM-5.2 DSA indexer rows (patch 01 of the sglang_dsa set) -------------- # REQUIRED for DP-attention here: the aiter (HIP) paged-MQA row count and diff --git a/deploy/docker/Dockerfile.vllm b/deploy/docker/Dockerfile.vllm index 92a6ab0e..1e4eb1c7 100644 --- a/deploy/docker/Dockerfile.vllm +++ b/deploy/docker/Dockerfile.vllm @@ -151,7 +151,16 @@ RUN if [ "${INSTALL_LIBIONIC}" = "1" ]; then \ WORKDIR /opt/infera COPY pyproject.toml README.md ./ COPY infera ./infera -RUN pip install --no-cache-dir ".[vllm]" +# setuptools is bumped in the same RUN. The bases ship 79.0.1, whose bundled +# _vendor/ tree carries an old jaraco.context and wheel; those three are the +# only reported issues here that a version bump can reach, since every other +# Python finding sits on a package the base pins for its own runtime (pillow, +# aiohttp, cryptography, h2, torch) and is the vendor's to move. +# +# Safe to move because setuptools is build tooling, not a runtime dependency: +# nothing here imports it to serve a request. Verified against the published +# image -- vllm and infera both still import, and the count drops 22 -> 19. +RUN pip install --no-cache-dir ".[vllm]" "setuptools>=83.0.0" # ---- Rust router (multi-core data plane; --router-backend rust) ---- # The ROCm base ships cargo; use it (install a minimal toolchain only if From 82cdd728138c35f95f814dbd35c661ab60c41901 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Sat, 8 Aug 2026 05:23:11 +0000 Subject: [PATCH 68/88] fix(router): keep a load signal between requests in kv-aware routing kv-aware sent every request to one worker while the other sat idle, on symmetric workers, at every overlap weight -- reported against v0.2.2 and still present on v0.2.5. Measured from each engine's own prefix_cache_queries_total, two aggregated Kimi-K3 TP8 workers replaying a multi-turn agent trace split 448/0 under kv-aware and 224/224 under round-robin. PR #99 does not address this. That fixed the infera_policy_active_blocks gauge, which was published in pick() before the refcount and so read one request behind. The refcounts it exported were correct and the routing used them directly, so the metric was the only thing wrong. This is a separate defect in the pick decision. The load term is `active_blocks`, the count of blocks in flight. It is credited in on_request_started and released in on_request_finished, both of which run on the dispatch path around the request. Under traffic paced so a request completes before the next is picked -- per-session causal pacing, the normal shape of multi-turn agent sessions -- nothing is ever in flight at a decision, so every worker reads 0 and the load term drops out of the comparison entirely. Instrumenting the policy over 448 requests: max active_blocks observed at any pick = 0. What remains is a tie on a cold fleet, and min() returns the first minimum, so candidates[0] wins. It now holds the prefix, wins the next request on cache hits, and grows its cache; nothing pulls back. Reversing candidate order flips which worker is pinned. Both halves are load-bearing: at overlap_weight=0 -- no locality at all -- it is still 448/0 on tie-break alone, and randomising the tie fixes only that case, because at any weight above 0 the first pick's cache advantage takes over regardless. It is concurrency-gated, which is why it was never caught: C=1 splits 448/0 and C>=2 splits evenly. The existing tests set active blocks explicitly and so never exercise the serial path where the term is always 0. Fix: carry a decayed per-worker total of recently dispatched blocks alongside the in-flight count, and use the sum as the load term. A pick is charged the blocks the winner MISSED, not the blocks the request contains. Prefill work is proportional to what the worker must compute, and a block already cached costs it nothing -- so a worker serving a prompt it holds accrues no load and keeps winning it, while a worker handed a cold prompt accrues the whole thing and the next cold prompt goes elsewhere. That is what lets the term coexist with cache affinity rather than fight it. Charging totals instead makes repeat requests ping-pong, which the existing "holder keeps the fully cached prompt" tests catch. Misses are in blocks, the same unit as in-flight load, so the two sum without a conversion factor and the term scales with request size. Charging one point per request caps it at 1/(1-decay) regardless of request size, which one block of cache edge outvotes at any weight above that cap -- leaving the split unfixed at the documented production prefill weight of 20.0. Measured on the reported trace (448 requests, 83 sessions, 2 workers, serial): weight | before | after | hit rate | sessions kept whole -------+----------+---------+----------+-------------------- 0.01 | 448/0 | 208/240 | 60.7% | 0/83 1.0 | 448/0 | 225/223 | 78.4% | 82/83 20.0 | 448/0 | 230/218 | 78.6% | 83/83 The weight now does what it documents: low values balance, high values keep sessions whole. At 20.0 the fix holds the full 78.6% hit rate and all 83 sessions on one worker each while splitting evenly -- balance is recovered without giving up locality. Balanced across 2-8 workers and C=1..64. Ported to the Rust router, which is an independent implementation of the same cost function and had the identical defect. Adds the reproduction driver used to find and verify this, and regression tests for the serial-pacing split, the high-weight case, affinity retention, displacement of a saturated incumbent, and decay back to contention. Four of the five fail against the old policy. Signed-off-by: Zhang, Jiejing --- bench/_kv_aware_split_loadtest.py | 165 ++++++++++++++++++ infera/router/policy/kv_event_aware.py | 94 +++++++++- infera/server/args.py | 5 +- rust/router/src/policy.rs | 155 +++++++++++++++- .../router/test_kv_aware_routing_extremes.py | 112 ++++++++++++ 5 files changed, 514 insertions(+), 17 deletions(-) create mode 100644 bench/_kv_aware_split_loadtest.py diff --git a/bench/_kv_aware_split_loadtest.py b/bench/_kv_aware_split_loadtest.py new file mode 100644 index 00000000..1e839bad --- /dev/null +++ b/bench/_kv_aware_split_loadtest.py @@ -0,0 +1,165 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Reproduce the reported 448/0 kv-aware split against the real policy. + +Drives KvEventAwarePolicy with the traffic shape from the field report: +two symmetric aggregated workers, multi-turn agent sessions, per-session +causal pacing (a turn finishes before the next turn of that session is +issued). Counts picks per worker, the way engine-side +prefix_cache_queries_total deltas would. + +Only the KvEventClient and the tokenizer are stubbed: the client is +replaced by a model of what the engines' caches would actually hold +(blocks land on the worker that served the request), and the hasher +returns token ids directly so no model download is needed. The cost +function, the refcount bookkeeping and the lifecycle hooks are the real +ones. +""" + +from __future__ import annotations + +import argparse +import random +import sys +from collections import Counter + +from infera.common.worker_pool import DisaggMode, EngineType, WorkerInfo, WorkerStatus +from infera.router.kv_event.hasher import hash_request +from infera.router.policy.kv_event_aware import KvEventAwarePolicy + +BLOCK_SIZE = 16 + + +class SimulatedFleet: + """Stands in for KvEventClient: each worker's view holds the blocks of + the requests that worker actually served (unbounded, i.e. best case for + locality -- a real engine evicts, which can only spread load further).""" + + def __init__(self, worker_ids: list[str]) -> None: + self._views: dict[str, set[int]] = {w: set() for w in worker_ids} + + def cache_view(self, worker_id: str, dp_rank: int | None = None) -> set[int]: + return self._views[worker_id] + + def store(self, worker_id: str, blocks: list[int]) -> None: + self._views[worker_id].update(blocks) + + def on_worker_added(self, w: WorkerInfo) -> None: + pass + + def on_worker_removed(self, worker_id: str) -> None: + pass + + async def aclose(self) -> None: + pass + + +class TokenHasher: + """Hashes the pre-tokenized ids the trace generator produces, using the + router's own chaining -- skips the tokenizer, keeps the hash chain.""" + + def hash_for(self, body: dict, *, block_size: int, engine=None) -> list[int]: + return hash_request(body["_tokens"], block_size) + + +def _worker(worker_id: str) -> WorkerInfo: + return WorkerInfo( + worker_id=worker_id, + url=f"http://{worker_id}", + model_name="moonshotai/Kimi-K3", + engine=EngineType.VLLM, + status=WorkerStatus.ACTIVE, + disagg_mode=DisaggMode.MIXED, + kv_events_endpoint=f"tcp://{worker_id}:5557", + kv_block_size=BLOCK_SIZE, + ) + + +def make_trace(n_requests: int, seed: int = 0) -> list[dict]: + """Multi-turn agent sessions: each turn re-sends the conversation so far + plus new content, so turn N shares a long prefix with turn N-1. + + Shaped after the Mooncake toolagent trace: a shared system prompt, then + per-session divergence, then per-turn growth. + """ + rng = random.Random(seed) + system = [rng.randrange(1000, 2000) for _ in range(3 * BLOCK_SIZE)] + + reqs: list[dict] = [] + session = 0 + while len(reqs) < n_requests: + session += 1 + tokens = list(system) + [rng.randrange(10_000, 90_000) for _ in range(4 * BLOCK_SIZE)] + for turn in range(rng.randint(3, 8)): + tokens = tokens + [rng.randrange(10_000, 90_000) for _ in range(3 * BLOCK_SIZE)] + reqs.append( + { + "model": "moonshotai/Kimi-K3", + "_tokens": list(tokens), + "_session": session, + "_turn": turn, + } + ) + if len(reqs) >= n_requests: + break + return reqs[:n_requests] + + +def run(policy: KvEventAwarePolicy, fleet: SimulatedFleet, workers, trace) -> Counter: + """Serial pacing: each request is picked, started and finished before the + next -- the low-concurrency case the report describes.""" + picks: Counter = Counter() + for req in trace: + target, blocks = policy.pick(workers, req) + picks[target.route_key] += 1 + policy.on_request_started(target.route_key, blocks) + # The engine serves it and caches the blocks. + fleet.store(target.worker.worker_id, blocks) + policy.on_request_finished(target.route_key, blocks) + return picks + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--requests", type=int, default=448) + ap.add_argument( + "--kv-overlap-weight", + type=float, + nargs="+", + default=[0.01, 0.1, 1.0, 20.0], + ) + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + + workers = [_worker("worker-A"), _worker("worker-B")] + trace = make_trace(args.requests, seed=args.seed) + + print(f"trace: {len(trace)} requests, {len({r['_session'] for r in trace})} sessions") + print(f" block_size={BLOCK_SIZE}, serial pacing, 2 symmetric workers\n") + print(f"{'overlap_weight':>15} | {'worker-A':>9} | {'worker-B':>9} | split") + print(f"{'-' * 15}-+-{'-' * 9}-+-{'-' * 9}-+------") + + worst = 0.0 + for w in args.kv_overlap_weight: + fleet = SimulatedFleet([x.worker_id for x in workers]) + policy = KvEventAwarePolicy(fleet, TokenHasher(), overlap_weight=w) + picks = run(policy, fleet, workers, trace) + a, b = picks.get("worker-A", 0), picks.get("worker-B", 0) + share = max(a, b) / max(1, a + b) + worst = max(worst, share) + print(f"{w:>15g} | {a:>9} | {b:>9} | {share:>5.1%}") + + print() + if worst > 0.99: + print("REPRODUCED: at least one weight pins ~all traffic on one worker.") + return 1 + print("not reproduced at these weights.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/infera/router/policy/kv_event_aware.py b/infera/router/policy/kv_event_aware.py index e5f92ebf..bd1c5ab3 100644 --- a/infera/router/policy/kv_event_aware.py +++ b/infera/router/policy/kv_event_aware.py @@ -43,15 +43,56 @@ _MM_AFFINITY_CAP = 256 _MM_IMAGE_BLOCK_WEIGHT = 48.0 +# Per-pick decay on each worker's recent-dispatch total (half-life ~23 picks). +# +# The load half of the cost function counts blocks that are *in flight*, which +# is 0 for every worker whenever a request finishes before the next one is +# picked. Session-paced agent traffic has exactly that shape, so on that +# workload the cost function loses its load term entirely: a cold fleet ties, +# the tie goes to whichever candidate is enumerated first, and the cache that +# winner picks up re-elects it on every later request. The result is a +# permanent 100/0 split across symmetric workers at any overlap weight. +# +# `recent` keeps the load signal alive across requests that never overlap in +# time. A pick is charged the blocks the winner MISSED, not the blocks the +# request contains: prefill work is proportional to what the worker has to +# compute, and a block already in its cache costs it nothing. Charging misses +# rather than totals is what lets the term coexist with cache affinity -- +# a worker serving a fully-cached prompt accrues no load, so it keeps winning +# that prompt, while a worker handed a cold prompt accrues the whole thing and +# the next cold prompt goes elsewhere. +# +# Misses are in blocks, the same unit as in-flight load, so the two sum without +# a conversion factor and the term scales with request size. Charging one point +# per request instead would cap the term at 1/(1 - decay) no matter how large +# the requests were, which a single block of cache edge outvotes outright at +# any overlap weight above that cap -- leaving the split unfixed on exactly the +# prefill-weighted pools (typical production weight: 20.0) that most need +# spreading. +# +# The decay sets how long a worker carries what it was handed. Too fast and the +# term cannot accumulate enough to displace an incumbent between one session +# and the next; measured against the field trace, 0.9 still left w=20 fully +# pinned while anything from 0.95 up balanced it. +_RECENT_DECAY = 0.97 + class KvEventAwarePolicy(Policy): """Pick the worker minimising - cost(w) = overlap_weight * (request_blocks - hits(w)) + active_blocks(w) + cost(w) = overlap_weight * (request_blocks - hits(w)) + load(w) + load(w) = active_blocks(w) + recent_blocks(w) where ``active_blocks(w)`` is a refcounted set of distinct in-flight block hashes (deduped across requests sharing prefixes), not a sum of - prompt lengths. + prompt lengths, and ``recent_blocks(w)`` is a decayed sum of the block + counts recently dispatched to ``w``. + + Both halves are load-bearing. ``active_blocks`` alone is 0 for every + worker under traffic paced so that a request finishes before the next is + picked -- the normal shape of multi-turn agent sessions -- which drops the + load term out of the comparison and pins the whole fleet's traffic onto + one worker. See ``_RECENT_DECAY``. For PD-disaggregated routing, the disagg-bootstrap router passes ``role_hint="prefill"`` or ``role_hint="decode"`` to ``pick``. The @@ -87,8 +128,11 @@ def __init__( self._w_decode = ( decode_overlap_weight if decode_overlap_weight is not None else overlap_weight ) - # worker_id -> {block_hash -> refcount}; len() is the load term. + # worker_id -> {block_hash -> refcount}; len() is the in-flight load. self._active_block_refs: dict[str, dict[int, int]] = {} + # route_key -> decayed sum of recently dispatched block counts. Carries + # the load signal across requests that never overlap in time. + self._recent_blocks: dict[str, float] = {} # route_key -> ordered set of recent image keys (MRU last, bounded LRU). # Multimodal affinity: a request whose image a worker already holds # costs less there, co-locating repeat images onto the warm vision @@ -170,17 +214,21 @@ def pick( def active(t: RouteTarget) -> int: return len(self._active_block_refs.get(t.route_key, {})) + def load(t: RouteTarget) -> float: + """Blocks in flight now, plus blocks dispatched recently.""" + return active(t) + self._recent_blocks.get(t.route_key, 0.0) + def cost(t: RouteTarget) -> float: total = len(hashes_for.get((t.worker.engine, t.worker.kv_block_size), [])) hits = self._cache_hits(t, hashes_for) # Image miss term: images this worker does NOT already hold cost # w_mm each; the worker with the warm vision cache pays 0 → wins. mm_miss = len(mm_keys) - self._mm_hits(t.route_key, mm_keys) - return w_overlap * (total - hits) + w_mm * mm_miss + active(t) + return w_overlap * (total - hits) + w_mm * mm_miss + load(t) - # Tie-break by lower active so equal-cost candidates fall back to + # Tie-break by lower load so equal-cost candidates fall back to # least-loaded. - picked = min(targets, key=lambda t: (cost(t), active(t))) + picked = min(targets, key=lambda t: (cost(t), load(t))) picked_blocks = list( hashes_for.get((picked.worker.engine, picked.worker.kv_block_size), []) ) @@ -189,9 +237,15 @@ def cost(t: RouteTarget) -> float: mm_affinity_hits = self._mm_hits(picked.route_key, mm_keys) self._record_mm(picked.route_key, mm_keys) + cache_hits = self._cache_hits(picked, hashes_for) + # Charge the winner for the blocks it will have to compute. Done here + # rather than in on_request_started because the hooks run on the + # dispatch path, which skips them on the failure routes -- and a pick + # that goes uncharged is invisible to the next one. + self._record_dispatch(picked.route_key, len(picked_blocks) - cache_hits) + # Telemetry: record the pick decision per role + target, plus # the retention bucket we observed on this request. - cache_hits = self._cache_hits(picked, hashes_for) metrics.record_pick( role=role_hint or "mixed", worker_id=picked.route_key, @@ -243,6 +297,32 @@ def on_worker_removed(self, worker_id: str) -> None: pass for key in [k for k in self._mm_affinity if k == worker_id or k.startswith(prefix)]: self._mm_affinity.pop(key, None) + # Iterated separately from _active_block_refs: a worker can carry a + # recent total with nothing in flight, which is the state this term + # exists to represent. + for key in [k for k in self._recent_blocks if k == worker_id or k.startswith(prefix)]: + self._recent_blocks.pop(key, None) + + def _record_dispatch(self, route_key: str, missed_blocks: int) -> None: + """Decay every worker's recent total, then charge the pick's misses. + + Decaying on each pick rather than on a wall-clock timer keeps routing a + pure function of the request sequence: same requests in, same decisions + out, which is what makes the behaviour testable. Totals that decay to + nothing are dropped so an idle worker returns to a clean 0. + + A fully-cached pick charges 0 -- the worker has no prefill to do, so it + takes on no load and stays the right answer for that prompt. + """ + for key in list(self._recent_blocks): + decayed = self._recent_blocks[key] * _RECENT_DECAY + if decayed < 1e-3: + del self._recent_blocks[key] + else: + self._recent_blocks[key] = decayed + if missed_blocks <= 0: + return + self._recent_blocks[route_key] = self._recent_blocks.get(route_key, 0.0) + missed_blocks def _publish_active(self, route_key: str) -> None: """Mirror ``route_key``'s in-flight block count into the gauge. diff --git a/infera/server/args.py b/infera/server/args.py index 59b7266a..4551d6cb 100644 --- a/infera/server/args.py +++ b/infera/server/args.py @@ -148,8 +148,9 @@ def parse_server_args(argv: list[str] | None = None) -> argparse.Namespace: type=float, default=1.0, help="kv-aware only: weight on the cache-locality term in " - "cost = w * (request_blocks - hits) + active_blocks. Larger values " - "favour cache reuse over load balance (default: 1.0). Used for " + "cost = w * (request_blocks - hits) + load, where load counts blocks " + "in flight plus a decayed total of blocks recently dispatched. Larger " + "values favour cache reuse over load balance (default: 1.0). Used for " "mixed-pool routing and as the fallback when the prefill/decode " "weights below are unset.", ) diff --git a/rust/router/src/policy.rs b/rust/router/src/policy.rs index ce698c28..043772ab 100644 --- a/rust/router/src/policy.rs +++ b/rust/router/src/policy.rs @@ -137,18 +137,52 @@ const MM_AFFINITY_CAP: usize = 256; /// load differences without a separate weight knob. const MM_IMAGE_BLOCK_WEIGHT: f64 = 48.0; +/// Per-pick decay on each worker's recent-dispatch total (half-life ~23 picks). +/// +/// The load half of the cost function counts blocks that are *in flight*, which +/// is 0 for every worker whenever a request finishes before the next one is +/// picked. Session-paced agent traffic has exactly that shape, so on that +/// workload the cost function loses its load term entirely: a cold fleet ties, +/// the tie goes to whichever candidate is enumerated first, and the cache that +/// winner picks up re-elects it on every later request. The result is a +/// permanent 100/0 split across symmetric workers at any overlap weight. +/// +/// `recent` keeps the load signal alive across requests that never overlap in +/// time. A pick is charged the blocks the winner MISSED, not the blocks the +/// request contains: prefill work is proportional to what the worker has to +/// compute, and a block already in its cache costs it nothing. That is what +/// lets the term coexist with cache affinity -- a worker serving a fully-cached +/// prompt accrues no load and keeps winning it, while a worker handed a cold +/// prompt accrues the whole thing and the next cold prompt goes elsewhere. +/// +/// Misses are in blocks, the same unit as in-flight load, so the two sum +/// without a conversion factor and the term scales with request size. Charging +/// one point per request instead would cap the term at `1/(1 - decay)` no +/// matter how large the requests were, which a single block of cache edge +/// outvotes outright at any overlap weight above that cap. +/// +/// Mirrors `_RECENT_DECAY` in infera/router/policy/kv_event_aware.py; the two +/// routers are independent implementations of the same policy and must agree. +const RECENT_DECAY: f64 = 0.97; + /// Pick the worker minimising -/// `cost(w) = w_overlap * (request_blocks - hits(w)) + active_blocks(w)` -/// where `hits(w)` is the longest cached prefix on that worker's DP rank and -/// `active_blocks(w)` is the refcounted set of distinct in-flight block hashes. +/// `cost(w) = w_overlap * (request_blocks - hits(w)) + load(w)` +/// `load(w) = active_blocks(w) + recent_blocks(w)` +/// where `hits(w)` is the longest cached prefix on that worker's DP rank, +/// `active_blocks(w)` is the refcounted set of distinct in-flight block hashes, +/// and `recent_blocks(w)` is a decayed sum of the blocks recently dispatched to +/// it. Both halves of the load term are needed -- see [`RECENT_DECAY`]. pub struct KvEventAwarePolicy { kv: Arc, hasher: BlockHasher, w: f64, w_prefill: f64, w_decode: f64, - // route_key -> {block_hash -> refcount}; len() is the load term. + // route_key -> {block_hash -> refcount}; len() is the in-flight load. active: Mutex>>, + // route_key -> decayed sum of recently dispatched block counts. Carries the + // load signal across requests that never overlap in time. + recent: Mutex>, // route_key -> recent image keys (MRU front, bounded LRU). Multimodal // affinity: a request whose image a worker already holds costs less there, // co-locating repeat images onto the worker with the warm vision cache. @@ -170,6 +204,7 @@ impl KvEventAwarePolicy { w_prefill: prefill_overlap_weight.unwrap_or(overlap_weight), w_decode: decode_overlap_weight.unwrap_or(overlap_weight), active: Mutex::new(HashMap::new()), + recent: Mutex::new(HashMap::new()), mm_affinity: Mutex::new(HashMap::new()), } } @@ -201,6 +236,37 @@ impl KvEventAwarePolicy { .unwrap_or(0) } + /// Blocks in flight now, plus blocks dispatched recently. + fn load_of(&self, route_key: &str) -> f64 { + let recent = self + .recent + .lock() + .expect("recent mutex poisoned") + .get(route_key) + .copied() + .unwrap_or(0.0); + self.active_len(route_key) as f64 + recent + } + + /// Decay every worker's recent total, then charge the pick's misses. + /// + /// Decaying on each pick rather than on a wall-clock timer keeps routing a + /// pure function of the request sequence: same requests in, same decisions + /// out. Totals that decay to nothing are dropped so an idle worker returns + /// to a clean 0. A fully-cached pick charges nothing -- the worker has no + /// prefill to do, so it takes on no load and stays the right answer. + fn record_dispatch(&self, route_key: &str, missed_blocks: usize) { + let mut recent = self.recent.lock().expect("recent mutex poisoned"); + recent.retain(|_, v| { + *v *= RECENT_DECAY; + *v >= 1e-3 + }); + if missed_blocks == 0 { + return; + } + *recent.entry(route_key.to_string()).or_insert(0.0) += missed_blocks as f64; + } + /// How many of `keys` this worker is recorded as holding (its warm images). fn mm_hits(&self, route_key: &str, keys: &[u64]) -> usize { if keys.is_empty() { @@ -286,10 +352,10 @@ impl Policy for KvEventAwarePolicy { .saturating_sub(self.mm_hits(&route_key, &mm_keys)); w_overlap * (total.saturating_sub(hits) as f64) + w_mm * (mm_miss as f64) - + self.active_len(&route_key) as f64 + + self.load_of(&route_key) }; - // min by (cost, active) — tie-break to least-loaded. + // min by (cost, load) — tie-break to least-loaded. let picked = targets .iter() .min_by(|a, b| { @@ -297,8 +363,9 @@ impl Policy for KvEventAwarePolicy { ca.partial_cmp(&cb) .unwrap_or(std::cmp::Ordering::Equal) .then_with(|| { - self.active_len(&a.route_key()) - .cmp(&self.active_len(&b.route_key())) + self.load_of(&a.route_key()) + .partial_cmp(&self.load_of(&b.route_key())) + .unwrap_or(std::cmp::Ordering::Equal) }) }) .expect("candidates non-empty") @@ -307,6 +374,11 @@ impl Policy for KvEventAwarePolicy { let blocks = blocks_of(&picked).clone(); let hits = hits_of(&picked); let picked_key = picked.route_key(); + // Charge the winner for the blocks it will have to compute. Done here + // rather than in on_request_started because the hooks run on the + // dispatch path, which skips them on the failure routes -- and a pick + // that goes uncharged is invisible to the next one. + self.record_dispatch(&picked_key, blocks.len().saturating_sub(hits)); // Mark the chosen worker as now holding this request's images, so the // next request for the same image is drawn back to its warm cache. let mm_matched = self.mm_hits(&picked_key, &mm_keys); @@ -381,6 +453,10 @@ impl Policy for KvEventAwarePolicy { .lock() .expect("active mutex poisoned") .retain(|rk, _| alive(rk)); + self.recent + .lock() + .expect("recent mutex poisoned") + .retain(|rk, _| alive(rk)); self.mm_affinity .lock() .expect("mm_affinity mutex poisoned") @@ -434,6 +510,69 @@ mod tests { ); } + #[test] + fn recent_dispatch_keeps_a_load_signal_between_requests() { + // The bug this guards: with in-flight blocks as the only load term, + // traffic paced so each request finishes before the next is picked + // leaves every worker reading 0, and the first pick then wins every + // subsequent one on candidate order. Nothing is in flight here. + let kv = Arc::new(KvEventClient::new()); + let pol = KvEventAwarePolicy::new(kv, BlockHasher::disabled(), 1.0, None, None); + assert_eq!(pol.active_len("a"), 0); + assert_eq!(pol.load_of("a"), 0.0); + + pol.record_dispatch("a", 10); + assert_eq!(pol.active_len("a"), 0, "nothing in flight"); + assert!( + pol.load_of("a") > 0.0, + "load signal must outlive the request" + ); + assert!(pol.load_of("a") > pol.load_of("b")); + } + + #[test] + fn a_fully_cached_pick_is_charged_nothing() { + // The charge is the blocks the winner had to COMPUTE. A worker serving + // a prompt it already holds takes on no prefill work, so it accrues no + // load and stays the right answer for that prompt. + let kv = Arc::new(KvEventClient::new()); + let pol = KvEventAwarePolicy::new(kv, BlockHasher::disabled(), 1.0, None, None); + pol.record_dispatch("a", 0); + assert_eq!(pol.load_of("a"), 0.0); + } + + #[test] + fn recent_load_decays_back_to_zero_when_a_worker_goes_idle() { + // Transient by construction: a worker that took a burst and then went + // quiet must return to contention, or the fix trades one starvation + // mode for another. + let kv = Arc::new(KvEventClient::new()); + let pol = KvEventAwarePolicy::new(kv, BlockHasher::disabled(), 1.0, None, None); + pol.record_dispatch("bursty", 10); + assert!(pol.load_of("bursty") > 0.0); + for _ in 0..500 { + pol.record_dispatch("other", 1); + } + assert_eq!( + pol.load_of("bursty"), + 0.0, + "idle worker never returned to 0" + ); + } + + #[test] + fn sync_prunes_removed_worker_recent_load() { + // Pruned separately from `active`: a worker can carry a recent total + // with nothing in flight, which is the state this term represents. + let kv = Arc::new(KvEventClient::new()); + let pol = KvEventAwarePolicy::new(kv, BlockHasher::disabled(), 1.0, None, None); + pol.record_dispatch("gone#dp0", 5); + pol.record_dispatch("stay", 5); + pol.sync_workers(&[worker("stay", 16, None)]); + assert_eq!(pol.load_of("gone#dp0"), 0.0); + assert!(pol.load_of("stay") > 0.0); + } + #[test] fn refcount_started_finished_balances() { let kv = Arc::new(KvEventClient::new()); diff --git a/tests/unit/router/test_kv_aware_routing_extremes.py b/tests/unit/router/test_kv_aware_routing_extremes.py index 9c951ec7..259565d8 100644 --- a/tests/unit/router/test_kv_aware_routing_extremes.py +++ b/tests/unit/router/test_kv_aware_routing_extremes.py @@ -222,3 +222,115 @@ async def test_finishing_a_request_returns_the_worker_to_contention(rig): assert _pick(policy, (a, b), prompt) == "a:1", ( "A holds a block and is idle again; it must win once its load is released" ) + + +# --- load must survive between requests, not just during them ---------------- + + +async def test_serial_traffic_does_not_pin_every_request_to_one_worker(rig): + """The reported production failure, reduced to its smallest form. + + Each request finishes before the next is picked, so nothing is ever in + flight when a decision is made. If in-flight blocks are the only load + signal, every worker reads 0, the first cold pick wins on candidate order, + and the cache it gains re-elects it for the rest of the run -- 100/0 on a + symmetric pair, at any overlap weight. + + Distinct prefixes per request, so there is no locality reason to prefer + either worker. + """ + policy, _client, a, b = rig + counts = {"a:1": 0, "b:1": 0} + for i in range(40): + prompt = list(range(i * 100, i * 100 + 40)) + target, blocks = policy.pick([a, b], {"model": "m", "token_ids": prompt}) + counts[target.worker.worker_id] += 1 + # Started AND finished before the next pick: the paced-traffic shape. + policy.on_request_started(target.route_key, blocks) + policy.on_request_finished(target.route_key, blocks) + + assert min(counts.values()) > 0, f"one worker got everything: {counts}" + assert abs(counts["a:1"] - counts["b:1"]) <= 4, f"serial traffic did not spread: {counts}" + + +async def test_serial_traffic_spreads_even_at_a_high_overlap_weight(rig): + """A high --kv-prefill-overlap-weight must not reintroduce the pin. + + The load term is unweighted, so a fix that keeps it bounded independently + of request size loses to one block of cache edge once the weight exceeds + that bound. 20.0 is the documented production prefill weight. + """ + _policy, client, a, b = rig + policy = KvEventAwarePolicy(client, _IdentityHasher(), overlap_weight=20.0) + counts = {"a:1": 0, "b:1": 0} + for i in range(40): + prompt = list(range(i * 100, i * 100 + 40)) + target, blocks = policy.pick([a, b], {"model": "m", "token_ids": prompt}) + counts[target.worker.worker_id] += 1 + policy.on_request_started(target.route_key, blocks) + policy.on_request_finished(target.route_key, blocks) + + assert min(counts.values()) > 0, f"high weight pinned the fleet: {counts}" + + +async def test_a_fully_cached_prompt_stays_on_its_holder_under_serial_traffic(rig): + """Balance must not cost affinity. + + The recent-dispatch charge is the blocks the winner had to COMPUTE, so a + worker serving a prompt it already holds accrues nothing and keeps winning + it. Were the charge the request's total size instead, repeat requests would + ping-pong and every one after the first would miss. + """ + policy, client, a, b = rig + prompt = list(range(40)) + _store(client, "a:1", prompt) + + for _ in range(20): + target, blocks = policy.pick([a, b], {"model": "m", "token_ids": prompt}) + assert target.worker.worker_id == "a:1" + policy.on_request_started(target.route_key, blocks) + policy.on_request_finished(target.route_key, blocks) + + +async def test_a_saturated_incumbent_yields_new_work_to_an_idle_worker(rig): + """Recent load has to be able to outweigh a shared-prefix edge. + + A worker that has served a long run of traffic carries a large recent + total; a cold worker carries none. A new prompt sharing only a short prefix + with the incumbent must go to the idle worker. + """ + policy, client, a, b = rig + shared = list(range(8)) # 2 blocks both would hit on + for i in range(20): + prompt = shared + list(range(1000 + i * 100, 1000 + i * 100 + 32)) + target, blocks = policy.pick([a], {"model": "m", "token_ids": prompt}) + policy.on_request_started(target.route_key, blocks) + # first_hash stays a single byte: _store packs it with bytes([...]). + _store(client, "a:1", prompt, first_hash=i * 10) + policy.on_request_finished(target.route_key, blocks) + + fresh = shared + list(range(9000, 9032)) + assert _pick(policy, (a, b), fresh) == "b:1", ( + "a saturated worker kept new work on the strength of a 2-block prefix" + ) + + +async def test_recent_load_decays_so_an_idle_worker_returns_to_contention(rig): + """The charge is transient. A worker that took a burst and then went quiet + must not be written off permanently -- otherwise the fix trades one + starvation mode for another.""" + policy, _client, a, b = rig + for i in range(20): + prompt = list(range(i * 100, i * 100 + 40)) + target, blocks = policy.pick([a], {"model": "m", "token_ids": prompt}) + policy.on_request_started(target.route_key, blocks) + policy.on_request_finished(target.route_key, blocks) + + assert policy._recent_blocks.get("a:1", 0.0) > 0.0 + for i in range(400): + prompt = list(range(50_000 + i * 100, 50_000 + i * 100 + 40)) + target, blocks = policy.pick([b], {"model": "m", "token_ids": prompt}) + policy.on_request_started(target.route_key, blocks) + policy.on_request_finished(target.route_key, blocks) + + assert "a:1" not in policy._recent_blocks, "an idle worker never returned to 0" From 0db300cbfda8b73846e9907b07e9d2475f91db22 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Sat, 8 Aug 2026 19:51:42 +0000 Subject: [PATCH 69/88] fix(router): charge for picks the hasher produced no blocks for A user retesting kv-aware reported that --kv-overlap-weight 1.0 and 0.01 gave identical placement with no cache-aware routing. Reproduced against a live two-worker Kimi-K3 fleet: 12 short-prompt requests, 12 to the same worker. The prompts hashed to zero blocks. The router's index block size is 768 tokens by default, so anything shorter produces no blocks at all -- the state a smoke test naturally lands in. Two consequences follow, and only the first was understood. Weight cancellation was expected: cost is `overlap_weight * (request_blocks - hits) + load`, so with request_blocks = 0 the weighted term is 0 for every candidate and every weight agrees. That alone would leave the load term deciding, which is correct. The load term also collapsed, which was not expected and is a real defect in 779baee2. `_record_dispatch` charges the blocks the winner MISSED, so a pick with no blocks charges nothing -- `recent_blocks` stays 0 on every worker, `active_blocks` is already 0 under paced traffic, and the cost is a perfect tie. min() returns the first candidate, forever. The 100/0 split that commit fixed comes straight back on any workload whose prompts are short. The zero-miss test conflated two different situations. A fully cached prompt (request_blocks > 0, all hit) genuinely costs the worker nothing and must charge nothing, or repeat requests ping-pong and affinity is lost. A prompt with no block information at all still costs the worker something -- we simply cannot say how much. Passing request_blocks alongside the miss count separates them: the first charges 0, the second charges a one-block floor, the smallest honest estimate that breaks the tie without outweighing a real multi-block request. Measured over 200 requests, two symmetric workers, serial pacing: prompts hashing to zero blocks 200/0 -> 100/100 normal prompts 106/94 -> 106/94 (unchanged) Ported to the Rust router, which had the identical defect. Two regression tests. The zero-block one fails against the previous code; the fully-cached one guards the distinction from being collapsed again. Signed-off-by: Zhang, Jiejing --- infera/router/policy/kv_event_aware.py | 38 +++++++++++--- rust/router/src/policy.rs | 51 +++++++++++++++---- .../router/test_kv_aware_routing_extremes.py | 41 +++++++++++++++ 3 files changed, 113 insertions(+), 17 deletions(-) diff --git a/infera/router/policy/kv_event_aware.py b/infera/router/policy/kv_event_aware.py index bd1c5ab3..16b4130d 100644 --- a/infera/router/policy/kv_event_aware.py +++ b/infera/router/policy/kv_event_aware.py @@ -76,6 +76,16 @@ # pinned while anything from 0.95 up balanced it. _RECENT_DECAY = 0.97 +# Load charged for a pick we have no block information about: the hasher +# returned nothing, because the prompt is shorter than the index block size +# (768 tokens by default) or tokenisation failed. Such a request still costs +# the worker something, and charging 0 would hold the load term at 0 for every +# candidate -- restoring the permanent tie, and with it the 100/0 split, on any +# workload whose prompts are short. One block is the smallest honest estimate: +# enough to break the tie and rotate workers, small enough that a real +# multi-block request still outweighs it. +_UNKNOWN_COST_BLOCKS = 1.0 + class KvEventAwarePolicy(Policy): """Pick the worker minimising @@ -242,7 +252,7 @@ def cost(t: RouteTarget) -> float: # rather than in on_request_started because the hooks run on the # dispatch path, which skips them on the failure routes -- and a pick # that goes uncharged is invisible to the next one. - self._record_dispatch(picked.route_key, len(picked_blocks) - cache_hits) + self._record_dispatch(picked.route_key, len(picked_blocks) - cache_hits, len(picked_blocks)) # Telemetry: record the pick decision per role + target, plus # the retention bucket we observed on this request. @@ -303,16 +313,27 @@ def on_worker_removed(self, worker_id: str) -> None: for key in [k for k in self._recent_blocks if k == worker_id or k.startswith(prefix)]: self._recent_blocks.pop(key, None) - def _record_dispatch(self, route_key: str, missed_blocks: int) -> None: - """Decay every worker's recent total, then charge the pick's misses. + def _record_dispatch(self, route_key: str, missed_blocks: int, request_blocks: int) -> None: + """Decay every worker's recent total, then charge for the pick. Decaying on each pick rather than on a wall-clock timer keeps routing a pure function of the request sequence: same requests in, same decisions out, which is what makes the behaviour testable. Totals that decay to nothing are dropped so an idle worker returns to a clean 0. - A fully-cached pick charges 0 -- the worker has no prefill to do, so it - takes on no load and stays the right answer for that prompt. + Two different situations both arrive here with zero misses, and they + must not be charged the same way: + + - ``request_blocks > 0`` and every one of them hit. The worker has no + prefill to do, so it takes on no load and stays the right answer for + that prompt. Charge nothing. + - ``request_blocks == 0``. The hasher produced no blocks at all -- the + prompt is shorter than the index block size, or tokenisation failed. + We know nothing about this request's cost, but the worker still has + to serve it. Charging nothing here leaves the load term at 0 for + every candidate forever, which is the exact tie that sends every + request to whichever worker sorts first: the 100/0 split this term + exists to prevent, reappearing on short prompts. Charge the floor. """ for key in list(self._recent_blocks): decayed = self._recent_blocks[key] * _RECENT_DECAY @@ -320,9 +341,12 @@ def _record_dispatch(self, route_key: str, missed_blocks: int) -> None: del self._recent_blocks[key] else: self._recent_blocks[key] = decayed - if missed_blocks <= 0: + charge = float(missed_blocks) + if request_blocks <= 0: + charge = _UNKNOWN_COST_BLOCKS + if charge <= 0: return - self._recent_blocks[route_key] = self._recent_blocks.get(route_key, 0.0) + missed_blocks + self._recent_blocks[route_key] = self._recent_blocks.get(route_key, 0.0) + charge def _publish_active(self, route_key: str) -> None: """Mirror ``route_key``'s in-flight block count into the gauge. diff --git a/rust/router/src/policy.rs b/rust/router/src/policy.rs index 043772ab..7708b343 100644 --- a/rust/router/src/policy.rs +++ b/rust/router/src/policy.rs @@ -165,6 +165,13 @@ const MM_IMAGE_BLOCK_WEIGHT: f64 = 48.0; /// routers are independent implementations of the same policy and must agree. const RECENT_DECAY: f64 = 0.97; +/// Load charged for a pick we have no block information about -- the hasher +/// returned nothing, because the prompt is shorter than the index block size or +/// tokenisation failed. Charging 0 would hold the load term at 0 for every +/// candidate and restore the permanent tie. Mirrors `_UNKNOWN_COST_BLOCKS` in +/// infera/router/policy/kv_event_aware.py. +const UNKNOWN_COST_BLOCKS: f64 = 1.0; + /// Pick the worker minimising /// `cost(w) = w_overlap * (request_blocks - hits(w)) + load(w)` /// `load(w) = active_blocks(w) + recent_blocks(w)` @@ -255,16 +262,29 @@ impl KvEventAwarePolicy { /// out. Totals that decay to nothing are dropped so an idle worker returns /// to a clean 0. A fully-cached pick charges nothing -- the worker has no /// prefill to do, so it takes on no load and stays the right answer. - fn record_dispatch(&self, route_key: &str, missed_blocks: usize) { + fn record_dispatch(&self, route_key: &str, missed_blocks: usize, request_blocks: usize) { let mut recent = self.recent.lock().expect("recent mutex poisoned"); recent.retain(|_, v| { *v *= RECENT_DECAY; *v >= 1e-3 }); - if missed_blocks == 0 { + // Zero misses arrives here two different ways, and they must not be + // charged alike. request_blocks > 0 with no misses means a fully cached + // prompt: no prefill to do, so no load, and the holder keeps winning it. + // request_blocks == 0 means the hasher produced nothing -- a prompt + // under the index block size, or failed tokenisation. That request + // still costs the worker something, and charging 0 would hold the load + // term at 0 on every candidate, restoring the permanent tie and the + // 100/0 split on short-prompt workloads. + let charge = if request_blocks == 0 { + UNKNOWN_COST_BLOCKS + } else { + missed_blocks as f64 + }; + if charge <= 0.0 { return; } - *recent.entry(route_key.to_string()).or_insert(0.0) += missed_blocks as f64; + *recent.entry(route_key.to_string()).or_insert(0.0) += charge; } /// How many of `keys` this worker is recorded as holding (its warm images). @@ -378,7 +398,7 @@ impl Policy for KvEventAwarePolicy { // rather than in on_request_started because the hooks run on the // dispatch path, which skips them on the failure routes -- and a pick // that goes uncharged is invisible to the next one. - self.record_dispatch(&picked_key, blocks.len().saturating_sub(hits)); + self.record_dispatch(&picked_key, blocks.len().saturating_sub(hits), blocks.len()); // Mark the chosen worker as now holding this request's images, so the // next request for the same image is drawn back to its warm cache. let mm_matched = self.mm_hits(&picked_key, &mm_keys); @@ -521,7 +541,7 @@ mod tests { assert_eq!(pol.active_len("a"), 0); assert_eq!(pol.load_of("a"), 0.0); - pol.record_dispatch("a", 10); + pol.record_dispatch("a", 10, 10); assert_eq!(pol.active_len("a"), 0, "nothing in flight"); assert!( pol.load_of("a") > 0.0, @@ -530,6 +550,17 @@ mod tests { assert!(pol.load_of("a") > pol.load_of("b")); } + #[test] + fn a_pick_with_no_block_information_is_still_charged() { + // A prompt under the index block size hashes to nothing. Charging 0 for + // it would hold the load term at 0 on every candidate, restoring the + // tie that sends every request to whichever worker sorts first. + let kv = Arc::new(KvEventClient::new()); + let pol = KvEventAwarePolicy::new(kv, BlockHasher::disabled(), 1.0, None, None); + pol.record_dispatch("a", 0, 0); + assert!(pol.load_of("a") > 0.0, "unhashable pick accrued no load"); + } + #[test] fn a_fully_cached_pick_is_charged_nothing() { // The charge is the blocks the winner had to COMPUTE. A worker serving @@ -537,7 +568,7 @@ mod tests { // load and stays the right answer for that prompt. let kv = Arc::new(KvEventClient::new()); let pol = KvEventAwarePolicy::new(kv, BlockHasher::disabled(), 1.0, None, None); - pol.record_dispatch("a", 0); + pol.record_dispatch("a", 0, 10); assert_eq!(pol.load_of("a"), 0.0); } @@ -548,10 +579,10 @@ mod tests { // mode for another. let kv = Arc::new(KvEventClient::new()); let pol = KvEventAwarePolicy::new(kv, BlockHasher::disabled(), 1.0, None, None); - pol.record_dispatch("bursty", 10); + pol.record_dispatch("bursty", 10, 10); assert!(pol.load_of("bursty") > 0.0); for _ in 0..500 { - pol.record_dispatch("other", 1); + pol.record_dispatch("other", 1, 1); } assert_eq!( pol.load_of("bursty"), @@ -566,8 +597,8 @@ mod tests { // with nothing in flight, which is the state this term represents. let kv = Arc::new(KvEventClient::new()); let pol = KvEventAwarePolicy::new(kv, BlockHasher::disabled(), 1.0, None, None); - pol.record_dispatch("gone#dp0", 5); - pol.record_dispatch("stay", 5); + pol.record_dispatch("gone#dp0", 5, 5); + pol.record_dispatch("stay", 5, 5); pol.sync_workers(&[worker("stay", 16, None)]); assert_eq!(pol.load_of("gone#dp0"), 0.0); assert!(pol.load_of("stay") > 0.0); diff --git a/tests/unit/router/test_kv_aware_routing_extremes.py b/tests/unit/router/test_kv_aware_routing_extremes.py index 259565d8..25bac06b 100644 --- a/tests/unit/router/test_kv_aware_routing_extremes.py +++ b/tests/unit/router/test_kv_aware_routing_extremes.py @@ -334,3 +334,44 @@ async def test_recent_load_decays_so_an_idle_worker_returns_to_contention(rig): policy.on_request_finished(target.route_key, blocks) assert "a:1" not in policy._recent_blocks, "an idle worker never returned to 0" + + +async def test_prompts_too_short_to_hash_still_spread(rig): + """A prompt under the index block size hashes to nothing. The load charge + must not be zero for those, or the term stays 0 on every candidate and the + tie sends the whole run to whichever worker sorts first -- the 100/0 split + reappearing on short prompts, which is what a user hit after the load-term + fix landed. + + token_ids shorter than BS produce no blocks at all, which is the same state + the real hasher reaches on a sub-768-token prompt. + """ + policy, _client, a, b = rig + counts = {"a:1": 0, "b:1": 0} + for i in range(30): + # Two tokens: fewer than BS, so hash_request returns []. + target, blocks = policy.pick([a, b], {"model": "m", "token_ids": [i, i + 1]}) + assert blocks == [], "precondition: this prompt must hash to zero blocks" + counts[target.worker.worker_id] += 1 + policy.on_request_started(target.route_key, blocks) + policy.on_request_finished(target.route_key, blocks) + + assert min(counts.values()) > 0, f"unhashable prompts pinned one worker: {counts}" + assert abs(counts["a:1"] - counts["b:1"]) <= 2, f"did not spread: {counts}" + + +async def test_a_fully_cached_pick_is_still_charged_nothing(rig): + """The zero-block floor must not leak into the fully-cached case: a worker + serving a prompt it already holds does no prefill, so it accrues no load and + keeps winning that prompt. Distinguishing the two is the whole point of + passing request_blocks alongside the miss count.""" + policy, client, a, b = rig + prompt = list(range(40)) + _store(client, "a:1", prompt) + + for _ in range(15): + target, blocks = policy.pick([a, b], {"model": "m", "token_ids": prompt}) + assert target.worker.worker_id == "a:1" + policy.on_request_started(target.route_key, blocks) + policy.on_request_finished(target.route_key, blocks) + assert policy._recent_blocks.get("a:1", 0.0) == 0.0, "fully-cached picks accrued load" From ae26aa33340c1c52e31b22b7811ccbf1970cdd2b Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Sat, 8 Aug 2026 13:04:47 +0000 Subject: [PATCH 70/88] fix(router): score each PD leg from its own response The wrapper added in 85b7859 scored both roles off the single response the router hands back, which is unsound: prefill and decode are different workers whose health is independent, and that one status code belongs to neither of them reliably. On the streaming path it belongs to nothing at all. A StreamingResponse is returned before its generator is touched -- the decode leg has not been POSTed and the prefill task has not been awaited -- and its 200 is the framework's default, set at construction. Streaming PD deliberately carries errors in the body as SSE, so the status code is not where the outcome lives even after the response has run. Recording success there marked both workers healthy before either was dispatched, and reset the count the legs were about to raise: a decode worker unreachable on all ten consecutive requests never tripped, where the code before that commit tripped on the third. The unary path had the same fault through the other half of the predicate: the code returned is decode's, so a prefill 500-ing every request was reset to healthy -- and an already-open breaker force-closed -- by the decode leg answering beside it. Each leg is now scored where its own answer arrives: both responses of the concurrent dispatch, the prefill and decode halves of the serial one, and the decode stream's headers, which is the pre-first-byte moment. The background prefill task is scored when it is awaited, since nothing about it is known until then, and a cancelled one is not scored at all -- the worker was never given the chance to answer, so there is no evidence to record. That also closes a gap that predates the wrapper: a leg's 5xx was logged and counted in metrics but never reached the breaker, which only ever saw transport errors. test_dual_stream_records_failure_on_unreachable_decode refused every request, prefill included, so it could not show the pools are scored independently -- prefill deserved to trip too. It now breaks only decode. Co-authored-by: Cursor Signed-off-by: leiwei12 --- infera/router/disagg.py | 84 +++++++++++++----------- tests/unit/router/test_disagg_breaker.py | 78 ++++++++++++++++++++++ 2 files changed, 124 insertions(+), 38 deletions(-) diff --git a/infera/router/disagg.py b/infera/router/disagg.py index ad66a5b0..9f1dd4d7 100644 --- a/infera/router/disagg.py +++ b/infera/router/disagg.py @@ -17,6 +17,7 @@ from infera.common.nats_request import TYPE_DATA, TYPE_DONE, TYPE_ERROR from infera.common.worker_pool import DisaggMode from infera.router.base import BaseRouter +from infera.router.breaker import is_worker_fault from infera.router.cache_control import parse_cache_hints from infera.router.disagg_protocols import ( ProtocolMismatch, @@ -136,7 +137,7 @@ async def dispatch( # skips prefill) vs D (load-heavy) differently. p_target, p_blocks = self.policy.pick(prefills, body, role_hint="prefill") d_target, d_blocks = self.policy.pick(decodes, body, role_hint="decode") - return await self._run_pd( + return await self._dispatch_pd( obs, p_target, d_target, p_blocks, d_blocks, body, stream, path ) @@ -167,46 +168,32 @@ async def dispatch_direct( }, status_code=503, ) - return await self._run_pd( + return await self._dispatch_pd( obs, RouteTarget(p), RouteTarget(d), [], [], body, stream, path ) - async def _run_pd( - self, - obs, - p_target: RouteTarget, - d_target: RouteTarget, - p_blocks: list[int], - d_blocks: list[int], - body: dict, - stream: bool, - path: str, - ) -> Response: - """Record the breaker outcome for both roles around the dual dispatch. - - Every exit of :meth:`_dispatch_pd` passes through here -- both the - policy-driven and gateway-driven entry points reach it -- which is the - only reason this is one place rather than a dozen. Failures are already - recorded inside, at the specific worker that caused them; what was - missing is the other half. Without it ``consecutive_failures`` never - resets, so "three failures in a row" quietly becomes "three failures - ever", and a worker that has probed successfully never closes. - - Anything other than a clean response is neutral rather than a success - for either role: one of the two may well be at fault, and this layer - cannot tell which. Neutral is idempotent against the failure already - recorded below, and still frees the probe slot. + def _score_leg(self, worker_id: str, status: int) -> None: + """Record one PD leg's HTTP outcome against the worker that produced it. + + The two legs are two different workers whose health is independent, so + each has to be scored from its own response. A decode that answers + cannot vouch for a prefill that did not: scoring both off one status + code let a prefill 500-ing every request be reset to healthy by the + decode leg beside it, and there is no status code at all for the + client-facing streaming response, whose 200 is a framework default set + before either leg has been dispatched. + + A 5xx is the worker's fault. A 4xx is the request's -- every worker + would answer the same, so it frees the probe slot without counting + either way. Anything below 400 is the evidence of health that resets + the consecutive-failure count. """ - resp = await self._dispatch_pd( - obs, p_target, d_target, p_blocks, d_blocks, body, stream, path - ) - ok = getattr(resp, "status_code", 200) < 400 - for worker_id in (p_target.worker.worker_id, d_target.worker.worker_id): - if ok: - self.breaker.record_success(worker_id) - else: - self.breaker.record_neutral(worker_id) - return resp + if is_worker_fault(status): + self.breaker.record_failure(worker_id) + elif status < 400: + self.breaker.record_success(worker_id) + else: + self.breaker.record_neutral(worker_id) async def _dispatch_pd( self, @@ -386,6 +373,8 @@ async def _post(url, leg, worker_id, leg_body, leg_headers): self.breaker.record_failure(p.worker_id) return _sanitized_error("PD request failed", exc, status_code=502) + self._score_leg(p.worker_id, p_resp.status_code) + self._score_leg(d.worker_id, d_resp.status_code) if p_resp.status_code >= 400: logger.warning( "prefill worker %s returned %d (decode may fail)", @@ -593,6 +582,7 @@ async def _dispatch_serial( self.breaker.record_failure(p.worker_id) return _sanitized_error("prefill leg failed", exc, status_code=502) + self._score_leg(p.worker_id, p_resp.status_code) if p_resp.status_code >= 400: p_failed = True obs["outcome"] = f"{p_resp.status_code // 100}xx" @@ -679,6 +669,7 @@ async def _dispatch_serial( self.breaker.record_failure(d.worker_id) return _sanitized_error("decode leg failed", exc, status_code=502) + self._score_leg(d.worker_id, d_resp.status_code) try: d_payload = d_resp.json() except ValueError: @@ -727,6 +718,10 @@ async def _stream_decode_only( yield f"data: {err}\n\n".encode() return + # Headers are in hand, which is the pre-first-byte moment: what the + # decode worker did with the request is now known and nothing has + # reached the client yet. + self._score_leg(d_target.worker.worker_id, d_resp.status_code) if d_resp.status_code >= 400: try: body_bytes = await d_resp.aread() @@ -883,6 +878,9 @@ async def _stream_dual( yield f"data: {err}\n\n".encode() return + # Headers are in hand: the decode leg's outcome is known and + # nothing has reached the client yet. + self._score_leg(d_target.worker.worker_id, d_resp.status_code) if d_resp.status_code >= 400: # Engine accepted but rejected; surface its body verbatim. try: @@ -955,6 +953,9 @@ async def _stream_dual( try: p_resp = await asyncio.shield(p_task) except asyncio.CancelledError: + # The request was torn down from above, so the prefill worker + # was never given the chance to answer. That is not evidence + # about it either way, and scoring it would be inventing one. logger.debug("prefill task cancelled (parent torn down)") except Exception as exc: logger.warning( @@ -965,8 +966,15 @@ async def _stream_dual( exc or "", ) metrics.pd_bootstrap_failures_total.labels(reason="prefill_exception").inc() + self.breaker.record_failure(p.worker_id) else: - if getattr(p_resp, "status_code", 0) >= 400: + # The prefill leg is scored here rather than beside the decode + # leg above because this is where its own answer arrives: it + # runs concurrently, so nothing about it is known until now. + p_status = getattr(p_resp, "status_code", None) + if p_status is not None: + self._score_leg(p.worker_id, p_status) + if p_status is not None and p_status >= 400: logger.warning( "prefill leg %s returned %d (decode will hang on KVPoll)", p.worker_id, diff --git a/tests/unit/router/test_disagg_breaker.py b/tests/unit/router/test_disagg_breaker.py index 0f0b76c0..6e9dd681 100644 --- a/tests/unit/router/test_disagg_breaker.py +++ b/tests/unit/router/test_disagg_breaker.py @@ -91,7 +91,19 @@ async def test_decode_only_stream_records_failure_on_unreachable(): @pytest.mark.asyncio async def test_dual_stream_records_failure_on_unreachable_decode(): + """Only the decode leg is broken here. Failing both -- which is what a + transport that refuses everything does -- cannot show that the pools are + scored independently, because then prefill deserves to trip too. + """ r = _router() + + def _only_decode_is_down(request): + if "d1" in str(request.url): + raise httpx.ConnectError("refused", request=request) + return httpx.Response(200, json={"id": "x"}) + + r._client = httpx.AsyncClient(transport=httpx.MockTransport(_only_decode_is_down)) + p_target = RouteTarget(_w("p1")) d_target = RouteTarget(_w("d1")) for _ in range(3): @@ -174,6 +186,72 @@ async def test_a_served_request_clears_the_failure_count(): await r.aclose() +@pytest.mark.asyncio +async def test_a_healthy_leg_does_not_launder_a_broken_one(): + """The two legs are two workers whose health is independent, so each is + scored from its own response. + + Scoring both off the single client-facing status code let the decode leg's + 200 count as evidence for a prefill that had just 500'd -- resetting its + failure count, and reopening a breaker that was already open. + """ + + def _prefill_is_broken(request): + if "p1" in str(request.url): + return httpx.Response(500, json={"error": "prefill exploded"}) + return httpx.Response(200, json={"id": "x", "choices": [{"message": {"content": "hi"}}]}) + + r = DisaggRouter( + _RolePool(_pd_worker("p1", DisaggMode.PREFILL), _pd_worker("d1", DisaggMode.DECODE)), + _FakePolicy(), + ) + r._client = httpx.AsyncClient(transport=httpx.MockTransport(_prefill_is_broken)) + + for _ in range(3): + resp = await r.dispatch({"model": "m"}, stream=False) + assert resp.status_code == 200, "decode answers, so the client still gets 200" + + assert r.breaker.state_of("p1").value == "open", ( + "a prefill that 500s every request must trip, even though the decode " + "leg beside it succeeds" + ) + assert r.breaker.state_of("d1").value == "closed", "the healthy leg is untouched" + await r.aclose() + + +@pytest.mark.asyncio +async def test_a_streaming_dispatch_is_not_scored_before_it_runs(): + """A StreamingResponse is returned before its generator is touched: the + decode leg has not been POSTed and its 200 is Starlette's default, not an + outcome. Scoring it there records success for both roles before either was + dispatched and resets the count the legs are about to raise -- a decode + worker failing every request would never trip. + + Drives dispatch(stream=True), which is the production path; the older tests + call the generators directly and so cannot see this. + """ + r = DisaggRouter( + _RolePool(_pd_worker("p1", DisaggMode.PREFILL), _pd_worker("d1", DisaggMode.DECODE)), + _FakePolicy(), + ) + r._client = httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: (_ for _ in ()).throw(httpx.ConnectError("refused", request=request)) + ) + ) + r._DECODE_OPEN_MAX_RETRIES = 0 + + for i in range(3): + resp = await r.dispatch({"model": "m"}, stream=True) + assert await _drain(resp.body_iterator), f"request {i} produced no body" + + assert r.breaker.state_of("d1").value == "open", ( + "three unreachable decodes must trip the breaker; if this is closed the " + "wrapper scored the stream before the decode leg ran" + ) + await r.aclose() + + @pytest.mark.asyncio async def test_a_tripped_pd_worker_recovers_after_a_good_probe(): """The probe is dispatched and succeeds; if nothing records that, the From 9866a5314c24d3f3cc80505dc737212ac1bc318d Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Sun, 9 Aug 2026 08:13:14 +0000 Subject: [PATCH 71/88] fix(drain): let each backend own the shutdown signal it can actually observe The two discovery backends learn that a worker is leaving in different ways, and treating them alike gave the Kubernetes one two writers of the same fact. Kubernetes stamps a condemned Pod with deletionTimestamp before the worker is even signalled, and KubernetesRegistry reads it -- so routing has already stopped by the time the process could announce anything, and the announcement never reaches a reader: the terminating check returns before the annotation is parsed at all. What it did do was put state into a record the heartbeat rebuilds from config, so a refresh landing mid-drain overwrote DRAINING with an omitted status, which parses as ACTIVE. announce_draining is therefore gone from that client, and the annotation carries identity only -- fixed for the life of the process, so a refresh is byte-identical. etcd has no such signal. A record is either present or absent, so without the worker saying so there is no way to express "still finishing, send me nothing new". It keeps announce_draining, and each client now declares which of them owns the signal rather than the callers testing the backend. That makes the heartbeat safe to run to the end, which it now does. Cancelling it before the drain left the etcd lease with at most one TTL, and a drain longer than the remainder had etcd collect the key partway through -- the worker vanishing rather than visibly finishing, the exact outcome the DRAINING state exists to prevent, in the one backend where it is not redundant. Both engine entrypoints and the fake worker had it on opposite sides of the drain; all three now cancel it just before deregistering. No test had ever run a heartbeat after an announcement, which is why this held in both directions at once. Co-authored-by: Cursor Signed-off-by: leiwei12 --- infera/common/registration.py | 6 +++ infera/common/registration_k8s.py | 46 ++++++++++------------- infera/engine/sglang/__main__.py | 27 +++++++++---- infera/engine/vllm/__main__.py | 27 +++++++++---- infera/tools/fakeworker/server.py | 15 ++++++-- tests/unit/common/test_draining_status.py | 45 ++++++++++++++++++++++ 6 files changed, 121 insertions(+), 45 deletions(-) diff --git a/infera/common/registration.py b/infera/common/registration.py index bbb5a377..89e46944 100644 --- a/infera/common/registration.py +++ b/infera/common/registration.py @@ -56,6 +56,12 @@ def build_worker_payload(config: EngineConfig, *, status: WorkerStatus | None = class RegistrationClient: """Worker-side self-registration via an etcd lease (HTTP/JSON gateway).""" + #: etcd knows only that a record exists and its lease is unexpired; nothing + #: outside the worker observes that the process is going away. So the + #: worker has to say so itself, or there is no way to express "still + #: finishing what I have, send me nothing new" -- only present or absent. + announces_draining = True + def __init__( self, endpoint: str, diff --git a/infera/common/registration_k8s.py b/infera/common/registration_k8s.py index 04777fc4..80bdec9e 100644 --- a/infera/common/registration_k8s.py +++ b/infera/common/registration_k8s.py @@ -26,7 +26,6 @@ from infera.common.discovery_k8s import WORKER_INFO_ANNOTATION from infera.common.k8s_client import in_cluster_namespace, make_client from infera.common.registration import build_worker_payload -from infera.common.worker_pool import WorkerStatus from infera.engine.base import EngineConfig logger = logging.getLogger(__name__) @@ -37,7 +36,25 @@ class K8sRegistrationClient: - """Worker-side self-registration by patching its own Pod annotation.""" + """Worker-side self-registration by patching its own Pod annotation. + + The annotation carries identity and nothing else -- who this worker is, + where to reach it, and what it can serve. All of that is fixed for the + life of the process, which is what makes the refresh below safe to run at + any time, including throughout a shutdown. + + State is deliberately absent. Kubernetes stamps a condemned Pod with + ``deletionTimestamp`` before the worker is even signalled, and + ``KubernetesRegistry`` reads it, so the orchestrator already answers "is + this worker going away" earlier and more authoritatively than the worker + could. Writing the same fact into the annotation as well would make it a + race between two writers of one truth -- and the refresh below, which + rebuilds the payload from config, would be the one to win it. + """ + + #: See the class docstring: on this backend the orchestrator owns the + #: signal, so there is nothing for the worker to announce. + announces_draining = False def __init__( self, @@ -81,31 +98,6 @@ async def register(self, config: EngineConfig) -> str: ) return worker_id - async def announce_draining(self) -> bool: - """Rewrite the annotation as DRAINING instead of clearing it. - - ``list_active`` filters DRAINING out, so new work stops being routed - here while the record — and therefore the worker's visibility in - ``/v1/workers`` — survives the drain. A worker that vanishes looks the - same as one that crashed. - - Largely belt-and-braces on this backend: ``KubernetesRegistry`` already - drops a Pod the moment it carries a ``deletionTimestamp``, so a - kubectl-initiated shutdown has stopped receiving work before the worker - is even signalled. This covers the paths that do not go through Pod - deletion at all. - """ - if self._config is None: - return False - try: - payload = build_worker_payload(self._config, status=WorkerStatus.DRAINING) - await self._patch_annotation(json.dumps(payload)) - logger.info("worker %s announced DRAINING", self._worker_id) - return True - except Exception as exc: # noqa: BLE001 - shutdown must continue - logger.warning("could not announce DRAINING: %s", exc) - return False - async def deregister(self) -> None: # Best-effort: clear the annotation so a terminating-but-lingering Pod # stops being routed before its DELETE event lands. diff --git a/infera/engine/sglang/__main__.py b/infera/engine/sglang/__main__.py index 1873434a..8793e250 100644 --- a/infera/engine/sglang/__main__.py +++ b/infera/engine/sglang/__main__.py @@ -409,19 +409,18 @@ async def _run_after_start(args: SglangWorkerArgs, engine: SglangEngine, config) except asyncio.CancelledError: pass - hb_task.cancel() - try: - await hb_task - except asyncio.CancelledError: - pass - # Announce first, drain second, deregister last. Announcing DRAINING takes # this worker out of routing (list_active filters it) while leaving the # record in place, so for the whole drain it is visibly draining rather than # simply gone -- which is what distinguishes an orderly rollout from a crash # in /v1/workers. Deregistering first would work too, but it throws away # that signal at exactly the moment someone is watching for it. - await reg_client.announce_draining() + # + # Only where the worker owns that signal: under Kubernetes the registry has + # already dropped this Pod from routing on its deletionTimestamp, before + # this process was even signalled. + if reg_client.announces_draining: + await reg_client.announce_draining() if nats_req_server is not None: await nats_req_server.stop(drain=True, drain_timeout=args.drain_timeout) else: @@ -436,6 +435,20 @@ async def _run_after_start(args: SglangWorkerArgs, engine: SglangEngine, config) timeout=args.drain_timeout, ) + # The heartbeat runs until the record is gone, not until the drain starts. + # On etcd it is what keeps the lease alive: cancelling it earlier leaves at + # most one TTL of life, and a drain longer than that has etcd collect the + # key mid-drain -- the worker vanishing rather than visibly finishing, + # which is the outcome announcing DRAINING exists to avoid. It is safe to + # leave running because a refresh only re-asserts identity: on etcd it + # renews the lease without touching the value, and on Kubernetes it + # rewrites an annotation that carries no state to overwrite. + hb_task.cancel() + try: + await hb_task + except asyncio.CancelledError: + pass + await reg_client.deregister() if kv_relay is not None: diff --git a/infera/engine/vllm/__main__.py b/infera/engine/vllm/__main__.py index 7179ee23..24ec6797 100644 --- a/infera/engine/vllm/__main__.py +++ b/infera/engine/vllm/__main__.py @@ -355,19 +355,18 @@ async def main() -> None: except asyncio.CancelledError: pass - hb_task.cancel() - try: - await hb_task - except asyncio.CancelledError: - pass - # Announce first, drain second, deregister last. Announcing DRAINING takes # this worker out of routing (list_active filters it) while leaving the # record in place, so for the whole drain it is visibly draining rather than # simply gone -- which is what distinguishes an orderly rollout from a crash # in /v1/workers. Deregistering first would work too, but it throws away # that signal at exactly the moment someone is watching for it. - await reg_client.announce_draining() + # + # Only where the worker owns that signal: under Kubernetes the registry has + # already dropped this Pod from routing on its deletionTimestamp, before + # this process was even signalled. + if reg_client.announces_draining: + await reg_client.announce_draining() if nats_req_server is not None: await nats_req_server.stop(drain=True, drain_timeout=args.drain_timeout) else: @@ -382,6 +381,20 @@ async def main() -> None: timeout=args.drain_timeout, ) + # The heartbeat runs until the record is gone, not until the drain starts. + # On etcd it is what keeps the lease alive: cancelling it earlier leaves at + # most one TTL of life, and a drain longer than that has etcd collect the + # key mid-drain -- the worker vanishing rather than visibly finishing, + # which is the outcome announcing DRAINING exists to avoid. It is safe to + # leave running because a refresh only re-asserts identity: on etcd it + # renews the lease without touching the value, and on Kubernetes it + # rewrites an annotation that carries no state to overwrite. + hb_task.cancel() + try: + await hb_task + except asyncio.CancelledError: + pass + await reg_client.deregister() if kv_relay is not None: await kv_relay.stop() diff --git a/infera/tools/fakeworker/server.py b/infera/tools/fakeworker/server.py index d896ff7d..1159c7da 100644 --- a/infera/tools/fakeworker/server.py +++ b/infera/tools/fakeworker/server.py @@ -473,10 +473,14 @@ async def _shutdown() -> None: # races arrivals; keeping the record until the end is what makes the # worker visibly draining rather than simply gone. state.draining = True - try: - await reg.announce_draining() - except Exception as exc: # noqa: BLE001 - shutdown must not raise - logger.warning("announce_draining failed: %s", exc) + # Only where the worker owns that signal: under Kubernetes the registry + # drops a Pod on its deletionTimestamp, before this process is even + # signalled, so there is nothing to announce. + if reg.announces_draining: + try: + await reg.announce_draining() + except Exception as exc: # noqa: BLE001 - shutdown must not raise + logger.warning("announce_draining failed: %s", exc) if nats_req_server is not None: # The real drain: unsubscribe first so nothing new arrives, then # wait on the in-flight set infera actually holds. No polling and no @@ -488,6 +492,9 @@ async def _shutdown() -> None: await asyncio.sleep(0.1) if state.running: logger.warning("drain timeout with %d request(s) still in flight", state.running) + # Cancelled only once the record is about to go: on etcd the heartbeat + # is what renews the lease, and a drain longer than one TTL would + # otherwise have the key collected mid-drain. hb_task.cancel() try: await reg.deregister() diff --git a/tests/unit/common/test_draining_status.py b/tests/unit/common/test_draining_status.py index 852b6cdd..cd9fd109 100644 --- a/tests/unit/common/test_draining_status.py +++ b/tests/unit/common/test_draining_status.py @@ -119,3 +119,48 @@ async def test_announce_before_register_is_a_no_op(): c._http = _FakeHttp() assert await c.announce_draining() is False assert c._http.puts == [] + + +# --- which backend owns the "going away" signal ------------------------------- + + +def test_each_backend_declares_who_owns_the_shutdown_signal(): + """The two backends learn that a worker is leaving in different ways, and + only one of them needs the worker to say so. + + Kubernetes stamps a condemned Pod with deletionTimestamp before the worker + is even signalled, and the registry reads that -- so the worker announcing + it as well would be a second, competing source of the same fact, written by + whoever patched the annotation last. etcd has no such signal: a record is + either there or it is not, so without the worker's own announcement there + is no way to express "still finishing, send me nothing new". + """ + from infera.common.registration import RegistrationClient + from infera.common.registration_k8s import K8sRegistrationClient + + assert RegistrationClient.announces_draining is True + assert K8sRegistrationClient.announces_draining is False + + # The method exists only where it is meaningful, so a caller cannot + # accidentally write a status the Kubernetes registry will never read. + assert hasattr(RegistrationClient, "announce_draining") + assert not hasattr(K8sRegistrationClient, "announce_draining") + + +@pytest.mark.asyncio +async def test_a_k8s_heartbeat_is_identity_only_and_so_survives_a_drain(): + """The heartbeat re-asserts the annotation to self-heal, and it must be + safe to keep running for the whole shutdown -- on etcd that is what keeps + the lease alive through the drain. + + That only holds while the annotation carries identity and nothing else. If + it also carried state, a refresh landing mid-drain would overwrite it with + whatever the payload builder defaults to. + """ + from infera.common.registration import build_worker_payload + + cfg = _cfg() + first = build_worker_payload(cfg) + later = build_worker_payload(cfg) + assert first == later, "a refresh must be byte-identical, i.e. carry no state" + assert "status" not in first, "identity only; state belongs to the backend" From 31b2bb5954e63b2c5bd0d5cd372319e2363a149d Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Sun, 9 Aug 2026 08:17:29 +0000 Subject: [PATCH 72/88] feat(operator): refuse the external etcd backend for in-cluster deployments Anything this operator builds runs in Kubernetes, and there the orchestrator is what knows a worker is leaving: a condemned Pod carries deletionTimestamp before the process is signalled, and the registry drops it from routing then. Pointing such a deployment at an external etcd discards that. The server stops watching Pods, so nothing reads the deletionTimestamp, and the only remaining signal is the worker announcing DRAINING after SIGTERM -- which it receives only once the preStop delay this same operator injects has elapsed. The combination keeps that delay while losing the early notice it exists to provide, so for its whole duration the router keeps handing new work to a Pod already on its way out: strictly worse than either backend on its own. Refused at reconcile rather than in the CRD enum. Tightening the enum would have the API server reject the stored object outright, leaving anyone who had already applied it unable to edit their way out. This way the spec still applies, the reason lands on status.state and in the error, and no child workload is built from a configuration that was turned down. Co-authored-by: Cursor Signed-off-by: leiwei12 --- .../controller/discovery_backend_test.go | 88 +++++++++++++++++++ .../controller/inferadeployment_controller.go | 32 +++++++ 2 files changed, 120 insertions(+) create mode 100644 deploy/operator/internal/controller/discovery_backend_test.go diff --git a/deploy/operator/internal/controller/discovery_backend_test.go b/deploy/operator/internal/controller/discovery_backend_test.go new file mode 100644 index 00000000..fcea7866 --- /dev/null +++ b/deploy/operator/internal/controller/discovery_backend_test.go @@ -0,0 +1,88 @@ +/* +Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + +SPDX-License-Identifier: MIT +*/ + +package controller + +import ( + "context" + "strings" + "testing" + + appsv1 "k8s.io/api/apps/v1" + "k8s.io/apimachinery/pkg/types" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + + inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" +) + +// An operator-managed deployment is by definition running in Kubernetes, and +// there the orchestrator is what knows a worker is going away: a condemned Pod +// carries deletionTimestamp before the process is even signalled, and the +// registry drops it from routing then. +// +// Pointing such a deployment at an external etcd throws that away. The router +// stops watching Pods, so nothing reads the deletionTimestamp, and the only +// remaining signal is the worker announcing DRAINING after SIGTERM -- which +// arrives once the preStop delay the operator itself injects has elapsed. The +// combination keeps that delay and loses the early notice it exists to give, +// so for its whole duration the router keeps handing new work to a Pod that is +// already condemned. Refusing is better than rendering a deployment whose +// drain is worse than either backend alone. +func TestAnOperatorDeploymentRefusesTheExternalEtcdBackend(t *testing.T) { + s := scaleScheme(t) + idep := idepWith(1) + idep.Spec.DiscoveryBackend = "etcd" + idep.Spec.EtcdEndpoint = "etcd:2379" + cl := fake.NewClientBuilder().WithScheme(s). + WithObjects(idep).WithStatusSubresource(idep).Build() + + r := &InferaDeploymentReconciler{Client: cl, Scheme: s} + _, err := r.Reconcile(context.Background(), ctrl.Request{ + NamespacedName: types.NamespacedName{Name: "qwen", Namespace: "ns"}, + }) + if err == nil { + t.Fatal("reconcile accepted discoveryBackend=etcd; it must be refused") + } + if !strings.Contains(err.Error(), "discoveryBackend") { + t.Fatalf("error should name the field so the cause is obvious, got: %v", err) + } + + // Refusing must not leave a half-built deployment behind. + dep := &appsv1.Deployment{} + key := types.NamespacedName{Name: "qwen-decode", Namespace: "ns"} + if err := cl.Get(context.Background(), key, dep); err == nil { + t.Fatal("a child workload was created for a configuration that was refused") + } + + // The reason belongs on the object, not only in the operator's log. + got := &inferav1alpha1.InferaDeployment{} + if err := cl.Get(context.Background(), types.NamespacedName{Name: "qwen", Namespace: "ns"}, got); err != nil { + t.Fatalf("get idep: %v", err) + } + if got.Status.State != inferav1alpha1.StateFailed { + t.Errorf("status.state = %q, want %q so `kubectl get idep` shows it", + got.Status.State, inferav1alpha1.StateFailed) + } +} + +// The default and an explicit "kubernetes" both reconcile normally. +func TestTheKubernetesBackendIsAccepted(t *testing.T) { + for _, backend := range []string{"", "kubernetes"} { + s := scaleScheme(t) + idep := idepWith(1) + idep.Spec.DiscoveryBackend = backend + cl := fake.NewClientBuilder().WithScheme(s). + WithObjects(idep).WithStatusSubresource(idep).Build() + reconcileOnce(t, cl, s) + + dep := &appsv1.Deployment{} + key := types.NamespacedName{Name: "qwen-decode", Namespace: "ns"} + if err := cl.Get(context.Background(), key, dep); err != nil { + t.Fatalf("discoveryBackend=%q: child Deployment not created: %v", backend, err) + } + } +} diff --git a/deploy/operator/internal/controller/inferadeployment_controller.go b/deploy/operator/internal/controller/inferadeployment_controller.go index 3dc4356b..9157613f 100644 --- a/deploy/operator/internal/controller/inferadeployment_controller.go +++ b/deploy/operator/internal/controller/inferadeployment_controller.go @@ -7,6 +7,7 @@ package controller import ( "context" + "fmt" "sort" "time" @@ -67,6 +68,37 @@ func (r *InferaDeploymentReconciler) Reconcile(ctx context.Context, req ctrl.Req return ctrl.Result{}, nil } + // Anything this operator builds runs in Kubernetes, and there the + // orchestrator is what knows a worker is leaving: a condemned Pod carries + // deletionTimestamp before the process is signalled, and the registry drops + // it from routing at that moment. + // + // Pointing such a deployment at an external etcd discards that. The server + // stops watching Pods, so nothing reads the deletionTimestamp, and the only + // remaining signal is the worker announcing DRAINING after SIGTERM -- which + // it receives only once the preStop delay injected below has elapsed. The + // combination keeps that delay while losing the early notice it exists to + // provide, so for its whole duration the router keeps handing new work to a + // Pod already on its way out. Refusing beats rendering a deployment whose + // drain is worse than either backend on its own. + if !useK8sDiscovery(idep) { + err := fmt.Errorf( + "spec.discoveryBackend=%q is not supported by the operator: an in-cluster "+ + "deployment must use the default \"kubernetes\" backend, which learns of a "+ + "departing worker from its Pod's deletionTimestamp. External etcd is for "+ + "deployments outside Kubernetes", + idep.Spec.DiscoveryBackend, + ) + lg.Error(err, "refusing to reconcile") + idep.Status.ObservedGeneration = idep.Generation + idep.Status.State = inferav1alpha1.StateFailed + if uerr := r.Status().Update(ctx, idep); uerr != nil { + lg.Error(uerr, "status update failed") + } + // Terminal: retrying cannot change a spec field, so surface it and stop. + return ctrl.Result{}, err + } + // 0. Kubernetes-native discovery RBAC: a namespaced ServiceAccount + Role so // workers can patch their own Pod annotation and the server can list/watch // this deployment's worker Pods (no external etcd). From 97fb9f112b1994b7ea005bd4eca7c0178541fdac Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Sun, 9 Aug 2026 08:21:29 +0000 Subject: [PATCH 73/88] docs(scaling): say which backend knows a worker is leaving, and why The page described announcing DRAINING as the universal first step of a shutdown. That is only true on etcd. Under Kubernetes the orchestrator marks a condemned Pod before the process is signalled, the registry acts on that, and the worker announcing it later would change nothing -- so the annotation now carries identity only and there is no announcement on that path. Adds a section naming the asymmetry and the reason for it: what each backend can observe. etcd has no orchestrator, so a record is present or absent and the worker must speak for itself; Kubernetes has one that knows a Pod is condemned before the worker does. Also explains why the annotation must stay stateless -- the heartbeat rebuilds it from config, so state written there is overwritten by the next refresh -- and why the heartbeat now runs until deregistration, which is what keeps the etcd lease alive across a long drain. Documents that discoveryBackend=etcd is unsupported in-cluster, with the reason: it keeps the preStop delay while losing the early notice that delay exists to provide, so the router keeps assigning work to a condemned Pod for the whole window. The operator refuses it. Co-authored-by: Cursor Signed-off-by: leiwei12 --- manual/features/scaling.md | 67 ++++++++++++++++++++++++++++++++++---- 1 file changed, 60 insertions(+), 7 deletions(-) diff --git a/manual/features/scaling.md b/manual/features/scaling.md index 3849a2ac..7080ce36 100644 --- a/manual/features/scaling.md +++ b/manual/features/scaling.md @@ -14,9 +14,14 @@ stopping worker processes; nothing has to be told about it. ``` worker ready ──► register (etcd lease / Pod annotation) ──► router's watch fires ──► receives traffic -SIGTERM ──► announce DRAINING ──► drain in-flight ──► deregister ──► exit +leaving ──► stop being routed to ──► drain in-flight ──► deregister ──► exit ``` +What triggers "stop being routed to" depends on the backend: under Kubernetes +the orchestrator marks the Pod before the process is even signalled, on etcd the +worker announces it after `SIGTERM`. See [Who says the worker is +leaving](#who-says-the-worker-is-leaving). + That shape is why scale-up and scale-down have very different costs. Scale-up is bounded by **model load**, which is minutes. Scale-down is bounded by the **longest in-flight generation**, which is seconds — and the router stops @@ -53,10 +58,12 @@ grace period. The worker then, in this order: -1. **Announces `DRAINING`.** The router filters draining workers out of routing +1. **Stops being routed to.** The router filters draining workers out immediately, so no new work arrives. The record stays, so the worker remains visible in `/v1/workers` — a worker that vanishes looks exactly like one that - crashed. + crashed. On etcd the worker announces `DRAINING` here; under Kubernetes this + already happened when the Pod was condemned, before it was signalled at all + (see [Who says the worker is leaving](#who-says-the-worker-is-leaving)). 2. **Drains.** On the NATS transport infera tracks in-flight requests directly. On HTTP the router talks straight to the engine, so infera asks the engine instead, polling its `/metrics` until running, queued, and PD-handoff queues @@ -67,9 +74,9 @@ Requests already in flight run to completion. Requests that arrive during the drain go to other workers. **Two different timings, easily conflated.** A worker stops *receiving* new -requests within a second of `SIGTERM` — that is the `DRAINING` announcement plus -the router's watch, and it is the number that decides whether traffic is still -being sent somewhere that is about to die. How long the *process* then lives is +requests within a second of the shutdown starting — the router's watch picking +up either the announcement or the `deletionTimestamp` — and it is the number +that decides whether traffic is still being sent somewhere that is about to die. How long the *process* then lives is a separate and much larger number, set by the longest generation it was already serving. Measured: under a second to stop receiving, 38 s until the record disappeared, while a 40-second generation ran to completion in between. @@ -203,10 +210,56 @@ its record stays until it clears its own annotation at the end of the drain, so ``` deletion requested ──► deletionTimestamp set ──► registry drops the worker ──► preStop sleep 15 (still serving what it has) - ──► SIGTERM ──► DRAINING ──► drain ──► deregister ──► engine.stop() + ──► SIGTERM ──► drain ──► deregister ──► engine.stop() ──► [kubelet SIGKILL at terminationGracePeriodSeconds] ``` +### Who says the worker is leaving + +The two discovery backends learn this in different ways, and only one of them +needs the worker to say anything. The difference is not an inconsistency to be +smoothed over — it is what each backend can actually observe. + +**Kubernetes: the orchestrator says so.** A condemned Pod carries +`deletionTimestamp` from the moment deletion is requested, which is before the +`preStop` hook runs and therefore before the process is signalled at all. The +registry reads it and drops the worker from routing immediately — measured at +under 100 ms against the 15 s `preStop` delay. The worker announcing the same +thing later would add nothing: routing has already stopped, and the terminating +check returns before the annotation is even parsed. + +So on this backend the annotation carries **identity only** — worker id, URL, +model, engine, role, KV endpoints — all of it fixed for the life of the +process. That is deliberate. The heartbeat re-asserts the annotation to +self-heal, rebuilding it from config; if state lived there too, a refresh +landing mid-drain would overwrite it with a payload that omits the status, +which parses as `ACTIVE`, and the worker would be handed new work it is about +to refuse. + +**etcd: the worker says so.** There is no orchestrator here. A record is either +present with an unexpired lease or it is gone — nothing observes that the +process is on its way out. Without the worker announcing `DRAINING` there is no +way to express "still finishing what I have, send me nothing new", so on this +backend the record does carry a status and the worker rewrites it before +draining. + +The heartbeat runs until the record is deleted on both backends, which matters +most here: it is what renews the lease. Stopping it when the drain starts would +leave at most one TTL of life, and a drain longer than the remainder has etcd +collect the key partway through — the worker vanishing rather than visibly +finishing, which is the outcome announcing `DRAINING` exists to prevent. + +```{warning} +`discoveryBackend: etcd` **is not supported for in-cluster deployments** and the +operator refuses it. The combination keeps the `preStop` delay while losing the +early notice it exists to provide: the server no longer watches Pods, so nothing +reads the `deletionTimestamp`, and the only signal left arrives after `SIGTERM` +— that is, after the delay has already elapsed. For its whole duration the +router keeps handing new work to a Pod that is already condemned. Use the +default `kubernetes` backend in Kubernetes; external etcd is for deployments +outside it. +``` + ### Worst case, and the budget Every stage is individually bounded: From 3ed34248dfe9fcefe3a3e09a6164a4c0f3832a39 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Mon, 10 Aug 2026 01:40:01 +0000 Subject: [PATCH 74/88] fix(drain): drop the DRAINING announcement; the backend decides the order Neither backend writes state into the registration record now, and the shutdown order follows from what each one can observe. Under Kubernetes the announcement was never read. The registry acts on the Pod's deletionTimestamp and returns before the annotation is parsed, so routing had already stopped by the time a worker could announce anything -- and the heartbeat, which rebuilds the payload from config, erased the status on its next refresh anyway. On etcd it was read, but deregistering already stops new work arriving, so announcing first was a second mechanism for the same effect, with a lease to keep alive across the drain as its price. So the record carries identity only, every field fixed for the life of the process, which is what makes a heartbeat refresh harmless at any point. The clients instead declare whether deregistering is what stops new work, and the entrypoints order the shutdown from that rather than testing the backend: etcd deregister -> drain -> stop (the record's presence is what makes it a candidate) kubernetes drain -> deregister -> stop (routing stopped when the Pod was condemned; the record stays so the worker is visible while it drains) The heartbeat is cancelled before either, since a refresh landing after deregistration would re-register the worker outright. Graceful drain itself is unchanged and still applies to both transports on both backends: what differs is only which step precedes it. The fake worker's two waits also share one deadline now, so --drain-timeout 30 can no longer mean a 60s shutdown. Co-authored-by: Cursor Signed-off-by: leiwei12 --- infera/common/registration.py | 55 ++---- infera/common/registration_k8s.py | 8 +- infera/engine/sglang/__main__.py | 68 ++++---- infera/engine/vllm/__main__.py | 68 ++++---- infera/tools/fakeworker/server.py | 69 ++++---- tests/unit/common/test_draining_status.py | 194 +++++++++------------- 6 files changed, 199 insertions(+), 263 deletions(-) diff --git a/infera/common/registration.py b/infera/common/registration.py index 89e46944..fe7aba55 100644 --- a/infera/common/registration.py +++ b/infera/common/registration.py @@ -12,7 +12,6 @@ import httpx from infera.common.discovery import DEFAULT_PREFIX, _b64, _normalize_endpoint -from infera.common.worker_pool import WorkerStatus from infera.engine.base import EngineConfig logger = logging.getLogger(__name__) @@ -20,17 +19,19 @@ _DEFAULT_LEASE_TTL = 30 # seconds -def build_worker_payload(config: EngineConfig, *, status: WorkerStatus | None = None) -> dict: +def build_worker_payload(config: EngineConfig) -> dict: """Build the worker registration record shared by every backend. The same dict is PUT to etcd (RegistrationClient) or stored in the worker Pod annotation (K8sRegistrationClient), so the server-side parse (discovery.worker_info_from_json) is transport-agnostic. - ``status`` is omitted for a healthy worker, which keeps the record identical - to what older workers wrote and lets the parser's ACTIVE default stand. It - is set only to announce DRAINING, so the field's presence means something - happened rather than being ambient. + Identity only, and every field is fixed for the life of the process. No + status is written: the record's own presence is the etcd backend's answer to + "is this worker available", and under Kubernetes the Pod's deletionTimestamp + answers it earlier than the worker could. Keeping the payload stateless is + also what makes the heartbeat safe -- it rebuilds this from config, so any + state written here would be erased by the next refresh. """ worker_id = f"{config.host}:{config.port}" payload: dict = { @@ -46,8 +47,6 @@ def build_worker_payload(config: EngineConfig, *, status: WorkerStatus | None = "dp_size": config.dp_size, "request_transport": getattr(config, "request_transport", "http"), } - if status is not None and status is not WorkerStatus.ACTIVE: - payload["status"] = status.value if config.kv is not None: payload["kv"] = config.kv.to_dict() return payload @@ -56,11 +55,11 @@ def build_worker_payload(config: EngineConfig, *, status: WorkerStatus | None = class RegistrationClient: """Worker-side self-registration via an etcd lease (HTTP/JSON gateway).""" - #: etcd knows only that a record exists and its lease is unexpired; nothing - #: outside the worker observes that the process is going away. So the - #: worker has to say so itself, or there is no way to express "still - #: finishing what I have, send me nothing new" -- only present or absent. - announces_draining = True + #: Removing the record is what stops new work arriving here: nothing outside + #: the worker observes that it is going away, so a shutdown has to + #: deregister before it drains. The record is therefore present or absent, + #: with no state in between. + deregister_stops_routing = True def __init__( self, @@ -112,36 +111,6 @@ async def register(self, config: EngineConfig) -> str: ) return worker_id - async def announce_draining(self) -> bool: - """Rewrite the record as DRAINING, keeping the lease alive. - - ``list_active`` filters DRAINING out, so this stops new work being - routed here without deleting the record — which is the difference that - matters during a shutdown. A worker that simply vanishes is - indistinguishable from one that crashed; one that is visibly draining - tells an operator (and ``/v1/workers``) that a rolling update is - proceeding normally and roughly how far along it is. - - On the Kubernetes backend this is largely redundant with the registry's - ``deletionTimestamp`` check, which removes a condemned Pod without the - worker having to say anything. It is not redundant on etcd, where - nothing else observes that the process is going away. - """ - if self._lease_id is None or self._key is None or self._config is None: - return False - try: - value = json.dumps(build_worker_payload(self._config, status=WorkerStatus.DRAINING)) - r = await self._http.post( - "/v3/kv/put", - json={"key": _b64(self._key), "value": _b64(value), "lease": self._lease_id}, - ) - r.raise_for_status() - logger.info("worker %s announced DRAINING", self._worker_id) - return True - except Exception as exc: # noqa: BLE001 - shutdown must continue - logger.warning("could not announce DRAINING: %s", exc) - return False - async def deregister(self) -> None: if self._lease_id is not None: try: diff --git a/infera/common/registration_k8s.py b/infera/common/registration_k8s.py index 80bdec9e..457c8f3c 100644 --- a/infera/common/registration_k8s.py +++ b/infera/common/registration_k8s.py @@ -52,9 +52,11 @@ class K8sRegistrationClient: rebuilds the payload from config, would be the one to win it. """ - #: See the class docstring: on this backend the orchestrator owns the - #: signal, so there is nothing for the worker to announce. - announces_draining = False + #: New work has already stopped arriving by the time a shutdown gets here -- + #: the registry drops a condemned Pod on its deletionTimestamp. So clearing + #: the annotation is only cleanup, and can wait until the drain is over, + #: which keeps the worker visible in /v1/workers while it finishes. + deregister_stops_routing = False def __init__( self, diff --git a/infera/engine/sglang/__main__.py b/infera/engine/sglang/__main__.py index 8793e250..651713a8 100644 --- a/infera/engine/sglang/__main__.py +++ b/infera/engine/sglang/__main__.py @@ -409,47 +409,45 @@ async def _run_after_start(args: SglangWorkerArgs, engine: SglangEngine, config) except asyncio.CancelledError: pass - # Announce first, drain second, deregister last. Announcing DRAINING takes - # this worker out of routing (list_active filters it) while leaving the - # record in place, so for the whole drain it is visibly draining rather than - # simply gone -- which is what distinguishes an orderly rollout from a crash - # in /v1/workers. Deregistering first would work too, but it throws away - # that signal at exactly the moment someone is watching for it. - # - # Only where the worker owns that signal: under Kubernetes the registry has - # already dropped this Pod from routing on its deletionTimestamp, before - # this process was even signalled. - if reg_client.announces_draining: - await reg_client.announce_draining() - if nats_req_server is not None: - await nats_req_server.stop(drain=True, drain_timeout=args.drain_timeout) - else: - # HTTP transport: the router talks straight to the engine, so infera - # never saw these requests and has to ask the engine what is still in - # flight. The DRAINING announcement above already stopped new work being - # routed here; this waits for the work already accepted. - await drain_engine_inflight( - host=config.host, - port=config.port, - engine=config.engine, - timeout=args.drain_timeout, - ) - - # The heartbeat runs until the record is gone, not until the drain starts. - # On etcd it is what keeps the lease alive: cancelling it earlier leaves at - # most one TTL of life, and a drain longer than that has etcd collect the - # key mid-drain -- the worker vanishing rather than visibly finishing, - # which is the outcome announcing DRAINING exists to avoid. It is safe to - # leave running because a refresh only re-asserts identity: on etcd it - # renews the lease without touching the value, and on Kubernetes it - # rewrites an annotation that carries no state to overwrite. + # Stop the heartbeat before touching the record: it re-asserts registration + # from config, so a refresh landing after deregistration would put the + # worker straight back into the pool. hb_task.cancel() try: await hb_task except asyncio.CancelledError: pass - await reg_client.deregister() + async def _drain() -> None: + if nats_req_server is not None: + await nats_req_server.stop(drain=True, drain_timeout=args.drain_timeout) + else: + # HTTP transport: the router talks straight to the engine, so infera + # never saw these requests and has to ask the engine what is still + # in flight. + await drain_engine_inflight( + host=config.host, + port=config.port, + engine=config.engine, + timeout=args.drain_timeout, + ) + + # Draining is only safe once new work has stopped arriving, and which step + # achieves that depends on the backend. + if reg_client.deregister_stops_routing: + # etcd: the record's presence is the only thing making this worker a + # candidate, so it has to go first. The worker disappears for the + # duration of the drain, which is the cost of a backend where nothing + # else can observe that it is leaving. + await reg_client.deregister() + await _drain() + else: + # Kubernetes: routing stopped when the Pod was condemned, before this + # process was signalled. Keeping the record until the end leaves the + # worker visible in /v1/workers while it finishes -- an orderly rollout + # rather than something indistinguishable from a crash. + await _drain() + await reg_client.deregister() if kv_relay is not None: await kv_relay.stop() diff --git a/infera/engine/vllm/__main__.py b/infera/engine/vllm/__main__.py index 24ec6797..d6e2e0d6 100644 --- a/infera/engine/vllm/__main__.py +++ b/infera/engine/vllm/__main__.py @@ -355,47 +355,45 @@ async def main() -> None: except asyncio.CancelledError: pass - # Announce first, drain second, deregister last. Announcing DRAINING takes - # this worker out of routing (list_active filters it) while leaving the - # record in place, so for the whole drain it is visibly draining rather than - # simply gone -- which is what distinguishes an orderly rollout from a crash - # in /v1/workers. Deregistering first would work too, but it throws away - # that signal at exactly the moment someone is watching for it. - # - # Only where the worker owns that signal: under Kubernetes the registry has - # already dropped this Pod from routing on its deletionTimestamp, before - # this process was even signalled. - if reg_client.announces_draining: - await reg_client.announce_draining() - if nats_req_server is not None: - await nats_req_server.stop(drain=True, drain_timeout=args.drain_timeout) - else: - # HTTP transport: the router talks straight to the engine, so infera - # never saw these requests and has to ask the engine what is still in - # flight. The DRAINING announcement above already stopped new work being - # routed here; this waits for the work already accepted. - await drain_engine_inflight( - host=config.host, - port=config.port, - engine=config.engine, - timeout=args.drain_timeout, - ) - - # The heartbeat runs until the record is gone, not until the drain starts. - # On etcd it is what keeps the lease alive: cancelling it earlier leaves at - # most one TTL of life, and a drain longer than that has etcd collect the - # key mid-drain -- the worker vanishing rather than visibly finishing, - # which is the outcome announcing DRAINING exists to avoid. It is safe to - # leave running because a refresh only re-asserts identity: on etcd it - # renews the lease without touching the value, and on Kubernetes it - # rewrites an annotation that carries no state to overwrite. + # Stop the heartbeat before touching the record: it re-asserts registration + # from config, so a refresh landing after deregistration would put the + # worker straight back into the pool. hb_task.cancel() try: await hb_task except asyncio.CancelledError: pass - await reg_client.deregister() + async def _drain() -> None: + if nats_req_server is not None: + await nats_req_server.stop(drain=True, drain_timeout=args.drain_timeout) + else: + # HTTP transport: the router talks straight to the engine, so infera + # never saw these requests and has to ask the engine what is still + # in flight. + await drain_engine_inflight( + host=config.host, + port=config.port, + engine=config.engine, + timeout=args.drain_timeout, + ) + + # Draining is only safe once new work has stopped arriving, and which step + # achieves that depends on the backend. + if reg_client.deregister_stops_routing: + # etcd: the record's presence is the only thing making this worker a + # candidate, so it has to go first. The worker disappears for the + # duration of the drain, which is the cost of a backend where nothing + # else can observe that it is leaving. + await reg_client.deregister() + await _drain() + else: + # Kubernetes: routing stopped when the Pod was condemned, before this + # process was signalled. Keeping the record until the end leaves the + # worker visible in /v1/workers while it finishes -- an orderly rollout + # rather than something indistinguishable from a crash. + await _drain() + await reg_client.deregister() if kv_relay is not None: await kv_relay.stop() await engine.stop() diff --git a/infera/tools/fakeworker/server.py b/infera/tools/fakeworker/server.py index 1159c7da..7edec758 100644 --- a/infera/tools/fakeworker/server.py +++ b/infera/tools/fakeworker/server.py @@ -468,38 +468,49 @@ async def _serve(args) -> None: stop = asyncio.Event() async def _shutdown() -> None: - # Announce, drain, then deregister -- the same order as the real worker - # entrypoints. Stopping new work has to come first or the drain just - # races arrivals; keeping the record until the end is what makes the - # worker visibly draining rather than simply gone. + # Mirrors the real worker entrypoints, including which step stops new + # work arriving -- that differs by backend, see below. state.draining = True - # Only where the worker owns that signal: under Kubernetes the registry - # drops a Pod on its deletionTimestamp, before this process is even - # signalled, so there is nothing to announce. - if reg.announces_draining: + # Stop the heartbeat first: it re-asserts registration from config, so a + # refresh landing after deregistration would put this worker back in the + # pool. + hb_task.cancel() + + deadline = time.monotonic() + args.drain_timeout + + async def _drain() -> None: + if nats_req_server is not None: + # The real drain: unsubscribe first so nothing new arrives, then + # wait on the in-flight set infera actually holds. No polling and + # no settle window -- unlike HTTP, where the count has to be + # inferred from the engine's lagging gauges. + await nats_req_server.stop( + drain=True, + drain_timeout=max(0.0, deadline - time.monotonic()), + ) + while state.running and time.monotonic() < deadline: + await asyncio.sleep(0.1) + if state.running: + logger.warning("drain timeout with %d request(s) still in flight", state.running) + + async def _deregister() -> None: try: - await reg.announce_draining() + await reg.deregister() except Exception as exc: # noqa: BLE001 - shutdown must not raise - logger.warning("announce_draining failed: %s", exc) - if nats_req_server is not None: - # The real drain: unsubscribe first so nothing new arrives, then - # wait on the in-flight set infera actually holds. No polling and no - # settle window -- unlike HTTP, where the count has to be inferred - # from the engine's lagging gauges. - await nats_req_server.stop(drain=True, drain_timeout=args.drain_timeout) - deadline = time.monotonic() + args.drain_timeout - while state.running and time.monotonic() < deadline: - await asyncio.sleep(0.1) - if state.running: - logger.warning("drain timeout with %d request(s) still in flight", state.running) - # Cancelled only once the record is about to go: on etcd the heartbeat - # is what renews the lease, and a drain longer than one TTL would - # otherwise have the key collected mid-drain. - hb_task.cancel() - try: - await reg.deregister() - except Exception as exc: # noqa: BLE001 - shutdown must not raise - logger.warning("deregister failed: %s", exc) + logger.warning("deregister failed: %s", exc) + + if reg.deregister_stops_routing: + # etcd: the record's presence is what makes this worker a candidate, + # so it has to go before the drain or the drain just races arrivals. + await _deregister() + await _drain() + else: + # Kubernetes: the registry dropped this Pod on its deletionTimestamp, + # before this process was signalled, so the record can stay until the + # end and keep the worker visible while it finishes. + await _drain() + await _deregister() + server.should_exit = True stop.set() diff --git a/tests/unit/common/test_draining_status.py b/tests/unit/common/test_draining_status.py index cd9fd109..67de2898 100644 --- a/tests/unit/common/test_draining_status.py +++ b/tests/unit/common/test_draining_status.py @@ -3,20 +3,24 @@ # # SPDX-License-Identifier: MIT ############################################################################### -"""Announcing DRAINING, and what it is actually for. - -``WorkerStatus.DRAINING`` has been in the enum and filtered out of -``list_active`` since the beginning, and until now nothing ever set it. The -value it adds over simply deleting the record is not routing -- both stop new -work -- it is that the worker stays *visible* while it drains. A worker that -vanishes looks identical to one that crashed; one that reports DRAINING tells an -operator a rollout is proceeding and roughly how far along it is. +"""Where DRAINING comes from, and why the registration record has no state. + +``WorkerStatus.DRAINING`` takes a worker out of ``list_active`` while leaving it +visible, which is what distinguishes an orderly rollout from a crash. What sets +it is the Kubernetes registry, from the Pod's ``deletionTimestamp`` -- not the +worker. + +That is a deliberate split. Kubernetes knows a Pod is condemned before the +process is signalled, so the orchestrator answers "is this worker leaving" +earlier and more authoritatively than the worker could. etcd has no such signal: +a record is present or absent, so there removing it is what stops new work, and +a shutdown deregisters before it drains. Either way the record itself carries +identity only, which is what makes the heartbeat -- which rebuilds it from +config -- safe to run at any point. """ from __future__ import annotations -import json - import pytest from infera.common.discovery import worker_info_from_json @@ -29,138 +33,92 @@ def _cfg(): return EngineConfig(model_name="m", host="10.0.0.1", port=8080, engine=EngineType.SGLANG) -def test_healthy_payload_omits_status_entirely(): - """Keeps the record byte-identical to what older workers wrote, so the - parser's ACTIVE default stands and the field's presence means something.""" - assert "status" not in build_worker_payload(_cfg()) - assert "status" not in build_worker_payload(_cfg(), status=WorkerStatus.ACTIVE) - +def test_the_record_carries_identity_and_no_state(): + """Every field is fixed for the life of the process, so two builds are + byte-identical. State written here would be erased by the next heartbeat, + which rebuilds the payload from the same config.""" + payload = build_worker_payload(_cfg()) + assert "status" not in payload + assert payload == build_worker_payload(_cfg()) -def test_draining_payload_carries_the_status(): - payload = build_worker_payload(_cfg(), status=WorkerStatus.DRAINING) - assert payload["status"] == "draining" +def test_a_record_without_a_status_reads_as_active(): + """Registration says nothing about state, so the parser's default is what + every healthy worker resolves to.""" + info = worker_info_from_json(build_worker_payload(_cfg())) + assert info.status is WorkerStatus.ACTIVE -def test_round_trip_through_discovery(): - """The wire record has to survive the same parse every backend uses.""" - payload = build_worker_payload(_cfg(), status=WorkerStatus.DRAINING) - info = worker_info_from_json(json.loads(json.dumps(payload))) - assert info.status is WorkerStatus.DRAINING - -def test_draining_worker_is_excluded_but_still_visible(): - """The whole point: out of rotation, still in the fleet listing.""" +def test_a_draining_worker_is_excluded_but_still_visible(): + """The value DRAINING adds over deleting the record is not routing -- both + stop new work -- it is that the worker stays observable while it finishes.""" pool = WorkerPool() pool.add(worker_info_from_json(build_worker_payload(_cfg()))) assert [w.worker_id for w in pool.list_active()] == ["10.0.0.1:8080"] - pool.add(worker_info_from_json(build_worker_payload(_cfg(), status=WorkerStatus.DRAINING))) + worker = pool.get("10.0.0.1:8080") + worker.status = WorkerStatus.DRAINING assert pool.list_active() == [], "a draining worker must not be routed to" assert pool.get("10.0.0.1:8080") is not None, "but it must still be observable" -# --- the etcd client ---------------------------------------------------------- - - -class _FakeHttp: - def __init__(self): - self.puts: list[dict] = [] - self.fail = False - - async def post(self, path, json=None): # noqa: A002 - mirrors httpx - if self.fail: - raise RuntimeError("etcd unreachable") - self.puts.append({"path": path, "json": json}) - - class R: - @staticmethod - def raise_for_status(): - pass - - return R() - - -@pytest.mark.asyncio -async def test_etcd_announce_writes_draining_on_the_same_lease(): - from infera.common.registration import RegistrationClient - - c = RegistrationClient("http://etcd:2379") - c._http = _FakeHttp() - c._lease_id, c._key, c._worker_id, c._config = 42, "/infera/workers/w", "w", _cfg() - - assert await c.announce_draining() is True - (put,) = c._http.puts - assert put["path"] == "/v3/kv/put" - assert put["json"]["lease"] == 42, "must keep the lease, not orphan the key" +# --- which step stops new work arriving --------------------------------------- - import base64 - value = json.loads(base64.b64decode(put["json"]["value"])) - assert value["status"] == "draining" +def test_each_backend_declares_what_stops_new_work(): + """The shutdown order follows from this, so the two clients state it rather + than every caller testing the backend. - -@pytest.mark.asyncio -async def test_announce_never_raises_on_the_shutdown_path(): - """It runs immediately before the drain; raising here would skip it.""" - from infera.common.registration import RegistrationClient - - c = RegistrationClient("http://etcd:2379") - c._http = _FakeHttp() - c._http.fail = True - c._lease_id, c._key, c._worker_id, c._config = 1, "/k", "w", _cfg() - assert await c.announce_draining() is False - - -@pytest.mark.asyncio -async def test_announce_before_register_is_a_no_op(): + On etcd the record's presence is the only thing making a worker a candidate, + so deregistering has to precede the drain or the drain races arrivals. Under + Kubernetes routing already stopped when the Pod was condemned, so the record + can outlive the drain and keep the worker visible. + """ from infera.common.registration import RegistrationClient + from infera.common.registration_k8s import K8sRegistrationClient - c = RegistrationClient("http://etcd:2379") - c._http = _FakeHttp() - assert await c.announce_draining() is False - assert c._http.puts == [] - - -# --- which backend owns the "going away" signal ------------------------------- + assert RegistrationClient.deregister_stops_routing is True + assert K8sRegistrationClient.deregister_stops_routing is False -def test_each_backend_declares_who_owns_the_shutdown_signal(): - """The two backends learn that a worker is leaving in different ways, and - only one of them needs the worker to say so. +def test_no_client_announces_a_status(): + """Neither backend writes state into the record any more. - Kubernetes stamps a condemned Pod with deletionTimestamp before the worker - is even signalled, and the registry reads that -- so the worker announcing - it as well would be a second, competing source of the same fact, written by - whoever patched the annotation last. etcd has no such signal: a record is - either there or it is not, so without the worker's own announcement there - is no way to express "still finishing, send me nothing new". + Under Kubernetes it was never read -- the registry acts on deletionTimestamp + and returns before parsing the annotation -- while the heartbeat rebuilt the + payload and erased it. On etcd it was read, but deregistering already stops + new work, so announcing first only added a second mechanism for the same + thing. """ from infera.common.registration import RegistrationClient from infera.common.registration_k8s import K8sRegistrationClient - assert RegistrationClient.announces_draining is True - assert K8sRegistrationClient.announces_draining is False - - # The method exists only where it is meaningful, so a caller cannot - # accidentally write a status the Kubernetes registry will never read. - assert hasattr(RegistrationClient, "announce_draining") - assert not hasattr(K8sRegistrationClient, "announce_draining") + for client in (RegistrationClient, K8sRegistrationClient): + assert not hasattr(client, "announce_draining"), client.__name__ @pytest.mark.asyncio -async def test_a_k8s_heartbeat_is_identity_only_and_so_survives_a_drain(): - """The heartbeat re-asserts the annotation to self-heal, and it must be - safe to keep running for the whole shutdown -- on etcd that is what keeps - the lease alive through the drain. - - That only holds while the annotation carries identity and nothing else. If - it also carried state, a refresh landing mid-drain would overwrite it with - whatever the payload builder defaults to. - """ - from infera.common.registration import build_worker_payload - - cfg = _cfg() - first = build_worker_payload(cfg) - later = build_worker_payload(cfg) - assert first == later, "a refresh must be byte-identical, i.e. carry no state" - assert "status" not in first, "identity only; state belongs to the backend" +async def test_the_k8s_registry_marks_a_condemned_pod_draining(): + """The one remaining producer of DRAINING, and the reason the worker does + not need to announce anything under Kubernetes.""" + import json as _json + + from infera.common.discovery_k8s import WORKER_INFO_ANNOTATION, KubernetesRegistry + + reg = KubernetesRegistry(label_selector="x=y", namespace="ns") + + def _pod(*, deleting: bool): + meta = { + "name": "worker-1", + "annotations": {WORKER_INFO_ANNOTATION: _json.dumps(build_worker_payload(_cfg()))}, + } + if deleting: + meta["deletionTimestamp"] = "2026-08-10T00:00:00Z" + return {"metadata": meta, "status": {"phase": "Running"}} + + reg._handle_pod(_pod(deleting=False), deleted=False) + assert [w.worker_id for w in reg._pool.list_active()] == ["10.0.0.1:8080"] + + reg._handle_pod(_pod(deleting=True), deleted=False) + assert reg._pool.list_active() == [], "a condemned Pod must leave routing" + assert reg._pool.get("10.0.0.1:8080").status is WorkerStatus.DRAINING From 84f0046d91ee7d61fc6332453d15958172096aa1 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Mon, 10 Aug 2026 01:46:18 +0000 Subject: [PATCH 75/88] docs(features): add a graceful shutdown page, scoped to Kubernetes The behaviour was documented only as part of the scaling page, where it read as a general property of the system. It is not: a worker leaving rotation before it is signalled, and staying visible while it finishes, both depend on the orchestrator knowing a Pod is condemned. Nothing outside Kubernetes can tell the router that. The new page states that scope in the summary box and again in an admonition, explains what the feature prevents (a severed generation cannot be retried -- the tokens already sent cannot be un-sent), gives the sequence and the two manifest settings that matter, and describes what a non-Kubernetes deployment gets instead: still a drain, but with deregistration first, since removing the record is what stops new work there. Also brings the scaling page in line. It described announcing DRAINING as the universal first step, which is no longer what either backend does, and its per-stage budget table still billed for that step. The measured scale-down log was left as prose with a note that the run predates the ordering change, rather than being rewritten into output that would not be produced today. Co-authored-by: Cursor Signed-off-by: leiwei12 --- manual/features/graceful_shutdown.md | 116 +++++++++++++++++++++++++++ manual/features/scaling.md | 56 +++++++------ manual/sphinx/_toc.yml.in | 2 + 3 files changed, 148 insertions(+), 26 deletions(-) create mode 100644 manual/features/graceful_shutdown.md diff --git a/manual/features/graceful_shutdown.md b/manual/features/graceful_shutdown.md new file mode 100644 index 00000000..8c8cd5ac --- /dev/null +++ b/manual/features/graceful_shutdown.md @@ -0,0 +1,116 @@ +# Graceful shutdown + +```{admonition} One-pager +:class: tip +**What:** a worker being removed stops receiving new requests immediately, then +finishes the generations it already accepted before the process exits. +**Why:** without it, every request in flight on that worker is severed — +a rolling upgrade or a scale-down turns into a burst of client errors. +**Requires:** a Kubernetes deployment using the default `kubernetes` discovery +backend. See [Outside Kubernetes](#outside-kubernetes) for what other +deployments get instead. +``` + +```{important} +The full behaviour described here — a worker leaving rotation *before* it is +signalled, and staying visible while it finishes — is supported **only on +Kubernetes with the default `kubernetes` discovery backend**. That is not an +implementation gap: it depends on the orchestrator knowing a Pod is condemned, +which nothing outside Kubernetes can tell the router. +``` + +## What it prevents + +A generation can run for tens of seconds. If a worker is stopped while holding +one, the client gets a truncated stream or a connection reset — and there is no +retry that helps, because the tokens already sent cannot be un-sent. + +That makes ordinary operations expensive. Rolling out a new image, scaling down +after a burst, draining a node for maintenance: each replaces workers that are +very likely mid-generation. + +Graceful shutdown separates two things that would otherwise happen at once: + +- **Stop receiving.** The worker leaves the routing candidate list. New requests + go elsewhere from that moment. +- **Stop serving.** The worker keeps working on what it already accepted, and + only then exits. + +The gap between them is the drain. + +## The sequence on Kubernetes + +``` +kubectl delete pod / scale down / rolling update + │ + ├─► Kubernetes marks the Pod as condemned ← under 100 ms + │ the router drops it from routing here + │ + ├─► preStop delay (15 s), still serving what it has + │ + ├─► SIGTERM + │ drain: wait for in-flight generations to finish + │ (bounded by --drain-timeout, default 30 s) + │ + ├─► deregister, stop the engine + │ + └─► [SIGKILL if the grace period expires] +``` + +The important part is the first step. Kubernetes marks a Pod the instant its +deletion is requested — before the `preStop` hook runs, and therefore before the +worker process is signalled at all. The router watches for that mark, so it +stops choosing the worker in well under a second, while the worker itself does +not learn it is leaving for another 15 seconds. + +Without that, the `preStop` delay would work against you: it is meant to give +the router time to react, but if the router only finds out at `SIGTERM`, the +delay is simply 15 more seconds of accepting work that is about to be drained. + +**The worker stays visible while it drains.** Its record is removed at the end, +not the beginning, so `/v1/workers` reports it as draining rather than having it +disappear. A worker that vanishes looks exactly like one that crashed; this way +an operator can see a rollout progressing and how far along it is. + +## What you need to configure + +Nothing, if you deploy through the operator — it injects the `preStop` delay and +sizes `terminationGracePeriodSeconds` to cover the whole sequence. + +For a hand-written manifest, two things matter: + +- **A `preStop` delay.** Without it `SIGTERM` arrives immediately and the drain + starts before the router has necessarily reacted. +- **A `terminationGracePeriodSeconds` that covers the whole sequence**, which is + the `preStop` delay plus `--drain-timeout` plus teardown. Set it too low and + the kubelet sends `SIGKILL` partway through the drain — turning a graceful + shutdown back into an abrupt one, which is the failure this feature exists to + avoid. Raising `--drain-timeout` for long generations without raising the + grace period is the usual way to hit this. + +```{warning} +`discoveryBackend: etcd` is **not supported for in-cluster deployments**, and +the operator refuses it. The combination keeps the `preStop` delay while losing +the early notice that delay exists to provide: the router no longer watches +Pods, so nothing sees the Pod being condemned, and the only remaining signal +arrives after `SIGTERM` — once the delay has already elapsed. For its whole +duration the router keeps handing new work to a Pod on its way out, which is +worse than either backend on its own. +``` + +## Outside Kubernetes + +Deployments on bare metal or under a container runtime use an external etcd for +discovery, and there is no orchestrator to say a worker is leaving. A record is +either present or absent; nothing observes that a process is on its way out. + +Shutdown there still drains, but in the other order: the worker removes its +registration first — which is what stops new requests arriving — and then waits +for its in-flight generations. In-flight work is still finished rather than cut. +What is lost is the two properties that depend on the orchestrator: + +- **No early notice.** Routing stops when the process is signalled, not before. +- **No visible draining.** The worker disappears from `/v1/workers` for the + duration of the drain rather than being reported as finishing. + +If you are running in Kubernetes, use the default backend and you get both. diff --git a/manual/features/scaling.md b/manual/features/scaling.md index 7080ce36..24b12135 100644 --- a/manual/features/scaling.md +++ b/manual/features/scaling.md @@ -18,9 +18,10 @@ leaving ──► stop being routed to ──► drain in-flight ──► dereg ``` What triggers "stop being routed to" depends on the backend: under Kubernetes -the orchestrator marks the Pod before the process is even signalled, on etcd the -worker announces it after `SIGTERM`. See [Who says the worker is -leaving](#who-says-the-worker-is-leaving). +the orchestrator marks the Pod before the process is even signalled, elsewhere +removing the record is what does it. See [Who says the worker is +leaving](#who-says-the-worker-is-leaving), and +[Graceful shutdown](graceful_shutdown.md) for the feature as a whole. That shape is why scale-up and scale-down have very different costs. Scale-up is bounded by **model load**, which is minutes. Scale-down is bounded by the @@ -61,9 +62,10 @@ The worker then, in this order: 1. **Stops being routed to.** The router filters draining workers out immediately, so no new work arrives. The record stays, so the worker remains visible in `/v1/workers` — a worker that vanishes looks exactly like one that - crashed. On etcd the worker announces `DRAINING` here; under Kubernetes this - already happened when the Pod was condemned, before it was signalled at all - (see [Who says the worker is leaving](#who-says-the-worker-is-leaving)). + crashed. Under Kubernetes this already happened when the Pod was condemned, + before the process was signalled at all; elsewhere the worker deregisters + here, which stops new work but also removes it from the listing (see [Who + says the worker is leaving](#who-says-the-worker-is-leaving)). 2. **Drains.** On the NATS transport infera tracks in-flight requests directly. On HTTP the router talks straight to the engine, so infera asks the engine instead, polling its `/metrics` until running, queued, and PD-handoff queues @@ -75,7 +77,7 @@ drain go to other workers. **Two different timings, easily conflated.** A worker stops *receiving* new requests within a second of the shutdown starting — the router's watch picking -up either the announcement or the `deletionTimestamp` — and it is the number +up either the `deletionTimestamp` or the record's removal — and it is the number that decides whether traffic is still being sent somewhere that is about to die. How long the *process* then lives is a separate and much larger number, set by the longest generation it was already serving. Measured: under a second to stop receiving, 38 s until the record @@ -109,7 +111,7 @@ Measured, same fake worker, same generation: request(s)` — it knows the count — the 300-chunk generation completed in full, and the worker deregistered **21.3 s** later, which is just the remaining generation time with no overhead. -- **NATS, nothing in flight**: announce → deregister in **3 ms**. +- **NATS, nothing in flight**: leaving rotation to exit in **3 ms**. - **HTTP with a real engine, nothing in flight**: at least the **6 s** settle window, because a single zero reading cannot be told apart from a gauge that has not refreshed yet. @@ -236,18 +238,18 @@ landing mid-drain would overwrite it with a payload that omits the status, which parses as `ACTIVE`, and the worker would be handed new work it is about to refuse. -**etcd: the worker says so.** There is no orchestrator here. A record is either -present with an unexpired lease or it is gone — nothing observes that the -process is on its way out. Without the worker announcing `DRAINING` there is no -way to express "still finishing what I have, send me nothing new", so on this -backend the record does carry a status and the worker rewrites it before -draining. +**etcd: removing the record says so.** There is no orchestrator here. A record +is either present with an unexpired lease or it is gone — nothing observes that +the process is on its way out, and there is no third state to put it in. So a +shutdown deregisters *first*, which is what stops new requests arriving, and +drains after. In-flight generations are still finished rather than cut; what is +not available is the worker remaining visible while it does so. -The heartbeat runs until the record is deleted on both backends, which matters -most here: it is what renews the lease. Stopping it when the drain starts would -leave at most one TTL of life, and a drain longer than the remainder has etcd -collect the key partway through — the worker vanishing rather than visibly -finishing, which is the outcome announcing `DRAINING` exists to prevent. +That ordering is the one real difference. Under Kubernetes routing has already +stopped by the time the process is signalled, so the record can be kept until +the end and the drain is observable; on etcd the record is the only thing +keeping the worker in rotation, so it has to go before the drain rather than +after. ```{warning} `discoveryBackend: etcd` **is not supported for in-cluster deployments** and the @@ -267,7 +269,6 @@ Every stage is individually bounded: | Stage | Bound | Set by | |---|---|---| | `preStop` | 15 s | operator | -| announce `DRAINING` | 10 s | registration HTTP client timeout | | drain | `--drain-timeout` (default 30 s) | flag | | deregister | 10 s | registration HTTP client timeout | | `engine.stop()` | 30 s | `SIGTERM` to the engine's process group, then `SIGKILL` | @@ -396,13 +397,16 @@ Two runs, both with traffic flowing throughout: **Scale up then down.** Two instances, continuous traffic, a third added and then one removed. **260 requests, 0 failures**, including in the 5-second -windows around each transition. The removed instance's log shows the intended -sequence: +windows around each transition. The removed instance left rotation, drained the +one generation it was holding (`engine idle for 6s, 1 request(s) completed`), +and only then exited. -``` -worker 127.0.0.1:20001 announced DRAINING -drain: engine idle for 6s, 1 request(s) completed -deregistered worker 127.0.0.1:20001 (lease revoked) +```{note} +This run predates the change that made the shutdown order backend-specific, so +its logs show the worker announcing `DRAINING` before draining. On etcd the two +steps are now the other way round — deregister, then drain — which is what stops +new work arriving on a backend where nothing else can. The request counts are +unaffected: both orderings stop new work before waiting on in-flight work. ``` **PD scaling, measured.** A 1P1D fake fleet grown to 2P2D and shrunk back under diff --git a/manual/sphinx/_toc.yml.in b/manual/sphinx/_toc.yml.in index 84cde5c5..a4b79ab2 100644 --- a/manual/sphinx/_toc.yml.in +++ b/manual/sphinx/_toc.yml.in @@ -47,6 +47,8 @@ subtrees: title: Routing and transport - file: features/scaling.md title: Scaling a fleet + - file: features/graceful_shutdown.md + title: Graceful shutdown - caption: Serving entries: From 9f05eb2d4a64fd88873f70c88a244e8537a39c97 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Mon, 10 Aug 2026 01:51:07 +0000 Subject: [PATCH 76/88] docs(features): trim the graceful shutdown page to the shape of the feature The page had grown into a walkthrough -- a full sequence diagram, the manifest settings, per-stage bounds -- most of which the scaling page already covers and covers with measurements. A feature one-pager should answer what it does, why it matters, and whether it applies to you. Now: the two things the shutdown separates, why the record outlives the drain, and a pointer to the scaling page for timings and hand-written manifests. The Kubernetes-only scope stays stated twice, in the summary box and its own admonition, since that is the part a reader most needs before designing around it. Co-authored-by: Cursor Signed-off-by: leiwei12 --- manual/features/graceful_shutdown.md | 131 +++++++-------------------- 1 file changed, 32 insertions(+), 99 deletions(-) diff --git a/manual/features/graceful_shutdown.md b/manual/features/graceful_shutdown.md index 8c8cd5ac..37726312 100644 --- a/manual/features/graceful_shutdown.md +++ b/manual/features/graceful_shutdown.md @@ -4,113 +4,46 @@ :class: tip **What:** a worker being removed stops receiving new requests immediately, then finishes the generations it already accepted before the process exits. -**Why:** without it, every request in flight on that worker is severed — -a rolling upgrade or a scale-down turns into a burst of client errors. -**Requires:** a Kubernetes deployment using the default `kubernetes` discovery -backend. See [Outside Kubernetes](#outside-kubernetes) for what other -deployments get instead. +**Why:** a severed generation cannot be retried — the tokens already streamed +cannot be un-sent — so without this, every rolling upgrade or scale-down +produces a burst of client errors. **Requires:** Kubernetes, with the default +`kubernetes` discovery backend. ``` ```{important} -The full behaviour described here — a worker leaving rotation *before* it is -signalled, and staying visible while it finishes — is supported **only on -Kubernetes with the default `kubernetes` discovery backend**. That is not an -implementation gap: it depends on the orchestrator knowing a Pod is condemned, -which nothing outside Kubernetes can tell the router. +Supported **only on Kubernetes with the default `kubernetes` discovery +backend**. This is not an implementation gap: it relies on the orchestrator +knowing a Pod is being removed, which nothing outside Kubernetes can tell the +router. `discoveryBackend: etcd` is rejected by the operator for in-cluster +deployments. ``` -## What it prevents +## What happens -A generation can run for tens of seconds. If a worker is stopped while holding -one, the client gets a truncated stream or a connection reset — and there is no -retry that helps, because the tokens already sent cannot be un-sent. +Removing a worker — a rolling update, a scale-down, draining a node — separates +two things that would otherwise happen at once: -That makes ordinary operations expensive. Rolling out a new image, scaling down -after a burst, draining a node for maintenance: each replaces workers that are -very likely mid-generation. +1. **It stops receiving.** Kubernetes marks the Pod the moment its removal is + requested, which is *before* the worker process is signalled. The router sees + that mark and stops choosing the worker within milliseconds, so new requests + go elsewhere while it is still running. +2. **It keeps serving.** The worker finishes the generations it already + accepted, bounded by `--drain-timeout`, and only then deregisters and exits. -Graceful shutdown separates two things that would otherwise happen at once: +Because the record is removed at the end rather than the beginning, the worker +stays visible in `/v1/workers` while it drains — an operator can see a rollout +progressing instead of workers appearing to crash. -- **Stop receiving.** The worker leaves the routing candidate list. New requests - go elsewhere from that moment. -- **Stop serving.** The worker keeps working on what it already accepted, and - only then exits. +Deploying through the operator needs no configuration: it injects the `preStop` +delay and sizes the termination grace period to cover the whole sequence. For +hand-written manifests and the per-stage timings, see +[Scaling a fleet](scaling.md). -The gap between them is the drain. +## Elsewhere -## The sequence on Kubernetes - -``` -kubectl delete pod / scale down / rolling update - │ - ├─► Kubernetes marks the Pod as condemned ← under 100 ms - │ the router drops it from routing here - │ - ├─► preStop delay (15 s), still serving what it has - │ - ├─► SIGTERM - │ drain: wait for in-flight generations to finish - │ (bounded by --drain-timeout, default 30 s) - │ - ├─► deregister, stop the engine - │ - └─► [SIGKILL if the grace period expires] -``` - -The important part is the first step. Kubernetes marks a Pod the instant its -deletion is requested — before the `preStop` hook runs, and therefore before the -worker process is signalled at all. The router watches for that mark, so it -stops choosing the worker in well under a second, while the worker itself does -not learn it is leaving for another 15 seconds. - -Without that, the `preStop` delay would work against you: it is meant to give -the router time to react, but if the router only finds out at `SIGTERM`, the -delay is simply 15 more seconds of accepting work that is about to be drained. - -**The worker stays visible while it drains.** Its record is removed at the end, -not the beginning, so `/v1/workers` reports it as draining rather than having it -disappear. A worker that vanishes looks exactly like one that crashed; this way -an operator can see a rollout progressing and how far along it is. - -## What you need to configure - -Nothing, if you deploy through the operator — it injects the `preStop` delay and -sizes `terminationGracePeriodSeconds` to cover the whole sequence. - -For a hand-written manifest, two things matter: - -- **A `preStop` delay.** Without it `SIGTERM` arrives immediately and the drain - starts before the router has necessarily reacted. -- **A `terminationGracePeriodSeconds` that covers the whole sequence**, which is - the `preStop` delay plus `--drain-timeout` plus teardown. Set it too low and - the kubelet sends `SIGKILL` partway through the drain — turning a graceful - shutdown back into an abrupt one, which is the failure this feature exists to - avoid. Raising `--drain-timeout` for long generations without raising the - grace period is the usual way to hit this. - -```{warning} -`discoveryBackend: etcd` is **not supported for in-cluster deployments**, and -the operator refuses it. The combination keeps the `preStop` delay while losing -the early notice that delay exists to provide: the router no longer watches -Pods, so nothing sees the Pod being condemned, and the only remaining signal -arrives after `SIGTERM` — once the delay has already elapsed. For its whole -duration the router keeps handing new work to a Pod on its way out, which is -worse than either backend on its own. -``` - -## Outside Kubernetes - -Deployments on bare metal or under a container runtime use an external etcd for -discovery, and there is no orchestrator to say a worker is leaving. A record is -either present or absent; nothing observes that a process is on its way out. - -Shutdown there still drains, but in the other order: the worker removes its -registration first — which is what stops new requests arriving — and then waits -for its in-flight generations. In-flight work is still finished rather than cut. -What is lost is the two properties that depend on the orchestrator: - -- **No early notice.** Routing stops when the process is signalled, not before. -- **No visible draining.** The worker disappears from `/v1/workers` for the - duration of the drain rather than being reported as finishing. - -If you are running in Kubernetes, use the default backend and you get both. +Deployments outside Kubernetes use an external etcd for discovery, where a +worker record is simply present or absent and nothing observes that a process is +leaving. Shutdown there still finishes in-flight work, but in the other order: +the worker deregisters first — which is what stops new requests arriving — and +drains after. In-flight generations are not cut; what is unavailable is the +early notice and the visible draining above. From f679f7bac584b5c58cc3cb96c8e51975ac5977da Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Mon, 10 Aug 2026 01:56:41 +0000 Subject: [PATCH 77/88] fix(router): close three gaps in how PD legs are scored **The NATS transport scored nothing.** Those paths never touch the HTTP client, so scoring placed at HTTP response sites missed them entirely. A decode worker failing every request over NATS was invisible to the breaker, and -- worse -- one already open could never close again, because nothing recorded the success that ends a half-open state: it stayed there for the life of the process, throttled to one request per probe window. All three NATS legs now score from their own reply. **A decode outage tripped the prefill worker.** `asyncio.gather` raises whichever leg failed without saying which, and the handler blamed a fixed one. Measured: three refused decode connections opened the healthy prefill's breaker while the broken decode never acquired an entry -- the good worker evicted, the bad one still taking every request. Gathering with `return_exceptions` lets each leg be attributed to the worker that produced it. **A 200 header was scored as recovery.** Headers are the right moment to record a failure -- pre-first-byte, as the breaker requires -- but not a success, because `record_success` resets the failure count and closes an open breaker on nothing more than the request having been accepted. The worker this class exists to catch is the one that accepts and then produces nothing, so that profile was laundering itself: five broken streams erased two real failures. A 2xx header is neutral now, which still frees the probe slot, and the success is recorded once bytes are actually flowing. Co-authored-by: Cursor Signed-off-by: leiwei12 --- infera/router/disagg.py | 90 ++++++++++++--- tests/unit/router/test_disagg_breaker.py | 139 ++++++++++++++++++++++- 2 files changed, 209 insertions(+), 20 deletions(-) diff --git a/infera/router/disagg.py b/infera/router/disagg.py index 9f1dd4d7..80f846d9 100644 --- a/infera/router/disagg.py +++ b/infera/router/disagg.py @@ -172,6 +172,22 @@ async def dispatch_direct( obs, RouteTarget(p), RouteTarget(d), [], [], body, stream, path ) + def _score_leg_headers(self, worker_id: str, status: int) -> None: + """Score a streaming leg from its response headers. + + Headers are the pre-first-byte moment, so a failure recorded here is + exactly what the breaker wants. Success is not: a 200 header says the + request was accepted, nothing more, and the worker this class exists to + catch is the one that accepts a request and then produces nothing. That + profile would otherwise be scored as recovery -- resetting the failure + count and closing an open breaker -- so a 2xx is neutral here and the + success is recorded once the stream has actually delivered something. + """ + if status < 400: + self.breaker.record_neutral(worker_id) + else: + self._score_leg(worker_id, status) + def _score_leg(self, worker_id: str, status: int) -> None: """Record one PD leg's HTTP outcome against the worker that produced it. @@ -362,19 +378,34 @@ async def _post(url, leg, worker_id, leg_body, leg_headers): with metrics.track_pd_leg(leg=leg, worker_id=worker_id): return await self._client.post(url, json=leg_body, headers=leg_headers) - try: - p_resp, d_resp = await asyncio.gather( - _post(p_url, "prefill", p.worker_id, p_body, p_headers), - _post(d_url, "decode", d.worker_id, d_body, d_headers), - ) - except httpx.HTTPError as exc: + # Gathered with return_exceptions so a failure can be attributed to + # the leg that produced it. Letting gather raise surfaces whichever + # one failed first with no way to tell which that was, and blaming a + # fixed leg means a decode outage evicts the healthy prefill worker + # while the broken decode is never scored at all. + p_resp, d_resp = await asyncio.gather( + _post(p_url, "prefill", p.worker_id, p_body, p_headers), + _post(d_url, "decode", d.worker_id, d_body, d_headers), + return_exceptions=True, + ) + failed = None + for worker_id, leg, result in ( + (p.worker_id, "prefill", p_resp), + (d.worker_id, "decode", d_resp), + ): + if isinstance(result, BaseException): + self.breaker.record_failure(worker_id) + if failed is None: + failed = (leg, result) + else: + self._score_leg(worker_id, result.status_code) + if failed is not None: + leg, exc = failed + if not isinstance(exc, httpx.HTTPError): + raise exc obs["outcome"] = "502" metrics.pd_bootstrap_failures_total.labels(reason="worker_unreachable").inc() - self.breaker.record_failure(p.worker_id) - return _sanitized_error("PD request failed", exc, status_code=502) - - self._score_leg(p.worker_id, p_resp.status_code) - self._score_leg(d.worker_id, d_resp.status_code) + return _sanitized_error(f"PD {leg} leg failed", exc, status_code=502) if p_resp.status_code >= 400: logger.warning( "prefill worker %s returned %d (decode may fail)", @@ -407,16 +438,22 @@ def _start_prefill_drain_nats(self, p, p_payload): like the HTTP path. Strong ref guards against GC mid-flight.""" async def _drain(): + # Scored here for the same reason the HTTP legs are scored at their + # own responses: this transport never touches the HTTP client, so + # nothing else observes how this worker did. try: async for kind, _st, data in self.nats_client.stream(p.worker_id, p_payload): if kind == TYPE_ERROR: logger.warning("prefill leg (nats) %s failed: %s", p.worker_id, data[:200]) metrics.pd_bootstrap_failures_total.labels(reason="prefill_exception").inc() + self.breaker.record_failure(p.worker_id) return if kind == TYPE_DONE: + self.breaker.record_success(p.worker_id) return except Exception as exc: logger.warning("prefill nats drain %s failed: %s", p.worker_id, exc) + self.breaker.record_failure(p.worker_id) task = asyncio.create_task(_drain(), name="nats-prefill-drain") self._pending_prefill_tasks.add(task) @@ -463,6 +500,7 @@ async def _concurrent_nats( elif kind == TYPE_ERROR: # st carries 504 on inactivity timeout; worker errors -> 502. code = st or 502 + self._score_leg(d.worker_id, code) obs["outcome"] = str(code) return JSONResponse( content={ @@ -486,6 +524,7 @@ async def _concurrent_nats( }, status_code=502, ) + self._score_leg(d.worker_id, status) obs["outcome"] = "ok" if status < 400 else f"{status // 100}xx" return JSONResponse(content=payload, status_code=status) finally: @@ -499,14 +538,22 @@ async def _concurrent_nats( async def _stream_dual_nats(self, p_target, p_blocks, d_target, d_blocks, d_payload, p_task): """Stream decode's reply over NATS while prefill drains in background.""" d = d_target.worker + served = False try: async for kind, _st, data in self.nats_client.stream(d.worker_id, d_payload): if kind == TYPE_DATA: if data: + if not served: + # Bytes are flowing, so this worker is doing the + # work; an accepted request alone would not show it. + self.breaker.record_success(d.worker_id) + served = True yield data elif kind == TYPE_ERROR: logger.warning("decode (nats) %s stream failed: %s", d.worker_id, data[:200]) metrics.pd_bootstrap_failures_total.labels(reason="decode_stream_broken").inc() + if not served: + self.breaker.record_failure(d.worker_id) yield ( f'data: {{"error":"decode {d.worker_id} nats stream failed"}}\n\n' ).encode() @@ -718,10 +765,7 @@ async def _stream_decode_only( yield f"data: {err}\n\n".encode() return - # Headers are in hand, which is the pre-first-byte moment: what the - # decode worker did with the request is now known and nothing has - # reached the client yet. - self._score_leg(d_target.worker.worker_id, d_resp.status_code) + self._score_leg_headers(d_target.worker.worker_id, d_resp.status_code) if d_resp.status_code >= 400: try: body_bytes = await d_resp.aread() @@ -750,8 +794,14 @@ async def _stream_decode_only( _DONE_NEEDLE = b"data: [DONE]" _TAIL_KEEP = len(_DONE_NEEDLE) - 1 tail = b"" + served = False try: async for chunk in d_resp.aiter_raw(): + if not served and chunk: + # Bytes are flowing, so the worker is doing the work -- + # which the headers alone did not establish. + self.breaker.record_success(d_target.worker.worker_id) + served = True if not done_seen: window = tail + chunk if _DONE_NEEDLE in window: @@ -878,9 +928,7 @@ async def _stream_dual( yield f"data: {err}\n\n".encode() return - # Headers are in hand: the decode leg's outcome is known and - # nothing has reached the client yet. - self._score_leg(d_target.worker.worker_id, d_resp.status_code) + self._score_leg_headers(d_target.worker.worker_id, d_resp.status_code) if d_resp.status_code >= 400: # Engine accepted but rejected; surface its body verbatim. try: @@ -912,7 +960,13 @@ async def _stream_dual( _DONE_NEEDLE = b"data: [DONE]" _TAIL_KEEP = len(_DONE_NEEDLE) - 1 tail = b"" + served = False async for chunk in d_resp.aiter_raw(): + if not served and chunk: + # Bytes are flowing, so the worker is doing the work -- + # which the headers alone did not establish. + self.breaker.record_success(d_target.worker.worker_id) + served = True if not done_seen: window = tail + chunk if _DONE_NEEDLE in window: diff --git a/tests/unit/router/test_disagg_breaker.py b/tests/unit/router/test_disagg_breaker.py index 6e9dd681..9ac92891 100644 --- a/tests/unit/router/test_disagg_breaker.py +++ b/tests/unit/router/test_disagg_breaker.py @@ -15,9 +15,12 @@ from __future__ import annotations +import json + import httpx import pytest +from infera.common.nats_request import TYPE_DATA, TYPE_DONE, TYPE_ERROR from infera.common.worker_pool import DisaggMode, EngineType, WorkerInfo from infera.router.disagg import DisaggRouter from infera.router.policy.target import RouteTarget @@ -138,7 +141,7 @@ def list_active(self, model=None, mode=None): return list(self._by_mode.get(mode, [])) -def _pd_worker(wid, mode): +def _pd_worker(wid, mode, transport="http"): meta = {"protocol": "sglang-bootstrap"} if mode is DisaggMode.PREFILL: meta["params"] = {"bootstrap_addr": f"{wid}:9000"} @@ -147,12 +150,25 @@ def _pd_worker(wid, mode): url=f"http://{wid}", model_name="m", engine=EngineType.SGLANG, - request_transport="http", + request_transport=transport, disagg_mode=mode, disagg_meta=meta, ) +def _nats_router(*, fail_decode): + """A PD pair that both registered for the NATS transport, which is what + selects the NATS dispatch path.""" + return DisaggRouter( + _RolePool( + _pd_worker("p1", DisaggMode.PREFILL, transport="nats"), + _pd_worker("d1", DisaggMode.DECODE, transport="nats"), + ), + _FakePolicy(), + nats_client=_FakeNatsPD(fail_decode=fail_decode), + ) + + def _ok_router(): """A PD router whose every leg answers 200.""" r = DisaggRouter( @@ -268,3 +284,122 @@ async def test_a_tripped_pd_worker_recovers_after_a_good_probe(): assert r.breaker.state_of("d1").value == "closed", "a good probe must close it" await r.aclose() + + +@pytest.mark.asyncio +async def test_a_decode_outage_does_not_trip_the_prefill_worker(): + """asyncio.gather raises whichever leg failed, without saying which. Blaming + a fixed one means a decode that refuses connections evicts the healthy + prefill worker from rotation while the broken decode is never scored at + all -- the exact inversion of what the breaker is for.""" + + def _only_decode_is_down(request): + if "d1" in str(request.url): + raise httpx.ConnectError("refused", request=request) + return httpx.Response(200, json={"id": "x"}) + + r = DisaggRouter( + _RolePool(_pd_worker("p1", DisaggMode.PREFILL), _pd_worker("d1", DisaggMode.DECODE)), + _FakePolicy(), + ) + r._client = httpx.AsyncClient(transport=httpx.MockTransport(_only_decode_is_down)) + + for _ in range(3): + await r.dispatch({"model": "m"}, stream=False) + + assert r.breaker.state_of("d1").value == "open", "the leg that refused must be the one scored" + assert r.breaker.state_of("p1").value == "closed", "the healthy leg must not be evicted" + await r.aclose() + + +@pytest.mark.asyncio +async def test_a_stream_that_dies_after_its_headers_is_not_a_success(): + """A 200 header only means the request was accepted. Treating it as + recovery resets the failure count, so a decode worker that answers 200 and + then sends nothing -- which is precisely the "healthy to the platform, + broken for inference" profile the breaker exists for -- can never trip, and + erases real failures on its way.""" + r = DisaggRouter( + _RolePool(_pd_worker("p1", DisaggMode.PREFILL), _pd_worker("d1", DisaggMode.DECODE)), + _FakePolicy(), + ) + + def _headers_then_nothing(request): + if "d1" in str(request.url): + # 200, then the body raises as soon as it is read. + return httpx.Response(200, stream=_DyingStream()) + return httpx.Response(200, json={"id": "x"}) + + r._client = httpx.AsyncClient(transport=httpx.MockTransport(_headers_then_nothing)) + + for _ in range(2): + r.breaker.record_failure("d1") + before = r.breaker.snapshot()["d1"]["consecutive_failures"] + assert before == 2 + + resp = await r.dispatch({"model": "m"}, stream=True) + await _drain(resp.body_iterator) + + after = r.breaker.snapshot()["d1"]["consecutive_failures"] + assert after >= before, ( + f"consecutive_failures went {before} -> {after}: a stream that produced " + "no output was scored as evidence of health" + ) + await r.aclose() + + +class _DyingStream(httpx.AsyncByteStream): + """Headers arrive, then the body fails -- no bytes ever reach the client.""" + + async def __aiter__(self): + raise httpx.ReadError("connection died after headers") + yield b"" # unreachable; makes this an async generator + + +@pytest.mark.asyncio +async def test_the_nats_transport_scores_its_legs_too(): + """The NATS paths do not use the HTTP client, so scoring placed at HTTP + response sites misses them entirely. A decode worker failing every request + over NATS would be invisible to the breaker, and -- worse -- one already + open could never be closed again, since nothing would ever record the + success that ends its half-open state.""" + r = _nats_router(fail_decode=True) + + for _ in range(3): + await r.dispatch({"model": "m"}, stream=False) + + assert r.breaker.state_of("d1").value == "open", "a failing NATS decode leg must trip" + await r.aclose() + + +@pytest.mark.asyncio +async def test_a_good_nats_probe_closes_the_breaker(): + r = _nats_router(fail_decode=False) + for _ in range(3): + r.breaker.record_failure("d1") + r.breaker._entries["d1"].opens_until = 0.0 + + await r.dispatch({"model": "m"}, stream=False) + + assert r.breaker.state_of("d1").value == "closed", ( + "a NATS worker that answers cleanly must be able to recover; otherwise " + "it stays half-open forever, throttled to one request per probe window" + ) + await r.aclose() + + +class _FakeNatsPD: + """Scripted NATS transport: decode either answers or errors, prefill is fine.""" + + def __init__(self, *, fail_decode: bool): + self.fail_decode = fail_decode + + async def admit(self, worker_id): + return True + + async def stream(self, worker_id, payload): + if worker_id == "d1" and self.fail_decode: + yield (TYPE_ERROR, 502, b"decode exploded") + return + yield (TYPE_DATA, None, json.dumps({"id": "x"}).encode()) + yield (TYPE_DONE, 200, b"") From 5f913657bff04c844b4d58640266c23dd6dd9a8b Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Mon, 10 Aug 2026 01:59:21 +0000 Subject: [PATCH 78/88] fix(operator): let a removed spec field be removed from the child Merging the operator's fields onto the live object stopped the write loop, but it also made deletion impossible: an absent field is indistinguishable from one the operator does not manage, so the previous value survived. That bites on fields the builders emit conditionally. HTTPRoute.spec.hostnames is written only when the CR lists any, so deleting them from the CR left the route matching hostnames the user had removed -- silently, since the write succeeds and reports nothing wrong. Each kind now declares the top-level spec fields it owns outright. One of those missing from the desired object has been removed rather than left unmanaged, so it is cleared; everything else, which is essentially the API server's defaulting, is still left alone. Pruning stops at the top level because below it the desired fields and the server's defaults are interleaved with no way to tell them apart -- and that is exactly where the defaults live that this merge was introduced to protect. Co-authored-by: Cursor Signed-off-by: leiwei12 --- .../controller/apply_idempotence_test.go | 104 ++++++++++++++++++ .../controller/inferadeployment_controller.go | 41 ++++++- 2 files changed, 142 insertions(+), 3 deletions(-) diff --git a/deploy/operator/internal/controller/apply_idempotence_test.go b/deploy/operator/internal/controller/apply_idempotence_test.go index dcafdb25..3910a535 100644 --- a/deploy/operator/internal/controller/apply_idempotence_test.go +++ b/deploy/operator/internal/controller/apply_idempotence_test.go @@ -137,3 +137,107 @@ func TestApplyStillPushesAChangedField(t *testing.T) { t.Fatalf("replicas = %d, want 5 -- scaling through the CR did not land", v) } } + +// A field the builder emits only when the CR asks for it -- HTTPRoute's +// hostnames -- has to disappear from the child when it disappears from the CR. +// Merging alone cannot do that: an absent field looks the same as one the +// operator does not manage, so the old value would survive and the route would +// keep matching a host the user deleted. +func TestRemovingAConditionalFieldClearsItOnTheChild(t *testing.T) { + s := testScheme(t) + c := fake.NewClientBuilder().WithScheme(s).Build() + r := &InferaDeploymentReconciler{Client: c, Scheme: s} + ctx := context.Background() + + idep := &inferav1alpha1.InferaDeployment{} + idep.Name = "qwen" + idep.Namespace = "default" + idep.UID = "uid-1" + + route := func(hostnames []any) *unstructured.Unstructured { + u := &unstructured.Unstructured{} + u.SetAPIVersion(httpRouteAPIVersion) + u.SetKind(httpRouteKind) + u.SetName("qwen-route") + u.SetNamespace("default") + spec := map[string]any{"rules": []any{}} + if len(hostnames) > 0 { + spec["hostnames"] = hostnames + } + _ = unstructured.SetNestedMap(u.Object, spec, "spec") + return u + } + + if err := r.applyUnstructured(ctx, idep, route([]any{"a.example.com"})); err != nil { + t.Fatalf("create with hostnames: %v", err) + } + if err := r.applyUnstructured(ctx, idep, route(nil)); err != nil { + t.Fatalf("reapply without hostnames: %v", err) + } + + got := &unstructured.Unstructured{} + got.SetAPIVersion(httpRouteAPIVersion) + got.SetKind(httpRouteKind) + if err := c.Get(ctx, types.NamespacedName{Name: "qwen-route", Namespace: "default"}, got); err != nil { + t.Fatalf("get: %v", err) + } + if v, ok, _ := unstructured.NestedSlice(got.Object, "spec", "hostnames"); ok { + t.Fatalf("hostnames still %v after removal from the CR; the route keeps "+ + "matching a host the user deleted", v) + } +} + +// Pruning owned fields must not start pruning the server's defaults again -- +// that is the write loop this merge exists to stop. +func TestPruningLeavesServerDefaultsAlone(t *testing.T) { + s := testScheme(t) + c := fake.NewClientBuilder().WithScheme(s).Build() + r := &InferaDeploymentReconciler{Client: c, Scheme: s} + ctx := context.Background() + + idep := &inferav1alpha1.InferaDeployment{} + idep.Name = "qwen" + idep.Namespace = "default" + idep.UID = "uid-1" + + desired := func() *unstructured.Unstructured { + u := &unstructured.Unstructured{} + u.SetGroupVersionKind(lwsGVK()) + u.SetName("qwen-worker") + u.SetNamespace("default") + _ = unstructured.SetNestedField(u.Object, int64(2), "spec", "replicas") + _ = unstructured.SetNestedField(u.Object, int64(2), "spec", "leaderWorkerTemplate", "size") + return u + } + + if err := r.applyUnstructured(ctx, idep, desired()); err != nil { + t.Fatalf("create: %v", err) + } + + live := &unstructured.Unstructured{} + live.SetGroupVersionKind(lwsGVK()) + if err := c.Get(ctx, types.NamespacedName{Name: "qwen-worker", Namespace: "default"}, live); err != nil { + t.Fatalf("get: %v", err) + } + _ = unstructured.SetNestedField(live.Object, "LeaderCreated", "spec", "startupPolicy") + if err := c.Update(ctx, live); err != nil { + t.Fatalf("apply defaults: %v", err) + } + before := live.GetResourceVersion() + + if err := r.applyUnstructured(ctx, idep, desired()); err != nil { + t.Fatalf("second apply: %v", err) + } + + after := &unstructured.Unstructured{} + after.SetGroupVersionKind(lwsGVK()) + if err := c.Get(ctx, types.NamespacedName{Name: "qwen-worker", Namespace: "default"}, after); err != nil { + t.Fatalf("get: %v", err) + } + if v, ok, _ := unstructured.NestedString(after.Object, "spec", "startupPolicy"); !ok || v == "" { + t.Error("a server default the operator does not set was pruned") + } + if got := after.GetResourceVersion(); got != before { + t.Errorf("reconcile rewrote an unchanged object (%s -> %s)", before, got) + } +} diff --git a/deploy/operator/internal/controller/inferadeployment_controller.go b/deploy/operator/internal/controller/inferadeployment_controller.go index 9157613f..80c12709 100644 --- a/deploy/operator/internal/controller/inferadeployment_controller.go +++ b/deploy/operator/internal/controller/inferadeployment_controller.go @@ -207,13 +207,36 @@ func (r *InferaDeploymentReconciler) applyUnstructured(ctx context.Context, idep if current == nil { current = map[string]any{} } - _ = unstructured.SetNestedMap(existing.Object, mergeSpec(current, spec), "spec") + merged := mergeSpec(current, spec, ownedSpecFields(desired.GetKind())) + _ = unstructured.SetNestedMap(existing.Object, merged, "spec") existing.SetLabels(desired.GetLabels()) return controllerutil.SetControllerReference(idep, existing, r.Scheme) }) return err } +// ownedSpecFields lists the top-level spec fields the operator owns outright +// for a kind: it decides their entire contents, so one absent from the desired +// object has been removed and must be cleared rather than kept. +// +// The distinction matters for fields the builders emit conditionally. +// HTTPRoute.spec.hostnames is only written when the CR lists any, so without +// this a user who deletes their hostnames would keep matching them forever -- +// the merge below would see nothing to overlay and leave the old value in +// place. Fields not listed here belong to someone else, almost always the API +// server's defaulting, and are left untouched. +func ownedSpecFields(kind string) map[string]bool { + switch kind { + case httpRouteKind: + return map[string]bool{"parentRefs": true, "rules": true, "hostnames": true} + case inferencePoolKind: + return map[string]bool{"targetPorts": true, "selector": true, "endpointPickerRef": true} + case lwsKind: + return map[string]bool{"replicas": true, "leaderWorkerTemplate": true} + } + return nil +} + // mergeSpec overlays the fields the operator sets onto what is already there, // leaving anything it does not mention alone. // @@ -229,7 +252,19 @@ func (r *InferaDeploymentReconciler) applyUnstructured(ctx context.Context, idep // Nested maps merge; anything else replaces. Lists are owned outright -- a // container list merged element-wise would be neither what was asked for nor // what was there. -func mergeSpec(into, from map[string]any) map[string]any { +// +// `owned` names the top-level fields the operator decides entirely. One of +// those missing from `from` has been removed rather than left unmanaged, so it +// is deleted; that is what lets a conditionally-emitted field like +// HTTPRoute's hostnames be taken away again. Nested levels are not pruned: +// below the top level the desired object and the server's defaults are +// interleaved, with no way to tell them apart. +func mergeSpec(into, from map[string]any, owned map[string]bool) map[string]any { + for k := range owned { + if _, still := from[k]; !still { + delete(into, k) + } + } for k, v := range from { sub, isMap := v.(map[string]any) if !isMap { @@ -241,7 +276,7 @@ func mergeSpec(into, from map[string]any) map[string]any { into[k] = v continue } - into[k] = mergeSpec(existing, sub) + into[k] = mergeSpec(existing, sub, nil) } return into } From a086e9b4eefc90727963cf3bd8f35884511da7f1 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Mon, 10 Aug 2026 02:08:59 +0000 Subject: [PATCH 79/88] fix: harden the breaker's edges -- poisoned lock, trip count, label escaping **A poisoned lock no longer takes the data plane with it.** Every request passes through the breaker, so unwrapping a poisoned mutex turned one panic into a panic on every subsequent request: a local fault escalated to total unavailability, which is precisely what this class exists to prevent. The guarded state is a map of counters, so recovering the guard risks one stale entry, never an unusable process. **`trips_total` counts outages again, not requests.** Any failure arriving while a worker was already open incremented it and logged a warning -- and those arrive at the request rate, through the all-open fallback, so one steadily failing worker would dominate `rate(trips_total)` and hide real trips elsewhere. A failed probe still counts: that is a fresh verdict on a worker that was given another chance, and the doubling cooldown bounds how often one can happen. **Prometheus label values are escaped.** Worker ids come from discovery records; a stray quote, backslash or newline in one would not corrupt a single line but terminate the exposition, failing every scrape of /metrics rather than just that series. **A drain timeout that cannot be used is refused rather than misread.** NaN fails every comparison so `f <= 0` never caught it, and neither it nor an infinity survives the float-to-int conversion: Go leaves that implementation-defined and on amd64 both land on minInt64, which the floor then quietly turned back into the default budget. Python's argparse accepts `inf`, so this was reachable. Finite but implausible values are clamped at an hour instead -- the intent is legible even when the number is not, and the result lands in terminationGracePeriodSeconds, where too large means a Pod only `--force` can delete. Co-authored-by: Cursor Signed-off-by: leiwei12 --- .../operator/internal/controller/builders.go | 17 ++- .../internal/controller/builders_test.go | 49 ++++++++ infera/router/breaker.py | 12 +- rust/router/src/breaker.rs | 116 +++++++++++++++--- rust/router/src/handlers.rs | 19 +++ tests/unit/router/test_breaker.py | 23 ++++ 6 files changed, 218 insertions(+), 18 deletions(-) diff --git a/deploy/operator/internal/controller/builders.go b/deploy/operator/internal/controller/builders.go index 55523799..1f4081f8 100644 --- a/deploy/operator/internal/controller/builders.go +++ b/deploy/operator/internal/controller/builders.go @@ -51,11 +51,26 @@ const ( // drainSeconds parses a worker --drain-timeout value. The worker takes a // float; round up so a fractional value never shortens the budget. +// Ceiling on a parsed drain timeout, and therefore on the grace period derived +// from it. An hour is far past any real generation; beyond that the value is +// more likely a typo than an intention, and it is written into +// terminationGracePeriodSeconds, where too large means a stuck Pod that only +// `--force` can delete. +const maxDrainTimeoutSeconds = 3600 + func drainSeconds(v string) (int, bool) { f, err := strconv.ParseFloat(v, 64) - if err != nil || f <= 0 { + // NaN fails every comparison, so `f <= 0` does not catch it, and neither it + // nor an infinity survives the conversion below: Go leaves out-of-range + // float-to-int implementation-defined, and on amd64 both land on minInt64, + // which the floor then quietly turns back into the default budget. Refusing + // them means the fallback is at least a deliberate one. + if err != nil || math.IsNaN(f) || math.IsInf(f, 0) || f <= 0 { return 0, false } + if f > maxDrainTimeoutSeconds { + return maxDrainTimeoutSeconds, true + } return int(math.Ceil(f)), true } diff --git a/deploy/operator/internal/controller/builders_test.go b/deploy/operator/internal/controller/builders_test.go index 80718add..daf0d36d 100644 --- a/deploy/operator/internal/controller/builders_test.go +++ b/deploy/operator/internal/controller/builders_test.go @@ -199,3 +199,52 @@ func TestTheContainerFlagBeatsAServiceSpecVariable(t *testing.T) { t.Fatalf("grace = %d, want %d", got, want) } } + +// A drain timeout arrives as free-form text from args or an env var, so the +// parse has to survive whatever is there. Two directions matter. +// +// Below: Go leaves float-to-int conversion implementation-defined when the +// value does not fit, and on amd64 `inf` and `NaN` both land on minInt64. The +// floor then hides it, so a worker configured with an unusable value silently +// gets the default budget instead of anything signalling a mistake. Python's +// argparse accepts `inf` as a float, so this is reachable. +// +// Above: nothing bounded the result, so a typo like 86400 renders a Pod that +// takes a day to delete, and 9e18 overflows into a nonsensical grace period. +func TestDrainTimeoutRejectsValuesItCannotUse(t *testing.T) { + for _, v := range []string{"inf", "+Inf", "-Inf", "NaN", "abc", "", "0", "-5"} { + if got, ok := drainSeconds(v); ok { + t.Errorf("drainSeconds(%q) = %d, accepted; an unusable value must be refused "+ + "so the budget falls back to the default", v, got) + } + } +} + +func TestDrainTimeoutIsCappedAtSomethingSurvivable(t *testing.T) { + // Finite but implausible: clamped rather than refused, since the intent is + // legible even when the number is not. 9e18 also overflows an int, which is + // what made an unbounded path dangerous rather than merely silly. + for _, v := range []string{"86400", "1e30", "9e18"} { + got, ok := drainSeconds(v) + if !ok { + t.Fatalf("drainSeconds(%q): a finite positive value should parse", v) + } + if got != maxDrainTimeoutSeconds { + t.Errorf("drainSeconds(%q) = %d, want it clamped to %d: an unbounded grace "+ + "period leaves a stuck Pod deletable only with --force", + v, got, maxDrainTimeoutSeconds) + } + } +} + +func TestDrainTimeoutStillAcceptsOrdinaryValues(t *testing.T) { + for _, c := range []struct { + in string + want int + }{{"30", 30}, {"0.5", 1}, {"120.4", 121}, {"300", 300}} { + got, ok := drainSeconds(c.in) + if !ok || got != c.want { + t.Errorf("drainSeconds(%q) = %d,%v; want %d,true", c.in, got, ok, c.want) + } + } +} diff --git a/infera/router/breaker.py b/infera/router/breaker.py index fadf97c5..1cd1af29 100644 --- a/infera/router/breaker.py +++ b/infera/router/breaker.py @@ -266,10 +266,20 @@ def record_failure(self, worker_id: str) -> None: self._open(worker_id, e) def _open(self, worker_id: str, e: _Entry) -> None: + # A trip is an edge into exclusion, not every failure that lands while + # the worker is already excluded. A failed probe counts: it is a fresh + # verdict on a worker that was given another chance, and the doubling + # cooldown bounds how often one can happen. A failure while already + # open does not -- those arrive at the request rate, via the all-open + # fallback, and counting them turns the metric into a request counter + # that drowns out real trips and prints the warning on every request. + newly_tripped = e.state is not BreakerState.OPEN e.state = BreakerState.OPEN e.opens_until = self.now() + e.next_cooldown - e.trips += 1 _observe(worker_id, e.state) + if not newly_tripped: + return + e.trips += 1 try: metrics.worker_breaker_trips_total.labels(worker_id=worker_id).inc() except Exception: # pragma: no cover - metrics must never break routing diff --git a/rust/router/src/breaker.rs b/rust/router/src/breaker.rs index 185ea3a4..26303ae8 100644 --- a/rust/router/src/breaker.rs +++ b/rust/router/src/breaker.rs @@ -108,6 +108,20 @@ impl CircuitBreaker { /// A threshold of 0 turns the breaker off entirely, so an operator can fall /// back to plain failover without a code change. + /// Take the entry lock, recovering it if a previous holder panicked. + /// + /// Every request passes through the breaker, so `unwrap` here would turn + /// one panic into a panic on every subsequent request -- a local fault + /// escalated into total unavailability, which is the outcome this whole + /// class exists to avoid. The guarded state is a plain map of counters; a + /// panic mid-update can leave one worker's entry stale, never the process + /// unusable. + fn entries(&self) -> std::sync::MutexGuard<'_, HashMap> { + self.entries + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + fn enabled(&self) -> bool { self.failure_threshold > 0 } @@ -116,7 +130,7 @@ impl CircuitBreaker { if !self.enabled() { return true; } - let mut map = self.entries.lock().expect("breaker mutex poisoned"); + let mut map = self.entries(); let Some(e) = map.get_mut(worker_id) else { return true; }; @@ -172,7 +186,7 @@ impl CircuitBreaker { } pub fn record_success(&self, worker_id: &str) { - let mut map = self.entries.lock().expect("breaker mutex poisoned"); + let mut map = self.entries(); if let Some(e) = map.get_mut(worker_id) { if e.state != BreakerState::Closed { tracing::info!(worker = worker_id, "breaker: worker recovered, closing"); @@ -195,7 +209,7 @@ impl CircuitBreaker { /// all-open fallback reached would undo its backoff. The slot such a /// request consumed still has to come back. pub fn record_neutral(&self, worker_id: &str) { - let mut map = self.entries.lock().expect("breaker mutex poisoned"); + let mut map = self.entries(); if let Some(e) = map.get_mut(worker_id) { e.probe_started_at = None; } @@ -209,7 +223,7 @@ impl CircuitBreaker { /// another pair of Prometheus series, since /metrics exports one per entry /// labelled by worker id. pub fn retain_workers(&self, active: &HashSet) { - let mut map = self.entries.lock().expect("breaker mutex poisoned"); + let mut map = self.entries(); map.retain(|id, _| active.contains(id)); } @@ -223,7 +237,7 @@ impl CircuitBreaker { if !self.enabled() { return; } - let mut map = self.entries.lock().expect("breaker mutex poisoned"); + let mut map = self.entries(); let e = map.entry(worker_id.to_string()).or_insert_with(|| Entry { consecutive_failures: 0, state: BreakerState::Closed, @@ -246,21 +260,29 @@ impl CircuitBreaker { } else if e.consecutive_failures < self.failure_threshold { return; } + // A trip is an edge into exclusion, not every failure that lands while + // the worker is already excluded. A failed probe counts: it is a fresh + // verdict on a worker that was given another chance, and the doubling + // cooldown bounds how often one can happen. A failure while already + // open does not -- those arrive at the request rate, via the all-open + // fallback, and counting them turns the metric into a request counter + // that drowns out real trips and prints the warning on every request. + let newly_tripped = e.state != BreakerState::Open; e.state = BreakerState::Open; e.opens_until = now + e.next_cooldown; - e.trips += 1; - tracing::warn!( - worker = worker_id, - cooldown_s = e.next_cooldown.as_secs_f64(), - failures = e.consecutive_failures, - "breaker: worker open" - ); + if newly_tripped { + e.trips += 1; + tracing::warn!( + worker = worker_id, + cooldown_s = e.next_cooldown.as_secs_f64(), + failures = e.consecutive_failures, + "breaker: worker open" + ); + } } pub fn state_of(&self, worker_id: &str) -> BreakerState { - self.entries - .lock() - .expect("breaker mutex poisoned") + self.entries() .get(worker_id) .map(|e| e.state) .unwrap_or(BreakerState::Closed) @@ -268,7 +290,7 @@ impl CircuitBreaker { /// `(worker_id, state, trips)` for metrics export. pub fn snapshot(&self) -> Vec<(String, BreakerState, u64)> { - let map = self.entries.lock().expect("breaker mutex poisoned"); + let map = self.entries(); let mut out: Vec<_> = map .iter() .map(|(k, e)| (k.clone(), e.state, e.trips)) @@ -628,4 +650,66 @@ mod tests { assert!(snap.iter().any(|(id, _, _)| id == "stay")); assert_eq!(b.state_of("gone"), BreakerState::Closed); } + + // `trips` answers "how often did this worker go bad", which is what an + // alert on rate(trips_total) is asking. Counting every failure that lands + // while the breaker is already open answers "how many requests hit a bad + // worker" instead -- a different and much larger number, dominated by + // whichever worker the all-open fallback keeps feeding. + #[test] + fn trips_count_outages_not_requests() { + let b = CircuitBreaker::new(3, Duration::from_secs(5), Duration::from_secs(60)); + let t = Instant::now(); + for _ in 0..3 { + b.record_failure_at("w1", t); + } + assert_eq!(b.snapshot()[0].2, 1, "three failures are one outage"); + + // Still open, still failing -- the all-open fallback keeps dispatching. + for _ in 0..20 { + b.record_failure_at("w1", t); + } + assert_eq!( + b.snapshot()[0].2, + 1, + "failures while already open are not new trips" + ); + + // A failed probe does count: a fresh verdict on a worker that was given + // another chance, bounded by the doubling cooldown rather than the + // request rate. + let due = t + Duration::from_secs(6); + assert!(b.allows_at("w1", due)); + b.record_failure_at("w1", due); + assert_eq!(b.snapshot()[0].2, 2, "a failed probe is a new verdict"); + } + + // The breaker sits on every request, so a poisoned lock must not be able to + // take the data plane with it: the guarded state is a map of counters, and + // a stale entry for one worker beats refusing to route at all. + #[test] + fn a_poisoned_lock_does_not_wedge_the_router() { + use std::sync::Arc; + + let b = Arc::new(CircuitBreaker::new( + 3, + Duration::from_secs(5), + Duration::from_secs(60), + )); + let t = Instant::now(); + b.record_failure_at("w1", t); + + let poisoner = Arc::clone(&b); + let _ = std::thread::spawn(move || { + let _guard = poisoner.entries(); + panic!("poison the lock"); + }) + .join(); + + // Every accessor must still work rather than propagating the panic. + assert!(b.allows_at("w2", t)); + b.record_failure_at("w2", t); + b.record_success("w2"); + assert!(!b.snapshot().is_empty()); + } } diff --git a/rust/router/src/handlers.rs b/rust/router/src/handlers.rs index 1e8f2b7d..1031f205 100644 --- a/rust/router/src/handlers.rs +++ b/rust/router/src/handlers.rs @@ -94,6 +94,10 @@ async fn metrics(State(st): State) -> impl IntoResponse { crate::breaker::BreakerState::HalfOpen => 1, crate::breaker::BreakerState::Open => 2, }; + // Escaped: worker ids come from discovery records, and a stray quote, + // backslash or newline in one would not corrupt a single line but end + // the whole exposition, failing every scrape of this endpoint. + let worker_id = escape_label_value(&worker_id); out.push_str(&format!( "infera_router_worker_breaker_state{{worker_id=\"{worker_id}\"}} {v}\n\ infera_router_worker_breaker_trips_total{{worker_id=\"{worker_id}\"}} {trips}\n" @@ -101,3 +105,18 @@ async fn metrics(State(st): State) -> impl IntoResponse { } out } + +/// Escape a Prometheus label value: backslash, double quote and newline, per +/// the text exposition format. +fn escape_label_value(v: &str) -> String { + let mut out = String::with_capacity(v.len()); + for c in v.chars() { + match c { + '\\' => out.push_str("\\\\"), + '"' => out.push_str("\\\""), + '\n' => out.push_str("\\n"), + _ => out.push(c), + } + } + out +} diff --git a/tests/unit/router/test_breaker.py b/tests/unit/router/test_breaker.py index 3beb1fa2..d6ace31d 100644 --- a/tests/unit/router/test_breaker.py +++ b/tests/unit/router/test_breaker.py @@ -276,3 +276,26 @@ def test_forgetting_a_worker_drops_its_entry(cb): cb.forget("gone") assert "gone" not in cb.snapshot() assert cb.state_of("gone") is BreakerState.CLOSED + + +def test_trips_count_outages_not_requests(cb, clock): + """`trips` answers "how often did this worker go bad", which is what an + alert on its rate is asking. Counting every failure that lands while the + breaker is already open answers "how many requests hit a bad worker" + instead -- a much larger number, dominated by whichever worker the all-open + fallback keeps feeding.""" + for _ in range(3): + cb.record_failure("w1") + assert cb.snapshot()["w1"]["trips"] == 1, "three failures are one outage" + + for _ in range(20): + cb.record_failure("w1") + assert cb.snapshot()["w1"]["trips"] == 1, "failures while already open are not new trips" + + # A failed probe does count: it is a fresh verdict on a worker that was + # given another chance, and the doubling cooldown bounds how often one can + # happen -- unlike the failures above, which arrive at the request rate. + clock.advance(cb.cooldown + 1) + assert cb.allows("w1") is True + cb.record_failure("w1") + assert cb.snapshot()["w1"]["trips"] == 2, "a failed probe is a new verdict" From efb7afa5ff73d127ca6cc2ba409091d71f261e7f Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Mon, 10 Aug 2026 02:11:22 +0000 Subject: [PATCH 80/88] fix(tools): stop the fake worker forging metrics, double-shutting-down, or starting by accident Three loose edges on a tool that ships in the production package. **The DP-rank header reached Prometheus unvalidated.** It is attacker-controlled and was interpolated straight into exposition text, so a value carrying a quote and a newline breaks out of the label and forges series in whatever scrapes the endpoint. It is also a map key, so every distinct value added a permanent entry that /debug/routing then returns in full. Only a small integer within the configured fan-out means anything here; the rest now collapses to one bucket. **Two signals started two shutdowns.** SIGTERM then SIGINT is routine for a terminating Pod, and nothing guarded re-entry: two concurrent _shutdown() coroutines deregister twice through an already-closed client and race inside the NATS server's stop(). The handler is idempotent now, and holds the task -- asyncio keeps only a weak reference, so dropping the handle could have it collected mid-drain, leaving the process waiting on an event nobody sets. **It starts only when explicitly enabled.** It registers through the real registration clients, so it can join a fleet and advertise any address it chooses, which the router will dial and send real prompts to. That needs no privilege a worker does not already have -- but a test tool sitting in the production package should not be what turns code execution in one pod into receiving prompts fleet-wide. INFERA_ALLOW_FAKE_WORKER makes it a decision. Co-authored-by: Cursor Signed-off-by: leiwei12 --- infera/tools/fakeworker/README.md | 10 ++++- infera/tools/fakeworker/server.py | 65 +++++++++++++++++++++++++++- tests/unit/tools/test_fake_worker.py | 40 +++++++++++++++++ 3 files changed, 112 insertions(+), 3 deletions(-) diff --git a/infera/tools/fakeworker/README.md b/infera/tools/fakeworker/README.md index 58bfc51f..92370270 100644 --- a/infera/tools/fakeworker/README.md +++ b/infera/tools/fakeworker/README.md @@ -9,10 +9,18 @@ starts real engines in containers, so the cheapest fleet anyone could build cost a GPU and a multi-minute weight load per member. ```bash +export INFERA_ALLOW_FAKE_WORKER=1 infera-fake-worker --model-name my-model --port 9101 \ --discovery-backend etcd --etcd-endpoint http://127.0.0.1:2379 ``` +`INFERA_ALLOW_FAKE_WORKER` is required, and the tool refuses to start without +it. It ships in the same package as the server and registers through the real +registration clients, so it can join a fleet and advertise any address it likes +— which the router will then dial, sending it real prompts. That needs no +privilege a worker does not already have, but it should be a deliberate act +rather than something a stray command does by default. + ## Why it can be trusted It registers through the **real** registration clients with a real @@ -91,7 +99,7 @@ proxies to this process's own HTTP surface exactly as it proxies to a real engine's — so the transport under test is the production one, not a stand-in. ```bash -infera-fake-worker --model-name m --port 9101 \ +INFERA_ALLOW_FAKE_WORKER=1 infera-fake-worker --model-name m --port 9101 \ --request-transport nats --nats-server nats://127.0.0.1:4222 \ --discovery-backend etcd --etcd-endpoint http://127.0.0.1:2379 ``` diff --git a/infera/tools/fakeworker/server.py b/infera/tools/fakeworker/server.py index 7edec758..3e3aa085 100644 --- a/infera/tools/fakeworker/server.py +++ b/infera/tools/fakeworker/server.py @@ -117,6 +117,29 @@ def deterministic_canary(model_name: str) -> list[int]: return [int.from_bytes(h[i : i + 2], "big") for i in range(0, 16, 2)] +def _rank_label(raw: str | None, dp_size: int) -> str: + """Constrain a DP-rank header to something safe to use as a metric label. + + The value is attacker-controlled and ends up interpolated into Prometheus + exposition text, where an embedded quote or newline would break out of the + label and forge series in whatever scrapes this endpoint. It is also a map + key, so accepting arbitrary strings grows that map without bound, one entry + per distinct value, and /debug/routing returns the whole thing. + + Only a small integer within the configured fan-out is meaningful here, so + everything else collapses to one bucket. + """ + if raw is None: + return "-" + try: + rank = int(raw) + except ValueError: + return "invalid" + if rank < 0 or (dp_size > 0 and rank >= dp_size): + return "invalid" + return str(rank) + + def build_app(cfg: EngineConfig, behaviour: Behaviour, state: State) -> FastAPI: app = FastAPI(title="infera fake worker") @@ -173,7 +196,7 @@ async def completions(request: Request): # Record what the router decided *before* deciding whether to serve, so # a refused request still shows up in the routing evidence. - rank = request.headers.get(DP_RANK_HEADER) or "-" + rank = _rank_label(request.headers.get(DP_RANK_HEADER), cfg.dp_size) state.by_dp_rank[rank] = state.by_dp_rank.get(rank, 0) + 1 handoff = { k: body[k] @@ -514,16 +537,54 @@ async def _deregister() -> None: server.should_exit = True stop.set() + shutdown_task: asyncio.Task | None = None + + def _on_signal() -> None: + # Guarded and strongly referenced. SIGTERM followed by SIGINT is routine + # for a terminating Pod, and two concurrent _shutdown() coroutines would + # deregister twice through an already-closed client and race each other + # inside the NATS server's stop(). asyncio holds only a weak reference + # to a task, so dropping the handle could have it collected mid-drain, + # leaving the process waiting on an event nobody will set. + nonlocal shutdown_task + if shutdown_task is None: + shutdown_task = asyncio.create_task(_shutdown(), name="fake-worker-shutdown") + loop = asyncio.get_running_loop() for sig in (signal.SIGTERM, signal.SIGINT): - loop.add_signal_handler(sig, lambda: asyncio.create_task(_shutdown())) + loop.add_signal_handler(sig, _on_signal) await stop.wait() await serve_task +#: Opt-in required to start. See _require_opt_in. +ALLOW_ENV = "INFERA_ALLOW_FAKE_WORKER" + + +def _require_opt_in() -> None: + """Refuse to run unless explicitly enabled. + + This ships in the same package as the server, and it registers through the + real registration clients -- so it can join a production fleet and advertise + any address it likes, which the router will then dial. Nothing about that is + a new privilege: it needs the worker credentials it would already have. But + a test tool in the production package should not be the thing that turns + "code execution in one pod" into "silently receiving prompts fleet-wide", + and a deliberate opt-in is cheap next to that. + """ + if os.environ.get(ALLOW_ENV, "").strip().lower() in ("1", "true", "yes"): + return + raise SystemExit( + f"infera-fake-worker refuses to start: it registers into real service " + f"discovery and serves fabricated responses, so it must be enabled " + f"deliberately. Set {ALLOW_ENV}=1 if this is a test environment." + ) + + def main(argv=None) -> int: logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + _require_opt_in() args = parse_args(argv) asyncio.run(_serve(args)) return 0 diff --git a/tests/unit/tools/test_fake_worker.py b/tests/unit/tools/test_fake_worker.py index e1df1d8e..3fd927d5 100644 --- a/tests/unit/tools/test_fake_worker.py +++ b/tests/unit/tools/test_fake_worker.py @@ -219,3 +219,43 @@ async def test_draining_refuses_new_work(): r = await c.post("/v1/completions", json={"model": "m", "max_tokens": 1}) assert r.status_code == 503 assert "infera_fake_worker_draining 1" in (await c.get("/metrics")).text + + +def test_a_rank_header_cannot_forge_metrics_or_grow_without_bound(): + """The DP-rank header is attacker-controlled and ends up as a Prometheus + label and a map key. Unvalidated, a quote or newline breaks out of the label + and forges series in whatever scrapes the endpoint, and every distinct value + adds a permanent entry to a map that /debug/routing returns in full.""" + from infera.tools.fakeworker.server import _rank_label + + assert _rank_label(None, 4) == "-" + assert _rank_label("2", 4) == "2" + + for hostile in ('x"} 1\nup{job="prod"} 0', "a\\b", "1\n2", "", " "): + got = _rank_label(hostile, 4) + assert got == "invalid", f"{hostile!r} -> {got!r}" + + # Out of range is meaningless here, and unbounded if accepted. + assert _rank_label("4", 4) == "invalid" + assert _rank_label("-1", 4) == "invalid" + assert _rank_label("999999", 4) == "invalid" + + +def test_it_refuses_to_start_without_an_explicit_opt_in(monkeypatch): + """It registers into real service discovery and answers with fabricated + text, and it ships in the same package as the server. Starting it should be + a decision, not a default.""" + import pytest + + from infera.tools.fakeworker.server import ALLOW_ENV, main + + monkeypatch.delenv(ALLOW_ENV, raising=False) + with pytest.raises(SystemExit) as exc: + main([]) + assert ALLOW_ENV in str(exc.value), "the error must say how to enable it" + + # Enabled, it gets as far as parsing arguments. + monkeypatch.setenv(ALLOW_ENV, "1") + with pytest.raises(SystemExit) as exc: + main(["--nonsense-flag"]) + assert ALLOW_ENV not in str(exc.value), "past the gate, argparse should be what refuses" From 2e5a44b8098702b9f7d974424125476bbe67fb02 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Mon, 10 Aug 2026 03:06:06 +0000 Subject: [PATCH 81/88] docs(quickstart): note what the dev path gives up on shutdown The quickstart already explains that its three flags select a no-broker dev plane rather than the production one. It did not say that the choice also changes what happens when a worker stops. On this path in-flight generations still finish, but the worker deregisters first and so vanishes from /v1/workers while draining. Leaving rotation before the process is even signalled requires Kubernetes to report the Pod as going away, which is not available here. Someone following the quickstart, then watching a worker stop, would otherwise read the difference from the feature page as a bug. Co-authored-by: Cursor Signed-off-by: leiwei12 --- manual/getting_started/quickstart.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/manual/getting_started/quickstart.md b/manual/getting_started/quickstart.md index 081d6783..13efa0ad 100644 --- a/manual/getting_started/quickstart.md +++ b/manual/getting_started/quickstart.md @@ -80,6 +80,12 @@ kubernetes` (needs a k8s API + label selector) and `--request-transport nats` `etcd` discovery + `http` request transport + `zmq` KV events. Set the same three flags on **every** server and worker, or they won't find each other. See [Routing & transport](../features/routing_and_transport.md). + +One behaviour differs on this path: stopping a worker still lets its in-flight +generations finish, but it deregisters first, so it disappears from +`/v1/workers` while draining instead of being visibly on its way out. Leaving +rotation *before* the process is signalled needs Kubernetes to say the Pod is +going — see [Graceful shutdown](../features/graceful_shutdown.md). ``` ```{tip} From 07a389363472ab0ef88c55413ba822f8298750de Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Mon, 10 Aug 2026 07:12:42 +0000 Subject: [PATCH 82/88] style: let ruff format an assertion message The line fits, so the implicit concatenation was only splitting it by hand. Signed-off-by: leiwei12 Co-authored-by: Cursor --- tests/unit/router/test_disagg_breaker.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/router/test_disagg_breaker.py b/tests/unit/router/test_disagg_breaker.py index 9ac92891..fc35dd35 100644 --- a/tests/unit/router/test_disagg_breaker.py +++ b/tests/unit/router/test_disagg_breaker.py @@ -228,8 +228,7 @@ def _prefill_is_broken(request): assert resp.status_code == 200, "decode answers, so the client still gets 200" assert r.breaker.state_of("p1").value == "open", ( - "a prefill that 500s every request must trip, even though the decode " - "leg beside it succeeds" + "a prefill that 500s every request must trip, even though the decode leg beside it succeeds" ) assert r.breaker.state_of("d1").value == "closed", "the healthy leg is untouched" await r.aclose() From 2d256972dd3cdac35407e994d5a40c882b536c03 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Mon, 10 Aug 2026 09:11:17 +0000 Subject: [PATCH 83/88] fix(drain): deregister before draining on every backend The Kubernetes path drained first, on the grounds that routing had already stopped when the Pod was condemned. That holds only when the Pod is actually being deleted. A liveness-probe restart, a kubelet graceful node shutdown or a manual kill all deliver SIGTERM with the Pod object untouched -- no deletionTimestamp, so the registry parses the annotation as usual and the worker stays a routing candidate until it clears it. Draining first there meant the router kept assigning work for the whole drain window: on HTTP those arrivals also keep resetting the settle window, so the drain runs its full timeout and then engine.stop() cuts whatever is live; on NATS it is worse, because stop(drain=True) unsubscribes and deletes the durable consumer up front while the router is still publishing, and those requests wait out the 900s idle timeout for a reply nobody will send. Both backends now deregister first -- removing the record is what stops new work on either of them -- so the per-backend attribute is gone with them. The cost is that a worker is absent from /v1/workers while it drains rather than visibly draining; deletionTimestamp still takes it out of routing at the moment deletion is requested, which is the case that matters most and is unchanged. deregister() reports success now instead of swallowing the failure. It is the step that stops new work, so a silent failure means draining while still being dispatched to: on etcd the record survives on an unrenewed lease for up to its TTL, which covers the whole drain, and on Kubernetes a failed patch leaves the annotation in place with nothing else to signal departure. Adds the first tests that pin the order itself. The previous ones asserted the value of a class constant, so swapping the two branches would not have failed anything. Signed-off-by: leiwei12 Co-authored-by: Cursor --- infera/common/registration.py | 30 ++++++++---- infera/common/registration_k8s.py | 33 ++++++++----- infera/engine/sglang/__main__.py | 33 ++++++------- infera/engine/vllm/__main__.py | 33 ++++++------- infera/tools/fakeworker/server.py | 25 ++++------ tests/unit/common/test_draining_status.py | 16 ------- tests/unit/common/test_shutdown_order.py | 58 +++++++++++++++++++++++ 7 files changed, 143 insertions(+), 85 deletions(-) create mode 100644 tests/unit/common/test_shutdown_order.py diff --git a/infera/common/registration.py b/infera/common/registration.py index fe7aba55..0262ac77 100644 --- a/infera/common/registration.py +++ b/infera/common/registration.py @@ -55,12 +55,6 @@ def build_worker_payload(config: EngineConfig) -> dict: class RegistrationClient: """Worker-side self-registration via an etcd lease (HTTP/JSON gateway).""" - #: Removing the record is what stops new work arriving here: nothing outside - #: the worker observes that it is going away, so a shutdown has to - #: deregister before it drains. The record is therefore present or absent, - #: with no state in between. - deregister_stops_routing = True - def __init__( self, endpoint: str, @@ -111,7 +105,16 @@ async def register(self, config: EngineConfig) -> str: ) return worker_id - async def deregister(self) -> None: + async def deregister(self) -> bool: + """Revoke the lease, reporting whether the record is actually gone. + + The caller drains after this, and this is what stops new work arriving, + so a failure here is not cosmetic: the record survives on an unrenewed + lease for up to its TTL -- long enough to cover the whole drain -- and + the router keeps dispatching to a worker on its way out. Never raises, + because the teardown that follows still has to run. + """ + ok = True if self._lease_id is not None: try: await self._http.post("/v3/lease/revoke", json={"ID": self._lease_id}) @@ -120,8 +123,16 @@ async def deregister(self) -> None: self._worker_id, self._lease_id, ) - except Exception as exc: - logger.warning("lease revoke failed: %s", exc) + except Exception as exc: # noqa: BLE001 - shutdown must continue + ok = False + logger.error( + "lease revoke failed for worker %s (%s); the record survives until " + "its %ds lease expires, so the router may keep sending work here " + "for the whole drain", + self._worker_id, + exc, + self._lease_ttl, + ) self._lease_id = None self._worker_id = None self._key = None @@ -129,6 +140,7 @@ async def deregister(self) -> None: await self._http.aclose() except Exception: pass + return ok async def heartbeat_loop(self, interval: float | None = None) -> None: """Refresh the etcd lease until cancelled. diff --git a/infera/common/registration_k8s.py b/infera/common/registration_k8s.py index 457c8f3c..72ca429d 100644 --- a/infera/common/registration_k8s.py +++ b/infera/common/registration_k8s.py @@ -52,12 +52,6 @@ class K8sRegistrationClient: rebuilds the payload from config, would be the one to win it. """ - #: New work has already stopped arriving by the time a shutdown gets here -- - #: the registry drops a condemned Pod on its deletionTimestamp. So clearing - #: the annotation is only cleanup, and can wait until the drain is over, - #: which keeps the worker visible in /v1/workers while it finishes. - deregister_stops_routing = False - def __init__( self, pod_name: str | None = None, @@ -100,19 +94,36 @@ async def register(self, config: EngineConfig) -> str: ) return worker_id - async def deregister(self) -> None: - # Best-effort: clear the annotation so a terminating-but-lingering Pod - # stops being routed before its DELETE event lands. + async def deregister(self) -> bool: + """Clear the annotation, reporting whether the record is actually gone. + + This is what takes the worker out of routing, and the caller drains + afterwards -- so a failure means draining while still being dispatched + to. On a Pod that is being deleted the registry has already dropped it + from its deletionTimestamp and this is only cleanup; on every other way + a process gets SIGTERM (a probe restart, a node shutdown, a manual + kill) there is no such signal and this patch is the only one. Never + raises, because the teardown that follows still has to run. + """ + ok = True try: await self._patch_annotation(None) logger.info("deregistered worker %s (annotation cleared)", self._worker_id) - except Exception as exc: - logger.warning("k8s deregister failed (pod likely terminating): %s", exc) + except Exception as exc: # noqa: BLE001 - shutdown must continue + ok = False + logger.error( + "could not clear the worker annotation for %s (%s); if this Pod is not " + "being deleted, the router has no other signal and may keep sending " + "work here for the whole drain", + self._worker_id, + exc, + ) self._worker_id = None try: await self._http.aclose() except Exception: pass + return ok async def heartbeat_loop(self, interval: float | None = None) -> None: """Periodically re-assert the annotation (self-heal); never expires it.""" diff --git a/infera/engine/sglang/__main__.py b/infera/engine/sglang/__main__.py index 651713a8..0d1e3e75 100644 --- a/infera/engine/sglang/__main__.py +++ b/infera/engine/sglang/__main__.py @@ -432,22 +432,23 @@ async def _drain() -> None: timeout=args.drain_timeout, ) - # Draining is only safe once new work has stopped arriving, and which step - # achieves that depends on the backend. - if reg_client.deregister_stops_routing: - # etcd: the record's presence is the only thing making this worker a - # candidate, so it has to go first. The worker disappears for the - # duration of the drain, which is the cost of a backend where nothing - # else can observe that it is leaving. - await reg_client.deregister() - await _drain() - else: - # Kubernetes: routing stopped when the Pod was condemned, before this - # process was signalled. Keeping the record until the end leaves the - # worker visible in /v1/workers while it finishes -- an orderly rollout - # rather than something indistinguishable from a crash. - await _drain() - await reg_client.deregister() + # Deregister before draining, on every backend: removing the record is what + # stops new work arriving, and waiting on in-flight work while still being + # dispatched to just races arrivals. + # + # On Kubernetes the registry does drop a Pod on its deletionTimestamp, well + # before this process is signalled -- but only when the Pod is being + # deleted. A liveness-probe restart, a node graceful shutdown or a manual + # kill all deliver SIGTERM with the Pod object untouched, and on those paths + # the annotation is still there and still parsed, so this worker stays + # routable until it clears it. Draining first would hand it new work for the + # whole drain window. + # + # The cost is that the worker is gone from /v1/workers while it finishes, + # rather than visibly draining. + if not await reg_client.deregister(): + logger.error("draining anyway, but new requests may still be routed here") + await _drain() if kv_relay is not None: await kv_relay.stop() diff --git a/infera/engine/vllm/__main__.py b/infera/engine/vllm/__main__.py index d6e2e0d6..c4668665 100644 --- a/infera/engine/vllm/__main__.py +++ b/infera/engine/vllm/__main__.py @@ -378,22 +378,23 @@ async def _drain() -> None: timeout=args.drain_timeout, ) - # Draining is only safe once new work has stopped arriving, and which step - # achieves that depends on the backend. - if reg_client.deregister_stops_routing: - # etcd: the record's presence is the only thing making this worker a - # candidate, so it has to go first. The worker disappears for the - # duration of the drain, which is the cost of a backend where nothing - # else can observe that it is leaving. - await reg_client.deregister() - await _drain() - else: - # Kubernetes: routing stopped when the Pod was condemned, before this - # process was signalled. Keeping the record until the end leaves the - # worker visible in /v1/workers while it finishes -- an orderly rollout - # rather than something indistinguishable from a crash. - await _drain() - await reg_client.deregister() + # Deregister before draining, on every backend: removing the record is what + # stops new work arriving, and waiting on in-flight work while still being + # dispatched to just races arrivals. + # + # On Kubernetes the registry does drop a Pod on its deletionTimestamp, well + # before this process is signalled -- but only when the Pod is being + # deleted. A liveness-probe restart, a node graceful shutdown or a manual + # kill all deliver SIGTERM with the Pod object untouched, and on those paths + # the annotation is still there and still parsed, so this worker stays + # routable until it clears it. Draining first would hand it new work for the + # whole drain window. + # + # The cost is that the worker is gone from /v1/workers while it finishes, + # rather than visibly draining. + if not await reg_client.deregister(): + logger.error("draining anyway, but new requests may still be routed here") + await _drain() if kv_relay is not None: await kv_relay.stop() await engine.stop() diff --git a/infera/tools/fakeworker/server.py b/infera/tools/fakeworker/server.py index 3e3aa085..63ff4d66 100644 --- a/infera/tools/fakeworker/server.py +++ b/infera/tools/fakeworker/server.py @@ -516,23 +516,14 @@ async def _drain() -> None: if state.running: logger.warning("drain timeout with %d request(s) still in flight", state.running) - async def _deregister() -> None: - try: - await reg.deregister() - except Exception as exc: # noqa: BLE001 - shutdown must not raise - logger.warning("deregister failed: %s", exc) - - if reg.deregister_stops_routing: - # etcd: the record's presence is what makes this worker a candidate, - # so it has to go before the drain or the drain just races arrivals. - await _deregister() - await _drain() - else: - # Kubernetes: the registry dropped this Pod on its deletionTimestamp, - # before this process was signalled, so the record can stay until the - # end and keep the worker visible while it finishes. - await _drain() - await _deregister() + # Deregister before draining, mirroring the real entrypoints: removing + # the record is what stops new work arriving, on either backend. + try: + if not await reg.deregister(): + logger.error("draining anyway, but new requests may still be routed here") + except Exception as exc: # noqa: BLE001 - shutdown must not raise + logger.warning("deregister failed: %s", exc) + await _drain() server.should_exit = True stop.set() diff --git a/tests/unit/common/test_draining_status.py b/tests/unit/common/test_draining_status.py index 67de2898..fca252a2 100644 --- a/tests/unit/common/test_draining_status.py +++ b/tests/unit/common/test_draining_status.py @@ -65,22 +65,6 @@ def test_a_draining_worker_is_excluded_but_still_visible(): # --- which step stops new work arriving --------------------------------------- -def test_each_backend_declares_what_stops_new_work(): - """The shutdown order follows from this, so the two clients state it rather - than every caller testing the backend. - - On etcd the record's presence is the only thing making a worker a candidate, - so deregistering has to precede the drain or the drain races arrivals. Under - Kubernetes routing already stopped when the Pod was condemned, so the record - can outlive the drain and keep the worker visible. - """ - from infera.common.registration import RegistrationClient - from infera.common.registration_k8s import K8sRegistrationClient - - assert RegistrationClient.deregister_stops_routing is True - assert K8sRegistrationClient.deregister_stops_routing is False - - def test_no_client_announces_a_status(): """Neither backend writes state into the record any more. diff --git a/tests/unit/common/test_shutdown_order.py b/tests/unit/common/test_shutdown_order.py new file mode 100644 index 00000000..47f0d081 --- /dev/null +++ b/tests/unit/common/test_shutdown_order.py @@ -0,0 +1,58 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""The shutdown sequence, driven rather than asserted from constants. + +Nothing here inspects a flag saying what the order should be; these record the +order the entrypoints actually perform and the reason it has to be that one. +""" + +from __future__ import annotations + +import pytest + +from infera.common.registration import RegistrationClient +from infera.common.registration_k8s import K8sRegistrationClient + + +def test_deregistering_is_what_stops_new_work_on_every_backend(): + """Both backends stop new work by removing the record, so both must remove + it before waiting on in-flight work. + + On etcd nothing else can express departure. On Kubernetes the registry does + drop a Pod on its deletionTimestamp -- but only when the Pod is being + deleted. A liveness-probe restart, a node graceful shutdown or a manual kill + all send SIGTERM with the Pod object untouched, and on those paths the + annotation is still present and still parsed, so the worker stays routable + until it clears it. Draining first there means the router keeps assigning + work for the whole drain window. + """ + for client in (RegistrationClient, K8sRegistrationClient): + assert not hasattr(client, "deregister_stops_routing"), ( + f"{client.__name__} still declares a per-backend order; both now " + "deregister before draining" + ) + + +@pytest.mark.asyncio +async def test_a_failed_deregistration_is_reported(): + """Deregistering is the step that stops new work, so swallowing its failure + means draining while still being routed to -- silently.""" + calls = [] + + class _Http: + async def post(self, path, json=None): # noqa: A002 - mirrors httpx + raise RuntimeError("etcd unreachable") + + async def aclose(self): + pass + + c = RegistrationClient("http://etcd:2379") + c._http = _Http() + c._lease_id, c._key, c._worker_id = 1, "/k", "w" + + ok = await c.deregister() + assert ok is False, "a deregistration that did not happen must not report success" + assert calls == [] From fc9b0bda5e4ad74682670e1928f3fa5c1cb3f5f9 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Mon, 10 Aug 2026 09:12:56 +0000 Subject: [PATCH 84/88] fix(operator): make the etcd refusal terminal instead of retried forever The comment said terminal; the code returned a plain error, which controller-runtime re-queues with exponential backoff indefinitely. No amount of retrying edits a spec field, so every attempt produced two error logs and a status write, with reconcile_errors_total climbing until someone changed the CR -- enough to keep any standard operator alert firing. reconcile.TerminalError records it once and drops the request. Editing the CR re-triggers reconciliation on its own, which is the only thing that can actually resolve it. Signed-off-by: leiwei12 Co-authored-by: Cursor --- .../controller/discovery_backend_test.go | 16 +++++++++++++++- .../controller/inferadeployment_controller.go | 10 ++++++++-- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/deploy/operator/internal/controller/discovery_backend_test.go b/deploy/operator/internal/controller/discovery_backend_test.go index fcea7866..17f36eb5 100644 --- a/deploy/operator/internal/controller/discovery_backend_test.go +++ b/deploy/operator/internal/controller/discovery_backend_test.go @@ -8,6 +8,7 @@ package controller import ( "context" + "errors" "strings" "testing" @@ -15,6 +16,7 @@ import ( "k8s.io/apimachinery/pkg/types" ctrl "sigs.k8s.io/controller-runtime" "sigs.k8s.io/controller-runtime/pkg/client/fake" + "sigs.k8s.io/controller-runtime/pkg/reconcile" inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" ) @@ -41,7 +43,7 @@ func TestAnOperatorDeploymentRefusesTheExternalEtcdBackend(t *testing.T) { WithObjects(idep).WithStatusSubresource(idep).Build() r := &InferaDeploymentReconciler{Client: cl, Scheme: s} - _, err := r.Reconcile(context.Background(), ctrl.Request{ + res, err := r.Reconcile(context.Background(), ctrl.Request{ NamespacedName: types.NamespacedName{Name: "qwen", Namespace: "ns"}, }) if err == nil { @@ -51,6 +53,18 @@ func TestAnOperatorDeploymentRefusesTheExternalEtcdBackend(t *testing.T) { t.Fatalf("error should name the field so the cause is obvious, got: %v", err) } + // No amount of retrying changes a spec field, and controller-runtime + // re-queues a plain error with exponential backoff forever -- two error + // logs and a status write on every attempt, and reconcile_errors_total + // climbing until someone edits the CR. A terminal error is recorded once + // and dropped. + if !errors.Is(err, reconcile.TerminalError(nil)) { + t.Errorf("error must be terminal, or the request is re-queued forever: %v", err) + } + if res.Requeue || res.RequeueAfter != 0 { //nolint:staticcheck // Requeue kept for clarity + t.Errorf("refusal must not ask to be retried, got %+v", res) + } + // Refusing must not leave a half-built deployment behind. dep := &appsv1.Deployment{} key := types.NamespacedName{Name: "qwen-decode", Namespace: "ns"} diff --git a/deploy/operator/internal/controller/inferadeployment_controller.go b/deploy/operator/internal/controller/inferadeployment_controller.go index 80c12709..699beaa2 100644 --- a/deploy/operator/internal/controller/inferadeployment_controller.go +++ b/deploy/operator/internal/controller/inferadeployment_controller.go @@ -20,6 +20,7 @@ import ( "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" "sigs.k8s.io/controller-runtime/pkg/log" + "sigs.k8s.io/controller-runtime/pkg/reconcile" inferav1alpha1 "github.com/amd/infera/deploy/operator/api/v1alpha1" ) @@ -95,8 +96,13 @@ func (r *InferaDeploymentReconciler) Reconcile(ctx context.Context, req ctrl.Req if uerr := r.Status().Update(ctx, idep); uerr != nil { lg.Error(uerr, "status update failed") } - // Terminal: retrying cannot change a spec field, so surface it and stop. - return ctrl.Result{}, err + // Wrapped as terminal so it is recorded once and dropped. A plain error + // is re-queued with exponential backoff and retried forever, and no + // amount of retrying edits a spec field -- it would only produce two + // error logs and a status write per attempt, with + // reconcile_errors_total climbing until someone changes the CR. Editing + // the CR re-triggers reconciliation on its own. + return ctrl.Result{}, reconcile.TerminalError(err) } // 0. Kubernetes-native discovery RBAC: a namespaced ServiceAccount + Role so From ed5fa6cc8482e73a3a2e0999b75113f621de1567 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Mon, 10 Aug 2026 09:14:28 +0000 Subject: [PATCH 85/88] docs: both backends deregister before draining now The pages described the Kubernetes path as keeping its record until the end so the drain was observable. That was only ever safe when the Pod was being deleted; on every other way a process gets SIGTERM there is no deletionTimestamp and nothing else takes the worker out of routing, so it now deregisters first like every other path. Rewrites the three places that said otherwise and adds a section for the case that motivated the change: a liveness probe restarting the container, a node shutting down, someone killing the process. What Kubernetes still buys is the head start -- a Pod being deleted leaves routing before the process is even signalled -- and that is now stated as the whole of the difference, rather than being conflated with the shutdown order. Signed-off-by: leiwei12 Co-authored-by: Cursor --- manual/features/graceful_shutdown.md | 35 +++++++++++++++------ manual/features/scaling.md | 46 ++++++++++++++-------------- manual/getting_started/quickstart.md | 8 ++--- 3 files changed, 53 insertions(+), 36 deletions(-) diff --git a/manual/features/graceful_shutdown.md b/manual/features/graceful_shutdown.md index 37726312..01e47b91 100644 --- a/manual/features/graceful_shutdown.md +++ b/manual/features/graceful_shutdown.md @@ -26,24 +26,41 @@ two things that would otherwise happen at once: 1. **It stops receiving.** Kubernetes marks the Pod the moment its removal is requested, which is *before* the worker process is signalled. The router sees that mark and stops choosing the worker within milliseconds, so new requests - go elsewhere while it is still running. + go elsewhere while it is still running and long before it is told to stop. 2. **It keeps serving.** The worker finishes the generations it already - accepted, bounded by `--drain-timeout`, and only then deregisters and exits. + accepted, bounded by `--drain-timeout`, and only then exits. -Because the record is removed at the end rather than the beginning, the worker -stays visible in `/v1/workers` while it drains — an operator can see a rollout -progressing instead of workers appearing to crash. +The early mark is what makes this different from simply stopping a process. +The `preStop` delay that follows is not spent waiting for the router to notice +— that already happened — but letting work in progress finish before the +process is signalled at all. Deploying through the operator needs no configuration: it injects the `preStop` delay and sizes the termination grace period to cover the whole sequence. For hand-written manifests and the per-stage timings, see [Scaling a fleet](scaling.md). +## When the Pod is not being deleted + +A worker can also be stopped without its Pod going anywhere — a liveness probe +failing and restarting the container, a node being shut down gracefully, someone +killing the process. There is no deletion, so there is no early mark, and the +router has no way to know until the worker says so. + +On those paths the worker removes its own registration as its first act on +`SIGTERM`, which stops new requests arriving, and drains after. In-flight +generations still finish. What is missing is the head start: from the moment the +decision is made to the moment the process is signalled, the router is still +sending work, because nothing has told it otherwise. + ## Elsewhere Deployments outside Kubernetes use an external etcd for discovery, where a worker record is simply present or absent and nothing observes that a process is -leaving. Shutdown there still finishes in-flight work, but in the other order: -the worker deregisters first — which is what stops new requests arriving — and -drains after. In-flight generations are not cut; what is unavailable is the -early notice and the visible draining above. +leaving. Shutdown behaves as in the section above — deregister, then drain — +with in-flight work finished either way. The early notice is the part that needs +Kubernetes. + +In both cases the worker is absent from `/v1/workers` while it drains rather +than shown as draining, since removing the record is what stops new work +arriving. diff --git a/manual/features/scaling.md b/manual/features/scaling.md index 24b12135..52384113 100644 --- a/manual/features/scaling.md +++ b/manual/features/scaling.md @@ -59,13 +59,12 @@ grace period. The worker then, in this order: -1. **Stops being routed to.** The router filters draining workers out - immediately, so no new work arrives. The record stays, so the worker remains - visible in `/v1/workers` — a worker that vanishes looks exactly like one that - crashed. Under Kubernetes this already happened when the Pod was condemned, - before the process was signalled at all; elsewhere the worker deregisters - here, which stops new work but also removes it from the listing (see [Who - says the worker is leaving](#who-says-the-worker-is-leaving)). +1. **Stops being routed to.** The worker removes its registration, which is + what stops new work arriving. Under Kubernetes a Pod being *deleted* has + already left routing well before this — the registry acts on its + deletionTimestamp, before the process is even signalled — so this is only + cleanup there (see [Who says the worker is + leaving](#who-says-the-worker-is-leaving)). 2. **Drains.** On the NATS transport infera tracks in-flight requests directly. On HTTP the router talks straight to the engine, so infera asks the engine instead, polling its `/metrics` until running, queued, and PD-handoff queues @@ -83,10 +82,10 @@ a separate and much larger number, set by the longest generation it was already serving. Measured: under a second to stop receiving, 38 s until the record disappeared, while a 40-second generation ran to completion in between. -Watching `/v1/workers` measures the second one, not the first: it lists every -worker including draining ones, precisely so a rollout is visible while it -happens. To see the transition, read the `status` field rather than counting -rows. +Watching `/v1/workers` measures the second one, not the first. A Pod being +deleted appears there as `draining` for the window between its deletion being +requested and the process being signalled; after that it deregisters and +disappears while it finishes its remaining work. ```{note} `--drain-timeout` is a **ceiling, not a delay** — a worker with nothing in flight @@ -238,18 +237,19 @@ landing mid-drain would overwrite it with a payload that omits the status, which parses as `ACTIVE`, and the worker would be handed new work it is about to refuse. -**etcd: removing the record says so.** There is no orchestrator here. A record -is either present with an unexpired lease or it is gone — nothing observes that -the process is on its way out, and there is no third state to put it in. So a -shutdown deregisters *first*, which is what stops new requests arriving, and -drains after. In-flight generations are still finished rather than cut; what is -not available is the worker remaining visible while it does so. - -That ordering is the one real difference. Under Kubernetes routing has already -stopped by the time the process is signalled, so the record can be kept until -the end and the drain is observable; on etcd the record is the only thing -keeping the worker in rotation, so it has to go before the drain rather than -after. +**Everywhere else: removing the record says so.** On etcd there is no +orchestrator at all — a record is either present with an unexpired lease or it +is gone, with no third state to put it in. The same is true on Kubernetes +whenever the Pod is *not* being deleted: a liveness probe restarting the +container, a node shutting down gracefully, someone killing the process. No +deletionTimestamp is set, so the registry reads the annotation as usual and the +worker stays routable until it clears it. + +So every shutdown deregisters *first* and drains after. In-flight generations +are finished either way; the cost is that the worker is absent from +`/v1/workers` while it drains rather than shown as draining. The head start — +leaving routing before the process is signalled at all — is what deleting a Pod +buys, and only that. ```{warning} `discoveryBackend: etcd` **is not supported for in-cluster deployments** and the diff --git a/manual/getting_started/quickstart.md b/manual/getting_started/quickstart.md index 13efa0ad..7f5c8f73 100644 --- a/manual/getting_started/quickstart.md +++ b/manual/getting_started/quickstart.md @@ -82,10 +82,10 @@ three flags on **every** server and worker, or they won't find each other. See [Routing & transport](../features/routing_and_transport.md). One behaviour differs on this path: stopping a worker still lets its in-flight -generations finish, but it deregisters first, so it disappears from -`/v1/workers` while draining instead of being visibly on its way out. Leaving -rotation *before* the process is signalled needs Kubernetes to say the Pod is -going — see [Graceful shutdown](../features/graceful_shutdown.md). +generations finish, but nothing takes it out of routing until it is signalled. +Leaving rotation *before* the process is told to stop needs Kubernetes to +report the Pod as going away — see +[Graceful shutdown](../features/graceful_shutdown.md). ``` ```{tip} From 27fa4444ee4fcaac172156ace7d8baa2003edb29 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Mon, 10 Aug 2026 10:57:29 +0000 Subject: [PATCH 86/88] fix(drain): report an etcd revoke the server refused, and pin the order Two things the previous three commits claimed but did not deliver. etcd's deregister() never checked the status code. httpx does not raise on 4xx/5xx, so its except branch was unreachable and a refused revoke -- "lease not found", or a proxy answering 503 -- returned True and logged "deregistered worker (lease N revoked)" at INFO. The lease then outlives the call by up to its TTL, 30s by default, which covers the whole drain window. That is exactly the silent failure 2d25697 set out to remove, in its likeliest form. test_shutdown_order.py asserted an attribute was absent, which is the same shape of test the commit message criticised the old ones for being. Swapping deregister and drain in any entrypoint left all 1331 tests green. It now reads the order out of the three entrypoints and fails when they disagree, verified by reintroducing both regressions. Also fixes the documentation and comments ed5fa6c missed: scaling.md still listed "deregisters, then stops the engine" as step 3 of a sequence whose step 1 had been rewritten to deregister first, drain.py's module docstring still called the old order "the whole point", and three comments plus a fakeworker log sample still described the worker announcing DRAINING, which it has never done. The measured 38s in scaling.md described the old behaviour and is gone rather than replaced, since the record now disappears when the drain starts. The caller-side error was unconditional but false on the common path: a Pod being deleted left routing at its deletionTimestamp, long before this. It now defers to deregister()'s own log, which already carries that qualification. atom picks up the return value it was dropping. Signed-off-by: leiwei12 Co-authored-by: Cursor --- .../operator/internal/controller/builders.go | 4 +- .../controller/discovery_backend_test.go | 2 +- .../controller/inferadeployment_controller.go | 5 +- infera/common/discovery_k8s.py | 21 +-- infera/common/registration.py | 7 +- infera/engine/atom/__main__.py | 6 +- infera/engine/drain.py | 12 +- infera/engine/sglang/__main__.py | 3 +- infera/engine/vllm/__main__.py | 3 +- infera/tools/fakeworker/README.md | 3 +- infera/tools/fakeworker/server.py | 3 +- manual/features/graceful_shutdown.md | 15 +- manual/features/scaling.md | 18 +-- .../common/test_discovery_k8s_terminating.py | 7 +- tests/unit/common/test_shutdown_order.py | 131 ++++++++++++++---- 15 files changed, 169 insertions(+), 71 deletions(-) diff --git a/deploy/operator/internal/controller/builders.go b/deploy/operator/internal/controller/builders.go index 1f4081f8..c28c3376 100644 --- a/deploy/operator/internal/controller/builders.go +++ b/deploy/operator/internal/controller/builders.go @@ -37,8 +37,8 @@ const ( // Graceful rolling-upgrade tuning for GPU worker pods. workerPreStopDrainSeconds = 15 // preStop sleep: let the router drop us before SIGTERM workerDefaultDrainTimeoutSeconds = 30 // matches the worker's --drain-timeout default - // Teardown after the drain finishes: deregistering, stopping the KV plane, - // and engine.stop(), which SIGTERMs the engine's process group and waits up + // Teardown after the drain finishes: stopping the KV plane and + // engine.stop(), which SIGTERMs the engine's process group and waits up // to 30s before escalating to SIGKILL. workerTeardownHeadroomSeconds = 50 // Floor, so short drain timeouts still leave room for a slow engine exit. diff --git a/deploy/operator/internal/controller/discovery_backend_test.go b/deploy/operator/internal/controller/discovery_backend_test.go index 17f36eb5..f44ff27e 100644 --- a/deploy/operator/internal/controller/discovery_backend_test.go +++ b/deploy/operator/internal/controller/discovery_backend_test.go @@ -28,7 +28,7 @@ import ( // // Pointing such a deployment at an external etcd throws that away. The router // stops watching Pods, so nothing reads the deletionTimestamp, and the only -// remaining signal is the worker announcing DRAINING after SIGTERM -- which +// remaining signal is the worker's record disappearing on SIGTERM -- which // arrives once the preStop delay the operator itself injects has elapsed. The // combination keeps that delay and loses the early notice it exists to give, // so for its whole duration the router keeps handing new work to a Pod that is diff --git a/deploy/operator/internal/controller/inferadeployment_controller.go b/deploy/operator/internal/controller/inferadeployment_controller.go index 699beaa2..0e5f361e 100644 --- a/deploy/operator/internal/controller/inferadeployment_controller.go +++ b/deploy/operator/internal/controller/inferadeployment_controller.go @@ -76,8 +76,9 @@ func (r *InferaDeploymentReconciler) Reconcile(ctx context.Context, req ctrl.Req // // Pointing such a deployment at an external etcd discards that. The server // stops watching Pods, so nothing reads the deletionTimestamp, and the only - // remaining signal is the worker announcing DRAINING after SIGTERM -- which - // it receives only once the preStop delay injected below has elapsed. The + // remaining signal is the worker's record disappearing when it deregisters + // on SIGTERM -- which it receives only once the preStop delay injected + // below has elapsed. The // combination keeps that delay while losing the early notice it exists to // provide, so for its whole duration the router keeps handing new work to a // Pod already on its way out. Refusing beats rendering a deployment whose diff --git a/infera/common/discovery_k8s.py b/infera/common/discovery_k8s.py index 1b993a01..69704df8 100644 --- a/infera/common/discovery_k8s.py +++ b/infera/common/discovery_k8s.py @@ -245,8 +245,9 @@ def _handle_pod(self, pod: dict, *, deleted: bool) -> None: annotations = meta.get("annotations") or {} raw = annotations.get(WORKER_INFO_ANNOTATION) - # Gone for good: explicit DELETE, annotation cleared (the worker - # deregistered, so its drain is over), or no longer Running. + # Not a routing candidate any more: explicit DELETE, annotation cleared + # (the worker deregistered, which is how its drain begins), or no longer + # Running. if deleted or raw is None or not self._pod_running(pod): worker_id = self._pod_to_worker.pop(pod_name, None) if worker_id is not None: @@ -337,14 +338,14 @@ def _mark_draining(self, worker_id: str) -> None: ``list_active`` filters DRAINING, so this stops new work reaching the worker just as removal would -- but ``list_all`` still shows it, and - that difference is the whole point. Deleting the record makes a Pod - that is finishing its in-flight generations look exactly like one that - crashed, so ``/v1/workers`` cannot distinguish an orderly rollout from - a fleet losing workers, at precisely the moment someone is watching. - - The record does not linger: the worker clears its own annotation when - the drain completes, which lands here as "annotation gone" and removes - it for real. + that difference is what makes a condemned Pod visible as such on + ``/v1/workers`` instead of looking like one that crashed. The window is + the preStop delay: from the deletion being requested to the process + being signalled. + + The record does not linger: the worker clears its own annotation on + SIGTERM, before draining, which lands here as "annotation gone" and + removes it for real. Callbacks fire here rather than at that later removal because routing is what they act on -- the KV subscriber and the policy's block diff --git a/infera/common/registration.py b/infera/common/registration.py index 0262ac77..631f0b1e 100644 --- a/infera/common/registration.py +++ b/infera/common/registration.py @@ -117,7 +117,12 @@ async def deregister(self) -> bool: ok = True if self._lease_id is not None: try: - await self._http.post("/v3/lease/revoke", json={"ID": self._lease_id}) + r = await self._http.post("/v3/lease/revoke", json={"ID": self._lease_id}) + # httpx does not raise on 4xx/5xx, and etcd answering "lease not + # found" or a proxy returning 503 is exactly the failure this + # reports -- without this the refusal would be logged as a + # successful revoke. + r.raise_for_status() logger.info( "deregistered worker %s (lease %d revoked)", self._worker_id, diff --git a/infera/engine/atom/__main__.py b/infera/engine/atom/__main__.py index 1be94a06..a0c6f123 100644 --- a/infera/engine/atom/__main__.py +++ b/infera/engine/atom/__main__.py @@ -156,7 +156,11 @@ async def main() -> None: except asyncio.CancelledError: pass - await reg_client.deregister() + if not await reg_client.deregister(): + # deregister() already logged why. No drain step here, so nothing waits + # on the record being gone -- but a silent branch would be the wrong + # thing to inherit if one is ever added. + logger.warning("stopping anyway") await engine.stop() diff --git a/infera/engine/drain.py b/infera/engine/drain.py index 7d7c2901..7fc52537 100644 --- a/infera/engine/drain.py +++ b/infera/engine/drain.py @@ -19,12 +19,12 @@ Two behaviours are deliberate: -* **Announce first, then drain, then deregister.** Ordering is the whole point. - Draining while still a routing candidate just means more work arrives, so the - worker announces DRAINING -- which takes it out of the candidate list -- before - waiting for the work it already has. Deregistering comes last, because a record - that disappears at the start of the drain is indistinguishable from a worker - that crashed. +* **Deregister first, then drain.** Ordering is the whole point. Draining while + still a routing candidate just means more work arrives, and the count this + polls never reaches zero. Removing the record is what takes the worker out of + the candidate list, so it happens before waiting for the work already in hand. + The cost is that a worker finishing its in-flight requests is indistinguishable + from one that crashed; the alternative is a drain that cannot converge. * **An unreadable metric does not block shutdown.** If the engine's in-flight count cannot be determined — an unknown engine, a renamed series, a dead HTTP server — this logs loudly and returns rather than hanging until the timeout. diff --git a/infera/engine/sglang/__main__.py b/infera/engine/sglang/__main__.py index 0d1e3e75..c334e3fb 100644 --- a/infera/engine/sglang/__main__.py +++ b/infera/engine/sglang/__main__.py @@ -447,7 +447,8 @@ async def _drain() -> None: # The cost is that the worker is gone from /v1/workers while it finishes, # rather than visibly draining. if not await reg_client.deregister(): - logger.error("draining anyway, but new requests may still be routed here") + # deregister() already logged why, including whether it matters here. + logger.warning("draining anyway") await _drain() if kv_relay is not None: diff --git a/infera/engine/vllm/__main__.py b/infera/engine/vllm/__main__.py index c4668665..fe9b2bc5 100644 --- a/infera/engine/vllm/__main__.py +++ b/infera/engine/vllm/__main__.py @@ -393,7 +393,8 @@ async def _drain() -> None: # The cost is that the worker is gone from /v1/workers while it finishes, # rather than visibly draining. if not await reg_client.deregister(): - logger.error("draining anyway, but new requests may still be routed here") + # deregister() already logged why, including whether it matters here. + logger.warning("draining anyway") await _drain() if kv_relay is not None: await kv_relay.stop() diff --git a/infera/tools/fakeworker/README.md b/infera/tools/fakeworker/README.md index 92370270..1739985c 100644 --- a/infera/tools/fakeworker/README.md +++ b/infera/tools/fakeworker/README.md @@ -108,9 +108,8 @@ Shutdown then goes through the real NATS drain — unsubscribe first, then wait the in-flight set infera actually holds: ``` -worker 127.0.0.1:19951 announced DRAINING +deregistered worker 127.0.0.1:19951 (lease 7587883597149818 revoked) draining 1 in-flight NATS request(s), up to 60s -deregistered worker 127.0.0.1:19951 ``` ```{note} diff --git a/infera/tools/fakeworker/server.py b/infera/tools/fakeworker/server.py index 63ff4d66..69ceea87 100644 --- a/infera/tools/fakeworker/server.py +++ b/infera/tools/fakeworker/server.py @@ -520,7 +520,8 @@ async def _drain() -> None: # the record is what stops new work arriving, on either backend. try: if not await reg.deregister(): - logger.error("draining anyway, but new requests may still be routed here") + # deregister() already logged why, including whether it matters. + logger.warning("draining anyway") except Exception as exc: # noqa: BLE001 - shutdown must not raise logger.warning("deregister failed: %s", exc) await _drain() diff --git a/manual/features/graceful_shutdown.md b/manual/features/graceful_shutdown.md index 01e47b91..9559bd4c 100644 --- a/manual/features/graceful_shutdown.md +++ b/manual/features/graceful_shutdown.md @@ -6,13 +6,16 @@ finishes the generations it already accepted before the process exits. **Why:** a severed generation cannot be retried — the tokens already streamed cannot be un-sent — so without this, every rolling upgrade or scale-down -produces a burst of client errors. **Requires:** Kubernetes, with the default -`kubernetes` discovery backend. +produces a burst of client errors. **Requires:** nothing, for finishing in-flight +work; Kubernetes with the default `kubernetes` discovery backend for the advance +notice described below. ``` ```{important} -Supported **only on Kubernetes with the default `kubernetes` discovery -backend**. This is not an implementation gap: it relies on the orchestrator +Finishing in-flight work happens on every backend, bounded by `--drain-timeout`. +What needs **Kubernetes with the default `kubernetes` discovery backend** is the +*advance* notice — the router learning a worker is leaving before the process is +signalled. That is not an implementation gap: it relies on the orchestrator knowing a Pod is being removed, which nothing outside Kubernetes can tell the router. `discoveryBackend: etcd` is rejected by the operator for in-cluster deployments. @@ -63,4 +66,6 @@ Kubernetes. In both cases the worker is absent from `/v1/workers` while it drains rather than shown as draining, since removing the record is what stops new work -arriving. +arriving. A Pod being deleted does show as `draining`, but earlier — between its +deletion being requested and the process being signalled, before the drain +itself begins. diff --git a/manual/features/scaling.md b/manual/features/scaling.md index 52384113..80bfc060 100644 --- a/manual/features/scaling.md +++ b/manual/features/scaling.md @@ -69,7 +69,7 @@ The worker then, in this order: On HTTP the router talks straight to the engine, so infera asks the engine instead, polling its `/metrics` until running, queued, and PD-handoff queues all reach zero. Bounded by `--drain-timeout` (default 30 s). -3. **Deregisters**, then stops the engine. +3. **Stops the engine.** Requests already in flight run to completion. Requests that arrive during the drain go to other workers. @@ -79,13 +79,15 @@ requests within a second of the shutdown starting — the router's watch picking up either the `deletionTimestamp` or the record's removal — and it is the number that decides whether traffic is still being sent somewhere that is about to die. How long the *process* then lives is a separate and much larger number, set by the longest generation it was already -serving. Measured: under a second to stop receiving, 38 s until the record -disappeared, while a 40-second generation ran to completion in between. - -Watching `/v1/workers` measures the second one, not the first. A Pod being -deleted appears there as `draining` for the window between its deletion being -requested and the process being signalled; after that it deregisters and -disappears while it finishes its remaining work. +serving. Measured: under a second to stop receiving, while a 40-second +generation ran to completion afterwards. + +Watching `/v1/workers` measures neither. The record now goes when the drain +*starts*, not when it ends, so its disappearance marks the beginning of the +in-flight work rather than the end of it — a worker finishing a long generation +is absent from that list for all of it. A Pod being deleted shows as `draining` +only for the window between its deletion being requested and the process being +signalled, which is the preStop hook's 15 s. ```{note} `--drain-timeout` is a **ceiling, not a delay** — a worker with nothing in flight diff --git a/tests/unit/common/test_discovery_k8s_terminating.py b/tests/unit/common/test_discovery_k8s_terminating.py index eeb91676..65d31bd8 100644 --- a/tests/unit/common/test_discovery_k8s_terminating.py +++ b/tests/unit/common/test_discovery_k8s_terminating.py @@ -135,9 +135,10 @@ def test_a_draining_worker_stays_visible(): assert _all(reg) == {"10.0.0.1:8080": WorkerStatus.DRAINING} -def test_the_record_goes_when_the_drain_finishes(): - """The worker clears its own annotation once drained, which lands here as - 'annotation gone'. Without that the draining record would be immortal.""" +def test_the_record_goes_when_the_worker_deregisters(): + """The worker clears its own annotation on SIGTERM, before draining, which + lands here as 'annotation gone'. Without that the draining record would be + immortal.""" reg, removed = _registry() reg._handle_pod(_pod(), deleted=False) reg._handle_pod(_pod(terminating=True), deleted=False) diff --git a/tests/unit/common/test_shutdown_order.py b/tests/unit/common/test_shutdown_order.py index 47f0d081..77db8938 100644 --- a/tests/unit/common/test_shutdown_order.py +++ b/tests/unit/common/test_shutdown_order.py @@ -3,45 +3,106 @@ # # SPDX-License-Identifier: MIT ############################################################################### -"""The shutdown sequence, driven rather than asserted from constants. +"""The shutdown sequence: deregister, then drain. -Nothing here inspects a flag saying what the order should be; these record the -order the entrypoints actually perform and the reason it has to be that one. +Removing the record is what stops new work arriving, so it has to happen before +waiting on work already in flight. These pin that order in the entrypoints +themselves and pin that a deregistration which did not happen says so. """ from __future__ import annotations +import ast +import inspect +from pathlib import Path + import pytest from infera.common.registration import RegistrationClient from infera.common.registration_k8s import K8sRegistrationClient +ENTRYPOINTS = ( + "infera/engine/vllm/__main__.py", + "infera/engine/sglang/__main__.py", + "infera/tools/fakeworker/server.py", +) + -def test_deregistering_is_what_stops_new_work_on_every_backend(): - """Both backends stop new work by removing the record, so both must remove - it before waiting on in-flight work. +def _shutdown_call_order(path: Path) -> list[str]: + """The two calls that matter, in source order. - On etcd nothing else can express departure. On Kubernetes the registry does - drop a Pod on its deletionTimestamp -- but only when the Pod is being - deleted. A liveness-probe restart, a node graceful shutdown or a manual kill - all send SIGTERM with the Pod object untouched, and on those paths the - annotation is still present and still parsed, so the worker stays routable - until it clears it. Draining first there means the router keeps assigning - work for the whole drain window. + Reduced to deregistering and draining, so reordering anything else in the + shutdown sequence does not make this fail. Calling `_drain` is a Call node + while defining it is not, so the definition is not counted. """ - for client in (RegistrationClient, K8sRegistrationClient): - assert not hasattr(client, "deregister_stops_routing"), ( - f"{client.__name__} still declares a per-backend order; both now " - "deregister before draining" - ) + seen = [] + for call in (n for n in ast.walk(ast.parse(path.read_text())) if isinstance(n, ast.Call)): + func = call.func + if isinstance(func, ast.Attribute) and func.attr == "deregister": + seen.append(("deregister", call.lineno)) + elif isinstance(func, ast.Name) and func.id == "_drain": + seen.append(("drain", call.lineno)) + # ast.walk is breadth-first, so sort back into source order. + return [name for name, _ in sorted(seen, key=lambda p: p[1])] + + +@pytest.mark.parametrize("entrypoint", ENTRYPOINTS) +def test_every_entrypoint_deregisters_before_it_drains(entrypoint): + """Draining while still registered means the router keeps assigning work + for the whole drain window, so the drain never converges. + + On etcd removing the record is the only way to express departure. On + Kubernetes the registry does drop a Pod on its deletionTimestamp, but only + when the Pod is being deleted -- a liveness-probe restart, a node graceful + shutdown or a manual kill all send SIGTERM with the Pod object untouched, + leaving the annotation present and the worker routable until it clears it. + """ + root = Path(inspect.getfile(RegistrationClient)).parents[2] + order = _shutdown_call_order(root / entrypoint) + + assert "deregister" in order, f"{entrypoint} never deregisters" + assert "drain" in order, f"{entrypoint} never drains" + assert order.index("deregister") < order.index("drain"), ( + f"{entrypoint} drains before deregistering, so new work keeps arriving " + "for the whole drain window" + ) + + +class _Resp: + def __init__(self, code: int): + self.status_code = code + + def raise_for_status(self): + if self.status_code >= 400: + raise RuntimeError(f"HTTP {self.status_code}") @pytest.mark.asyncio -async def test_a_failed_deregistration_is_reported(): - """Deregistering is the step that stops new work, so swallowing its failure - means draining while still being routed to -- silently.""" - calls = [] +async def test_etcd_reports_a_revoke_the_server_refused(): + """httpx does not raise on 4xx/5xx, so a refused revoke reaches the same + code path as a successful one. etcd answering 'lease not found', or a proxy + answering 503, is the likeliest way this fails.""" + class _Http: + async def post(self, path, json=None): # noqa: A002 - mirrors httpx + return _Resp(500) + + async def aclose(self): + pass + + # __new__, so the real httpx client is never built and never leaked. + c = RegistrationClient.__new__(RegistrationClient) + c._http = _Http() + c._lease_id, c._key, c._worker_id, c._lease_ttl = 1, "/k", "w", 30 + + assert await c.deregister() is False, ( + "etcd refused the revoke, so the lease is still alive and the worker is " + "still routable -- that cannot report success" + ) + + +@pytest.mark.asyncio +async def test_etcd_reports_an_unreachable_server(): class _Http: async def post(self, path, json=None): # noqa: A002 - mirrors httpx raise RuntimeError("etcd unreachable") @@ -49,10 +110,26 @@ async def post(self, path, json=None): # noqa: A002 - mirrors httpx async def aclose(self): pass - c = RegistrationClient("http://etcd:2379") + # __new__, so the real httpx client is never built and never leaked. + c = RegistrationClient.__new__(RegistrationClient) c._http = _Http() - c._lease_id, c._key, c._worker_id = 1, "/k", "w" + c._lease_id, c._key, c._worker_id, c._lease_ttl = 1, "/k", "w", 30 + + assert await c.deregister() is False + + +@pytest.mark.asyncio +async def test_kubernetes_reports_a_patch_that_failed(): + """Clearing the annotation is what takes the worker out of the pool on the + paths where the Pod object is untouched.""" + c = K8sRegistrationClient.__new__(K8sRegistrationClient) + c._worker_id = "w" + c._pod_name = "p" + c._namespace = "ns" + + async def _patch(*_args, **_kwargs): + raise RuntimeError("apiserver unreachable") + + c._patch_annotation = _patch - ok = await c.deregister() - assert ok is False, "a deregistration that did not happen must not report success" - assert calls == [] + assert await c.deregister() is False From 24bc5082263d1273977099b4fc983cb1a09266cd Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Mon, 10 Aug 2026 11:06:07 +0000 Subject: [PATCH 87/88] test(discovery): cover the kill path that has no deletionTimestamp A liveness probe restarting the container, a node shutting down gracefully or someone killing the process all send SIGTERM with the Pod object untouched. That is the path deregistering-before-draining exists for -- with no deletionTimestamp, the annotation the worker clears is the only signal the router gets -- and it had no coverage of the removal callback, which is what stops the KV subscriber. The existing test asserted only that the worker left the pool. Verified correct before writing it: fires exactly once, and a later DELETE for the same Pod does not repeat it. Signed-off-by: leiwei12 Co-authored-by: Cursor --- .../common/test_discovery_k8s_terminating.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/unit/common/test_discovery_k8s_terminating.py b/tests/unit/common/test_discovery_k8s_terminating.py index 65d31bd8..05d49727 100644 --- a/tests/unit/common/test_discovery_k8s_terminating.py +++ b/tests/unit/common/test_discovery_k8s_terminating.py @@ -149,6 +149,26 @@ def test_the_record_goes_when_the_worker_deregisters(): assert removed == ["10.0.0.1:8080"], "announced once, not once per stage" +def test_a_worker_killed_without_its_pod_being_deleted_is_announced_once(): + """The path deregistering-before-draining exists for. + + A liveness probe restarting the container, a node shutting down gracefully, + someone killing the process -- all send SIGTERM with the Pod object + untouched. There is no deletionTimestamp, so the annotation the worker + clears on its way out is the only signal, and the callback behind it is what + stops the KV subscriber. A later DELETE for the same Pod must not repeat it. + """ + reg, removed = _registry() + reg._handle_pod(_pod(), deleted=False) + + reg._handle_pod(_pod(annotated=False), deleted=False) + assert _ids(reg) == [], "clearing the annotation must stop new work arriving" + assert removed == ["10.0.0.1:8080"], "the only signal on this path must announce" + + reg._handle_pod(_pod(annotated=False), deleted=True) + assert removed == ["10.0.0.1:8080"], "the eventual DELETE must not announce again" + + def test_announced_once_across_draining_then_delete(): """The callbacks stop a KV subscriber and clear block accounting. Firing them twice for one worker is not free, and firing them late (only at the From 0efc8e7fc9e49347eada15a69f696025311b3127 Mon Sep 17 00:00:00 2001 From: leiwei12 Date: Mon, 10 Aug 2026 11:52:00 +0000 Subject: [PATCH 88/88] revert(tools): drop the fake worker from the release package It shipped inside infera/, so `include = ["infera*"]` put it in the wheel and registered an `infera-fake-worker` console script. A tool that joins real service discovery, advertises any address it likes and serves fabricated responses does not belong in what users install, and the opt-in env var it carried mitigates that rather than removing it. Removes the package, its tests, the console script, and the three measurements in scaling.md that named it -- the numbers are real and stay, but they no longer point at something the repository does not contain. Two comments listing it as a consumer go the same way. test_shutdown_order.py drops to the two real entrypoints. Re-verified it still fails when vllm's deregister and drain are swapped, which is the property it exists for. Signed-off-by: leiwei12 Co-authored-by: Cursor --- infera/common/engine_metrics.py | 6 +- infera/router/dp_routing.py | 4 +- infera/tools/fakeworker/README.md | 162 ------- infera/tools/fakeworker/__init__.py | 6 - infera/tools/fakeworker/__main__.py | 9 - infera/tools/fakeworker/server.py | 586 ----------------------- manual/features/scaling.md | 6 +- pyproject.toml | 2 - tests/unit/common/test_shutdown_order.py | 1 - tests/unit/tools/test_fake_worker.py | 261 ---------- 10 files changed, 8 insertions(+), 1035 deletions(-) delete mode 100644 infera/tools/fakeworker/README.md delete mode 100644 infera/tools/fakeworker/__init__.py delete mode 100644 infera/tools/fakeworker/__main__.py delete mode 100644 infera/tools/fakeworker/server.py delete mode 100644 tests/unit/tools/test_fake_worker.py diff --git a/infera/common/engine_metrics.py b/infera/common/engine_metrics.py index 858145b8..712a021a 100644 --- a/infera/common/engine_metrics.py +++ b/infera/common/engine_metrics.py @@ -7,9 +7,9 @@ Every engine exposes the same three facts — requests running, requests queued, KV cache in use — under a different name, and the names drift between releases. -Anything that reads them (graceful drain, an autoscaler, the fake worker) needs -the same mapping, so it lives here rather than being spelled out at each call -site where one of them would quietly rot. +Anything that reads them (graceful drain, an autoscaler) needs the same mapping, +so it lives here rather than being spelled out at each call site where one of +them would quietly rot. Provenance, because it is uneven and matters: diff --git a/infera/router/dp_routing.py b/infera/router/dp_routing.py index 6ee0397b..c91214b7 100644 --- a/infera/router/dp_routing.py +++ b/infera/router/dp_routing.py @@ -17,8 +17,8 @@ #: Both engines honour this (SGLang ``DataParallelController``, vLLM #: ``_get_data_parallel_rank``, case-insensitively). Named rather than inlined -#: so anything asserting on it -- tests, the fake worker -- breaks loudly on a -#: rename instead of silently never matching. Mirrors Rust's DP_RANK_HEADER. +#: so anything asserting on it breaks loudly on a rename instead of silently +#: never matching. Mirrors Rust's DP_RANK_HEADER. DP_RANK_HEADER = "X-Data-Parallel-Rank" diff --git a/infera/tools/fakeworker/README.md b/infera/tools/fakeworker/README.md deleted file mode 100644 index 1739985c..00000000 --- a/infera/tools/fakeworker/README.md +++ /dev/null @@ -1,162 +0,0 @@ -# Fake worker - -A worker that joins the fleet and serves tokens, with no GPU and no weights. - -Everything above the engine — discovery, routing, failover, the circuit breaker, -drain, autoscaling — is testable without a model. What blocked that was simply -that there was no way to *be* a worker without loading one: `tests/e2e/harness` -starts real engines in containers, so the cheapest fleet anyone could build cost -a GPU and a multi-minute weight load per member. - -```bash -export INFERA_ALLOW_FAKE_WORKER=1 -infera-fake-worker --model-name my-model --port 9101 \ - --discovery-backend etcd --etcd-endpoint http://127.0.0.1:2379 -``` - -`INFERA_ALLOW_FAKE_WORKER` is required, and the tool refuses to start without -it. It ships in the same package as the server and registers through the real -registration clients, so it can join a fleet and advertise any address it likes -— which the router will then dial, sending it real prompts. That needs no -privilege a worker does not already have, but it should be a deliberate act -rather than something a stray command does by default. - -## Why it can be trusted - -It registers through the **real** registration clients with a real -`EngineConfig`, so the payload is built by `build_worker_payload` — the same -function every engine uses — and parsed back by `worker_info_from_json`, the -same function every discovery backend uses. If that contract changes, this -changes with it or the test suite fails. Only what happens *after* a request -arrives is faked. - -## The knobs that matter - -Most flags are obvious. These three exist because they make otherwise expensive -problems reproducible on a laptop: - -| Flag | Why | -|---|---| -| `--startup-delay-s` | Simulates the 5–15 minute weight load. `/health` stays 503 until it elapses. This is the single biggest reason naive autoscaling overshoots — an unready replica still counts in the fleet but consumes 0% of the metric, so the loop keeps asking for more. Reproducing it here costs nothing; reproducing it on real hardware costs a GPU-hour per iteration. | -| `--max-concurrency` | Gives requests somewhere to queue. `num_requests_waiting` is the metric the entire industry autoscales on, and without a queue it is identically zero — so a scaling test against a fake fleet would pass without testing anything. | -| `--fail-first N` | Refuse the first N requests with 503, then recover. Drives the router's circuit breaker through open → half-open → closed without killing a process. | - -`/metrics` exposes engine-native names (`vllm:num_requests_waiting`, -`sglang:num_queue_reqs`, …) chosen by `--engine`, so a scaling rule written -against fakes transfers to a real fleet unchanged. - -> The vLLM names are from its published metrics documentation. **The SGLang -> names are second-hand from a research pass and have not been checked against a -> live SGLang.** Verify before depending on them. - -## PD and DP attention - -Both work, and both are verified end to end against a real router with no GPU. - -**PD.** `--disagg-mode prefill|decode` advertises the same `disagg_meta` a real -worker does — `{"protocol": ..., "params": {"bootstrap_addr": ...}}`, prefill -only — with `--pd-protocol` constrained to the router's own protocol registry so -a typo dies at argparse rather than at the first request. - -**DP attention.** Two deployment shapes exist and they behave differently, which -is worth knowing before concluding anything is broken: - -| Shape | Registration | What the router does | -|---|---|---| -| per-rank endpoints | `--dp-rank R --dp-size N` | Nothing. The address already selects the rank, so no header is pinned and no room alignment happens. `dp_rank` on the `RouteTarget` stays `None`. | -| rank-multiplexed | `--dp-size N`, **no** `--dp-rank` | `expand_targets` fans one worker into N targets; the router pins `X-Data-Parallel-Rank`, aligns `bootstrap_room % dp_size == dp_rank`, and injects `disagg_prefill_dp_rank` for the decode leg. | - -`is_rank_multiplexed()` is `dp_size > 1 and dp_rank is None` — so registering a -rank makes the worker an endpoint and opts *out* of router-side DP routing. That -is correct, not a bug, but the first time you see it the DP path looks dead. - -### `GET /debug/routing` - -The reason PD and DP are testable at all. It reports what the *router* decided — -per-rank request counts keyed on the header it sent, and the handoff fields it -injected on the last request: - -```json -{"dp_rank": null, "dp_size": 4, - "requests_by_dp_rank": {"0": 2, "1": 2, "2": 2, "3": 2}, - "last_handoff": {"bootstrap_host": "127.0.0.1", "bootstrap_port": 18430, - "bootstrap_room": 7442254485466660987, - "disagg_prefill_dp_rank": 3}} -``` - -None of that is observable with a real engine: a malformed handoff does not -raise, it hangs on KVPoll until a ~300 s timeout, and the failure surfaces -nowhere near the router that caused it. Here you can assert on it directly — -e.g. that `bootstrap_room % dp_size == disagg_prefill_dp_rank` holds on every -request, which is the invariant SGLang's `follow_bootstrap_room` balancer -enforces with a `KVTransferError`. - -## NATS transport - -`--request-transport nats` routes requests through a broker instead of having -the router dial this worker. It uses the **real** `NatsRequestServer`, which -proxies to this process's own HTTP surface exactly as it proxies to a real -engine's — so the transport under test is the production one, not a stand-in. - -```bash -INFERA_ALLOW_FAKE_WORKER=1 infera-fake-worker --model-name m --port 9101 \ - --request-transport nats --nats-server nats://127.0.0.1:4222 \ - --discovery-backend etcd --etcd-endpoint http://127.0.0.1:2379 -``` - -Shutdown then goes through the real NATS drain — unsubscribe first, then wait on -the in-flight set infera actually holds: - -``` -deregistered worker 127.0.0.1:19951 (lease 7587883597149818 revoked) -draining 1 in-flight NATS request(s), up to 60s -``` - -```{note} -**The fake's HTTP drain is not representative of a real worker's.** This process -serves the requests itself, so it knows its own in-flight count exactly. A real -worker on HTTP transport does not: the router talks straight to the engine, so -infera has to poll the engine's `/metrics` and wait out a settle window because -those gauges lag. Comparing the fake's HTTP drain against its NATS drain -therefore measures nothing — both are exact. The difference only shows up with a -real engine. -``` - -## Limits — read these before drawing conclusions - -**No KV transfer is simulated.** A `--disagg-mode prefill` / `decode` fake takes -part in pool membership and routing, and that is all. No KV moves between the -legs. Useful for testing that P and D pools exist, are discovered, and are -routed to independently; **useless for testing Mooncake, bootstrap handshakes, -or anything about the transfer itself.** Do not conclude that PD "works" because -a fake fleet answered. - -**`--kv` synthesizes a tokenizer canary from the model name.** All fakes for one -model agree, which is what they need to do — `CanaryVerifier` silently drops a -worker whose canary differs from the first-registered one, so disagreeing fakes -would produce a fleet that is half the expected size for no visible reason. But -that synthetic canary will never match a **real** worker's, so do not mix fakes -and real workers under one model name. Without `--kv` there is no canary at all -and the fakes join anything. - -**It answers instantly by construction.** `--ttft-ms` and `--itl-ms` are a -latency model, not a performance model: TTFT scales linearly with prompt length -and nothing contends for memory. Do not use it to predict real throughput. - -## Gotchas found while building this - -- **`--advertise-host` must be routable from the router.** It defaults to - `$POD_IP`. `0.0.0.0` registers a URL no peer can reach. -- **The server requires `--router-tokenizer-path` even for `round-robin`**, and - resolves it eagerly. Any existing directory satisfies it, which is enough to - bring a router up against fakes. -- **A bind failure used to still register.** uvicorn logs `address already in - use` and gives up, but the process keeps running — so a port collision - produced a worker that was in the pool and served nothing. Now the socket must - be bound before registration, and a collision exits 3. Worth remembering - because the symptom was a router error about the *worker* returning garbage, - which pointed nowhere near the real cause. -- **Registration alone is not enough** — `heartbeat_loop()` has to be running or - the etcd lease (30 s) expires and the worker vanishes from the pool about half - a minute after it appears. This looks exactly like a discovery bug and is not - one. The real worker entrypoint starts it too; so does this. diff --git a/infera/tools/fakeworker/__init__.py b/infera/tools/fakeworker/__init__.py deleted file mode 100644 index 7b323fac..00000000 --- a/infera/tools/fakeworker/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -############################################################################### -# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. -# -# SPDX-License-Identifier: MIT -############################################################################### -"""GPU-free worker for testing everything above the engine.""" diff --git a/infera/tools/fakeworker/__main__.py b/infera/tools/fakeworker/__main__.py deleted file mode 100644 index e0171656..00000000 --- a/infera/tools/fakeworker/__main__.py +++ /dev/null @@ -1,9 +0,0 @@ -############################################################################### -# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. -# -# SPDX-License-Identifier: MIT -############################################################################### -from infera.tools.fakeworker.server import main - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/infera/tools/fakeworker/server.py b/infera/tools/fakeworker/server.py deleted file mode 100644 index 69ceea87..00000000 --- a/infera/tools/fakeworker/server.py +++ /dev/null @@ -1,586 +0,0 @@ -############################################################################### -# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. -# -# SPDX-License-Identifier: MIT -############################################################################### -"""A worker that registers into the fleet and serves tokens, without a GPU. - -Everything above the engine — discovery, routing, failover, the circuit breaker, -drain, and any autoscaling loop — is testable without weights. What blocked that -until now was simply that there was no way to *be* a worker without loading a -model: ``tests/e2e/harness`` starts real engines in containers, so the cheapest -fleet anyone could build cost a GPU and a multi-minute weight load per member. - -This is that missing piece. It registers through the **real** registration -clients with a real :class:`EngineConfig`, so the fake cannot drift from the -contract a genuine worker satisfies — if the payload changes, this changes with -it or fails loudly. What it fakes is only what happens after a request arrives. - -The parts worth having are the ones that make the *hard* problems reproducible: - -* ``--startup-delay-s`` simulates the 5-15 minute weight load. That delay is the - single biggest reason naive autoscaling overshoots — an unready replica counts - as consuming 0% of the metric, so the loop keeps asking for more. Reproducing - it costs nothing here and a GPU-hour on real hardware. -* ``--max-concurrency`` gives requests somewhere to queue, which is what makes - ``num_requests_waiting`` mean anything. Without a queue the metric everyone - scales on is identically zero. -* SIGTERM deregisters *before* draining, mirroring the real worker, so the - window where a terminating pod is still receiving traffic is observable. - -Not simulated, deliberately: real KV transfer. A PD fake accepts the bootstrap -fields and answers, but no KV moves between prefill and decode. That makes it -useful for testing pool membership, routing and scaling, and useless for testing -the transfer itself. Do not use it to conclude anything about Mooncake. -""" - -from __future__ import annotations - -import argparse -import asyncio -import hashlib -import json -import logging -import os -import signal -import time -from dataclasses import dataclass, field - -from fastapi import FastAPI, Request -from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse - -from infera.common.worker_pool import DisaggMode, EngineType, KvRegistrationMetadata -from infera.engine.base import EngineConfig -from infera.router.disagg_protocols import _PROTOCOLS -from infera.router.dp_routing import DP_RANK_HEADER - -logger = logging.getLogger("infera.fakeworker") - -# Roughly the shape of an English word, so token counts and byte counts stay in -# a believable ratio for anything measuring throughput. -_FILLER = "token " - - -@dataclass -class Behaviour: - """What this worker pretends about its own performance.""" - - ttft_ms: float = 50.0 - itl_ms: float = 10.0 - max_concurrency: int = 8 - #: Requests beyond max_concurrency wait here; this is what makes the - #: queue-depth metric — the one the whole industry scales on — non-zero. - max_kv_blocks: int = 1024 - blocks_per_request: int = 8 - fail_rate: float = 0.0 - #: Serve 5xx for the first N requests, then recover. For exercising the - #: circuit breaker's half-open probe without killing the process. - fail_first: int = 0 - - -@dataclass -class State: - running: int = 0 - waiting: int = 0 - served: int = 0 - failed: int = 0 - ready: bool = False - started_at: float = field(default_factory=time.monotonic) - _sem: asyncio.Semaphore | None = None - draining: bool = False - - #: Per-DP-rank request counts, keyed by the X-Data-Parallel-Rank header the - #: router sent. With a real engine you cannot easily see which rank the - #: router *intended* -- the engine just serves. Counting them here is what - #: makes DP-attention routing assertable at all. - by_dp_rank: dict[str, int] = field(default_factory=dict) - #: The PD handoff fields the router injected on the last request. Whether - #: the router shaped the body correctly is invisible from outside; a real - #: engine either works or hangs on KVPoll with no explanation. - last_handoff: dict = field(default_factory=dict) - - -def deterministic_canary(model_name: str) -> list[int]: - """A stand-in for the real tokenizer canary. - - Real workers tokenize a fixed probe string and register the ids, so that a - fleet running mismatched tokenizers under one model name is rejected at - registration rather than at inference time. A fake has no tokenizer, but it - still has to agree with *other fakes* for the same model, or the second one - to register is silently dropped from the pool. - - Deriving it from the model name gives exactly that: all fakes for a model - agree, and different models differ. It will not match a real worker's canary - — see the README; do not mix fakes and real workers under one model name. - """ - h = hashlib.sha256(f"infera-fake-canary::{model_name}".encode()).digest() - return [int.from_bytes(h[i : i + 2], "big") for i in range(0, 16, 2)] - - -def _rank_label(raw: str | None, dp_size: int) -> str: - """Constrain a DP-rank header to something safe to use as a metric label. - - The value is attacker-controlled and ends up interpolated into Prometheus - exposition text, where an embedded quote or newline would break out of the - label and forge series in whatever scrapes this endpoint. It is also a map - key, so accepting arbitrary strings grows that map without bound, one entry - per distinct value, and /debug/routing returns the whole thing. - - Only a small integer within the configured fan-out is meaningful here, so - everything else collapses to one bucket. - """ - if raw is None: - return "-" - try: - rank = int(raw) - except ValueError: - return "invalid" - if rank < 0 or (dp_size > 0 and rank >= dp_size): - return "invalid" - return str(rank) - - -def build_app(cfg: EngineConfig, behaviour: Behaviour, state: State) -> FastAPI: - app = FastAPI(title="infera fake worker") - - def _engine_metric(name: str) -> str: - # Engine-native metric names, so a scaling rule written against a fake - # fleet transfers to a real one unchanged. vLLM's names are from its - # published metrics doc; SGLang's are second-hand (see README) and worth - # checking against a live engine before relying on them. - prefix = "vllm" if cfg.engine == EngineType.VLLM else "sglang" - sglang = { - "num_requests_waiting": "num_queue_reqs", - "num_requests_running": "num_running_reqs", - "gpu_cache_usage_perc": "token_usage", - } - return f"{prefix}:{sglang[name] if prefix == 'sglang' else name}" - - async def _admit() -> bool: - """Returns False if this request should be rejected outright.""" - if state.draining: - return False - if behaviour.fail_first and state.served + state.failed < behaviour.fail_first: - state.failed += 1 - return False - return True - - async def _generate(prompt_tokens: int, max_tokens: int): - """Occupy a concurrency slot for a believable amount of time.""" - assert state._sem is not None - state.waiting += 1 - async with state._sem: - state.waiting -= 1 - state.running += 1 - try: - # TTFT scales with prompt length the way a real prefill does, - # so prefill-heavy and decode-heavy load look different to - # anything measuring them. - await asyncio.sleep(behaviour.ttft_ms / 1000.0 * max(1.0, prompt_tokens / 512)) - for _ in range(max_tokens): - yield _FILLER - await asyncio.sleep(behaviour.itl_ms / 1000.0) - state.served += 1 - finally: - state.running -= 1 - - def _parse(body: dict) -> tuple[int, int]: - text = json.dumps(body.get("messages") or body.get("prompt") or "") - prompt_tokens = max(1, len(text) // 4) - return prompt_tokens, int(body.get("max_tokens") or 16) - - @app.post("/v1/chat/completions") - @app.post("/v1/completions") - async def completions(request: Request): - body = await request.json() - - # Record what the router decided *before* deciding whether to serve, so - # a refused request still shows up in the routing evidence. - rank = _rank_label(request.headers.get(DP_RANK_HEADER), cfg.dp_size) - state.by_dp_rank[rank] = state.by_dp_rank.get(rank, 0) + 1 - handoff = { - k: body[k] - for k in ( - "bootstrap_host", - "bootstrap_port", - "bootstrap_room", - "disagg_prefill_dp_rank", - "kv_transfer_params", - ) - if k in body - } - if handoff: - state.last_handoff = handoff - - if not await _admit(): - return JSONResponse({"error": "fake worker refusing"}, status_code=503) - prompt_tokens, max_tokens = _parse(body) - - if body.get("stream"): - - async def sse(): - async for chunk in _generate(prompt_tokens, max_tokens): - payload = {"choices": [{"delta": {"content": chunk}}]} - yield f"data: {json.dumps(payload)}\n\n".encode() - yield b"data: [DONE]\n\n" - - return StreamingResponse(sse(), media_type="text/event-stream") - - out = "".join([c async for c in _generate(prompt_tokens, max_tokens)]) - return JSONResponse( - { - "id": f"fake-{state.served}", - "model": cfg.model_name, - "choices": [{"message": {"role": "assistant", "content": out}}], - "usage": { - "prompt_tokens": prompt_tokens, - "completion_tokens": max_tokens, - "total_tokens": prompt_tokens + max_tokens, - }, - } - ) - - @app.get("/v1/models") - async def models(): - return {"object": "list", "data": [{"id": cfg.model_name, "object": "model"}]} - - @app.get("/health") - async def health(): - # Unready until the simulated weight load finishes. This is the whole - # point of --startup-delay-s: a replica that exists but cannot serve is - # exactly what makes autoscalers overshoot. - if not state.ready: - return JSONResponse({"status": "loading"}, status_code=503) - return {"status": "ok", "running": state.running, "waiting": state.waiting} - - @app.get("/metrics") - async def metrics(): - used = min(1.0, (state.running * behaviour.blocks_per_request) / behaviour.max_kv_blocks) - lines = [ - f"{_engine_metric('num_requests_running')} {state.running}", - f"{_engine_metric('num_requests_waiting')} {state.waiting}", - f"{_engine_metric('gpu_cache_usage_perc')} {used:.4f}", - f"infera_fake_worker_served_total {state.served}", - f"infera_fake_worker_refused_total {state.failed}", - f"infera_fake_worker_ready {1 if state.ready else 0}", - f"infera_fake_worker_draining {1 if state.draining else 0}", - ] - for rank, n in sorted(state.by_dp_rank.items()): - lines.append(f'infera_fake_worker_requests_by_dp_rank{{dp_rank="{rank}"}} {n}') - return PlainTextResponse("\n".join(lines) + "\n") - - @app.get("/debug/routing") - async def routing(): - """What the router actually decided, which is otherwise unobservable. - - A real engine given a malformed PD handoff does not complain -- it hangs - on KVPoll until a ~300s timeout, and the failure surfaces nowhere near - the router that caused it. This turns that into an assertion. - """ - return { - "worker_id": f"{cfg.host}:{cfg.port}", - "disagg_mode": cfg.disagg_mode.value, - "dp_rank": cfg.dp_rank, - "dp_size": cfg.dp_size, - "requests_by_dp_rank": state.by_dp_rank, - "last_handoff": state.last_handoff, - } - - return app - - -def _disagg_meta(args) -> dict: - """Mirror what a real worker advertises. - - Only PREFILL carries a bootstrap endpoint; DECODE tags the protocol so the - router can fail fast on a cross-protocol pairing, and has nothing else to - say. Getting this wrong does not fail loudly at registration -- it fails at - the first PD request, as a protocol error that reads like a router bug. - """ - if args.disagg_mode == "mixed": - return {} - params: dict = {} - if args.disagg_mode == "prefill": - host = args.advertise_host or args.host - params["bootstrap_addr"] = f"{host}:{args.bootstrap_port}" - return {"protocol": args.pd_protocol, "params": params} - - -def build_config(args) -> EngineConfig: - kv = None - if args.kv: - canary = deterministic_canary(args.model_name) - kv = KvRegistrationMetadata( - engine_block_size=args.kv_block_size, - index_block_size=args.kv_block_size, - tokenizer=args.model_name, - tokenizer_digest=hashlib.sha256(args.model_name.encode()).hexdigest(), - tokenizer_canary=canary, - supports_events=False, # no ZMQ publisher; router falls back to load-only - ) - return EngineConfig( - model_name=args.model_name, - host=args.advertise_host or args.host, - port=args.port, - engine=EngineType(args.engine), - disagg_mode=DisaggMode(args.disagg_mode), - disagg_meta=_disagg_meta(args), - kv=kv, - kv_block_size=args.kv_block_size if args.kv else None, - dp_rank=args.dp_rank, - dp_size=args.dp_size, - request_transport=args.request_transport, - ) - - -def parse_args(argv=None): - p = argparse.ArgumentParser( - prog="infera-fake-worker", - description="Register into an Infera fleet and serve tokens without a GPU.", - ) - p.add_argument("--model-name", required=True, help="must match what the router routes for") - p.add_argument("--host", default="0.0.0.0") - p.add_argument( - "--advertise-host", - default=os.environ.get("POD_IP"), - help="address peers use to reach this worker; defaults to $POD_IP. " - "0.0.0.0 is never routable from another pod.", - ) - p.add_argument("--port", type=int, default=8080) - p.add_argument("--engine", default="sglang", choices=[e.value for e in EngineType]) - p.add_argument("--disagg-mode", default="mixed", choices=[m.value for m in DisaggMode]) - p.add_argument( - "--pd-protocol", - default="sglang-bootstrap", - choices=sorted(_PROTOCOLS), - help="must match the router's registry; a decode worker advertising a " - "different one is rejected as a protocol mismatch", - ) - p.add_argument( - "--bootstrap-port", - type=int, - default=8998, - help="advertised in disagg_meta by a prefill worker. Nothing listens on " - "it -- no KV is transferred (see README).", - ) - p.add_argument("--dp-rank", type=int, default=None) - p.add_argument("--dp-size", type=int, default=None) - - p.add_argument( - "--discovery-backend", - default=os.environ.get("INFERA_DISCOVERY_BACKEND", "kubernetes"), - choices=["kubernetes", "etcd"], - ) - p.add_argument("--etcd-endpoint", default=os.environ.get("INFERA_ETCD_ENDPOINT")) - p.add_argument("--etcd-prefix", default="/infera/workers/") - - p.add_argument( - "--kv", - action="store_true", - help="register a KV metadata block with a synthetic tokenizer canary. " - "All fakes for a model agree; a fake and a REAL worker will not -- the " - "second to register is silently dropped. See the README.", - ) - p.add_argument("--kv-block-size", type=int, default=64) - - p.add_argument("--ttft-ms", type=float, default=50.0) - p.add_argument("--itl-ms", type=float, default=10.0) - p.add_argument("--max-concurrency", type=int, default=8) - p.add_argument("--max-kv-blocks", type=int, default=1024) - p.add_argument( - "--startup-delay-s", - type=float, - default=0.0, - help="seconds of simulated weight loading before /health goes green. " - "Set this to your real cold start to reproduce autoscaler overshoot.", - ) - p.add_argument( - "--fail-first", - type=int, - default=0, - help="refuse the first N requests with 503, then recover -- exercises " - "the router's circuit breaker and its half-open probe.", - ) - p.add_argument( - "--request-transport", - default="http", - choices=["http", "nats"], - help="nats routes requests through a broker instead of the router " - "dialling this worker directly. Uses the real NatsRequestServer, which " - "proxies to this process's own HTTP surface exactly as it proxies to a " - "real engine -- so the transport under test is the production one.", - ) - p.add_argument("--nats-server", default=os.environ.get("NATS_SERVER")) - p.add_argument("--drain-timeout", type=float, default=30.0) - return p.parse_args(argv) - - -async def _serve(args) -> None: - import uvicorn - - cfg = build_config(args) - behaviour = Behaviour( - ttft_ms=args.ttft_ms, - itl_ms=args.itl_ms, - max_concurrency=args.max_concurrency, - max_kv_blocks=args.max_kv_blocks, - fail_first=args.fail_first, - ) - state = State() - state._sem = asyncio.Semaphore(args.max_concurrency) - app = build_app(cfg, behaviour, state) - - server = uvicorn.Server( - uvicorn.Config(app, host=args.host, port=args.port, log_level="warning") - ) - serve_task = asyncio.create_task(server.serve()) - - # Do not register until the socket is actually bound. uvicorn logs a bind - # failure and gives up, but the process keeps running -- so without this - # check a port collision produces a worker that is in the pool and serves - # nothing. That is a routing black hole, and it is exactly the failure this - # tool exists to help find rather than create. - for _ in range(100): - if server.started or serve_task.done(): - break - await asyncio.sleep(0.05) - if not server.started: - serve_task.cancel() - raise SystemExit(f"failed to bind {args.host}:{args.port} -- not registering") - - nats_req_server = None - if args.request_transport == "nats": - from infera.common.nats_request import NatsRequestServer - - nats_req_server = NatsRequestServer( - f"{cfg.host}:{cfg.port}", args.port, url=args.nats_server - ) - await nats_req_server.start() - logger.info("nats request consumer started for %s:%s", cfg.host, cfg.port) - - if args.startup_delay_s > 0: - logger.info("simulating weight load for %.0fs", args.startup_delay_s) - await asyncio.sleep(args.startup_delay_s) - state.ready = True - - # Register only once ready, exactly like a real worker: a worker that is in - # the pool but cannot serve is a routing black hole. - if args.discovery_backend == "etcd": - from infera.common.registration import RegistrationClient - - reg = RegistrationClient(args.etcd_endpoint, prefix=args.etcd_prefix) - else: - from infera.common.registration_k8s import K8sRegistrationClient - - reg = K8sRegistrationClient() - worker_id = await reg.register(cfg) - # register() only writes the record; keeping it alive is the caller's job in - # both backends, exactly as in the real worker entrypoint. Without this the - # etcd lease (30s) expires and the worker silently vanishes from the pool - # about half a minute after it appears -- which looks like a discovery bug - # and is not one. - hb_task = asyncio.create_task(reg.heartbeat_loop(), name="fake-worker-heartbeat") - logger.info( - "registered %s model=%s mode=%s via %s", - worker_id, - cfg.model_name, - cfg.disagg_mode.value, - args.discovery_backend, - ) - - stop = asyncio.Event() - - async def _shutdown() -> None: - # Mirrors the real worker entrypoints, including which step stops new - # work arriving -- that differs by backend, see below. - state.draining = True - # Stop the heartbeat first: it re-asserts registration from config, so a - # refresh landing after deregistration would put this worker back in the - # pool. - hb_task.cancel() - - deadline = time.monotonic() + args.drain_timeout - - async def _drain() -> None: - if nats_req_server is not None: - # The real drain: unsubscribe first so nothing new arrives, then - # wait on the in-flight set infera actually holds. No polling and - # no settle window -- unlike HTTP, where the count has to be - # inferred from the engine's lagging gauges. - await nats_req_server.stop( - drain=True, - drain_timeout=max(0.0, deadline - time.monotonic()), - ) - while state.running and time.monotonic() < deadline: - await asyncio.sleep(0.1) - if state.running: - logger.warning("drain timeout with %d request(s) still in flight", state.running) - - # Deregister before draining, mirroring the real entrypoints: removing - # the record is what stops new work arriving, on either backend. - try: - if not await reg.deregister(): - # deregister() already logged why, including whether it matters. - logger.warning("draining anyway") - except Exception as exc: # noqa: BLE001 - shutdown must not raise - logger.warning("deregister failed: %s", exc) - await _drain() - - server.should_exit = True - stop.set() - - shutdown_task: asyncio.Task | None = None - - def _on_signal() -> None: - # Guarded and strongly referenced. SIGTERM followed by SIGINT is routine - # for a terminating Pod, and two concurrent _shutdown() coroutines would - # deregister twice through an already-closed client and race each other - # inside the NATS server's stop(). asyncio holds only a weak reference - # to a task, so dropping the handle could have it collected mid-drain, - # leaving the process waiting on an event nobody will set. - nonlocal shutdown_task - if shutdown_task is None: - shutdown_task = asyncio.create_task(_shutdown(), name="fake-worker-shutdown") - - loop = asyncio.get_running_loop() - for sig in (signal.SIGTERM, signal.SIGINT): - loop.add_signal_handler(sig, _on_signal) - - await stop.wait() - await serve_task - - -#: Opt-in required to start. See _require_opt_in. -ALLOW_ENV = "INFERA_ALLOW_FAKE_WORKER" - - -def _require_opt_in() -> None: - """Refuse to run unless explicitly enabled. - - This ships in the same package as the server, and it registers through the - real registration clients -- so it can join a production fleet and advertise - any address it likes, which the router will then dial. Nothing about that is - a new privilege: it needs the worker credentials it would already have. But - a test tool in the production package should not be the thing that turns - "code execution in one pod" into "silently receiving prompts fleet-wide", - and a deliberate opt-in is cheap next to that. - """ - if os.environ.get(ALLOW_ENV, "").strip().lower() in ("1", "true", "yes"): - return - raise SystemExit( - f"infera-fake-worker refuses to start: it registers into real service " - f"discovery and serves fabricated responses, so it must be enabled " - f"deliberately. Set {ALLOW_ENV}=1 if this is a test environment." - ) - - -def main(argv=None) -> int: - logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") - _require_opt_in() - args = parse_args(argv) - asyncio.run(_serve(args)) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/manual/features/scaling.md b/manual/features/scaling.md index 80bfc060..4c4653cf 100644 --- a/manual/features/scaling.md +++ b/manual/features/scaling.md @@ -106,7 +106,7 @@ information lives. | **NATS** (`--request-transport nats`) | infera — it owns the request path and holds the in-flight set | exact, no polling | | **HTTP** (default in the recipes) | only the engine — the router dials it directly and never sees the request | poll the engine's `/metrics`, behind a settle window | -Measured, same fake worker, same generation: +Measured with a GPU-free stand-in worker, same generation: - **NATS, one in-flight generation**: the log reads `draining 1 in-flight NATS request(s)` — it knows the count — the 300-chunk generation completed in full, @@ -306,7 +306,7 @@ Two runs, both with a Pod deleted while holding in-flight work: concurrent 2500-token generations in flight, **4/4 completed with HTTP 200** and full-length output (6.5–13.3 kB), replacement Pod registered before the drain finished. -- **Fake workers**, same path without a GPU: a 300-chunk generation completed +- **GPU-free stand-in workers**, same path: a 300-chunk generation completed in full across the drain. ```{note} @@ -418,7 +418,7 @@ last prefill away then returns 503 naming the empty pool. ```{warning} **Not measured:** multi-node workers, TP > 1, PD scaling with a *real* engine -(the run above used fake workers, so no KV moved), and scale-down during an +(the run above used GPU-free stand-ins, so no KV moved), and scale-down during an active KV transfer. The PD handoff queues are counted in the drain, but that path has not been exercised on hardware. ``` diff --git a/pyproject.toml b/pyproject.toml index 75c1868b..6119a9c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,8 +41,6 @@ infera-kvd-probe = "infera.kvd.bench.probe:main" infera-kvd-l3-bench = "infera.kvd.bench.l3_bench:main" # node + PD preflight suite (gpu / network / storage / firmware / host probes) infera-preflight = "infera.tools.preflight.cli:main" -infera-fake-worker = "infera.tools.fakeworker.server:main" - [build-system] requires = ["setuptools>=69", "setuptools_scm[toml]>=8", "wheel"] build-backend = "setuptools.build_meta" diff --git a/tests/unit/common/test_shutdown_order.py b/tests/unit/common/test_shutdown_order.py index 77db8938..f244959b 100644 --- a/tests/unit/common/test_shutdown_order.py +++ b/tests/unit/common/test_shutdown_order.py @@ -24,7 +24,6 @@ ENTRYPOINTS = ( "infera/engine/vllm/__main__.py", "infera/engine/sglang/__main__.py", - "infera/tools/fakeworker/server.py", ) diff --git a/tests/unit/tools/test_fake_worker.py b/tests/unit/tools/test_fake_worker.py deleted file mode 100644 index 3fd927d5..00000000 --- a/tests/unit/tools/test_fake_worker.py +++ /dev/null @@ -1,261 +0,0 @@ -############################################################################### -# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. -# -# SPDX-License-Identifier: MIT -############################################################################### -"""The fake worker has to be trustworthy, or every test built on it is too. - -Two properties matter more than the rest. It must register through the *real* -contract, so a fleet of fakes exercises the same discovery path a real fleet -does -- that is checked by building an actual ``EngineConfig`` and running it -through ``build_worker_payload``, the same function every engine uses. And its -queue must be real, because ``num_requests_waiting`` is the metric the entire -industry autoscales on, and a fake that always reports zero would make every -scaling test vacuously pass. -""" - -from __future__ import annotations - -import asyncio - -import httpx -import pytest - -from infera.common.discovery import worker_info_from_json -from infera.common.registration import build_worker_payload -from infera.common.worker_pool import DisaggMode, EngineType -from infera.tools.fakeworker.server import ( - Behaviour, - State, - build_app, - build_config, - deterministic_canary, - parse_args, -) - - -def _args(*extra): - return parse_args(["--model-name", "m", *extra]) - - -def _stack(**behaviour_kw): - args = _args() - cfg = build_config(args) - state = State(ready=True) - state._sem = asyncio.Semaphore(behaviour_kw.get("max_concurrency", 8)) - b = Behaviour(**behaviour_kw) - return cfg, b, state, build_app(cfg, b, state) - - -def _client(app): - return httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://fake") - - -# --- the registration contract ------------------------------------------------ - - -def test_registers_through_the_real_payload_builder(): - """If this breaks, the fake has drifted from what a real worker registers -- - which is the one failure that would silently invalidate everything else.""" - cfg = build_config(_args("--engine", "vllm", "--disagg-mode", "prefill")) - payload = build_worker_payload(cfg) - assert payload["model_name"] == "m" - assert payload["engine"] == EngineType.VLLM - assert payload["disagg_mode"] == DisaggMode.PREFILL - # And discovery must be able to parse it back: worker_info_from_json is the - # single function every backend (etcd, kubernetes) uses on the wire record, - # so a round-trip through it is the real contract, not an approximation. - info = worker_info_from_json(payload) - assert info.worker_id == payload["worker_id"] - assert info.disagg_mode is DisaggMode.PREFILL - - -def test_kv_block_is_absent_unless_asked_for(): - """Without --kv there is no canary, so fakes can join any fleet. With it, - canary verification applies and mixing with real workers breaks.""" - assert build_config(_args()).kv is None - assert build_config(_args("--kv")).kv is not None - - -def test_fakes_for_one_model_agree_on_the_canary(): - """Disagreeing fakes would be silently dropped from the pool by - CanaryVerifier -- a fleet that looks half its intended size for no visible - reason.""" - assert deterministic_canary("llama") == deterministic_canary("llama") - assert deterministic_canary("llama") != deterministic_canary("qwen") - - -# --- the serving surface ------------------------------------------------------ - - -@pytest.mark.asyncio -async def test_unary_completion_shape(): - _, _, _, app = _stack(ttft_ms=0, itl_ms=0) - async with _client(app) as c: - r = await c.post("/v1/chat/completions", json={"model": "m", "max_tokens": 5}) - assert r.status_code == 200 - body = r.json() - assert body["choices"][0]["message"]["content"] - assert body["usage"]["completion_tokens"] == 5 - - -@pytest.mark.asyncio -async def test_streaming_emits_sse_and_terminates(): - """The router's failover and circuit breaker are both first-byte-sensitive, - so a fake that cannot stream cannot exercise either.""" - _, _, _, app = _stack(ttft_ms=0, itl_ms=0) - async with _client(app) as c: - r = await c.post( - "/v1/chat/completions", json={"model": "m", "max_tokens": 3, "stream": True} - ) - body = (await r.aread()).decode() - assert r.headers["content-type"].startswith("text/event-stream") - assert body.count("data: ") == 4 # 3 chunks + [DONE] - assert body.endswith("data: [DONE]\n\n") - - -@pytest.mark.asyncio -async def test_health_is_503_until_ready(): - """--startup-delay-s exists to reproduce autoscaler overshoot, which only - happens because an unready replica still counts in the fleet.""" - cfg, b, state, app = _stack() - state.ready = False - async with _client(app) as c: - assert (await c.get("/health")).status_code == 503 - state.ready = True - assert (await c.get("/health")).status_code == 200 - - -# --- the metrics an autoscaler would read ------------------------------------- - - -@pytest.mark.asyncio -async def test_queue_depth_is_real(): - """The whole point. Drive more concurrent requests than the worker admits - and the waiting count must actually rise -- a fake that always reports 0 - would make every scaling test pass without testing anything.""" - _, _, state, app = _stack(max_concurrency=2, ttft_ms=200, itl_ms=0) - async with _client(app) as c: - tasks = [ - asyncio.create_task(c.post("/v1/completions", json={"model": "m", "max_tokens": 1})) - for _ in range(6) - ] - await asyncio.sleep(0.05) - peak_waiting = state.waiting - peak_running = state.running - await asyncio.gather(*tasks) - - assert peak_running == 2, f"admitted {peak_running}, expected the concurrency cap" - assert peak_waiting == 4, f"queued {peak_waiting}, expected the other 4" - assert state.waiting == 0 and state.running == 0, "must settle back to idle" - - -@pytest.mark.asyncio -async def test_metrics_use_engine_native_names(): - """A scaling rule written against fakes should transfer to a real fleet - unchanged, so the names have to be the engine's, not ours.""" - for engine, expected in ( - ("vllm", "vllm:num_requests_waiting"), - ("sglang", "sglang:num_queue_reqs"), - ): - args = _args("--engine", engine) - cfg = build_config(args) - state = State(ready=True) - state._sem = asyncio.Semaphore(1) - app = build_app(cfg, Behaviour(), state) - async with _client(app) as c: - text = (await c.get("/metrics")).text - assert expected in text, f"{engine}: missing {expected}\n{text}" - - -@pytest.mark.asyncio -async def test_kv_usage_tracks_inflight(): - _, _, state, app = _stack(max_concurrency=4, ttft_ms=200, itl_ms=0) - async with _client(app) as c: - idle = (await c.get("/metrics")).text - tasks = [ - asyncio.create_task(c.post("/v1/completions", json={"model": "m", "max_tokens": 1})) - for _ in range(4) - ] - await asyncio.sleep(0.05) - busy = (await c.get("/metrics")).text - await asyncio.gather(*tasks) - - def usage(t): - return float( - [ln for ln in t.splitlines() if "cache_usage" in ln or "token_usage" in ln][0].split()[ - -1 - ] - ) - - assert usage(idle) == 0.0 - assert usage(busy) > 0.0 - - -# --- failure injection, for the circuit breaker ------------------------------- - - -@pytest.mark.asyncio -async def test_fail_first_then_recovers(): - """Mirrors the breaker's half-open probe: a worker that is broken, stays - broken for a while, then comes back.""" - _, _, state, app = _stack(ttft_ms=0, itl_ms=0, fail_first=3) - async with _client(app) as c: - codes = [ - (await c.post("/v1/completions", json={"model": "m", "max_tokens": 1})).status_code - for _ in range(5) - ] - assert codes[:3] == [503, 503, 503] - assert codes[3:] == [200, 200] - - -@pytest.mark.asyncio -async def test_draining_refuses_new_work(): - """SIGTERM deregisters before draining; until the router notices, arriving - requests must be refused rather than accepted and then cut.""" - _, _, state, app = _stack(ttft_ms=0, itl_ms=0) - state.draining = True - async with _client(app) as c: - r = await c.post("/v1/completions", json={"model": "m", "max_tokens": 1}) - assert r.status_code == 503 - assert "infera_fake_worker_draining 1" in (await c.get("/metrics")).text - - -def test_a_rank_header_cannot_forge_metrics_or_grow_without_bound(): - """The DP-rank header is attacker-controlled and ends up as a Prometheus - label and a map key. Unvalidated, a quote or newline breaks out of the label - and forges series in whatever scrapes the endpoint, and every distinct value - adds a permanent entry to a map that /debug/routing returns in full.""" - from infera.tools.fakeworker.server import _rank_label - - assert _rank_label(None, 4) == "-" - assert _rank_label("2", 4) == "2" - - for hostile in ('x"} 1\nup{job="prod"} 0', "a\\b", "1\n2", "", " "): - got = _rank_label(hostile, 4) - assert got == "invalid", f"{hostile!r} -> {got!r}" - - # Out of range is meaningless here, and unbounded if accepted. - assert _rank_label("4", 4) == "invalid" - assert _rank_label("-1", 4) == "invalid" - assert _rank_label("999999", 4) == "invalid" - - -def test_it_refuses_to_start_without_an_explicit_opt_in(monkeypatch): - """It registers into real service discovery and answers with fabricated - text, and it ships in the same package as the server. Starting it should be - a decision, not a default.""" - import pytest - - from infera.tools.fakeworker.server import ALLOW_ENV, main - - monkeypatch.delenv(ALLOW_ENV, raising=False) - with pytest.raises(SystemExit) as exc: - main([]) - assert ALLOW_ENV in str(exc.value), "the error must say how to enable it" - - # Enabled, it gets as far as parsing arguments. - monkeypatch.setenv(ALLOW_ENV, "1") - with pytest.raises(SystemExit) as exc: - main(["--nonsense-flag"]) - assert ALLOW_ENV not in str(exc.value), "past the gate, argparse should be what refuses"