diff --git a/internal/execution/guarded.go b/internal/execution/guarded.go index 072d546..b0694d4 100644 --- a/internal/execution/guarded.go +++ b/internal/execution/guarded.go @@ -875,9 +875,11 @@ func runLocal( } startedAt := time.Now().UTC() + vramStop := samplePeakVRAM(ctx, 500*time.Millisecond) out, peakRAMMB, runErr := runWithReservationHeartbeat(ledger, execID, func() (string, int64, error) { return runLocalWithOutput(ctx, command, env, stdoutWriter, stderrWriter) }) + peakVRAMMB := vramStop() elapsed := time.Since(startedAt) resp.Output = out resp.ExitCode = exitCode(runErr) @@ -887,6 +889,7 @@ func runLocal( runtimeChanged = true recordExecutionOutcome(st, reqs, resp, runErr, elapsed, peakRAMMB) resp.PeakRAMMB = peakRAMMB + resp.PeakVRAMMB = peakVRAMMB resp.WallTimeMS = durationMilliseconds(elapsed) return resp, runErr } @@ -896,6 +899,7 @@ func runLocal( recordExecutionOutcome(st, reqs, resp, nil, elapsed, peakRAMMB) resp.OK = true resp.PeakRAMMB = peakRAMMB + resp.PeakVRAMMB = peakVRAMMB resp.WallTimeMS = durationMilliseconds(elapsed) return resp, nil } diff --git a/internal/execution/vram.go b/internal/execution/vram.go new file mode 100644 index 0000000..84a4f0d --- /dev/null +++ b/internal/execution/vram.go @@ -0,0 +1,89 @@ +package execution + +import ( + "bufio" + "context" + "os/exec" + "strconv" + "strings" + "sync" + "time" +) + +// samplePeakVRAM polls nvidia-smi at a fixed interval while run is executing +// and returns the peak total GPU memory used across all GPUs, in MiB. If +// nvidia-smi is unavailable or every sample fails, it returns 0 and a nil +// error — VRAM is best-effort telemetry, not a load-bearing signal. +// +// The caller starts the sampler before launching the workload, then calls +// stop() once the workload has finished; the returned peak is the max total +// observed across all successful samples taken between start and stop. +// +// On machines without NVIDIA GPUs (e.g. Apple Silicon, Intel-only nodes), +// nvidia-smi will not be on PATH and this returns 0 cleanly. +func samplePeakVRAM(ctx context.Context, interval time.Duration) (stop func() int64) { + var mu sync.Mutex + peak := int64(0) + done := make(chan struct{}) + + go func() { + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-done: + return + case <-ctx.Done(): + return + case <-ticker.C: + used, ok := queryTotalVRAMUsed(ctx) + if !ok { + continue + } + mu.Lock() + if used > peak { + peak = used + } + mu.Unlock() + } + } + }() + + return func() int64 { + close(done) + mu.Lock() + defer mu.Unlock() + return peak + } +} + +// queryTotalVRAMUsed runs nvidia-smi once and sums memory.used across all +// GPUs. Returns (totalMiB, true) on success, (0, false) if nvidia-smi is +// missing, fails, or emits unparseable output. +func queryTotalVRAMUsed(ctx context.Context) (int64, bool) { + cmd := exec.CommandContext(ctx, "nvidia-smi", + "--query-gpu=memory.used", + "--format=csv,noheader,nounits", + ) + out, err := cmd.Output() + if err != nil { + return 0, false + } + var total int64 + scan := bufio.NewScanner(strings.NewReader(string(out))) + for scan.Scan() { + line := strings.TrimSpace(scan.Text()) + if line == "" { + continue + } + mib, err := strconv.ParseInt(line, 10, 64) + if err != nil { + return 0, false + } + total += mib + } + if err := scan.Err(); err != nil { + return 0, false + } + return total, true +} diff --git a/internal/execution/vram_test.go b/internal/execution/vram_test.go new file mode 100644 index 0000000..4076280 --- /dev/null +++ b/internal/execution/vram_test.go @@ -0,0 +1,54 @@ +package execution + +import ( + "context" + "testing" + "time" +) + +// TestSamplePeakVRAMReturnsZeroWhenNvidiaSmiMissing verifies the sampler +// degrades cleanly on machines without nvidia-smi (Apple Silicon, Intel-only +// nodes). It does this by pointing PATH at an empty directory so the binary +// can't be found. +func TestSamplePeakVRAMReturnsZeroWhenNvidiaSmiMissing(t *testing.T) { + t.Setenv("PATH", "/nonexistent/empty-path") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + stop := samplePeakVRAM(ctx, 20*time.Millisecond) + // Give the ticker a couple of beats to attempt (and fail) a sample. + time.Sleep(80 * time.Millisecond) + peak := stop() + if peak != 0 { + t.Fatalf("expected 0 peak when nvidia-smi missing, got %d", peak) + } +} + +// TestSamplePeakVRAMStopsCleanly verifies that calling stop() returns a +// non-negative int64 and doesn't block or panic, even if the goroutine never +// took a successful sample. +func TestSamplePeakVRAMStopsCleanly(t *testing.T) { + t.Setenv("PATH", "/nonexistent/empty-path") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + stop := samplePeakVRAM(ctx, 50*time.Millisecond) + peak := stop() + if peak < 0 { + t.Fatalf("expected non-negative peak, got %d", peak) + } + // Calling stop twice must not panic on double-close of the channel. + // (We don't call it twice here — the contract is single-stop — but + // we verify the returned value is sane.) +} + +// TestQueryTotalVRAMUsedHandlesMissingBinary verifies the query helper +// returns (0, false) when nvidia-smi isn't on PATH, not an error. +func TestQueryTotalVRAMUsedHandlesMissingBinary(t *testing.T) { + t.Setenv("PATH", "/nonexistent/empty-path") + total, ok := queryTotalVRAMUsed(context.Background()) + if ok { + t.Fatalf("expected ok=false when nvidia-smi missing, got total=%d", total) + } + if total != 0 { + t.Fatalf("expected total=0 when missing, got %d", total) + } +}