Repository navigation
fix(audio): take the generation lock before loading, run on the model this request loaded - #134
Merged
Merged
Conversation
… this request loaded /v1/audio/transcriptions and /v1/audio/speech loaded (and cold-swapped) the model first and queued for the generation lock second, then ran on whatever the slot held once the lock came free. A load therefore ran beside an LLM generation, and a request parked at the lock could have its model replaced, or the slot emptied, by the next request. Same shape as the embedder/reranker race fixed in #133. Both handlers now acquire the lock before the load and release it when the load fails. AudioEngine gains transcribe(model:...) / synthesize(model:...), which bind the model their own load produced instead of re-reading the slot; the two-step loadSTT/transcribe pair stays for the GUI's sequential use. A malformed model id is answered 400 before the lock, like the lookup errors on /v1/embeddings; the validators are shared with the loaders so they cannot drift. The handlers talk to the engine through a small AudioBackend protocol so the lock contract and the reply shapes are tested against a scripted backend, without weights or Metal: a load waits behind a held lock and releases it on failure, malformed ids do not wait, text/verbose_json and wav/pcm replies, and the pcm sample-rate 400.
The lock tests showed that the backend was not called while another holder had the lock and that the lock was free afterwards, but not that the handler itself held the lock during the call: a handler that acquired and released before loading passed them. The scripted backend now tries to take the lock from inside each call, gives it straight back if it gets it, and records the outcome; the four lock tests assert it never got it.
… call, read the lock state in tests Review of the first round found that taking the lock before the load also put a model's first download under it: a multi-GB speech model fetched on first use would have stalled every chat, embeddings and rerank request for minutes, with nothing in the logs saying why. The handlers now fetch the files first (AudioEngine.prepareSTT/prepareTTS, upstream's own resolve-or- download step, files only) outside the lock, and only load and run under it. A repo the Hub cannot deliver fails with 500 load_failed before the lock is touched. Concurrent fetches of one repo share a single download, because upstream clears a directory it finds incomplete and two downloads would delete each other's files. The app's view models cancel a superseded task and start the next one at once, so the superseded call can still be inside the engine; the GUI's AudioServicing now goes through the model-bound transcribe(model:) and synthesize(model:) too, and the slot-reading transcribe(audioURL:) and synthesize(text:) are gone. A swap drains MLX's buffer cache between the two models, as the embedder and reranker swaps do (#130), which the "same gate" comment claimed and the code did not do. The SpeechModelBox doc now says what actually serializes synthesis (the generation lock on the server; nothing in the app yet). The Metal CI job runs the suite in parallel on a slow runner, where a request round trip took close to 20 s and two "does not wait for the lock" tests (one from #133) failed on a fixed 2 s bound. The tests now read the server's lock state directly instead of racing an acquire against a timer, and the bounded waits are 60 s with the lock held throughout, so only a real wait-for-lock regression can fail them.
…off the actor, stop a cancelled request before it loads Second review round. The fetch dedupe moved out of AudioEngine into AudioSnapshotFetcher, a process-wide actor: the app's engine and the server's share one cache directory, and a user picking another file while a model was still downloading started a second download of the same repo, which upstream answers by clearing the directory under the first. The app's AudioServicing now fetches before it loads, like the server. prepareSTT/prepareTTS are nonisolated so a request's download no longer waits behind a forward pass running on the engine; they touch no slot and no MLX state. transcribe(model:) and synthesize(model:) check for cancellation before the load and before the inference, so the app's superseded request stops instead of loading a second copy. A load that fails after a resident model was dropped drains the cache again, as the embedder/reranker swaps do. Tests: the two lock tests asserted that the fetch had reached the backend within a fixed 300 ms, the same kind of bound that failed on the CI runners (both jobs on the previous push); they now poll for the arrival with the lock held, up to 60 s, and only then check that nothing loaded. The shared-fetch test holds the fake fetch at a gate until the fetcher reports the second caller has joined, so it no longer depends on a 400 ms window. The background-post helpers raise their flag on a transport error too, so a timeout is reported as what it is. CHANGELOG: the known-gap sentence said the second copy some loaders download was "still" fetched under the lock; in v0.9.0 nothing was. It now names the six speech-to-text families that ignore the cache directory, Kokoro's grapheme-to-phoneme fetch, Whisper's fallback tokenizer fetch, and that syntheses can still overlap on the app's engine.
…transcription checkpoint, fix the Kokoro note Third review round, nits. A request whose task was cancelled while the backend ran threw CancellationError out of the backend and was reported as 500 audio_failed, while a cancellation while waiting for the lock is 500 cancelled; the mapping now gives both the same code. The transcription path's cancellation checkpoint gets its own test (a pre-cancelled call with a malformed id: the check throws before validation would, so no network and no confusion with the validation error). The embeddings tests get the two leftovers from the previous round: the post helper raises its flag on a transport error too, and the lock-free check's doc notes it is owner- agnostic. CHANGELOG: Kokoro's grapheme-to-phoneme fetch happens in synthesis, which ran under the lock in v0.9.0 as well; it is no longer listed among the downloads that moved under the lock. "Stops before it loads or runs" now says what happens to a request superseded while its load is already running: it stops before the forward pass.
This was referenced Oct 4, 2026
magicnight
added a commit
that referenced
this pull request
Oct 9, 2026
Six speech-to-text families (voxtral, cohere, canary, wav2vec2/mms, lasr, moonshine) resolved their model a second time inside the loader without the cache AudioEngine passes, and on the text-to-speech side Chatterbox and OmniVoice did the same for their whole model, and the Llama and Qwen3 voices and Echo for the codec they fetch from a second repo, so those downloads landed in ~/.cache/huggingface, and since #134 inside the generation lock (#135). The controlled fork now passes the cache through those loaders (the STT half is upstream issue Blaizzy/mlx-audio-swift#279) and carries upstream's pending PR Blaizzy/mlx-audio-swift#257, without which a snapshot fetched by the first resolve would have counted as complete for a loader that asks for more patterns, such as cohere's tokenizer.model or SenseVoice's am.mvn. That PR recorded the patterns a load asked for; the Hub client serves the cached snapshot offline, so it could certify files never downloaded. The fork certifies by the Hub's listing instead, fetches one by one only the listed files not on disk at the revision the snapshot came from (or main when the app's Hub cache does not record that revision or the Hub no longer has the commit), refuses a listed path that cannot be written under the model directory, bounds the listing in time, backs off a repo whose listing ran out of time or took seconds to fail, a 404 included (a failure that comes back at once is asked again next time, and neither a cancelled caller nor a fetch with nothing to certify books one), certifies a fresh fetch from what landed on disk, drops a manifest found beside a snapshot it could not use, and records nothing offline. The pin moves from d138532 to da4cd79. Upstream's v0.1.5 builds against mlx-swift-lm 3.32.3 but has neither fix, so the fork stays for the cache fixes alone. A gated smoke loads a model through AudioEngine with HF_HUB_CACHE pointed at an empty directory and checks that it stays empty, that the files a loader's own patterns name land in the app cache, and that the weights the prefetch wrote are not replaced; it loads a speech model too, and checks the codec it fetches from a second repo lands in the app cache as well.
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
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.
The audio handlers had the same race the embedder/reranker had before #133.
/v1/audio/transcriptionsand/v1/audio/speechloaded (and cold-swapped) the model first, queued for the generation lock second, and then ran on whatever the slot held once the lock came free. Two consequences: a model load ran beside an LLM generation, and a request parked at the lock could have its model replaced, or the slot emptied, by the next request, so it transcribed with the wrong model or got a 500 off an empty slot. Present since v0.9.0.What changes:
AudioEngine.prepareSTT/prepareTTS, upstream's own resolve-or-download step, files only) outside the lock; take the lock; load and run; release. The fetch is outside the lock on purpose: audio models are pulled from the Hub on first use, a multi-GB download can take minutes, and the first version of this branch would have had every chat, embeddings and rerank request queued behind it. A repo the Hub cannot deliver (typo, gated, offline) now fails with 500load_failedbefore the lock is touched. The fetch is alsononisolated, so a request's download does not wait behind a forward pass running on the engine actor.AudioSnapshotFetcher.shared): upstream clears a directory it finds incomplete, so two downloads of one repo delete each other's files, and the app's engine and the server's share one cache directory. The app'sAudioServicingfetches first too, so a user who picks another file while a model is still downloading joins the download instead of restarting it.transcribe(model:)andsynthesize(model:)check for cancellation before the load and before the inference, so the app's superseded request stops at the next checkpoint instead of loading a second copy (one superseded mid-load finishes that load and stops before the forward pass); on the server a request cancelled while the backend runs is 500cancelled, the same code as a cancellation while waiting for the lock, and the lock is released as on any other failure. Two syntheses can still overlap on the app's engine, which has no lock of its own yet (follow-up).AudioEnginegetstranscribe(model:...)andsynthesize(model:...), which bind the model their own load produced instead of reading the slot again. The slot-readingtranscribe(audioURL:)/synthesize(text:)are removed;loadSTT/loadTTSstay as preload API. The app'sAudioServicinggoes through the model-bound calls too: its view models cancel a superseded task and start the next one at once, and the superseded call can still be inside the engine. The app target builds./v1/embeddings; the validators are shared with the loaders so the server and the engine cannot disagree.AudioBackendprotocol.AudioEngineis the only production conformer; the protocol exists so the tests can script the backend.invalidAudioModelID→ 400,modelLoadFailed→ 500load_failed,CancellationError→ 500cancelled, anything else → 500audio_failednamed after the operation. Wire codes are unchanged for every case that existed; an error of unknown type is nowaudio_failedrather thanload_failed, which is honest since every fetch and load error is wrapped.SpeechModelBoxdoc now says what serializes synthesis: the box'ssynthesizeis a nonisolated async method and runs off the actor, so only the generation lock does on the server, and nothing does in the app yet (follow-up).Tests, all under bare
swift test(no weights, no Hub, no Metal):text/verbose_jsonandwav/pcmreplies with the model id, language, temperature and the staged upload reaching the backend;pcmat a non-24 kHz rate is a 400 and the lock is free afterwards. The backend records, from inside each call, whether the server's lock was held (generationLockIsHeld, a new internal read), so a handler that acquires and releases before loading is caught too.prepare*reject a malformed id locally and leave nothing resident; two concurrent prepares for one repo make one fetch and a different repo makes its own, with the fake fetch held at a gate until the fetcher reports the second caller has joined, so no timing is involved; a Hub failure becomesmodelLoadFailed; a pre-cancelled caller stops before the load on both paths (speech against a local TTS directory with a bogusconfig.json, transcription with a malformed id whose validation error cannot be confused with the cancellation; neither touches the network, with or without the check). The fetch and the fetcher are injectable for these.Mutation check, each guard broken in turn and the audio + engine suites run against it; all 18 failed as they must:
…LoadWaitsForTheGenerationLock…(answered while locked)lockHeldDuringCalls == [false]lockIsFreeafter the 500aMalformedModelIDDoesNotWaitForTheLockprepareCallsempty while locked,lockHeldDuringPrepare == [true, true], the fetch-failure testmodelLoadFailedfolded intoaudio_failedspeechPCMAtTheWrongSampleRateIs400concurrentPreparesForOneRepoShareASingleFetchprepareSTT/ from the load paththeModelBoundEntryPointsRejectAMalformedIDLocallyand the loader testsaCancelledCallerStopsBeforeTheLoad,aCancelledTranscriptionStopsBeforeTheLoadCancellationErrorno longer mapped tocancelledaCancellationIs500CancelledCI note: the Metal job on the first push failed two "does not wait for the lock" tests, one of them #133's
embeddingsLookupAndKindErrorsDoNotWaitForTheLock, on fixed 2 s bounds; the second push failed both lock tests on a fixed 300 ms bound of my own. That job runs all 990 tests in parallel on a slow runner, where a request round trip took close to 20 s. Both test files now read the server's lock state directly instead of racing an acquire against a timer, every wait for something that must happen polls up to 60 s with the lock held throughout, and the only fixed sleeps left are followed by "has not happened yet" assertions, which a slow runner can only pass. The full suite passes locally underxcodebuild testwith the same flags as the job.Not covered by a test: that the real engine's load from the fetched files uses the bound model rather than the slot, and the second cancellation checkpoint (after the load, before the forward pass) on either path. Both need a loadable checkpoint; we still have no audio checkpoint verified on our hardware (README says so), and the binding is a few lines the review can read.
Known gap, in the CHANGELOG: the fetch covers what upstream's own first step downloads, so whatever a loader fetches on its own now happens under the lock because the load does. Six upstream speech-to-text families (voxtral, cohere, canary, wav2vec2/mms, lasr, moonshine) ignore the cache directory they are given and download a second full copy into the shared Hugging Face cache, and Whisper fetches tokenizer files from the matching
openai/whisper-*repo when its snapshot lacks them; in v0.9.0 those ran outside the lock, with everything else. Kokoro's first synthesis fetches its grapheme-to-phoneme model under the lock, as it did in v0.9.0, since synthesis always ran there. Fixing the first two means passing the cache directory through in the mlx-audio-swift fork, or fetching those pieces before the lock too; follow-up.Other follow-ups the review surfaced, none introduced here: the download is no longer stopped when a client disconnects mid-fetch (it used to run inside the load; whether Hummingbird cancels the handler on disconnect is unverified either way); model ids differing only in case get separate fetch entries but one directory on APFS; the app's engine still has nothing serializing two syntheses.