feat(failover): run one prober per cluster and pin warm targets on workers

Assisted-by: Claude:claude-opus-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

core/application/application.go

@@ -17,6 +17,7 @@ import (
 	"github.com/mudler/LocalAI/core/services/cloudproxy/mitm"
 	"github.com/mudler/LocalAI/core/services/facerecognition"
 	"github.com/mudler/LocalAI/core/services/failover"
+	"github.com/mudler/LocalAI/core/services/failover/distsync"
 	"github.com/mudler/LocalAI/core/services/galleryop"
 	"github.com/mudler/LocalAI/core/services/monitoring"
 	"github.com/mudler/LocalAI/core/services/nodes"
@@ -94,6 +95,10 @@ type Application struct {
 	// Distributed mode services (nil when not in distributed mode)
 	distributed *DistributedServices
 
+	// failoverSync shares failover state between frontends; nil in
+	// standalone mode or when it could not start.
+	failoverSync *distsync.Sync
+
 	// Upgrade checker (background service for detecting backend upgrades)
 	upgradeChecker *UpgradeChecker
 
@@ -515,6 +520,13 @@ func (a *Application) IsDistributed() bool {
 func (a *Application) Shutdown() error {
 	var err error
 	a.shutdownOnce.Do(func() {
+		// Before distributed shutdown: the sync's subscriptions live on the
+		// NATS connection that closes there.
+		if a.failoverSync != nil {
+			if closeErr := a.failoverSync.Close(); closeErr != nil {
+				xlog.Warn("failover: closing state sync", "error", closeErr)
+			}
+		}
 		a.distributed.Shutdown()
 		if a.modelLoader != nil {
 			err = a.modelLoader.StopAllGRPC()

core/application/distributed.go

@@ -77,7 +77,9 @@ func (ds *DistributedServices) Shutdown() {
 // Returns nil if distributed mode is not enabled.
 // configLoader is used by the SmartRouter to compute concurrency-group
 // anti-affinity at placement time (#9659); it may be nil in tests.
-func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoader *config.ModelConfigLoader) (*DistributedServices, error) {
+// pinned, when set, replaces configLoader as the source of models the router
+// and reconciler must keep loaded (it adds warm failover targets).
+func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoader *config.ModelConfigLoader, pinned nodes.PinnedModelResolver) (*DistributedServices, error) {
 	if !cfg.Distributed.Enabled {
 		return nil, nil
 	}
@@ -383,6 +385,9 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
 		conflictResolver = configLoader
 		pinnedResolver = configLoader
 	}
+	if pinned != nil {
+		pinnedResolver = pinned
+	}
 	modelCleanup := nodes.NewModelCleanupService(registry, remoteUnloader)
 	router := nodes.NewSmartRouter(registry, nodes.SmartRouterOptions{
 		Unloader:         remoteUnloader,

core/application/failover.go

@@ -24,8 +24,15 @@ var preloadModelByName = backend.PreloadModelByName
 // 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.
+//
+// In distributed mode only the probe leader preloads: every frontend would
+// otherwise ask the workers for the same load. Followers still pin, so their
+// own watchdog never evicts a warm target they happen to hold.
 func (a *Application) applyFailoverWarmTargets(warm []string) {
 	a.SyncPinnedModelsToWatchdog()
+	if a.IsDistributed() && !a.failoverManager.IsLeader() {
+		return
+	}
 	go func() {
 		for _, name := range warm {
 			if _, err := preloadModelByName(a.ApplicationConfig().Context, a.ModelConfigLoader(), a.ModelLoader(), a.ApplicationConfig(), name); err != nil {

core/application/failover_distributed.go

@@ -0,0 +1,69 @@
+package application
+
+import (
+	"context"
+
+	"github.com/mudler/LocalAI/core/services/advisorylock"
+	"github.com/mudler/LocalAI/core/services/failover"
+	"github.com/mudler/LocalAI/core/services/failover/distsync"
+	"github.com/mudler/LocalAI/core/services/nodes"
+	"github.com/mudler/xlog"
+	"gorm.io/gorm"
+)
+
+// failoverLeaderGate elects the frontend that probes failover targets and
+// decides chains, so N frontends do not probe every target N times or
+// disagree on which target is active. A lock error counts as "not leader":
+// skipping one tick is safe, two leaders are not.
+func failoverLeaderGate(db *gorm.DB) failover.LeaderGate {
+	return func(ctx context.Context, fn func()) bool {
+		ok, err := advisorylock.TryWithLockCtx(ctx, db, advisorylock.KeyFailoverProber, func() error {
+			fn()
+			return nil
+		})
+		if err != nil {
+			xlog.Warn("failover: could not take the prober leader lock", "error", err)
+			return false
+		}
+		return ok
+	}
+}
+
+// failoverPinnedResolver adds warm failover targets to the config-pinned
+// models, so the SmartRouter and ReplicaReconciler keep them loaded on the
+// workers the same way the local watchdog keeps them loaded in standalone mode.
+type failoverPinnedResolver struct {
+	base nodes.PinnedModelResolver
+	fm   *failover.Manager
+}
+
+func (r *failoverPinnedResolver) GetPinnedModelNames() []string {
+	var pinned []string
+	if r.base != nil {
+		pinned = r.base.GetPinnedModelNames()
+	}
+	return failover.MergePinned(pinned, r.fm.WarmTargets())
+}
+
+// startFailoverDistributed makes the failover manager cluster-aware: one
+// frontend (the advisory-lock holder) probes and decides, and pins, target
+// health and chain state are shared over NATS. A sync failure is logged and
+// the manager keeps probing on its own, as in standalone mode.
+func (a *Application) startFailoverDistributed(ctx context.Context) {
+	db := a.distributedDB()
+	pins, err := distsync.NewPinStore(db)
+	if err != nil {
+		xlog.Error("failover: pins will not persist, could not prepare the pin store", "error", err)
+		pins = nil // distsync.New treats a nil store as "no durable pins"
+	}
+	s, err := distsync.New(ctx, a.distributed.Nats, pins, a.failoverManager)
+	if err != nil {
+		xlog.Error("failover: state will not be shared between frontends", "error", err)
+		return
+	}
+	a.failoverSync = s
+	// Gate only once state is shared: a follower learns target health and
+	// chain decisions solely from the leader's publishes, so gating without
+	// the sync would leave every follower's chains frozen.
+	a.failoverManager.SetLeaderGate(failoverLeaderGate(db))
+}

core/application/failover_distributed_test.go

@@ -0,0 +1,102 @@
+package application
+
+import (
+	"context"
+	"time"
+
+	"github.com/mudler/LocalAI/core/config"
+	"github.com/mudler/LocalAI/core/services/failover"
+	"github.com/mudler/LocalAI/pkg/model"
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+)
+
+// staticPins is a nodes.PinnedModelResolver with a fixed list.
+type staticPins []string
+
+func (s staticPins) GetPinnedModelNames() []string { return s }
+
+// failoverSource is a failover.ConfigSource over a fixed set of configs.
+type failoverSource map[string]config.ModelConfig
+
+func (s failoverSource) GetModelConfig(name string) (config.ModelConfig, bool) {
+	c, ok := s[name]
+	return c, ok
+}
+
+func (s failoverSource) GetAllModelsConfigs() []config.ModelConfig {
+	out := make([]config.ModelConfig, 0, len(s))
+	for _, c := range s {
+		out = append(out, c)
+	}
+	return out
+}
+
+// warmChainSource has one chain whose two local targets are warm.
+func warmChainSource() failoverSource {
+	return failoverSource{
+		"b": {Name: "b", Backend: "llama-cpp"},
+		"c": {Name: "c", Backend: "llama-cpp"},
+		"chain": {Name: "chain", Failover: &config.FailoverConfig{Targets: []config.FailoverTarget{
+			{Model: "b", Warm: true}, {Model: "c", Warm: true},
+		}}},
+	}
+}
+
+var _ = Describe("failoverPinnedResolver", func() {
+	It("merges config pins and warm failover targets without duplicates", func() {
+		fm := failover.New(warmChainSource())
+		fm.Sync()
+		Expect(fm.WarmTargets()).To(Equal([]string{"b", "c"}))
+
+		r := &failoverPinnedResolver{base: staticPins{"pinned-a", "b"}, fm: fm}
+		Expect(r.GetPinnedModelNames()).To(Equal([]string{"pinned-a", "b", "c"}))
+	})
+})
+
+var _ = Describe("failoverLeaderGate", func() {
+	It("lets only one frontend lead at a time on the same database", func() {
+		// Not PostgreSQL, so advisorylock falls back to its in-process lock,
+		// which has the same try-lock semantics as pg_try_advisory_lock.
+		db, err := gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{})
+		Expect(err).ToNot(HaveOccurred())
+		first, second := failoverLeaderGate(db), failoverLeaderGate(db)
+		ctx := context.Background()
+
+		var secondInside, secondRan bool
+		Expect(first(ctx, func() {
+			secondInside = second(ctx, func() { secondRan = true })
+		})).To(BeTrue())
+		Expect(secondInside).To(BeFalse(), "the second gate must not lead while the first holds the lock")
+		Expect(secondRan).To(BeFalse())
+
+		Expect(second(ctx, func() { secondRan = true })).To(BeTrue())
+		Expect(secondRan).To(BeTrue())
+	})
+})
+
+var _ = Describe("applyFailoverWarmTargets in distributed mode", func() {
+	It("pins but does not preload on a frontend that is not the probe leader", func() {
+		preloaded := make(chan string, 1)
+		orig := preloadModelByName
+		preloadModelByName = func(_ context.Context, _ *config.ModelConfigLoader, _ *model.ModelLoader, _ *config.ApplicationConfig, name string) ([]string, error) {
+			preloaded <- name
+			return nil, nil
+		}
+		DeferCleanup(func() { preloadModelByName = orig })
+
+		// A gate that never grants leadership: this frontend is a follower.
+		fm := failover.New(warmChainSource(),
+			failover.WithLeaderGate(func(context.Context, func()) bool { return false }))
+		app := &Application{
+			applicationConfig: &config.ApplicationConfig{Context: context.Background()},
+			distributed:       &DistributedServices{},
+			failoverManager:   fm,
+		}
+
+		app.applyFailoverWarmTargets([]string{"b"})
+		Consistently(preloaded, 200*time.Millisecond).ShouldNot(Receive(), "only the probe leader preloads warm targets")
+	})
+})

core/application/startup.go

@@ -286,12 +286,16 @@ func New(opts ...config.AppOption) (*Application, error) {
 	// the model configs are loaded, so it is declared out here.
 	var revisionStore modeladmin.RevisionStore
 
-	distSvc, err := initDistributed(options, application.authDB, application.ModelConfigLoader())
+	distSvc, err := initDistributed(options, application.authDB, application.ModelConfigLoader(),
+		&failoverPinnedResolver{base: application.ModelConfigLoader(), fm: application.failoverManager})
 	if err != nil {
 		return nil, fmt.Errorf("distributed mode initialization failed: %w", err)
 	}
 	if distSvc != nil {
 		application.distributed = distSvc
+		// Before failoverManager.Run starts below: the gate and sync must be
+		// in place for its first tick.
+		application.startFailoverDistributed(options.Context)
 		// Wire remote model unloader so ShutdownModel works for remote nodes
 		// Uses NATS to tell serve-backend nodes to Free + kill their backend process
 		application.modelLoader.SetRemoteUnloader(distSvc.Unloader)

core/services/advisorylock/keys.go

@@ -12,4 +12,7 @@ const (
 	KeySchemaMigrate        int64 = 105
 	KeyBackendUpgradeCheck  int64 = 106
 	KeyStateReconciler      int64 = 107
+	// KeyFailoverProber elects the one frontend that probes failover
+	// targets, decides chains and preloads warm targets.
+	KeyFailoverProber       int64 = 108
 )

core/services/failover/schedule.go

@@ -46,11 +46,14 @@ func (m *Manager) Run(ctx context.Context) {
 // leader decision.
 func (m *Manager) Tick(ctx context.Context) {
 	m.Sync()
-	if m.gate == nil {
+	m.mu.Lock()
+	gate := m.gate // SetLeaderGate may replace it after construction
+	m.mu.Unlock()
+	if gate == nil {
 		m.lead(ctx)
 		return
 	}
-	if m.gate(ctx, func() { m.lead(ctx) }) {
+	if gate(ctx, func() { m.lead(ctx) }) {
 		return
 	}
 	m.mu.Lock()

core/services/failover/statesync.go

@@ -48,6 +48,18 @@ type LeaderGate func(ctx context.Context, fn func()) bool
 // grants leadership. Without it the manager is always the leader.
 func WithLeaderGate(g LeaderGate) Option { return func(m *Manager) { m.gate = g } }
 
+// SetLeaderGate is WithLeaderGate for a manager that already exists: the
+// application builds the manager before distributed init, where the gate's
+// database becomes known. Leadership is reset to match New, so the first tick
+// that wins the gate counts as becoming leader (warm set redelivered, state
+// republished at once).
+func (m *Manager) SetLeaderGate(g LeaderGate) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.gate = g
+	m.leader = g == nil
+}
+
 // SetStateSync attaches the sync layer and hydrates pins from it. Pins are
 // read outside the lock because the store may need I/O.
 func (m *Manager) SetStateSync(s StateSync) {

core/services/failover/statesync_test.go

@@ -140,6 +140,18 @@ var _ = Describe("Manager state sync", func() {
 		Expect(st.Pinned).ToNot(BeNil())
 	})
 
+	It("SetLeaderGate gates a manager built without one", func() {
+		// Production builds the manager before distributed init, so the gate
+		// arrives through the setter rather than the option.
+		p := &fakeProber{fail: map[string]error{}}
+		m := New(src, WithClock(clock), WithProber(p))
+		Expect(m.IsLeader()).To(BeTrue())
+		m.SetLeaderGate(gateFor(false))
+		Expect(m.IsLeader()).To(BeFalse(), "a gated manager is not the leader until the gate grants it")
+		m.Tick(ctx)
+		Consistently(func() int { return len(p.take()) }, 200*time.Millisecond).Should(Equal(0), "a follower must not probe")
+	})
+
 	It("only the leader probes", func() {
 		pa, pb := &fakeProber{fail: map[string]error{}}, &fakeProber{fail: map[string]error{}}
 		a = New(src, WithClock(clock), WithProber(pa), WithLeaderGate(gateFor(true)))

docs/content/features/model-failover.md

@@ -167,7 +167,8 @@ curl -X DELETE http://localhost:8080/api/failover/assistant-llm/pin
 ```
 
 While a chain is pinned, only the pinned target serves it. Health checks
-continue. A restart removes the pin.
+continue. On a single LocalAI instance, a restart removes the pin. In
+[distributed mode](#distributed-mode), pins persist.
 
 ## Assistant and MCP
 
@@ -175,10 +176,25 @@ The LocalAI Assistant and `local-ai mcp-server` offer `list_failover_chains`,
 `pin_failover_target` and `unpin_failover_target`. Create and edit chains with
 the model config tools, like any other model.
 
+## Distributed mode
+
+In [distributed mode]({{%relref "features/distributed-mode" %}}), all frontends
+share one failover state:
+
+- Pins apply to the whole cluster. LocalAI stores them in the database, so
+  they persist across restarts. A pin set on one frontend applies on all.
+- Target health and the active target of each chain are shared over NATS.
+  All frontends send a chain's requests to the same target.
+- One frontend, the probe leader, runs the health checks, decides fail-over
+  and fail-back, and loads warm targets. A PostgreSQL advisory lock selects
+  the leader. If the leader stops, another frontend takes over.
+- Warm targets stay loaded on the workers. The router and the replica
+  reconciler treat them like pinned models and do not evict them.
+- A frontend that starts late gets the current state within 10 seconds,
+  because the leader sends its full state again every 10 seconds.
+
 ## Limits
 
-- Failover state is kept in memory by each LocalAI instance. Several frontends
-  in distributed mode each keep their own view.
 - Chains do not nest.
 - See also [model aliases]({{%relref "features/model-aliases" %}}) and the
   [realtime API]({{%relref "features/openai-realtime" %}}).