From 59932cf07ceb086dd610519bf55eea9714192d37 Mon Sep 17 00:00:00 2001 From: brycehenson Date: Mon, 17 Aug 2026 04:25:34 +0000 Subject: [PATCH 01/19] Add model auto-unload controls --- README.md | 10 ++- api/src/core/config.py | 4 + api/src/inference/model_manager.py | 128 +++++++++++++++++++++++++++-- api/src/routers/development.py | 47 +++++++++++ api/tests/test_model_unload.py | 110 ++++++++++++++++++++++++- docker/gpu/docker-compose.yml | 5 +- docs/configuration.md | 4 +- 7 files changed, 295 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 939092ed..543436ca 100644 --- a/README.md +++ b/README.md @@ -698,7 +698,11 @@ Key Performance Metrics: ### Model Unload / VRAM Reclaim -`POST /dev/unload` frees the model from VRAM and reloads lazily on the next request. Reclaim scales with load (the activation pool, not just weights) but plateaus: chunks cap at 450 tokens. Long-form = ~30 paragraphs. Same setup as above. +`POST /dev/unload` frees the model from VRAM and reloads lazily on the next request. Set `ALLOW_DEV_UNLOAD=true` to expose the lifecycle controls: `GET /dev/model`, `POST /dev/unload`, and `POST /dev/reload`. + +For shared-GPU hosts, set `MODEL_AUTO_UNLOAD_ENABLED=true` and tune `MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS` to unload automatically after the model has been idle. In-flight generation keeps the model loaded; the next request reloads it automatically, or `POST /dev/reload` can pre-warm it. + +Reclaim scales with load (the activation pool, not just weights) but plateaus: chunks cap at 450 tokens. Long-form = ~30 paragraphs. Same setup as above.

Short workload @@ -759,7 +763,9 @@ System state and resource usage, for debugging exhaustion or performance issues. - `/debug/threads` - Get thread information and stack traces - `/debug/storage` - Disk usage per mounted partition - `/debug/system` - Get system information (CPU, memory, GPU) -- `POST /dev/unload` - Release model from VRAM; reloads lazily on next request. Off by default; set `ALLOW_DEV_UNLOAD=true` to enable +- `/dev/model` - Get model load state and auto-unload timing +- `POST /dev/unload` - Release model from VRAM; reloads lazily on next request +- `POST /dev/reload` - Load the model immediately after an unload or before traffic arrives Stability: the `/v1/*` OpenAI-compatible routes are the stable API. `/dev/*` and `/debug/*` are operational helpers, and may change or move behind flags between minor releases. diff --git a/api/src/core/config.py b/api/src/core/config.py index 9d7f77ba..8c5ea36b 100644 --- a/api/src/core/config.py +++ b/api/src/core/config.py @@ -40,6 +40,10 @@ class Settings(BaseSettings): False # Whether to allow saving combined voices locally ) allow_dev_unload: bool = False # Whether to expose the POST /dev/unload endpoint + model_auto_unload_enabled: bool = False # Whether to unload the model after an idle timeout + model_auto_unload_timeout_seconds: float = ( + 300.0 # Idle seconds before unloading when MODEL_AUTO_UNLOAD_ENABLED=true + ) enable_debug_endpoints: bool = ( False # Whether to expose /debug/* host and process introspection routes ) diff --git a/api/src/inference/model_manager.py b/api/src/inference/model_manager.py index b1468717..8b025909 100644 --- a/api/src/inference/model_manager.py +++ b/api/src/inference/model_manager.py @@ -1,6 +1,7 @@ """Kokoro V1 model management.""" import asyncio +import time from typing import Optional import torch @@ -29,6 +30,9 @@ def __init__(self, config: Optional[ModelConfig] = None): self._backend: Optional[KokoroV1] = None # Explicitly type as KokoroV1 self._device: Optional[str] = None self._lock = asyncio.Lock() + self._active_requests = 0 + self._last_used_at: Optional[float] = None + self._idle_unload_task: Optional[asyncio.Task] = None def _determine_device(self) -> str: """Determine device based on settings.""" @@ -102,11 +106,12 @@ async def initialize_with_warmup(self, voice_manager) -> tuple[str, str, int]: raise RuntimeError(f"Warmup failed: {e}") async def ensure_backend(self) -> None: - """Reload the backend if it was unloaded via /dev/unload.""" + """Reload the backend if it was unloaded.""" if self._backend: return async with self._lock: if not self._backend: + self._cancel_idle_unload_timer() await self.initialize() await self.load_model(self._config.pytorch_kokoro_v1_file) @@ -137,30 +142,110 @@ async def load_model(self, path: str) -> None: try: await self._backend.load_model(path) + self._last_used_at = time.monotonic() except FileNotFoundError as e: raise e except Exception as e: raise RuntimeError(f"Failed to load model: {e}") + def _auto_unload_timeout(self) -> float: + return max(0.0, float(settings.model_auto_unload_timeout_seconds)) + + def _auto_unload_enabled(self) -> bool: + return settings.model_auto_unload_enabled and self._auto_unload_timeout() > 0 + + def _cancel_idle_unload_timer(self) -> None: + try: + current_task = asyncio.current_task() + except RuntimeError: + current_task = None + if ( + self._idle_unload_task + and not self._idle_unload_task.done() + and self._idle_unload_task is not current_task + ): + self._idle_unload_task.cancel() + self._idle_unload_task = None + + def _unload_backend_locked(self) -> bool: + if self._backend is None: + return False + self._backend.unload() + self._backend = None + self._last_used_at = time.monotonic() + return True + + def _schedule_idle_unload_timer_locked(self) -> None: + self._cancel_idle_unload_timer() + if ( + not self._auto_unload_enabled() + or self._backend is None + or self._active_requests > 0 + ): + return + + timeout = self._auto_unload_timeout() + self._idle_unload_task = asyncio.create_task(self._idle_unload_after(timeout)) + + async def _idle_unload_after(self, timeout: float) -> None: + try: + await asyncio.sleep(timeout) + async with self._lock: + if ( + not self._auto_unload_enabled() + or self._backend is None + or self._active_requests > 0 + or self._last_used_at is None + ): + return + + idle_for = time.monotonic() - self._last_used_at + if idle_for < self._auto_unload_timeout(): + self._schedule_idle_unload_timer_locked() + return + + unloaded = self._unload_backend_locked() + + if unloaded: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + logger.info("Model auto-unloaded after idle timeout") + except asyncio.CancelledError: + pass + + async def _begin_request(self) -> None: + async with self._lock: + self._active_requests += 1 + self._cancel_idle_unload_timer() + + async def _end_request(self) -> None: + async with self._lock: + self._active_requests = max(0, self._active_requests - 1) + self._last_used_at = time.monotonic() + self._schedule_idle_unload_timer_locked() + async def generate(self, *args, **kwargs): """Generate audio using initialized backend. Raises: RuntimeError: If generation fails """ - await self.ensure_backend() - assert self._backend is not None, "ensure_backend left no backend" - + await self._begin_request() try: + await self.ensure_backend() + assert self._backend is not None, "ensure_backend left no backend" async for chunk in self._backend.generate(*args, **kwargs): if settings.default_volume_multiplier != 1.0: chunk.audio *= settings.default_volume_multiplier yield chunk except Exception as e: raise RuntimeError(f"Generation failed: {e}") + finally: + await self._end_request() def unload_all(self) -> None: """Unload model and free resources.""" + self._cancel_idle_unload_timer() if self._backend: self._backend.unload() self._backend = None @@ -168,13 +253,42 @@ def unload_all(self) -> None: async def unload(self) -> None: """Release model from GPU memory. Reloads automatically on next request.""" async with self._lock: - if self._backend is not None: - self._backend.unload() - self._backend = None + self._cancel_idle_unload_timer() + self._unload_backend_locked() if torch.cuda.is_available(): torch.cuda.empty_cache() logger.info("Model unloaded from GPU memory") + async def reload(self) -> None: + """Reload the model immediately.""" + async with self._lock: + self._cancel_idle_unload_timer() + self._unload_backend_locked() + await self.initialize() + await self.load_model(self._config.pytorch_kokoro_v1_file) + self._schedule_idle_unload_timer_locked() + logger.info("Model reloaded") + + def status(self) -> dict: + """Return model lifecycle state for API responses.""" + timeout = self._auto_unload_timeout() + idle_for = None + unload_in = None + if self._last_used_at is not None: + idle_for = max(0.0, time.monotonic() - self._last_used_at) + if self._auto_unload_enabled() and self._backend is not None: + unload_in = max(0.0, timeout - idle_for) + return { + "backend": self.current_backend, + "device": self._device, + "loaded": self._backend is not None, + "active_requests": self._active_requests, + "auto_unload_enabled": settings.model_auto_unload_enabled, + "auto_unload_timeout_seconds": timeout, + "idle_seconds": idle_for, + "seconds_until_auto_unload": unload_in, + } + @property def current_backend(self) -> str: """Get current backend type.""" diff --git a/api/src/routers/development.py b/api/src/routers/development.py index bc2e1eab..7a23ad59 100644 --- a/api/src/routers/development.py +++ b/api/src/routers/development.py @@ -505,3 +505,50 @@ async def unload_model( except Exception as e: logger.error(f"Error unloading model: {e}") raise HTTPException(status_code=500, detail={"error": str(e)}) + + +@router.post("/dev/reload") +async def reload_model( + tts_service: TTSService = Depends(get_tts_service), +): + """Reload the model immediately. + + Normal inference also reloads lazily after an unload; this endpoint is useful + when you want to pre-warm the model before the next request. + """ + if not settings.allow_dev_unload: + raise HTTPException( + status_code=403, + detail={"error": "The /dev/reload endpoint is disabled"}, + ) + try: + if tts_service.model_manager is None: + raise HTTPException( + status_code=503, detail={"error": "Model manager not initialized"} + ) + await tts_service.model_manager.reload() + return JSONResponse( + {"status": "loaded", "model": tts_service.model_manager.status()} + ) + except HTTPException: + raise + except Exception as e: + logger.error(f"Error reloading model: {e}") + raise HTTPException(status_code=500, detail={"error": str(e)}) + + +@router.get("/dev/model") +async def model_status( + tts_service: TTSService = Depends(get_tts_service), +): + """Return model load state and auto-unload settings.""" + if not settings.allow_dev_unload: + raise HTTPException( + status_code=403, + detail={"error": "The /dev/model endpoint is disabled"}, + ) + if tts_service.model_manager is None: + raise HTTPException( + status_code=503, detail={"error": "Model manager not initialized"} + ) + return JSONResponse(tts_service.model_manager.status()) diff --git a/api/tests/test_model_unload.py b/api/tests/test_model_unload.py index 05da4353..a7faee37 100644 --- a/api/tests/test_model_unload.py +++ b/api/tests/test_model_unload.py @@ -1,4 +1,4 @@ -"""Tests for ModelManager.unload(), lazy reinit in generate(), and POST /dev/unload.""" +"""Tests for model unload, auto-unload, lazy reload, and dev lifecycle endpoints.""" import asyncio from contextlib import contextmanager @@ -36,6 +36,8 @@ async def _override(): def _enable_dev_unload(monkeypatch): """Enable the /dev/unload gate for the endpoint tests in this module.""" monkeypatch.setattr(settings, "allow_dev_unload", True) + monkeypatch.setattr(settings, "model_auto_unload_enabled", False) + monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 300.0) # --------------------------------------------------------------------------- @@ -46,6 +48,7 @@ def _enable_dev_unload(monkeypatch): def test_manager_init_creates_lock(): manager = ModelManager() assert isinstance(manager._lock, asyncio.Lock) + assert manager._active_requests == 0 @pytest.mark.asyncio @@ -185,6 +188,85 @@ async def fake_generate(*args, **kwargs): assert len(chunks) == 1 +@pytest.mark.asyncio +async def test_generate_schedules_idle_unload_when_enabled(monkeypatch): + """Finished generation schedules model unload after the configured idle period.""" + monkeypatch.setattr(settings, "model_auto_unload_enabled", True) + monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 0.01) + + manager = ModelManager() + mock_backend = MagicMock() + audio_chunk = AudioChunk(np.zeros(10, dtype=np.float32)) + + async def fake_generate(*args, **kwargs): + yield audio_chunk + + mock_backend.generate = fake_generate + manager._backend = mock_backend + + with patch("api.src.inference.model_manager.torch") as mock_torch: + mock_torch.cuda.is_available.return_value = False + chunks = [] + async for chunk in manager.generate("hello", ("voice", "/path/voice.pt")): + chunks.append(chunk) + await asyncio.sleep(0.03) + + assert len(chunks) == 1 + mock_backend.unload.assert_called_once() + assert manager._backend is None + + +@pytest.mark.asyncio +async def test_active_request_blocks_idle_unload(monkeypatch): + """The idle timer does not unload while generation is still active.""" + monkeypatch.setattr(settings, "model_auto_unload_enabled", True) + monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 0.01) + + manager = ModelManager() + mock_backend = MagicMock() + + async def slow_generate(*args, **kwargs): + await asyncio.sleep(0.03) + yield AudioChunk(np.zeros(10, dtype=np.float32)) + + mock_backend.generate = slow_generate + manager._backend = mock_backend + + with patch("api.src.inference.model_manager.torch") as mock_torch: + mock_torch.cuda.is_available.return_value = False + chunks = [] + async for chunk in manager.generate("hello", ("voice", "/path/voice.pt")): + assert manager._active_requests == 1 + mock_backend.unload.assert_not_called() + chunks.append(chunk) + + assert len(chunks) == 1 + assert manager._active_requests == 0 + assert manager._idle_unload_task is not None + manager._cancel_idle_unload_timer() + + +def test_status_reports_model_lifecycle_state(monkeypatch): + monkeypatch.setattr(settings, "model_auto_unload_enabled", True) + monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 30.0) + manager = ModelManager() + manager._backend = MagicMock() + manager._device = "cuda" + manager._last_used_at = 10.0 + + with patch("api.src.inference.model_manager.time.monotonic", return_value=15.0): + status = manager.status() + + assert status["backend"] == "kokoro_v1" + assert status["device"] == "cuda" + assert status["loaded"] is True + assert status["active_requests"] == 0 + assert status["auto_unload_enabled"] is True + assert status["auto_unload_timeout_seconds"] == 30.0 + assert status["idle_seconds"] == 5.0 + assert status["seconds_until_auto_unload"] == 25.0 + + # --------------------------------------------------------------------------- # POST /dev/unload endpoint tests # --------------------------------------------------------------------------- @@ -261,3 +343,29 @@ def test_unload_endpoint_500_on_exception(): assert response.status_code == 500 assert "GPU exploded" in response.json()["detail"]["error"] + + +def test_reload_endpoint_returns_status(): + mock_manager = MagicMock() + mock_manager.reload = AsyncMock() + mock_manager.status.return_value = {"loaded": True} + service = _mock_service(manager=mock_manager) + + with override_tts_service(service): + response = client.post("/dev/reload") + + assert response.status_code == 200 + assert response.json() == {"status": "loaded", "model": {"loaded": True}} + mock_manager.reload.assert_called_once() + + +def test_model_status_endpoint_returns_status(): + mock_manager = MagicMock() + mock_manager.status.return_value = {"loaded": False} + service = _mock_service(manager=mock_manager) + + with override_tts_service(service): + response = client.get("/dev/model") + + assert response.status_code == 200 + assert response.json() == {"loaded": False} diff --git a/docker/gpu/docker-compose.yml b/docker/gpu/docker-compose.yml index 1a917138..ae83b166 100644 --- a/docker/gpu/docker-compose.yml +++ b/docker/gpu/docker-compose.yml @@ -28,7 +28,9 @@ services: - PYTHONUNBUFFERED=1 - API_LOG_LEVEL=DEBUG - DOWNLOAD_MODEL=true - # - ALLOW_DEV_UNLOAD=true + - ALLOW_DEV_UNLOAD=true + - MODEL_AUTO_UNLOAD_ENABLED=true + - MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS=300 # - ENABLE_DEBUG_ENDPOINTS=true deploy: resources: @@ -37,4 +39,3 @@ services: - driver: nvidia count: all capabilities: [gpu] - diff --git a/docs/configuration.md b/docs/configuration.md index a3d637c2..80914f3a 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -153,7 +153,9 @@ Names are the field names from `api/src/core/config.py`, uppercased. Unrecognize | Variable | Default | | |---|---|---| | `ENABLE_DEBUG_ENDPOINTS` | `false` | Expose `/debug/*` host and process introspection | -| `ALLOW_DEV_UNLOAD` | `false` | Expose `POST /dev/unload` | +| `ALLOW_DEV_UNLOAD` | `false` | Expose `/dev/model`, `POST /dev/unload`, and `POST /dev/reload` | +| `MODEL_AUTO_UNLOAD_ENABLED` | `false` | Unload the model after the idle timeout | +| `MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS` | `300.0` | Idle seconds before auto-unload when enabled | ## Logging From 543f485c6f63bc873fc15cf11eb5bd3a629755a7 Mon Sep 17 00:00:00 2001 From: brycehenson Date: Mon, 17 Aug 2026 04:36:48 +0000 Subject: [PATCH 02/19] Schedule auto-unload after model load --- api/src/inference/model_manager.py | 1 + api/tests/test_model_unload.py | 44 ++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/api/src/inference/model_manager.py b/api/src/inference/model_manager.py index 8b025909..2a0adcf1 100644 --- a/api/src/inference/model_manager.py +++ b/api/src/inference/model_manager.py @@ -143,6 +143,7 @@ async def load_model(self, path: str) -> None: try: await self._backend.load_model(path) self._last_used_at = time.monotonic() + self._schedule_idle_unload_timer_locked() except FileNotFoundError as e: raise e except Exception as e: diff --git a/api/tests/test_model_unload.py b/api/tests/test_model_unload.py index a7faee37..04abb485 100644 --- a/api/tests/test_model_unload.py +++ b/api/tests/test_model_unload.py @@ -216,6 +216,50 @@ async def fake_generate(*args, **kwargs): assert manager._backend is None +@pytest.mark.asyncio +async def test_load_model_schedules_idle_unload_when_enabled(monkeypatch): + """Startup-style model loads also schedule unload without waiting for traffic.""" + monkeypatch.setattr(settings, "model_auto_unload_enabled", True) + monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 0.01) + + manager = ModelManager() + mock_backend = MagicMock() + mock_backend.load_model = AsyncMock() + manager._backend = mock_backend + + with patch("api.src.inference.model_manager.torch") as mock_torch: + mock_torch.cuda.is_available.return_value = False + await manager.load_model("/path/model.pt") + await asyncio.sleep(0.03) + + mock_backend.load_model.assert_called_once_with("/path/model.pt") + mock_backend.unload.assert_called_once() + assert manager._backend is None + + +@pytest.mark.asyncio +async def test_load_model_does_not_schedule_idle_unload_during_active_request( + monkeypatch, +): + """Lazy loads during generation wait for request completion before scheduling.""" + monkeypatch.setattr(settings, "model_auto_unload_enabled", True) + monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 0.01) + + manager = ModelManager() + mock_backend = MagicMock() + mock_backend.load_model = AsyncMock() + manager._backend = mock_backend + manager._active_requests = 1 + + await manager.load_model("/path/model.pt") + await asyncio.sleep(0.03) + + mock_backend.load_model.assert_called_once_with("/path/model.pt") + mock_backend.unload.assert_not_called() + assert manager._backend is mock_backend + assert manager._idle_unload_task is None + + @pytest.mark.asyncio async def test_active_request_blocks_idle_unload(monkeypatch): """The idle timer does not unload while generation is still active.""" From ec7a1b41cb4d6b55880d58a91197d4c1f7d50a68 Mon Sep 17 00:00:00 2001 From: brycehenson Date: Mon, 17 Aug 2026 06:01:05 +0000 Subject: [PATCH 03/19] unload to cpu --- README.md | 2 +- api/src/core/config.py | 1 + api/src/inference/kokoro_v1.py | 62 +++++++++++++-- api/src/inference/model_manager.py | 116 +++++++++++++++++++++++++---- api/tests/test_kokoro_v1.py | 41 ++++++++++ api/tests/test_model_unload.py | 63 ++++++++++++++++ docker/gpu/docker-compose.yml | 1 + docs/configuration.md | 1 + notes.md | 85 +++++++++++++++++++++ 9 files changed, 350 insertions(+), 22 deletions(-) create mode 100644 notes.md diff --git a/README.md b/README.md index 543436ca..d0102e58 100644 --- a/README.md +++ b/README.md @@ -700,7 +700,7 @@ Key Performance Metrics: `POST /dev/unload` frees the model from VRAM and reloads lazily on the next request. Set `ALLOW_DEV_UNLOAD=true` to expose the lifecycle controls: `GET /dev/model`, `POST /dev/unload`, and `POST /dev/reload`. -For shared-GPU hosts, set `MODEL_AUTO_UNLOAD_ENABLED=true` and tune `MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS` to unload automatically after the model has been idle. In-flight generation keeps the model loaded; the next request reloads it automatically, or `POST /dev/reload` can pre-warm it. +For shared-GPU hosts, set `MODEL_AUTO_UNLOAD_ENABLED=true` and tune `MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS` to unload automatically after the model has been idle. In-flight generation keeps the model loaded; the next request reloads it automatically, or `POST /dev/reload` can pre-warm it. Set `MODEL_UNLOAD_STRATEGY=cpu_cache` to keep model weights in system RAM for faster reload while still clearing GPU memory; the default `destroy` strategy releases model objects completely. Reclaim scales with load (the activation pool, not just weights) but plateaus: chunks cap at 450 tokens. Long-form = ~30 paragraphs. Same setup as above. diff --git a/api/src/core/config.py b/api/src/core/config.py index 8c5ea36b..fc3f3fc6 100644 --- a/api/src/core/config.py +++ b/api/src/core/config.py @@ -44,6 +44,7 @@ class Settings(BaseSettings): model_auto_unload_timeout_seconds: float = ( 300.0 # Idle seconds before unloading when MODEL_AUTO_UNLOAD_ENABLED=true ) + model_unload_strategy: str = "destroy" # "destroy" or "cpu_cache" enable_debug_endpoints: bool = ( False # Whether to expose /debug/* host and process introspection routes ) diff --git a/api/src/inference/kokoro_v1.py b/api/src/inference/kokoro_v1.py index cecc2584..e9b8ef3e 100644 --- a/api/src/inference/kokoro_v1.py +++ b/api/src/inference/kokoro_v1.py @@ -89,6 +89,7 @@ def __init__(self): # Strictly respect settings.use_gpu self._device = settings.get_device() self._model: Optional[KModel] = None + self._model_cpu_cached = False self._pipelines: Dict[str, KPipeline] = {} # Store pipelines by lang_code self._voice_cache: Dict[str, torch.Tensor] = {} # Cache voice tensors by path @@ -144,12 +145,53 @@ async def load_model(self, path: str) -> None: self._model = self._model.cuda() else: self._model = self._model.cpu() + self._model_cpu_cached = False except FileNotFoundError: raise except Exception as e: raise RuntimeError(f"Failed to load Kokoro model: {e}") + def _move_model_to_device(self) -> None: + """Move a CPU-cached model back to the configured inference device.""" + if self._model is None or not self._model_cpu_cached: + return + + logger.info(f"Moving CPU-cached Kokoro model back to {self._device}") + if self._device == "mps": + self._model = self._model.to(torch.device("mps")) + elif self._device == "cuda": + self._model = self._model.cuda() + torch.cuda.synchronize() + else: + self._model = self._model.cpu() + self._model_cpu_cached = False + logger.info(f"CPU-cached Kokoro model restored to {self._device}") + + def restore_to_device(self) -> None: + """Restore a CPU-cached model to the configured inference device.""" + self._move_model_to_device() + + def _clear_runtime_caches(self) -> None: + """Release cached objects that can hold device tensors.""" + for pipeline in self._pipelines.values(): + del pipeline + self._pipelines.clear() + self._voice_cache.clear() + + def _offload_model_to_cpu(self) -> bool: + """Move the model out of VRAM while retaining weights in system RAM.""" + if self._model is None or self._device not in {"cuda", "mps"}: + return False + + logger.info("Moving Kokoro model to CPU cache") + self._model = self._model.cpu() + self._model_cpu_cached = True + self._clear_runtime_caches() + self._clear_memory() + logger.info("Kokoro model offloaded to CPU cache and device caches cleared") + return True + def _get_pipeline(self, lang_code: str) -> KPipeline: """Get or create pipeline for language code. @@ -197,6 +239,7 @@ async def generate_from_tokens( raise RuntimeError("Model not loaded") try: + self._move_model_to_device() # Memory management for GPU if self._device == "cuda": if self._check_memory(): @@ -295,6 +338,7 @@ async def generate( if not self.is_loaded: raise RuntimeError("Model not loaded") try: + self._move_model_to_device() # Memory management for GPU if self._device == "cuda": if self._check_memory(): @@ -455,18 +499,26 @@ def _clear_memory(self) -> None: if hasattr(torch.mps, "empty_cache"): torch.mps.empty_cache() - def unload(self) -> None: + def unload(self, strategy: str = "destroy") -> None: """Unload model and free resources.""" + if strategy == "cpu_cache" and self._offload_model_to_cpu(): + return + + logger.info("Destroying Kokoro model backend state") if self._model is not None: del self._model self._model = None - for pipeline in self._pipelines.values(): - del pipeline - self._pipelines.clear() - self._voice_cache.clear() + self._model_cpu_cached = False + self._clear_runtime_caches() if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.synchronize() + logger.info("Kokoro model backend state destroyed") + + @property + def is_cpu_cached(self) -> bool: + """Check if model weights are retained in CPU RAM after device unload.""" + return self._model_cpu_cached @property def is_loaded(self) -> bool: diff --git a/api/src/inference/model_manager.py b/api/src/inference/model_manager.py index 2a0adcf1..f2b5317a 100644 --- a/api/src/inference/model_manager.py +++ b/api/src/inference/model_manager.py @@ -141,9 +141,16 @@ async def load_model(self, path: str) -> None: raise RuntimeError("Backend not initialized") try: + logger.info( + f"Loading model onto {self._device or 'configured device'} from {path}" + ) await self._backend.load_model(path) self._last_used_at = time.monotonic() self._schedule_idle_unload_timer_locked() + logger.info( + f"Model loaded; unload_strategy={self._unload_strategy()} " + f"auto_unload_enabled={settings.model_auto_unload_enabled}" + ) except FileNotFoundError as e: raise e except Exception as e: @@ -155,6 +162,25 @@ def _auto_unload_timeout(self) -> float: def _auto_unload_enabled(self) -> bool: return settings.model_auto_unload_enabled and self._auto_unload_timeout() > 0 + def _backend_is_loaded(self) -> bool: + if self._backend is None: + return False + return getattr(self._backend, "is_loaded", True) is not False + + def _backend_is_cpu_cached(self) -> bool: + if self._backend is None: + return False + return getattr(self._backend, "is_cpu_cached", False) is True + + def _unload_strategy(self) -> str: + strategy = settings.model_unload_strategy.strip().lower() + if strategy not in {"destroy", "cpu_cache"}: + logger.warning( + f"Unknown MODEL_UNLOAD_STRATEGY={settings.model_unload_strategy!r}; using destroy" + ) + return "destroy" + return strategy + def _cancel_idle_unload_timer(self) -> None: try: current_task = asyncio.current_task() @@ -165,27 +191,49 @@ def _cancel_idle_unload_timer(self) -> None: and not self._idle_unload_task.done() and self._idle_unload_task is not current_task ): + logger.debug("Cancelling pending model auto-unload timer") self._idle_unload_task.cancel() self._idle_unload_task = None def _unload_backend_locked(self) -> bool: if self._backend is None: + logger.info("Model unload requested, but no backend is loaded") return False - self._backend.unload() - self._backend = None + strategy = self._unload_strategy() + logger.info(f"Unloading model with strategy={strategy}") + self._backend.unload(strategy=strategy) + if strategy == "destroy" or not self._backend_is_loaded(): + self._backend = None + logger.info("Model unloaded and backend destroyed") + elif self._backend_is_cpu_cached(): + logger.info("Model offloaded from GPU and retained in CPU cache") + else: + logger.info("Model unload completed with backend still available") self._last_used_at = time.monotonic() return True def _schedule_idle_unload_timer_locked(self) -> None: self._cancel_idle_unload_timer() - if ( - not self._auto_unload_enabled() - or self._backend is None - or self._active_requests > 0 - ): + if not self._auto_unload_enabled(): + logger.debug("Model auto-unload timer not scheduled: disabled") + return + if self._backend is None: + logger.debug("Model auto-unload timer not scheduled: no backend loaded") + return + if self._backend_is_cpu_cached(): + logger.debug("Model auto-unload timer not scheduled: model is CPU cached") + return + if self._active_requests > 0: + logger.debug( + f"Model auto-unload timer not scheduled: {self._active_requests} active request(s)" + ) return timeout = self._auto_unload_timeout() + logger.info( + f"Scheduling model auto-unload in {timeout:.1f}s " + f"with strategy={self._unload_strategy()}" + ) self._idle_unload_task = asyncio.create_task(self._idle_unload_after(timeout)) async def _idle_unload_after(self, timeout: float) -> None: @@ -198,31 +246,50 @@ async def _idle_unload_after(self, timeout: float) -> None: or self._active_requests > 0 or self._last_used_at is None ): + logger.debug( + "Model auto-unload skipped when timer fired: " + f"enabled={self._auto_unload_enabled()} " + f"backend_loaded={self._backend is not None} " + f"active_requests={self._active_requests} " + f"last_used_at_set={self._last_used_at is not None}" + ) return idle_for = time.monotonic() - self._last_used_at if idle_for < self._auto_unload_timeout(): + logger.debug( + f"Model auto-unload timer fired early after {idle_for:.1f}s idle; rescheduling" + ) self._schedule_idle_unload_timer_locked() return + logger.info(f"Model idle for {idle_for:.1f}s; auto-unloading now") unloaded = self._unload_backend_locked() if unloaded: if torch.cuda.is_available(): torch.cuda.empty_cache() - logger.info("Model auto-unloaded after idle timeout") + logger.info( + f"Model auto-unload completed with strategy={self._unload_strategy()}" + ) except asyncio.CancelledError: - pass + logger.debug("Model auto-unload timer cancelled") async def _begin_request(self) -> None: async with self._lock: self._active_requests += 1 self._cancel_idle_unload_timer() + logger.debug( + f"Model request started; active_requests={self._active_requests}" + ) async def _end_request(self) -> None: async with self._lock: self._active_requests = max(0, self._active_requests - 1) self._last_used_at = time.monotonic() + logger.debug( + f"Model request finished; active_requests={self._active_requests}" + ) self._schedule_idle_unload_timer_locked() async def generate(self, *args, **kwargs): @@ -253,22 +320,33 @@ def unload_all(self) -> None: async def unload(self) -> None: """Release model from GPU memory. Reloads automatically on next request.""" + logger.info("Manual model unload requested") async with self._lock: self._cancel_idle_unload_timer() self._unload_backend_locked() if torch.cuda.is_available(): torch.cuda.empty_cache() - logger.info("Model unloaded from GPU memory") + logger.info("Manual model unload completed") async def reload(self) -> None: """Reload the model immediately.""" + logger.info("Manual model reload requested") async with self._lock: self._cancel_idle_unload_timer() - self._unload_backend_locked() - await self.initialize() - await self.load_model(self._config.pytorch_kokoro_v1_file) + if ( + self._unload_strategy() == "cpu_cache" + and self._backend is not None + and self._backend_is_loaded() + ): + logger.info("Restoring model from CPU cache") + self._backend.restore_to_device() + self._last_used_at = time.monotonic() + else: + self._unload_backend_locked() + await self.initialize() + await self.load_model(self._config.pytorch_kokoro_v1_file) self._schedule_idle_unload_timer_locked() - logger.info("Model reloaded") + logger.info("Manual model reload completed") def status(self) -> dict: """Return model lifecycle state for API responses.""" @@ -277,13 +355,19 @@ def status(self) -> dict: unload_in = None if self._last_used_at is not None: idle_for = max(0.0, time.monotonic() - self._last_used_at) - if self._auto_unload_enabled() and self._backend is not None: + if ( + self._auto_unload_enabled() + and self._backend_is_loaded() + and not self._backend_is_cpu_cached() + ): unload_in = max(0.0, timeout - idle_for) return { "backend": self.current_backend, "device": self._device, - "loaded": self._backend is not None, + "loaded": self._backend_is_loaded() and not self._backend_is_cpu_cached(), + "cpu_cached": self._backend_is_loaded() and self._backend_is_cpu_cached(), "active_requests": self._active_requests, + "unload_strategy": self._unload_strategy(), "auto_unload_enabled": settings.model_auto_unload_enabled, "auto_unload_timeout_seconds": timeout, "idle_seconds": idle_for, diff --git a/api/tests/test_kokoro_v1.py b/api/tests/test_kokoro_v1.py index 4aa62184..a30a4d1b 100644 --- a/api/tests/test_kokoro_v1.py +++ b/api/tests/test_kokoro_v1.py @@ -68,6 +68,47 @@ def test_unload_with_pipelines(kokoro_backend): assert kokoro_backend._voice_cache == {} # Voice tensors should be released +def test_cpu_cache_unload_moves_model_to_cpu_and_clears_runtime_caches( + kokoro_backend, +): + """CPU-cache unload keeps model weights but clears device-backed runtime state.""" + cuda_model = MagicMock() + cpu_model = MagicMock() + cuda_model.cpu.return_value = cpu_model + kokoro_backend._model = cuda_model + kokoro_backend._device = "cuda" + kokoro_backend._pipelines = {"a": MagicMock()} + kokoro_backend._voice_cache = {"af_heart.pt:cuda": MagicMock()} + + with patch.object(kokoro_backend, "_clear_memory") as mock_clear: + kokoro_backend.unload(strategy="cpu_cache") + + cuda_model.cpu.assert_called_once() + mock_clear.assert_called_once() + assert kokoro_backend._model is cpu_model + assert kokoro_backend.is_loaded + assert kokoro_backend.is_cpu_cached + assert kokoro_backend._pipelines == {} + assert kokoro_backend._voice_cache == {} + + +def test_restore_to_device_moves_cpu_cached_model_to_cuda(kokoro_backend): + cpu_model = MagicMock() + cuda_model = MagicMock() + cpu_model.cuda.return_value = cuda_model + kokoro_backend._model = cpu_model + kokoro_backend._device = "cuda" + kokoro_backend._model_cpu_cached = True + + with patch("api.src.inference.kokoro_v1.torch") as mock_torch: + kokoro_backend.restore_to_device() + + cpu_model.cuda.assert_called_once() + mock_torch.cuda.synchronize.assert_called_once() + assert kokoro_backend._model is cuda_model + assert not kokoro_backend.is_cpu_cached + + @pytest.mark.asyncio async def test_generate_validation(kokoro_backend): """Test generation validation.""" diff --git a/api/tests/test_model_unload.py b/api/tests/test_model_unload.py index 04abb485..c31ae1e5 100644 --- a/api/tests/test_model_unload.py +++ b/api/tests/test_model_unload.py @@ -38,6 +38,7 @@ def _enable_dev_unload(monkeypatch): monkeypatch.setattr(settings, "allow_dev_unload", True) monkeypatch.setattr(settings, "model_auto_unload_enabled", False) monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 300.0) + monkeypatch.setattr(settings, "model_unload_strategy", "destroy") # --------------------------------------------------------------------------- @@ -65,6 +66,45 @@ async def test_unload_clears_backend(): assert manager._backend is None +@pytest.mark.asyncio +async def test_unload_cpu_cache_keeps_backend(monkeypatch): + monkeypatch.setattr(settings, "model_unload_strategy", "cpu_cache") + + manager = ModelManager() + mock_backend = MagicMock() + mock_backend.is_loaded = True + mock_backend.is_cpu_cached = True + manager._backend = mock_backend + + with patch("api.src.inference.model_manager.torch") as mock_torch: + mock_torch.cuda.is_available.return_value = False + await manager.unload() + + mock_backend.unload.assert_called_once_with(strategy="cpu_cache") + assert manager._backend is mock_backend + + +@pytest.mark.asyncio +async def test_reload_restores_cpu_cached_backend(monkeypatch): + monkeypatch.setattr(settings, "model_unload_strategy", "cpu_cache") + + manager = ModelManager() + mock_backend = MagicMock() + mock_backend.is_loaded = True + mock_backend.is_cpu_cached = True + manager._backend = mock_backend + + with ( + patch.object(manager, "initialize", new_callable=AsyncMock) as mock_init, + patch.object(manager, "load_model", new_callable=AsyncMock) as mock_load, + ): + await manager.reload() + + mock_backend.restore_to_device.assert_called_once() + mock_init.assert_not_called() + mock_load.assert_not_called() + + @pytest.mark.asyncio async def test_unload_when_already_none_is_noop(): manager = ModelManager() @@ -305,12 +345,35 @@ def test_status_reports_model_lifecycle_state(monkeypatch): assert status["device"] == "cuda" assert status["loaded"] is True assert status["active_requests"] == 0 + assert status["unload_strategy"] == "destroy" + assert status["cpu_cached"] is False assert status["auto_unload_enabled"] is True assert status["auto_unload_timeout_seconds"] == 30.0 assert status["idle_seconds"] == 5.0 assert status["seconds_until_auto_unload"] == 25.0 +def test_status_reports_cpu_cached_state(monkeypatch): + monkeypatch.setattr(settings, "model_unload_strategy", "cpu_cache") + monkeypatch.setattr(settings, "model_auto_unload_enabled", True) + monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 30.0) + + manager = ModelManager() + mock_backend = MagicMock() + mock_backend.is_loaded = True + mock_backend.is_cpu_cached = True + manager._backend = mock_backend + manager._last_used_at = 10.0 + + with patch("api.src.inference.model_manager.time.monotonic", return_value=15.0): + status = manager.status() + + assert status["loaded"] is False + assert status["cpu_cached"] is True + assert status["unload_strategy"] == "cpu_cache" + assert status["seconds_until_auto_unload"] is None + + # --------------------------------------------------------------------------- # POST /dev/unload endpoint tests # --------------------------------------------------------------------------- diff --git a/docker/gpu/docker-compose.yml b/docker/gpu/docker-compose.yml index ae83b166..c4d37107 100644 --- a/docker/gpu/docker-compose.yml +++ b/docker/gpu/docker-compose.yml @@ -31,6 +31,7 @@ services: - ALLOW_DEV_UNLOAD=true - MODEL_AUTO_UNLOAD_ENABLED=true - MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS=300 + - MODEL_UNLOAD_STRATEGY=cpu_cache # - ENABLE_DEBUG_ENDPOINTS=true deploy: resources: diff --git a/docs/configuration.md b/docs/configuration.md index 80914f3a..b1e1675b 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -156,6 +156,7 @@ Names are the field names from `api/src/core/config.py`, uppercased. Unrecognize | `ALLOW_DEV_UNLOAD` | `false` | Expose `/dev/model`, `POST /dev/unload`, and `POST /dev/reload` | | `MODEL_AUTO_UNLOAD_ENABLED` | `false` | Unload the model after the idle timeout | | `MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS` | `300.0` | Idle seconds before auto-unload when enabled | +| `MODEL_UNLOAD_STRATEGY` | `destroy` | `destroy` releases model objects; `cpu_cache` keeps model weights in system RAM for faster reload | ## Logging diff --git a/notes.md b/notes.md new file mode 100644 index 00000000..dabb4e10 --- /dev/null +++ b/notes.md @@ -0,0 +1,85 @@ +# Kokoro GPU Unload Notes + +## Question + +Could reload after GPU unload be faster if the model stays resident in system RAM +and is moved back to CUDA on demand, instead of destroying and reconstructing the +backend from disk? + +## Benchmark Setup + +- Ran entirely inside the running container: `kokoro-tts-gpu-kokoro-tts-1` +- GPU: NVIDIA GeForce RTX 2070 SUPER +- PyTorch: `2.8.0+cu126` +- Model: `/app/api/src/models/v1_0/kokoro-v1_0.pth` +- Voice: `af_heart` +- Test text: short unload/reload sentence +- The API model was unloaded before benchmarking to free VRAM. + +No project files were changed for the benchmark. + +## Current API Behavior + +Measured through the live API using `POST /dev/unload` followed by +`POST /v1/audio/speech`. + +| Case | Time | +|---|---:| +| Cold request after unload, average | `1.835s` | +| Warm request, average | `0.135s` | +| Reload penalty, average | `1.700s` | + +Interpretation: with the current destroy/reload mechanism, the next short request +after unload pays about `+1.7s`. + +## Direct Model Load Benchmark + +Compared fresh model construction plus CUDA load against keeping the model object +in CPU RAM and moving it back to CUDA. + +| Case | Time | +|---|---:| +| Fresh model load, average | `1.692s` | +| CPU RAM to CUDA, average | `0.193s` | +| Estimated time saved | `1.499s` | + +Interpretation: most of the reload time is CPU-side reconstruction/loading, not +the CUDA transfer itself. + +## Backend-Level CPU Cache Benchmark + +This benchmark kept the backend/model/pipeline/voice state alive in CPU RAM, +then moved the model back to CUDA before generating. + +| Case | Time | +|---|---:| +| Fresh backend load plus cold generation, average | `2.247s` | +| CPU-cache reload plus generation, average | `0.334s` | +| Estimated time saved | `1.913s` | + +Interpretation: in an optimistic implementation, reload plus first generation was +roughly `6x` to `7x` faster than the current fresh backend path. + +## Takeaway + +CPU-RAM caching looks worth implementing. A likely design is an unload strategy +that moves the loaded model to CPU and clears CUDA memory, while preserving enough +backend state to avoid reconstructing the model from disk on the next request. + +Potential strategy setting: + +```env +MODEL_UNLOAD_STRATEGY=destroy +# or +MODEL_UNLOAD_STRATEGY=cpu_cache +``` + +Expected tradeoff: + +- `destroy`: maximum system RAM release, slower reload. +- `cpu_cache`: keeps more system RAM in use, much faster reload, still reclaims + most model VRAM. + +Important caveat: the backend-level benchmark is a best case because it preserved +pipeline and voice cache state. If an implementation preserves only model weights +but rebuilds pipeline or voice state, the speedup may be smaller. From 5cb79f52d1ae57a1a452a7e9441e2a39b9eb0c75 Mon Sep 17 00:00:00 2001 From: brycehenson Date: Mon, 17 Aug 2026 06:44:50 +0000 Subject: [PATCH 04/19] add timeout time in log message --- api/src/inference/model_manager.py | 6 +++++- api/tests/test_model_unload.py | 22 ++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/api/src/inference/model_manager.py b/api/src/inference/model_manager.py index f2b5317a..b06c2b95 100644 --- a/api/src/inference/model_manager.py +++ b/api/src/inference/model_manager.py @@ -159,6 +159,9 @@ async def load_model(self, path: str) -> None: def _auto_unload_timeout(self) -> float: return max(0.0, float(settings.model_auto_unload_timeout_seconds)) + def _format_seconds(self, seconds: float) -> str: + return f"{seconds:g}s" + def _auto_unload_enabled(self) -> bool: return settings.model_auto_unload_enabled and self._auto_unload_timeout() > 0 @@ -270,7 +273,8 @@ async def _idle_unload_after(self, timeout: float) -> None: if torch.cuda.is_available(): torch.cuda.empty_cache() logger.info( - f"Model auto-unload completed with strategy={self._unload_strategy()}" + "Model auto-unloaded after idle timeout of " + f"{self._format_seconds(self._auto_unload_timeout())}" ) except asyncio.CancelledError: logger.debug("Model auto-unload timer cancelled") diff --git a/api/tests/test_model_unload.py b/api/tests/test_model_unload.py index c31ae1e5..8480b740 100644 --- a/api/tests/test_model_unload.py +++ b/api/tests/test_model_unload.py @@ -256,6 +256,28 @@ async def fake_generate(*args, **kwargs): assert manager._backend is None +@pytest.mark.asyncio +async def test_idle_auto_unload_log_includes_configured_timeout(monkeypatch): + monkeypatch.setattr(settings, "model_auto_unload_enabled", True) + monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 30.0) + + manager = ModelManager() + mock_backend = MagicMock() + manager._backend = mock_backend + manager._last_used_at = 0.0 + + with ( + patch("api.src.inference.model_manager.time.monotonic", return_value=31.0), + patch("api.src.inference.model_manager.torch") as mock_torch, + patch("api.src.inference.model_manager.logger.info") as mock_log_info, + ): + mock_torch.cuda.is_available.return_value = False + await manager._idle_unload_after(0) + + mock_backend.unload.assert_called_once() + mock_log_info.assert_any_call("Model auto-unloaded after idle timeout of 30s") + + @pytest.mark.asyncio async def test_load_model_schedules_idle_unload_when_enabled(monkeypatch): """Startup-style model loads also schedule unload without waiting for traffic.""" From e8e4abcd56cfb76f6fea381972b88b978a8c64a6 Mon Sep 17 00:00:00 2001 From: brycehenson Date: Mon, 17 Aug 2026 04:25:34 +0000 Subject: [PATCH 05/19] Add model auto-unload controls --- README.md | 10 ++- api/src/core/config.py | 4 + api/src/inference/model_manager.py | 128 +++++++++++++++++++++++++++-- api/src/routers/development.py | 47 +++++++++++ api/tests/test_model_unload.py | 110 ++++++++++++++++++++++++- docker/gpu/docker-compose.yml | 5 +- docs/configuration.md | 4 +- 7 files changed, 295 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 939092ed..543436ca 100644 --- a/README.md +++ b/README.md @@ -698,7 +698,11 @@ Key Performance Metrics: ### Model Unload / VRAM Reclaim -`POST /dev/unload` frees the model from VRAM and reloads lazily on the next request. Reclaim scales with load (the activation pool, not just weights) but plateaus: chunks cap at 450 tokens. Long-form = ~30 paragraphs. Same setup as above. +`POST /dev/unload` frees the model from VRAM and reloads lazily on the next request. Set `ALLOW_DEV_UNLOAD=true` to expose the lifecycle controls: `GET /dev/model`, `POST /dev/unload`, and `POST /dev/reload`. + +For shared-GPU hosts, set `MODEL_AUTO_UNLOAD_ENABLED=true` and tune `MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS` to unload automatically after the model has been idle. In-flight generation keeps the model loaded; the next request reloads it automatically, or `POST /dev/reload` can pre-warm it. + +Reclaim scales with load (the activation pool, not just weights) but plateaus: chunks cap at 450 tokens. Long-form = ~30 paragraphs. Same setup as above.

Short workload @@ -759,7 +763,9 @@ System state and resource usage, for debugging exhaustion or performance issues. - `/debug/threads` - Get thread information and stack traces - `/debug/storage` - Disk usage per mounted partition - `/debug/system` - Get system information (CPU, memory, GPU) -- `POST /dev/unload` - Release model from VRAM; reloads lazily on next request. Off by default; set `ALLOW_DEV_UNLOAD=true` to enable +- `/dev/model` - Get model load state and auto-unload timing +- `POST /dev/unload` - Release model from VRAM; reloads lazily on next request +- `POST /dev/reload` - Load the model immediately after an unload or before traffic arrives Stability: the `/v1/*` OpenAI-compatible routes are the stable API. `/dev/*` and `/debug/*` are operational helpers, and may change or move behind flags between minor releases. diff --git a/api/src/core/config.py b/api/src/core/config.py index 9d7f77ba..8c5ea36b 100644 --- a/api/src/core/config.py +++ b/api/src/core/config.py @@ -40,6 +40,10 @@ class Settings(BaseSettings): False # Whether to allow saving combined voices locally ) allow_dev_unload: bool = False # Whether to expose the POST /dev/unload endpoint + model_auto_unload_enabled: bool = False # Whether to unload the model after an idle timeout + model_auto_unload_timeout_seconds: float = ( + 300.0 # Idle seconds before unloading when MODEL_AUTO_UNLOAD_ENABLED=true + ) enable_debug_endpoints: bool = ( False # Whether to expose /debug/* host and process introspection routes ) diff --git a/api/src/inference/model_manager.py b/api/src/inference/model_manager.py index b1468717..8b025909 100644 --- a/api/src/inference/model_manager.py +++ b/api/src/inference/model_manager.py @@ -1,6 +1,7 @@ """Kokoro V1 model management.""" import asyncio +import time from typing import Optional import torch @@ -29,6 +30,9 @@ def __init__(self, config: Optional[ModelConfig] = None): self._backend: Optional[KokoroV1] = None # Explicitly type as KokoroV1 self._device: Optional[str] = None self._lock = asyncio.Lock() + self._active_requests = 0 + self._last_used_at: Optional[float] = None + self._idle_unload_task: Optional[asyncio.Task] = None def _determine_device(self) -> str: """Determine device based on settings.""" @@ -102,11 +106,12 @@ async def initialize_with_warmup(self, voice_manager) -> tuple[str, str, int]: raise RuntimeError(f"Warmup failed: {e}") async def ensure_backend(self) -> None: - """Reload the backend if it was unloaded via /dev/unload.""" + """Reload the backend if it was unloaded.""" if self._backend: return async with self._lock: if not self._backend: + self._cancel_idle_unload_timer() await self.initialize() await self.load_model(self._config.pytorch_kokoro_v1_file) @@ -137,30 +142,110 @@ async def load_model(self, path: str) -> None: try: await self._backend.load_model(path) + self._last_used_at = time.monotonic() except FileNotFoundError as e: raise e except Exception as e: raise RuntimeError(f"Failed to load model: {e}") + def _auto_unload_timeout(self) -> float: + return max(0.0, float(settings.model_auto_unload_timeout_seconds)) + + def _auto_unload_enabled(self) -> bool: + return settings.model_auto_unload_enabled and self._auto_unload_timeout() > 0 + + def _cancel_idle_unload_timer(self) -> None: + try: + current_task = asyncio.current_task() + except RuntimeError: + current_task = None + if ( + self._idle_unload_task + and not self._idle_unload_task.done() + and self._idle_unload_task is not current_task + ): + self._idle_unload_task.cancel() + self._idle_unload_task = None + + def _unload_backend_locked(self) -> bool: + if self._backend is None: + return False + self._backend.unload() + self._backend = None + self._last_used_at = time.monotonic() + return True + + def _schedule_idle_unload_timer_locked(self) -> None: + self._cancel_idle_unload_timer() + if ( + not self._auto_unload_enabled() + or self._backend is None + or self._active_requests > 0 + ): + return + + timeout = self._auto_unload_timeout() + self._idle_unload_task = asyncio.create_task(self._idle_unload_after(timeout)) + + async def _idle_unload_after(self, timeout: float) -> None: + try: + await asyncio.sleep(timeout) + async with self._lock: + if ( + not self._auto_unload_enabled() + or self._backend is None + or self._active_requests > 0 + or self._last_used_at is None + ): + return + + idle_for = time.monotonic() - self._last_used_at + if idle_for < self._auto_unload_timeout(): + self._schedule_idle_unload_timer_locked() + return + + unloaded = self._unload_backend_locked() + + if unloaded: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + logger.info("Model auto-unloaded after idle timeout") + except asyncio.CancelledError: + pass + + async def _begin_request(self) -> None: + async with self._lock: + self._active_requests += 1 + self._cancel_idle_unload_timer() + + async def _end_request(self) -> None: + async with self._lock: + self._active_requests = max(0, self._active_requests - 1) + self._last_used_at = time.monotonic() + self._schedule_idle_unload_timer_locked() + async def generate(self, *args, **kwargs): """Generate audio using initialized backend. Raises: RuntimeError: If generation fails """ - await self.ensure_backend() - assert self._backend is not None, "ensure_backend left no backend" - + await self._begin_request() try: + await self.ensure_backend() + assert self._backend is not None, "ensure_backend left no backend" async for chunk in self._backend.generate(*args, **kwargs): if settings.default_volume_multiplier != 1.0: chunk.audio *= settings.default_volume_multiplier yield chunk except Exception as e: raise RuntimeError(f"Generation failed: {e}") + finally: + await self._end_request() def unload_all(self) -> None: """Unload model and free resources.""" + self._cancel_idle_unload_timer() if self._backend: self._backend.unload() self._backend = None @@ -168,13 +253,42 @@ def unload_all(self) -> None: async def unload(self) -> None: """Release model from GPU memory. Reloads automatically on next request.""" async with self._lock: - if self._backend is not None: - self._backend.unload() - self._backend = None + self._cancel_idle_unload_timer() + self._unload_backend_locked() if torch.cuda.is_available(): torch.cuda.empty_cache() logger.info("Model unloaded from GPU memory") + async def reload(self) -> None: + """Reload the model immediately.""" + async with self._lock: + self._cancel_idle_unload_timer() + self._unload_backend_locked() + await self.initialize() + await self.load_model(self._config.pytorch_kokoro_v1_file) + self._schedule_idle_unload_timer_locked() + logger.info("Model reloaded") + + def status(self) -> dict: + """Return model lifecycle state for API responses.""" + timeout = self._auto_unload_timeout() + idle_for = None + unload_in = None + if self._last_used_at is not None: + idle_for = max(0.0, time.monotonic() - self._last_used_at) + if self._auto_unload_enabled() and self._backend is not None: + unload_in = max(0.0, timeout - idle_for) + return { + "backend": self.current_backend, + "device": self._device, + "loaded": self._backend is not None, + "active_requests": self._active_requests, + "auto_unload_enabled": settings.model_auto_unload_enabled, + "auto_unload_timeout_seconds": timeout, + "idle_seconds": idle_for, + "seconds_until_auto_unload": unload_in, + } + @property def current_backend(self) -> str: """Get current backend type.""" diff --git a/api/src/routers/development.py b/api/src/routers/development.py index bc2e1eab..7a23ad59 100644 --- a/api/src/routers/development.py +++ b/api/src/routers/development.py @@ -505,3 +505,50 @@ async def unload_model( except Exception as e: logger.error(f"Error unloading model: {e}") raise HTTPException(status_code=500, detail={"error": str(e)}) + + +@router.post("/dev/reload") +async def reload_model( + tts_service: TTSService = Depends(get_tts_service), +): + """Reload the model immediately. + + Normal inference also reloads lazily after an unload; this endpoint is useful + when you want to pre-warm the model before the next request. + """ + if not settings.allow_dev_unload: + raise HTTPException( + status_code=403, + detail={"error": "The /dev/reload endpoint is disabled"}, + ) + try: + if tts_service.model_manager is None: + raise HTTPException( + status_code=503, detail={"error": "Model manager not initialized"} + ) + await tts_service.model_manager.reload() + return JSONResponse( + {"status": "loaded", "model": tts_service.model_manager.status()} + ) + except HTTPException: + raise + except Exception as e: + logger.error(f"Error reloading model: {e}") + raise HTTPException(status_code=500, detail={"error": str(e)}) + + +@router.get("/dev/model") +async def model_status( + tts_service: TTSService = Depends(get_tts_service), +): + """Return model load state and auto-unload settings.""" + if not settings.allow_dev_unload: + raise HTTPException( + status_code=403, + detail={"error": "The /dev/model endpoint is disabled"}, + ) + if tts_service.model_manager is None: + raise HTTPException( + status_code=503, detail={"error": "Model manager not initialized"} + ) + return JSONResponse(tts_service.model_manager.status()) diff --git a/api/tests/test_model_unload.py b/api/tests/test_model_unload.py index 05da4353..a7faee37 100644 --- a/api/tests/test_model_unload.py +++ b/api/tests/test_model_unload.py @@ -1,4 +1,4 @@ -"""Tests for ModelManager.unload(), lazy reinit in generate(), and POST /dev/unload.""" +"""Tests for model unload, auto-unload, lazy reload, and dev lifecycle endpoints.""" import asyncio from contextlib import contextmanager @@ -36,6 +36,8 @@ async def _override(): def _enable_dev_unload(monkeypatch): """Enable the /dev/unload gate for the endpoint tests in this module.""" monkeypatch.setattr(settings, "allow_dev_unload", True) + monkeypatch.setattr(settings, "model_auto_unload_enabled", False) + monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 300.0) # --------------------------------------------------------------------------- @@ -46,6 +48,7 @@ def _enable_dev_unload(monkeypatch): def test_manager_init_creates_lock(): manager = ModelManager() assert isinstance(manager._lock, asyncio.Lock) + assert manager._active_requests == 0 @pytest.mark.asyncio @@ -185,6 +188,85 @@ async def fake_generate(*args, **kwargs): assert len(chunks) == 1 +@pytest.mark.asyncio +async def test_generate_schedules_idle_unload_when_enabled(monkeypatch): + """Finished generation schedules model unload after the configured idle period.""" + monkeypatch.setattr(settings, "model_auto_unload_enabled", True) + monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 0.01) + + manager = ModelManager() + mock_backend = MagicMock() + audio_chunk = AudioChunk(np.zeros(10, dtype=np.float32)) + + async def fake_generate(*args, **kwargs): + yield audio_chunk + + mock_backend.generate = fake_generate + manager._backend = mock_backend + + with patch("api.src.inference.model_manager.torch") as mock_torch: + mock_torch.cuda.is_available.return_value = False + chunks = [] + async for chunk in manager.generate("hello", ("voice", "/path/voice.pt")): + chunks.append(chunk) + await asyncio.sleep(0.03) + + assert len(chunks) == 1 + mock_backend.unload.assert_called_once() + assert manager._backend is None + + +@pytest.mark.asyncio +async def test_active_request_blocks_idle_unload(monkeypatch): + """The idle timer does not unload while generation is still active.""" + monkeypatch.setattr(settings, "model_auto_unload_enabled", True) + monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 0.01) + + manager = ModelManager() + mock_backend = MagicMock() + + async def slow_generate(*args, **kwargs): + await asyncio.sleep(0.03) + yield AudioChunk(np.zeros(10, dtype=np.float32)) + + mock_backend.generate = slow_generate + manager._backend = mock_backend + + with patch("api.src.inference.model_manager.torch") as mock_torch: + mock_torch.cuda.is_available.return_value = False + chunks = [] + async for chunk in manager.generate("hello", ("voice", "/path/voice.pt")): + assert manager._active_requests == 1 + mock_backend.unload.assert_not_called() + chunks.append(chunk) + + assert len(chunks) == 1 + assert manager._active_requests == 0 + assert manager._idle_unload_task is not None + manager._cancel_idle_unload_timer() + + +def test_status_reports_model_lifecycle_state(monkeypatch): + monkeypatch.setattr(settings, "model_auto_unload_enabled", True) + monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 30.0) + manager = ModelManager() + manager._backend = MagicMock() + manager._device = "cuda" + manager._last_used_at = 10.0 + + with patch("api.src.inference.model_manager.time.monotonic", return_value=15.0): + status = manager.status() + + assert status["backend"] == "kokoro_v1" + assert status["device"] == "cuda" + assert status["loaded"] is True + assert status["active_requests"] == 0 + assert status["auto_unload_enabled"] is True + assert status["auto_unload_timeout_seconds"] == 30.0 + assert status["idle_seconds"] == 5.0 + assert status["seconds_until_auto_unload"] == 25.0 + + # --------------------------------------------------------------------------- # POST /dev/unload endpoint tests # --------------------------------------------------------------------------- @@ -261,3 +343,29 @@ def test_unload_endpoint_500_on_exception(): assert response.status_code == 500 assert "GPU exploded" in response.json()["detail"]["error"] + + +def test_reload_endpoint_returns_status(): + mock_manager = MagicMock() + mock_manager.reload = AsyncMock() + mock_manager.status.return_value = {"loaded": True} + service = _mock_service(manager=mock_manager) + + with override_tts_service(service): + response = client.post("/dev/reload") + + assert response.status_code == 200 + assert response.json() == {"status": "loaded", "model": {"loaded": True}} + mock_manager.reload.assert_called_once() + + +def test_model_status_endpoint_returns_status(): + mock_manager = MagicMock() + mock_manager.status.return_value = {"loaded": False} + service = _mock_service(manager=mock_manager) + + with override_tts_service(service): + response = client.get("/dev/model") + + assert response.status_code == 200 + assert response.json() == {"loaded": False} diff --git a/docker/gpu/docker-compose.yml b/docker/gpu/docker-compose.yml index 1a917138..ae83b166 100644 --- a/docker/gpu/docker-compose.yml +++ b/docker/gpu/docker-compose.yml @@ -28,7 +28,9 @@ services: - PYTHONUNBUFFERED=1 - API_LOG_LEVEL=DEBUG - DOWNLOAD_MODEL=true - # - ALLOW_DEV_UNLOAD=true + - ALLOW_DEV_UNLOAD=true + - MODEL_AUTO_UNLOAD_ENABLED=true + - MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS=300 # - ENABLE_DEBUG_ENDPOINTS=true deploy: resources: @@ -37,4 +39,3 @@ services: - driver: nvidia count: all capabilities: [gpu] - diff --git a/docs/configuration.md b/docs/configuration.md index a3d637c2..80914f3a 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -153,7 +153,9 @@ Names are the field names from `api/src/core/config.py`, uppercased. Unrecognize | Variable | Default | | |---|---|---| | `ENABLE_DEBUG_ENDPOINTS` | `false` | Expose `/debug/*` host and process introspection | -| `ALLOW_DEV_UNLOAD` | `false` | Expose `POST /dev/unload` | +| `ALLOW_DEV_UNLOAD` | `false` | Expose `/dev/model`, `POST /dev/unload`, and `POST /dev/reload` | +| `MODEL_AUTO_UNLOAD_ENABLED` | `false` | Unload the model after the idle timeout | +| `MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS` | `300.0` | Idle seconds before auto-unload when enabled | ## Logging From e222136a72152d387c19fd0e14a61d786e2fc29a Mon Sep 17 00:00:00 2001 From: brycehenson Date: Mon, 17 Aug 2026 04:36:48 +0000 Subject: [PATCH 06/19] Schedule auto-unload after model load --- api/src/inference/model_manager.py | 1 + api/tests/test_model_unload.py | 44 ++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/api/src/inference/model_manager.py b/api/src/inference/model_manager.py index 8b025909..2a0adcf1 100644 --- a/api/src/inference/model_manager.py +++ b/api/src/inference/model_manager.py @@ -143,6 +143,7 @@ async def load_model(self, path: str) -> None: try: await self._backend.load_model(path) self._last_used_at = time.monotonic() + self._schedule_idle_unload_timer_locked() except FileNotFoundError as e: raise e except Exception as e: diff --git a/api/tests/test_model_unload.py b/api/tests/test_model_unload.py index a7faee37..04abb485 100644 --- a/api/tests/test_model_unload.py +++ b/api/tests/test_model_unload.py @@ -216,6 +216,50 @@ async def fake_generate(*args, **kwargs): assert manager._backend is None +@pytest.mark.asyncio +async def test_load_model_schedules_idle_unload_when_enabled(monkeypatch): + """Startup-style model loads also schedule unload without waiting for traffic.""" + monkeypatch.setattr(settings, "model_auto_unload_enabled", True) + monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 0.01) + + manager = ModelManager() + mock_backend = MagicMock() + mock_backend.load_model = AsyncMock() + manager._backend = mock_backend + + with patch("api.src.inference.model_manager.torch") as mock_torch: + mock_torch.cuda.is_available.return_value = False + await manager.load_model("/path/model.pt") + await asyncio.sleep(0.03) + + mock_backend.load_model.assert_called_once_with("/path/model.pt") + mock_backend.unload.assert_called_once() + assert manager._backend is None + + +@pytest.mark.asyncio +async def test_load_model_does_not_schedule_idle_unload_during_active_request( + monkeypatch, +): + """Lazy loads during generation wait for request completion before scheduling.""" + monkeypatch.setattr(settings, "model_auto_unload_enabled", True) + monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 0.01) + + manager = ModelManager() + mock_backend = MagicMock() + mock_backend.load_model = AsyncMock() + manager._backend = mock_backend + manager._active_requests = 1 + + await manager.load_model("/path/model.pt") + await asyncio.sleep(0.03) + + mock_backend.load_model.assert_called_once_with("/path/model.pt") + mock_backend.unload.assert_not_called() + assert manager._backend is mock_backend + assert manager._idle_unload_task is None + + @pytest.mark.asyncio async def test_active_request_blocks_idle_unload(monkeypatch): """The idle timer does not unload while generation is still active.""" From 26edc3bfbeefa9fb376269ae206b73afa9a2bc5b Mon Sep 17 00:00:00 2001 From: brycehenson Date: Mon, 17 Aug 2026 06:46:07 +0000 Subject: [PATCH 07/19] Log model auto-unload timeout --- api/src/inference/model_manager.py | 8 +++++++- api/tests/test_model_unload.py | 22 ++++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/api/src/inference/model_manager.py b/api/src/inference/model_manager.py index 2a0adcf1..ca0695e1 100644 --- a/api/src/inference/model_manager.py +++ b/api/src/inference/model_manager.py @@ -152,6 +152,9 @@ async def load_model(self, path: str) -> None: def _auto_unload_timeout(self) -> float: return max(0.0, float(settings.model_auto_unload_timeout_seconds)) + def _format_seconds(self, seconds: float) -> str: + return f"{seconds:g}s" + def _auto_unload_enabled(self) -> bool: return settings.model_auto_unload_enabled and self._auto_unload_timeout() > 0 @@ -210,7 +213,10 @@ async def _idle_unload_after(self, timeout: float) -> None: if unloaded: if torch.cuda.is_available(): torch.cuda.empty_cache() - logger.info("Model auto-unloaded after idle timeout") + logger.info( + "Model auto-unloaded after idle timeout of " + f"{self._format_seconds(self._auto_unload_timeout())}" + ) except asyncio.CancelledError: pass diff --git a/api/tests/test_model_unload.py b/api/tests/test_model_unload.py index 04abb485..ee52f8dc 100644 --- a/api/tests/test_model_unload.py +++ b/api/tests/test_model_unload.py @@ -216,6 +216,28 @@ async def fake_generate(*args, **kwargs): assert manager._backend is None +@pytest.mark.asyncio +async def test_idle_auto_unload_log_includes_configured_timeout(monkeypatch): + monkeypatch.setattr(settings, "model_auto_unload_enabled", True) + monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 30.0) + + manager = ModelManager() + mock_backend = MagicMock() + manager._backend = mock_backend + manager._last_used_at = 0.0 + + with ( + patch("api.src.inference.model_manager.time.monotonic", return_value=31.0), + patch("api.src.inference.model_manager.torch") as mock_torch, + patch("api.src.inference.model_manager.logger.info") as mock_log_info, + ): + mock_torch.cuda.is_available.return_value = False + await manager._idle_unload_after(0) + + mock_backend.unload.assert_called_once() + mock_log_info.assert_any_call("Model auto-unloaded after idle timeout of 30s") + + @pytest.mark.asyncio async def test_load_model_schedules_idle_unload_when_enabled(monkeypatch): """Startup-style model loads also schedule unload without waiting for traffic.""" From 330de4377bea75f792cba809b38074b2d15e6823 Mon Sep 17 00:00:00 2001 From: brycehenson Date: Mon, 17 Aug 2026 06:54:21 +0000 Subject: [PATCH 08/19] simplify readme --- README.md | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/README.md b/README.md index 543436ca..6bd85402 100644 --- a/README.md +++ b/README.md @@ -698,10 +698,7 @@ Key Performance Metrics: ### Model Unload / VRAM Reclaim -`POST /dev/unload` frees the model from VRAM and reloads lazily on the next request. Set `ALLOW_DEV_UNLOAD=true` to expose the lifecycle controls: `GET /dev/model`, `POST /dev/unload`, and `POST /dev/reload`. - -For shared-GPU hosts, set `MODEL_AUTO_UNLOAD_ENABLED=true` and tune `MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS` to unload automatically after the model has been idle. In-flight generation keeps the model loaded; the next request reloads it automatically, or `POST /dev/reload` can pre-warm it. - +`POST /dev/unload` frees the model from VRAM and reloads lazily on the next request. To automatically unload the model after an idle timeout, set MODEL_AUTO_UNLOAD_ENABLED=true and adjust MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS. Reclaim scales with load (the activation pool, not just weights) but plateaus: chunks cap at 450 tokens. Long-form = ~30 paragraphs. Same setup as above.

From eb6d91169de32e109314d0db5552336ea0f220a4 Mon Sep 17 00:00:00 2001 From: brycehenson Date: Mon, 17 Aug 2026 07:14:31 +0000 Subject: [PATCH 09/19] reduce readme diff --- README.md | 8 +++++--- docker/gpu/docker-compose.yml | 5 ++--- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 6bd85402..7ee012a5 100644 --- a/README.md +++ b/README.md @@ -698,8 +698,8 @@ Key Performance Metrics: ### Model Unload / VRAM Reclaim -`POST /dev/unload` frees the model from VRAM and reloads lazily on the next request. To automatically unload the model after an idle timeout, set MODEL_AUTO_UNLOAD_ENABLED=true and adjust MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS. -Reclaim scales with load (the activation pool, not just weights) but plateaus: chunks cap at 450 tokens. Long-form = ~30 paragraphs. Same setup as above. +`POST /dev/unload` frees the model from VRAM and reloads lazily on the next request. Reclaim scales with load (the activation pool, not just weights) but plateaus: chunks cap at 450 tokens. Long-form = ~30 paragraphs. Same setup as above. +

Short workload @@ -713,6 +713,8 @@ Reclaim scales with load (the activation pool, not just weights) but plateaus: c Floor is host + CUDA context. Reproduce with `uv run --extra benchmarks assorted_checks/benchmarks/benchmark_model_unload.py` from `examples/`. +To automatically unload the model after an idle timeout, set MODEL_AUTO_UNLOAD_ENABLED=true and adjust MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS. + ### Transcription roundtrip (WER/CER) End-to-end roundtrip: synthesize with Kokoro, transcribe the result back with [`faster-whisper`](https://github.com/SYSTRAN/faster-whisper), compare to the source text. Scripts and data live under `examples/assorted_checks/test_transcription/`. @@ -761,7 +763,7 @@ System state and resource usage, for debugging exhaustion or performance issues. - `/debug/storage` - Disk usage per mounted partition - `/debug/system` - Get system information (CPU, memory, GPU) - `/dev/model` - Get model load state and auto-unload timing -- `POST /dev/unload` - Release model from VRAM; reloads lazily on next request +- `POST /dev/unload` - Release model from VRAM; reloads lazily on next request. Off by default; set `ALLOW_DEV_UNLOAD=true` to enable - `POST /dev/reload` - Load the model immediately after an unload or before traffic arrives Stability: the `/v1/*` OpenAI-compatible routes are the stable API. `/dev/*` and `/debug/*` are operational helpers, and may change or move behind flags between minor releases. diff --git a/docker/gpu/docker-compose.yml b/docker/gpu/docker-compose.yml index ae83b166..1a917138 100644 --- a/docker/gpu/docker-compose.yml +++ b/docker/gpu/docker-compose.yml @@ -28,9 +28,7 @@ services: - PYTHONUNBUFFERED=1 - API_LOG_LEVEL=DEBUG - DOWNLOAD_MODEL=true - - ALLOW_DEV_UNLOAD=true - - MODEL_AUTO_UNLOAD_ENABLED=true - - MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS=300 + # - ALLOW_DEV_UNLOAD=true # - ENABLE_DEBUG_ENDPOINTS=true deploy: resources: @@ -39,3 +37,4 @@ services: - driver: nvidia count: all capabilities: [gpu] + From e910a7612d1fcd0307592d282aa0b866437146c1 Mon Sep 17 00:00:00 2001 From: brycehenson Date: Mon, 17 Aug 2026 07:15:24 +0000 Subject: [PATCH 10/19] remove added newline --- README.md | 1 - 1 file changed, 1 deletion(-) diff --git a/README.md b/README.md index 7ee012a5..32f92c81 100644 --- a/README.md +++ b/README.md @@ -700,7 +700,6 @@ Key Performance Metrics: `POST /dev/unload` frees the model from VRAM and reloads lazily on the next request. Reclaim scales with load (the activation pool, not just weights) but plateaus: chunks cap at 450 tokens. Long-form = ~30 paragraphs. Same setup as above. -

Short workload Long-form workload From d621ff09abc3ed2282cacba2be5d09bbef2e932a Mon Sep 17 00:00:00 2001 From: brycehenson Date: Mon, 17 Aug 2026 08:13:39 +0000 Subject: [PATCH 11/19] handle USE_GPU=true MODEL_UNLOAD_STRATEGY=cpu_cache --- api/src/inference/model_manager.py | 5 +++++ api/tests/test_model_unload.py | 33 ++++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/api/src/inference/model_manager.py b/api/src/inference/model_manager.py index b06c2b95..eefc0339 100644 --- a/api/src/inference/model_manager.py +++ b/api/src/inference/model_manager.py @@ -182,6 +182,11 @@ def _unload_strategy(self) -> str: f"Unknown MODEL_UNLOAD_STRATEGY={settings.model_unload_strategy!r}; using destroy" ) return "destroy" + if strategy == "cpu_cache" and not settings.use_gpu: + logger.warning( + "MODEL_UNLOAD_STRATEGY=cpu_cache requires USE_GPU=true; using destroy" + ) + return "destroy" return strategy def _cancel_idle_unload_timer(self) -> None: diff --git a/api/tests/test_model_unload.py b/api/tests/test_model_unload.py index 8480b740..a83bc57c 100644 --- a/api/tests/test_model_unload.py +++ b/api/tests/test_model_unload.py @@ -84,6 +84,39 @@ async def test_unload_cpu_cache_keeps_backend(monkeypatch): assert manager._backend is mock_backend +def test_cpu_cache_strategy_warns_and_uses_destroy_without_gpu(monkeypatch): + monkeypatch.setattr(settings, "use_gpu", False) + monkeypatch.setattr(settings, "model_unload_strategy", "cpu_cache") + + manager = ModelManager() + + with patch("api.src.inference.model_manager.logger.warning") as mock_warning: + assert manager._unload_strategy() == "destroy" + + mock_warning.assert_called_once_with( + "MODEL_UNLOAD_STRATEGY=cpu_cache requires USE_GPU=true; using destroy" + ) + + +@pytest.mark.asyncio +async def test_unload_cpu_cache_destroys_backend_without_gpu(monkeypatch): + monkeypatch.setattr(settings, "use_gpu", False) + monkeypatch.setattr(settings, "model_unload_strategy", "cpu_cache") + + manager = ModelManager() + mock_backend = MagicMock() + mock_backend.is_loaded = True + mock_backend.is_cpu_cached = True + manager._backend = mock_backend + + with patch("api.src.inference.model_manager.torch") as mock_torch: + mock_torch.cuda.is_available.return_value = False + await manager.unload() + + mock_backend.unload.assert_called_once_with(strategy="destroy") + assert manager._backend is None + + @pytest.mark.asyncio async def test_reload_restores_cpu_cached_backend(monkeypatch): monkeypatch.setattr(settings, "model_unload_strategy", "cpu_cache") From 241f98d2ebc1b1647eab8f2954803ea2ee47895c Mon Sep 17 00:00:00 2001 From: brycehenson Date: Thu, 20 Aug 2026 05:54:00 +0000 Subject: [PATCH 12/19] Refine model auto-unload handling Centralize request timeout accounting into ModelManager.hold(). Use it for generate_from_phonemes(). Remove MODEL_AUTO_UNLOAD_ENABLED. Replace it with MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS=0 is the default/off state. --- README.md | 2 +- api/src/core/config.py | 3 +- api/src/inference/model_manager.py | 29 +++++++----- api/src/services/tts_service.py | 71 +++++++++++++++--------------- api/tests/test_model_unload.py | 9 +--- docs/configuration.md | 3 +- 6 files changed, 58 insertions(+), 59 deletions(-) diff --git a/README.md b/README.md index 32f92c81..079b9020 100644 --- a/README.md +++ b/README.md @@ -712,7 +712,7 @@ Key Performance Metrics: Floor is host + CUDA context. Reproduce with `uv run --extra benchmarks assorted_checks/benchmarks/benchmark_model_unload.py` from `examples/`. -To automatically unload the model after an idle timeout, set MODEL_AUTO_UNLOAD_ENABLED=true and adjust MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS. +To automatically unload the model after an idle timeout, set `MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS` to a positive number of seconds. The default `0` disables auto-unload. ### Transcription roundtrip (WER/CER) diff --git a/api/src/core/config.py b/api/src/core/config.py index 8c5ea36b..8bc40c02 100644 --- a/api/src/core/config.py +++ b/api/src/core/config.py @@ -40,9 +40,8 @@ class Settings(BaseSettings): False # Whether to allow saving combined voices locally ) allow_dev_unload: bool = False # Whether to expose the POST /dev/unload endpoint - model_auto_unload_enabled: bool = False # Whether to unload the model after an idle timeout model_auto_unload_timeout_seconds: float = ( - 300.0 # Idle seconds before unloading when MODEL_AUTO_UNLOAD_ENABLED=true + 0.0 # Idle seconds before unloading; 0 disables auto-unload ) enable_debug_endpoints: bool = ( False # Whether to expose /debug/* host and process introspection routes diff --git a/api/src/inference/model_manager.py b/api/src/inference/model_manager.py index ca0695e1..f81f975d 100644 --- a/api/src/inference/model_manager.py +++ b/api/src/inference/model_manager.py @@ -2,6 +2,7 @@ import asyncio import time +from contextlib import asynccontextmanager from typing import Optional import torch @@ -156,7 +157,7 @@ def _format_seconds(self, seconds: float) -> str: return f"{seconds:g}s" def _auto_unload_enabled(self) -> bool: - return settings.model_auto_unload_enabled and self._auto_unload_timeout() > 0 + return self._auto_unload_timeout() > 0 def _cancel_idle_unload_timer(self) -> None: try: @@ -231,24 +232,30 @@ async def _end_request(self) -> None: self._last_used_at = time.monotonic() self._schedule_idle_unload_timer_locked() + @asynccontextmanager + async def hold(self): + await self._begin_request() + try: + yield + finally: + await self._end_request() + async def generate(self, *args, **kwargs): """Generate audio using initialized backend. Raises: RuntimeError: If generation fails """ - await self._begin_request() try: - await self.ensure_backend() - assert self._backend is not None, "ensure_backend left no backend" - async for chunk in self._backend.generate(*args, **kwargs): - if settings.default_volume_multiplier != 1.0: - chunk.audio *= settings.default_volume_multiplier - yield chunk + async with self.hold(): + await self.ensure_backend() + assert self._backend is not None, "ensure_backend left no backend" + async for chunk in self._backend.generate(*args, **kwargs): + if settings.default_volume_multiplier != 1.0: + chunk.audio *= settings.default_volume_multiplier + yield chunk except Exception as e: raise RuntimeError(f"Generation failed: {e}") - finally: - await self._end_request() def unload_all(self) -> None: """Unload model and free resources.""" @@ -290,7 +297,7 @@ def status(self) -> dict: "device": self._device, "loaded": self._backend is not None, "active_requests": self._active_requests, - "auto_unload_enabled": settings.model_auto_unload_enabled, + "auto_unload_enabled": self._auto_unload_enabled(), "auto_unload_timeout_seconds": timeout, "idle_seconds": idle_for, "seconds_until_auto_unload": unload_in, diff --git a/api/src/services/tts_service.py b/api/src/services/tts_service.py index 465dbe39..bd3adf3e 100644 --- a/api/src/services/tts_service.py +++ b/api/src/services/tts_service.py @@ -563,44 +563,45 @@ async def generate_from_phonemes( """ start_time = time.time() try: - await self.model_manager.ensure_backend() - backend = self.model_manager.get_backend() - voice_name, voice_path = await self._get_voices_path(voice) - - if isinstance(backend, KokoroV1): - # For Kokoro V1, use generate_from_tokens with raw phonemes - result = None - # Use provided lang_code or determine from voice name - pipeline_lang_code = lang_code if lang_code else voice[:1].lower() - logger.info( - f"Using lang_code '{pipeline_lang_code}' for voice '{voice_name}' in phoneme pipeline" - ) + async with self.model_manager.hold(): + await self.model_manager.ensure_backend() + backend = self.model_manager.get_backend() + voice_name, voice_path = await self._get_voices_path(voice) - try: - # Use backend's pipeline management - for r in backend._get_pipeline( - pipeline_lang_code - ).generate_from_tokens( - tokens=phonemes, # Pass raw phonemes string - voice=voice_path, - speed=speed, - ): - if r.audio is not None: - result = r - break - except Exception as e: - logger.error(f"Failed to generate from phonemes: {e}") - raise RuntimeError(f"Phoneme generation failed: {e}") + if isinstance(backend, KokoroV1): + # For Kokoro V1, use generate_from_tokens with raw phonemes + result = None + # Use provided lang_code or determine from voice name + pipeline_lang_code = lang_code if lang_code else voice[:1].lower() + logger.info( + f"Using lang_code '{pipeline_lang_code}' for voice '{voice_name}' in phoneme pipeline" + ) - if result is None or result.audio is None: - raise ValueError("No audio generated") + try: + # Use backend's pipeline management + for r in backend._get_pipeline( + pipeline_lang_code + ).generate_from_tokens( + tokens=phonemes, # Pass raw phonemes string + voice=voice_path, + speed=speed, + ): + if r.audio is not None: + result = r + break + except Exception as e: + logger.error(f"Failed to generate from phonemes: {e}") + raise RuntimeError(f"Phoneme generation failed: {e}") - processing_time = time.time() - start_time - return result.audio.numpy(), processing_time - else: - raise ValueError( - "Phoneme generation only supported with Kokoro V1 backend" - ) + if result is None or result.audio is None: + raise ValueError("No audio generated") + + processing_time = time.time() - start_time + return result.audio.numpy(), processing_time + else: + raise ValueError( + "Phoneme generation only supported with Kokoro V1 backend" + ) except Exception as e: logger.error(f"Error in phoneme audio generation: {str(e)}") diff --git a/api/tests/test_model_unload.py b/api/tests/test_model_unload.py index ee52f8dc..6767a306 100644 --- a/api/tests/test_model_unload.py +++ b/api/tests/test_model_unload.py @@ -36,8 +36,7 @@ async def _override(): def _enable_dev_unload(monkeypatch): """Enable the /dev/unload gate for the endpoint tests in this module.""" monkeypatch.setattr(settings, "allow_dev_unload", True) - monkeypatch.setattr(settings, "model_auto_unload_enabled", False) - monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 300.0) + monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 0.0) # --------------------------------------------------------------------------- @@ -191,7 +190,6 @@ async def fake_generate(*args, **kwargs): @pytest.mark.asyncio async def test_generate_schedules_idle_unload_when_enabled(monkeypatch): """Finished generation schedules model unload after the configured idle period.""" - monkeypatch.setattr(settings, "model_auto_unload_enabled", True) monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 0.01) manager = ModelManager() @@ -218,7 +216,6 @@ async def fake_generate(*args, **kwargs): @pytest.mark.asyncio async def test_idle_auto_unload_log_includes_configured_timeout(monkeypatch): - monkeypatch.setattr(settings, "model_auto_unload_enabled", True) monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 30.0) manager = ModelManager() @@ -241,7 +238,6 @@ async def test_idle_auto_unload_log_includes_configured_timeout(monkeypatch): @pytest.mark.asyncio async def test_load_model_schedules_idle_unload_when_enabled(monkeypatch): """Startup-style model loads also schedule unload without waiting for traffic.""" - monkeypatch.setattr(settings, "model_auto_unload_enabled", True) monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 0.01) manager = ModelManager() @@ -264,7 +260,6 @@ async def test_load_model_does_not_schedule_idle_unload_during_active_request( monkeypatch, ): """Lazy loads during generation wait for request completion before scheduling.""" - monkeypatch.setattr(settings, "model_auto_unload_enabled", True) monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 0.01) manager = ModelManager() @@ -285,7 +280,6 @@ async def test_load_model_does_not_schedule_idle_unload_during_active_request( @pytest.mark.asyncio async def test_active_request_blocks_idle_unload(monkeypatch): """The idle timer does not unload while generation is still active.""" - monkeypatch.setattr(settings, "model_auto_unload_enabled", True) monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 0.01) manager = ModelManager() @@ -313,7 +307,6 @@ async def slow_generate(*args, **kwargs): def test_status_reports_model_lifecycle_state(monkeypatch): - monkeypatch.setattr(settings, "model_auto_unload_enabled", True) monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 30.0) manager = ModelManager() manager._backend = MagicMock() diff --git a/docs/configuration.md b/docs/configuration.md index 80914f3a..8dd3d73d 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -154,8 +154,7 @@ Names are the field names from `api/src/core/config.py`, uppercased. Unrecognize |---|---|---| | `ENABLE_DEBUG_ENDPOINTS` | `false` | Expose `/debug/*` host and process introspection | | `ALLOW_DEV_UNLOAD` | `false` | Expose `/dev/model`, `POST /dev/unload`, and `POST /dev/reload` | -| `MODEL_AUTO_UNLOAD_ENABLED` | `false` | Unload the model after the idle timeout | -| `MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS` | `300.0` | Idle seconds before auto-unload when enabled | +| `MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS` | `0.0` | Idle seconds before auto-unload; `0` disables auto-unload | ## Logging From 3d9efd70ebd56724c8774f6feab82cdef1ab358b Mon Sep 17 00:00:00 2001 From: brycehenson Date: Thu, 20 Aug 2026 06:57:35 +0000 Subject: [PATCH 13/19] changelog --- CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 486c93c4..2a5d23ae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,9 @@ Notable changes to this project will be documented in this file. Per-PR attribution and contributor credits are published automatically on the corresponding GitHub release page; this file is the curated, human-readable summary. ## [Unreleased] +### Added +- Optional model auto-unload after an idle timeout (`MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS`, default off) to release VRAM. Reloads on the next request. `/dev/model` reports load/idle state and `POST /dev/reload` pre-warms the model, both behind `ALLOW_DEV_UNLOAD`. + ### Fixed - Native Windows installs (`start-cpu.ps1` etc) no longer need a C++ toolchain: `pyopenjtalk-plus` (a drop-in fork with prebuilt Windows wheels) replaces `pyopenjtalk` on win32 only (#508, proposed by @siliconfps). Needs a recent `uv`. Linux, macOS, and Docker are unchanged. From 9af63961f3fd168b06ee274d9341e0892fd6d364 Mon Sep 17 00:00:00 2001 From: brycehenson Date: Thu, 20 Aug 2026 07:16:15 +0000 Subject: [PATCH 14/19] changelog, remove notes, change option name to "move_to_cpu" (from "cpu_cache") --- CHANGELOG.md | 1 + README.md | 2 +- api/src/core/config.py | 2 +- api/src/inference/kokoro_v1.py | 2 +- api/src/inference/model_manager.py | 8 +-- api/tests/test_kokoro_v1.py | 6 +-- api/tests/test_model_unload.py | 22 ++++---- docker/gpu/docker-compose.yml | 2 +- docs/configuration.md | 2 +- notes.md | 85 ------------------------------ 10 files changed, 24 insertions(+), 108 deletions(-) delete mode 100644 notes.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 2a5d23ae..c86ad31f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ Per-PR attribution and contributor credits are published automatically on the co ## [Unreleased] ### Added - Optional model auto-unload after an idle timeout (`MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS`, default off) to release VRAM. Reloads on the next request. `/dev/model` reports load/idle state and `POST /dev/reload` pre-warms the model, both behind `ALLOW_DEV_UNLOAD`. +- `MODEL_UNLOAD_STRATEGY=move_to_cpu` moves model weights from GPU to system RAM on unload, so reloads are faster than the default `destroy` strategy which loads from disk. ### Fixed - Native Windows installs (`start-cpu.ps1` etc) no longer need a C++ toolchain: `pyopenjtalk-plus` (a drop-in fork with prebuilt Windows wheels) replaces `pyopenjtalk` on win32 only (#508, proposed by @siliconfps). Needs a recent `uv`. Linux, macOS, and Docker are unchanged. diff --git a/README.md b/README.md index 5bc9962c..d42ad527 100644 --- a/README.md +++ b/README.md @@ -700,7 +700,7 @@ Key Performance Metrics: `POST /dev/unload` frees the model from VRAM and reloads lazily on the next request. Set `ALLOW_DEV_UNLOAD=true` to expose the lifecycle controls: `GET /dev/model`, `POST /dev/unload`, and `POST /dev/reload`. -For shared-GPU hosts, set `MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS` above `0` to unload automatically after the model has been idle. In-flight generation keeps the model loaded; the next request reloads it automatically, or `POST /dev/reload` can pre-warm it. Set `MODEL_UNLOAD_STRATEGY=cpu_cache` to keep model weights in system RAM for faster reload while still clearing GPU memory; the default `destroy` strategy releases model objects completely. +For shared-GPU hosts, set `MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS` above `0` to unload automatically after the model has been idle. In-flight generation keeps the model loaded; the next request reloads it automatically, or `POST /dev/reload` can pre-warm it. Set `MODEL_UNLOAD_STRATEGY=move_to_cpu` to move model weights from GPU to system RAM on unload for faster reload while still clearing GPU memory. The default `destroy` strategy releases model objects completely and will load from disk. Reclaim scales with load (the activation pool, not just weights) but plateaus: chunks cap at 450 tokens. Long-form = ~30 paragraphs. Same setup as above. diff --git a/api/src/core/config.py b/api/src/core/config.py index 225ed309..b6027e89 100644 --- a/api/src/core/config.py +++ b/api/src/core/config.py @@ -43,7 +43,7 @@ class Settings(BaseSettings): model_auto_unload_timeout_seconds: float = ( 0.0 # Idle seconds before unloading; 0 disables auto-unload ) - model_unload_strategy: str = "destroy" # "destroy" or "cpu_cache" + model_unload_strategy: str = "destroy" # "destroy" or "move_to_cpu" enable_debug_endpoints: bool = ( False # Whether to expose /debug/* host and process introspection routes ) diff --git a/api/src/inference/kokoro_v1.py b/api/src/inference/kokoro_v1.py index e9b8ef3e..4f6bb827 100644 --- a/api/src/inference/kokoro_v1.py +++ b/api/src/inference/kokoro_v1.py @@ -501,7 +501,7 @@ def _clear_memory(self) -> None: def unload(self, strategy: str = "destroy") -> None: """Unload model and free resources.""" - if strategy == "cpu_cache" and self._offload_model_to_cpu(): + if strategy == "move_to_cpu" and self._offload_model_to_cpu(): return logger.info("Destroying Kokoro model backend state") diff --git a/api/src/inference/model_manager.py b/api/src/inference/model_manager.py index 2e0e5793..dcea1a04 100644 --- a/api/src/inference/model_manager.py +++ b/api/src/inference/model_manager.py @@ -178,14 +178,14 @@ def _backend_is_cpu_cached(self) -> bool: def _unload_strategy(self) -> str: strategy = settings.model_unload_strategy.strip().lower() - if strategy not in {"destroy", "cpu_cache"}: + if strategy not in {"destroy", "move_to_cpu"}: logger.warning( f"Unknown MODEL_UNLOAD_STRATEGY={settings.model_unload_strategy!r}; using destroy" ) return "destroy" - if strategy == "cpu_cache" and not settings.use_gpu: + if strategy == "move_to_cpu" and not settings.use_gpu: logger.warning( - "MODEL_UNLOAD_STRATEGY=cpu_cache requires USE_GPU=true; using destroy" + "MODEL_UNLOAD_STRATEGY=move_to_cpu requires USE_GPU=true; using destroy" ) return "destroy" return strategy @@ -350,7 +350,7 @@ async def reload(self) -> None: async with self._lock: self._cancel_idle_unload_timer() if ( - self._unload_strategy() == "cpu_cache" + self._unload_strategy() == "move_to_cpu" and self._backend is not None and self._backend_is_loaded() ): diff --git a/api/tests/test_kokoro_v1.py b/api/tests/test_kokoro_v1.py index a30a4d1b..02c52476 100644 --- a/api/tests/test_kokoro_v1.py +++ b/api/tests/test_kokoro_v1.py @@ -68,10 +68,10 @@ def test_unload_with_pipelines(kokoro_backend): assert kokoro_backend._voice_cache == {} # Voice tensors should be released -def test_cpu_cache_unload_moves_model_to_cpu_and_clears_runtime_caches( +def test_move_to_cpu_unload_moves_model_to_cpu_and_clears_runtime_caches( kokoro_backend, ): - """CPU-cache unload keeps model weights but clears device-backed runtime state.""" + """move_to_cpu unload keeps model weights but clears device-backed runtime state.""" cuda_model = MagicMock() cpu_model = MagicMock() cuda_model.cpu.return_value = cpu_model @@ -81,7 +81,7 @@ def test_cpu_cache_unload_moves_model_to_cpu_and_clears_runtime_caches( kokoro_backend._voice_cache = {"af_heart.pt:cuda": MagicMock()} with patch.object(kokoro_backend, "_clear_memory") as mock_clear: - kokoro_backend.unload(strategy="cpu_cache") + kokoro_backend.unload(strategy="move_to_cpu") cuda_model.cpu.assert_called_once() mock_clear.assert_called_once() diff --git a/api/tests/test_model_unload.py b/api/tests/test_model_unload.py index e8d62825..efc51089 100644 --- a/api/tests/test_model_unload.py +++ b/api/tests/test_model_unload.py @@ -66,9 +66,9 @@ async def test_unload_clears_backend(): @pytest.mark.asyncio -async def test_unload_cpu_cache_keeps_backend(monkeypatch): +async def test_unload_move_to_cpu_keeps_backend(monkeypatch): monkeypatch.setattr(settings, "use_gpu", True) - monkeypatch.setattr(settings, "model_unload_strategy", "cpu_cache") + monkeypatch.setattr(settings, "model_unload_strategy", "move_to_cpu") manager = ModelManager() mock_backend = MagicMock() @@ -80,13 +80,13 @@ async def test_unload_cpu_cache_keeps_backend(monkeypatch): mock_torch.cuda.is_available.return_value = False await manager.unload() - mock_backend.unload.assert_called_once_with(strategy="cpu_cache") + mock_backend.unload.assert_called_once_with(strategy="move_to_cpu") assert manager._backend is mock_backend -def test_cpu_cache_strategy_warns_and_uses_destroy_without_gpu(monkeypatch): +def test_move_to_cpu_strategy_warns_and_uses_destroy_without_gpu(monkeypatch): monkeypatch.setattr(settings, "use_gpu", False) - monkeypatch.setattr(settings, "model_unload_strategy", "cpu_cache") + monkeypatch.setattr(settings, "model_unload_strategy", "move_to_cpu") manager = ModelManager() @@ -94,14 +94,14 @@ def test_cpu_cache_strategy_warns_and_uses_destroy_without_gpu(monkeypatch): assert manager._unload_strategy() == "destroy" mock_warning.assert_called_once_with( - "MODEL_UNLOAD_STRATEGY=cpu_cache requires USE_GPU=true; using destroy" + "MODEL_UNLOAD_STRATEGY=move_to_cpu requires USE_GPU=true; using destroy" ) @pytest.mark.asyncio -async def test_unload_cpu_cache_destroys_backend_without_gpu(monkeypatch): +async def test_unload_move_to_cpu_destroys_backend_without_gpu(monkeypatch): monkeypatch.setattr(settings, "use_gpu", False) - monkeypatch.setattr(settings, "model_unload_strategy", "cpu_cache") + monkeypatch.setattr(settings, "model_unload_strategy", "move_to_cpu") manager = ModelManager() mock_backend = MagicMock() @@ -120,7 +120,7 @@ async def test_unload_cpu_cache_destroys_backend_without_gpu(monkeypatch): @pytest.mark.asyncio async def test_reload_restores_cpu_cached_backend(monkeypatch): monkeypatch.setattr(settings, "use_gpu", True) - monkeypatch.setattr(settings, "model_unload_strategy", "cpu_cache") + monkeypatch.setattr(settings, "model_unload_strategy", "move_to_cpu") manager = ModelManager() mock_backend = MagicMock() @@ -405,7 +405,7 @@ def test_status_reports_model_lifecycle_state(monkeypatch): def test_status_reports_cpu_cached_state(monkeypatch): monkeypatch.setattr(settings, "use_gpu", True) - monkeypatch.setattr(settings, "model_unload_strategy", "cpu_cache") + monkeypatch.setattr(settings, "model_unload_strategy", "move_to_cpu") monkeypatch.setattr(settings, "model_auto_unload_timeout_seconds", 30.0) manager = ModelManager() @@ -420,7 +420,7 @@ def test_status_reports_cpu_cached_state(monkeypatch): assert status["loaded"] is False assert status["cpu_cached"] is True - assert status["unload_strategy"] == "cpu_cache" + assert status["unload_strategy"] == "move_to_cpu" assert status["seconds_until_auto_unload"] is None diff --git a/docker/gpu/docker-compose.yml b/docker/gpu/docker-compose.yml index bbef5710..41352dea 100644 --- a/docker/gpu/docker-compose.yml +++ b/docker/gpu/docker-compose.yml @@ -30,7 +30,7 @@ services: - DOWNLOAD_MODEL=true # - ALLOW_DEV_UNLOAD=true # - MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS=300 - # - MODEL_UNLOAD_STRATEGY=cpu_cache + # - MODEL_UNLOAD_STRATEGY=move_to_cpu # - ENABLE_DEBUG_ENDPOINTS=true deploy: resources: diff --git a/docs/configuration.md b/docs/configuration.md index 4bceb50b..d39e374e 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -155,7 +155,7 @@ Names are the field names from `api/src/core/config.py`, uppercased. Unrecognize | `ENABLE_DEBUG_ENDPOINTS` | `false` | Expose `/debug/*` host and process introspection | | `ALLOW_DEV_UNLOAD` | `false` | Expose `/dev/model`, `POST /dev/unload`, and `POST /dev/reload` | | `MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS` | `0.0` | Idle seconds before auto-unload; `0` disables auto-unload | -| `MODEL_UNLOAD_STRATEGY` | `destroy` | `destroy` releases model objects; `cpu_cache` keeps model weights in system RAM for faster reload | +| `MODEL_UNLOAD_STRATEGY` | `destroy` | `destroy` releases model objects; `move_to_cpu` moves model weights from GPU to system RAM for faster reload | ## Logging diff --git a/notes.md b/notes.md deleted file mode 100644 index dabb4e10..00000000 --- a/notes.md +++ /dev/null @@ -1,85 +0,0 @@ -# Kokoro GPU Unload Notes - -## Question - -Could reload after GPU unload be faster if the model stays resident in system RAM -and is moved back to CUDA on demand, instead of destroying and reconstructing the -backend from disk? - -## Benchmark Setup - -- Ran entirely inside the running container: `kokoro-tts-gpu-kokoro-tts-1` -- GPU: NVIDIA GeForce RTX 2070 SUPER -- PyTorch: `2.8.0+cu126` -- Model: `/app/api/src/models/v1_0/kokoro-v1_0.pth` -- Voice: `af_heart` -- Test text: short unload/reload sentence -- The API model was unloaded before benchmarking to free VRAM. - -No project files were changed for the benchmark. - -## Current API Behavior - -Measured through the live API using `POST /dev/unload` followed by -`POST /v1/audio/speech`. - -| Case | Time | -|---|---:| -| Cold request after unload, average | `1.835s` | -| Warm request, average | `0.135s` | -| Reload penalty, average | `1.700s` | - -Interpretation: with the current destroy/reload mechanism, the next short request -after unload pays about `+1.7s`. - -## Direct Model Load Benchmark - -Compared fresh model construction plus CUDA load against keeping the model object -in CPU RAM and moving it back to CUDA. - -| Case | Time | -|---|---:| -| Fresh model load, average | `1.692s` | -| CPU RAM to CUDA, average | `0.193s` | -| Estimated time saved | `1.499s` | - -Interpretation: most of the reload time is CPU-side reconstruction/loading, not -the CUDA transfer itself. - -## Backend-Level CPU Cache Benchmark - -This benchmark kept the backend/model/pipeline/voice state alive in CPU RAM, -then moved the model back to CUDA before generating. - -| Case | Time | -|---|---:| -| Fresh backend load plus cold generation, average | `2.247s` | -| CPU-cache reload plus generation, average | `0.334s` | -| Estimated time saved | `1.913s` | - -Interpretation: in an optimistic implementation, reload plus first generation was -roughly `6x` to `7x` faster than the current fresh backend path. - -## Takeaway - -CPU-RAM caching looks worth implementing. A likely design is an unload strategy -that moves the loaded model to CPU and clears CUDA memory, while preserving enough -backend state to avoid reconstructing the model from disk on the next request. - -Potential strategy setting: - -```env -MODEL_UNLOAD_STRATEGY=destroy -# or -MODEL_UNLOAD_STRATEGY=cpu_cache -``` - -Expected tradeoff: - -- `destroy`: maximum system RAM release, slower reload. -- `cpu_cache`: keeps more system RAM in use, much faster reload, still reclaims - most model VRAM. - -Important caveat: the backend-level benchmark is a best case because it preserved -pipeline and voice cache state. If an implementation preserves only model weights -but rebuilds pipeline or voice state, the speedup may be smaller. From 190b522dbc7444ba640beb1fb36b408333dd5d85 Mon Sep 17 00:00:00 2001 From: brycehenson Date: Thu, 20 Aug 2026 08:00:44 +0000 Subject: [PATCH 15/19] created benchmark script for MODEL_UNLOAD_STRATEGY. Reduce readme diff --- README.md | 7 +- .../benchmark_model_unload_strategies.py | 301 ++++++++++++++++++ 2 files changed, 305 insertions(+), 3 deletions(-) create mode 100644 examples/assorted_checks/benchmarks/benchmark_model_unload_strategies.py diff --git a/README.md b/README.md index d42ad527..e81df70d 100644 --- a/README.md +++ b/README.md @@ -698,10 +698,8 @@ Key Performance Metrics: ### Model Unload / VRAM Reclaim -`POST /dev/unload` frees the model from VRAM and reloads lazily on the next request. Set `ALLOW_DEV_UNLOAD=true` to expose the lifecycle controls: `GET /dev/model`, `POST /dev/unload`, and `POST /dev/reload`. - -For shared-GPU hosts, set `MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS` above `0` to unload automatically after the model has been idle. In-flight generation keeps the model loaded; the next request reloads it automatically, or `POST /dev/reload` can pre-warm it. Set `MODEL_UNLOAD_STRATEGY=move_to_cpu` to move model weights from GPU to system RAM on unload for faster reload while still clearing GPU memory. The default `destroy` strategy releases model objects completely and will load from disk. +`POST /dev/unload` frees the model from VRAM and reloads lazily on the next request. Reclaim scales with load (the activation pool, not just weights) but plateaus: chunks cap at 450 tokens. Long-form = ~30 paragraphs. Same setup as above.

@@ -716,7 +714,10 @@ Reclaim scales with load (the activation pool, not just weights) but plateaus: c Floor is host + CUDA context. Reproduce with `uv run --extra benchmarks assorted_checks/benchmarks/benchmark_model_unload.py` from `examples/`. +`POST /dev/reload` reloads the model and `GET /dev/model` reports model load state and auto-unload settings. Set `ALLOW_DEV_UNLOAD=true` to expose these controls. + To automatically unload the model after an idle timeout, set `MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS` to a positive number of seconds. The default `0` disables auto-unload. +Set `MODEL_UNLOAD_STRATEGY=move_to_cpu` to move model weights from GPU to system RAM on unload for faster reload while still clearing GPU memory. The default `destroy` strategy releases model objects completely and will load from disk. ### Transcription roundtrip (WER/CER) diff --git a/examples/assorted_checks/benchmarks/benchmark_model_unload_strategies.py b/examples/assorted_checks/benchmarks/benchmark_model_unload_strategies.py new file mode 100644 index 00000000..bade4ab1 --- /dev/null +++ b/examples/assorted_checks/benchmarks/benchmark_model_unload_strategies.py @@ -0,0 +1,301 @@ +#!/usr/bin/env python3 +"""Measure unload and reload timings against a running Kokoro service. + +The service must already be running with ALLOW_DEV_UNLOAD=true. This script +does not start or stop containers; it follows the convention of the other +benchmark scripts in this directory and only calls HTTP endpoints. + +Usage, from repo root with a service already running: + uv run --extra benchmarks \ + examples/assorted_checks/benchmarks/benchmark_model_unload_strategies.py \ + --trials 10 + +To benchmark both strategies manually from repo root, build the image once, +run the service with one strategy, run this benchmark, then repeat with the +other strategy: + + docker build -f docker/gpu/Dockerfile.optimized -t kokoro-benchmark-gpu . + + docker rm -f kokoro-benchmark 2>/dev/null || true + docker run -d \ + --name kokoro-benchmark \ + --gpus '"device=1"' \ + -p 8880:8880 \ + --env-file .env \ + -e PYTHONPATH=/app:/app/api \ + -e USE_GPU=true \ + -e PYTHONUNBUFFERED=1 \ + -e API_LOG_LEVEL=DEBUG \ + -e DOWNLOAD_MODEL=true \ + -e ALLOW_DEV_UNLOAD=true \ + -e MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS=0 \ + -e MODEL_UNLOAD_STRATEGY=move_to_cpu \ + -v "$PWD/api:/app/api" \ + -v "$PWD/web:/app/web" \ + --user 1001:1001 \ + kokoro-benchmark-gpu + uv run --extra benchmarks \ + examples/assorted_checks/benchmarks/benchmark_model_unload_strategies.py \ + --trials 10 --strategy move_to_cpu --output-prefix model_unload_move_to_cpu + + docker rm -f kokoro-benchmark + docker run -d \ + --name kokoro-benchmark \ + --gpus '"device=1"' \ + -p 8880:8880 \ + --env-file .env \ + -e PYTHONPATH=/app:/app/api \ + -e USE_GPU=true \ + -e PYTHONUNBUFFERED=1 \ + -e API_LOG_LEVEL=DEBUG \ + -e DOWNLOAD_MODEL=true \ + -e ALLOW_DEV_UNLOAD=true \ + -e MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS=0 \ + -e MODEL_UNLOAD_STRATEGY=destroy \ + -v "$PWD/api:/app/api" \ + -v "$PWD/web:/app/web" \ + --user 1001:1001 \ + kokoro-benchmark-gpu + uv run --extra benchmarks \ + examples/assorted_checks/benchmarks/benchmark_model_unload_strategies.py \ + --trials 10 --strategy destroy --output-prefix model_unload_destroy + + docker rm -f kokoro-benchmark + +""" +import argparse +import csv +import json +import os +import statistics +import time +from datetime import datetime +from typing import Any + +import requests + + +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +DEFAULT_URL = "http://127.0.0.1:8880" + + +def save_json_results(results: dict[str, Any], output_file: str) -> None: + os.makedirs(os.path.dirname(output_file), exist_ok=True) + with open(output_file, "w", encoding="utf-8") as handle: + json.dump(results, handle, indent=2) + + +def write_benchmark_stats(stats: list[dict[str, Any]], output_file: str) -> None: + os.makedirs(os.path.dirname(output_file), exist_ok=True) + with open(output_file, "w", encoding="utf-8") as handle: + for section in stats: + handle.write(f"=== {section['title']} ===\n\n") + for label, value in section["stats"].items(): + if isinstance(value, float): + handle.write(f"{label}: {value:.2f}\n") + else: + handle.write(f"{label}: {value}\n") + handle.write("\n") + + +def request( + method: str, base_url: str, path: str, timeout: int = 600 +) -> tuple[bytes, float]: + start = time.perf_counter() + response = requests.request(method, f"{base_url}{path}", timeout=timeout) + response.raise_for_status() + return response.content, time.perf_counter() - start + + +def model_status(base_url: str) -> dict: + body, _ = request("GET", base_url, "/dev/model", timeout=60) + return json.loads(body.decode("utf-8")) + + +def run_trial(strategy: str, trial: int, args: argparse.Namespace) -> dict: + print(f"\n=== {strategy} trial {trial}/{args.trials} ===") + + print(" unload") + _, unload_seconds = request( + "POST", args.url, "/dev/unload", timeout=args.endpoint_timeout + ) + time.sleep(args.settle) + + print(" reload") + _, reload_seconds = request( + "POST", args.url, "/dev/reload", timeout=args.endpoint_timeout + ) + time.sleep(args.settle) + + row = { + "strategy": strategy, + "trial": trial, + "initial_load_seconds": args.initial_load_seconds, + "unload_seconds": unload_seconds, + "reload_seconds": reload_seconds, + } + print( + " -> unload={unload:.3f}s reload={reload:.3f}s".format( + unload=unload_seconds, + reload=reload_seconds, + ) + ) + return row + + +def stats_for(values: list[float]) -> dict: + return { + "average": statistics.mean(values), + "standard_deviation": statistics.stdev(values) if len(values) > 1 else 0.0, + } + + +def summarize(results: list[dict]) -> list[dict]: + summary = [] + for strategy in sorted({row["strategy"] for row in results}): + rows = [row for row in results if row["strategy"] == strategy] + initial = [ + row["initial_load_seconds"] + for row in rows + if row["initial_load_seconds"] is not None + ] + unload = [row["unload_seconds"] for row in rows] + reload = [row["reload_seconds"] for row in rows] + summary.append( + { + "strategy": strategy, + "trials": len(rows), + "initial_load_seconds": stats_for(initial) if initial else None, + "unload_seconds": stats_for(unload), + "reload_seconds": stats_for(reload), + } + ) + return summary + + +def write_trials_csv(results: list[dict], output_file: str) -> None: + os.makedirs(os.path.dirname(output_file), exist_ok=True) + with open(output_file, "w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=list(results[0].keys())) + writer.writeheader() + writer.writerows(results) + + +def fmt(value: float | None) -> str: + if value is None: + return "n/a" + return f"{value:.3f}s" + + +def write_report(summary: list[dict], results: list[dict], output_file: str) -> None: + os.makedirs(os.path.dirname(output_file), exist_ok=True) + with open(output_file, "w", encoding="utf-8") as handle: + handle.write("# Model Unload Strategy Timing\n\n") + handle.write(f"Generated: {datetime.now().isoformat()}\n\n") + handle.write("Standard deviation is sample standard deviation across trials.\n\n") + handle.write( + "| strategy | initial load avg | initial load stddev | unload avg | " + "unload stddev | reload avg | reload stddev |\n" + ) + handle.write("|---|---:|---:|---:|---:|---:|---:|\n") + for row in summary: + initial = row["initial_load_seconds"] or {} + unload = row["unload_seconds"] + reload = row["reload_seconds"] + handle.write( + f"| {row['strategy']} | {fmt(initial.get('average'))} | " + f"{fmt(initial.get('standard_deviation'))} | " + f"{fmt(unload['average'])} | {fmt(unload['standard_deviation'])} | " + f"{fmt(reload['average'])} | {fmt(reload['standard_deviation'])} |\n" + ) + + handle.write("\n## Trial Data\n\n") + handle.write("| strategy | trial | initial load | unload | reload |\n") + handle.write("|---|---:|---:|---:|---:|\n") + for row in results: + handle.write( + f"| {row['strategy']} | {row['trial']} | " + f"{fmt(row['initial_load_seconds'])} | " + f"{fmt(row['unload_seconds'])} | {fmt(row['reload_seconds'])} |\n" + ) + + +def write_stats_file(summary: list[dict], output_file: str) -> None: + stats = [] + for row in summary: + values = {} + for label in ("initial_load_seconds", "unload_seconds", "reload_seconds"): + section = row[label] + if section is None: + continue + values[f"{label}_average"] = section["average"] + values[f"{label}_standard_deviation"] = section["standard_deviation"] + stats.append( + {"title": f"Model Unload Strategy - {row['strategy']}", "stats": values} + ) + write_benchmark_stats(stats, output_file) + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--url", default=DEFAULT_URL) + ap.add_argument("--trials", type=int, default=10) + ap.add_argument("--strategy", help="expected strategy; defaults to /dev/model") + ap.add_argument("--initial-load-seconds", type=float) + ap.add_argument("--settle", type=float, default=1.0) + ap.add_argument("--endpoint-timeout", type=int, default=600) + ap.add_argument("--output-json") + ap.add_argument("--output-prefix", default="model_unload_strategy") + ap.add_argument("--no-report", action="store_true") + args = ap.parse_args() + + status = model_status(args.url) + strategy = args.strategy or status.get("unload_strategy") + if not strategy: + raise RuntimeError("could not determine strategy from /dev/model") + if status.get("unload_strategy") != strategy: + raise RuntimeError( + f"expected unload_strategy={strategy!r}, got {status.get('unload_strategy')!r}" + ) + + results = [run_trial(strategy, trial, args) for trial in range(1, args.trials + 1)] + summary = summarize(results) + payload = { + "timestamp": datetime.now().isoformat(), + "url": args.url, + "model": status, + "results": results, + "summary": summary, + } + + output_data_dir = os.path.join(SCRIPT_DIR, "output_data") + json_path = args.output_json or os.path.join( + output_data_dir, f"{args.output_prefix}_results.json" + ) + save_json_results(payload, json_path) + + if not args.no_report: + write_trials_csv( + results, + os.path.join(output_data_dir, f"{args.output_prefix}_trials.csv"), + ) + write_report( + summary, + results, + os.path.join(output_data_dir, f"{args.output_prefix}_report.md"), + ) + write_stats_file( + summary, + os.path.join(output_data_dir, f"{args.output_prefix}_stats.txt"), + ) + + print("\ndone.") + print(f"- {json_path}") + if not args.no_report: + print(f"- {os.path.join(output_data_dir, f'{args.output_prefix}_trials.csv')}") + print(f"- {os.path.join(output_data_dir, f'{args.output_prefix}_report.md')}") + print(f"- {os.path.join(output_data_dir, f'{args.output_prefix}_stats.txt')}") + + +if __name__ == "__main__": + main() From 692bd26c34768d1521d36026f3f801cd2f88e3b2 Mon Sep 17 00:00:00 2001 From: brycehenson Date: Thu, 20 Aug 2026 08:03:46 +0000 Subject: [PATCH 16/19] extra space in readme --- README.md | 1 - 1 file changed, 1 deletion(-) diff --git a/README.md b/README.md index e81df70d..f1ac4b29 100644 --- a/README.md +++ b/README.md @@ -698,7 +698,6 @@ Key Performance Metrics: ### Model Unload / VRAM Reclaim - `POST /dev/unload` frees the model from VRAM and reloads lazily on the next request. Reclaim scales with load (the activation pool, not just weights) but plateaus: chunks cap at 450 tokens. Long-form = ~30 paragraphs. Same setup as above. From 0a41a3c839cb72d149dfe96625327eb1281ac0d7 Mon Sep 17 00:00:00 2001 From: brycehenson Date: Thu, 20 Aug 2026 08:22:32 +0000 Subject: [PATCH 17/19] remove initial_load_seconds from benchmark_model_unload_strategies.py --- CHANGELOG.md | 2 +- .../benchmark_model_unload_strategies.py | 33 +++++-------------- 2 files changed, 9 insertions(+), 26 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c86ad31f..2ede87fc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,7 +7,7 @@ Per-PR attribution and contributor credits are published automatically on the co ## [Unreleased] ### Added - Optional model auto-unload after an idle timeout (`MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS`, default off) to release VRAM. Reloads on the next request. `/dev/model` reports load/idle state and `POST /dev/reload` pre-warms the model, both behind `ALLOW_DEV_UNLOAD`. -- `MODEL_UNLOAD_STRATEGY=move_to_cpu` moves model weights from GPU to system RAM on unload, so reloads are faster than the default `destroy` strategy which loads from disk. +- `MODEL_UNLOAD_STRATEGY=move_to_cpu` moves model weights from GPU to system RAM on unload, so reloads are faster than the default `destroy` strategy, which reloads from disk. ### Fixed - Native Windows installs (`start-cpu.ps1` etc) no longer need a C++ toolchain: `pyopenjtalk-plus` (a drop-in fork with prebuilt Windows wheels) replaces `pyopenjtalk` on win32 only (#508, proposed by @siliconfps). Needs a recent `uv`. Linux, macOS, and Docker are unchanged. diff --git a/examples/assorted_checks/benchmarks/benchmark_model_unload_strategies.py b/examples/assorted_checks/benchmarks/benchmark_model_unload_strategies.py index bade4ab1..747523fd 100644 --- a/examples/assorted_checks/benchmarks/benchmark_model_unload_strategies.py +++ b/examples/assorted_checks/benchmarks/benchmark_model_unload_strategies.py @@ -21,7 +21,6 @@ --name kokoro-benchmark \ --gpus '"device=1"' \ -p 8880:8880 \ - --env-file .env \ -e PYTHONPATH=/app:/app/api \ -e USE_GPU=true \ -e PYTHONUNBUFFERED=1 \ @@ -61,7 +60,7 @@ --trials 10 --strategy destroy --output-prefix model_unload_destroy docker rm -f kokoro-benchmark - + """ import argparse import csv @@ -130,7 +129,6 @@ def run_trial(strategy: str, trial: int, args: argparse.Namespace) -> dict: row = { "strategy": strategy, "trial": trial, - "initial_load_seconds": args.initial_load_seconds, "unload_seconds": unload_seconds, "reload_seconds": reload_seconds, } @@ -154,18 +152,12 @@ def summarize(results: list[dict]) -> list[dict]: summary = [] for strategy in sorted({row["strategy"] for row in results}): rows = [row for row in results if row["strategy"] == strategy] - initial = [ - row["initial_load_seconds"] - for row in rows - if row["initial_load_seconds"] is not None - ] unload = [row["unload_seconds"] for row in rows] reload = [row["reload_seconds"] for row in rows] summary.append( { "strategy": strategy, "trials": len(rows), - "initial_load_seconds": stats_for(initial) if initial else None, "unload_seconds": stats_for(unload), "reload_seconds": stats_for(reload), } @@ -193,29 +185,23 @@ def write_report(summary: list[dict], results: list[dict], output_file: str) -> handle.write("# Model Unload Strategy Timing\n\n") handle.write(f"Generated: {datetime.now().isoformat()}\n\n") handle.write("Standard deviation is sample standard deviation across trials.\n\n") - handle.write( - "| strategy | initial load avg | initial load stddev | unload avg | " - "unload stddev | reload avg | reload stddev |\n" - ) - handle.write("|---|---:|---:|---:|---:|---:|---:|\n") + handle.write("| strategy | unload avg | unload stddev | reload avg | reload stddev |\n") + handle.write("|---|---:|---:|---:|---:|\n") for row in summary: - initial = row["initial_load_seconds"] or {} unload = row["unload_seconds"] reload = row["reload_seconds"] handle.write( - f"| {row['strategy']} | {fmt(initial.get('average'))} | " - f"{fmt(initial.get('standard_deviation'))} | " - f"{fmt(unload['average'])} | {fmt(unload['standard_deviation'])} | " + f"| {row['strategy']} | {fmt(unload['average'])} | " + f"{fmt(unload['standard_deviation'])} | " f"{fmt(reload['average'])} | {fmt(reload['standard_deviation'])} |\n" ) handle.write("\n## Trial Data\n\n") - handle.write("| strategy | trial | initial load | unload | reload |\n") - handle.write("|---|---:|---:|---:|---:|\n") + handle.write("| strategy | trial | unload | reload |\n") + handle.write("|---|---:|---:|---:|\n") for row in results: handle.write( f"| {row['strategy']} | {row['trial']} | " - f"{fmt(row['initial_load_seconds'])} | " f"{fmt(row['unload_seconds'])} | {fmt(row['reload_seconds'])} |\n" ) @@ -224,10 +210,8 @@ def write_stats_file(summary: list[dict], output_file: str) -> None: stats = [] for row in summary: values = {} - for label in ("initial_load_seconds", "unload_seconds", "reload_seconds"): + for label in ("unload_seconds", "reload_seconds"): section = row[label] - if section is None: - continue values[f"{label}_average"] = section["average"] values[f"{label}_standard_deviation"] = section["standard_deviation"] stats.append( @@ -241,7 +225,6 @@ def main(): ap.add_argument("--url", default=DEFAULT_URL) ap.add_argument("--trials", type=int, default=10) ap.add_argument("--strategy", help="expected strategy; defaults to /dev/model") - ap.add_argument("--initial-load-seconds", type=float) ap.add_argument("--settle", type=float, default=1.0) ap.add_argument("--endpoint-timeout", type=int, default=600) ap.add_argument("--output-json") From 02ecef7797d574b88340fb3a6b4aeec0367bcc63 Mon Sep 17 00:00:00 2001 From: brycehenson Date: Thu, 20 Aug 2026 08:24:17 +0000 Subject: [PATCH 18/19] remove newline --- README.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/README.md b/README.md index f1ac4b29..339efea6 100644 --- a/README.md +++ b/README.md @@ -698,8 +698,7 @@ Key Performance Metrics: ### Model Unload / VRAM Reclaim -`POST /dev/unload` frees the model from VRAM and reloads lazily on the next request. -Reclaim scales with load (the activation pool, not just weights) but plateaus: chunks cap at 450 tokens. Long-form = ~30 paragraphs. Same setup as above. +`POST /dev/unload` frees the model from VRAM and reloads lazily on the next request. Reclaim scales with load (the activation pool, not just weights) but plateaus: chunks cap at 450 tokens. Long-form = ~30 paragraphs. Same setup as above.

Short workload From 39ed88b2f26ef8fd15892b5e59efcbb6bc5bd1bb Mon Sep 17 00:00:00 2001 From: brycehenson Date: Thu, 20 Aug 2026 09:00:54 +0000 Subject: [PATCH 19/19] readme Debug Endpoints, clarify that ALLOW_DEV_UNLOAD=true is needed. config.py inline comment on allow_dev_unload --- README.md | 4 ++-- api/src/core/config.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 079b9020..46556cad 100644 --- a/README.md +++ b/README.md @@ -761,9 +761,9 @@ System state and resource usage, for debugging exhaustion or performance issues. - `/debug/threads` - Get thread information and stack traces - `/debug/storage` - Disk usage per mounted partition - `/debug/system` - Get system information (CPU, memory, GPU) -- `/dev/model` - Get model load state and auto-unload timing +- `/dev/model` - Get model load state and auto-unload timing. Off by default; set `ALLOW_DEV_UNLOAD=true` to enable - `POST /dev/unload` - Release model from VRAM; reloads lazily on next request. Off by default; set `ALLOW_DEV_UNLOAD=true` to enable -- `POST /dev/reload` - Load the model immediately after an unload or before traffic arrives +- `POST /dev/reload` - Load the model into VRAM. Off by default; set `ALLOW_DEV_UNLOAD=true` to enable Stability: the `/v1/*` OpenAI-compatible routes are the stable API. `/dev/*` and `/debug/*` are operational helpers, and may change or move behind flags between minor releases. diff --git a/api/src/core/config.py b/api/src/core/config.py index 8bc40c02..3f6887f0 100644 --- a/api/src/core/config.py +++ b/api/src/core/config.py @@ -39,7 +39,7 @@ class Settings(BaseSettings): allow_local_voice_saving: bool = ( False # Whether to allow saving combined voices locally ) - allow_dev_unload: bool = False # Whether to expose the POST /dev/unload endpoint + allow_dev_unload: bool = False # Whether to expose /dev/model, POST /dev/unload, and POST /dev/reload model_auto_unload_timeout_seconds: float = ( 0.0 # Idle seconds before unloading; 0 disables auto-unload )