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.
This commit is contained in:
Jamie Pine
2026-07-19 17:45:52 -07:00
parent f2cf2a729d
commit 0070c04bcf
10 changed files with 286 additions and 66 deletions
+4 -4
View File
@@ -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())
+39
View File
@@ -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()
+5 -7
View File
@@ -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())
+7 -8
View File
@@ -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: