bee/audit/internal/platform/nvidia_bandwidth_socket_test.go

package platform

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

func fakeNvidiaSmiBusIDs(t *testing.T, csv string) {
	t.Helper()
	old := satExecCommand
	satExecCommand = func(name string, args ...string) *exec.Cmd {
		if name == "nvidia-smi" {
			return exec.Command("printf", csv)
		}
		return exec.Command(name, args...)
	}
	t.Cleanup(func() { satExecCommand = old })
}

func fakeNUMANodes(t *testing.T, byBDF map[string]string) {
	t.Helper()
	old := satReadFile
	satReadFile = func(path string) ([]byte, error) {
		bdf := filepath.Base(filepath.Dir(path))
		if node, ok := byBDF[bdf]; ok {
			return []byte(node), nil
		}
		return nil, os.ErrNotExist
	}
	t.Cleanup(func() { satReadFile = old })
}

func TestGPUNUMANodesResolvesFromPCIBusID(t *testing.T) {
	fakeNvidiaSmiBusIDs(t, "0, 00000000:05:00.0\n1, 00000000:F4:00.0\n")
	fakeNUMANodes(t, map[string]string{
		"0000:05:00.0": "0\n",
		"0000:F4:00.0": "1\n",
	})

	nodes, err := gpuNUMANodes([]int{0, 1})
	if err != nil {
		t.Fatalf("gpuNUMANodes error: %v", err)
	}
	if nodes[0] != 0 || nodes[1] != 1 {
		t.Fatalf("nodes=%v want {0:0, 1:1}", nodes)
	}
}

func TestGPUNUMANodesSkipsUnresolvableNode(t *testing.T) {
	fakeNvidiaSmiBusIDs(t, "0, 00000000:05:00.0\n1, 00000000:06:00.0\n")
	fakeNUMANodes(t, map[string]string{
		"0000:05:00.0": "0\n",
		// GPU 1's numa_node deliberately missing, and node -1 (no affinity).
		"0000:06:00.0": "-1\n",
	})

	nodes, err := gpuNUMANodes([]int{0, 1})
	if err != nil {
		t.Fatalf("gpuNUMANodes error: %v", err)
	}
	if _, ok := nodes[1]; ok {
		t.Fatalf("nodes=%v want GPU 1 absent (node -1 means no affinity)", nodes)
	}
	if nodes[0] != 0 {
		t.Fatalf("nodes[0]=%d want 0", nodes[0])
	}
}

func TestGPUBandwidthSocketGroupsSplitsBySocket(t *testing.T) {
	fakeNvidiaSmiBusIDs(t, "0, 00000000:05:00.0\n1, 00000000:06:00.0\n2, 00000000:76:00.0\n3, 00000000:77:00.0\n4, 00000000:F4:00.0\n5, 00000000:F5:00.0\n")
	fakeNUMANodes(t, map[string]string{
		"0000:05:00.0": "0\n",
		"0000:06:00.0": "0\n",
		"0000:76:00.0": "0\n",
		"0000:77:00.0": "0\n",
		"0000:F4:00.0": "1\n",
		"0000:F5:00.0": "1\n",
	})

	groups := gpuBandwidthSocketGroups([]int{0, 1, 2, 3, 4, 5}, nil)
	if len(groups) != 2 {
		t.Fatalf("groups=%v want 2 groups", groups)
	}
	if joinIndexList(groups[0]) != "0,1,2,3" {
		t.Fatalf("groups[0]=%v want 0,1,2,3", groups[0])
	}
	if joinIndexList(groups[1]) != "4,5" {
		t.Fatalf("groups[1]=%v want 4,5", groups[1])
	}
}

func TestGPUBandwidthSocketGroupsFallsBackToSingleGroup(t *testing.T) {
	t.Run("single NUMA node", func(t *testing.T) {
		fakeNvidiaSmiBusIDs(t, "0, 00000000:05:00.0\n1, 00000000:06:00.0\n")
		fakeNUMANodes(t, map[string]string{
			"0000:05:00.0": "0\n",
			"0000:06:00.0": "0\n",
		})
		groups := gpuBandwidthSocketGroups([]int{0, 1}, nil)
		if len(groups) != 1 || joinIndexList(groups[0]) != "0,1" {
			t.Fatalf("groups=%v want single group [0,1]", groups)
		}
	})

	t.Run("unresolvable NUMA node", func(t *testing.T) {
		fakeNvidiaSmiBusIDs(t, "0, 00000000:05:00.0\n1, 00000000:06:00.0\n")
		fakeNUMANodes(t, map[string]string{
			"0000:05:00.0": "0\n",
			// GPU 1 missing entirely.
		})
		groups := gpuBandwidthSocketGroups([]int{0, 1}, nil)
		if len(groups) != 1 || joinIndexList(groups[0]) != "0,1" {
			t.Fatalf("groups=%v want single fallback group [0,1]", groups)
		}
	})

	t.Run("nvidia-smi command failure", func(t *testing.T) {
		old := satExecCommand
		satExecCommand = func(name string, args ...string) *exec.Cmd {
			return exec.Command("false")
		}
		t.Cleanup(func() { satExecCommand = old })

		groups := gpuBandwidthSocketGroups([]int{0, 1}, nil)
		if len(groups) != 1 || joinIndexList(groups[0]) != "0,1" {
			t.Fatalf("groups=%v want single fallback group [0,1]", groups)
		}
	})
}

func TestNormalizeNvidiaBDF(t *testing.T) {
	cases := map[string]string{
		"00000000:05:00.0": "0000:05:00.0",
		"0000:05:00.0":     "0000:05:00.0",
		"garbage":          "garbage",
	}
	for in, want := range cases {
		if got := normalizeNvidiaBDF(in); got != want {
			t.Fatalf("normalizeNvidiaBDF(%q)=%q want %q", in, got, want)
		}
	}
}

func TestRunNvidiaBandwidthPackSplitsPerSocketThenAll(t *testing.T) {
	fakeNvidiaSmiBusIDs(t, "0, 00000000:05:00.0\n1, 00000000:06:00.0\n2, 00000000:F4:00.0\n3, 00000000:F5:00.0\n")
	fakeNUMANodes(t, map[string]string{
		"0000:05:00.0": "0\n",
		"0000:06:00.0": "0\n",
		"0000:F4:00.0": "1\n",
		"0000:F5:00.0": "1\n",
	})

	dir := t.TempDir()
	s := &System{}
	_, err := s.RunNvidiaBandwidthPack(nil, dir, []int{0, 1, 2, 3}, nil)
	if err != nil {
		t.Fatalf("RunNvidiaBandwidthPack error: %v", err)
	}

	entries, err := os.ReadDir(dir)
	if err != nil {
		t.Fatalf("ReadDir: %v", err)
	}
	if len(entries) != 1 {
		t.Fatalf("want exactly one run dir, got %v", entries)
	}
	runDir := filepath.Join(dir, entries[0].Name())

	wantFiles := []string{
		"00-nvidia-smi-persistence-mode.log",
		"01-nvidia-smi-q.log",
		"02-dcgmi-discovery.log",
		"03-dcgmi-nvbandwidth-socket0.log",
		"04-dcgmi-nvbandwidth-socket1.log",
		"05-dcgmi-nvbandwidth-all.log",
		"06-nvidia-smi-after.log",
	}
	for _, name := range wantFiles {
		if _, err := os.Stat(filepath.Join(runDir, name)); err != nil {
			t.Fatalf("missing expected job output %s: %v", name, err)
		}
	}
}

func TestRunNvidiaBandwidthPackSinglePassWhenOneSocket(t *testing.T) {
	fakeNvidiaSmiBusIDs(t, "0, 00000000:05:00.0\n1, 00000000:06:00.0\n")
	fakeNUMANodes(t, map[string]string{
		"0000:05:00.0": "0\n",
		"0000:06:00.0": "0\n",
	})

	dir := t.TempDir()
	s := &System{}
	_, err := s.RunNvidiaBandwidthPack(nil, dir, []int{0, 1}, nil)
	if err != nil {
		t.Fatalf("RunNvidiaBandwidthPack error: %v", err)
	}

	entries, err := os.ReadDir(dir)
	if err != nil {
		t.Fatalf("ReadDir: %v", err)
	}
	runDir := filepath.Join(dir, entries[0].Name())

	if _, err := os.Stat(filepath.Join(runDir, "03-dcgmi-nvbandwidth.log")); err != nil {
		t.Fatalf("missing single-pass job output: %v", err)
	}
	if _, err := os.Stat(filepath.Join(runDir, "03-dcgmi-nvbandwidth-socket0.log")); err == nil {
		t.Fatalf("did not expect a per-socket split for a single-socket system")
	}
}