feat(hardware): read free VRAM and the UMD CUDA banner

cmd/flynn/modelfit.go

@@ -12,13 +12,10 @@ import (
 	"github.com/ionalpha/flynn/internal/hardware"
 )
 
-// detectBox and probeFreeVRAM are the two reads of the real machine behind `flynn models
-// --fit`. They are variables so the fit view can be driven from a described box, which is
-// how its ranking is tested on a machine with a different GPU.
-var (
-	detectBox     = hardware.Detect
-	probeFreeVRAM = hardware.FreeVRAM
-)
+// detectBox is the read of the real machine behind `flynn models --fit`. It is a variable so
+// the fit view can be driven from a described box, which is how its ranking is tested on a
+// machine with a different GPU.
+var detectBox = hardware.Detect
 
 // fitMachine is the machine a fit is judged against and how it was determined.
 type fitMachine struct {
@@ -50,7 +47,7 @@ func resolveMachine(vramGB, ramGB float64) fitMachine {
 	case box.HasGPU():
 		m.VRAMFreeBytes = box.VRAMBytes
 		desc := fmt.Sprintf("%s, %s VRAM", box.GPUName, humanBytes(box.VRAMBytes))
-		if free, ok := probeFreeVRAM(ctx); ok && free >= 0 && free < box.VRAMBytes {
+		if free := box.VRAMFreeBytes; free > 0 && free < box.VRAMBytes {
 			m.VRAMFreeBytes = free
 			desc = fmt.Sprintf("%s, %s free of %s VRAM, %s in use", box.GPUName, humanBytes(free), humanBytes(box.VRAMBytes), humanBytes(box.VRAMBytes-free))
 		}

cmd/flynn/modelfit_test.go

@@ -43,13 +43,11 @@ func TestPlacementColNamesEachPlacement(t *testing.T) {
 }
 
 // stubMachine makes the fit view see a described box instead of the one running the test.
-// freeBytes < 0 means the free-memory probe is unavailable.
-func stubMachine(t *testing.T, box hardware.Box, freeBytes int64) {
+func stubMachine(t *testing.T, box hardware.Box) {
 	t.Helper()
-	oldBox, oldFree := detectBox, probeFreeVRAM
-	t.Cleanup(func() { detectBox, probeFreeVRAM = oldBox, oldFree })
+	old := detectBox
+	t.Cleanup(func() { detectBox = old })
 	detectBox = func(context.Context) hardware.Box { return box }
-	probeFreeVRAM = func(context.Context) (int64, bool) { return freeBytes, freeBytes >= 0 }
 }
 
 func fitOutput(t *testing.T, args ...string) string {
@@ -66,7 +64,7 @@ func fitOutput(t *testing.T, args ...string) string {
 // list as a labelled split, and the largest model that stays on the card is named beside it.
 func TestFitViewOnTheFieldTestBox(t *testing.T) {
 	const gb = 1 << 30
-	stubMachine(t, hardware.Box{GPUName: "RTX 5080 Laptop", VRAMBytes: 16 * gb, RAMBytes: 16 * 1_000_000_000}, 16*gb-1_100_000_000)
+	stubMachine(t, hardware.Box{GPUName: "RTX 5080 Laptop", VRAMBytes: 16 * gb, VRAMFreeBytes: 16*gb - 1_100_000_000, RAMBytes: 16 * 1_000_000_000})
 	out := fitOutput(t)
 	for _, want := range []string{
 		"RTX 5080 Laptop", "in use", "context 32k, KV cache f16",
@@ -86,44 +84,39 @@ func TestFitViewTable(t *testing.T) {
 	cases := []struct {
 		name  string
 		box   hardware.Box
-		free  int64
 		args  []string
 		want  string
 		wants []string
 	}{
 		{
 			name: "12GB card, 64GB RAM",
-			box:  hardware.Box{GPUName: "RTX 4070", VRAMBytes: 12 * gb, RAMBytes: 64 * gb},
-			free: 11 * gb,
+			box:  hardware.Box{GPUName: "RTX 4070", VRAMBytes: 12 * gb, VRAMFreeBytes: 11 * gb, RAMBytes: 64 * gb},
 			want: "recommended local model: qwen3.6:35b-a3b (gpu+ram",
 		},
 		{
 			name: "CPU only",
 			box:  hardware.Box{RAMBytes: 16 * gb},
-			free: -1,
 			want: "recommended local model: ",
 			wants: []string{
 				"system RAM", "usable",
 			},
 		},
 		{
 			name: "24GB card holds the 35B whole",
-			box:  hardware.Box{GPUName: "RTX 4090", VRAMBytes: 24 * gb, RAMBytes: 32 * gb},
-			free: 23 * gb,
+			box:  hardware.Box{GPUName: "RTX 4090", VRAMBytes: 24 * gb, VRAMFreeBytes: 23 * gb, RAMBytes: 32 * gb},
 			want: "recommended local model: qwen3.6:35b-a3b (gpu)",
 		},
 		{
 			name:  "a q8_0 cache at 131k tokens leaves the 35B tight where f16 leaves it over budget",
-			box:   hardware.Box{GPUName: "RTX 4090", VRAMBytes: 24 * gb},
-			free:  23 * gb,
+			box:   hardware.Box{GPUName: "RTX 4090", VRAMBytes: 24 * gb, VRAMFreeBytes: 23 * gb},
 			args:  []string{"--ctx", "131072", "--kv", "q8"},
 			want:  "KV cache q8_0",
 			wants: []string{"tight"},
 		},
 	}
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			stubMachine(t, tc.box, tc.free)
+			stubMachine(t, tc.box)
 			out := fitOutput(t, tc.args...)
 			for _, want := range append([]string{tc.want}, tc.wants...) {
 				if !strings.Contains(out, want) {
@@ -135,7 +128,7 @@ func TestFitViewTable(t *testing.T) {
 }
 
 func TestFitViewRejectsBadSettings(t *testing.T) {
-	stubMachine(t, hardware.Box{RAMBytes: 16_000_000_000}, -1)
+	stubMachine(t, hardware.Box{RAMBytes: 16_000_000_000})
 	for _, args := range [][]string{{"--kv", "f32"}, {"--ctx", "0"}, {"--ctx", "-5"}} {
 		var b strings.Builder
 		if err := runModels(append([]string{"--local", "--fit"}, args...), "", &b); err == nil {
@@ -146,7 +139,7 @@ func TestFitViewRejectsBadSettings(t *testing.T) {
 
 // TestFitViewWithNothingDetected says so rather than ranking against zero memory.
 func TestFitViewWithNothingDetected(t *testing.T) {
-	stubMachine(t, hardware.Box{}, -1)
+	stubMachine(t, hardware.Box{})
 	out := fitOutput(t)
 	if !strings.Contains(out, "could not detect") {
 		t.Fatalf("expected the undetected message:\n%s", out)
@@ -157,7 +150,7 @@ func TestFitViewWithNothingDetected(t *testing.T) {
 // same answer on a laptop with 16 GB of RAM and a workstation with 256 GB, and --ram opts
 // into a split.
 func TestGivenVRAMIgnoresThisMachinesRAM(t *testing.T) {
-	stubMachine(t, hardware.Box{GPUName: "big", VRAMBytes: 80_000_000_000, RAMBytes: 256_000_000_000}, -1)
+	stubMachine(t, hardware.Box{GPUName: "big", VRAMBytes: 80_000_000_000, RAMBytes: 256_000_000_000})
 	without := fitOutput(t, "--vram", "16")
 	if strings.Contains(without, "gpu+ram (") || !strings.Contains(without, "no system RAM") {
 		t.Fatalf("--vram alone must not borrow this machine's RAM:\n%s", without)
@@ -168,12 +161,27 @@ func TestGivenVRAMIgnoresThisMachinesRAM(t *testing.T) {
 	}
 }
 
-// TestFreeVRAMProbeAbsentFallsBackToTotal: with no probe the detected total is used and
-// the line does not claim anything is in use.
-func TestFreeVRAMProbeAbsentFallsBackToTotal(t *testing.T) {
-	stubMachine(t, hardware.Box{GPUName: "X", VRAMBytes: 16_000_000_000, RAMBytes: 8_000_000_000}, -1)
-	m := resolveMachine(0, 0)
-	if m.VRAMFreeBytes != 16_000_000_000 || strings.Contains(m.description, "in use") {
-		t.Fatalf("machine = %+v", m)
+// TestFitJudgesVRAMAgainstTheFreeFigure: the fit uses the free memory Detect reports, and
+// falls back to the total when the figure is unreported or not below the total, so the
+// line only claims memory is in use when the probe showed it.
+func TestFitJudgesVRAMAgainstTheFreeFigure(t *testing.T) {
+	const gb = 1_000_000_000
+	for _, tc := range []struct {
+		name  string
+		box   hardware.Box
+		want  int64
+		inUse bool
+	}{
+		{"free figure used", hardware.Box{GPUName: "X", VRAMBytes: 16 * gb, VRAMFreeBytes: 11 * gb, RAMBytes: 8 * gb}, 11 * gb, true},
+		{"unreported falls back to the total", hardware.Box{GPUName: "X", VRAMBytes: 16 * gb, RAMBytes: 8 * gb}, 16 * gb, false},
+		{"figure above the total is ignored", hardware.Box{GPUName: "X", VRAMBytes: 16 * gb, VRAMFreeBytes: 17 * gb, RAMBytes: 8 * gb}, 16 * gb, false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			stubMachine(t, tc.box)
+			m := resolveMachine(0, 0)
+			if m.VRAMFreeBytes != tc.want || strings.Contains(m.description, "in use") != tc.inUse {
+				t.Fatalf("machine = %+v, description %q", m, m.description)
+			}
+		})
 	}
 }

internal/hardware/hardware.go

@@ -23,6 +23,11 @@ type Box struct {
 	GPUName string
 	// VRAMBytes is the detected GPU memory in bytes, 0 if none was found.
 	VRAMBytes int64
+	// VRAMFreeBytes is the GPU memory not in use when the probe ran, in bytes, 0 if the
+	// driver did not report it or reported none free. The desktop and other programs already
+	// hold part of the card, so a model sized to VRAMBytes can fail to load; a fit is judged
+	// against this figure when it is set.
+	VRAMFreeBytes int64
 	// RAMBytes is the total system memory in bytes, 0 if it could not be read. A
 	// CPU-only run is bounded by this, not by VRAM, so it decides fit on the common
 	// machine that has no usable GPU.
@@ -90,7 +95,7 @@ func Detect(ctx context.Context) Box {
 	var b Box
 	if out, ok := runNvidiaSMI(ctx); ok {
 		if gpu, ok := parseNvidiaSMI(out); ok {
-			b.GPUName, b.VRAMBytes = gpu.GPUName, gpu.VRAMBytes
+			b.GPUName, b.VRAMBytes, b.VRAMFreeBytes = gpu.GPUName, gpu.VRAMBytes, gpu.VRAMFreeBytes
 			b.ComputeCapability, b.DriverVersion = gpu.ComputeCapability, gpu.DriverVersion
 			b.GPUArch = archForComputeCap(b.ComputeCapability)
 		}
@@ -120,20 +125,22 @@ var lookPath = exec.LookPath
 // runNvidiaSMI queries the NVIDIA management tool for the per-GPU fields, returning false
 // when the tool is missing or errors (no GPU, no driver, not installed). Memory and name
 // have always been queried; compute capability and driver version are added so the box's
-// architecture and CUDA reach are known, not just its size.
+// architecture and CUDA reach are known, not just its size. Free memory is the last column,
+// so the positional fields before it keep their order.
 func runNvidiaSMI(ctx context.Context) (string, bool) {
 	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
 	defer cancel()
 	out, err := runCommand(ctx, "nvidia-smi",
-		"--query-gpu=memory.total,name,compute_cap,driver_version", "--format=csv,noheader,nounits")
+		"--query-gpu=memory.total,name,compute_cap,driver_version,memory.free", "--format=csv,noheader,nounits")
 	if err != nil {
 		return "", false
 	}
 	return string(out), true
 }
 
 // runNvidiaSMIHeader runs nvidia-smi with no query, whose banner carries the maximum CUDA
-// version the driver supports ("CUDA Version: 12.4"). That number is not available from
+// version the driver supports ("CUDA Version: 12.4", or "CUDA UMD Version: 13.3" on current
+// drivers). That number is not available from
 // the per-GPU query, so it is read from the default output. Returns false when the tool
 // is absent or errors.
 func runNvidiaSMIHeader(ctx context.Context) (string, bool) {
@@ -147,7 +154,7 @@ func runNvidiaSMIHeader(ctx context.Context) (string, bool) {
 }
 
 // parseNvidiaSMI reads the first usable GPU row of the query output
-// ("<MiB>, <name>, <compute_cap>, <driver_version>") into a Box. Memory is reported in
+// ("<MiB>, <name>, <compute_cap>, <driver_version>, <free MiB>") into a Box. Memory is reported in
 // mebibytes with the nounits format, so it is scaled to bytes. Older nvidia-smi builds
 // or a trimmed query may print fewer columns, so the extra fields are read only when
 // present. It returns false when no row yields a memory total.
@@ -172,6 +179,13 @@ func parseNvidiaSMI(out string) (Box, bool) {
 		if len(cols) > 3 {
 			b.DriverVersion = strings.TrimSpace(cols[3])
 		}
+		if len(cols) > 4 {
+			// A free figure above the total, or a non-positive one, is not a reading this
+			// probe can use, so the field stays unset rather than describing the card wrongly.
+			if free, err := strconv.ParseInt(strings.TrimSpace(cols[4]), 10, 64); err == nil && free > 0 && free <= mib {
+				b.VRAMFreeBytes = free * 1024 * 1024
+			}
+		}
 		return b, true
 	}
 	return Box{}, false
@@ -232,16 +246,28 @@ func archForComputeCap(cc string) string {
 	}
 }
 
+// cudaLabels are the banner labels the maximum CUDA version is printed under. Older drivers
+// print "CUDA Version"; current ones print "CUDA UMD Version", the user-mode driver's figure,
+// which is the same number for this purpose. Neither label is a substring of the other, so
+// the order does not change which one matches.
+var cudaLabels = []string{"CUDA Version:", "CUDA UMD Version:"}
+
 // parseCUDAVersion reads the maximum supported CUDA version out of the nvidia-smi banner,
-// which prints it as "CUDA Version: 12.4" near the driver version. It returns an empty
-// string when the banner does not carry it.
+// which prints it as "CUDA Version: 12.4" or "CUDA UMD Version: 13.3" near the driver
+// version. It returns an empty string when the banner carries neither label.
 func parseCUDAVersion(out string) string {
-	const marker = "CUDA Version:"
-	i := strings.Index(out, marker)
-	if i < 0 {
-		return ""
+	for _, label := range cudaLabels {
+		if i := strings.Index(out, label); i >= 0 {
+			return cudaVersionAfter(out[i+len(label):])
+		}
 	}
-	rest := strings.TrimSpace(out[i+len(marker):])
+	return ""
+}
+
+// cudaVersionAfter reads the dotted version that follows a banner label, or returns empty when
+// the text after the label is not a version.
+func cudaVersionAfter(rest string) string {
+	rest = strings.TrimSpace(rest)
 	end := strings.IndexFunc(rest, func(r rune) bool {
 		return r != '.' && (r < '0' || r > '9')
 	})