From 918acf7e6fbde57156daf747a4c98bbad03bf79a Mon Sep 17 00:00:00 2001 From: mimeding <264272563+mimeding@users.noreply.github.com> Date: Tue, 7 Jul 2026 16:14:28 -0400 Subject: [PATCH] Harden load-time MTP warmup: drain deadline, no solo-gate queuing, lease pin MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Post-merge hardening for #1903 warmup, fixing three verified defects in MLXBatchAdapter.warmupNativeMTPAtLoad and its finishLoadedContainer call site: - No timeout: "for await _ in prepared.stream {}" could hang forever on a wedged generation while finishLoadedContainer held the process-wide cold-load slot. The drain now races a 30-second deadline (a 2-token greedy warmup completes in <2 s on any supported hardware; 30 s means something is wedged). On timeout: cancel prepared.genTask, reset the registry warm flag, log a notice, and return — the load proceeds without MTP warm and the first request falls back to AR warmup. - Global solo-gate contention: at maxBatchSize == 1 the warmup's generate queued on the process-wide SoloGenerationGate, so under manualMultiModel loading model A during model B's long generation stalled A's entire load. SoloGenerationGate gains a non-queuing tryAcquire(modelName:) -> Lease?, and generate gains a soloGateBehavior parameter (.waitForTurn default / .skipIfBusy). The warmup path uses .skipIfBusy: when the gate is held, generate throws SoloGateBusyError before consumeNativeMTPColdWarmup runs, so the warmup is skipped without consuming the warm flag and the first real request performs the AR warmup as before. - Lease-less warmup window: finishLoadedContainer publishes modelCache and removes the loading record before awaiting the warmup, so an unload / strict-single-model eviction could tear the container down under the in-flight warmup generation. The model is now pinned with ModelLease.shared.acquire(name) across the warmup and released on every path (warmupNativeMTPAtLoad never throws; the timeout path returns normally). Also corrects the escape-hatch doc comment ("defaults write ai.osaurus ...") to the app's actual bundle id domain com.dinoki.osaurus. New SoloGateWarmupHardeningTests covers tryAcquire semantics (free -> lease, busy -> nil, failed probe leaves queued waiters unaffected) and the drain deadline via an injected never-ending stream; existing NativeMTPWarmupFlagTests stay green. Co-Authored-By: Claude Fable 5 --- .../OsaurusCore/Services/ModelRuntime.swift | 9 ++ .../ModelRuntime/MLXBatchAdapter.swift | 127 ++++++++++++++++-- .../SoloGateWarmupHardeningTests.swift | 104 ++++++++++++++ 3 files changed, 232 insertions(+), 8 deletions(-) create mode 100644 Packages/OsaurusCore/Tests/Service/SoloGateWarmupHardeningTests.swift diff --git a/Packages/OsaurusCore/Services/ModelRuntime.swift b/Packages/OsaurusCore/Services/ModelRuntime.swift index 33494928f..3389f1515 100644 --- a/Packages/OsaurusCore/Services/ModelRuntime.swift +++ b/Packages/OsaurusCore/Services/ModelRuntime.swift @@ -672,6 +672,14 @@ public actor ModelRuntime { // greedy generation, so the first user request decodes with MTP. // Failure is non-fatal: the request path's cold-warmup rule remains // as the fallback. + // + // Pin the model for the duration of the warmup. `modelCache[name]` + // is already published and the loading record removed above, so an + // unload / strict-single-model eviction racing this window would + // otherwise see a zero lease count and tear the container down + // under the in-flight warmup generation. The lease is released on + // every exit path (`warmupNativeMTPAtLoad` never throws). + await ModelLease.shared.acquire(name) await MLXBatchAdapter.warmupNativeMTPAtLoad( modelName: name, container: holder.container, @@ -679,6 +687,7 @@ public actor ModelRuntime { runtime: getConfig(), maxBatchSize: InferenceFeatureFlags.mlxBatchEngineMaxBatchSize ) + await ModelLease.shared.release(name) genLog.info( "loadContainer: loaded \(name, privacy: .public) isVLM=\(holder.isVLM, privacy: .public)" diff --git a/Packages/OsaurusCore/Services/ModelRuntime/MLXBatchAdapter.swift b/Packages/OsaurusCore/Services/ModelRuntime/MLXBatchAdapter.swift index b4d17ac91..6d10f1121 100644 --- a/Packages/OsaurusCore/Services/ModelRuntime/MLXBatchAdapter.swift +++ b/Packages/OsaurusCore/Services/ModelRuntime/MLXBatchAdapter.swift @@ -266,6 +266,18 @@ struct MLXBatchAdapter { return Lease(gate: self) } + /// Non-queuing probe: returns a lease when the gate is free, `nil` + /// when another generation holds it. Never enqueues a waiter, so a + /// `nil` result has no effect on the gate or on queued `acquire` + /// callers. Used by the load-time MTP warmup, which must never wait + /// behind an unrelated model's generation while the process-wide + /// cold-load slot is held. + func tryAcquire(modelName: String) -> Lease? { + guard !busy else { return nil } + busy = true + return Lease(gate: self) + } + private func release() { guard busy else { return } if !waiters.isEmpty { @@ -538,6 +550,12 @@ struct MLXBatchAdapter { await soloGate.acquire(modelName: modelName) } + /// Non-queuing variant of `acquireSoloLease` — see + /// `SoloGenerationGate.tryAcquire(modelName:)`. + func tryAcquireSoloLease(for modelName: String) async -> SoloGenerationGate.Lease? { + await soloGate.tryAcquire(modelName: modelName) + } + func consumeNativeMTPColdWarmup(modelName: String, requested: Bool) -> Bool { guard requested else { return false } if nativeMTPWarmModels.contains(modelName) { @@ -568,9 +586,39 @@ struct MLXBatchAdapter { /// Escape hatch in case a family surfaces a first-generation issue that /// the two-token warmup does not absorb: - /// defaults write ai.osaurus ai.osaurus.mtp.disableLoadWarmup -bool true + /// defaults write com.dinoki.osaurus ai.osaurus.mtp.disableLoadWarmup -bool true static let mtpLoadWarmupDisabledKey = "ai.osaurus.mtp.disableLoadWarmup" + /// Deadline for draining the warmup generation. A 2-token greedy warmup + /// completes in well under 2 s on any supported hardware; 30 s means the + /// generation is wedged (hung engine actor, Metal stall). The caller + /// (`finishLoadedContainer`) holds the process-wide cold-load slot, so + /// an unbounded drain would block every subsequent model load. + static let mtpLoadWarmupDrainDeadlineNanoseconds: UInt64 = 30 * 1_000_000_000 + + /// Drain `stream`, giving up after `nanoseconds`. Returns `true` when + /// the stream finished, `false` on timeout. `AsyncStream` iteration is + /// cancellation-aware, so `cancelAll()` promptly unblocks the drain + /// child when the deadline child wins. + static func drainWithDeadline( + _ stream: AsyncStream, + nanoseconds: UInt64 + ) async -> Bool { + await withTaskGroup(of: Bool.self) { group in + group.addTask { + for await _ in stream {} + return true + } + group.addTask { + try? await Task.sleep(nanoseconds: nanoseconds) + return false + } + let finished = await group.next() ?? false + group.cancelAll() + return finished + } + } + /// Runs the native-MTP cold warmup at model-load time instead of on the /// user's first request. The registry's cold-warmup rule forces the /// first generation per model into plain AR mode; without this, that @@ -595,6 +643,13 @@ struct MLXBatchAdapter { let startedAt = CFAbsoluteTimeGetCurrent() do { + // `.skipIfBusy`: at maxBatchSize == 1 a plain `generate` would + // QUEUE on the process-wide SoloGenerationGate. Under + // manualMultiModel that means loading model A during model B's + // long generation would stall A's entire load (the caller holds + // the cold-load slot) just to run a hidden 2-token warmup. If the + // gate is busy, `generate` throws before the warm flag is + // consumed and we skip the warmup entirely. let prepared = try await generate( modelName: modelName, container: container, @@ -605,13 +660,39 @@ struct MLXBatchAdapter { stopSequences: [], draftStrategy: draftStrategy, runtime: runtime, - maxBatchSize: maxBatchSize + maxBatchSize: maxBatchSize, + soloGateBehavior: .skipIfBusy + ) + // Bound the drain: an unbounded `for await` here would hang + // `finishLoadedContainer` — and with it the process-wide + // cold-load slot — forever on a wedged generation. See + // `mtpLoadWarmupDrainDeadlineNanoseconds` for the 30 s rationale. + let finished = await drainWithDeadline( + prepared.stream, + nanoseconds: mtpLoadWarmupDrainDeadlineNanoseconds ) - for await _ in prepared.stream {} + guard finished else { + prepared.genTask.cancel() + // The generation consumed the warm flag before wedging; + // reset it so the first real request runs the AR warmup. + await Registry.shared.resetNativeMTPWarmup(modelName: modelName) + batchAdapterLog.notice( + "native MTP load warmup timed out after 30s for \(modelName, privacy: .public) — cancelled; load proceeds without MTP warm, first request falls back to AR warmup" + ) + return + } let elapsedMs = Int((CFAbsoluteTimeGetCurrent() - startedAt) * 1000) batchAdapterLog.info( "native MTP load warmup: completed for \(modelName, privacy: .public) in \(elapsedMs, privacy: .public)ms; first user request decodes with MTP" ) + } catch is SoloGateBusyError { + // Another model's generation holds the solo gate. The throw + // happens before `consumeNativeMTPColdWarmup`, so the warm flag + // is untouched — the first real request performs the AR warmup + // exactly as it did before load-time warmup existed. + batchAdapterLog.notice( + "native MTP load warmup: solo gate busy for \(modelName, privacy: .public) — deferring to first-request AR warmup" + ) } catch { // The failed generation may already have consumed the warm flag; // reset it so the request path's AR warmup applies unchanged. @@ -1040,6 +1121,24 @@ struct MLXBatchAdapter { } } + /// How `generate` treats the process-wide `SoloGenerationGate` when + /// `maxBatchSize == 1`. + enum SoloGateBehavior: Sendable { + /// Queue behind the current holder (regular request path). + case waitForTurn + /// Probe without queuing and throw `SoloGateBusyError` when another + /// generation holds the gate. Load-time warmup path only: the + /// warmup must never wait behind an unrelated model's generation + /// while `finishLoadedContainer` holds the process-wide cold-load + /// slot. + case skipIfBusy + } + + /// Thrown by `generate(soloGateBehavior: .skipIfBusy)` when the solo + /// gate is held. Thrown before the native-MTP warm flag is consumed, + /// so callers can skip the warmup without un-warming anything. + struct SoloGateBusyError: Error {} + /// Tokenize the chat + tools, fetch (or create) the per-model /// `BatchEngine`, and submit one request via `engine.generate`. Returns /// the resulting `Generation` stream wrapped with cancellation plumbing. @@ -1054,7 +1153,8 @@ struct MLXBatchAdapter { stopSequences: [String], draftStrategy: MLXLMCommon.DraftStrategy?, runtime: RuntimeConfig, - maxBatchSize: Int + maxBatchSize: Int, + soloGateBehavior: SoloGateBehavior = .waitForTurn ) async throws -> PreparedStream { let trace = generation.ttftTrace trace?.mark("batch_prepare_start") @@ -1067,10 +1167,21 @@ struct MLXBatchAdapter { PrefillDebugLog.shared.log( "==GEN GENERATE-ENTER model=\(modelName) maxBatch=\(maxBatchSize)" ) - let soloLease = - maxBatchSize == 1 - ? await Registry.shared.acquireSoloLease(for: modelName) - : nil + let soloLease: SoloGenerationGate.Lease? + if maxBatchSize == 1 { + switch soloGateBehavior { + case .waitForTurn: + soloLease = await Registry.shared.acquireSoloLease(for: modelName) + case .skipIfBusy: + guard let probed = await Registry.shared.tryAcquireSoloLease(for: modelName) + else { + throw SoloGateBusyError() + } + soloLease = probed + } + } else { + soloLease = nil + } if Task.isCancelled { if let soloLease { await soloLease.release() } throw CancellationError() diff --git a/Packages/OsaurusCore/Tests/Service/SoloGateWarmupHardeningTests.swift b/Packages/OsaurusCore/Tests/Service/SoloGateWarmupHardeningTests.swift new file mode 100644 index 000000000..fc87ca1f6 --- /dev/null +++ b/Packages/OsaurusCore/Tests/Service/SoloGateWarmupHardeningTests.swift @@ -0,0 +1,104 @@ +// +// SoloGateWarmupHardeningTests.swift +// osaurusTests +// +// Post-merge hardening coverage for the load-time MTP warmup (#1903): +// +// - `SoloGenerationGate.tryAcquire` semantics: a probe returns a lease +// when the gate is free, `nil` when busy, and a failed probe never +// enqueues — queued `acquire` waiters are unaffected. +// - `drainWithDeadline`: the warmup's stream drain must come back at the +// deadline on a wedged generation (the caller holds the process-wide +// cold-load slot), and must report completion when the stream ends. +// The wedged case is exercised with an injected never-ending stream — +// `drainWithDeadline` takes the `AsyncStream` directly, so no real +// engine is needed. +// + +import Foundation +@preconcurrency import MLXLMCommon +import Testing + +@testable import OsaurusCore + +struct SoloGateWarmupHardeningTests { + + // MARK: - SoloGenerationGate.tryAcquire + + @Test func tryAcquireReturnsLeaseWhenFreeAndNilWhenBusy() async { + let gate = MLXBatchAdapter.SoloGenerationGate() + + let lease = await gate.tryAcquire(modelName: "model-a") + #expect(lease != nil) + + // While held, a probe fails without queuing. + #expect(await gate.tryAcquire(modelName: "model-b") == nil) + + await lease?.release() + + // Free again: the probe succeeds. + let again = await gate.tryAcquire(modelName: "model-b") + #expect(again != nil) + await again?.release() + } + + @Test func blockingAcquireStillWorksAfterFailedProbe() async { + let gate = MLXBatchAdapter.SoloGenerationGate() + + let holder = await gate.acquire(modelName: "model-a") + #expect(await gate.tryAcquire(modelName: "model-b") == nil) + await holder.release() + + // The failed probe must not have corrupted the gate state. + let next = await gate.acquire(modelName: "model-b") + await next.release() + } + + @Test func failedProbeDoesNotDisturbQueuedWaiters() async { + let gate = MLXBatchAdapter.SoloGenerationGate() + let holder = await gate.acquire(modelName: "model-a") + + // Queue a blocking waiter behind the holder. + let waiter = Task { + let lease = await gate.acquire(modelName: "model-b") + await lease.release() + return true + } + // Let the waiter park (best-effort; the assertions below hold on + // either interleaving). + try? await Task.sleep(nanoseconds: 50_000_000) + + // A probe while busy fails and must not enqueue. + #expect(await gate.tryAcquire(modelName: "model-c") == nil) + + // Releasing the holder hands the gate to the queued waiter, which + // completes normally — the failed probe left no trace. + await holder.release() + #expect(await waiter.value) + } + + // MARK: - Warmup drain deadline + + @Test func drainReportsCompletionWhenStreamEnds() async { + let stream = AsyncStream { continuation in + continuation.finish() + } + let finished = await MLXBatchAdapter.drainWithDeadline( + stream, nanoseconds: 1_000_000_000) + #expect(finished) + } + + @Test func drainTimesOutOnWedgedStream() async { + // Injected never-ending stream: the continuation is kept alive and + // never finished (dropping it would auto-finish the stream), so + // iteration would block forever without the deadline. + let (stream, continuation) = AsyncStream.makeStream() + let start = Date() + let finished = await MLXBatchAdapter.drainWithDeadline( + stream, nanoseconds: 50_000_000) + #expect(!finished) + // Must come back at the deadline, not hang. + #expect(Date().timeIntervalSince(start) < 5) + continuation.finish() + } +}