nvcurve/nvcurve/hal/gpu.py

"""GPU discovery and initialization."""

import contextlib
import ctypes
import logging
import sys
from typing import Any

from ..nvapi.bootstrap import query_interface
from ..nvapi.constants import FUNC
from ..nvapi.types import GpuInfo

log = logging.getLogger("nvcurve.hal.gpu")


def init_nvapi() -> None:
    """Initialize NvAPI. Must be called before any GPU operations."""
    init_fn = query_interface(FUNC["Initialize"], nargs=0)
    if not init_fn or init_fn() != 0:
        raise RuntimeError("NvAPI_Initialize failed")


def enumerate_gpus() -> tuple[ctypes.Array, int]:
    """Return (gpu_handles_array, count). Exits if no GPUs found."""
    gpus = (ctypes.c_void_p * 64)()
    ngpu = ctypes.c_int32()
    enum_fn = query_interface(FUNC["EnumPhysicalGPUs"])
    if enum_fn is None:
        raise RuntimeError("NvAPI function EnumPhysicalGPUs not available")
    enum_fn(ctypes.byref(gpus), ctypes.byref(ngpu))
    if ngpu.value == 0:
        print("No NVIDIA GPUs found")
        sys.exit(1)
    return gpus, ngpu.value


def get_gpu_name(gpu) -> str:
    """Return the full name string for a GPU handle."""
    name_buf = ctypes.create_string_buffer(256)
    fn = query_interface(FUNC["GetFullName"])
    if fn is None:
        raise RuntimeError("NvAPI function GetFullName not available")
    fn(gpu, name_buf)
    return name_buf.value.decode(errors="replace")


def discover_gpus() -> list[GpuInfo]:
    """Initialize NvAPI and NVML, and return a list of GpuInfo for all physical GPUs."""
    init_nvapi()
    gpus, count = enumerate_gpus()
    infos = []

    try:
        import pynvml as _pynvml

        _pynvml.nvmlInit()
    except Exception:
        _pynvml = None
    # Aliased as Any so attribute access is not flagged when the import failed.
    pynvml: Any = _pynvml

    for i in range(count):
        name = get_gpu_name(gpus[i])
        uuid = None
        pci_bus_id = None
        if pynvml is not None:
            try:
                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
                raw_uuid = pynvml.nvmlDeviceGetUUID(handle)
                # NVML might return bytes
                if isinstance(raw_uuid, bytes):
                    uuid = raw_uuid.decode("utf-8", errors="ignore")
                elif raw_uuid is not None:
                    uuid = str(raw_uuid)
                pci_info = pynvml.nvmlDeviceGetPciInfo(handle)
                # Parse something like "00000000:01:00.0" -> bus is 1.
                # PCI bus numbers are hex by convention (pynvml's field is an
                # int; the str/bytes branches are defensive).
                bus = pci_info.bus
                if isinstance(bus, bytes):
                    pci_bus_id = int(bus.decode("utf-8", errors="ignore"), 16)
                elif isinstance(bus, str):
                    pci_bus_id = int(bus, 16)
                else:
                    pci_bus_id = int(bus)
            except Exception as exc:
                log.debug("NVML query for GPU %d failed: %s", i, exc)
        infos.append(GpuInfo(name=name, index=i, uuid=uuid, pci_bus_id=pci_bus_id))

    if pynvml is not None:
        with contextlib.suppress(Exception):
            pynvml.nvmlShutdown()

    return infos


def get_gpu(index: int = 0):
    """Initialize NvAPI, enumerate GPUs, and return the handle for `index`.

    Also returns the GPU name as a convenience.
    Returns (handle, name).
    """
    init_nvapi()
    gpus, count = enumerate_gpus()
    if index >= count:
        print(f"GPU index {index} out of range (found {count} GPU(s))")
        sys.exit(1)
    gpu = gpus[index]
    name = get_gpu_name(gpu)
    return gpu, name