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() + } +}