voculinux/src/vocalinux/speech_recognition/model_ladder.py

"""Memory-aware whisper.cpp model ladder for Voculinux (ticket #20).

The default model is not one-size-fits-all on a 4 GB dGPU: the ladder picks
the rung that fits free memory at load time instead of failing halfway
through a Turbo load.

Rungs (quantized Turbo first, small floor last)::

    large-v3-turbo-q8_0 -> large-v3-turbo-q5_0 -> small.en / small

The thresholds are PROVISIONAL. They are the rule
``scripts/measure_perf_baseline.py::recommend_thresholds`` applies (measured
footprint plus the 512 MiB desktop headroom) to the transcribe peaks
recorded for ticket #19 (1132 MiB for q8_0, 846 MiB for q5_0). Ticket #19
is still open and its measurement is not recorded in this repository, so
these numbers are the calibration to replace, not evidence. A re-measure
changes the constants here, and ``tests/test_model_ladder.py`` pins them so
that change is deliberate.

Backend order follows the same readiness rule: CUDA ranks before Vulkan
only when a complete architecture-matching CUDA toolkit is installed
(``nvidia-smi`` sees a GPU, compute capability >= 7.0, and nvcc answers on
PATH or a complete toolkit root under ``CUDA_HOME``, ``CUDA_PATH`` or
``/usr/local/cuda*``); otherwise the order is Vulkan, then CPU. The order only
names backends the installed build compiled in. A driver without the toolkit is
exactly what made the old VRAM-aware recommendation unreachable.
"""

import logging
import os
import platform
import shutil
import subprocess
from collections.abc import Mapping
from pathlib import Path

logger = logging.getLogger(__name__)

DEFAULT_RUNG: str = "large-v3-turbo-q8_0"
LOW_MEMORY_RUNG: str = "large-v3-turbo-q5_0"
FLOOR_RUNG_ENGLISH: str = "small.en"
FLOOR_RUNG_MULTILINGUAL: str = "small"

LADDER_RUNGS: tuple[str, ...] = (DEFAULT_RUNG, LOW_MEMORY_RUNG)

#: MiB of headroom kept for the desktop and other GPU clients. Same value
#: the baseline protocol adds to measured footprints.
HEADROOM_MIB: float = 512.0

#: Provisional transcribe peaks (MiB) from ticket #19, not yet recorded in the repo.
_MEASURED_Q8_0_PEAK_MIB: float = 1132.0
_MEASURED_Q5_0_PEAK_MIB: float = 846.0

#: Free memory at or above this keeps the default q8_0 rung.
PREFER_Q8_0_ABOVE_MIB: float = _MEASURED_Q8_0_PEAK_MIB + HEADROOM_MIB
#: Free memory at or above this keeps the q5_0 rung; below it is the floor.
PREFER_Q5_0_ABOVE_MIB: float = _MEASURED_Q5_0_PEAK_MIB + HEADROOM_MIB

#: Minimum CUDA compute capability (major, minor) the bundled backend
#: targets. Older architectures cannot run its kernels.
MIN_CUDA_COMPUTE_CAPABILITY: tuple[int, int] = (7, 0)


#: Backends whose kernels run on a GPU, so their model must fit free VRAM on a
#: discrete card. A CPU build keeps its model in system RAM.
GPU_BACKENDS: tuple[str, ...] = ("vulkan", "cuda")

#: Environment variables that may name a CUDA toolkit root, in probe order.
_CUDA_ROOT_ENV_VARS: tuple[str, ...] = ("CUDA_HOME", "CUDA_PATH")
#: Directory whose ``cuda*`` entries are probed as toolkit roots. Tests redirect it.
_CUDA_USR_LOCAL_ROOT: Path = Path("/usr/local")
#: Toolkit ``targets/`` directory for each host machine. Voculinux publishes
#: only glibc x86_64, so any other machine has no matching toolkit target.
_CUDA_TARGET_FOR_MACHINE: dict[str, str] = {"x86_64": "x86_64-linux"}


def _floor_rung(language_is_english: bool) -> str:
    """Return the floor rung for the language: English-only pairs with ``.en``."""
    return FLOOR_RUNG_ENGLISH if language_is_english else FLOOR_RUNG_MULTILINGUAL


def select_ladder_rung(language_is_english: bool, free_mib: float) -> str:
    """Return the ladder rung that fits ``free_mib`` MiB of free memory.

    Pure function of its inputs so the policy is testable without hardware.
    """
    if free_mib >= PREFER_Q8_0_ABOVE_MIB:
        return DEFAULT_RUNG
    if free_mib >= PREFER_Q5_0_ABOVE_MIB:
        return LOW_MEMORY_RUNG
    return _floor_rung(language_is_english)


def resolve_startup_rung(
    pinned_variant: str,
    language_is_english: bool,
    free_mib: float,
) -> str:
    """Return the rung to load: a manual pin wins verbatim, otherwise the ladder.

    A manual selection is never overridden here; only unpinned loads derive
    their rung from free memory.
    """
    if pinned_variant:
        return pinned_variant
    return select_ladder_rung(language_is_english, free_mib)


def next_rung_down(rung: str, language_is_english: bool) -> str | None:
    """Return the rung one step below ``rung``, or None at the floor.

    The floor itself never steps down: a second consecutive failure is a
    real error and must surface instead of looping.
    """
    if rung == DEFAULT_RUNG:
        return LOW_MEMORY_RUNG
    if rung == LOW_MEMORY_RUNG:
        return _floor_rung(language_is_english)
    return None


def _run_probe(argv: list[str], timeout: int = 10) -> str | None:
    """Run a probe command, returning stdout or None when unavailable.

    Host binaries inherit a stripped host environment so probes keep working
    inside the AppImage bundle (see ``utils/host_process.py``).
    """
    from ..utils.host_process import host_env

    if shutil.which(argv[0]) is None:
        return None
    try:
        completed = subprocess.run(
            argv, capture_output=True, text=True, timeout=timeout, env=host_env()
        )
    except (OSError, subprocess.SubprocessError) as error:
        logger.debug("Probe %s failed: %s", argv[0], error)
        return None
    if completed.returncode != 0:
        return None
    return completed.stdout


def _parse_first_free_mib(smi_output: str) -> float | None:
    """Parse the first GPU's free memory (MiB) from nvidia-smi CSV output."""
    for line in smi_output.splitlines():
        fields = [field.strip() for field in line.split(",")]
        if len(fields) >= 1:
            try:
                return float(fields[0])
            except ValueError:
                continue
    return None


def _parse_first_compute_cap(smi_output: str) -> tuple[int, int] | None:
    """Parse the first GPU's compute capability from nvidia-smi CSV output."""
    for line in smi_output.splitlines():
        fields = [field.strip() for field in line.split(",")]
        if len(fields) >= 2:
            try:
                major, _, minor = fields[1].partition(".")
                return (int(major), int((minor or "0").strip()))
            except ValueError:
                continue
    return None


def cuda_toolkit_ready_from_probe(
    has_nvidia_gpu: bool,
    nvcc_ok: bool,
    compute_cap: tuple[int, int] | None,
) -> bool:
    """Return whether a complete architecture-matching CUDA toolkit is present.

    Pure function of its inputs: an NVIDIA GPU alone (driver only) is not a
    toolkit, and a toolkit that cannot target this architecture is not
    usable either.
    """
    if not has_nvidia_gpu or not nvcc_ok or compute_cap is None:
        return False
    return compute_cap >= MIN_CUDA_COMPUTE_CAPABILITY


def _toolkit_is_complete(root: Path, machine: str) -> bool:
    """Return whether ``root`` holds an nvcc, the runtime header and libcudart.

    ``machine`` must have a toolkit target: the libraries are looked up in that
    target's ``lib`` directory (and ``lib64`` on x86_64, where it is a link to
    the same place in the runfile layout).
    """
    target = _CUDA_TARGET_FOR_MACHINE.get(machine)
    if target is None:
        return False
    nvcc = root / "bin" / "nvcc"
    if not (nvcc.is_file() and os.access(nvcc, os.X_OK)):
        return False
    if not (root / "include" / "cuda_runtime.h").is_file():
        return False
    lib_dirs = [root / "targets" / target / "lib"]
    if machine == "x86_64":
        lib_dirs.append(root / "lib64")
    return any(lib.is_dir() and any(lib.glob("libcudart.so*")) for lib in lib_dirs)


def _cuda_toolkit_roots(environ: Mapping[str, str]) -> list[Path]:
    """Return candidate toolkit roots: CUDA_HOME, CUDA_PATH, then /usr/local/cuda*."""
    roots: list[Path] = []
    for name in _CUDA_ROOT_ENV_VARS:
        value = environ.get(name, "").strip()
        if value:
            roots.append(Path(value))
    if _CUDA_USR_LOCAL_ROOT.is_dir():
        roots.extend(
            sorted(
                (path for path in _CUDA_USR_LOCAL_ROOT.glob("cuda*") if path.is_dir()),
                reverse=True,
            )
        )
    return roots


def _toolkit_root_is_complete(environ: Mapping[str, str], machine: str) -> bool:
    """Return whether any candidate root is a complete toolkit for ``machine``."""
    checked: set[Path] = set()
    for root in _cuda_toolkit_roots(environ):
        try:
            resolved = root.resolve()
        except OSError:
            continue
        if resolved in checked:
            continue
        checked.add(resolved)
        if _toolkit_is_complete(resolved, machine):
            return True
    return False


def _cuda_toolkit_installed() -> bool:
    """Return whether a CUDA toolkit is usable: nvcc on PATH, or a complete root."""
    if _run_probe(["nvcc", "--version"]) is not None:
        return True
    return _toolkit_root_is_complete(os.environ, platform.machine())


def cuda_toolkit_ready() -> bool:
    """Probe the host for a complete architecture-matching CUDA toolkit."""
    smi = _run_probe(
        ["nvidia-smi", "--query-gpu=memory.free,compute_cap", "--format=csv,noheader,nounits"]
    )
    if not smi:
        return False
    return cuda_toolkit_ready_from_probe(
        has_nvidia_gpu=True,
        nvcc_ok=_cuda_toolkit_installed(),
        compute_cap=_parse_first_compute_cap(smi),
    )


def free_vram_mib() -> float | None:
    """Return the first NVIDIA GPU's free VRAM in MiB, or None without one."""
    smi = _run_probe(["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"])
    if not smi:
        return None
    return _parse_first_free_mib(smi)


def free_system_ram_mib() -> float | None:
    """Return free system RAM in MiB (psutil, falling back to /proc/meminfo)."""
    try:
        import psutil

        return float(psutil.virtual_memory().available) / (1024**2)
    except (ImportError, AttributeError, OSError) as error:
        logger.debug("psutil free RAM unavailable: %s", error)
    try:
        with open("/proc/meminfo", encoding="utf-8") as handle:
            for line in handle:
                if line.startswith("MemAvailable:"):
                    return float(line.split()[1]) / 1024.0
    except OSError as error:
        logger.debug("/proc/meminfo unavailable: %s", error)
    return None


def free_memory_for_backend(gpu_backend: str) -> tuple[float | None, str]:
    """Return ``(free MiB, source)`` for the memory the installed build uses.

    ``gpu_backend`` is what the installed whisper.cpp build actually has
    (``vulkan``, ``cuda`` or ``cpu``). A GPU build reads free VRAM on a
    discrete NVIDIA card; integrated graphics and CPU builds read free system
    RAM. ``source`` is ``"VRAM"`` or ``"RAM"`` so callers can say which.
    """
    if gpu_backend in GPU_BACKENDS:
        vram = free_vram_mib()
        if vram is not None:
            return vram, "VRAM"
    return free_system_ram_mib(), "RAM"


def free_memory_mib() -> float | None:
    """Return free memory in MiB: discrete-GPU VRAM when present, else RAM."""
    vram = free_vram_mib()
    if vram is not None:
        return vram
    return free_system_ram_mib()


def preferred_backend_order(build_backend: str | None = None) -> tuple[str, ...]:
    """Return the backend attempt order for this host and installed build.

    ``build_backend`` is the GPU backend the installed pywhispercpp build
    compiled in (``vulkan``, ``cuda`` or ``cpu``). The runtime can only use
    those, so the order names only them: CUDA first when a complete
    architecture-matching toolkit is installed, then Vulkan, then CPU. A
    Vulkan-only build never lists CUDA. ``None`` (build unknown) keeps the
    unfiltered order: CUDA, Vulkan, CPU with a toolkit, otherwise Vulkan, CPU.
    """
    if build_backend is None:
        if cuda_toolkit_ready():
            return ("cuda", "vulkan", "cpu")
        return ("vulkan", "cpu")
    order: list[str] = []
    if build_backend == "cuda" and cuda_toolkit_ready():
        order.append("cuda")
    if build_backend == "vulkan":
        order.append("vulkan")
    order.append("cpu")
    return tuple(order)


def describe_environment() -> str:
    """Return a one-line summary of the ladder inputs for logging."""
    vram = free_vram_mib()
    if vram is not None:
        return f"free VRAM {vram:.0f} MiB, CUDA toolkit ready: {cuda_toolkit_ready()}"
    ram = free_system_ram_mib()
    if ram is not None:
        return f"no dGPU, free RAM {ram:.0f} MiB"
    return "memory unreadable; assuming the floor"


def is_ladder_rung(variant: str) -> bool:
    """Return whether ``variant`` is a rung the loader may step between."""
    return variant in LADDER_RUNGS or variant in (FLOOR_RUNG_ENGLISH, FLOOR_RUNG_MULTILINGUAL)


__all__ = [
    "DEFAULT_RUNG",
    "LOW_MEMORY_RUNG",
    "FLOOR_RUNG_ENGLISH",
    "FLOOR_RUNG_MULTILINGUAL",
    "LADDER_RUNGS",
    "HEADROOM_MIB",
    "PREFER_Q8_0_ABOVE_MIB",
    "PREFER_Q5_0_ABOVE_MIB",
    "MIN_CUDA_COMPUTE_CAPABILITY",
    "GPU_BACKENDS",
    "select_ladder_rung",
    "resolve_startup_rung",
    "next_rung_down",
    "is_ladder_rung",
    "cuda_toolkit_ready_from_probe",
    "cuda_toolkit_ready",
    "free_memory_for_backend",
    "free_memory_mib",
    "preferred_backend_order",
    "describe_environment",
]