Unify NVIDIA GPU recovery paths

audit/internal/platform/benchmark.go

@@ -105,6 +105,7 @@ var (
 	benchmarkSkippedPattern    = regexp.MustCompile(`^([a-z0-9_]+)(?:\[\d+\])?=SKIPPED (.+)$`)
 	benchmarkIterationsPattern = regexp.MustCompile(`^([a-z0-9_]+)_iterations=(\d+)$`)
 	benchmarkGeteuid           = os.Geteuid
+	benchmarkResetNvidiaGPU    = resetNvidiaGPU
 	benchmarkSleep             = time.Sleep
 )
 
@@ -249,6 +250,35 @@ func setBenchmarkPowerLimit(ctx context.Context, verboseLog string, gpuIndex, po
 	return nil
 }
 
+func resetBenchmarkGPU(ctx context.Context, verboseLog string, gpuIndex int, logFunc func(string)) error {
+	if logFunc != nil {
+		logFunc(fmt.Sprintf("power benchmark pre-flight: GPU %d reset via shared NVIDIA recover path", gpuIndex))
+	}
+	out, err := benchmarkResetNvidiaGPU(gpuIndex)
+	appendSATVerboseLog(verboseLog,
+		fmt.Sprintf("[%s] start power-preflight-gpu-%d-reset.log", time.Now().UTC().Format(time.RFC3339), gpuIndex),
+		"cmd: bee-nvidia-recover reset-gpu "+strconv.Itoa(gpuIndex),
+	)
+	if trimmed := strings.TrimSpace(out); trimmed != "" && logFunc != nil {
+		for _, line := range strings.Split(trimmed, "\n") {
+			line = strings.TrimSpace(line)
+			if line != "" {
+				logFunc(line)
+			}
+		}
+	}
+	rc := 0
+	if err != nil {
+		rc = 1
+	}
+	appendSATVerboseLog(verboseLog,
+		fmt.Sprintf("[%s] finish power-preflight-gpu-%d-reset.log", time.Now().UTC().Format(time.RFC3339), gpuIndex),
+		fmt.Sprintf("rc: %d", rc),
+		"",
+	)
+	return err
+}
+
 func resetBenchmarkGPUs(ctx context.Context, verboseLog string, gpuIndices []int, logFunc func(string)) []int {
 	if len(gpuIndices) == 0 {
 		return nil
@@ -266,8 +296,7 @@ func resetBenchmarkGPUs(ctx context.Context, verboseLog string, gpuIndices []int
 	}
 	var failed []int
 	for _, idx := range gpuIndices {
-		name := fmt.Sprintf("power-preflight-gpu-%d-reset.log", idx)
-		if _, err := runSATCommandCtx(ctx, verboseLog, name, []string{"nvidia-smi", "-i", strconv.Itoa(idx), "-r"}, nil, logFunc); err != nil {
+		if err := resetBenchmarkGPU(ctx, verboseLog, idx, logFunc); err != nil {
 			failed = append(failed, idx)
 			if logFunc != nil {
 				logFunc(fmt.Sprintf("power benchmark pre-flight: GPU %d reset failed: %v", idx, err))
@@ -4440,8 +4469,7 @@ func (s *System) RunNvidiaPowerBench(ctx context.Context, baseDir string, opts N
 		_ = os.MkdirAll(singleDir, 0755)
 		singleInfo := cloneBenchmarkGPUInfoMap(infoByIndex)
 		if failed := resetBenchmarkGPUs(ctx, verboseLog, []int{idx}, logFunc); len(failed) > 0 {
-			result.Findings = append(result.Findings,
-				fmt.Sprintf("GPU %d reset pre-flight did not complete before its first power test; throttle counters may contain stale state.", idx))
+			return "", fmt.Errorf("power benchmark pre-flight: failed to reset GPU %d; benchmark aborted to keep measurements clean", idx)
 		}
 		logFunc(fmt.Sprintf("power calibration: GPU %d single-card baseline", idx))
 		singlePowerStopCh := make(chan struct{})

audit/internal/platform/benchmark_test.go

@@ -2,7 +2,7 @@ package platform
 
 import (
 	"context"
-	"os"
+	"fmt"
 	"os/exec"
 	"path/filepath"
 	"strings"
@@ -188,18 +188,16 @@ func TestBenchmarkCalibrationThrottleReasonIgnoresPowerReasons(t *testing.T) {
 }
 
 func TestResetBenchmarkGPUsSkipsWithoutRoot(t *testing.T) {
-	t.Parallel()
-
 	oldGeteuid := benchmarkGeteuid
-	oldExec := satExecCommand
+	oldReset := benchmarkResetNvidiaGPU
 	benchmarkGeteuid = func() int { return 1000 }
-	satExecCommand = func(name string, args ...string) *exec.Cmd {
-		t.Fatalf("unexpected command: %s %v", name, args)
-		return nil
+	benchmarkResetNvidiaGPU = func(int) (string, error) {
+		t.Fatal("unexpected reset call")
+		return "", nil
 	}
 	t.Cleanup(func() {
 		benchmarkGeteuid = oldGeteuid
-		satExecCommand = oldExec
+		benchmarkResetNvidiaGPU = oldReset
 	})
 
 	var logs []string
@@ -215,44 +213,52 @@ func TestResetBenchmarkGPUsSkipsWithoutRoot(t *testing.T) {
 }
 
 func TestResetBenchmarkGPUsResetsEachGPU(t *testing.T) {
-	t.Parallel()
-
-	dir := t.TempDir()
-	script := filepath.Join(dir, "nvidia-smi")
-	argsLog := filepath.Join(dir, "args.log")
-	if err := os.WriteFile(script, []byte("#!/bin/sh\nprintf '%s\\n' \"$*\" >> "+argsLog+"\nprintf 'ok\\n'\n"), 0755); err != nil {
-		t.Fatalf("write script: %v", err)
-	}
-
 	oldGeteuid := benchmarkGeteuid
 	oldSleep := benchmarkSleep
-	oldLookPath := satLookPath
+	oldReset := benchmarkResetNvidiaGPU
 	benchmarkGeteuid = func() int { return 0 }
 	benchmarkSleep = func(time.Duration) {}
-	satLookPath = func(file string) (string, error) {
-		if file == "nvidia-smi" {
-			return script, nil
-		}
-		return exec.LookPath(file)
+	var calls []int
+	benchmarkResetNvidiaGPU = func(index int) (string, error) {
+		calls = append(calls, index)
+		return "ok\n", nil
 	}
 	t.Cleanup(func() {
 		benchmarkGeteuid = oldGeteuid
 		benchmarkSleep = oldSleep
-		satLookPath = oldLookPath
+		benchmarkResetNvidiaGPU = oldReset
 	})
 
-	failed := resetBenchmarkGPUs(context.Background(), filepath.Join(dir, "verbose.log"), []int{2, 5}, nil)
+	failed := resetBenchmarkGPUs(context.Background(), filepath.Join(t.TempDir(), "verbose.log"), []int{2, 5}, nil)
 	if len(failed) != 0 {
 		t.Fatalf("failed=%v want no failures", failed)
 	}
-	raw, err := os.ReadFile(argsLog)
-	if err != nil {
-		t.Fatalf("read args log: %v", err)
+	if got, want := fmt.Sprint(calls), "[2 5]"; got != want {
+		t.Fatalf("calls=%v want %s", calls, want)
 	}
-	got := strings.Fields(string(raw))
-	want := []string{"-i", "2", "-r", "-i", "5", "-r"}
-	if strings.Join(got, " ") != strings.Join(want, " ") {
-		t.Fatalf("args=%v want %v", got, want)
+}
+
+func TestResetBenchmarkGPUsTracksFailuresFromSharedReset(t *testing.T) {
+	oldGeteuid := benchmarkGeteuid
+	oldSleep := benchmarkSleep
+	oldReset := benchmarkResetNvidiaGPU
+	benchmarkGeteuid = func() int { return 0 }
+	benchmarkSleep = func(time.Duration) {}
+	benchmarkResetNvidiaGPU = func(index int) (string, error) {
+		if index == 5 {
+			return "busy\n", exec.ErrNotFound
+		}
+		return "ok\n", nil
+	}
+	t.Cleanup(func() {
+		benchmarkGeteuid = oldGeteuid
+		benchmarkSleep = oldSleep
+		benchmarkResetNvidiaGPU = oldReset
+	})
+
+	failed := resetBenchmarkGPUs(context.Background(), filepath.Join(t.TempDir(), "verbose.log"), []int{2, 5}, nil)
+	if got, want := fmt.Sprint(failed), "[5]"; got != want {
+		t.Fatalf("failed=%v want %s", failed, want)
 	}
 }
 

audit/internal/platform/nvidia_recover.go

@@ -3,6 +3,8 @@ package platform
 import (
 	"fmt"
 	"os/exec"
+	"strconv"
+	"strings"
 	"time"
 )
 
@@ -28,3 +30,22 @@ func runNvidiaRecover(args ...string) (string, error) {
 	raw, err := exec.Command("sudo", helperArgs...).CombinedOutput()
 	return string(raw), err
 }
+
+func resetNvidiaGPU(index int) (string, error) {
+	if index < 0 {
+		return "", fmt.Errorf("gpu index must be >= 0")
+	}
+	out, err := runNvidiaRecover("reset-gpu", strconv.Itoa(index))
+	if strings.TrimSpace(out) == "" && err == nil {
+		out = "GPU reset completed.\n"
+	}
+	return out, err
+}
+
+func restartNvidiaDrivers() (string, error) {
+	out, err := runNvidiaRecover("restart-drivers")
+	if strings.TrimSpace(out) == "" && err == nil {
+		out = "NVIDIA drivers restarted.\n"
+	}
+	return out, err
+}

audit/internal/platform/sat.go

@@ -404,14 +404,7 @@ func normalizeNvidiaBusID(v string) string {
 }
 
 func (s *System) ResetNvidiaGPU(index int) (string, error) {
-	if index < 0 {
-		return "", fmt.Errorf("gpu index must be >= 0")
-	}
-	out, err := runNvidiaRecover("reset-gpu", strconv.Itoa(index))
-	if strings.TrimSpace(out) == "" && err == nil {
-		out = "GPU reset completed.\n"
-	}
-	return out, err
+	return resetNvidiaGPU(index)
 }
 
 // RunNCCLTests runs nccl-tests all_reduce_perf across the selected NVIDIA GPUs.

audit/internal/platform/services.go

@@ -62,7 +62,7 @@ func (s *System) ServiceState(name string) string {
 
 func (s *System) ServiceDo(name string, action ServiceAction) (string, error) {
 	if name == "bee-nvidia" && action == ServiceRestart {
-		return runNvidiaRecover("restart-drivers")
+		return restartNvidiaDrivers()
 	}
 	// bee-web runs as the bee user; sudo is required to control system services.
 	// /etc/sudoers.d/bee grants bee NOPASSWD:ALL.