mirror of
https://github.com/mudler/LocalAI.git
synced 2026-04-16 12:59:33 -04:00
feat: add backend upgrade detection and execution logic
Add CheckBackendUpgrades() to compare installed backend versions/digests against gallery entries, and UpgradeBackend() to perform atomic upgrades with backup-based rollback on failure. Includes Agent A's data model changes (Version/URI/Digest fields, GetImageDigest).
This commit is contained in:
Ettore Di Giacinto
223
core/gallery/upgrade.go
Normal file
223
core/gallery/upgrade.go
Normal file
@@ -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
|
||||
}
|
||||
219
core/gallery/upgrade_test.go
Normal file
219
core/gallery/upgrade_test.go
Normal file
@@ -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"))
|
||||
})
|
||||
})
|
||||
})
|
||||
Reference in New Issue
Block a user