import asyncio
import logging
import math
import os
import queue
import re
import sys
import threading
import time
from concurrent.futures import Executor, Future, ThreadPoolExecutor
from utils.containment import contain_system_exit
_torch = None
_OmniVoice = None
def _lazy_torch():
global _torch
if _torch is None:
import torch as _t
_torch = _t
return _torch
def _missing_module_is_omnivoice(exc: ModuleNotFoundError) -> bool:
"""True when *exc* says the ``omnivoice`` package itself is not importable.
``ModuleNotFoundError`` is raised for two very different situations along
this import, and only one of them is fixable by putting the source tree on
``sys.path`` (#1415):
* ``omnivoice`` (or a submodule of it) is genuinely absent — a missing or
broken editable install, which the #564 fallback repairs; ``exc.name``
names the omnivoice package.
* something ``omnivoice`` imports is absent or broken — a torch /
torchaudio / torchvision mismatch, or transformers' lazy module refusing
an attribute whose backing import failed
("Could not import module 'AutoFeatureExtractor'", which carries no
``name`` at all). Nothing about ``sys.path`` is wrong here.
Treating the second as the first re-imported from the same broken
environment, failed identically, and logged that the editable install was
missing — a confident diagnosis of the wrong component.
``exc.name`` is the authority, and its absence is decisive rather than
unknown: the stdlib always sets it, so a ModuleNotFoundError without one
was raised by hand — which is exactly what transformers' lazy module does.
"""
name = getattr(exc, "name", None)
if not name:
return False
return name == "omnivoice" or name.startswith("omnivoice.")
def _lazy_omnivoice():
global _OmniVoice
if _OmniVoice is None:
try:
from omnivoice.models.omnivoice import OmniVoice as _OV
except ModuleNotFoundError as exc:
if not _missing_module_is_omnivoice(exc):
raise
from core.omnivoice_path import ensure_omnivoice_importable
_backend_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
ensure_omnivoice_importable(_backend_dir, logger)
from omnivoice.models.omnivoice import OmniVoice as _OV
_OmniVoice = _OV
return _OmniVoice
from core.config import IDLE_TIMEOUT_SECONDS, CPU_POOL_WORKERS
logger = logging.getLogger("omnivoice.model")
_GPU_VRAM_PER_JOB_GB = 5.0
_GPU_WORKER_CAP = 4
class WorkerStopIteration(RuntimeError):
"""A pool worker raised a bare ``StopIteration``.
asyncio refuses to put ``StopIteration`` into a Future — ``_copy_future_
state`` raises ``TypeError: StopIteration interacts badly with generators
and cannot be raised into a Future`` *inside the event loop's callback*, so
the ``run_in_executor`` future is never completed and the awaiting caller
waits **forever**. Not a theoretical edge: verified on the bundled CPython
3.11, and the failure has no error, no event and no timeout — a render just
stops, which is indistinguishable to the user from a wedged app.
Generator-driven engines reach it on ordinary bad input: VoxCPM's
``next_and_close`` is a bare ``next(gen)``, so a generator that ends without
yielding (text the model normalises away to nothing, for instance) raises
exactly this out of ``backend.generate`` (#1321 class).
Translating it to a RuntimeError at the pool boundary — the one place every
dispatch funnels through — turns a silent hang into a normal failure that
the existing per-chapter / per-job error handling reports. Subclasses
RuntimeError so every `except Exception` site upstream keeps working.
"""
def _guard_stopiteration(fn):
"""Wrap `fn` so a bare StopIteration can never escape into a Future."""
def _guarded(*a, **kw):
try:
return fn(*a, **kw)
except StopIteration as e:
raise WorkerStopIteration(
"the engine stopped without producing a result (StopIteration) — "
"its generator ended before yielding anything, which usually means "
"it could not handle this input"
) from e
return _guarded
class _GuardedCpuPool(ThreadPoolExecutor):
"""CPU pool with the same StopIteration guard as the GPU pool."""
def submit(self, fn, /, *args, **kwargs):
return super().submit(_guard_stopiteration(fn), *args, **kwargs)
_gpu_pool_singleton: "_ResilientGpuPool | None" = None
_cpu_pool = _GuardedCpuPool(max_workers=CPU_POOL_WORKERS)
def _workers_for_free_vram(free_gb: float) -> int:
"""GPU worker count for a given free-VRAM figure: free // per-job budget,
floored at 1 and capped at _GPU_WORKER_CAP. Pure so the sizing policy is
unit-tested without a GPU (the #567 crash hinged on this returning >1 on
8 GB cards)."""
return max(1, min(_GPU_WORKER_CAP, int(free_gb // _GPU_VRAM_PER_JOB_GB)))
def _pick_gpu_workers() -> int:
"""Pick a sensible GPU worker count from the runtime environment.
Resolution order:
1. OMNIVOICE_GPU_WORKERS env var (explicit user override, clamped 1..16).
2. CUDA / ROCm: free VRAM // per-job budget, capped at 4.
3. MPS / CPU / unknown: 1.
Designed to fail safe — any exception → 1 worker, never propagated.
"""
override = os.environ.get("OMNIVOICE_GPU_WORKERS")
if override:
try:
n = int(override)
return max(1, min(16, n))
except ValueError:
logger.warning("OMNIVOICE_GPU_WORKERS=%r is not an integer; ignoring", override)
try:
torch = _lazy_torch()
if hasattr(torch, "cuda") and torch.cuda.is_available():
free_bytes, _total = torch.cuda.mem_get_info()
free_gb = free_bytes / (1024 ** 3)
workers = _workers_for_free_vram(free_gb)
logger.info(
"GPU pool sized to %d worker(s) — %.1f GB free / %.1f GB per job (cap %d)",
workers, free_gb, _GPU_VRAM_PER_JOB_GB, _GPU_WORKER_CAP,
)
return workers
if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
logger.info("GPU pool: MPS detected, using 1 worker (shared system memory)")
return 1
except Exception as e:
logger.warning("GPU worker probe failed (%s); defaulting to 1", e)
return 1
_GPU_POOL_THREAD_PREFIX = "gpu-pool"
def _build_gpu_pool() -> ThreadPoolExecutor:
workers = _pick_gpu_workers()
return ThreadPoolExecutor(
max_workers=workers, thread_name_prefix=_GPU_POOL_THREAD_PREFIX)
def running_on_gpu_pool() -> bool:
"""True iff the calling thread is a gpu-pool worker (already holds a slot).
Routes that dispatch backend work via run_on_gpu_pool_guarded are already on
a pool worker; re-acquiring a slot there would self-deadlock on a 1-worker
pool (MPS). Used by SubprocessBackend.generate()'s on-pool skip and by
_heal_tts_placement.
"""
return threading.current_thread().name.startswith(_GPU_POOL_THREAD_PREFIX)
class _ResilientGpuPool(Executor):
"""A stable, self-healing wrapper around the GPU `ThreadPoolExecutor`.
The crash this fixes (#589 #599): `_reset_gpu_pool()` shuts the pool down on
a model-load timeout, but consumers that captured the executor *object* at
import time (`from services.model_manager import _gpu_pool` at module level —
generation, dub_generate, dub_core, dub_translate, openai_compat) kept
submitting to the dead pool and got `RuntimeError: cannot schedule new
futures after shutdown` on the next generate/dub/translate.
Making `_gpu_pool` a single long-lived wrapper whose *inner* pool is swapped
means those references never go stale: every `submit()` resolves the live
pool, and a submit that races a shutdown rebuilds once and retries. Building
the inner pool stays lazy so we still size workers after torch's device
probe (the reason for the original `__getattr__` indirection).
"""
def __init__(self):
self._pool: "ThreadPoolExecutor | None" = None
self._lock = threading.Lock()
self._stats_lock = threading.Lock()
self._queued = 0
self._running = 0
self._avg_job_s = 0.0
def _live_pool(self) -> ThreadPoolExecutor:
pool = self._pool
if pool is None:
with self._lock:
if self._pool is None:
self._pool = _build_gpu_pool()
pool = self._pool
return pool
def _submit_live(self, fn, /, *args, **kwargs):
try:
return self._live_pool().submit(fn, *args, **kwargs)
except RuntimeError as e:
if "shutdown" not in str(e).lower():
raise
with self._lock:
self._pool = _build_gpu_pool()
pool = self._pool
return pool.submit(fn, *args, **kwargs)
def submit(self, fn, /, *args, **kwargs):
token = {"counted": False}
def _tracked(*a, **kw):
with self._stats_lock:
token["counted"] = True
self._queued -= 1
self._running += 1
t0 = time.monotonic()
try:
return _guard_stopiteration(fn)(*a, **kw)
finally:
elapsed = time.monotonic() - t0
with self._stats_lock:
self._running -= 1
self._avg_job_s = (
elapsed if self._avg_job_s <= 0
else 0.7 * self._avg_job_s + 0.3 * elapsed
)
with self._stats_lock:
self._queued += 1
try:
fut = self._submit_live(_tracked, *args, **kwargs)
except BaseException:
with self._stats_lock:
if not token["counted"]:
token["counted"] = True
self._queued -= 1
raise
def _drain(_f, token=token):
with self._stats_lock:
if not token["counted"]:
token["counted"] = True
self._queued -= 1
fut.add_done_callback(_drain)
return fut
def stats(self) -> dict:
"""Live queue depth / worker occupancy — the input to admission control."""
with self._stats_lock:
queued, running, avg = self._queued, self._running, self._avg_job_s
pool = self._pool
workers = getattr(pool, "_max_workers", None) or 1
return {"queued": queued, "running": running,
"workers": workers, "avg_job_s": avg}
def reset(self) -> None:
"""Abandon the current worker pool; the next submit builds a fresh one.
Deliberately **not** ``cancel_futures=True`` (#1190/#1202): that killed
innocent peers — a queued job belonging to a *different* request was
cancelled because *this* request timed out, and surfaced to that caller
as a bare ``CancelledError``. ``shutdown(wait=False)`` only refuses NEW
submissions; work already in the old pool's queue still drains on the
old pool's workers, so peers complete normally while new work goes to
the fresh pool.
Honesty about what this reclaims: **nothing**. Python cannot kill the
thread wedged in the timed-out job — it keeps running (and keeps its
VRAM) until it finishes on its own. Dropping the pool only stops NEW
work from queueing behind it; it does not restore the device. That is
why the timeout guidance no longer claims capacity was restored.
"""
with self._lock:
pool, self._pool = self._pool, None
if pool is not None:
try:
pool.shutdown(wait=False)
except Exception:
pass
def shutdown(self, wait=True, *, cancel_futures=False):
with self._lock:
pool, self._pool = self._pool, None
if pool is not None:
pool.shutdown(wait=wait, cancel_futures=cancel_futures)
def _get_gpu_pool() -> "_ResilientGpuPool":
"""Internal accessor for the GPU pool singleton. Same object as the
module-level `_gpu_pool` attribute, but resolvable from inside this module
(Python's module `__getattr__` only fires for lookups from *outside*).
"""
global _gpu_pool_singleton
if _gpu_pool_singleton is None:
_gpu_pool_singleton = _ResilientGpuPool()
return _gpu_pool_singleton
def __getattr__(name: str):
"""Lazy module attribute — initialises `_gpu_pool` on first access so we
can probe the device after torch finishes its lazy import. Without this
we'd be forced to commit to max_workers=1 at module import time, before
knowing whether CUDA is even available.
"""
if name == "_gpu_pool":
return _get_gpu_pool()
raise AttributeError(f"module 'services.model_manager' has no attribute {name!r}")
_GENERATE_TIMEOUT_EXPLICIT = "OMNIVOICE_GENERATE_TIMEOUT_S" in os.environ
GPU_JOB_TIMEOUT_S = float(os.environ.get("OMNIVOICE_GENERATE_TIMEOUT_S", "300.0"))
_CONFIGURED_GPU_JOB_TIMEOUT_S = GPU_JOB_TIMEOUT_S
_CPU_GENERATE_TIMEOUT_EXPLICIT = "OMNIVOICE_CPU_GENERATE_TIMEOUT_S" in os.environ
CPU_JOB_TIMEOUT_S = float(os.environ.get("OMNIVOICE_CPU_GENERATE_TIMEOUT_S", "600.0"))
_CONFIGURED_CPU_JOB_TIMEOUT_S = CPU_JOB_TIMEOUT_S
GPU_QUEUE_TIMEOUT_S = float(os.environ.get("OMNIVOICE_GPU_QUEUE_TIMEOUT_S", "1800.0"))
MODEL_LOAD_HEARTBEAT_GRACE_S = float(
os.environ.get("OMNIVOICE_MODEL_LOAD_HEARTBEAT_GRACE_S", "30.0"))
MODEL_LOAD_EXTRA_TIMEOUT_S = float(
os.environ.get("OMNIVOICE_MODEL_LOAD_TIMEOUT_S", "1800.0"))
GENERATE_PROGRESS_GRACE_S = float(
os.environ.get("OMNIVOICE_GENERATE_PROGRESS_GRACE_S", "300.0"))
_MODEL_LOAD_ACTIVITY: dict = {}
def report_model_load_activity() -> None:
"""Record that the CURRENT THREAD's job is making model-load progress.
Called by engine code that can prove liveness — e.g. SubprocessBackend
each time a sidecar progress frame arrives during a cold load. The
guarded waiter uses it to extend the execution deadline (bounded by
MODEL_LOAD_EXTRA_TIMEOUT_S) instead of abandoning a healthy download.
"""
_MODEL_LOAD_ACTIVITY[threading.get_ident()] = (
time.monotonic(), MODEL_LOAD_HEARTBEAT_GRACE_S,
)
def report_generate_progress() -> None:
"""Record that the CURRENT THREAD's job finished a unit of synthesis.
Same contract as the load heartbeat, different evidence: a multi-chunk
render that just completed chunk 7 of 20 is demonstrably working, however
slow it is. Without this, a long text on a modest GPU hit the 300s
execution budget mid-render and was abandoned as "too heavy for the
available compute" — with most of its chunks already rendered, and no way
for the user to tell that from a genuine wedge (#1338/#1348/#1391).
Carries a longer freshness window than the load heartbeat because chunks
are coarse: see GENERATE_PROGRESS_GRACE_S.
"""
_MODEL_LOAD_ACTIVITY[threading.get_ident()] = (
time.monotonic(), GENERATE_PROGRESS_GRACE_S,
)
class GpuJobTimeoutError(TimeoutError):
"""A GPU-pool job **that actually started executing** overran its bound.
Only raised once a worker picked the job up, so the message's "too heavy
for the available compute" reading is truthful. Queue wait is bounded
separately and surfaces as :class:`GpuPoolBusyError`.
"""
class GpuPoolBusyError(TimeoutError):
"""The GPU pool is saturated — the job never started, so nothing was lost.
Retryable verbatim: no compute was spent, no partial state exists. Carries
``retry_after`` (seconds) so HTTP callers can emit a real ``Retry-After``
and scripted clients can back off instead of hammering a busy backend.
"""
def __init__(self, message: str, *, retry_after: float = 30.0):
super().__init__(message)
self.retry_after = max(1, int(round(retry_after)))
def generate_timeout_s(
text: "str | None", *, engine: object = None, execution_device: "str | None" = None,
min_vram_gb: float = 0.0, hardware_family: "str | None" = None,
vram_gb: "float | None" = None,
_include_sidecar_grace: bool = True,
) -> float:
"""THE wall-clock execution budget for one synthesis job, scaled to input.
Single source of truth for every TTS dispatch (#1190/#1202). The
length-scaled budget landed in v0.3.22 but was wired into only two call
sites in generation.py's classic path — the streaming path the UI tries
FIRST, plus /v1/audio/speech, batch, dub and archetype previews, all still
used the flat 300s, which is why 0.3.22 users kept seeing "exceeded 300s"
on long inputs. Lives here (not in a router) so every router shares it
without importing generation.py.
Policy: floor at the configured OMNIVOICE_GENERATE_TIMEOUT_S (accelerated
hosts) or OMNIVOICE_CPU_GENERATE_TIMEOUT_S (CPU hosts — the latter wins
for CPU whenever it is itself explicit, even if the former also is; see
the #1787 comment on the module-level constants), plus 1s per 40
characters past a 1200-character free allowance — generous enough for
CPU-class hardware, still bounded (a wedged job is caught in minutes, not
hours).
#1804: "accelerated" is not one performance class. A card with less VRAM
than the engine declares it needs pages to system RAM over PCIe and renders
SLOWER than the same machine's CPU would — yet, judged by device family
alone, it was handed HALF the CPU budget. That inversion is what three 4 GB
reporters hit (#1226 GTX 1650 Ti, #1222 Quadro P2000, #1804 GTX 1650), all
on the engine that declares a 6 GB floor. Every layer already knew: routing
raises a caveat, the preflight toast warns, and the timeout message names
the card. Only the budget ignored it. So an under-provisioned accelerator
now floors at the CPU budget — the class of hardware it actually performs
like. ``min_vram_gb`` is the engine's declared floor; callers that pass
``engine`` get it read off the engine automatically. Native runtimes pass
an explicit ``vram_gb=0`` when their dedicated-memory probe failed; that
unknown capacity gets the same conservative CPU-class budget without
claiming the card is under-provisioned in user-facing diagnostics.
"""
base = GPU_JOB_TIMEOUT_S
explicit_budget = _GENERATE_TIMEOUT_EXPLICIT or GPU_JOB_TIMEOUT_S != _CONFIGURED_GPU_JOB_TIMEOUT_S
try:
from core.device_caps import detect_host_caps
caps = detect_host_caps()
family = execution_device or caps.family
if not min_vram_gb and engine is not None:
min_vram_gb = float(getattr(engine, "min_vram_gb", 0.0) or 0.0)
if execution_device is None and engine is not None:
from services.engine_routing import runtime_compute_profile
profile = runtime_compute_profile(engine, caps)
family = profile["effective_device"]
min_vram_gb = profile["min_vram_gb"]
hardware_family = profile.get("runtime_hardware_family")
vram_gb = profile.get("runtime_vram_gb")
universal_override = (
_GENERATE_TIMEOUT_EXPLICIT
or GPU_JOB_TIMEOUT_S != _CONFIGURED_GPU_JOB_TIMEOUT_S
)
cpu_explicit = (
_CPU_GENERATE_TIMEOUT_EXPLICIT
or CPU_JOB_TIMEOUT_S != _CONFIGURED_CPU_JOB_TIMEOUT_S
)
if family == "cpu" and (cpu_explicit or not universal_override):
base = CPU_JOB_TIMEOUT_S
explicit_budget = cpu_explicit
elif not universal_override and family in (
"cuda", "rocm", "vulkan", "xpu",
):
from services.engine_routing import under_provisioned_vram
runtime_family = hardware_family or family
unknown_dedicated_vram = (
min_vram_gb > 0
and runtime_family in ("cuda", "rocm", "xpu", "vulkan")
and vram_gb is not None
and float(vram_gb or 0.0) <= 0
)
if unknown_dedicated_vram or under_provisioned_vram(
caps, min_vram_gb, family=hardware_family, vram_gb=vram_gb,
):
base = max(base, CPU_JOB_TIMEOUT_S)
except Exception:
pass
sidecar_grace = 0.0
if not explicit_budget and engine is not None and hasattr(engine, "recv_timeout_s"):
try:
sidecar_timeout = float(engine.recv_timeout_s)
if math.isfinite(sidecar_timeout) and sidecar_timeout > 0:
base = max(base, sidecar_timeout)
sidecar_grace = 5.0 if _include_sidecar_grace else 0.0
except (TypeError, ValueError):
pass
return base + (max(0, len(text or "") - 1200) / 40.0) + sidecar_grace
def _retry_after_estimate(stats: dict) -> float:
"""Seconds a caller should wait before retrying, from live pool state.
Queue depth ahead of you, divided by workers, times a recent job's wall
time. Bounded to 5..300s so the hint is always usable (and never zero on a
cold pool with no timing history yet)."""
base = stats.get("avg_job_s") or 0.0
if base <= 0:
base = 30.0
workers = max(1, int(stats.get("workers") or 1))
waves = (int(stats.get("queued") or 0) + 1) / workers
return max(5.0, min(300.0, base * waves))
def gpu_pool_stats(executor=None) -> dict:
"""Live pool occupancy, or a permissive default for executors that don't
track it (plain ThreadPoolExecutor in tests / injected executors)."""
ex = executor if executor is not None else _get_gpu_pool()
fn = getattr(ex, "stats", None)
if callable(fn):
try:
return fn()
except Exception:
pass
return {"queued": 0, "running": 0, "workers": 1, "avg_job_s": 0.0}
def check_gpu_admission(*, what: str = "GPU job", executor=None) -> None:
"""Admission control at SUBMIT (#1190/#1202) — raise before queueing when
the pool is already backed up.
Policy: refuse when ``queued >= workers`` — every worker is busy AND a full
wave of jobs is *already waiting* ahead of this one. Deliberately NOT the
stricter "no worker is free": on the 1-worker hosts this bug hurts most,
that would reject the ordinary second concurrent request the desktop UI
issues routinely and which completes fine today. The looser rule still
catches the case that matters — a scripted client fanning out N requests at
a pool that can only serialize them — and turns a silent multi-minute wait
into an immediate, honest "retry in N seconds".
"""
stats = gpu_pool_stats(executor)
if stats.get("queued", 0) < max(1, int(stats.get("workers") or 1)):
return
retry_after = _retry_after_estimate(stats)
raise GpuPoolBusyError(
f"{what} was not accepted: the local GPU worker pool is saturated "
f"({stats.get('running', 0)} running, {stats.get('queued', 0)} already "
f"queued on {stats.get('workers', 1)} worker(s)). Nothing was started, "
f"so this request is safe to retry as-is in about "
f"{int(retry_after)}s. To raise throughput, run fewer "
f"concurrent requests, or set OMNIVOICE_GPU_WORKERS if the machine has "
f"spare VRAM.",
retry_after=retry_after,
)
def _log_safe(what: str) -> str:
"""Backward-compatible alias for the shared logging seam."""
from core.logging_utils import log_safe
return log_safe(what, limit=120)
def _swallow_abandoned(fut) -> None:
"""Consume the result of a future we stopped awaiting, so an abandoned
wedged job can't emit "Future exception was never retrieved" noise."""
try:
if not fut.cancelled():
fut.exception()
except (asyncio.CancelledError, Exception):
pass
async def run_on_gpu_pool_guarded(fn, *, what: str = "GPU job",
timeout: "float | None" = None,
executor=None,
queue_timeout: "float | None" = None,
min_vram_gb: float = 0.0,
on_abandon=None):
"""Run blocking ``fn`` on the GPU pool, bounding **execution** — not the
wait for a free worker.
Two clocks (#1190/#1202):
* ``queue_timeout`` (generous, ``GPU_QUEUE_TIMEOUT_S``) covers the time the
job sits in the pool queue. Exceeding it raises :class:`GpuPoolBusyError`
— the job is cancelled out of the queue before it ever runs, so no
compute is wasted and the caller can retry verbatim.
* ``timeout`` (``GPU_JOB_TIMEOUT_S`` by default) starts only when a worker
actually picks the job up. Exceeding *that* is a genuinely wedged/too-slow
job → :class:`GpuJobTimeoutError` + pool ``reset()``.
Previously both were one clock started at submit: ``run_in_executor``
returns immediately, so a job queued behind a busy 1-worker pool burned its
entire budget waiting and then reported "too heavy for the available
compute" without having executed one instruction.
``fn`` must be a zero-arg callable — wrap args with ``functools.partial``.
Executors without ``reset`` (a plain ThreadPoolExecutor in tests) still get
both bounds; only the reset step is skipped.
``min_vram_gb`` is the declared VRAM floor of the engine this job belongs
to (``TTSBackend.min_vram_gb``); it only shapes the timeout MESSAGE. Left
at 0 — the default, and correct for every non-TTS job on this pool
(reference transcribe, watermarking, dub steps) — the under-provisioned-GPU
wording is never used, because nothing measured says it applies (#1226).
``on_abandon`` is called once, after a job whose caller stopped waiting can
no longer access its inputs. A queued job that is cancelled before it
starts calls it immediately; a running thread calls it from ``_job``'s
finalizer. Normal completion never calls it. This lets request-owned temp
files outlive abandoned workers without delaying ordinary requests (#1668).
"""
from services.inference_cancellation import InferenceCancellation
cancellation = InferenceCancellation()
loop = asyncio.get_running_loop()
ex = executor if executor is not None else _get_gpu_pool()
timeout = GPU_JOB_TIMEOUT_S if timeout is None else float(timeout)
queue_timeout = GPU_QUEUE_TIMEOUT_S if queue_timeout is None else float(queue_timeout)
started = asyncio.Event()
_inner = contain_system_exit(fn, what)
_ident_box: dict = {}
_abandon_lock = threading.Lock()
_abandon_state = {
"requested": False,
"finished": False,
"callback_called": False,
}
def _fire_abandon_callback() -> None:
if on_abandon is None:
return
with _abandon_lock:
if _abandon_state["callback_called"]:
return
_abandon_state["callback_called"] = True
try:
on_abandon()
except Exception:
logger.exception("%s abandon cleanup failed", _log_safe(what))
def _job():
_ident_box["ident"] = threading.get_ident()
try:
loop.call_soon_threadsafe(started.set)
except RuntimeError:
pass
try:
with cancellation.activate():
return _inner()
finally:
_MODEL_LOAD_ACTIVITY.pop(threading.get_ident(), None)
with _abandon_lock:
_abandon_state["finished"] = True
abandoned = _abandon_state["requested"]
if abandoned:
_fire_abandon_callback()
from core.render_trace import bind as bind_render_trace
concurrent_fut = ex.submit(bind_render_trace(_job))
fut = asyncio.wrap_future(concurrent_fut, loop=loop)
def _abandon() -> None:
cancellation.cancel()
cancelled_before_start = concurrent_fut.cancel()
with _abandon_lock:
_abandon_state["requested"] = True
finished = _abandon_state["finished"]
fut.cancel()
if cancelled_before_start or finished:
_fire_abandon_callback()
waiter = asyncio.ensure_future(started.wait())
try:
done, _pending = await asyncio.wait(
{waiter, fut}, timeout=queue_timeout,
return_when=asyncio.FIRST_COMPLETED,
)
except asyncio.CancelledError:
_abandon()
fut.add_done_callback(_swallow_abandoned)
raise
finally:
waiter.cancel()
if not done:
_abandon()
fut.add_done_callback(_swallow_abandoned)
stats = gpu_pool_stats(ex)
logger.warning(
"%s waited %.0fs for a free GPU worker and was never started "
"(%d queued / %d running) — reporting pool saturation (#1190).",
_log_safe(what), queue_timeout,
stats.get("queued", 0), stats.get("running", 0),
)
raise GpuPoolBusyError(
f"{what} waited {queue_timeout:.0f}s for a free GPU worker and "
f"never started, so nothing was computed and the request is safe "
f"to retry as-is. The backend is alive but every worker is busy "
f"with earlier jobs. Run fewer concurrent requests, or raise "
f"OMNIVOICE_GPU_WORKERS if the machine has spare VRAM.",
retry_after=_retry_after_estimate(stats),
)
_t0 = time.monotonic()
_soft_deadline = _t0 + timeout
_hard_deadline = _soft_deadline + MODEL_LOAD_EXTRA_TIMEOUT_S
_extended = False
class _ExecutionDeadlineExceeded(Exception):
"""Internal deadline, distinct from a worker's own TimeoutError."""
try:
while True:
_now = time.monotonic()
if _now < _soft_deadline:
_slice = min(_soft_deadline - _now, 5.0)
else:
_beat = _MODEL_LOAD_ACTIVITY.get(_ident_box.get("ident"))
_last, _grace = _beat if _beat else (None, 0.0)
if (_last is None
or _now - _last > _grace
or _now >= _hard_deadline):
if concurrent_fut.done():
return await fut
raise _ExecutionDeadlineExceeded()
if not _extended:
_extended = True
logger.info(
"%s reached its %.0fs execution budget while still "
"making progress — extending while heartbeats continue "
"(grace %.0fs, cap +%.0fs) (#1367/#1391).",
_log_safe(what), timeout, _grace,
MODEL_LOAD_EXTRA_TIMEOUT_S,
)
_slice = max(0.05, min(
(_last + _grace) - _now,
_hard_deadline - _now,
5.0,
))
_done, _ = await asyncio.wait({fut}, timeout=_slice)
if _done:
return fut.result()
except asyncio.CancelledError:
_abandon()
fut.add_done_callback(_swallow_abandoned)
raise
except _ExecutionDeadlineExceeded as timeout_exc:
_abandon()
fut.add_done_callback(_swallow_abandoned)
stacks = log_gpu_pool_worker_stacks(what, timeout, executor=ex)
_reset = getattr(ex, "reset", None)
if callable(_reset):
try:
_reset()
logger.warning(
"%s exceeded %.0fs of EXECUTION time — abandoned the "
"GPU-pool worker; it keeps running (and holding the "
"device) until it finishes on its own (#730/#1190).",
_log_safe(what), timeout,
)
except Exception:
logger.exception("GPU pool reset after %s timeout failed",
_log_safe(what))
raise GpuJobTimeoutError(
_timeout_guidance(
what, timeout, min_vram_gb, wedged=_stack_shows_a_wedge(stacks),
)
) from timeout_exc
_WEDGE_STACK_DEPTH = 25
def _live_pool_thread_idents(executor) -> "set | None":
"""Thread idents belonging to ``executor``'s CURRENT inner pool, or None
when they can't be established.
Needed because a wedged worker survives ``reset()`` — it cannot be
cancelled, so it keeps running under the same ``gpu-pool`` name the
replacement pool also uses. Without this, the second timeout in a session
logs the stale thread alongside the live one with nothing to tell them
apart, and the stale stack is the more misleading of the two: it names an
operation that is no longer the one that just failed (greptile).
``ThreadPoolExecutor._threads`` is private but has been the storage for its
worker set since 3.2 and is stable across every version we support; None
here is a soft degrade to "label nothing", never an error.
"""
pool = getattr(executor, "_pool", executor)
threads = getattr(pool, "_threads", None)
if not threads:
return None
try:
return {t.ident for t in threads if t.ident is not None}
except Exception:
return None
def log_gpu_pool_worker_stacks(what: str, timeout: float, executor=None) -> str:
"""Log where every GPU-pool worker is currently executing. Never raises.
The gap this closes (#1338/#1329/#1348): when a job overran its execution
budget we logged *that* it had, reset the pool, and returned a message
about the machine being too slow — with no record of what the abandoned
thread was actually doing. So every report of this class arrived
undiagnosable, and the only way forward was to ask the user to reproduce it
under a debugger. On an RTX 3060 rendering one sentence, "too heavy for the
available compute" is almost certainly the wrong story, and nothing in the
log could contradict it.
``sys._current_frames()`` reads the frame of every live thread, including
one wedged inside a C call — which is exactly the case here, since the
worker cannot be cancelled and keeps running after we abandon it. Filtered
to gpu-pool workers so the log names the stuck job, not the web server.
Returns the formatted text (also for tests); empty when nothing matched.
"""
try:
import sys as _sys
import threading as _threading
import traceback as _traceback
names = {
t.ident: t.name for t in _threading.enumerate()
if t.ident is not None and t.name.startswith(_GPU_POOL_THREAD_PREFIX)
}
if not names:
return ""
live = _live_pool_thread_idents(executor) if executor is not None else None
frames = _sys._current_frames()
blocks = []
for ident, name in sorted(names.items(), key=lambda kv: kv[1]):
frame = frames.get(ident)
if frame is None:
continue
if live is None:
label = name
elif ident in live:
label = f"{name} (current pool)"
else:
label = (
f"{name} (STALE — a worker abandoned by an earlier timeout, "
f"still running; not the job that just failed)"
)
stack = "".join(_traceback.format_stack(frame, limit=_WEDGE_STACK_DEPTH))
blocks.append(f"--- {label} ---\n{stack.rstrip()}")
if not blocks:
return ""
try:
from core.failure import sanitize as _sanitize
text = _sanitize("\n".join(blocks))
except Exception:
logger.exception("Could not sanitize GPU-pool worker stacks; "
"omitting them rather than logging raw paths")
return ""
logger.warning(
"%s exceeded %.0fs — stack of every GPU-pool worker at the moment "
"it was abandoned. The deepest frame is where it is stuck; if that "
"is inside the model rather than a data copy, this is a hang and "
"not an under-provisioned machine (#1338):\n%s",
_log_safe(what), timeout, text,
)
return text
except Exception:
logger.exception("Could not capture GPU-pool worker stacks")
return ""
_WEDGE_STDLIB_FILES = (
"/threading.py", "\\threading.py",
"/asyncio/locks.py", "\\asyncio\\locks.py",
"/concurrent/futures/_base.py", "\\concurrent\\futures\\_base.py",
"/queue.py", "\\queue.py",
)
_WEDGE_FUNCTIONS = frozenset({
"acquire", "wait", "result", "get", "join", "_wait_for_tstate_lock",
})
_FRAME_HEAD = re.compile(r'^\s*File "(?P<file>.+)", line \d+, in (?P<func>\S+)\s*$')
def _stack_shows_a_wedge(stacks: "str | None") -> bool:
"""True when the abandoned worker's DEEPEST frame is a blocking wait.
The message this feeds is the one users actually read, and for years it
said the same thing whatever happened: "too heavy for the available
compute". That is a specific, testable claim, and when the worker is
parked on a lock it is simply false — nothing was computed, so nothing was
too heavy. #1416 and #1419 both arrived as "my machine is too slow"
reports from people whose jobs never ran at all (a cold load waiting on a
lock owned by another event loop, #1417), and #1329 is the same wedge seen
from the dub loop. Every one of them was sent to look at their hardware.
Only the last frame counts, and it must be a blocking primitive in a
standard-library module. Both halves matter (CodeRabbit): a compute job's
*callers* routinely include a lock it has already left, so scanning the
whole stack would flag nearly everything; and an application function
named ``wait`` or ``result`` is not evidence of anything, so the function
name alone is not enough either.
Reads the text :func:`log_gpu_pool_worker_stacks` already captured — no
second stack walk, and no cost at all on the healthy path.
Conservative: unknown or unparseable stacks return False and keep the old
wording. Claiming a hang we cannot see would be the same mistake pointing
the other way.
"""
if not stacks:
return False
deepest = None
for line in str(stacks).splitlines():
m = _FRAME_HEAD.match(line)
if m:
deepest = m
if deepest is None:
return False
func = deepest.group("func")
if func not in _WEDGE_FUNCTIONS:
return False
path = deepest.group("file").replace("\\", "/")
return any(
path.endswith(tail.replace("\\", "/")) for tail in _WEDGE_STDLIB_FILES
)
def _timeout_guidance(
what: str, timeout: float, min_vram_gb: float = 0.0, *, wedged: bool = False,
) -> str:
"""Device-aware timeout message (#896): a CPU-only host must never be told
to "set the engine to CPU" or blamed on VRAM — on CPU the job is simply
compute-bound. GPU hosts keep the VRAM-contention guidance.
Honesty fix (#1190/#1202): this used to promise "Capacity was restored
automatically". It was not. Python cannot kill the abandoned worker
thread — it runs to completion still holding its VRAM, so an immediate
retry contends with the zombie and is *more* likely to fail, which is
exactly how one slow chunk cascaded into a whole failed batch. The message
now says what actually happens and gives both interactive and scripted
callers something to do about it.
"""
family = "cuda"
device_name, vram_gb = "", 0.0
_caps = None
try:
from core.device_caps import detect_host_caps
_caps = detect_host_caps()
family = _caps.family
device_name, vram_gb = _caps.device_name, _caps.vram_gb
except Exception:
pass
if wedged:
return (
f"{what} was abandoned after {timeout:.0f}s without doing any "
"work — it spent the whole time waiting on an internal lock, not "
"computing. This is a bug in VoiceStudio, not a limit of your "
"machine, so shorter text or a lighter engine won't help. "
"Restart the backend to clear it (Settings → Logs → Backend has "
"the stack trace that was captured), and please report it with "
"that log at https://github.com/debpalash/VoiceStudio/issues — "
"the trace names exactly where it stopped."
)
common = (
f"{what} ran for more than {timeout:.0f}s of actual compute time and "
"was abandoned — the backend is running, but this job was too heavy "
"for the available compute. The abandoned job cannot be killed: it "
"keeps running and keeps holding the device until it finishes on its "
"own, so an immediate retry competes with it. Wait for the current "
"job to drain (or restart the backend) before retrying; "
)
if family == "cpu":
return common + (
"this machine renders on CPU, where long generations are "
"compute-bound. For a durable fix try shorter text or a lighter "
"engine (OmniVoice GGUF and Supertonic-3 are CPU-tuned). If you "
"expect very long single generations, raise "
"the compute-time budget in Settings → Performance & Device."
)
from services.engine_routing import under_provisioned_vram
if under_provisioned_vram(_caps, min_vram_gb):
return common + (
f"{device_name or 'this GPU'} has {vram_gb:.1f} GB of VRAM and "
f"this engine wants about {min_vram_gb:.0f} GB — generations here "
f"are slow enough to hit the limit even with nothing else loaded. "
f"The durable fix is a lighter engine (OmniVoice GGUF and "
f"Supertonic-3 are tuned for small/no GPU) or shorter text; "
f"Flush caches / Unload the resident model (top toolbar or "
f"Model Catalogue) frees what little headroom there is. (Raise "
f"the compute-time budget in Settings → Performance & Device if "
f"you'd rather let long "
f"generations run.)"
)
return common + (
"most often the GPU is VRAM-starved (a resident model and this job "
"contend for memory). For a durable fix, Flush caches / Unload the "
"resident model (top toolbar or Model Catalogue) before retrying, "
"try shorter text, a lighter engine, or set the engine to CPU in "
"Model Catalogue. (Raise the compute-time budget in "
"Settings → Performance & Device for very "
"long single generations.)"
)
_WATERMARK_STOP = object()
class _WatermarkExecutor(Executor):
"""Single daemon worker with a bounded shutdown contract.
``ThreadPoolExecutor`` uses non-daemon workers that Python joins at exit,
so ``wait=False`` still delays process exit while ``wait=True`` can hang
lifespan teardown forever. AudioSeal loading is not cooperatively
cancellable; a daemon worker plus a bounded join is the only thread-based
contract that both preserves in-process model warm-up and guarantees exit.
"""
def __init__(self) -> None:
self._items: queue.Queue = queue.Queue()
self._lock = threading.Lock()
self._shutdown = False
self._thread: threading.Thread | None = None
def submit(self, fn, /, *args, **kwargs) -> Future:
future: Future = Future()
with self._lock:
if self._shutdown:
raise RuntimeError("cannot schedule new futures after shutdown")
if self._thread is None:
self._thread = threading.Thread(
target=self._run,
name="watermark_0",
daemon=True,
)
self._thread.start()
self._items.put((future, fn, args, kwargs))
return future
def _run(self) -> None:
while True:
item = self._items.get()
if item is _WATERMARK_STOP:
return
future, fn, args, kwargs = item
if not future.set_running_or_notify_cancel():
continue
try:
future.set_result(fn(*args, **kwargs))
except (Exception, SystemExit, KeyboardInterrupt) as exc:
future.set_exception(exc)
def is_stopped(self) -> bool:
"""Whether shutdown has completed and this executor can be replaced."""
with self._lock:
return self._shutdown and (
self._thread is None or not self._thread.is_alive()
)
def is_shutdown(self) -> bool:
with self._lock:
return self._shutdown
def shutdown(
self,
wait: bool = True,
*,
cancel_futures: bool = False,
timeout: float | None = None,
) -> bool:
with self._lock:
self._shutdown = True
thread = self._thread
if cancel_futures:
while True:
try:
item = self._items.get_nowait()
except queue.Empty:
break
if item is not _WATERMARK_STOP:
item[0].cancel()
self._items.put(_WATERMARK_STOP)
if wait and thread is not None:
thread.join(timeout=timeout)
return thread is None or not thread.is_alive()
_watermark_pool_singleton: "_WatermarkExecutor | None" = None
_watermark_pool_lock = threading.Lock()
_watermark_pool_accepting = True
def begin_watermark_pool_lifecycle() -> None:
"""Open watermark submissions for a newly-started app lifespan."""
global _watermark_pool_accepting, _watermark_pool_singleton
with _watermark_pool_lock:
if (
_watermark_pool_singleton is not None
and _watermark_pool_singleton.is_stopped()
):
_watermark_pool_singleton = None
_watermark_pool_accepting = (
_watermark_pool_singleton is None
or not _watermark_pool_singleton.is_shutdown()
)
def get_watermark_pool() -> _WatermarkExecutor:
"""Dedicated 1-worker pool for provenance marking. Built lazily so hosts
with watermarking disabled never spawn the thread.
The executor is captured and returned UNDER the lock: reading the global
again after an unlocked null-check could race shutdown_watermark_pool's
reset and hand out None (CodeRabbit, PR #1577)."""
global _watermark_pool_accepting, _watermark_pool_singleton
with _watermark_pool_lock:
if not _watermark_pool_accepting:
if (
_watermark_pool_singleton is not None
and _watermark_pool_singleton.is_stopped()
):
_watermark_pool_singleton = None
_watermark_pool_accepting = True
else:
raise RuntimeError("watermark executor is shutting down")
if (
_watermark_pool_singleton is not None
and _watermark_pool_singleton.is_stopped()
):
_watermark_pool_singleton = None
if _watermark_pool_singleton is None:
_watermark_pool_singleton = _WatermarkExecutor()
return _watermark_pool_singleton
def shutdown_watermark_pool(*, timeout: float = 20.0) -> None:
"""Drain the watermark pool at app shutdown (PR #1577).
Refuse queued work and wait for the active operation: Python cannot kill
a thread inside AudioSeal loading, so returning early would let model
initialization continue during interpreter teardown. The draining pool
remains published until its worker stops, preventing concurrent producers
from creating a replacement that escapes this shutdown. A process that
keeps running after lifespan shutdown (the test suite does exactly this)
gets a fresh pool once the old worker has actually stopped."""
global _watermark_pool_accepting, _watermark_pool_singleton
with _watermark_pool_lock:
_watermark_pool_accepting = False
pool = _watermark_pool_singleton
if pool is not None:
stopped = pool.shutdown(
wait=True,
cancel_futures=True,
timeout=max(0.0, float(timeout)),
)
if stopped:
with _watermark_pool_lock:
if _watermark_pool_singleton is pool:
_watermark_pool_singleton = None
else:
logger.warning(
"Watermark worker exceeded the %.1fs shutdown deadline; "
"abandoning its daemon thread",
timeout,
)
model = None
_model_lock = asyncio.Lock()
_model_load_thread_lock = threading.Lock()
_last_used = time.time()
_loading_detail: dict = {
"sub_stage": None,
"detail": "",
"error": None,
"progress": None,
}
def _configure_rocm_if_needed(torch):
"""Auto-set HSA_OVERRIDE_GFX_VERSION for AMD GPUs on ROCm.
ROCm-enabled PyTorch reports `torch.cuda.is_available() == True` but
some consumer AMD GPUs have GFX IDs the installed build wasn't compiled
for. Setting HSA_OVERRIDE_GFX_VERSION lets them run with the closest
supported architecture.
The override is applied **only when the native gfx is genuinely absent
from this build's arch list**. Newer ROCm wheels support parts that used
to need remapping (gfx1151/Strix Halo is native from ROCm 7.x), and
overriding a natively-supported GPU forces it onto foreign kernels for no
reason — so the map is a fallback, not an unconditional rewrite.
"""
from core.device_caps import (
ROCM_GFX_OVERRIDES,
build_arch_list,
hsa_override_for,
)
if os.environ.get("HSA_OVERRIDE_GFX_VERSION"):
return
try:
device_name = torch.cuda.get_device_name(0).lower()
if not any(kw in device_name for kw in ("amd", "radeon", "instinct")):
return
props = torch.cuda.get_device_properties(0)
gcn_arch = getattr(props, "gcnArchName", "") or ""
gfx_id = gcn_arch.split(":")[0].strip().lower()
target = ROCM_GFX_OVERRIDES.get(gfx_id)
if not target:
return
arch_list = {a.split(":")[0].strip().lower() for a in build_arch_list(torch)}
if not arch_list:
logger.debug(
"ROCm: no arch list from this torch build; leaving "
"HSA_OVERRIDE_GFX_VERSION unset for %s (%s)", device_name, gfx_id,
)
return
if gfx_id in arch_list:
logger.info("ROCm: %s (%s) is natively supported by this build; "
"no HSA_OVERRIDE_GFX_VERSION needed", device_name, gfx_id)
return
if target not in arch_list:
logger.warning(
"ROCm: %s (%s) is unsupported by this build and its remap "
"target %s is missing too — not setting "
"HSA_OVERRIDE_GFX_VERSION.", device_name, gfx_id, target,
)
return
override = hsa_override_for(target)
os.environ["HSA_OVERRIDE_GFX_VERSION"] = override
logger.info("ROCm: auto-set HSA_OVERRIDE_GFX_VERSION=%s (%s) for %s (%s)",
override, target, device_name, gfx_id)
except Exception as e:
logger.debug("ROCm GFX auto-config skipped: %s", e)
def check_device_compatibility():
"""Check if PyTorch supports the current GPU's architecture.
Returns (compatible, warning_message). Compatible is True if OK or
no discrete GPU is present. The arch comparison itself lives in
``core.device_caps.arch_unsupported()`` — shared with the probe, and
CUDA/ROCm-aware (a ROCm build lists ``gfx…``, not ``sm_…`` — #1228).
"""
from core.device_caps import arch_unsupported
torch = _lazy_torch()
if not torch.cuda.is_available():
return True, None
mismatch = arch_unsupported(torch)
if mismatch is None:
return True, None
device_arch, arch_list = mismatch
try:
device_name = torch.cuda.get_device_name(0)
except Exception:
device_name = "GPU"
if getattr(getattr(torch, "version", None), "hip", None) is not None:
return False, (
f"{device_name} ({device_arch}) is not supported by this ROCm "
f"PyTorch build. Supported architectures: {', '.join(arch_list)}. "
f"Set HSA_OVERRIDE_GFX_VERSION to the closest supported target "
f"(e.g. 11.0.0 for a gfx11xx card) or install a ROCm build that "
f"lists {device_arch}."
)
return False, (
f"{device_name} ({device_arch}) is not supported by this PyTorch build. "
f"Supported architectures: {', '.join(arch_list)}. "
f"Install a build that covers it: pip install --force-reinstall torch "
f"--index-url https://download.pytorch.org/whl/cu128"
)
def get_best_device():
"""Detect the best available compute device.
Priority: CUDA/ROCm > Intel XPU > DirectML > MPS > CPU
The *family* decision delegates to ``core.device_caps.detect_host_caps()``
(the single source of truth) so the probe and this loader can never
disagree. This function keeps the side-effects the probe deliberately
avoids: the ROCm ``HSA_OVERRIDE_GFX_VERSION`` env override and the
DirectML device-string return (DirectML is not a torch device family, so
the probe reports it as ``cpu`` — we still resolve the real device string
here for Windows DirectML users). The string contract is unchanged:
``"cuda"`` / ``"xpu"`` / a DirectML device string / ``"mps"`` / ``"cpu"``.
"""
from core.device_caps import detect_host_caps
torch = _lazy_torch()
family = detect_host_caps().family
if family in ("cuda", "rocm"):
_configure_rocm_if_needed(torch)
compatible, warning = check_device_compatibility()
if not compatible:
logger.warning(warning)
if not _env_flag("OMNIVOICE_FORCE_CUDA"):
logger.warning(
"Falling back to CPU: this GPU is unsupported by the installed "
"PyTorch build (set OMNIVOICE_FORCE_CUDA=1 to force CUDA anyway)."
)
return "cpu"
return "cuda"
if family == "xpu":
try:
logger.info("Using Intel XPU device: %s", torch.xpu.get_device_name(0))
except Exception:
logger.info("Using Intel XPU device")
return "xpu"
if family == "mps":
return "mps"
if family == "cpu":
try:
import torch_directml
if torch_directml.device_count() > 0:
logger.info("Using DirectML device (GPU %d)", 0)
return str(torch_directml.device(0))
except ImportError:
pass
return "cpu"
_COMPILE_ERR_MODULE_PREFIXES = ("torch._dynamo", "torch._inductor", "torch.fx", "triton")
_COMPILE_ERR_TB_MARKERS = ("/_dynamo/", "/_inductor/", "/triton/", "torch/fx/")
_COMPILE_ERR_MSG_MARKERS = (
"dynamo", "inductor", "triton", "cudagraph",
"symbolically trace", "torch.compile", "fx graph",
)
def _is_compile_runtime_failure(exc: BaseException) -> bool:
"""True when an exception originates in the torch.compile stack (Dynamo /
Inductor / Triton / FX / CUDA-graph trees) rather than in the model itself.
#278: an independent compile-stack failure can surface during generation
as an AssertionError out of torch/_inductor/cudagraph_trees.py. An
architecture missing from the running torch build's arch list is rejected
earlier by should_torch_compile(), before this runtime fallback applies.
#278 also quotes "Detected that you are using FX to symbolically trace a
dynamo-optimized function"; Dynamo raises that on any device, CPU included,
so it is a compile-stack error to catch here but never an arch signal.
Walks the exception chain and checks (a) the exception type's module,
(b) the message, (c) the traceback file paths — the cudagraph case is a
bare AssertionError, so the traceback check is load-bearing.
"""
import traceback as _tb
seen: set[int] = set()
cur: BaseException | None = exc
while cur is not None and id(cur) not in seen:
seen.add(id(cur))
mod = type(cur).__module__ or ""
if mod.startswith(_COMPILE_ERR_MODULE_PREFIXES):
return True
msg = str(cur).lower()
if any(marker in msg for marker in _COMPILE_ERR_MSG_MARKERS):
return True
try:
for frame in _tb.extract_tb(cur.__traceback__):
filename = (frame.filename or "").replace("\\", "/")
if any(marker in filename for marker in _COMPILE_ERR_TB_MARKERS):
return True
except Exception as traceback_scan_error:
logging.debug(
"Skipping traceback marker scan while classifying compile runtime failure: %s",
traceback_scan_error,
)
if cur.__cause__ is not None:
cur = cur.__cause__
elif not cur.__suppress_context__:
cur = cur.__context__
else:
cur = None
return False
def _install_compile_fallback(_model) -> None:
"""Wrap ``model.generate`` so a torch.compile failure at inference time
falls back to the eager (uncompiled) model instead of failing the
generation (#278).
All TTS paths (generate, archetype previews, dub, stream, batch) funnel
through ``model.generate``, so this is the single choke point. On a
compile-stack failure we: log a clear warning, restore the eager module
(``OptimizedModule._orig_mod``), disable compile for the rest of the
session via ``engine_env.mark_compile_runtime_failure``, reset dynamo
state, and retry the call once eagerly. Non-compile errors (real OOM,
validation, …) propagate unchanged — fully backward compatible for users
whose torch.compile works.
"""
orig_generate = _model.generate
def _generate_with_compile_fallback(*args, **kwargs):
try:
return orig_generate(*args, **kwargs)
except Exception as exc:
compiled = getattr(_model, "llm", None)
eager = getattr(compiled, "_orig_mod", None)
if eager is None or not _is_compile_runtime_failure(exc):
raise
logger.warning(
"torch.compile runtime failure during generation (%s: %s) — "
"falling back to the eager model and disabling torch.compile "
"for this session. Generation is being retried without it.",
type(exc).__name__, exc,
)
from services import engine_env
engine_env.mark_compile_runtime_failure(f"{type(exc).__name__}: {exc}")
_model.llm = eager
try:
torch = _lazy_torch()
torch._dynamo.reset()
except Exception as reset_exc:
logger.debug(
"Non-fatal: failed to reset torch._dynamo state after compile failure (%s: %s). "
"Continuing with eager fallback.",
type(reset_exc).__name__,
reset_exc,
)
try:
return orig_generate(*args, **kwargs)
except Exception as eager_exc:
raise eager_exc from None
_model.generate = _generate_with_compile_fallback
def _is_flashinfer_runtime_failure(exc: BaseException) -> bool:
"""True when an exception originates in the FlashInfer fast path (the
flashinfer package, our omnivoice_flashinfer patch module, or CUDA-graph
capture/replay) rather than in the model or the request itself. Same
chain/traceback walk as ``_is_compile_runtime_failure``."""
import traceback as _tb
tb_markers = ("/flashinfer/", "omnivoice_flashinfer")
msg_markers = ("flashinfer", "cuda graph", "cudagraph")
seen: set[int] = set()
cur: BaseException | None = exc
while cur is not None and id(cur) not in seen:
seen.add(id(cur))
mod = type(cur).__module__ or ""
if mod.startswith("flashinfer"):
return True
msg = str(cur).lower()
if any(marker in msg for marker in msg_markers):
return True
try:
for frame in _tb.extract_tb(cur.__traceback__):
filename = (frame.filename or "").replace("\\", "/")
if any(marker in filename for marker in tb_markers):
return True
except Exception:
pass
if cur.__cause__ is not None:
cur = cur.__cause__
elif not cur.__suppress_context__:
cur = cur.__context__
else:
cur = None
return False
def _unapply_flashinfer(_model) -> None:
"""Restore the standard execution path on a FlashInfer-patched model.
``apply_flashinfer`` works entirely through *instance-level* state —
MethodType-bound ``forward``/``_generate_iterative`` overrides and
``_fi_*`` attributes — so deleting those attributes restores the class
implementations exactly. The attention implementation is restored to the
one captured before apply (``_fi_orig_attn_impl`` — could be
flash_attention_2, not just sdpa), and use_cache is re-enabled."""
llm = getattr(_model, "llm", None)
orig_attn = getattr(_model, "_fi_orig_attn_impl", None) or "sdpa"
if llm is not None:
for module in llm.modules():
if "forward" in vars(module):
del module.forward
for attr in ("_fi_w_qkv", "_fi_qkv_split", "_fi_rope_theta", "_fi_w_gate_up"):
if attr in vars(module):
delattr(module, attr)
try:
llm.set_attn_implementation(orig_attn)
except Exception:
logger.exception(
"failed to restore %s attention after FlashInfer", orig_attn
)
llm.config.use_cache = True
for attr in (
"_fi_orig_attn_impl",
"_generate_iterative",
"_fi_runner",
"_fi_graph_cache",
"_fi_enable_cuda_graph",
"_fi_graph_buckets",
"_fi_overhead_budget",
):
if attr in vars(_model):
delattr(_model, attr)
def _install_flashinfer_fallback(_model) -> None:
"""Wrap ``model.generate`` so a FlashInfer failure at inference time falls
back to the standard path instead of failing the generation — the same
contract as ``_install_compile_fallback`` (#278): an optimization must
never turn a working generation into an error."""
orig_generate = _model.generate
def _generate_with_flashinfer_fallback(*args, **kwargs):
try:
return orig_generate(*args, **kwargs)
except Exception as exc:
if not _is_flashinfer_runtime_failure(exc):
raise
logger.warning(
"FlashInfer runtime failure during generation (%s: %s) — "
"restoring the standard path and disabling FlashInfer for "
"this session. Generation is being retried without it.",
type(exc).__name__, exc,
)
from services import engine_env
engine_env.mark_flashinfer_runtime_failure(
f"{type(exc).__name__}: {exc}"
)
_unapply_flashinfer(_model)
_model.generate = orig_generate
try:
return orig_generate(*args, **kwargs)
except Exception as plain_exc:
raise plain_exc from None
_model.generate = _generate_with_flashinfer_fallback
_TORCH_COMPILE_MODE = "reduce-overhead"
_CUDAGRAPH_COMPILE_MODES = frozenset({"reduce-overhead", "max-autotune"})
_CUDAGRAPH_MIN_CAPABILITY = (8, 0)
_FORCE_CUDAGRAPH_ENV = "OMNIVOICE_FORCE_CUDAGRAPH"
def _resolve_compile_mode() -> str:
"""The ``torch.compile`` mode to use on this GPU (#2135).
Returns the configured cudagraph mode on Ampere+, and the non-cudagraph
``"default"`` on older architectures where graph capture has been observed
to abort the process. Fails *safe* (→ "default") only when we positively
identify a pre-Ampere device; any probe error keeps the configured mode so
a weird torch build doesn't silently lose the optimization.
"""
if _TORCH_COMPILE_MODE not in _CUDAGRAPH_COMPILE_MODES:
return _TORCH_COMPILE_MODE
if os.environ.get(_FORCE_CUDAGRAPH_ENV, "").strip().lower() in {"1", "true", "yes", "on"}:
logger.warning(
"%s=1 — keeping torch.compile mode %r on a GPU where CUDA-graph "
"capture is not known-good (#2135).",
_FORCE_CUDAGRAPH_ENV, _TORCH_COMPILE_MODE,
)
return _TORCH_COMPILE_MODE
try:
import torch
if not torch.cuda.is_available():
return _TORCH_COMPILE_MODE
capability = torch.cuda.get_device_capability(0)
except Exception:
logger.debug("compile-mode capability probe failed; keeping %r",
_TORCH_COMPILE_MODE, exc_info=True)
return _TORCH_COMPILE_MODE
if tuple(capability) >= _CUDAGRAPH_MIN_CAPABILITY:
return _TORCH_COMPILE_MODE
try:
device_name = torch.cuda.get_device_name(0)
except Exception:
device_name = "this GPU"
logger.info(
"torch.compile mode %r downgraded to 'default' on %s (sm_%d%d): CUDA-graph "
"capture below sm_%d%d has been seen to abort the backend process (#2135). "
"Compiled kernels are still used. Set %s=1 to override.",
_TORCH_COMPILE_MODE, device_name, capability[0], capability[1],
_CUDAGRAPH_MIN_CAPABILITY[0], _CUDAGRAPH_MIN_CAPABILITY[1],
_FORCE_CUDAGRAPH_ENV,
)
return "default"
_compiled_inference_executor: "ThreadPoolExecutor | None" = None
_compiled_inference_thread_ident: "int | None" = None
def _get_compiled_inference_executor() -> ThreadPoolExecutor:
"""The single-thread executor that owns ALL inference on a compiled model.
Created lazily the first time a model is compiled with a cudagraph mode;
reused across model reloads (idle unload → reload keeps the same thread,
which is fine — a fresh compile simply captures its graphs there too).
The worker is spun up eagerly so its thread ident is known for the
re-entrancy guard in `_install_compile_thread_affinity`.
"""
global _compiled_inference_executor, _compiled_inference_thread_ident
if _compiled_inference_executor is None:
_compiled_inference_executor = ThreadPoolExecutor(
max_workers=1, thread_name_prefix="compiled-infer",
)
_compiled_inference_thread_ident = _compiled_inference_executor.submit(
threading.get_ident
).result()
return _compiled_inference_executor
def _install_compile_thread_affinity(_model) -> None:
"""Pin every ``model.generate`` call to the dedicated compile thread (#315).
Wraps ``model.generate`` (the single choke point all TTS paths funnel
through — generate, archetype previews, dub, stream, batch) so the call
body always runs on `_get_compiled_inference_executor()`'s one thread.
That makes the thread that *captures* the CUDA graph on the first render
and the thread that *replays* it on every later render the same thread,
deterministically, regardless of which `_gpu_pool` worker dispatched it.
Installed AFTER `_install_compile_fallback`, so the call-time order is:
caller thread → hop to the dedicated thread → eager-fallback wrapper →
real generate (the #278 classification/retry also runs on the dedicated
thread, with native tracebacks). The hop is a no-op when already on the
dedicated thread — a 1-worker executor submitting to itself would
deadlock, so the re-entrancy guard is load-bearing.
"""
executor = _get_compiled_inference_executor()
inner_generate = _model.generate
def _generate_on_compile_thread(*args, **kwargs):
if threading.get_ident() == _compiled_inference_thread_ident:
return inner_generate(*args, **kwargs)
return executor.submit(inner_generate, *args, **kwargs).result()
_model.generate = _generate_on_compile_thread
def _set_loading(sub_stage: str, detail: str = "", error: str | None = None, progress: float | None = None):
"""Update the loading detail dict atomically."""
_loading_detail["sub_stage"] = sub_stage
_loading_detail["detail"] = detail
_loading_detail["error"] = error
_loading_detail["progress"] = progress
try:
from core import event_bus
event_bus.emit("model_status", {"sub_stage": sub_stage})
except Exception:
logger.debug("Could not publish model status", exc_info=True)
def _env_flag(name: str, default: bool = False) -> bool:
value = os.environ.get(name)
if value is None:
return default
return value.strip().lower() in {"1", "true", "yes", "on"}
def should_preload_tts_asr() -> bool:
"""Whether VoiceStudio.from_pretrained should attach PyTorch Whisper.
The default is intentionally false. On Apple Silicon, eager TTS + ASR
loading can overcommit unified memory and leave desktop startup stuck
at the model-loading stage. ASR backends still load on demand.
"""
return _env_flag("OMNIVOICE_PRELOAD_TTS_ASR")
def _is_incomplete_cache_error(exc: BaseException) -> bool:
"""True when `exc` is the truncated-HF-cache class (#352 / #581 / #1273).
transformers raises an OSError when the on-disk snapshot has config and
tokenizer files but no weight shard — the signature of an interrupted
download. We match on the message (stable across transformers 4.x/5.x)
rather than the error type, since the same OSError type covers unrelated
I/O failures.
There are TWO wordings, and this used to match only the first, so a
half-written repo whose *subfolder* failed to load (#1273:
"Error no file named model.safetensors, … found in directory
…/snapshots/<rev>/audio_tokenizer") got neither the automatic repair nor
an actionable message — just a raw 500. `core.failure` owns the phrase
list so the heal and the error text can't drift apart."""
from core.failure import is_incomplete_cache_message
return is_incomplete_cache_message(str(exc))
def _is_corrupt_model_file_error(exc: BaseException) -> bool:
"""True when a model weight or config file cannot be parsed.
The other half of the interrupted-download class (#1406). transformers
only raises the "does not appear to have a file named …" signature when
the shard is *absent*; a shard that stops mid-file, gets truncated by
antivirus, or is actually a saved HTML error page opens fine and then
fails inside safetensors:
Error while deserializing header: header too large
That is a ``SafetensorError`` from a Rust extension — not an ``OSError``,
so it never reached the recovery ladder and surfaced as a raw 500 on every
generation (the reporter hit it from voice design *and* from a gallery
preview, which is what a shared broken shard looks like).
The whole exception chain is checked, not just the outermost message:
transformers wraps the tensor library's error in its own before it gets
here, and matching only the surface would miss every wrapped case."""
from core.failure import is_corrupt_model_file_message
return any(is_corrupt_model_file_message(str(e)) for e in _exception_chain(exc))
def _is_corrupt_weights_error(exc: BaseException) -> bool:
"""Backward-compatible wrapper for the original #1406 helper name."""
return _is_corrupt_model_file_error(exc)
def _hf_offline() -> bool:
"""Respect HF's offline switches so repair never makes a network call the
user opted out of. `snapshot_download` would itself raise offline, but
checking up front lets us skip straight to the actionable message."""
return _env_flag("HF_HUB_OFFLINE") or _env_flag("TRANSFORMERS_OFFLINE")
_LINK_REPAIR_ATTEMPTED: set[str] = set()
_FORCED_REDOWNLOAD_ATTEMPTED: set[str] = set()
def _selfheal_broken_snapshot_links(checkpoint: str) -> bool:
"""Rung 0 of cache recovery: delete-and-restore broken snapshot entries.
Returns True only when broken entries were found, removed AND restored —
i.e. retrying the load is worth it. At most one attempt per repo per
process. Never raises; when it returns False the legacy resume/force
ladder still runs."""
if checkpoint in _LINK_REPAIR_ATTEMPTED:
return False
_LINK_REPAIR_ATTEMPTED.add(checkpoint)
if os.path.isdir(checkpoint):
return False
try:
from services.hf_cache_repair import repair_repo_cache
summary = repair_repo_cache(checkpoint)
except Exception as repair_err:
logger.warning("Snapshot-link self-heal for %s errored: %s",
checkpoint, repair_err)
return False
if summary.get("removed") and summary.get("ok"):
logger.warning(
"Model cache for %s had %d broken file link(s) — repaired "
"automatically (%s), retrying the load.",
checkpoint, summary["removed"],
summary.get("outcome") or "healed",
)
return True
if summary.get("found"):
logger.warning(
"Model cache for %s has %d broken file link(s) that could not be "
"auto-repaired (%s).",
checkpoint, summary["found"], summary.get("error") or "unknown",
)
return False
def _manual_cache_delete_hint(checkpoint: str) -> str:
"""Names the exact on-disk folder to delete when every auto-repair rung
failed — "delete the model" is only actionable if the user can find it.
Empty for local-directory checkpoints (they don't live in the hub cache)."""
try:
if os.path.isdir(checkpoint):
return ""
from services.hf_cache_repair import repo_cache_dir
return (
f" If the problem persists, quit VoiceStudio, delete "
f"{repo_cache_dir(checkpoint)} and restart — the model "
"re-downloads automatically."
)
except Exception:
return ""
_last_repair_error: str = ""
def _repair_failure_detail() -> str:
"""One sanitized clause naming why auto-repair failed, or "" (#886).
Feeds user-facing messages (the generate 500 detail / model status), so it
goes through core.failure.sanitize — and because the cause text is now part
of the surfaced error, the shared HF-mirror hint (#874) fires on it when
the repair failed against an unreachable configured mirror."""
if not _last_repair_error:
return ""
try:
from core.failure import sanitize
cause = sanitize(_last_repair_error)
except Exception:
cause = _last_repair_error
return f" Auto-repair failed with: {cause}."
def _repair_model_cache(checkpoint: str, *, force: bool = False) -> bool:
"""Re-fetch a checkpoint's missing files in place and report success.
An interrupted download leaves the cache missing only some files;
`snapshot_download` resumes/fills exactly those (already-present, correctly
sized blobs are skipped by hash, so a near-complete cache repairs in
seconds and a complete one would no-op). Returns False — leaving the caller
to surface the actionable delete-and-reinstall message — when repair is
impossible (offline) or the re-fetch itself fails (no network, gated repo,
full disk). Never raises; repair is best-effort.
``force=True`` passes ``force_download`` so the re-fetch replaces files that
are *present but corrupt* — a truncated/garbled blob that still has the right
size won't be re-fetched by the default resume (#739). It re-downloads the
whole snapshot, so it's the last resort the load path only reaches after a
plain resume-repair didn't fix the cache."""
global _last_repair_error
_last_repair_error = ""
if _hf_offline():
logger.warning(
"Model cache for %s is incomplete but HF offline mode is set — "
"cannot auto-repair.", checkpoint,
)
_last_repair_error = (
"Hugging Face offline mode is enabled (HF_HUB_OFFLINE/TRANSFORMERS_OFFLINE)"
)
return False
try:
from huggingface_hub import snapshot_download
except Exception as imp_err:
logger.warning("Cannot import snapshot_download to repair cache: %s", imp_err)
_last_repair_error = f"{type(imp_err).__name__}: {imp_err}"
return False
try:
from services.hf_cache_repair import hf_cache_home
from services.hf_revisions import installed_revision
cache_root = hf_cache_home()
revision = installed_revision(checkpoint, cache_root)
except (OSError, ValueError) as revision_err:
_last_repair_error = str(revision_err)
logger.warning("Refusing unpinned model repair for %s: %s", checkpoint, revision_err)
return False
dl_kwargs: dict = {
"repo_id": checkpoint,
"revision": revision,
"cache_dir": cache_root,
}
try:
from services import endpoint_race
endpoint = endpoint_race.effective_endpoint()
except Exception:
endpoint = os.environ.get("HF_ENDPOINT")
if endpoint:
dl_kwargs["endpoint"] = endpoint
if force:
dl_kwargs["force_download"] = True
if os.name == "nt":
dl_kwargs["local_dir_use_symlinks"] = False
def _attempt() -> None:
"""One snapshot_download, tolerating an hf_hub that rejects the optional
symlink knob. Lets real failures (network, gated repo, disk) propagate."""
try:
snapshot_download(**dl_kwargs)
except TypeError:
dl_kwargs.pop("local_dir_use_symlinks", None)
snapshot_download(**dl_kwargs)
try:
retries = max(1, int(os.environ.get("OMNIVOICE_MODEL_REPAIR_RETRIES", "3")))
except ValueError:
retries = 3
try:
backoff = max(0.0, float(os.environ.get("OMNIVOICE_MODEL_REPAIR_BACKOFF_S", "2")))
except ValueError:
backoff = 2.0
logger.info(
"Auto-repairing incomplete model cache for %s (up to %d attempt(s)) …",
checkpoint, retries,
)
for attempt in range(1, retries + 1):
try:
_attempt()
logger.info("Auto-repair of %s completed; retrying model load.", checkpoint)
return True
except Exception as e:
logger.warning(
"Auto-repair of %s attempt %d/%d failed: %s",
checkpoint, attempt, retries, e,
)
_last_repair_error = f"{type(e).__name__}: {e}"
if attempt < retries:
try:
from services import endpoint_race
if endpoint_race.reselect_after_failure(checkpoint, str(e)):
new_ep = endpoint_race.effective_endpoint()
if new_ep:
dl_kwargs["endpoint"] = new_ep
else:
dl_kwargs.pop("endpoint", None)
logger.info(
"Auto-repair of %s: endpoint failover — retrying on %s",
checkpoint, new_ep or "https://huggingface.co",
)
except Exception:
pass
if backoff:
time.sleep(backoff * attempt)
return False
_DEFAULT_OMNIVOICE_CHECKPOINT = "k2-fsa/OmniVoice"
def resolve_omnivoice_checkpoint() -> str:
"""Resolve the VoiceStudio TTS checkpoint from ``OMNIVOICE_MODEL``, self-healing
a misconfigured value.
A valid checkpoint is either a HuggingFace repo id (``org/repo`` — contains a
``/``) or an existing local directory. A bare token like ``"omnivoice"`` — a
TTS *engine id* that leaked into ``OMNIVOICE_MODEL`` (e.g. a stale pref/env) —
is neither, and would crash model load with *"omnivoice is not a local folder
and is not a valid model identifier listed on huggingface.co/models"* (#693).
Fall back to the default rather than 500 on every launch.
"""
checkpoint = os.environ.get("OMNIVOICE_MODEL", _DEFAULT_OMNIVOICE_CHECKPOINT).strip()
if not checkpoint:
return _DEFAULT_OMNIVOICE_CHECKPOINT
if checkpoint == "test":
return checkpoint
if "/" in checkpoint or "\\" in checkpoint or os.path.isabs(checkpoint):
return checkpoint
logger.warning(
"OMNIVOICE_MODEL=%r is not a HuggingFace repo id (org/repo) or a local "
"path — falling back to %s (#693).",
checkpoint, _DEFAULT_OMNIVOICE_CHECKPOINT,
)
return _DEFAULT_OMNIVOICE_CHECKPOINT
_INTERPRETER_SHUTDOWN_MSG = "cannot schedule new futures after interpreter shutdown"
_SCHEDULE_AFTER_SHUTDOWN_MSG = "cannot schedule new futures after"
class ModelLoadInterruptedByShutdown(RuntimeError):
"""A model load cut short because the backend is shutting down (#1174).
Benign by definition — the load didn't *fail*, the process is exiting.
``_load_model_sync`` raises this instead of the raw executor error so no
caller (preload task, request handler, log formatter) can dress an
expected teardown up as a crash: no ERROR log, no ``/model/status``
phantom error, no exit-code-poisoning traceback.
"""
_shutting_down = threading.Event()
def begin_shutdown() -> None:
"""Graceful shutdown started: in-flight/queued model loads are now benign
cancellations, and new loads must not start (#1174)."""
_shutting_down.set()
def reset_shutdown_flag() -> None:
"""New run starting — arm model loads again (lifespan startup)."""
_shutting_down.clear()
def is_shutting_down() -> bool:
return _shutting_down.is_set()
def _exception_chain(exc: "BaseException | None"):
"""Yield ``exc`` and every ``__cause__``/``__context__`` ancestor once
(cycle-safe). transformers' lazy-import + materialization machinery wraps
the original error several layers deep."""
seen: set[int] = set()
while exc is not None and id(exc) not in seen:
seen.add(id(exc))
yield exc
exc = exc.__cause__ or exc.__context__
def _is_interpreter_shutdown_error(exc: "BaseException | None") -> bool:
"""True when `exc` (or anything in its cause/context chain) is — or
carries the text of — the ``RuntimeError`` a ``ThreadPoolExecutor`` raises
once Python has begun interpreter shutdown, i.e. the operation was
interrupted by the process exiting, not by a real fault.
Two match modes, both required:
- the live exception object: ``RuntimeError`` whose message mentions
``interpreter shutdown`` anywhere in the chain;
- the *stringified* form: transformers ≥5 aggregates materializer-worker
errors into NEW exceptions whose message embeds the original traceback
as text (``log_conversion_errors`` formats it into
``loading_info.conversion_errors`` → ``SkipParameters`` → summary
raise), which changes the type AND severs the cause chain — the exact
miss behind the "Model loading failed: cannot schedule new futures
after interpreter shutdown" ERROR logged during pytest teardown
(#1174). Matching the full CPython phrase inside any message keeps
that conclusive without loosening the plain-pool case.
"""
for e in _exception_chain(exc):
if isinstance(e, RuntimeError) and "interpreter shutdown" in str(e):
return True
if _INTERPRETER_SHUTDOWN_MSG in str(e):
return True
return False
def _is_schedule_after_shutdown_error(exc: "BaseException | None") -> bool:
"""Any executor 'cannot schedule new futures after …' rejection, either
variant, live or stringified. Only consulted while ``_shutting_down`` is
set: during app shutdown even the plain single-pool variant is benign
(our own ``_reset_gpu_pool()``/executor teardown caused it). Outside
shutdown the plain variant stays the #589-class real fault and must NOT
be silenced."""
return any(_SCHEDULE_AFTER_SHUTDOWN_MSG in str(e) for e in _exception_chain(exc))
def _load_model_sync():
global model
if _shutting_down.is_set():
logger.info("Model load skipped: backend is shutting down.")
raise ModelLoadInterruptedByShutdown("model load skipped: backend shutting down")
from utils.hf_progress import register_listener, unregister_listener
def _on_hf_progress(ev):
pct = ev.get("pct", 0.0)
filename = ev.get("filename", "")
phase = ev.get("phase", "")
if pct > 0:
pct_int = min(round(pct * 100), 99)
detail = _loading_detail.get("detail", "")
base = detail.split(" —")[0].split(" (")[0]
_loading_detail["progress"] = pct_int
_loading_detail["detail"] = f"{base} — {pct_int}%"
lid = register_listener(_on_hf_progress)
try:
_set_loading("importing", "Importing PyTorch & VoiceStudio runtime…")
logger.info("Importing PyTorch & VoiceStudio runtime…")
torch = _lazy_torch()
VoiceStudio = _lazy_omnivoice()
device = get_best_device()
checkpoint = resolve_omnivoice_checkpoint()
_set_loading("loading_weights", f"Loading TTS weights on {device}…")
logger.info("Loading VoiceStudio model on device: %s", device)
preload_asr = should_preload_tts_asr()
if preload_asr:
logger.info("Preloading PyTorch Whisper after TTS model load.")
else:
logger.info("Skipping PyTorch Whisper preload; ASR will load on demand.")
def _load():
return VoiceStudio.from_pretrained(
checkpoint, device_map=device, dtype=torch.float16, load_asr=False,
)
def _recover_corrupt_weights(exc: BaseException):
"""Re-fetch weights that are on disk but unparseable (#1406).
Deliberately a FORCED re-download rather than the resume ladder
below: a resume trusts a blob that is already the expected size
and would never re-fetch the one that is actually wrong.
"""
repair_checkpoint = checkpoint
for nested_exc in _exception_chain(exc):
repository_id = getattr(nested_exc, "repository_id", None)
if repository_id == "eustlb/higgs-audio-v2-tokenizer":
repair_checkpoint = repository_id
break
asset_label = (
"audio tokenizer"
if repair_checkpoint != checkpoint
else "TTS model"
)
if repair_checkpoint in _FORCED_REDOWNLOAD_ATTEMPTED:
raise RuntimeError(
f"The {asset_label} files for {repair_checkpoint} are damaged and a "
"re-download did not fix them. Open the engine's Weights list in Model Catalogue, "
"delete the VoiceStudio TTS model, and install it again."
f"{_manual_cache_delete_hint(repair_checkpoint)}"
) from exc
_FORCED_REDOWNLOAD_ATTEMPTED.add(repair_checkpoint)
logger.warning(
"%s files for %s are present but unparseable (%s) — a "
"download that stopped mid-file, or a file altered on disk "
"after it arrived. Re-fetching them.",
asset_label,
repair_checkpoint,
exc,
)
_set_loading("loading_weights", "Model files are damaged — re-downloading…")
if not _repair_model_cache(repair_checkpoint, force=True):
raise RuntimeError(
f"The {asset_label} files for {repair_checkpoint} are damaged — a "
"download that stopped part-way, or a file changed on "
"disk after it arrived — and could not be re-downloaded "
f"automatically.{_repair_failure_detail()} Open Settings "
"→ Models, delete the VoiceStudio TTS model, and install "
f"it again.{_manual_cache_delete_hint(repair_checkpoint)}"
) from exc
_set_loading("loading_weights", f"Loading TTS weights on {device}…")
try:
return _load()
except Exception as exc2:
if not _is_corrupt_weights_error(exc2):
raise
raise RuntimeError(
f"The {asset_label} files for {repair_checkpoint} are still damaged "
"after being re-downloaded. Open the engine's Weights list in Model Catalogue, "
"delete the VoiceStudio TTS model, and install it again."
f"{_manual_cache_delete_hint(repair_checkpoint)}"
) from exc2
try:
_model = _load()
except OSError as e:
if _is_corrupt_weights_error(e):
_model = _recover_corrupt_weights(e)
elif not _is_incomplete_cache_error(e):
raise
else:
_model = None
if _selfheal_broken_snapshot_links(checkpoint):
_set_loading(
"loading_weights",
"Model cache had broken file links — repaired "
"automatically, retrying…",
)
try:
_model = _load()
except OSError as e_link:
if not _is_incomplete_cache_error(e_link):
raise
logger.warning(
"Load still failing after snapshot-link repair of %s — "
"falling back to resume repair.", checkpoint,
)
e = e_link
_model = None
if _model is None:
_set_loading("loading_weights", "Repairing incomplete model cache…")
if not _repair_model_cache(checkpoint):
raise RuntimeError(
f"The TTS model cache for {checkpoint} is incomplete "
"(weights missing — usually an interrupted download)."
f"{_repair_failure_detail()} "
"Open the engine's Weights list in Model Catalogue, delete the VoiceStudio TTS model, "
f"and install it again.{_manual_cache_delete_hint(checkpoint)}"
) from e
_set_loading("loading_weights", f"Loading TTS weights on {device}…")
try:
_model = _load()
except OSError as e2:
if _is_corrupt_weights_error(e2):
_model = _recover_corrupt_weights(e2)
elif _is_incomplete_cache_error(e2):
_set_loading("loading_weights", "Re-downloading model files…")
if _repair_model_cache(checkpoint, force=True):
try:
_model = _load()
except OSError as e3:
raise RuntimeError(
f"The TTS model cache for {checkpoint} is incomplete "
"and could not be auto-repaired. Open Model "
"Catalogue → Models, delete the VoiceStudio TTS model, and install "
f"it again.{_manual_cache_delete_hint(checkpoint)}"
) from e3
else:
raise RuntimeError(
f"The TTS model cache for {checkpoint} is incomplete and "
f"could not be auto-repaired.{_repair_failure_detail()} "
"Open the engine's Weights list in Model Catalogue, delete the VoiceStudio TTS model, "
f"and install it again.{_manual_cache_delete_hint(checkpoint)}"
) from e2
else:
raise RuntimeError(
f"The TTS model cache for {checkpoint} is incomplete and "
"could not be auto-repaired. Open the engine's Weights list in Model Catalogue, delete "
"the VoiceStudio TTS model, and install it again."
f"{_manual_cache_delete_hint(checkpoint)}"
) from e2
except Exception as e_corrupt:
if not _is_corrupt_weights_error(e_corrupt):
raise
_model = _recover_corrupt_weights(e_corrupt)
if preload_asr:
try:
_model.load_asr_model()
except Exception as asr_exc:
if not _is_corrupt_model_file_error(asr_exc):
raise
raise RuntimeError(
"The transcription model's files are damaged. Open "
"the engine's Weights list in Model Catalogue, delete the transcription (ASR) model, "
"and install it again; or set OMNIVOICE_PRELOAD_TTS_ASR=0 "
"to stop preloading it alongside TTS."
) from asr_exc
flashinfer_applied = False
try:
from services.engine_env import (
mark_flashinfer_runtime_failure,
should_flashinfer,
)
fi_mode = should_flashinfer(device)
if fi_mode != "off":
_set_loading("compiling", "Applying FlashInfer kernels…")
try:
from omnivoice.models.omnivoice_flashinfer import apply_flashinfer
_model._fi_orig_attn_impl = getattr(
_model.llm.config, "_attn_implementation", "sdpa"
)
apply_flashinfer(_model, enable_cuda_graph=(fi_mode == "graph"))
except Exception as fi_exc:
mark_flashinfer_runtime_failure(
f"{type(fi_exc).__name__}: {fi_exc}"
)
_unapply_flashinfer(_model)
else:
flashinfer_applied = True
_install_flashinfer_fallback(_model)
_install_compile_thread_affinity(_model)
logger.info(
"FlashInfer applied (mode=%s) — torch.compile skipped "
"for this load.", fi_mode,
)
except Exception:
logger.exception("FlashInfer opt-in check failed; continuing without")
try:
from services.engine_env import should_torch_compile
if not flashinfer_applied and should_torch_compile(device):
_set_loading("compiling", "Compiling model (torch.compile)…")
compile_mode = _resolve_compile_mode()
try:
_model.llm = torch.compile(_model.llm, mode=compile_mode)
except Exception as compile_exc:
from services.engine_env import mark_compile_runtime_failure
mark_compile_runtime_failure(f"{type(compile_exc).__name__}: {compile_exc}")
logger.warning(
"torch.compile failed (%s) — continuing with the eager model.",
compile_exc,
)
else:
_install_compile_fallback(_model)
if compile_mode in _CUDAGRAPH_COMPILE_MODES:
_install_compile_thread_affinity(_model)
logger.info(
"torch.compile mode %r uses CUDA graphs — compiled-model "
"inference pinned to a single dedicated thread (#315).",
compile_mode,
)
logger.info("torch.compile applied (mode=%r).", compile_mode)
except Exception as e:
logger.info("torch.compile skipped: %s", e)
try:
setattr(_model, "_voicestudio_checkpoint", checkpoint)
setattr(
_model,
"_voicestudio_loaded_at",
time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
)
except Exception:
logger.debug("Could not attach resident model identity", exc_info=True)
_set_loading("ready", "Model ready", progress=100)
logger.info("VoiceStudio model loaded successfully.")
return _model
except ModelLoadInterruptedByShutdown:
raise
except Exception as exc:
if _is_interpreter_shutdown_error(exc) or (
_shutting_down.is_set() and _is_schedule_after_shutdown_error(exc)
):
logger.info(
"Model load aborted: shutdown during load — benign, not a failure."
)
raise ModelLoadInterruptedByShutdown("shutdown during load") from exc
try:
from core.failure import build_failure
_f = build_failure(exc, stage="model-load", include_diagnostic=False)
err_msg = _f["reason"] + (f" — {_f['hint']}" if _f.get("hint") else "")
except Exception:
err_msg = str(exc)
_set_loading("error", "Model loading failed", error=err_msg)
logger.error("Model loading failed: %s", str(exc), exc_info=exc)
raise
finally:
unregister_listener(lid)
def _model_load_timeout() -> float:
"""Overall ceiling (seconds) for a single model load/download attempt.
Backstop for any hang the HF per-read socket timeouts don't catch
(a wedged torch.compile, a deadlock, etc.). Generous by default so a
legitimate cold multi-GB download on a slow link still completes;
overridable via OMNIVOICE_MODEL_LOAD_TIMEOUT for very slow networks.
"""
try:
return max(30.0, float(os.environ.get("OMNIVOICE_MODEL_LOAD_TIMEOUT", "1200")))
except (ValueError, TypeError):
return 1200.0
def _reset_gpu_pool() -> None:
"""Recover from a wedged/timed-out load by abandoning the GPU worker pool.
The resilient wrapper is kept (its identity is shared by every importer);
only its inner `ThreadPoolExecutor` is dropped, so the next submit builds a
fresh worker. This is what stops stale references from raising "cannot
schedule new futures after shutdown" after a reset (#589 #599).
"""
if _gpu_pool_singleton is not None:
_gpu_pool_singleton.reset()
async def _load_model_with_timeout():
"""Run the blocking model load on the GPU pool, bounded by a deadline.
Raises RuntimeError on timeout (and resets the poisoned pool) so callers
surface an actionable error instead of hanging indefinitely.
This is the shared load boundary for BOTH get_model() and the startup
preload_model() — the memory reclaim must live here, or a memory-tight
machine gets protected on demand loads but OS-killed during the startup
preload (review finding on the original placement in get_model()).
"""
_make_room_before_tts_load()
loop = asyncio.get_running_loop()
timeout = _model_load_timeout()
try:
return await asyncio.wait_for(
loop.run_in_executor(_get_gpu_pool(), _load_model_sync),
timeout=timeout,
)
except asyncio.TimeoutError as exc:
_set_loading("error", "Model load timed out", error="timeout")
_reset_gpu_pool()
logger.error("Model load exceeded %ss; resetting GPU pool.", timeout)
raise RuntimeError(
f"Model loading timed out after {int(timeout)}s — usually a network "
"stall downloading the model (proxy, firewall, or antivirus). Check "
"your connection or set a Hugging Face mirror in Settings, then retry."
) from exc
async def get_model():
global model, _last_used
_last_used = time.time()
if model is not None:
await _heal_tts_placement()
await asyncio.get_running_loop().run_in_executor(None, make_room_before_generate)
return model
if running_on_gpu_pool():
if model is None:
with _model_load_thread_lock:
if model is None:
from core.run_sentinel import touch_activity
touch_activity("model_load", "omnivoice-tts")
_make_room_before_tts_load()
model = _load_model_sync()
return model
async with _model_lock:
if model is None:
from core.run_sentinel import touch_activity
touch_activity("model_load", "omnivoice-tts")
model = await _load_model_with_timeout()
return model
def _make_room_before_tts_load() -> None:
"""Evict-then-load: free what we already own before a tight TTS load.
The audit's top gap: on a 16 GB unified-memory box a plain TTS load could
still be OS-killed — the dub path frees memory before *ASR* loads
(offload_tts_for_asr, #1119), but nothing freed memory before a *TTS*
load, and a warm dictation model (~2 GB) is routinely the difference.
Deliberately NOT admission control: refusing a load on an estimate would
brick machines that would actually cope (the #1111 decision — advisory
only). This only releases things the app already reclaims on idle anyway
(the capture-ASR model, engine instances, allocator caches), just *now*
instead of after the idle timeout — and only when free memory is actually
tight, so a roomy machine pays nothing.
"""
try:
from services.memory_budget import available_memory
free_gb = (available_memory() or {}).get("ram_available_gb")
if free_gb is None or free_gb >= _UNIFIED_OFFLOAD_HEADROOM_GB:
return
logger.info(
"Memory tight before TTS load (%.1f GB free), releasing idle "
"models first.", free_gb,
)
_release_idle_tts_memory("load")
except Exception:
logger.debug("pre-load memory reclaim skipped", exc_info=True)
def _release_idle_tts_memory(stage):
"""Drop capture-ASR, TTS side caches, and allocator caches. Best-effort;
never raises (a cleanup failure must not break the load/generate that called
it). Shared by the cold-load and warm-generate make-room paths so the
eviction recipe cannot drift between them (#730/#1190)."""
try:
try:
from services.asr_backend import release_idle_capture_backend
release_idle_capture_backend(0.0)
except Exception:
logger.debug("capture-ASR pre-%s release failed", stage, exc_info=True)
release_tts_side_caches()
free_vram()
except Exception:
logger.debug("pre-%s memory reclaim skipped", stage, exc_info=True)
def _should_make_room_for_generate():
"""Decide whether to free idle GPU memory before a generate (#730/#1190).
Modes (OMNIVOICE_FREE_VRAM_BEFORE_GENERATE):
auto (default): free when free system RAM is below the unified headroom,
mirroring _make_room_before_tts_load. A roomy machine pays nothing.
always: free before every generate (small per-call cost from gc.collect +
cache drop).
never: opt out.
"""
mode = os.environ.get("OMNIVOICE_FREE_VRAM_BEFORE_GENERATE", "auto").strip().lower()
if mode == "never":
return False
if mode == "always":
return True
try:
from services.memory_budget import available_memory
free_gb = (available_memory() or {}).get("ram_available_gb")
if free_gb is not None and free_gb < _UNIFIED_OFFLOAD_HEADROOM_GB:
return True
except Exception:
logger.debug("make_room memory probe failed", exc_info=True)
return False
def make_room_before_generate():
"""Free idle GPU memory before a warm, heavy generate (#730/#1190).
The cold LOAD path already evicts (``_make_room_before_tts_load`` runs inside
``_load_model_with_timeout``), but the warm path (model already resident,
``get_model`` returns early at the cache check) skipped it. A long generate
on a VRAM-tight MPS box then contended with capture-ASR and the clone-prompt
side cache until it exceeded the execution budget and was abandoned, which is
exactly how one slow synth cascaded into a stuck, device-holding backend.
This runs the same fail-safe eviction the load path uses, just before a
generate the policy says is likely to starve.
Deliberately NOT admission control and NOT a device reclaim. It only drops
things the app already releases on idle, just now instead of later, so a
roomy machine or a short synth pays nothing. It cannot kill an already
abandoned worker; only a crash-isolated subprocess engine can (see
services.subprocess_backend).
"""
if not _should_make_room_for_generate():
return
_release_idle_tts_memory("generate")
def _checkpoint_in_local_cache(checkpoint: str) -> bool:
"""True when ``checkpoint`` is loadable with NO network: an existing local
directory, or a COMPLETE HF cache snapshot. ``snapshot_download(...,
local_files_only=True)`` never constructs an HTTP session, so a broken
proxy env (#959: ``ALL_PROXY``/``HTTPS_PROXY=socks5://`` without socksio)
can't false-negative this probe. Never raises."""
if os.path.isdir(checkpoint):
return True
try:
from huggingface_hub import snapshot_download
snapshot_download(checkpoint, local_files_only=True)
return True
except Exception:
return False
def _headless_worker() -> bool:
"""True when this process serves remote work and has no local UI."""
try:
from worker.agent import worker_mode_enabled
except Exception:
return False
return worker_mode_enabled()
async def preload_model():
"""Background model warm-up — call from lifespan startup.
Loads the TTS model on the GPU pool thread so the first /generate
call is near-instant instead of waiting 4-6s for weight loading.
Non-blocking: if models aren't installed yet, silently exits.
"""
global model, _last_used
if model is not None:
return
try:
from core.device_caps import detect_host_caps
if detect_host_caps().family == "mps":
logger.info(
"Native TTS preload skipped: OmniVoice uses crash isolation on this host."
)
return
except Exception:
logger.debug("effective TTS preload selection failed", exc_info=True)
if _headless_worker():
logger.info(
"Preload skipped: this process is running as a remote worker, so the "
"model loads on first request and is released when it goes idle."
)
return
try:
checkpoint = resolve_omnivoice_checkpoint()
if not _checkpoint_in_local_cache(checkpoint):
logger.info(
"Preload skipped: %s is not installed locally — the model "
"will load (and download if requested) on first use.",
checkpoint,
)
return
logger.info("Preloading TTS model in background…")
_last_used = time.time()
async with _model_lock:
if model is None:
model = await _load_model_with_timeout()
logger.info("Preload complete — model ready.")
except ModelLoadInterruptedByShutdown:
logger.info("Model preload stopped: shutdown during load — benign.")
except Exception as e:
logger.warning("Model preload failed (non-fatal): %s", e, exc_info=e)
try:
from core.failure import build_failure
from core.failure import describe_exception
reason = " | ".join(
describe_exception(exc) for exc in _exception_chain(e)
) or describe_exception(e)
failure = build_failure(
reason, stage="model-preload", include_diagnostic=False,
)
detail = failure.get("hint") or failure.get("reason") or str(e)
except Exception:
detail = (
"The TTS model could not be loaded. Settings → Logs → Backend "
"has the full error."
)
_set_loading("error", detail, error=detail)
def get_model_status():
is_loaded = model is not None
try:
is_loading = (not is_loaded) and _model_lock.locked()
except Exception:
is_loading = False
status = "loading" if is_loading else ("ready" if is_loaded else "idle")
checkpoint = None
loaded_at = None
if is_loaded:
checkpoint = getattr(model, "_voicestudio_checkpoint", None)
loaded_at = getattr(model, "_voicestudio_loaded_at", None)
if not checkpoint:
try:
checkpoint = resolve_omnivoice_checkpoint()
except Exception:
logger.debug("Could not resolve resident model identity", exc_info=True)
result = {
"loaded": is_loaded,
"loading": is_loading,
"status": status,
}
if checkpoint is not None:
result["checkpoint"] = checkpoint
if loaded_at is not None:
result["loaded_at"] = loaded_at
sub = _loading_detail.get("sub_stage")
err = _loading_detail.get("error")
if sub and (is_loading or is_loaded or err):
result["sub_stage"] = sub
result["detail"] = _loading_detail.get("detail", "")
progress = _loading_detail.get("progress")
if progress is not None:
result["progress"] = progress
if err:
result["error"] = err
return result
def _resolve_idle_timeout() -> float:
"""In-process model idle timeout in seconds (MM2-05): prefs store → env →
core.config default, env winning. Resolved per-tick so a settings change
takes effect without a restart."""
try:
from core import prefs
return float(prefs.resolve(
"idle_timeout_seconds",
env="OMNIVOICE_IDLE_TIMEOUT_S",
default=IDLE_TIMEOUT_SECONDS,
))
except (TypeError, ValueError, ImportError):
return float(IDLE_TIMEOUT_SECONDS)
async def idle_worker():
torch = _lazy_torch()
while True:
await asyncio.sleep(30)
idle_timeout = _resolve_idle_timeout()
async with _model_lock:
if model is not None and time.time() - _last_used > idle_timeout:
logger.info("Idle timeout reached. Unloading VoiceStudio model to free VRAM.")
unload_shared_model()
try:
from services.asr_backend import release_idle_capture_backend
if release_idle_capture_backend(idle_timeout):
free_vram()
except Exception:
logger.warning("idle capture-ASR release failed", exc_info=True)
try:
from services.watermark import release_idle_models
release_idle_models(idle_timeout)
except Exception:
logger.warning("idle watermark-model release failed", exc_info=True)
def release_tts_side_caches():
"""Drop caches keyed to the TTS model, for when the model itself is released.
The voice-clone prompt cache (services.tts_backend) holds encoded reference
tensors belonging to *this* model instance. If the model is unloaded but the
prompts survive, an "unload" no longer means unload (#1119) — they sit in the
very memory the unload was reclaiming (``_offload_unified_memory`` drops the
model precisely to hand that RAM to the ASR model).
Previously only ``OmniVoiceBackend.unload()`` cleared them, which sufficed
while the cache was adapter-only. The native ``/generate`` path now populates
it too, and that path unloads through *here*, never through the adapter.
Reached through ``sys.modules`` rather than an import, deliberately:
``tts_backend`` already imports this module, so importing it back would close
a real cycle — and doing it at *import* time (e.g. a registration hook) drags
``core.config`` in earlier than it is today, which perturbs DATA_DIR binding.
A plain lookup has neither problem, and is exactly right besides: if the module
was never imported, it has no cache to clear.
Best-effort by construction — cache hygiene must never be able to break an
unload, because a failed unload is how the backend gets OOM-killed.
"""
mod = sys.modules.get("services.tts_backend")
if mod is None:
return
try:
mod.clear_clone_prompt_cache()
except Exception:
logger.debug("clone-prompt cache clear failed during unload", exc_info=True)
def _clear_cublas_workspaces(torch) -> None:
"""Drop cuBLAS's per-handle workspaces before emptying the cache.
Measured on a 4090: after unloading the model, ``empty_cache()`` left
803 MB reserved with 8.5 MB allocated. A segment dump explained it —
**one** 803 MB segment, 794.7 MB of it inactive-but-split, pinned by a
single live 8,519,680-byte block. That number is cuBLAS's default
workspace. It is taken from the caching allocator on first use, it lands
inside whatever segment the model load had just grown, and it is held for
the life of the cuBLAS handle — so one 8.5 MB block kept three quarters of
a gigabyte from ever going back to the driver, no matter how many times
the user pressed Flush Memory.
Clearing the workspaces first lets the whole segment go. The next cuBLAS
call re-allocates one, which is why this belongs here (on the unload
paths) and not on any hot path.
Private API, so it is optional by construction: a torch build without it
keeps today's behaviour rather than failing an unload.
"""
clear = getattr(getattr(torch, "_C", None), "_cuda_clearCublasWorkspaces", None)
if clear is None:
return
try:
clear()
except Exception:
logger.debug("clearing cuBLAS workspaces failed", exc_info=True)
def free_vram():
"""Release cached GPU memory on any accelerator (CUDA, MPS, XPU, NPU)."""
torch = _lazy_torch()
import gc
gc.collect()
if torch.cuda.is_available():
_clear_cublas_workspaces(torch)
torch.cuda.empty_cache()
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
torch.mps.empty_cache()
elif hasattr(torch, "xpu") and torch.xpu.is_available():
torch.xpu.empty_cache()
elif hasattr(torch, "npu") and torch.npu.is_available():
torch.npu.empty_cache()
def unload_shared_model() -> bool:
"""Drop the shared VoiceStudio model and actually give the memory back.
The order is the entire point of this function. Clearing the reference has
to come FIRST, then the allocator caches. ``free_vram()`` run while
``model`` is still bound releases nothing: the weights are still reachable,
so ``gc.collect()`` keeps them and ``empty_cache()`` only returns blocks
the allocator already considered free. The reference drops a moment later,
the weights go back into torch's cache, and nobody ever hands that cache to
the driver — so the unload is logged, the engine is dropped from the
registry, and ``nvidia-smi`` does not move.
Six call sites open-coded this pair and one of them had it inverted — the
one the engine-registry sweep reaches, which is the sweep a headless worker
node runs. A worker therefore sat on 3.6 GB indefinitely while reporting
the engine released, and every other path looked fine (#1495). One helper,
so there is one ordering and nowhere left to get it wrong.
Takes no lock of its own: the sync engine-registry path
(``OmniVoiceBackend.unload``) cannot await one, and callers that do hold
``_model_lock`` simply keep holding it across the call. Assignment is
GIL-atomic, so the worst a race costs is a redundant reload. Idempotent —
returns False when nothing was resident.
"""
global model
if model is None:
return False
model = None
release_tts_side_caches()
free_vram()
return True
def _has_dedicated_vram():
"""Check if the current device has limited dedicated VRAM that needs offloading."""
torch = _lazy_torch()
if torch.cuda.is_available():
return True
if hasattr(torch, "xpu") and torch.xpu.is_available():
return True
if hasattr(torch, "npu") and torch.npu.is_available():
return True
return False
_UNIFIED_OFFLOAD_HEADROOM_GB = float(
os.environ.get("OMNIVOICE_UNIFIED_OFFLOAD_HEADROOM_GB", "6.0")
)
def _offload_unified_memory() -> bool:
"""Release the TTS model on a unified-memory host when RAM is tight.
Returns True when the model was actually released. Never raises — a failure
to make room must not abort the transcription that asked for it."""
global model
try:
from services.memory_budget import available_memory
free_gb = available_memory().get("ram_available_gb")
if free_gb is not None and free_gb > _UNIFIED_OFFLOAD_HEADROOM_GB:
return False
logger.info(
"Unified memory tight (%s GB free) — releasing the TTS model so ASR has room "
"(it reloads on the next generation).",
"unknown" if free_gb is None else f"{free_gb:.1f}",
)
unload_shared_model()
return True
except Exception as e:
logger.warning("unified-memory TTS offload failed (continuing): %s", e)
return False
def offload_tts_for_asr():
"""Move TTS model to CPU to free VRAM for ASR (WhisperX large-v3).
On a 7-8 GB laptop GPU the TTS model (~2.4 GB) and WhisperX large-v3
(~3 GB) plus the VAD model can't coexist. Offloading the TTS model to
CPU before transcription prevents CUDA OOM, then restore_tts_after_asr()
moves it back.
Works on CUDA (NVIDIA + ROCm) and Intel XPU.
"""
global model
torch = _lazy_torch()
if model is None:
return
if not _has_dedicated_vram():
_offload_unified_memory()
return
try:
if torch.cuda.is_available():
free_mem = torch.cuda.mem_get_info()[0]
if free_mem > 8 * 1024 ** 3:
return
except Exception:
pass
try:
logger.info("Offloading TTS model to CPU to free VRAM for ASR...")
model.to("cpu")
free_vram()
logger.info("TTS model offloaded. VRAM freed for ASR.")
except Exception as e:
logger.warning("TTS offload failed: %s", e)
def restore_tts_after_asr():
"""Move TTS model back to the GPU after ASR completes."""
global model
torch = _lazy_torch()
if model is None:
return
if not _has_dedicated_vram():
return
try:
device = get_best_device()
if device in ("cuda", "xpu"):
logger.info("Restoring TTS model to %s...", device)
model.to(device)
free_vram()
except Exception as e:
logger.warning("TTS restore to %s failed: %s", get_best_device(), e)
def _first_param_device(obj):
"""Device the weights of ``obj`` actually live on, or None if undeterminable.
The TTS runtime is a wrapper object, not necessarily an ``nn.Module``, so
fall back to the first sub-module that owns parameters. Never raises.
"""
try:
params = getattr(obj, "parameters", None)
if callable(params):
for p in params():
return p.device
except Exception:
pass
try:
torch = _lazy_torch()
for v in vars(obj).values():
if isinstance(v, torch.nn.Module):
for p in v.parameters():
return p.device
except Exception:
pass
return None
def _stranded_tts_target():
"""Target device string when the loaded TTS model is stranded off it, else None.
Ordered cheapest-first so the hot path (model already on the accelerator)
costs a single parameter probe: anything not sitting on CPU is by
definition not stranded, because the only thing that moves the model is
``offload_tts_for_asr()`` and it only ever moves it to CPU.
"""
m = model
if m is None:
return None
dev = _first_param_device(m)
if dev is None or getattr(dev, "type", None) != "cpu":
return None
if not _has_dedicated_vram():
return None
try:
target = get_best_device()
except Exception:
return None
return target if target in ("cuda", "xpu") else None
def ensure_tts_on_device() -> bool:
"""Move the TTS model back onto its target device if it was stranded on CPU.
Returns True when a move actually happened. Never raises — a failed move
just leaves the model on CPU, which is exactly the pre-fix behaviour
(slow), never a failed generation.
"""
target = _stranded_tts_target()
m = model
if target is None or m is None:
return False
try:
logger.warning(
"TTS model found stranded on CPU (an ASR offload was never restored) — "
"moving it back to %s; generation would otherwise run 10-50x slower (#1191).",
target,
)
m.to(target)
free_vram()
return True
except Exception as e:
logger.warning("TTS placement self-heal to %s failed (staying on CPU): %s", target, e)
return False
async def _heal_tts_placement() -> None:
"""Async wrapper for :func:`ensure_tts_on_device` used by ``get_model()``.
The cheap mismatch probe runs inline; the rare actual move is dispatched to
the **GPU pool** so it serializes against in-flight inference — moving a
shared model's weights underneath a running ``generate()`` is the one way
this could make things worse than the bug it fixes. The pool that can
strand a model is always 1-worker (``offload_tts_for_asr`` only fires below
8 GB free VRAM, and ``_workers_for_free_vram`` gives such a host a single
worker), so occupying a slot is genuine mutual exclusion there.
"""
if _stranded_tts_target() is None:
return
if running_on_gpu_pool():
ensure_tts_on_device()
return
async with _model_lock:
if _stranded_tts_target() is None:
return
try:
await asyncio.get_running_loop().run_in_executor(
_get_gpu_pool(), ensure_tts_on_device
)
except Exception as e:
logger.warning("TTS placement self-heal could not run: %s", e)
_diar_pipeline = None
DIARIZATION_ERR_NO_TOKEN = "NO_TOKEN"
DIARIZATION_ERR_LICENSE = "PYANNOTE_LICENSE_REQUIRED"
DIARIZATION_ERR_LOAD = "LOAD_FAILED"
DIARIZATION_ERR_MISSING = "MODEL_MISSING"
def _classify_diarization_error(exc: BaseException) -> str:
"""Map a pyannote/HF-hub exception to one of the diarization error
sentinels above.
The 401/403 path is the canonical "user hasn't accepted the model
license on huggingface.co" symptom — both `Pipeline.from_pretrained`
and `huggingface_hub` raise distinct exception classes for it
depending on the installed versions, so we sniff on both the class
name and the stringified message rather than importing the
`HfHubHTTPError` symbol directly (which is not stable across
huggingface_hub majors).
"""
name = type(exc).__name__.lower()
msg = str(exc).lower()
if "localentrynotfounderror" in name or isinstance(exc, FileNotFoundError):
return DIARIZATION_ERR_MISSING
if (
"401" in msg
or "403" in msg
or "unauthorized" in msg
or "gated" in msg
or "accept" in msg and ("license" in msg or "terms" in msg or "user conditions" in msg)
or "gatedrepoerror" in name
or "repositorynotfounderror" in name and "gated" in msg
):
return DIARIZATION_ERR_LICENSE
return DIARIZATION_ERR_LOAD
def _ensure_pyannote_hf_token_compat():
"""pyannote-audio 3.x calls huggingface_hub.hf_hub_download / snapshot_download
with the ``use_auth_token`` kwarg, which huggingface_hub 1.x removed (only
``token`` remains) — raising ``hf_hub_download() got an unexpected keyword
argument 'use_auth_token'`` and breaking diarization (#167).
Wrap those functions to translate the deprecated kwarg. We patch
huggingface_hub itself BEFORE pyannote is imported, so pyannote's
``from huggingface_hub import hf_hub_download`` binds the wrapped fn; we
also patch any already-imported pyannote submodule that bound it directly.
Idempotent (guarded by an attribute marker).
"""
import functools
import sys as _sys
import huggingface_hub as _hf
def _wrap(orig):
if orig is None or getattr(orig, "_ov_uat_shim", False):
return orig
@functools.wraps(orig)
def _wrapped(*args, **kwargs):
if "use_auth_token" in kwargs:
kwargs.setdefault("token", kwargs.pop("use_auth_token"))
return orig(*args, **kwargs)
_wrapped._ov_uat_shim = True
return _wrapped
for _name in ("hf_hub_download", "snapshot_download"):
if hasattr(_hf, _name):
setattr(_hf, _name, _wrap(getattr(_hf, _name)))
for _modname, _mod in list(_sys.modules.items()):
if _modname.startswith("pyannote.") and _mod is not None:
for _name in ("hf_hub_download", "snapshot_download"):
if hasattr(_mod, _name):
setattr(_mod, _name, _wrap(getattr(_mod, _name)))
def get_diarization_pipeline(return_error: bool = False):
"""Load (or return the cached) pyannote speaker-diarization-3.1 pipeline.
Default return: the pipeline instance, or `None` if anything went
wrong (no token, license not accepted, model load crashed). Existing
callers (dub_core legacy `_transcribe`) rely on the `None` sentinel.
When `return_error=True`, returns a 2-tuple
`(pipeline | None, error_sentinel | None)` where `error_sentinel` is
one of the `DIARIZATION_ERR_*` constants. This shape is what the
streaming `_diarize` path uses to emit a structured SSE warning with
a docs deeplink — issue #78.
"""
global _diar_pipeline
from services.diarization_runtime import SORTFORMER, selected_backend
if selected_backend() == SORTFORMER:
try:
from services.diarization_native import NativeSortformer
pipeline = NativeSortformer()
return (pipeline, None) if return_error else pipeline
except Exception as exc:
logger.exception("Could not prepare native Sortformer")
return (None, _classify_diarization_error(exc)) if return_error else None
if _diar_pipeline is not None:
return (_diar_pipeline, None) if return_error else _diar_pipeline
from services import token_resolver
resolved = token_resolver.resolve()
hf_token = resolved.token if resolved else False
try:
torch = _lazy_torch()
_ensure_pyannote_hf_token_compat()
try:
from services.asr_backend import WhisperXBackend
WhisperXBackend._allow_vad_pickle_globals()
except Exception as _glob_e:
logger.debug("pyannote safe-globals allowlist skipped: %s", _glob_e)
from pyannote.audio import Pipeline
logger.info("Loading Pyannote Diarization Pipeline...")
from services.diarization_local import local_pipeline_config
with local_pipeline_config() as config_path:
pipeline = Pipeline.from_pretrained(config_path, use_auth_token=hf_token)
if pipeline is None:
raise RuntimeError("The installed diarisation pipeline could not be loaded")
device = get_best_device()
if device in ("cuda",):
pipeline.to(torch.device(device))
_diar_pipeline = pipeline
logger.info("Pyannote Diarization Pipeline loaded on %s.", device)
return (_diar_pipeline, None) if return_error else _diar_pipeline
except Exception as e:
err_class = _classify_diarization_error(e)
if resolved is None and err_class == DIARIZATION_ERR_MISSING:
err_class = DIARIZATION_ERR_NO_TOKEN
logger.exception(
"Failed to load Pyannote pipeline (class=%s)", err_class,
)
return (None, err_class) if return_error else None
def unload_diarization_pipeline() -> bool:
"""Release a resident pyannote pipeline after the runtime changes."""
global _diar_pipeline
pipeline = _diar_pipeline
_diar_pipeline = None
if pipeline is None:
return False
del pipeline
try:
free_vram()
except Exception:
logger.debug("Could not clear accelerator cache after diarisation unload", exc_info=True)
return True