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