VoiceStudio/backend/services/model_manager.py

import os
import re
import sys
import time
import asyncio
import logging
import queue
import threading
from concurrent.futures import Executor, Future, ThreadPoolExecutor

from utils.containment import contain_system_exit

# ── Lazy imports ─────────────────────────────────────────────────────
# torch and VoiceStudio are heavy (~2-3s import on Apple Silicon).
# Deferring them until first use cuts cold start from ~4s to ~1.5s,
# so health/status endpoints respond immediately on boot.

_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:
            # The class is OmniVoice — a library identifier, not product
            # branding. The VoiceStudio rename must not touch it (checkpoint
            # configs reference the class name via transformers architectures).
            from omnivoice.models.omnivoice import OmniVoice as _OV
        except ModuleNotFoundError as exc:
            if not _missing_module_is_omnivoice(exc):
                # Something in omnivoice's OWN import chain is missing — not
                # omnivoice itself (#1415). transformers' lazy module raises
                # ModuleNotFoundError for any attribute whose backing import
                # failed ("Could not import module 'AutoFeatureExtractor'"),
                # and a missing torchaudio/torchvision raises it by name. The
                # source-tree fallback below cannot fix any of those: it
                # re-imports from the same broken environment and fails
                # identically, having logged that the *editable install* is
                # broken — which sent the reporter, and us, after the wrong
                # thing. Let it through with its own cause intact; classify()
                # already names it TRANSFORMERS_IMPORT and hints at the real
                # remedy.
                raise
            # The venv's editable install is missing/broken (#564). main.py wires
            # the source fallback at startup, but resolve it here too so the
            # model-load path self-heals and logs the paths it searched.
            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")

# Per-TTS-job VRAM headroom estimate. VoiceStudio's forward + autoregressive
# decode peaks around 1.6 GB, but the interactive clone path co-loads WhisperX
# large-v3 ASR (~3 GB) to transcribe the reference, so a *concurrent* clone job
# is realistically ~5 GB. The old 2.5 GB budget over-committed: an 8 GB card
# (~7 GB free) got 2 workers, and two concurrent clone jobs blew past VRAM into
# a sticky CUDA "illegal memory access" that aborts the whole backend process —
# the wave of "Can't reach the local backend" crash reports on 8 GB GPUs
# (#567/#570/#571/#580/#582/#583/#584). Budgeting 5 GB serializes to 1 worker on
# ≤10 GB cards (no contention → no crash) while 16/24 GB cards still parallelize.
# Power users override with OMNIVOICE_GPU_WORKERS.
_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


# thread_name_prefix for the GPU pool, centralised so the "am I on a gpu-pool
# worker?" predicates (running_on_gpu_pool below; SubprocessBackend.generate's
# on-pool skip) cannot drift from the pool's actual prefix. A drift would
# silently re-introduce the 1-worker self-deadlock this couples against.
_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()
        # ── Queue accounting (#1190/#1202) ───────────────────────────────
        # `queued` = submitted but not yet picked up by a worker; `running` =
        # executing right now. Admission control (check_gpu_admission) and the
        # Retry-After estimate both read these, so a scripted client learns the
        # pool is saturated at SUBMIT instead of after a 300s silent wait.
        self._stats_lock = threading.Lock()
        self._queued = 0
        self._running = 0
        self._avg_job_s = 0.0  # EMA of completed job wall time

    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:
            # "cannot schedule new futures after shutdown": the inner pool was
            # reset (or torn down) under us. Rebuild once and retry so a stale
            # caller self-heals instead of 500-ing. (Interpreter-shutdown races
            # re-raise on the retry — we don't loop.)
            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):
        # Every dispatch (guarded or raw run_in_executor) funnels through here,
        # so wrapping the callable is the one place that sees queue→run→done
        # for the whole pool.
        token = {"counted": False}

        def _tracked(*a, **kw):
            with self._stats_lock:
                token["counted"] = True
                self._queued -= 1
                self._running += 1
            t0 = time.monotonic()
            try:
                # A bare StopIteration here would never reach the caller — it
                # hangs the awaiting future instead (see WorkerStopIteration).
                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):
            # A job cancelled before a worker picked it up never runs _tracked;
            # release its queue slot here so the depth can't drift upward.
            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}")


# ── GPU-job timeout guard (#730 class; residual #850/#802/#755 …) ─────
# A blocking GPU job that wedges on a Windows+CUDA hang keeps occupying its
# worker forever — run_in_executor can't cancel the thread. With a 1–2 worker
# pool that starves *every* other request, so the next user action surfaces as
# the misleading "Can't reach the local backend" even though the process is
# alive. ASR/dub/model-load already bound+reset on hang (run_transcribe_guarded,
# _reset_pool_on_wedge, _load_model_with_timeout); the TTS **generate** paths
# (generation.py, tts_stream.py) were the last unguarded dispatch — and the
# residual on-main reports all fail on generate:start (audio). This is the same
# guard generalised so every GPU dispatch shares one recovery path.
_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 synthesis is healthy but substantially slower than accelerated inference.
# Keep a separate, bounded floor so a short render on CPU is not abandoned at
# the GPU-oriented five-minute deadline (#1588).
#
# #1787 review fix: an explicit OMNIVOICE_CPU_GENERATE_TIMEOUT_S must ALWAYS
# govern CPU dispatches, even when OMNIVOICE_GENERATE_TIMEOUT_S is ALSO
# explicit. Before this flag existed, `universal_override` below treated any
# explicit GENERATE_TIMEOUT_S as authoritative for CPU too, so the Settings
# panel's "CPU budget" row could be saved and silently never apply whenever
# the "Accelerated" row was also set — the exact defect (a control that looks
# like it works and doesn't) issue #1787 exists to remove. Setting ONLY
# OMNIVOICE_GENERATE_TIMEOUT_S keeps its historical "universal" behavior
# unchanged (test_explicit_universal_generate_timeout_wins_on_cpu) — nobody
# who already relies on that single-var override loses it. The only case that
# changes is the previously-undocumented, previously-broken combination of
# setting BOTH: the more specific (CPU) value now wins for CPU jobs, matching
# what a user who filled in both Settings rows was told would happen.
_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

# Queue-wait budget — a SEPARATE, deliberately generous clock (#1190/#1202).
# The execution bound above must never be spent waiting in line: a job queued
# behind a busy 1-worker pool used to burn its whole 300s budget without
# executing a single instruction and then be told it was "too heavy for the
# available compute". Waiting long is normal on a 1-worker host (that is what
# serialization means); waiting *forever* is not, so the queue still has a
# bound — crossing it means saturation, which is a retryable 503, not a
# too-heavy job.
GPU_QUEUE_TIMEOUT_S = float(os.environ.get("OMNIVOICE_GPU_QUEUE_TIMEOUT_S", "1800.0"))

# ── model-load heartbeats (#1367) ────────────────────────────────────────────
# A first-use generate on a subprocess engine DOWNLOADS the model inside the
# job, and the sidecar proves the download is healthy by emitting a progress
# frame every ~5s. The execution clock above ignored that: a slow connection
# blew the 300s budget mid-download and the user was told their hardware was
# too slow, while the sidecar's own watchdog was happily fed. These three make
# the two clocks agree — a job is only "wedged" when it is SILENT.
#
# How long a heartbeat stays fresh. Sidecars emit every ~5s (_HEARTBEAT_S in
# each engine's main.py); 30s tolerates a stall between frames without keeping
# a genuinely dead load alive for long.
MODEL_LOAD_HEARTBEAT_GRACE_S = float(
    os.environ.get("OMNIVOICE_MODEL_LOAD_HEARTBEAT_GRACE_S", "30.0"))
# Cap on the EXTRA time heartbeats can buy beyond the normal execution budget.
# Without a cap, a load that heartbeats but never finishes would hold its
# worker forever. 1800s of extension ≈ a 5 GB model at ~2.5 MB/s on top of the
# 300s base — beyond that, telling the user is better than silently waiting.
MODEL_LOAD_EXTRA_TIMEOUT_S = float(
    os.environ.get("OMNIVOICE_MODEL_LOAD_TIMEOUT_S", "1800.0"))

# How long a SYNTHESIS heartbeat stays fresh. Much longer than the load grace
# on purpose: the finest progress signal a generate has is "a chunk finished",
# and one chunk of a long text on a modest GPU can legitimately take minutes
# (#1391: an RTX 2060 SUPER with 5.7 GB free). Judging that by the 30s
# sidecar-frame grace would call every slow-but-healthy render wedged, which is
# the bug. At this grace the distinction is the honest one: a job that has not
# finished a single chunk in a whole base budget really has stopped.
GENERATE_PROGRESS_GRACE_S = float(
    os.environ.get("OMNIVOICE_GENERATE_PROGRESS_GRACE_S", "300.0"))

#: thread ident -> (monotonic time of its last heartbeat, how long it stays
#: fresh). Written by report_model_load_activity() / report_generate_progress()
#: from pool-worker threads, read by the guarded waiter, cleared when the job
#: ends. Plain dict: CPython dict ops are atomic enough for a small tuple, and
#: a torn read only costs one 5s wait slice.
_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,
) -> 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
    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
        )
        # An explicit (env-set, or runtime-changed the same way tests do)
        # CPU budget is more specific than the universal override and always
        # wins for CPU dispatches — see the #1787 comment above.
        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
        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,
            ):
                # `max`, never a plain assignment: an operator who raised the
                # accelerated budget above the CPU one must not have it cut.
                base = max(base, CPU_JOB_TIMEOUT_S)
    except Exception:
        # Device probing is advisory here; the configured universal bound is
        # still safe when a platform probe is unavailable during startup.
        pass
    return base + (max(0, len(text or "") - 1200) / 40.0)


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:  # noqa: BLE001 — telemetry must never break a request
            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):  # noqa: BLE001 — cleanup only
        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).
    """
    loop = asyncio.get_running_loop()
    ex = executor if executor is not None else _get_gpu_pool()
    # Resolved at CALL time, not def time, so monkeypatching/reloading the
    # module constant reaches every call site (the old default bound at def).
    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)
    # The worker thread's ident, published by _job so the waiter can read this
    # job's model-load heartbeats (#1367). A dict, not a nonlocal: the closure
    # runs on a pool thread while the waiter reads from the event loop.
    _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:  # noqa: BLE001 — cleanup cannot hide the pool result
            logger.exception("%s abandon cleanup failed", _log_safe(what))

    def _job():
        # First thing the worker does: tell the awaiting coroutine the
        # execution clock may start. call_soon_threadsafe is the only
        # loop-safe way to touch an asyncio primitive from a pool thread.
        _ident_box["ident"] = threading.get_ident()
        try:
            loop.call_soon_threadsafe(started.set)
        except RuntimeError:
            pass  # loop already closed (caller vanished) — still run the job
        try:
            return _inner()
        finally:
            # Idents are reused by the OS; a stale heartbeat under this ident
            # must not vouch for some future job on the same thread.
            _MODEL_LOAD_ACTIVITY.pop(threading.get_ident(), None)
            with _abandon_lock:
                _abandon_state["finished"] = True
                abandoned = _abandon_state["requested"]
            if abandoned:
                _fire_abandon_callback()

    concurrent_fut = ex.submit(_job)
    fut = asyncio.wrap_future(concurrent_fut, loop=loop)

    def _abandon() -> None:
        # Keep the concurrent future so we can distinguish a job cancelled out
        # of the queue from a thread that Python cannot stop once it has begun.
        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:
        # Phase 1 — queue wait. Watch the future too, so a job that fails or is
        # cancelled while still queued resolves here instead of hanging.
        done, _pending = await asyncio.wait(
            {waiter, fut}, timeout=queue_timeout,
            return_when=asyncio.FIRST_COMPLETED,
        )
    except asyncio.CancelledError:
        # Caller went away (client disconnect). We stop awaiting the job, so
        # make sure its eventual result/exception is consumed rather than
        # logged as "Future exception was never retrieved".
        _abandon()
        fut.add_done_callback(_swallow_abandoned)
        raise
    finally:
        waiter.cancel()

    if not done:
        # Never picked up: cancel it out of the queue (a not-yet-started
        # concurrent future cancels cleanly) and report saturation, NOT a
        # too-heavy job.
        _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),
        )

    # Phase 2 — execution. The clock starts here: this job owns a worker.
    #
    # Not a single wait_for (#1367): a first-use generate on a subprocess
    # engine downloads its model inside the job, and the sidecar proves the
    # download is healthy with progress frames the backend forwards via
    # report_model_load_activity(). Sliced waiting lets the deadline extend
    # while those heartbeats stay fresh — bounded by MODEL_LOAD_EXTRA_TIMEOUT_S
    # — so a slow connection is no longer reported as too-slow hardware. A job
    # that goes SILENT still dies at the original deadline (± one slice).
    _t0 = time.monotonic()
    _soft_deadline = _t0 + timeout
    _hard_deadline = _soft_deadline + MODEL_LOAD_EXTRA_TIMEOUT_S
    _extended = False
    try:
        while True:
            _now = time.monotonic()
            if _now < _soft_deadline:
                _slice = min(_soft_deadline - _now, 5.0)
            else:
                # Soft budget exhausted. Keep waiting ONLY on the strength of a
                # fresh heartbeat from this job's worker thread — a model-load
                # progress frame, or a completed synthesis chunk. Each carries
                # its own freshness window (loads report every ~5s; chunks are
                # minutes apart on slow hardware).
                _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):
                    raise asyncio.TimeoutError()
                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,
                    )
                # Wake at the next decision point (heartbeat expiry or the
                # cap), not a fixed 5s — a fixed slice overshoots both.
                _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:
        # Caller went away mid-execution. The old wait_for cancelled the
        # wrapper itself; asyncio.wait does not, so do both halves here or the
        # eventual result is logged as "Future exception was never retrieved".
        _abandon()
        fut.add_done_callback(_swallow_abandoned)
        raise
    except asyncio.TimeoutError as timeout_exc:
        # Parity with the old wait_for semantics: cancel the asyncio wrapper;
        # the worker thread keeps going regardless. Consume whatever it
        # eventually produces.
        _abandon()
        fut.add_done_callback(_swallow_abandoned)
        # Capture the stacks BEFORE reset(): reset() replaces the executor, and
        # once the wedged thread is no longer a pool worker we can no longer
        # tell it apart from any other thread in the process.
        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


#: Frames to keep per wedged worker. Deep enough to cross the engine adapter
#: into the model's own call stack, shallow enough that a 1-worker and an
#: 8-worker host both produce a log a human will actually read.
_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)  # unwrap _ResilientGpuPool
    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:  # noqa: BLE001
        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 ""
        # Stack frames carry absolute source paths, and on a user's machine
        # those start with their home directory — i.e. their account name. This
        # log lands in backend.log, which goes into diagnostic bundles and
        # prefilled bug reports, so it must be sanitized like every other
        # surfaced text (CWE-532; CodeRabbit). core.failure.sanitize also
        # redacts HF tokens and *TOKEN*/*KEY*/*SECRET* env values, which a
        # frame's local-variable-free repr should never contain — but "should
        # never" is not a reason to log it unredacted.
        try:
            from core.failure import sanitize as _sanitize
            text = _sanitize("\n".join(blocks))
        except Exception:  # noqa: BLE001 — never lose the diagnostic to this
            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:  # noqa: BLE001 — diagnostics must never mask the timeout
        logger.exception("Could not capture GPU-pool worker stacks")
        return ""


#: Standard-library modules whose blocking primitives a wedged worker parks in.
#: Matched on the *file* of the deepest frame, so a user function that happens
#: to be named ``wait`` or ``result`` cannot be mistaken for one of these.
_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",
)

#: Blocking entry points within those modules. A thread sitting in one of these
#: is waiting on another thread, by definition — there is no slow-but-working
#: interpretation of it.
_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"  # conservative default: GPU wording if the probe fails
    device_name, vram_gb = "", 0.0
    _caps = None  # a failed probe stays None; under_provisioned_vram() reads it safely
    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:  # noqa: BLE001 — guidance must never mask the timeout
        pass
    if wedged:
        # The worker spent the whole budget parked on a lock. None of the
        # hardware advice below applies — shorter text and a lighter engine
        # cannot speed up a job that never started (#1416/#1419/#1329).
        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."
        )
    # #1226/#1222: two users on 4 GB cards were told, generically, that the GPU
    # "is VRAM-starved" — true, but it read as a transient contention problem
    # they could flush their way out of, when their card was simply too small
    # for the engine they had selected. Say so instead — but ONLY when the
    # caller passed the engine's measured floor and the host is a dedicated-
    # VRAM family below it. This function serves every GPU-pool job (reference
    # transcribe, watermarking, dub steps, CPU-only engines on a GPU host), so
    # a threshold applied without knowing whose job it is would confidently
    # misdiagnose most of them. And on MPS `vram_gb` is a unified-memory
    # heuristic (RAM/2), not a dedicated pool to compare against.
    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 pool (#1169 load, split out in #1190) ──────────────────────
# AudioSeal's generator is loaded with `AudioSeal.load_generator(...)` and
# never moved to an accelerator: `embed_watermark` is CPU work on CPU tensors.
# Running it on the GPU pool therefore reserves a *GPU* worker for a job that
# uses no VRAM — and since #1169 routed every producer (including per-chunk
# stream previews) through mark_synthetic, on an 8 GB host (exactly 1 GPU
# worker) each watermark embed serialized directly ahead of the next generate,
# doubling the effective queue depth of a streamed multi-chunk render.
# Giving it its own tiny pool removes that head-of-line blocking with no VRAM
# risk, because the work was never on the device to begin with.
_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  # type: ignore
_model_lock = asyncio.Lock()

#: Process-wide exclusion for a cold load that runs INLINE on a GPU-pool
#: worker (#1417). `_model_lock` cannot serve there — it is an asyncio.Lock
#: bound to the server loop, and that path arrives on a bootstrap loop from
#: another thread. A threading.Lock is loop-agnostic, so the two together
#: guarantee only one cold load is ever in flight whichever route reached it.
_model_load_thread_lock = threading.Lock()
_last_used = time.time()
# Idle timeout is resolved per-tick in _resolve_idle_timeout() (MM2-05) from
# prefs/env/core.config — no module-level duplicate of IDLE_TIMEOUT_SECONDS.

# ── Loading sub-stage tracker ────────────────────────────────────────
# Updated by _load_model_sync() so get_model_status() can report
# granular progress to the frontend pill.
_loading_detail: dict = {
    "sub_stage": None,   # importing | loading_weights | compiling | ready | error
    "detail": "",        # human-readable description
    "error": None,       # error message string if failed
    "progress": None,    # 0-100 percentage (None = indeterminate)
}

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  # User already set it manually
    try:
        device_name = torch.cuda.get_device_name(0).lower()
        # Only AMD GPUs need this — skip NVIDIA
        if not any(kw in device_name for kw in ("amd", "radeon", "instinct")):
            return
        # Try to read the GFX version from the device properties
        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:
            # Metadata unavailable — an UNKNOWN build, not a confirmed
            # mismatch. Remapping on a guess could push a natively-supported
            # GPU onto foreign kernels, so fail open and change nothing.
            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:
            # The remap target isn't in this build either — setting the
            # override would only change WHICH kernel is missing. Leave it
            # unset so check_device_compatibility() reports the real mismatch
            # and the CPU fallback engages.
            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

    # ── NVIDIA CUDA or AMD ROCm (both present through torch.cuda) ─────
    if family in ("cuda", "rocm"):
        _configure_rocm_if_needed(torch)
        compatible, warning = check_device_compatibility()
        if not compatible:
            logger.warning(warning)
            # #756: the GPU's compute capability isn't in this torch build's arch
            # list, so CUDA kernels can't launch ("no kernel image is available
            # for execution") — every generate would 500. Too-old (Pascal sm_61)
            # and too-new (Blackwell sm_120 on pre-cu128 wheels) both land here.
            # Fall back to CPU so the app WORKS (slowly) instead of dead-ending;
            # OMNIVOICE_FORCE_CUDA=1 overrides for users who installed a matching
            # torch and know the arch_list probe is wrong for their setup.
            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"

    # ── Intel Arc / discrete GPU via IPEX ────────────────────────────
    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"

    # ── Apple Silicon MPS ────────────────────────────────────────────
    # Checked BEFORE DirectML to mirror the probe's family-priority order
    # (cuda > rocm > xpu > mps; DirectML is not a torch family) so the loader
    # and detect_host_caps() never disagree on a host that somehow exposes both.
    if family == "mps":
        return "mps"

    # ── DirectML — universal Windows GPU (probe reports this as "cpu") ─
    # Reached only when no torch family was detected (family == "cpu"), which is
    # exactly the DirectML case — the probe classifies DirectML hosts as cpu.
    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:
            # DirectML is optional; an absent package leaves CPU available.
            pass

    # Other families need an explicitly compatible loader (e.g. NPU sidecars).
    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: on GPU architectures Triton doesn't support yet (e.g. Blackwell
    sm_120), the compiled model dies mid-generation with errors like
    "Detected that you are using FX to symbolically trace a dynamo-optimized
    function" or an AssertionError out of torch/_inductor/cudagraph_trees.py.
    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,
            )
        # Follow the chain, honoring `raise ... from None` (the eager-retry
        # path suppresses the original compile error so a genuine eager
        # failure isn't misclassified as a compile failure).
        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:
                # `from None` so a genuine eager failure (e.g. a real OOM)
                # isn't chained to — and misclassified as — the compile error.
                raise eager_exc from None

    _model.generate = _generate_with_compile_fallback


# ── FlashInfer runtime fallback (upstream k2-fsa port) ──────────────────────


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 BEFORE exposing the eager path: while the teardown
            # mutates modules, _model.generate still routes through the
            # thread-affinity wrapper, so a concurrent render queues behind
            # this call instead of racing the half-restored model (Greptile,
            # #1565 round 2). Only a fully restored model is published.
            _unapply_flashinfer(_model)
            _model.generate = orig_generate
            try:
                return orig_generate(*args, **kwargs)
            except Exception as plain_exc:
                # `from None`: a genuine standard-path failure must not be
                # chained to — and misread as — the FlashInfer error.
                raise plain_exc from None

    _model.generate = _generate_with_flashinfer_fallback


# ── #315: thread affinity for cudagraph-compiled models ─────────────────────
# `torch.compile(mode="reduce-overhead")` captures CUDA graphs, and captured
# graph state is **thread-local** (torch/_inductor/cudagraph_trees keys its
# tree manager off the capturing thread). The `_gpu_pool` runs up to
# `_GPU_WORKER_CAP` threads, so render #1 captures the graph on worker A and a
# later render dispatched to worker B replays against mismatched cudagraph
# state — silently corrupting the audio (static / slowed playback, no
# exception, so the #278 eager fallback never fires). Fix: every call into a
# cudagraph-compiled model executes on ONE dedicated thread; uncompiled
# models (CPU / MPS / Windows-no-Triton / compile-disabled) keep the full pool.

_TORCH_COMPILE_MODE = "reduce-overhead"
# Compile modes that enable CUDA graphs under the hood — these need the
# single-thread affinity below. "default" / "max-autotune-no-cudagraphs"
# would not.
_CUDAGRAPH_COMPILE_MODES = frozenset({"reduce-overhead", "max-autotune"})

# ── #2135: CUDA-graph capture needs Ampere or newer ─────────────────────────
# On a Turing T4 (sm_75) the cudagraph mode above took the whole backend
# process down on the first generate — no Python traceback, no HTTP response,
# just a dead PID (the native capture aborts below the interpreter, so neither
# the #278 eager fallback nor any `except` can see it). The graph *capture* is
# the risky part, not Inductor: dropping to the non-cudagraph "default" mode
# keeps the compiled kernels (and most of the speedup) while removing the
# crash surface. Ampere (sm_80) is the floor because that is where the app has
# actual passing evidence; anything older takes the conservative path.
_CUDAGRAPH_MIN_CAPABILITY = (8, 0)
# Escape hatch in the other direction, for operators benchmarking on old GPUs.
_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
    # Model state is a declared real-time event. Emit only at these explicit
    # lifecycle transitions; high-frequency Hugging Face byte progress updates
    # write the dict directly and remain covered by the active one-second poll.
    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")


# ── Broken-snapshot-link self-heal ───────────────────────────────────
# A sibling of the incomplete-cache class above: the blobs are FULLY
# downloaded, but the snapshots/<rev>/ entries pointing at them are dangling
# symlinks (0 KB) or zero-byte stand-ins — blob-naming mismatches between
# download modes, interrupted renames, or antivirus interference all produce
# this state (reported on Windows, where the NTFS links show as 0 KB, but the
# heal is generic). os.path.isfile() on a dangling link is False, so
# transformers raises the same "does not appear to have a file named …"
# signature even though the bytes are on disk. The resume repair below can't
# fix it (snapshot_download may trust/short-circuit on the existing broken
# entry), so rung 0 of the recovery ladder deletes exactly the broken entries
# and restores them — see services.hf_cache_repair.

# Repos this process already attempted the link self-heal for — the retry
# after a repair may only happen ONCE per repo per process, so a cache that
# stays broken can't loop repair↔retry.
_LINK_REPAIR_ATTEMPTED: set[str] = set()

#: Repos whose weights we have already force-re-downloaded this process
#: (#1406). Without it, a shard that stays unparseable after a full re-fetch
#: would pull the whole model again on EVERY generate request — one bad file
#: turning into unbounded traffic. Same once-per-repo-per-process contract as
#: the snapshot-link repair above (CodeRabbit).
_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  # a local-directory checkpoint doesn't use the hub cache
    try:
        from services.hf_cache_repair import repair_repo_cache
        summary = repair_repo_cache(checkpoint)
    except Exception as repair_err:  # repair must never break the ladder
        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 ""


# Why the LAST _repair_model_cache run failed ("" when it succeeded / hasn't
# run). #886: the "could not be auto-repaired" message used to drop the cause
# entirely, so a mirror outage, offline mode, or a full disk all read the same.
_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:  # pragma: no cover - huggingface_hub is a hard dep
        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,
    }
    # Explicit endpoint (HF_ENDPOINT / pref) wins; otherwise the automatic
    # endpoint selection's cached pick applies (services.endpoint_race).
    try:
        from services import endpoint_race
        endpoint = endpoint_race.effective_endpoint()
    except Exception:  # endpoint resolution must never break the repair
        endpoint = os.environ.get("HF_ENDPOINT")
    if endpoint:
        dl_kwargs["endpoint"] = endpoint
    if force:
        # Replace present-but-corrupt blobs that resume would trust by size.
        dl_kwargs["force_download"] = True
    if os.name == "nt":
        # Match the install path (download.py): avoid symlinks on Windows.
        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)  # nosec B615 -- installed immutable revision
        except TypeError:
            # Older/newer huggingface_hub may not accept local_dir_use_symlinks
            # on a cache-only call — retry without the optional knob.
            dl_kwargs.pop("local_dir_use_symlinks", None)
            snapshot_download(**dl_kwargs)  # nosec B615 -- installed immutable revision

    # Bounded retries (#739): an incomplete cache *is* an interrupted download, so
    # a single transient blip mid-repair shouldn't drop the user back to a manual
    # delete-and-reinstall. snapshot_download resumes between attempts (present,
    # correctly-sized blobs are skipped by hash), so each retry continues where
    # the last left off — cheap and idempotent. Counts/backoff are env-tunable
    # for restricted networks and kept fast (backoff=0) in tests.
    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:
                # Endpoint failover (auto mode only, once per repo per
                # process — same guard pattern as the snapshot-link rung): a
                # network-classified repair failure re-races the endpoints so
                # the next attempt retries on the winner instead of burning
                # every retry on a dead host. Explicit user endpoints are
                # never switched.
                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:  # failover must never break the ladder
                    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":
        # Test-suite sentinel (tests/conftest.py sets OMNIVOICE_MODEL=test):
        # return it verbatim. Self-healing it to the real default — "test"
        # is a bare token like the #693 engine-id leak — would hand every
        # app-booting test the real 2.3 GB k2-fsa/OmniVoice checkpoint,
        # which is exactly the download the sentinel exists to prevent. A
        # real load against "test" fails fast with a clear HF error instead.
        return checkpoint
    # Honor a HF repo id (org/repo) or an EXPLICIT local path (absolute, or with
    # a path separator). A bare token like "omnivoice" must NOT be treated as a
    # local dir even if a cwd-relative folder happens to share its name — that
    # is exactly the engine-id leak (#693), so self-heal to the default.
    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


#: CPython's exact executor-rejection message once interpreter shutdown began
#: (concurrent/futures/thread.py). The "interpreter" word is what separates a
#: process teardown from an ordinary single-pool reset ("…after shutdown").
_INTERPRETER_SHUTDOWN_MSG = "cannot schedule new futures after interpreter shutdown"
#: Prefix shared by BOTH executor-rejection variants (interpreter + plain pool).
_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.
    """


# Flipped by main.py's lifespan: set the moment graceful shutdown starts,
# cleared on startup (in-process relaunches: TestClient boots, the
# --health-check thread). While set, executor-rejection errors during a load
# — including the plain single-pool "…after shutdown" variant our own
# _reset_gpu_pool() causes — are classified as a benign cancelled-load
# instead of the #589-class real fault.
_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():
        # The graceful shutdown began before this queued load got a worker
        # (e.g. 1-worker MPS pool with a capture-ASR warmup ahead of it).
        # Don't start a multi-GB import/load the process is about to abandon
        # — bail before torch is even imported (#1174).
        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

    # Register a listener that updates _loading_detail with real-time
    # download/weight-loading percentages from hf_hub_download tqdm bars.
    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)  # cap at 99 until fully done
            detail = _loading_detail.get("detail", "")
            # Append percentage to the existing detail label
            base = detail.split(" —")[0].split(" (")[0]  # strip old suffix
            _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:
                # Already re-fetched this repo once this process and it is
                # still unparseable. Re-downloading again would be the same
                # gigabytes for the same result, once per generate request.
                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:
            # #352 / #581: a truncated HF cache surfaces here as "does not
            # appear to have a file named pytorch_model.bin or
            # model.safetensors". Instead of dead-ending the user with a
            # manual delete-and-reinstall instruction, try to self-repair: an
            # interrupted download leaves the cache missing only some files,
            # and snapshot_download() resumes/fills exactly those (a complete
            # cache never reaches this branch, so the fast path is untouched).
            if _is_corrupt_weights_error(e):
                # Present-but-unparseable wearing an OSError (#1406) —
                # transformers wraps a tensor-library failure in one. The
                # resume ladder below is the wrong repair (it would trust the
                # bad blob), so divert before the missing-shard check drops
                # this as unrecognised and 500s.
                _model = _recover_corrupt_weights(e)
            elif not _is_incomplete_cache_error(e):
                raise
            else:
                # Rung 0: broken snapshot links — the blobs are on disk but the
                # snapshot entries don't resolve (dangling symlinks / zero-byte
                # stand-ins). Delete exactly the broken entries, restore, and
                # retry the load ONCE (guarded per repo per process). A cache
                # without broken links falls straight through to the resume
                # ladder below.
                _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:
                        # Resume-repair ran but the cache is still unusable. The usual
                        # cause beyond "repo genuinely lacks weights" is a blob that's
                        # present with the right size but corrupt — snapshot_download's
                        # resume trusts it and never re-fetches it (#739). Force a full
                        # re-download (replaces corrupt blobs) and retry once more before
                        # falling back to the manual delete-and-reinstall message.
                        if _is_corrupt_weights_error(e2):
                            # The resume filled the missing files, then exposed a
                            # present-but-damaged blob. A second resume would trust
                            # that blob, so switch to the forced corruption repair.
                            _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:
            # safetensors raises SafetensorError from a Rust extension and
            # torch raises UnpicklingError — neither is an OSError, so the
            # ladder above never saw them and the load 500'd with a raw
            # "Error while deserializing header: header too large" (#1406).
            # Anything that is not this class re-raises untouched, so no
            # unrelated failure is swallowed by the broad clause.
            if not _is_corrupt_weights_error(e_corrupt):
                raise
            _model = _recover_corrupt_weights(e_corrupt)

        if preload_asr:
            # Keep ASR outside `from_pretrained`: if its separate HF cache is
            # corrupt, it must never be mistaken for the TTS checkpoint and
            # trigger a second multi-GB TTS load/re-download (CodeRabbit).
            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 opt-in (upstream k2-fsa port): packed CFG attention +
        # fused kernels, ~2x on upstream's benchmarks. Applied INSTEAD of
        # torch.compile — both rewrite the llm's execution and they do not
        # compose. Best-effort: any apply failure latches the session off and
        # the standard path continues untouched.
        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

                    # Captured BEFORE apply so unapply (either the failure
                    # branch below or the generate-time fallback) restores
                    # the true prior implementation.
                    _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:  # noqa: BLE001 — perf opt, never fatal
                    mark_flashinfer_runtime_failure(
                        f"{type(fi_exc).__name__}: {fi_exc}"
                    )
                    # apply_flashinfer mutates the model as it goes — a
                    # failure partway leaves half-patched modules that would
                    # crash the next render (Greptile, #1565). Restore fully.
                    _unapply_flashinfer(_model)
                else:
                    flashinfer_applied = True
                    _install_flashinfer_fallback(_model)
                    # BOTH modes pin inference to one thread. Graph mode for
                    # the #315 reason (captured CUDA-graph state is
                    # thread-local); eager mode because the FlashInfer
                    # attention wrapper and packed position ids are planned
                    # per generation in module state — two _gpu_pool workers
                    # interleaving plan() and run() would corrupt each
                    # other's layout (CodeRabbit/Greptile, #1565).
                    _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:
            # plan-02 (#65): gate on Triton availability (+ user setting), not
            # just device==cuda. Triton has no Windows wheel, so the old
            # cuda-only check OOM'd on Windows+CUDA; should_torch_compile()
            # falls back to eager there.
            from services.engine_env import should_torch_compile

            if not flashinfer_applied and should_torch_compile(device):
                _set_loading("compiling", "Compiling model (torch.compile)…")
                # #2135: resolved per-GPU — pre-Ampere drops to the
                # non-cudagraph mode rather than risking a native abort.
                compile_mode = _resolve_compile_mode()
                try:
                    _model.llm = torch.compile(_model.llm, mode=compile_mode)
                except Exception as compile_exc:
                    # #278: compile is an optimization, never a point of
                    # failure — keep the eager model and remember the failure
                    # so later loads this session skip compile up front.
                    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:
                    # Compilation is lazy: Dynamo/Inductor/Triton can still
                    # blow up on the first *forward* (e.g. unsupported new GPU
                    # archs, #278). Wrap generate so that falls back to eager
                    # instead of failing the generation.
                    _install_compile_fallback(_model)
                    if compile_mode in _CUDAGRAPH_COMPILE_MODES:
                        # #315: reduce-overhead uses CUDA graphs, whose
                        # captured state is thread-local. Pin all inference to
                        # one dedicated thread so a later render dispatched to
                        # a different _gpu_pool worker can't replay a graph it
                        # didn't capture (static / slowed audio from the 2nd
                        # render onward).
                        _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)

        # Bind status identity to the object that actually finished loading.
        # Resolving preferences later can name a newly-selected checkpoint
        # while the previous one is still resident, and process-global load
        # metadata can be overwritten by a loader that completed after its
        # caller timed out. Instance metadata keeps /model/status honest.
        try:
            setattr(_model, "_voicestudio_checkpoint", checkpoint)
            setattr(
                _model,
                "_voicestudio_loaded_at",
                time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            )
        except Exception:
            # OmniVoice is an ordinary nn.Module and accepts attributes, but
            # a future slotted/proxied model must still be usable. Status falls
            # back to the effective configured checkpoint below.
            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:
        # A model load interrupted by *interpreter/process shutdown* is not a
        # real fault — the backend is on its way out (uvicorn stopping, a failed
        # port bind, or the user closing the app mid-load). transformers
        # materializes weights in its OWN thread pool, which raises "cannot
        # schedule new futures after interpreter shutdown" on the way down.
        # Likewise once main.py's lifespan flipped `begin_shutdown()`, even the
        # plain single-pool rejection is benign — our own _reset_gpu_pool()
        # caused it. Convert those to ModelLoadInterruptedByShutdown (logged
        # calmly at INFO) instead of dressing an expected teardown up as a
        # crash: otherwise the backend-crash report fills with a scary
        # traceback for what is a normal shutdown, /model/status flips to a
        # phantom error, and on some shutdown paths the escaping RuntimeError
        # poisons the process exit code (#1174: SIGTERM mid-load → exit 1 →
        # the desktop shell toasts "the backend crashed").
        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
        # Surface an ACTIONABLE, sanitized error in /model/status (it's shown in
        # the first-run System Check). build_failure classifies the cause and
        # attaches a fix hint — e.g. a corrupted transformers install
        # ([Errno 2] … modeling_*.py) now says "reinstall transformers" instead
        # of an unhelpful raw path + "try restarting" — and strips the home dir.
        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:  # never let failure-formatting mask the real error
            err_msg = str(exc)
        _set_loading("error", "Model loading failed", error=err_msg)
        # #1000 class: transformers' lazy-import machinery wraps ANY disruption
        # to an inner import (including one interrupted by process teardown)
        # in a generic "Could not import module X. Are this object's
        # requirements defined correctly?" — logging only str(exc) discarded
        # the real cause in __cause__/__context__ and made a shutdown race
        # look like a broken install. exc_info surfaces the full chain.
        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:
        # Placement self-heal (#1191). The ASR offload/restore pair below is a
        # *balanced-call* contract, and any unbalanced path (abort, terminal
        # error, client disconnect) used to leave the TTS model resident on CPU
        # — where it stayed for EVERY later generation until the idle unload
        # fired, at 10-50x the latency. Verifying placement here makes the
        # contract unnecessary: a future unbalanced offload can no longer
        # strand the model, because the next generation moves it back.
        await _heal_tts_placement()
        # Free idle GPU memory before this warm generate reuses the resident
        # model. The cold-load path already evicts (_make_room_before_tts_load);
        # this closes the WARM path for every native TTS generate (/generate, WS
        # TTS, dub, batch, audiobook), not just a couple of routes. No-op on a
        # roomy machine. Off the event loop because the eviction does gc.collect
        # + cache drop + ASR teardown that can block for hundreds of ms.
        await asyncio.get_running_loop().run_in_executor(None, make_room_before_generate)
        return model

    if running_on_gpu_pool():
        # Same reasoning as _heal_tts_placement below, applied to the COLD
        # path it never covered (#1417). We are on a pool worker, reached from
        # OmniVoiceBackend._ensure_loaded(), which bootstraps a *fresh* event
        # loop with asyncio.run(). `_model_lock` is bound to the server loop,
        # so awaiting it here either raises outright:
        #
        #   RuntimeError: <asyncio.locks.Lock …> is bound to a different event loop
        #
        # (the reported 500 on /v1/audio/speech) or deadlocks, depending on
        # which loop touched the lock first.
        #
        # The load must also run INLINE, in this very thread. Going through
        # `_load_model_with_timeout()` would hand `_load_model_sync` back to
        # `_get_gpu_pool()` — the pool we are currently occupying — and MPS
        # pins that pool to a single worker, so it would wait on itself. That
        # is the same deadlock wearing a different hat (CodeRabbit, #1418).
        #
        # Exclusion comes from `_model_load_thread_lock` rather than the GPU
        # slot: holding a slot is not exclusion when the pool has more than
        # one worker, which CUDA hosts do.
        if model is None:
            with _model_load_thread_lock:
                if model is None:  # another thread loaded it while we waited
                    from core.run_sentinel import touch_activity
                    touch_activity("model_load", "omnivoice-tts")
                    # Same reclaim `_load_model_with_timeout` performs; a
                    # memory-tight machine needs it on this path too.
                    _make_room_before_tts_load()
                    model = _load_model_sync()
        return model

    async with _model_lock:
        if model is None:
            # Crash forensics (#1164): a cold TTS model load is where memory
            # exhaustion (OS OOM kill) most often lands — record that one
            # started so an unclean death is attributable by the next run.
            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:  # noqa: BLE001 -- making room must never break loading
        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)  # 0s idle = release if unleased
        except Exception:  # noqa: BLE001 -- best-effort, never blocks the caller
            logger.debug("capture-ASR pre-%s release failed", stage, exc_info=True)
        release_tts_side_caches()
        free_vram()
    except Exception:  # noqa: BLE001 -- a cleanup failure must never break the caller
        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:  # noqa: BLE001 -- a probe failure must never block a generate
        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  # noqa: PLC0415
    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  # already loaded

    # On MPS the configured ``omnivoice`` id resolves to a crash-isolated
    # sidecar. Warming the native singleton here would put the same fatal MPS
    # allocator risk back into the API process before the isolated engine is
    # ever asked to synthesize.
    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:  # noqa: BLE001 -- preload selection must stay best-effort
        logger.debug("effective TTS preload selection failed", exc_info=True)

    # A machine lending its GPU has no local user to warm the model FOR. This
    # preload exists to make the first /generate feel instant for the person
    # sitting in front of the app; on a headless node there is nobody sitting
    # there, so it is several GB of VRAM held from boot against a request that
    # may never come — and the idle sweep cannot reclaim it, because the sweep
    # owns the worker executor's engines and this is the default local model.
    # Observed on hardware: a node that had run nothing still sat at 2.4 GB.
    #
    # A machine that is BOTH a desktop app and a worker keeps the warm-up:
    # there is a real user there, and the whole point stands.
    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:
        # Warm-up is gated on LOCAL availability only — never a Hub API
        # probe. The old `model_info(checkpoint)` probe proved the repo
        # exists on huggingface.co, NOT that this machine has it installed,
        # so on any networked machine with an uninstalled model (fresh
        # install, empty-cache CI/test run) every app boot silently pulled
        # the multi-GB checkpoint in a background thread the moment lifespan
        # started — violating this function's "if models aren't installed
        # yet, silently exits" contract. The cache-only check also never
        # constructs an HTTP session, so the #959 class (broken
        # ALL_PROXY/HTTPS_PROXY=socks5:// env raising at client
        # construction) can't false-negative it, and startup stays free of
        # network calls (local-first). Uses the same resolver as the load
        # path (#693) so a leaked engine id can't skew the probe.
        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:
        # Expected teardown (#1174): the backend was shut down while the
        # preload was still loading weights. Info, no traceback — a WARNING
        # with a stack here is exactly the crash-shaped noise the
        # classification exists to prevent.
        logger.info("Model preload stopped: shutdown during load — benign.")
    except Exception as e:
        # See the matching exc_info note on the _load_model_sync handler above
        # (#1000 class) — the full chain, not just str(e), is what actually
        # distinguishes a real dependency problem from a shutdown-interrupted
        # import.
        logger.warning("Model preload failed (non-fatal): %s", e, exc_info=e)
        # Non-fatal must not mean invisible (#1415). A broken dependency in the
        # model's import chain fails here and nowhere else until the user tries
        # to generate — so the app starts clean, reports itself healthy, and
        # simply produces nothing, which is how the reporter's environment
        # looked. Record it on the status the UI already reads, with the
        # classified remedy attached; the next successful load clears it.
        try:
            from core.failure import build_failure

            from core.failure import describe_exception

            # The whole chain, not just the surface: transformers reports a
            # broken dependency as a lazy-attribute error and keeps the real
            # cause in __cause__, so classifying the outermost message alone
            # loses the only part that names a remedy.
            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:  # noqa: BLE001 — never lose the warning to this
            # NOT str(e): the whole point of build_failure is that it sanitizes,
            # and an exception message routinely carries absolute paths — i.e.
            # the user's account name — which this string is about to publish
            # through /model/status (CWE-532; CodeRabbit). A fixed message that
            # points at the log beats leaking one into the API.
            detail = (
                "The TTS model could not be loaded. Settings → Logs → Backend "
                "has the full error."
            )
        # `sub_stage` is a public API enum and the frontend keys failure state
        # off `error`. Keep the human-readable word "failed" in the detail,
        # not in the state machine (#1695).
        _set_loading("error", detail, error=detail)

def get_model_status():
    is_loaded = model is not None
    # asyncio.Lock exposes .locked() on all supported Python versions; wrap in try for safety.
    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:
                # Status is a recovery surface. A broken preferences layer
                # must not turn a resident-model query into a 500.
                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
    # Attach sub-stage detail only while it describes the current resident/load
    # state, or when a failure must remain actionable. A completed model can be
    # unloaded while the last successful "ready" detail remains in memory.
    # Publishing that stale detail alongside status=idle/loaded=false gives
    # clients two contradictory readiness states.
    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()
        # The capture/dictation ASR was never idle-released — so once a user
        # dictated, its model stayed resident for the life of the process while
        # the TTS model dutifully freed its 3.8 GB. On a 16 GB Mac that left the
        # backend sitting at ~6.2 GB idle, which is what tipped it into the
        # memory pressure that gets it killed mid-generate (#1076/#1092/#1093/
        # #1101). Give it the same bargain the TTS model already makes. Held
        # off while a live dictation stream has a lease, so nothing is unloaded
        # mid-sentence.
        try:
            from services.asr_backend import release_idle_capture_backend

            if release_idle_capture_backend(idle_timeout):
                free_vram()
        except Exception:  # noqa: BLE001 — the reaper must never kill idle_worker
            logger.warning("idle capture-ASR release failed", exc_info=True)
        # Same bargain for the AudioSeal watermark models, which loaded on the
        # first embed and were never released. Deliberately only here and not
        # in the make-room paths: watermarking runs immediately *after* a
        # generate, so evicting it just before one would only buy a reload.
        try:
            from services.watermark import release_idle_models

            release_idle_models(idle_timeout)
        except Exception:  # noqa: BLE001 — the reaper must never kill idle_worker
            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:  # noqa: BLE001
        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:  # noqa: BLE001 — freeing memory must never raise
        logger.debug("clearing cuBLAS workspaces failed", exc_info=True)


def free_vram():
    """Release cached GPU memory on any accelerator (CUDA, MPS, XPU)."""
    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()


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
    return False



# Free RAM below which the TTS model is released before ASR loads on a
# unified-memory machine. WhisperX large-v3 needs ~3 GB plus VAD and overhead,
# so a box with less than this much headroom cannot hold both — and on a Mac the
# loser is the whole backend process (the OS kills it). Tunable for bigger boxes.
_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  # plenty of room — keep the model warm, pay no reload
        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:  # noqa: BLE001
        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():
        # UNIFIED MEMORY (Apple Silicon / CPU). Moving the model "to CPU" frees
        # nothing here — it is the same physical RAM — which is why this used to
        # bail out entirely. But the conclusion was wrong: the fix on unified
        # memory isn't to MOVE the model, it's to RELEASE it.
        #
        # Holding the ~3.8 GB TTS model resident while WhisperX large-v3 (~3 GB)
        # loads on top of it is what gets the backend OOM-killed mid-dub on a
        # 16 GB Mac (#1119) — the transcribe stream just dies. Unload it and the
        # room is real. get_model() lazily reloads on the next TTS use, so the
        # only cost is that reload, and only when memory was actually tight.
        _offload_unified_memory()
        return
    try:
        # Check if there's enough free VRAM to skip offloading
        if torch.cuda.is_available():
            free_mem = torch.cuda.mem_get_info()[0]
            if free_mem > 8 * 1024 ** 3:  # > 8 GB free → skip offload
                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():
        # Nothing to restore on unified memory: offload UNLOADED the model, and
        # get_model() reloads it lazily on the next TTS call. Reloading it here
        # would just re-occupy the RAM we freed, right when the dub still has
        # translation and synthesis ahead of it.
        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:  # noqa: BLE001 — a probe must never break generation
        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:  # noqa: BLE001
        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():
        # Unified memory / CPU-only: the offload RELEASES the model rather than
        # moving it, and CPU is the legitimate home here. Nothing to heal.
        return None
    try:
        target = get_best_device()
    except Exception:  # noqa: BLE001
        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:  # noqa: BLE001
        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():
        # Reached from a GPU-pool thread — OmniVoiceBackend._ensure_loaded()
        # bootstraps a fresh loop with asyncio.run(get_model()) from inside
        # generate(). We already hold the GPU slot, so we already have the
        # exclusion the move needs; awaiting our own pool (or the model lock
        # held by the loop that is waiting on us) would deadlock. Move inline.
        ensure_tts_on_device()
        return
    async with _model_lock:
        if _stranded_tts_target() is None:
            return  # another caller healed it while we waited
        try:
            await asyncio.get_running_loop().run_in_executor(
                _get_gpu_pool(), ensure_tts_on_device
            )
        except Exception as e:  # noqa: BLE001
            logger.warning("TTS placement self-heal could not run: %s", e)

_diar_pipeline = None

# Sentinel error classes used by callers (dub_core) to decide whether to
# emit a structured SSE warning with a docs deeplink. Kept as module-level
# constants so tests can pin them — they cross the SSE wire and the
# frontend's errorDocsMap classifies on the same strings.
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

    # Phase 1 AUTH-01: 3-source resolver (App → Env → HF-CLI). Per
    # Pitfall #1 in 01-RESEARCH.md — exactly one place in the backend
    # reads HF tokens, and that place is `token_resolver.resolve()`.
    from services import token_resolver
    resolved = token_resolver.resolve()
    # Access is checked during explicit installation. An already-installed
    # local bundle remains usable after a token expires or is removed.
    hf_token = resolved.token if resolved else False
    try:
        torch = _lazy_torch()
        _ensure_pyannote_hf_token_compat()  # #167: use_auth_token -> token
        # PyTorch 2.6 flipped torch.load's default to weights_only=True, whose
        # secure unpickler rejects the pyannote checkpoint's metadata globals
        # (torch_version.TorchVersion, omegaconf nodes, …) — surfacing as
        # "Weights only load failed / Unsupported global" and breaking
        # diarization on torch>=2.6 even after the license is accepted (#270).
        # Reuse the exact allowlist the WhisperX VAD load registers so the
        # secure load path succeeds; it is idempotent and per-process.
        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()
        # Pyannote supports CUDA and CPU; route XPU/DirectML to CPU
        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)
        # Without a token, a missing local bundle means the user must connect
        # Hugging Face before the explicit install can run. Once installed,
        # local_pipeline_config succeeds and diarisation remains fully local.
        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