diff --git a/core/application/failover.go b/core/application/failover.go index 3194e294b..f0b5c2c2b 100644 --- a/core/application/failover.go +++ b/core/application/failover.go @@ -5,13 +5,28 @@ import ( "github.com/mudler/xlog" ) +// preloadModelByName is a seam over backend.PreloadModelByName so tests can +// substitute a controllable loader instead of touching real models/disk. +var preloadModelByName = backend.PreloadModelByName + // applyFailoverWarmTargets pins warm failover targets in the watchdog and // loads them, so a switch does not wait for a cold load. +// +// SyncPinnedModelsToWatchdog runs synchronously: it is a cheap in-memory +// update, and the pin must land before the watchdog can evict a target that +// is about to become (or stay) a chain's active path. Preloading is not +// cheap — it can download or load a multi-GB model — and this callback runs +// on the failover manager's single scheduler goroutine (Sync, called from +// Tick, called from Run), before that tick's probes fire. A slow or hung +// load here would freeze probing and fail-back for every chain, so it runs +// in its own goroutine instead of blocking the scheduler loop. func (a *Application) applyFailoverWarmTargets(warm []string) { a.SyncPinnedModelsToWatchdog() - for _, name := range warm { - if _, err := backend.PreloadModelByName(a.ApplicationConfig().Context, a.ModelConfigLoader(), a.ModelLoader(), a.ApplicationConfig(), name); err != nil { - xlog.Warn("failover: could not preload warm target", "model", name, "error", err) + go func() { + for _, name := range warm { + if _, err := preloadModelByName(a.ApplicationConfig().Context, a.ModelConfigLoader(), a.ModelLoader(), a.ApplicationConfig(), name); err != nil { + xlog.Warn("failover: could not preload warm target", "model", name, "error", err) + } } - } + }() } diff --git a/core/application/failover_test.go b/core/application/failover_test.go new file mode 100644 index 000000000..a9057425a --- /dev/null +++ b/core/application/failover_test.go @@ -0,0 +1,41 @@ +package application + +import ( + "context" + "time" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/model" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("applyFailoverWarmTargets", func() { + It("returns promptly even when the preload call blocks", func() { + // Guards against a regression to a synchronous preload loop: onWarm + // runs on the failover manager's single scheduler goroutine, so a + // blocking loader here must not block the caller. + started := make(chan struct{}) + release := make(chan struct{}) + orig := preloadModelByName + preloadModelByName = func(ctx context.Context, cl *config.ModelConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig, name string) ([]string, error) { + close(started) + <-release // never released within the test's timeout + return nil, nil + } + DeferCleanup(func() { preloadModelByName = orig }) + DeferCleanup(func() { close(release) }) + + app := &Application{applicationConfig: &config.ApplicationConfig{Context: context.Background()}} + + callReturned := make(chan struct{}) + go func() { + defer GinkgoRecover() + app.applyFailoverWarmTargets([]string{"warm-a"}) + close(callReturned) + }() + + Eventually(callReturned, time.Second).Should(BeClosed(), "applyFailoverWarmTargets must not wait on the preload goroutine") + Eventually(started, time.Second).Should(BeClosed(), "the preload goroutine should still run in the background") + }) +})