llmgw/internal/agent/discovery/gpu_nvml.go

//go:build nvml

// This file is the NVML-backed GPU discovery path, built only with the `nvml`
// tag because the bindings require cgo. Without the tag the agent falls back to
// nvidia-smi, which keeps the server side — and development machines without a
// C toolchain or a driver — building cleanly (plan §3.2).
//
// What NVML adds over nvidia-smi is NVLink topology. Nothing else exposes which
// GPUs share an NVLink domain, and that is what keeps a tensor-parallel group
// off the slow path inside a node (plan §6.2, constraint 3).

package discovery

import (
	"context"
	"fmt"
	"strings"

	"github.com/NVIDIA/go-nvml/pkg/nvml"
)

// nvmlAvailable reports that this build can use NVML.
const nvmlAvailable = true

// DiscoverGPUsNVML enumerates GPUs through NVML.
//
// It returns an error rather than falling back on its own: the caller decides
// whether to degrade to nvidia-smi, because on a node that is supposed to have
// working NVML a silent downgrade hides a broken driver installation.
func DiscoverGPUsNVML(_ context.Context) ([]GPU, error) {
	if ret := nvml.Init(); ret != nvml.SUCCESS {
		return nil, fmt.Errorf("nvml init: %s", nvml.ErrorString(ret))
	}
	defer func() { _ = nvml.Shutdown() }()

	driver, _ := nvml.SystemGetDriverVersion()
	cuda := ""
	if v, ret := nvml.SystemGetCudaDriverVersion(); ret == nvml.SUCCESS {
		// NVML reports 12040 for 12.4.
		cuda = fmt.Sprintf("%d.%d", v/1000, (v%1000)/10)
	}

	count, ret := nvml.DeviceGetCount()
	if ret != nvml.SUCCESS {
		return nil, fmt.Errorf("nvml device count: %s", nvml.ErrorString(ret))
	}

	// Bus id to index, so NVLink peers can be reported as indices. NVML gives
	// the remote end of a link as a PCI address, which is not what placement
	// works in.
	byBusID := make(map[string]int, count)
	gpus := make([]GPU, 0, count)

	for i := range count {
		dev, ret := nvml.DeviceGetHandleByIndex(i)
		if ret != nvml.SUCCESS {
			return nil, fmt.Errorf("nvml device %d: %s", i, nvml.ErrorString(ret))
		}

		g := GPU{Index: i, NUMANode: -1, DriverVersion: driver, CUDAVersion: cuda}
		g.UUID, _ = dev.GetUUID()
		g.Name, _ = dev.GetName()

		if mem, ret := dev.GetMemoryInfo(); ret == nvml.SUCCESS {
			g.MemoryTotalMB = int(mem.Total / (1 << 20))
			g.MemoryUsedMB = int(mem.Used / (1 << 20))
		}
		if major, minor, ret := dev.GetCudaComputeCapability(); ret == nvml.SUCCESS {
			g.ComputeCapability = fmt.Sprintf("%d.%d", major, minor)
		}
		if pci, ret := dev.GetPciInfo(); ret == nvml.SUCCESS {
			g.PCIBusID = busIDString(pci.BusId[:])
			byBusID[normalizeBusID(g.PCIBusID)] = i
		}
		if t, ret := dev.GetTemperature(nvml.TEMPERATURE_GPU); ret == nvml.SUCCESS {
			g.Health.TemperatureC = int(t)
		}
		// Volatile counters reset on driver reload, which is what we want: an
		// error from three months and one reboot ago is not a reason to refuse
		// a card today.
		if n, ret := dev.GetTotalEccErrors(nvml.MEMORY_ERROR_TYPE_UNCORRECTED, nvml.VOLATILE_ECC); ret == nvml.SUCCESS {
			g.Health.ECCErrors = int(n)
		}

		assessHealth(&g)
		gpus = append(gpus, g)
	}

	// A second pass, because a link's remote end may be a GPU enumerated later.
	for i := range gpus {
		dev, ret := nvml.DeviceGetHandleByIndex(gpus[i].Index)
		if ret != nvml.SUCCESS {
			continue
		}
		gpus[i].NVLinkPeers = nvlinkPeers(dev, byBusID, gpus[i].Index)
	}
	return gpus, nil
}

// nvlinkPeers returns the indices of GPUs reachable from dev over NVLink.
func nvlinkPeers(dev nvml.Device, byBusID map[string]int, self int) []int {
	seen := map[int]struct{}{}
	var out []int

	for link := range nvml.NVLINK_MAX_LINKS {
		pci, ret := dev.GetNvLinkRemotePciInfo(link)
		if ret != nvml.SUCCESS {
			// Unpopulated links return an error; that is the normal case for
			// most of the 36 possible links, so it is not worth reporting.
			continue
		}
		peer, ok := byBusID[normalizeBusID(busIDString(pci.BusId[:]))]
		if !ok || peer == self {
			// The remote end can be an NVSwitch rather than a GPU, in which
			// case it is not in the map.
			continue
		}
		if _, dup := seen[peer]; dup {
			continue
		}
		seen[peer] = struct{}{}
		out = append(out, peer)
	}
	return out
}

// busIDString turns NVML's fixed-size C char array into a Go string.
func busIDString(b []int8) string {
	buf := make([]byte, 0, len(b))
	for _, c := range b {
		if c == 0 {
			break
		}
		buf = append(buf, byte(c))
	}
	return string(buf)
}

// normalizeBusID makes bus ids comparable: NVML reports them uppercase from one
// call and lowercase from another, and sometimes with a shortened domain.
func normalizeBusID(s string) string {
	s = strings.ToLower(strings.TrimSpace(s))
	// "0000:1b:00.0" and "00000000:1B:00.0" are the same device.
	if i := strings.Index(s, ":"); i > 0 {
		s = strings.TrimLeft(s[:i], "0") + s[i:]
	}
	return strings.TrimPrefix(s, ":")
}