bee/audit/internal/platform/nvidia_config_check.go

package platform

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

// RunNvidiaConfigCheckPack runs a read-only check of NVIDIA GPU
// configuration and NVLink topology that DCGM diag does not cover: DCGM's
// diagnostic levels stress-test compute/memory/power/bandwidth, but never
// assert that persistent config (ECC, MIG, power limit) matches the factory
// default, and NVIDIA's own DGX BasePOD deployment guide calls out
// topology/NVLink validation as a distinct manual step outside DCGM diag.
// Confidential Computing readiness — previously its own standalone SAT
// target — is folded in here as one more read-only fact about GPU state
// rather than a dedicated test, since none of these are stress tests; they
// are one-shot state/config reads that all answer "is this GPU configured
// the way we expect", not "does it survive load".
func (s *System) RunNvidiaConfigCheckPack(ctx context.Context, baseDir string, logFunc func(string)) (string, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if baseDir == "" {
		baseDir = "/var/log/bee-sat"
	}
	ts := time.Now().UTC().Format("20060102-150405")
	runDir := filepath.Join(baseDir, "nvidia-config-"+ts)
	if err := os.MkdirAll(runDir, 0755); err != nil {
		return "", err
	}
	verboseLog := filepath.Join(runDir, "verbose.log")

	status := NvidiaConfigCheckStatus{CollectedAt: time.Now().UTC()}

	// -- GPU config compliance (ECC/MIG/power limit) -----------------------
	if settings, err := s.ListNvidiaGPUSettings(); err != nil {
		status.Notes = append(status.Notes, "nvidia-smi GPU settings unavailable (no driver, or no GPU present): "+err.Error())
	} else {
		for _, g := range settings {
			f := NvidiaGPUConfigFinding{
				Index: g.Index, Name: g.Name,
				ECCCurrent: g.ECCCurrent, MIGCurrent: g.MIGCurrent, MIGPending: g.MIGPending,
				PowerLimitW: g.PowerLimitW, PowerDefaultLimitW: g.PowerDefaultLimitW,
				Issues: evaluateNvidiaGPUConfig(g),
			}
			status.GPUs = append(status.GPUs, f)
			for _, issue := range f.Issues {
				status.Warnings = append(status.Warnings, fmt.Sprintf("GPU %d (%s): %s", g.Index, g.Name, issue))
			}
			if strings.EqualFold(g.CCState, "on") {
				status.CCState = "ON"
			}
		}
	}

	// -- NVLink topology (bonded pairs, active links, error counters) ------
	topoOut, _ := runSATCommandCtx(ctx, verboseLog, "nvidia-smi-topo-m", []string{"nvidia-smi", "topo", "-m"}, nil, logFunc)
	_ = os.WriteFile(filepath.Join(runDir, "01-nvidia-smi-topo-m.log"), topoOut, 0644)
	statusOut, _ := runSATCommandCtx(ctx, verboseLog, "nvidia-smi-nvlink-s", []string{"nvidia-smi", "nvlink", "-s"}, nil, logFunc)
	_ = os.WriteFile(filepath.Join(runDir, "02-nvidia-smi-nvlink-s.log"), statusOut, 0644)
	errOut, _ := runSATCommandCtx(ctx, verboseLog, "nvidia-smi-nvlink-e", []string{"nvidia-smi", "nvlink", "-e"}, nil, logFunc)
	_ = os.WriteFile(filepath.Join(runDir, "03-nvidia-smi-nvlink-e.log"), errOut, 0644)

	pairs := parseNvidiaNVLinkBondedPairs(string(topoOut))
	linkStatus := parseNvidiaNVLinkStatus(string(statusOut))
	linkErrors := parseNvidiaNVLinkErrors(string(errOut))
	for _, pair := range pairs {
		f := evaluateNvidiaNVLinkPair(pair, linkStatus, linkErrors)
		status.NVLinkPairs = append(status.NVLinkPairs, f)
		for _, issue := range f.Issues {
			status.Warnings = append(status.Warnings, fmt.Sprintf("NVLink GPU%d<->GPU%d: %s", f.GPUA, f.GPUB, issue))
		}
	}

	// -- Confidential Computing readiness (informational only) -------------
	ccOut, ccErr := runSATCommandCtx(ctx, verboseLog, "nvidia-smi-conf-compute-q", []string{"nvidia-smi", "conf-compute", "-q"}, nil, logFunc)
	_ = os.WriteFile(filepath.Join(runDir, "04-nvidia-smi-conf-compute-q.log"), ccOut, 0644)
	if ccErr == nil {
		fields := parseConfComputeFields(ccOut)
		if status.CCState == "" {
			status.CCState = fields["CC State"]
		}
		status.CPUCCCapability = fields["CPU CC Capabilities"]
		status.GPUCCCapability = fields["GPU CC Capabilities"]
	}
	dmesgOut, _ := runSATCommandCtx(ctx, verboseLog, "dmesg", []string{"dmesg"}, nil, nil)
	ccDmesgLines := filterConfComputeDmesgLines(dmesgOut)
	_ = os.WriteFile(filepath.Join(runDir, "05-dmesg-cc-relevant.log"), []byte(strings.Join(ccDmesgLines, "\n")+"\n"), 0644)
	lowerDmesg := strings.ToLower(strings.Join(ccDmesgLines, "\n"))
	status.HostAMDSEVSNPActive = strings.Contains(lowerDmesg, "sev-snp enabled")
	status.HostIntelTDXActive = strings.Contains(lowerDmesg, "tdx module") && strings.Contains(lowerDmesg, "module initialized") ||
		strings.Contains(lowerDmesg, "virt/tdx: module initialized")
	for _, path := range []string{"/sys/module/kvm_amd/parameters/sev_snp"} {
		if out, err := runSATCommandCtx(ctx, verboseLog, "sysfs-sev-snp", []string{"cat", path}, nil, nil); err == nil &&
			strings.EqualFold(strings.TrimSpace(string(out)), "Y") {
			status.HostAMDSEVSNPActive = true
		}
	}
	gpuCanRunCC := strings.EqualFold(strings.TrimSpace(status.GPUCCCapability), "CC Capable")
	cpuCapReported := strings.TrimSpace(status.CPUCCCapability)
	cpuCanRunCC := status.HostAMDSEVSNPActive || status.HostIntelTDXActive ||
		(cpuCapReported != "" && !strings.EqualFold(cpuCapReported, "NONE"))
	status.CCReady = cpuCanRunCC && gpuCanRunCC

	summary := renderNvidiaConfigCheckSummary(status)
	if err := os.WriteFile(filepath.Join(runDir, "summary.txt"), []byte(summary), 0644); err != nil {
		return "", err
	}
	report := renderNvidiaConfigCheckReport(status)
	if err := os.WriteFile(filepath.Join(runDir, "nvidia-config-report.txt"), []byte(report), 0644); err != nil {
		return "", err
	}

	return runDir, nil
}

// NvidiaConfigCheckStatus is the result of RunNvidiaConfigCheckPack.
type NvidiaConfigCheckStatus struct {
	CollectedAt time.Time `json:"collected_at"`

	GPUs        []NvidiaGPUConfigFinding  `json:"gpus,omitempty"`
	NVLinkPairs []NvidiaNVLinkPairFinding `json:"nvlink_pairs,omitempty"`

	// Confidential Computing readiness — informational only, does not gate
	// overall_status: an unconfigured/NOT_READY CC state is a
	// feature-readiness fact on most fleets (CC is opt-in), not a fault.
	CCReady         bool   `json:"cc_ready"`
	CCState         string `json:"cc_state,omitempty"`
	CPUCCCapability string `json:"cpu_cc_capability,omitempty"`
	GPUCCCapability string `json:"gpu_cc_capability,omitempty"`

	HostAMDSEVSNPActive bool `json:"host_amd_sev_snp_active"`
	HostIntelTDXActive  bool `json:"host_intel_tdx_active"`

	// Warnings drive overall_status: any entry here fails the check.
	Warnings []string `json:"warnings,omitempty"`
	Notes    []string `json:"notes,omitempty"`
}

// NvidiaGPUConfigFinding is one GPU's config-compliance result.
type NvidiaGPUConfigFinding struct {
	Index              int      `json:"index"`
	Name               string   `json:"name"`
	ECCCurrent         string   `json:"ecc_current"`
	MIGCurrent         string   `json:"mig_current"`
	MIGPending         string   `json:"mig_pending"`
	PowerLimitW        float64  `json:"power_limit_w"`
	PowerDefaultLimitW float64  `json:"power_default_limit_w"`
	Issues             []string `json:"issues,omitempty"`
}

// NvidiaNVLinkPairFinding is one NVLink-bonded GPU pair's link-count/error
// validation result.
type NvidiaNVLinkPairFinding struct {
	GPUA          int      `json:"gpu_a"`
	GPUB          int      `json:"gpu_b"`
	ExpectedLinks int      `json:"expected_links"` // from "nvidia-smi topo -m"'s NVx cell
	ActiveLinks   int      `json:"active_links"`   // from "nvidia-smi nvlink -s"
	TotalLinks    int      `json:"total_links"`
	ErrorLinks    int      `json:"error_links"` // links with a nonzero replay/recovery/CRC counter
	Issues        []string `json:"issues,omitempty"`
}

// evaluateNvidiaGPUConfig flags config drift that DCGM diag does not check:
// ECC disabled (NVIDIA data-center GPUs ship with ECC on), a MIG mode
// change stuck pending a reset/reboot, and a power limit capped meaningfully
// below the card's own factory default (a silent throttle that a stress
// test would only surface as "lower than expected performance", with no
// clear cause).
func evaluateNvidiaGPUConfig(g NvidiaGPUSetting) []string {
	var issues []string
	if strings.EqualFold(g.ECCCurrent, "disabled") {
		issues = append(issues, "ECC is disabled")
	}
	if g.MIGCurrent != "" && g.MIGPending != "" && !strings.EqualFold(g.MIGCurrent, g.MIGPending) {
		issues = append(issues, fmt.Sprintf("MIG mode change pending (current=%s, pending=%s) — apply with a GPU reset or reboot", g.MIGCurrent, g.MIGPending))
	}
	if g.PowerDefaultLimitW > 0 && g.PowerLimitW > 0 {
		deficit := (g.PowerDefaultLimitW - g.PowerLimitW) / g.PowerDefaultLimitW
		if deficit > 0.05 {
			issues = append(issues, fmt.Sprintf("power limit %.0fW is %.0f%% below factory default %.0fW", g.PowerLimitW, deficit*100, g.PowerDefaultLimitW))
		}
	}
	return issues
}

// evaluateNvidiaNVLinkPair compares a bonded pair's actual link state
// against what the topology matrix says should be there.
func evaluateNvidiaNVLinkPair(pair nvidiaNVLinkBondedPair, status map[int][]nvidiaNVLinkPort, errors map[int]map[int][3]int64) NvidiaNVLinkPairFinding {
	f := NvidiaNVLinkPairFinding{GPUA: pair.gpuA, GPUB: pair.gpuB, ExpectedLinks: pair.links}

	active, total, errLinks := 0, 0, 0
	for _, gpu := range []int{pair.gpuA, pair.gpuB} {
		for _, port := range status[gpu] {
			total++
			if port.active {
				active++
			}
		}
		for _, counters := range errors[gpu] {
			if counters[0] != 0 || counters[1] != 0 || counters[2] != 0 {
				errLinks++
			}
		}
	}
	f.ActiveLinks, f.TotalLinks, f.ErrorLinks = active, total, errLinks

	if total > 0 && active < total {
		f.Issues = append(f.Issues, fmt.Sprintf("%d/%d NVLinks inactive on a bonded pair", total-active, total))
	}
	if errLinks > 0 {
		f.Issues = append(f.Issues, fmt.Sprintf("%d link(s) reporting replay/recovery/CRC errors", errLinks))
	}
	return f
}

// ---------------------------------------------------------------------------
// NVLink parsing — deliberately self-contained rather than importing
// internal/collector or internal/webui's equivalent parsers (isNVLinkBridgeCandidate
// et al. and parseGPUPairAdjacency/parseTopoNVLinkStatus/parseTopoNVLinkErrors,
// respectively), matching this codebase's existing precedent of small,
// package-local duplication over cross-package coupling for narrow parsing
// helpers (see webui/page_topo.go's isNICDeviceClassDev/isRAIDControllerClass).
// ---------------------------------------------------------------------------

type nvidiaNVLinkBondedPair struct {
	gpuA, gpuB int
	links      int
}

var (
	nvidiaNVLinkANSIRe = regexp.MustCompile("\x1b\\[[0-9;]*[A-Za-z]")
	nvidiaNVLinkNVRe   = regexp.MustCompile(`(?i)^NV(\d+)$`)
)

// parseNvidiaNVLinkBondedPairs returns every GPU pair with a nonzero NVLink
// bond count from a "nvidia-smi topo -m" matrix, deduplicated (A,B) == (B,A).
func parseNvidiaNVLinkBondedPairs(raw string) []nvidiaNVLinkBondedPair {
	lines := strings.Split(nvidiaNVLinkANSIRe.ReplaceAllString(raw, ""), "\n")
	headerIdx := -1
	var gpuColIndices []int
	for i, line := range lines {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "GPU0") {
			parts := strings.Fields(trimmed)
			for j, col := range parts {
				if strings.HasPrefix(col, "GPU") {
					gpuColIndices = append(gpuColIndices, j)
				}
			}
			if len(gpuColIndices) >= 2 {
				headerIdx = i
			}
			break
		}
	}
	if headerIdx < 0 {
		return nil
	}
	colIdxToGPU := make(map[int]int, len(gpuColIndices))
	for gpuIdx, colIdx := range gpuColIndices {
		colIdxToGPU[colIdx] = gpuIdx
	}

	seen := map[[2]int]bool{}
	var pairs []nvidiaNVLinkBondedPair
	for _, line := range lines[headerIdx+1:] {
		trimmed := strings.TrimSpace(line)
		if !strings.HasPrefix(trimmed, "GPU") {
			continue
		}
		cells := strings.Fields(trimmed)
		if len(cells) == 0 {
			continue
		}
		rowGPU, err := strconv.Atoi(strings.TrimPrefix(cells[0], "GPU"))
		if err != nil {
			continue
		}
		for colIdx, colGPU := range colIdxToGPU {
			if colGPU == rowGPU {
				continue
			}
			dataIdx := colIdx + 1
			if dataIdx >= len(cells) {
				continue
			}
			m := nvidiaNVLinkNVRe.FindStringSubmatch(cells[dataIdx])
			if len(m) != 2 {
				continue
			}
			nv, err := strconv.Atoi(m[1])
			if err != nil || nv <= 0 {
				continue
			}
			a, b := rowGPU, colGPU
			if a > b {
				a, b = b, a
			}
			key := [2]int{a, b}
			if seen[key] {
				continue
			}
			seen[key] = true
			pairs = append(pairs, nvidiaNVLinkBondedPair{gpuA: a, gpuB: b, links: nv})
		}
	}
	return pairs
}

type nvidiaNVLinkPort struct {
	active bool
}

var (
	nvidiaNVLinkGPUHeaderRe = regexp.MustCompile(`^GPU (\d+):`)
	nvidiaNVLinkSpeedLineRe = regexp.MustCompile(`^Link (\d+):\s*[\d.]+\s*GB/s`)
	nvidiaNVLinkInactiveRe  = regexp.MustCompile(`^Link (\d+):\s*<inactive>`)
	nvidiaNVLinkErrorLineRe = regexp.MustCompile(`^Link (\d+):\s*(Replay|Recovery|CRC) Errors:\s*(\d+)`)
)

// parseNvidiaNVLinkStatus parses "nvidia-smi nvlink -s" output into per-GPU
// per-link active/inactive state.
func parseNvidiaNVLinkStatus(raw string) map[int][]nvidiaNVLinkPort {
	result := map[int][]nvidiaNVLinkPort{}
	currentGPU := -1
	for _, line := range strings.Split(raw, "\n") {
		trimmed := strings.TrimSpace(line)
		if m := nvidiaNVLinkGPUHeaderRe.FindStringSubmatch(trimmed); m != nil {
			currentGPU, _ = strconv.Atoi(m[1])
			continue
		}
		if currentGPU < 0 {
			continue
		}
		if nvidiaNVLinkInactiveRe.MatchString(trimmed) {
			result[currentGPU] = append(result[currentGPU], nvidiaNVLinkPort{active: false})
			continue
		}
		if nvidiaNVLinkSpeedLineRe.MatchString(trimmed) {
			result[currentGPU] = append(result[currentGPU], nvidiaNVLinkPort{active: true})
		}
	}
	return result
}

// parseNvidiaNVLinkErrors parses "nvidia-smi nvlink -e" output into, per GPU
// then link index, [replay, recovery, crc] error counts.
func parseNvidiaNVLinkErrors(raw string) map[int]map[int][3]int64 {
	result := map[int]map[int][3]int64{}
	currentGPU := -1
	for _, line := range strings.Split(raw, "\n") {
		trimmed := strings.TrimSpace(line)
		if m := nvidiaNVLinkGPUHeaderRe.FindStringSubmatch(trimmed); m != nil {
			currentGPU, _ = strconv.Atoi(m[1])
			continue
		}
		if currentGPU < 0 {
			continue
		}
		m := nvidiaNVLinkErrorLineRe.FindStringSubmatch(trimmed)
		if m == nil {
			continue
		}
		linkIdx, _ := strconv.Atoi(m[1])
		count, _ := strconv.ParseInt(m[3], 10, 64)
		if result[currentGPU] == nil {
			result[currentGPU] = map[int][3]int64{}
		}
		c := result[currentGPU][linkIdx]
		switch m[2] {
		case "Replay":
			c[0] = count
		case "Recovery":
			c[1] = count
		case "CRC":
			c[2] = count
		}
		result[currentGPU][linkIdx] = c
	}
	return result
}

// ---------------------------------------------------------------------------
// Confidential Computing (folded in — see RunNvidiaConfigCheckPack doc)
// ---------------------------------------------------------------------------

// parseConfComputeFields parses the indented "Key : Value" block emitted by
// `nvidia-smi conf-compute -q`, e.g.:
//
//	CC State                : OFF
//	Multi-GPU Mode          : Protected PCIe
//	CPU CC Capabilities     : INTEL TDX
//	GPU CC Capabilities     : CC Capable
//	CC GPUs Ready State     : Not Ready
func parseConfComputeFields(out []byte) map[string]string {
	fields := map[string]string{}
	for _, line := range strings.Split(string(out), "\n") {
		idx := strings.Index(line, ":")
		if idx < 0 {
			continue
		}
		key := strings.TrimSpace(line[:idx])
		val := strings.TrimSpace(line[idx+1:])
		if key == "" || val == "" {
			continue
		}
		fields[key] = val
	}
	return fields
}

// filterConfComputeDmesgLines returns the dmesg lines relevant to CPU
// Confidential Computing support (AMD SEV/SEV-ES/SEV-SNP, Intel TDX).
func filterConfComputeDmesgLines(dmesgOut []byte) []string {
	var lines []string
	for _, raw := range bytes.Split(dmesgOut, []byte("\n")) {
		lower := strings.ToLower(string(raw))
		if strings.Contains(lower, "sev") || strings.Contains(lower, "tdx") {
			lines = append(lines, string(raw))
		}
	}
	return lines
}

// ---------------------------------------------------------------------------
// Report rendering
// ---------------------------------------------------------------------------

func renderNvidiaConfigCheckSummary(status NvidiaConfigCheckStatus) string {
	var b strings.Builder
	fmt.Fprintf(&b, "run_at_utc=%s\n", status.CollectedAt.Format(time.RFC3339))
	fmt.Fprintf(&b, "gpu_count=%d\n", len(status.GPUs))
	eccIssues, migIssues, powerIssues := 0, 0, 0
	for _, g := range status.GPUs {
		for _, issue := range g.Issues {
			switch {
			case strings.Contains(issue, "ECC"):
				eccIssues++
			case strings.Contains(issue, "MIG"):
				migIssues++
			case strings.Contains(issue, "power limit"):
				powerIssues++
			}
		}
	}
	fmt.Fprintf(&b, "gpu_ecc_issues=%d\n", eccIssues)
	fmt.Fprintf(&b, "gpu_mig_pending_issues=%d\n", migIssues)
	fmt.Fprintf(&b, "gpu_power_limit_issues=%d\n", powerIssues)
	fmt.Fprintf(&b, "nvlink_pairs_checked=%d\n", len(status.NVLinkPairs))
	nvlinkIssues := 0
	for _, p := range status.NVLinkPairs {
		if len(p.Issues) > 0 {
			nvlinkIssues++
		}
	}
	fmt.Fprintf(&b, "nvlink_pairs_with_issues=%d\n", nvlinkIssues)
	fmt.Fprintf(&b, "cc_ready=%t\n", status.CCReady)
	fmt.Fprintf(&b, "cc_state=%s\n", status.CCState)
	fmt.Fprintf(&b, "cpu_cc_capability=%s\n", status.CPUCCCapability)
	fmt.Fprintf(&b, "gpu_cc_capability=%s\n", status.GPUCCCapability)
	if len(status.Warnings) == 0 {
		fmt.Fprintln(&b, "overall_status=OK")
	} else {
		fmt.Fprintln(&b, "overall_status=FAILED")
	}
	return b.String()
}

func renderNvidiaConfigCheckReport(status NvidiaConfigCheckStatus) string {
	var b strings.Builder
	line := strings.Repeat("=", 80)
	b.WriteString(line + "\n")
	b.WriteString("NVIDIA GPU Configuration & NVLink Topology Check\n")
	b.WriteString(line + "\n\n")

	b.WriteString("-- GPU Configuration ----------------------------------------------------------\n")
	for _, g := range status.GPUs {
		verdict := "OK"
		if len(g.Issues) > 0 {
			verdict = "ISSUES FOUND"
		}
		fmt.Fprintf(&b, "  GPU %d (%s): %s\n", g.Index, g.Name, verdict)
		fmt.Fprintf(&b, "    ECC: %s | MIG: %s (pending %s) | Power: %.0fW / %.0fW default\n",
			nonEmptyOr(g.ECCCurrent, "unknown"), nonEmptyOr(g.MIGCurrent, "unknown"), nonEmptyOr(g.MIGPending, "unknown"),
			g.PowerLimitW, g.PowerDefaultLimitW)
		for _, issue := range g.Issues {
			fmt.Fprintf(&b, "    - %s\n", issue)
		}
	}
	b.WriteString("\n-- NVLink Topology --------------------------------------------------------------\n")
	if len(status.NVLinkPairs) == 0 {
		b.WriteString("  No NVLink-bonded GPU pairs found.\n")
	}
	for _, p := range status.NVLinkPairs {
		verdict := "OK"
		if len(p.Issues) > 0 {
			verdict = "ISSUES FOUND"
		}
		fmt.Fprintf(&b, "  GPU%d <-> GPU%d: %d/%d links active (topology expects %d): %s\n",
			p.GPUA, p.GPUB, p.ActiveLinks, p.TotalLinks, p.ExpectedLinks, verdict)
		for _, issue := range p.Issues {
			fmt.Fprintf(&b, "    - %s\n", issue)
		}
	}

	b.WriteString("\n-- Confidential Computing (informational) ---------------------------------------\n")
	verdict := "NOT READY"
	if status.CCReady {
		verdict = "READY"
	}
	fmt.Fprintf(&b, "  Verdict: %s\n", verdict)
	fmt.Fprintf(&b, "  CC State: %s | CPU CC: %s | GPU CC: %s\n",
		nonEmptyOr(status.CCState, "unknown"), nonEmptyOr(status.CPUCCCapability, "unknown"), nonEmptyOr(status.GPUCCCapability, "unknown"))

	if len(status.Notes) > 0 {
		b.WriteString("\n-- Notes -------------------------------------------------------------------------\n")
		for _, n := range status.Notes {
			fmt.Fprintf(&b, "  - %s\n", n)
		}
	}

	fmt.Fprintf(&b, "\nCollected : %s\n", status.CollectedAt.Format("2006-01-02 15:04:05 UTC"))
	b.WriteString(line + "\n")
	return b.String()
}

func nonEmptyOr(v, fallback string) string {
	if strings.TrimSpace(v) == "" {
		return fallback
	}
	return v
}