fix(server): bound resident model memory with LRU eviction and idle TTL sweeper

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.

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)
 

mlx_audio/tests/test_server.py

@@ -1,3 +1,4 @@
+import asyncio
 import functools
 import io
 import json
@@ -945,4 +946,132 @@ 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