Repository navigation
fix(engine): drain a cold auxiliary or speech load that fails after allocating - #146
Conversation
…llocating The server's reranker and embedder slots and the audio engine drained a failed load only when they had just released a resident model, so a cold load that failed after allocating (weights that do not match the model, say) left what it had read in MLX's cache, and nothing in the server drained it until a later swap, unload or stop, or the next plain generation's own clear. They now drain under the chat engine's rule (#136): whenever a model has loaded somewhere in the process, so Metal is known to be up. A load that fails after allocating before any model has loaded in the process is still not drained, and a load that fails before touching MLX still never initializes Metal. A gated smoke loads an embedder, then asks for a reranker whose weights were rewritten with one tensor renamed, which fails after the weights were read and allocated, and checks the cache is near empty; then the same through the embedder slot.
|
Local pipeline on d38f957 (
Gated, same head: Cursor Bugbot has reported "usage limit reached" since 2026-10-06 and did not review this PR. |
1 participant
A follow-up to #144 (noted in its review): the server's two auxiliary slots drained a failed load only when they had just released a resident engine (
hadResident), whileMLXSwiftEngine.loadsince #140 also drains when a model has loaded anywhere in the process (EngineMemory.hasLoadedAModel), the signal that Metal is up. A cold reranker or embedder load that fails after allocating, weights that do not match the model for instance (mlx-swift-lm'sloadWeightsreads and evaluates every tensor beforeupdate(parameters:verify:)rejects the set), left what it had read in MLX's cache, and nothing in the server drained it until a later swap, unload or stop, or the next plain generation's own clear. The audio engine's STT and TTS loads had the same shape (there a tokenizer or codec that fails after the weights were evaluated).What changes:
ensureEmbedderLoaded,ensureRerankerLoaded, the STT and TTS loads inAudioEngine) drain whenhadResident || EngineMemory.hasLoadedAModel(the audio pair with theirreleasedResident), the chat engine's rule. A load that fails after allocating before any model has loaded in the process is still not drained, since Metal may not be up, and a load that fails before touching MLX never initializes Metal, which the SPM test job (no metallib) relies on.Tests:
MACMLX_RUN_RERANK_SMOKE=1):testAColdLoadThatFailsAfterAllocatingIsDrainedbuilds a copy of ms-marco-MiniLM-L-6-v2 and one of BAAI/bge-small-zh-v1.5 withmodel.safetensorsrewritten through MLX with one tensor renamed, loads the real embedder first (a model has loaded, so Metal is up, and the reranker slot is empty), asks for the broken reranker cold (a 500, nothing installed), checks through MLX's peak that the failed load allocated more than the bound before it failed (so the bound cannot pass vacuously if a later mlx-swift-lm verifies keys before evaluating), and checks MLX's cache is under 64 MiB afterwards, where the undrained read would leave the checkpoint's 91 MB; then unloads the embedder through the door and does the same through the embedder slot (96 MB). The audio pair is not driven by a test: STT loads only by repo id through the app's Hub cache, and no speech checkpoint is kept on this machine. A TTS case is possible (a local copy of a Qwen3 voice with a brokentokenizer.jsonfails in its post-load hook after the model is evaluated) but needs a gated download; the change there is the same one-clause rule, covered by review.hadResident/releasedResident: the reranker's and the embedder's are caught by the smoke (93.1 MB and 95.8 MB left in MLX's cache when the catch does not drain); the two audio mutants survive, as said above, and are covered by review. With the fix the smoke measured 0 B cached after each failed load (the embedder's 95.8 MB still active after the first, 20 B after the second).