diff --git a/CHANGELOG.md b/CHANGELOG.md index 486c93c4..2ede87fc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,10 @@ 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`. +- `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/README.md b/README.md index 939092ed..46bcb73a 100644 --- a/README.md +++ b/README.md @@ -712,6 +712,11 @@ Key Performance Metrics: 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) 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/`. @@ -759,7 +764,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. 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 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 9d7f77ba..a0a2b7f2 100644 --- a/api/src/core/config.py +++ b/api/src/core/config.py @@ -39,7 +39,11 @@ 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 + ) + 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 cecc2584..4f6bb827 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 == "move_to_cpu" 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 b1468717..dcea1a04 100644 --- a/api/src/inference/model_manager.py +++ b/api/src/inference/model_manager.py @@ -1,6 +1,8 @@ """Kokoro V1 model management.""" import asyncio +import time +from contextlib import asynccontextmanager from typing import Optional import torch @@ -29,6 +31,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 +107,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) @@ -136,44 +142,253 @@ 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={self._auto_unload_enabled()}" + ) 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 _format_seconds(self, seconds: float) -> str: + return f"{seconds:g}s" + + def _auto_unload_enabled(self) -> bool: + return 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", "move_to_cpu"}: + logger.warning( + f"Unknown MODEL_UNLOAD_STRATEGY={settings.model_unload_strategy!r}; using destroy" + ) + return "destroy" + if strategy == "move_to_cpu" and not settings.use_gpu: + logger.warning( + "MODEL_UNLOAD_STRATEGY=move_to_cpu requires USE_GPU=true; using destroy" + ) + return "destroy" + return strategy + + 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 + ): + 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 + 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(): + 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: + 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 + ): + 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 of " + f"{self._format_seconds(self._auto_unload_timeout())}" + ) + except asyncio.CancelledError: + 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() + + @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.ensure_backend() - assert self._backend is not None, "ensure_backend left no backend" - try: - 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}") def unload_all(self) -> None: """Unload model and free resources.""" + self._cancel_idle_unload_timer() if self._backend: self._backend.unload() self._backend = 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: - 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") + 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() + if ( + self._unload_strategy() == "move_to_cpu" + 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("Manual model reload completed") + + 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_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_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": self._auto_unload_enabled(), + "auto_unload_timeout_seconds": timeout, + "idle_seconds": idle_for, + "seconds_until_auto_unload": unload_in, + } @property def current_backend(self) -> str: 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/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_kokoro_v1.py b/api/tests/test_kokoro_v1.py index 4aa62184..02c52476 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_move_to_cpu_unload_moves_model_to_cpu_and_clears_runtime_caches( + kokoro_backend, +): + """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 + 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="move_to_cpu") + + 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 05da4353..efc51089 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_timeout_seconds", 0.0) + monkeypatch.setattr(settings, "model_unload_strategy", "destroy") # --------------------------------------------------------------------------- @@ -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 @@ -62,6 +65,80 @@ async def test_unload_clears_backend(): assert manager._backend is None +@pytest.mark.asyncio +async def test_unload_move_to_cpu_keeps_backend(monkeypatch): + monkeypatch.setattr(settings, "use_gpu", True) + monkeypatch.setattr(settings, "model_unload_strategy", "move_to_cpu") + + 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="move_to_cpu") + assert manager._backend is mock_backend + + +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", "move_to_cpu") + + 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=move_to_cpu requires USE_GPU=true; using destroy" + ) + + +@pytest.mark.asyncio +async def test_unload_move_to_cpu_destroys_backend_without_gpu(monkeypatch): + monkeypatch.setattr(settings, "use_gpu", False) + monkeypatch.setattr(settings, "model_unload_strategy", "move_to_cpu") + + 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, "use_gpu", True) + monkeypatch.setattr(settings, "model_unload_strategy", "move_to_cpu") + + 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() @@ -185,6 +262,168 @@ 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_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_idle_auto_unload_log_includes_configured_timeout(monkeypatch): + 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.""" + 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_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.""" + 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_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["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, "use_gpu", True) + monkeypatch.setattr(settings, "model_unload_strategy", "move_to_cpu") + 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"] == "move_to_cpu" + assert status["seconds_until_auto_unload"] is None + + # --------------------------------------------------------------------------- # POST /dev/unload endpoint tests # --------------------------------------------------------------------------- @@ -261,3 +500,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..41352dea 100644 --- a/docker/gpu/docker-compose.yml +++ b/docker/gpu/docker-compose.yml @@ -29,6 +29,8 @@ services: - API_LOG_LEVEL=DEBUG - DOWNLOAD_MODEL=true # - ALLOW_DEV_UNLOAD=true + # - MODEL_AUTO_UNLOAD_TIMEOUT_SECONDS=300 + # - MODEL_UNLOAD_STRATEGY=move_to_cpu # - 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..d39e374e 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_TIMEOUT_SECONDS` | `0.0` | Idle seconds before auto-unload; `0` disables auto-unload | +| `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/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..747523fd --- /dev/null +++ b/examples/assorted_checks/benchmarks/benchmark_model_unload_strategies.py @@ -0,0 +1,284 @@ +#!/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 \ + -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, + "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] + unload = [row["unload_seconds"] for row in rows] + reload = [row["reload_seconds"] for row in rows] + summary.append( + { + "strategy": strategy, + "trials": len(rows), + "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 | unload avg | unload stddev | reload avg | reload stddev |\n") + handle.write("|---|---:|---:|---:|---:|\n") + for row in summary: + unload = row["unload_seconds"] + reload = row["reload_seconds"] + handle.write( + 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 | unload | reload |\n") + handle.write("|---|---:|---:|---:|\n") + for row in results: + handle.write( + f"| {row['strategy']} | {row['trial']} | " + 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 ("unload_seconds", "reload_seconds"): + section = row[label] + 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("--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()