Harden #1903's load-time MTP warmup: drain deadline, gate-busy skip, lease pinning

Packages/OsaurusCore/Services/ModelRuntime.swift

@@ -672,13 +672,22 @@ 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,
             draftStrategy: holder.draftStrategy,
             runtime: getConfig(),
             maxBatchSize: InferenceFeatureFlags.mlxBatchEngineMaxBatchSize
         )
+        await ModelLease.shared.release(name)
 
         genLog.info(
             "loadContainer: loaded \(name, privacy: .public) isVLM=\(holder.isVLM, privacy: .public)"

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<Generation>,
+        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()

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