diff --git a/cmd/flynn/modelfit.go b/cmd/flynn/modelfit.go index 2a5c8712..140ab256 100644 --- a/cmd/flynn/modelfit.go +++ b/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)) } diff --git a/cmd/flynn/modelfit_test.go b/cmd/flynn/modelfit_test.go index e97074be..926dc7f7 100644 --- a/cmd/flynn/modelfit_test.go +++ b/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,21 +84,18 @@ 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", @@ -108,14 +103,12 @@ func TestFitViewTable(t *testing.T) { }, { 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"}, @@ -123,7 +116,7 @@ func TestFitViewTable(t *testing.T) { } 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) + } + }) } } diff --git a/internal/hardware/freevram.go b/internal/hardware/freevram.go deleted file mode 100644 index ee5fd514..00000000 --- a/internal/hardware/freevram.go +++ /dev/null @@ -1,41 +0,0 @@ -package hardware - -import ( - "context" - "strconv" - "strings" - "time" -) - -// FreeVRAM reports how much memory of the first NVIDIA GPU is unused right now. The -// desktop compositor and other programs already hold part of the card, so a model sized to -// [Box.VRAMBytes] can still fail to load; this is the figure a fit is judged against. It -// reports false when the vendor tool is absent or prints nothing usable, and the caller -// then falls back to the total. -func FreeVRAM(ctx context.Context) (int64, bool) { - ctx, cancel := context.WithTimeout(ctx, 3*time.Second) - defer cancel() - out, err := runCommand(ctx, "nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits") - if err != nil { - return 0, false - } - return parseFreeVRAM(string(out)) -} - -// parseFreeVRAM reads the first row of the memory.free query, a count of mebibytes. A -// negative or non-numeric row ("N/A" on a driver that hides the figure) is unusable, and -// is not skipped in favour of a later GPU's row, which would describe a different card. -func parseFreeVRAM(out string) (int64, bool) { - for _, line := range strings.Split(out, "\n") { - line = strings.TrimSpace(line) - if line == "" { - continue - } - mib, err := strconv.ParseInt(line, 10, 64) - if err != nil || mib < 0 { - return 0, false - } - return mib * 1024 * 1024, true - } - return 0, false -} diff --git a/internal/hardware/freevram_test.go b/internal/hardware/freevram_test.go deleted file mode 100644 index c89b7846..00000000 --- a/internal/hardware/freevram_test.go +++ /dev/null @@ -1,57 +0,0 @@ -package hardware - -import ( - "context" - "testing" -) - -func TestParseFreeVRAM(t *testing.T) { - for _, tc := range []struct { - name string - in string - want int64 - ok bool - }{ - {"one gpu", "15360\n", 15360 * 1024 * 1024, true}, - {"first row of two, with blank lines and CRLF", "\n 1024 \r\n2048\n", 1024 * 1024 * 1024, true}, - {"empty output", "", 0, false}, - {"hidden figure", "N/A\n", 0, false}, - {"negative", "-5\n", 0, false}, - {"unusable first row is not skipped", "N/A\n2048\n", 0, false}, - } { - t.Run(tc.name, func(t *testing.T) { - got, ok := parseFreeVRAM(tc.in) - if got != tc.want || ok != tc.ok { - t.Errorf("parseFreeVRAM(%q) = %d, %v; want %d, %v", tc.in, got, ok, tc.want, tc.ok) - } - }) - } -} - -// TestFreeVRAMRunsTheQueryAndReadsIt drives FreeVRAM through the probe seam: a driver that -// answers, one that prints nothing usable, a tool that is missing, and a tool that exits -// non-zero. Only the first gives a figure, and the query asked for is memory.free alone. -func TestFreeVRAMRunsTheQueryAndReadsIt(t *testing.T) { - for _, tc := range []struct { - name string - host fakeHost - want int64 - ok bool - }{ - {"answers", fakeHost{cmdOut: map[string]string{"nvidia-smi": "14000\n"}}, 14000 * 1024 * 1024, true}, - {"prints N/A", fakeHost{cmdOut: map[string]string{"nvidia-smi": "N/A\n"}}, 0, false}, - {"not installed", fakeHost{}, 0, false}, - {"exits non-zero", fakeHost{cmdErr: map[string]bool{"nvidia-smi": true}}, 0, false}, - } { - t.Run(tc.name, func(t *testing.T) { - tc.host.install(t) - got, ok := FreeVRAM(context.Background()) - if got != tc.want || ok != tc.ok { - t.Errorf("FreeVRAM() = %d, %v; want %d, %v", got, ok, tc.want, tc.ok) - } - if len(tc.host.calls) != 1 || tc.host.calls[0] != "nvidia-smi" { - t.Errorf("ran %v, want one nvidia-smi call", tc.host.calls) - } - }) - } -} diff --git a/internal/hardware/hardware.go b/internal/hardware/hardware.go index bda20b32..40e12058 100644 --- a/internal/hardware/hardware.go +++ b/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,12 +125,13 @@ 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 } @@ -133,7 +139,8 @@ func runNvidiaSMI(ctx context.Context) (string, bool) { } // 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 -// (", , , ") into a Box. Memory is reported in +// (", , , , ") 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') }) diff --git a/internal/hardware/hardware_test.go b/internal/hardware/hardware_test.go index 0027c2e7..5822d0cb 100644 --- a/internal/hardware/hardware_test.go +++ b/internal/hardware/hardware_test.go @@ -70,6 +70,39 @@ func TestParseNvidiaSMICapability(t *testing.T) { } } +// TestParseNvidiaSMIFreeVRAM reads the free-memory column the query appends after the +// existing fields. Each case is a figure the driver can print: a reading, a hidden one, a +// query that predates the column, and rows a reader must not trust. +func TestParseNvidiaSMIFreeVRAM(t *testing.T) { + const mib = 1024 * 1024 + cases := []struct { + name string + out string + wantFree int64 + }{ + {"free is the fifth column", "24564, NVIDIA GeForce RTX 4090, 8.9, 550.54.14, 20000\n", 20000 * mib}, + {"crlf row", "24564, NVIDIA GeForce RTX 4090, 8.9, 550.54.14, 20000\r\n", 20000 * mib}, + {"free equal to total", "8192, A, 8.9, 550.54.14, 8192", 8192 * mib}, + {"hidden figure", "24564, A, 8.9, 550.54.14, [N/A]\n", 0}, + {"query without the column", "16384, A, 8.9, 550.54.14\n", 0}, + {"zero free is not reported", "16384, A, 8.9, 550.54.14, 0\n", 0}, + {"negative free", "16384, A, 8.9, 550.54.14, -5\n", 0}, + {"above the total is dropped", "8192, A, 8.9, 550.54.14, 9000\n", 0}, + {"first row's free, not a later gpu's", "16384, A, 8.9, 550.54.14, 1000\n8192, B, 8.6, 550.54.14, 5000\n", 1000 * mib}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + b, ok := parseNvidiaSMI(tc.out) + if !ok { + t.Fatalf("should parse %q", tc.out) + } + if b.VRAMFreeBytes != tc.wantFree { + t.Fatalf("VRAMFreeBytes = %d, want %d", b.VRAMFreeBytes, tc.wantFree) + } + }) + } +} + func TestArchForComputeCap(t *testing.T) { cases := map[string]string{ "": "", @@ -100,6 +133,10 @@ func TestParseCUDAVersion(t *testing.T) { {"CUDA Version: 12 ", "12"}, {"no cuda banner here", ""}, {"CUDA Version: N/A", ""}, + // Drivers 610 and later print the user-mode driver's figure under a different label. + {"| NVIDIA-SMI 610.88 KMD Version: 610.88 CUDA UMD Version: 13.3 |", "13.3"}, + {"CUDA UMD Version: 13.3\n", "13.3"}, + {"CUDA UMD Version: N/A", ""}, } for _, tc := range cases { if got := parseCUDAVersion(tc.in); got != tc.want { diff --git a/internal/hardware/probe_test.go b/internal/hardware/probe_test.go index 4104f7d6..8e679e79 100644 --- a/internal/hardware/probe_test.go +++ b/internal/hardware/probe_test.go @@ -231,6 +231,25 @@ func TestEngineRuntimeMatchIsCaseInsensitive(t *testing.T) { } } +// TestDetectReadsFreeVRAMAndUMDBanner checks both fields come from probes the box already +// runs: the free memory from the per-GPU query, and the CUDA version from a banner that uses +// the "CUDA UMD Version" label. +func TestDetectReadsFreeVRAMAndUMDBanner(t *testing.T) { + const umdBanner = "| NVIDIA-SMI 610.88 KMD Version: 610.88 CUDA UMD Version: 13.3 |" + h := &fakeHost{cmdOut: map[string]string{ + "nvidia-smi": "16376, NVIDIA GeForce RTX 5080 Laptop, 12.0, 610.88, 11264\n" + umdBanner, + }} + h.install(t) + + b := Detect(context.Background()) + if b.VRAMBytes != 16376*1024*1024 || b.VRAMFreeBytes != 11264*1024*1024 { + t.Fatalf("VRAMBytes=%d VRAMFreeBytes=%d, want the total and free figures from the query", b.VRAMBytes, b.VRAMFreeBytes) + } + if b.CUDAVersion != "13.3" { + t.Fatalf("CUDAVersion = %q, want 13.3 from the UMD banner", b.CUDAVersion) + } +} + // TestRunNvidiaSMIReportsAbsence checks the two nvidia-smi probes surface a missing or // failing tool as (\"\", false) rather than an empty success, which is what keeps the // GPU fields unset instead of zeroed.