cpu_backend.py

"""Evidence that this process runs on CPU only, instead of a label that says so.

The status used to publish `gpu_enabled: False` and `n_gpu_layers: 0` as
hardcoded literals. Nothing checked anything, so if the installed wheel had CUDA
compiled in and a load site ever asked for offload, the status would still have
said "cpu-only".

This lives in its own module so a test can execute it. `app.py` pulls in gradio
and the whole Minecraft stack, which a unit test should not need to satisfy.

IMPORTANT -- why the binary query is opt-in
-------------------------------------------
`llama_cpp.llama_supports_gpu_offload()` must NOT be called before
`llama_backend_init()`. Measured on this project: calling it first raises
`OSError [WinError -1073741795]`, i.e. NTSTATUS 0xC000001D
STATUS_ILLEGAL_INSTRUCTION -- the call runs code that the CPU cannot execute
before the backend is registered. On Linux the equivalent is SIGILL, which
kills the process and cannot be caught by try/except, so an unconditional call
at import would have taken the Space down.

The library listing below needs no native call and is enough on its own: if the
wheel ships no CUDA/Vulkan/Metal/ROCm backend, it structurally cannot offload a
layer. The binary query stays available behind LLAMA_BINARY_PROBE=1 for a
deliberate check after a model has been loaded.
"""

from __future__ import annotations

import os
from pathlib import Path

# Every load site must ask for zero offloaded layers. Omitting the argument is
# not the same as passing zero: it lets the library decide.
REQUIRED_OFFLOAD_ARGUMENTS = {
    "n_gpu_layers": 0,
    "offload_kqv": False,
}

GPU_LIBRARY_TOKENS = ("cuda", "cublas", "vulkan", "metal", "hip", "rocm")
LIBRARY_SUFFIXES = (".so", ".dll", ".dylib")

BINARY_PROBE_ENABLED = os.environ.get("LLAMA_BINARY_PROBE", "0") == "1"


def backend_libraries() -> list[str]:
    """Names of the native backend libraries actually shipped with llama_cpp.

    Pure filesystem inspection: no native call, so it is safe at any time.
    """
    try:
        import llama_cpp
    except ImportError:
        return []
    base = Path(llama_cpp.__file__).resolve().parent
    return sorted(
        path.name
        for path in base.rglob("*")
        if path.is_file() and path.suffix.lower() in LIBRARY_SUFFIXES
    )


def binary_offload_support():
    """Ask the linked binary whether it can offload. Risky before backend init.

    Returns True, False, or a string explaining why the answer is unavailable.
    Never raises.
    """
    try:
        from llama_cpp import llama_cpp as _backend

        return bool(_backend.llama_supports_gpu_offload())
    except BaseException as exc:  # noqa: BLE001 - a probe must not take the process down
        return f"unavailable: {type(exc).__name__}"


def cpu_only_report(probe_binary: bool | None = None) -> dict:
    """Report whether this process can use a GPU, rather than asserting it cannot."""
    if probe_binary is None:
        probe_binary = BINARY_PROBE_ENABLED

    libraries = backend_libraries()
    cuda_libraries = sorted(
        name for name in libraries if any(token in name.lower() for token in GPU_LIBRARY_TOKENS)
    )

    report = {
        "gpu_offload_supported": binary_offload_support() if probe_binary else "not probed",
        "gpu_binary_probe_enabled": probe_binary,
        "gpu_backend_libraries": libraries,
        "cuda_libraries_present": cuda_libraries,
        "n_gpu_layers_requested": REQUIRED_OFFLOAD_ARGUMENTS["n_gpu_layers"],
    }
    # CPU-only is structural when the wheel ships no GPU backend at all, and
    # enforced when every load site additionally asks for zero layers.
    report["cpu_only"] = (
        not cuda_libraries and report["n_gpu_layers_requested"] == 0
    )
    return report


def cpu_only_summary(report: dict) -> str:
    return (
        f"cpu_only={report['cpu_only']} "
        f"gpu_libs={report['gpu_backend_libraries']} "
        f"cuda_libs={report['cuda_libraries_present'] or 'none'} "
        f"n_gpu_layers={report['n_gpu_layers_requested']} "
        f"binary_probe={report['gpu_offload_supported']}"
    )