Add cancel() and isBusy to InferenceEngine protocol

stikves · stikves · commit f1a8bf3d3678 · 2026-06-12T13:09:59.000-07:00
diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift
@@ -39,6 +39,7 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
 
     nonisolated(unsafe) private var engine: EngineImpl
     private let engineInUse = Atomic<Bool>(false)
+    private let generationTask = Mutex<Task<Void, Never>?>(nil)
     let config: ModelConfig
 
     init(
@@ -104,7 +105,7 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
         let stopReasonStore = StopReasonStore()
         let (base, outputContinuation) =
             AsyncThrowingStream<InferenceOutput, any Error>.makeStream()
-        Task {
+        let task = Task {
             self.acquireEngine()
             defer { self.releaseEngine() }
             do {
@@ -139,15 +140,28 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
                 outputContinuation.finish(throwing: error)
             }
         }
+        generationTask.withLock { $0 = task }
         return GenerationSequence(base: base, stopReasonStore: stopReasonStore)
     }
 
+    var isBusy: Bool { engineInUse.load(ordering: .acquiring) }
+
+    func cancel() async {
+        let task = generationTask.withLock { v in
+            v?.cancel()
+            return v
+        }
+        await task?.value
+        await engine.computeStream.currentWorkCompleted()
+        generationTask.withLock { $0 = nil }
+    }
+
     /// Wait for any in-flight generate() Task to return the engine.
     private func drain() {
         var attempts = 0
         while engineInUse.load(ordering: .acquiring) {
             attempts += 1
-            if attempts > 5000 {
+            if attempts > Self.drainTimeoutIterations {
                 fatalError("Engine not returned after drain() — tokenSequence Task stuck?")
             }
             Thread.sleep(forTimeInterval: 0.001)
diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift
@@ -73,6 +73,7 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
 
     // Track in-flight generation for drain (same pattern as pipelined engine)
     private let generating = Mutex(false)
+    private let cancelRequested = Mutex(false)
 
     // MARK: - Init
 
@@ -347,12 +348,31 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
 
     // MARK: - Lifecycle
 
+    public var isBusy: Bool { generating.withLock { $0 } }
+
+    public func cancel() async {
+        cancelRequested.withLock { $0 = true }
+        var attempts = 0
+        while generating.withLock({ $0 }) {
+            attempts += 1
+            if attempts > Self.drainTimeoutIterations {
+                CLILogger.log(
+                    "WARNING: Sequential engine cancel() drain timeout after \(Self.drainTimeoutIterations) iterations",
+                    component: "CoreAISequential"
+                )
+                break
+            }
+            await Task.yield()
+        }
+        cancelRequested.withLock { $0 = false }
+    }
+
     /// Wait for any in-flight generate() Task to finish.
     private func drain() {
         var attempts = 0
         while generating.withLock({ $0 }) {
             attempts += 1
-            if attempts > 5000 {
+            if attempts > Self.drainTimeoutIterations {
                 fatalError("Sequential engine drain() timeout — generation Task stuck?")
             }
             Thread.sleep(forTimeInterval: 0.001)
@@ -550,6 +570,12 @@ extension CoreAISequentialEngine.GenerationSequence {
             do {
                 try Task.checkCancellation()
 
+                if engine.cancelRequested.withLock({ $0 }) {
+                    stopReasonStore.set(.cancelled)
+                    finishAndRelease()
+                    return nil
+                }
+
                 guard engine.processedTokenCount < inputTokens.count else {
                     throw InferenceRuntimeError.invalidState("No new tokens to process")
                 }
diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift
@@ -6,6 +6,7 @@
 import CoreAI
 import CoreAIShared
 import Foundation
+import Synchronization
 
 /// Static-shape inference engine using Core AI models.
 public final class StaticShapeEngine: InferenceEngine, @unchecked Sendable {
@@ -47,6 +48,10 @@ public final class StaticShapeEngine: InferenceEngine, @unchecked Sendable {
     // Number of tokens already processed in the current sequence.
     private var processedTokenCount: Int = 0
 
+    // Track in-flight generation for drain/cancel
+    private let generating = Mutex(false)
+    private let cancelRequested = Mutex(false)
+
     // MARK: - Initialization
 
     public init(configuration: ModelConfig, preparedModel: PreparedModel) async throws {
@@ -531,7 +536,39 @@ public final class StaticShapeEngine: InferenceEngine, @unchecked Sendable {
 
     // MARK: - Lifecycle
 
+    public var isBusy: Bool { generating.withLock { $0 } }
+
+    public func cancel() async {
+        cancelRequested.withLock { $0 = true }
+        var attempts = 0
+        while generating.withLock({ $0 }) {
+            attempts += 1
+            if attempts > Self.drainTimeoutIterations {
+                CLILogger.log(
+                    "WARNING: StaticShape engine cancel() drain timeout after \(Self.drainTimeoutIterations) iterations",
+                    component: "StaticShape"
+                )
+                break
+            }
+            await Task.yield()
+        }
+        cancelRequested.withLock { $0 = false }
+    }
+
+    /// Wait for any in-flight generate() iteration to finish.
+    private func drain() {
+        var attempts = 0
+        while generating.withLock({ $0 }) {
+            attempts += 1
+            if attempts > Self.drainTimeoutIterations {
+                fatalError("StaticShape engine drain() timeout — generation stuck?")
+            }
+            Thread.sleep(forTimeInterval: 0.001)
+        }
+    }
+
     public func reset() {
+        drain()
         let resetSpan = InstrumentsProfiler.beginReset(engine: "StaticShape")
         processedTokenCount = 0
         resetSpan.end()
@@ -595,6 +632,8 @@ extension StaticShapeEngine.GenerationSequence {
 
         private var inputTokens: [StaticShapeEngine.TokenId]
         private var step: Int = 0
+        private var didAcquireLock: Bool = false
+        private var finished: Bool = false
 
         init(
             engine: StaticShapeEngine,
@@ -620,15 +659,29 @@ extension StaticShapeEngine.GenerationSequence {
         }
 
         public mutating func next() async throws -> InferenceOutput? {
+            if finished { return nil }
+
+            if !didAcquireLock {
+                engine.generating.withLock { $0 = true }
+                didAcquireLock = true
+            }
+
             guard step < maxTokens else {
                 // Natural exhaustion. Don't clobber a reason a decoder set (e.g. `.eos`).
                 stopReasonStore.setIfUnset(.maxTokens)
+                finishAndRelease()
                 return nil
             }
 
             do {
                 try Task.checkCancellation()
 
+                if engine.cancelRequested.withLock({ $0 }) {
+                    stopReasonStore.set(.cancelled)
+                    finishAndRelease()
+                    return nil
+                }
+
                 // When forced, we still need the forward pass (for logits + KV cache update)
                 // but skip the sampler — the next token is predetermined.
                 let (logits, sampledToken) = try await engine.inference(
@@ -647,11 +700,22 @@ extension StaticShapeEngine.GenerationSequence {
                 )
             } catch is CancellationError {
                 stopReasonStore.set(.cancelled)
+                finishAndRelease()
                 throw CancellationError()
             } catch {
                 stopReasonStore.set(.error)
+                finishAndRelease()
                 throw error
             }
         }
+
+        private mutating func finishAndRelease() {
+            guard !finished else { return }
+            finished = true
+            if didAcquireLock {
+                engine.generating.withLock { $0 = false }
+                didAcquireLock = false
+            }
+        }
     }
 }
diff --git a/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift b/swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift
@@ -106,6 +106,14 @@ public protocol InferenceEngine: Sendable {
 
     // MARK: - Lifecycle
 
+    /// Whether the engine is currently running a generation.
+    var isBusy: Bool { get }
+
+    /// Cancel in-flight generation and wait for the engine to become idle.
+    ///
+    /// Best-effort: logs a warning if the drain times out but does not trap.
+    func cancel() async
+
     /// Reset internal state including KV cache.
     func reset() async throws
 
@@ -156,6 +164,19 @@ extension InferenceConfiguration {
 
 // MARK: - Default Implementations
 
+extension InferenceEngine {
+    /// Maximum iterations for drain loops before timing out.
+    /// - `cancel()` logs a warning and breaks (best-effort).
+    /// - `reset()` traps (hard lifecycle boundary).
+    static var drainTimeoutIterations: Int { 5000 }
+
+    /// Default: engine is not busy.
+    public var isBusy: Bool { false }
+
+    /// Default no-op cancel.
+    public func cancel() async {}
+}
+
 extension InferenceEngine {
     /// Default: supportsLogits is false. Engines that can return per-step
     /// logits (sequential, static-shape) override this to true.
diff --git a/swift/Sources/CoreAILanguageModels/LanguageModel/CoreAILanguageModel.swift b/swift/Sources/CoreAILanguageModels/LanguageModel/CoreAILanguageModel.swift
@@ -290,6 +290,7 @@ public struct CoreAILanguageModel: LanguageModel {
             let maxTokens = request.generationOptions.maximumResponseTokens ?? 512
 
             // Reset engine state for new generation
+            await engine.cancel()
             try await engine.reset()
 
             // FoundationModels now threads entry identity itself based on event
diff --git a/swift/Tests/LanguageModelsTests/TestUtilities.swift b/swift/Tests/LanguageModelsTests/TestUtilities.swift
@@ -158,6 +158,10 @@ class MockEngine: InferenceEngine, @unchecked Sendable {
         }
     }
 
+    var isBusy: Bool { false }
+
+    func cancel() async {}
+
     func reset() async throws {
         resetCalled = true
         inferenceCallCount = 0
diff --git a/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift b/swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift
@@ -288,3 +288,42 @@ struct GenerateMultiCallTests {
         #expect(count == 0)
     }
 }
+
+// MARK: - Cancel API Tests
+
+@Suite("CancelAPI")
+struct CancelAPITests {
+    @Test("MockEngine isBusy returns false by default")
+    func mockEngineNotBusy() {
+        let engine = MockEngine(tokens: [1, 2, 3])
+        #expect(engine.isBusy == false)
+    }
+
+    @Test("MockEngine cancel() is a no-op")
+    func mockEngineCancelNoOp() async {
+        let engine = MockEngine(tokens: [1, 2, 3])
+        await engine.cancel()
+        #expect(engine.isBusy == false)
+    }
+
+    @Test("drainTimeoutIterations default value")
+    func drainTimeoutDefault() {
+        #expect(MockEngine.drainTimeoutIterations == 5000)
+    }
+
+    @Test("cancel() before generate does not crash")
+    func cancelBeforeGenerate() async {
+        let engine = MockEngine(tokens: [1, 2, 3])
+        await engine.cancel()
+        // Should be able to generate normally after cancel
+        var tokens: [Int32] = []
+        for try await output in try engine.generate(
+            with: [1],
+            samplingConfiguration: .greedy,
+            inferenceOptions: InferenceOptions(maxTokens: 3)
+        ) {
+            tokens.append(output.tokenId)
+        }
+        #expect(tokens.count == 3)
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -158,6 +158,10 @@ class MockEngine: InferenceEngine, @unchecked Sendable {`
`158`	`158`	`}`
`159`	`159`	`}`
`160`	`160`
	`161`	`+ var isBusy: Bool { false }`
	`162`	`+`
	`163`	`+ func cancel() async {}`
	`164`	`+`
`161`	`165`	`func reset() async throws {`
`162`	`166`	`resetCalled = true`
`163`	`167`	`inferenceCallCount = 0`