Harden #1903's load-time MTP warmup: drain deadline, gate-busy skip, lease pinning#1929
Draft
mimeding wants to merge 1 commit into
Draft
Harden #1903's load-time MTP warmup: drain deadline, gate-busy skip, lease pinning#1929mimeding wants to merge 1 commit into
mimeding wants to merge 1 commit into
Conversation
…ase pin Post-merge hardening for osaurus-ai#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 <noreply@anthropic.com>
mimeding
marked this pull request as draft
July 18, 2026 22:53
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Post-merge hardening for #1903's load-time MTP warmup, remediating three defects found in adversarial review of the merged code:
AsyncStreamiteration is cancellation-responsive); on timeout the generation task is cancelled, the warm flag reset, and the load proceeds without MTP warm. A 2-token warmup completes in <2 s on any supported hardware; 30 s means wedged.maxBatchSize == 1the warmup queued on the global SoloGenerationGate while the cold-load slot was held: under multi-model, loading model A during model B's long generation stalled A's entire load (and any other pending load) for B's full response.SoloGenerationGate.tryAcquire(never enqueues) + asoloGateBehavior: .skipIfBusyseam ongenerate— a busy gate skips the warmup entirely, deferring to the original first-request AR warmup without consuming the warm flag (the throw happens before consumption by construction). The seam shape avoids both the probe-then-call deadlock and the release-then-call TOCTOU.modelCachepublication and loading-record removal but before the caller's lease, so unload/strict-evict could tear the container down mid-warmup.finishLoadedContainernow holds aModelLeaseacross the warmup.Also corrects the kill-switch doc command to the app's actual defaults domain (
com.dinoki.osaurus— the previously documentedai.osaurusdomain is never read, making the documented rollback inert).Why
#1903 is merged and its failure paths are production paths. All three defects share a trigger class this codebase documents extensively (wedged Metal generations, multi-model concurrency) and had no recovery story.
Measurement
8 tests green:
tryAcquiresemantics (free→lease, busy→nil, no gate-state disturbance, queued waiters unaffected) and the drain deadline (completes on finished stream; times out at deadline on an injected never-ending stream). The pre-existing warmup flag tests stay green; behavior for the common single-model case is unchanged (gate free → warmup runs exactly as today).Risk & rollback
Strictly narrows failure modes: every new path (timeout, busy-skip) degrades to #1903's original fallback (first-request AR warmup). Revert = one commit.
🤖 Generated with Claude Code