mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-18 22:30:40 -07:00
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:
@@ -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())
|
||||
|
||||
@@ -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()
|
||||
@@ -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())
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user