diff --git a/core/gallery/backends_test.go b/core/gallery/backends_test.go index 848840191..39a0521c7 100644 --- a/core/gallery/backends_test.go +++ b/core/gallery/backends_test.go @@ -105,10 +105,11 @@ var _ = Describe("Gallery Backends", func() { Name: "meta-backend", }, CapabilitiesMap: map[string]string{ - "nvidia": "nvidia-backend", - "amd": "amd-backend", - "intel": "intel-backend", - "metal": "metal-backend", + "nvidia": "nvidia-backend", + "amd": "amd-backend", + "intel": "intel-backend", + "metal": "metal-backend", + "default": "default-backend", }, } @@ -133,7 +134,14 @@ var _ = Describe("Gallery Backends", func() { URI: testImage, } - backends := GalleryElements[*GalleryBackend]{nvidiaBackend, amdBackend, metalBackend} + defaultBackend := &GalleryBackend{ + Metadata: Metadata{ + Name: "default-backend", + }, + URI: testImage, + } + + backends := GalleryElements[*GalleryBackend]{nvidiaBackend, amdBackend, metalBackend, defaultBackend} if runtime.GOOS == "darwin" && runtime.GOARCH == "arm64" { metal := &system.SystemState{} @@ -142,15 +150,26 @@ var _ = Describe("Gallery Backends", func() { } else { // Test with NVIDIA system state - nvidiaSystemState := &system.SystemState{GPUVendor: "nvidia"} + nvidiaSystemState := &system.SystemState{GPUVendor: "nvidia", VRAM: 1000000000000} bestBackend := metaBackend.FindBestBackendFromMeta(nvidiaSystemState, backends) Expect(bestBackend).To(Equal(nvidiaBackend)) // Test with AMD system state - amdSystemState := &system.SystemState{GPUVendor: "amd"} + amdSystemState := &system.SystemState{GPUVendor: "amd", VRAM: 1000000000000} bestBackend = metaBackend.FindBestBackendFromMeta(amdSystemState, backends) Expect(bestBackend).To(Equal(amdBackend)) + // Test with default system state (not enough VRAM) + defaultSystemState := &system.SystemState{GPUVendor: "amd"} + bestBackend = metaBackend.FindBestBackendFromMeta(defaultSystemState, backends) + Expect(bestBackend).To(Equal(defaultBackend)) + + // Test with default system state + defaultSystemState = &system.SystemState{GPUVendor: "default"} + bestBackend = metaBackend.FindBestBackendFromMeta(defaultSystemState, backends) + Expect(bestBackend).To(Equal(defaultBackend)) + + backends = GalleryElements[*GalleryBackend]{nvidiaBackend, amdBackend, metalBackend} // Test with unsupported GPU vendor unsupportedSystemState := &system.SystemState{GPUVendor: "unsupported"} bestBackend = metaBackend.FindBestBackendFromMeta(unsupportedSystemState, backends) @@ -201,7 +220,7 @@ var _ = Describe("Gallery Backends", func() { Expect(err).NotTo(HaveOccurred()) // Test with NVIDIA system state - nvidiaSystemState := &system.SystemState{GPUVendor: "nvidia"} + nvidiaSystemState := &system.SystemState{GPUVendor: "nvidia", VRAM: 1000000000000} err = InstallBackendFromGallery([]config.Gallery{gallery}, nvidiaSystemState, "meta-backend", tempDir, nil, true) Expect(err).NotTo(HaveOccurred()) @@ -275,7 +294,7 @@ var _ = Describe("Gallery Backends", func() { Expect(err).NotTo(HaveOccurred()) // Test with NVIDIA system state - nvidiaSystemState := &system.SystemState{GPUVendor: "nvidia"} + nvidiaSystemState := &system.SystemState{GPUVendor: "nvidia", VRAM: 1000000000000} err = InstallBackendFromGallery([]config.Gallery{gallery}, nvidiaSystemState, "meta-backend", tempDir, nil, true) Expect(err).NotTo(HaveOccurred()) @@ -352,7 +371,7 @@ var _ = Describe("Gallery Backends", func() { Expect(err).NotTo(HaveOccurred()) // Test with NVIDIA system state - nvidiaSystemState := &system.SystemState{GPUVendor: "nvidia"} + nvidiaSystemState := &system.SystemState{GPUVendor: "nvidia", VRAM: 1000000000000} err = InstallBackendFromGallery([]config.Gallery{gallery}, nvidiaSystemState, "meta-backend", tempDir, nil, true) Expect(err).NotTo(HaveOccurred()) diff --git a/pkg/system/capabilities.go b/pkg/system/capabilities.go index 19897dacb..056e569fa 100644 --- a/pkg/system/capabilities.go +++ b/pkg/system/capabilities.go @@ -5,12 +5,15 @@ import ( "runtime" "strings" + "github.com/jaypipes/ghw/pkg/gpu" "github.com/mudler/LocalAI/pkg/xsysinfo" "github.com/rs/zerolog/log" ) type SystemState struct { GPUVendor string + gpus []*gpu.GraphicsCard + VRAM uint64 } const ( @@ -91,24 +94,32 @@ func (s *SystemState) getSystemCapabilities() string { } log.Info().Str("Capability", s.GPUVendor).Msgf("Capability automatically detected, set %s to override", capabilityEnv) + // If vram is less than 4GB, let's default to CPU but warn the user that they can override that via env + if s.VRAM <= 4*1024*1024*1024 { + log.Warn().Msgf("VRAM is less than 4GB, defaulting to CPU. Set %s to override", capabilityEnv) + return defaultCapability + } + return s.GPUVendor } func GetSystemState() (*SystemState, error) { - gpuVendor, _ := detectGPUVendor() + // Detection is best-effort here, we don't want to fail if it fails + gpus, _ := xsysinfo.GPUs() + log.Debug().Any("gpus", gpus).Msg("GPUs") + gpuVendor, _ := detectGPUVendor(gpus) log.Debug().Str("gpuVendor", gpuVendor).Msg("GPU vendor") + vram, _ := xsysinfo.TotalAvailableVRAM() + log.Debug().Any("vram", vram).Msg("Total available VRAM") return &SystemState{ GPUVendor: gpuVendor, + gpus: gpus, + VRAM: vram, }, nil } -func detectGPUVendor() (string, error) { - gpus, err := xsysinfo.GPUs() - if err != nil { - return "", err - } - +func detectGPUVendor(gpus []*gpu.GraphicsCard) (string, error) { for _, gpu := range gpus { if gpu.DeviceInfo != nil { if gpu.DeviceInfo.Vendor != nil {