From 0070c04bcfd49df06bb26f3316fc73c83045d40b Mon Sep 17 00:00:00 2001 From: Jamie Pine Date: Sun, 19 Jul 2026 17:40:03 -0700 Subject: [PATCH] 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 | 6 +- backend/backends/__init__.py | 26 ++++- backend/backends/mlx_backend.py | 84 +++++++++----- backend/backends/qwen_llm_backend.py | 28 ++++- backend/routes/models.py | 6 +- backend/services/llm.py | 8 +- backend/services/mlx_thread.py | 39 +++++++ backend/services/transcribe.py | 12 +- backend/services/tts.py | 15 ++- backend/tests/test_mlx_thread_affinity.py | 128 ++++++++++++++++++++++ 10 files changed, 286 insertions(+), 66 deletions(-) create mode 100644 backend/services/mlx_thread.py create mode 100644 backend/tests/test_mlx_thread_affinity.py diff --git a/backend/app.py b/backend/app.py index c63a2a6c..79ca9eaf 100644 --- a/backend/app.py +++ b/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") diff --git a/backend/backends/__init__.py b/backend/backends/__init__.py index 2437a87b..abaf70c1 100644 --- a/backend/backends/__init__.py +++ b/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 diff --git a/backend/backends/mlx_backend.py b/backend/backends/mlx_backend.py index 9692e59b..92ba80ee 100644 --- a/backend/backends/mlx_backend.py +++ b/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) diff --git a/backend/backends/qwen_llm_backend.py b/backend/backends/qwen_llm_backend.py index e77a3540..3df20824 100644 --- a/backend/backends/qwen_llm_backend.py +++ b/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, diff --git a/backend/routes/models.py b/backend/routes/models.py index 7cbb7b04..3f3655e0 100644 --- a/backend/routes/models.py +++ b/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("/", "--")) diff --git a/backend/services/llm.py b/backend/services/llm.py index e89c9f8c..0eff3e7f 100644 --- a/backend/services/llm.py +++ b/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()) diff --git a/backend/services/mlx_thread.py b/backend/services/mlx_thread.py new file mode 100644 index 00000000..435a33a2 --- /dev/null +++ b/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() diff --git a/backend/services/transcribe.py b/backend/services/transcribe.py index e400dbd5..4104bb3b 100644 --- a/backend/services/transcribe.py +++ b/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()) diff --git a/backend/services/tts.py b/backend/services/tts.py index d4f90ff3..e5b05675 100644 --- a/backend/services/tts.py +++ b/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: diff --git a/backend/tests/test_mlx_thread_affinity.py b/backend/tests/test_mlx_thread_affinity.py new file mode 100644 index 00000000..cf298bd4 --- /dev/null +++ b/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