Skip to content

Commit aa3bbf6

Browse files
authored
Add --clear-coreai-cache flag to clear Core AI specialization cache before model load (#127)
Co-authored-by: Tao Jia <tjia1818@users.noreply.github.com>
1 parent 367ad52 commit aa3bbf6

12 files changed

Lines changed: 296 additions & 79 deletions

File tree

Package.swift

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,7 @@ let package = Package(
131131
name: "llm-runner",
132132
dependencies: [
133133
"CoreAILanguageModels",
134+
"CoreAIShared",
134135
.product(name: "ArgumentParser", package: "swift-argument-parser"),
135136
],
136137
path: "swift/Sources/Tools/llm-runner",
@@ -154,6 +155,7 @@ let package = Package(
154155
name: "object-detector",
155156
dependencies: [
156157
"CoreAIObjectDetector",
158+
"CoreAIShared",
157159
.product(name: "ArgumentParser", package: "swift-argument-parser"),
158160
],
159161
path: "swift/Sources/Tools/object-detector",
@@ -165,6 +167,7 @@ let package = Package(
165167
name: "diffusion-runner",
166168
dependencies: [
167169
"CoreAIDiffusionPipeline",
170+
"CoreAIShared",
168171
.product(name: "ArgumentParser", package: "swift-argument-parser"),
169172
],
170173
path: "swift/Sources/Tools/diffusion-runner",
@@ -176,6 +179,7 @@ let package = Package(
176179
name: "speech-runner",
177180
dependencies: [
178181
"CoreAISpeech",
182+
"CoreAIShared",
179183
.product(name: "ArgumentParser", package: "swift-argument-parser"),
180184
],
181185
path: "swift/Sources/Tools/speech-runner",
@@ -213,6 +217,7 @@ let package = Package(
213217
name: "LanguageModelsTests",
214218
dependencies: [
215219
"CoreAILanguageModels",
220+
"CoreAIShared",
216221
"TestUtilities",
217222
.product(name: "Transformers", package: "swift-transformers"),
218223
],

swift/Sources/CoreAILanguageModels/Profiling/PerformanceMetrics.swift

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
// Use of this source code is governed by a BSD-3-clause license that can
44
// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
55

6+
import CoreAIShared
67
import Foundation
78

89
/// A protocol abstracting time measurement for testability.

swift/Sources/CoreAILanguageModels/Profiling/Timing.swift

Lines changed: 0 additions & 34 deletions
This file was deleted.
Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
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 Foundation
7+
8+
// MARK: - Duration Time Conversions
9+
10+
/// Convenience conversions from `Duration` to floating-point time units, shared
11+
/// across the package's command-line tools and profiling code for elapsed-time reporting.
12+
///
13+
/// Example usage:
14+
/// ```swift
15+
/// let start = ContinuousClock.now
16+
/// // ... do work ...
17+
/// let elapsed = (ContinuousClock.now - start).inSeconds
18+
/// print("Elapsed: \(elapsed)s")
19+
/// ```
20+
extension Duration {
21+
/// Duration in seconds as a `Double`.
22+
package var inSeconds: Double {
23+
let (secs, attoseconds) = components
24+
return Double(secs) + Double(attoseconds) / 1e18
25+
}
26+
27+
/// Duration in milliseconds as a `Double`.
28+
package var inMilliseconds: Double {
29+
inSeconds * 1000.0
30+
}
31+
}

swift/Sources/CoreAIShared/Runtime/ModelStructure.swift

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,86 @@ public struct PreparedModel: Sendable {
131131
return url
132132
}
133133

134+
// MARK: - Cache Inspection
135+
136+
/// File extensions that identify a Core AI model asset (source or compiled).
137+
private static let assetExtensions: Set<String> = ["aimodel", "aimodelc"]
138+
139+
/// Enumerates the Core AI model asset(s) reachable from `url`.
140+
///
141+
/// A model directory (e.g. an LLM bundle) contains one or more asset components alongside
142+
/// other files (tokenizer, metadata); we can't assume specific component filenames, so scan
143+
/// the directory for every `.aimodel`/`.aimodelc` entry. If `url` is itself an asset, it is
144+
/// returned as the sole component. This filename-agnostic approach stays correct as new model
145+
/// families add differently-named components.
146+
///
147+
/// - Parameter url: Either a bundle directory containing asset components, or a single asset.
148+
/// - Returns: The asset URLs to operate on, sorted for stable output. Empty only if `url` is a
149+
/// directory with no asset components.
150+
public static func modelAssetURLs(at url: URL) throws -> [URL] {
151+
// A path ending in a known asset extension IS the asset (asset bundles are themselves
152+
// directories, so this must be checked before treating `url` as a container to scan).
153+
if assetExtensions.contains(url.pathExtension) {
154+
return [url]
155+
}
156+
let entries = try FileManager.default.contentsOfDirectory(
157+
at: url,
158+
includingPropertiesForKeys: nil
159+
)
160+
return
161+
entries
162+
.filter { assetExtensions.contains($0.pathExtension) }
163+
.sorted { $0.path < $1.path }
164+
}
165+
166+
/// Clears the Core AI specialization cache for every model asset reachable from `url`,
167+
/// forcing re-specialization on the next load.
168+
///
169+
/// Discovers components via ``modelAssetURLs(at:)`` — pass either a bundle directory or a
170+
/// single asset. Used by CLI tools implementing `--clear-coreai-cache`.
171+
///
172+
/// - Parameter url: A bundle directory containing asset components, or a single asset.
173+
/// - Returns: The asset URLs whose cache entries were cleared.
174+
@discardableResult
175+
public static func clearCache(at url: URL) throws -> [URL] {
176+
let assetURLs = try modelAssetURLs(at: url)
177+
for assetURL in assetURLs {
178+
let coreaiURL = resolveCoreAIModelURL(from: assetURL)
179+
try AIModelCache.default.deleteEntries(for: coreaiURL)
180+
}
181+
return assetURLs
182+
}
183+
184+
/// Reports whether the default Core AI cache already holds a specialized asset for `url`
185+
/// under the given `options`.
186+
///
187+
/// This only inspects the cache via `AIModelCache.model(for:options:)`; it never triggers
188+
/// specialization. Returns `false` if no entry exists, or if an entry exists but fails to load.
189+
///
190+
/// - Important: `options` must match the options the loader will use for `url`, otherwise a
191+
/// real cached specialization won't be found. Callers that load via `prepare(at:)` should use
192+
/// the ``isCached(at:)`` overload; callers that load via `AIModel(contentsOf:)` or a custom
193+
/// `SpecializationOptions` must pass the same value here.
194+
public static func isCached(at url: URL, options: SpecializationOptions) -> Bool {
195+
let coreaiURL = resolveCoreAIModelURL(from: url)
196+
do {
197+
return try AIModelCache.default.model(for: coreaiURL, options: options) != nil
198+
} catch {
199+
return false
200+
}
201+
}
202+
203+
/// Reports whether the default Core AI cache already holds a specialized asset for `url`,
204+
/// using the same structure-derived `SpecializationOptions` that ``prepare(at:)`` uses.
205+
///
206+
/// Use this only for models loaded through ``prepare(at:)``. For other loaders, use
207+
/// ``isCached(at:options:)`` with the matching options.
208+
public static func isCached(at url: URL) -> Bool {
209+
let coreaiURL = resolveCoreAIModelURL(from: url)
210+
let options = probeStructure(at: coreaiURL).specializationOptions
211+
return isCached(at: coreaiURL, options: options)
212+
}
213+
134214
// MARK: - Asset Preparation
135215

136216
/// Prepares a Core AI model asset by loading via `AIModel` and detecting its structure.

swift/Sources/Tools/benchmark/BenchmarkMain.swift

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,9 @@
66
// Based on mlx-lm benchmark (https://github.com/ml-explore/mlx-lm)
77

88
import ArgumentParser
9+
import CoreAI
910
import CoreAILanguageModels
11+
import CoreAIShared
1012
import Foundation
1113

1214
@main
@@ -40,6 +42,12 @@ struct LLMBenchmark: AsyncParsableCommand {
4042
@Option(name: .customLong("output-json"), help: "Write summary JSON to file")
4143
var outputJson: String?
4244

45+
@Flag(
46+
name: .customLong("clear-coreai-cache"),
47+
help: "Clear Core AI cached specialization for this model before loading (forces re-specialization)"
48+
)
49+
var clearCoreAICache: Bool = false
50+
4351
func validate() throws {
4452
if promptTokens < 1 { throw ValidationError("--prompt-tokens must be >= 1") }
4553
if generationTokens < 1 { throw ValidationError("--generation-tokens must be >= 1") }
@@ -57,6 +65,17 @@ struct LLMBenchmark: AsyncParsableCommand {
5765
let bundle = try LanguageBundle(from: model)
5866
let vocabSize = bundle.vocabSize
5967

68+
let modelURL = try bundle.requireModelURL(for: ModelBundle.ComponentKey.main)
69+
70+
if clearCoreAICache {
71+
let cleared = try PreparedModel.clearCache(at: bundle.bundlePath)
72+
print("\n🗑️ Cleared specialization cache for \(bundle.name) (\(cleared.count) component(s))")
73+
}
74+
75+
// Detect an existing cached specialization before loading so we can annotate the load
76+
// time below. This only inspects the cache; it never triggers specialization.
77+
let cacheHit = PreparedModel.isCached(at: modelURL)
78+
6079
let engineConfig = ModelConfig(
6180
name: bundle.name,
6281
tokenizer: bundle.tokenizer,
@@ -66,11 +85,13 @@ struct LLMBenchmark: AsyncParsableCommand {
6685
function: bundle.language.functionMap?.name(for: "main") ?? "main"
6786
)
6887
let configData = try JSONEncoder().encode(engineConfig)
69-
print("\n⏳ Preparing AI asset...")
88+
print("\n⏳ Preparing AI asset...", terminator: "")
89+
fflush(stdout)
7090
let engine = try await EngineFactory.createEngine(
7191
config: configData,
72-
modelURL: try bundle.requireModelURL(for: ModelBundle.ComponentKey.main)
92+
modelURL: modelURL
7393
)
94+
print(cacheHit ? " done (cache hit)" : " done")
7495

7596
let prompt = randomPrompt(vocabSize: vocabSize, count: promptTokens, seed: seed)
7697
let sampling = SamplingConfiguration(temperature: 0)

swift/Sources/Tools/diffusion-runner/DiffusionRunnerMain.swift

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import ArgumentParser
77
import CoreAI
88
import CoreAIDiffusionPipeline
9+
import CoreAIShared
910
import CoreGraphics
1011
import Foundation
1112
import ImageIO
@@ -73,16 +74,27 @@ struct DiffusionRunner: AsyncParsableCommand {
7374
help: "Path to pipeline trace dir — use Python's noise + embeddings instead of generating")
7475
var traceInputsDir: String?
7576

77+
@Flag(
78+
name: .customLong("clear-coreai-cache"),
79+
help: "Clear Core AI cached specialization for this model before loading (forces re-specialization)"
80+
)
81+
var clearCoreAICache: Bool = false
82+
7683
func run() async throws {
77-
let modelURL = URL(fileURLWithPath: model)
84+
let bundleURL = URL(fileURLWithPath: model)
85+
86+
if clearCoreAICache {
87+
let cleared = try PreparedModel.clearCache(at: bundleURL)
88+
print("🗑️ Cleared specialization cache for \(bundleURL.lastPathComponent) (\(cleared.count) component(s))")
89+
}
7890

7991
if let parityDir = parityTestDir {
80-
try await runParityTest(modelURL: modelURL, dataDir: URL(fileURLWithPath: parityDir))
92+
try await runParityTest(modelURL: bundleURL, dataDir: URL(fileURLWithPath: parityDir))
8193
return
8294
}
8395

8496
if let traceDir = traceInputsDir {
85-
try await runWithTraceInputs(modelURL: modelURL, traceDir: URL(fileURLWithPath: traceDir))
97+
try await runWithTraceInputs(modelURL: bundleURL, traceDir: URL(fileURLWithPath: traceDir))
8698
return
8799
}
88100

@@ -96,7 +108,7 @@ struct DiffusionRunner: AsyncParsableCommand {
96108
}
97109

98110
// Determine pipeline type and dispatch
99-
let resolvedDescriptor = try PipelineDescriptor.resolve(at: modelURL, config: configSource)
111+
let resolvedDescriptor = try PipelineDescriptor.resolve(at: bundleURL, config: configSource)
100112
let isFlux2 = resolvedDescriptor.type == .flux2
101113
let isSd3 = resolvedDescriptor.type == .stableDiffusion3
102114

@@ -131,7 +143,7 @@ struct DiffusionRunner: AsyncParsableCommand {
131143
)
132144

133145
if isFlux2 {
134-
let pipeline = try await Flux2Pipeline(from: modelURL, config: configSource, mode: decodeResolution)
146+
let pipeline = try await Flux2Pipeline(from: bundleURL, config: configSource, mode: decodeResolution)
135147

136148
print("Generating (FLUX.2): \"\(prompt)\"")
137149
print("Steps: \(effectiveSteps), Guidance: \(effectiveGuidance), Seed: \(seed)")
@@ -156,7 +168,7 @@ struct DiffusionRunner: AsyncParsableCommand {
156168
try saveImage(image, to: outputURL)
157169
print("Saved: \(output)")
158170
} else if isSd3 {
159-
let pipeline = try await SD3Pipeline(from: modelURL, config: configSource)
171+
let pipeline = try await SD3Pipeline(from: bundleURL, config: configSource)
160172

161173
print("Generating (SD 3.x): \"\(prompt)\"")
162174
print("Steps: \(effectiveSteps), Guidance: \(effectiveGuidance), Seed: \(seed)")
@@ -181,7 +193,7 @@ struct DiffusionRunner: AsyncParsableCommand {
181193
try saveImage(image, to: outputURL)
182194
print("Saved: \(output)")
183195
} else {
184-
let pipeline = try await StableDiffusionPipeline.load(from: modelURL, config: configSource)
196+
let pipeline = try await StableDiffusionPipeline.load(from: bundleURL, config: configSource)
185197

186198
print("Generating: \"\(prompt)\"")
187199
print("Steps: \(effectiveSteps), Guidance: \(effectiveGuidance), Seed: \(seed)")
@@ -791,10 +803,3 @@ struct DiffusionRunner: AsyncParsableCommand {
791803
a.map(abs).max() ?? 0
792804
}
793805
}
794-
795-
extension Duration {
796-
var inSeconds: Double {
797-
let (secs, attoseconds) = self.components
798-
return Double(secs) + Double(attoseconds) / 1e18
799-
}
800-
}

swift/Sources/Tools/image-segmenter/ImageSegmentationRunnerMain.swift

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
55

66
import ArgumentParser
7+
import CoreAI
78
import CoreAIImageSegmenter
89
import CoreAIShared
910
import CoreGraphics
@@ -81,6 +82,12 @@ struct ImageSegmenterCLI: AsyncParsableCommand {
8182
@Flag(name: .long, help: "Print verbose progress information.")
8283
var verbose: Bool = false
8384

85+
@Flag(
86+
name: .customLong("clear-coreai-cache"),
87+
help: "Clear Core AI cached specialization for this model before loading (forces re-specialization)"
88+
)
89+
var clearCoreAICache: Bool = false
90+
8491
@Option(
8592
name: .customLong("parity-test"),
8693
help: """
@@ -160,8 +167,26 @@ struct ImageSegmenterCLI: AsyncParsableCommand {
160167
CLILogger.level = 1
161168
}
162169

170+
let bundle = try ModelBundle(from: model)
171+
let modelURL = try bundle.requireModelURL(for: ModelBundle.ComponentKey.main)
172+
173+
if clearCoreAICache {
174+
let cleared = try PreparedModel.clearCache(at: bundle.bundlePath)
175+
print("🗑️ Cleared specialization cache for \(bundle.name) (\(cleared.count) component(s))")
176+
}
177+
178+
// Detect an existing cached specialization before loading so we can annotate the load time
179+
// below. Only inspects the cache; never specializes. `ImageSegmenter` loads via
180+
// `PreparedModel.prepare`, so the structure-derived options overload matches.
181+
let cacheHit = PreparedModel.isCached(at: modelURL)
182+
163183
if verbose { print("Creating image segmenter...") }
184+
print("⏳ Preparing AI asset...", terminator: "")
185+
fflush(stdout)
186+
let loadStart = ContinuousClock.now
164187
let runner = try await ImageSegmenter(resourcesAt: model)
188+
let loadElapsed = ContinuousClock.now - loadStart
189+
print(" done in \(String(format: "%.3f", loadElapsed.inSeconds))s\(cacheHit ? " (cache hit)" : "")")
165190

166191
let cgImage = try loadCGImage(from: imagePath)
167192
if verbose { print("Loaded image: \(cgImage.width)×\(cgImage.height)") }

0 commit comments

Comments
 (0)