Skip to content

Commit 1a754e1

Browse files
stikvesstikves
committed
Add cancel() and isBusy with GenerationToken pattern
1 parent d83f0b0 commit 1a754e1

7 files changed

Lines changed: 329 additions & 29 deletions

File tree

‎swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift‎

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,12 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
4141
private let engineInUse = Atomic<Bool>(false)
4242
let config: ModelConfig
4343

44+
// Generation lifecycle
45+
private let _activeToken = Mutex<GenerationToken?>(nil)
46+
private let _generationTask = Mutex<Task<Void, Never>?>(nil)
47+
48+
var isBusy: Bool { _activeToken.withLock { $0 != nil } }
49+
4450
init(
4551
config: ModelConfig,
4652
preparedModel: PreparedModel,
@@ -104,9 +110,20 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
104110
let stopReasonStore = StopReasonStore()
105111
let (base, outputContinuation) =
106112
AsyncThrowingStream<InferenceOutput, any Error>.makeStream()
107-
Task {
113+
114+
let token = GenerationToken()
115+
_activeToken.withLock { $0 = token }
116+
117+
let task = Task {
108118
self.acquireEngine()
109-
defer { self.releaseEngine() }
119+
defer {
120+
self.releaseEngine()
121+
// Only clear if this generation still owns both slots
122+
if self._activeToken.withLock({ $0 === token }) {
123+
self._activeToken.withLock { $0 = nil }
124+
self._generationTask.withLock { $0 = nil }
125+
}
126+
}
110127
do {
111128
let (tokenStream, tokenContinuation) =
112129
AsyncThrowingStream<InferenceEngine.TokenId, any Error>.makeStream()
@@ -139,6 +156,7 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
139156
outputContinuation.finish(throwing: error)
140157
}
141158
}
159+
_generationTask.withLock { $0 = task }
142160
return GenerationSequence(base: base, stopReasonStore: stopReasonStore)
143161
}
144162

@@ -154,8 +172,20 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
154172
}
155173
}
156174

175+
func cancel() async throws {
176+
let task: Task<Void, Never>? = _generationTask.withLock { task in
177+
task?.cancel()
178+
defer { task = nil }
179+
return task
180+
}
181+
_activeToken.withLock { $0?.cancel(); $0 = nil }
182+
await task?.value
183+
}
184+
157185
func reset() {
158186
drain()
187+
_activeToken.withLock { $0?.cancel(); $0 = nil }
188+
_generationTask.withLock { $0 = nil }
159189
guard tryAcquireEngine() else { return }
160190
defer { releaseEngine() }
161191
engine.reset()

‎swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift‎

Lines changed: 41 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -71,8 +71,16 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
7171
// Track processed tokens for incremental inference
7272
private var processedTokenCount: Int = 0
7373

74-
// Track in-flight generation for drain (same pattern as pipelined engine)
75-
private let generating = Mutex(false)
74+
// Track in-flight generation via token (replaces simple bool lock)
75+
private let _activeToken = Mutex<GenerationToken?>(nil)
76+
77+
public var isBusy: Bool { _activeToken.withLock { $0 != nil } }
78+
79+
/// Clear the engine's active token if it matches the given token.
80+
/// Called by the iterator when generation finishes or is cancelled.
81+
func clearTokenIfActive(_ token: GenerationToken) {
82+
_activeToken.withLock { if $0 === token { $0 = nil } }
83+
}
7684

7785
// MARK: - Init
7886

@@ -337,11 +345,14 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
337345
samplingConfiguration: SamplingConfiguration,
338346
inferenceOptions: InferenceOptions
339347
) throws -> GenerationSequence {
340-
GenerationSequence(
348+
let token = GenerationToken()
349+
_activeToken.withLock { $0 = token }
350+
return GenerationSequence(
341351
engine: self,
342352
input: input,
343353
samplingConfiguration: samplingConfiguration,
344-
inferenceOptions: inferenceOptions
354+
inferenceOptions: inferenceOptions,
355+
generationToken: token
345356
)
346357
}
347358

@@ -350,7 +361,7 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
350361
/// Wait for any in-flight generate() Task to finish.
351362
private func drain() {
352363
var attempts = 0
353-
while generating.withLock({ $0 }) {
364+
while _activeToken.withLock({ $0 != nil }) {
354365
attempts += 1
355366
if attempts > 5000 {
356367
fatalError("Sequential engine drain() timeout — generation Task stuck?")
@@ -359,8 +370,12 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
359370
}
360371
}
361372

373+
public func cancel() async throws {
374+
_activeToken.withLock { $0?.cancel(); $0 = nil }
375+
}
376+
362377
public func reset() {
363-
drain()
378+
_activeToken.withLock { $0?.cancel(); $0 = nil }
364379
let resetSpan = InstrumentsProfiler.beginReset(engine: "CoreAIClean")
365380
processedTokenCount = 0
366381
zeroFill(&keyCache)
@@ -464,6 +479,7 @@ extension CoreAISequentialEngine {
464479
let input: [CoreAISequentialEngine.TokenId]
465480
let samplingConfiguration: SamplingConfiguration
466481
let inferenceOptions: InferenceOptions
482+
let generationToken: GenerationToken
467483

468484
/// Shared with the iterator so the caller can read why generation ended.
469485
let stopReasonStore = StopReasonStore()
@@ -480,7 +496,8 @@ extension CoreAISequentialEngine {
480496
input: input,
481497
samplingConfiguration: samplingConfiguration,
482498
inferenceOptions: inferenceOptions,
483-
stopReasonStore: stopReasonStore
499+
stopReasonStore: stopReasonStore,
500+
generationToken: generationToken
484501
)
485502
}
486503
}
@@ -497,24 +514,26 @@ extension CoreAISequentialEngine.GenerationSequence {
497514
private let forcedContinuation: [CoreAISequentialEngine.TokenId]?
498515
private let maxTokens: Int
499516
private let stopReasonStore: StopReasonStore
517+
private let generationToken: GenerationToken
500518

501519
private var inputTokens: [CoreAISequentialEngine.TokenId]
502520
private var step: Int = 0
503-
private var didAcquireLock: Bool = false
504521
private var finished: Bool = false
505522

506523
init(
507524
engine: CoreAISequentialEngine,
508525
input: [CoreAISequentialEngine.TokenId],
509526
samplingConfiguration: SamplingConfiguration,
510527
inferenceOptions: InferenceOptions,
511-
stopReasonStore: StopReasonStore
528+
stopReasonStore: StopReasonStore,
529+
generationToken: GenerationToken
512530
) {
513531
self.engine = engine
514532
self.samplingConfiguration = samplingConfiguration
515533
self.returnsLogits = inferenceOptions.includeLogits
516534
self.forcedContinuation = inferenceOptions.forcedContinuation
517535
self.stopReasonStore = stopReasonStore
536+
self.generationToken = generationToken
518537
self.inputTokens = input
519538
if let forced = inferenceOptions.forcedContinuation {
520539
self.maxTokens = forced.count
@@ -527,17 +546,16 @@ extension CoreAISequentialEngine.GenerationSequence {
527546
}
528547

529548
deinit {
530-
if didAcquireLock {
531-
engine.generating.withLock { $0 = false }
532-
}
549+
engine.clearTokenIfActive(generationToken)
533550
}
534551

535552
public func next() async throws -> InferenceOutput? {
536553
if finished { return nil }
537554

538-
if !didAcquireLock {
539-
engine.generating.withLock { $0 = true }
540-
didAcquireLock = true
555+
if generationToken.isCancelled {
556+
stopReasonStore.set(.cancelled)
557+
finishAndRelease()
558+
return nil
541559
}
542560

543561
guard step < maxTokens else {
@@ -572,6 +590,13 @@ extension CoreAISequentialEngine.GenerationSequence {
572590
logitBuffer = lastLogits
573591
}
574592

593+
// Check cancellation after inference step
594+
if generationToken.isCancelled {
595+
stopReasonStore.set(.cancelled)
596+
finishAndRelease()
597+
return nil
598+
}
599+
575600
let nextToken: Int32
576601
if let forced = forcedContinuation {
577602
nextToken = forced[step]
@@ -603,10 +628,7 @@ extension CoreAISequentialEngine.GenerationSequence {
603628
return
604629
}
605630
finished = true
606-
if didAcquireLock {
607-
engine.generating.withLock { $0 = false }
608-
didAcquireLock = false
609-
}
631+
engine.clearTokenIfActive(generationToken)
610632
}
611633
}
612634
}

‎swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift‎

Lines changed: 53 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import CoreAI
77
import CoreAIShared
88
import Foundation
9+
import Synchronization
910

1011
/// Static-shape inference engine using Core AI models.
1112
public final class StaticShapeEngine: InferenceEngine, @unchecked Sendable {
@@ -47,6 +48,16 @@ public final class StaticShapeEngine: InferenceEngine, @unchecked Sendable {
4748
// Number of tokens already processed in the current sequence.
4849
private var processedTokenCount: Int = 0
4950

51+
// Track in-flight generation via token
52+
private let _activeToken = Mutex<GenerationToken?>(nil)
53+
54+
public var isBusy: Bool { _activeToken.withLock { $0 != nil } }
55+
56+
/// Clear the engine's active token if it matches the given token.
57+
func clearTokenIfActive(_ token: GenerationToken) {
58+
_activeToken.withLock { if $0 === token { $0 = nil } }
59+
}
60+
5061
// MARK: - Initialization
5162

5263
public init(configuration: ModelConfig, preparedModel: PreparedModel) async throws {
@@ -317,11 +328,14 @@ public final class StaticShapeEngine: InferenceEngine, @unchecked Sendable {
317328
samplingConfiguration: SamplingConfiguration,
318329
inferenceOptions: InferenceOptions
319330
) throws -> GenerationSequence {
320-
GenerationSequence(
331+
let token = GenerationToken()
332+
_activeToken.withLock { $0 = token }
333+
return GenerationSequence(
321334
engine: self,
322335
input: input,
323336
samplingConfiguration: samplingConfiguration,
324-
inferenceOptions: inferenceOptions
337+
inferenceOptions: inferenceOptions,
338+
generationToken: token
325339
)
326340
}
327341

@@ -531,7 +545,12 @@ public final class StaticShapeEngine: InferenceEngine, @unchecked Sendable {
531545

532546
// MARK: - Lifecycle
533547

548+
public func cancel() async throws {
549+
_activeToken.withLock { $0?.cancel(); $0 = nil }
550+
}
551+
534552
public func reset() {
553+
_activeToken.withLock { $0?.cancel(); $0 = nil }
535554
let resetSpan = InstrumentsProfiler.beginReset(engine: "StaticShape")
536555
processedTokenCount = 0
537556
resetSpan.end()
@@ -559,6 +578,7 @@ extension StaticShapeEngine {
559578
let input: [TokenId]
560579
let samplingConfiguration: SamplingConfiguration
561580
let inferenceOptions: InferenceOptions
581+
let generationToken: GenerationToken
562582

563583
/// Shared with the iterator so the caller can read why generation ended.
564584
let stopReasonStore = StopReasonStore()
@@ -575,7 +595,8 @@ extension StaticShapeEngine {
575595
input: input,
576596
samplingConfiguration: samplingConfiguration,
577597
inferenceOptions: inferenceOptions,
578-
stopReasonStore: stopReasonStore
598+
stopReasonStore: stopReasonStore,
599+
generationToken: generationToken
579600
)
580601
}
581602
}
@@ -592,22 +613,26 @@ extension StaticShapeEngine.GenerationSequence {
592613
private let forcedContinuation: [StaticShapeEngine.TokenId]?
593614
private let maxTokens: Int
594615
private let stopReasonStore: StopReasonStore
616+
private let generationToken: GenerationToken
595617

596618
private var inputTokens: [StaticShapeEngine.TokenId]
597619
private var step: Int = 0
620+
private var finished: Bool = false
598621

599622
init(
600623
engine: StaticShapeEngine,
601624
input: [StaticShapeEngine.TokenId],
602625
samplingConfiguration: SamplingConfiguration,
603626
inferenceOptions: InferenceOptions,
604-
stopReasonStore: StopReasonStore
627+
stopReasonStore: StopReasonStore,
628+
generationToken: GenerationToken
605629
) {
606630
self.engine = engine
607631
self.samplingConfiguration = samplingConfiguration
608632
self.returnsLogits = inferenceOptions.includeLogits
609633
self.forcedContinuation = inferenceOptions.forcedContinuation
610634
self.stopReasonStore = stopReasonStore
635+
self.generationToken = generationToken
611636
self.inputTokens = input
612637
if let forced = inferenceOptions.forcedContinuation {
613638
self.maxTokens = forced.count
@@ -620,9 +645,18 @@ extension StaticShapeEngine.GenerationSequence {
620645
}
621646

622647
public mutating func next() async throws -> InferenceOutput? {
648+
if finished { return nil }
649+
650+
if generationToken.isCancelled {
651+
stopReasonStore.set(.cancelled)
652+
finishAndRelease()
653+
return nil
654+
}
655+
623656
guard step < maxTokens else {
624657
// Natural exhaustion. Don't clobber a reason a decoder set (e.g. `.eos`).
625658
stopReasonStore.setIfUnset(.maxTokens)
659+
finishAndRelease()
626660
return nil
627661
}
628662

@@ -637,6 +671,13 @@ extension StaticShapeEngine.GenerationSequence {
637671
returnsLogits: returnsLogits || forcedContinuation != nil
638672
)
639673

674+
// Check cancellation after inference step
675+
if generationToken.isCancelled {
676+
stopReasonStore.set(.cancelled)
677+
finishAndRelease()
678+
return nil
679+
}
680+
640681
let nextToken = forcedContinuation?[step] ?? sampledToken
641682
inputTokens.append(nextToken)
642683
step += 1
@@ -647,11 +688,19 @@ extension StaticShapeEngine.GenerationSequence {
647688
)
648689
} catch is CancellationError {
649690
stopReasonStore.set(.cancelled)
691+
finishAndRelease()
650692
throw CancellationError()
651693
} catch {
652694
stopReasonStore.set(.error)
695+
finishAndRelease()
653696
throw error
654697
}
655698
}
699+
700+
private mutating func finishAndRelease() {
701+
guard !finished else { return }
702+
finished = true
703+
engine.clearTokenIfActive(generationToken)
704+
}
656705
}
657706
}

‎swift/Sources/CoreAILanguageModels/InferenceEngines/GenerationToken.swift‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
// Copyright 2026 Apple Inc.
2+
//
3+
// Use of this source code is governed by a BSD-3-clause license that can
4+
// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
5+
6+
import Synchronization
7+
8+
/// A token representing an active generation session.
9+
///
10+
/// Created by `generate()`, held by the iterator. The engine retains a
11+
/// reference to the active token and can cancel it at any time. The iterator
12+
/// checks `isCancelled` on each `next()` call.
13+
public final class GenerationToken: Sendable {
14+
private let _cancelled = Mutex(false)
15+
16+
public var isCancelled: Bool { _cancelled.withLock { $0 } }
17+
18+
public func cancel() { _cancelled.withLock { $0 = true } }
19+
}

0 commit comments

Comments
 (0)