diff --git a/docs/guides/web-ui-api-server.md b/docs/guides/web-ui-api-server.md index 5250ae59..c6aefc53 100644 --- a/docs/guides/web-ui-api-server.md +++ b/docs/guides/web-ui-api-server.md @@ -34,9 +34,12 @@ The API will be available at `http://localhost:8000` and the Studio UI at `http: | `--realtime-transcription-delay-ms` | `null` | Transcription latency/quality knob for models that support it (e.g. `voxtral_realtime`) | | `--vad-model` | `mlx-community/silero-vad` | Streaming VAD model used for server-side turn detection (`server_vad`) on `/v1/realtime` | | `--tts-max-batch-size` | `8` | Maximum compatible TTS speech requests per continuous batch session | +| `--max-resident-models` | `0` | Max models kept loaded at once; loading another evicts the least-recently-used (LRU). `0` disables the bound (unbounded) | +| `--model-idle-ttl-seconds` | `0` | Unload models unused for this many seconds (0 disables the idle sweeper) | The two realtime flags also read from `MLX_AUDIO_REALTIME_MODEL` and `MLX_AUDIO_REALTIME_TRANSCRIPTION_DELAY_MS` if present; the CLI flags take precedence. `--vad-model` likewise reads from `MLX_AUDIO_VAD_MODEL`. The TTS batching flag also reads from `MLX_AUDIO_TTS_MAX_BATCH_SIZE`; the CLI flag takes precedence. +The memory bounds also read from `MLX_AUDIO_MAX_RESIDENT_MODELS` and `MLX_AUDIO_MODEL_IDLE_TTL_SECONDS`; the CLI flags take precedence. ### CORS Configuration @@ -160,6 +163,14 @@ curl -X POST "http://localhost:8000/v1/models?model_name=mlx-community/Kokoro-82 curl -X DELETE "http://localhost:8000/v1/models?model_name=mlx-community/Kokoro-82M-bf16" ``` +Loaded models are evicted automatically when `--max-resident-models` is set: once +the bound is reached, loading another model evicts the least-recently-used one +(`0`, the default, leaves the resident set unbounded — recommended when one server +mixes STT and TTS, which would otherwise evict each other's model), and +`--model-idle-ttl-seconds` unloads models unused for that long. Set either to suit +multi-voice / multi-model workloads so memory stays bounded without manual +`DELETE` calls. + ### Real-Time WebSocket Transcription The server exposes two WebSocket endpoints for live transcription. Both accept 16-bit signed little-endian PCM audio; they differ in wire protocol and intended consumers. diff --git a/mlx_audio/server.py b/mlx_audio/server.py index 299b0243..bbe6c875 100644 --- a/mlx_audio/server.py +++ b/mlx_audio/server.py @@ -8,11 +8,13 @@ import argparse import asyncio import base64 +import gc import inspect import io import json import os import subprocess +import threading import time import uuid import webbrowser @@ -91,27 +93,91 @@ def sanitize_for_json(obj: Any) -> Any: class ModelProvider: - def __init__(self): + def __init__( + self, + max_resident_models: int = 0, + idle_ttl_seconds: float = 0.0, + ): self.models: Dict[str, Dict[str, Any]] = {} - self.lock = asyncio.Lock() + self._last_used: Dict[str, float] = {} + # 0 = unbounded (never evict); >= 1 bounds the resident set and evicts + # least-recently-used when a new model would exceed it. + self.max_resident_models = max(0, max_resident_models) + self.idle_ttl_seconds = idle_ttl_seconds + self._lock = threading.Lock() + self._sweeper_task: Optional[asyncio.Task] = None def load_model(self, model_name: str): - if model_name not in self.models: - self.models[model_name] = load_model(model_name) - - return self.models[model_name] + with self._lock: + if model_name not in self.models: + if ( + self.max_resident_models > 0 + and len(self.models) >= self.max_resident_models + ): + self._evict_least_recently_used() + self.models[model_name] = load_model(model_name) + + self._last_used[model_name] = time.monotonic() + return self.models[model_name] async def remove_model(self, model_name: str) -> bool: - async with self.lock: + with self._lock: if model_name in self.models: - del self.models[model_name] + self._evict(model_name) return True return False async def get_available_models(self): - async with self.lock: + with self._lock: return list(self.models.keys()) + def _evict(self, model_name: str) -> None: + """Drop a loaded model and release its MLX memory (call with the lock held).""" + self.models.pop(model_name, None) + self._last_used.pop(model_name, None) + # Collect cyclical references (e.g. model <-> processor graphs) before + # clearing the MLX buffer cache, otherwise those buffers stay alive and + # mx.clear_cache() cannot release them. + gc.collect() + mx.clear_cache() + + def _evict_least_recently_used(self) -> Optional[str]: + """Evict the model unused for the longest time, returning its name.""" + if not self._last_used: + return None + victim = min(self._last_used, key=self._last_used.get) + self._evict(victim) + return victim + + def start_sweeper(self) -> None: + """Start the background idle-TTL sweeper (no-op when TTL is disabled).""" + if self.idle_ttl_seconds > 0 and self._sweeper_task is None: + self._sweeper_task = asyncio.create_task(self._idle_sweeper()) + + def stop_sweeper(self) -> None: + if self._sweeper_task is not None: + self._sweeper_task.cancel() + self._sweeper_task = None + + async def _idle_sweeper(self) -> None: + while True: + await asyncio.sleep(self.idle_ttl_seconds) + self.evict_stale(time.monotonic()) + + def evict_stale(self, now: float) -> List[str]: + """Evict every model unused for at least ``idle_ttl_seconds``.""" + if self.idle_ttl_seconds <= 0: + return [] + with self._lock: + stale = [ + name + for name, last_used in self._last_used.items() + if now - last_used >= self.idle_ttl_seconds + ] + for name in stale: + self._evict(name) + return stale + app = FastAPI() @@ -152,9 +218,11 @@ def setup_cors(app: FastAPI, allowed_origins: List[str]): @asynccontextmanager async def app_lifespan(app: FastAPI): del app + model_provider.start_sweeper() try: yield finally: + model_provider.stop_sweeper() global INFERENCE_BROKER if INFERENCE_BROKER is not None: INFERENCE_BROKER.stop_and_join() @@ -208,8 +276,26 @@ class SeparationResponse(BaseModel): sample_rate: int -# Initialize the ModelProvider -model_provider = ModelProvider() +# Initialize the ModelProvider. Memory bounds are configurable so multi-voice +# / multi-model servers can bound resident model memory without restarting. +def _env_int(name: str, default: int) -> int: + try: + return int(os.environ.get(name, default)) + except (TypeError, ValueError): + return default + + +def _env_float(name: str, default: float) -> float: + try: + return float(os.environ.get(name, default)) + except (TypeError, ValueError): + return default + + +model_provider = ModelProvider( + max_resident_models=_env_int("MLX_AUDIO_MAX_RESIDENT_MODELS", 0), + idle_ttl_seconds=_env_float("MLX_AUDIO_MODEL_IDLE_TTL_SECONDS", 0.0), +) REALTIME_INFERENCE_LOCK = asyncio.Lock() INFERENCE_BROKER: Optional[InferenceBroker] = None @@ -2118,6 +2204,27 @@ def main(): "Overrides $MLX_AUDIO_TTS_MAX_BATCH_SIZE." ), ) + parser.add_argument( + "--max-resident-models", + type=int, + default=None, + help=( + "Maximum number of models kept loaded in memory at once; loading " + "another evicts the least-recently-used one (LRU). 0 disables the " + "LRU bound (unbounded). " + "Overrides $MLX_AUDIO_MAX_RESIDENT_MODELS (default: 0)." + ), + ) + parser.add_argument( + "--model-idle-ttl-seconds", + type=float, + default=None, + help=( + "Unload models unused for this many seconds via a background sweeper " + "(0 disables the idle sweeper). " + "Overrides $MLX_AUDIO_MODEL_IDLE_TTL_SECONDS (default: 0)." + ), + ) args = parser.parse_args() if args.realtime_model: @@ -2130,6 +2237,21 @@ def main(): os.environ["MLX_AUDIO_VAD_MODEL"] = args.vad_model if args.tts_max_batch_size is not None: os.environ["MLX_AUDIO_TTS_MAX_BATCH_SIZE"] = str(args.tts_max_batch_size) + # The provider is constructed at import time, so apply the memory bounds + # directly (CLI flags take precedence over the environment). + model_provider.max_resident_models = max( + 0, + ( + args.max_resident_models + if args.max_resident_models is not None + else _env_int("MLX_AUDIO_MAX_RESIDENT_MODELS", 0) + ), + ) + model_provider.idle_ttl_seconds = ( + args.model_idle_ttl_seconds + if args.model_idle_ttl_seconds is not None + else _env_float("MLX_AUDIO_MODEL_IDLE_TTL_SECONDS", 0.0) + ) setup_cors(app, args.allowed_origins) diff --git a/mlx_audio/tests/test_server.py b/mlx_audio/tests/test_server.py index c8e071ad..624c2157 100644 --- a/mlx_audio/tests/test_server.py +++ b/mlx_audio/tests/test_server.py @@ -1,3 +1,4 @@ +import asyncio import functools import io import json @@ -945,4 +946,130 @@ async def test_stream_inference_results_reraises_error_by_default(): pass assert exc_info.value is error - assert handle.cancelled + + +# -- ModelProvider memory eviction ------------------------------------------ + + +def _make_provider(max_resident=1, idle_ttl=0.0): + from mlx_audio.server import ModelProvider + + return ModelProvider(max_resident_models=max_resident, idle_ttl_seconds=idle_ttl) + + +def test_model_provider_lru_evicts_least_recently_used(monkeypatch): + loaded = [] + monkeypatch.setattr( + "mlx_audio.server.load_model", + lambda name: loaded.append(name) or f"model-{name}", + ) + provider = _make_provider(max_resident=2) + + provider.load_model("a") + provider.load_model("b") + provider.load_model("a") # a becomes the most recently used + provider.load_model("c") # evicts b (LRU), keeps a + + assert set(provider.models) == {"a", "c"} + assert set(provider._last_used.keys()) == {"a", "c"} + + +def test_model_provider_evicts_when_capacity_exceeded(monkeypatch): + loaded = [] + monkeypatch.setattr( + "mlx_audio.server.load_model", + lambda name: loaded.append(name) or f"model-{name}", + ) + provider = _make_provider(max_resident=1) + + provider.load_model("a") + provider.load_model("b") + + assert list(provider.models.keys()) == ["b"] + assert loaded == ["a", "b"] + + +async def test_model_provider_remove_model_clears_memory(monkeypatch): + cleared = [] + monkeypatch.setattr("mlx_audio.server.load_model", lambda name: f"model-{name}") + monkeypatch.setattr("mlx_audio.server.mx.clear_cache", lambda: cleared.append(1)) + provider = _make_provider() + + provider.load_model("a") + assert "a" in provider.models + + removed = await provider.remove_model("a") + + assert removed is True + assert provider.models == {} + assert cleared == [1] + assert "a" not in provider._last_used + + +def test_model_provider_evict_collects_garbage_before_clearing_cache(monkeypatch): + order = [] + monkeypatch.setattr("mlx_audio.server.load_model", lambda name: f"model-{name}") + monkeypatch.setattr( + "mlx_audio.server.gc.collect", lambda: order.append("gc.collect") + ) + monkeypatch.setattr( + "mlx_audio.server.mx.clear_cache", lambda: order.append("mx.clear_cache") + ) + provider = _make_provider() + + provider.load_model("a") + provider.load_model("b") + + assert order == ["gc.collect", "mx.clear_cache"] + + +async def test_model_provider_remove_model_missing_returns_false(): + provider = _make_provider() + assert await provider.remove_model("missing") is False + + +async def test_model_provider_idle_sweeper_evicts_stale_models(monkeypatch): + loaded = [] + monkeypatch.setattr( + "mlx_audio.server.load_model", + lambda name: loaded.append(name) or f"model-{name}", + ) + clock = {"now": 0.0} + monkeypatch.setattr("mlx_audio.server.time.monotonic", lambda: clock["now"]) + provider = _make_provider(idle_ttl=0.05) + provider.start_sweeper() + try: + provider.load_model("a") + clock["now"] = 0.06 + stale = provider.evict_stale(clock["now"]) + assert stale == ["a"] + assert "a" not in provider.models + + provider.load_model("b") + clock["now"] = 0.10 # b is only 0.04s old — still fresh + assert provider.evict_stale(clock["now"]) == [] + assert "b" in provider.models + + clock["now"] = 0.20 # b is now 0.14s old — stale + assert provider.evict_stale(clock["now"]) == ["b"] + assert "b" not in provider.models + finally: + provider.stop_sweeper() + + +def test_model_provider_sweeper_is_noop_when_ttl_disabled(monkeypatch): + provider = _make_provider(idle_ttl=0.0) + provider.start_sweeper() + assert provider._sweeper_task is None + provider.stop_sweeper() + + +def test_model_provider_zero_max_resident_is_unbounded(monkeypatch): + monkeypatch.setattr("mlx_audio.server.load_model", lambda name: f"model-{name}") + provider = _make_provider(max_resident=0) + assert provider.max_resident_models == 0 + + for i in range(5): + provider.load_model(f"m{i}") + + assert len(provider.models) == 5