fix(stt): close idle-unload races — strong model ref, single long-lived watcher

Review pass on the idle-unload feature found two material concurrency
bugs; both fixed here with a regression guard:

1. Unload-vs-use null deref (HIGH): _transcribe_local re-read the
   module global _local_model at the transcribe call site. An idle
   unload firing between the model load and transcribe() evaluated
   None.transcribe → AttributeError → user-visible 'Local
   transcription failed'. The window was real: the idle timer was only
   touched AFTER a successful transcription, so a voice note arriving
   exactly as the timeout expired raced the watcher directly.
   Fix: bind a strong local reference under the model lock and use it
   for the whole transcription (the watcher can null the global at any
   time; this in-flight call keeps its instance — the generator holds
   self, so no use-after-free). Also touch the idle timer at the START
   of transcription so a long in-flight transcribe can't be counted as
   idle time. The CUDA-fallback retry path gets the same treatment
   (locked global write, local ref use).

2. Watcher replacement race + response-path join (MEDIUM/HIGH): the
   old design stopped and re-started the watcher after EVERY
   transcription with an unlocked set/join(5)/clear/start sequence on
   shared globals. Two concurrent voice messages could interleave to
   leave TWO live watchers (one with a stale, shorter timeout — a
   raised unload_after_idle_seconds could still unload on the old
   value), and the join(timeout=5) sat on the user-visible response
   path (a watcher blocked on _local_model_lock during a concurrent
   multi-second model load stalls the reply up to 5s).
   Fix: single long-lived watcher under a management lock — started
   only when none is alive (per-transcription cost: one lock + one
   is_alive check), re-reads the configured timeout from config every
   cycle (config edits now apply within one 30s interval, without
   waiting for the next voice message — previously undocumented), and
   stands down without unloading when the timeout is set to 0
   mid-idle.

Tests: 17 now — idempotent start (same thread, no churn), config
re-read + stand-down-when-disabled, and the race guard
(unload firing mid-transcription must not fail the in-flight call).
The race guard is mutation-verified: reverting the fix (re-reading
the global at the call site) makes it fail with the exact NoneType
error; the fixed code passes.
This commit is contained in:
kshitij 2026-08-07 18:33:55 +05:30 committed by kshitij
parent 7b006ea6e8
commit 72c63aa586
2 changed files with 173 additions and 47 deletions

View File

@ -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"

View File

@ -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(