Vetting/agent/tests/gpu.go

package tests

import (
	"context"
	"encoding/json"
	"fmt"
	"os/exec"
	"strings"
	"time"
)

// GPU enumerates VGA / 3D PCI devices. No devices → skip cleanly (a
// CPU-only server passes this stage by virtue of having nothing to
// stress). Devices present → try nvidia-smi for NVIDIA cards, else
// accept PCI presence.
func GPU(ctx context.Context, d Deps) Outcome {
	pciStart := time.Now()
	devices := listGPUPCI(ctx)
	pciEnd := time.Now()
	if len(devices) == 0 {
		d.Info("GPU: no VGA/3D PCI devices found — skipping stage")
		return Outcome{
			Passed:  true,
			Summary: "skipped (no GPU present)",
			Extras:  map[string]any{"skipped": true, "reason": "no_gpu_present"},
		}
	}
	d.Info("GPU: found " + joinDevices(devices))

	nvStart := time.Now()
	nvidia := nvidiaSmiList(ctx)
	nvEnd := time.Now()
	extras := map[string]any{
		"pci_devices": devices,
		"skipped":     false,
	}
	if len(nvidia) > 0 {
		extras["nvidia"] = nvidia
		d.Info("GPU: nvidia-smi reports: " + strings.Join(nvidia, ", "))
	}

	// Sub-step rows: one per enumerated PCI device, plus (optionally) one
	// per NVIDIA card when nvidia-smi sees anything. PCI enumeration runs
	// once for all devices — we bracket that single invocation by
	// pciStart/pciEnd and attribute the window to each device row so the
	// UI can still slice the log per row by time.
	var subs []SubStepReport
	for i, dev := range devices {
		summary, _ := json.Marshal(map[string]any{"pci": dev, "ordinal": i})
		subs = append(subs, SubStepReport{
			Name:        fmt.Sprintf("pci #%d", i),
			Passed:      true,
			StartedAt:   pciStart,
			CompletedAt: pciEnd,
			SummaryJSON: summary,
		})
	}
	for i, line := range nvidia {
		summary, _ := json.Marshal(map[string]any{"nvidia_smi": line})
		subs = append(subs, SubStepReport{
			Name:        fmt.Sprintf("nvidia #%d", i),
			Passed:      true,
			StartedAt:   nvStart,
			CompletedAt: nvEnd,
			SummaryJSON: summary,
		})
	}

	return Outcome{
		Passed:   true,
		Summary:  formatCount(len(devices), "GPU present"),
		Extras:   extras,
		SubSteps: subs,
	}
}

// listGPUPCI shells out to lspci. Returns human-readable strings, one
// per VGA/3D device. If lspci isn't available we return nil and the
// caller treats it as "no GPU" which auto-skips.
func listGPUPCI(ctx context.Context) []string {
	cmd := exec.CommandContext(ctx, "lspci", "-mm")
	out, err := cmd.Output()
	if err != nil {
		return nil
	}
	var devs []string
	for _, line := range strings.Split(string(out), "\n") {
		l := strings.ToLower(line)
		if strings.Contains(l, "vga compatible controller") || strings.Contains(l, "3d controller") {
			devs = append(devs, strings.TrimSpace(line))
		}
	}
	return devs
}

// nvidiaSmiList returns each card's "<name>, <pci bus>" line; empty
// slice when nvidia-smi isn't installed or fails.
func nvidiaSmiList(ctx context.Context) []string {
	cmd := exec.CommandContext(ctx, "nvidia-smi", "-L")
	out, err := cmd.Output()
	if err != nil {
		return nil
	}
	var lines []string
	for _, l := range strings.Split(string(out), "\n") {
		l = strings.TrimSpace(l)
		if l != "" {
			lines = append(lines, l)
		}
	}
	return lines
}

func joinDevices(devs []string) string {
	if len(devs) == 0 {
		return ""
	}
	if len(devs) == 1 {
		return devs[0]
	}
	return devs[0] + " (+" + strings.TrimSpace(formatCount(len(devs)-1, "more")) + ")"
}