From d5ef809f3de794297f471797921fd22bbecb78f9 Mon Sep 17 00:00:00 2001 From: AXIS Contributor Date: Mon, 27 Jul 2026 13:42:16 -0400 Subject: [PATCH 1/2] feat(execution): sample local VRAM peak via nvidia-smi during guarded execution MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fills the 1A gap from the v0.14.7 optimization spec audit: PeakVRAMMB was previously only populated from *remote* execution responses; local execution recorded zero VRAM even on GPU hosts. Adds a background sampler (internal/execution/vram.go) that polls nvidia-smi --query-gpu=memory.used at a 500ms interval while the local workload runs, sums across all GPUs, and returns the max total observed. runLocal starts the sampler right before runWithReservationHeartbeat and stops it immediately after, setting resp.PeakVRAMMB on both success and failure paths. Degrades cleanly to 0 on hosts without nvidia-smi (Apple Silicon, Intel-only, or PATH without the binary) — the sampler's query returns (0, false) and the goroutine stops without error. VRAM is best-effort telemetry, not load-bearing. Verified: - Unit tests pass on the no-GPU path (PATH=/nonexistent). - Live smoke test on cranium (2x NVIDIA GPUs) captured 17896 MiB peak over a 500ms window. - Full execution + mcp suites pass; go build ./... clean. --- internal/execution/guarded.go | 4 ++ internal/execution/vram.go | 89 +++++++++++++++++++++++++++++++++ internal/execution/vram_test.go | 54 ++++++++++++++++++++ 3 files changed, 147 insertions(+) create mode 100644 internal/execution/vram.go create mode 100644 internal/execution/vram_test.go 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..849c02c --- /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 +} \ No newline at end of file diff --git a/internal/execution/vram_test.go b/internal/execution/vram_test.go new file mode 100644 index 0000000..e9720bb --- /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) + } +} \ No newline at end of file From 11a0f2550a7a679fa6025e7dab45c38324525e46 Mon Sep 17 00:00:00 2001 From: AXIS Contributor Date: Mon, 27 Jul 2026 14:52:07 -0400 Subject: [PATCH 2/2] style: gofmt --- internal/execution/vram.go | 2 +- internal/execution/vram_test.go | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/internal/execution/vram.go b/internal/execution/vram.go index 849c02c..84a4f0d 100644 --- a/internal/execution/vram.go +++ b/internal/execution/vram.go @@ -86,4 +86,4 @@ func queryTotalVRAMUsed(ctx context.Context) (int64, bool) { return 0, false } return total, true -} \ No newline at end of file +} diff --git a/internal/execution/vram_test.go b/internal/execution/vram_test.go index e9720bb..4076280 100644 --- a/internal/execution/vram_test.go +++ b/internal/execution/vram_test.go @@ -51,4 +51,4 @@ func TestQueryTotalVRAMUsedHandlesMissingBinary(t *testing.T) { if total != 0 { t.Fatalf("expected total=0 when missing, got %d", total) } -} \ No newline at end of file +}