llmgw/internal/agent/discovery/gpu_smi.go

package discovery

import (
	"context"
	"encoding/csv"
	"fmt"
	"os/exec"
	"strconv"
	"strings"
)

// smiQueryFields are the columns requested from nvidia-smi, in order.
var smiQueryFields = []string{
	"index",
	"uuid",
	"name",
	"memory.total",
	"memory.used",
	"compute_cap",
	"pci.bus_id",
	"driver_version",
	"temperature.gpu",
	"ecc.errors.uncorrected.volatile.total",
	// Last, and optional. An nvidia-smi too old to know this field returns a
	// shorter row, and every column before it still lines up — which is why it
	// goes on the end rather than anywhere else, and why the parser requires
	// only the mandatory prefix.
	"utilization.gpu",
}

// mandatorySMIFields is how many columns a row must have to be usable.
//
// Everything before utilization.gpu. A driver that does not know the last field
// gives a short row, and rejecting those would make an llmgw upgrade silently
// stop discovering GPUs on an older node — a far worse outcome than not knowing
// how busy a card is.
const mandatorySMIFields = 10

// DiscoverGPUsSMI enumerates GPUs by shelling out to nvidia-smi.
//
// This is the cgo-free path (plan §3.2): the NVML bindings need cgo, which
// would force the whole binary — server included — to be built with a C
// toolchain. The query gives everything except topology, which Discover
// reads from `nvidia-smi topo -m` afterwards.
func DiscoverGPUsSMI(ctx context.Context, binary string) ([]GPU, error) {
	if binary == "" {
		// Not configured and not installed is a node without an NVIDIA card,
		// not a failed discovery. Reported as an error it also cut discovery
		// short before InfiniBand was looked at, on every CPU node that had
		// not been told by hand to skip GPUs.
		if _, err := exec.LookPath("nvidia-smi"); err != nil {
			return nil, nil
		}
		binary = "nvidia-smi"
	}
	out, err := exec.CommandContext(ctx, binary,
		"--query-gpu="+strings.Join(smiQueryFields, ","),
		"--format=csv,noheader,nounits",
	).Output()
	if err != nil {
		return nil, fmt.Errorf("running %s: %w", binary, err)
	}
	cards, err := parseSMI(string(out))
	if err != nil || len(cards) == 0 {
		return cards, err
	}
	// The query lists whole cards, MIG mode or not. A card in MIG mode is
	// replaced by its instances, which only -L names and only the full
	// report sizes; asked only when some card could be in MIG mode, since
	// the report is the slowest thing nvidia-smi does.
	if !anyMIGCapable(cards) {
		return cards, nil
	}
	listed, report := readMIG(ctx, binary)
	return expandMIG(cards, listed, report), nil
}

// anyMIGCapable reports whether any card could be in MIG mode: MIG began
// with compute capability 8.0 (the A100). The consumer cards of the same
// generations have none, but telling them apart by name is guesswork, and
// asking one costs only two more nvidia-smi runs at registration.
func anyMIGCapable(cards []GPU) bool {
	for _, c := range cards {
		major, _, _ := strings.Cut(c.ComputeCapability, ".")
		if n, err := strconv.Atoi(major); err == nil && n >= 8 {
			return true
		}
	}
	return false
}

func parseSMI(out string) ([]GPU, error) {
	r := csv.NewReader(strings.NewReader(out))
	r.TrimLeadingSpace = true
	r.FieldsPerRecord = -1 // tolerate driver versions with extra columns

	rows, err := r.ReadAll()
	if err != nil {
		return nil, fmt.Errorf("parsing nvidia-smi output: %w", err)
	}

	gpus := make([]GPU, 0, len(rows))
	for _, row := range rows {
		if len(row) < mandatorySMIFields {
			continue
		}
		idx, err := strconv.Atoi(strings.TrimSpace(row[0]))
		if err != nil {
			continue
		}

		g := GPU{
			Index:             idx,
			UUID:              strings.TrimSpace(row[1]),
			Name:              strings.TrimSpace(row[2]),
			ComputeCapability: strings.TrimSpace(row[5]),
			PCIBusID:          strings.TrimSpace(row[6]),
			DriverVersion:     strings.TrimSpace(row[7]),
			NUMANode:          -1,
		}
		g.MemoryTotalMB, _ = parseCSVInt(row[3])
		g.MemoryUsedMB, _ = parseCSVInt(row[4])
		g.Health.TemperatureC, _ = parseCSVInt(row[8])
		g.Health.ECCErrors, _ = parseCSVInt(row[9])
		if len(row) > 10 {
			g.UtilizationPercent, _ = parseCSVInt(row[10])
		}

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

// ForeignMemoryToleranceMB is how much allocated memory is written off as
// driver overhead rather than a foreign process.
const ForeignMemoryToleranceMB = 512

// foreignMemorySuffix ends the problem a card with another tenant on it is
// kept out of placement for. One phrase, so that the control plane can find
// it again and drop it when the tenant leaves.
const foreignMemorySuffix = " MiB already allocated by another process"

// ForeignMemoryProblem is the problem for mb of memory held by something
// that is not one of this agent's engines, or "" when that is too little to
// count.
func ForeignMemoryProblem(mb int) string {
	if mb <= ForeignMemoryToleranceMB {
		return ""
	}
	return fmt.Sprintf("%d%s", mb, foreignMemorySuffix)
}

// IsForeignMemoryProblem reports whether a problem is ForeignMemoryProblem's.
func IsForeignMemoryProblem(p string) bool { return strings.HasSuffix(p, foreignMemorySuffix) }

// assessHealth flags conditions that should keep a GPU out of placement.
func assessHealth(g *GPU) {
	if g.Health.ECCErrors > 0 {
		g.Health.Problems = append(g.Health.Problems,
			fmt.Sprintf("%d uncorrected ECC errors", g.Health.ECCErrors))
	}
	// Memory already in use means something else is on the card. Starting an
	// engine there would either fail to allocate or quietly evict the other
	// process's throughput.
	//
	// Only a first reading: this runs before the agent knows which processes
	// are its own, so an engine still running from before a restart reads as
	// foreign here. Every heartbeat measures it again, without the agent's
	// own engines, and the control plane replaces this with that.
	if p := ForeignMemoryProblem(g.MemoryUsedMB); p != "" {
		g.Health.Problems = append(g.Health.Problems, p)
	}
}