diff --git a/infera/common/nats_request.py b/infera/common/nats_request.py index 0894d951..9f5d71cf 100644 --- a/infera/common/nats_request.py +++ b/infera/common/nats_request.py @@ -15,13 +15,20 @@ Wire protocol (one request -> N reply messages on a fresh inbox): request (server -> ``infera.req.``, reply=): - JSON {"path": str, "stream": bool, "headers": {..}|null, "body": {..}} + JSON {"path": str, "stream": bool, "headers": {..}|null, "body": {..}, + "migratable": bool} reply (worker -> ), framed by the ``rs-type`` header: data : payload = raw response bytes (an SSE chunk, or the full JSON body) done : payload = b"" , header ``rs-status`` = HTTP status code error: payload = utf-8 error text (transport/proxy failure) +``migratable`` says the router can continue this generation on another worker +(see infera.router.migration), which is what lets a draining worker hand it back +early instead of holding the shutdown open. It is a promise about the *router*, +so the worker must not assume it: without it, a request is drained the slow way +and severing it early would just be an error the client did not have to see. + The worker side proxies to its own local engine HTTP (127.0.0.1:), so the engine itself is unchanged; the consumer is a thin task started alongside it (like the KV relay). @@ -85,6 +92,12 @@ TYPE_DONE = "done" TYPE_ERROR = "error" +# The error payload a worker sends when it hands a generation back at the start +# of a drain rather than holding the shutdown open for it. The router reads it +# to tell "this worker is leaving on purpose" apart from a worker that broke, +# which are the same event to the client but not to an operator reading metrics. +DRAINING_NOTICE = b"infera: worker draining" + # Throttle knob (single variable, default OFF). When > 0, the per-instance # request path is JetStream-backed and the router refuses to dispatch to a # worker whose backlog (num_pending + num_ack_pending on its request consumer) @@ -454,6 +467,11 @@ def __init__( # In-flight proxy tasks keyed by reply inbox, so a cancel signal can # abort the matching request's engine call. self._inflight: dict[str, asyncio.Task] = {} + # Inboxes whose router can continue the generation on another worker. + self._migratable: set[str] = set() + # Inboxes already told the worker is draining, so the cancellation that + # follows does not report itself a second time. + self._handed_back: set[str] = set() async def start(self) -> None: self._nc = await _connect(self._url, "infera-worker-req") @@ -541,8 +559,17 @@ async def stop(self, *, drain: bool = False, drain_timeout: float = 0.0) -> None _done, pending = await asyncio.wait(inflight, timeout=remaining) if pending: logger.warning( - "drain timeout; cancelling %d unfinished request(s)", len(pending) + "drain timeout; %d request(s) did not finish in time", len(pending) ) + # 2b. Whatever is still running was about to be cancelled. Hand the + # resumable part back to the router instead: migration is the + # alternative to cutting these, not a shortcut past the wait above. + # Moving a generation costs the next worker a re-read of everything + # produced so far, so it is worth doing only once the request has been + # given the window it was promised -- or, with no window configured, + # once it is clear there will not be one. + if drain: + await self._hand_back_migratable() # 3. Cancel whatever is left (all of it on the non-drain path). for task in list(self._inflight.values()): if not task.done(): @@ -611,6 +638,39 @@ async def _await_queued(self, deadline: float) -> None: 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 _hand_back_migratable(self) -> None: + """Give the router back what it can finish elsewhere, instead of cutting it. + + Called once the drain window has run out, on the requests that outlived + it. Everything here was going to be cancelled a moment later, so the + choice this makes is not whether these generations survive the drain -- + it is whether the client learns that they did not. + + The notice goes out *before* the task is cancelled: cancellation also + sends an error, but a generic one, and the router would then treat a + planned handover as a worker that broke. + """ + handed = [i for i in self._migratable if not self._done(i)] + if not handed: + return + logger.info("drain: handing %d resumable request(s) back to the router", len(handed)) + for inbox in handed: + self._handed_back.add(inbox) + try: + await self._reply(inbox, TYPE_ERROR, DRAINING_NOTICE) + except Exception as exc: # noqa: BLE001 - shutdown must continue + logger.debug("could not hand back %s: %s", inbox[-12:], exc) + task = self._inflight.get(inbox) + if task is not None and not task.done(): + # Stops the engine generating for a stream nobody reads now that + # the router has been told to take it elsewhere. + task.cancel() + self._migratable.clear() + + def _done(self, inbox: str) -> bool: + task = self._inflight.get(inbox) + return task is None or task.done() + async def _reply( self, inbox: str, rtype: str, data: bytes = b"", status: int | None = None ) -> None: @@ -631,6 +691,13 @@ async def _on_request(self, msg) -> None: # (which tears down the engine connection -> engine stops generating). task = asyncio.create_task(self._proxy(inbox, msg), name=f"nats-req-{inbox[-12:]}") self._inflight[inbox] = task + # Remembered for drain: only a request the router said it can continue + # elsewhere may be handed back early. + try: + if json.loads(msg.data).get("migratable"): + self._migratable.add(inbox) + except Exception: # noqa: BLE001 - a body _proxy will reject anyway + pass async def _on_cancel(self, msg) -> None: # nats-py requires subscription callbacks to be coroutines. @@ -661,12 +728,15 @@ async def _proxy(self, inbox: str, msg) -> None: except asyncio.CancelledError: # Router gave up (timeout / client disconnect). The engine connection # is torn down by exiting the stream context; best-effort error reply - # (the router may have already unsubscribed). - logger.info("NATS request aborted (cancelled): %s", inbox[-12:]) - try: - await self._reply(inbox, TYPE_ERROR, b"request cancelled") - except Exception: - pass + # (the router may have already unsubscribed). A request handed back + # for drain was cancelled by us and has already been told why, so it + # must not get a second, contradictory error frame. + if inbox not in self._handed_back: + logger.info("NATS request aborted (cancelled): %s", inbox[-12:]) + try: + await self._reply(inbox, TYPE_ERROR, b"request cancelled") + except Exception: + pass except Exception as exc: logger.warning("NATS request proxy failed: %s", exc) try: @@ -675,6 +745,8 @@ async def _proxy(self, inbox: str, msg) -> None: pass finally: self._inflight.pop(inbox, None) + self._migratable.discard(inbox) + self._handed_back.discard(inbox) # Ack only after the request is fully proxied so the backlog gauge # (num_ack_pending) reflects genuinely in-flight work. await self._ack(msg) diff --git a/infera/engine/sglang/args.py b/infera/engine/sglang/args.py index ebd55e60..433ec063 100644 --- a/infera/engine/sglang/args.py +++ b/infera/engine/sglang/args.py @@ -135,8 +135,10 @@ def parse_sglang_args(argv: list[str] | None = None) -> SglangWorkerArgs: default=float(__import__("os").environ.get("INFERA_DRAIN_TIMEOUT", "30") or 30), help="Graceful shutdown: on SIGTERM the worker stops accepting new NATS " "requests and lets in-flight generations finish for up to this many " - "seconds before cancelling leftovers (rolling-upgrade drain). Default 30; " - "0 = cancel in-flight immediately. Overrides $INFERA_DRAIN_TIMEOUT.", + "seconds (rolling-upgrade drain). How long one generation is worth " + "waiting for: whatever is still running at the deadline is cancelled, " + "or handed back to the router when it enabled --migration-limit. " + "Default 30; 0 = do not wait. Overrides $INFERA_DRAIN_TIMEOUT.", ) parser.add_argument( "--advertise-host", diff --git a/infera/engine/vllm/args.py b/infera/engine/vllm/args.py index 158f84b1..b38eb279 100644 --- a/infera/engine/vllm/args.py +++ b/infera/engine/vllm/args.py @@ -275,8 +275,10 @@ def parse_vllm_args(argv: list[str] | None = None) -> VllmWorkerArgs: default=float(__import__("os").environ.get("INFERA_DRAIN_TIMEOUT", "30") or 30), help="Graceful shutdown: on SIGTERM the worker stops accepting new NATS " "requests and lets in-flight generations finish for up to this many " - "seconds before cancelling leftovers (rolling-upgrade drain). Default 30; " - "0 = cancel in-flight immediately. Overrides $INFERA_DRAIN_TIMEOUT.", + "seconds (rolling-upgrade drain). How long one generation is worth " + "waiting for: whatever is still running at the deadline is cancelled, " + "or handed back to the router when it enabled --migration-limit. " + "Default 30; 0 = do not wait. Overrides $INFERA_DRAIN_TIMEOUT.", ) parser.add_argument( "--advertise-host", diff --git a/infera/router/auto.py b/infera/router/auto.py index 0db4a494..456ed51f 100644 --- a/infera/router/auto.py +++ b/infera/router/auto.py @@ -39,13 +39,16 @@ class AutoRouter(BaseRouter): a dumb dispatcher. """ - def __init__(self, *args, **kwargs) -> None: + def __init__(self, *args, migration_limit: int = 0, **kwargs) -> None: super().__init__(*args, **kwargs) self._mixed = MixedRouter( self.pool, self.policy, nats_client=self.nats_client, request_max_retries=self.request_max_retries, + # Only mixed workers can carry a generation elsewhere; the PD path + # has a second leg whose state would also have to move. + migration_limit=migration_limit, # 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. diff --git a/infera/router/migration.py b/infera/router/migration.py new file mode 100644 index 00000000..b1eacc13 --- /dev/null +++ b/infera/router/migration.py @@ -0,0 +1,375 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Carrying a half-finished generation from one worker to another. + +A worker that goes away mid-generation -- drained, evicted, crashed -- takes an +unfinished response with it. Retrying the request from the top is not equivalent: +the client has already been sent tokens, so a fresh generation would either +repeat them or contradict what it already read. + +What moves instead is the generation so far, appended to the prompt with the +token budget reduced by what it cost, so the next worker continues rather than +restarts. The client sees one uninterrupted stream and never learns a worker +changed underneath it. + +**Exactly, when the engine allows it.** Given the token ids the engine actually +sampled (see infera.router.token_ids), the continuation carries those ids and +the prompt's, and the next worker resumes from the identical sequence. Without +them the decoded text is carried instead and re-encoded downstream, which can +shift a token boundary: the same words, not provably the same tokens. Exactness +is preferred wherever it is available and the fallback is silent, because an +approximate continuation is still far better than a severed stream. + +**What this is not.** The KV cache does not move, so the new worker re-reads the +carried prefix -- kv-aware routing usually lands on a worker that already holds +some of it, which is what keeps that affordable. And no continuation is +byte-identical to what the original worker would have produced: sampling state +does not survive the move, and exact ids do not change that. A caller who needs +reproducible output for a fixed seed should not enable migration. + +**Requires the NATS transport.** Over HTTP the router hands the connection +straight to the engine and never sees a frame boundary, so there is nothing to +accumulate from. +""" + +from __future__ import annotations + +import json +import logging +from dataclasses import dataclass + +from infera.router.token_ids import deltas_from_chunk, prompt_from_chunk, strip_token_ids + +logger = logging.getLogger(__name__) + +# The sentinel that ends an OpenAI stream. Never carried: it belongs to the +# stream the client is reading, not to the generation being moved. +_DONE_SENTINEL = b"data: [DONE]" +_DATA_PREFIX = b"data: " + +# Only completions accepts a pre-tokenized prompt; chat has no such entry, so an +# exact continuation of a chat request has to be issued against this path. +COMPLETIONS_PATH = "/v1/completions" + +# Request fields whose meaning would not survive being reissued as completions. +# Tools are the obvious one -- a completions request cannot emit a tool call, so +# converting one that might need to would quietly remove the ability halfway +# through the answer. +_CHAT_ONLY_FIELDS = ("tools", "functions", "tool_choice", "function_call", "response_format") + + +def as_chat_chunk(raw: bytes) -> bytes: + """Re-shape a completions chunk into the chat form the client is reading. + + An exact chat continuation has to be issued against the completions + endpoint, which answers in a different shape: ``choices[].text`` where the + client expects ``choices[].delta.content``. Without this the migration would + be plainly visible as the response changing format mid-stream. + + Anything unrecognised is passed through untouched. A chunk this cannot + convert is a chunk it does not understand, and mangling it would be worse + than letting it through. + """ + if not raw: + return raw + out: list[bytes] = [] + rewritten = False + for line in raw.split(b"\n"): + stripped = line.strip() + if not stripped.startswith(_DATA_PREFIX) or stripped.startswith(_DONE_SENTINEL): + out.append(line) + continue + try: + obj = json.loads(stripped[len(_DATA_PREFIX) :]) + except ValueError: + out.append(line) + continue + if not isinstance(obj, dict) or not isinstance(obj.get("choices"), list): + out.append(line) + continue + obj["object"] = "chat.completion.chunk" + for choice in obj["choices"]: + if isinstance(choice, dict) and "text" in choice: + choice["delta"] = {"content": choice.pop("text") or ""} + # `logprobs` here describes completions tokens and has a + # different shape in chat; dropping beats mistranslating. + choice.pop("logprobs", None) + out.append(_DATA_PREFIX + json.dumps(obj, separators=(",", ":")).encode()) + rewritten = True + return b"\n".join(out) if rewritten else raw + + +def _asks_for_several_completions(body: dict) -> bool: + """Whether this request produces more than one generation. + + `n` and `best_of` are the direct forms. A list-valued `prompt` is the + indirect one: completions accepts a batch, and each entry is its own + generation. A batch of exactly one is still a batch -- the reply is indexed + per prompt -- so length is not what decides it. + + The exception is a pre-tokenized prompt, `[int, ...]`, which is a single + generation expressed as token ids rather than a batch of anything. + """ + for key in ("n", "best_of"): + value = body.get(key) + if isinstance(value, int) and value > 1: + return True + prompt = body.get("prompt") + if isinstance(prompt, list) and not _is_token_array(prompt): + return True + return False + + +def _is_token_array(prompt: list) -> bool: + return bool(prompt) and all(isinstance(x, int) and not isinstance(x, bool) for x in prompt) + + +@dataclass(frozen=True) +class Continuation: + """The request that makes another worker finish this generation.""" + + body: dict + path: str + exact: bool + + +class MigrationState: + """What has been produced so far, and what it costs to carry it. + + Fed every chunk on its way to the client, so what is accumulated is what the + client actually received -- not what a worker claims to have sent. That + distinction is the point: after a migration the two must agree, or the + client sees a seam. + """ + + def __init__(self, body: dict, *, limit: int, path: str = COMPLETIONS_PATH) -> None: + self._original_body = body + self._path = path + self._text: list[str] = [] + self._chunks_with_text = 0 + self._output_ids: list[int] = [] + self._prompt_ids: list[int] | None = None + # Cleared the moment a chunk carries text the engine did not account for + # in ids. Continuing from ids that cover only part of what the client + # read would drop the rest, so the whole id path is abandoned instead. + self._ids_cover_output = True + self.migrations_left = limit + # Set once anything unparseable arrives. A generation that cannot be + # reconstructed exactly must not be migrated at all: continuing from a + # partial prefix would silently drop output the client already read. + self.poisoned = False + # A request asking for several completions has no single prefix to + # carry: what is accumulated is one choice's output, and every other one + # would resume from it. Rejected up front rather than on the first + # chunk, since the shape of the request already says so. + if _asks_for_several_completions(body): + self._poison("more than one completion was asked for") + + @property + def produced_text(self) -> str: + return "".join(self._text) + + @property + def produced_tokens(self) -> int: + """Tokens generated so far. + + Exact when the engine reported ids. Otherwise the number of chunks that + carried text, which every engine here emits one token at a time. Only + used to reduce the remaining budget, where being off by a little changes + the length of the answer and nothing else. + """ + if self._has_exact_output(): + return len(self._output_ids) + return self._chunks_with_text + + def is_exact(self) -> bool: + """Whether the continuation can be issued as token ids. + + Needs the prompt's ids as well as the output's: resuming from exact + output ids appended to a re-encoded prompt would just move the ambiguity + from one end of the sequence to the other. + """ + return bool(self._prompt_ids) and self._has_exact_output() + + def _has_exact_output(self) -> bool: + return self._ids_cover_output and bool(self._output_ids) + + def observe(self, chunk: bytes) -> bytes: + """Record one chunk, and return what should go to the client. + + The ids are asked for by the router, not the caller, so they are taken + out on the way past. That happens whatever the state of this object: + giving up on migrating a request is the router's own business, and must + not change the shape of what the caller receives. + """ + if not chunk: + return chunk + out_lines: list[bytes] = [] + rewritten = False + for line in chunk.split(b"\n"): + stripped = line.strip() + if not stripped.startswith(_DATA_PREFIX) or stripped.startswith(_DONE_SENTINEL): + out_lines.append(line) + continue + try: + obj = json.loads(stripped[len(_DATA_PREFIX) :]) + except ValueError: + # Not JSON: this is not a stream we know how to reconstruct, and + # not one to rewrite either -- it passes through as it arrived. + self._poison("chunk is not JSON") + out_lines.append(line) + continue + if not self.poisoned: + self._record(obj) + if strip_token_ids(obj): + # Rebuild only the lines that carried ids; everything else keeps + # the engine's own bytes. + out_lines.append(_DATA_PREFIX + json.dumps(obj, separators=(",", ":")).encode()) + rewritten = True + else: + out_lines.append(line) + return b"\n".join(out_lines) if rewritten else chunk + + def _record(self, obj: object) -> None: + if isinstance(obj, dict) and len(obj.get("choices") or []) > 1: + # Not asked for, but arrived anyway. Accumulating the first choice + # would give every other one a prefix belonging to it. + self._poison("the engine returned more than one choice") + return + if self._prompt_ids is None: + self._prompt_ids = prompt_from_chunk(obj) + ids = deltas_from_chunk(obj) + if not ids and self._carries_more_than_text(obj): + # Output that does not survive being carried as text. Resuming would + # replay it: the client holds half a tool call and would be sent a + # whole one after it. Exact ids cover this -- they are the tokens + # behind whatever the parser produced -- so this only stops the + # text path. + self._poison("output is not plain text (tool call or reasoning)") + return + if ids: + self._output_ids.extend(ids) + delta = self._delta_of(obj) + if delta is None: + return + if not ids and not isinstance(self._original_body.get("prompt", ""), str): + # A pre-tokenized prompt can be extended with ids and nothing else. + # Text output the engine did not account for therefore leaves this + # request with no way to be continued at all. + self._poison("a pre-tokenized prompt cannot be extended with text") + return + self._text.append(delta) + self._chunks_with_text += 1 + if not ids: + # Text the ids do not account for. Whatever the engine is doing -- + # not reporting ids, or reporting them only sometimes -- the id + # sequence is no longer a faithful record of what was produced. + self._ids_cover_output = False + + @staticmethod + def _carries_more_than_text(obj: object) -> bool: + """Whether a chunk holds output that concatenated text cannot represent. + + Tool calls arrive under their own key rather than as content, and a + reasoning parser moves the model's thinking out of it. Either way the + text the client received is not the whole of what was produced, so a + continuation built from that text alone is missing part of the answer. + """ + if not isinstance(obj, dict): + return False + choices = obj.get("choices") + if not isinstance(choices, list) or not choices: + return False + first = choices[0] + if not isinstance(first, dict): + return False + delta = first.get("delta") + if not isinstance(delta, dict): + return False + return bool(delta.get("tool_calls") or delta.get("reasoning_content")) + + @staticmethod + def _delta_of(obj: object) -> str | None: + """The text a chunk adds, for chat and completions alike.""" + if not isinstance(obj, dict): + return None + choices = obj.get("choices") + if not isinstance(choices, list) or not choices: + return None + first = choices[0] + if not isinstance(first, dict): + return None + # Chat completions put it under `delta`, completions under `text`. + delta = first.get("delta") + if isinstance(delta, dict): + content = delta.get("content") + return content if isinstance(content, str) and content else None + text = first.get("text") + return text if isinstance(text, str) and text else None + + def _poison(self, why: str) -> None: + self.poisoned = True + logger.warning("migration disabled for this request: %s", why) + + def can_migrate(self) -> bool: + return not self.poisoned and self.migrations_left > 0 + + def next_continuation(self) -> Continuation: + """The request that makes another worker continue this generation.""" + self.migrations_left -= 1 + if self.is_exact() and self._convertible(): + return self._exact_continuation() + return self._text_continuation() + + def _convertible(self) -> bool: + """Whether this request can be reissued against the completions path. + + A completions request cannot emit a tool call or honour a chat response + format, so a request that might need either keeps the text path: a + continuation that silently loses a capability is worse than one whose + token boundaries might differ. + """ + if self._path.endswith("/completions") and not self._path.endswith("/chat/completions"): + return True + return not any(self._original_body.get(f) for f in _CHAT_ONLY_FIELDS) + + def _exact_continuation(self) -> Continuation: + body = self._base_body() + body.pop("messages", None) + body["prompt"] = list(self._prompt_ids or []) + self._output_ids + return Continuation(body=body, path=COMPLETIONS_PATH, exact=True) + + def _text_continuation(self) -> Continuation: + body = self._base_body() + carried = self.produced_text + if body.get("messages") is not None: + # Chat: an assistant turn holding what has been said so far. Engines + # continue such a turn rather than starting a new one, which is + # exactly the semantics needed here. + messages = list(body["messages"]) + messages.append({"role": "assistant", "content": carried}) + body["messages"] = messages + else: + original = body.get("prompt", "") + if not isinstance(original, str): + # Guarded against at construction; asserted here because + # formatting silently produces a repr rather than failing. + raise AssertionError("text continuation of a non-string prompt") + body["prompt"] = f"{original}{carried}" + return Continuation(body=body, path=self._path, exact=False) + + def _base_body(self) -> dict: + body = dict(self._original_body) + # Asked for by the router for its own use; the next worker is told + # separately whether they are still wanted. + body.pop("return_token_ids", None) + for key in ("max_tokens", "max_completion_tokens"): + budget = body.get(key) + if isinstance(budget, int): + # Never below 1: a request for zero tokens is rejected outright, + # which would turn a migration into an error the client sees. + body[key] = max(1, budget - self.produced_tokens) + return body diff --git a/infera/router/mixed.py b/infera/router/mixed.py index 9f2fe4fa..65617de0 100644 --- a/infera/router/mixed.py +++ b/infera/router/mixed.py @@ -13,13 +13,20 @@ from fastapi import Response from fastapi.responses import JSONResponse, StreamingResponse -from infera.common.nats_request import TYPE_DATA, TYPE_DONE, TYPE_ERROR +from infera.common.nats_request import ( + DRAINING_NOTICE, + 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 +from infera.router.migration import MigrationState, as_chat_chunk +from infera.router.token_ids import supports_streaming_ids from infera.server import metrics logger = logging.getLogger(__name__) @@ -45,8 +52,11 @@ class MixedRouter(BaseRouter): (the client already holds partial output). """ - def __init__(self, *args, **kwargs) -> None: + def __init__(self, *args, migration_limit: int = 0, **kwargs) -> None: super().__init__(*args, **kwargs) + # How many times one generation may be carried to another worker. Zero + # disables it; see infera.router.migration for what carrying costs. + self._migration_limit = max(0, int(migration_limit or 0)) # Bound connect time so unreachable workers fail fast; leave read open # for arbitrarily long generations. Bump connection limits well above # httpx defaults (100) so we can sustain high-concurrency benchmarks @@ -140,6 +150,17 @@ async def _attempt(self, target, blocks, body, hints, path, stream, obs) -> Resp forwarded_body.pop("_infera_cache_hints", None) forwarded_body.pop("_infera_request_id", None) + # Ask for the sampled token ids only when a migration could actually use + # them: they cost the engine work and the client never sees them, since + # they are taken back out before the stream is forwarded. + if ( + stream + and self._migration_limit > 0 + and supports_streaming_ids(worker.engine, path) + and "return_token_ids" not in forwarded_body + ): + forwarded_body["return_token_ids"] = True + use_nats = self.nats_client is not None and worker.request_transport == "nats" self.policy.on_request_started(target.route_key, blocks) @@ -181,29 +202,83 @@ async def _attempt_stream( if kind == TYPE_DATA: obs["outcome"] = "ok" # committed once first byte is in hand + state = self._migration_state(forwarded_body, use_nats, path) + + def passthrough(raw): + # observe() also takes the router-only token ids back out, so + # what is yielded is always what the caller asked for. + return state.observe(raw) if state is not None else raw + + def to_client_shape(raw): + return as_chat_chunk(passthrough(raw)) + async def generate(): + current, cur_target, cur_blocks = agen, target, blocks + transform = passthrough try: if data0: - yield data0 - async for k, _st, d in agen: - if k == TYPE_DATA: - if d: - yield d - elif k == TYPE_ERROR: - logger.warning( - "stream from worker %s failed mid-stream: %s", - worker.worker_id, - d[:200], + yield transform(data0) + while True: + async for k, _st, d in current: + if k == TYPE_DATA: + if d: + yield transform(d) + elif k == TYPE_ERROR: + # A drain notice is a planned handover, not a + # fault: the worker is leaving and expects us to + # take the generation elsewhere. + draining = d.startswith(DRAINING_NOTICE) + if draining: + reason = "worker_draining" + logger.info( + "worker %s is draining; taking its stream elsewhere", + cur_target.worker.worker_id, + ) + else: + reason = "stream_broken" + logger.warning( + "stream from worker %s failed mid-stream: %s", + cur_target.worker.worker_id, + d[:200], + ) + break + else: # done + return + else: + # The generator ended without saying why, which is + # the shape a crashed worker leaves behind. + reason = "stream_broken" + + # The stream broke. Whatever the client already read has + # to be honoured, so the only options are to carry the + # generation to another worker or to end it visibly. + resumed = None + if state is not None and state.can_migrate(): + resumed = await self._resume_elsewhere( + state, cur_target, path, obs, reason ) + elif state is not None: + # Migration was possible for this request and is not + # any more, which is worth counting separately from + # a deployment that never enabled it. + metrics.migrations_failed_total.labels( + reason="poisoned" if state.poisoned else "limit" + ).inc() + if resumed is None: + what = "is shutting down" if reason == "worker_draining" else "failed" yield ( - f'data: {{"error":"worker {worker.worker_id} ' - f'stream failed mid-stream"}}\n\n' + f'data: {{"error":"worker {cur_target.worker.worker_id} ' + f'{what} mid-stream"}}\n\n' ).encode() return - else: # done - return + current, next_target, next_blocks, reshape = resumed + transform = to_client_shape if reshape else passthrough + # The previous attempt's load accounting ends here; the + # new one is already counted by _resume_elsewhere. + self.policy.on_request_finished(cur_target.route_key, cur_blocks) + cur_target, cur_blocks = next_target, next_blocks finally: - self.policy.on_request_finished(target.route_key, blocks) + self.policy.on_request_finished(cur_target.route_key, cur_blocks) return StreamingResponse(generate(), media_type="text/event-stream") @@ -223,6 +298,99 @@ async def generate(): ) ) + def _migration_state(self, forwarded_body, use_nats, path): + """State for carrying this generation elsewhere, or None if it cannot be. + + Only over NATS: on the HTTP path the router hands the connection to the + engine and never sees a frame, so there is nothing to accumulate and + nothing to resume from. + """ + if not use_nats or self._migration_limit <= 0: + return None + return MigrationState(forwarded_body, limit=self._migration_limit, path=path) + + async def _resume_elsewhere(self, state, failed_target, path, obs, reason): + """Continue this generation on a different worker. + + Returns the new event stream and its target, or None when nobody else + can take it -- in which case the caller ends the stream visibly rather + than leaving the client waiting on a generation that stopped. + """ + cont = state.next_continuation() + body = cont.body + # An exact chat continuation is issued against the completions endpoint, + # so what is sent and what the client is reading can differ. + send_path, client_path = cont.path, path + model = body.get("model") + candidates = [ + w + for w in self.pool.list_active(model=model, mode=DisaggMode.MIXED) + if w.worker_id != failed_target.worker.worker_id + ] + # The failed worker is excluded even if it is the only one: it just + # dropped this stream, and handing the request back to it is how a + # migration loop starts. + candidates = self.breaker.filter(candidates) + if not candidates: + logger.warning("cannot migrate: no other worker serves model=%r", model) + metrics.migrations_failed_total.labels(reason="no_candidate").inc() + return None + + target, blocks = self.policy.pick(candidates, body) + worker = target.worker + if worker.request_transport != "nats" or self.nats_client is None: + # The accumulated state is only resumable over NATS. + logger.warning("cannot migrate: worker %s is not on nats", worker.worker_id) + metrics.migrations_failed_total.labels(reason="not_nats").inc() + return None + + self.policy.on_request_started(target.route_key, blocks) + agen = self._normalized_stream( + worker, + f"{worker.url}{send_path}", + body, + dp_rank_header(target), + send_path, + use_nats=True, + ) + # Peek: a worker that cannot take it must not consume the migration + # budget silently, and the caller needs a stream that is already + # producing before it commits to it. + try: + kind, _st, first = await agen.__anext__() + except StopAsyncIteration: + kind, first = TYPE_DONE, b"" + if kind != TYPE_DATA: + await agen.aclose() + self.policy.on_request_finished(target.route_key, blocks) + logger.warning("migration to %s failed before first byte", worker.worker_id) + metrics.migrations_failed_total.labels(reason="no_first_byte").inc() + return None + + logger.info( + "migrated a live generation from %s to %s after %d token(s) (%s)", + failed_target.worker.worker_id, + worker.worker_id, + state.produced_tokens, + "exact token ids" if cont.exact else "carried text", + ) + metrics.migrations_total.labels(reason=reason).inc() + obs["outcome"] = "ok" + + async def resumed(): + # The peeked frame is replayed unobserved: the caller runs every + # frame through the same transform, and recording it here as well + # would count its tokens twice. + if first: + yield TYPE_DATA, None, first + async for item in agen: + yield item + + # Replies arrive in the shape of whatever endpoint was used, which is + # not necessarily the one the client is reading. + reshape = send_path != client_path + return resumed(), target, blocks, reshape + async def _attempt_unary( self, worker, url, forwarded_body, dp_headers, path, use_nats, obs ) -> Response: @@ -320,7 +488,15 @@ async def _normalized_stream( """Unify HTTP and NATS streaming into ``(kind, status, data)`` events where ``kind`` is one of TYPE_DATA / TYPE_DONE / TYPE_ERROR.""" if use_nats: - payload = {"path": path, "stream": True, "headers": dp_headers, "body": forwarded_body} + payload = { + "path": path, + "stream": True, + "headers": dp_headers, + "body": forwarded_body, + # Lets a draining worker return this stream immediately instead + # of holding its shutdown open; only true when we can resume it. + "migratable": self._migration_limit > 0, + } async for kind, st, data in self.nats_client.stream(worker.worker_id, payload): yield (kind, st, data) return diff --git a/infera/router/token_ids.py b/infera/router/token_ids.py new file mode 100644 index 00000000..83eacfc7 --- /dev/null +++ b/infera/router/token_ids.py @@ -0,0 +1,134 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Reading the token ids an engine reports alongside the text it streams. + +Migration continues a generation on another worker. Handing over the decoded +text works, but re-encoding it is not guaranteed to reproduce the ids the model +actually sampled -- tokenizers are not injective on the way back, so a boundary +can shift and the second worker resumes from a slightly different sequence. +Asking the engine for the ids removes that step entirely. + +Both supported engines can report them, and neither agrees on where: + + vLLM chunk["choices"][i]["token_ids"] -- the deltas in this chunk + chunk["prompt_token_ids"] -- the prompt, first chunk only + SGLang chunk["sglext"]["completion_token_ids"] -- one list per choice + chunk["sglext"]["prompt_token_ids"] + +Both shapes are accepted wherever they appear, because the field has moved +between engine releases and a router pinned to one layout would silently stop +being exact. Silently is the problem: every reader here returns None rather +than a guess, and a caller that gets None falls back to carrying text, which is +approximate but never wrong about what was produced. +""" + +from __future__ import annotations + +from infera.common.worker_pool import EngineType + + +def deltas_from_chunk(obj: object) -> list[int] | None: + """The ids generated in this chunk, or None if the engine did not say. + + None is not an error. It is the ordinary state for an engine that was not + asked for ids, or was asked and does not support it on this endpoint. + """ + if not isinstance(obj, dict): + return None + choices = obj.get("choices") + if isinstance(choices, list) and choices and isinstance(choices[0], dict): + ids = _int_list(choices[0].get("token_ids")) + if ids is not None: + return ids + ext = obj.get("sglext") + if isinstance(ext, dict): + per_choice = ext.get("completion_token_ids") + # One list per choice; only the first is ours (n > 1 is rejected + # upstream for migratable requests). + if isinstance(per_choice, list) and per_choice: + return _int_list(per_choice[0]) + return _int_list(ext.get("token_ids")) + return None + + +def prompt_from_chunk(obj: object) -> list[int] | None: + """The prompt ids, which engines attach to the first chunk only. + + This is what makes an exact continuation possible at all: without it the + router would have to re-encode the prompt, reintroducing on the input side + exactly the ambiguity the output ids were fetched to avoid. + """ + if not isinstance(obj, dict): + return None + ids = _int_list(obj.get("prompt_token_ids")) + if ids is not None: + return ids + choices = obj.get("choices") + if isinstance(choices, list) and choices and isinstance(choices[0], dict): + ids = _int_list(choices[0].get("prompt_token_ids")) + if ids is not None: + return ids + ext = obj.get("sglext") + if isinstance(ext, dict): + return _int_list(ext.get("prompt_token_ids")) + return None + + +def strip_token_ids(obj: dict) -> bool: + """Remove the id fields from a chunk on its way to the client. + + The router asks for these; the caller did not. Leaving them in would put a + field in the response that the same request does not produce when migration + is off, which is a difference an operator's config should not make. + + Returns whether anything was removed, so a caller can skip re-serialising a + chunk that does not need it. + """ + touched = obj.pop("prompt_token_ids", None) is not None + touched |= obj.pop("sglext", None) is not None + choices = obj.get("choices") + if isinstance(choices, list): + for choice in choices: + if isinstance(choice, dict): + touched |= choice.pop("token_ids", None) is not None + touched |= choice.pop("prompt_token_ids", None) is not None + return touched + + +def supports_streaming_ids(engine: EngineType | None, path: str) -> bool: + """Whether asking this engine for ids on this endpoint is safe. + + SGLang rejects the request outright -- not ignores it -- when ids are asked + for on streaming chat, so sending it anyway would turn every migratable + chat request into a 400. Its completions endpoint is fine, and vLLM + supports both. + + ATOM is left out: it is not known to accept the field, and an engine that + errors on an unknown parameter would fail requests that work today. + """ + if engine == EngineType.VLLM: + return True + if engine == EngineType.SGLANG: + return not path.endswith("/chat/completions") + return False + + +def _int_list(value: object) -> list[int] | None: + """A list of ids, or None for anything else. + + Rejects rather than filters: a partially-parsed sequence would be a + plausible-looking prefix of the truth, and continuing from it drops output + the client already read. + """ + if not isinstance(value, list): + return None + out: list[int] = [] + for item in value: + # bool is an int in Python and never a token id. + if isinstance(item, bool) or not isinstance(item, int): + return None + out.append(item) + return out diff --git a/infera/server/__main__.py b/infera/server/__main__.py index 75057355..e4b33ced 100644 --- a/infera/server/__main__.py +++ b/infera/server/__main__.py @@ -260,7 +260,15 @@ def on_worker_removed(worker_id: str) -> None: nats_client=nats_request_client, request_max_retries=args.request_max_retries, breaker=breaker, + migration_limit=args.migration_limit, ) + if args.migration_limit: + logger.info( + "request migration enabled (limit=%d): a generation whose worker " + "goes away mid-stream continues on another one", + args.migration_limit, + ) + scaler = None if args.enable_scaling_api: # Resolved lazily on first call: the deployment is read from this Pod's diff --git a/infera/server/args.py b/infera/server/args.py index 7a059f74..207306c4 100644 --- a/infera/server/args.py +++ b/infera/server/args.py @@ -214,9 +214,22 @@ def parse_server_args(argv: list[str] | None = None) -> argparse.Namespace: help="Bounded failover: number of ALTERNATE mixed workers to try if a " "dispatch fails BEFORE any response data has reached the client " "(unreachable / NATS error / idle-timeout-before-first-token / 429 " - "backlog). Mid-stream failures are never retried. Default 1; 0 disables. " + "backlog). Mid-stream failures are not retried here; see " + "--migration-limit for those. Default 1; 0 disables. " "Overrides $INFERA_REQUEST_MAX_RETRIES.", ) + parser.add_argument( + "--migration-limit", + type=int, + default=int(os.environ.get("INFERA_MIGRATION_LIMIT", "0") or 0), + help="Carry a streaming generation to another worker when its own goes " + "away mid-stream, instead of ending the response with an error. The " + "text produced so far is appended to the prompt so the client sees one " + "uninterrupted stream. Requires the NATS transport, applies to mixed " + "(non-PD) workers, and does not preserve sampling state -- the " + "continuation is not byte-identical to what the original worker would " + "have produced. Default 0 (disabled). Overrides $INFERA_MIGRATION_LIMIT.", + ) parser.add_argument( "--breaker-failure-threshold", type=int, diff --git a/infera/server/metrics.py b/infera/server/metrics.py index b130bcd5..6f8ef227 100644 --- a/infera/server/metrics.py +++ b/infera/server/metrics.py @@ -112,6 +112,20 @@ registry=REGISTRY, ) +migrations_total = Counter( + "infera_migrations_total", + "Live generations carried from one worker to another, by what prompted it.", + labelnames=("reason",), # reason โˆˆ {stream_broken, worker_draining} + registry=REGISTRY, +) + +migrations_failed_total = Counter( + "infera_migrations_failed_total", + "Generations that could not be carried, and so ended visibly to the client.", + labelnames=("reason",), # reason โˆˆ {no_candidate, not_nats, no_first_byte, limit, poisoned} + registry=REGISTRY, +) + pd_bootstrap_failures_total = Counter( "infera_pd_bootstrap_failures_total", "PD bootstrap protocol failures (missing bootstrap_addr, P unreachable, etc.).", diff --git a/manual/features/feature_matrix.md b/manual/features/feature_matrix.md index 7c6f97fe..1acc3833 100644 --- a/manual/features/feature_matrix.md +++ b/manual/features/feature_matrix.md @@ -15,9 +15,13 @@ engine. **Legend:** โœ… supported ยท ๐Ÿšง work in progress ยท blank = not suppor | **KV-Aware Routing + DP-Attention** | โœ… | โœ… | โœ… | [KV-Aware Routing][kv] | | **Tiered KV Cache Offload (kvd)** | โœ… | ๐Ÿšง | ๐Ÿšง | [KV Cache Offload][kvd] | | **Multimodal (image / audio / video)** | | | | | +| **Request Migration** | โœ… | โœ… | โœ… | [Request Migration][mig] | KV-cache offload (`kvd`), including AIC GPU-Direct, is **vLLM-only** today. +Request migration does not depend on the engine, but does require the NATS +request transport and mixed (non-PD) workers; it is off unless enabled. [pd]: ./pd_disaggregation.md [kv]: ./kv_aware_routing.md [kvd]: ./kv_cache_offload.md +[mig]: ./graceful_shutdown.md#request-migration diff --git a/manual/features/graceful_shutdown.md b/manual/features/graceful_shutdown.md index 9559bd4c..dbeae890 100644 --- a/manual/features/graceful_shutdown.md +++ b/manual/features/graceful_shutdown.md @@ -6,66 +6,111 @@ 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:** nothing, for finishing in-flight +produces a burst of client errors. **Requires:** nothing, to finish in-flight work; Kubernetes with the default `kubernetes` discovery backend for the advance -notice described below. +notice below. ``` -```{important} -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. -``` - -## What happens - Removing a worker โ€” a rolling update, a scale-down, draining a node โ€” separates 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 and long before it is told to stop. +1. **It stops receiving.** The router stops choosing the worker as soon as its + removal is requested, which is before the worker itself is told to stop, 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 exits. -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 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. 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. +delay and sizes the termination grace period to cover the whole sequence, from +the drain timeout you set. For hand-written manifests and the per-stage timings, +see [Scaling a fleet](scaling.md). + +```{important} +Finishing in-flight work happens on every backend. 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. +``` + +## Request migration + +`--drain-timeout` is how long one generation is worth waiting for. Most finish +well inside it; a long one may not, and something has to happen at the deadline. +By default it is cut, and the client reads that as a failure. + +Request migration is the alternative: the generation is finished on a different +worker, so the client reads one uninterrupted stream and never learns a worker +changed underneath it. It covers the opposite case too โ€” a worker that crashes, +is evicted, or drops off the network, where there was no advance notice and no +drain window to run out of. + +It is off by default: + +```bash +infera-server --migration-limit 1 # or $INFERA_MIGRATION_LIMIT=1 +``` + +The limit is how many times one generation may be moved. `1` covers a worker +dying without letting a request wander the fleet during a broader outage. + +Migration does not shorten the drain: a generation is moved only after it has +had the window it was promised, so a Pod exits at the same moment either way. +What changes is whether the client sees an error. + +```{important} +The continuation is **not byte-identical** to what the original worker would +have produced: sampling state does not move with the request. Output stays +coherent and the seam is not visible to a reader, but a caller who needs +reproducible output for a fixed seed should leave this off. +``` + +### What it applies to + +Streaming requests on mixed (non-PD) workers using the NATS request transport. +Everything else keeps the behaviour it had: PD streams and HTTP-transport +workers end with an error as before, and non-streaming requests are already +covered by the ordinary failover in `--request-max-retries`, which re-runs them +cleanly rather than stitching one together. + +Accuracy depends on the engine. Where it reports the token ids it sampled โ€” vLLM +on both endpoints, SGLang on completions โ€” the continuation resumes from exactly +those; otherwise the decoded text is carried instead, which reads the same but +may not tokenize identically. The choice is automatic, and nothing fails because +the ids were unavailable. + +Requests asking for more than one completion โ€” `n`, `best_of`, or a batch of +prompts โ€” are excluded: there is no single generation to carry. + +A request can also stop being migratable partway through, and then ends with an +error rather than an approximation: streams the router cannot parse, tool calls +or reasoning content on the text path, since they do not appear in the text the +client receives, and a pre-tokenized prompt whose engine stops reporting ids, +which leaves nothing that can be extended. + +### Observing it + +`infera_migrations_total{reason}` counts generations moved, separating a rollout +doing its job (`worker_draining`) from a worker dying (`stream_broken`). +`infera_migrations_failed_total{reason}` counts the ones that could not be; a +rising `no_candidate` usually means the model has too few replicas for a +migration to have anywhere to go. + +## Where the advance notice is missing + +In-flight work is always finished. The head start is what varies, and two cases +do not get it: + +**A worker stopped without its Pod being removed** โ€” a liveness probe restarting +the container, a node shutting down, someone killing the process. Nothing +observes a deletion that never happens, so the router keeps sending work until +the worker itself drops out of the pool. + +**Deployments outside Kubernetes**, which discover through an external etcd +where a worker record is simply present or absent, with nothing to observe that +a process is leaving. + +In both cases a draining worker is absent from `/v1/workers` rather than shown +as draining. A Pod being deleted does show as `draining`, in the window between +its removal being requested and the drain beginning. diff --git a/tests/unit/common/test_drain_handback.py b/tests/unit/common/test_drain_handback.py new file mode 100644 index 00000000..4a9b7f84 --- /dev/null +++ b/tests/unit/common/test_drain_handback.py @@ -0,0 +1,229 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Handing resumable generations back when the drain window runs out. + +A generation that outlives the drain used to be cancelled, which the client read +as a failure. Now the ones the router can resume are handed back to it instead +and finish on another worker. + +This happens *after* the wait, never in place of it. Moving a generation costs +the next worker a re-read of everything produced so far, and a request that +would have finished on its own within the window should simply be allowed to. + +The load-bearing constraint is the other half: a request the router *cannot* +resume must still be left alone. Cutting one early would turn a shutdown the +client never noticed into an error it did. +""" + +from __future__ import annotations + +import asyncio +import json + +import pytest + +from infera.common.nats_request import ( + DRAINING_NOTICE, + HDR_TYPE, + TYPE_ERROR, + NatsRequestServer, +) + + +class _Msg: + def __init__(self, payload: dict, inbox: str): + self.data = json.dumps(payload).encode() + self.reply = inbox + self.headers = None + + async def ack(self): + pass + + +class _Conn: + """Records replies, so a test can see what the router would have received.""" + + def __init__(self): + self.replies: list[tuple[str, str, bytes]] = [] + + async def publish(self, subject, data=b"", headers=None): + self.replies.append((subject, (headers or {}).get(HDR_TYPE), data)) + + async def drain(self): + pass + + def new_inbox(self): + return "_INBOX.test" + + def types_for(self, inbox: str) -> list[str]: + return [t for (subj, t, _d) in self.replies if subj == inbox] + + def payloads_for(self, inbox: str) -> list[bytes]: + return [d for (subj, _t, d) in self.replies if subj == inbox] + + +def _server(conn) -> NatsRequestServer: + srv = NatsRequestServer.__new__(NatsRequestServer) + srv._nc = conn + srv._inflight = {} + srv._migratable = set() + srv._handed_back = set() + srv._sub = None + srv._cancel_sub = None + srv._js = None + srv._worker_id = "w1" + srv._max_duration = 0 + srv._http = None + return srv + + +async def _never_ends(): + await asyncio.Event().wait() + + +def _accept(srv, inbox: str, *, migratable: bool) -> asyncio.Task: + """Register an in-flight request the way _on_request would.""" + task = asyncio.create_task(_never_ends()) + srv._inflight[inbox] = task + if json.loads(_Msg({"migratable": migratable}, inbox).data).get("migratable"): + srv._migratable.add(inbox) + return task + + +@pytest.mark.asyncio +async def test_a_resumable_generation_is_returned_at_once(): + conn = _Conn() + srv = _server(conn) + task = _accept(srv, "inbox-a", migratable=True) + + await srv._hand_back_migratable() + await asyncio.sleep(0) + + assert conn.types_for("inbox-a") == [TYPE_ERROR] + assert conn.payloads_for("inbox-a") == [DRAINING_NOTICE] + assert task.cancelled() or task.cancelling(), "the engine must stop generating" + + +@pytest.mark.asyncio +async def test_a_generation_nobody_can_resume_is_left_alone(): + """The whole feature is opt-in from the router's side. Without that promise + the worker must drain the slow way, or it converts a shutdown the client + never saw into a visible failure.""" + conn = _Conn() + srv = _server(conn) + task = _accept(srv, "inbox-b", migratable=False) + + await srv._hand_back_migratable() + await asyncio.sleep(0) + + assert conn.replies == [], "nothing may be sent for a request that cannot move" + assert not task.done() + task.cancel() + + +@pytest.mark.asyncio +async def test_only_the_resumable_half_is_returned(): + conn = _Conn() + srv = _server(conn) + movable = _accept(srv, "yes", migratable=True) + staying = _accept(srv, "no", migratable=False) + + await srv._hand_back_migratable() + await asyncio.sleep(0) + + assert conn.types_for("yes") == [TYPE_ERROR] + assert conn.types_for("no") == [] + assert movable.cancelled() or movable.cancelling() + assert not staying.done() + staying.cancel() + + +@pytest.mark.asyncio +async def test_the_handover_is_announced_once(): + """The cancellation that follows also reports an error. Two frames for one + event would have the router read a planned handover as a second failure.""" + conn = _Conn() + srv = _server(conn) + _accept(srv, "inbox-c", migratable=True) + + await srv._hand_back_migratable() + assert "inbox-c" in srv._handed_back, "the cancel path must know to stay quiet" + + # What _proxy does when the cancellation lands. + if "inbox-c" not in srv._handed_back: + await srv._reply("inbox-c", TYPE_ERROR, b"request cancelled") + assert conn.types_for("inbox-c") == [TYPE_ERROR] + + +@pytest.mark.asyncio +async def test_a_finished_request_is_not_disturbed(): + conn = _Conn() + srv = _server(conn) + done = asyncio.create_task(asyncio.sleep(0)) + await done + srv._inflight["gone"] = done + srv._migratable.add("gone") + + await srv._hand_back_migratable() + assert conn.replies == [] + + +@pytest.mark.asyncio +async def test_nothing_in_flight_is_a_no_op(): + conn = _Conn() + srv = _server(conn) + await srv._hand_back_migratable() + assert conn.replies == [] + + +@pytest.mark.asyncio +async def test_the_wait_comes_first_and_the_handover_after(): + """A request that finishes inside the drain window costs nothing. Handing + it over early would buy the next worker a re-read for no reason.""" + conn = _Conn() + srv = _server(conn) + finished = asyncio.create_task(asyncio.sleep(0)) + await finished + srv._inflight["quick"] = finished + srv._migratable.add("quick") + still_going = _accept(srv, "slow", migratable=True) + + await srv.stop(drain=True, drain_timeout=0.01) + await asyncio.sleep(0) + + assert conn.types_for("quick") == [], "it finished on its own; nothing to hand back" + assert conn.payloads_for("slow") == [DRAINING_NOTICE] + assert still_going.cancelled() or still_going.cancelling() + + +@pytest.mark.asyncio +async def test_a_zero_window_still_hands_over_rather_than_cutting(): + """`--drain-timeout 0` means leave now, not sever what could have lived. + The handover is a local publish, so it costs nothing to honour that.""" + conn = _Conn() + srv = _server(conn) + task = _accept(srv, "inbox-z", migratable=True) + + await srv.stop(drain=True, drain_timeout=0) + await asyncio.sleep(0) + + assert conn.payloads_for("inbox-z") == [DRAINING_NOTICE] + assert task.cancelled() or task.cancelling() + + +@pytest.mark.asyncio +async def test_an_abrupt_stop_hands_nothing_over(): + """Not a drain: this is the emergency path, where the caller has asked for + everything to stop at once rather than for an orderly exit.""" + conn = _Conn() + srv = _server(conn) + task = _accept(srv, "inbox-a", migratable=True) + + await srv.stop() + await asyncio.sleep(0) + + assert conn.replies == [] + assert task.cancelled() or task.cancelling() diff --git a/tests/unit/router/test_failover.py b/tests/unit/router/test_failover.py index 65a2b99f..e731ad18 100644 --- a/tests/unit/router/test_failover.py +++ b/tests/unit/router/test_failover.py @@ -121,7 +121,7 @@ async def test_stream_no_retry_after_first_byte(): resp = await r.dispatch({"model": "m"}, stream=True) body = await _drain_stream(resp) assert b"partial" in body # first chunk delivered - assert b"stream failed mid-stream" in body # error surfaced inline + assert b"failed mid-stream" in body # error surfaced inline assert b"SHOULD-NOT-BE-USED" not in body assert nats.streamed == ["w1"] # committed to w1, no failover await r.aclose() diff --git a/tests/unit/router/test_migration.py b/tests/unit/router/test_migration.py new file mode 100644 index 00000000..7c3cc275 --- /dev/null +++ b/tests/unit/router/test_migration.py @@ -0,0 +1,307 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Carrying a generation to another worker without the client noticing. + +The whole feature rests on one property: the text handed to the next worker is +exactly what the client already received. If it is short, the client reads a +gap; if it is long, it reads a repetition. Both are worse than the failure +migration exists to hide, so most of this file is about that equality. +""" + +from __future__ import annotations + +import json + +from infera.router.migration import MigrationState + + +def chat_chunk(content: str) -> bytes: + return f"data: {json.dumps({'choices': [{'delta': {'content': content}}]})}\n\n".encode() + + +def completion_chunk(text: str) -> bytes: + return f"data: {json.dumps({'choices': [{'text': text}]})}\n\n".encode() + + +def test_the_carried_text_is_what_the_client_received(): + st = MigrationState({"messages": [{"role": "user", "content": "hi"}]}, limit=1) + for piece in ("Hello", ",", " world"): + st.observe(chat_chunk(piece)) + assert st.produced_text == "Hello, world" + + +def test_completions_carry_the_same_way(): + st = MigrationState({"prompt": "once"}, limit=1) + st.observe(completion_chunk(" upon")) + st.observe(completion_chunk(" a time")) + assert st.produced_text == " upon a time" + + +def test_several_chunks_in_one_frame_are_all_counted(): + # A slow reader coalesces frames; the transport is free to deliver them + # together and the accumulated text must not depend on that. + st = MigrationState({"prompt": ""}, limit=1) + st.observe(completion_chunk("a") + completion_chunk("b") + completion_chunk("c")) + assert st.produced_text == "abc" + assert st.produced_tokens == 3 + + +def test_the_done_sentinel_is_not_carried(): + """`[DONE]` belongs to the stream the client is reading, not to the + generation. Carried into a prompt it would become model input.""" + st = MigrationState({"prompt": ""}, limit=1) + st.observe(completion_chunk("x")) + st.observe(b"data: [DONE]\n\n") + assert st.produced_text == "x" + + +def test_a_chat_continuation_appends_an_assistant_turn(): + body = {"messages": [{"role": "user", "content": "count"}], "max_tokens": 10} + st = MigrationState(body, limit=1) + for piece in ("one", " two"): + st.observe(chat_chunk(piece)) + + nxt = st.next_continuation().body + assert nxt["messages"][-1] == {"role": "assistant", "content": "one two"} + # Two tokens spent, eight left: the next worker finishes the answer rather + # than producing a second one of full length. + assert nxt["max_tokens"] == 8 + # The original is untouched -- a failed migration must be able to fall back. + assert body["max_tokens"] == 10 + assert len(body["messages"]) == 1 + + +def test_a_completion_continuation_extends_the_prompt(): + st = MigrationState({"prompt": "once", "max_tokens": 5}, limit=1) + st.observe(completion_chunk(" upon")) + nxt = st.next_continuation().body + assert nxt["prompt"] == "once upon" + assert nxt["max_tokens"] == 4 + + +def test_the_budget_never_reaches_zero(): + """A request for zero tokens is rejected, which would turn a migration the + client cannot see into an error it can.""" + st = MigrationState({"prompt": "", "max_tokens": 2}, limit=1) + for _ in range(5): + st.observe(completion_chunk("x")) + assert st.next_continuation().body["max_tokens"] == 1 + + +def test_a_request_without_a_budget_stays_without_one(): + st = MigrationState({"prompt": ""}, limit=1) + st.observe(completion_chunk("x")) + assert "max_tokens" not in st.next_continuation().body + + +def test_an_unparseable_chunk_disables_migration(): + """Continuing from a prefix that could not be fully reconstructed would + drop output the client already read -- worse than not migrating.""" + st = MigrationState({"prompt": ""}, limit=1) + st.observe(completion_chunk("good")) + st.observe(b"data: {not json\n\n") + assert st.poisoned + assert not st.can_migrate() + + +def test_chunks_after_poisoning_are_ignored(): + st = MigrationState({"prompt": ""}, limit=1) + st.observe(b"data: {not json\n\n") + st.observe(completion_chunk("later")) + assert st.produced_text == "" + + +def test_the_migration_limit_is_enforced(): + st = MigrationState({"prompt": ""}, limit=2) + assert st.can_migrate() + st.next_continuation() + assert st.can_migrate() + st.next_continuation() + assert not st.can_migrate(), "a request must not migrate forever" + + +def test_a_zero_limit_disables_migration(): + st = MigrationState({"prompt": ""}, limit=0) + assert not st.can_migrate() + + +def exact_chunk(content: str, ids: list[int], prompt_ids: list[int] | None = None) -> bytes: + """A chunk from an engine that was asked to report token ids.""" + obj = {"choices": [{"delta": {"content": content}, "token_ids": ids}]} + if prompt_ids is not None: + obj["prompt_token_ids"] = prompt_ids + return f"data: {json.dumps(obj)}\n\n".encode() + + +def test_exact_ids_are_carried_instead_of_text(): + """The point of the whole id path: the next worker resumes from the + sequence the model sampled, not from a re-encoding of the words.""" + st = MigrationState({"prompt": "hi", "max_tokens": 10}, limit=1, path="/v1/completions") + st.observe(exact_chunk("Hel", [100], prompt_ids=[1, 2])) + st.observe(exact_chunk("lo", [200])) + + assert st.is_exact() + cont = st.next_continuation() + assert cont.exact + assert cont.body["prompt"] == [1, 2, 100, 200] + assert cont.body["max_tokens"] == 8, "two real tokens, not two chunks" + + +def test_the_token_count_is_the_real_one_when_ids_are_known(): + """Without ids this is chunks-with-text, which is only a good guess. A + chunk carrying several tokens makes the two differ.""" + st = MigrationState({"prompt": "", "max_tokens": 20}, limit=1, path="/v1/completions") + st.observe(exact_chunk("a b c", [1, 2, 3], prompt_ids=[9])) + assert st.produced_tokens == 3 + assert st.next_continuation().body["max_tokens"] == 17 + + +def test_ids_without_a_prompt_are_not_enough(): + """Appending exact output ids to a re-encoded prompt just moves the + ambiguity to the other end of the sequence.""" + st = MigrationState({"prompt": "hi"}, limit=1, path="/v1/completions") + st.observe(exact_chunk("out", [5])) # engine never sent prompt ids + assert not st.is_exact() + assert st.next_continuation().exact is False + + +def test_text_the_ids_do_not_cover_abandons_the_id_path(): + """Half the output accounted for is worse than none: continuing from those + ids would drop the text they omit, which the client has already read.""" + st = MigrationState({"prompt": ""}, limit=1, path="/v1/completions") + st.observe(exact_chunk("counted", [1], prompt_ids=[0])) + st.observe(completion_chunk("unaccounted")) + assert not st.is_exact() + cont = st.next_continuation() + assert cont.exact is False + assert cont.body["prompt"] == "countedunaccounted", "falls back to the full text" + + +def test_an_exact_chat_continuation_moves_to_the_completions_path(): + """Chat has no pre-tokenized entry, so exactness costs a change of + endpoint; the caller is told so it can convert the replies back.""" + st = MigrationState( + {"messages": [{"role": "user", "content": "hi"}]}, + limit=1, + path="/v1/chat/completions", + ) + st.observe(exact_chunk("part", [7], prompt_ids=[1, 2])) + + cont = st.next_continuation() + assert cont.exact + assert cont.path == "/v1/completions" + assert cont.body["prompt"] == [1, 2, 7] + assert "messages" not in cont.body + + +def test_a_chat_request_with_tools_keeps_the_text_path(): + """A completions request cannot emit a tool call. Losing that halfway + through an answer is worse than a token boundary that might differ.""" + st = MigrationState( + { + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"type": "function", "function": {"name": "f"}}], + }, + limit=1, + path="/v1/chat/completions", + ) + st.observe(exact_chunk("part", [7], prompt_ids=[1, 2])) + + cont = st.next_continuation() + assert cont.exact is False + assert cont.path == "/v1/chat/completions" + assert cont.body["messages"][-1]["content"] == "part" + + +def test_structured_output_also_keeps_the_text_path(): + st = MigrationState( + { + "messages": [{"role": "user", "content": "hi"}], + "response_format": {"type": "json_object"}, + }, + limit=1, + path="/v1/chat/completions", + ) + st.observe(exact_chunk("{", [7], prompt_ids=[1])) + assert st.next_continuation().exact is False + + +def test_the_ids_are_taken_out_before_the_client_sees_them(): + """Enabling migration must not change the shape of the response.""" + st = MigrationState({"prompt": ""}, limit=1, path="/v1/completions") + out = st.observe(exact_chunk("hi", [42], prompt_ids=[1])) + + assert b"token_ids" not in out + assert b"prompt_token_ids" not in out + assert json.loads(out.split(b"data: ")[1])["choices"][0]["delta"]["content"] == "hi" + assert st.produced_text == "hi", "still recorded, just not forwarded" + + +def test_a_chunk_without_ids_is_forwarded_untouched(): + """No ids means nothing to strip, and the engine's own bytes go through + without a re-serialisation that could perturb them.""" + st = MigrationState({"prompt": ""}, limit=1, path="/v1/completions") + original = completion_chunk("hi") + assert st.observe(original) is original + + +def test_the_router_only_field_is_not_passed_on(): + st = MigrationState({"prompt": "", "return_token_ids": True}, limit=1, path="/v1/completions") + st.observe(exact_chunk("hi", [1], prompt_ids=[0])) + assert "return_token_ids" not in st.next_continuation().body + + +def tool_call_chunk() -> bytes: + delta = {"tool_calls": [{"index": 0, "function": {"name": "f", "arguments": '{"a"'}}]} + return f"data: {json.dumps({'choices': [{'delta': delta}]})}\n\n".encode() + + +def test_a_tool_call_stops_the_text_path(): + """Tool calls arrive under their own key, not as content, so carried text + would omit them entirely -- and the client, already holding half a call, + would be sent a whole one after the migration.""" + st = MigrationState({"messages": []}, limit=1, path="/v1/chat/completions") + st.observe(chat_chunk("Let me check")) + st.observe(tool_call_chunk()) + + assert st.poisoned + assert not st.can_migrate() + + +def test_hidden_reasoning_stops_the_text_path(): + """A reasoning parser moves the model's thinking out of `content`. Carrying + only what the client saw would resume from an answer with its reasoning + removed, which is missing output rather than imprecise output.""" + st = MigrationState({"messages": []}, limit=1, path="/v1/chat/completions") + st.observe(b'data: {"choices": [{"delta": {"reasoning_content": "hmm..."}}]}\n\n') + + assert st.poisoned + + +def test_exact_ids_survive_a_tool_call(): + """The ids are the tokens behind whatever the parser emitted, so they carry + the call itself. Only the text path is defeated by it.""" + st = MigrationState({"messages": []}, limit=1, path="/v1/chat/completions") + st.observe(exact_chunk("Let me check", [10], prompt_ids=[1])) + obj = { + "choices": [{"delta": {"tool_calls": [{"index": 0}]}, "token_ids": [20, 21]}], + } + st.observe(f"data: {json.dumps(obj)}\n\n".encode()) + + assert not st.poisoned + assert st.is_exact() + assert st.next_continuation().body["prompt"] == [1, 10, 20, 21] + + +def test_keepalives_and_role_frames_add_nothing(): + # An opening chat frame carries `role` and no content; comment lines are + # heartbeats. Neither is output, and counting them would shorten the answer. + st = MigrationState({"prompt": ""}, limit=1) + st.observe(b'data: {"choices": [{"delta": {"role": "assistant"}}]}\n\n') + st.observe(b": keepalive\n\n") + st.observe(completion_chunk("real")) + assert st.produced_text == "real" + assert st.produced_tokens == 1 diff --git a/tests/unit/router/test_migration_dispatch.py b/tests/unit/router/test_migration_dispatch.py new file mode 100644 index 00000000..426c1d77 --- /dev/null +++ b/tests/unit/router/test_migration_dispatch.py @@ -0,0 +1,476 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Migration as the client experiences it. + +`test_migration.py` covers the accounting; this drives the dispatch path, where +the property that matters is what reaches the client. A migration the client can +detect -- a gap, a repetition, an error frame -- has failed at its only job, +however correct the bookkeeping was. +""" + +from __future__ import annotations + +import json + +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.mixed import MixedRouter +from infera.router.policy.target import RouteTarget + + +def _w(wid: str) -> WorkerInfo: + return WorkerInfo( + worker_id=wid, + url=f"http://{wid}", + model_name="m", + engine=EngineType.SGLANG, + request_transport="nats", + disagg_mode=DisaggMode.MIXED, + ) + + +class _Pool: + def __init__(self, workers): + self._workers = workers + + def list_active(self, model=None, mode=None): + return list(self._workers) + + +class _Policy: + """Round-robins, so a migration lands somewhere other than the failure.""" + + def __init__(self): + self.picks: list[str] = [] + + def pick(self, candidates, body, role_hint=None): + target = RouteTarget(candidates[0]) + self.picks.append(target.worker.worker_id) + return target, [] + + def on_request_started(self, route_key, blocks): + pass + + def on_request_finished(self, route_key, blocks): + pass + + +def chunk(text: str) -> bytes: + return f"data: {json.dumps({'choices': [{'delta': {'content': text}}]})}\n\n".encode() + + +class _ScriptedNats: + """Replays a scripted reply per worker, recording the bodies it was sent.""" + + def __init__(self, scripts: dict[str, list[tuple]]): + self.scripts = scripts + self.bodies: dict[str, dict] = {} + self.sent: dict[str, dict] = {} + + async def admit(self, worker_id): + return True + + async def stream(self, worker_id, payload): + self.sent[worker_id] = payload + self.bodies[worker_id] = payload["body"] + assert "path" in payload + for item in self.scripts.get(worker_id, []): + yield item + + +def _router(nats, workers, *, limit: int) -> MixedRouter: + return MixedRouter(_Pool(workers), _Policy(), nats_client=nats, migration_limit=limit) + + +async def _collect(response) -> bytes: + out = b"" + async for piece in response.body_iterator: + out += piece if isinstance(piece, bytes) else piece.encode() + return out + + +@pytest.mark.asyncio +async def test_a_broken_stream_continues_on_another_worker(): + """The client reads one uninterrupted answer across a worker failure.""" + nats = _ScriptedNats( + { + # w1 produces two tokens, then its stream breaks. + "w1": [ + (TYPE_DATA, None, chunk("Hello")), + (TYPE_DATA, None, chunk(", ")), + (TYPE_ERROR, None, b"worker vanished"), + ], + # w2 continues, and finishes. + "w2": [ + (TYPE_DATA, None, chunk("world")), + (TYPE_DONE, 200, b""), + ], + } + ) + r = _router(nats, [_w("w1"), _w("w2")], limit=1) + resp = await r.dispatch( + {"model": "m", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 10}, + stream=True, + ) + body = await _collect(resp) + + assert b"Hello" in body and b", " in body and b"world" in body + assert b"error" not in body, "the client must not see the failure" + # Nothing is repeated: the second worker continued rather than restarted. + assert body.count(b"Hello") == 1 + await r.aclose() + + +@pytest.mark.asyncio +async def test_the_second_worker_is_asked_to_continue_not_to_restart(): + """What the next worker receives is the original request plus what the + client already read, with the budget reduced by what it cost.""" + nats = _ScriptedNats( + { + "w1": [ + (TYPE_DATA, None, chunk("one")), + (TYPE_DATA, None, chunk(" two")), + (TYPE_ERROR, None, b"gone"), + ], + "w2": [(TYPE_DATA, None, chunk(" three")), (TYPE_DONE, 200, b"")], + } + ) + r = _router(nats, [_w("w1"), _w("w2")], limit=1) + resp = await r.dispatch( + {"model": "m", "messages": [{"role": "user", "content": "count"}], "max_tokens": 9}, + stream=True, + ) + await _collect(resp) + + sent = nats.bodies["w2"] + assert sent["messages"][-1] == {"role": "assistant", "content": "one two"} + assert sent["max_tokens"] == 7, "two tokens were already spent" + await r.aclose() + + +@pytest.mark.asyncio +async def test_a_failure_the_client_never_sees_is_still_visible_when_unmigratable(): + """With nowhere to go the stream ends with an error rather than silently + stopping -- a truncated answer with no explanation is worse.""" + nats = _ScriptedNats({"w1": [(TYPE_DATA, None, chunk("partial")), (TYPE_ERROR, None, b"gone")]}) + r = _router(nats, [_w("w1")], limit=1) # no second worker + resp = await r.dispatch( + {"model": "m", "messages": [{"role": "user", "content": "hi"}]}, stream=True + ) + body = await _collect(resp) + + assert b"partial" in body + assert b"error" in body + await r.aclose() + + +@pytest.mark.asyncio +async def test_migration_is_off_by_default(): + """It changes what a worker is asked to produce, so it is opt-in.""" + nats = _ScriptedNats( + {"w1": [(TYPE_DATA, None, chunk("x")), (TYPE_ERROR, None, b"gone")], "w2": []} + ) + r = _router(nats, [_w("w1"), _w("w2")], limit=0) + resp = await r.dispatch( + {"model": "m", "messages": [{"role": "user", "content": "hi"}]}, stream=True + ) + body = await _collect(resp) + + assert b"error" in body + assert "w2" not in nats.bodies, "nothing may be dispatched when migration is off" + await r.aclose() + + +def exact_chunk(content: str, ids: list[int], prompt_ids: list[int] | None = None) -> bytes: + obj = {"choices": [{"delta": {"content": content}, "token_ids": ids}]} + if prompt_ids is not None: + obj["prompt_token_ids"] = prompt_ids + return f"data: {json.dumps(obj)}\n\n".encode() + + +def completion_chunk(text: str, ids: list[int], prompt_ids: list[int] | None = None) -> bytes: + obj = {"choices": [{"text": text, "token_ids": ids}], "object": "text_completion"} + if prompt_ids is not None: + obj["prompt_token_ids"] = prompt_ids + return f"data: {json.dumps(obj)}\n\n".encode() + + +def _vllm(wid: str) -> WorkerInfo: + return WorkerInfo( + worker_id=wid, + url=f"http://{wid}", + model_name="m", + engine=EngineType.VLLM, + request_transport="nats", + disagg_mode=DisaggMode.MIXED, + ) + + +@pytest.mark.asyncio +async def test_an_exact_continuation_resumes_from_the_sampled_ids(): + """The second worker is handed the ids the model actually produced, not a + re-encoding of the words, so no token boundary can shift.""" + nats = _ScriptedNats( + { + "w1": [ + (TYPE_DATA, None, exact_chunk("Hel", [100], prompt_ids=[1, 2])), + (TYPE_DATA, None, exact_chunk("lo", [200])), + (TYPE_ERROR, None, b"gone"), + ], + "w2": [ + (TYPE_DATA, None, completion_chunk(" there", [300])), + (TYPE_DONE, 200, b""), + ], + } + ) + r = _router(nats, [_vllm("w1"), _vllm("w2")], limit=1) + resp = await r.dispatch( + {"model": "m", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 10}, + stream=True, + ) + body = await _collect(resp) + + sent = nats.sent["w2"] + assert sent["body"]["prompt"] == [1, 2, 100, 200], "prompt ids + everything sampled" + assert sent["path"] == "/v1/completions", "the only endpoint taking token ids" + assert sent["body"]["max_tokens"] == 8, "two real tokens spent" + assert b"Hel" in body and b"lo" in body and b" there" in body + assert b"error" not in body + await r.aclose() + + +@pytest.mark.asyncio +async def test_the_client_keeps_reading_chat_across_an_exact_migration(): + """The continuation is fetched from the completions endpoint, which answers + in a different shape. The client asked for chat and must keep getting it.""" + nats = _ScriptedNats( + { + "w1": [ + (TYPE_DATA, None, exact_chunk("start", [1], prompt_ids=[9])), + (TYPE_ERROR, None, b"gone"), + ], + "w2": [(TYPE_DATA, None, completion_chunk(" end", [2])), (TYPE_DONE, 200, b"")], + } + ) + r = _router(nats, [_vllm("w1"), _vllm("w2")], limit=1) + resp = await r.dispatch( + {"model": "m", "messages": [{"role": "user", "content": "hi"}]}, stream=True + ) + body = await _collect(resp) + + frames = [json.loads(ln[6:]) for ln in body.split(b"\n") if ln.startswith(b"data: ")] + assert all("delta" in f["choices"][0] for f in frames), "every frame is chat-shaped" + assert all("text" not in f["choices"][0] for f in frames) + assert [f["choices"][0]["delta"]["content"] for f in frames] == ["start", " end"] + await r.aclose() + + +@pytest.mark.asyncio +async def test_the_token_ids_never_reach_the_client(): + """The router asks for them for its own use. Leaving them in would make the + response depend on an operator's migration setting.""" + nats = _ScriptedNats( + {"w1": [(TYPE_DATA, None, exact_chunk("hi", [5], prompt_ids=[1])), (TYPE_DONE, 200, b"")]} + ) + r = _router(nats, [_vllm("w1")], limit=1) + body = await _collect( + await r.dispatch( + {"model": "m", "messages": [{"role": "user", "content": "hi"}]}, stream=True + ) + ) + assert b"token_ids" not in body + assert b"hi" in body + await r.aclose() + + +@pytest.mark.asyncio +async def test_ids_are_requested_only_when_a_migration_could_use_them(): + nats = _ScriptedNats({"w1": [(TYPE_DATA, None, chunk("x")), (TYPE_DONE, 200, b"")]}) + r = _router(nats, [_vllm("w1")], limit=1) + await _collect( + await r.dispatch({"model": "m", "messages": [{"role": "u", "content": "h"}]}, stream=True) + ) + assert nats.sent["w1"]["body"]["return_token_ids"] is True + await r.aclose() + + off = _ScriptedNats({"w1": [(TYPE_DATA, None, chunk("x")), (TYPE_DONE, 200, b"")]}) + r_off = _router(off, [_vllm("w1")], limit=0) + await _collect( + await r_off.dispatch( + {"model": "m", "messages": [{"role": "u", "content": "h"}]}, stream=True + ) + ) + assert "return_token_ids" not in off.sent["w1"]["body"], ( + "an engine must not be asked for work no migration will use" + ) + await r_off.aclose() + + +@pytest.mark.asyncio +async def test_sglang_chat_streams_are_never_asked_for_ids(): + """SGLang rejects the request outright, so asking would turn every + migratable chat request into a 400.""" + nats = _ScriptedNats({"w1": [(TYPE_DATA, None, chunk("x")), (TYPE_DONE, 200, b"")]}) + r = _router(nats, [_w("w1")], limit=1) # _w builds an SGLANG worker + await _collect( + await r.dispatch({"model": "m", "messages": [{"role": "u", "content": "h"}]}, stream=True) + ) + assert "return_token_ids" not in nats.sent["w1"]["body"] + await r.aclose() + + +@pytest.mark.asyncio +async def test_an_engine_that_reports_no_ids_still_migrates_on_text(): + """The exact path is an optimisation. Losing it must cost precision, not + the migration itself.""" + nats = _ScriptedNats( + { + "w1": [(TYPE_DATA, None, chunk("half")), (TYPE_ERROR, None, b"gone")], + "w2": [(TYPE_DATA, None, chunk(" done")), (TYPE_DONE, 200, b"")], + } + ) + r = _router(nats, [_w("w1"), _w("w2")], limit=1) + body = await _collect( + await r.dispatch( + {"model": "m", "messages": [{"role": "user", "content": "hi"}]}, stream=True + ) + ) + sent = nats.sent["w2"] + assert sent["path"] == "/v1/chat/completions", "no ids means no endpoint change" + assert sent["body"]["messages"][-1] == {"role": "assistant", "content": "half"} + assert b"half" in body and b" done" in body and b"error" not in body + await r.aclose() + + +@pytest.mark.asyncio +async def test_a_draining_worker_hands_its_stream_over(): + """A worker leaving on purpose looks like a broken one to the client -- the + stream continues either way -- but not to an operator, so the two are + counted separately.""" + from infera.common.nats_request import DRAINING_NOTICE + from infera.server import metrics + + before = metrics.migrations_total.labels(reason="worker_draining")._value.get() + nats = _ScriptedNats( + { + "w1": [(TYPE_DATA, None, chunk("half")), (TYPE_ERROR, None, DRAINING_NOTICE)], + "w2": [(TYPE_DATA, None, chunk(" done")), (TYPE_DONE, 200, b"")], + } + ) + r = _router(nats, [_w("w1"), _w("w2")], limit=1) + resp = await r.dispatch( + {"model": "m", "messages": [{"role": "user", "content": "hi"}]}, stream=True + ) + body = await _collect(resp) + + assert b"half" in body and b" done" in body + assert b"error" not in body + after = metrics.migrations_total.labels(reason="worker_draining")._value.get() + assert after == before + 1, "a planned handover is not counted as a fault" + await r.aclose() + + +@pytest.mark.asyncio +async def test_the_worker_is_told_whether_the_stream_can_be_resumed(): + """The worker cannot know on its own: handing a stream back early is only + safe because the router promised to continue it.""" + nats = _ScriptedNats({"w1": [(TYPE_DATA, None, chunk("x")), (TYPE_DONE, 200, b"")]}) + r = _router(nats, [_w("w1")], limit=1) + await _collect( + await r.dispatch( + {"model": "m", "messages": [{"role": "user", "content": "hi"}]}, stream=True + ) + ) + assert nats.sent["w1"]["migratable"] is True + await r.aclose() + + nats_off = _ScriptedNats({"w1": [(TYPE_DATA, None, chunk("x")), (TYPE_DONE, 200, b"")]}) + r_off = _router(nats_off, [_w("w1")], limit=0) + await _collect( + await r_off.dispatch( + {"model": "m", "messages": [{"role": "user", "content": "hi"}]}, stream=True + ) + ) + assert nats_off.sent["w1"]["migratable"] is False, ( + "a worker must not shed streams this router cannot resume" + ) + await r_off.aclose() + + +@pytest.mark.asyncio +async def test_a_tool_call_ends_the_stream_rather_than_replaying_it(): + """The client holds half a tool call. Resuming would send it a whole one + after that, so the stream ends visibly instead.""" + from infera.server import metrics + + before = metrics.migrations_failed_total.labels(reason="poisoned")._value.get() + tool = json.dumps( + {"choices": [{"delta": {"tool_calls": [{"index": 0, "function": {"name": "f"}}]}}]} + ) + nats = _ScriptedNats( + { + "w1": [ + (TYPE_DATA, None, chunk("checking")), + (TYPE_DATA, None, f"data: {tool}\n\n".encode()), + (TYPE_ERROR, None, b"gone"), + ], + "w2": [(TYPE_DATA, None, chunk("whatever")), (TYPE_DONE, 200, b"")], + } + ) + r = _router(nats, [_vllm("w1"), _vllm("w2")], limit=1) + body = await _collect( + await r.dispatch( + {"model": "m", "messages": [{"role": "user", "content": "hi"}]}, stream=True + ) + ) + + assert b"error" in body + assert "w2" not in nats.sent, "a request that cannot be rebuilt must not be resumed" + after = metrics.migrations_failed_total.labels(reason="poisoned")._value.get() + assert after == before + 1, "the lost capability has to be visible to an operator" + await r.aclose() + + +@pytest.mark.asyncio +async def test_a_draining_worker_that_cannot_be_replaced_says_so(): + """Ending with 'failed' would report a planned shutdown as a fault.""" + from infera.common.nats_request import DRAINING_NOTICE + + nats = _ScriptedNats( + {"w1": [(TYPE_DATA, None, chunk("half")), (TYPE_ERROR, None, DRAINING_NOTICE)]} + ) + r = _router(nats, [_vllm("w1")], limit=1) # nowhere else to go + body = await _collect( + await r.dispatch( + {"model": "m", "messages": [{"role": "user", "content": "hi"}]}, stream=True + ) + ) + assert b"shutting down" in body + assert b"failed" not in body + await r.aclose() + + +@pytest.mark.asyncio +async def test_the_limit_bounds_how_often_one_generation_moves(): + """Two workers that both break, with a budget of one move: the second + failure ends the stream instead of starting a third attempt.""" + nats = _ScriptedNats( + { + "w1": [(TYPE_DATA, None, chunk("a")), (TYPE_ERROR, None, b"gone")], + "w2": [(TYPE_DATA, None, chunk("b")), (TYPE_ERROR, None, b"gone too")], + } + ) + r = _router(nats, [_w("w1"), _w("w2")], limit=1) + resp = await r.dispatch( + {"model": "m", "messages": [{"role": "user", "content": "hi"}]}, stream=True + ) + body = await _collect(resp) + + assert b"a" in body and b"b" in body + assert b"error" in body, "the budget is spent; the failure now reaches the client" + await r.aclose() diff --git a/tests/unit/router/test_migration_regressions.py b/tests/unit/router/test_migration_regressions.py new file mode 100644 index 00000000..57c0aa89 --- /dev/null +++ b/tests/unit/router/test_migration_regressions.py @@ -0,0 +1,139 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Requests whose shape the carried prefix cannot represent. + +Every case here produced output that was wrong rather than absent, which is the +worse failure: a severed stream is visibly a failure, while a continuation built +from the wrong prefix reads as the model losing the plot. +""" + +from __future__ import annotations + +import json + +from infera.router.migration import MigrationState + + +def completion_chunk(text: str) -> bytes: + return f"data: {json.dumps({'choices': [{'text': text}]})}\n\n".encode() + + +def exact_chunk(content: str, ids: list[int], prompt_ids: list[int] | None = None) -> bytes: + obj = {"choices": [{"delta": {"content": content}, "token_ids": ids}]} + if prompt_ids is not None: + obj["prompt_token_ids"] = prompt_ids + return f"data: {json.dumps(obj)}\n\n".encode() + + +def test_a_pre_tokenized_prompt_is_never_carried_as_text(): + """`prompt` may be a token array, which engines accept verbatim. Formatting + one into a string yields its Python repr -- the next worker would be asked + to continue the literal text "[1, 2, 3]hello".""" + st = MigrationState({"prompt": [1, 2, 3]}, limit=1, path="/v1/completions") + st.observe(completion_chunk("hello")) # no ids: the text path is all there is + + assert not st.can_migrate(), "a token array cannot be extended with text" + + +def test_a_batch_prompt_is_never_migrated(): + """A list of prompts is several generations at once. There is no single + prefix to carry, and formatting the list would produce its repr.""" + st = MigrationState({"prompt": ["a", "b"]}, limit=1, path="/v1/completions") + st.observe(completion_chunk("hello")) + + assert not st.can_migrate() + + +def test_a_pre_tokenized_prompt_still_migrates_exactly(): + """The id path builds `prompt_ids + output_ids`, which is correct for a + token array -- it is only the text path that cannot express one.""" + st = MigrationState({"prompt": [1, 2, 3]}, limit=1, path="/v1/completions") + st.observe(exact_chunk("hello", [42], prompt_ids=[1, 2, 3])) + + assert st.can_migrate() + cont = st.next_continuation() + assert cont.exact + assert cont.body["prompt"] == [1, 2, 3, 42] + + +def test_more_than_one_choice_is_never_migrated(): + """Only the first choice is accumulated, so every other one would resume + from a prefix belonging to the first: not a gap, but the wrong content.""" + st = MigrationState({"prompt": "hi", "n": 2}, limit=1, path="/v1/completions") + st.observe(completion_chunk("hello")) + + assert not st.can_migrate() + + +def test_several_choices_in_a_chunk_are_noticed_even_without_n(): + """`best_of`, or an engine that returns several choices for its own + reasons, reaches the same place by a different route.""" + obj = {"choices": [{"index": 0, "text": "AAA"}, {"index": 1, "text": "BBB"}]} + st = MigrationState({"prompt": "hi"}, limit=1, path="/v1/completions") + st.observe(f"data: {json.dumps(obj)}\n\n".encode()) + + assert not st.can_migrate() + assert st.produced_text == "", "a prefix covering one choice of several is not a prefix" + + +def test_several_completions_are_rejected_on_the_exact_path_too(): + """Exact ids do not help here: they are accumulated for one choice, so a + second one would resume from the first one's tokens.""" + st = MigrationState({"prompt": "hi", "n": 2}, limit=1, path="/v1/completions") + obj = {"choices": [{"text": "A", "token_ids": [1]}], "prompt_token_ids": [9]} + st.observe(f"data: {json.dumps(obj)}\n\n".encode()) + + assert not st.can_migrate() + + +def test_a_batch_of_token_arrays_is_still_a_batch(): + st = MigrationState({"prompt": [[1, 2], [3, 4]]}, limit=1, path="/v1/completions") + assert not st.can_migrate() + + +def test_an_empty_prompt_list_is_not_read_as_a_token_array(): + """Nothing sensible can be carried, and treating it as a single generation + would put it back on the text path it cannot take.""" + st = MigrationState({"prompt": []}, limit=1, path="/v1/completions") + assert not st.can_migrate() + + +def test_best_of_is_rejected_like_n(): + st = MigrationState({"prompt": "hi", "best_of": 3}, limit=1, path="/v1/completions") + assert not st.can_migrate() + + +def test_ids_are_still_stripped_after_migration_is_ruled_out(): + """The ids are requested by the router and the client never asked for them. + Whether a request is still migratable is the router's business; it must not + change the shape of what the caller receives.""" + st = MigrationState({"messages": []}, limit=1, path="/v1/chat/completions") + st.observe(exact_chunk("hi", [1], prompt_ids=[0])) + + tool = {"choices": [{"delta": {"tool_calls": [{"index": 0}]}}]} + st.observe(f"data: {json.dumps(tool)}\n\n".encode()) + assert st.poisoned + + after = st.observe(exact_chunk("more", [12])) + assert b"token_ids" not in after, "ids leaked once the request stopped being migratable" + assert b"more" in after + + +def test_ids_are_stripped_even_after_an_unparseable_chunk(): + """The same holds for the other way a request stops being migratable.""" + st = MigrationState({"messages": []}, limit=1, path="/v1/chat/completions") + st.observe(b"data: {not json\n\n") + assert st.poisoned + + after = st.observe(exact_chunk("more", [12])) + assert b"token_ids" not in after + + +def test_an_unparseable_chunk_reaches_the_client_unchanged(): + """What the router cannot read, it must not rewrite.""" + st = MigrationState({"messages": []}, limit=1, path="/v1/chat/completions") + raw = b"data: {not json\n\n" + assert st.observe(raw) == raw diff --git a/tests/unit/router/test_token_ids.py b/tests/unit/router/test_token_ids.py new file mode 100644 index 00000000..7649f8f1 --- /dev/null +++ b/tests/unit/router/test_token_ids.py @@ -0,0 +1,108 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Reading token ids out of engine chunks, and refusing to guess. + +Two engines report ids in two different places, and the field has moved between +releases. What matters more than covering every shape is the failure mode: an +unrecognised chunk must produce None so the caller falls back to carrying text, +never a partial list that looks like the truth. +""" + +from __future__ import annotations + +from infera.common.worker_pool import EngineType +from infera.router.token_ids import ( + deltas_from_chunk, + prompt_from_chunk, + strip_token_ids, + supports_streaming_ids, +) + + +def test_vllm_reports_deltas_per_choice(): + chunk = {"choices": [{"delta": {"content": "hi"}, "token_ids": [15339, 1917]}]} + assert deltas_from_chunk(chunk) == [15339, 1917] + + +def test_vllm_reports_the_prompt_at_the_top_level(): + assert prompt_from_chunk({"prompt_token_ids": [1, 2, 3], "choices": []}) == [1, 2, 3] + + +def test_sglang_reports_deltas_under_its_own_extension(): + chunk = {"choices": [{"text": "hi"}], "sglext": {"completion_token_ids": [[4, 5]]}} + assert deltas_from_chunk(chunk) == [4, 5] + + +def test_sglang_reports_the_prompt_under_the_same_extension(): + assert prompt_from_chunk({"sglext": {"prompt_token_ids": [7, 8]}}) == [7, 8] + + +def test_a_chunk_without_ids_reads_as_absent(): + """The ordinary case for an engine that was never asked. Absent is not an + error -- the caller carries text instead.""" + assert deltas_from_chunk({"choices": [{"delta": {"content": "hi"}}]}) is None + assert prompt_from_chunk({"choices": [{"delta": {"content": "hi"}}]}) is None + + +def test_a_malformed_id_list_is_rejected_whole(): + """Filtering out the bad entries would leave a plausible prefix, and + resuming from a prefix drops output the client already read.""" + assert deltas_from_chunk({"choices": [{"token_ids": [1, "two", 3]}]}) is None + assert deltas_from_chunk({"choices": [{"token_ids": [1, None]}]}) is None + assert deltas_from_chunk({"choices": [{"token_ids": "12"}]}) is None + + +def test_booleans_are_not_token_ids(): + # bool is an int in Python; a JSON true here means the field is not ids. + assert deltas_from_chunk({"choices": [{"token_ids": [True, False]}]}) is None + + +def test_nonsense_input_is_survived(): + for junk in (None, [], "text", 42, {"choices": "no"}, {"choices": [None]}): + assert deltas_from_chunk(junk) is None + assert prompt_from_chunk(junk) is None + + +def test_an_empty_delta_list_is_not_absent(): + """A chunk that generated nothing is different from one that did not say -- + the first still means the engine is reporting ids.""" + assert deltas_from_chunk({"choices": [{"token_ids": []}]}) == [] + + +def test_stripping_leaves_the_response_as_the_client_expects(): + chunk = { + "choices": [{"delta": {"content": "hi"}, "token_ids": [1], "prompt_token_ids": [2]}], + "prompt_token_ids": [2, 3], + "sglext": {"completion_token_ids": [[1]]}, + } + assert strip_token_ids(chunk) is True + assert chunk == {"choices": [{"delta": {"content": "hi"}}]} + + +def test_stripping_reports_when_there_was_nothing_to_strip(): + """Lets the caller skip re-serialising a chunk it did not change.""" + chunk = {"choices": [{"delta": {"content": "hi"}}]} + assert strip_token_ids(chunk) is False + assert chunk == {"choices": [{"delta": {"content": "hi"}}]} + + +def test_sglang_is_not_asked_for_ids_on_streaming_chat(): + """It rejects the request rather than ignoring the field, so asking would + turn every migratable chat request into a 400.""" + assert supports_streaming_ids(EngineType.SGLANG, "/v1/chat/completions") is False + assert supports_streaming_ids(EngineType.SGLANG, "/v1/completions") is True + + +def test_vllm_is_asked_on_both_endpoints(): + assert supports_streaming_ids(EngineType.VLLM, "/v1/chat/completions") is True + assert supports_streaming_ids(EngineType.VLLM, "/v1/completions") is True + + +def test_an_unknown_engine_is_not_asked(): + """An engine that errors on an unknown parameter would fail requests that + work today, so silence is the safe default.""" + assert supports_streaming_ids(EngineType.ATOM, "/v1/completions") is False + assert supports_streaming_ids(None, "/v1/completions") is False