fix(gpu-detection): default to CPU if there is less than 4GB of GPU available

Signed-off-by: Ettore Di Giacinto <[email protected]>

Author: mudler

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())
 
@@ -272,7 +291,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())
 
@@ -344,7 +363,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())
 

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 {