dashboard/app/collect/gpu.py

import glob
import re
import shutil
import subprocess

from app.sample import Sample

_name_cache: str | None = None


def _read(path: str) -> str | None:
    """Read a sysfs file, returning its stripped contents.

    Args:
        path: path under /sys.

    Returns:
        The file contents, or None if it cannot be read.
    """
    try:
        with open(path) as f:
            return f.read().strip()
    except OSError:
        return None


def shorten(name: str) -> str:
    """Shorten a raw GPU device name (lspci / lact) for display.

    Strips a trailing "(rev ...)" marker, then reformats by bracket
    group: a name like "Renoir [Radeon Vega Series / ...]" becomes
    "Renoir (Radeon Vega Series)"; a name with two or more groups (typical
    for unbound PCI IDs, e.g. "[1002] Device [1586]") becomes
    "first-group middle-text (last-group)"; anything else is truncated to
    50 characters.

    Args:
        name: raw device name from lspci or lact.

    Returns:
        A display-friendly name.
    """
    name = re.sub(r"\s*\(rev.*\)$", "", name).strip()
    groups = re.findall(r"\[([^\]]+)\]", name)
    if len(groups) >= 2:
        brand = groups[0]
        series = groups[-1].split(" / ")[0]
        model = name.split("]", 1)[1].split("[", 1)[0].strip()
        return f"{brand} {model} ({series})".strip()
    if len(groups) == 1:
        series = groups[0].split(" / ")[0]
        model = name.split("[", 1)[0].strip()
        return f"{model} ({series})".strip()
    return name[:50]


def _gpu_name() -> str:
    """Resolve the display GPU name, cached for the process lifetime.

    Runs `lspci` once and takes the first VGA / 3D-controller device name,
    shortened with shorten(). Falls back to "GPU" if lspci is missing or
    no matching device line is found.

    Returns:
        The display name to put on the overview card and Sample.
    """
    global _name_cache
    if _name_cache is None:
        _name_cache = "GPU"
        if shutil.which("lspci"):
            try:
                out = subprocess.run(
                    ["lspci"], capture_output=True, text=True, timeout=5, check=False
                ).stdout
                for line in out.splitlines():
                    if "VGA" in line or "3D controller" in line:
                        _name_cache = shorten(line.split(":", 2)[-1].strip())
                        break
            except (OSError, subprocess.SubprocessError):
                pass
    return _name_cache


def _amd(s: Sample) -> bool:
    """Fill GPU fields from AMD sysfs (amdgpu driver).

    Reads gpu_busy_percent, mem_info_vram_used/total, and hwmon
    temp1_input (millidegrees) from each /sys/class/drm/card*/device.
    Busy percent is averaged across cards, VRAM summed, temperature is the
    hottest card. The display name comes from _gpu_name().

    Args:
        s: sample to fill.

    Returns:
        True if at least one card reported a busy percent, else False
        (leaving s untouched).
    """
    devices = sorted(glob.glob("/sys/class/drm/card[0-9]*/device/gpu_busy_percent"))
    if not devices:
        return False
    busy_sum = 0
    count = 0
    vram_used = 0
    vram_total = 0
    temps: list[float] = []
    for busy_path in devices:
        dev = busy_path.rsplit("/", 1)[0]
        try:
            busy_sum += int(_read(busy_path) or 0)
            count += 1
        except ValueError:
            continue
        vram_used += int(_read(f"{dev}/mem_info_vram_used") or 0)
        vram_total += int(_read(f"{dev}/mem_info_vram_total") or 0)
        for hwmon in glob.glob(f"{dev}/hwmon/hwmon*"):
            t = _read(f"{hwmon}/temp1_input")
            if t:
                try:
                    temps.append(int(t) / 1000.0)
                except ValueError:
                    pass
    if count == 0:
        return False
    s.gpu = round(busy_sum / count, 1)
    s.vram_used = vram_used
    s.vram_total = vram_total
    s.vram_pct = round(vram_used / vram_total * 100, 1) if vram_total else None
    s.gpu_temp = max(temps) if temps else None
    s.gpu_name = _gpu_name()
    return True


def _nvidia(s: Sample) -> bool:
    """Fill GPU fields by querying nvidia-smi.

    Runs `nvidia-smi --query-gpu=...` (5 s timeout) and parses the
    CSV: busy percent averaged across GPUs, VRAM summed (MiB converted to
    bytes), temperature the hottest GPU, name from the first line.

    Args:
        s: sample to fill.

    Returns:
        True if nvidia-smi exists and returned usable data, else False.
    """
    if not shutil.which("nvidia-smi"):
        return False
    try:
        out = subprocess.run(
            [
                "nvidia-smi",
                "--query-gpu=utilization.gpu,memory.used,memory.total,temperature.gpu,name",
                "--format=csv,noheader,nounits",
            ],
            capture_output=True,
            text=True,
            timeout=5,
            check=True,
        ).stdout
    except (OSError, subprocess.SubprocessError):
        return False
    lines = [l for l in out.splitlines() if l.strip()]
    if not lines:
        return False
    busy = used = total = 0
    temp = 0
    for line in lines:
        parts = [p.strip() for p in line.split(",")]
        try:
            busy += int(parts[0])
            used += int(parts[1])
            total += int(parts[2])
            temp = max(temp, int(parts[3]))
        except ValueError:
            continue
    name = lines[0].split(",")[-1].strip()
    vram_used = used * 1024 * 1024
    vram_total = total * 1024 * 1024
    s.gpu = round(busy / len(lines), 1)
    s.vram_used = vram_used
    s.vram_total = vram_total
    s.vram_pct = round(vram_used / vram_total * 100, 1) if vram_total else None
    s.gpu_temp = float(temp)
    s.gpu_name = name
    return True


def fill(s: Sample) -> None:
    """Fill the gpu / vram / gpu_temp / gpu_name fields of a Sample.

    Tries the AMD sysfs path first (no subprocess), then nvidia-smi.
    If neither applies, the fields keep their Sample defaults.

    Args:
        s: sample to fill.
    """
    _ = _amd(s) or _nvidia(s)