llmgw/internal/agent/discovery/topo_test.go

package discovery

import (
	"reflect"
	"strings"
	"testing"
)

// row joins cells the way nvidia-smi prints them: tab-separated, each padded.
func row(cells ...string) string { return strings.Join(cells, "\t") }

// dgxH100 is `nvidia-smi topo -m` on an eight-GPU NVSwitch machine with two
// adapters, as newer drivers print it: NICn columns and a legend naming them.
var dgxH100 = strings.Join([]string{
	row("\x1b[4m", "GPU0", "GPU1", "GPU2", "GPU3", "NIC0", "NIC1", "CPU Affinity", "NUMA Affinity", "GPU NUMA ID\x1b[0m"),
	row("GPU0", " X ", "NV18", "NV18", "NV18", "PXB", "SYS", "0-55,112-167", "0", "N/A"),
	row("GPU1", "NV18", " X ", "NV18", "NV18", "PXB", "SYS", "0-55,112-167", "0", "N/A"),
	row("GPU2", "NV18", "NV18", " X ", "NV18", "SYS", "PXB", "56-111,168-223", "1", "N/A"),
	row("GPU3", "NV18", "NV18", "NV18", " X ", "SYS", "PIX", "56-111,168-223", "1", "N/A"),
	row("NIC0", "PXB", "PXB", "SYS", "SYS", " X ", "SYS"),
	row("NIC1", "SYS", "SYS", "PXB", "PIX", "SYS", " X "),
	"",
	"Legend:",
	"",
	"  X    = Self",
	"  SYS  = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)",
	"  NV#  = Connection traversing a bonded set of # NVLinks",
	"",
	"NIC Legend:",
	"",
	"  NIC0: mlx5_0",
	"  NIC1: mlx5_1",
	"",
}, "\n")

func TestNVLinkAndAdapterDistanceAreReadFromTopo(t *testing.T) {
	topo, err := parseTopo(dgxH100)
	if err != nil {
		t.Fatal(err)
	}
	if got := topo.NVLink[0]; !reflect.DeepEqual(got, map[int]int{1: 18, 2: 18, 3: 18}) {
		t.Errorf("GPU0's NVLink peers = %v", got)
	}
	// The legend turns NIC1 into the adapter's own name.
	if got := topo.NICPath[3]; !reflect.DeepEqual(got, map[string]string{"mlx5_0": "SYS", "mlx5_1": "PIX"}) {
		t.Errorf("GPU3's adapters = %v", got)
	}
	if topo.NUMA[2] != 1 || topo.NUMA[0] != 0 {
		t.Errorf("NUMA = %v", topo.NUMA)
	}

	gpus := []GPU{{Index: 0}, {Index: 3}}
	topo.Apply(gpus)
	if !reflect.DeepEqual(gpus[0].NVLinkPeers, []int{1, 2, 3}) || gpus[1].NUMANode != 1 ||
		gpus[1].NICPaths["mlx5_1"] != "PIX" {
		t.Errorf("applied = %+v", gpus)
	}
}

func TestOlderDriversNameAdaptersInTheHeader(t *testing.T) {
	// No legend: the columns are the adapters' own names, and NVLink between
	// pairs only, as on an A100 PCIe box with bridges.
	out := strings.Join([]string{
		row("", "GPU0", "GPU1", "GPU2", "GPU3", "mlx5_0", "CPU Affinity", "NUMA Affinity"),
		row("GPU0", " X ", "NV12", "SYS", "SYS", "NODE", "0-31", "0"),
		row("GPU1", "NV12", " X ", "SYS", "SYS", "NODE", "0-31", "0"),
		row("GPU2", "SYS", "SYS", " X ", "NV12", "SYS", "32-63", "1"),
		row("GPU3", "SYS", "SYS", "NV12", " X ", "SYS", "32-63", "1"),
		row("mlx5_0", "NODE", "NODE", "SYS", "SYS", " X "),
		"",
	}, "\n")
	topo, err := parseTopo(out)
	if err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(topo.NVLink[2], map[int]int{3: 12}) || len(topo.NVLink[0]) != 1 {
		t.Errorf("NVLink = %v", topo.NVLink)
	}
	if topo.NICPath[0]["mlx5_0"] != "NODE" {
		t.Errorf("adapters = %v", topo.NICPath)
	}
}

func TestAPCIeMachineHasNoNVLink(t *testing.T) {
	out := strings.Join([]string{
		row("", "GPU0", "GPU1", "CPU Affinity", "NUMA Affinity"),
		row("GPU0", " X ", "PHB", "0-15", "0"),
		row("GPU1", "PHB", " X ", "0-15", "0"),
		"",
	}, "\n")
	topo, err := parseTopo(out)
	if err != nil {
		t.Fatal(err)
	}
	if len(topo.NVLink) != 0 {
		t.Errorf("NVLink on a PCIe machine: %v", topo.NVLink)
	}
}

func TestNVMLPeersAreKeptOverTopo(t *testing.T) {
	topo, _ := parseTopo(dgxH100)
	gpus := []GPU{{Index: 0, NVLinkPeers: []int{1}}}
	topo.Apply(gpus)
	if !reflect.DeepEqual(gpus[0].NVLinkPeers, []int{1}) {
		t.Errorf("NVML's answer was replaced: %v", gpus[0].NVLinkPeers)
	}
}

func TestSomethingThatIsNotAMatrixIsAnError(t *testing.T) {
	if _, err := parseTopo("NVIDIA-SMI has failed because it couldn't communicate with the NVIDIA driver."); err == nil {
		t.Error("expected an error")
	}
}