Files
VoiceStudio/tests/test_model_load_timeout.py
c965a7fbdd fix(backend): self-healing GPU pool so a reset can't strand requests (#589/#599) (#601)
`_reset_gpu_pool()` fires on a model-load timeout to recover a wedged worker —
it shut the ThreadPoolExecutor down and rebuilt a fresh one on next access. But
several request handlers (generation, dub_generate, dub_core, dub_translate,
openai_compat) did a *module-level* `from services.model_manager import
_gpu_pool`, capturing the executor object at import time. After a reset those
references pointed at the dead pool, so the next generate/dub/transcribe/
translate raised `RuntimeError: cannot schedule new futures after shutdown` —
surfacing as a 500 or "Can't reach the local backend" (#589 #599).

Make `_gpu_pool` a single long-lived `_ResilientGpuPool` wrapper (a
concurrent.futures.Executor) whose *inner* ThreadPoolExecutor is swapped:
- every submit() resolves the live pool, and a submit that races a shutdown
  rebuilds once and retries, so a stale captured reference self-heals;
- `_reset_gpu_pool()` now drops only the inner pool (fresh worker on retry)
  while preserving the wrapper identity every importer holds;
- pool sizing stays lazy, so we still probe the device after torch's lazy
  import (the reason for the original __getattr__ indirection).

Fixes the whole class — all importers share one wrapper, module-level or
function-level. Regression tests cover stale-ref-survives-reset, identity
stability, submit-after-inner-shutdown self-heal, and asyncio.run_in_executor
compatibility; updated the load-timeout test to the new reset semantics.

Co-authored-by: mergetest <test@local>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-22 04:42:16 +05:30

67 lines
2.5 KiB
Python

"""Regression: model load/download must never hang forever (Windows demo-voice
"runs indefinitely, no audio, no error" report).
`get_model()` wraps the blocking load in an `asyncio.wait_for` deadline; on
timeout it drops the poisoned GPU pool and raises a clear RuntimeError instead
of leaving `/generate` (and the UI spinner) wedged forever.
"""
from __future__ import annotations
import asyncio
import sys
import threading
import pytest
@pytest.fixture
def model_manager(monkeypatch):
for mod_name in ("core.config", "services.model_manager"):
if getattr(sys.modules.get(mod_name), "__file__", None) is None:
sys.modules.pop(mod_name, None)
import services.model_manager as mm
monkeypatch.setattr(mm, "model", None)
return mm
def test_model_load_timeout_respects_env(model_manager, monkeypatch):
mm = model_manager
monkeypatch.delenv("OMNIVOICE_MODEL_LOAD_TIMEOUT", raising=False)
assert mm._model_load_timeout() == 1200.0
monkeypatch.setenv("OMNIVOICE_MODEL_LOAD_TIMEOUT", "5000")
assert mm._model_load_timeout() == 5000.0
monkeypatch.setenv("OMNIVOICE_MODEL_LOAD_TIMEOUT", "not-a-number")
assert mm._model_load_timeout() == 1200.0
monkeypatch.setenv("OMNIVOICE_MODEL_LOAD_TIMEOUT", "1") # below the safety floor
assert mm._model_load_timeout() == 30.0
def test_get_model_times_out_and_resets_pool(model_manager, monkeypatch):
mm = model_manager
# Isolated lock so we never reuse one bound to a previous test's loop.
monkeypatch.setattr(mm, "_model_lock", asyncio.Lock())
monkeypatch.setattr(mm, "_model_load_timeout", lambda: 0.3)
release = threading.Event()
def _hang(): # simulates a wedged download that never returns
release.wait(2.0)
return object()
monkeypatch.setattr(mm, "_load_model_sync", _hang)
assert mm._get_gpu_pool() is not None # a pool exists before the timeout
try:
with pytest.raises(RuntimeError, match="timed out"):
asyncio.run(mm.get_model())
assert mm.model is None # no half-loaded model
# The resilient wrapper survives (importers hold its identity); only its
# inner worker pool is dropped, so a retry builds a fresh worker.
assert mm._gpu_pool_singleton is not None
assert mm._gpu_pool_singleton._pool is None
assert not mm._model_lock.locked() # lock released for a retry
finally:
release.set() # let the orphaned worker exit immediately