Skip to content

fix(audio): take the generation lock before loading, run on the model this request loaded - #134

Merged
magicnight merged 5 commits into
mainfrom
fix/audio-swap-under-lock
Oct 4, 2026
Merged

magicnight merged 5 commits into
mainfrom
fix/audio-swap-under-lock

Conversation

@magicnight

@magicnight magicnight commented Oct 4, 2026 •

Copy link
Copy Markdown
Owner

The audio handlers had the same race the embedder/reranker had before #133. /v1/audio/transcriptions and /v1/audio/speech loaded (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:

  • Each handler now runs in three steps: fetch the model's files (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 500 load_failed before the lock is touched. The fetch is also nonisolated, so a request's download does not wait behind a forward pass running on the engine actor.
  • Concurrent fetches of one repo share a single download, process-wide (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's AudioServicing fetches first too, so a user who picks another file while a model is still downloading joins the download instead of restarting it.
  • transcribe(model:) and synthesize(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 500 cancelled, 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).
  • AudioEngine gets transcribe(model:...) and synthesize(model:...), which bind the model their own load produced instead of reading the slot again. The slot-reading transcribe(audioURL:) / synthesize(text:) are removed; loadSTT / loadTTS stay as preload API. The app's AudioServicing goes 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.
  • A swap drains MLX's buffer cache between dropping one model and loading the next, and again when the replacement's load fails, as the embedder and reranker swaps do (Reranker and embedder engines are outside the model pool budget, and a swap keeps both models loaded #130); only when a model was resident, so a failed first load never touches MLX.
  • A malformed model id is answered 400 before any of this, like the lookup errors on /v1/embeddings; the validators are shared with the loaders so the server and the engine cannot disagree.
  • The handlers reach the engine through a small internal AudioBackend protocol. AudioEngine is the only production conformer; the protocol exists so the tests can script the backend.
  • Failure mapping covers the whole call: invalidAudioModelID → 400, modelLoadFailed → 500 load_failed, CancellationError → 500 cancelled, anything else → 500 audio_failed named after the operation. Wire codes are unchanged for every case that existed; an error of unknown type is now audio_failed rather than load_failed, which is honest since every fetch and load error is wrapped.
  • The SpeechModelBox doc now says what serializes synthesis: the box's synthesize is 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):

  • against a scripted backend: the fetch runs without waiting for a held lock while the load waits behind it, on both routes; a failed load releases the lock; a failed fetch is answered with the lock held and never loads; a malformed id on either route is answered with the lock held and never reaches the backend; text / verbose_json and wav / pcm replies with the model id, language, temperature and the staged upload reaching the backend; pcm at 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.
  • on the engine: the model-bound entry points and 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 becomes modelLoadFailed; a pre-cancelled caller stops before the load on both paths (speech against a local TTS directory with a bogus config.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:

mutation (each route where it applies) caught by
load moved ahead of the acquire …LoadWaitsForTheGenerationLock… (answered while locked)
lock released right after the acquire, before the load the in-call probe: lockHeldDuringCalls == [false]
release dropped on the load-failure path lockIsFree after the 500
id validation moved behind the lock aMalformedModelIDDoesNotWaitForTheLock
fetch moved under the lock prepareCalls empty while locked, lockHeldDuringPrepare == [true, true], the fetch-failure test
modelLoadFailed folded into audio_failed the load-failure tests and the mapping tests
pcm sample-rate guard removed speechPCMAtTheWrongSampleRateIs400
fetch dedupe removed concurrentPreparesForOneRepoShareASingleFetch
validation dropped from prepareSTT / from the load path theModelBoundEntryPointsRejectAMalformedIDLocally and the loader tests
cancellation check before the load removed, either path aCancelledCallerStopsBeforeTheLoad, aCancelledTranscriptionStopsBeforeTheLoad
CancellationError no longer mapped to cancelled aCancellationIs500Cancelled

CI 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 under xcodebuild test with 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.

… 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.
@magicnight
magicnight merged commit a109693 into main Oct 4, 2026
4 checks passed
@magicnight
magicnight deleted the fix/audio-swap-under-lock branch October 4, 2026 11:26
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant