llmgw/internal/agent/discovery/discover_test.go

package discovery

import (
	"context"
	"log/slog"
	"os"
	"path/filepath"
	"runtime"
	"testing"
)

func fakeSMI(t *testing.T) string {
	t.Helper()
	stub := filepath.Join(t.TempDir(), "nvidia-smi")
	script := "#!/bin/sh\ncat <<'EOF'\n" + smiSample + "EOF\n"
	if err := os.WriteFile(stub, []byte(script), 0o755); err != nil {
		t.Fatal(err)
	}
	return stub
}

func TestDiscoverUsesSMIByDefault(t *testing.T) {
	inv, err := Discover(t.Context(), Options{
		NvidiaSMI: fakeSMI(t),
		IBRoot:    filepath.Join(t.TempDir(), "absent"),
	}, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatal(err)
	}
	if inv.Source != "nvidia-smi" {
		t.Errorf("source=%q", inv.Source)
	}
	if len(inv.GPUs) != 2 {
		t.Errorf("got %d GPUs", len(inv.GPUs))
	}
}

func TestDiscoverFallsBackWhenNVMLUnavailable(t *testing.T) {
	// Asking for NVML in a build without it must still produce an inventory:
	// the fallback is the whole point of the build tag.
	inv, err := Discover(t.Context(), Options{
		PreferNVML: true,
		NvidiaSMI:  fakeSMI(t),
		IBRoot:     filepath.Join(t.TempDir(), "absent"),
	}, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatal(err)
	}
	if len(inv.GPUs) == 0 {
		t.Fatal("no GPUs discovered through the fallback")
	}
	if NVMLSupported() {
		// On an NVML build without a driver the NVML attempt fails and the
		// fallback still has to work.
		if inv.Source != "nvml" && inv.Source != "nvidia-smi" {
			t.Errorf("source=%q", inv.Source)
		}
	} else if inv.Source != "nvidia-smi" {
		t.Errorf("a build without NVML reported source=%q", inv.Source)
	}
}

func TestDiscoverSkipGPUs(t *testing.T) {
	inv, err := Discover(t.Context(), Options{
		SkipGPUs: true,
		IBRoot:   fakeSysfs(t),
	}, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatal(err)
	}
	if len(inv.GPUs) != 0 {
		t.Errorf("GPUs were discovered despite SkipGPUs")
	}
	if len(inv.HCAs) != 3 {
		t.Errorf("got %d HCAs", len(inv.HCAs))
	}
	if inv.Source != "none" {
		t.Errorf("source=%q", inv.Source)
	}
}

func TestDiscoverReportsGPUFailure(t *testing.T) {
	// A node meant to have GPUs whose discovery fails must say so rather than
	// reporting an empty, healthy-looking inventory.
	_, err := Discover(context.Background(), Options{
		NvidiaSMI: filepath.Join(t.TempDir(), "absent"),
		IBRoot:    filepath.Join(t.TempDir(), "absent"),
	}, slog.New(slog.DiscardHandler))
	if err == nil {
		t.Error("a failed GPU discovery was reported as success")
	}
}

func TestNVMLSupportedMatchesBuild(t *testing.T) {
	// Documents which backend this binary has, and keeps both files compiling.
	t.Logf("NVML backend compiled in: %v", NVMLSupported())
	if !NVMLSupported() {
		if _, err := DiscoverGPUsNVML(t.Context()); err != ErrNoNVML {
			t.Errorf("expected ErrNoNVML, got %v", err)
		}
	}
}

func TestMemoryIsReportedEvenWhenGPUDiscoveryFails(t *testing.T) {
	// A node with no nvidia-smi is exactly the node whose memory matters, since
	// a CPU profile is what it will be asked to run. Reading memory after GPU
	// discovery meant that node reported none: the GPU error returns early and
	// takes a fact that has nothing to do with cards with it. Seen on node05,
	// which has 32 GB and was reporting "-".
	inv, err := Discover(t.Context(), Options{NvidiaSMI: "/nonexistent/nvidia-smi"}, nil)
	if err == nil {
		t.Fatal("GPU discovery was expected to fail with no nvidia-smi")
	}
	// On Linux this is a real number; elsewhere there is no /proc/meminfo and 0
	// is the honest answer. Either way it must not be lost to the GPU error.
	if runtime.GOOS == "linux" && inv.MemoryTotalMB <= 0 {
		t.Errorf("memory was not reported alongside a GPU failure: %d", inv.MemoryTotalMB)
	}
}