bee/audit/internal/platform/nvidia_pcie_bandwidth.go

package platform

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"bee/audit/internal/collector"
)

// nvidiaPCIeBandwidthFinding is one GPU's post-load PCIe link-speed result.
type nvidiaPCIeBandwidthFinding struct {
	Index       int
	BDF         string
	BeforeSpeed string
	AfterSpeed  string
	MaxSpeed    string
	Width       int
	MaxWidth    int
	Degraded    bool
	Supported   bool
	Sampled     bool
}

// captureNvidiaPCIeLinkBaseline records link capabilities immediately before
// the existing nvbandwidth SAT. The post-load half is completed by
// finishNvidiaPCIeLinkCheck as soon as that same SAT returns, so one real PCIe
// traffic pass provides both the DCGM verdict and the negotiated-link verdict.
func captureNvidiaPCIeLinkBaseline(gpuIndices []int) ([]nvidiaPCIeBandwidthFinding, error) {
	bdfByIndex, err := gpuIndexToBDF(gpuIndices)
	if err != nil {
		return nil, fmt.Errorf("resolve GPU BDFs: %w", err)
	}
	indices := make([]int, 0, len(bdfByIndex))
	for idx := range bdfByIndex {
		indices = append(indices, idx)
	}
	sort.Ints(indices)

	findings := make([]nvidiaPCIeBandwidthFinding, 0, len(indices))
	for _, idx := range indices {
		bdf := bdfByIndex[idx]
		before, beforeOK := readPCIeSysfsString(bdf, "current_link_speed")
		max, maxOK := readPCIeSysfsString(bdf, "max_link_speed")
		maxWidth, widthOK := readPCIeSysfsInt(bdf, "max_link_width")
		findings = append(findings, nvidiaPCIeBandwidthFinding{
			Index: idx, BDF: bdf, BeforeSpeed: before, MaxSpeed: max, MaxWidth: maxWidth,
			Supported: beforeOK && maxOK && widthOK,
		})
	}
	return findings, nil
}

func finishNvidiaPCIeLinkCheck(runDir string, findings []nvidiaPCIeBandwidthFinding) error {
	report := renderNvidiaPCIeBandwidthReport(findings)
	if err := os.WriteFile(filepath.Join(runDir, "nvidia-pcie-link-under-load-report.txt"), []byte(report), 0644); err != nil {
		return err
	}
	return appendNvidiaPCIeLinkSummary(filepath.Join(runDir, "summary.txt"), findings)
}

func sampleNvidiaPCIeLinkAfterLoad(findings []nvidiaPCIeBandwidthFinding) {
	for i := range findings {
		bdf := findings[i].BDF
		after, afterOK := readPCIeSysfsString(bdf, "current_link_speed")
		width, widthOK := readPCIeSysfsInt(bdf, "current_link_width")
		findings[i].AfterSpeed = after
		findings[i].Width = width
		findings[i].Sampled = true
		findings[i].Supported = findings[i].Supported && afterOK && widthOK
		findings[i].Degraded = findings[i].Supported &&
			(after != findings[i].MaxSpeed || width != findings[i].MaxWidth)
	}
}

// gpuIndexToBDF resolves each of gpuIndices to its PCI BDF via nvidia-smi.
// An empty gpuIndices resolves every GPU nvidia-smi reports.
func gpuIndexToBDF(gpuIndices []int) (map[int]string, error) {
	out, err := satExecCommand("nvidia-smi", "--query-gpu=index,pci.bus_id", "--format=csv,noheader,nounits").Output()
	if err != nil {
		return nil, fmt.Errorf("nvidia-smi: %w", err)
	}
	want := make(map[int]struct{}, len(gpuIndices))
	for _, idx := range gpuIndices {
		want[idx] = struct{}{}
	}
	filterByIndex := len(want) > 0

	result := make(map[int]string)
	for _, line := range strings.Split(strings.TrimSpace(string(out)), "\n") {
		fields := strings.SplitN(line, ",", 2)
		if len(fields) != 2 {
			continue
		}
		idx, err := strconv.Atoi(strings.TrimSpace(fields[0]))
		if err != nil {
			continue
		}
		if filterByIndex {
			if _, ok := want[idx]; !ok {
				continue
			}
		}
		result[idx] = normalizeNvidiaBDF(strings.TrimSpace(fields[1]))
	}
	return result, nil
}

func readPCIeSysfsString(bdf, attr string) (string, bool) {
	raw, err := satReadFile(filepath.Join("/sys/bus/pci/devices", bdf, attr))
	if err != nil {
		return "", false
	}
	v := strings.TrimSpace(string(raw))
	if v == "" {
		return "", false
	}
	return collector.NormalizePCILinkSpeed(v), true
}

func readPCIeSysfsInt(bdf, attr string) (int, bool) {
	raw, err := satReadFile(filepath.Join("/sys/bus/pci/devices", bdf, attr))
	if err != nil {
		return 0, false
	}
	v, err := strconv.Atoi(strings.TrimSpace(string(raw)))
	if err != nil || v < 0 {
		return 0, false
	}
	return v, true
}

func appendNvidiaPCIeLinkSummary(path string, findings []nvidiaPCIeBandwidthFinding) error {
	raw, err := os.ReadFile(path)
	if err != nil {
		return err
	}
	var extra strings.Builder
	fmt.Fprintf(&extra, "pcie_gpu_count=%d\n", len(findings))
	degraded := 0
	supported := 0
	var reasons []string
	for _, f := range findings {
		status := "UNSUPPORTED"
		if f.Supported && f.Sampled {
			supported++
			status = statusLabel(!f.Degraded)
		}
		fmt.Fprintf(&extra, "pcie_gpu%d_status=%s\n", f.Index, status)
		if f.Degraded {
			degraded++
			reasons = append(reasons, fmt.Sprintf("GPU%d (%s): under load at %s x%d, capable of %s x%d",
				f.Index, f.BDF, f.AfterSpeed, f.Width, f.MaxSpeed, f.MaxWidth))
		}
	}
	fmt.Fprintf(&extra, "pcie_degraded=%d\n", degraded)
	linkStatus := "UNSUPPORTED"
	if degraded > 0 {
		linkStatus = "FAILED"
	} else if supported == len(findings) && supported > 0 {
		linkStatus = "OK"
	}
	fmt.Fprintf(&extra, "pcie_link_under_load_status=%s\n", linkStatus)
	if len(reasons) > 0 {
		fmt.Fprintf(&extra, "pcie_link_under_load_detail=%s\n", strings.Join(reasons, "; "))
	}
	if linkStatus == "FAILED" {
		lines := strings.Split(string(raw), "\n")
		for i := range lines {
			if strings.HasPrefix(lines[i], "overall_status=") {
				lines[i] = "overall_status=FAILED"
				break
			}
		}
		raw = []byte(strings.Join(lines, "\n"))
	}
	if len(raw) > 0 && raw[len(raw)-1] != '\n' {
		raw = append(raw, '\n')
	}
	raw = append(raw, extra.String()...)
	return os.WriteFile(path, raw, 0644)
}

func statusLabel(ok bool) string {
	if ok {
		return "OK"
	}
	return "FAILED"
}

func renderNvidiaPCIeBandwidthReport(findings []nvidiaPCIeBandwidthFinding) string {
	var b strings.Builder
	line := strings.Repeat("=", 80)
	b.WriteString(line + "\n")
	b.WriteString("NVIDIA GPU PCIe Link-Under-Load Check\n")
	b.WriteString(line + "\n\n")
	for _, f := range findings {
		verdict := "UNSUPPORTED"
		if f.Supported && f.Sampled && !f.Degraded {
			verdict = "OK"
		} else if f.Degraded {
			verdict = "DEGRADED"
		}
		fmt.Fprintf(&b, "GPU%d (%s): %s\n", f.Index, f.BDF, verdict)
		fmt.Fprintf(&b, "    before=%s after=%s max=%s width=%d/%d\n",
			f.BeforeSpeed, f.AfterSpeed, f.MaxSpeed, f.Width, f.MaxWidth)
	}
	return b.String()
}