stackpilot/backend/services/gpu_service.py

"""GPU detection and Compose YAML injection helpers.

Detection sources:
  * NVIDIA  — `nvidia-smi` query (when the toolkit is present)
  * AMD/Intel — scan /dev/dri/renderD* + /sys/class/drm/*/device/vendor
"""
from __future__ import annotations

import glob
import os
import subprocess
from dataclasses import asdict, dataclass
from typing import Literal, Optional

Vendor = Literal["nvidia", "amd", "intel"]

# PCI vendor IDs found in /sys/class/drm/*/device/vendor
_PCI_VENDORS = {
    "0x10de": "nvidia",
    "0x1002": "amd",
    "0x8086": "intel",
}


@dataclass
class GPUInfo:
    vendor: Vendor
    index: int
    name: str
    uuid: Optional[str] = None
    device_path: Optional[str] = None  # e.g. /dev/dri/renderD128
    driver: str = ""
    vram_mb: Optional[int] = None

    def to_dict(self) -> dict:
        return asdict(self)


def _detect_nvidia() -> list[GPUInfo]:
    gpus: list[GPUInfo] = []
    try:
        out = subprocess.run(
            [
                "nvidia-smi",
                "--query-gpu=index,name,uuid,memory.total,driver_version",
                "--format=csv,noheader,nounits",
            ],
            capture_output=True,
            text=True,
            timeout=8,
        )
    except (FileNotFoundError, subprocess.SubprocessError):
        return gpus
    if out.returncode != 0:
        return gpus
    for line in out.stdout.strip().splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) < 5:
            continue
        idx, name, uuid, mem, driver = parts[:5]
        try:
            index = int(idx)
        except ValueError:
            index = 0
        try:
            vram = int(float(mem))
        except ValueError:
            vram = None
        gpus.append(
            GPUInfo(
                vendor="nvidia",
                index=index,
                name=name,
                uuid=uuid,
                driver=driver,
                vram_mb=vram,
            )
        )
    return gpus


def _read(path: str) -> str:
    try:
        with open(path, "r", encoding="utf-8") as fh:
            return fh.read().strip()
    except OSError:
        return ""


def _detect_dri() -> list[GPUInfo]:
    """Detect AMD/Intel render nodes via /dev/dri + sysfs."""
    gpus: list[GPUInfo] = []
    render_nodes = sorted(glob.glob("/dev/dri/renderD*"))
    for i, node in enumerate(render_nodes):
        base = os.path.basename(node)  # renderD128
        vendor_id = _read(f"/sys/class/drm/{base}/device/vendor").lower()
        vendor = _PCI_VENDORS.get(vendor_id)
        if vendor == "nvidia":
            # NVIDIA is reported via nvidia-smi; skip its DRI node here.
            continue
        if not vendor:
            continue
        # Try to read a human-ish name.
        device_id = _read(f"/sys/class/drm/{base}/device/device")
        name = {
            "amd": "AMD GPU",
            "intel": "Intel iGPU",
        }.get(vendor, "GPU")
        if device_id:
            name = f"{name} ({device_id})"
        gpus.append(
            GPUInfo(
                vendor=vendor,  # type: ignore[arg-type]
                index=i,
                name=name,
                device_path=node,
                driver=vendor,
            )
        )
    return gpus


def detect_gpus() -> list[GPUInfo]:
    return _detect_nvidia() + _detect_dri()


# --------------------------------------------------------------------------- #
# Injection helpers — mutate a single service dict in place
# --------------------------------------------------------------------------- #


def inject_nvidia(
    service: dict,
    *,
    device_ids: Optional[list[str]] = None,
    count: Optional[int] = None,
    capabilities: Optional[list[str]] = None,
) -> dict:
    """Add an NVIDIA device reservation under deploy.resources."""
    caps = capabilities or ["gpu"]
    reservation: dict = {"driver": "nvidia", "capabilities": caps}
    if device_ids:
        reservation["device_ids"] = device_ids
    elif count is not None:
        reservation["count"] = count
    else:
        reservation["count"] = "all"

    deploy = service.setdefault("deploy", {})
    resources = deploy.setdefault("resources", {})
    reservations = resources.setdefault("reservations", {})
    devices = reservations.setdefault("devices", [])
    # Replace any existing nvidia reservation.
    devices[:] = [d for d in devices if d.get("driver") != "nvidia"]
    devices.append(reservation)
    return service


def inject_dri(
    service: dict,
    *,
    vendor: Vendor,
    add_render_group: bool = True,
    add_video_group: bool = False,
    set_libva: bool = False,
) -> dict:
    """Pass through /dev/dri and (optionally) add render/video groups.

    For Intel QSV, set LIBVA_DRIVER_NAME=iHD.
    """
    devices = service.setdefault("devices", [])
    if "/dev/dri:/dev/dri" not in devices:
        devices.append("/dev/dri:/dev/dri")

    groups = service.setdefault("group_add", [])
    if add_render_group and "render" not in groups:
        groups.append("render")
    if add_video_group and "video" not in groups:
        groups.append("video")

    if set_libva and vendor == "intel":
        _ensure_env(service, "LIBVA_DRIVER_NAME", "iHD")
    return service


def remove_gpu(service: dict) -> dict:
    """Strip GPU-related config from a service."""
    deploy = service.get("deploy", {})
    resources = deploy.get("resources", {})
    reservations = resources.get("reservations", {})
    if "devices" in reservations:
        reservations["devices"] = [
            d for d in reservations["devices"] if d.get("driver") != "nvidia"
        ]
        if not reservations["devices"]:
            reservations.pop("devices", None)
        if not reservations:
            resources.pop("reservations", None)
        if not resources:
            deploy.pop("resources", None)
        if not deploy:
            service.pop("deploy", None)
    # DRI passthrough
    if "devices" in service:
        service["devices"] = [
            d for d in service["devices"] if d != "/dev/dri:/dev/dri"
        ]
        if not service["devices"]:
            service.pop("devices", None)
    if "group_add" in service:
        service["group_add"] = [
            g for g in service["group_add"] if g not in ("render", "video")
        ]
        if not service["group_add"]:
            service.pop("group_add", None)
    return service


def _ensure_env(service: dict, key: str, value: str) -> None:
    env = service.get("environment")
    if env is None:
        service["environment"] = [f"{key}={value}"]
        return
    if isinstance(env, dict):
        env[key] = value
        return
    # list form
    prefix = f"{key}="
    env = [e for e in env if not (isinstance(e, str) and e.startswith(prefix))]
    env.append(f"{key}={value}")
    service["environment"] = env