[Frontend] Abort engine request on client disconnect (free leaked KV)#1562
Open
yhl-amd wants to merge 2 commits into
Open
[Frontend] Abort engine request on client disconnect (free leaked KV)#1562yhl-amd wants to merge 2 commits into
yhl-amd wants to merge 2 commits into
Conversation
…st leak The server never reacted to a client hanging up mid-request, and the non-streaming path never released finished requests. Under retrying clients (opencode / Claude Code) this piled up "pending requests" without bound and wasted GPU on output nobody reads. Two independent bugs, both verified with controlled curl tests: 1. No disconnect cancellation. Neither the streaming generators nor the non-stream handlers aborted the engine sequence when the client went away, so an abandoned request ran to max_tokens. Add an engine abort path (Sequence.aborted -> Scheduler reuses the normal finish/deallocate stop path; EngineUtilityHandler.abort_request; CoreManager.abort_request) and wire it into the API layer. Streaming generators (stream_chat_response / _fanout) now run their body under try/finally so cleanup_streaming_request (which aborts the seq) fires on GeneratorExit. For non-stream requests Starlette does NOT cancel the handler on disconnect, so add `_run_nonstream_with_disconnect()` which races the generation against `raw_request.is_disconnected()` and cancels on disconnect. `chat_completions` and `completions` gain a `raw_request` parameter for this. 2. Non-stream request leak. generate_async / _multimodal / _fanout only removed the request from io_processor.requests on a setup exception, never on normal completion, so every completed non-stream request leaked a Sequence (pending requests grew forever). Wrap their token loops in try/finally that always pops the request (and aborts on early exit). Streaming already pops via cleanup_streaming_request. Also add `"choices": []` to the streaming usage chunk so AI-SDK clients (opencode) stop rejecting the final usage frame. Verified (DeepSeek-V4, TP8): 6 sequential non-stream completions keep pending flat (was 1->8); 3 mid-generation disconnects -> abort found, GPU drops immediately, pending flat (was +3 leak, GPU pinned to max_tokens); normal streaming + non-stream answers unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
3 tasks
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
yhl-amd
force-pushed
the
feat/api-abort-on-disconnect
branch
from
77348b7 to
40c6c03
Compare
…ct-abort # Conflicts: # atom/entrypoints/openai/api_server.py
yhl-amd
force-pushed
the
feat/api-abort-on-disconnect
branch
from
40c6c03 to
7a0492f
Compare
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
1 participant
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.
Summary
Non-streaming API handlers were not cancelled when the client hung up (Starlette only cancels
StreamingResponse, not plain handlers), so the engine kept generating and the sequence's KV blocks leaked until it hitmax_tokens. This adds a client-disconnect abort path.Changes
EngineCoreMgr.abort_requestbroadcasts anabort_requestutility command;_handle_abort_requestmarks the seqaborted; the scheduler finishes it at the next step via the normal stop path (frees KV, emits a finishedRequestOutput). AddsSequence.aborted._run_nonstream_with_disconnectrunsgenerate_asyncin a task and pollsrequest.is_disconnected(); on disconnect it cancels the task, whose teardown aborts + pops the request. Wired into/v1/chat/completionsand/v1/completions(return HTTP 499).generate_async/generate_async_multimodal/generate_async_fanout/cleanup_streaming_requestabort + pop on early exit.Test plan
🤖 Generated with Claude Code