FabledCurator/agent/fc_agent/accel.py

"""Which accelerator each runtime actually got — reported once, at startup.

The agent has two GPU runtimes and both fall back to the CPU without raising:
torch when the driver is too old for its CUDA build, and onnxruntime (the imgutils
detector + CCIP models) when its CUDA provider cannot load its libraries. A
fallback shows up only as slower work, and nothing reported it. On 2026-09-24 the
image turned out to be running a CUDA-13 torch and onnxruntime on a CUDA-12 base
(#1451), and whether the ONNX half was on the GPU could not be answered from
anything the agent had ever logged.

Also the fix for the likeliest way the ONNX half misses: onnxruntime-gpu's CUDA
provider finds libcudart/cuBLAS/cuDNN only on the loader path, and in this image
they live in the `nvidia-*` pip packages torch installs. `preload_dlls()` (ORT
1.21+) loads them from there, so the provider resolves them by soname.

Stdlib-only at import, so the unit suite can import it — torch and onnxruntime
are imported inside the functions.
"""

from __future__ import annotations

import ctypes
import importlib
import logging
from pathlib import Path

log = logging.getLogger("fc_agent.accel")

# Filled by report(); /status carries it so the page can show it too.
LAST: dict = {}


def torch_status(imp=importlib.import_module) -> dict:
    try:
        torch = imp("torch")
    except Exception as e:
        return {"device": "unavailable", "error": str(e)}
    out = {"version": torch.__version__, "cuda_build": torch.version.cuda}
    if torch.cuda.is_available():
        out["device"] = "cuda"
        out["gpu"] = torch.cuda.get_device_name(0)
    else:
        out["device"] = "cpu"
    return out


def onnx_status(imp=importlib.import_module, load=ctypes.CDLL) -> dict:
    try:
        ort = imp("onnxruntime")
    except Exception as e:
        return {"device": "unavailable", "error": str(e)}
    out = {"version": ort.__version__, "providers": list(ort.get_available_providers())}
    if "CUDAExecutionProvider" not in out["providers"]:
        out["device"] = "cpu"
        return out
    preload = getattr(ort, "preload_dlls", None)
    if preload is not None:
        try:
            preload()
        except Exception as e:
            out["preload_error"] = str(e)
    # "Available" only means the build HAS the provider. Loading its library is
    # what resolves libcudart/cuBLAS/cuDNN — the step that fails when they are
    # missing, and the one a session would otherwise fail silently on.
    capi = Path(ort.__file__).parent / "capi"
    try:
        load(str(capi / "libonnxruntime_providers_shared.so"), mode=ctypes.RTLD_GLOBAL)
        load(str(capi / "libonnxruntime_providers_cuda.so"))
    except OSError as e:
        out["device"] = "cpu"
        out["error"] = str(e)
        return out
    # Loading proves the libraries resolve, NOT that a GPU can be used: on
    # 2026-09-24 this reported "onnx on GPU" beside torch failing cuInit with
    # "CUDA unknown error" (a driver update awaiting a reboot). Asking the CUDA
    # runtime for a device initialises the driver the provider would use.
    error = _cuda_device_error(load)
    out["device"] = "cpu" if error else "cuda"
    if error:
        out["error"] = error
    return out


def _cuda_device_error(load=ctypes.CDLL) -> str | None:
    """None when the CUDA runtime can reach a device, else why it cannot."""
    try:
        cudart = load("libcudart.so.13")
    except OSError as e:
        return str(e)
    count = ctypes.c_int(0)
    rc = cudart.cudaGetDeviceCount(ctypes.byref(count))
    if rc != 0:
        cudart.cudaGetErrorString.restype = ctypes.c_char_p
        return f"cudaGetDeviceCount: {cudart.cudaGetErrorString(rc).decode()} ({rc})"
    return None if count.value > 0 else "no CUDA device visible"


def report() -> dict:
    """Check both runtimes, log the result, and keep it for /status."""
    LAST.clear()
    LAST.update(torch=torch_status(), onnx=onnx_status())
    for name, s in LAST.items():
        if s.get("device") == "cuda":
            log.info("accel: %s on GPU (%s)", name, s)
        else:
            log.warning("accel: %s is NOT on the GPU — work runs on the CPU (%s)", name, s)
    return dict(LAST)