diff --git a/tests/tools/test_stt_idle_unload.py b/tests/tools/test_stt_idle_unload.py index 70c61e30387fa..d98110ddb9b3a 100644 --- a/tests/tools/test_stt_idle_unload.py +++ b/tests/tools/test_stt_idle_unload.py @@ -16,9 +16,16 @@ Contract under test: 3. ``_touch_transcription_time`` resets the idle timer. 4. ``_unload_local_model`` is safe to call when the model is already None. 5. Config resolution handles garbage values gracefully. +6. Watcher start is idempotent (single long-lived thread, no churn on the + transcription response path) and the loop re-reads config each cycle. +7. RACE GUARD: an unload firing mid-transcription must not fail the + in-flight transcription — ``_transcribe_local`` binds a strong local + reference to the model instance instead of re-reading the global. """ +import struct import time +import wave from unittest.mock import MagicMock, patch import pytest @@ -127,10 +134,13 @@ class TestIdleUnloadWatcher: """The watcher unloads the model after the configured idle period.""" mock_model = MagicMock(name="whisper_model") # Use a very short timeout and patch the check interval to 0.01s - # so the test runs in < 1 second. + # so the test runs in < 1 second. The watcher re-reads config each + # cycle, so patch _load_stt_config to keep the timeout active. with patch.object(tt, "_local_model", mock_model), \ patch.object(tt, "_local_model_name", "base"), \ patch.object(tt, "_IDLE_UNLOAD_CHECK_INTERVAL", 0.01), \ + patch.object(tt, "_load_stt_config", + return_value={"local": {"unload_after_idle_seconds": 1}}), \ patch.object(tt, "_last_transcription_time", time.monotonic() - 100): _start_idle_unload_watcher(timeout_seconds=1) # Wait for the watcher to fire @@ -152,10 +162,12 @@ class TestIdleUnloadWatcher: tt._local_model_name = "base" tt._IDLE_UNLOAD_CHECK_INTERVAL = 0.01 tt._last_transcription_time = time.monotonic() - _start_idle_unload_watcher(timeout_seconds=100) - # Give it a few check cycles — model must survive - time.sleep(0.1) - assert tt._local_model is not None + with patch.object(tt, "_load_stt_config", + return_value={"local": {"unload_after_idle_seconds": 100}}): + _start_idle_unload_watcher(timeout_seconds=100) + # Give it a few check cycles — model must survive + time.sleep(0.1) + assert tt._local_model is not None finally: tt._local_model = original_model tt._local_model_name = original_name @@ -171,19 +183,101 @@ class TestIdleUnloadWatcher: time.sleep(0.05) # No crash, no hang — watcher detected _local_model is None and exited - def test_watcher_stopped_on_new_start(self): - """Starting a new watcher stops the previous one.""" + def test_start_is_idempotent_while_watcher_alive(self): + """A second start while a watcher is alive is a no-op (single + long-lived watcher — no stop/join/restart churn on the hot path).""" mock_model = MagicMock(name="whisper_model") with patch.object(tt, "_local_model", mock_model), \ patch.object(tt, "_local_model_name", "base"), \ - patch.object(tt, "_IDLE_UNLOAD_CHECK_INTERVAL", 0.01), \ + patch.object(tt, "_IDLE_UNLOAD_CHECK_INTERVAL", 0.5), \ + patch.object(tt, "_load_stt_config", + return_value={"local": {"unload_after_idle_seconds": 100}}), \ patch.object(tt, "_last_transcription_time", time.monotonic()): _start_idle_unload_watcher(timeout_seconds=100) first_thread = tt._idle_unload_thread assert first_thread is not None and first_thread.is_alive() _start_idle_unload_watcher(timeout_seconds=200) - second_thread = tt._idle_unload_thread - assert second_thread is not first_thread - # First thread should have been stopped - time.sleep(0.05) - assert not first_thread.is_alive() + assert tt._idle_unload_thread is first_thread # same thread, no churn + tt._idle_unload_stop.set() # clean up + first_thread.join(timeout=2) + + def test_watcher_rereads_config_and_stands_down_when_disabled(self): + """Setting unload_after_idle_seconds to 0 mid-idle stops the watcher + without unloading — config edits apply within one check interval.""" + mock_model = MagicMock(name="whisper_model") + original_model = tt._local_model + original_name = tt._local_model_name + original_interval = tt._IDLE_UNLOAD_CHECK_INTERVAL + original_ts = tt._last_transcription_time + try: + tt._local_model = mock_model + tt._local_model_name = "base" + tt._IDLE_UNLOAD_CHECK_INTERVAL = 0.01 + tt._last_transcription_time = time.monotonic() - 1000 # long idle + with patch.object(tt, "_load_stt_config", + return_value={"local": {"unload_after_idle_seconds": 0}}): + _start_idle_unload_watcher(timeout_seconds=1) + thread = tt._idle_unload_thread + assert thread is not None + thread.join(timeout=2) + assert not thread.is_alive() # stood down... + assert tt._local_model is not None # ...without unloading + finally: + tt._local_model = original_model + tt._local_model_name = original_name + tt._IDLE_UNLOAD_CHECK_INTERVAL = original_interval + tt._last_transcription_time = original_ts + + +# ============================================================================ +# Race guard: unload mid-transcription must not break the in-flight call +# ============================================================================ + + +class TestUnloadDuringTranscriptionRace: + def test_transcribe_survives_concurrent_unload(self, tmp_path): + """_transcribe_local binds a strong local ref under the lock; an + idle unload nulling the module global mid-transcription must not + produce 'NoneType has no attribute transcribe'.""" + wav_path = tmp_path / "a.wav" + n = 16000 + with wave.open(str(wav_path), "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(16000) + wf.writeframes(struct.pack(f"<{n}h", *([0] * n))) + + seg = MagicMock() + seg.text = " hello" + seg.no_speech_prob = 0.0 + seg.avg_logprob = -0.2 + info = MagicMock() + info.language = "en" + info.duration = 1.0 + + mock_model = MagicMock(name="whisper_model") + mock_model.transcribe.return_value = (iter([seg]), info) + + real_kwargs_builder = tt.build_local_transcribe_kwargs + + def kwargs_then_unload(*args, **kwargs): + # Simulate the watcher firing in the window BETWEEN the model + # load and the transcribe call (build_local_transcribe_kwargs + # runs exactly there): the module global goes away while this + # transcription is still in flight. + out = real_kwargs_builder(*args, **kwargs) + _unload_local_model() + assert tt._local_model is None + return out + + with patch.object(tt, "_HAS_FASTER_WHISPER", True), \ + patch.object(tt, "_load_stt_config", return_value={"local": {}}), \ + patch.object(tt, "build_local_transcribe_kwargs", + side_effect=kwargs_then_unload), \ + patch.object(tt, "_local_model", mock_model), \ + patch.object(tt, "_local_model_name", "base"): + from tools.transcription_tools import _transcribe_local + result = _transcribe_local(str(wav_path), "base") + + assert result["success"] is True, result.get("error") + assert result["transcript"] == "hello" diff --git a/tools/transcription_tools.py b/tools/transcription_tools.py index f60fb1d7f1479..e31127384973e 100644 --- a/tools/transcription_tools.py +++ b/tools/transcription_tools.py @@ -143,12 +143,15 @@ _local_model_lock = threading.Lock() # The model singleton above is loaded once and never released — hundreds of MB # of RAM/VRAM sit idle between voice messages. On long-running gateway # processes (especially with local LLMs competing for the same GPU) this is -# wasteful. A lightweight daemon thread checks _last_transcription_time and -# unloads the model after a configurable idle period. The next voice message -# reloads it transparently. +# wasteful. A single long-lived daemon thread checks _last_transcription_time +# and unloads the model after a configurable idle period, then exits. The next +# voice message reloads the model and restarts the watcher transparently. _last_transcription_time: float = 0.0 _idle_unload_thread: Optional[threading.Thread] = None _idle_unload_stop = threading.Event() +# Serializes watcher start checks so two concurrent transcriptions can't +# both observe "no watcher alive" and spawn duplicates. +_idle_unload_mgmt_lock = threading.Lock() _IDLE_UNLOAD_CHECK_INTERVAL = 30 # seconds between idle checks @@ -1513,38 +1516,56 @@ def _unload_local_model() -> None: def _start_idle_unload_watcher(timeout_seconds: int) -> None: - """Start (or replace) the background idle-unload thread. + """Ensure the idle-unload watcher thread is running. - The thread checks every ``_IDLE_UNLOAD_CHECK_INTERVAL`` seconds whether - ``_last_transcription_time`` is older than ``timeout_seconds``. If so, it - unloads the model and exits. The next transcription restarts it. + A single long-lived watcher: started only when none is alive, so the + per-transcription cost is one lock + one ``is_alive()`` check — no + stop/join/restart churn on the response path. The loop re-reads the + configured timeout from config every cycle, so changing + ``stt.local.unload_after_idle_seconds`` takes effect within one check + interval without a restart. After unloading (or when the timeout is set + to 0/never, or the model is already gone) the thread exits; the next + transcription restarts it. + + ``timeout_seconds`` seeds the first cycle so a just-written config is + honored even if a concurrent config read would race. """ global _idle_unload_thread - # Stop any existing watcher — a new timeout may have been configured - _idle_unload_stop.set() - if _idle_unload_thread is not None and _idle_unload_thread.is_alive(): - _idle_unload_thread.join(timeout=5) - _idle_unload_stop.clear() + with _idle_unload_mgmt_lock: + if _idle_unload_thread is not None and _idle_unload_thread.is_alive(): + return - def _watch(): - while not _idle_unload_stop.is_set(): - if _idle_unload_stop.wait(_IDLE_UNLOAD_CHECK_INTERVAL): - break - if _local_model is None: - break - idle_for = time.monotonic() - _last_transcription_time - if idle_for >= timeout_seconds: - _unload_local_model() - break + def _watch(initial_timeout=timeout_seconds): + timeout = initial_timeout + while not _idle_unload_stop.is_set(): + if _idle_unload_stop.wait(_IDLE_UNLOAD_CHECK_INTERVAL): + break + if _local_model is None: + break + # Re-read the timeout each cycle: config edits apply without + # waiting for the next voice message. + try: + timeout = _get_idle_unload_seconds( + _load_stt_config().get("local") or {} + ) + except Exception: # noqa: BLE001 - keep the seed value + timeout = initial_timeout + if timeout <= 0: + break # unload disabled mid-flight — stand down + idle_for = time.monotonic() - _last_transcription_time + if idle_for >= timeout: + _unload_local_model() + break - _idle_unload_thread = threading.Thread( - target=_watch, name="hermes-stt-idle-unload", daemon=True - ) - _idle_unload_thread.start() + _idle_unload_stop.clear() + _idle_unload_thread = threading.Thread( + target=_watch, name="hermes-stt-idle-unload", daemon=True + ) + _idle_unload_thread.start() def _touch_transcription_time() -> None: - """Record that a transcription just completed (resets the idle timer).""" + """Record transcription activity (resets the idle timer).""" global _last_transcription_time _last_transcription_time = time.monotonic() @@ -1729,10 +1750,18 @@ def _transcribe_local(file_path: str, model_name: str) -> Dict[str, Any]: try: local_cfg = _load_stt_config().get("local") or {} + # Reset the idle timer BEFORE loading/transcribing so the idle-unload + # watcher can't count a long in-flight transcription as idle time and + # unload mid-use. + _touch_transcription_time() # Lazy-load the model (downloads on first use, ~150 MB for 'base'). # Double-checked lock: concurrent voice messages must not both # download/load the model (#24767). - if _local_model is None or _local_model_name != model_name: + # ``model`` is a strong local reference bound under the lock: the idle + # watcher may null the module global at any time, but this + # transcription keeps using the instance it grabbed. + model = _local_model + if model is None or _local_model_name != model_name: with _local_model_lock: if _local_model is None or _local_model_name != model_name: logger.info("Loading faster-whisper model '%s' (first load downloads the model)...", model_name) @@ -1747,7 +1776,10 @@ def _transcribe_local(file_path: str, model_name: str) -> Dict[str, Any]: compute_type=local_cfg.get("compute_type", "auto"), ) _local_model_name = model_name + model = _local_model + if model is None: # defensive: load failed without raising + return {"success": False, "transcript": "", "error": "Local whisper model failed to load"} # Shared hardened kwargs: VAD filter (default on), no cross-window # conditioning, language/initial_prompt resolution — one owner for # every local faster-whisper call site. @@ -1756,7 +1788,7 @@ def _transcribe_local(file_path: str, model_name: str) -> Dict[str, Any]: transcribe_kwargs = build_local_transcribe_kwargs(stt_config) try: - segments, info = _local_model.transcribe(file_path, **transcribe_kwargs) + segments, info = model.transcribe(file_path, **transcribe_kwargs) transcript = _join_confident_segments(segments, local_config) except Exception as exc: # CUDA runtime libs sometimes only fail at dlopen-on-first-use, @@ -1771,12 +1803,12 @@ def _transcribe_local(file_path: str, model_name: str) -> Dict[str, Any]: "evicting cached model and retrying on CPU (int8).", exc, ) - _local_model = None - _local_model_name = None from faster_whisper import WhisperModel - _local_model = WhisperModel(model_name, device="cpu", compute_type="int8") - _local_model_name = model_name - segments, info = _local_model.transcribe(file_path, **transcribe_kwargs) + model = WhisperModel(model_name, device="cpu", compute_type="int8") + with _local_model_lock: + _local_model = model + _local_model_name = model_name + segments, info = model.transcribe(file_path, **transcribe_kwargs) transcript = _join_confident_segments(segments, local_config) logger.info(