refactor: modularize audit and harden build validation
This commit is contained in:
Mikhail Chusavitin
@@ -0,0 +1,396 @@
|
||||
package platform
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/csv"
|
||||
"fmt"
|
||||
"math"
|
||||
"os"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
func normalizeNvidiaBenchmarkOptionsForBenchmark(opts NvidiaBenchmarkOptions) NvidiaBenchmarkOptions {
|
||||
switch strings.TrimSpace(strings.ToLower(opts.Profile)) {
|
||||
case NvidiaBenchmarkProfileStability:
|
||||
opts.Profile = NvidiaBenchmarkProfileStability
|
||||
case NvidiaBenchmarkProfileOvernight:
|
||||
opts.Profile = NvidiaBenchmarkProfileOvernight
|
||||
default:
|
||||
opts.Profile = NvidiaBenchmarkProfileStandard
|
||||
}
|
||||
if opts.SizeMB < 0 {
|
||||
opts.SizeMB = 0
|
||||
}
|
||||
opts.ServerPowerSource = normalizeBenchmarkPowerSource(opts.ServerPowerSource)
|
||||
opts.GPUIndices = dedupeSortedIndices(opts.GPUIndices)
|
||||
opts.ExcludeGPUIndices = dedupeSortedIndices(opts.ExcludeGPUIndices)
|
||||
return opts
|
||||
}
|
||||
|
||||
func resolveBenchmarkProfile(profile string) benchmarkProfileSpec {
|
||||
switch strings.TrimSpace(strings.ToLower(profile)) {
|
||||
case NvidiaBenchmarkProfileStability:
|
||||
return benchmarkProfileSpec{Name: NvidiaBenchmarkProfileStability, BaselineSec: 30, WarmupSec: 120, SteadySec: 3600, NCCLSec: 300, CooldownSec: 0}
|
||||
case NvidiaBenchmarkProfileOvernight:
|
||||
return benchmarkProfileSpec{Name: NvidiaBenchmarkProfileOvernight, BaselineSec: 60, WarmupSec: 180, SteadySec: 27000, NCCLSec: 600, CooldownSec: 0}
|
||||
default:
|
||||
return benchmarkProfileSpec{Name: NvidiaBenchmarkProfileStandard, BaselineSec: 15, WarmupSec: 45, SteadySec: 480, NCCLSec: 180, CooldownSec: 0}
|
||||
}
|
||||
}
|
||||
|
||||
// benchmarkGPUInfoQuery describes a nvidia-smi --query-gpu field set to try.
|
||||
// Fields are tried in order; the first successful query wins. Extended fields
|
||||
// (attribute.multiprocessor_count, power.default_limit) are not supported on
|
||||
// all driver versions, so we fall back to the base set if the full query fails.
|
||||
// The minimal fallback omits clock fields entirely — clocks.max.* returns
|
||||
// exit status 2 on some GPU generations (e.g. Blackwell); missing data is
|
||||
// then recovered from nvidia-smi -q.
|
||||
var benchmarkGPUInfoQueries = []struct {
|
||||
fields string
|
||||
extended bool // whether this query includes optional extended fields
|
||||
minimal bool // clock fields omitted; max clocks must be filled separately
|
||||
}{
|
||||
{
|
||||
fields: "index,uuid,name,pci.bus_id,vbios_version,power.limit,clocks.max.graphics,clocks.max.memory,clocks.base.graphics,attribute.multiprocessor_count,power.default_limit",
|
||||
extended: true,
|
||||
},
|
||||
{
|
||||
fields: "index,uuid,name,pci.bus_id,vbios_version,power.limit,clocks.max.graphics,clocks.max.memory,clocks.base.graphics",
|
||||
extended: false,
|
||||
},
|
||||
{
|
||||
fields: "index,uuid,name,pci.bus_id,vbios_version,power.limit",
|
||||
minimal: true,
|
||||
},
|
||||
}
|
||||
|
||||
// enrichGPUInfoWithNvidiaSMIQ fills benchmark GPU metadata from nvidia-smi -q
|
||||
// for fields that may be missing from --query-gpu on some driver versions.
|
||||
func enrichGPUInfoWithNvidiaSMIQ(infoByIndex map[int]benchmarkGPUInfo, nvsmiQ []byte) {
|
||||
if len(infoByIndex) == 0 || len(nvsmiQ) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Build bus_id → index map for matching verbose sections to GPU indices.
|
||||
busToBenchIdx := make(map[string]int, len(infoByIndex))
|
||||
for idx, info := range infoByIndex {
|
||||
if info.BusID != "" {
|
||||
// nvidia-smi -q uses "GPU 00000000:4E:00.0" (8-digit domain),
|
||||
// while --query-gpu returns the same format; normalise to lower.
|
||||
busToBenchIdx[strings.ToLower(strings.TrimSpace(info.BusID))] = idx
|
||||
}
|
||||
}
|
||||
|
||||
// Split the verbose output into per-GPU sections on "^GPU " lines.
|
||||
gpuSectionRe := regexp.MustCompile(`(?m)^GPU\s+([\dA-Fa-f:\.]+)`)
|
||||
maxGfxRe := regexp.MustCompile(`(?i)Max Clocks[\s\S]*?Graphics\s*:\s*(\d+)\s*MHz`)
|
||||
maxMemRe := regexp.MustCompile(`(?i)Max Clocks[\s\S]*?Memory\s*:\s*(\d+)\s*MHz`)
|
||||
defaultPwrRe := regexp.MustCompile(`(?i)Default Power Limit\s*:\s*([0-9.]+)\s*W`)
|
||||
currentPwrRe := regexp.MustCompile(`(?i)Current Power Limit\s*:\s*([0-9.]+)\s*W`)
|
||||
minPwrRe := regexp.MustCompile(`(?i)Min Power Limit\s*:\s*([0-9.]+)\s*W`)
|
||||
maxPwrRe := regexp.MustCompile(`(?i)Max Power Limit\s*:\s*([0-9.]+)\s*W`)
|
||||
smCountRe := regexp.MustCompile(`(?i)Multiprocessor Count\s*:\s*(\d+)`)
|
||||
shutdownTempRe := regexp.MustCompile(`(?i)GPU Shutdown Temp\s*:\s*(\d+)\s*C`)
|
||||
slowdownTempRe := regexp.MustCompile(`(?i)GPU Slowdown Temp\s*:\s*(\d+)\s*C`)
|
||||
|
||||
sectionStarts := gpuSectionRe.FindAllSubmatchIndex(nvsmiQ, -1)
|
||||
for i, loc := range sectionStarts {
|
||||
busID := strings.ToLower(string(nvsmiQ[loc[2]:loc[3]]))
|
||||
benchIdx, ok := busToBenchIdx[busID]
|
||||
if !ok {
|
||||
// Bus IDs from verbose output may have a different domain prefix;
|
||||
// try suffix match on the slot portion (XX:XX.X).
|
||||
for k, v := range busToBenchIdx {
|
||||
if strings.HasSuffix(k, busID) || strings.HasSuffix(busID, k) {
|
||||
benchIdx = v
|
||||
ok = true
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
end := len(nvsmiQ)
|
||||
if i+1 < len(sectionStarts) {
|
||||
end = sectionStarts[i+1][0]
|
||||
}
|
||||
section := nvsmiQ[loc[0]:end]
|
||||
|
||||
info := infoByIndex[benchIdx]
|
||||
|
||||
if info.MaxGraphicsClockMHz == 0 {
|
||||
if m := maxGfxRe.FindSubmatch(section); m != nil {
|
||||
if v, err := strconv.ParseFloat(string(m[1]), 64); err == nil {
|
||||
info.MaxGraphicsClockMHz = v
|
||||
}
|
||||
}
|
||||
}
|
||||
if info.MaxMemoryClockMHz == 0 {
|
||||
if m := maxMemRe.FindSubmatch(section); m != nil {
|
||||
if v, err := strconv.ParseFloat(string(m[1]), 64); err == nil {
|
||||
info.MaxMemoryClockMHz = v
|
||||
}
|
||||
}
|
||||
}
|
||||
if info.DefaultPowerLimitW == 0 {
|
||||
if m := defaultPwrRe.FindSubmatch(section); m != nil {
|
||||
if v, err := strconv.ParseFloat(string(m[1]), 64); err == nil && v > 0 {
|
||||
info.DefaultPowerLimitW = v
|
||||
}
|
||||
}
|
||||
}
|
||||
if info.PowerLimitW == 0 {
|
||||
if m := currentPwrRe.FindSubmatch(section); m != nil {
|
||||
if v, err := strconv.ParseFloat(string(m[1]), 64); err == nil && v > 0 {
|
||||
info.PowerLimitW = v
|
||||
}
|
||||
}
|
||||
}
|
||||
if info.MinPowerLimitW == 0 {
|
||||
if m := minPwrRe.FindSubmatch(section); m != nil {
|
||||
if v, err := strconv.ParseFloat(string(m[1]), 64); err == nil && v > 0 {
|
||||
info.MinPowerLimitW = v
|
||||
}
|
||||
}
|
||||
}
|
||||
if info.MaxPowerLimitW == 0 {
|
||||
if m := maxPwrRe.FindSubmatch(section); m != nil {
|
||||
if v, err := strconv.ParseFloat(string(m[1]), 64); err == nil && v > 0 {
|
||||
info.MaxPowerLimitW = v
|
||||
}
|
||||
}
|
||||
}
|
||||
if info.MultiprocessorCount == 0 {
|
||||
if m := smCountRe.FindSubmatch(section); m != nil {
|
||||
if v, err := strconv.Atoi(string(m[1])); err == nil && v > 0 {
|
||||
info.MultiprocessorCount = v
|
||||
}
|
||||
}
|
||||
}
|
||||
if info.ShutdownTempC == 0 {
|
||||
if m := shutdownTempRe.FindSubmatch(section); m != nil {
|
||||
if v, err := strconv.ParseFloat(string(m[1]), 64); err == nil && v > 0 {
|
||||
info.ShutdownTempC = v
|
||||
}
|
||||
}
|
||||
}
|
||||
if info.SlowdownTempC == 0 {
|
||||
if m := slowdownTempRe.FindSubmatch(section); m != nil {
|
||||
if v, err := strconv.ParseFloat(string(m[1]), 64); err == nil && v > 0 {
|
||||
info.SlowdownTempC = v
|
||||
}
|
||||
}
|
||||
}
|
||||
infoByIndex[benchIdx] = info
|
||||
}
|
||||
}
|
||||
|
||||
func queryBenchmarkGPUInfo(gpuIndices []int) (map[int]benchmarkGPUInfo, error) {
|
||||
var lastErr error
|
||||
for _, q := range benchmarkGPUInfoQueries {
|
||||
args := []string{
|
||||
"--query-gpu=" + q.fields,
|
||||
"--format=csv,noheader,nounits",
|
||||
}
|
||||
if len(gpuIndices) > 0 {
|
||||
args = append([]string{"--id=" + joinIndexList(gpuIndices)}, args...)
|
||||
}
|
||||
out, err := satExecCommand("nvidia-smi", args...).Output()
|
||||
if err != nil {
|
||||
lastErr = fmt.Errorf("nvidia-smi gpu info (%s): %w", q.fields[:min(len(q.fields), 40)], err)
|
||||
continue
|
||||
}
|
||||
|
||||
r := csv.NewReader(strings.NewReader(string(out)))
|
||||
r.TrimLeadingSpace = true
|
||||
r.FieldsPerRecord = -1
|
||||
rows, err := r.ReadAll()
|
||||
if err != nil {
|
||||
lastErr = fmt.Errorf("parse nvidia-smi gpu info: %w", err)
|
||||
continue
|
||||
}
|
||||
|
||||
minFields := 6
|
||||
if !q.minimal {
|
||||
minFields = 9
|
||||
}
|
||||
infoByIndex := make(map[int]benchmarkGPUInfo, len(rows))
|
||||
for _, row := range rows {
|
||||
if len(row) < minFields {
|
||||
continue
|
||||
}
|
||||
idx, err := strconv.Atoi(strings.TrimSpace(row[0]))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
info := benchmarkGPUInfo{
|
||||
Index: idx,
|
||||
UUID: strings.TrimSpace(row[1]),
|
||||
Name: strings.TrimSpace(row[2]),
|
||||
BusID: strings.TrimSpace(row[3]),
|
||||
VBIOS: strings.TrimSpace(row[4]),
|
||||
PowerLimitW: parseBenchmarkFloat(row[5]),
|
||||
}
|
||||
if !q.minimal {
|
||||
info.MaxGraphicsClockMHz = parseBenchmarkFloat(row[6])
|
||||
info.MaxMemoryClockMHz = parseBenchmarkFloat(row[7])
|
||||
if len(row) >= 9 {
|
||||
info.BaseGraphicsClockMHz = parseBenchmarkFloat(row[8])
|
||||
}
|
||||
if q.extended {
|
||||
if len(row) >= 10 {
|
||||
info.MultiprocessorCount = int(parseBenchmarkFloat(row[9]))
|
||||
}
|
||||
if len(row) >= 11 {
|
||||
info.DefaultPowerLimitW = parseBenchmarkFloat(row[10])
|
||||
}
|
||||
}
|
||||
}
|
||||
infoByIndex[idx] = info
|
||||
}
|
||||
return infoByIndex, nil
|
||||
}
|
||||
return nil, lastErr
|
||||
}
|
||||
|
||||
func applyBenchmarkNormalization(ctx context.Context, verboseLog string, gpuIndices []int, infoByIndex map[int]benchmarkGPUInfo, result *NvidiaBenchmarkResult) []benchmarkRestoreAction {
|
||||
if os.Geteuid() != 0 {
|
||||
result.Normalization.Status = "partial"
|
||||
result.Normalization.Notes = append(result.Normalization.Notes, "benchmark normalization skipped: root privileges are required for persistence mode and clock locks")
|
||||
for _, idx := range gpuIndices {
|
||||
result.Normalization.GPUs = append(result.Normalization.GPUs, BenchmarkNormalizationGPU{
|
||||
Index: idx,
|
||||
Notes: []string{"normalization skipped: root privileges are required"},
|
||||
})
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
var restore []benchmarkRestoreAction
|
||||
for _, idx := range gpuIndices {
|
||||
rec := BenchmarkNormalizationGPU{Index: idx}
|
||||
if _, err := runSATCommandCtx(ctx, verboseLog, fmt.Sprintf("normalize-gpu-%d-pm", idx), []string{"nvidia-smi", "-i", strconv.Itoa(idx), "-pm", "1"}, nil, nil); err != nil {
|
||||
rec.PersistenceMode = "failed"
|
||||
rec.Notes = append(rec.Notes, "failed to enable persistence mode")
|
||||
result.Normalization.Status = "partial"
|
||||
} else {
|
||||
rec.PersistenceMode = "applied"
|
||||
}
|
||||
|
||||
if info, ok := infoByIndex[idx]; ok && info.MaxGraphicsClockMHz > 0 {
|
||||
target := int(math.Round(info.MaxGraphicsClockMHz))
|
||||
if out, err := runSATCommandCtx(ctx, verboseLog, fmt.Sprintf("normalize-gpu-%d-lgc", idx), []string{"nvidia-smi", "-i", strconv.Itoa(idx), "-lgc", strconv.Itoa(target)}, nil, nil); err != nil {
|
||||
rec.GPUClockLockStatus = "failed"
|
||||
rec.Notes = append(rec.Notes, "graphics clock lock failed: "+strings.TrimSpace(string(out)))
|
||||
result.Normalization.Status = "partial"
|
||||
} else {
|
||||
rec.GPUClockLockStatus = "applied"
|
||||
rec.GPUClockLockMHz = float64(target)
|
||||
idxCopy := idx
|
||||
restore = append(restore, benchmarkRestoreAction{name: fmt.Sprintf("gpu-%d-rgc", idxCopy), fn: func() {
|
||||
_, _ = runSATCommandCtx(context.Background(), verboseLog, fmt.Sprintf("restore-gpu-%d-rgc", idxCopy), []string{"nvidia-smi", "-i", strconv.Itoa(idxCopy), "-rgc"}, nil, nil)
|
||||
}})
|
||||
}
|
||||
} else {
|
||||
rec.GPUClockLockStatus = "skipped"
|
||||
rec.Notes = append(rec.Notes, "graphics clock lock skipped: gpu inventory unavailable or MaxGraphicsClockMHz=0")
|
||||
result.Normalization.Status = "partial"
|
||||
}
|
||||
|
||||
if info, ok := infoByIndex[idx]; ok && info.MaxMemoryClockMHz > 0 {
|
||||
target := int(math.Round(info.MaxMemoryClockMHz))
|
||||
out, err := runSATCommandCtx(ctx, verboseLog, fmt.Sprintf("normalize-gpu-%d-lmc", idx), []string{"nvidia-smi", "-i", strconv.Itoa(idx), "-lmc", strconv.Itoa(target)}, nil, nil)
|
||||
switch {
|
||||
case err == nil:
|
||||
rec.MemoryClockLockStatus = "applied"
|
||||
rec.MemoryClockLockMHz = float64(target)
|
||||
idxCopy := idx
|
||||
restore = append(restore, benchmarkRestoreAction{name: fmt.Sprintf("gpu-%d-rmc", idxCopy), fn: func() {
|
||||
_, _ = runSATCommandCtx(context.Background(), verboseLog, fmt.Sprintf("restore-gpu-%d-rmc", idxCopy), []string{"nvidia-smi", "-i", strconv.Itoa(idxCopy), "-rmc"}, nil, nil)
|
||||
}})
|
||||
case strings.Contains(strings.ToLower(string(out)), "deferred") || strings.Contains(strings.ToLower(string(out)), "not supported"):
|
||||
rec.MemoryClockLockStatus = "unsupported"
|
||||
rec.Notes = append(rec.Notes, "memory clock lock unsupported on this GPU/driver path")
|
||||
result.Normalization.Status = "partial"
|
||||
default:
|
||||
rec.MemoryClockLockStatus = "failed"
|
||||
rec.Notes = append(rec.Notes, "memory clock lock failed: "+strings.TrimSpace(string(out)))
|
||||
result.Normalization.Status = "partial"
|
||||
}
|
||||
}
|
||||
|
||||
result.Normalization.GPUs = append(result.Normalization.GPUs, rec)
|
||||
}
|
||||
return restore
|
||||
}
|
||||
|
||||
func collectBenchmarkSamples(ctx context.Context, durationSec int, gpuIndices []int) ([]GPUMetricRow, error) {
|
||||
if durationSec <= 0 {
|
||||
return nil, nil
|
||||
}
|
||||
deadline := time.Now().Add(time.Duration(durationSec) * time.Second)
|
||||
var rows []GPUMetricRow
|
||||
start := time.Now()
|
||||
for {
|
||||
if ctx.Err() != nil {
|
||||
return rows, ctx.Err()
|
||||
}
|
||||
samples, err := sampleBenchmarkTelemetry(gpuIndices)
|
||||
if err == nil {
|
||||
elapsed := time.Since(start).Seconds()
|
||||
for i := range samples {
|
||||
samples[i].ElapsedSec = elapsed
|
||||
}
|
||||
rows = append(rows, samples...)
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
break
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return rows, ctx.Err()
|
||||
case <-time.After(time.Second):
|
||||
}
|
||||
}
|
||||
return rows, nil
|
||||
}
|
||||
|
||||
func runBenchmarkCommandWithMetrics(ctx context.Context, verboseLog, name string, cmd []string, env []string, gpuIndices []int, logFunc func(string)) ([]byte, []GPUMetricRow, error) {
|
||||
stopCh := make(chan struct{})
|
||||
doneCh := make(chan struct{})
|
||||
var metricRows []GPUMetricRow
|
||||
start := time.Now()
|
||||
|
||||
go func() {
|
||||
defer close(doneCh)
|
||||
ticker := time.NewTicker(time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-stopCh:
|
||||
return
|
||||
case <-ticker.C:
|
||||
samples, err := sampleBenchmarkTelemetry(gpuIndices)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
elapsed := time.Since(start).Seconds()
|
||||
for i := range samples {
|
||||
samples[i].ElapsedSec = elapsed
|
||||
}
|
||||
metricRows = append(metricRows, samples...)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
out, err := runSATCommandCtx(ctx, verboseLog, name, cmd, env, logFunc)
|
||||
close(stopCh)
|
||||
<-doneCh
|
||||
|
||||
return out, metricRows, err
|
||||
}
|
||||
Reference in New Issue
Block a user