c965a7fbdd
`_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>
77 lines
2.9 KiB
Python
77 lines
2.9 KiB
Python
"""Regression: the GPU pool must survive a reset without breaking long-lived
|
|
references (#589 #599 — "cannot schedule new futures after shutdown").
|
|
|
|
`_reset_gpu_pool()` fires on a model-load timeout. Before the fix, consumers
|
|
that did a *module-level* `from services.model_manager import _gpu_pool`
|
|
(generation, dub_generate, dub_core, dub_translate, openai_compat) captured the
|
|
ThreadPoolExecutor object once — so after a reset they kept submitting to the
|
|
shut-down pool and every generate/dub 500'd with "cannot schedule new futures
|
|
after shutdown". The resilient wrapper keeps a stable identity and rebuilds its
|
|
inner pool on demand, so those references self-heal.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
|
|
@pytest.fixture
|
|
def mm(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
|
|
# Keep the pool tiny + device-probe-free regardless of the host.
|
|
monkeypatch.setattr(_mm, "_pick_gpu_workers", lambda: 1)
|
|
# Start from a clean singleton so tests don't share a wrapper.
|
|
monkeypatch.setattr(_mm, "_gpu_pool_singleton", None)
|
|
return _mm
|
|
|
|
|
|
def test_stale_reference_survives_reset(mm):
|
|
# Mimic a module-level `from services.model_manager import _gpu_pool`.
|
|
captured = mm._gpu_pool # triggers __getattr__ → wrapper
|
|
assert captured.submit(lambda: 7).result(timeout=5) == 7
|
|
|
|
mm._reset_gpu_pool() # the load-timeout recovery path
|
|
|
|
# The SAME captured reference must still work — no "cannot schedule new
|
|
# futures after shutdown".
|
|
assert captured.submit(lambda: 11).result(timeout=5) == 11
|
|
|
|
|
|
def test_reset_keeps_wrapper_identity_drops_inner_pool(mm):
|
|
pool = mm._get_gpu_pool()
|
|
pool.submit(lambda: None).result(timeout=5) # force-build the inner pool
|
|
assert pool._pool is not None
|
|
|
|
mm._reset_gpu_pool()
|
|
assert mm._get_gpu_pool() is pool # stable identity
|
|
assert pool._pool is None # inner worker pool dropped
|
|
|
|
pool.submit(lambda: None).result(timeout=5) # rebuilds transparently
|
|
assert pool._pool is not None
|
|
|
|
|
|
def test_submit_after_inner_shutdown_self_heals(mm):
|
|
pool = mm._get_gpu_pool()
|
|
# Simulate the exact failure: a stale inner pool that's been shut down.
|
|
pool.submit(lambda: None).result(timeout=5)
|
|
pool._pool.shutdown(wait=True)
|
|
# Without the retry this raises RuntimeError("cannot schedule new futures
|
|
# after shutdown"); the wrapper rebuilds and succeeds.
|
|
assert pool.submit(lambda: 42).result(timeout=5) == 42
|
|
|
|
|
|
def test_wrapper_usable_with_asyncio_run_in_executor(mm):
|
|
import asyncio
|
|
|
|
pool = mm._get_gpu_pool()
|
|
|
|
async def _go():
|
|
loop = asyncio.get_running_loop()
|
|
return await loop.run_in_executor(pool, lambda: 5)
|
|
|
|
assert asyncio.run(_go()) == 5
|