Add idle model unload support

.github/workflows/tests/test_api_validation.py

@@ -6,6 +6,7 @@
 import tempfile
 import types
 import unittest
+from unittest import mock
 
 
 class _HTTPException(Exception):
@@ -218,6 +219,20 @@ def transcribe(self, _path, **kwargs):
         return iter([_FakeSegment()]), _FakeInfo()
 
 
+class _CapturedWhisperModel:
+    instances = []
+
+    def __init__(self, model_name, **kwargs):
+        self.model_name = model_name
+        self.kwargs = kwargs
+        self.calls = []
+        type(self).instances.append(self)
+
+    def transcribe(self, _path, **kwargs):
+        self.calls.append(kwargs)
+        return iter([_FakeSegment()]), _FakeInfo()
+
+
 class BeamPropagationTests(unittest.TestCase):
     def setUp(self):
         self.old_model = api_server._model
@@ -301,5 +316,134 @@ async def collect_stream():
         self.assertTrue(any("transcript.text.done" in frame for frame in frames))
 
 
+class IdleUnloadFeatureTests(unittest.TestCase):
+    def setUp(self):
+        self.old_model = api_server._model
+        self.old_model_name = api_server._model_name
+        self.old_model_config = api_server._model_config
+        self.old_last_model_used_at = api_server._last_model_used_at
+        self.old_active_inferences = api_server._active_inferences
+        self.old_idle_unload_seconds = api_server._idle_unload_seconds
+        self.old_beam_size = api_server._beam_size
+        self.old_max_request_beam = api_server._max_request_beam
+        self.old_word_timestamps = api_server._word_timestamps
+        self.old_max_upload_bytes = api_server._max_upload_bytes
+        self.old_diarization_enabled = api_server._diarization_enabled
+
+    def tearDown(self):
+        api_server._model = self.old_model
+        api_server._model_name = self.old_model_name
+        api_server._model_config = self.old_model_config
+        api_server._last_model_used_at = self.old_last_model_used_at
+        api_server._active_inferences = self.old_active_inferences
+        api_server._idle_unload_seconds = self.old_idle_unload_seconds
+        api_server._beam_size = self.old_beam_size
+        api_server._max_request_beam = self.old_max_request_beam
+        api_server._word_timestamps = self.old_word_timestamps
+        api_server._max_upload_bytes = self.old_max_upload_bytes
+        api_server._diarization_enabled = self.old_diarization_enabled
+
+    def test_load_model_reads_idle_unload_seconds_and_caches_config(self):
+        env = {
+            "WHISPER_MODEL": "medium",
+            "WHISPER_DEVICE": "cuda",
+            "WHISPER_COMPUTE_TYPE": "float16",
+            "WHISPER_THREADS": "6",
+            "HF_HOME": "/cache/whisper",
+            "WHISPER_LOCAL_ONLY": "1",
+            "WHISPER_BEAM": "8",
+            "WHISPER_MAX_REQUEST_BEAM": "12",
+            "WHISPER_WORD_TIMESTAMPS": "true",
+            "WHISPER_MAX_UPLOAD_MB": "256",
+            "WHISPER_IDLE_UNLOAD_SECONDS": "900",
+        }
+        fake_fw = types.ModuleType("faster_whisper")
+        fake_fw.WhisperModel = _CapturedWhisperModel
+
+        with mock.patch.dict(sys.modules, {"faster_whisper": fake_fw}):
+            with mock.patch.object(api_server, "_load_model_from_config", autospec=True) as load_config:
+                with mock.patch.dict(os.environ, env, clear=False):
+                    api_server._load_model()
+
+        load_config.assert_called_once_with()
+        self.assertEqual(api_server._idle_unload_seconds, 900)
+        self.assertEqual(api_server._beam_size, 8)
+        self.assertEqual(api_server._max_request_beam, 12)
+        self.assertTrue(api_server._word_timestamps)
+        self.assertEqual(api_server._max_upload_bytes, 256 * 1024 * 1024)
+        self.assertEqual(
+            api_server._model_config,
+            {
+                "model_name": "medium",
+                "device": "cuda",
+                "compute_type": "float16",
+                "threads": 6,
+                "cache_dir": "/cache/whisper",
+                "local_files_only": True,
+            },
+        )
+
+    def test_idle_unload_recovers_on_next_request(self):
+        fake_fw = types.ModuleType("faster_whisper")
+        fake_fw.WhisperModel = _CapturedWhisperModel
+        _CapturedWhisperModel.instances.clear()
+
+        api_server._model = None
+        api_server._model_name = None
+        api_server._model_config = {
+            "model_name": "base",
+            "device": "cpu",
+            "compute_type": "int8",
+            "threads": 2,
+            "cache_dir": "/cache/whisper",
+            "local_files_only": False,
+        }
+        api_server._last_model_used_at = 0
+        api_server._active_inferences = 0
+        api_server._beam_size = 5
+        api_server._max_request_beam = 10
+        api_server._word_timestamps = False
+
+        with mock.patch.dict(sys.modules, {"faster_whisper": fake_fw}):
+            response = asyncio.run(api_server._handle_audio(
+                task="transcribe",
+                file=_FakeUpload(),
+                model="whisper-1",
+                language=None,
+                prompt=None,
+                response_format="json",
+                temperature=0,
+                stream=None,
+                beam=None,
+            ))
+
+        self.assertEqual(response.content, {"text": "hello"})
+        self.assertEqual(len(_CapturedWhisperModel.instances), 1)
+        instance = _CapturedWhisperModel.instances[0]
+        self.assertEqual(instance.model_name, "base")
+        self.assertEqual(instance.kwargs["download_root"], "/cache/whisper")
+        self.assertIs(api_server._model, instance)
+        self.assertEqual(api_server._model_name, "base")
+        self.assertGreater(api_server._last_model_used_at, 0)
+        self.assertEqual(instance.calls[0]["beam_size"], 5)
+
+    def test_release_model_clears_cached_model_and_cuda_cache(self):
+        empty_cache_calls = []
+        fake_torch = types.ModuleType("torch")
+        fake_torch.cuda = types.SimpleNamespace(
+            is_available=lambda: True,
+            empty_cache=lambda: empty_cache_calls.append(True),
+        )
+
+        api_server._model = object()
+        api_server._model_name = "base"
+
+        with mock.patch.dict(sys.modules, {"torch": fake_torch}):
+            api_server._release_model("idle for 600s")
+
+        self.assertIsNone(api_server._model)
+        self.assertTrue(empty_cache_calls)
+
+
 if __name__ == "__main__":
     unittest.main()

README.md

@@ -183,6 +183,7 @@ This Docker image uses the following variables, that can be declared in an `env`
 | `WHISPER_BEAM` | Beam size for transcription and translation decoding. Higher values may improve accuracy at the cost of speed. Use `1` for fastest (greedy) decoding. | `5` |
 | `WHISPER_MAX_REQUEST_BEAM` | Maximum beam size allowed for the per-request `beam` override. Set to `0` to disable this limit. | `10` |
 | `WHISPER_MAX_UPLOAD_MB` | Maximum uploaded audio file size in MB. Requests above this limit return HTTP 413. Set to `0` to disable the limit. | `1024` |
+| `WHISPER_IDLE_UNLOAD_SECONDS` | Unload the active model after this many idle seconds to reduce RAM/VRAM use. Set to `0` to keep the model loaded. The next transcription or translation request reloads the model and may be slower. | `0` |
 | `WHISPER_LOCAL_ONLY` | When set to any non-empty value (e.g. `true`), disables all HuggingFace model downloads. For offline or air-gapped deployments with pre-cached models. | *(not set)* |
 | `WHISPER_WORD_TIMESTAMPS` | When set to `true`, enables word-level timestamps globally for all requests. The `verbose_json` output will include a top-level `words` array with per-word timing and confidence. Can also be enabled per-request via `timestamp_granularities[]=word`. | *(not set)* |
 | `WHISPER_DIARIZATION` | Set to `true` to enable speaker diarization. Identifies who is speaking in each segment. Uses [sherpa-onnx](https://github.com/k2-fsa/sherpa-onnx) with pyannote segmentation-3.0 ONNX models (~45 MB, auto-downloaded on first use). Not supported in streaming mode. | *(not set)* |

api_server.py

@@ -13,6 +13,7 @@
 """
 
 import asyncio
+import gc
 import json
 import logging
 import os
@@ -44,6 +45,11 @@
 
 _model = None       # WhisperModel instance
 _model_name = None  # name as loaded (e.g. "base")
+_model_config = None  # cached load config for lazy reloads
+_last_model_used_at = time.monotonic()
+_active_inferences = 0
+_idle_unload_seconds = 0
+_idle_monitor_task = None
 _beam_size = 5      # beam size used for transcription
 _max_request_beam = 10  # maximum per-request beam override; 0 disables the cap
 _word_timestamps = False  # default for word-level timestamps
@@ -118,7 +124,7 @@ def _load_diarizer() -> None:
 
 def _load_model() -> None:
     """Import and initialise the faster-whisper model from environment config."""
-    global _model, _model_name, _beam_size, _max_request_beam, _word_timestamps, _max_upload_bytes
+    global _model, _model_name, _model_config, _beam_size, _max_request_beam, _word_timestamps, _max_upload_bytes, _idle_unload_seconds
 
     from faster_whisper import WhisperModel  # deferred — keeps import fast
 
@@ -132,6 +138,7 @@ def _load_model() -> None:
     _max_request_beam = _env_int("WHISPER_MAX_REQUEST_BEAM", 10)
     _word_timestamps = os.environ.get("WHISPER_WORD_TIMESTAMPS", "").strip().lower() == "true"
     max_upload_mb    = _env_int("WHISPER_MAX_UPLOAD_MB", 1024)
+    _idle_unload_seconds = _env_int("WHISPER_IDLE_UNLOAD_SECONDS", 0)
     if _max_request_beam < 0:
         logger.error(
             "Invalid value for WHISPER_MAX_REQUEST_BEAM: %d (expected 0 or greater); using default 10",
@@ -146,28 +153,95 @@ def _load_model() -> None:
         max_upload_mb = 1024
     _max_upload_bytes = max_upload_mb * 1024 * 1024
 
+    _model_config = {
+        "model_name": model_name,
+        "device": device,
+        "compute_type": compute_type,
+        "threads": threads,
+        "cache_dir": cache_dir,
+        "local_files_only": local_files_only,
+    }
+    _load_model_from_config()
+
+
+def _load_model_from_config() -> None:
+    """Load the configured model if it is not currently resident."""
+    global _model, _model_name, _last_model_used_at
+    if _model is not None:
+        return
+    if not _model_config:
+        _load_model()
+        return
+    from faster_whisper import WhisperModel
+    cfg = _model_config
     logger.info(
-        "Loading model '%s' | device=%s compute_type=%s threads=%d beam=%d max_request_beam=%d word_ts=%s max_upload_mb=%d local_only=%s cache=%s",
-        model_name, device, compute_type, threads, _beam_size, _max_request_beam, _word_timestamps, max_upload_mb, local_files_only, cache_dir,
+        "Loading model '%s' | device=%s compute_type=%s threads=%d local_only=%s cache=%s",
+        cfg["model_name"], cfg["device"], cfg["compute_type"], cfg["threads"], cfg["local_files_only"], cfg["cache_dir"],
     )
     t0 = time.monotonic()
     _model = WhisperModel(
-        model_name,
-        device=device,
-        compute_type=compute_type,
-        cpu_threads=threads,
-        download_root=cache_dir,
-        local_files_only=local_files_only,
+        cfg["model_name"],
+        device=cfg["device"],
+        compute_type=cfg["compute_type"],
+        cpu_threads=cfg["threads"],
+        download_root=cfg["cache_dir"],
+        local_files_only=cfg["local_files_only"],
     )
-    _model_name = model_name
-    logger.info("Model '%s' ready in %.1fs", model_name, time.monotonic() - t0)
+    _model_name = cfg["model_name"]
+    _last_model_used_at = time.monotonic()
+    logger.info("Model '%s' ready in %.1fs", _model_name, time.monotonic() - t0)
+
+
+def _release_model(reason: str) -> None:
+    """Unload the resident model and ask CUDA-capable libraries to free caches."""
+    global _model
+    if _model is None:
+        return
+    logger.info("Unloading model '%s' (%s)", _model_name, reason)
+    _model = None
+    gc.collect()
+    try:
+        import torch
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    except Exception as exc:  # noqa: BLE001
+        logger.debug("CUDA cache cleanup skipped: %s", exc)
+
+
+async def _idle_unload_monitor() -> None:
+    """Periodically unload the model after WHISPER_IDLE_UNLOAD_SECONDS."""
+    if _idle_unload_seconds <= 0:
+        return
+    logger.info("Idle model unload enabled after %ds", _idle_unload_seconds)
+    while True:
+        await asyncio.sleep(max(5, min(60, _idle_unload_seconds // 2 or 5)))
+        if _model is None or _active_inferences > 0:
+            continue
+        if time.monotonic() - _last_model_used_at < _idle_unload_seconds:
+            continue
+        if not _inference_lock.acquire(blocking=False):
+            continue
+        try:
+            if _active_inferences == 0 and _model is not None and time.monotonic() - _last_model_used_at >= _idle_unload_seconds:
+                _release_model(f"idle for {_idle_unload_seconds}s")
+        finally:
+            _inference_lock.release()
 
 
 @asynccontextmanager
 async def _lifespan(app: FastAPI):
+    global _idle_monitor_task
     _load_model()
     _load_diarizer()
-    yield
+    if _idle_unload_seconds > 0:
+        _idle_monitor_task = asyncio.create_task(_idle_unload_monitor())
+    try:
+        yield
+    finally:
+        if _idle_monitor_task:
+            _idle_monitor_task.cancel()
+        with _inference_lock:
+            _release_model("shutdown")
 
 
 # ---------------------------------------------------------------------------
@@ -374,8 +448,11 @@ async def _stream_sse(
     seg_queue: asyncio.Queue = asyncio.Queue()
 
     def _run() -> None:
+        global _active_inferences, _last_model_used_at
         with _inference_lock:
+            _active_inferences += 1
             try:
+                _load_model_from_config()
                 segs_gen, _ = _model.transcribe(
                     tmp_path,
                     language=lang,
@@ -390,6 +467,8 @@ def _run() -> None:
             except Exception as exc:  # noqa: BLE001
                 loop.call_soon_threadsafe(seg_queue.put_nowait, exc)
             finally:
+                _last_model_used_at = time.monotonic()
+                _active_inferences = max(0, _active_inferences - 1)
                 loop.call_soon_threadsafe(seg_queue.put_nowait, None)  # sentinel
 
     loop.run_in_executor(None, _run)
@@ -477,8 +556,8 @@ async def _handle_audio(
     Shared implementation for transcription and translation endpoints.
     ``task`` is either ``"transcribe"`` or ``"translate"``.
     """
-    if _model is None:
-        raise HTTPException(status_code=503, detail="Model is not loaded yet. Please retry.")
+    # Model may have been unloaded by the idle monitor; it is reloaded below
+    # under _inference_lock immediately before transcription.
 
     _validate_temperature(temperature)
     request_beam = _resolve_request_beam(beam)
@@ -586,17 +665,24 @@ async def _handle_audio(
     try:
         try:
             with _inference_lock:
-                segments_gen, info = _model.transcribe(
-                    tmp_path,
-                    language=lang,
-                    task=task,
-                    initial_prompt=prompt or None,
-                    temperature=temperature,
-                    beam_size=request_beam,
-                    word_timestamps=wt_flag,
-                    vad_filter=True,
-                )
-                segments = list(segments_gen)  # consume the generator before the temp file is removed
+                global _active_inferences, _last_model_used_at
+                _active_inferences += 1
+                try:
+                    _load_model_from_config()
+                    segments_gen, info = _model.transcribe(
+                        tmp_path,
+                        language=lang,
+                        task=task,
+                        initial_prompt=prompt or None,
+                        temperature=temperature,
+                        beam_size=request_beam,
+                        word_timestamps=wt_flag,
+                        vad_filter=True,
+                    )
+                    segments = list(segments_gen)  # consume the generator before the temp file is removed
+                finally:
+                    _last_model_used_at = time.monotonic()
+                    _active_inferences = max(0, _active_inferences - 1)
 
         except HTTPException:
             raise