fix(mlx): run all MLX operations on a single dedicated thread — fixes "There is no Stream(gpu, N) in current thread"

[flowbrix]

Symptom

On macOS (Apple Silicon, MLX backend), Qwen3-TTS voice-clone (ICL) generation fails with HTTP 500 on every request:

RuntimeError: There is no Stream(gpu, 1) in current thread.

The same error hits the MLX STT path (voicebox.transcribe via HTTP/MCP), with varying stream numbers (Stream(gpu, 4), …).

Fixes #675
Fixes #699

Root cause

mlx_lm creates a thread-local GPU stream at module import time — mlx_lm/generate.py:226 in mlx-lm 0.31.x:

generation_stream = mx.new_thread_local_stream(mx.default_device())

Any MLX array or op that captures that stream can only be evaluated on the thread that created it. backend/backends/mlx_backend.py dispatched model loading and generation through asyncio.to_thread, which uses a rotating worker pool — so the import, the load, and each generation could land on different threads. As soon as a generation ran on a thread other than the importing one, MLX raised There is no Stream(gpu, N) in current thread and the request failed. (mlx_audio has module-level streams on its paths too, so the STT backend is affected the same way.)

This is why behavior looked inconsistent across entry points (#675): whichever call path happened to reuse the importing thread worked; everything else failed.

Fix

Route all MLX work — module imports, model loading, generation, transcription — through one shared single-worker ThreadPoolExecutor (_MLX_EXECUTOR + _run_on_mlx_thread helper), replacing the four asyncio.to_thread call sites in mlx_backend.py. Everything MLX now lives on the same dedicated thread for the lifetime of the process.

The serialization is free in practice: MLX is single-GPU and generations were already serialized by the backend.

Verification

On an M3 Max (macOS 26, mlx-lm 0.31.3):

• Before: Qwen3-TTS 0.6B cloned-voice generation via /generate/stream → HTTP 500 (There is no Stream(gpu, 1) in current thread), 100% reproducible.
• After: same request succeeds; warm RTF ~0.9. Repeated/concurrent requests stay stable (no thread-affinity dependence on which handler thread serves the request).
• Kokoro TTS and Whisper STT paths re-tested, unchanged.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes
  • Improved MLX-based speech-to-text and text-to-speech reliability by running model loading and inference through a dedicated execution thread.
  • Reduced the risk of concurrency-related issues during audio processing.

[coderabbitai]

📝 Walkthrough

Walkthrough

MLX backend model loading, TTS generation, and STT transcription now run through a shared single-worker executor instead of asyncio.to_thread, keeping MLX operations on one dedicated thread.

Changes

MLX thread dispatch

Layer / File(s)	Summary
MLX executor and dispatcher backend/backends/mlx_backend.py	Adds a single-worker MLX executor and async helper for dispatching blocking operations.
TTS and STT operation dispatch backend/backends/mlx_backend.py	Routes TTS loading/generation and STT loading/transcription through the dedicated dispatcher.

Estimated code review effort: 2 (Simple) | ~10 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Backend
  participant MLXDispatcher
  participant DedicatedThread
  Backend->>MLXDispatcher: Dispatch model loading or inference
  MLXDispatcher->>DedicatedThread: Execute synchronous MLX helper
  DedicatedThread-->>Backend: Return result

Loading

Suggested reviewers: jamiepine

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly describes the main change: routing MLX work to one dedicated thread to fix the stream-context error.
Linked Issues check	✅ Passed	The change matches issues #675 and #699 by moving MLX loading and inference paths onto a shared single-worker executor.
Out of Scope Changes check	✅ Passed	The PR appears scoped to the MLX thread-affinity fix in mlx_backend.py with no unrelated code changes.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

backend/backends/mlx_backend.py (1)

25-28: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Preserve contextvars to maintain asyncio.to_thread's behavior.

asyncio.to_thread natively copies contextvars, which are frequently used by web frameworks and logging libraries to track request IDs and tracing contexts across async tasks. When manually switching to loop.run_in_executor(), these context variables are dropped.

You can explicitly copy the context to perfectly mirror the original to_thread behavior and prevent potential tracing regressions.

♻️ Proposed refactor to preserve context

+import contextvars
+import functools
+
 async def _run_on_mlx_thread(fn, *args):
     """Run a blocking MLX operation on the dedicated MLX thread."""
-    return await asyncio.get_running_loop().run_in_executor(_MLX_EXECUTOR, fn, *args)
+    loop = asyncio.get_running_loop()
+    ctx = contextvars.copy_context()
+    func_call = functools.partial(ctx.run, fn, *args)
+    return await loop.run_in_executor(_MLX_EXECUTOR, func_call)

🤖 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/mlx_backend.py` around lines 25 - 28, Update
_run_on_mlx_thread to capture the current contextvars context before submitting
the blocking operation, then execute fn with its arguments inside that copied
context on _MLX_EXECUTOR, preserving asyncio.to_thread behavior.

🤖 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.

Nitpick comments:
In `@backend/backends/mlx_backend.py`:
- Around line 25-28: Update _run_on_mlx_thread to capture the current
contextvars context before submitting the blocking operation, then execute fn
with its arguments inside that copied context on _MLX_EXECUTOR, preserving
asyncio.to_thread behavior.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 37c93f4e-9826-4a07-ad9c-34f4eb06e4db

📥 Commits

Reviewing files that changed from the base of the PR and between f2cf2a7 and b038d17.

📒 Files selected for processing (1)

• backend/backends/mlx_backend.py