diff --git a/core/gallery/backend_types.go b/core/gallery/backend_types.go index f57e7ffc..39167048 100644 --- a/core/gallery/backend_types.go +++ b/core/gallery/backend_types.go @@ -3,9 +3,10 @@ package gallery import "github.com/mudler/LocalAI/core/config" 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 +15,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 a2df466d..8217fa4f 100644 --- a/core/gallery/backends.go +++ b/core/gallery/backends.go @@ -6,12 +6,31 @@ import ( "path/filepath" "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 ( + aliasFile = "alias" + metaFile = "meta" + runFile = "run.sh" +) + +func findBestBackendFromMeta(backend *GalleryBackend, systemState *system.SystemState, backends GalleryElements[*GalleryBackend]) *GalleryBackend { + 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 +38,38 @@ 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) + } + + // Then, let's create an meta file to point to the best backend + metaFile := filepath.Join(metaBackendPath, metaFile) + if err := os.WriteFile(metaFile, []byte(bestBackend.Name), 0644); err != nil { + return fmt.Errorf("failed to write meta file %q: %v", metaFile, err) + } + + return nil } return InstallBackend(basePath, backend, downloadStatus) @@ -32,6 +82,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) @@ -50,7 +104,7 @@ func InstallBackend(basePath string, config *GalleryBackend, downloadStatus func if config.Alias != "" { // Write an alias file inside - aliasFile := filepath.Join(backendPath, "alias") + aliasFile := filepath.Join(backendPath, aliasFile) if err := os.WriteFile(aliasFile, []byte(config.Alias), 0644); err != nil { return fmt.Errorf("failed to write alias file %q: %v", aliasFile, err) } @@ -60,9 +114,52 @@ func InstallBackend(basePath string, config *GalleryBackend, downloadStatus func } 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 + } + + for _, backend := range backends { + if backend.IsDir() { + aliasFile := filepath.Join(basePath, backend.Name(), aliasFile) + alias, err := os.ReadFile(aliasFile) + if err != nil { + return err + } + if string(alias) == name { + backendDirectory = filepath.Join(basePath, backend.Name()) + break + } + } + } + + if backendDirectory == "" { + return fmt.Errorf("no backend found with name %q", name) + } + } + + // If it's a meta, delete also associated backend + metaFile := filepath.Join(backendDirectory, metaFile) + if _, err := os.Stat(metaFile); err == nil { + meta, err := os.ReadFile(metaFile) + if err != nil { + return err + } + metaBackendDirectory := filepath.Join(basePath, string(meta)) + 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", string(meta)) + } + os.RemoveAll(metaBackendDirectory) + } + + return os.RemoveAll(backendDirectory) } func ListSystemBackends(basePath string) (map[string]string, error) { @@ -75,10 +172,10 @@ 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") + aliasFile := filepath.Join(basePath, backend.Name(), aliasFile) if _, err := os.Stat(aliasFile); err == nil { // read the alias file, and use it as key alias, err := os.ReadFile(aliasFile) diff --git a/core/gallery/backends_test.go b/core/gallery/backends_test.go index 44b4fea5..1f78d760 100644 --- a/core/gallery/backends_test.go +++ b/core/gallery/backends_test.go @@ -5,8 +5,14 @@ import ( "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 +41,203 @@ 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 alias + metaBackendPath := filepath.Join(tempDir, "meta-backend") + err := os.MkdirAll(metaBackendPath, 0750) + Expect(err).NotTo(HaveOccurred()) + + // Create alias file pointing to concrete backend + aliasFilePath := filepath.Join(metaBackendPath, "meta") + err = os.WriteFile(aliasFilePath, []byte("concrete-backend"), 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") diff --git a/core/system/capabilities.go b/core/system/capabilities.go new file mode 100644 index 00000000..b2fd0459 --- /dev/null +++ b/core/system/capabilities.go @@ -0,0 +1,40 @@ +package system + +import ( + "strings" + + "github.com/mudler/LocalAI/pkg/xsysinfo" +) + +type SystemState struct { + GPUVendor string +} + +func GetSystemState() (*SystemState, error) { + gpuVendor, err := detectGPUVendor() + if err != nil { + return nil, err + } + return &SystemState{GPUVendor: gpuVendor}, nil +} + +func detectGPUVendor() (string, error) { + gpus, err := xsysinfo.GPUs() + if err != nil { + return "", err + } + + for _, gpu := range gpus { + if strings.ToUpper(gpu.DeviceInfo.Vendor.Name) == "NVIDIA" { + return "nvidia", nil + } + if strings.ToUpper(gpu.DeviceInfo.Vendor.Name) == "AMD" { + return "amd", nil + } + if strings.ToUpper(gpu.DeviceInfo.Vendor.Name) == "INTEL" { + return "intel", nil + } + } + + return "", nil +}