LocalAI/core/startup/model_preload.go

package startup

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"slices"
	"time"

	"github.com/google/uuid"
	"github.com/mudler/LocalAI/core/config"
	"github.com/mudler/LocalAI/core/gallery"
	"github.com/mudler/LocalAI/core/gallery/importers"
	"github.com/mudler/LocalAI/core/services/galleryop"
	"github.com/mudler/LocalAI/pkg/model"
	"github.com/mudler/LocalAI/pkg/system"
	"github.com/mudler/LocalAI/pkg/utils"
	"github.com/mudler/xlog"
)

// InstallModels will preload models from the given list of URLs and galleries
// It will download the model if it is not already present in the model path
// It will also try to resolve if the model is an embedded model YAML configuration
func InstallModels(ctx context.Context, galleryService *galleryop.GalleryService, galleries, backendGalleries []config.Gallery, systemState *system.SystemState, modelLoader *model.ModelLoader, enforceScan, autoloadBackendGalleries, requireBackendIntegrity bool, downloadStatus func(string, string, string, float64), models ...string) error {
	return InstallModelsWithOptions(ctx, galleryService, galleries, backendGalleries, systemState, modelLoader, enforceScan, autoloadBackendGalleries, requireBackendIntegrity, downloadStatus, nil, models...)
}

// InstallModelsWithOptions is InstallModels with extra install options, which
// is how the CLI passes a variant pin. It is a separate entry point rather than
// a variadic tail on InstallModels because that signature already ends in a
// variadic model list.
func InstallModelsWithOptions(ctx context.Context, galleryService *galleryop.GalleryService, galleries, backendGalleries []config.Gallery, systemState *system.SystemState, modelLoader *model.ModelLoader, enforceScan, autoloadBackendGalleries, requireBackendIntegrity bool, downloadStatus func(string, string, string, float64), extraOptions []gallery.InstallOption, models ...string) error {
	// create an error that groups all errors
	var err error
	installOptions := slices.Clone(extraOptions)
	if galleryService != nil {
		installOptions = append(installOptions, gallery.WithArtifactMaterializer(galleryService.ModelArtifactMaterializer()))
	}
	for _, url := range models {
		// Check if it's a model gallery, or print a warning
		e, found := installModel(ctx, galleries, backendGalleries, url, systemState, modelLoader, downloadStatus, enforceScan, autoloadBackendGalleries, requireBackendIntegrity, installOptions...)
		if e != nil && found {
			xlog.Error("[startup] failed installing model", "error", err, "model", url)
			err = errors.Join(err, e)
		} else if !found {
			xlog.Debug("[startup] model not found in the gallery", "model", url)

			if galleryService == nil {
				return fmt.Errorf("cannot start autoimporter, not sure how to handle this uri")
			}

			// TODO: we should just use the discoverModelConfig here and default to this.
			modelConfig, discoverErr := importers.DiscoverModelConfig(url, json.RawMessage{})
			if discoverErr != nil {
				xlog.Error("[startup] failed to discover model config", "error", discoverErr, "model", url)
				err = errors.Join(discoverErr, fmt.Errorf("failed to discover model config: %w", err))
				continue
			}

			uuid, uuidErr := uuid.NewUUID()
			if uuidErr != nil {
				err = errors.Join(uuidErr, fmt.Errorf("failed to generate UUID: %w", uuidErr))
				continue
			}

			galleryService.ModelGalleryChannel <- galleryop.ManagementOp[gallery.GalleryModel, gallery.ModelConfig]{
				Req: gallery.GalleryModel{
					Overrides: map[string]any{},
				},
				ID:                 uuid.String(),
				GalleryElementName: modelConfig.Name,
				GalleryElement:     &modelConfig,
				BackendGalleries:   backendGalleries,
			}

			var status *galleryop.OpStatus
			poll := time.NewTicker(50 * time.Millisecond)
			defer poll.Stop()
			for {
				status = galleryService.GetStatus(uuid.String())
				if status != nil && status.Processed {
					break
				}
				select {
				case <-ctx.Done():
					return ctx.Err()
				case <-poll.C:
				}
			}

			if status.Error != nil {
				xlog.Error("[startup] failed to import model", "error", status.Error, "model", modelConfig.Name, "url", url)
				return status.Error
			}

			xlog.Info("[startup] imported model", "model", modelConfig.Name, "url", url)
		}
	}
	return err
}

func installModel(ctx context.Context, galleries, backendGalleries []config.Gallery, modelName string, systemState *system.SystemState, modelLoader *model.ModelLoader, downloadStatus func(string, string, string, float64), enforceScan, autoloadBackendGalleries, requireBackendIntegrity bool, options ...gallery.InstallOption) (error, bool) {
	models, err := gallery.AvailableGalleryModels(galleries, systemState)
	if err != nil {
		return err, false
	}

	model := gallery.FindGalleryElement(models, modelName)
	if model == nil {
		return err, false
	}

	if downloadStatus == nil {
		downloadStatus = utils.DisplayDownloadFunction
	}

	xlog.Info("installing model", "model", modelName, "license", model.License)
	err = gallery.InstallModelFromGallery(ctx, galleries, backendGalleries, systemState, modelLoader, modelName, gallery.GalleryModel{}, downloadStatus, enforceScan, autoloadBackendGalleries, requireBackendIntegrity, options...)
	if err != nil {
		return err, true
	}

	return nil, true
}