Skip to content

[Bugfix] Cancel inference on client disconnect + fix non-stream request leak#1441

Open
yhl-amd wants to merge 2 commits into
ROCm:mainfrom
yhl-amd:fix/client-disconnect-abort
Open

[Bugfix] Cancel inference on client disconnect + fix non-stream request leak#1441
yhl-amd wants to merge 2 commits into
ROCm:mainfrom
yhl-amd:fix/client-disconnect-abort

Conversation

@yhl-amd

@yhl-amd yhl-amd commented Jul 2, 2026

Copy link
Copy Markdown
Contributor

Summary

Two independent bugs in the OpenAI server, both surfaced by retrying clients (opencode / Claude Code) which pile up unbounded pending requests and waste GPU on output nobody reads.

1. No disconnect cancellation. Neither the streaming generators nor the non-stream handlers aborted the engine sequence when the client hung up, 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 (broadcast to all DP ranks).
  • Streaming generators (stream_chat_response / _fanout) now run under try/finally so cleanup_streaming_request (which aborts the seq) fires on GeneratorExit.
  • Starlette does not cancel the handler on disconnect for non-stream, so add _run_nonstream_with_disconnect() which races generation against raw_request.is_disconnected() and cancels on disconnect. chat_completions / completions gain a raw_request param.

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 → every completed non-stream request leaked a Sequence (pending grew forever). Wrap the token loops in try/finally that always pops (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.

Why

Under an agentic client that retries / cancels mid-generation, the pre-fix server leaks a Sequence per completed non-stream request and keeps abandoned requests running to max_tokens. pending requests grows without bound and GPU stays pinned on discarded work.

Test plan

Verified on DeepSeek-V4, TP8:

  • 6 sequential non-stream completions → pending requests stays flat (was 1→8).
  • 3 mid-generation client disconnects → abort_request ... found=True, GPU utilization drops immediately, pending flat (was +3 leak with GPU pinned to max_tokens).
  • Normal streaming + non-stream answers unchanged.

Reproduce:

PORT=9700; M=<served-model-id>
# (1) non-stream leak: fire 6, pending must stay flat
for i in $(seq 6); do
  curl -s localhost:$PORT/v1/chat/completions -H 'Content-Type: application/json' \
    -d '{"model":"'"$M"'","messages":[{"role":"user","content":"hi"}],"max_tokens":8}' >/dev/null
done
# (2) disconnect cancel: start long gen, drop after 2s -> log shows abort found, GPU drops
timeout 2 curl -s localhost:$PORT/v1/chat/completions -H 'Content-Type: application/json' \
  -d '{"model":"'"$M"'","messages":[{"role":"user","content":"write a very long essay"}],"max_tokens":4000}' >/dev/null
# (3) regression: normal stream + non-stream still correct
curl -s localhost:$PORT/v1/chat/completions -H 'Content-Type: application/json' \
  -d '{"model":"'"$M"'","messages":[{"role":"user","content":"2+2?"}],"max_tokens":10}'

Split out of #1428 (which now carries only the DeepSeek-V4 DSML tool-call changes) so this fix can be reviewed on its own.

🤖 Generated with Claude Code

@yhl-amd

yhl-amd commented Jul 2, 2026

Copy link
Copy Markdown
Contributor Author

Split out of #1428 so the client-disconnect / non-stream-leak fix can be reviewed on its own. #1428 now carries only the DeepSeek-V4 DSML tool-call changes.

…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>
@yhl-amd
yhl-amd force-pushed the fix/client-disconnect-abort branch from 5d00023 to 40c6c03 Compare July 12, 2026 15:47
Address review: fold the parallel `Sequence.aborted` boolean into the
existing `SequenceStatus` enum so a sequence's terminal intent has a
single source of truth.

- sequence: add SequenceStatus.ABORTED (kept distinct from FINISHED so an
  aborted running seq still rides one cleanup pass; is_finished() stays
  False until then); drop Sequence.aborted.
- engine_utility: _handle_abort_request sets seq.status = ABORTED.
- scheduler (running): finish check reads seq.status == ABORTED.
- scheduler (waiting): intercept ABORTED seqs when popped from `waiting`,
  BEFORE the waiting->running promotion overwrites status with RUNNING and
  loses the abort intent. Such seqs hold no KV and need no forward pass, so
  finish them outright via `_rejected` (mirrors the unschedulable exit).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant