feat(gallery): automatically install missing backends along models (#5736)

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

Author: mudler

core/application/startup.go

@@ -57,7 +57,7 @@ func New(opts ...config.AppOption) (*Application, error) {
 		}
 	}
 
-	if err := pkgStartup.InstallModels(options.Galleries, options.ModelPath, options.EnforcePredownloadScans, nil, options.ModelsURL...); err != nil {
+	if err := pkgStartup.InstallModels(options.Galleries, options.BackendGalleries, options.ModelPath, options.BackendsPath, options.EnforcePredownloadScans, options.AutoloadBackendGalleries, nil, options.ModelsURL...); err != nil {
 		log.Error().Err(err).Msg("error installing models")
 	}
 
@@ -86,13 +86,13 @@ func New(opts ...config.AppOption) (*Application, error) {
 	}
 
 	if options.PreloadJSONModels != "" {
-		if err := services.ApplyGalleryFromString(options.ModelPath, options.PreloadJSONModels, options.EnforcePredownloadScans, options.Galleries); err != nil {
+		if err := services.ApplyGalleryFromString(options.ModelPath, options.BackendsPath, options.EnforcePredownloadScans, options.AutoloadBackendGalleries, options.Galleries, options.BackendGalleries, options.PreloadJSONModels); err != nil {
 			return nil, err
 		}
 	}
 
 	if options.PreloadModelsFromPath != "" {
-		if err := services.ApplyGalleryFromFile(options.ModelPath, options.PreloadModelsFromPath, options.EnforcePredownloadScans, options.Galleries); err != nil {
+		if err := services.ApplyGalleryFromFile(options.ModelPath, options.BackendsPath, options.EnforcePredownloadScans, options.AutoloadBackendGalleries, options.Galleries, options.BackendGalleries, options.PreloadModelsFromPath); err != nil {
 			return nil, err
 		}
 	}

core/backend/llm.go

@@ -42,9 +42,10 @@ func ModelInference(ctx context.Context, s string, messages []schema.Message, im
 		if _, err := os.Stat(modelFile); os.IsNotExist(err) {
 			utils.ResetDownloadTimers()
 			// if we failed to load the model, we try to download it
-			err := gallery.InstallModelFromGallery(o.Galleries, modelFile, loader.ModelPath, gallery.GalleryModel{}, utils.DisplayDownloadFunction, o.EnforcePredownloadScans)
+			err := gallery.InstallModelFromGallery(o.Galleries, o.BackendGalleries, modelFile, loader.ModelPath, o.BackendsPath, gallery.GalleryModel{}, utils.DisplayDownloadFunction, o.EnforcePredownloadScans, o.AutoloadBackendGalleries)
 			if err != nil {
-				return nil, err
+				log.Error().Err(err).Msgf("failed to install model %q from gallery", modelFile)
+				//return nil, err
 			}
 		}
 	}

core/cli/models.go

@@ -16,17 +16,20 @@ import (
 )
 
 type ModelsCMDFlags struct {
-	Galleries  string `env:"LOCALAI_GALLERIES,GALLERIES" help:"JSON list of galleries" group:"models" default:"${galleries}"`
-	ModelsPath string `env:"LOCALAI_MODELS_PATH,MODELS_PATH" type:"path" default:"${basepath}/models" help:"Path containing models used for inferencing" group:"storage"`
+	Galleries        string `env:"LOCALAI_GALLERIES,GALLERIES" help:"JSON list of galleries" group:"models" default:"${galleries}"`
+	BackendGalleries string `env:"LOCALAI_BACKEND_GALLERIES,BACKEND_GALLERIES" help:"JSON list of backend galleries" group:"backends" default:"${backends}"`
+	ModelsPath       string `env:"LOCALAI_MODELS_PATH,MODELS_PATH" type:"path" default:"${basepath}/models" help:"Path containing models used for inferencing" group:"storage"`
+	BackendsPath     string `env:"LOCALAI_BACKENDS_PATH,BACKENDS_PATH" type:"path" default:"${basepath}/backends" help:"Path containing backends used for inferencing" group:"storage"`
 }
 
 type ModelsList struct {
 	ModelsCMDFlags `embed:""`
 }
 
 type ModelsInstall struct {
-	DisablePredownloadScan bool     `env:"LOCALAI_DISABLE_PREDOWNLOAD_SCAN" help:"If true, disables the best-effort security scanner before downloading any files." group:"hardening" default:"false"`
-	ModelArgs              []string `arg:"" optional:"" name:"models" help:"Model configuration URLs to load"`
+	DisablePredownloadScan   bool     `env:"LOCALAI_DISABLE_PREDOWNLOAD_SCAN" help:"If true, disables the best-effort security scanner before downloading any files." group:"hardening" default:"false"`
+	AutoloadBackendGalleries bool     `env:"LOCALAI_AUTOLOAD_BACKEND_GALLERIES" help:"If true, automatically loads backend galleries" group:"backends" default:"true"`
+	ModelArgs                []string `arg:"" optional:"" name:"models" help:"Model configuration URLs to load"`
 
 	ModelsCMDFlags `embed:""`
 }
@@ -62,6 +65,11 @@ func (mi *ModelsInstall) Run(ctx *cliContext.Context) error {
 		log.Error().Err(err).Msg("unable to load galleries")
 	}
 
+	var backendGalleries []config.Gallery
+	if err := json.Unmarshal([]byte(mi.BackendGalleries), &backendGalleries); err != nil {
+		log.Error().Err(err).Msg("unable to load backend galleries")
+	}
+
 	for _, modelName := range mi.ModelArgs {
 
 		progressBar := progressbar.NewOptions(
@@ -100,7 +108,7 @@ func (mi *ModelsInstall) Run(ctx *cliContext.Context) error {
 			log.Info().Str("model", modelName).Str("license", model.License).Msg("installing model")
 		}
 
-		err = startup.InstallModels(galleries, mi.ModelsPath, !mi.DisablePredownloadScan, progressCallback, modelName)
+		err = startup.InstallModels(galleries, backendGalleries, mi.ModelsPath, mi.BackendsPath, !mi.DisablePredownloadScan, mi.AutoloadBackendGalleries, progressCallback, modelName)
 		if err != nil {
 			return err
 		}

core/cli/run.go

@@ -30,13 +30,14 @@ type RunCMD struct {
 	LocalaiConfigDir             string        `env:"LOCALAI_CONFIG_DIR" type:"path" default:"${basepath}/configuration" help:"Directory for dynamic loading of certain configuration files (currently api_keys.json and external_backends.json)" group:"storage"`
 	LocalaiConfigDirPollInterval time.Duration `env:"LOCALAI_CONFIG_DIR_POLL_INTERVAL" help:"Typically the config path picks up changes automatically, but if your system has broken fsnotify events, set this to an interval to poll the LocalAI Config Dir (example: 1m)" group:"storage"`
 	// The alias on this option is there to preserve functionality with the old `--config-file` parameter
-	ModelsConfigFile    string   `env:"LOCALAI_MODELS_CONFIG_FILE,CONFIG_FILE" aliases:"config-file" help:"YAML file containing a list of model backend configs" group:"storage"`
-	BackendGalleries    string   `env:"LOCALAI_BACKEND_GALLERIES,BACKEND_GALLERIES" help:"JSON list of backend galleries" group:"backends" default:"${backends}"`
-	Galleries           string   `env:"LOCALAI_GALLERIES,GALLERIES" help:"JSON list of galleries" group:"models" default:"${galleries}"`
-	AutoloadGalleries   bool     `env:"LOCALAI_AUTOLOAD_GALLERIES,AUTOLOAD_GALLERIES" group:"models"`
-	PreloadModels       string   `env:"LOCALAI_PRELOAD_MODELS,PRELOAD_MODELS" help:"A List of models to apply in JSON at start" group:"models"`
-	Models              []string `env:"LOCALAI_MODELS,MODELS" help:"A List of model configuration URLs to load" group:"models"`
-	PreloadModelsConfig string   `env:"LOCALAI_PRELOAD_MODELS_CONFIG,PRELOAD_MODELS_CONFIG" help:"A List of models to apply at startup. Path to a YAML config file" group:"models"`
+	ModelsConfigFile         string   `env:"LOCALAI_MODELS_CONFIG_FILE,CONFIG_FILE" aliases:"config-file" help:"YAML file containing a list of model backend configs" group:"storage"`
+	BackendGalleries         string   `env:"LOCALAI_BACKEND_GALLERIES,BACKEND_GALLERIES" help:"JSON list of backend galleries" group:"backends" default:"${backends}"`
+	Galleries                string   `env:"LOCALAI_GALLERIES,GALLERIES" help:"JSON list of galleries" group:"models" default:"${galleries}"`
+	AutoloadGalleries        bool     `env:"LOCALAI_AUTOLOAD_GALLERIES,AUTOLOAD_GALLERIES" group:"models" default:"true"`
+	AutoloadBackendGalleries bool     `env:"LOCALAI_AUTOLOAD_BACKEND_GALLERIES,AUTOLOAD_BACKEND_GALLERIES" group:"backends" default:"true"`
+	PreloadModels            string   `env:"LOCALAI_PRELOAD_MODELS,PRELOAD_MODELS" help:"A List of models to apply in JSON at start" group:"models"`
+	Models                   []string `env:"LOCALAI_MODELS,MODELS" help:"A List of model configuration URLs to load" group:"models"`
+	PreloadModelsConfig      string   `env:"LOCALAI_PRELOAD_MODELS_CONFIG,PRELOAD_MODELS_CONFIG" help:"A List of models to apply at startup. Path to a YAML config file" group:"models"`
 
 	F16         bool `name:"f16" env:"LOCALAI_F16,F16" help:"Enable GPU acceleration" group:"performance"`
 	Threads     int  `env:"LOCALAI_THREADS,THREADS" short:"t" help:"Number of threads used for parallel computation. Usage of the number of physical cores in the system is suggested" group:"performance"`
@@ -192,6 +193,10 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
 		opts = append(opts, config.EnableGalleriesAutoload)
 	}
 
+	if r.AutoloadBackendGalleries {
+		opts = append(opts, config.EnableBackendGalleriesAutoload)
+	}
+
 	if r.PreloadBackendOnly {
 		_, err := application.New(opts...)
 		return err

core/config/application_config.go

@@ -55,7 +55,7 @@ type ApplicationConfig struct {
 
 	ExternalGRPCBackends map[string]string
 
-	AutoloadGalleries bool
+	AutoloadGalleries, AutoloadBackendGalleries bool
 
 	SingleBackend           bool
 	ParallelBackendRequests bool
@@ -192,6 +192,10 @@ var EnableGalleriesAutoload = func(o *ApplicationConfig) {
 	o.AutoloadGalleries = true
 }
 
+var EnableBackendGalleriesAutoload = func(o *ApplicationConfig) {
+	o.AutoloadBackendGalleries = true
+}
+
 func WithExternalBackend(name string, uri string) AppOption {
 	return func(o *ApplicationConfig) {
 		if o.ExternalGRPCBackends == nil {

core/gallery/backends.go

@@ -71,7 +71,22 @@ func findBestBackendFromMeta(backend *GalleryBackend, systemState *system.System
 }
 
 // Installs a model from the gallery
-func InstallBackendFromGallery(galleries []config.Gallery, systemState *system.SystemState, 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), force bool) error {
+	if !force {
+		// check if we already have the backend installed
+		backends, err := ListSystemBackends(basePath)
+		if err != nil {
+			return err
+		}
+		if _, ok := backends[name]; ok {
+			return nil
+		}
+	}
+
+	if name == "" {
+		return fmt.Errorf("backend name is empty")
+	}
+
 	log.Debug().Interface("galleries", galleries).Str("name", name).Msg("Installing backend from gallery")
 
 	backends, err := AvailableBackends(galleries, basePath)

core/gallery/backends_test.go

@@ -42,13 +42,13 @@ var _ = Describe("Gallery Backends", func() {
 
 	Describe("InstallBackendFromGallery", func() {
 		It("should return error when backend is not found", func() {
-			err := InstallBackendFromGallery(galleries, nil, "non-existent", tempDir, nil)
+			err := InstallBackendFromGallery(galleries, nil, "non-existent", tempDir, nil, true)
 			Expect(err).To(HaveOccurred())
 			Expect(err.Error()).To(ContainSubstring("no backend found with name \"non-existent\""))
 		})
 
 		It("should install backend from gallery", func() {
-			err := InstallBackendFromGallery(galleries, nil, "test-backend", tempDir, nil)
+			err := InstallBackendFromGallery(galleries, nil, "test-backend", tempDir, nil, true)
 			Expect(err).ToNot(HaveOccurred())
 			Expect(filepath.Join(tempDir, "test-backend", "run.sh")).To(BeARegularFile())
 		})
@@ -181,7 +181,7 @@ var _ = Describe("Gallery Backends", func() {
 
 			// Test with NVIDIA system state
 			nvidiaSystemState := &system.SystemState{GPUVendor: "nvidia"}
-			err = InstallBackendFromGallery([]config.Gallery{gallery}, nvidiaSystemState, "meta-backend", tempDir, nil)
+			err = InstallBackendFromGallery([]config.Gallery{gallery}, nvidiaSystemState, "meta-backend", tempDir, nil, true)
 			Expect(err).NotTo(HaveOccurred())
 
 			metaBackendPath := filepath.Join(tempDir, "meta-backend")

core/gallery/models.go

@@ -10,6 +10,7 @@ import (
 	"dario.cat/mergo"
 	"github.com/mudler/LocalAI/core/config"
 	lconfig "github.com/mudler/LocalAI/core/config"
+	"github.com/mudler/LocalAI/core/system"
 	"github.com/mudler/LocalAI/pkg/downloader"
 	"github.com/mudler/LocalAI/pkg/utils"
 
@@ -69,7 +70,9 @@ type PromptTemplate struct {
 }
 
 // Installs a model from the gallery
-func InstallModelFromGallery(galleries []config.Gallery, name string, basePath string, req GalleryModel, downloadStatus func(string, string, string, float64), enforceScan bool) error {
+func InstallModelFromGallery(
+	modelGalleries, backendGalleries []config.Gallery,
+	name string, basePath, backendBasePath string, req GalleryModel, downloadStatus func(string, string, string, float64), enforceScan, automaticallyInstallBackend bool) error {
 
 	applyModel := func(model *GalleryModel) error {
 		name = strings.ReplaceAll(name, string(os.PathSeparator), "__")
@@ -119,14 +122,26 @@ func InstallModelFromGallery(galleries []config.Gallery, name string, basePath s
 			return err
 		}
 
-		if err := InstallModel(basePath, installName, &config, model.Overrides, downloadStatus, enforceScan); err != nil {
+		installedModel, err := InstallModel(basePath, installName, &config, model.Overrides, downloadStatus, enforceScan)
+		if err != nil {
 			return err
 		}
 
+		if automaticallyInstallBackend && installedModel.Backend != "" {
+			systemState, err := system.GetSystemState()
+			if err != nil {
+				return err
+			}
+
+			if err := InstallBackendFromGallery(backendGalleries, systemState, installedModel.Backend, backendBasePath, downloadStatus, false); err != nil {
+				return err
+			}
+		}
+
 		return nil
 	}
 
-	models, err := AvailableGalleryModels(galleries, basePath)
+	models, err := AvailableGalleryModels(modelGalleries, basePath)
 	if err != nil {
 		return err
 	}
@@ -139,11 +154,11 @@ func InstallModelFromGallery(galleries []config.Gallery, name string, basePath s
 	return applyModel(model)
 }
 
-func InstallModel(basePath, nameOverride string, config *ModelConfig, configOverrides map[string]interface{}, downloadStatus func(string, string, string, float64), enforceScan bool) error {
+func InstallModel(basePath, nameOverride string, config *ModelConfig, configOverrides map[string]interface{}, downloadStatus func(string, string, string, float64), enforceScan bool) (*lconfig.BackendConfig, error) {
 	// Create base path if it doesn't exist
 	err := os.MkdirAll(basePath, 0750)
 	if err != nil {
-		return fmt.Errorf("failed to create base path: %v", err)
+		return nil, fmt.Errorf("failed to create base path: %v", err)
 	}
 
 	if len(configOverrides) > 0 {
@@ -155,7 +170,7 @@ func InstallModel(basePath, nameOverride string, config *ModelConfig, configOver
 		log.Debug().Msgf("Checking %q exists and matches SHA", file.Filename)
 
 		if err := utils.VerifyPath(file.Filename, basePath); err != nil {
-			return err
+			return nil, err
 		}
 
 		// Create file path
@@ -165,32 +180,32 @@ func InstallModel(basePath, nameOverride string, config *ModelConfig, configOver
 			scanResults, err := downloader.HuggingFaceScan(downloader.URI(file.URI))
 			if err != nil && errors.Is(err, downloader.ErrUnsafeFilesFound) {
 				log.Error().Str("model", config.Name).Strs("clamAV", scanResults.ClamAVInfectedFiles).Strs("pickles", scanResults.DangerousPickles).Msg("Contains unsafe file(s)!")
-				return err
+				return nil, err
 			}
 		}
 		uri := downloader.URI(file.URI)
 		if err := uri.DownloadFile(filePath, file.SHA256, i, len(config.Files), downloadStatus); err != nil {
-			return err
+			return nil, err
 		}
 	}
 
 	// Write prompt template contents to separate files
 	for _, template := range config.PromptTemplates {
 		if err := utils.VerifyPath(template.Name+".tmpl", basePath); err != nil {
-			return err
+			return nil, err
 		}
 		// Create file path
 		filePath := filepath.Join(basePath, template.Name+".tmpl")
 
 		// Create parent directory
 		err := os.MkdirAll(filepath.Dir(filePath), 0750)
 		if err != nil {
-			return fmt.Errorf("failed to create parent directory for prompt template %q: %v", template.Name, err)
+			return nil, fmt.Errorf("failed to create parent directory for prompt template %q: %v", template.Name, err)
 		}
 		// Create and write file content
 		err = os.WriteFile(filePath, []byte(template.Content), 0600)
 		if err != nil {
-			return fmt.Errorf("failed to write prompt template %q: %v", template.Name, err)
+			return nil, fmt.Errorf("failed to write prompt template %q: %v", template.Name, err)
 		}
 
 		log.Debug().Msgf("Prompt template %q written", template.Name)
@@ -202,9 +217,11 @@ func InstallModel(basePath, nameOverride string, config *ModelConfig, configOver
 	}
 
 	if err := utils.VerifyPath(name+".yaml", basePath); err != nil {
-		return err
+		return nil, err
 	}
 
+	backendConfig := lconfig.BackendConfig{}
+
 	// write config file
 	if len(configOverrides) != 0 || len(config.ConfigFile) != 0 {
 		configFilePath := filepath.Join(basePath, name+".yaml")
@@ -213,33 +230,33 @@ func InstallModel(basePath, nameOverride string, config *ModelConfig, configOver
 		configMap := make(map[string]interface{})
 		err = yaml.Unmarshal([]byte(config.ConfigFile), &configMap)
 		if err != nil {
-			return fmt.Errorf("failed to unmarshal config YAML: %v", err)
+			return nil, fmt.Errorf("failed to unmarshal config YAML: %v", err)
 		}
 
 		configMap["name"] = name
 
 		if err := mergo.Merge(&configMap, configOverrides, mergo.WithOverride); err != nil {
-			return err
+			return nil, err
 		}
 
 		// Write updated config file
 		updatedConfigYAML, err := yaml.Marshal(configMap)
 		if err != nil {
-			return fmt.Errorf("failed to marshal updated config YAML: %v", err)
+			return nil, fmt.Errorf("failed to marshal updated config YAML: %v", err)
 		}
 
-		backendConfig := lconfig.BackendConfig{}
 		err = yaml.Unmarshal(updatedConfigYAML, &backendConfig)
 		if err != nil {
-			return fmt.Errorf("failed to unmarshal updated config YAML: %v", err)
+			return nil, fmt.Errorf("failed to unmarshal updated config YAML: %v", err)
 		}
+
 		if !backendConfig.Validate() {
-			return fmt.Errorf("failed to validate updated config YAML")
+			return nil, fmt.Errorf("failed to validate updated config YAML")
 		}
 
 		err = os.WriteFile(configFilePath, updatedConfigYAML, 0600)
 		if err != nil {
-			return fmt.Errorf("failed to write updated config file: %v", err)
+			return nil, fmt.Errorf("failed to write updated config file: %v", err)
 		}
 
 		log.Debug().Msgf("Written config file %s", configFilePath)
@@ -249,14 +266,12 @@ func InstallModel(basePath, nameOverride string, config *ModelConfig, configOver
 	modelFile := filepath.Join(basePath, galleryFileName(name))
 	data, err := yaml.Marshal(config)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	log.Debug().Msgf("Written gallery file %s", modelFile)
 
-	return os.WriteFile(modelFile, data, 0600)
-
-	//return nil
+	return &backendConfig, os.WriteFile(modelFile, data, 0600)
 }
 
 func galleryFileName(name string) string {