llmgw/internal/agent/discovery/topo.go

package discovery

import (
	"context"
	"fmt"
	"os/exec"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

// Topology is what `nvidia-smi topo -m` says about how a node's GPUs and
// network adapters are wired together.
//
// Read from nvidia-smi rather than NVML so the portable, cgo-free binary knows
// it too. Without it that binary saw no NVLink at all: every GPU reported no
// peers, and a tensor-parallel group was chosen by index on machines where
// NVLink decides whether it runs at full speed.
type Topology struct {
	// NVLink is, per GPU, the GPUs it reaches over NVLink and how many links
	// are bonded to each.
	NVLink map[int]map[int]int
	// NICPath is, per GPU, how far each network adapter is, by its HCA name:
	// PIX, PXB, PHB, NODE or SYS, nearest first. A GPU should talk to the
	// fabric through the adapter nearest it; one across the socket costs a
	// trip over the processor interconnect on every transfer.
	NICPath map[int]map[string]string
	// NUMA is the NUMA node each GPU is attached to, where reported.
	NUMA map[int]int
}

// PathRank orders the relations nvidia-smi reports, nearest first. Lower is
// closer; an unknown relation ranks last.
func PathRank(rel string) int {
	switch strings.ToUpper(rel) {
	case "PIX":
		return 0
	case "PXB":
		return 1
	case "PHB":
		return 2
	case "NODE":
		return 3
	case "SYS":
		return 4
	}
	return 5
}

// DiscoverTopology runs `nvidia-smi topo -m`. A node without nvidia-smi has
// no topology to report, which is not an error.
func DiscoverTopology(ctx context.Context, binary string) (Topology, error) {
	if binary == "" {
		if _, err := exec.LookPath("nvidia-smi"); err != nil {
			return Topology{}, nil
		}
		binary = "nvidia-smi"
	}
	out, err := exec.CommandContext(ctx, binary, "topo", "-m").Output()
	if err != nil {
		return Topology{}, fmt.Errorf("running %s topo -m: %w", binary, err)
	}
	return parseTopo(string(out))
}

// ParseTopology reads `nvidia-smi topo -m` output.
func ParseTopology(out string) (Topology, error) { return parseTopo(out) }

var (
	ansi       = regexp.MustCompile(`\x1b\[[0-9;]*m`)
	nvLinks    = regexp.MustCompile(`^NV(\d+)$`)
	nicLegend  = regexp.MustCompile(`^\s*(NIC\d+)\s*:\s*(\S+)\s*$`)
	gpuColumn  = regexp.MustCompile(`^GPU(\d+)$`)
	numaColumn = "NUMA Affinity"
)

// parseTopo reads the matrix. Columns are tab-separated; the header names
// each one — GPUn, then the adapters, as NICn with a legend below on newer
// drivers or by HCA name on older ones — then CPU and NUMA affinity.
func parseTopo(out string) (Topology, error) {
	out = ansi.ReplaceAllString(out, "")
	lines := strings.Split(out, "\n")

	legend := map[string]string{}
	for _, l := range lines {
		if m := nicLegend.FindStringSubmatch(l); m != nil {
			legend[m[1]] = m[2]
		}
	}

	var header []string
	start := -1
	for i, l := range lines {
		cells := splitTabs(l)
		if len(cells) > 1 && cells[0] == "" && gpuColumn.MatchString(cells[1]) {
			header, start = cells, i
			break
		}
	}
	if start < 0 {
		return Topology{}, fmt.Errorf("no GPU matrix in nvidia-smi topo output")
	}

	t := Topology{NVLink: map[int]map[int]int{}, NICPath: map[int]map[string]string{}, NUMA: map[int]int{}}
	for _, l := range lines[start+1:] {
		cells := splitTabs(l)
		if len(cells) < 2 {
			// The matrix ends at the first blank line, before the legend.
			if strings.TrimSpace(l) == "" {
				break
			}
			continue
		}
		m := gpuColumn.FindStringSubmatch(cells[0])
		if m == nil {
			continue // an adapter's row: its view of the GPUs is the same matrix
		}
		gpu, _ := strconv.Atoi(m[1])
		for c := 1; c < len(cells) && c < len(header); c++ {
			col, v := header[c], strings.TrimSpace(cells[c])
			switch {
			case gpuColumn.MatchString(col):
				peer, _ := strconv.Atoi(gpuColumn.FindStringSubmatch(col)[1])
				if n := nvLinks.FindStringSubmatch(v); n != nil && peer != gpu {
					links, _ := strconv.Atoi(n[1])
					if t.NVLink[gpu] == nil {
						t.NVLink[gpu] = map[int]int{}
					}
					t.NVLink[gpu][peer] = links
				}
			case col == numaColumn:
				if n, err := strconv.Atoi(v); err == nil {
					t.NUMA[gpu] = n
				}
			case col == "CPU Affinity" || col == "GPU NUMA ID" || col == "":
			default:
				// An adapter: NICn named by the legend, or an HCA by name.
				hca := col
				if name, ok := legend[col]; ok {
					hca = name
				}
				if PathRank(v) < 5 {
					if t.NICPath[gpu] == nil {
						t.NICPath[gpu] = map[string]string{}
					}
					t.NICPath[gpu][hca] = strings.ToUpper(v)
				}
			}
		}
	}
	if len(t.NVLink) == 0 && len(t.NICPath) == 0 && len(t.NUMA) == 0 {
		return t, fmt.Errorf("the nvidia-smi topo matrix had no rows")
	}
	return t, nil
}

func splitTabs(l string) []string {
	l = strings.TrimRight(l, "\r")
	cells := strings.Split(l, "\t")
	for i := range cells {
		cells[i] = strings.TrimSpace(cells[i])
	}
	return cells
}

// Apply puts the topology on the GPUs it describes. What a GPU already knows
// — NVLink peers from NVML, which is the more exact source — is kept.
//
// The matrix is of whole cards. A MIG instance takes its card's NUMA node and
// adapters, and none of its NVLink: instances cannot reach each other or any
// other card over it, and peers listed for one would draw a tensor-parallel
// group onto it.
func (t Topology) Apply(gpus []GPU) {
	for i := range gpus {
		g := &gpus[i]
		card := g.CardIndex()
		if len(g.NVLinkPeers) == 0 && !g.IsMIG() {
			for peer := range t.NVLink[card] {
				g.NVLinkPeers = append(g.NVLinkPeers, peer)
			}
			sort.Ints(g.NVLinkPeers)
		}
		if n, ok := t.NUMA[card]; ok {
			g.NUMANode = n
		}
		if paths := t.NICPath[card]; len(paths) > 0 {
			g.NICPaths = paths
		}
	}
}