llmgw/internal/agent/discovery/gpu.go

package discovery

import (
	"encoding/json"
	"fmt"
	"math"
	"regexp"
	"slices"
	"strconv"
	"strings"
)

// GPU is one accelerator on this node.
type GPU struct {
	Index int    `json:"index"`
	UUID  string `json:"uuid"`
	Name  string `json:"name"`
	// MemoryTotalMB is the total framebuffer size.
	MemoryTotalMB int `json:"memory_total_mb"`
	// MemoryUsedMB is what is already allocated. Non-zero on an idle node means
	// something other than our engines is using the card (plan §5.1).
	MemoryUsedMB int `json:"memory_used_mb"`
	// ComputeCapability gates quantization choices: an FP8 profile must not
	// land on a card below 8.9 (plan §6.2).
	ComputeCapability string `json:"compute_capability"`
	NUMANode          int    `json:"numa_node"`
	PCIBusID          string `json:"pci_bus_id,omitempty"`
	DriverVersion     string `json:"driver_version,omitempty"`
	CUDAVersion       string `json:"cuda_version,omitempty"`
	// NVLinkPeers are the indices of GPUs reachable over NVLink. Within a node
	// a TP group should stay inside one NVLink domain.
	NVLinkPeers []int `json:"nvlink_peers,omitempty"`
	// Health carries anything that should keep the GPU out of placement.
	Health GPUHealth `json:"health"`
}

// GPUHealth records conditions that make a GPU unsafe to schedule onto.
type GPUHealth struct {
	// ECCErrors is the count of uncorrected ECC errors.
	ECCErrors int `json:"ecc_errors"`
	// TemperatureC is the current die temperature.
	TemperatureC int `json:"temperature_c"`
	// Problems lists human-readable reasons to avoid this GPU.
	Problems []string `json:"problems,omitempty"`
}

// OK reports whether the GPU is fit for placement.
func (h GPUHealth) OK() bool { return len(h.Problems) == 0 }

// nonAlnum collapses everything that is not a letter or digit.
var nonAlnum = regexp.MustCompile(`[^a-z0-9]+`)

// Class derives the GPU class used for placement: same model, same memory.
//
// Two cards only belong in the same tensor-parallel group if they are the same
// class, so the derivation has to be stable across driver versions and naming
// quirks ("NVIDIA H100 80GB HBM3" and "H100-SXM5-80GB" are the same class).
func (g GPU) Class() string {
	name := strings.ToLower(g.Name)
	name = strings.TrimPrefix(name, "nvidia ")
	name = strings.TrimPrefix(name, "tesla ")

	// Keep the model token: the first word that contains a digit, which is how
	// NVIDIA names every accelerator (h100, a100, l40s, rtx4090).
	var model string
	for _, tok := range nonAlnum.Split(name, -1) {
		if tok == "" || !strings.ContainsFunc(tok, func(r rune) bool { return r >= '0' && r <= '9' }) {
			continue
		}
		// Skip memory and form-factor tokens, which are not the model.
		if strings.HasSuffix(tok, "gb") || tok == "80gb" || tok == "sxm4" || tok == "sxm5" || tok == "pcie" {
			continue
		}
		model = tok
		break
	}
	if model == "" {
		model = nonAlnum.ReplaceAllString(name, "-")
	}

	// Round to the nearest GiB: vendors report 81559 MiB for an "80GB" card.
	gb := int(math.Round(float64(g.MemoryTotalMB) / 1024))
	return fmt.Sprintf("%s-%dg", model, gb)
}

// MarshalJSON adds the derived class, so a consumer of the JSON does not have
// to reimplement the derivation to know which GPUs are interchangeable.
func (g GPU) MarshalJSON() ([]byte, error) {
	type alias GPU // prevents recursing back into this method
	return json.Marshal(struct {
		alias
		Class string `json:"class"`
	}{alias(g), g.Class()})
}

// Inventory is everything discovery found on this node.
type Inventory struct {
	GPUs []GPU `json:"gpus"`
	HCAs []HCA `json:"hcas"`
	// Source names the backend that produced the GPU list, so a consumer can
	// tell "no GPUs" from "discovery did not work".
	Source string `json:"source"`
}

// Classes returns the distinct GPU classes present, sorted.
func (inv Inventory) Classes() []string {
	seen := map[string]struct{}{}
	for _, g := range inv.GPUs {
		seen[g.Class()] = struct{}{}
	}
	out := make([]string, 0, len(seen))
	for c := range seen {
		out = append(out, c)
	}
	slices.Sort(out)
	return out
}

// parseCSVInt reads an integer from nvidia-smi's CSV output, tolerating the
// "[N/A]" and "Not Supported" values it emits for unsupported queries.
func parseCSVInt(s string) (int, bool) {
	s = strings.TrimSpace(s)
	s = strings.TrimSuffix(s, " MiB")
	s = strings.TrimSuffix(s, " W")
	if s == "" || strings.HasPrefix(s, "[") || strings.EqualFold(s, "Not Supported") {
		return 0, false
	}
	v, err := strconv.Atoi(s)
	if err != nil {
		return 0, false
	}
	return v, true
}