diff --git a/Sources/DFlash/DFlashRuntime.swift b/Sources/DFlash/DFlashRuntime.swift index 3a00c8f..08fd94f 100644 --- a/Sources/DFlash/DFlashRuntime.swift +++ b/Sources/DFlash/DFlashRuntime.swift @@ -465,6 +465,10 @@ public enum DFlashRuntime { var prefetchedBlockLen: Int? while generatedTokenIDs.count < maxNewTokens { + // The consumer may be gone (stop sequence hit, client disconnected): the + // wrapper task is cancelled, and this loop must stop computing, not just + // stop yielding (#224). + if Task.isCancelled { break } let remaining = maxNewTokens - generatedTokenIDs.count let blockLen = max(1, min(effectiveBlockTokens, remaining)) diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index a4bb8da..9aeef3c 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -2321,6 +2321,73 @@ func collectBody(_ request: Request) async throws -> Data { return Data(bodyBytes) } +/// The end-of-sequence set the standard path stops on (`buildStopTokenIds` in +/// mlx-swift-lm): the configuration's ids, the tokenizer's EOS and the extra EOS tokens +/// resolved through the tokenizer. DFlash takes it as a list. +func stopTokenIDs( + eosTokenIDs: Set, tokenizerEOS: Int?, extraEOSTokens: Set, + tokenID: (String) -> Int? +) -> [Int] { + var ids = eosTokenIDs + if let tokenizerEOS { ids.insert(tokenizerEOS) } + for token in extraEOSTokens { + if let id = tokenID(token) { ids.insert(id) } + } + return ids.sorted() +} + +/// Mutable detokenizer shared with the bridge task in `dflashGenerationStream`. +private final class DetokenizerBox: @unchecked Sendable { + var detokenizer: NaiveStreamingDetokenizer + init(_ d: NaiveStreamingDetokenizer) { self.detokenizer = d } +} + +/// Bridges a DFlash event stream to the `Generation` stream the handlers consume. +/// +/// When the consumer stops early — a text-level stop sequence, a client disconnect — the +/// returned stream is torn down, which cancels the bridge task; that ends the DFlash +/// stream and, through its own `onTermination`, the DFlash generation task. Before this +/// the bridge was a bare `Task` nobody cancelled, so the DFlash loop ran on to +/// `max_tokens` after the response had ended and the next request's generation +/// overlapped it (#224). +func dflashGenerationStream( + _ events: AsyncStream, + tokenLimit: Int, + nextText: @escaping @Sendable (Int) -> String? +) -> AsyncStream { + AsyncStream { continuation in + let bridge = Task { + for await event in events { + if Task.isCancelled { break } + switch event { + case .token(let tokenID, _, _, _): + if let chunk = nextText(tokenID) { + continuation.yield(.chunk(chunk, tokenId: tokenID)) + } + case .prefill, .prefillProgress: + break + case .summary(let summary): + print("[SwiftLM] DFlash summary: \(summary.generationTokens) tokens, \(String(format: "%.1f", summary.tokensPerSecond)) tok/s, acceptance=\(String(format: "%.1f%%", summary.acceptanceRatio * 100)), \(summary.cyclesCompleted) cycles") + // The SSE/non-streaming handlers emit finish_reason, usage and + // `[DONE]` from `.info`. Without it, a DFlash run that ends on EOS + // or max_tokens (rather than a textual stop sequence) closes the + // stream with no `[DONE]` sentinel. + let prefillSec = summary.phaseTimingsUs.prefill / 1_000_000.0 + continuation.yield(.info(GenerateCompletionInfo( + promptTokenCount: summary.promptTokenCount, + generationTokenCount: summary.generationTokens, + promptTime: prefillSec, + generationTime: summary.elapsedUs / 1_000_000.0 - prefillSec, + stopReason: summary.generationTokens >= tokenLimit ? .length : .stop + ))) + } + } + continuation.finish() + } + continuation.onTermination = { _ in bridge.cancel() } + } +} + // ── Chat Completions Handler ───────────────────────────────────────────────── func handleChatCompletion( @@ -2522,50 +2589,25 @@ func handleChatCompletion( fflush(stdout) // Convert DFlashEvent stream to Generation stream with proper streaming detokenizer let dflashTokenizer = await container.tokenizer + let dflashConfiguration = await container.configuration let dflashStream = DFlashRuntime.generate( targetModel: targetModel, draftModel: dflashDraft, promptTokens: promptTokens, maxNewTokens: tokenLimit, - blockTokens: dflashBlockSize + blockTokens: dflashBlockSize, + // The end-of-sequence set the standard path stops on. Without it DFlash ran + // to max_tokens past `<|im_end|>` (#224). + stopTokenIDs: stopTokenIDs( + eosTokenIDs: dflashConfiguration.eosTokenIds, + tokenizerEOS: dflashTokenizer.eosTokenId, + extraEOSTokens: dflashConfiguration.extraEOSTokens, + tokenID: dflashTokenizer.convertTokenToId) ) - - // Use a class wrapper so the detokenizer can be mutated inside the closure - final class DetokenizerBox: @unchecked Sendable { - var detokenizer: NaiveStreamingDetokenizer - init(_ d: NaiveStreamingDetokenizer) { self.detokenizer = d } - } let box = DetokenizerBox(NaiveStreamingDetokenizer(tokenizer: dflashTokenizer)) - - let genStream = AsyncStream { continuation in - Task { - for await event in dflashStream { - switch event { - case .token(let tokenID, _, _, _): - box.detokenizer.append(token: tokenID) - if let chunk = box.detokenizer.next() { - continuation.yield(.chunk(chunk, tokenId: tokenID)) - } - case .prefill, .prefillProgress: - break - case .summary(let summary): - print("[SwiftLM] DFlash summary: \(summary.generationTokens) tokens, \(String(format: "%.1f", summary.tokensPerSecond)) tok/s, acceptance=\(String(format: "%.1f%%", summary.acceptanceRatio * 100)), \(summary.cyclesCompleted) cycles") - // The SSE/non-streaming handlers emit finish_reason, usage and - // `[DONE]` from `.info`. Without it, a DFlash run that ends on EOS - // or max_tokens (rather than a textual stop sequence) closes the - // stream with no `[DONE]` sentinel. - let prefillSec = summary.phaseTimingsUs.prefill / 1_000_000.0 - continuation.yield(.info(GenerateCompletionInfo( - promptTokenCount: summary.promptTokenCount, - generationTokenCount: summary.generationTokens, - promptTime: prefillSec, - generationTime: summary.elapsedUs / 1_000_000.0 - prefillSec, - stopReason: summary.generationTokens >= tokenLimit ? .length : .stop - ))) - } - } - continuation.finish() - } + let genStream = dflashGenerationStream(dflashStream, tokenLimit: tokenLimit) { tokenID in + box.detokenizer.append(token: tokenID) + return box.detokenizer.next() } return (genStream, nil) diff --git a/tests/SwiftLMTests/DFlashStreamBridgeTests.swift b/tests/SwiftLMTests/DFlashStreamBridgeTests.swift new file mode 100644 index 0000000..2f6daf0 --- /dev/null +++ b/tests/SwiftLMTests/DFlashStreamBridgeTests.swift @@ -0,0 +1,106 @@ +import XCTest +import Foundation +@testable import DFlash +import MLXLMCommon +@testable import SwiftLM + +// MARK: - The DFlash → Generation bridge (issue #224) +// +// The bridge used to be a bare `Task` that nobody cancelled. When the SSE handler ended +// a response early (a text-level stop sequence, a client disconnect) the handler released +// its slot, but the bridge kept draining the DFlash stream and the DFlash loop kept +// generating up to `max_tokens`; the next request then ran a second generation on the +// same model and MLX aborted. These pin the two halves of the fix that need no model: +// the stop-token set handed to DFlash, and cancellation reaching the DFlash stream. +final class DFlashStreamBridgeTests: XCTestCase { + + // MARK: stopTokenIDs(…) — same union the standard path stops on + + func testStopTokenIDsUnionsConfigurationTokenizerAndExtraTokens() { + let ids = stopTokenIDs( + eosTokenIDs: [151645, 151643], + tokenizerEOS: 151645, + extraEOSTokens: ["<|endoftext|>", "<|im_end|>", "<|not_in_vocab|>"], + tokenID: { ["<|endoftext|>": 151643, "<|im_end|>": 151645][$0] }) + XCTAssertEqual(ids, [151643, 151645], "deduplicated; unknown extra tokens are dropped") + } + + func testStopTokenIDsWithoutTokenizerEOS() { + XCTAssertEqual(stopTokenIDs(eosTokenIDs: [2], tokenizerEOS: nil, extraEOSTokens: [], tokenID: { _ in nil }), [2]) + XCTAssertEqual(stopTokenIDs(eosTokenIDs: [], tokenizerEOS: nil, extraEOSTokens: [], tokenID: { _ in nil }), []) + } + + // MARK: dflashGenerationStream(…) — tearing down the consumer cancels the producer + + /// A DFlash-shaped stream that never ends on its own, and reports whether the + /// consumer side tore it down. + private final class EndlessEvents: @unchecked Sendable { + let terminated = XCTestExpectation(description: "DFlash stream terminated") + private(set) var yielded = 0 + let stream: AsyncStream + + init() { + let terminated = self.terminated + var counter = 0 + stream = AsyncStream { continuation in + let producer = Task { + while !Task.isCancelled { + counter += 1 + continuation.yield(.token(tokenID: counter, generatedTokens: counter, acceptanceRatio: 0, cyclesCompleted: counter)) + try? await Task.sleep(for: .milliseconds(5)) + } + continuation.finish() + } + continuation.onTermination = { _ in + producer.cancel() + terminated.fulfill() + } + } + } + } + + func testStreamTornDownAfterConsumerStopsEarly() async throws { + let events = EndlessEvents() + var received: [String] = [] + do { + // Scoped like the handler's `for await` over the generation stream: once the + // consumer breaks out and the stream value goes away, the bridge must be + // torn down and the DFlash stream with it. + let generation = dflashGenerationStream(events.stream, tokenLimit: 1_000) { "t\($0) " } + for await item in generation { + if case .chunk(let text, _) = item { received.append(text) } + if received.count == 3 { break } + } + } + XCTAssertEqual(received, ["t1 ", "t2 ", "t3 "]) + await fulfillment(of: [events.terminated], timeout: 2) + } + + func testBridgeMapsTokensAndSummary() async throws { + let summary = DFlashSummary( + elapsedUs: 2_000_000, promptTokenCount: 7, generatedTokenIDs: [1, 2, 3], + acceptedFromDraft: 1, acceptanceRatio: 0.33, blockTokens: 16, cyclesCompleted: 2, + phaseTimingsUs: .init(prefill: 500_000, draft: 0, verify: 0, replay: 0)) + let (events, continuation) = AsyncStream.makeStream() + continuation.yield(.prefill(promptTokenCount: 7, prefillUs: 500_000)) + continuation.yield(.token(tokenID: 1, generatedTokens: 1, acceptanceRatio: 0, cyclesCompleted: 1)) + continuation.yield(.token(tokenID: 2, generatedTokens: 2, acceptanceRatio: 0, cyclesCompleted: 1)) + continuation.yield(.summary(summary)) + continuation.finish() + + var chunks: [String] = [] + var info: GenerateCompletionInfo? + // The detokenizer seam may hold text back (nil) and release it on a later token. + for await item in dflashGenerationStream(events, tokenLimit: 3, nextText: { $0 == 1 ? nil : "one two" }) { + switch item { + case .chunk(let text, _): chunks.append(text) + case .info(let i): info = i + default: break + } + } + XCTAssertEqual(chunks, ["one two"]) + XCTAssertEqual(info?.generationTokenCount, 3) + XCTAssertEqual(info?.promptTokenCount, 7) + XCTAssertEqual(info?.stopReason, .length, "3 generated against a limit of 3 is a length stop") + } +} diff --git a/tests/test-dflash.sh b/tests/test-dflash.sh index 92a4a6d..7d674e5 100755 --- a/tests/test-dflash.sh +++ b/tests/test-dflash.sh @@ -189,6 +189,26 @@ else fail "Speculative decoding path not activated (missing log line)" fi +# ── Test 4: A stream that ends early must not leave a generation running ── +# Issue #224: the text-level stop ended the SSE stream and released the slot, but the +# DFlash loop ran on to max_tokens; the next request then overlapped it and MLX aborted. +# A one-word answer against a large budget leaves plenty of orphaned work for the next +# request to collide with. +log "Test 4: early-ended stream does not leave a generation running" +EARLY=$(curl -sf -N --max-time 120 -X POST "$URL/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d "{\"model\":\"$MAIN_MODEL\",\"stream\":true,\"max_tokens\":200,\"messages\":[{\"role\":\"user\",\"content\":\"Reply with the single word: ok\"}]}" \ + 2>/dev/null || true) +AFTER=$(curl -sf --max-time 120 -X POST "$URL/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d "{\"model\":\"$MAIN_MODEL\",\"max_tokens\":10,\"messages\":[{\"role\":\"user\",\"content\":\"Say the word: done\"}]}" 2>/dev/null || echo "") +AFTER_CONTENT=$(echo "$AFTER" | jq -r '.choices[0].message.content // empty' 2>/dev/null || echo "") +if echo "$EARLY" | grep -q "data: \[DONE\]" && [ -n "$AFTER_CONTENT" ] && curl -sf "$URL/health" >/dev/null 2>&1; then + pass "request after an early-ended stream completed and the server is still up" +else + fail "request after an early-ended stream failed (server alive: $(curl -sf "$URL/health" >/dev/null 2>&1 && echo yes || echo no))" +fi + # ── Test 5: Multiple sequential requests (stability) ──────────────── log "Test 5: Sequential request stability (3 requests)"