From c48271188d1397e426a4a81b6d55deaaa25dd051 Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Mon, 15 Jun 2026 13:46:06 -0700 Subject: [PATCH] Fix pipelined engine reset(): cancel before drain MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The producer task holds the engine lock until it finishes. If the consumer stops early (e.g. EOS detected), the producer may still be running. drain() waits for the lock to be released, but cancellation was only happening AFTER drain — creating a deadlock that triggers the 5-second fatalError timeout. Fix: cancel the GenerationToken and Task before calling drain(). The producer checks isCancelled on its next loop iteration, releases the lock, and drain() can proceed. Fixes #41 --- .../InferenceEngines/CoreAIPipelinedEngine.swift | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift index bd2e84f..4ddc8f2 100644 --- a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift +++ b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift @@ -186,12 +186,17 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable { } func reset() { - drain() + // Cancel active generation BEFORE draining — otherwise drain() waits + // forever for a producer that will never release the engine. _activeToken.withLock { $0?.cancel() $0 = nil } - _generationTask.withLock { $0 = nil } + _generationTask.withLock { + $0?.cancel() + $0 = nil + } + drain() guard tryAcquireEngine() else { return } defer { releaseEngine() } engine.reset()