fix(mlx): keep unload_model off the MLX worker so it cannot stall the event loop

Routing unload through the single MLX worker and waiting on the result
blocked the caller (the FastAPI event loop via /models/unload) for the
remainder of any in-flight generation — measured 3.7 s stall on a 16 s
clip on an M2 Ultra, unbounded for long texts. Dropping the model
reference inline is thread-safe (MLX frees buffers through its global
allocator) and is what main did before the thread-affinity change; the
generation in flight keeps its own reference and the next generate()
reloads on the worker.
This commit is contained in:
jamiepine
2026-10-04 00:01:47 +00:00
committed by capy-ai-staging[bot]
parent 99d05b917c
commit cf984885e9
+11 -21
View File
@@ -133,21 +133,16 @@ class MLXTTSBackend:
def unload_model(self): def unload_model(self):
"""Unload the model to free memory. """Unload the model to free memory.
Safe to call from any thread (e.g. the FastAPI event loop, from the Runs inline on the calling thread rather than on the MLX worker:
/models/unload routes): the actual teardown is submitted to the MLX dropping the Python reference is thread-safe (MLX frees buffers
worker thread and awaited synchronously here, so Metal resources are through its global allocator, no stream needed), and routing it
always released on the same OS thread that created them. Do not call through the single worker would block the caller — usually the
this from within a callable already running ON the MLX worker thread FastAPI event loop, via /models/unload — until any in-flight
(e.g. from _reload_sync) — use _unload_model_sync directly there, or generation on that worker finishes. A generation still running keeps
this would deadlock the single-worker executor waiting on itself. its own reference to the model, so it completes normally and the
next generate() reloads via _reload_and_generate_sync.
""" """
# Submit unconditionally rather than checking self.model here first: self._unload_model_sync()
# that check would race against a load already queued on the worker
# thread (this call could see None, skip, and leave a model that
# finishes loading a moment later still resident). The loaded check
# belongs inside _unload_model_sync, where it runs atomically with
# the teardown itself.
_mlx_executor.submit(self._unload_model_sync).result()
def _unload_model_sync(self): def _unload_model_sync(self):
if self.model is not None: if self.model is not None:
@@ -375,13 +370,8 @@ class MLXSTTBackend:
logger.info("MLX Whisper model %s loaded successfully", model_size) logger.info("MLX Whisper model %s loaded successfully", model_size)
def unload_model(self): def unload_model(self):
"""Unload the model to free memory. """Unload the model to free memory (inline; see MLXTTSBackend.unload_model)."""
self._unload_model_sync()
Safe to call from any thread — see MLXTTSBackend.unload_model for why,
including why this submits unconditionally instead of checking
self.model first.
"""
_mlx_executor.submit(self._unload_model_sync).result()
def _unload_model_sync(self): def _unload_model_sync(self):
if self.model is not None: if self.model is not None: