llmgw/internal/agent/discovery/gpu_test.go

package discovery

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

const smiSample = `0, GPU-11111111-2222-3333-4444-555555555555, NVIDIA H100 80GB HBM3, 81559, 4, 9.0, 00000000:1B:00.0, 550.54.15, 34, 0
1, GPU-66666666-7777-8888-9999-000000000000, NVIDIA H100 80GB HBM3, 81559, 0, 9.0, 00000000:1C:00.0, 550.54.15, 36, 0
`

func TestParseSMI(t *testing.T) {
	gpus, err := parseSMI(smiSample)
	if err != nil {
		t.Fatal(err)
	}
	if len(gpus) != 2 {
		t.Fatalf("expected 2 GPUs, got %d", len(gpus))
	}

	g := gpus[0]
	if g.Index != 0 || g.UUID != "GPU-11111111-2222-3333-4444-555555555555" {
		t.Errorf("identity wrong: %+v", g)
	}
	if g.MemoryTotalMB != 81559 {
		t.Errorf("MemoryTotalMB=%d", g.MemoryTotalMB)
	}
	if g.ComputeCapability != "9.0" {
		t.Errorf("ComputeCapability=%q; FP8 placement depends on it", g.ComputeCapability)
	}
	if g.DriverVersion != "550.54.15" {
		t.Errorf("DriverVersion=%q", g.DriverVersion)
	}
	if !g.Health.OK() {
		t.Errorf("a healthy GPU was flagged: %v", g.Health.Problems)
	}
}

func TestParseSMIUnsupportedFields(t *testing.T) {
	// Consumer cards report [N/A] for ECC; that must not fail the whole parse.
	const row = `0, GPU-abc, NVIDIA GeForce RTX 4090, 24564, 0, 8.9, 00000000:01:00.0, 550.54.15, 45, [N/A]
`
	gpus, err := parseSMI(row)
	if err != nil {
		t.Fatal(err)
	}
	if len(gpus) != 1 {
		t.Fatalf("expected 1 GPU, got %d", len(gpus))
	}
	if gpus[0].Health.ECCErrors != 0 || !gpus[0].Health.OK() {
		t.Errorf("[N/A] ECC was misread: %+v", gpus[0].Health)
	}
}

func TestHealthFlagsForeignMemory(t *testing.T) {
	// Memory in use on an otherwise idle card means another process is there.
	const row = `0, GPU-abc, NVIDIA H100 80GB HBM3, 81559, 40000, 9.0, 00000000:1B:00.0, 550.54.15, 34, 0
`
	gpus, _ := parseSMI(row)
	if gpus[0].Health.OK() {
		t.Error("40 GiB allocated by another process was not flagged")
	}
}

func TestHealthFlagsECC(t *testing.T) {
	const row = `0, GPU-abc, NVIDIA H100 80GB HBM3, 81559, 0, 9.0, 00000000:1B:00.0, 550.54.15, 34, 3
`
	gpus, _ := parseSMI(row)
	if gpus[0].Health.OK() {
		t.Error("uncorrected ECC errors were not flagged")
	}
}

func TestGPUClass(t *testing.T) {
	// The same card must derive the same class however the driver names it,
	// because a tensor-parallel group may not mix classes.
	cases := []struct {
		name   string
		memMB  int
		expect string
	}{
		{"NVIDIA H100 80GB HBM3", 81559, "h100-80g"},
		{"NVIDIA H100-SXM5-80GB", 81559, "h100-80g"},
		{"Tesla A100-SXM4-40GB", 40960, "a100-40g"},
		{"NVIDIA A100 80GB PCIe", 81920, "a100-80g"},
		// Sizes come from what the driver reports, in GiB. That does not always
		// match the marketing name: an L40S is sold as 48GB but reports 44.99
		// GiB, because the marketing figure is decimal GB. Being consistent
		// matters more than matching the box, since the number is what decides
		// whether a model fits.
		{"NVIDIA L40S", 46068, "l40s-45g"},
		// The model token is what distinguishes a 4090 from a 4080, so "rtx"
		// is not enough on its own.
		{"NVIDIA GeForce RTX 4090", 24564, "4090-24g"},
	}
	for _, c := range cases {
		got := GPU{Name: c.name, MemoryTotalMB: c.memMB}.Class()
		if got != c.expect {
			t.Errorf("%q (%d MiB) -> %q, want %q", c.name, c.memMB, got, c.expect)
		}
	}
}

func TestInventoryClasses(t *testing.T) {
	inv := Inventory{GPUs: []GPU{
		{Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559},
		{Name: "NVIDIA H100-SXM5-80GB", MemoryTotalMB: 81559},
		{Name: "Tesla A100-SXM4-40GB", MemoryTotalMB: 40960},
	}}
	got := inv.Classes()
	if len(got) != 2 || got[0] != "a100-40g" || got[1] != "h100-80g" {
		t.Errorf("classes=%v", got)
	}
}

// TestDiscoverGPUsSMI runs the real code path against a stub binary, so the
// command construction is covered too, not just the parser.
func TestDiscoverGPUsSMI(t *testing.T) {
	dir := t.TempDir()
	stub := filepath.Join(dir, "nvidia-smi")
	script := "#!/bin/sh\ncat <<'EOF'\n" + smiSample + "EOF\n"
	if err := os.WriteFile(stub, []byte(script), 0o755); err != nil {
		t.Fatal(err)
	}

	gpus, err := DiscoverGPUsSMI(t.Context(), stub)
	if err != nil {
		t.Fatal(err)
	}
	if len(gpus) != 2 {
		t.Fatalf("expected 2 GPUs, got %d", len(gpus))
	}
}

func TestDiscoverGPUsSMIMissingBinary(t *testing.T) {
	// A node with no NVIDIA driver must produce a clear error rather than a
	// panic or an empty success.
	_, err := DiscoverGPUsSMI(context.Background(), filepath.Join(t.TempDir(), "absent"))
	if err == nil {
		t.Error("expected an error when nvidia-smi is missing")
	}
}