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

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

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto
2025-08-14 16:57:33 +02:00
committed by GitHub
parent b52bfaf1b3
commit 19c92c70c5
2 changed files with 47 additions and 17 deletions

View File

@@ -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())

View File

@@ -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 {