llmgw/internal/agent/discovery/pcie.go

package discovery

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

// A card in the wrong slot: an x16 card in a slot wired x8 or x4, behind a
// riser, or in a slot of an older PCIe generation than the card's, links
// slower than it was built to, and nothing says so. A model on the card
// still decodes at the card's own memory speed, but it loads slower, a
// tensor-parallel replica over PCIe is slower, and so is anything that
// spills to the host. The driver says how the link stands and what the
// card could do; the difference is reported where the node's warnings are.

// PCIeLink is how a card is linked to its host.
type PCIeLink struct {
	// Gen is the fastest generation the card and the slot together allow —
	// not the current one, which drops while the card idles, to save power.
	// GPUGen is the card's own fastest, 0 when the driver does not say.
	Gen    int `json:"gen,omitempty"`
	GPUGen int `json:"gpu_gen,omitempty"`
	// Width is the lanes in use now, WidthMax the card's most.
	Width    int `json:"width,omitempty"`
	WidthMax int `json:"width_max,omitempty"`
}

// laneGBs is one lane's usable rate in GB/s, by generation.
var laneGBs = map[int]float64{1: 0.25, 2: 0.5, 3: 0.985, 4: 1.969, 5: 3.938, 6: 7.877}

// GBs is what a link of gen and width carries, in GB/s each way.
func pcieGBs(gen, width int) float64 { return laneGBs[gen] * float64(width) }

// Problem says how a card's link falls short of what the card could do,
// empty when it does not or the driver did not say.
func (p PCIeLink) Problem() string {
	gen := p.Gen
	if p.GPUGen > 0 && gen == 0 {
		gen = p.GPUGen
	}
	slowGen := p.GPUGen > 0 && p.Gen > 0 && p.Gen < p.GPUGen
	narrow := p.WidthMax > 0 && p.Width > 0 && p.Width < p.WidthMax
	if !slowGen && !narrow {
		return ""
	}
	could := pcieGBs(max(p.GPUGen, gen), p.WidthMax)
	has := pcieGBs(gen, p.Width)
	var why []string
	if narrow {
		why = append(why, fmt.Sprintf("x%d of its x%d lanes", p.Width, p.WidthMax))
	}
	if slowGen {
		why = append(why, fmt.Sprintf("PCIe Gen%d of its Gen%d", p.Gen, p.GPUGen))
	}
	return fmt.Sprintf("linked at %s: %.1f GB/s where it could have %.1f. "+
		"Check the slot (wired fewer lanes, or an older generation), a riser, or the BIOS's PCIe setting",
		strings.Join(why, " and "), has, could)
}

// pcieFields are asked apart from the rest, and the card's own fastest
// generation apart again: a driver that does not know a field refuses the
// whole query, and these must not cost a node its cards.
var pcieFields = [][]string{
	{"index", "pcie.link.gen.max", "pcie.link.width.current", "pcie.link.width.max", "pcie.link.gen.gpumax"},
	{"index", "pcie.link.gen.max", "pcie.link.width.current", "pcie.link.width.max"},
}

// readPCIe fills each card's link from nvidia-smi, leaving it empty when
// the driver says nothing of it.
func readPCIe(ctx context.Context, binary string, cards []GPU) {
	for _, fields := range pcieFields {
		out, err := exec.CommandContext(ctx, binary, "--query-gpu="+strings.Join(fields, ","),
			"--format=csv,noheader,nounits").Output()
		if err != nil {
			continue
		}
		links := parsePCIe(string(out))
		for i := range cards {
			if l, ok := links[cards[i].Index]; ok && cards[i].MIG == nil {
				cards[i].PCIe = l
			}
		}
		break
	}
	for i := range cards {
		if cards[i].MIG == nil {
			cards[i].PCIe = withSysfs(cards[i].PCIe, "/sys/bus/pci/devices", cards[i].PCIBusID)
		}
	}
}

// withSysfs takes the card's own lane count, and its own fastest
// generation, from the kernel: the driver's width.max is what card and slot
// allow together, so a x16 card in a slot wired x8 reads x8 of x8 there,
// while the device's max_link_width is the card's, 16.
func withSysfs(l PCIeLink, root, busID string) PCIeLink {
	dev := sysfsBusID(busID)
	if dev == "" {
		return l
	}
	read := func(name string) string {
		b, err := os.ReadFile(filepath.Join(root, dev, name)) //nolint:gosec // under /sys
		if err != nil {
			return ""
		}
		return strings.TrimSpace(string(b))
	}
	if n, err := strconv.Atoi(read("max_link_width")); err == nil && n > 0 {
		l.WidthMax = n
	}
	if n, err := strconv.Atoi(read("current_link_width")); err == nil && n > 0 {
		l.Width = n
	}
	if g := genOfSpeed(read("max_link_speed")); g > 0 && l.GPUGen == 0 {
		l.GPUGen = g
	}
	return l
}

// sysfsBusID is nvidia-smi's 00000000:3B:00.0 as the kernel names it,
// 0000:3b:00.0.
func sysfsBusID(id string) string {
	id = strings.ToLower(strings.TrimSpace(id))
	dom, rest, ok := strings.Cut(id, ":")
	if !ok || len(dom) < 4 {
		return ""
	}
	return dom[len(dom)-4:] + ":" + rest
}

// genOfSpeed is the PCIe generation of a link speed as sysfs writes it.
func genOfSpeed(s string) int {
	switch {
	case strings.HasPrefix(s, "2.5 GT/s"):
		return 1
	case strings.HasPrefix(s, "5.0 GT/s"), strings.HasPrefix(s, "5 GT/s"):
		return 2
	case strings.HasPrefix(s, "8.0 GT/s"), strings.HasPrefix(s, "8 GT/s"):
		return 3
	case strings.HasPrefix(s, "16.0 GT/s"), strings.HasPrefix(s, "16 GT/s"):
		return 4
	case strings.HasPrefix(s, "32.0 GT/s"), strings.HasPrefix(s, "32 GT/s"):
		return 5
	case strings.HasPrefix(s, "64.0 GT/s"), strings.HasPrefix(s, "64 GT/s"):
		return 6
	}
	return 0
}

func parsePCIe(out string) map[int]PCIeLink {
	r := csv.NewReader(strings.NewReader(out))
	r.TrimLeadingSpace = true
	r.FieldsPerRecord = -1
	rows, _ := r.ReadAll()
	num := func(s string) int {
		n, _ := strconv.Atoi(strings.TrimSpace(s))
		return n
	}
	links := map[int]PCIeLink{}
	for _, row := range rows {
		if len(row) < 4 {
			continue
		}
		idx, err := strconv.Atoi(strings.TrimSpace(row[0]))
		if err != nil {
			continue
		}
		l := PCIeLink{Gen: num(row[1]), Width: num(row[2]), WidthMax: num(row[3])}
		if len(row) > 4 {
			l.GPUGen = num(row[4])
		}
		links[idx] = l
	}
	return links
}