package discover import ( "io" "log/slog" "os" "path/filepath" "testing" "github.com/ollama/ollama/logutil" "github.com/ollama/ollama/ml" ) func TestLlamaServerDiscovery(t *testing.T) { originalProbe := probeLlamaServerVulkanDevices probeLlamaServerVulkanDevices = func(_ []string) ([]vulkanPhysicalDevice, error) { return nil, errWindowsVulkanProbeUnsupported } t.Cleanup(func() { probeLlamaServerVulkanDevices = originalProbe }) t.Run("output only trace", func(t *testing.T) { original := slog.Default() t.Cleanup(func() { slog.SetDefault(original) }) slog.SetDefault(logutil.NewLogger(io.Discard, slog.LevelDebug)) if got := llamaServerDiscoveryOutput(t.Context()); got != io.Discard { t.Fatal("debug logging should discard raw llama-server discovery output") } slog.SetDefault(logutil.NewLogger(io.Discard, logutil.LevelTrace)) if got := llamaServerDiscoveryOutput(t.Context()); got == io.Discard { t.Fatal("trace logging should emit raw llama-server discovery output") } }) t.Run("parse devices", func(t *testing.T) { type wantDevice struct { name string library string totalMiB uint64 compute string driver string gfxTarget string checkIntegrated bool integrated bool } tests := []struct { name string output string libDirs []string want []wantDevice }{ { name: "NVIDIA CUDA", output: `load_backend: loaded CUDA backend from /lib/ollama/cuda_v12/libggml-cuda.so Available devices: NVIDIA GeForce RTX 4090: NVIDIA CUDA (24564 MiB, 23592 MiB free) `, libDirs: []string{"/lib/ollama", "/lib/ollama/cuda_v12"}, want: []wantDevice{{ name: "NVIDIA GeForce RTX 4090", library: "CUDA", totalMiB: 24564, driver: "12.0", }}, }, { name: "Metal", output: `Available devices: Metal: Apple M3 Max (98304 MiB, 98303 MiB free) `, want: []wantDevice{{ name: "Metal", library: "Metal", totalMiB: 98304, }}, }, { name: "ROCm with gfx target", output: ` Device 0: AMD Radeon RX 6700 XT, gfx1031 (0x1031), VMM: no, Wave Size: 32, VRAM: 12272 MiB Available devices: ROCm0: AMD Radeon RX 6700 XT (12272 MiB, 12248 MiB free) `, libDirs: []string{"/lib/ollama", "/lib/ollama/rocm_v7_2"}, want: []wantDevice{{ name: "ROCm0", library: "ROCm", totalMiB: 12272, compute: "gfx1031", gfxTarget: "gfx1031", }}, }, { name: "multi GPU", output: `Available devices: CUDA0: NVIDIA GeForce RTX 4090 (24564 MiB, 23592 MiB free) CUDA1: NVIDIA GeForce RTX 3060 (12288 MiB, 11500 MiB free) `, libDirs: []string{"/lib/ollama", "/lib/ollama/cuda_v12"}, want: []wantDevice{ {name: "CUDA0", library: "CUDA", totalMiB: 24564}, {name: "CUDA1", library: "CUDA", totalMiB: 12288}, }, }, { name: "Vulkan UMA", output: `ggml_vulkan: 0 = Intel(R) Graphics (Intel open-source Mesa driver) | uma: 1 | fp16: 1 | bf16: 0 | warp size: 32 | shared memory: 65536 | int dot: 1 | matrix cores: none Available devices: Vulkan0: Intel(R) Graphics (16384 MiB, 12288 MiB free) `, libDirs: []string{"/lib/ollama", "/lib/ollama/vulkan"}, want: []wantDevice{{ name: "Vulkan0", library: "Vulkan", totalMiB: 16384, checkIntegrated: true, integrated: true, }}, }, { name: "Vulkan without UMA metadata", output: `Available devices: Vulkan0: AMD Radeon(TM) Graphics (32768 MiB, 31000 MiB free) `, libDirs: []string{"/lib/ollama", "/lib/ollama/vulkan"}, want: []wantDevice{{ name: "Vulkan0", library: "Vulkan", totalMiB: 32768, checkIntegrated: true, }}, }, { name: "CUDA device filtered by compiled archs", output: `ggml_cuda_init: found 1 CUDA devices (Total VRAM: 6063 MiB): Device 0: NVIDIA GeForce GTX 1060 6GB, compute capability 6.1, VMM: yes, VRAM: 6063 MiB load_backend: loaded CUDA backend from /lib/ollama/cuda_v13/libggml-cuda.so system_info: n_threads = 4 | CUDA : ARCHS = 750,800,860,890,900,1000,1030,1100,1200,1210 | Available devices: CUDA0: NVIDIA GeForce GTX 1060 6GB (6063 MiB, 5900 MiB free) `, libDirs: []string{"/lib/ollama", "/lib/ollama/cuda_v13"}, }, { name: "CUDA device kept by compiled archs", output: `ggml_cuda_init: found 1 CUDA devices (Total VRAM: 16379 MiB): Device 0: NVIDIA GeForce RTX 4060 Ti, compute capability 8.9, VMM: yes, VRAM: 16379 MiB system_info: n_threads = 16 | CUDA : ARCHS = 750,800,860,890,900,1000,1030,1100,1200,1210 | Available devices: CUDA0: NVIDIA GeForce RTX 4060 Ti (16379 MiB, 14900 MiB free) `, want: []wantDevice{{ name: "CUDA0", library: "CUDA", totalMiB: 16379, compute: "8.9", }}, }, { name: "CUDA without compiled archs fails open", output: `ggml_cuda_init: found 1 CUDA devices (Total VRAM: 6063 MiB): Device 0: NVIDIA GeForce GTX 1060 6GB, compute capability 6.1, VMM: yes, VRAM: 6063 MiB Available devices: CUDA0: NVIDIA GeForce GTX 1060 6GB (6063 MiB, 5900 MiB free) `, want: []wantDevice{{ name: "CUDA0", library: "CUDA", totalMiB: 6063, compute: "6.1", }}, }, { name: "CUDA without compute capability fails open", output: `system_info: n_threads = 4 | CUDA : ARCHS = 750,800 | Available devices: CUDA0: Some Future GPU (8192 MiB, 8000 MiB free) `, want: []wantDevice{{ name: "CUDA0", library: "CUDA", totalMiB: 8192, }}, }, { name: "CUDA mixed arch support", output: `ggml_cuda_init: found 2 CUDA devices: Device 0: NVIDIA GeForce GTX 1060, compute capability 6.1, VMM: yes, VRAM: 6063 MiB Device 1: NVIDIA GeForce RTX 4060 Ti, compute capability 8.9, VMM: yes, VRAM: 16379 MiB system_info: n_threads = 8 | CUDA : ARCHS = 750,800,860,890 | Available devices: CUDA0: NVIDIA GeForce GTX 1060 (6063 MiB, 5900 MiB free) CUDA1: NVIDIA GeForce RTX 4060 Ti (16379 MiB, 14900 MiB free) `, want: []wantDevice{{ name: "CUDA1", library: "CUDA", totalMiB: 16379, compute: "8.9", }}, }, { name: "ROCm gfx target with xnack suffix", output: `ggml_cuda_init: found 2 ROCm devices (Total VRAM: 32736 MiB): Device 0: AMD Radeon RX 6800, gfx1030 (0x1030), VMM: no, Wave Size: 32, VRAM: 16368 MiB Device 1: AMD Radeon Pro VII, gfx906:sramecc+:xnack- (0x906), VMM: no, Wave Size: 64, VRAM: 16368 MiB Available devices: ROCm0: AMD Radeon RX 6800 (16368 MiB, 16342 MiB free) ROCm1: AMD Radeon Pro VII (16368 MiB, 16348 MiB free) `, want: []wantDevice{ {name: "ROCm0", library: "ROCm", totalMiB: 16368, compute: "gfx1030", gfxTarget: "gfx1030"}, {name: "ROCm1", library: "ROCm", totalMiB: 16368, compute: "gfx906", gfxTarget: "gfx906"}, }, }, { name: "unknown library", output: `Available devices: Future0: Mystery Accelerator (8192 MiB, 8000 MiB free) `, want: []wantDevice{{ name: "Future0", library: "Mystery Accelerator", totalMiB: 8192, }}, }, { name: "no devices", output: "Available devices:\n", }, { name: "empty output", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { if tt.libDirs == nil { tt.libDirs = []string{"/lib/ollama"} } devices := parseLlamaServerDevices(tt.output, tt.libDirs) if len(devices) != len(tt.want) { t.Fatalf("got %d devices, want %d", len(devices), len(tt.want)) } for i, want := range tt.want { got := devices[i] if want.name != "" && got.Name != want.name { t.Errorf("device %d name = %q, want %q", i, got.Name, want.name) } if want.library != "" && got.Library != want.library { t.Errorf("device %d library = %q, want %q", i, got.Library, want.library) } if want.totalMiB > 0 && got.TotalMemory != want.totalMiB*1024*1024 { t.Errorf("device %d total memory = %d, want %d MiB", i, got.TotalMemory, want.totalMiB) } if want.compute != "" && got.Compute() != want.compute { t.Errorf("device %d compute = %q, want %q", i, got.Compute(), want.compute) } if want.driver != "" && got.Driver() != want.driver { t.Errorf("device %d driver = %q, want %q", i, got.Driver(), want.driver) } if want.gfxTarget != "" && got.GFXTarget != want.gfxTarget { t.Errorf("device %d gfx target = %q, want %q", i, got.GFXTarget, want.gfxTarget) } if want.checkIntegrated && got.Integrated != want.integrated { t.Errorf("device %d integrated = %v, want %v", i, got.Integrated, want.integrated) } } }) } }) t.Run("parse fixtures", func(t *testing.T) { type wantDevice struct { name string library string totalMiB uint64 compute string gfxTarget string integrated bool } tests := []struct { name string output string libDirs []string want []wantDevice wantSkip string }{ { name: "cuda mixed archs filters unsupported device", output: `ggml_cuda_init: found 2 CUDA devices: Device 0: NVIDIA GeForce GTX 1060, compute capability 6.1, VMM: yes, VRAM: 6063 MiB Device 1: NVIDIA GeForce RTX 4060 Ti, compute capability 8.9, VMM: yes, VRAM: 16379 MiB system_info: n_threads = 8 | CUDA : ARCHS = 750,800,860,890 | Available devices: CUDA0: NVIDIA GeForce GTX 1060 (6063 MiB, 5900 MiB free) CUDA1: NVIDIA GeForce RTX 4060 Ti (16379 MiB, 14900 MiB free) `, want: []wantDevice{{ name: "CUDA1", library: "CUDA", totalMiB: 16379, compute: "8.9", }}, }, { name: "rocm gfx targets preserve suffix-free compute", output: `ggml_cuda_init: found 2 ROCm devices (Total VRAM: 32736 MiB): Device 0: AMD Radeon RX 6800, gfx1030 (0x1030), VMM: no, Wave Size: 32, VRAM: 16368 MiB Device 1: AMD Radeon Pro VII, gfx906:sramecc+:xnack- (0x906), VMM: no, Wave Size: 64, VRAM: 16368 MiB Available devices: ROCm0: AMD Radeon RX 6800 (16368 MiB, 16342 MiB free) ROCm1: AMD Radeon Pro VII (16368 MiB, 16348 MiB free) `, want: []wantDevice{ {name: "ROCm0", library: "ROCm", totalMiB: 16368, compute: "gfx1030", gfxTarget: "gfx1030"}, {name: "ROCm1", library: "ROCm", totalMiB: 16368, compute: "gfx906", gfxTarget: "gfx906"}, }, }, { name: "vulkan uma marks integrated", output: `ggml_vulkan: 0 = Intel(R) Graphics (Intel open-source Mesa driver) | uma: 1 | fp16: 1 | bf16: 0 | warp size: 32 | shared memory: 65536 | int dot: 1 | matrix cores: none Available devices: Vulkan0: Intel(R) Graphics (16384 MiB, 12288 MiB free) `, want: []wantDevice{{ name: "Vulkan0", library: "Vulkan", totalMiB: 16384, integrated: true, }}, }, { name: "windows vulkan without uma stays unclassified", output: `load_backend: loaded Vulkan backend from C:\ollama\lib\ollama\vulkan\ggml-vulkan.dll Available devices: Vulkan0: AMD Radeon(TM) Graphics (32768 MiB, 31000 MiB free) Vulkan1: AMD Radeon RX 7900 XTX (24564 MiB, 23000 MiB free) `, want: []wantDevice{ {name: "Vulkan0", library: "Vulkan", totalMiB: 32768}, {name: "Vulkan1", library: "Vulkan", totalMiB: 24564}, }, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { libDirs := tt.libDirs if libDirs == nil { libDirs = []string{"/lib/ollama"} } got := parseLlamaServerDevices(tt.output, libDirs) if len(got) != len(tt.want) { t.Fatalf("got %d devices, want %d", len(got), len(tt.want)) } for i, want := range tt.want { if got[i].Name != want.name { t.Fatalf("device %d name = %q, want %q", i, got[i].Name, want.name) } if got[i].Library != want.library { t.Fatalf("device %d library = %q, want %q", i, got[i].Library, want.library) } if got[i].TotalMemory != want.totalMiB*1024*1024 { t.Fatalf("device %d total memory = %d, want %d MiB", i, got[i].TotalMemory, want.totalMiB) } if want.compute != "" && got[i].Compute() != want.compute { t.Fatalf("device %d compute = %q, want %q", i, got[i].Compute(), want.compute) } if want.gfxTarget != "" && got[i].GFXTarget != want.gfxTarget { t.Fatalf("device %d gfx target = %q, want %q", i, got[i].GFXTarget, want.gfxTarget) } if got[i].Integrated != want.integrated { t.Fatalf("device %d integrated = %v, want %v", i, got[i].Integrated, want.integrated) } } }) } }) t.Run("skips mismatched Vulkan native metadata", func(t *testing.T) { output := `Available devices: Vulkan0: Intel(R) UHD Graphics 770 (32768 MiB, 31000 MiB free) ` nativeDevices := []nativeProbeDevice{{ Library: "Vulkan", Index: 0, IndexMatchesBackend: true, Description: "NVIDIA GeForce RTX 4060 Ti", DeviceID: "0000:05:00.0", IntegratedKnown: true, TotalMemory: 16107 * 1024 * 1024, }} devices := parseLlamaServerDevicesWithNative(output, "", []string{"/lib/ollama", "/lib/ollama/vulkan"}, nativeDevices) if len(devices) != 1 { t.Fatalf("got %d devices, want 1", len(devices)) } if devices[0].PCIID != "" { t.Fatalf("PCIID = %q, want empty", devices[0].PCIID) } if devices[0].Integrated { t.Fatal("Integrated = true, want false") } }) t.Run("native probe uma lines do not key into llama-server device order", func(t *testing.T) { // llama-server enumerates [Intel iGPU, NVIDIA dGPU]; the native probe // enumerates the same devices in the opposite order. Its ggml_vulkan // uma lines must not overwrite llama-server's index-keyed UMA map, // otherwise the classification inverts (#16667). output := `ggml_vulkan: 0 = Intel(R) RaptorLake-S Mobile Graphics Controller (Intel Corporation) | uma: 1 | fp16: 1 | bf16: 0 | warp size: 32 | shared memory: 65536 | int dot: 1 | matrix cores: none ggml_vulkan: 1 = NVIDIA GeForce RTX 4080 Laptop GPU (NVIDIA) | uma: 0 | fp16: 1 | bf16: 0 | warp size: 32 | shared memory: 49152 | int dot: 1 | matrix cores: KHR_coopmat Available devices: Vulkan0: Intel(R) RaptorLake-S Mobile Graphics Controller (32550 MiB, 31800 MiB free) Vulkan1: NVIDIA GeForce RTX 4080 Laptop GPU (12282 MiB, 11000 MiB free) ` nativeOutput := `ggml_vulkan: 0 = NVIDIA GeForce RTX 4080 Laptop GPU (NVIDIA) | uma: 0 | fp16: 1 | bf16: 0 | warp size: 32 | shared memory: 49152 | int dot: 1 | matrix cores: KHR_coopmat ggml_vulkan: 1 = Intel(R) RaptorLake-S Mobile Graphics Controller (Intel Corporation) | uma: 1 | fp16: 1 | bf16: 0 | warp size: 32 | shared memory: 65536 | int dot: 1 | matrix cores: none ` devices := parseLlamaServerDevicesWithNative(output, nativeOutput, []string{"/lib/ollama", "/lib/ollama/vulkan"}, nil) if len(devices) != 2 { t.Fatalf("got %d devices, want 2", len(devices)) } if !devices[0].Integrated { t.Fatal("device 0 (Intel iGPU) Integrated = false, want true") } if devices[1].Integrated { t.Fatal("device 1 (NVIDIA dGPU) Integrated = true, want false") } }) t.Run("cuda runtime version", func(t *testing.T) { dir := t.TempDir() if err := os.WriteFile(filepath.Join(dir, "libcudart.so.12.8.90"), nil, 0o644); err != nil { t.Fatal(err) } major, minor, ok := cudaRuntimeVersion([]string{dir}) if !ok || major != 12 || minor != 8 { t.Fatalf("cudaRuntimeVersion = %d.%d, %v, want 12.8, true", major, minor, ok) } major, minor, ok = cudaRuntimeVersion([]string{filepath.Join(t.TempDir(), "cuda_v13")}) if !ok || major != 13 || minor != 0 { t.Fatalf("cudaRuntimeVersion fallback = %d.%d, %v, want 13.0, true", major, minor, ok) } }) t.Run("refine windows vulkan devices", func(t *testing.T) { makeDevices := func() []ml.DeviceInfo { return []ml.DeviceInfo{ {DeviceID: ml.DeviceID{ID: "0", Library: "Vulkan"}, Description: "AMD Radeon(TM) Graphics"}, {DeviceID: ml.DeviceID{ID: "1", Library: "Vulkan"}, Description: "AMD Radeon RX 7900 XTX"}, {DeviceID: ml.DeviceID{ID: "0", Library: "CUDA"}, Description: "NVIDIA GeForce RTX 4090"}, } } tests := []struct { name string devices []ml.DeviceInfo probed []vulkanPhysicalDevice want []bool applied bool }{ { name: "fills missing integrated bit", probed: []vulkanPhysicalDevice{ {Name: "AMD Radeon(TM) Graphics", Integrated: true}, {Name: "AMD Radeon RX 7900 XTX", Integrated: false}, }, want: []bool{true, false, false}, applied: true, }, { name: "matches names when order differs", probed: []vulkanPhysicalDevice{ {Name: "AMD Radeon RX 7900 XTX", Integrated: false}, {Name: "AMD Radeon(TM) Graphics", Integrated: true}, }, want: []bool{true, false, false}, applied: true, }, { name: "skips when names do not line up", probed: []vulkanPhysicalDevice{ {Name: "Wrong GPU", Integrated: true}, {Name: "AMD Radeon RX 7900 XTX", Integrated: false}, }, want: []bool{false, false, false}, }, { name: "skips when probe finds fewer devices", probed: []vulkanPhysicalDevice{{Name: "AMD Radeon(TM) Graphics", Integrated: true}}, want: []bool{false, false, false}, }, { name: "matches subset when probe enumerates extra devices", devices: []ml.DeviceInfo{ {DeviceID: ml.DeviceID{ID: "0", Library: "Vulkan"}, Description: "Intel(R) RaptorLake-S Mobile Graphics Controller"}, {DeviceID: ml.DeviceID{ID: "1", Library: "Vulkan"}, Description: "NVIDIA GeForce RTX 4080 Laptop GPU"}, }, probed: []vulkanPhysicalDevice{ {Name: "NVIDIA GeForce RTX 4080 Laptop GPU", Integrated: false}, {Name: "Intel(R) RaptorLake-S Mobile Graphics Controller", Integrated: true}, {Name: "Microsoft Direct3D12 (NVIDIA GeForce RTX 4080 Laptop GPU)", Integrated: false}, {Name: "Microsoft Direct3D12 (Intel(R) RaptorLake-S Mobile Graphics Controller)", Integrated: true}, {Name: "llvmpipe (LLVM 17.0.6, 256 bits)", Integrated: false}, }, want: []bool{true, false}, applied: true, }, { name: "skips ambiguous duplicate names with conflicting types", devices: []ml.DeviceInfo{ {DeviceID: ml.DeviceID{ID: "0", Library: "Vulkan"}, Description: "AMD Radeon Graphics"}, }, probed: []vulkanPhysicalDevice{ {Name: "AMD Radeon Graphics", Integrated: true}, {Name: "AMD Radeon Graphics", Integrated: false}, }, want: []bool{false}, }, { name: "matches duplicate names with agreeing types", devices: []ml.DeviceInfo{ {DeviceID: ml.DeviceID{ID: "0", Library: "Vulkan"}, Description: "AMD Radeon RX 7900 XTX"}, {DeviceID: ml.DeviceID{ID: "1", Library: "Vulkan"}, Description: "AMD Radeon RX 7900 XTX"}, }, probed: []vulkanPhysicalDevice{ {Name: "AMD Radeon RX 7900 XTX", Integrated: false}, {Name: "AMD Radeon RX 7900 XTX", Integrated: false}, }, want: []bool{false, false}, applied: true, }, { name: "overwrites stale classification", devices: []ml.DeviceInfo{ {DeviceID: ml.DeviceID{ID: "0", Library: "Vulkan"}, Description: "AMD Radeon(TM) Graphics", Integrated: true}, {DeviceID: ml.DeviceID{ID: "1", Library: "Vulkan"}, Description: "AMD Radeon RX 7900 XTX"}, }, probed: []vulkanPhysicalDevice{ {Name: "AMD Radeon(TM) Graphics", Integrated: false}, {Name: "AMD Radeon RX 7900 XTX", Integrated: false}, }, want: []bool{false, false}, applied: true, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { devices := tt.devices if devices == nil { devices = makeDevices() } applied := applyWindowsVulkanRefinement(devices, tt.probed) if applied != tt.applied { t.Fatalf("applied = %v, want %v", applied, tt.applied) } got := devices if len(got) != len(tt.want) { t.Fatalf("got %d devices, want %d", len(got), len(tt.want)) } for i, want := range tt.want { if got[i].Integrated != want { t.Fatalf("device %d integrated = %v, want %v", i, got[i].Integrated, want) } } }) } }) }