From c6446cf7bfb7cea76408013b614d4b2c530eaa03 Mon Sep 17 00:00:00 2001 From: Joshua Langberg Date: Tue, 15 Sep 2026 18:26:23 -0400 Subject: [PATCH 1/2] Clear the MLX cache on the first generated token (#620) TokenIterator.next() cleared the buffer cache when tokenCount % 256 == 0, but checked after incrementing tokenCount, so the first clear came at token 256. mlx-lm checks before incrementing and also clears on the first token, right after prefill. A generation shorter than 256 tokens never cleared, so repeated short requests grew the cache without bound. Co-authored-by: Joshua Langberg --- Libraries/MLXLMCommon/Evaluate.swift | 6 +-- .../TokenIteratorClearCacheTests.swift | 40 +++++++++++++++++++ 2 files changed, 43 insertions(+), 3 deletions(-) create mode 100644 Tests/MLXLMTests/TokenIteratorClearCacheTests.swift diff --git a/Libraries/MLXLMCommon/Evaluate.swift b/Libraries/MLXLMCommon/Evaluate.swift index 065a458119..160d99d733 100644 --- a/Libraries/MLXLMCommon/Evaluate.swift +++ b/Libraries/MLXLMCommon/Evaluate.swift @@ -949,15 +949,15 @@ public struct TokenIterator: TokenIteratorProtocol { // way in its generation loop. asyncEval([token] + cache.flatMap { $0.state }) - tokenCount += 1 - // Periodically return freed buffers that cannot be reused (odd or // monotonically growing sizes accumulate in the pool otherwise). - // Matches mlx-lm's clear cadence. + // Matches mlx-lm, which also clears on the first token, after prefill. if tokenCount % 256 == 0 { MLX.Memory.clearCache() } + tokenCount += 1 + return previousY.tokens.item(Int.self) } } diff --git a/Tests/MLXLMTests/TokenIteratorClearCacheTests.swift b/Tests/MLXLMTests/TokenIteratorClearCacheTests.swift new file mode 100644 index 0000000000..8368f077cf --- /dev/null +++ b/Tests/MLXLMTests/TokenIteratorClearCacheTests.swift @@ -0,0 +1,40 @@ +// Copyright © 2026 Apple Inc. + +import MLX +import MLXLLM +import MLXLMCommon +import MLXNN +import XCTest + +final class TokenIteratorClearCacheTests: XCTestCase { + + private func makeTinyModel() -> LlamaModel { + let config = LlamaConfiguration( + hiddenSize: 64, hiddenLayers: 4, intermediateSize: 128, + attentionHeads: 8, rmsNormEps: 0.00001, vocabularySize: 100, kvHeads: 4) + let model = LlamaModel(config) + quantize(model: model, groupSize: 64, bits: 4) + eval(model) + return model + } + + // Short generations never reach token 256, so the first token must clear the cache too, + // as mlx-lm does. Otherwise each request leaves its buffers in the cache. + func testFirstTokenClearsBufferCache() throws { + let model = makeTinyModel() + var iterator = try TokenIterator( + input: LMInput(tokens: MLXArray([1, 2, 3, 4, 5])), model: model, + parameters: GenerateParameters(maxTokens: 8)) + + let seeded = 256 * 1024 * 1024 + do { + let buffer = MLXArray.zeros([seeded], dtype: .uint8) + 1 + eval(buffer) + } + XCTAssertGreaterThanOrEqual(Memory.cacheMemory, seeded) + + _ = iterator.next() + + XCTAssertLessThan(Memory.cacheMemory, seeded / 4) + } +} From ee673d6a71d76e67b532dc7eaf91d92edc3bb8bb Mon Sep 17 00:00:00 2001 From: Christopher Webb <207731778+thechriswebb@users.noreply.github.com> Date: Tue, 22 Sep 2026 14:45:38 -0500 Subject: [PATCH 2/2] Extract the model cache and stop whole-cache eviction abandoning an in-flight load (#603) * Move the model cache into its own file - Moves the constraint kind, the tokenizer bias, and the cache actor out of the model file, changing only the access the cross-file use requires: `private actor ModelCache` becomes `actor ModelCache`, because a `private` top-level type is visible only inside its own file. No other content changes. The model cache grew large enough inside MLXLanguageModel.swift to warrant its own file, separate from the model's own responsibilities. * Catch only the fork-unavailable error when a constraint clone fails - Narrows the freshly-compiled constraint's clone attempt from a blanket `try?` to a catch on `GrammarError.forkFailed`, matching the cached-template clone site above it, which already narrowed this way. Every other clone failure now reaches the caller instead of being silently treated as the expected one. A blanket `try?` treated every clone failure as the one expected case, xgrammar lacking `Fork()`. `GrammarConstraint.clone()` has exactly one throw site today, so this narrowing changes no observable behavior yet; it only stops a future second failure mode from being swallowed the same way. * Cancel every registered load from evictAll() - Cancels every registered load task, which remove(modelID:) already did for one model. - Corrects the eviction documentation, which claimed an in-flight load completes. - Adds a test that evictAll() cancels an in-flight load, and the same probe against evict(). - Bounds both new tests with a time limit, because each parks its loader and a regression would hang. evictAll() dropped every load registration without cancelling, so nothing in the process could reach a load that was still running. A caller awaiting preload() or respond() now fails, and callers coalesced onto one load fail together. Cancellation is cooperative. A load already inside loadWeights runs to completion and allocates, and that limit applies to evict() too. * Rename the grammar tokenizer members away from an xgrammar abbreviation - Renames the grammar tokenizer members and their storage away from an abbreviation of the library name: `xgTokenizers` becomes `grammarTokenizers`, `makeXGTokenizer` becomes `makeGrammarTokenizer`, and `hasCachedXGTokenizer` becomes `hasCachedGrammarTokenizer`, at every call site. - Renames the same local variables at every call site: `xgTokenizer` and `xgTok` become `grammarTokenizer`, and the `vocab` local inside `makeGrammarTokenizer` becomes `vocabulary` to avoid shadowing the `vocab:` parameter label once its neighbor is renamed. The cache's tokenizer members were named after xgrammar's own abbreviation rather than the type they store, which the style guide's initialism rule flags. * Name the download progress throughput property in full - Renames the observable throughput property to throughputBytesPerSecond, dropping the abbreviation. - Drops public from the two reporting entry points, leaving the observable state and the shared instance public. Every caller of the two entry points is in this module, and the tests reach them through @testable import, so neither needs to be public. No release contains this type: the newest tag is 3.31.4 from June, and the whole MLXFoundationModels library arrived in July, so no consumer can depend on either name. * Track download progress for each model - Keys the progress state on the model ID, with one published value per model and accessors that add every entry. - Reports the end of a load from a defer, so a failed or cancelled load clears its own entry. - Sends every report through one stream and one main-actor consumer, which applies them in the order the loader sent them. - Gives every load its own identifier, so a report from a cancelled load cannot change the load that replaced it. - Removes isActive and modelName, and adds activeDownloads and download(forModelID:). One fraction and one identifier cannot describe two loads, and the cache loads models concurrently, so two downloads overwrote each other's state. reportCompleted() ran only after a successful load, so a cancelled load left the state active for the life of the process. Each report also created its own main-actor task, and those tasks run in no fixed order, so a progress report could apply after the end and leave a finished download on screen. Eviction cancels a load and forgets it at once, so a replacement can start while the cancelled load still transfers bytes, and both shared one entry. No release contains this type, so none of these changes can break a consumer. --- .../GenerableRoundTripTests.swift | 6 +- .../Grammar/EmitStopSignalTests.swift | 6 +- ...ForwardTokenizationDisagreementTests.swift | 4 +- .../GuidedGenerationBenchmarkTests.swift | 14 +- .../GuidedGenerationTests.swift | 12 +- .../HardReserveStressTests.swift | 6 +- .../MaxTokenTruncationTests.swift | 12 +- .../MultiModelGuidedGenerationTests.swift | 16 +- .../PrewarmGrammarTests.swift | 4 +- .../MLXDownloadProgress.swift | 322 +++++++++----- .../MLXLanguageModel.swift | 340 ++------------- .../MLXFoundationModels/ModelCache.swift | 298 +++++++++++++ .../MLXDownloadProgressTests.swift | 405 ++++++++++++++++++ .../ModelCacheEvictionTests.swift | 104 +++++ 14 files changed, 1111 insertions(+), 438 deletions(-) create mode 100644 Libraries/MLXFoundationModels/ModelCache.swift create mode 100644 Tests/MLXFoundationModelsTests/MLXDownloadProgressTests.swift diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GenerableRoundTripTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GenerableRoundTripTests.swift index 8d6663d932..20c58ec4a0 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GenerableRoundTripTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GenerableRoundTripTests.swift @@ -280,13 +280,13 @@ struct GenerableRoundTripTests { kvBits: Int? = nil ) async throws -> String { try await container.perform { context in - let xgTokenizer = try await MLXLanguageModel.makeXGTokenizer( + let grammarTokenizer = try await MLXLanguageModel.makeGrammarTokenizer( modelID: modelID, tokenizer: context.tokenizer ) let constraint = try GrammarConstraint( - tokenizer: xgTokenizer, + tokenizer: grammarTokenizer, jsonSchema: jsonSchema, fastForward: true, hostTokenizer: context.tokenizer @@ -320,7 +320,7 @@ struct GenerableRoundTripTests { context: context, constraint: constraint, maxTokens: maxTokens, - vocabSize: Int(xgTokenizer.vocabSize), + vocabSize: Int(grammarTokenizer.vocabSize), kvBits: kvBits, completionReserve: completionReserve, hardReserve: hardReserve, diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/EmitStopSignalTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/EmitStopSignalTests.swift index 16502f0549..e5b6a04de7 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/EmitStopSignalTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/EmitStopSignalTests.swift @@ -53,12 +53,12 @@ struct EmitStopSignalTests { let container = try await loadTestModelContainer(id: TestFixtures.defaultModelID) try await container.perform { context in - let xgTokenizer = try await MLXLanguageModel.makeXGTokenizer( + let grammarTokenizer = try await MLXLanguageModel.makeGrammarTokenizer( modelID: TestFixtures.defaultModelID, tokenizer: context.tokenizer ) let constraint = try GrammarConstraint( - tokenizer: xgTokenizer, + tokenizer: grammarTokenizer, jsonSchema: schema, fastForward: true, hostTokenizer: context.tokenizer @@ -83,7 +83,7 @@ struct EmitStopSignalTests { context: context, constraint: constraint, maxTokens: 128, - vocabSize: Int(xgTokenizer.vocabSize) + vocabSize: Int(grammarTokenizer.vocabSize) ) { _ in callCount += 1 if firstFalseAt != nil { diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/FastForwardTokenizationDisagreementTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/FastForwardTokenizationDisagreementTests.swift index d2e79aab82..195be23410 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/FastForwardTokenizationDisagreementTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/Grammar/FastForwardTokenizationDisagreementTests.swift @@ -162,7 +162,7 @@ struct FastForwardTokenizationDisagreementTests { } try await qwenContainer.perform { qwenContext in - let xgTokenizer = try GrammarTokenizer( + let grammarTokenizer = try GrammarTokenizer( vocab: seeds.vocab, vocabType: seeds.vocabType, eosTokenId: seeds.eosTokenId @@ -173,7 +173,7 @@ struct FastForwardTokenizationDisagreementTests { // Gemma-bound matcher does not have in its current mask. let grammar = "root ::= \"\(Self.forcedPayload)\"\n" let constraint = try GrammarConstraint( - tokenizer: xgTokenizer, + tokenizer: grammarTokenizer, grammar: grammar, fastForward: true, hostTokenizer: qwenContext.tokenizer diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationBenchmarkTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationBenchmarkTests.swift index a258f668df..a042d065e1 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationBenchmarkTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationBenchmarkTests.swift @@ -322,13 +322,13 @@ struct GuidedGenerationBenchmarkTests { let container = try await loadTestModelContainer(id: TestFixtures.defaultModelID) let modelID = TestFixtures.defaultModelID - let (xgTokenizer, hostTokenizer): (GrammarTokenizer, any Tokenizer) = + let (grammarTokenizer, hostTokenizer): (GrammarTokenizer, any Tokenizer) = try await container.perform { context in - let xg = try await MLXLanguageModel.makeXGTokenizer( + let tokenizer = try await MLXLanguageModel.makeGrammarTokenizer( modelID: modelID, tokenizer: context.tokenizer ) - return (xg, context.tokenizer) + return (tokenizer, context.tokenizer) } let schema = """ @@ -349,7 +349,7 @@ struct GuidedGenerationBenchmarkTests { for _ in 0 ..< iterations { let start = ContinuousClock.now let constraint = try GrammarConstraint( - tokenizer: xgTokenizer, + tokenizer: grammarTokenizer, jsonSchema: schema, fastForward: true, hostTokenizer: hostTokenizer @@ -447,12 +447,12 @@ struct GuidedGenerationBenchmarkTests { modelID: String = TestFixtures.defaultModelID ) async throws -> RunResult { try await container.perform { context in - let xgTokenizer = try await MLXLanguageModel.makeXGTokenizer( + let grammarTokenizer = try await MLXLanguageModel.makeGrammarTokenizer( modelID: modelID, tokenizer: context.tokenizer ) let constraint = try GrammarConstraint( - tokenizer: xgTokenizer, + tokenizer: grammarTokenizer, jsonSchema: Self.benchmarkSchema, fastForward: true, hostTokenizer: context.tokenizer @@ -472,7 +472,7 @@ struct GuidedGenerationBenchmarkTests { context: context, constraint: constraint, maxTokens: Self.benchmarkMaxTokens, - vocabSize: Int(xgTokenizer.vocabSize) + vocabSize: Int(grammarTokenizer.vocabSize) ) { text in charCount += text.count deltaCount += 1 diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationTests.swift index 78af325940..82060f8db8 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/GuidedGenerationTests.swift @@ -324,12 +324,12 @@ struct GuidedGenerationTests { let container = try await loadTestModelContainer(id: modelID) let raw: String = try await container.perform { context in - let xgTokenizer = try await MLXLanguageModel.makeXGTokenizer( + let grammarTokenizer = try await MLXLanguageModel.makeGrammarTokenizer( modelID: modelID, tokenizer: context.tokenizer ) let constraint = try GrammarConstraint( - tokenizer: xgTokenizer, + tokenizer: grammarTokenizer, jsonSchema: schema, fastForward: true, hostTokenizer: context.tokenizer @@ -364,7 +364,7 @@ struct GuidedGenerationTests { context: context, constraint: constraint, maxTokens: maxTokens, - vocabSize: Int(xgTokenizer.vocabSize), + vocabSize: Int(grammarTokenizer.vocabSize), completionReserve: reserve, hardReserve: hardReserve, closingBias: closingBias, @@ -490,12 +490,12 @@ struct GuidedGenerationTests { let schema = Self.unboundedSchema let raw: String = try await container.perform { context in - let xgTokenizer = try await MLXLanguageModel.makeXGTokenizer( + let grammarTokenizer = try await MLXLanguageModel.makeGrammarTokenizer( modelID: modelID, tokenizer: context.tokenizer ) let constraint = try GrammarConstraint( - tokenizer: xgTokenizer, + tokenizer: grammarTokenizer, jsonSchema: schema, fastForward: true, hostTokenizer: context.tokenizer @@ -536,7 +536,7 @@ struct GuidedGenerationTests { context: context, constraint: constraint, maxTokens: 256, - vocabSize: Int(xgTokenizer.vocabSize), + vocabSize: Int(grammarTokenizer.vocabSize), completionReserve: reserve, hardReserve: hardReserve, closingBias: closingBias, diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/HardReserveStressTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/HardReserveStressTests.swift index 40f7691e13..89dfbf2f5c 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/HardReserveStressTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/HardReserveStressTests.swift @@ -158,12 +158,12 @@ struct HardReserveStressTests { let container = try await loadTestModelContainer(id: modelID) let raw: String = try await container.perform { context in - let xgTokenizer = try await MLXLanguageModel.makeXGTokenizer( + let grammarTokenizer = try await MLXLanguageModel.makeGrammarTokenizer( modelID: modelID, tokenizer: context.tokenizer ) let constraint = try GrammarConstraint( - tokenizer: xgTokenizer, + tokenizer: grammarTokenizer, jsonSchema: schema, fastForward: true, hostTokenizer: context.tokenizer @@ -208,7 +208,7 @@ struct HardReserveStressTests { context: context, constraint: constraint, maxTokens: maxTokens, - vocabSize: Int(xgTokenizer.vocabSize), + vocabSize: Int(grammarTokenizer.vocabSize), completionReserve: softReserve, hardReserve: hardReserve, closingBias: closingBias, diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/MaxTokenTruncationTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/MaxTokenTruncationTests.swift index 05cd656d82..94ef012297 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/MaxTokenTruncationTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/MaxTokenTruncationTests.swift @@ -43,13 +43,13 @@ struct MaxTokenTruncationTests { } """ - let xgTokenizer = try await MLXLanguageModel.makeXGTokenizer( + let grammarTokenizer = try await MLXLanguageModel.makeGrammarTokenizer( modelID: TestFixtures.defaultModelID, tokenizer: context.tokenizer ) let constraint = try GrammarConstraint( - tokenizer: xgTokenizer, + tokenizer: grammarTokenizer, jsonSchema: complexSchema, fastForward: true, hostTokenizer: context.tokenizer @@ -68,7 +68,7 @@ struct MaxTokenTruncationTests { context: context, constraint: constraint, maxTokens: 5, - vocabSize: Int(xgTokenizer.vocabSize) + vocabSize: Int(grammarTokenizer.vocabSize) ) { _ in true } } } @@ -122,7 +122,7 @@ struct MaxTokenTruncationTests { let container = try await loadTestModelContainer(id: TestFixtures.defaultModelID) try await container.perform { context in - let xgTokenizer = try await MLXLanguageModel.makeXGTokenizer( + let grammarTokenizer = try await MLXLanguageModel.makeGrammarTokenizer( modelID: TestFixtures.defaultModelID, tokenizer: context.tokenizer ) @@ -138,7 +138,7 @@ struct MaxTokenTruncationTests { """ let constraint = try GrammarConstraint( - tokenizer: xgTokenizer, + tokenizer: grammarTokenizer, jsonSchema: arraySchema, fastForward: true, hostTokenizer: context.tokenizer @@ -157,7 +157,7 @@ struct MaxTokenTruncationTests { context: context, constraint: constraint, maxTokens: 3, - vocabSize: Int(xgTokenizer.vocabSize) + vocabSize: Int(grammarTokenizer.vocabSize) ) { _ in true } } } diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/MultiModelGuidedGenerationTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/MultiModelGuidedGenerationTests.swift index 1fa8947a37..738d1b3de2 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/MultiModelGuidedGenerationTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/MultiModelGuidedGenerationTests.swift @@ -120,12 +120,12 @@ struct MultiModelGuidedGenerationTests { """ let raw: String = try await container.perform { context in - let xgTokenizer = try await MLXLanguageModel.makeXGTokenizer( + let grammarTokenizer = try await MLXLanguageModel.makeGrammarTokenizer( modelID: modelID, tokenizer: context.tokenizer ) let constraint = try GrammarConstraint( - tokenizer: xgTokenizer, + tokenizer: grammarTokenizer, jsonSchema: schema, fastForward: true, hostTokenizer: context.tokenizer @@ -155,7 +155,7 @@ struct MultiModelGuidedGenerationTests { context: context, constraint: constraint, maxTokens: 1024, - vocabSize: Int(xgTokenizer.vocabSize), + vocabSize: Int(grammarTokenizer.vocabSize), completionReserve: reserve, closingBias: closingBias, whitespaceBias: whitespaceBias, @@ -246,12 +246,12 @@ struct MultiModelGuidedGenerationTests { """ let raw: String = try await container.perform { context in - let xgTokenizer = try await MLXLanguageModel.makeXGTokenizer( + let grammarTokenizer = try await MLXLanguageModel.makeGrammarTokenizer( modelID: modelID, tokenizer: context.tokenizer ) let constraint = try GrammarConstraint( - tokenizer: xgTokenizer, + tokenizer: grammarTokenizer, jsonSchema: schema, fastForward: true, hostTokenizer: context.tokenizer @@ -284,7 +284,7 @@ struct MultiModelGuidedGenerationTests { context: context, constraint: constraint, maxTokens: 4096, - vocabSize: Int(xgTokenizer.vocabSize), + vocabSize: Int(grammarTokenizer.vocabSize), completionReserve: reserve, closingBias: closingBias, whitespaceBias: whitespaceBias, @@ -438,12 +438,12 @@ struct MultiModelGuidedGenerationTests { let modelID = Self.gemmaModelID let container = try await loadTestModelContainer(id: modelID) try await container.perform { context in - let xgTokenizer = try await MLXLanguageModel.makeXGTokenizer( + let grammarTokenizer = try await MLXLanguageModel.makeGrammarTokenizer( modelID: modelID, tokenizer: context.tokenizer ) let constraint = try GrammarConstraint( - tokenizer: xgTokenizer, + tokenizer: grammarTokenizer, jsonSchema: schema, fastForward: true, hostTokenizer: context.tokenizer diff --git a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/PrewarmGrammarTests.swift b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/PrewarmGrammarTests.swift index 50226193e8..8a6a224d80 100644 --- a/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/PrewarmGrammarTests.swift +++ b/IntegrationTesting/IntegrationTestingTests/MLXFoundationModelsIntegration/GuidedGeneration/PrewarmGrammarTests.swift @@ -13,7 +13,7 @@ import FoundationModels struct PrewarmGrammarTests { @Test - func prewarmCreatesXGTokenizer() async throws { + func prewarmCreatesGrammarTokenizer() async throws { guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } let model = makeTestModel(TestFixtures.defaultModelID) let executor = try makeMLXExecutor(for: model) @@ -26,7 +26,7 @@ struct PrewarmGrammarTests { // Assert the genuine cache hit, not merely that a later respond works // (a guided respond succeeds with or without warmup — only the seam // proves warmUp did the pre-creation). - let cached = await MLXLanguageModel.hasCachedXGTokenizer(modelID: model.modelID) + let cached = await MLXLanguageModel.hasCachedGrammarTokenizer(modelID: model.modelID) #expect(cached, "warmUp should pre-create the GrammarTokenizer") // And a guided generation still succeeds end-to-end after warmUp. diff --git a/Libraries/MLXFoundationModels/MLXDownloadProgress.swift b/Libraries/MLXFoundationModels/MLXDownloadProgress.swift index 4affdce89a..f53ef2211d 100644 --- a/Libraries/MLXFoundationModels/MLXDownloadProgress.swift +++ b/Libraries/MLXFoundationModels/MLXDownloadProgress.swift @@ -9,9 +9,36 @@ import Foundation +/// One reported change to a model's load. +/// +/// Carries only `Sendable` primitives, so a producer on any thread builds one +/// without an isolation check. +enum DownloadProgressUpdate: Sendable { + case started(modelID: String, loadID: UUID) + case progress(modelID: String, loadID: UUID, fraction: Double, completed: Int64, total: Int64) + case ended(modelID: String, loadID: UUID) +} + +/// The shared instance's inbox. +/// +/// A producer yields synchronously from whatever thread it is on, and +/// `AsyncStream` delivers in yield order. Order matters for one model's own +/// updates: an end applied before its start would leave the entry stranded. The +/// buffering policy is stated rather than defaulted, because a policy that drops +/// elements would break both the order and the start-to-end pairing. +private let sharedProgressInbox = AsyncStream.makeStream( + bufferingPolicy: .unbounded) + +/// Reads `MLXDownloadProgress.shared` once, which is what starts the drain. +/// Nothing else does, so without this the inbox buffers every update for the +/// life of the process. A global initializer runs lazily and exactly once. +private let startProgressDrainOnce: Void = { + Task { @MainActor in _ = MLXDownloadProgress.shared } +}() + /// Observable download progress for MLX model loading. /// -/// Tracks whether a model is being downloaded/loaded and reports progress. +/// Tracks one entry per model the adapter is downloading, and aggregates them. /// Shared singleton so any view in the app can observe download state. /// /// Usage: @@ -20,8 +47,10 @@ import Foundation /// var downloadProgress = MLXDownloadProgress.shared /// /// var body: some View { -/// if downloadProgress.isActive { -/// ProgressView(value: downloadProgress.fractionCompleted) +/// ForEach(downloadProgress.activeDownloads) { download in +/// ProgressView(value: download.fractionCompleted) { +/// Text(download.modelID) +/// } /// } /// } /// } @@ -30,124 +59,215 @@ import Foundation @Observable public final class MLXDownloadProgress { - /// Shared singleton instance. - public static let shared = MLXDownloadProgress() + /// One model's download progress. + /// + /// A value, so a consumer reads it once and holds it across an actor hop. + public struct Download: Sendable, Identifiable { - /// Whether a model is currently being downloaded or loaded. - public private(set) var isActive = false + /// The model identifier, as `ModelConfiguration.name` gives it. A + /// repository ID or a path fragment, never a display name. + public let modelID: String - /// Download progress from 0.0 to 1.0. - public private(set) var fractionCompleted: Double = 0 + /// The model identifier, which is what distinguishes one download here. + public var id: String { modelID } - /// The model identifier being downloaded, if any. - public private(set) var modelName: String? + /// Progress from 0.0 to 1.0. + public let fractionCompleted: Double - /// When the current download started. nil when inactive. - /// Consumers can compute elapsed time as `Date.now.timeIntervalSince(startedAt)`. - public private(set) var startedAt: Date? + /// Bytes downloaded so far, from `Progress.completedUnitCount`. + public let completedBytes: Int64 - /// Bytes downloaded so far for the current download. Derived from the - /// underlying `Progress.completedUnitCount`. - public private(set) var completedBytes: Int64 = 0 + /// Total bytes, from `Progress.totalUnitCount`. + public let totalBytes: Int64 - /// Total bytes for the current download. Derived from the underlying - /// `Progress.totalUnitCount`. May be 0 before the first progress report. - public private(set) var totalBytes: Int64 = 0 + /// When the load now running first reported progress for this model. + public let startedAt: Date + + /// Rolling average throughput over the last five seconds of samples. nil + /// until two samples span more than a tenth of a second. + public let throughputBytesPerSecond: Double? + } + + /// Shared singleton instance, and the only instance that drains + /// `sharedProgressInbox`. + public static let shared: MLXDownloadProgress = { + let progress = MLXDownloadProgress() + progress.consume(sharedProgressInbox.stream) + return progress + }() + + /// One entry per model being loaded. An entry whose `download` is nil has + /// started and has not reported yet, so it is not a download. + private var states: [String: State] = [:] + + private struct State { + /// Which load owns this entry. Eviction cancels a load and forgets it at + /// once, so a replacement load can start while the cancelled one is still + /// unwinding a network read. Every update the cancelled load still sends + /// carries its own identifier, and this is what rejects it. + let loadID: UUID + var download: Download? + var samples: [(time: Date, bytes: Int64)] = [] + } - /// Rolling average throughput in bytes per second, computed over the - /// most recent ~5 seconds of progress samples. nil until we have at - /// least two samples spanning a meaningful window. + /// Width of the throughput window. Short enough that a stall shows within a + /// few seconds, long enough to smooth the jitter in chunk arrivals. + private static let throughputWindow: TimeInterval = 5 + + /// Internal, so a test builds an isolated instance and calls ``apply(_:)`` on + /// it. An instance other than ``shared`` drains no inbox, so nothing reaches + /// it except a direct ``apply(_:)`` call. + init() {} + + // MARK: - Reading + + /// Unsorted, for the aggregates. ``activeDownloads`` is the published order. + private var reportedDownloads: [Download] { + states.values.compactMap(\.download) + } + + /// Every model downloading now, earliest first. Empty when none is. + public var activeDownloads: [Download] { + reportedDownloads.sorted { ($0.startedAt, $0.modelID) < ($1.startedAt, $1.modelID) } + } + + /// This model's download, or nil when it is not downloading or has not + /// reported progress yet. + public func download(forModelID modelID: String) -> Download? { + states[modelID]?.download + } + + /// Combined progress across every downloading model, weighted by bytes. 0 + /// before the first report carries a total. + public var fractionCompleted: Double { + let total = totalBytes + guard total > 0 else { return 0 } + return Double(completedBytes) / Double(total) + } + + /// Bytes downloaded so far, across every downloading model. + public var completedBytes: Int64 { + reportedDownloads.reduce(0) { $0 + $1.completedBytes } + } + + /// Total bytes, across every downloading model. + public var totalBytes: Int64 { + reportedDownloads.reduce(0) { $0 + $1.totalBytes } + } + + /// When the earliest download still running first reported, and nil when none + /// is running. Consumers can compute elapsed time as + /// `Date.now.timeIntervalSince(startedAt)`. + public var startedAt: Date? { + reportedDownloads.map(\.startedAt).min() + } + + /// Combined throughput across every downloading model, and nil when no model + /// has enough samples yet. /// - /// Rolling (not cumulative) so a stall shows up immediately as the - /// number dropping toward 0 -- consumers can show "still moving" vs - /// "stuck" without needing a separate indicator. - public private(set) var throughputBytesPerSec: Double? - - /// Width of the throughput rolling window. Short enough that stalls - /// are visible within a few seconds; long enough to smooth out the - /// natural jitter in HF chunk arrivals. - private let throughputWindow: TimeInterval = 5.0 - - /// Samples used to compute rolling throughput. Pruned to - /// `throughputWindow` on every `reportProgress` call. - private var samples: [(time: Date, bytes: Int64)] = [] - - private init() {} - - /// Nonisolated entry point for `reportProgress` so callers from sendable - /// closures (e.g. the cache loader's `progressHandler`) don't have to - /// hop to the main actor just to read `.shared`. The instance method is - /// already nonisolated; this shim only forwards. - nonisolated public static func report(progress: Progress, modelID: String) { - Task { @MainActor in - shared.reportProgress(progress, modelID: modelID) - } + /// Rolling (not cumulative) so a stall shows up immediately as the number + /// dropping toward 0 -- consumers can show "still moving" vs "stuck" without + /// needing a separate indicator. + public var throughputBytesPerSecond: Double? { + let rates = reportedDownloads.compactMap(\.throughputBytesPerSecond) + return rates.isEmpty ? nil : rates.reduce(0, +) } - /// Nonisolated entry point for `reportCompleted`. Same rationale as - /// ``report(progress:modelID:)``. - nonisolated public static func reportCompleted() { - Task { @MainActor in - shared.reportCompleted() - } + // MARK: - Reporting + + /// Reports that a load for `modelID` started. `loadID` identifies this load + /// attempt, and every later call for it must carry the same value. + /// + /// Pair every call with ``reportEnded(modelID:loadID:)``, or that model's + /// progress never clears. + nonisolated static func reportStarted(modelID: String, loadID: UUID) { + _ = startProgressDrainOnce + sharedProgressInbox.continuation.yield(.started(modelID: modelID, loadID: loadID)) + } + + /// Reports one progress change for the load `loadID` identifies. + /// + /// Nonisolated so a producer inside a sendable closure, for example the cache + /// loader's `progressHandler`, reaches the inbox without a hop to the main + /// actor. A hop here would lose the order the reports were made in. + nonisolated static func report(progress: Progress, modelID: String, loadID: UUID) { + _ = startProgressDrainOnce + sharedProgressInbox.continuation.yield( + .progress( + modelID: modelID, + loadID: loadID, + fraction: progress.fractionCompleted, + completed: progress.completedUnitCount, + total: progress.totalUnitCount)) + } + + /// Reports that the load `loadID` identifies ended, whether it succeeded, + /// failed, or was cancelled. + nonisolated static func reportEnded(modelID: String, loadID: UUID) { + _ = startProgressDrainOnce + sharedProgressInbox.continuation.yield(.ended(modelID: modelID, loadID: loadID)) } - nonisolated func reportProgress(_ progress: Progress, modelID: String) { - let fraction = progress.fractionCompleted - // Don't show the progress UI for already-cached models (immediate 100%) - guard fraction < 1.0 else { return } - let completed = progress.completedUnitCount - let total = progress.totalUnitCount + // MARK: - Applying + + /// Drains `updates` for the life of this instance. Only ``shared`` calls it. + /// A second consumer would take updates the shared instance needs. + private func consume(_ updates: AsyncStream) { Task { @MainActor in - if self.startedAt == nil { - self.startedAt = Date() - self.samples.removeAll() + for await update in updates { + self.apply(update) } - self.isActive = true - self.fractionCompleted = fraction - self.modelName = modelID - self.completedBytes = completed - self.totalBytes = total - self.appendSampleAndRecompute(bytes: completed) } } - nonisolated func reportCompleted() { - Task { @MainActor in - self.isActive = false - self.fractionCompleted = 1.0 - self.modelName = nil - self.startedAt = nil - self.completedBytes = 0 - self.totalBytes = 0 - self.throughputBytesPerSec = nil - self.samples.removeAll() + /// Folds one update into the published state, in the order it was yielded. + /// + /// An update from a load that no longer owns its model's entry is dropped, so + /// a load that was cancelled cannot disturb the load that replaced it. + func apply(_ update: DownloadProgressUpdate) { + switch update { + case .started(let modelID, let loadID): + // Replaces any earlier entry, so the newest load owns the model and + // brings its own start time and samples. Nothing is published for it + // until it reports. + states[modelID] = State(loadID: loadID) + + case .progress(let modelID, let loadID, let fraction, let completed, let total): + // A model already on disk reports 1.0 at once. Publishing it would + // flash a full progress bar for one frame. A finishing download's + // final 1.0 is dropped for the same reason, and its entry goes away + // when its load ends. + guard fraction < 1 else { return } + guard var state = states[modelID], state.loadID == loadID else { return } + + let now = Date.now + state.samples.append((time: now, bytes: completed)) + let cutoff = now.addingTimeInterval(-Self.throughputWindow) + state.samples.removeAll { $0.time < cutoff } + state.download = Download( + modelID: modelID, + fractionCompleted: fraction, + completedBytes: completed, + totalBytes: total, + startedAt: state.download?.startedAt ?? now, + throughputBytesPerSecond: Self.throughput(across: state.samples)) + states[modelID] = state + + case .ended(let modelID, let loadID): + guard states[modelID]?.loadID == loadID else { return } + states[modelID] = nil } } - /// Append the latest byte count, prune samples outside the rolling - /// window, and recompute throughput. Requires at least 2 samples - /// spanning a non-trivial time interval to produce a meaningful rate. - private func appendSampleAndRecompute(bytes: Int64) { - let now = Date() - samples.append((time: now, bytes: bytes)) - let cutoff = now.addingTimeInterval(-throughputWindow) - samples.removeAll { $0.time < cutoff } - - guard let oldest = samples.first, - let newest = samples.last, - samples.count >= 2 - else { - throughputBytesPerSec = nil - return - } - let dt = newest.time.timeIntervalSince(oldest.time) - guard dt > 0.1 else { - throughputBytesPerSec = nil - return - } - let db = newest.bytes - oldest.bytes - throughputBytesPerSec = Double(db) / dt + /// Rate across the retained samples. nil below two samples, and nil when they + /// span a tenth of a second or less, because a shorter span gives a + /// meaningless rate. + private static func throughput(across samples: [(time: Date, bytes: Int64)]) -> Double? { + guard samples.count >= 2, let oldest = samples.first, let newest = samples.last + else { return nil } + let seconds = newest.time.timeIntervalSince(oldest.time) + guard seconds > 0.1 else { return nil } + return Double(newest.bytes - oldest.bytes) / seconds } } diff --git a/Libraries/MLXFoundationModels/MLXLanguageModel.swift b/Libraries/MLXFoundationModels/MLXLanguageModel.swift index 1392828172..d58ea466f7 100644 --- a/Libraries/MLXFoundationModels/MLXLanguageModel.swift +++ b/Libraries/MLXFoundationModels/MLXLanguageModel.swift @@ -18,281 +18,6 @@ import MLX import os.log import MLXGuidedGeneration -// MARK: - Constraint Cache Kind - -/// Selects which xgrammar constructor a cached template was compiled -/// with. Used by the constraint cache so a JSON-schema source and a -/// structural-tag source can never alias even if their text collides. -enum ConstraintKind { - case json - case structuralTag -} - -// MARK: - Tokenizer Bias Cache Entry - -/// Tokenizer-derived logit biases, cached per model. Both arrays are pure -/// functions of the tokenizer, so they are identical for a model's lifetime. -/// `@unchecked Sendable`: every field is `let` and read-only after construction -/// (the arrays are only *added* to logits in `GuidedGenerationLoop`, never -/// mutated), and the entry is shared across actors via `ModelCache` — the same -/// pattern as `GrammarTokenizer`/`GrammarConstraint` in `XGrammarBridge.swift`. -final class TokenizerBias: @unchecked Sendable { - let closing: MLXArray - let whitespace: MLXArray - let whitespaceTokenIDs: Set - - init(closing: MLXArray, whitespace: MLXArray, whitespaceTokenIDs: Set) { - self.closing = closing - self.whitespace = whitespace - self.whitespaceTokenIDs = whitespaceTokenIDs - } -} - -// MARK: - Model Cache Actor - -/// Thread-safe model cache using Swift actor isolation. -/// Prevents race conditions when multiple concurrent requests try to load the model. -/// Supports caching multiple models by their identifiers. -private actor ModelCache { - /// Class wrapper around `Task` so actor-reentrancy supersession guards can - /// use `===` identity comparison. `Task` is a value type; a wrapper lets us - /// detect whether `evictAll()` replaced a loading entry mid-flight. - private final class LoadTask { - let task: Task - init(_ task: Task) { self.task = task } - } - - private var containers: [String: ModelContainer] = [:] - private var loadingTasks: [String: LoadTask] = [:] - /// In-flight loads tagged as a warmup of an already-present model, which - /// must NOT surface as `.downloading` (there is no user-facing download). - /// A subset of `loadingTasks`' keys. See `load` and `isDownloading`. - private var suppressedLoadIDs: Set = [] - private var xgTokenizers: [String: GrammarTokenizer] = [:] - /// Cached compiled constraint templates keyed by (modelID, schemaJSON). - /// Clone from template instead of recompiling the grammar each request. - private var constraintTemplates: [String: GrammarConstraint] = [:] - /// Cached per-model logit biases (closing + whitespace). Pure functions of - /// the tokenizer, so computed once per model and reused across requests. - private var tokenizerBiases: [String: TokenizerBias] = [:] - /// Most recent load error per model. Cleared on a subsequent successful - /// load. Surfaced through `MLXLanguageModel.availability` so callers can - /// distinguish "never tried" from "tried and failed". - private var lastErrors: [String: any Error] = [:] - - /// Gets the cached model container for the given model ID, loading it if necessary. - /// Concurrent callers for the same model will share the same loading task, preventing duplicate loads. - /// - /// The `loader` closure carries the transport types (downloader, tokenizer - /// loader). Keeping them out of the cache means the cache itself stays - /// agnostic of how a container is acquired -- first caller wins; later - /// callers reuse the cached container regardless of which loader they - /// brought along. - func load( - modelID: String, - suppressDownloadingState: Bool = false, - loader: @Sendable @escaping () async throws -> ModelContainer - ) async throws -> ModelContainer { - if let cached = containers[modelID] { - return cached - } - - if let existingLoadTask = loadingTasks[modelID] { - // Coalesced onto an in-flight load: the first caller's - // classification (downloading vs. suppressed) stands — we do not - // re-tag. This collision is benign because the suppress decision is - // conditioned on disk-presence: a warmup and a genuine download for - // a not-yet-present model both classify as downloading, so they - // agree; when the model IS present, `availability` resolves to - // `.available` regardless of the in-flight load. - return try await existingLoadTask.task.value - } - - let loadTask = LoadTask( - Task { - try await loader() - }) - loadingTasks[modelID] = loadTask - // Tag a warmup-of-an-already-present model out of the `.downloading` - // signal (computed by the caller as warmup AND modelExistsOnDisk()). - if suppressDownloadingState { - suppressedLoadIDs.insert(modelID) - } - - do { - let loaded = try await loadTask.task.value - // Supersession guard: `evict()`/`evictAll()` may have removed this - // load while it was suspended (actor reentrancy). If we are no longer - // the registered task, hand the awaiter its container but do NOT - // re-populate the cache — ARC frees the weights when the awaiter - // releases it. - guard loadingTasks[modelID] === loadTask else { return loaded } - containers[modelID] = loaded - loadingTasks[modelID] = nil - suppressedLoadIDs.remove(modelID) - lastErrors[modelID] = nil - return loaded - } catch { - // Same guard on the failure path: a superseded load must not re-add a - // stale lastErrors entry for a model nobody holds. - if loadingTasks[modelID] === loadTask { - loadingTasks[modelID] = nil - suppressedLoadIDs.remove(modelID) - lastErrors[modelID] = error - } - throw error - } - } - - /// Whether a *genuine download* is in flight for the given model: a load - /// task is running and it was not tagged as a warmup of an already-present - /// model. Drives `availability`'s `.downloading` state, so a background - /// warmup of an already-downloaded model does not spuriously report - /// `.downloading`. (A warmup that triggers a real fetch is not tagged and - /// does report here.) - func isDownloading(modelID: String) -> Bool { - loadingTasks[modelID] != nil && !suppressedLoadIDs.contains(modelID) - } - - /// The most recent load error for the given model, if a previous attempt - /// failed and no successful load has happened since. - func lastError(modelID: String) -> (any Error)? { - lastErrors[modelID] - } - - /// Gets or creates a cached GrammarTokenizer for the given model. - func makeXGTokenizer( - modelID: String, - tokenizer: any Tokenizer - ) throws -> GrammarTokenizer { - if let cached = xgTokenizers[modelID] { - return cached - } - let vocab = TokenizerVocabExtractor.extractForGrammar(from: tokenizer) - let xgTok = try GrammarTokenizer( - vocab: vocab.vocab, - vocabType: vocab.vocabType, - eosTokenId: Int32(tokenizer.eosTokenId ?? 0) - ) - xgTokenizers[modelID] = xgTok - return xgTok - } - - /// Whether an `GrammarTokenizer` is already cached for the given model. - /// Used by `MLXLanguageModel.hasCachedXGTokenizer` so tests can assert - /// that `warmUp()` pre-created it (a genuine cache hit) rather than only - /// that a later guided respond happens to succeed. - func hasCachedXGTokenizer(modelID: String) -> Bool { - xgTokenizers[modelID] != nil - } - - /// Gets or creates the cached tokenizer-derived logit biases for a model. - func makeTokenizerBias( - modelID: String, - tokenizer: any Tokenizer - ) -> TokenizerBias { - if let cached = tokenizerBiases[modelID] { - return cached - } - let closing = ClosingTokenBias.compute( - tokenizer: tokenizer, - eosTokenId: tokenizer.eosTokenId - ) - let (whitespace, whitespaceTokenIDs) = WhitespaceTokenBias.compute( - tokenizer: tokenizer - ) - let bias = TokenizerBias( - closing: closing, - whitespace: whitespace, - whitespaceTokenIDs: whitespaceTokenIDs - ) - tokenizerBiases[modelID] = bias - return bias - } - - /// Gets a fresh constraint by cloning a cached template, or compiles and caches one first. - /// - /// Grammar compilation is expensive (~5-20ms). By caching the compiled template - /// and cloning it (~0.1ms), repeated requests with the same schema skip recompilation. - /// When Fork() is unavailable (xgrammar < v0.1.34), the clone attempt fails gracefully - /// and each request compiles a fresh constraint instead. - func makeConstraint( - modelID: String, - kind: ConstraintKind, - source: String, - tokenizer: GrammarTokenizer, - hostTokenizer: any Tokenizer, - fastForward: Bool - ) throws -> GrammarConstraint { - let cacheKey = "\(modelID):\(kind):\(source)" - if let template = constraintTemplates[cacheKey] { - do { - return try template.clone() - } catch GrammarError.forkFailed { - constraintTemplates.removeValue(forKey: cacheKey) - } - } - let constraint: GrammarConstraint - switch kind { - case .json: - constraint = try GrammarConstraint( - tokenizer: tokenizer, - jsonSchema: source, - fastForward: fastForward, - hostTokenizer: hostTokenizer - ) - case .structuralTag: - constraint = try GrammarConstraint( - tokenizer: tokenizer, - structuralTag: source, - fastForward: fastForward, - hostTokenizer: hostTokenizer - ) - } - if let cloned = try? constraint.clone() { - constraintTemplates[cacheKey] = constraint - return cloned - } - return constraint - } - - /// Evicts all cached state: model containers, tokenizers, constraint - /// templates, and per-model tokenizer biases. No GPU-stream synchronization - /// is required — in-flight callers retain their own `ModelContainer` and - /// free it via ARC on completion. - func evictAll() { - containers.removeAll() - loadingTasks.removeAll() - suppressedLoadIDs.removeAll() - xgTokenizers.removeAll() - constraintTemplates.removeAll() - tokenizerBiases.removeAll() - lastErrors.removeAll() - } - - /// Evicts a single model's state across every per-model cache: its container, - /// xgrammar tokenizer, all compiled constraint templates, tokenizer bias, - /// last load error, the suppressed-download tag, and any in-flight load - /// registration. - /// Best-effort cancels an in-flight load (the load path is not - /// cancellation-aware today, so this is a no-op safety net); the - /// load-completion guard in `load()` is what prevents a superseded load - /// from re-populating after removal. - func remove(modelID: String) { - // `loadingTasks` holds a `LoadTask` box; cancel the wrapped `Task`. - loadingTasks[modelID]?.task.cancel() - loadingTasks.removeValue(forKey: modelID) - suppressedLoadIDs.remove(modelID) - containers.removeValue(forKey: modelID) - xgTokenizers.removeValue(forKey: modelID) - constraintTemplates = constraintTemplates.filter { - !$0.key.hasPrefix("\(modelID):") - } - tokenizerBiases.removeValue(forKey: modelID) - lastErrors.removeValue(forKey: modelID) - } -} - // MARK: - MLXLanguageModel /// A language model implementation that uses MLX for local inference. @@ -362,6 +87,10 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { /// Loads the model container for a configuration, forwarding download /// progress. Injected so this module carries no HuggingFace or /// swift-transformers dependency; the HuggingFace wiring lives in callers. + /// + /// This closure must honor cancellation while it is suspended. ``evict()`` + /// and ``evictAll()`` cancel the load task tied to this closure. A closure + /// that ignores cancellation keeps running and keeps its caller waiting. public typealias ContainerLoader = @Sendable ( _ configuration: ModelConfiguration, @@ -410,20 +139,26 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { // Configure the buffer pool once per process rather than on every // load, so a consumer's own `Memory.cacheLimit` survives our loads. _ = Self.configureGPUCacheOnce - let container = try await load(configuration) { progress in - MLXDownloadProgress.report(progress: progress, modelID: configuration.name) + // One identifier per load attempt. Eviction cancels a load and + // forgets it at once, so a replacement can start while this one is + // still unwinding, and the progress observable needs to tell them + // apart. + let loadID = UUID() + MLXDownloadProgress.reportStarted(modelID: configuration.name, loadID: loadID) + defer { MLXDownloadProgress.reportEnded(modelID: configuration.name, loadID: loadID) } + return try await load(configuration) { progress in + MLXDownloadProgress.report( + progress: progress, modelID: configuration.name, loadID: loadID) } - MLXDownloadProgress.reportCompleted() - return container } } /// Gets or creates a cached GrammarTokenizer for the given model. - static func makeXGTokenizer( + static func makeGrammarTokenizer( modelID: String, tokenizer: any Tokenizer ) async throws -> GrammarTokenizer { - try await cache.makeXGTokenizer(modelID: modelID, tokenizer: tokenizer) + try await cache.makeGrammarTokenizer(modelID: modelID, tokenizer: tokenizer) } /// Gets the cached per-model tokenizer-derived logit biases (closing + @@ -457,18 +192,27 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { /// Whether the shared cache already holds an `GrammarTokenizer` for the model. /// Internal test seam (not public API): lets `PrewarmGrammarTests` confirm /// `warmUp()` pre-created the tokenizer. - static func hasCachedXGTokenizer(modelID: String) async -> Bool { - await cache.hasCachedXGTokenizer(modelID: modelID) + static func hasCachedGrammarTokenizer(modelID: String) async -> Bool { + await cache.hasCachedGrammarTokenizer(modelID: modelID) } /// Evicts every cached model, tokenizer, constraint template, and per-model /// tokenizer bias, freeing the GPU memory held by model weights. Subsequent - /// requests reload from the on-disk cache. + /// requests reload from the on-disk cache. Cancels every registered + /// in-flight load task, the same way ``evict()`` does for one model. + /// + /// Safe to call during in-flight `respond()`/`warmUp()` work that already holds + /// a `ModelContainer`: that work keeps its own strong reference and synchronizes + /// the GPU on exit, so dropping the cache's reference cannot free weights out + /// from under a live kernel — the weights free via ARC once that work returns. /// - /// Safe to call during in-flight `respond()`/`warmUp()` work: each holds its - /// own strong reference to the `ModelContainer` and synchronizes the GPU on - /// exit, so dropping the cache's reference cannot free weights out from under - /// a live kernel — the weights free via ARC once that work returns. + /// Work still inside its load phase is the other case. This call cancels every + /// registered load task, so if the loader honors cancellation, an awaiting + /// `preload()` or `respond()` fails instead of completing. The loader decides + /// the error type, so it may not be `CancellationError`. The cache coalesces + /// concurrent callers for one model onto a single load task, so one such + /// failure reaches every caller waiting on that load, including callers that + /// never asked for an eviction. public static func evictAll() async { await cache.evictAll() } @@ -479,8 +223,10 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { /// /// Safe to call during an in-flight `respond()`: that call retains its own /// `ModelContainer` and finishes normally; the weights free via ARC once it - /// returns. Evicting a model whose load is still in flight removes it cleanly - /// — the in-flight load completes but does not re-populate the cache. + /// returns. Evicting a model with an in-flight load cancels that load's task. + /// If the loader honors cancellation, the awaiting `preload()` or `respond()` + /// call fails instead of completing. The loader decides the error type, so it + /// may not be `CancellationError`. public func evict() async { await Self.cache.remove(modelID: modelID) } @@ -636,7 +382,7 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { // per-request schema/tool grammar that prewarm doesn't possess — a // pre-built constraint would land under a key no real respond() reads. let tokenizer = await container.tokenizer - _ = try await Self.makeXGTokenizer( + _ = try await Self.makeGrammarTokenizer( modelID: modelID, tokenizer: tokenizer) // Force Metal shader JIT with a minimal 1-token generate, run inside @@ -1263,7 +1009,7 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { tools: requiredToolDefinitions ) - let xgTokenizer = try await MLXLanguageModel.makeXGTokenizer( + let grammarTokenizer = try await MLXLanguageModel.makeGrammarTokenizer( modelID: modelID, tokenizer: context.tokenizer ) @@ -1271,7 +1017,7 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { modelID: modelID, kind: .structuralTag, source: toolCallingGrammar, - tokenizer: xgTokenizer, + tokenizer: grammarTokenizer, hostTokenizer: context.tokenizer, fastForward: true ) @@ -1356,7 +1102,7 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { context: context, constraint: constraint, maxTokens: phase2MaxTokens, - vocabSize: Int(xgTokenizer.vocabSize), + vocabSize: Int(grammarTokenizer.vocabSize), completionReserve: completionReserve, hardReserve: hardReserve, closingBias: closingBias, @@ -1651,14 +1397,14 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { context: ModelContext, channel: LanguageModelExecutorGenerationChannel ) async throws { - let xgTokenizer = try await MLXLanguageModel.makeXGTokenizer( + let grammarTokenizer = try await MLXLanguageModel.makeGrammarTokenizer( modelID: modelID, tokenizer: context.tokenizer) let constraint = try await MLXLanguageModel.makeConstraint( modelID: modelID, kind: .json, source: schemaJSON, - tokenizer: xgTokenizer, + tokenizer: grammarTokenizer, hostTokenizer: context.tokenizer, fastForward: true) let maxTokens = requestedMaxTokens ?? Self.defaultMaxTokens @@ -1690,7 +1436,7 @@ public struct MLXLanguageModel: FoundationModels.LanguageModel, Sendable { context: context, constraint: constraint, maxTokens: maxTokens, - vocabSize: Int(xgTokenizer.vocabSize), + vocabSize: Int(grammarTokenizer.vocabSize), completionReserve: completionReserve, hardReserve: hardReserve, closingBias: bias.closing, diff --git a/Libraries/MLXFoundationModels/ModelCache.swift b/Libraries/MLXFoundationModels/ModelCache.swift new file mode 100644 index 0000000000..161680239f --- /dev/null +++ b/Libraries/MLXFoundationModels/ModelCache.swift @@ -0,0 +1,298 @@ +// Copyright © 2026 Apple Inc. + +#if FoundationModelsIntegration +#if canImport(FoundationModels, _version: 2) + +import Foundation +import MLX +import MLXGuidedGeneration +import MLXLMCommon + +// MARK: - Constraint Cache Kind + +/// Selects which xgrammar constructor a cached template was compiled +/// with. Used by the constraint cache so a JSON-schema source and a +/// structural-tag source can never alias even if their text collides. +enum ConstraintKind { + case json + case structuralTag +} + +// MARK: - Tokenizer Bias Cache Entry + +/// Tokenizer-derived logit biases, cached per model. Both arrays are pure +/// functions of the tokenizer, so they are identical for a model's lifetime. +/// `@unchecked Sendable`: every field is `let` and read-only after construction +/// (the arrays are only *added* to logits in `GuidedGenerationLoop`, never +/// mutated), and the entry is shared across actors via `ModelCache` — the same +/// pattern as `GrammarTokenizer`/`GrammarConstraint` in `XGrammarBridge.swift`. +final class TokenizerBias: @unchecked Sendable { + let closing: MLXArray + let whitespace: MLXArray + let whitespaceTokenIDs: Set + + init(closing: MLXArray, whitespace: MLXArray, whitespaceTokenIDs: Set) { + self.closing = closing + self.whitespace = whitespace + self.whitespaceTokenIDs = whitespaceTokenIDs + } +} + +// MARK: - Model Cache Actor + +/// Thread-safe model cache using Swift actor isolation. +/// Prevents race conditions when multiple concurrent requests try to load the model. +/// Supports caching multiple models by their identifiers. +actor ModelCache { + /// Class wrapper around `Task` so actor-reentrancy supersession guards can + /// use `===` identity comparison. `Task` is a value type; a wrapper lets us + /// detect whether `evictAll()` replaced a loading entry mid-flight. + private final class LoadTask { + let task: Task + init(_ task: Task) { self.task = task } + } + + private var containers: [String: ModelContainer] = [:] + private var loadingTasks: [String: LoadTask] = [:] + /// In-flight loads tagged as a warmup of an already-present model, which + /// must NOT surface as `.downloading` (there is no user-facing download). + /// A subset of `loadingTasks`' keys. See `load` and `isDownloading`. + private var suppressedLoadIDs: Set = [] + private var grammarTokenizers: [String: GrammarTokenizer] = [:] + /// Cached compiled constraint templates keyed by (modelID, schemaJSON). + /// Clone from template instead of recompiling the grammar each request. + private var constraintTemplates: [String: GrammarConstraint] = [:] + /// Cached per-model logit biases (closing + whitespace). Pure functions of + /// the tokenizer, so computed once per model and reused across requests. + private var tokenizerBiases: [String: TokenizerBias] = [:] + /// Most recent load error per model. Cleared on a subsequent successful + /// load. Surfaced through `MLXLanguageModel.availability` so callers can + /// distinguish "never tried" from "tried and failed". + private var lastErrors: [String: any Error] = [:] + + /// Gets the cached model container for the given model ID, loading it if necessary. + /// Concurrent callers for the same model will share the same loading task, preventing duplicate loads. + /// + /// The `loader` closure carries the transport types (downloader, tokenizer + /// loader). Keeping them out of the cache means the cache itself stays + /// agnostic of how a container is acquired -- first caller wins; later + /// callers reuse the cached container regardless of which loader they + /// brought along. + /// + /// The loader must honor cancellation while it is suspended: `evictAll()` and + /// `remove(modelID:)` cancel the registered load task, and a loader that + /// ignores cancellation keeps running and keeps its awaiter waiting. + func load( + modelID: String, + suppressDownloadingState: Bool = false, + loader: @Sendable @escaping () async throws -> ModelContainer + ) async throws -> ModelContainer { + if let cached = containers[modelID] { + return cached + } + + if let existingLoadTask = loadingTasks[modelID] { + // Coalesced onto an in-flight load: the first caller's + // classification (downloading vs. suppressed) stands — we do not + // re-tag. This collision is benign because the suppress decision is + // conditioned on disk-presence: a warmup and a genuine download for + // a not-yet-present model both classify as downloading, so they + // agree; when the model IS present, `availability` resolves to + // `.available` regardless of the in-flight load. + return try await existingLoadTask.task.value + } + + let loadTask = LoadTask( + Task { + try await loader() + }) + loadingTasks[modelID] = loadTask + // Tag a warmup-of-an-already-present model out of the `.downloading` + // signal (computed by the caller as warmup AND modelExistsOnDisk()). + if suppressDownloadingState { + suppressedLoadIDs.insert(modelID) + } + + do { + let loaded = try await loadTask.task.value + // Supersession guard: `evict()`/`evictAll()` may have removed this + // load while it was suspended (actor reentrancy). If we are no longer + // the registered task, hand the awaiter its container but do NOT + // re-populate the cache — ARC frees the weights when the awaiter + // releases it. + guard loadingTasks[modelID] === loadTask else { return loaded } + containers[modelID] = loaded + loadingTasks[modelID] = nil + suppressedLoadIDs.remove(modelID) + lastErrors[modelID] = nil + return loaded + } catch { + // Same guard on the failure path: a superseded load must not re-add a + // stale lastErrors entry for a model nobody holds. + if loadingTasks[modelID] === loadTask { + loadingTasks[modelID] = nil + suppressedLoadIDs.remove(modelID) + lastErrors[modelID] = error + } + throw error + } + } + + /// Whether a *genuine download* is in flight for the given model: a load + /// task is running and it was not tagged as a warmup of an already-present + /// model. Drives `availability`'s `.downloading` state, so a background + /// warmup of an already-downloaded model does not spuriously report + /// `.downloading`. (A warmup that triggers a real fetch is not tagged and + /// does report here.) + func isDownloading(modelID: String) -> Bool { + loadingTasks[modelID] != nil && !suppressedLoadIDs.contains(modelID) + } + + /// The most recent load error for the given model, if a previous attempt + /// failed and no successful load has happened since. + func lastError(modelID: String) -> (any Error)? { + lastErrors[modelID] + } + + /// Gets or creates a cached GrammarTokenizer for the given model. + func makeGrammarTokenizer( + modelID: String, + tokenizer: any Tokenizer + ) throws -> GrammarTokenizer { + if let cached = grammarTokenizers[modelID] { + return cached + } + let vocabulary = TokenizerVocabExtractor.extractForGrammar(from: tokenizer) + let grammarTokenizer = try GrammarTokenizer( + vocab: vocabulary.vocab, + vocabType: vocabulary.vocabType, + eosTokenId: Int32(tokenizer.eosTokenId ?? 0) + ) + grammarTokenizers[modelID] = grammarTokenizer + return grammarTokenizer + } + + /// Whether an `GrammarTokenizer` is already cached for the given model. + /// Used by `MLXLanguageModel.hasCachedGrammarTokenizer` so tests can assert + /// that `warmUp()` pre-created it (a genuine cache hit) rather than only + /// that a later guided respond happens to succeed. + func hasCachedGrammarTokenizer(modelID: String) -> Bool { + grammarTokenizers[modelID] != nil + } + + /// Gets or creates the cached tokenizer-derived logit biases for a model. + func makeTokenizerBias( + modelID: String, + tokenizer: any Tokenizer + ) -> TokenizerBias { + if let cached = tokenizerBiases[modelID] { + return cached + } + let closing = ClosingTokenBias.compute( + tokenizer: tokenizer, + eosTokenId: tokenizer.eosTokenId + ) + let (whitespace, whitespaceTokenIDs) = WhitespaceTokenBias.compute( + tokenizer: tokenizer + ) + let bias = TokenizerBias( + closing: closing, + whitespace: whitespace, + whitespaceTokenIDs: whitespaceTokenIDs + ) + tokenizerBiases[modelID] = bias + return bias + } + + /// Gets a fresh constraint by cloning a cached template, or compiles and caches one first. + /// + /// Grammar compilation is expensive (~5-20ms). By caching the compiled template + /// and cloning it (~0.1ms), repeated requests with the same schema skip recompilation. + /// When Fork() is unavailable (xgrammar < v0.1.34), the clone attempt fails gracefully + /// and each request compiles a fresh constraint instead. Any other clone failure + /// reaches the caller. + func makeConstraint( + modelID: String, + kind: ConstraintKind, + source: String, + tokenizer: GrammarTokenizer, + hostTokenizer: any Tokenizer, + fastForward: Bool + ) throws -> GrammarConstraint { + let cacheKey = "\(modelID):\(kind):\(source)" + if let template = constraintTemplates[cacheKey] { + do { + return try template.clone() + } catch GrammarError.forkFailed { + constraintTemplates.removeValue(forKey: cacheKey) + } + } + let constraint: GrammarConstraint + switch kind { + case .json: + constraint = try GrammarConstraint( + tokenizer: tokenizer, + jsonSchema: source, + fastForward: fastForward, + hostTokenizer: hostTokenizer + ) + case .structuralTag: + constraint = try GrammarConstraint( + tokenizer: tokenizer, + structuralTag: source, + fastForward: fastForward, + hostTokenizer: hostTokenizer + ) + } + do { + let clone = try constraint.clone() + constraintTemplates[cacheKey] = constraint + return clone + } catch GrammarError.forkFailed { + return constraint + } + } + + /// Evicts all cached state: model containers, tokenizers, constraint + /// templates, and per-model tokenizer biases. Cancels every registered load + /// task, as `remove(modelID:)` does for a single model. No GPU-stream + /// synchronization is required — in-flight callers retain their own + /// `ModelContainer` and free it via ARC on completion. + func evictAll() { + containers.removeAll() + for loadTask in loadingTasks.values { + loadTask.task.cancel() + } + loadingTasks.removeAll() + suppressedLoadIDs.removeAll() + grammarTokenizers.removeAll() + constraintTemplates.removeAll() + tokenizerBiases.removeAll() + lastErrors.removeAll() + } + + /// Evicts a single model's state across every per-model cache: its container, + /// xgrammar tokenizer, all compiled constraint templates, tokenizer bias, + /// last load error, the suppressed-download tag, and any in-flight load + /// registration. + /// + /// Cancels an in-flight load and drops its registration. A loader that honors + /// cancellation stops early. The load-completion guard in `load()` is + /// what prevents any superseded load from re-populating after removal. + func remove(modelID: String) { + // `loadingTasks` holds a `LoadTask` box; cancel the wrapped `Task`. + loadingTasks[modelID]?.task.cancel() + loadingTasks.removeValue(forKey: modelID) + suppressedLoadIDs.remove(modelID) + containers.removeValue(forKey: modelID) + grammarTokenizers.removeValue(forKey: modelID) + constraintTemplates = constraintTemplates.filter { + !$0.key.hasPrefix("\(modelID):") + } + tokenizerBiases.removeValue(forKey: modelID) + lastErrors.removeValue(forKey: modelID) + } +} + +#endif // canImport(FoundationModels) +#endif // FoundationModelsIntegration diff --git a/Tests/MLXFoundationModelsTests/MLXDownloadProgressTests.swift b/Tests/MLXFoundationModelsTests/MLXDownloadProgressTests.swift new file mode 100644 index 0000000000..6e7601bacc --- /dev/null +++ b/Tests/MLXFoundationModelsTests/MLXDownloadProgressTests.swift @@ -0,0 +1,405 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import MLX +import MLXLMCommon +import MLXNN +import Testing + +@testable import MLXFoundationModels + +#if FoundationModelsIntegration && canImport(FoundationModels, _version: 2) + +/// Unit tests for the main-actor fold. Each builds its own instance, so none of +/// them reads or writes `MLXDownloadProgress.shared`, and the suite needs no +/// serialization against its neighbors. +@Suite("MLXDownloadProgress state") +struct MLXDownloadProgressStateTests { + + @Test @MainActor + func aReportForAModelThatNeverStartedIsIgnored() { + let progress = MLXDownloadProgress() + progress.apply( + .progress(modelID: "org/a", loadID: UUID(), fraction: 0.5, completed: 5, total: 10)) + + #expect(progress.activeDownloads.isEmpty) + } + + @Test @MainActor + func aStartedModelAppearsOnlyAfterItsFirstReport() { + let progress = MLXDownloadProgress() + let load = UUID() + progress.apply(.started(modelID: "org/a", loadID: load)) + + #expect(progress.activeDownloads.isEmpty) + + progress.apply( + .progress(modelID: "org/a", loadID: load, fraction: 0.5, completed: 5, total: 10)) + + #expect(progress.activeDownloads.map(\.modelID) == ["org/a"]) + #expect(progress.download(forModelID: "org/a")?.completedBytes == 5) + #expect(progress.download(forModelID: "org/a")?.totalBytes == 10) + } + + @Test @MainActor + func aModelAlreadyOnDiskNeverAppears() { + let progress = MLXDownloadProgress() + let load = UUID() + progress.apply(.started(modelID: "org/a", loadID: load)) + progress.apply( + .progress(modelID: "org/a", loadID: load, fraction: 1, completed: 10, total: 10)) + + #expect(progress.activeDownloads.isEmpty) + } + + @Test @MainActor + func twoModelsReportIndependentlyAndTheAggregatesSumThem() { + let progress = MLXDownloadProgress() + let loadA = UUID() + let loadB = UUID() + progress.apply(.started(modelID: "org/a", loadID: loadA)) + progress.apply( + .progress(modelID: "org/a", loadID: loadA, fraction: 0.5, completed: 5, total: 10)) + progress.apply(.started(modelID: "org/b", loadID: loadB)) + progress.apply( + .progress(modelID: "org/b", loadID: loadB, fraction: 0.25, completed: 30, total: 120)) + + #expect(progress.activeDownloads.map(\.modelID) == ["org/a", "org/b"]) + #expect(progress.download(forModelID: "org/a")?.fractionCompleted == 0.5) + #expect(progress.download(forModelID: "org/b")?.fractionCompleted == 0.25) + #expect(progress.completedBytes == 35) + #expect(progress.totalBytes == 130) + #expect(progress.fractionCompleted == 35.0 / 130.0) + } + + @Test @MainActor + func endingOneModelLeavesTheOtherDownloading() { + let progress = MLXDownloadProgress() + let loadA = UUID() + let loadB = UUID() + progress.apply(.started(modelID: "org/a", loadID: loadA)) + progress.apply( + .progress(modelID: "org/a", loadID: loadA, fraction: 0.5, completed: 5, total: 10)) + progress.apply(.started(modelID: "org/b", loadID: loadB)) + progress.apply( + .progress(modelID: "org/b", loadID: loadB, fraction: 0.25, completed: 30, total: 120)) + progress.apply(.ended(modelID: "org/a", loadID: loadA)) + + #expect(progress.activeDownloads.map(\.modelID) == ["org/b"]) + #expect(progress.download(forModelID: "org/a") == nil) + #expect(progress.completedBytes == 30) + } + + @Test @MainActor + func aReportThatArrivesAfterTheEndIsIgnored() { + let progress = MLXDownloadProgress() + let load = UUID() + progress.apply(.started(modelID: "org/a", loadID: load)) + progress.apply( + .progress(modelID: "org/a", loadID: load, fraction: 0.5, completed: 5, total: 10)) + progress.apply(.ended(modelID: "org/a", loadID: load)) + progress.apply( + .progress(modelID: "org/a", loadID: load, fraction: 0.75, completed: 7, total: 10)) + + #expect(progress.activeDownloads.isEmpty) + } + + @Test @MainActor + func aCancelledLoadDoesNotEndTheLoadThatReplacedIt() { + let progress = MLXDownloadProgress() + let cancelled = UUID() + let replacement = UUID() + progress.apply(.started(modelID: "org/a", loadID: cancelled)) + progress.apply( + .progress(modelID: "org/a", loadID: cancelled, fraction: 0.4, completed: 4, total: 10)) + // An eviction cancels the first load and forgets it at once, so the + // replacement starts while the first is still unwinding and reports its + // end afterwards. + progress.apply(.started(modelID: "org/a", loadID: replacement)) + progress.apply(.ended(modelID: "org/a", loadID: cancelled)) + progress.apply( + .progress( + modelID: "org/a", loadID: replacement, fraction: 0.1, completed: 1, total: 10)) + + #expect(progress.download(forModelID: "org/a")?.fractionCompleted == 0.1) + + progress.apply(.ended(modelID: "org/a", loadID: replacement)) + + #expect(progress.activeDownloads.isEmpty) + } + + @Test @MainActor + func aReportFromACancelledLoadCannotMoveTheReplacementBackward() { + let progress = MLXDownloadProgress() + let cancelled = UUID() + let replacement = UUID() + progress.apply(.started(modelID: "org/a", loadID: cancelled)) + progress.apply( + .progress(modelID: "org/a", loadID: cancelled, fraction: 0.4, completed: 4, total: 10)) + progress.apply(.started(modelID: "org/a", loadID: replacement)) + progress.apply( + .progress( + modelID: "org/a", loadID: replacement, fraction: 0.1, completed: 1, total: 10)) + let replacementStart = progress.download(forModelID: "org/a")?.startedAt + // The cancelled load's last report, arriving after its replacement began. + progress.apply( + .progress(modelID: "org/a", loadID: cancelled, fraction: 0.4, completed: 4, total: 10)) + + let download = progress.download(forModelID: "org/a") + #expect(download?.fractionCompleted == 0.1) + #expect(download?.completedBytes == 1) + #expect(download?.startedAt == replacementStart) + // The cancelled load's byte count never enters the replacement's samples, + // so no rate is computed across two loads. + #expect(download?.throughputBytesPerSecond == nil) + } + + @Test @MainActor + func aReplacementLoadPublishesNothingUntilItReports() { + let progress = MLXDownloadProgress() + let cancelled = UUID() + let replacement = UUID() + progress.apply(.started(modelID: "org/a", loadID: cancelled)) + progress.apply( + .progress(modelID: "org/a", loadID: cancelled, fraction: 0.4, completed: 4, total: 10)) + progress.apply(.started(modelID: "org/a", loadID: replacement)) + + #expect(progress.activeDownloads.isEmpty) + #expect(progress.download(forModelID: "org/a") == nil) + } + + @Test @MainActor + func nothingDownloadingReportsEmptyAggregates() { + let progress = MLXDownloadProgress() + + #expect(progress.activeDownloads.isEmpty) + #expect(progress.fractionCompleted == 0) + #expect(progress.completedBytes == 0) + #expect(progress.totalBytes == 0) + #expect(progress.startedAt == nil) + #expect(progress.throughputBytesPerSecond == nil) + } +} + +/// Waits until the shared progress state satisfies `condition`. +/// +/// A producer only hands its update to the inbox, so a test cannot read the state +/// on the line after a load returns. Each turn waits for the next change to what +/// `condition` reads, and never sleeps. A condition that never holds hangs here, +/// which the caller's time limit turns into a failure. +/// +/// Runs on the main actor and does not suspend between testing the condition and +/// installing the tracking, so a change cannot land in that gap and leave this +/// waiting for a change that already happened. +@MainActor +private func waitForSharedProgress(toSatisfy condition: @MainActor () -> Bool) async { + while !condition() { + await withCheckedContinuation { continuation in + withObservationTracking { + _ = condition() + } onChange: { + continuation.resume() + } + } + } +} + +/// Satisfies `LanguageModel`'s requirements and nothing more. A test that only +/// loads a container never generates a token, so these bodies never run. Do not +/// treat this as a model. It has no parameters, and `callAsFunction` returns a +/// fixed zero array whose dimensions mean nothing. +private final class StubLanguageModel: Module, LanguageModel, KVCacheDimensionProvider { + var kvHeads: [Int] { [] } + + func prepare( + _ input: LMInput, cache: [KVCache], state _: LMOutput.State?, prefill _: PrefillParameters + ) throws -> PrepareResult { + .tokens(input.text) + } + + func callAsFunction(_ inputs: MLXArray, cache: [KVCache]?) -> MLXArray { + MLXArray.zeros([1, 1, 1]) + } +} + +/// Satisfies `UserInputProcessor`'s one requirement. A test that only loads a +/// container prepares no input, so this body never runs. +private struct StubInputProcessor: UserInputProcessor { + func prepare(input: UserInput) throws -> LMInput { + LMInput(tokens: MLXArray([0])) + } +} + +/// Holds a loader open, so a test observes the progress a load reported before +/// that load ends. +private actor ProgressLoadGate { + private var reportedAlready = false + private var reportedContinuation: CheckedContinuation? + private var releasedAlready = false + private var releaseContinuation: CheckedContinuation? + + func signalReported() { + reportedAlready = true + reportedContinuation?.resume() + reportedContinuation = nil + } + + func waitUntilReported() async { + if reportedAlready { return } + await withCheckedContinuation { reportedContinuation = $0 } + } + + func waitForRelease() async { + if releasedAlready { return } + await withCheckedContinuation { releaseContinuation = $0 } + } + + func release() { + releasedAlready = true + releaseContinuation?.resume() + releaseContinuation = nil + } +} + +/// Thrown when a parked loader resumes without being cancelled. Only the +/// cancellation handler releases it, so this error means the test's own gate +/// misbehaved rather than the cache. +private struct ParkedLoaderResumedWithoutCancellation: Error {} + +/// Nested under the serialized `FoundationModelsCacheTests` parent, declared in +/// ModelCacheEvictionTests.swift, because these tests drive a real +/// `MLXLanguageModel` load and read the process-global +/// `MLXDownloadProgress.shared`. +extension FoundationModelsCacheTests { + + @Suite("MLXDownloadProgress load lifecycle") + struct DownloadProgressLifecycle { + + struct LoadFailure: Error {} + + @Test( + "a failed load clears the progress state for its model", + .timeLimit(.minutes(1))) + func failedLoadClearsProgress() async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + + let id = "org/throwing-loader-\(UUID().uuidString)" + let gate = ProgressLoadGate() + let model = MLXLanguageModel( + configuration: ModelConfiguration(id: id), + capabilities: [], + weightsLocation: { _ in URL(fileURLWithPath: "/tmp") }, + load: { _, progressHandler in + let fakeProgress = Progress(totalUnitCount: 10) + fakeProgress.completedUnitCount = 5 + progressHandler(fakeProgress) + await gate.signalReported() + await gate.waitForRelease() + throw LoadFailure() + }) + + let loadTask = Task { try await model.loadContainer() } + await gate.waitUntilReported() + await waitForSharedProgress { + MLXDownloadProgress.shared.download(forModelID: id) != nil + } + + await gate.release() + await #expect(throws: LoadFailure.self) { + try await loadTask.value + } + + await waitForSharedProgress { + MLXDownloadProgress.shared.download(forModelID: id) == nil + } + } + + @Test( + "a cancelled load clears the progress state for its model", + .timeLimit(.minutes(1))) + func cancelledLoadClearsProgress() async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + + let id = "org/cancelled-loader-\(UUID().uuidString)" + let gate = ProgressLoadGate() + let model = MLXLanguageModel( + configuration: ModelConfiguration(id: id), + capabilities: [], + weightsLocation: { _ in URL(fileURLWithPath: "/tmp") }, + load: { _, progressHandler in + let fakeProgress = Progress(totalUnitCount: 10) + fakeProgress.completedUnitCount = 5 + progressHandler(fakeProgress) + await withTaskCancellationHandler { + await gate.waitForRelease() + } onCancel: { + // A cancellation handler cannot await, so hop onto the + // gate to release the parked loader. + Task { await gate.release() } + } + try Task.checkCancellation() + throw ParkedLoaderResumedWithoutCancellation() + }) + + let loadTask = Task { try await model.loadContainer() } + await waitForSharedProgress { + MLXDownloadProgress.shared.download(forModelID: id) != nil + } + + await model.evict() + await #expect(throws: CancellationError.self) { + try await loadTask.value + } + + await waitForSharedProgress { + MLXDownloadProgress.shared.download(forModelID: id) == nil + } + } + + @Test( + "a successful load clears the progress state for its model", + .timeLimit(.minutes(1))) + func successfulLoadClearsProgress() async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + + let id = "org/succeeding-loader-\(UUID().uuidString)" + let gate = ProgressLoadGate() + let model = MLXLanguageModel( + configuration: ModelConfiguration(id: id), + capabilities: [], + weightsLocation: { _ in URL(fileURLWithPath: "/tmp") }, + load: { configuration, progressHandler in + let fakeProgress = Progress(totalUnitCount: 10) + fakeProgress.completedUnitCount = 5 + progressHandler(fakeProgress) + await gate.signalReported() + await gate.waitForRelease() + return ModelContainer( + context: ModelContext( + configuration: configuration, + model: StubLanguageModel(), + processor: StubInputProcessor(), + tokenizer: ByteTokenizer())) + }) + + let loadTask = Task { try await model.loadContainer() } + await gate.waitUntilReported() + await waitForSharedProgress { + MLXDownloadProgress.shared.download(forModelID: id) != nil + } + + await gate.release() + _ = try await loadTask.value + + await waitForSharedProgress { + MLXDownloadProgress.shared.download(forModelID: id) == nil + } + + // A successful load caches the container under `id`. Drop it so it + // does not outlive this test in the process-global cache. + await model.evict() + } + } +} + +#endif // FoundationModelsIntegration && canImport(FoundationModels, _version: 2) diff --git a/Tests/MLXFoundationModelsTests/ModelCacheEvictionTests.swift b/Tests/MLXFoundationModelsTests/ModelCacheEvictionTests.swift index 9474f1a4d2..d9e32ca665 100644 --- a/Tests/MLXFoundationModelsTests/ModelCacheEvictionTests.swift +++ b/Tests/MLXFoundationModelsTests/ModelCacheEvictionTests.swift @@ -165,6 +165,54 @@ extension FoundationModelsCacheTests { let lastError = await MLXLanguageModel.lastLoadErrorInCache(modelID: id) #expect(lastError == nil, "a superseded load must not re-populate cache state") } + + @Test( + "evict() cancels the model's registered load task", + .timeLimit(.minutes(1))) + func evictCancelsRegisteredLoad() async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + + let probe = CancellationProbe() + let model = MLXLanguageModel( + configuration: ModelConfiguration(id: "org/cancel-evict-\(UUID().uuidString)"), + capabilities: [], + weightsLocation: { _ in URL(fileURLWithPath: "/no/such/path") }, + load: makeCancellationProbingLoader(probe: probe)) + + let loadTask = Task { try await model.preload() } + await probe.waitUntilParked() + + await model.evict() + + await #expect(throws: CancellationError.self) { + try await loadTask.value + } + #expect(await probe.wasCancelled) + } + + @Test( + "evictAll() cancels every registered load task", + .timeLimit(.minutes(1))) + func evictAllCancelsRegisteredLoads() async throws { + guard #available(iOS 27.0, macOS 27.0, visionOS 27.0, *) else { return } + + let probe = CancellationProbe() + let model = MLXLanguageModel( + configuration: ModelConfiguration(id: "org/cancel-evictall-\(UUID().uuidString)"), + capabilities: [], + weightsLocation: { _ in URL(fileURLWithPath: "/no/such/path") }, + load: makeCancellationProbingLoader(probe: probe)) + + let loadTask = Task { try await model.preload() } + await probe.waitUntilParked() + + await MLXLanguageModel.evictAll() + + await #expect(throws: CancellationError.self) { + try await loadTask.value + } + #expect(await probe.wasCancelled) + } } } @@ -173,6 +221,62 @@ extension FoundationModelsCacheTests { // Hoisted to file scope so the eviction tests above share one definition. Kept // file-`private` so they don't collide with AvailabilityTests.swift's own stubs. +/// Builds a container loader that parks until the cache cancels its load task. +/// The loader never produces a container: the test awaits its failure, so the +/// probe reads back without any timing assumption. +@available(iOS 27.0, macOS 27.0, visionOS 27.0, *) +private func makeCancellationProbingLoader( + probe: CancellationProbe +) -> MLXLanguageModel.ContainerLoader { + { _, _ in + await withTaskCancellationHandler { + await probe.parkUntilCancelled() + } onCancel: { + // A cancellation handler cannot await, so hop onto the probe actor to + // record the cancellation and resume the parked continuation. + Task { await probe.recordCancellation() } + } + try Task.checkCancellation() + throw ParkedLoaderResumedWithoutCancellation() + } +} + +/// Records the cancellation of a registered load task, and parks the loader until +/// that cancellation arrives. +private actor CancellationProbe { + private(set) var wasCancelled = false + private var isParked = false + private var parkContinuation: CheckedContinuation? + private var parkedWaiter: CheckedContinuation? + + /// Suspends the loader until ``recordCancellation()`` runs. + func parkUntilCancelled() async { + if wasCancelled { return } + isParked = true + parkedWaiter?.resume() + parkedWaiter = nil + await withCheckedContinuation { parkContinuation = $0 } + } + + /// Suspends until the loader parks, so a test evicts only after the cache has + /// registered the load task. + func waitUntilParked() async { + if isParked { return } + await withCheckedContinuation { parkedWaiter = $0 } + } + + func recordCancellation() { + wasCancelled = true + parkContinuation?.resume() + parkContinuation = nil + } +} + +/// Thrown if the parked loader resumes without a cancellation. Only +/// `CancellationProbe.recordCancellation()` resumes it, so this error means the +/// probe misbehaved rather than the cache. +private struct ParkedLoaderResumedWithoutCancellation: Error {} + /// Coordinates the in-flight window: the downloader signals when a load has entered /// (so the load task is registered), then parks until the test releases it. private actor LoadGate {