Studio: free the llama-server slot when a chat stream reaches [DONE] - #7564
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7b5017a9a6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| async for chunk in iterator: | ||
| yield chunk |
There was a problem hiding this comment.
If the ASGI send() call stalls on the [DONE] frame, or another consumer stops iterating as soon as it receives that sentinel, execution remains suspended at yield chunk and never reaches this release; the outer finally also cannot run while the response is wedged. On a one-slot backend this preserves the unbounded admission starvation the change is intended to prevent. Detect and release the safe terminal chunk before yielding it; the tool-stream loop at lines 9714–9717 has the same ordering.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in f7fa02b: the release now happens before the yield. Confirmed the ordering matters, and for a broader reason than a wedge: Starlette's stream_response suspends the body iterator across the whole of await send(...), uvicorn's send() awaits flow.drain() once the socket is write-paused, and Starlette never aclose()s a body iterator, so a client that stops reading parks the generator on the yield and everything after it is deferred to GC.
Same ordering applied to the tool-stream loop.
| Delta chunks are one JSON line ending in ``}``, so only the two terminal | ||
| emitters match: the normal end of stream and ``_openai_stream_error_sse``. | ||
| """ | ||
| return isinstance(chunk, str) and chunk.rstrip().endswith(_SSE_DONE_LINE) |
There was a problem hiding this comment.
When local chunk processing fails while the synchronous GGUF generator is still active, _openai_stream_error_sse produces a chunk ending in [DONE], so this predicate causes the outer loop to release the admission lease before the inner generator resumes into its finally, sets the cancel event, drains any worker, and closes gen. A following request can therefore enter llama-server while the failed request still occupies its slot, exceeding the configured parallelism and potentially wedging both requests. Only a successful stream-completion sentinel should trigger the early release; error sentinels must retain the existing cleanup/finally release.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in f7fa02b: the predicate is now exact equality against the bare success sentinel, so the error form from _openai_stream_error_sse (routes/inference.py:434) no longer frees the slot. Confirmed your reachability argument: the error chunk is yielded from the inner generator's except block at :10167, ahead of the finally that sets the cancel event, drains the worker and closes gen, so the sync generator can still be parked on a yield inside _open_stream's httpx client.
Your comment also led to a third case: the cancel path breaks the read loop at :10062 and falls through to stream_completed = True and the plain sentinel at :10156 with gen never driven to StopIteration, so exact equality alone was not enough either. Both loops now also gate on a flag set only when the sync generator returned on its own. Regression tests for all three are in tests/test_gguf_stream_slot_release_ordering.py.
|
@codex review |
for more information, see https://pre-commit.ci
|
@codex review |
|
Codex Review: Didn't find any major issues. Keep it up! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
@codex review |
1 similar comment
|
@codex review |
|
Codex Review: Didn't find any major issues. 🚀 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
Pushed
The dicts were only working around rebinding across the closure. Both flags are bound in Worth noting for review: omitting Tests: 355 passed across the slot-release, admission, stream-cancel and tool-passthrough suites. The one test that targeted the helper directly is removed; the behaviour it guarded is already covered at route level by |
|
@codex review |
|
Codex Review: Didn't find any major issues. 🚀 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
1 participant
The
Chat UI Testsjob has been timing out intermittently across many unrelated branches. Digging into the artifacts of one failure, it is a real Studio bug, not a Playwright flake: a chat request that has finished generating keeps its llama-server slot until the HTTP response is torn down, so the next chat request queues behind a generation that is already over. The admission queue has no default timeout (DEFAULT_ADMISSION_QUEUE_TIMEOUT_S = None), so that wait is unbounded.What the artifacts show
Correlating the Studio request log with the llama-server log from run 30326911000 (offsets line up one-to-one across all six generations):
/v1/chat/completionsstartsllama-server freed its only slot 180 seconds before Studio let go of it. Turn 7 never reached llama-server at all: it sat in
_openai_admission_wait_stream_chunksemitting keepalives until the browser was torn down. The event loop served eight other requests in between, so nothing was blocked, andn_slots = 1in that job.Turn 6 did reach
[DONE]. The screenshot shows its answer rendered and the context counter at 117, and that counter is written from the SSE usage chunk after the streaming loop exits, which only happens on[DONE]or EOF. The body stayed open for 180 s, so it was[DONE].Why nothing recovered
[DONE](chat-api.ts), so the request stays open until navigation._SameTaskStreamingResponseoverrides__call__and detects disconnects only bysend()raisingOSError. uvicorn advertises ASGIspec_version 2.3and itssend()returns silently when disconnected, and Starlette gates thatOSErrorpath onspec_version >= (2, 4). So under uvicorn the branch cannot fire, and the stocklisten_for_disconnecttask group that would have fired was replaced._DEFAULT_STREAM_STALL_TIMEOUT_Sbounds gaps between upstream llama-server reads and stops applying once[DONE]is parsed;_openai_compat_stream_stall_timeout()covers only the passthrough proxy;_LOCAL_TOOL_STREAM_STALL_KEEPALIVE_Sis a cadence, not a cap. Hence 180 s (the Playwright budget), never 120 s.Changes
1. Release the admission lease when the decode is genuinely over. The lease was released only in the stream's outer
finally, which does not run until the ASGI body is torn down. Three conditions now gate the early release, all three needed:stream_responsesuspends the body iterator across the whole ofawait send(...), uvicorn'ssend()awaitsflow.drain()once the socket is write-paused, and Starlette neveraclose()s a body iterator. So a client that stops reading parks the generator on theyieldand anything after it is deferred to GC._openai_stream_error_sse(:434) also ends in[DONE], but it is yielded from the inner generator'sexceptblock ahead of thefinallythat sets the cancel event, drains the worker and closesgen, so the sync generator can still be parked inside_open_stream's httpx client.:10062and falls through tostream_completed = Trueand the plain sentinel at:10156without ever drivinggentoStopIteration, and thefinallythen skipsgen.close(). Pressing Stop is ordinary traffic, so the sentinel alone does not prove llama-server is finished.release()is idempotent, so the existingfinallystays a correct backstop. Applied to bothadmitted_gguf_stream_chunksandadmitted_gguf_tool_stream, which share the queue.2. Bound the disconnect-watcher stop.
_stop_local_disconnect_cancel_watcherranawait watcherwith no timeout, inside the stream'sfinally.asyncio.waitnever cancels its argument and never re-raises, so the watcher can be abandoned after a timeout; it is a poll loop holding no resources.This is the amplification, and it is what makes the failure user-visible. Anyone on a 1-slot backend who sends a second message after a response wedges gets "Generating..." forever with no error, and two further consequences of the response never terminating: the
ActiveGenerationtracker stays open, andLlamaKeepWarmMiddleware._finish()never fires, so idle auto-unload stops working.I have not found what wedges the teardown in the first place. Candidates in that
finallyare narrow, but I could not reproduce it: 7 clean end-to-end runs of the CI Playwright test locally (including with the stack pinned to 2 cores), 7 sequential and 7 held-open HTTP turns, and several thousand synthetic uvicorn requests all came back clean. It needs the 4-vCPU runner under CPU inference load. Decoupling the slot means the wedge no longer takes the next request down with it.Tests
tests/test_gguf_stream_slot_release.pydrives the real ASGI route with_stop_local_disconnect_cancel_watchermonkeypatched to hang, putting the response in exactly the state the CI job was left in, and areceivethat never disconnects, matching the browser. It asserts the slot is free once the[DONE]frame has been sent and that a second caller is admitted immediately.tests/test_gguf_stream_slot_release_ordering.pycovers the three gating conditions, each against the real route:test_slot_is_free_before_the_done_frame_reaches_sendtest_error_sentinel_keeps_the_slot_until_the_generator_is_closedtest_cancelled_stream_keeps_the_slot_until_the_generator_is_closedAll four route-level tests fail without the corresponding source change and pass with it, verified by reverting.
test_gguf_stream_slot_release.py+..._ordering.py+ admission, stream-cancel and tool-passthrough suites: 344 passedmainhere (missing optional deps and network cases).