wehub-resource-sync
271 lines
7.3 KiB
Go
271 lines
7.3 KiB
Go
package ml
|
|
|
|
import (
|
|
"bytes"
|
|
"log/slog"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestMergeEnvWithRunnerEnvOverrides(t *testing.T) {
|
|
devices := []DeviceInfo{
|
|
{
|
|
DeviceID: DeviceID{Library: "Metal", ID: "0"},
|
|
RunnerEnvOverrides: map[string]string{"GGML_METAL_TENSOR_DISABLE": "1"},
|
|
},
|
|
}
|
|
|
|
env := GetDevicesEnv(devices)
|
|
|
|
if got, want := env["GGML_METAL_TENSOR_DISABLE"], "1"; got != want {
|
|
t.Fatalf("GGML_METAL_TENSOR_DISABLE = %q, want %q", got, want)
|
|
}
|
|
}
|
|
|
|
func TestGetDevicesEnvWarnsOnConflictingOverrides(t *testing.T) {
|
|
var logs bytes.Buffer
|
|
oldLogger := slog.Default()
|
|
slog.SetDefault(slog.New(slog.NewTextHandler(&logs, &slog.HandlerOptions{Level: slog.LevelDebug})))
|
|
t.Cleanup(func() {
|
|
slog.SetDefault(oldLogger)
|
|
})
|
|
|
|
devices := []DeviceInfo{
|
|
{
|
|
DeviceID: DeviceID{Library: "Metal", ID: "0"},
|
|
RunnerEnvOverrides: map[string]string{"TEST_OVERRIDE": "one"},
|
|
},
|
|
{
|
|
DeviceID: DeviceID{Library: "Metal", ID: "1"},
|
|
RunnerEnvOverrides: map[string]string{"TEST_OVERRIDE": "two"},
|
|
},
|
|
}
|
|
|
|
env := GetDevicesEnv(devices)
|
|
|
|
if got, want := env["TEST_OVERRIDE"], "two"; got != want {
|
|
t.Fatalf("TEST_OVERRIDE = %q, want %q", got, want)
|
|
}
|
|
|
|
if !strings.Contains(logs.String(), "conflicting device environment override") {
|
|
t.Fatalf("expected warning log, got %q", logs.String())
|
|
}
|
|
}
|
|
|
|
func TestGetDevicesEnvFiltersVisibleDevices(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
gpus []DeviceInfo
|
|
key string
|
|
want string
|
|
}{
|
|
{
|
|
name: "single CUDA",
|
|
gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA", ID: "3"}}},
|
|
key: "CUDA_VISIBLE_DEVICES",
|
|
want: "3",
|
|
},
|
|
{
|
|
name: "multiple CUDA",
|
|
gpus: []DeviceInfo{
|
|
{DeviceID: DeviceID{Library: "CUDA", ID: "3"}},
|
|
{DeviceID: DeviceID{Library: "CUDA", ID: "4"}},
|
|
},
|
|
key: "CUDA_VISIBLE_DEVICES",
|
|
want: "3,4",
|
|
},
|
|
{
|
|
name: "multiple remapped CUDA",
|
|
gpus: []DeviceInfo{
|
|
{DeviceID: DeviceID{Library: "CUDA", ID: "0"}, FilterID: "0"},
|
|
{DeviceID: DeviceID{Library: "CUDA", ID: "1"}, FilterID: "2"},
|
|
},
|
|
key: "CUDA_VISIBLE_DEVICES",
|
|
want: "0,2",
|
|
},
|
|
{
|
|
name: "single Vulkan",
|
|
gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "Vulkan", ID: "0"}, FilterID: "1"}},
|
|
key: "GGML_VK_VISIBLE_DEVICES",
|
|
want: "1",
|
|
},
|
|
{
|
|
name: "multiple Vulkan",
|
|
gpus: []DeviceInfo{
|
|
{DeviceID: DeviceID{Library: "Vulkan", ID: "0"}, FilterID: "1"},
|
|
{DeviceID: DeviceID{Library: "Vulkan", ID: "1"}, FilterID: "0"},
|
|
},
|
|
key: "GGML_VK_VISIBLE_DEVICES",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
env := GetDevicesEnv(tt.gpus)
|
|
if got := env[tt.key]; got != tt.want {
|
|
t.Fatalf("%s = %q, want %q", tt.key, got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestDeviceCompareVulkanDuplicates(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
a DeviceInfo
|
|
b DeviceInfo
|
|
want DeviceComparison
|
|
}{
|
|
{
|
|
name: "vulkan duplicate with missing pci",
|
|
a: DeviceInfo{
|
|
DeviceID: DeviceID{Library: "CUDA", ID: "0"},
|
|
Description: "NVIDIA GeForce RTX 4060 Ti",
|
|
PCIID: "0000:01:00.0",
|
|
TotalMemory: 16 * 1024 * 1024 * 1024,
|
|
},
|
|
b: DeviceInfo{
|
|
DeviceID: DeviceID{Library: "Vulkan", ID: "2"},
|
|
Description: "NVIDIA GeForce RTX 4060 Ti",
|
|
TotalMemory: 16107 * 1024 * 1024,
|
|
},
|
|
want: DuplicateDevice,
|
|
},
|
|
{
|
|
name: "vulkan radv suffix duplicate",
|
|
a: DeviceInfo{
|
|
DeviceID: DeviceID{Library: "ROCm", ID: "0"},
|
|
Description: "AMD Radeon RX 6400",
|
|
PCIID: "0000:03:00.0",
|
|
TotalMemory: 4 * 1024 * 1024 * 1024,
|
|
},
|
|
b: DeviceInfo{
|
|
DeviceID: DeviceID{Library: "Vulkan", ID: "0"},
|
|
Description: "AMD Radeon RX 6400 (RADV NAVI24)",
|
|
TotalMemory: 4096 * 1024 * 1024,
|
|
},
|
|
want: DuplicateDevice,
|
|
},
|
|
{
|
|
name: "different vulkan igpu",
|
|
a: DeviceInfo{
|
|
DeviceID: DeviceID{Library: "CUDA", ID: "0"},
|
|
Description: "NVIDIA GeForce RTX 4060 Ti",
|
|
PCIID: "0000:01:00.0",
|
|
TotalMemory: 16 * 1024 * 1024 * 1024,
|
|
},
|
|
b: DeviceInfo{
|
|
DeviceID: DeviceID{Library: "Vulkan", ID: "0"},
|
|
Description: "Intel(R) UHD Graphics 770",
|
|
TotalMemory: 32 * 1024 * 1024 * 1024,
|
|
},
|
|
want: UniqueDevice,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
if got := tt.a.Compare(tt.b); got != tt.want {
|
|
t.Fatalf("Compare = %v, want %v", got, tt.want)
|
|
}
|
|
if got := tt.b.Compare(tt.a); got != tt.want {
|
|
t.Fatalf("reverse Compare = %v, want %v", got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestFlashAttentionSupported(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
gpus []DeviceInfo
|
|
want bool
|
|
}{
|
|
{
|
|
name: "cuda compute 5.0 unsupported",
|
|
gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 5, ComputeMinor: 0}},
|
|
},
|
|
{
|
|
name: "cuda compute 6.0 supported",
|
|
gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 6, ComputeMinor: 0}},
|
|
want: true,
|
|
},
|
|
{
|
|
name: "cuda compute 6.1 supported",
|
|
gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 6, ComputeMinor: 1}},
|
|
want: true,
|
|
},
|
|
{
|
|
name: "cuda compute 6.2 supported",
|
|
gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 6, ComputeMinor: 2}},
|
|
want: true,
|
|
},
|
|
{
|
|
name: "cuda compute 7.2 unsupported",
|
|
gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 7, ComputeMinor: 2}},
|
|
},
|
|
{
|
|
name: "cuda compute 7.0 supported",
|
|
gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 7, ComputeMinor: 0}},
|
|
want: true,
|
|
},
|
|
{
|
|
name: "cuda compute 7.5 supported",
|
|
gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 7, ComputeMinor: 5}},
|
|
want: true,
|
|
},
|
|
{
|
|
name: "cuda unknown driver supported by compute",
|
|
gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA"}, ComputeMajor: 8, ComputeMinor: 9}},
|
|
want: true,
|
|
},
|
|
{
|
|
name: "cuda unknown compute unsupported",
|
|
gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: -1, ComputeMinor: -1}},
|
|
},
|
|
{
|
|
name: "mixed cuda unsupported",
|
|
gpus: []DeviceInfo{
|
|
{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 8, ComputeMinor: 9},
|
|
{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 5, ComputeMinor: 0},
|
|
},
|
|
},
|
|
{
|
|
name: "metal supported",
|
|
gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "Metal"}}},
|
|
want: true,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
if got := FlashAttentionSupported(tt.gpus); got != tt.want {
|
|
t.Fatalf("FlashAttentionSupported = %t, want %t", got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestFlashAttentionSupportedWarnsWhenCUDADriverUnknown(t *testing.T) {
|
|
var logs bytes.Buffer
|
|
oldLogger := slog.Default()
|
|
slog.SetDefault(slog.New(slog.NewTextHandler(&logs, &slog.HandlerOptions{Level: slog.LevelWarn})))
|
|
t.Cleanup(func() {
|
|
slog.SetDefault(oldLogger)
|
|
})
|
|
|
|
gpus := []DeviceInfo{{
|
|
DeviceID: DeviceID{Library: "CUDA"},
|
|
Description: "NVIDIA RTX",
|
|
ComputeMajor: 8,
|
|
ComputeMinor: 9,
|
|
}}
|
|
|
|
if !FlashAttentionSupported(gpus) {
|
|
t.Fatal("expected unknown-driver modern CUDA GPU to allow flash attention")
|
|
}
|
|
if !strings.Contains(logs.String(), "CUDA driver version unavailable") {
|
|
t.Fatalf("expected unknown driver warning, got %q", logs.String())
|
|
}
|
|
}
|