diff --git a/core/application/upgrade_checker.go b/core/application/upgrade_checker.go index e18548a4e1b2..74fcf500267d 100644 --- a/core/application/upgrade_checker.go +++ b/core/application/upgrade_checker.go @@ -2,6 +2,7 @@ package application import ( "context" + "errors" "sync" "time" @@ -222,9 +223,14 @@ func (uc *UpgradeChecker) runCheck(ctx context.Context) { }, nil) } else { err = gallery.UpgradeBackend(ctx, uc.systemState, uc.modelLoader, - uc.galleries, name, nil, uc.appConfig.RequireBackendIntegrity) + uc.galleries, name, nil, uc.appConfig.RequireBackendIntegrity, + gallery.WithSkipIfBackendBusy()) } - if err != nil { + if errors.Is(err, gallery.ErrBackendOperationInProgress) { + xlog.Debug("Skipping auto-upgrade because backend is already being upgraded", + "backend", name) + continue + } else if err != nil { xlog.Error("Failed to auto-upgrade backend", "backend", name, "error", err) } else { diff --git a/core/gallery/backend_operation.go b/core/gallery/backend_operation.go new file mode 100644 index 000000000000..4d0ceea005bd --- /dev/null +++ b/core/gallery/backend_operation.go @@ -0,0 +1,106 @@ +package gallery + +import ( + "context" + "errors" + "path/filepath" + "sync" +) + +// ErrBackendOperationInProgress is returned when a non-blocking backend +// upgrade finds another upgrade already using the same backend path. +var ErrBackendOperationInProgress = errors.New("backend operation already in progress") + +// UpgradeOption customizes UpgradeBackend behavior. +type UpgradeOption func(*upgradeOptions) + +type upgradeOptions struct { + skipIfBusy bool +} + +// WithSkipIfBackendBusy makes UpgradeBackend return +// ErrBackendOperationInProgress instead of waiting for an operation already +// upgrading the same resolved backend path. Background auto-upgrades use this +// so an explicit upgrade remains the owner of the in-flight work. +func WithSkipIfBackendBusy() UpgradeOption { + return func(options *upgradeOptions) { + options.skipIfBusy = true + } +} + +type backendOperationEntry struct { + token chan struct{} + refs int +} + +type backendOperationCoordinator struct { + mu sync.Mutex + entries map[string]*backendOperationEntry +} + +func newBackendOperationCoordinator() *backendOperationCoordinator { + return &backendOperationCoordinator{ + entries: make(map[string]*backendOperationEntry), + } +} + +var backendOperations = newBackendOperationCoordinator() + +// acquire serializes upgrades by their normalized backend path. refs counts +// both the holder and waiters so an entry cannot be removed and replaced while +// a waiter still refers to it. +func (c *backendOperationCoordinator) acquire(ctx context.Context, path string, skipIfBusy bool) (func(), error) { + key := filepath.Clean(path) + + c.mu.Lock() + entry := c.entries[key] + if entry == nil { + entry = &backendOperationEntry{token: make(chan struct{}, 1)} + entry.token <- struct{}{} + c.entries[key] = entry + } + entry.refs++ + c.mu.Unlock() + + acquired := false + if skipIfBusy { + select { + case <-entry.token: + acquired = true + default: + } + } else { + select { + case <-entry.token: + acquired = true + case <-ctx.Done(): + } + } + + if !acquired { + c.dropReference(key, entry) + if skipIfBusy { + return nil, ErrBackendOperationInProgress + } + return nil, ctx.Err() + } + + var once sync.Once + release := func() { + once.Do(func() { + entry.token <- struct{}{} + c.dropReference(key, entry) + }) + } + return release, nil +} + +func (c *backendOperationCoordinator) dropReference(key string, entry *backendOperationEntry) { + c.mu.Lock() + defer c.mu.Unlock() + + entry.refs-- + if entry.refs == 0 && c.entries[key] == entry { + delete(c.entries, key) + } +} diff --git a/core/gallery/backend_operation_internal_test.go b/core/gallery/backend_operation_internal_test.go new file mode 100644 index 000000000000..00a7de41ddcf --- /dev/null +++ b/core/gallery/backend_operation_internal_test.go @@ -0,0 +1,178 @@ +package gallery + +import ( + "context" + "errors" + "os" + "path/filepath" + "testing" + "time" + + "github.com/mudler/LocalAI/pkg/system" +) + +func waitForBackendOperationReferences(t *testing.T, coordinator *backendOperationCoordinator, path string, want int) { + t.Helper() + key := filepath.Clean(path) + deadline := time.Now().Add(time.Second) + for { + coordinator.mu.Lock() + entry := coordinator.entries[key] + got := 0 + if entry != nil { + got = entry.refs + } + coordinator.mu.Unlock() + if got == want { + return + } + if time.Now().After(deadline) { + t.Fatalf("backend operation references = %d, want %d", got, want) + } + time.Sleep(time.Millisecond) + } +} + +func TestBackendOperationCoordinatorSerializesByPath(t *testing.T) { + coordinator := newBackendOperationCoordinator() + ctx := context.Background() + path := filepath.Join(t.TempDir(), "backend") + + releaseFirst, err := coordinator.acquire(ctx, path, false) + if err != nil { + t.Fatalf("acquire first operation: %v", err) + } + + // A different backend remains independent. + releaseOther, err := coordinator.acquire(ctx, path+"-other", true) + if err != nil { + t.Fatalf("acquire different backend: %v", err) + } + releaseOther() + + // A background operation on the same backend backs off immediately. + if _, err := coordinator.acquire(ctx, path, true); !errors.Is(err, ErrBackendOperationInProgress) { + t.Fatalf("non-blocking acquire error = %v, want ErrBackendOperationInProgress", err) + } + + acquired := make(chan error, 1) + go func() { + release, err := coordinator.acquire(ctx, path, false) + if err == nil { + release() + } + acquired <- err + }() + + waitForBackendOperationReferences(t, coordinator, path, 2) + select { + case err := <-acquired: + t.Fatalf("second operation completed before release: %v", err) + default: + } + + releaseFirst() + select { + case err := <-acquired: + if err != nil { + t.Fatalf("second operation failed after release: %v", err) + } + case <-time.After(time.Second): + t.Fatal("second operation did not acquire after release") + } + + coordinator.mu.Lock() + defer coordinator.mu.Unlock() + if len(coordinator.entries) != 0 { + t.Fatalf("coordinator retained %d idle entries", len(coordinator.entries)) + } +} + +func TestBackendOperationCoordinatorDropsCanceledWaiter(t *testing.T) { + coordinator := newBackendOperationCoordinator() + path := filepath.Join(t.TempDir(), "backend") + release, err := coordinator.acquire(context.Background(), path, false) + if err != nil { + t.Fatalf("acquire first operation: %v", err) + } + + canceled, cancel := context.WithCancel(context.Background()) + cancel() + if _, err := coordinator.acquire(canceled, path, false); !errors.Is(err, context.Canceled) { + t.Fatalf("canceled acquire error = %v, want context.Canceled", err) + } + release() + + // A canceled waiter must not leave the path permanently busy. + reacquired, err := coordinator.acquire(context.Background(), path, true) + if err != nil { + t.Fatalf("reacquire after cancellation: %v", err) + } + reacquired() +} + +func TestUpgradeBackendUsesResolvedConcretePath(t *testing.T) { + backendsPath := t.TempDir() + state, err := system.GetSystemState(system.WithBackendPath(backendsPath)) + if err != nil { + t.Fatalf("get system state: %v", err) + } + + concreteName := "concrete-backend" + concretePath := filepath.Join(backendsPath, concreteName) + if err := os.MkdirAll(concretePath, 0750); err != nil { + t.Fatalf("create concrete backend: %v", err) + } + if err := os.WriteFile(filepath.Join(concretePath, runFile), []byte("#!/bin/sh\n"), 0755); err != nil { + t.Fatalf("write concrete run file: %v", err) + } + if err := writeBackendMetadata(concretePath, &BackendMetadata{Name: concreteName, Version: "1"}); err != nil { + t.Fatalf("write concrete metadata: %v", err) + } + + metaName := "meta-backend" + metaPath := filepath.Join(backendsPath, metaName) + if err := os.MkdirAll(metaPath, 0750); err != nil { + t.Fatalf("create meta backend: %v", err) + } + if err := writeBackendMetadata(metaPath, &BackendMetadata{Name: metaName, MetaBackendFor: concreteName}); err != nil { + t.Fatalf("write meta metadata: %v", err) + } + + release, err := backendOperations.acquire(context.Background(), concretePath, false) + if err != nil { + t.Fatalf("hold concrete backend operation: %v", err) + } + t.Cleanup(release) + + err = UpgradeBackend( + context.Background(), state, nil, nil, metaName, nil, false, + WithSkipIfBackendBusy(), + ) + if !errors.Is(err, ErrBackendOperationInProgress) { + release() + t.Fatalf("auto-upgrade error = %v, want ErrBackendOperationInProgress", err) + } + + waitResult := make(chan error, 1) + go func() { + waitResult <- UpgradeBackend(context.Background(), state, nil, nil, metaName, nil, false) + }() + waitForBackendOperationReferences(t, backendOperations, concretePath, 2) + select { + case err := <-waitResult: + release() + t.Fatalf("manual upgrade returned while concrete backend was busy: %v", err) + default: + } + + release() + select { + case err := <-waitResult: + if err == nil || errors.Is(err, ErrBackendOperationInProgress) { + t.Fatalf("manual upgrade error after release = %v, want normal gallery lookup failure", err) + } + case <-time.After(time.Second): + t.Fatal("manual upgrade did not continue after concrete backend was released") + } +} diff --git a/core/gallery/upgrade.go b/core/gallery/upgrade.go index 7abc65a4358b..ed9f5f341d87 100644 --- a/core/gallery/upgrade.go +++ b/core/gallery/upgrade.go @@ -240,7 +240,7 @@ func summarizeNodeDrift(nodes []NodeBackendRef) (majority struct{ version, diges // 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), requireIntegrity bool) error { +func UpgradeBackend(ctx context.Context, systemState *system.SystemState, modelLoader *model.ModelLoader, galleries []config.Gallery, backendName string, downloadStatus func(string, string, string, float64), requireIntegrity bool, opts ...UpgradeOption) error { // Look up the installed backend installedBackends, err := ListSystemBackends(systemState) if err != nil { @@ -259,7 +259,32 @@ func UpgradeBackend(ctx context.Context, systemState *system.SystemState, modelL // 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, requireIntegrity) + return UpgradeBackend(ctx, systemState, modelLoader, galleries, installed.Metadata.MetaBackendFor, downloadStatus, requireIntegrity, opts...) + } + + options := upgradeOptions{} + for _, opt := range opts { + if opt != nil { + opt(&options) + } + } + + backendPath := filepath.Join(systemState.Backend.BackendsPath, backendName) + release, err := backendOperations.acquire(ctx, backendPath, options.skipIfBusy) + if err != nil { + return fmt.Errorf("backend %q: %w", backendName, err) + } + defer release() + + // Refresh installed state after waiting. Another upgrade may have replaced + // metadata while this call was blocked on the same backend path. + installedBackends, err = ListSystemBackends(systemState) + if err != nil { + return fmt.Errorf("failed to refresh installed backends: %w", err) + } + installed, ok = installedBackends.Get(backendName) + if !ok { + return fmt.Errorf("backend %q: %w", backendName, ErrBackendNotFound) } // Find the gallery entry. Unfiltered for the same reason as the check @@ -285,7 +310,6 @@ func UpgradeBackend(ctx context.Context, systemState *system.SystemState, modelL return fmt.Errorf("upgrade %q: %w", backendName, err) } - backendPath := filepath.Join(systemState.Backend.BackendsPath, backendName) tmpPath := backendPath + ".upgrade-tmp" backupPath := backendPath + ".backup" diff --git a/core/services/galleryop/managers_local.go b/core/services/galleryop/managers_local.go index 70fc018e000b..b798b049886f 100644 --- a/core/services/galleryop/managers_local.go +++ b/core/services/galleryop/managers_local.go @@ -105,12 +105,20 @@ func (b *LocalBackendManager) ListBackends() (gallery.SystemBackends, error) { return gallery.ListSystemBackends(b.systemState) } -// UpgradeBackend ignores op.ID and op.TargetNodeID: a single-node install -// reports progress through the local progressCb already, and there is only -// one node to target. Both fields only matter for distributed per-node -// streaming/scoping (see DistributedBackendManager.UpgradeBackend). +// UpgradeBackend uses an empty op.ID to identify the UpgradeChecker's +// background operation, which backs off instead of waiting when the same +// backend is busy. A single-node upgrade reports progress through the local +// progressCb already, and TargetNodeID only matters for distributed per-node +// scoping. func (b *LocalBackendManager) UpgradeBackend(ctx context.Context, op *ManagementOp[gallery.GalleryBackend, any], progressCb ProgressCallback) error { - return gallery.UpgradeBackend(ctx, b.systemState, b.modelLoader, b.backendGalleries, op.GalleryElementName, progressCb, b.requireBackendIntegrity) + var opts []gallery.UpgradeOption + if op.ID == "" { + opts = append(opts, gallery.WithSkipIfBackendBusy()) + } + return gallery.UpgradeBackend( + ctx, b.systemState, b.modelLoader, b.backendGalleries, + op.GalleryElementName, progressCb, b.requireBackendIntegrity, opts..., + ) } func (b *LocalBackendManager) CheckUpgrades(ctx context.Context) (map[string]gallery.UpgradeInfo, error) {