fix(mlx): serialize accelerator lifecycle and inference

Route MLX load, inference, unload, reset, cache cleanup, and shutdown through a single worker. Add affinity and concurrent-unload regression coverage.\n\nVerified: 17 related backend tests; frontend CI; cargo check.

backend/app.py

@@ -353,15 +353,15 @@ async def _run_shutdown() -> None:
     """Unload models on lifespan exit."""
     logger.info("Voicebox server shutting down...")
     try:
-        tts.unload_tts_model()
+        await tts.unload_tts_model()
     except Exception:
         logger.exception("Failed to unload TTS model")
     try:
-        transcribe.unload_whisper_model()
+        await transcribe.unload_whisper_model()
     except Exception:
         logger.exception("Failed to unload Whisper model")
     try:
-        llm.unload_llm_model()
+        await llm.unload_llm_model()
     except Exception:
         logger.exception("Failed to unload LLM model")
 

backend/backends/__init__.py

@@ -547,7 +547,21 @@ async def ensure_model_cached_or_raise(engine: str, model_size: str = "default")
             )
 
 
-def unload_model_by_config(config: ModelConfig) -> bool:
+async def unload_backend(backend) -> None:
+    """Free a backend's model, serialized onto the MLX worker when it has one.
+
+    MLX backends expose an async ``unload`` that runs the free on the dedicated
+    MLX thread so it can't collide with an in-flight load/generate. Other
+    backends only carry the synchronous ``unload_model``.
+    """
+    unload = getattr(backend, "unload", None)
+    if unload is not None:
+        await unload()
+    else:
+        backend.unload_model()
+
+
+async def unload_model_by_config(config: ModelConfig) -> bool:
     """Unload a model given its config. Returns True if it was loaded, False otherwise."""
     from . import get_tts_backend_for_engine
     from ..services import tts, transcribe, llm as llm_service
@@ -555,7 +569,7 @@ def unload_model_by_config(config: ModelConfig) -> bool:
     if config.engine == "whisper":
         whisper_model = transcribe.get_whisper_model()
         if whisper_model.is_loaded() and whisper_model.model_size == config.model_size:
-            transcribe.unload_whisper_model()
+            await unload_backend(whisper_model)
             return True
         return False
 
@@ -563,7 +577,7 @@ def unload_model_by_config(config: ModelConfig) -> bool:
         backend = llm_service.get_llm_model()
         loaded_size = getattr(backend, "_current_model_size", None) or getattr(backend, "model_size", None)
         if backend.is_loaded() and loaded_size == config.model_size:
-            backend.unload_model()
+            await unload_backend(backend)
             return True
         return False
 
@@ -571,7 +585,7 @@ def unload_model_by_config(config: ModelConfig) -> bool:
         tts_model = tts.get_tts_model()
         loaded_size = getattr(tts_model, "_current_model_size", None) or getattr(tts_model, "model_size", None)
         if tts_model.is_loaded() and loaded_size == config.model_size:
-            tts.unload_tts_model()
+            await unload_backend(tts_model)
             return True
         return False
 
@@ -579,14 +593,14 @@ def unload_model_by_config(config: ModelConfig) -> bool:
         backend = get_tts_backend_for_engine(config.engine)
         loaded_size = getattr(backend, "_current_model_size", None) or getattr(backend, "model_size", None)
         if backend.is_loaded() and loaded_size == config.model_size:
-            backend.unload_model()
+            await unload_backend(backend)
             return True
         return False
 
     # All other TTS engines
     backend = get_tts_backend_for_engine(config.engine)
     if backend.is_loaded():
-        backend.unload_model()
+        await unload_backend(backend)
         return True
     return False
 

backend/backends/mlx_backend.py

@@ -3,7 +3,6 @@ MLX backend implementation for TTS and STT using mlx-audio.
 """
 
 from typing import Optional, List, Tuple
-import asyncio
 import logging
 import numpy as np
 from pathlib import Path
@@ -19,6 +18,7 @@ ensure_original_qwen_config_cached()
 
 from . import TTSBackend, STTBackend, LANGUAGE_CODE_TO_NAME, WHISPER_HF_REPOS
 from .base import is_model_cached, combine_voice_prompts as _combine_voice_prompts, model_load_progress
+from ..services.mlx_thread import run_on_mlx_thread, clear_mlx_cache
 from ..utils.cache import get_cache_key, get_cached_voice_prompt, cache_voice_prompt
 
 
@@ -63,6 +63,22 @@ class MLXTTSBackend:
             weight_extensions=(".safetensors", ".bin", ".npz"),
         )
 
+    def _ensure_loaded_sync(self, model_size: Optional[str]):
+        """Load the model if the requested size isn't already resident.
+
+        Runs on the MLX worker thread so it stays serialized with generation.
+        """
+        if model_size is None:
+            model_size = self.model_size
+
+        if self.model is not None and self._current_model_size == model_size:
+            return
+
+        if self.model is not None and self._current_model_size != model_size:
+            self.unload_model()
+
+        self._load_model_sync(model_size)
+
     async def load_model_async(self, model_size: Optional[str] = None):
         """
         Lazy load the MLX TTS model.
@@ -70,23 +86,15 @@ class MLXTTSBackend:
         Args:
             model_size: Model size to load (1.7B or 0.6B)
         """
-        if model_size is None:
-            model_size = self.model_size
-
-        # If already loaded with correct size, return
-        if self.model is not None and self._current_model_size == model_size:
-            return
-
-        # Unload existing model if different size requested
-        if self.model is not None and self._current_model_size != model_size:
-            self.unload_model()
-
-        # Run blocking load in thread pool
-        await asyncio.to_thread(self._load_model_sync, model_size)
+        await run_on_mlx_thread(self._ensure_loaded_sync, model_size)
 
     # Alias for compatibility
     load_model = load_model_async
 
+    async def unload(self):
+        """Free the model, serialized onto the MLX worker thread."""
+        await run_on_mlx_thread(self.unload_model)
+
     def _load_model_sync(self, model_size: str):
         """Synchronous model loading."""
         model_path = self._get_model_path(model_size)
@@ -110,6 +118,7 @@ class MLXTTSBackend:
             del self.model
             self.model = None
             self._current_model_size = None
+            clear_mlx_cache()
             logger.info("MLX TTS model unloaded")
 
     async def create_voice_prompt(
@@ -187,8 +196,6 @@ class MLXTTSBackend:
         Returns:
             Tuple of (audio_array, sample_rate)
         """
-        await self.load_model_async(None)
-
         logger.info("Generating audio for text: %s", text)
 
         def _generate_sync():
@@ -258,8 +265,13 @@ class MLXTTSBackend:
 
             return audio, sample_rate
 
-        # Run blocking inference in thread pool
-        audio, sample_rate = await asyncio.to_thread(_generate_sync)
+        # Load-if-needed and inference run as one job on the MLX worker so a
+        # concurrent unload or different-size load can't land between them.
+        def _load_and_generate():
+            self._ensure_loaded_sync(None)
+            return _generate_sync()
+
+        audio, sample_rate = await run_on_mlx_thread(_load_and_generate)
 
         return audio, sample_rate
 
@@ -279,12 +291,10 @@ class MLXSTTBackend:
         hf_repo = WHISPER_HF_REPOS.get(model_size, f"openai/whisper-{model_size}")
         return is_model_cached(hf_repo, weight_extensions=(".safetensors", ".bin", ".npz"))
 
-    async def load_model_async(self, model_size: Optional[str] = None):
-        """
-        Lazy load the MLX Whisper model.
+    def _ensure_loaded_sync(self, model_size: Optional[str]):
+        """Load the model if the requested size isn't already resident.
 
-        Args:
-            model_size: Model size (tiny, base, small, medium, large)
+        Runs on the MLX worker thread so it stays serialized with transcription.
         """
         if model_size is None:
             model_size = self.model_size
@@ -292,12 +302,24 @@ class MLXSTTBackend:
         if self.model is not None and self.model_size == model_size:
             return
 
-        # Run blocking load in thread pool
-        await asyncio.to_thread(self._load_model_sync, model_size)
+        self._load_model_sync(model_size)
+
+    async def load_model_async(self, model_size: Optional[str] = None):
+        """
+        Lazy load the MLX Whisper model.
+
+        Args:
+            model_size: Model size (tiny, base, small, medium, large)
+        """
+        await run_on_mlx_thread(self._ensure_loaded_sync, model_size)
 
     # Alias for compatibility
     load_model = load_model_async
 
+    async def unload(self):
+        """Free the model, serialized onto the MLX worker thread."""
+        await run_on_mlx_thread(self.unload_model)
+
     def _load_model_sync(self, model_size: str):
         """Synchronous model loading."""
         progress_model_name = f"whisper-{model_size}"
@@ -319,6 +341,7 @@ class MLXSTTBackend:
         if self.model is not None:
             del self.model
             self.model = None
+            clear_mlx_cache()
             logger.info("MLX Whisper model unloaded")
 
     async def transcribe(
@@ -338,8 +361,6 @@ class MLXSTTBackend:
         Returns:
             Transcribed text
         """
-        await self.load_model_async(model_size)
-
         def _transcribe_sync():
             """Run synchronous transcription in thread pool."""
             # MLX Whisper transcription using generate method
@@ -363,5 +384,10 @@ class MLXSTTBackend:
             else:
                 return str(result).strip()
 
-        # Run blocking transcription in thread pool
-        return await asyncio.to_thread(_transcribe_sync)
+        # Load-if-needed and transcription run as one job on the MLX worker so
+        # a concurrent unload or load can't land between them.
+        def _load_and_transcribe():
+            self._ensure_loaded_sync(model_size)
+            return _transcribe_sync()
+
+        return await run_on_mlx_thread(_load_and_transcribe)

backend/backends/qwen_llm_backend.py

@@ -19,6 +19,7 @@ from .base import (
     manual_seed,
     model_load_progress,
 )
+from ..services.mlx_thread import run_on_mlx_thread, clear_mlx_cache
 from ..utils.hf_offline_patch import force_offline_if_cached
 
 logger = logging.getLogger(__name__)
@@ -202,7 +203,11 @@ class MLXQwenLLMBackend:
             weight_extensions=(".safetensors", ".bin", ".npz"),
         )
 
-    async def load_model(self, model_size: Optional[str] = None) -> None:
+    def _ensure_loaded_sync(self, model_size: Optional[str]) -> None:
+        """Load the model if the requested size isn't already resident.
+
+        Runs on the MLX worker thread so it stays serialized with generation.
+        """
         if model_size is None:
             model_size = self.model_size
 
@@ -212,7 +217,14 @@ class MLXQwenLLMBackend:
         if self.model is not None and self._current_model_size != model_size:
             self.unload_model()
 
-        await asyncio.to_thread(self._load_model_sync, model_size)
+        self._load_model_sync(model_size)
+
+    async def load_model(self, model_size: Optional[str] = None) -> None:
+        await run_on_mlx_thread(self._ensure_loaded_sync, model_size)
+
+    async def unload(self) -> None:
+        """Free the model, serialized onto the MLX worker thread."""
+        await run_on_mlx_thread(self.unload_model)
 
     def _load_model_sync(self, model_size: str) -> None:
         from mlx_lm import load as mlx_load
@@ -243,6 +255,7 @@ class MLXQwenLLMBackend:
         self.model = None
         self.tokenizer = None
         self._current_model_size = None
+        clear_mlx_cache()
         logger.info("Qwen3 (MLX) unloaded")
 
     async def generate(
@@ -254,10 +267,13 @@ class MLXQwenLLMBackend:
         model_size: Optional[str] = None,
         examples: Optional[list[tuple[str, str]]] = None,
     ) -> str:
-        await self.load_model(model_size)
-        return await asyncio.to_thread(
-            self._generate_sync, prompt, system, max_tokens, temperature, examples
-        )
+        # Load-if-needed and inference run as one job on the MLX worker so a
+        # concurrent unload or different-size load can't land between them.
+        def _load_and_generate() -> str:
+            self._ensure_loaded_sync(model_size)
+            return self._generate_sync(prompt, system, max_tokens, temperature, examples)
+
+        return await run_on_mlx_thread(_load_and_generate)
 
     def _generate_sync(
         self,

backend/routes/models.py

@@ -66,7 +66,7 @@ async def unload_model():
     from ..services import tts
 
     try:
-        tts.unload_tts_model()
+        await tts.unload_tts_model()
         return {"message": "Model unloaded successfully"}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -82,7 +82,7 @@ async def unload_model_by_name(model_name: str):
         raise HTTPException(status_code=400, detail=f"Unknown model: {model_name}")
 
     try:
-        was_loaded = unload_model_by_config(config)
+        was_loaded = await unload_model_by_config(config)
         if not was_loaded:
             return {"message": f"Model {model_name} is not loaded"}
         return {"message": f"Model {model_name} unloaded successfully"}
@@ -454,7 +454,7 @@ async def delete_model(model_name: str):
     hf_repo_id = config.hf_repo_id
 
     try:
-        unload_model_by_config(config)
+        await unload_model_by_config(config)
 
         cache_dir = hf_constants.HF_HUB_CACHE
         repo_cache_dir = Path(cache_dir) / ("models--" + hf_repo_id.replace("/", "--"))

backend/services/llm.py

@@ -2,7 +2,7 @@
 LLM inference module - delegates to backend abstraction layer.
 """
 
-from ..backends import get_llm_backend, LLMBackend
+from ..backends import LLMBackend, get_llm_backend, unload_backend
 
 
 def get_llm_model() -> LLMBackend:
@@ -10,6 +10,6 @@ def get_llm_model() -> LLMBackend:
     return get_llm_backend()
 
 
-def unload_llm_model() -> None:
-    """Unload LLM model to free memory."""
-    get_llm_backend().unload_model()
+async def unload_llm_model() -> None:
+    """Unload LLM model to free memory, serialized onto the MLX worker."""
+    await unload_backend(get_llm_backend())

backend/services/mlx_thread.py

@@ -0,0 +1,39 @@
+"""Single dedicated worker thread for all MLX GPU work.
+
+MLX's Metal command encoder/stream is thread-local: it binds to whichever
+thread first touches the GPU device. ``asyncio.to_thread()`` uses the event
+loop's default executor, which hands successive calls to different worker
+threads — a model loaded on one thread and generated on another raises
+"There is no Stream(gpu, N) in current thread" (issue #699).
+
+Routing every MLX load, generate, transcribe and unload through this one
+worker keeps them on a single thread. Because the pool has a single worker,
+submitted jobs also run to completion one at a time in submission order, so a
+load-then-infer pair submitted as one job cannot be interleaved with an unload
+or a different-size load from another request.
+"""
+
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+
+_mlx_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="mlx-worker")
+
+
+def run_on_mlx_thread(func, *args):
+    """Run ``func(*args)`` on the single dedicated MLX worker thread."""
+    loop = asyncio.get_running_loop()
+    return loop.run_in_executor(_mlx_executor, func, *args)
+
+
+def clear_mlx_cache() -> None:
+    """Return MLX's cached unified memory to the OS after a model is freed.
+
+    Must run on the MLX worker thread (call it from an unload that is already
+    routed through ``run_on_mlx_thread``). ``clear_cache`` moved out of the
+    ``mlx.core.metal`` namespace in newer MLX, so resolve it from either.
+    """
+    import mlx.core as mx
+
+    clear = getattr(mx, "clear_cache", None) or getattr(getattr(mx, "metal", None), "clear_cache", None)
+    if clear is not None:
+        clear()

backend/services/transcribe.py

@@ -2,21 +2,19 @@
 STT (Speech-to-Text) module - delegates to backend abstraction layer.
 """
 
-from typing import Optional
-from ..backends import get_stt_backend, STTBackend
+from ..backends import STTBackend, get_stt_backend, unload_backend
 
 
 def get_whisper_model() -> STTBackend:
     """
     Get STT backend instance (MLX or PyTorch based on platform).
-    
+
     Returns:
         STT backend instance
     """
     return get_stt_backend()
 
 
-def unload_whisper_model():
-    """Unload Whisper model to free memory."""
-    backend = get_stt_backend()
-    backend.unload_model()
+async def unload_whisper_model():
+    """Unload Whisper model to free memory, serialized onto the MLX worker."""
+    await unload_backend(get_stt_backend())

backend/services/tts.py

@@ -2,28 +2,27 @@
 TTS inference module - delegates to backend abstraction layer.
 """
 
-from typing import Optional
-import numpy as np
 import io
+
+import numpy as np
 import soundfile as sf
 
-from ..backends import get_tts_backend, TTSBackend
+from ..backends import TTSBackend, get_tts_backend, unload_backend
 
 
 def get_tts_model() -> TTSBackend:
     """
     Get TTS backend instance (MLX or PyTorch based on platform).
-    
+
     Returns:
         TTS backend instance
     """
     return get_tts_backend()
 
 
-def unload_tts_model():
-    """Unload TTS model to free memory."""
-    backend = get_tts_backend()
-    backend.unload_model()
+async def unload_tts_model():
+    """Unload TTS model to free memory, serialized onto the MLX worker."""
+    await unload_backend(get_tts_backend())
 
 
 def audio_to_wav_bytes(audio: np.ndarray, sample_rate: int) -> bytes:

backend/tests/test_mlx_thread_affinity.py

@@ -0,0 +1,128 @@
+"""Regression tests for MLX single-thread serialization.
+
+MLX's Metal stream is thread-local, so every load/generate/unload must run on
+one dedicated worker thread (issue #699), and a load+infer pair must run as one
+atomic job so a concurrent unload or different-size load can't land between the
+load and the inference that reads the model.
+
+These drive the real async orchestration on ``MLXQwenLLMBackend`` with the
+heavy mlx-lm calls faked, so they exercise the shipped code paths without
+needing MLX installed.
+"""
+
+import asyncio
+import threading
+import time
+
+import pytest
+
+from backend.backends.qwen_llm_backend import MLXQwenLLMBackend
+from backend.services import llm as llm_service
+from backend.services.mlx_thread import run_on_mlx_thread
+
+
+@pytest.mark.asyncio
+async def test_run_on_mlx_thread_uses_a_single_worker():
+    idents = set()
+
+    def record():
+        idents.add(threading.get_ident())
+
+    await asyncio.gather(*(run_on_mlx_thread(record) for _ in range(12)))
+
+    assert len(idents) == 1, "MLX work must stay pinned to one worker thread"
+    assert idents.pop() != threading.get_ident(), "MLX work must not run on the event loop thread"
+
+
+def _install_fakes(backend, worker_threads):
+    """Replace the heavy sync internals with fakes that record their thread.
+
+    ``_load_model_sync`` and ``_generate_sync`` sleep briefly so that, if the
+    load and inference of one request were ever split into separate jobs, a
+    second request could interleave and be observed.
+    """
+
+    def fake_load(model_size):
+        worker_threads.add(threading.get_ident())
+        time.sleep(0.02)
+        backend.model = {"size": model_size}
+        backend._current_model_size = model_size
+        backend.model_size = model_size
+
+    def fake_unload():
+        worker_threads.add(threading.get_ident())
+        backend.model = None
+        backend._current_model_size = None
+
+    def fake_generate(prompt, system, max_tokens, temperature, examples=None):
+        worker_threads.add(threading.get_ident())
+        # Capture the resident model, do "work", then confirm it wasn't
+        # swapped or freed underneath us — that is exactly the interleave the
+        # atomic load+infer job is meant to prevent.
+        resident = backend.model
+        assert resident is not None, "model was freed mid-generation"
+        time.sleep(0.02)
+        assert backend.model is resident, "model was swapped mid-generation"
+        return resident["size"]
+
+    backend._load_model_sync = fake_load
+    backend.unload_model = fake_unload
+    backend._generate_sync = fake_generate
+
+
+@pytest.mark.asyncio
+async def test_concurrent_generate_does_not_cross_models():
+    backend = MLXQwenLLMBackend()
+    worker_threads = set()
+    _install_fakes(backend, worker_threads)
+
+    small, large = await asyncio.gather(
+        backend.generate("a", model_size="0.6B"),
+        backend.generate("b", model_size="4B"),
+    )
+
+    assert small == "0.6B"
+    assert large == "4B"
+    assert len(worker_threads) == 1, "load and generate must share the one MLX thread"
+
+
+@pytest.mark.asyncio
+async def test_unload_cannot_free_model_mid_generation():
+    backend = MLXQwenLLMBackend()
+    worker_threads = set()
+    _install_fakes(backend, worker_threads)
+
+    await backend.load_model("0.6B")
+
+    # An unload issued while a generation is in flight must serialize behind it
+    # on the worker rather than free the model out from under it.
+    size, _ = await asyncio.gather(
+        backend.generate("a", model_size="0.6B"),
+        backend.unload(),
+    )
+
+    assert size == "0.6B"
+    assert backend.model is None, "unload should still take effect once generation completes"
+    assert len(worker_threads) == 1
+
+
+@pytest.mark.asyncio
+async def test_service_path_unload_serializes_with_generation(monkeypatch):
+    # The service unload helpers (tts/stt/llm) all route through unload_backend,
+    # which must serialize on the MLX worker rather than free the model on the
+    # event-loop thread mid-generation.
+    backend = MLXQwenLLMBackend()
+    worker_threads = set()
+    _install_fakes(backend, worker_threads)
+    monkeypatch.setattr(llm_service, "get_llm_backend", lambda: backend)
+
+    await backend.load_model("0.6B")
+
+    size, _ = await asyncio.gather(
+        backend.generate("a", model_size="0.6B"),
+        llm_service.unload_llm_model(),
+    )
+
+    assert size == "0.6B"
+    assert backend.model is None
+    assert len(worker_threads) == 1