thoxmesh_node/gpu.py

"""Graceful GPU detection.

Every GPU question in this package routes through here, because the naive form
of the check is a crash rather than a false:

    subprocess.run(["nvidia-smi"], capture_output=True).returncode == 0

``subprocess.run`` **raises** ``FileNotFoundError`` when the binary is not on
PATH. On a CPU Colab runtime nvidia-smi is not installed, so that expression does
not evaluate to ``False`` — it terminates the process. That exact line shipped in
the Colab notebook and killed every CPU session before the node could start.

The rules this module follows:

* Absent tooling, broken tooling and a hung probe all mean **"no usable GPU"**,
  never an exception. A node that cannot see a GPU is a perfectly good CPU node.
* Detection is cached, because it shells out and is called from both startup and
  configuration.
* Nothing here imports torch. The node serves GGUF through llama.cpp and has no
  torch dependency; adding one just to ask about CUDA would be a large install
  for a question ``nvidia-smi`` already answers.
"""

from __future__ import annotations

import logging
import os
import shutil
import subprocess
from dataclasses import dataclass

logger = logging.getLogger("thox.gpu")

#: Seconds to wait for nvidia-smi. A wedged driver can hang it indefinitely.
_PROBE_TIMEOUT_SECONDS = 60

#: Layers to offload when a GPU is present and no explicit count is configured.
#: -1 means "all layers" in llama.cpp; it clamps itself to what fits in VRAM.
GPU_LAYERS_ALL = -1

_cached: "GpuInfo | None" = None


@dataclass(frozen=True)
class GpuInfo:
    """What we could determine about GPU availability."""

    available: bool
    name: str = ""
    detail: str = ""

    def describe(self) -> str:
        """One line suitable for a startup log or a notebook cell."""
        if self.available:
            return f"GPU available: {self.name or 'unknown device'}"
        return f"No usable GPU ({self.detail or 'nvidia-smi not found'}); running CPU-only"


def probe_gpu(*, use_cache: bool = True) -> GpuInfo:
    """Determine GPU availability without ever raising.

    Ordering matters: ``shutil.which`` is checked before ``subprocess.run`` so
    the missing-binary case never reaches a call that would raise. The
    try/except then covers a binary that exists but cannot execute — a driver or
    library mismatch surfaces as ``OSError`` on exec, and a wedged driver as a
    timeout.
    """
    global _cached
    if use_cache and _cached is not None:
        return _cached

    info: GpuInfo
    binary = shutil.which("nvidia-smi")
    if binary is None:
        info = GpuInfo(available=False, detail="nvidia-smi not on PATH")
    else:
        try:
            completed = subprocess.run(
                [binary, "--query-gpu=name", "--format=csv,noheader"],
                capture_output=True,
                text=True,
                timeout=_PROBE_TIMEOUT_SECONDS,
            )
        except (OSError, subprocess.SubprocessError) as exc:
            # Binary present but unusable: driver/library mismatch, or it hung.
            info = GpuInfo(available=False, detail=f"nvidia-smi failed: {exc}")
        else:
            if completed.returncode != 0:
                stderr = (completed.stderr or "").strip().splitlines()
                info = GpuInfo(
                    available=False,
                    detail=f"nvidia-smi exited {completed.returncode}"
                    + (f": {stderr[0]}" if stderr else ""),
                )
            else:
                name = (completed.stdout or "").strip().splitlines()
                info = GpuInfo(
                    available=bool(name),
                    name=name[0].strip() if name else "",
                    detail="" if name else "nvidia-smi reported no devices",
                )

    _cached = info
    logger.info("%s", info.describe())
    return info


def has_gpu() -> bool:
    """``True`` when a usable GPU was detected. Never raises."""
    return probe_gpu().available


def resolve_gpu_layers(configured: str | int | None) -> int:
    """Translate the configured layer count into a number for llama.cpp.

    ``"auto"`` (the default) offloads everything when a GPU is present and
    nothing when it is not, so the same configuration is correct on both Colab
    runtime types. An explicit number is always honoured — including a non-zero
    value on a machine where detection failed, since an operator who sets it
    knows something the probe does not.
    """
    if configured is None or configured == "":
        configured = "auto"
    if isinstance(configured, str) and configured.strip().lower() == "auto":
        return GPU_LAYERS_ALL if has_gpu() else 0
    try:
        return int(configured)
    except (TypeError, ValueError):
        logger.warning("invalid gpu layer count %r; falling back to CPU-only (0)", configured)
        return 0


def reset_cache() -> None:
    """Clear the cached probe. For tests, and after installing drivers."""
    global _cached
    _cached = None


def environment_report() -> dict[str, object]:
    """Everything the ``doctor`` command needs to explain this machine.

    Deliberately returns data rather than printing, so it can be asserted on in
    tests and rendered differently by callers.
    """
    info = probe_gpu()
    return {
        "gpu_available": info.available,
        "gpu_name": info.name,
        "gpu_detail": info.detail,
        "nvidia_smi": shutil.which("nvidia-smi"),
        "cpu_count": os.cpu_count(),
    }