llmgw/internal/agent/discovery/ib.go

// Package discovery inspects a node's GPUs, InfiniBand adapters and the
// topology between them (plan §5.1 and §5.2).
package discovery

import (
	"errors"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"slices"
	"strconv"
	"strings"
)

// SysfsRoot is the default place InfiniBand devices appear.
const SysfsRoot = "/sys/class/infiniband"

// LinkLayer distinguishes real InfiniBand from RoCE.
type LinkLayer string

const (
	// LinkInfiniBand is native InfiniBand.
	LinkInfiniBand LinkLayer = "InfiniBand"
	// LinkEthernet is RoCE — Ethernet carrying RDMA.
	LinkEthernet LinkLayer = "Ethernet"
)

// Port is one port of an HCA.
type Port struct {
	Num int `json:"num"`
	// State is the port's logical state; only ACTIVE ports carry traffic.
	State string `json:"state"`
	// PhysState is the physical link state, useful for telling "cable unplugged"
	// from "no subnet manager".
	PhysState string `json:"phys_state"`
	// RateGbps is the negotiated rate.
	RateGbps  int       `json:"rate_gbps"`
	LinkLayer LinkLayer `json:"link_layer"`
	// GIDPrefix is the subnet prefix from gids/0, normalized to lowercase hex
	// groups. It is the only automatic fabric hint available, and an unreliable
	// one: OpenSM's default is fe80:: everywhere (plan §5.2).
	GIDPrefix string `json:"gid_prefix"`
	LID       int    `json:"lid"`
	SMLID     int    `json:"sm_lid"`
	// Counters are the port's error and traffic counters, by their sysfs
	// names, as far as the driver has them. The data counters are converted
	// to bytes; sysfs counts them in four-byte words.
	Counters map[string]uint64 `json:"counters,omitempty"`
	// RoCEv2GID is the GID table index NCCL should use on a RoCE port: the
	// first RoCE v2 entry holding an IPv4 address, else the first RoCE v2
	// entry. -1 on InfiniBand, and on a RoCE port with no v2 entry. NCCL's
	// default of 0 is a RoCE v1 link-local entry on most setups, which does
	// not route, and a job on it hangs.
	RoCEv2GID int `json:"rocev2_gid"`
	// PortGUID is the port's own identifier, the interface half of its GID.
	PortGUID string `json:"port_guid,omitempty"`
	// SM is the subnet manager this port sees as master, from sminfo. Nil on
	// RoCE, which has none, and when infiniband-diags is not installed.
	SM *SMInfo `json:"sm,omitempty"`
	// Network is, on RoCE, the IP network of the RoCE v2 address: the only
	// thing that says two RoCE ports are on the same fabric.
	Network string `json:"network,omitempty"`
}

// SMInfo is a subnet manager as a port sees it.
type SMInfo struct {
	GUID     string `json:"guid"`
	LID      int    `json:"lid"`
	Priority int    `json:"priority"`
	// State is SMINFO_MASTER, SMINFO_STANDBY, ...
	State string `json:"state"`
}

// FabricKey is what identifies the physical fabric this port is on: the
// master subnet manager on InfiniBand, since every port on one fabric sees
// the same one and ports on separate fabrics cannot; the IP network on RoCE.
// Empty when neither is known, and for a port that is not up.
func (p Port) FabricKey() string {
	if !p.Active() {
		return ""
	}
	switch {
	case p.LinkLayer == LinkEthernet && p.Network != "":
		return "roce:" + p.Network
	case p.SM != nil && p.SM.GUID != "":
		return "sm:" + p.SM.GUID
	}
	return ""
}

// PortCounters are the counters worth watching, from each port's counters/
// directory. The errors are the ones that say a cable, a transceiver or a
// switch port is going bad before a job fails on it; the data counters say
// whether the fabric is carrying anything at all.
var PortCounters = []string{
	"symbol_error", "link_error_recovery", "link_downed", "port_rcv_errors",
	"port_rcv_remote_physical_errors", "port_xmit_discards",
	"local_link_integrity_errors", "excessive_buffer_overrun_errors",
	"port_xmit_data", "port_rcv_data",
}

// dataCounter reports whether a counter counts data in four-byte words.
func dataCounter(name string) bool { return name == "port_xmit_data" || name == "port_rcv_data" }

// Active reports whether the port can carry traffic.
func (p Port) Active() bool { return strings.EqualFold(p.State, "ACTIVE") }

// HCA is one host channel adapter.
type HCA struct {
	Name     string `json:"name"`
	Ports    []Port `json:"ports"`
	PCIAddr  string `json:"pci_addr,omitempty"`
	NUMANode int    `json:"numa_node"`
	// FirmwareVersion is reported by the driver; worth capturing because
	// NCCL behaviour differs between firmware revisions.
	FirmwareVersion string `json:"firmware_version,omitempty"`
	// NodeGUID identifies the adapter on its fabric: it is the name a
	// switch's view of the fabric, and every other tool, knows it by.
	NodeGUID string `json:"node_guid,omitempty"`
}

// ActivePorts returns the ports that can carry traffic.
func (h HCA) ActivePorts() []Port {
	var out []Port
	for _, p := range h.Ports {
		if p.Active() {
			out = append(out, p)
		}
	}
	return out
}

// DiscoverIB enumerates the InfiniBand adapters under root.
//
// A missing root is not an error: plenty of nodes have no IB at all, and the
// scheduler treats them as single-node islands (plan §2).
func DiscoverIB(root string) ([]HCA, error) {
	if root == "" {
		root = SysfsRoot
	}
	entries, err := os.ReadDir(root)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return nil, nil
		}
		return nil, fmt.Errorf("reading %s: %w", root, err)
	}

	var out []HCA
	for _, e := range entries {
		hca, err := readHCA(filepath.Join(root, e.Name()), e.Name())
		if err != nil {
			// One unreadable adapter must not hide the rest; the node may still
			// be usable through its other HCAs.
			continue
		}
		out = append(out, hca)
	}
	slices.SortFunc(out, func(a, b HCA) int { return strings.Compare(a.Name, b.Name) })
	return out, nil
}

func readHCA(dir, name string) (HCA, error) {
	h := HCA{Name: name, NUMANode: -1}
	h.FirmwareVersion = readTrimmed(filepath.Join(dir, "fw_ver"))
	h.NodeGUID = normalizeGUID(readTrimmed(filepath.Join(dir, "node_guid")))

	// device is a symlink into the PCI tree; its target's base name is the
	// bus address, and numa_node lives beside it.
	if target, err := filepath.EvalSymlinks(filepath.Join(dir, "device")); err == nil {
		h.PCIAddr = filepath.Base(target)
		if n, err := strconv.Atoi(readTrimmed(filepath.Join(target, "numa_node"))); err == nil {
			h.NUMANode = n
		}
	}

	portDirs, err := os.ReadDir(filepath.Join(dir, "ports"))
	if err != nil {
		return h, fmt.Errorf("%s: no ports: %w", name, err)
	}
	for _, pd := range portDirs {
		num, err := strconv.Atoi(pd.Name())
		if err != nil {
			continue
		}
		h.Ports = append(h.Ports, readPort(filepath.Join(dir, "ports", pd.Name()), num))
	}
	slices.SortFunc(h.Ports, func(a, b Port) int { return a.Num - b.Num })
	return h, nil
}

func readPort(dir string, num int) Port {
	p := Port{
		Num:       num,
		State:     parseEnumeratedState(readTrimmed(filepath.Join(dir, "state"))),
		PhysState: parseEnumeratedState(readTrimmed(filepath.Join(dir, "phys_state"))),
		RateGbps:  parseRate(readTrimmed(filepath.Join(dir, "rate"))),
		LinkLayer: LinkLayer(readTrimmed(filepath.Join(dir, "link_layer"))),
		GIDPrefix: parseGIDPrefix(readTrimmed(filepath.Join(dir, "gids", "0"))),
		LID:       parseMaybeHex(readTrimmed(filepath.Join(dir, "lid"))),
		SMLID:     parseMaybeHex(readTrimmed(filepath.Join(dir, "sm_lid"))),
	}
	p.RoCEv2GID = -1
	p.PortGUID = portGUID(readTrimmed(filepath.Join(dir, "gids", "0")))
	if p.LinkLayer == LinkEthernet {
		p.RoCEv2GID = rocev2GID(dir)
		if p.RoCEv2GID >= 0 {
			idx := strconv.Itoa(p.RoCEv2GID)
			p.Network = RoCENetwork(readTrimmed(filepath.Join(dir, "gid_attrs", "ndevs", idx)),
				readTrimmed(filepath.Join(dir, "gids", idx)))
		}
	}
	for _, name := range PortCounters {
		raw := readTrimmed(filepath.Join(dir, "counters", name))
		v, err := strconv.ParseUint(raw, 10, 64)
		if err != nil {
			continue // not every driver has every counter
		}
		if dataCounter(name) {
			v *= 4
		}
		if p.Counters == nil {
			p.Counters = map[string]uint64{}
		}
		p.Counters[name] = v
	}
	return p
}

// parseEnumeratedState turns sysfs's "4: ACTIVE" into "ACTIVE".
func parseEnumeratedState(s string) string {
	if _, rest, found := strings.Cut(s, ":"); found {
		return strings.TrimSpace(rest)
	}
	return s
}

// parseRate turns "100 Gb/sec (4X EDR)" into 100.
func parseRate(s string) int {
	f := strings.Fields(s)
	if len(f) == 0 {
		return 0
	}
	// The rate can be fractional ("2.5 Gb/sec") on old hardware; truncate.
	v, err := strconv.ParseFloat(f[0], 64)
	if err != nil {
		return 0
	}
	return int(v)
}

// parseGIDPrefix keeps the first four groups of a GID, which is the subnet
// prefix; the remaining four are the interface id.
func parseGIDPrefix(gid string) string {
	parts := strings.Split(strings.ToLower(gid), ":")
	if len(parts) < 4 {
		return ""
	}
	return strings.Join(parts[:4], ":")
}

// parseMaybeHex accepts both "0x1a" and "26"; sysfs is inconsistent between
// drivers and kernel versions.
func parseMaybeHex(s string) int {
	s = strings.TrimSpace(s)
	if s == "" {
		return 0
	}
	if after, ok := strings.CutPrefix(strings.ToLower(s), "0x"); ok {
		v, err := strconv.ParseInt(after, 16, 64)
		if err != nil {
			return 0
		}
		return int(v)
	}
	v, err := strconv.Atoi(s)
	if err != nil {
		return 0
	}
	return v
}

func readTrimmed(path string) string {
	b, err := os.ReadFile(path)
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(b))
}

// rocev2GID finds the GID index a RoCE port's traffic should use.
func rocev2GID(dir string) int {
	first := -1
	for i := 0; i < 64; i++ {
		typ := readTrimmed(filepath.Join(dir, "gid_attrs", "types", strconv.Itoa(i)))
		if !strings.EqualFold(typ, "RoCE v2") {
			continue
		}
		gid := strings.ToLower(readTrimmed(filepath.Join(dir, "gids", strconv.Itoa(i))))
		if strings.HasPrefix(gid, "0000:0000:0000:0000:0000:ffff:") {
			return i
		}
		if first < 0 {
			first = i
		}
	}
	return first
}

// normalizeGUID writes a GUID one way: sysfs says "0c42:a103:0045:bd6c",
// sminfo and ibnetdiscover say "0x0c42a1030045bd6c".
func normalizeGUID(g string) string {
	g = strings.ToLower(strings.TrimSpace(g))
	g = strings.TrimPrefix(g, "0x")
	g = strings.ReplaceAll(g, ":", "")
	if g == "" || strings.Trim(g, "0") == "" {
		return ""
	}
	for len(g) < 16 {
		g = "0" + g
	}
	return "0x" + g
}

// NormalizeGUID is normalizeGUID for other packages.
func NormalizeGUID(g string) string { return normalizeGUID(g) }

// portGUID is the interface identifier of a GID: its last four groups.
func portGUID(gid string) string {
	parts := strings.Split(gid, ":")
	if len(parts) != 8 {
		return ""
	}
	return normalizeGUID(strings.Join(parts[4:], ""))
}

// RoCENetwork is the IP network an IPv4-mapped RoCE GID's address sits in,
// read from the interface that carries it. Replaced in tests.
var RoCENetwork = func(ndev, gid string) string {
	ip := gidIPv4(gid)
	if ip == nil || ndev == "" {
		return ""
	}
	ifc, err := net.InterfaceByName(ndev)
	if err != nil {
		return ""
	}
	addrs, _ := ifc.Addrs()
	for _, a := range addrs {
		if n, ok := a.(*net.IPNet); ok && n.IP.Equal(ip) {
			return (&net.IPNet{IP: ip.Mask(n.Mask), Mask: n.Mask}).String()
		}
	}
	return ""
}

// gidIPv4 is the IPv4 address in an IPv4-mapped GID, nil otherwise.
func gidIPv4(gid string) net.IP {
	g := strings.ToLower(gid)
	const mapped = "0000:0000:0000:0000:0000:ffff:"
	if !strings.HasPrefix(g, mapped) {
		return nil
	}
	hex := strings.ReplaceAll(strings.TrimPrefix(g, mapped), ":", "")
	if len(hex) != 8 {
		return nil
	}
	b := make(net.IP, 4)
	for i := range 4 {
		v, err := strconv.ParseUint(hex[i*2:i*2+2], 16, 8)
		if err != nil {
			return nil
		}
		b[i] = byte(v)
	}
	return b
}