diff --git a/core/gallery/upgrade.go b/core/gallery/upgrade.go new file mode 100644 index 000000000..aec75852c --- /dev/null +++ b/core/gallery/upgrade.go @@ -0,0 +1,223 @@ +package gallery + +import ( + "context" + "fmt" + "os" + "path/filepath" + "time" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/downloader" + "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/oci" + "github.com/mudler/LocalAI/pkg/system" + "github.com/mudler/xlog" + cp "github.com/otiai10/copy" +) + +// UpgradeInfo holds details about an available backend upgrade. +type UpgradeInfo struct { + BackendName string `json:"backend_name"` + InstalledVersion string `json:"installed_version"` + AvailableVersion string `json:"available_version"` + InstalledDigest string `json:"installed_digest,omitempty"` + AvailableDigest string `json:"available_digest,omitempty"` +} + +// CheckBackendUpgrades compares installed backends against gallery entries +// and returns a map of backend names to UpgradeInfo for those that have +// newer versions or different OCI digests available. +func CheckBackendUpgrades(ctx context.Context, galleries []config.Gallery, systemState *system.SystemState) (map[string]UpgradeInfo, error) { + galleryBackends, err := AvailableBackends(galleries, systemState) + if err != nil { + return nil, fmt.Errorf("failed to list available backends: %w", err) + } + + installedBackends, err := ListSystemBackends(systemState) + if err != nil { + return nil, fmt.Errorf("failed to list installed backends: %w", err) + } + + result := make(map[string]UpgradeInfo) + + for _, installed := range installedBackends { + // Skip system backends — they are managed outside the gallery + if installed.IsSystem { + continue + } + if installed.Metadata == nil { + continue + } + + // Find matching gallery entry by metadata name + galleryEntry := FindGalleryElement(galleryBackends, installed.Metadata.Name) + if galleryEntry == nil { + continue + } + + installedVersion := installed.Metadata.Version + galleryVersion := galleryEntry.Version + + // If both sides have versions, compare them + if galleryVersion != "" && installedVersion != "" { + if galleryVersion != installedVersion { + result[installed.Metadata.Name] = UpgradeInfo{ + BackendName: installed.Metadata.Name, + InstalledVersion: installedVersion, + AvailableVersion: galleryVersion, + } + } + // Versions match — no upgrade needed + continue + } + + // If either version is empty, fall back to OCI digest comparison + if installed.Metadata.Digest != "" && downloader.URI(galleryEntry.URI).LooksLikeOCI() { + remoteDigest, err := oci.GetImageDigest(galleryEntry.URI, "", nil, nil) + if err != nil { + xlog.Warn("Failed to get remote OCI digest for upgrade check", "backend", installed.Metadata.Name, "error", err) + continue + } + if remoteDigest != installed.Metadata.Digest { + result[installed.Metadata.Name] = UpgradeInfo{ + BackendName: installed.Metadata.Name, + InstalledDigest: installed.Metadata.Digest, + AvailableDigest: remoteDigest, + } + } + } + // No version info and no digest to compare — skip + } + + return result, nil +} + +// UpgradeBackend upgrades a single backend to the latest gallery version using +// an atomic swap with backup-based rollback on failure. +func UpgradeBackend(ctx context.Context, systemState *system.SystemState, modelLoader *model.ModelLoader, galleries []config.Gallery, backendName string, downloadStatus func(string, string, string, float64)) error { + // Look up the installed backend + installedBackends, err := ListSystemBackends(systemState) + if err != nil { + return fmt.Errorf("failed to list installed backends: %w", err) + } + + installed, ok := installedBackends.Get(backendName) + if !ok { + return fmt.Errorf("backend %q: %w", backendName, ErrBackendNotFound) + } + + if installed.IsSystem { + return fmt.Errorf("system backend %q cannot be upgraded via gallery", backendName) + } + + // If this is a meta backend, recursively upgrade the concrete backend it points to + if installed.Metadata != nil && installed.Metadata.MetaBackendFor != "" { + xlog.Info("Meta backend detected, upgrading concrete backend", "meta", backendName, "concrete", installed.Metadata.MetaBackendFor) + return UpgradeBackend(ctx, systemState, modelLoader, galleries, installed.Metadata.MetaBackendFor, downloadStatus) + } + + // Find the gallery entry + galleryBackends, err := AvailableBackends(galleries, systemState) + if err != nil { + return fmt.Errorf("failed to list available backends: %w", err) + } + + galleryEntry := FindGalleryElement(galleryBackends, backendName) + if galleryEntry == nil { + return fmt.Errorf("no gallery entry found for backend %q", backendName) + } + + backendPath := filepath.Join(systemState.Backend.BackendsPath, backendName) + tmpPath := backendPath + ".upgrade-tmp" + backupPath := backendPath + ".backup" + + // Clean up any stale tmp/backup dirs from prior attempts + os.RemoveAll(tmpPath) + os.RemoveAll(backupPath) + + // Step 1: Download the new backend into the tmp directory + if err := os.MkdirAll(tmpPath, 0750); err != nil { + return fmt.Errorf("failed to create upgrade tmp dir: %w", err) + } + + uri := downloader.URI(galleryEntry.URI) + if uri.LooksLikeDir() { + if err := cp.Copy(string(uri), tmpPath); err != nil { + os.RemoveAll(tmpPath) + return fmt.Errorf("failed to copy backend from directory: %w", err) + } + } else { + if err := uri.DownloadFileWithContext(ctx, tmpPath, "", 1, 1, downloadStatus); err != nil { + os.RemoveAll(tmpPath) + return fmt.Errorf("failed to download backend: %w", err) + } + } + + // Step 2: Validate — check that run.sh exists in the new content + newRunFile := filepath.Join(tmpPath, runFile) + if _, err := os.Stat(newRunFile); os.IsNotExist(err) { + os.RemoveAll(tmpPath) + return fmt.Errorf("upgrade validation failed: run.sh not found in new backend") + } + + // Step 3: Atomic swap — rename current to backup, then tmp to current + if err := os.Rename(backendPath, backupPath); err != nil { + os.RemoveAll(tmpPath) + return fmt.Errorf("failed to move current backend to backup: %w", err) + } + + if err := os.Rename(tmpPath, backendPath); err != nil { + // Restore backup on failure + xlog.Error("Failed to move new backend into place, restoring backup", "error", err) + if restoreErr := os.Rename(backupPath, backendPath); restoreErr != nil { + xlog.Error("Failed to restore backup", "error", restoreErr) + } + os.RemoveAll(tmpPath) + return fmt.Errorf("failed to move new backend into place: %w", err) + } + + // Step 4: Write updated metadata, preserving alias from old metadata + var oldAlias string + if installed.Metadata != nil { + oldAlias = installed.Metadata.Alias + } + + newMetadata := &BackendMetadata{ + Name: backendName, + Version: galleryEntry.Version, + URI: galleryEntry.URI, + InstalledAt: time.Now().Format(time.RFC3339), + Alias: oldAlias, + } + + if galleryEntry.Gallery.URL != "" { + newMetadata.GalleryURL = galleryEntry.Gallery.URL + } + + // Record OCI digest if applicable (non-fatal on failure) + if uri.LooksLikeOCI() { + digest, digestErr := oci.GetImageDigest(galleryEntry.URI, "", nil, nil) + if digestErr != nil { + xlog.Warn("Failed to get OCI image digest after upgrade", "uri", galleryEntry.URI, "error", digestErr) + } else { + newMetadata.Digest = digest + } + } + + if err := writeBackendMetadata(backendPath, newMetadata); err != nil { + // Metadata write failure is not worth rolling back the entire upgrade + xlog.Error("Failed to write metadata after upgrade", "error", err) + } + + // Step 5: Re-register backends so the model loader picks up any changes + if err := RegisterBackends(systemState, modelLoader); err != nil { + xlog.Warn("Failed to re-register backends after upgrade", "error", err) + } + + // Step 6: Remove backup + os.RemoveAll(backupPath) + + xlog.Info("Backend upgraded successfully", "backend", backendName, "version", galleryEntry.Version) + return nil +} diff --git a/core/gallery/upgrade_test.go b/core/gallery/upgrade_test.go new file mode 100644 index 000000000..f65b4276b --- /dev/null +++ b/core/gallery/upgrade_test.go @@ -0,0 +1,219 @@ +package gallery_test + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "time" + + "github.com/mudler/LocalAI/core/config" + . "github.com/mudler/LocalAI/core/gallery" + "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/system" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gopkg.in/yaml.v3" +) + +var _ = Describe("Upgrade Detection and Execution", func() { + var ( + tempDir string + backendsPath string + galleryPath string + systemState *system.SystemState + galleries []config.Gallery + ) + + // installBackendWithVersion creates a fake installed backend directory with + // the given name, version, and optional run.sh content. + installBackendWithVersion := func(name, version string, runContent ...string) { + dir := filepath.Join(backendsPath, name) + Expect(os.MkdirAll(dir, 0750)).To(Succeed()) + + content := "#!/bin/sh\necho ok" + if len(runContent) > 0 { + content = runContent[0] + } + Expect(os.WriteFile(filepath.Join(dir, "run.sh"), []byte(content), 0755)).To(Succeed()) + + metadata := BackendMetadata{ + Name: name, + Version: version, + InstalledAt: time.Now().Format(time.RFC3339), + } + data, err := json.MarshalIndent(metadata, "", " ") + Expect(err).NotTo(HaveOccurred()) + Expect(os.WriteFile(filepath.Join(dir, "metadata.json"), data, 0644)).To(Succeed()) + } + + // writeGalleryYAML writes a gallery YAML file with the given backends. + writeGalleryYAML := func(backends []GalleryBackend) { + data, err := yaml.Marshal(backends) + Expect(err).NotTo(HaveOccurred()) + Expect(os.WriteFile(galleryPath, data, 0644)).To(Succeed()) + } + + BeforeEach(func() { + var err error + tempDir, err = os.MkdirTemp("", "upgrade-test-*") + Expect(err).NotTo(HaveOccurred()) + + backendsPath = tempDir + + galleryPath = filepath.Join(tempDir, "gallery.yaml") + + // Write a default empty gallery + writeGalleryYAML([]GalleryBackend{}) + + galleries = []config.Gallery{ + { + Name: "test-gallery", + URL: "file://" + galleryPath, + }, + } + + systemState, err = system.GetSystemState( + system.WithBackendPath(backendsPath), + ) + Expect(err).NotTo(HaveOccurred()) + }) + + AfterEach(func() { + os.RemoveAll(tempDir) + }) + + Describe("CheckBackendUpgrades", func() { + It("should detect upgrade when gallery version differs from installed version", func() { + // Install a backend at v1.0.0 + installBackendWithVersion("my-backend", "1.0.0") + + // Gallery advertises v2.0.0 + writeGalleryYAML([]GalleryBackend{ + { + Metadata: Metadata{ + Name: "my-backend", + }, + URI: filepath.Join(tempDir, "some-source"), + Version: "2.0.0", + }, + }) + + upgrades, err := CheckBackendUpgrades(context.Background(), galleries, systemState) + Expect(err).NotTo(HaveOccurred()) + Expect(upgrades).To(HaveKey("my-backend")) + Expect(upgrades["my-backend"].InstalledVersion).To(Equal("1.0.0")) + Expect(upgrades["my-backend"].AvailableVersion).To(Equal("2.0.0")) + }) + + It("should NOT flag upgrade when versions match", func() { + installBackendWithVersion("my-backend", "2.0.0") + + writeGalleryYAML([]GalleryBackend{ + { + Metadata: Metadata{ + Name: "my-backend", + }, + URI: filepath.Join(tempDir, "some-source"), + Version: "2.0.0", + }, + }) + + upgrades, err := CheckBackendUpgrades(context.Background(), galleries, systemState) + Expect(err).NotTo(HaveOccurred()) + Expect(upgrades).To(BeEmpty()) + }) + + It("should skip backends without version info and without OCI digest", func() { + // Install without version + installBackendWithVersion("my-backend", "") + + // Gallery also without version + writeGalleryYAML([]GalleryBackend{ + { + Metadata: Metadata{ + Name: "my-backend", + }, + URI: filepath.Join(tempDir, "some-source"), + }, + }) + + upgrades, err := CheckBackendUpgrades(context.Background(), galleries, systemState) + Expect(err).NotTo(HaveOccurred()) + Expect(upgrades).To(BeEmpty()) + }) + }) + + Describe("UpgradeBackend", func() { + It("should replace backend directory and update metadata", func() { + // Install v1 + installBackendWithVersion("my-backend", "1.0.0", "#!/bin/sh\necho v1") + + // Create a source directory with v2 content + srcDir := filepath.Join(tempDir, "v2-source") + Expect(os.MkdirAll(srcDir, 0750)).To(Succeed()) + Expect(os.WriteFile(filepath.Join(srcDir, "run.sh"), []byte("#!/bin/sh\necho v2"), 0755)).To(Succeed()) + + // Gallery points to the v2 source dir + writeGalleryYAML([]GalleryBackend{ + { + Metadata: Metadata{ + Name: "my-backend", + }, + URI: srcDir, + Version: "2.0.0", + }, + }) + + ml := model.NewModelLoader(systemState) + err := UpgradeBackend(context.Background(), systemState, ml, galleries, "my-backend", nil) + Expect(err).NotTo(HaveOccurred()) + + // Verify run.sh was updated + content, err := os.ReadFile(filepath.Join(backendsPath, "my-backend", "run.sh")) + Expect(err).NotTo(HaveOccurred()) + Expect(string(content)).To(Equal("#!/bin/sh\necho v2")) + + // Verify metadata was updated + metaData, err := os.ReadFile(filepath.Join(backendsPath, "my-backend", "metadata.json")) + Expect(err).NotTo(HaveOccurred()) + var meta BackendMetadata + Expect(json.Unmarshal(metaData, &meta)).To(Succeed()) + Expect(meta.Version).To(Equal("2.0.0")) + Expect(meta.Name).To(Equal("my-backend")) + }) + + It("should restore backup on failure", func() { + // Install v1 + installBackendWithVersion("my-backend", "1.0.0", "#!/bin/sh\necho v1") + + // Gallery points to a nonexistent path (no run.sh will be found) + nonExistentDir := filepath.Join(tempDir, "does-not-exist") + writeGalleryYAML([]GalleryBackend{ + { + Metadata: Metadata{ + Name: "my-backend", + }, + URI: nonExistentDir, + Version: "2.0.0", + }, + }) + + ml := model.NewModelLoader(systemState) + err := UpgradeBackend(context.Background(), systemState, ml, galleries, "my-backend", nil) + Expect(err).To(HaveOccurred()) + + // Verify v1 is still intact + content, err := os.ReadFile(filepath.Join(backendsPath, "my-backend", "run.sh")) + Expect(err).NotTo(HaveOccurred()) + Expect(string(content)).To(Equal("#!/bin/sh\necho v1")) + + // Verify metadata still says v1 + metaData, err := os.ReadFile(filepath.Join(backendsPath, "my-backend", "metadata.json")) + Expect(err).NotTo(HaveOccurred()) + var meta BackendMetadata + Expect(json.Unmarshal(metaData, &meta)).To(Succeed()) + Expect(meta.Version).To(Equal("1.0.0")) + }) + }) +})