llmgw/internal/agent/discovery/ib_test.go

package discovery

import (
	"os"
	"path/filepath"
	"testing"
)

// fakeSysfs builds an InfiniBand tree shaped like the real one.
func fakeSysfs(t *testing.T) string {
	t.Helper()
	root := t.TempDir()

	write := func(path, content string) {
		t.Helper()
		full := filepath.Join(root, path)
		if err := os.MkdirAll(filepath.Dir(full), 0o755); err != nil {
			t.Fatal(err)
		}
		if err := os.WriteFile(full, []byte(content+"\n"), 0o644); err != nil {
			t.Fatal(err)
		}
	}

	// mlx5_0: one active IB port.
	write("ib/mlx5_0/fw_ver", "20.35.1012")
	write("ib/mlx5_0/ports/1/state", "4: ACTIVE")
	write("ib/mlx5_0/ports/1/phys_state", "5: LinkUp")
	write("ib/mlx5_0/ports/1/rate", "200 Gb/sec (4X HDR)")
	write("ib/mlx5_0/ports/1/link_layer", "InfiniBand")
	write("ib/mlx5_0/ports/1/gids/0", "FEC0:0000:0000:000A:0002:C903:00A0:1234")
	write("ib/mlx5_0/ports/1/lid", "0x1a")
	write("ib/mlx5_0/ports/1/sm_lid", "1")
	write("ib/mlx5_0/ports/1/counters/symbol_error", "3")
	write("ib/mlx5_0/ports/1/counters/link_downed", "1")
	write("ib/mlx5_0/ports/1/counters/port_xmit_data", "1000")

	// mlx5_1: two ports, the second one down.
	write("ib/mlx5_1/ports/1/state", "4: ACTIVE")
	write("ib/mlx5_1/ports/1/rate", "200 Gb/sec (4X HDR)")
	write("ib/mlx5_1/ports/1/link_layer", "InfiniBand")
	write("ib/mlx5_1/ports/1/gids/0", "fec0:0000:0000:000a:0002:c903:00a0:5678")
	write("ib/mlx5_1/ports/2/state", "1: DOWN")
	write("ib/mlx5_1/ports/2/phys_state", "3: Disabled")
	write("ib/mlx5_1/ports/2/rate", "0")
	write("ib/mlx5_1/ports/2/link_layer", "InfiniBand")

	// mlx5_2: RoCE rather than native IB.
	write("ib/mlx5_2/ports/1/state", "4: ACTIVE")
	write("ib/mlx5_2/ports/1/rate", "100 Gb/sec (4X EDR)")
	write("ib/mlx5_2/ports/1/link_layer", "Ethernet")
	write("ib/mlx5_2/ports/1/gids/0", "fe80:0000:0000:0000:0000:0000:0000:0001")
	write("ib/mlx5_2/ports/1/gid_attrs/types/0", "IB/RoCE v1")
	write("ib/mlx5_2/ports/1/gids/1", "fe80:0000:0000:0000:0000:0000:0000:0001")
	write("ib/mlx5_2/ports/1/gid_attrs/types/1", "RoCE v2")
	write("ib/mlx5_2/ports/1/gids/3", "0000:0000:0000:0000:0000:ffff:0a00:0001")
	write("ib/mlx5_2/ports/1/gid_attrs/types/3", "RoCE v2")
	write("ib/mlx5_2/ports/1/gid_attrs/ndevs/3", "ens1f0")

	// A PCI device with a NUMA node, linked the way the driver does it.
	pci := filepath.Join(root, "pci", "0000:0b:00.0")
	if err := os.MkdirAll(pci, 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(pci, "numa_node"), []byte("1\n"), 0o644); err != nil {
		t.Fatal(err)
	}
	if err := os.Symlink(pci, filepath.Join(root, "ib", "mlx5_0", "device")); err != nil {
		t.Fatal(err)
	}

	return filepath.Join(root, "ib")
}

func TestDiscoverIB(t *testing.T) {
	hcas, err := DiscoverIB(fakeSysfs(t))
	if err != nil {
		t.Fatal(err)
	}
	if len(hcas) != 3 {
		t.Fatalf("expected 3 HCAs, got %d", len(hcas))
	}
	// Sorted by name, so the order is predictable for the rest of the test.
	if hcas[0].Name != "mlx5_0" || hcas[2].Name != "mlx5_2" {
		t.Fatalf("HCAs are not sorted by name: %v %v", hcas[0].Name, hcas[2].Name)
	}

	h := hcas[0]
	if h.FirmwareVersion != "20.35.1012" {
		t.Errorf("fw_ver=%q", h.FirmwareVersion)
	}
	if h.PCIAddr != "0000:0b:00.0" {
		t.Errorf("PCIAddr=%q", h.PCIAddr)
	}
	if h.NUMANode != 1 {
		t.Errorf("NUMANode=%d; GPU-to-HCA affinity depends on this", h.NUMANode)
	}

	p := h.Ports[0]
	if p.State != "ACTIVE" || !p.Active() {
		t.Errorf("state=%q", p.State)
	}
	if p.PhysState != "LinkUp" {
		t.Errorf("phys_state=%q", p.PhysState)
	}
	if p.RateGbps != 200 {
		t.Errorf("rate=%d", p.RateGbps)
	}
	if p.LinkLayer != LinkInfiniBand {
		t.Errorf("link_layer=%q", p.LinkLayer)
	}
	if p.GIDPrefix != "fec0:0000:0000:000a" {
		t.Errorf("GIDPrefix=%q; it must be lowercased and cut to four groups", p.GIDPrefix)
	}
	if p.LID != 26 {
		t.Errorf("LID=%d, expected 0x1a to parse as 26", p.LID)
	}
	if p.SMLID != 1 {
		t.Errorf("SMLID=%d, expected a decimal value to parse too", p.SMLID)
	}
}

func TestActivePortsExcludesDown(t *testing.T) {
	hcas, err := DiscoverIB(fakeSysfs(t))
	if err != nil {
		t.Fatal(err)
	}
	h := hcas[1] // mlx5_1
	if len(h.Ports) != 2 {
		t.Fatalf("expected 2 ports, got %d", len(h.Ports))
	}
	active := h.ActivePorts()
	if len(active) != 1 || active[0].Num != 1 {
		t.Errorf("a DOWN port was reported as active: %+v", active)
	}
}

func TestRoCEIsDistinguished(t *testing.T) {
	// RoCE is usable but behaves differently from native IB; the placement
	// logic needs to be able to tell them apart.
	hcas, err := DiscoverIB(fakeSysfs(t))
	if err != nil {
		t.Fatal(err)
	}
	if got := hcas[2].Ports[0].LinkLayer; got != LinkEthernet {
		t.Errorf("mlx5_2 link_layer=%q, expected Ethernet (RoCE)", got)
	}
}

func TestMissingRootIsNotAnError(t *testing.T) {
	// A node without InfiniBand is perfectly normal; it just forms its own
	// single-node island.
	hcas, err := DiscoverIB(filepath.Join(t.TempDir(), "nonexistent"))
	if err != nil {
		t.Fatalf("a node without IB should not be an error: %v", err)
	}
	if len(hcas) != 0 {
		t.Errorf("expected no HCAs, got %d", len(hcas))
	}
}

func TestUnreadableAdapterDoesNotHideOthers(t *testing.T) {
	root := fakeSysfs(t)
	// An adapter with no ports directory at all.
	if err := os.MkdirAll(filepath.Join(root, "mlx5_broken"), 0o755); err != nil {
		t.Fatal(err)
	}

	hcas, err := DiscoverIB(root)
	if err != nil {
		t.Fatal(err)
	}
	if len(hcas) != 3 {
		t.Errorf("a broken adapter changed the result: got %d HCAs", len(hcas))
	}
}

func TestParseHelpers(t *testing.T) {
	t.Run("rate", func(t *testing.T) {
		for in, want := range map[string]int{
			"200 Gb/sec (4X HDR)": 200,
			"100 Gb/sec (4X EDR)": 100,
			"2.5 Gb/sec (1X SDR)": 2,
			"":                    0,
			"garbage":             0,
		} {
			if got := parseRate(in); got != want {
				t.Errorf("parseRate(%q)=%d, want %d", in, got, want)
			}
		}
	})

	t.Run("state", func(t *testing.T) {
		for in, want := range map[string]string{
			"4: ACTIVE": "ACTIVE",
			"1: DOWN":   "DOWN",
			"5: LinkUp": "LinkUp",
			"ACTIVE":    "ACTIVE",
			"":          "",
		} {
			if got := parseEnumeratedState(in); got != want {
				t.Errorf("parseEnumeratedState(%q)=%q, want %q", in, got, want)
			}
		}
	})

	t.Run("hex", func(t *testing.T) {
		for in, want := range map[string]int{
			"0x1a": 26, "0X1A": 26, "26": 26, "": 0, "zzz": 0,
		} {
			if got := parseMaybeHex(in); got != want {
				t.Errorf("parseMaybeHex(%q)=%d, want %d", in, got, want)
			}
		}
	})

	t.Run("gid", func(t *testing.T) {
		if got := parseGIDPrefix("FE80:0000:0000:0000:0002:C903:00A0:1234"); got != "fe80:0000:0000:0000" {
			t.Errorf("parseGIDPrefix=%q", got)
		}
		if got := parseGIDPrefix("short:gid"); got != "" {
			t.Errorf("a malformed GID should yield an empty prefix, got %q", got)
		}
	})
}

func TestAPortsCountersAreRead(t *testing.T) {
	hcas, err := DiscoverIB(fakeSysfs(t))
	if err != nil {
		t.Fatal(err)
	}
	c := hcas[0].Ports[0].Counters
	// Data in bytes: sysfs counts four-byte words.
	if c["symbol_error"] != 3 || c["link_downed"] != 1 || c["port_xmit_data"] != 4000 {
		t.Errorf("counters = %v", c)
	}
	if _, ok := c["port_rcv_errors"]; ok {
		t.Error("a counter the driver does not have should be absent, not zero")
	}
}

func TestRoCEFindsItsRoutableGID(t *testing.T) {
	// NCCL's default GID index is a v1 link-local entry on most hosts, which
	// does not route: the job hangs. The v2 entry holding an IPv4 address is
	// the one to use, over an earlier v2 link-local one.
	hcas, err := DiscoverIB(fakeSysfs(t))
	if err != nil {
		t.Fatal(err)
	}
	for _, h := range hcas {
		p := h.Ports[0]
		switch h.Name {
		case "mlx5_2":
			if p.RoCEv2GID != 3 {
				t.Errorf("RoCE GID index = %d, want 3", p.RoCEv2GID)
			}
		default:
			if p.RoCEv2GID != -1 {
				t.Errorf("%s is InfiniBand and has a GID index %d", h.Name, p.RoCEv2GID)
			}
		}
	}
}

func TestTheSubnetManagerAPortSeesIsRead(t *testing.T) {
	sm, err := ParseSMInfo("sminfo: sm lid 1 sm guid 0x98039B03009fcfd6, activity count 5452 priority 15 state 3 SMINFO_MASTER\n")
	if err != nil {
		t.Fatal(err)
	}
	if sm.GUID != "0x98039b03009fcfd6" || sm.LID != 1 || sm.Priority != 15 || sm.State != "SMINFO_MASTER" {
		t.Errorf("sm = %+v", sm)
	}
	if _, err := ParseSMInfo("ibwarn: [1234] mad_rpc_open_port: can't open UMAD port"); err == nil {
		t.Error("an error from sminfo should not parse")
	}
}

func TestAPortsFabricIsItsSubnetManagerOrItsNetwork(t *testing.T) {
	// Two ports on one InfiniBand fabric see one master SM, and on separate
	// fabrics cannot; RoCE has no SM, and is its IP network.
	hcas, err := DiscoverIB(fakeSysfs(t))
	if err != nil {
		t.Fatal(err)
	}
	if hcas[0].NodeGUID != "" {
		t.Errorf("no node_guid file, yet %q", hcas[0].NodeGUID)
	}
	err = AttachSM(hcas, func(hca string, port int) (*SMInfo, error) {
		return &SMInfo{GUID: "0x00000000000000aa", LID: 1, State: "SMINFO_MASTER"}, nil
	})
	if err != nil {
		t.Fatal(err)
	}
	if got := hcas[0].Ports[0].FabricKey(); got != "sm:0x00000000000000aa" {
		t.Errorf("IB key = %q", got)
	}
	if got := hcas[1].Ports[1].FabricKey(); got != "" {
		t.Errorf("a port that is down is on no fabric, got %q", got)
	}
	if hcas[2].Ports[0].SM != nil {
		t.Error("RoCE was asked for a subnet manager")
	}
	if got := hcas[0].Ports[0].PortGUID; got != "0x0002c90300a01234" {
		t.Errorf("port GUID = %q", got)
	}
}

func TestARoCEPortsNetworkIsItsFabric(t *testing.T) {
	old := RoCENetwork
	t.Cleanup(func() { RoCENetwork = old })
	RoCENetwork = func(ndev, gid string) string {
		if ip := gidIPv4(gid); ip != nil && ip.String() == "10.0.0.1" {
			return "10.0.0.0/24"
		}
		return ""
	}
	hcas, err := DiscoverIB(fakeSysfs(t))
	if err != nil {
		t.Fatal(err)
	}
	roce := hcas[2].Ports[0]
	if roce.Network != "10.0.0.0/24" || roce.FabricKey() != "roce:10.0.0.0/24" {
		t.Errorf("RoCE network %q key %q", roce.Network, roce.FabricKey())
	}
}

// smirDump is `saquery SMIR` on a fabric with a master and a standby, in the
// layout infiniband-diags prints.
const smirDump = "SMInfoRecord dump:\n\t\tRID\n\t\tLID...................1\n\t\tSMInfo dump:\n" +
	"\t\tGUID..................0x248a07030030f0c0\n\t\tSM_Key................0x0000000000000000\n" +
	"\t\tActCount..............386773\n\t\tPriority..............15\n\t\tSMState...............3\n" +
	"SMInfoRecord dump:\n\t\tRID\n\t\tLID...................12\n\t\tSMInfo dump:\n" +
	"\t\tGUID..................0x248a07030030f1d4\n\t\tSM_Key................0x0000000000000000\n" +
	"\t\tActCount..............12\n\t\tPriority..............14\n\t\tSMState...............2\n"

func TestEverySubnetManagerOnAFabricIsListed(t *testing.T) {
	recs, err := ParseSMIR(smirDump)
	if err != nil {
		t.Fatal(err)
	}
	if len(recs) != 2 {
		t.Fatalf("records = %+v", recs)
	}
	if recs[0].GUID != "0x248a07030030f0c0" || recs[0].State != "master" || recs[0].Priority != 15 || recs[0].LID != 1 {
		t.Errorf("master = %+v", recs[0])
	}
	if recs[1].State != "standby" || recs[1].LID != 12 {
		t.Errorf("standby = %+v", recs[1])
	}
	if _, err := ParseSMIR("ibwarn: can't open UMAD port"); err == nil {
		t.Error("an error should not parse")
	}
}