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"
	"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"`
}

// 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"`
}

// 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"))

	// 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"))),
	}
	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))
}