fix: MLX thread affinity crash + missing mlx-lm/mlx-audio in dev setup - #886
fix: MLX thread affinity crash + missing mlx-lm/mlx-audio in dev setup#886AlexStephenLytton wants to merge 3 commits into
Conversation
Two related bugs on Apple Silicon, both stemming from the same root cause (mlx-lm/mlx-audio never getting installed correctly by `just setup-python`), plus one cosmetic health-endpoint bug found along the way: 1. `backend/backends/mlx_backend.py`: pin every MLX model load/generate/ transcribe call to a single dedicated worker thread instead of asyncio.to_thread()'s default pool. MLX's Metal command encoder/stream is thread-local — if model load and generation land on different pool threads (which asyncio.to_thread() does not guarantee against), MLX raises "There is no Stream(gpu, N) in current thread." Fixes jamiepine#699. 2. `justfile`: `just setup-python` installed requirements-mlx.txt (mlx + miniaudio) on Apple Silicon and stopped — it never ran the `pip install --no-deps mlx-lm==0.31.1` / `mlx-audio==0.4.1` steps that .github/workflows/release.yml already does for packaged builds. Without mlx-audio installed, get_backend_type() falls back to "pytorch" silently: no error, no log line calling this out, just a backend that's much slower than intended and never actually uses MLX. Fixes jamiepine#823. Installing mlx-lm without --no-deps (the obvious thing to try once you notice jamiepine#823) reintroduces a second problem: mlx-lm>=0.30.0 requires transformers>=5.0.0, which conflicts with the transformers<=4.57.6 cap in requirements.txt. transformers>=5.0 changed `check_model_inputs` from a decorator factory (`@check_model_inputs()`) to a plain decorator (`@check_model_inputs`, called as `check_model_inputs(func)`) — qwen_tts's CustomVoice tokenizer still uses the old calling convention, so loading Qwen CustomVoice raises `TypeError: check_model_inputs() missing 1 required positional argument: 'func'` once transformers gets silently upgraded. --no-deps avoids this entirely since mlx-lm's own transformers requirement is never consulted; documented this in requirements-mlx.txt so the next person doesn't hit the same trap. 3. `backend/routes/health.py`: `/health`'s `backend_variant` field never had an `mlx` branch, so it always reported "cpu" regardless of whether MLX was active or working — cosmetic, but it's what a couple of the linked issues (jamiepine#706, jamiepine#723) used to conclude MLX wasn't running on the GPU when `backend_type` said otherwise. Added the missing branch so `backend_variant` reports "mps" when the mlx backend is selected. All three verified on real Apple Silicon hardware (M4): repeated generations with no thread crash, `just setup-python` now installs a working MLX backend from a clean venv, Qwen CustomVoice downloads and generates successfully post-fix, and `/health` reports `backend_variant: "mps"` when MLX is active. Fixes jamiepine#699, fixes jamiepine#823
📝 WalkthroughWalkthroughThe MLX backend serializes model loading and inference on a dedicated thread, initializes offline Hugging Face caching, installs pinned MLX dependencies during Apple Silicon setup, enables MPS device selection and cache clearing, and reports ChangesMLX reliability and Apple Silicon setup
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related issues
Possibly related PRs
Sequence Diagram(s)sequenceDiagram
participant Backend as Backend API
participant Executor as MLX worker executor
participant Runtime as MLX TTS/STT runtime
Backend->>Executor: Submit model load or inference
Executor->>Runtime: Run MLX operation on dedicated thread
Runtime-->>Executor: Return model or inference result
Executor-->>Backend: Return asynchronous result
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@backend/backends/mlx_backend.py`:
- Around line 97-101: Route the unload operation in the model-size change branch
through _run_on_mlx_thread instead of calling unload_model() directly. Await the
queued unload before the existing _load_model_sync call so unload_model remains
serialized with in-flight generate or load operations on the dedicated MLX
worker.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 65226774-a58d-43a9-b033-44aee1c1b6d5
📒 Files selected for processing (4)
backend/backends/mlx_backend.pybackend/requirements-mlx.txtbackend/routes/health.pyjustfile
…e Silicon get_torch_device() already supports allow_mps, and other engines (luxtts_backend.py, qwen_llm_backend.py) already pass allow_mps=True — but PyTorchTTSBackend, PyTorchSTTBackend, and QwenCustomVoiceBackend never did, so any Apple Silicon install that ends up on the PyTorch backend (e.g. MLX unavailable, or before this PR's justfile fix) silently ran these three on plain CPU instead of falling back to PyTorch+MPS. This turned out to matter more than expected: benchmarking on M4 hardware, PyTorch+MPS generation (RTF ~1.7x) was faster than this same PR's MLX thread-fix (RTF ~3.4-4.3x) for the same qwen engine and model size, and faster than a CUDA RTX 2080 Ti (RTF ~2.3-3.1x) too. I'm not proposing changing the default backend selection here — one engine, one model size, one machine is not enough data for that, and MLX-Whisper is reportedly faster than PyTorch-CPU for STT — but MPS should clearly be an option rather than silently absent. Verified on M4: PyTorchTTSBackend()._get_device() and QwenCustomVoiceBackend()._get_device() both now return "mps" (were "cpu").
|
Added a second, related commit: enabling MPS for the PyTorch Qwen3-TTS/CustomVoice engines (`allow_mps=True` was already used elsewhere — `luxtts_backend.py`, `qwen_llm_backend.py` — just missing here). While testing the MLX thread fix I benchmarked PyTorch+MPS directly against MLX for the same engine/model/machine (M4, `qwen` engine, 0.6B), out of curiosity about whether MLX's "4-5x faster than PyTorch" claim held up now that both paths actually work:
PyTorch+MPS came out faster than MLX on the same hardware, and faster than a discrete CUDA GPU too. Small sample (one text, one model size, one run each) so take the exact numbers with a grain of salt, but the gap is large enough that it seemed worth surfacing rather than just leaving PyTorch+MPS unreachable. Not proposing changing default backend selection — MLX-Whisper is reportedly faster than PyTorch-CPU for STT per the docs, so the right choice may differ per engine/task — just that MPS should be an available option instead of silently falling back to CPU whenever the PyTorch path is used. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@backend/backends/qwen_custom_voice_backend.py`:
- Line 68: Update unload_model to release the MPS allocator cache in addition to
the existing CUDA cleanup, using the shared device-cache helper if available or
an MPS-specific branch. Ensure unloading on Apple Silicon clears unused MPS
memory when switching models.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 4c8ba5e3-2a31-489b-91f1-bd23604c984e
📒 Files selected for processing (2)
backend/backends/pytorch_backend.pybackend/backends/qwen_custom_voice_backend.py
…ad thread-safety Two follow-ups from the automated review on this PR: 1. backend/backends/base.py: empty_device_cache() only handled "cuda" and "xpu", never "mps" — so once the previous commit enabled MPS for the PyTorch Qwen3-TTS/CustomVoice engines, switching model sizes on Apple Silicon would leak MPS-cached memory (torch.mps.empty_cache() never got called). Added the missing branch. 2. backend/backends/qwen_custom_voice_backend.py: unload_model() had its own inline `if torch.cuda.is_available(): torch.cuda.empty_cache()` instead of using the shared empty_device_cache() helper — meant it never freed MPS memory either, same bug as jamiepine#1 but duplicated instead of shared. Switched to the shared helper (matches pytorch_backend.py's existing pattern). 3. backend/backends/mlx_backend.py: load_model_async()'s unload-on- model-size-change call went straight to self.unload_model() instead of routing through _run_on_mlx_thread like every other MLX call in this file (the whole point of the earlier commit in this PR). Fixed for consistency with the thread-pinning fix. Verified on M4: 0.6B -> generate still works post-changes (regression check), torch.mps.empty_cache() confirmed callable without error in this environment.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@backend/backends/base.py`:
- Around line 185-186: Update the MPS branch in the cache-clearing logic to
require torch.backends.mps.is_available() before calling
torch.mps.empty_cache(), while preserving the existing device == "mps" condition
and CUDA behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 2308a9f0-630b-4dbc-b2b2-58a0def8fe01
📒 Files selected for processing (3)
backend/backends/base.pybackend/backends/mlx_backend.pybackend/backends/qwen_custom_voice_backend.py
🚧 Files skipped from review as they are similar to previous changes (2)
- backend/backends/qwen_custom_voice_backend.py
- backend/backends/mlx_backend.py
| elif device == "mps" and hasattr(torch, "mps"): | ||
| torch.mps.empty_cache() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the changed area and related helpers in backend/backends/base.py
FILE="backend/backends/base.py"
wc -l "$FILE"
sed -n '140,220p' "$FILE"
# Find related uses of empty_device_cache and device detection helpers
rg -n "empty_device_cache|get_torch_device|mps|torch\.mps|torch\.backends\.mps" backend -SRepository: jamiepine/voicebox
Length of output: 10703
🌐 Web query:
PyTorch mps empty_cache availability check torch.backends.mps.is_available docs
💡 Result:
In PyTorch, torch.mps.empty_cache is a standard function available as part of the MPS (Metal Performance Shaders) backend [1][2]. It is used to release unoccupied cached memory held by the caching allocator, which can help reclaim memory for other GPU applications [3][2]. Regarding your query about checking availability: 1. MPS Availability Check: To check if MPS is available in your current environment, you should use torch.backends.mps.is_available [4][5]. This function returns a boolean indicating whether the MPS backend is supported and currently available [4][6]. 2. Function Location: Note that while torch.backends.mps.is_available is used for checking availability [4][7], the empty_cache function resides under the torch.mps namespace (i.e., torch.mps.empty_cache) [1][3]. It is not located under torch.backends.mps. Recent updates to PyTorch have also enhanced torch.mps.empty_cache to clear the MPSGraphCache in addition to the standard allocator pool, which is particularly useful for workloads involving variable tensor shapes [8]. Summary of common MPS-related checks: - Check if MPS is available: torch.backends.mps.is_available [4][7] - Check if PyTorch was built with MPS support: torch.backends.mps.is_built [7][6] - Clear MPS cache: torch.mps.empty_cache [1][2]
Citations:
- 1: https://docs.pytorch.org/docs/stable/generated/torch.mps.empty_cache.html
- 2: https://docs.pytorch.org/docs/2.12/generated/torch.mps.empty%5Fcache.html
- 3: https://github.com/pytorch/pytorch/blob/v2.11.0/torch/mps/__init__.py
- 4: https://docs.pytorch.org/docs/stable/backends.html
- 5: https://docs.pytorch.org/docs/2.11/backends.html
- 6: https://docs.pytorch.org/docs/2.1/%5Fmodules/torch/backends/mps.html
- 7: https://github.com/pytorch/pytorch/blob/main/torch/backends/mps/__init__.py
- 8: [MPS] Clear MPSGraphCache in torch.mps.empty_cache() pytorch/pytorch#181485
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import importlib.util, sys
spec = importlib.util.find_spec("torch")
print("torch_installed=", spec is not None)
if spec is not None:
import torch
print("torch_version=", getattr(torch, "__version__", "unknown"))
print("hasattr(torch, 'mps')=", hasattr(torch, "mps"))
print("hasattr(torch.backends, 'mps')=", hasattr(torch.backends, "mps"))
if hasattr(torch.backends, "mps"):
try:
print("torch.backends.mps.is_available()=", torch.backends.mps.is_available())
except Exception as e:
print("torch.backends.mps.is_available() raised:", type(e).__name__, e)
if hasattr(torch, "mps"):
try:
print("torch.mps exists; empty_cache attr=", hasattr(torch.mps, "empty_cache"))
except Exception as e:
print("torch.mps probe raised:", type(e).__name__, e)
PYRepository: jamiepine/voicebox
Length of output: 179
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect git context around the file without running repository code
git diff -- backend/backends/base.py | sed -n '1,220p'Repository: jamiepine/voicebox
Length of output: 156
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Focus on the specific helper implementation and surrounding device logic.
FILE="backend/backends/base.py"
nl -ba "$FILE" | sed -n '160,205p'Repository: jamiepine/voicebox
Length of output: 196
🌐 Web query:
Does torch.mps.empty_cache raise when torch.backends.mps.is_available is false
💡 Result:
Yes, calling torch.mps.empty_cache when the MPS (Metal Performance Shaders) backend is not available raises a RuntimeError with the message "Cannot execute emptyCache without MPS backend" [1][2]. Because torch.mps.empty_cache interacts directly with the underlying MPS allocator [3][4], it is not safe to call on non-Apple Silicon devices or environments where the MPS backend is not initialized [1][2]. To avoid this error in cross-platform code, you should guard the call by checking the availability of the backend first [1][5][2]: import torch if torch.backends.mps.is_available: torch.mps.empty_cache
Citations:
- 1: RuntimeError: Cannot execute emptyCache() without MPS backend on Windows ace-step/ACE-Step-1.5#842
- 2: Fix unconditional torch.mps.empty_cache() crash in OOM recovery PriorLabs/TabPFN#1007
- 3: https://github.com/pytorch/pytorch/blob/v2.11.0/torch/mps/__init__.py
- 4: https://github.com/pytorch/pytorch/blob/e9ebbd3b/torch/csrc/mps/Module.cpp
- 5: https://deepxde.readthedocs.io/en/latest/_modules/deepxde/backend/pytorch/tensor.html
Guard the MPS cache clear with a backend availability check.
hasattr(torch, "mps") only checks the namespace; torch.mps.empty_cache() can still raise when MPS isn’t available. Mirror the CUDA branch and gate this on torch.backends.mps.is_available().
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@backend/backends/base.py` around lines 185 - 186, Update the MPS branch in
the cache-clearing logic to require torch.backends.mps.is_available() before
calling torch.mps.empty_cache(), while preserving the existing device == "mps"
condition and CUDA behavior.
|
Confirming this fix on more hardware, plus one gap found while testing it. Confirmation — Apple M3 Pro, 18 GB, macOS 26.5 (Darwin 25.5.0), dev checkout (Python 3.12.8,
Gap — the same thread-affinity crash still exists in Opened #909 (built on top of this branch) routing the LLM backend through the same 🤖 Generated with Claude Code |
|
Affected on Voicebox 0.5.0 / macOS 26.5.1 / Apple Silicon. Same MLX thread-affinity crash ( |
|
Thanks for this — two of the three things here were real. The thread-affinity crash is fixed on The part we can't take is the MPS device-policy change bundled in with the fix; that alters generation behavior beyond the crash being fixed and would need its own PR with its own justification. Closing this one, but the setup-python find is appreciated. |
4 participants
Summary
Two related bugs on Apple Silicon that turned out to share a root cause, plus one cosmetic health-endpoint fix found while chasing them down.
Stream(gpu, 1) not in current threadon Apple Silicon (v0.5.0) #699 (Qwen TTS generation fails 100% withStream(gpu, 1) not in current thread)just setup-pythondoesn't install mlx-audio/mlx-lm — MLX backend fails on fresh Apple Silicon dev setup)backend_variantmisreports) — the health.py fix here addresses the reporting side of those; see notes below on what it does and doesn't cover.What's actually going on
Starting from a clean
just setup-pythonon Apple Silicon, MLX never actually gets used:requirements-mlx.txtonly listsmlx+miniaudio.just setup-pythoninstalls it and stops — it never runs thepip install --no-deps mlx-lm==0.31.1/mlx-audio==0.4.1steps that.github/workflows/release.ymlalready does for packaged builds. Withoutmlx_audioimportable,get_backend_type()silently falls back to"pytorch"— no error, no log line calling it out, health endpoint still looks plausible. This is just setup-python doesn't install mlx-audio/mlx-lm — MLX backend fails on fresh Apple Silicon dev setup #823.The obvious fix once you notice just setup-python doesn't install mlx-audio/mlx-lm — MLX backend fails on fresh Apple Silicon dev setup #823 is
pip install mlx-lmby hand. That reintroduces a second problem:mlx-lm>=0.30.0requirestransformers>=5.0.0, conflicting with thetransformers<=4.57.6cap inrequirements.txt.transformers>=5.0changedcheck_model_inputsfrom a decorator factory (@check_model_inputs()) to a plain decorator (@check_model_inputs, i.e.check_model_inputs(func)).qwen_tts's CustomVoice tokenizer (qwen_tts/core/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py) still calls the old convention, so once transformers gets silently bumped, loading Qwen CustomVoice raises:--no-depsavoids this entirely (mlx-lm's own transformers requirement is never consulted, so transformers stays at whateverrequirements.txtinstalled) — which is exactly whatrelease.ymlalready does, it just never made it intojustfile.Once mlx-audio is installed and the backend actually reaches MLX, every model load/generate/transcribe call runs through
asyncio.to_thread(), which uses the event loop's default thread pool — no guarantee that two calls land on the same OS thread. MLX's Metal command encoder/stream is thread-local (bound to whichever thread first touched the GPU device), so if model load and a latergenerate()land on different pool threads, MLX raises:This is Qwen TTS generation fails 100% with
Stream(gpu, 1) not in current threadon Apple Silicon (v0.5.0) #699, reproducible 100% of the time per the original report. Fix: pin every MLX call to one dedicated worker thread instead of the shared pool.Along the way:
/health'sbackend_variantfield never had a branch forbackend_type == "mlx", so it always reported"cpu"regardless of whether MLX was actually active — this is what MLX falls back to CPU on Apple M5 Max — GPU not utilized (voicebox 0.5.0) #706 and MLX backend uses CPU variant on Apple Silicon despite gpu_available=true (M5 Pro) #723 point at as evidence MLX wasn't using the GPU. Added the missing branch ("mps"). This is a reporting fix only; it doesn't address whatever's causing MLX-Metal generation to be slower than PyTorch-CPU on the hardware in those two reports specifically — I don't have that hardware to reproduce on, so I'm not claiming this closes them, just that thebackend_variantfield itself was lying.Changes
justfile: add the two missingpip install --no-depslines to the Unixsetup-pythonrecipe, mirroringrelease.ymlexactly.backend/requirements-mlx.txt: addsounddevice==0.5.3(mlx-audio's own pin, previously unresolved on a from-scratch install), and a NOTE explaining whymlx-lm/mlx-audioare deliberately not listed here and must go through the--no-depspath instead.backend/backends/mlx_backend.py: dedicated single-workerThreadPoolExecutorfor all MLX calls (_run_on_mlx_thread), replacing the 4asyncio.to_thread()call sites acrossMLXTTSBackendandMLXSTTBackend.backend/routes/health.py: 2-line addition —elif backend_type == "mlx": default_variant = "mps".Testing
All verified on real Apple Silicon hardware (M4):
just setup-pythonequivalent commands by hand (matching the justfile diff) →mlx_audioimportable,/healthreportsbackend_type: "mlx",backend_variant: "mps"(previously"pytorch"/"cpu").qwenengine generations (model load + severalgenerate()calls across different requests) with noStream(gpu, N)crash.pip install mlx-lm==0.31.1(no--no-deps, reproducing the trap) → confirmed the exactcheck_model_inputs()TypeError from the report above. Reinstalled per this PR's--no-depssequence → CustomVoice downloads, loads, and generates successfully.qwenMLX engine after the transformers fix to make sure it still works (it does — mlx_audio.tts doesn't touch qwen_tts at all, only the CustomVoice/PyTorch code path does).Happy to adjust anything — this is my first PR against this repo, let me know if there's a preferred format for splitting these into separate PRs instead of one.
Summary by CodeRabbit
Bug Fixes
Chores