From a6d9988e8462808a205c3e645ec9989285b48dcf Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 24 Jun 2025 17:08:27 +0200 Subject: [PATCH] feat(backend gallery): add meta packages (#5696) * feat(backend gallery): add meta packages So we can have meta packages such as "vllm" that automatically installs the corresponding package depending on the GPU that is being currently detected in the system. Signed-off-by: Ettore Di Giacinto * feat: use a metadata file Signed-off-by: Ettore Di Giacinto --------- Signed-off-by: Ettore Di Giacinto --- core/gallery/backend_types.go | 25 +++- core/gallery/backends.go | 190 ++++++++++++++++++++++++--- core/gallery/backends_test.go | 230 +++++++++++++++++++++++++++++++-- core/http/app.go | 5 +- core/services/backends.go | 5 +- core/services/gallery.go | 13 +- core/system/capabilities.go | 49 +++++++ pkg/startup/backend_preload.go | 7 +- 8 files changed, 488 insertions(+), 36 deletions(-) create mode 100644 core/system/capabilities.go diff --git a/core/gallery/backend_types.go b/core/gallery/backend_types.go index f57e7ffc5..9c8bda933 100644 --- a/core/gallery/backend_types.go +++ b/core/gallery/backend_types.go @@ -2,10 +2,25 @@ package gallery import "github.com/mudler/LocalAI/core/config" +// BackendMetadata represents the metadata stored in a JSON file for each installed backend +type BackendMetadata struct { + // Alias is an optional alternative name for the backend + Alias string `json:"alias,omitempty"` + // MetaBackendFor points to the concrete backend if this is a meta backend + MetaBackendFor string `json:"meta_backend_for,omitempty"` + // Name is the original name from the gallery + Name string `json:"name,omitempty"` + // GalleryURL is the URL of the gallery this backend came from + GalleryURL string `json:"gallery_url,omitempty"` + // InstalledAt is the timestamp when the backend was installed + InstalledAt string `json:"installed_at,omitempty"` +} + type GalleryBackend struct { - Metadata `json:",inline" yaml:",inline"` - Alias string `json:"alias,omitempty" yaml:"alias,omitempty"` - URI string `json:"uri,omitempty" yaml:"uri,omitempty"` + Metadata `json:",inline" yaml:",inline"` + Alias string `json:"alias,omitempty" yaml:"alias,omitempty"` + URI string `json:"uri,omitempty" yaml:"uri,omitempty"` + CapabilitiesMap map[string]string `json:"capabilities,omitempty" yaml:"capabilities,omitempty"` } type GalleryBackends []*GalleryBackend @@ -14,6 +29,10 @@ func (m *GalleryBackend) SetGallery(gallery config.Gallery) { m.Gallery = gallery } +func (m *GalleryBackend) IsMeta() bool { + return len(m.CapabilitiesMap) > 0 +} + func (m *GalleryBackend) SetInstalled(installed bool) { m.Installed = installed } diff --git a/core/gallery/backends.go b/core/gallery/backends.go index a2df466d9..1b703e700 100644 --- a/core/gallery/backends.go +++ b/core/gallery/backends.go @@ -1,17 +1,79 @@ package gallery import ( + "encoding/json" "fmt" "os" "path/filepath" + "time" "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/system" "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/oci" + "github.com/rs/zerolog/log" ) +const ( + metadataFile = "metadata.json" + runFile = "run.sh" +) + +// readBackendMetadata reads the metadata JSON file for a backend +func readBackendMetadata(backendPath string) (*BackendMetadata, error) { + metadataPath := filepath.Join(backendPath, metadataFile) + + // If metadata file doesn't exist, return nil (for backward compatibility) + if _, err := os.Stat(metadataPath); os.IsNotExist(err) { + return nil, nil + } + + data, err := os.ReadFile(metadataPath) + if err != nil { + return nil, fmt.Errorf("failed to read metadata file %q: %v", metadataPath, err) + } + + var metadata BackendMetadata + if err := json.Unmarshal(data, &metadata); err != nil { + return nil, fmt.Errorf("failed to unmarshal metadata file %q: %v", metadataPath, err) + } + + return &metadata, nil +} + +// writeBackendMetadata writes the metadata JSON file for a backend +func writeBackendMetadata(backendPath string, metadata *BackendMetadata) error { + metadataPath := filepath.Join(backendPath, metadataFile) + + data, err := json.MarshalIndent(metadata, "", " ") + if err != nil { + return fmt.Errorf("failed to marshal metadata: %v", err) + } + + if err := os.WriteFile(metadataPath, data, 0644); err != nil { + return fmt.Errorf("failed to write metadata file %q: %v", metadataPath, err) + } + + return nil +} + +func findBestBackendFromMeta(backend *GalleryBackend, systemState *system.SystemState, backends GalleryElements[*GalleryBackend]) *GalleryBackend { + if systemState == nil { + return nil + } + + realBackend := backend.CapabilitiesMap[systemState.GPUVendor] + if realBackend == "" { + return nil + } + + return backends.FindByName(realBackend) +} + // Installs a model from the gallery -func InstallBackendFromGallery(galleries []config.Gallery, name string, basePath string, downloadStatus func(string, string, string, float64)) error { +func InstallBackendFromGallery(galleries []config.Gallery, systemState *system.SystemState, name string, basePath string, downloadStatus func(string, string, string, float64)) error { + log.Debug().Interface("galleries", galleries).Str("name", name).Msg("Installing backend from gallery") + backends, err := AvailableBackends(galleries, basePath) if err != nil { return err @@ -19,7 +81,44 @@ func InstallBackendFromGallery(galleries []config.Gallery, name string, basePath backend := FindGalleryElement(backends, name, basePath) if backend == nil { - return fmt.Errorf("no model found with name %q", name) + return fmt.Errorf("no backend found with name %q", name) + } + + if backend.IsMeta() { + log.Debug().Interface("systemState", systemState).Str("name", name).Msg("Backend is a meta backend") + + // Then, let's try to find the best backend based on the capabilities map + bestBackend := findBestBackendFromMeta(backend, systemState, backends) + if bestBackend == nil { + return fmt.Errorf("no backend found with capabilities %q", backend.CapabilitiesMap) + } + + log.Debug().Str("name", name).Str("bestBackend", bestBackend.Name).Msg("Installing backend from meta backend") + + // Then, let's install the best backend + if err := InstallBackend(basePath, bestBackend, downloadStatus); err != nil { + return err + } + + // we need now to create a path for the meta backend, with the alias to the installed ones so it can be used to remove it + metaBackendPath := filepath.Join(basePath, name) + if err := os.MkdirAll(metaBackendPath, 0750); err != nil { + return fmt.Errorf("failed to create meta backend path %q: %v", metaBackendPath, err) + } + + // Create metadata for the meta backend + metaMetadata := &BackendMetadata{ + MetaBackendFor: bestBackend.Name, + Name: name, + GalleryURL: backend.Gallery.URL, + InstalledAt: time.Now().Format(time.RFC3339), + } + + if err := writeBackendMetadata(metaBackendPath, metaMetadata); err != nil { + return fmt.Errorf("failed to write metadata for meta backend %q: %v", name, err) + } + + return nil } return InstallBackend(basePath, backend, downloadStatus) @@ -32,6 +131,10 @@ func InstallBackend(basePath string, config *GalleryBackend, downloadStatus func return fmt.Errorf("failed to create base path: %v", err) } + if config.IsMeta() { + return fmt.Errorf("meta backends cannot be installed directly") + } + name := config.Name img, err := oci.GetImage(config.URI, "", nil, nil) @@ -48,21 +151,73 @@ func InstallBackend(basePath string, config *GalleryBackend, downloadStatus func return fmt.Errorf("failed to extract image %q: %v", config.URI, err) } + // Create metadata for the backend + metadata := &BackendMetadata{ + Name: name, + GalleryURL: config.Gallery.URL, + InstalledAt: time.Now().Format(time.RFC3339), + } + if config.Alias != "" { - // Write an alias file inside - aliasFile := filepath.Join(backendPath, "alias") - if err := os.WriteFile(aliasFile, []byte(config.Alias), 0644); err != nil { - return fmt.Errorf("failed to write alias file %q: %v", aliasFile, err) - } + metadata.Alias = config.Alias + } + + if err := writeBackendMetadata(backendPath, metadata); err != nil { + return fmt.Errorf("failed to write metadata for backend %q: %v", name, err) } return nil } func DeleteBackendFromSystem(basePath string, name string) error { - backendFile := filepath.Join(basePath, name) + backendDirectory := filepath.Join(basePath, name) - return os.RemoveAll(backendFile) + // check if the backend dir exists + if _, err := os.Stat(backendDirectory); os.IsNotExist(err) { + // if doesn't exist, it might be an alias, so we need to check if we have a matching alias in + // all the backends in the basePath + backends, err := os.ReadDir(basePath) + if err != nil { + return err + } + foundBackend := false + + for _, backend := range backends { + if backend.IsDir() { + metadata, err := readBackendMetadata(filepath.Join(basePath, backend.Name())) + if err != nil { + return err + } + if metadata != nil && metadata.Alias == name { + backendDirectory = filepath.Join(basePath, backend.Name()) + foundBackend = true + break + } + } + } + + // If no backend found, return successfully (idempotent behavior) + if !foundBackend { + return fmt.Errorf("no backend found with name %q", name) + } + } + + // If it's a meta backend, delete also associated backend + metadata, err := readBackendMetadata(backendDirectory) + if err != nil { + return err + } + + if metadata != nil && metadata.MetaBackendFor != "" { + metaBackendDirectory := filepath.Join(basePath, metadata.MetaBackendFor) + log.Debug().Str("backendDirectory", metaBackendDirectory).Msg("Deleting meta backend") + if _, err := os.Stat(metaBackendDirectory); os.IsNotExist(err) { + return fmt.Errorf("meta backend %q not found", metadata.MetaBackendFor) + } + os.RemoveAll(metaBackendDirectory) + } + + return os.RemoveAll(backendDirectory) } func ListSystemBackends(basePath string) (map[string]string, error) { @@ -75,17 +230,16 @@ func ListSystemBackends(basePath string) (map[string]string, error) { for _, backend := range backends { if backend.IsDir() { - runFile := filepath.Join(basePath, backend.Name(), "run.sh") + runFile := filepath.Join(basePath, backend.Name(), runFile) backendsNames[backend.Name()] = runFile - aliasFile := filepath.Join(basePath, backend.Name(), "alias") - if _, err := os.Stat(aliasFile); err == nil { - // read the alias file, and use it as key - alias, err := os.ReadFile(aliasFile) - if err != nil { - return nil, err - } - backendsNames[string(alias)] = runFile + // Check for alias in metadata + metadata, err := readBackendMetadata(filepath.Join(basePath, backend.Name())) + if err != nil { + return nil, err + } + if metadata != nil && metadata.Alias != "" { + backendsNames[metadata.Alias] = runFile } } } diff --git a/core/gallery/backends_test.go b/core/gallery/backends_test.go index 44b4fea53..864ed3b5f 100644 --- a/core/gallery/backends_test.go +++ b/core/gallery/backends_test.go @@ -1,12 +1,19 @@ package gallery import ( + "encoding/json" "os" "path/filepath" "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/system" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "gopkg.in/yaml.v2" +) + +const ( + testImage = "quay.io/mudler/tests:localai-backend-test" ) var _ = Describe("Gallery Backends", func() { @@ -35,18 +42,209 @@ var _ = Describe("Gallery Backends", func() { Describe("InstallBackendFromGallery", func() { It("should return error when backend is not found", func() { - err := InstallBackendFromGallery(galleries, "non-existent", tempDir, nil) + err := InstallBackendFromGallery(galleries, nil, "non-existent", tempDir, nil) Expect(err).To(HaveOccurred()) - Expect(err.Error()).To(ContainSubstring("no model found with name")) + Expect(err.Error()).To(ContainSubstring("no backend found with name \"non-existent\"")) }) It("should install backend from gallery", func() { - err := InstallBackendFromGallery(galleries, "test-backend", tempDir, nil) + err := InstallBackendFromGallery(galleries, nil, "test-backend", tempDir, nil) Expect(err).ToNot(HaveOccurred()) Expect(filepath.Join(tempDir, "test-backend", "run.sh")).To(BeARegularFile()) }) }) + Describe("Meta Backends", func() { + It("should identify meta backends correctly", func() { + metaBackend := &GalleryBackend{ + Metadata: Metadata{ + Name: "meta-backend", + }, + CapabilitiesMap: map[string]string{ + "nvidia": "nvidia-backend", + "amd": "amd-backend", + "intel": "intel-backend", + }, + } + + Expect(metaBackend.IsMeta()).To(BeTrue()) + + regularBackend := &GalleryBackend{ + Metadata: Metadata{ + Name: "regular-backend", + }, + URI: testImage, + } + + Expect(regularBackend.IsMeta()).To(BeFalse()) + + emptyMetaBackend := &GalleryBackend{ + Metadata: Metadata{ + Name: "empty-meta-backend", + }, + CapabilitiesMap: map[string]string{}, + } + + Expect(emptyMetaBackend.IsMeta()).To(BeFalse()) + + nilMetaBackend := &GalleryBackend{ + Metadata: Metadata{ + Name: "nil-meta-backend", + }, + CapabilitiesMap: nil, + } + + Expect(nilMetaBackend.IsMeta()).To(BeFalse()) + }) + + It("should find best backend from meta based on system capabilities", func() { + metaBackend := &GalleryBackend{ + Metadata: Metadata{ + Name: "meta-backend", + }, + CapabilitiesMap: map[string]string{ + "nvidia": "nvidia-backend", + "amd": "amd-backend", + "intel": "intel-backend", + }, + } + + nvidiaBackend := &GalleryBackend{ + Metadata: Metadata{ + Name: "nvidia-backend", + }, + URI: testImage, + } + + amdBackend := &GalleryBackend{ + Metadata: Metadata{ + Name: "amd-backend", + }, + URI: testImage, + } + + backends := GalleryElements[*GalleryBackend]{nvidiaBackend, amdBackend} + + // Test with NVIDIA system state + nvidiaSystemState := &system.SystemState{GPUVendor: "nvidia"} + bestBackend := findBestBackendFromMeta(metaBackend, nvidiaSystemState, backends) + Expect(bestBackend).To(Equal(nvidiaBackend)) + + // Test with AMD system state + amdSystemState := &system.SystemState{GPUVendor: "amd"} + bestBackend = findBestBackendFromMeta(metaBackend, amdSystemState, backends) + Expect(bestBackend).To(Equal(amdBackend)) + + // Test with unsupported GPU vendor + unsupportedSystemState := &system.SystemState{GPUVendor: "unsupported"} + bestBackend = findBestBackendFromMeta(metaBackend, unsupportedSystemState, backends) + Expect(bestBackend).To(BeNil()) + }) + + It("should handle meta backend deletion correctly", func() { + metaBackend := &GalleryBackend{ + Metadata: Metadata{ + Name: "meta-backend", + }, + CapabilitiesMap: map[string]string{ + "nvidia": "nvidia-backend", + "amd": "amd-backend", + "intel": "intel-backend", + }, + } + + nvidiaBackend := &GalleryBackend{ + Metadata: Metadata{ + Name: "nvidia-backend", + }, + URI: testImage, + } + + amdBackend := &GalleryBackend{ + Metadata: Metadata{ + Name: "amd-backend", + }, + URI: testImage, + } + + gallery := config.Gallery{ + Name: "test-gallery", + URL: "file://" + filepath.Join(tempDir, "backend-gallery.yaml"), + } + + galleryBackend := GalleryBackends{amdBackend, nvidiaBackend, metaBackend} + + dat, err := yaml.Marshal(galleryBackend) + Expect(err).NotTo(HaveOccurred()) + err = os.WriteFile(filepath.Join(tempDir, "backend-gallery.yaml"), dat, 0644) + Expect(err).NotTo(HaveOccurred()) + + // Test with NVIDIA system state + nvidiaSystemState := &system.SystemState{GPUVendor: "nvidia"} + err = InstallBackendFromGallery([]config.Gallery{gallery}, nvidiaSystemState, "meta-backend", tempDir, nil) + Expect(err).NotTo(HaveOccurred()) + + metaBackendPath := filepath.Join(tempDir, "meta-backend") + Expect(metaBackendPath).To(BeADirectory()) + + concreteBackendPath := filepath.Join(tempDir, "nvidia-backend") + Expect(concreteBackendPath).To(BeADirectory()) + + allBackends, err := ListSystemBackends(tempDir) + Expect(err).NotTo(HaveOccurred()) + Expect(allBackends).To(HaveKey("meta-backend")) + Expect(allBackends).To(HaveKey("nvidia-backend")) + + // Delete meta backend by name + err = DeleteBackendFromSystem(tempDir, "meta-backend") + Expect(err).NotTo(HaveOccurred()) + + // Verify meta backend directory is deleted + Expect(metaBackendPath).NotTo(BeADirectory()) + + // Verify concrete backend directory is deleted + Expect(concreteBackendPath).NotTo(BeADirectory()) + }) + + It("should list meta backends correctly in system backends", func() { + // Create a meta backend directory with metadata + metaBackendPath := filepath.Join(tempDir, "meta-backend") + err := os.MkdirAll(metaBackendPath, 0750) + Expect(err).NotTo(HaveOccurred()) + + // Create metadata file pointing to concrete backend + metadata := &BackendMetadata{ + MetaBackendFor: "concrete-backend", + Name: "meta-backend", + InstalledAt: "2023-01-01T00:00:00Z", + } + metadataData, err := json.Marshal(metadata) + Expect(err).NotTo(HaveOccurred()) + err = os.WriteFile(filepath.Join(metaBackendPath, "metadata.json"), metadataData, 0644) + Expect(err).NotTo(HaveOccurred()) + + // Create the concrete backend directory with run.sh + concreteBackendPath := filepath.Join(tempDir, "concrete-backend") + err = os.MkdirAll(concreteBackendPath, 0750) + Expect(err).NotTo(HaveOccurred()) + err = os.WriteFile(filepath.Join(concreteBackendPath, "run.sh"), []byte("#!/bin/bash"), 0755) + Expect(err).NotTo(HaveOccurred()) + + // List system backends + backends, err := ListSystemBackends(tempDir) + Expect(err).NotTo(HaveOccurred()) + + // Should include both the meta backend name and concrete backend name + Expect(backends).To(HaveKey("meta-backend")) + Expect(backends).To(HaveKey("concrete-backend")) + + // meta-backend should point to its own run.sh + Expect(backends["meta-backend"]).To(Equal(filepath.Join(tempDir, "meta-backend", "run.sh"))) + // concrete-backend should point to its own run.sh + Expect(backends["concrete-backend"]).To(Equal(filepath.Join(tempDir, "concrete-backend", "run.sh"))) + }) + }) + Describe("InstallBackend", func() { It("should create base path if it doesn't exist", func() { newPath := filepath.Join(tempDir, "new-path") @@ -73,10 +271,17 @@ var _ = Describe("Gallery Backends", func() { err := InstallBackend(tempDir, &backend, nil) Expect(err).ToNot(HaveOccurred()) - Expect(filepath.Join(tempDir, "test-backend", "alias")).To(BeARegularFile()) - content, err := os.ReadFile(filepath.Join(tempDir, "test-backend", "alias")) + Expect(filepath.Join(tempDir, "test-backend", "metadata.json")).To(BeARegularFile()) + + // Read and verify metadata + metadataData, err := os.ReadFile(filepath.Join(tempDir, "test-backend", "metadata.json")) Expect(err).ToNot(HaveOccurred()) - Expect(string(content)).To(ContainSubstring("test-alias")) + var metadata BackendMetadata + err = json.Unmarshal(metadataData, &metadata) + Expect(err).ToNot(HaveOccurred()) + Expect(metadata.Alias).To(Equal("test-alias")) + Expect(metadata.Name).To(Equal("test-backend")) + Expect(filepath.Join(tempDir, "test-backend", "run.sh")).To(BeARegularFile()) // Check that the alias was recognized @@ -103,7 +308,7 @@ var _ = Describe("Gallery Backends", func() { It("should not error when backend doesn't exist", func() { err := DeleteBackendFromSystem(tempDir, "non-existent") - Expect(err).NotTo(HaveOccurred()) + Expect(err).To(HaveOccurred()) }) }) @@ -134,8 +339,15 @@ var _ = Describe("Gallery Backends", func() { err := os.MkdirAll(backendPath, 0750) Expect(err).NotTo(HaveOccurred()) - // Create alias file - err = os.WriteFile(filepath.Join(backendPath, "alias"), []byte(alias), 0644) + // Create metadata file with alias + metadata := &BackendMetadata{ + Alias: alias, + Name: backendName, + InstalledAt: "2023-01-01T00:00:00Z", + } + metadataData, err := json.Marshal(metadata) + Expect(err).NotTo(HaveOccurred()) + err = os.WriteFile(filepath.Join(backendPath, "metadata.json"), metadataData, 0644) Expect(err).NotTo(HaveOccurred()) backends, err := ListSystemBackends(tempDir) diff --git a/core/http/app.go b/core/http/app.go index ce8ce1641..aba00ff7e 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -205,7 +205,10 @@ func API(application *application.Application) (*fiber.App, error) { utils.LoadConfig(application.ApplicationConfig().ConfigsDir, openai.AssistantsFileConfigFile, &openai.AssistantFiles) galleryService := services.NewGalleryService(application.ApplicationConfig(), application.ModelLoader()) - galleryService.Start(application.ApplicationConfig().Context, application.BackendLoader()) + err = galleryService.Start(application.ApplicationConfig().Context, application.BackendLoader()) + if err != nil { + return nil, err + } requestExtractor := middleware.NewRequestExtractor(application.BackendLoader(), application.ModelLoader(), application.ApplicationConfig()) diff --git a/core/services/backends.go b/core/services/backends.go index b83ed8dd4..7a52a1b10 100644 --- a/core/services/backends.go +++ b/core/services/backends.go @@ -2,12 +2,13 @@ package services import ( "github.com/mudler/LocalAI/core/gallery" + "github.com/mudler/LocalAI/core/system" "github.com/mudler/LocalAI/pkg/utils" "github.com/rs/zerolog/log" ) -func (g *GalleryService) backendHandler(op *GalleryOp[gallery.GalleryBackend]) error { +func (g *GalleryService) backendHandler(op *GalleryOp[gallery.GalleryBackend], systemState *system.SystemState) error { utils.ResetDownloadTimers() g.UpdateStatus(op.ID, &GalleryOpStatus{Message: "processing", Progress: 0}) @@ -23,7 +24,7 @@ func (g *GalleryService) backendHandler(op *GalleryOp[gallery.GalleryBackend]) e g.modelLoader.DeleteExternalBackend(op.GalleryElementName) } else { log.Warn().Msgf("installing backend %s", op.GalleryElementName) - err = gallery.InstallBackendFromGallery(g.appConfig.BackendGalleries, op.GalleryElementName, g.appConfig.BackendsPath, progressCallback) + err = gallery.InstallBackendFromGallery(g.appConfig.BackendGalleries, systemState, op.GalleryElementName, g.appConfig.BackendsPath, progressCallback) if err == nil { err = gallery.RegisterBackends(g.appConfig.BackendsPath, g.modelLoader) } diff --git a/core/services/gallery.go b/core/services/gallery.go index 0c33d2435..b8306b30e 100644 --- a/core/services/gallery.go +++ b/core/services/gallery.go @@ -7,7 +7,9 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/gallery" + "github.com/mudler/LocalAI/core/system" "github.com/mudler/LocalAI/pkg/model" + "github.com/rs/zerolog/log" ) type GalleryService struct { @@ -50,7 +52,7 @@ func (g *GalleryService) GetAllStatus() map[string]*GalleryOpStatus { return g.statuses } -func (g *GalleryService) Start(c context.Context, cl *config.BackendConfigLoader) { +func (g *GalleryService) Start(c context.Context, cl *config.BackendConfigLoader) error { // updates the status with an error var updateError func(id string, e error) if !g.appConfig.OpaqueErrors { @@ -63,13 +65,18 @@ func (g *GalleryService) Start(c context.Context, cl *config.BackendConfigLoader } } + systemState, err := system.GetSystemState() + if err != nil { + log.Error().Err(err).Msg("failed to get system state") + } + go func() { for { select { case <-c.Done(): return case op := <-g.BackendGalleryChannel: - err := g.backendHandler(&op) + err := g.backendHandler(&op, systemState) if err != nil { updateError(op.ID, err) } @@ -82,4 +89,6 @@ func (g *GalleryService) Start(c context.Context, cl *config.BackendConfigLoader } } }() + + return nil } diff --git a/core/system/capabilities.go b/core/system/capabilities.go new file mode 100644 index 000000000..f1dba5104 --- /dev/null +++ b/core/system/capabilities.go @@ -0,0 +1,49 @@ +package system + +import ( + "strings" + + "github.com/mudler/LocalAI/pkg/xsysinfo" + "github.com/rs/zerolog/log" +) + +type SystemState struct { + GPUVendor string +} + +func GetSystemState() (*SystemState, error) { + gpuVendor, _ := detectGPUVendor() + log.Debug().Str("gpuVendor", gpuVendor).Msg("GPU vendor") + + return &SystemState{ + GPUVendor: gpuVendor, + }, nil +} + +func detectGPUVendor() (string, error) { + gpus, err := xsysinfo.GPUs() + if err != nil { + return "", err + } + + for _, gpu := range gpus { + if gpu.DeviceInfo != nil { + if gpu.DeviceInfo.Vendor != nil { + gpuVendorName := strings.ToUpper(gpu.DeviceInfo.Vendor.Name) + if gpuVendorName == "NVIDIA" { + return "nvidia", nil + } + if gpuVendorName == "AMD" { + return "amd", nil + } + if gpuVendorName == "INTEL" { + return "intel", nil + } + return "nvidia", nil + } + } + + } + + return "", nil +} diff --git a/pkg/startup/backend_preload.go b/pkg/startup/backend_preload.go index 17403c0ce..cbc37ca05 100644 --- a/pkg/startup/backend_preload.go +++ b/pkg/startup/backend_preload.go @@ -7,10 +7,15 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/gallery" + "github.com/mudler/LocalAI/core/system" ) func InstallExternalBackends(galleries []config.Gallery, backendPath string, downloadStatus func(string, string, string, float64), backends ...string) error { var errs error + systemState, err := system.GetSystemState() + if err != nil { + return fmt.Errorf("failed to get system state: %w", err) + } for _, backend := range backends { switch { case strings.HasPrefix(backend, "oci://"): @@ -22,7 +27,7 @@ func InstallExternalBackends(galleries []config.Gallery, backendPath string, dow errs = errors.Join(err, fmt.Errorf("error installing backend %s", backend)) } default: - err := gallery.InstallBackendFromGallery(galleries, backend, backendPath, downloadStatus) + err := gallery.InstallBackendFromGallery(galleries, systemState, backend, backendPath, downloadStatus) if err != nil { errs = errors.Join(err, fmt.Errorf("error installing backend %s", backend)) }