bee/audit/internal/platform/gpu_metrics.go

package platform

import (
	"bytes"
	"fmt"
	"math"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"
)

// GPUMetricRow is one telemetry sample from nvidia-smi during a stress test.
type GPUMetricRow struct {
	Stage                 string  `json:"stage,omitempty"`
	StageStartSec         float64 `json:"stage_start_sec,omitempty"`
	StageEndSec           float64 `json:"stage_end_sec,omitempty"`
	ElapsedSec            float64 `json:"elapsed_sec"`
	GPUIndex              int     `json:"index"`
	TempC                 float64 `json:"temp_c"`
	UsagePct              float64 `json:"usage_pct"`
	MemUsagePct           float64 `json:"mem_usage_pct"`
	PowerW                float64 `json:"power_w"`
	ClockMHz              float64 `json:"clock_mhz"`
	MemClockMHz           float64 `json:"mem_clock_mhz"`
	FanAvgRPM             float64 `json:"fan_avg_rpm,omitempty"`
	FanDutyCyclePct       float64 `json:"fan_duty_cycle_pct,omitempty"`
	FanDutyCycleAvailable bool    `json:"fan_duty_cycle_available,omitempty"`
}

// sampleGPUMetrics runs nvidia-smi once and returns current metrics for each GPU.
func sampleGPUMetrics(gpuIndices []int) ([]GPUMetricRow, error) {
	args := []string{
		"--query-gpu=index,temperature.gpu,utilization.gpu,utilization.memory,power.draw,clocks.current.graphics,clocks.current.memory",
		"--format=csv,noheader,nounits",
	}
	if len(gpuIndices) > 0 {
		ids := make([]string, len(gpuIndices))
		for i, idx := range gpuIndices {
			ids[i] = strconv.Itoa(idx)
		}
		args = append([]string{"--id=" + strings.Join(ids, ",")}, args...)
	}
	out, err := exec.Command("nvidia-smi", args...).Output()
	if err != nil {
		return nil, err
	}
	var rows []GPUMetricRow
	for _, line := range strings.Split(strings.TrimSpace(string(out)), "\n") {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		parts := strings.Split(line, ", ")
		if len(parts) < 7 {
			continue
		}
		idx, _ := strconv.Atoi(strings.TrimSpace(parts[0]))
		rows = append(rows, GPUMetricRow{
			GPUIndex:    idx,
			TempC:       parseGPUFloat(parts[1]),
			UsagePct:    parseGPUFloat(parts[2]),
			MemUsagePct: parseGPUFloat(parts[3]),
			PowerW:      parseGPUFloat(parts[4]),
			ClockMHz:    parseGPUFloat(parts[5]),
			MemClockMHz: parseGPUFloat(parts[6]),
		})
	}
	return rows, nil
}

func parseGPUFloat(s string) float64 {
	s = strings.TrimSpace(s)
	if s == "N/A" || s == "[Not Supported]" || s == "" {
		return 0
	}
	v, _ := strconv.ParseFloat(s, 64)
	return v
}

// SampleGPUMetrics runs nvidia-smi once and returns current metrics for each GPU.
func SampleGPUMetrics(gpuIndices []int) ([]GPUMetricRow, error) {
	return sampleGPUMetrics(gpuIndices)
}

// sampleAMDGPUMetrics queries rocm-smi for live GPU metrics.
func sampleAMDGPUMetrics() ([]GPUMetricRow, error) {
	out, err := runROCmSMI("--showtemp", "--showuse", "--showpower", "--showmemuse", "--csv")
	if err != nil {
		return nil, err
	}
	lines := strings.Split(strings.TrimSpace(string(out)), "\n")
	if len(lines) < 2 {
		return nil, fmt.Errorf("rocm-smi: insufficient output")
	}

	// Parse header to find column indices by name.
	headers := strings.Split(lines[0], ",")
	colIdx := func(keywords ...string) int {
		for i, h := range headers {
			hl := strings.ToLower(strings.TrimSpace(h))
			for _, kw := range keywords {
				if strings.Contains(hl, kw) {
					return i
				}
			}
		}
		return -1
	}
	idxTemp := colIdx("sensor edge", "temperature (c)", "temp")
	idxUse := colIdx("gpu use (%)")
	idxMem := colIdx("vram%", "memory allocated")
	idxPow := colIdx("average graphics package power", "power (w)")

	var rows []GPUMetricRow
	for _, line := range lines[1:] {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		parts := strings.Split(line, ",")
		idx := len(rows)
		row := GPUMetricRow{GPUIndex: idx}
		get := func(i int) float64 {
			if i < 0 || i >= len(parts) {
				return 0
			}
			v := strings.TrimSpace(parts[i])
			if strings.EqualFold(v, "n/a") {
				return 0
			}
			return parseGPUFloat(v)
		}
		row.TempC = get(idxTemp)
		row.UsagePct = get(idxUse)
		row.MemUsagePct = get(idxMem)
		row.PowerW = get(idxPow)
		rows = append(rows, row)
	}
	if len(rows) == 0 {
		return nil, fmt.Errorf("rocm-smi: no GPU rows parsed")
	}
	return rows, nil
}

// WriteGPUMetricsCSV writes collected rows as a CSV file.
func WriteGPUMetricsCSV(path string, rows []GPUMetricRow) error {
	var b bytes.Buffer
	b.WriteString("stage,elapsed_sec,gpu_index,temperature_c,usage_pct,mem_usage_pct,power_w,clock_mhz,mem_clock_mhz,fan_avg_rpm,fan_duty_cycle_pct,fan_duty_cycle_available\n")
	for _, r := range rows {
		dutyAvail := 0
		if r.FanDutyCycleAvailable {
			dutyAvail = 1
		}
		fmt.Fprintf(&b, "%s,%.1f,%d,%.1f,%.1f,%.1f,%.1f,%.0f,%.0f,%.0f,%.1f,%d\n",
			strconv.Quote(strings.TrimSpace(r.Stage)), r.ElapsedSec, r.GPUIndex, r.TempC, r.UsagePct, r.MemUsagePct, r.PowerW, r.ClockMHz, r.MemClockMHz, r.FanAvgRPM, r.FanDutyCyclePct, dutyAvail)
	}
	return os.WriteFile(path, b.Bytes(), 0644)
}

type gpuMetricStageSpan struct {
	Name  string
	Start float64
	End   float64
}

// WriteGPUMetricsHTML writes a standalone HTML file with one SVG chart per GPU.
func WriteGPUMetricsHTML(path string, rows []GPUMetricRow) error {
	// Group by GPU index preserving order.
	seen := make(map[int]bool)
	var order []int
	gpuMap := make(map[int][]GPUMetricRow)
	for _, r := range rows {
		if !seen[r.GPUIndex] {
			seen[r.GPUIndex] = true
			order = append(order, r.GPUIndex)
		}
		gpuMap[r.GPUIndex] = append(gpuMap[r.GPUIndex], r)
	}

	stageSpans := buildGPUMetricStageSpans(rows)
	stageColorByName := make(map[string]string, len(stageSpans))
	for i, span := range stageSpans {
		stageColorByName[span.Name] = gpuMetricStagePalette[i%len(gpuMetricStagePalette)]
	}

	var legend strings.Builder
	if len(stageSpans) > 0 {
		legend.WriteString(`<div class="stage-legend">`)
		for _, span := range stageSpans {
			fmt.Fprintf(&legend, `<span class="stage-chip"><span class="stage-swatch" style="background:%s"></span>%s</span>`,
				stageColorByName[span.Name], gpuHTMLEscape(span.Name))
		}
		legend.WriteString(`</div>`)
	}

	var svgs strings.Builder
	for _, gpuIdx := range order {
		svgs.WriteString(drawGPUChartSVG(gpuMap[gpuIdx], gpuIdx, stageSpans, stageColorByName))
		svgs.WriteString("\n")
	}

	ts := time.Now().UTC().Format("2006-01-02 15:04:05 UTC")
	html := fmt.Sprintf(`<!DOCTYPE html>
<html><head>
<meta charset="utf-8">
<title>GPU Stress Test Metrics</title>
<style>
:root{--bg:#fff;--surface:#fff;--surface-2:#f9fafb;--border:rgba(34,36,38,.15);--border-lite:rgba(34,36,38,.1);--ink:rgba(0,0,0,.87);--muted:rgba(0,0,0,.6)}
*{box-sizing:border-box}
body{font:14px/1.5 Lato,"Helvetica Neue",Arial,Helvetica,sans-serif;background:var(--bg);color:var(--ink);margin:0}
.page{padding:24px}
.card{background:var(--surface);border:1px solid var(--border);border-radius:4px;box-shadow:0 1px 2px rgba(34,36,38,.15);overflow:hidden}
.card-head{padding:11px 16px;background:var(--surface-2);border-bottom:1px solid var(--border);font-weight:700;font-size:13px}
.card-body{padding:16px}
h1{font-size:22px;margin:0 0 6px}
p{color:var(--muted);font-size:13px;margin:0 0 16px}
.stage-legend{display:flex;flex-wrap:wrap;gap:10px;margin:0 0 16px}
.stage-chip{display:inline-flex;align-items:center;gap:8px;padding:4px 10px;border-radius:999px;background:var(--surface-2);border:1px solid var(--border-lite);font-size:12px}
.stage-swatch{display:inline-block;width:12px;height:12px;border-radius:999px}
.chart-block{margin-top:16px}
</style>
</head><body>
<div class="page">
<div class="card">
<div class="card-head">GPU Stress Test Metrics</div>
<div class="card-body">
<h1>GPU Stress Test Metrics</h1>
<p>Generated %s</p>
%s
<div class="chart-block">%s</div>
</div>
</div>
</div>
</body></html>`, ts, legend.String(), svgs.String())

	return os.WriteFile(path, []byte(html), 0644)
}

// drawGPUChartSVG generates a self-contained SVG chart for one GPU.
func drawGPUChartSVG(rows []GPUMetricRow, gpuIdx int, stageSpans []gpuMetricStageSpan, stageColorByName map[string]string) string {
	// Layout
	const W, H = 960, 520
	const plotX1 = 120 // usage axis / chart left border
	const plotX2 = 840 // power axis / chart right border
	const plotY1 = 70  // top
	const plotY2 = 465 // bottom  (PH = 395)
	const PW = plotX2 - plotX1
	const PH = plotY2 - plotY1
	// Outer axes
	const tempAxisX = 60   // temp axis line
	const clockAxisX = 900 // clock axis line

	colors := [4]string{"#e74c3c", "#3498db", "#2ecc71", "#f39c12"}
	seriesLabel := [4]string{
		fmt.Sprintf("GPU %d Temp (°C)", gpuIdx),
		fmt.Sprintf("GPU %d Usage (%%)", gpuIdx),
		fmt.Sprintf("GPU %d Power (W)", gpuIdx),
		fmt.Sprintf("GPU %d Clock (MHz)", gpuIdx),
	}
	axisLabel := [4]string{"Temperature (°C)", "GPU Usage (%)", "Power (W)", "Clock (MHz)"}

	// Extract series
	t := make([]float64, len(rows))
	vals := [4][]float64{}
	for i := range vals {
		vals[i] = make([]float64, len(rows))
	}
	for i, r := range rows {
		t[i] = r.ElapsedSec
		vals[0][i] = r.TempC
		vals[1][i] = r.UsagePct
		vals[2][i] = r.PowerW
		vals[3][i] = r.ClockMHz
	}

	tMin, tMax := gpuMinMax(t)
	type axisScale struct {
		ticks    []float64
		min, max float64
	}
	var axes [4]axisScale
	for i := 0; i < 4; i++ {
		mn, mx := gpuMinMax(vals[i])
		tks := gpuNiceTicks(mn, mx, 8)
		axes[i] = axisScale{ticks: tks, min: tks[0], max: tks[len(tks)-1]}
	}

	xv := func(tv float64) float64 {
		if tMax == tMin {
			return float64(plotX1)
		}
		return float64(plotX1) + (tv-tMin)/(tMax-tMin)*float64(PW)
	}
	yv := func(v float64, ai int) float64 {
		a := axes[ai]
		if a.max == a.min {
			return float64(plotY1 + PH/2)
		}
		return float64(plotY2) - (v-a.min)/(a.max-a.min)*float64(PH)
	}

	var b strings.Builder

	fmt.Fprintf(&b, `<svg xmlns="http://www.w3.org/2000/svg" width="%d" height="%d"`+
		` style="background:#fff;border-radius:8px;display:block;margin:0 auto 24px;`+
		`box-shadow:0 2px 12px rgba(0,0,0,.12)">`+"\n", W, H)

	// Title
	fmt.Fprintf(&b, `<text x="%d" y="22" text-anchor="middle" font-family="sans-serif"`+
		` font-size="14" font-weight="bold" fill="#333">GPU Stress Test Metrics — GPU %d</text>`+"\n",
		plotX1+PW/2, gpuIdx)

	// Horizontal grid (align to temp axis ticks)
	b.WriteString(`<g stroke="#e0e0e0" stroke-width="0.5">` + "\n")
	for _, tick := range axes[0].ticks {
		y := yv(tick, 0)
		if y < float64(plotY1) || y > float64(plotY2) {
			continue
		}
		fmt.Fprintf(&b, `<line x1="%d" y1="%.1f" x2="%d" y2="%.1f"/>`+"\n",
			plotX1, y, plotX2, y)
	}
	// Vertical grid
	xTicks := gpuNiceTicks(tMin, tMax, 10)
	for _, tv := range xTicks {
		x := xv(tv)
		if x < float64(plotX1) || x > float64(plotX2) {
			continue
		}
		fmt.Fprintf(&b, `<line x1="%.1f" y1="%d" x2="%.1f" y2="%d"/>`+"\n",
			x, plotY1, x, plotY2)
	}
	b.WriteString("</g>\n")

	// Stage backgrounds
	for _, span := range stageSpans {
		x1 := xv(span.Start)
		x2 := xv(span.End)
		if x2 < x1 {
			x1, x2 = x2, x1
		}
		if x2-x1 < 1 {
			x2 = x1 + 1
		}
		color := stageColorByName[span.Name]
		fmt.Fprintf(&b, `<rect x="%.1f" y="%d" width="%.1f" height="%d" fill="%s" fill-opacity="0.18"/>`+"\n",
			x1, plotY1, x2-x1, PH, color)
		fmt.Fprintf(&b, `<text x="%.1f" y="%d" font-family="sans-serif" font-size="10" fill="#444" text-anchor="middle">%s</text>`+"\n",
			x1+(x2-x1)/2, plotY1+12, gpuHTMLEscape(span.Name))
	}

	// Chart border
	fmt.Fprintf(&b, `<rect x="%d" y="%d" width="%d" height="%d"`+
		` fill="none" stroke="#333" stroke-width="1"/>`+"\n",
		plotX1, plotY1, PW, PH)

	// X axis ticks and labels
	b.WriteString(`<g font-family="sans-serif" font-size="11" fill="#333" text-anchor="middle">` + "\n")
	for _, tv := range xTicks {
		x := xv(tv)
		if x < float64(plotX1) || x > float64(plotX2) {
			continue
		}
		fmt.Fprintf(&b, `<text x="%.1f" y="%d">%s</text>`+"\n", x, plotY2+18, gpuFormatTick(tv))
		fmt.Fprintf(&b, `<line x1="%.1f" y1="%d" x2="%.1f" y2="%d" stroke="#333" stroke-width="1"/>`+"\n",
			x, plotY2, x, plotY2+4)
	}
	b.WriteString("</g>\n")
	fmt.Fprintf(&b, `<text x="%d" y="%d" font-family="sans-serif" font-size="13"`+
		` fill="#333" text-anchor="middle">Time (seconds)</text>`+"\n",
		plotX1+PW/2, plotY2+38)

	// Y axes: [tempAxisX, plotX1, plotX2, clockAxisX]
	axisLineX := [4]int{tempAxisX, plotX1, plotX2, clockAxisX}
	axisRight := [4]bool{false, false, true, true}
	// Label x positions (for rotated vertical text)
	axisLabelX := [4]int{10, 68, 868, 950}

	for i := 0; i < 4; i++ {
		ax := axisLineX[i]
		right := axisRight[i]
		color := colors[i]

		// Axis line
		fmt.Fprintf(&b, `<line x1="%d" y1="%d" x2="%d" y2="%d"`+
			` stroke="%s" stroke-width="1"/>`+"\n",
			ax, plotY1, ax, plotY2, color)

		// Ticks and tick labels
		fmt.Fprintf(&b, `<g font-family="sans-serif" font-size="10" fill="%s">`+"\n", color)
		for _, tick := range axes[i].ticks {
			y := yv(tick, i)
			if y < float64(plotY1) || y > float64(plotY2) {
				continue
			}
			dx := -5
			textX := ax - 8
			anchor := "end"
			if right {
				dx = 5
				textX = ax + 8
				anchor = "start"
			}
			fmt.Fprintf(&b, `<line x1="%d" y1="%.1f" x2="%d" y2="%.1f"`+
				` stroke="%s" stroke-width="1"/>`+"\n",
				ax, y, ax+dx, y, color)
			fmt.Fprintf(&b, `<text x="%d" y="%.1f" text-anchor="%s" dy="4">%s</text>`+"\n",
				textX, y, anchor, gpuFormatTick(tick))
		}
		b.WriteString("</g>\n")

		// Axis label (rotated)
		lx := axisLabelX[i]
		fmt.Fprintf(&b, `<text transform="translate(%d,%d) rotate(-90)"`+
			` font-family="sans-serif" font-size="12" fill="%s" text-anchor="middle">%s</text>`+"\n",
			lx, plotY1+PH/2, color, axisLabel[i])
	}

	// Data lines
	for i := 0; i < 4; i++ {
		var pts strings.Builder
		for j := range rows {
			x := xv(t[j])
			y := yv(vals[i][j], i)
			if j == 0 {
				fmt.Fprintf(&pts, "%.1f,%.1f", x, y)
			} else {
				fmt.Fprintf(&pts, " %.1f,%.1f", x, y)
			}
		}
		fmt.Fprintf(&b, `<polyline points="%s" fill="none" stroke="%s" stroke-width="1.5"/>`+"\n",
			pts.String(), colors[i])
	}

	// Legend
	const legendY = 42
	for i := 0; i < 4; i++ {
		lx := plotX1 + i*(PW/4) + 10
		fmt.Fprintf(&b, `<line x1="%d" y1="%d" x2="%d" y2="%d"`+
			` stroke="%s" stroke-width="2"/>`+"\n",
			lx, legendY, lx+20, legendY, colors[i])
		fmt.Fprintf(&b, `<text x="%d" y="%d" font-family="sans-serif" font-size="12" fill="#333">%s</text>`+"\n",
			lx+25, legendY+4, seriesLabel[i])
	}

	b.WriteString("</svg>\n")
	return b.String()
}

func gpuMinMax(vals []float64) (float64, float64) {
	if len(vals) == 0 {
		return 0, 1
	}
	mn, mx := vals[0], vals[0]
	for _, v := range vals[1:] {
		if v < mn {
			mn = v
		}
		if v > mx {
			mx = v
		}
	}
	return mn, mx
}

func gpuNiceTicks(mn, mx float64, targetCount int) []float64 {
	if mn == mx {
		mn -= 1
		mx += 1
	}
	r := mx - mn
	step := math.Pow(10, math.Floor(math.Log10(r/float64(targetCount))))
	for _, f := range []float64{1, 2, 5, 10} {
		if r/(f*step) <= float64(targetCount)*1.5 {
			step = f * step
			break
		}
	}
	lo := math.Floor(mn/step) * step
	hi := math.Ceil(mx/step) * step
	var ticks []float64
	for v := lo; v <= hi+step*0.001; v += step {
		ticks = append(ticks, math.Round(v*1e9)/1e9)
	}
	return ticks
}

func gpuFormatTick(v float64) string {
	if v == math.Trunc(v) {
		return strconv.Itoa(int(v))
	}
	return strconv.FormatFloat(v, 'f', 1, 64)
}

var gpuMetricStagePalette = []string{
	"#d95c5c",
	"#2185d0",
	"#21ba45",
	"#f2c037",
	"#6435c9",
	"#00b5ad",
	"#a5673f",
}

func buildGPUMetricStageSpans(rows []GPUMetricRow) []gpuMetricStageSpan {
	var spans []gpuMetricStageSpan
	for _, row := range rows {
		name := strings.TrimSpace(row.Stage)
		if name == "" {
			name = "run"
		}
		start := row.StageStartSec
		end := row.StageEndSec
		if end <= start {
			start = row.ElapsedSec
			end = row.ElapsedSec
		}
		if len(spans) == 0 || spans[len(spans)-1].Name != name {
			spans = append(spans, gpuMetricStageSpan{Name: name, Start: start, End: end})
			continue
		}
		if start < spans[len(spans)-1].Start {
			spans[len(spans)-1].Start = start
		}
		if end > spans[len(spans)-1].End {
			spans[len(spans)-1].End = end
		}
	}
	for i := range spans {
		if spans[i].End <= spans[i].Start {
			spans[i].End = spans[i].Start + 1
		}
	}
	return spans
}

var gpuHTMLReplacer = strings.NewReplacer(
	"&", "&amp;",
	"<", "&lt;",
	">", "&gt;",
	`"`, "&quot;",
	"'", "&#39;",
)

func gpuHTMLEscape(s string) string {
	return gpuHTMLReplacer.Replace(s)
}