Repository navigation
fix: stop a non-streaming generation when the client disconnects - #250
Conversation
A non-streaming request whose client had gone away kept generating to max_tokens or EOS, holding its slot the whole time. With the default --parallel 1 every later request queued behind it, and a client that times out and retries (the OpenAI SDKs' default) left one full generation per attempt. The streaming path did not have this problem: its next write fails and tears down the body stream, which cancels the consumer. A non-streaming response writes nothing until the end, so nothing told it the client had left. Hummingbird runs the handler inline and enables remote half-closure, so a closed socket shows up on the pipeline as ChannelEvent.inputClosed (a reset as channelInactive), and nothing in the server watched for either. - ConnectionRequestContext keeps the connection's channel; the two POST routes pass it to their handlers. - ClientDisconnectWatch sits at the head of the pipeline for the span of one non-streaming request and fires once on inputClosed or channelInactive, forwarding both. It is removed when the handler returns, so a keep-alive connection closing between requests is not taken for an abandoned one. - cancellingOnClientDisconnect runs the non-streaming generation in a task the watch cancels. Cancellation reaches the producer through the AsyncStream as it does for SSE; the handler falls out of its loop and releases the slot through the existing path. Both the chat and the /v1/completions non-streaming branches use it. Streaming is unchanged, and the handlers still run without a channel (as the unit tests call them). Tests: ClientDisconnectTests (EmbeddedChannel) cover the watch and the wrapper; contract Test 14 abandons a non-streaming request and checks the slot is released with only a few tokens generated.
|
Thanks for #250. One related gap that fits this PR (it's the tool-call-only save issue from #245): In both Suggested fix: drop the case .info(let info):
if let onPrefillDone { await onPrefillDone() }
A test that would cover it: with Qwen2.5-0.5B-Instruct-4bit and a Note: I checked this by reading the code only; I have not run it. I have not checked how Happy to open this as a follow-up PR on top of #250 once it merges, or you can fold it in. Whichever you prefer. |
Both chat handlers save the prompt cache when the producer's `.info` arrives, and only did so once a text chunk had been seen. A turn the model answers with nothing but a tool call never yields a chunk: the parser holds the whole `<tool_call>` body and emits one `.toolCall`. So the save was skipped and the next turn of the agent loop (tool result, follow-up) re-prefilled the whole conversation. This predates the move of the save to `.info`: before that it sat inside the first-chunk block, which such a turn never entered either. Drop the first-chunk guard at `.info` in both handlers. `.info` is the producer's last yield, so prefill has long finished and the cache is no longer being written; the save slices the KV state to the prompt length, as it does after a text turn. Early exits (stop sequence, disconnect) still never reach `.info` and still skip the save. Tests: ToolCallPromptCacheTests drive both handlers with a synthetic stream (a tool-call-only turn saves once, text turns still save once, a stream cut before `.info` does not save); contract Test 15 sends a tool-call turn, non-streaming and streaming, and checks that the follow-up's prompt-cache hit reuses the whole prompt of that turn. Raised in SharpAI#245; suggested for this PR by @solderzzc on SharpAI#250.
|
Folded in as a second commit.
PR description updated to cover both. |
Closes #249.
A non-streaming request whose client disconnects kept generating to
max_tokensor EOS, holding its slot the whole time. With the default--parallel 1every later request queued behind it, and a client that times out and retries (the OpenAI SDKs' default) stacked up one full generation per attempt. The streaming path did not have this problem: its next write fails and tears down the body stream, which cancels the consumer. A non-streaming response writes nothing until the end, so nothing ever told it the client had left.What changes
ConnectionRequestContext);BasicRequestContextdrops it. The two POST routes pass the channel to their handlers.ClientDisconnectWatch, a small pipeline handler installed at the head of the pipeline for the span of one non-streaming request, fires once when the client's side closes. Hummingbird enables remote half-closure, so a closed socket arrives asChannelEvent.inputClosed; a reset arrives aschannelInactive. Both are forwarded unchanged. The watch is removed when the handler returns, so a keep-alive connection closing between requests is not mistaken for an abandoned one.cancellingOnClientDisconnectruns the non-streaming generation (prefill plus the.chunk/.infoloop) in a task that the watch cancels. Cancellation flows through theAsyncStreamto the producer the same way it does for SSE: the model stops after the token in flight, the handler falls out of its loop and releases the slot through the existing path, and the response it builds goes to the closed socket. Both the chat and the/v1/completionsnon-streaming branches use it. Nothing changes for streaming requests, and the handlers keep working without a channel (the unit tests call them that way).Like the streaming path, an abandoned request does not reach
.info, so it does not save the prompt cache. (The second commit below changes only what happens once.infois reached.)Tests
tests/SwiftLMTests/ClientDisconnectTests.swift: the watch fires once oninputClosed, once onchannelInactive, ignores other events and forwards what it sees (EmbeddedChannel); a non-streaming handler fed an endless synthetic stream returns, releases its slot and leavesrequests_activeat 0 after the watch fires; a normal completion removes the watch; the wrapper runs plainly without a channel.tests/test-contract.shTest 14: a non-streaming request curl abandons after half a second must leaverequests_activeat 0 within 5 s and add fewer than 400 tokens toswiftlm_tokens_generated_total. On main it adds the whole answer.Prefill is not interruptible, so a client that leaves during a long prefill still holds the slot until the first decode step; that is the same as on the streaming path.
What I ran (M1 Max, Xcode 27)
swift test --filter SwiftLMTests: 337 passed, 0 failed (main: 326, plus the 11 new across both commits).tests/test-contract.shwith the CI model: 24 passed, 0 failed, 2 skipped (the usual capability skips). Test 14 fails on main, where the abandoned request generated 737 tokens, and passes on the branch with 132. Test 15 fails without the second commit and passes with it (numbers below)./v1/completionsrequest 0.03 s, and the OpenAI SDK timeout-and-retry case 17.5 s → 0.0 s, with each attempt stopping after the tokens generated before the client left instead of running tomax_tokens.Second commit: a tool-call-only turn saves the prompt cache (#245)
Both chat handlers save the prompt cache when the producer's
.infoarrives, and only did so once a text chunk had been seen. A turn the model answers with nothing but a tool call never yields a chunk, so the save was skipped and the next turn of the agent loop re-prefilled the whole conversation. The!firstTokenguard at.infois dropped in both handlers:.infois the producer's last yield, so prefill has finished and the cache is no longer being written, and the save slices the KV state to the prompt length as it does after a text turn. Early exits (stop sequence, disconnect) still skip the save.tests/SwiftLMTests/ToolCallPromptCacheTests.swift: a tool-call-only turn saves once on both handlers, text turns still save once, a stream cut before.infodoes not save.tests/test-contract.shTest 15: a tool-call turn (non-streaming, then streaming) followed by its tool result; the follow-up's prompt-cache hit must reuse at least the whole prompt of the tool-call turn. Before the change the follow-up reused 4 and 7 cached tokens of prompts of 176 and 177; after it, all 176 and 177.