Skip to content

Commit 342fd81

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

7 files changed

Lines changed: 343 additions & 29 deletions

File tree

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

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,13 @@ 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 activeGenerationToken: GenerationToken? { _activeToken.withLock { $0 } }
49+
var isBusy: Bool { _activeToken.withLock { $0 != nil } }
50+
4451
init(
4552
config: ModelConfig,
4653
preparedModel: PreparedModel,
@@ -104,9 +111,17 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
104111
let stopReasonStore = StopReasonStore()
105112
let (base, outputContinuation) =
106113
AsyncThrowingStream<InferenceOutput, any Error>.makeStream()
107-
Task {
114+
115+
let token = GenerationToken()
116+
_activeToken.withLock { $0 = token }
117+
118+
let task = Task {
108119
self.acquireEngine()
109-
defer { self.releaseEngine() }
120+
defer {
121+
self.releaseEngine()
122+
self._activeToken.withLock { if $0 === token { $0 = nil } }
123+
self._generationTask.withLock { $0 = nil }
124+
}
110125
do {
111126
let (tokenStream, tokenContinuation) =
112127
AsyncThrowingStream<InferenceEngine.TokenId, any Error>.makeStream()
@@ -139,6 +154,7 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
139154
outputContinuation.finish(throwing: error)
140155
}
141156
}
157+
_generationTask.withLock { $0 = task }
142158
return GenerationSequence(base: base, stopReasonStore: stopReasonStore)
143159
}
144160

@@ -154,8 +170,20 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
154170
}
155171
}
156172

173+
func cancel() async throws {
174+
let task: Task<Void, Never>? = _generationTask.withLock { task in
175+
task?.cancel()
176+
defer { task = nil }
177+
return task
178+
}
179+
_activeToken.withLock { $0?.cancel(); $0 = nil }
180+
await task?.value
181+
}
182+
157183
func reset() {
158184
drain()
185+
_activeToken.withLock { $0?.cancel(); $0 = nil }
186+
_generationTask.withLock { $0 = nil }
159187
guard tryAcquireEngine() else { return }
160188
defer { releaseEngine() }
161189
engine.reset()

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

Lines changed: 42 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -71,8 +71,17 @@ 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 activeGenerationToken: GenerationToken? { _activeToken.withLock { $0 } }
78+
public var isBusy: Bool { _activeToken.withLock { $0 != nil } }
79+
80+
/// Clear the engine's active token if it matches the given token.
81+
/// Called by the iterator when generation finishes or is cancelled.
82+
func clearTokenIfActive(_ token: GenerationToken) {
83+
_activeToken.withLock { if $0 === token { $0 = nil } }
84+
}
7685

7786
// MARK: - Init
7887

@@ -337,11 +346,14 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
337346
samplingConfiguration: SamplingConfiguration,
338347
inferenceOptions: InferenceOptions
339348
) throws -> GenerationSequence {
340-
GenerationSequence(
349+
let token = GenerationToken()
350+
_activeToken.withLock { $0 = token }
351+
return GenerationSequence(
341352
engine: self,
342353
input: input,
343354
samplingConfiguration: samplingConfiguration,
344-
inferenceOptions: inferenceOptions
355+
inferenceOptions: inferenceOptions,
356+
generationToken: token
345357
)
346358
}
347359

@@ -350,7 +362,7 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
350362
/// Wait for any in-flight generate() Task to finish.
351363
private func drain() {
352364
var attempts = 0
353-
while generating.withLock({ $0 }) {
365+
while _activeToken.withLock({ $0 != nil }) {
354366
attempts += 1
355367
if attempts > 5000 {
356368
fatalError("Sequential engine drain() timeout — generation Task stuck?")
@@ -359,8 +371,12 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
359371
}
360372
}
361373

374+
public func cancel() async throws {
375+
_activeToken.withLock { $0?.cancel(); $0 = nil }
376+
}
377+
362378
public func reset() {
363-
drain()
379+
_activeToken.withLock { $0?.cancel(); $0 = nil }
364380
let resetSpan = InstrumentsProfiler.beginReset(engine: "CoreAIClean")
365381
processedTokenCount = 0
366382
zeroFill(&keyCache)
@@ -464,6 +480,7 @@ extension CoreAISequentialEngine {
464480
let input: [CoreAISequentialEngine.TokenId]
465481
let samplingConfiguration: SamplingConfiguration
466482
let inferenceOptions: InferenceOptions
483+
let generationToken: GenerationToken
467484

468485
/// Shared with the iterator so the caller can read why generation ended.
469486
let stopReasonStore = StopReasonStore()
@@ -480,7 +497,8 @@ extension CoreAISequentialEngine {
480497
input: input,
481498
samplingConfiguration: samplingConfiguration,
482499
inferenceOptions: inferenceOptions,
483-
stopReasonStore: stopReasonStore
500+
stopReasonStore: stopReasonStore,
501+
generationToken: generationToken
484502
)
485503
}
486504
}
@@ -497,24 +515,26 @@ extension CoreAISequentialEngine.GenerationSequence {
497515
private let forcedContinuation: [CoreAISequentialEngine.TokenId]?
498516
private let maxTokens: Int
499517
private let stopReasonStore: StopReasonStore
518+
private let generationToken: GenerationToken
500519

501520
private var inputTokens: [CoreAISequentialEngine.TokenId]
502521
private var step: Int = 0
503-
private var didAcquireLock: Bool = false
504522
private var finished: Bool = false
505523

506524
init(
507525
engine: CoreAISequentialEngine,
508526
input: [CoreAISequentialEngine.TokenId],
509527
samplingConfiguration: SamplingConfiguration,
510528
inferenceOptions: InferenceOptions,
511-
stopReasonStore: StopReasonStore
529+
stopReasonStore: StopReasonStore,
530+
generationToken: GenerationToken
512531
) {
513532
self.engine = engine
514533
self.samplingConfiguration = samplingConfiguration
515534
self.returnsLogits = inferenceOptions.includeLogits
516535
self.forcedContinuation = inferenceOptions.forcedContinuation
517536
self.stopReasonStore = stopReasonStore
537+
self.generationToken = generationToken
518538
self.inputTokens = input
519539
if let forced = inferenceOptions.forcedContinuation {
520540
self.maxTokens = forced.count
@@ -527,17 +547,16 @@ extension CoreAISequentialEngine.GenerationSequence {
527547
}
528548

529549
deinit {
530-
if didAcquireLock {
531-
engine.generating.withLock { $0 = false }
532-
}
550+
engine.clearTokenIfActive(generationToken)
533551
}
534552

535553
public func next() async throws -> InferenceOutput? {
536554
if finished { return nil }
537555

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

543562
guard step < maxTokens else {
@@ -572,6 +591,13 @@ extension CoreAISequentialEngine.GenerationSequence {
572591
logitBuffer = lastLogits
573592
}
574593

594+
// Check cancellation after inference step
595+
if generationToken.isCancelled {
596+
stopReasonStore.set(.cancelled)
597+
finishAndRelease()
598+
return nil
599+
}
600+
575601
let nextToken: Int32
576602
if let forced = forcedContinuation {
577603
nextToken = forced[step]
@@ -603,10 +629,7 @@ extension CoreAISequentialEngine.GenerationSequence {
603629
return
604630
}
605631
finished = true
606-
if didAcquireLock {
607-
engine.generating.withLock { $0 = false }
608-
didAcquireLock = false
609-
}
632+
engine.clearTokenIfActive(generationToken)
610633
}
611634
}
612635
}

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

Lines changed: 54 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,17 @@ 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 activeGenerationToken: GenerationToken? { _activeToken.withLock { $0 } }
55+
public var isBusy: Bool { _activeToken.withLock { $0 != nil } }
56+
57+
/// Clear the engine's active token if it matches the given token.
58+
func clearTokenIfActive(_ token: GenerationToken) {
59+
_activeToken.withLock { if $0 === token { $0 = nil } }
60+
}
61+
5062
// MARK: - Initialization
5163

5264
public init(configuration: ModelConfig, preparedModel: PreparedModel) async throws {
@@ -317,11 +329,14 @@ public final class StaticShapeEngine: InferenceEngine, @unchecked Sendable {
317329
samplingConfiguration: SamplingConfiguration,
318330
inferenceOptions: InferenceOptions
319331
) throws -> GenerationSequence {
320-
GenerationSequence(
332+
let token = GenerationToken()
333+
_activeToken.withLock { $0 = token }
334+
return GenerationSequence(
321335
engine: self,
322336
input: input,
323337
samplingConfiguration: samplingConfiguration,
324-
inferenceOptions: inferenceOptions
338+
inferenceOptions: inferenceOptions,
339+
generationToken: token
325340
)
326341
}
327342

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

532547
// MARK: - Lifecycle
533548

549+
public func cancel() async throws {
550+
_activeToken.withLock { $0?.cancel(); $0 = nil }
551+
}
552+
534553
public func reset() {
554+
_activeToken.withLock { $0?.cancel(); $0 = nil }
535555
let resetSpan = InstrumentsProfiler.beginReset(engine: "StaticShape")
536556
processedTokenCount = 0
537557
resetSpan.end()
@@ -559,6 +579,7 @@ extension StaticShapeEngine {
559579
let input: [TokenId]
560580
let samplingConfiguration: SamplingConfiguration
561581
let inferenceOptions: InferenceOptions
582+
let generationToken: GenerationToken
562583

563584
/// Shared with the iterator so the caller can read why generation ended.
564585
let stopReasonStore = StopReasonStore()
@@ -575,7 +596,8 @@ extension StaticShapeEngine {
575596
input: input,
576597
samplingConfiguration: samplingConfiguration,
577598
inferenceOptions: inferenceOptions,
578-
stopReasonStore: stopReasonStore
599+
stopReasonStore: stopReasonStore,
600+
generationToken: generationToken
579601
)
580602
}
581603
}
@@ -592,22 +614,26 @@ extension StaticShapeEngine.GenerationSequence {
592614
private let forcedContinuation: [StaticShapeEngine.TokenId]?
593615
private let maxTokens: Int
594616
private let stopReasonStore: StopReasonStore
617+
private let generationToken: GenerationToken
595618

596619
private var inputTokens: [StaticShapeEngine.TokenId]
597620
private var step: Int = 0
621+
private var finished: Bool = false
598622

599623
init(
600624
engine: StaticShapeEngine,
601625
input: [StaticShapeEngine.TokenId],
602626
samplingConfiguration: SamplingConfiguration,
603627
inferenceOptions: InferenceOptions,
604-
stopReasonStore: StopReasonStore
628+
stopReasonStore: StopReasonStore,
629+
generationToken: GenerationToken
605630
) {
606631
self.engine = engine
607632
self.samplingConfiguration = samplingConfiguration
608633
self.returnsLogits = inferenceOptions.includeLogits
609634
self.forcedContinuation = inferenceOptions.forcedContinuation
610635
self.stopReasonStore = stopReasonStore
636+
self.generationToken = generationToken
611637
self.inputTokens = input
612638
if let forced = inferenceOptions.forcedContinuation {
613639
self.maxTokens = forced.count
@@ -620,9 +646,18 @@ extension StaticShapeEngine.GenerationSequence {
620646
}
621647

622648
public mutating func next() async throws -> InferenceOutput? {
649+
if finished { return nil }
650+
651+
if generationToken.isCancelled {
652+
stopReasonStore.set(.cancelled)
653+
finishAndRelease()
654+
return nil
655+
}
656+
623657
guard step < maxTokens else {
624658
// Natural exhaustion. Don't clobber a reason a decoder set (e.g. `.eos`).
625659
stopReasonStore.setIfUnset(.maxTokens)
660+
finishAndRelease()
626661
return nil
627662
}
628663

@@ -637,6 +672,13 @@ extension StaticShapeEngine.GenerationSequence {
637672
returnsLogits: returnsLogits || forcedContinuation != nil
638673
)
639674

675+
// Check cancellation after inference step
676+
if generationToken.isCancelled {
677+
stopReasonStore.set(.cancelled)
678+
finishAndRelease()
679+
return nil
680+
}
681+
640682
let nextToken = forcedContinuation?[step] ?? sampledToken
641683
inputTokens.append(nextToken)
642684
step += 1
@@ -647,11 +689,19 @@ extension StaticShapeEngine.GenerationSequence {
647689
)
648690
} catch is CancellationError {
649691
stopReasonStore.set(.cancelled)
692+
finishAndRelease()
650693
throw CancellationError()
651694
} catch {
652695
stopReasonStore.set(.error)
696+
finishAndRelease()
653697
throw error
654698
}
655699
}
700+
701+
private mutating func finishAndRelease() {
702+
guard !finished else { return }
703+
finished = true
704+
engine.clearTokenIfActive(generationToken)
705+
}
656706
}
657707
}

‎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)