fix(server): guarantee AbortMsg delivery on stream cancellation - #222
Open
Artemowka22 wants to merge 1 commit into
Open
fix(server): guarantee AbortMsg delivery on stream cancellation#222Artemowka22 wants to merge 1 commit into
Artemowka22 wants to merge 1 commit into
Conversation
stream_with_cancellation reacted to a client disconnect with an unowned asyncio.create_task(abort_user(uid)): request teardown could outrun delivery and a failure inside the task degraded to a never-retrieved-exception warning, so the scheduler kept decoding for a client that was gone. Await the abort inline behind asyncio.shield (a second cancellation cannot kill the delivery task), and make abort_user claim the uid first — exactly one AbortMsg even if cancellation runs twice, and none at all when the stream already finished normally. Found via the freetoken-mlx downstream audit (docs/AUDIT.md, defect 2).
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
FrontendManager.stream_with_cancellationreacts to a client disconnect with a bareasyncio.create_task(self.abort_user(uid))and re-raises. The task has no owner: request teardown can complete before the abort coroutine ever runs, and an exception inside it degrades to a "Task exception was never retrieved" warning. The scheduler then keeps decoding for a client that is gone — under a longmax_tokensthis wastes the GPU for minutes and pins KV pages.Fix (ported from a downstream audit — agisota/freetoken-mlx, docs/AUDIT.md, defect 2):
abort_useridempotent — exactly oneAbortMsgper uid even if the handler runs twice (disconnect + server-side cancel).Test plan
pytest tests/server/test_stream_cancellation.py(new): cancellation delivers exactly one AbortMsg before teardown finishes; double cancellation stays single-shot; normal completion sends no aborttests/server/suite passes