From 0af09eae78270d8f5a3eb8bfe9b28bc6241ddf2d Mon Sep 17 00:00:00 2001 From: Valentine Ubani Mayaki Date: Fri, 17 Jul 2026 11:27:30 +0200 Subject: [PATCH 1/2] Add idle model unload support --- README.md | 1 + api_server.py | 136 ++++++++++++++++++++++++++++++++++++++++---------- run.sh | 7 +++ 3 files changed, 119 insertions(+), 25 deletions(-) diff --git a/README.md b/README.md index 72ec0c6..1842eab 100644 --- a/README.md +++ b/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)* | diff --git a/api_server.py b/api_server.py index 3a6f233..ae48bfb 100644 --- a/api_server.py +++ b/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 diff --git a/run.sh b/run.sh index 46c1de6..fbf9338 100644 --- a/run.sh +++ b/run.sh @@ -66,6 +66,8 @@ WHISPER_MAX_REQUEST_BEAM=$(nospaces "$WHISPER_MAX_REQUEST_BEAM") WHISPER_MAX_REQUEST_BEAM=$(noquotes "$WHISPER_MAX_REQUEST_BEAM") WHISPER_MAX_UPLOAD_MB=$(nospaces "$WHISPER_MAX_UPLOAD_MB") WHISPER_MAX_UPLOAD_MB=$(noquotes "$WHISPER_MAX_UPLOAD_MB") +WHISPER_IDLE_UNLOAD_SECONDS=$(nospaces "$WHISPER_IDLE_UNLOAD_SECONDS") +WHISPER_IDLE_UNLOAD_SECONDS=$(noquotes "$WHISPER_IDLE_UNLOAD_SECONDS") WHISPER_LOCAL_ONLY=$(nospaces "$WHISPER_LOCAL_ONLY") WHISPER_LOCAL_ONLY=$(noquotes "$WHISPER_LOCAL_ONLY") WHISPER_WORD_TIMESTAMPS=$(nospaces "$WHISPER_WORD_TIMESTAMPS") @@ -92,6 +94,7 @@ _USER_COMPUTE_TYPE="$WHISPER_COMPUTE_TYPE" [ -z "$WHISPER_BEAM" ] && WHISPER_BEAM=5 [ -z "$WHISPER_MAX_REQUEST_BEAM" ] && WHISPER_MAX_REQUEST_BEAM=10 [ -z "$WHISPER_MAX_UPLOAD_MB" ] && WHISPER_MAX_UPLOAD_MB=1024 +[ -z "$WHISPER_IDLE_UNLOAD_SECONDS" ] && WHISPER_IDLE_UNLOAD_SECONDS=0 [ -z "$WHISPER_DIARIZE_NUM_SPEAKERS" ] && WHISPER_DIARIZE_NUM_SPEAKERS=-1 [ -z "$WHISPER_DIARIZE_THRESHOLD" ] && WHISPER_DIARIZE_THRESHOLD=0.5 @@ -157,6 +160,9 @@ fi if ! printf '%s' "$WHISPER_MAX_UPLOAD_MB" | grep -Eq '^(0|[1-9][0-9]*)$'; then exiterr "WHISPER_MAX_UPLOAD_MB must be 0 (unlimited) or a positive integer." fi +if ! printf '%s' "$WHISPER_IDLE_UNLOAD_SECONDS" | grep -Eq '^(0|[1-9][0-9]*)$'; then + exiterr "WHISPER_IDLE_UNLOAD_SECONDS must be 0 (disabled) or a positive integer." +fi # Validate diarization speaker counts (-1 for auto-detect, or a positive integer) if ! printf '%s' "$WHISPER_DIARIZE_NUM_SPEAKERS" | grep -Eq '^(-1|[1-9][0-9]*)$'; then @@ -302,6 +308,7 @@ export WHISPER_LOG_LEVEL export WHISPER_BEAM export WHISPER_MAX_REQUEST_BEAM export WHISPER_MAX_UPLOAD_MB +export WHISPER_IDLE_UNLOAD_SECONDS export WHISPER_LOCAL_ONLY export WHISPER_WORD_TIMESTAMPS export WHISPER_DIARIZATION From 760d9d9a1946b754729e4eeea54afb352d867e4d Mon Sep 17 00:00:00 2001 From: Valentine Ubani Mayaki Date: Fri, 17 Jul 2026 17:50:21 +0200 Subject: [PATCH 2/2] Add idle unload tests --- .../workflows/tests/test_api_validation.py | 144 ++++++++++++++++++ 1 file changed, 144 insertions(+) diff --git a/.github/workflows/tests/test_api_validation.py b/.github/workflows/tests/test_api_validation.py index 01389a2..97e49bc 100644 --- a/.github/workflows/tests/test_api_validation.py +++ b/.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()