feat(router): continue a streaming generation on another worker when its own goes away - #119
Open
weilei0120 wants to merge 5 commits into
Open
feat(router): continue a streaming generation on another worker when its own goes away#119weilei0120 wants to merge 5 commits into
weilei0120 wants to merge 5 commits into
Conversation
… streams Continuing a generation on another worker means handing it what was produced so far. Decoded text works, but re-encoding it is not guaranteed to reproduce the ids the model sampled -- tokenizers are not injective in that direction, so a boundary can shift and the next worker resumes from a slightly different sequence. The engines can report the ids themselves, which removes that step. Neither agrees on where: vLLM puts the deltas on the choice and the prompt at the top level, SGLang uses its own `sglext` object, and the field has moved between releases. Both shapes are accepted wherever they appear, because a router pinned to one layout would silently stop being exact. Silently is the problem, so every reader returns None rather than a guess and a malformed id list is rejected whole rather than filtered: a partially-parsed sequence is a plausible-looking prefix of the truth, and resuming from one drops output the client already read. SGLang rejects -- not ignores -- a streaming chat request that asks for ids, so that combination is never asked. ATOM is left out entirely for the same class of reason: an engine that errors on an unknown parameter would fail requests that work today. Signed-off-by: leiwei12 <lei.wei@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
…finish A generation still running when the drain window closes was cancelled, which the client read as a failure. Where the router can continue it elsewhere, the worker now returns it instead, and the client reads one uninterrupted stream. 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, so a request that would have finished on its own inside the window is simply allowed to; `--drain-timeout` becomes how long a generation is worth waiting for, and migration decides what happens to the ones that exceed it. A timeout of 0 asks to leave immediately and is honoured -- nothing is waited for, but what can be resumed is still handed back, since doing so is a local publish that delays nothing. The worker cannot decide this alone. Severing a stream early is only safe because the router promised to continue it, so the request carries `migratable` and only marked ones are shed. With migration disabled nothing is marked and the drain behaves exactly as before, which is what makes this safe to land. The notice precedes the cancellation rather than replacing it: cancelling also sends an error, but a generic one, and the router would read a planned handover as a worker that broke. It also suppresses the second frame, so one event does not produce two contradictory errors, and the bookkeeping is cleared in the proxy's finally block so a finished request leaves nothing behind. Signed-off-by: leiwei12 <lei.wei@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
A worker that goes away mid-stream -- crashed, evicted, drained -- used to end the response with an error, because the tokens already sent cannot be un-sent and a retry would produce a different answer the client would read as a repetition or a contradiction. What moves instead is the generation so far, put back in front of the model with the budget reduced by what it cost, so the next worker finishes the sentence rather than starting a new one. Off by default; `--migration-limit 1` enables it. It changes what a worker is asked to produce, and the continuation is not byte-identical to what the original would have written -- sampling state does not survive the move -- so a caller needing reproducible output for a fixed seed should leave it off. Carried exactly wherever the engine reports its token ids, since re-encoding decoded text can shift a boundary. Chat has no pre-tokenized entry, so an exact chat continuation is reissued against completions and the replies are reshaped back; that conversion is skipped, and the text form used, when the request carries tools or a response format, because a completions request cannot honour them and losing a capability halfway through an answer is worse than a boundary that might differ. The ids are asked for only when a migration could use them and are removed before forwarding, so enabling this never changes the shape of what a client receives. Everything that cannot be faithfully rebuilt stops being migratable rather than being approximated: a chunk that is not JSON, ids that cover only part of the output, and tool calls or hidden reasoning, which arrive outside `content` and would be replayed whole after the client already holds half of one. Those streams end visibly, as they did before this change. What is accumulated is what the client actually received rather than what a worker claims to have sent -- after a migration the two must agree, or there is a seam -- and the migration itself is counted by why it happened, so a planned handover reads differently from a fault to whoever is watching. Signed-off-by: leiwei12 <lei.wei@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
…viour Migration answers a question this page already raises — what happens to a generation still running when the drain window closes — so it is a section here rather than a page of its own, next to the timeout that governs it. It also covers the case this page otherwise cannot: a worker that leaves without notice. Kept to what changes a reader's decisions: that it exists, that it is off by default, what turns it on, that the continuation is not byte-identical, which requests it applies to, and how to see it working. The rest of the page is trimmed on the same principle. It used to explain which mark the router reads and when, what the preStop delay is and is not spent on, and the order in which a worker deregisters and drains — none of which changes what a reader configures or expects, and all of which dates quickly. The two sections on paths without advance notice documented one thing from two angles and are now one. Signed-off-by: leiwei12 <lei.wei@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
weilei0120
requested review from
JohnQinAMD,
jiejingzhangamd,
limou102 and
xiaobochen-amd
as code owners
August 14, 2026 06:36
Contributor
There was a problem hiding this comment.
Pull request overview
Adds opt-in request migration for NATS-streamed generations so a client can keep reading uninterrupted when a worker disappears mid-stream (crash/eviction) or is draining (graceful shutdown), with best-effort exact prefix preservation via engine-reported token IDs.
Changes:
- Introduces migration state tracking + continuation dispatch in the mixed router, including optional token-id requesting/stripping and chat-shape reshaping when resuming via
/v1/completions. - Extends the NATS request protocol and worker drain behavior to “hand back” migratable in-flight streams after the drain window, plus new migration success/failure metrics.
- Adds unit tests and documentation for migration behavior, constraints, and observability; wires
--migration-limitthrough server/router construction.
Reviewed changes
Copilot reviewed 17 out of 17 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/router/test_token_ids.py | Adds coverage for token-id extraction/stripping and engine endpoint gating. |
| tests/unit/router/test_migration.py | Adds MigrationState accounting and continuation-shaping tests. |
| tests/unit/router/test_migration_dispatch.py | Adds end-to-end dispatch-path migration tests (broken stream + draining handover). |
| tests/unit/router/test_failover.py | Updates assertion for revised mid-stream error text. |
| tests/unit/common/test_drain_handback.py | Adds tests for worker drain handback of migratable requests. |
| manual/features/graceful_shutdown.md | Documents request migration behavior, scope, and metrics. |
| manual/features/feature_matrix.md | Adds request migration to the feature matrix with constraints note. |
| infera/server/metrics.py | Adds infera_migrations_total and infera_migrations_failed_total. |
| infera/server/args.py | Adds --migration-limit server flag with env override. |
| infera/server/main.py | Wires migration_limit into router construction and logs enablement. |
| infera/router/token_ids.py | Implements token-id parsing/stripping and “safe to ask” policy. |
| infera/router/mixed.py | Implements streaming migration + continuation dispatch + migratable NATS flagging. |
| infera/router/migration.py | Implements MigrationState and completions→chat reshaping helper. |
| infera/router/auto.py | Passes migration_limit through to MixedRouter (PD excluded). |
| infera/engine/vllm/args.py | Updates drain-timeout help to mention migration handback. |
| infera/engine/sglang/args.py | Updates drain-timeout help to mention migration handback. |
| infera/common/nats_request.py | Adds migratable request field + drain handback mechanism and notice constant. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+271
to
+272
| def can_migrate(self) -> bool: | ||
| return not self.poisoned and self.migrations_left > 0 |
Each 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. A `prompt` given as a token array -- which the engines accept verbatim, and which this branch verified they do -- was formatted into the carried text with an f-string, yielding its Python repr. The next worker was asked to continue the literal characters "[1, 2, 3]hello". Reachable wherever an engine does not report ids, which is every SGLang build that ignores `return_token_ids`. A batch of prompts had the same defect, and both are now refused rather than formatted. A request for several completions accumulated only the first choice, so every other one would have resumed from a prefix belonging to it -- content in the wrong place rather than content missing. `n`, `best_of` and prompt batches are rejected from the outset, and an engine returning several choices unasked is caught as it arrives. Exact ids do not rescue this: they are accumulated per choice too. Giving up on migrating a request also stopped the token ids being taken back out, so an engine reporting them on only some chunks could leak them to a caller that never asked -- against the guarantee that enabling migration does not change the shape of a response. Stripping no longer depends on whether the request is still migratable; only the accumulating does. Signed-off-by: leiwei12 <lei.wei@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this solves
A worker that goes away mid-stream — crashed, evicted, or rolled — used to end
the response with an error. That failure cannot be handled by retrying: the
client has already read half an answer, and re-running the request produces a
different one, which reads as a repetition or a contradiction.
This carries the generation to another worker instead, so the client reads one
uninterrupted stream and never learns a worker changed underneath it.
Off by default.
--migration-limit 1enables it.How it behaves
Two triggers, one mechanism. A worker that crashes breaks the stream and the
router picks it up; a worker being drained hands back what it could not
finish inside
--drain-timeoutrather than cutting it. They differ only in themetric label, so a rollout doing its job reads differently from a worker dying.
The prefix is carried exactly wherever the engine reports the token ids it
sampled, since re-encoding decoded text can shift a boundary. Where it does not,
the text is carried instead — approximate, but far better than a severed stream,
and the fallback is automatic.
Three decisions worth reviewing
The worker cannot shed a stream on its own. Cutting one early is only safe
because the router promised to continue it, so the router marks each request as
resumable and a worker only sheds marked ones. With migration disabled nothing
is marked and the drain behaves exactly as before — that is what makes this safe
to land in a deployment that has not opted in.
Migration happens after the drain window, not instead of it. Moving a
generation costs the next worker a re-read of everything produced so far, which
is not worth paying for a request that had a second left to run.
--drain-timeoutkeeps its meaning — how long a generation is worth waiting for — and migration
only decides what happens to the ones that exceed it. Pods exit at the same
moment they did before.
Anything that cannot be faithfully rebuilt stops being migratable rather
than being approximated: a chunk that is not JSON, ids covering only part of the
output, and tool calls or reasoning content, which arrive outside the text the
client receives and would otherwise be replayed after it already holds half of
one. Those streams end visibly, as they did before this change.
Not covered
PD-disaggregated requests (both legs and the transfer between them would have to
move), the HTTP transport (the router never sees a frame boundary), and
non-streaming requests, which
--request-max-retriesalready handles better byre-running them cleanly.
The continuation is not byte-identical to what the original worker would
have produced — sampling state does not move — and exact ids do not change that;
they make the prefix faithful, not the result reproducible. Documented next to
the flag.
Test plan
path; 1402 unit tests pass overall. The 6 remaining failures are a
pre-existing baseline (sglang/gaie not installed), confirmed unchanged by
running them against a stashed tree.
with a generation interrupted the way a rolling upgrade interrupts one.
deleted the Pod serving a live stream. Log reads
migrated a live generation ... after 1083 token(s) (exact token ids); the client received3000 chunks, zero error frames, one 0.05s gap. Verified the ids never
reach the client and that
promptas a token array is accepted verbatim.return_token_idsentirely — silently, on every endpoint — so it exercisedthe text fallback. Same outcome for the client:
migrated ... after 1128 token(s) (carried text), zero error frames, output complete and in order.The SGLang result is worth flagging: the fallback is the path that build will
always take, and it works, but "exact" is a vLLM-only benefit until SGLang ships
return_token_idson streaming endpoints.Made with Cursor