diff --git a/core/application/application.go b/core/application/application.go index f5454aba1..b15d848a0 100644 --- a/core/application/application.go +++ b/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() diff --git a/core/application/distributed.go b/core/application/distributed.go index 4c6214e57..6ba8234cc 100644 --- a/core/application/distributed.go +++ b/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, diff --git a/core/application/failover.go b/core/application/failover.go index a513c1f10..ffe5a8e51 100644 --- a/core/application/failover.go +++ b/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 { diff --git a/core/application/failover_distributed.go b/core/application/failover_distributed.go new file mode 100644 index 000000000..7c1b63b9f --- /dev/null +++ b/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)) +} diff --git a/core/application/failover_distributed_test.go b/core/application/failover_distributed_test.go new file mode 100644 index 000000000..b4cf84bf0 --- /dev/null +++ b/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") + }) +}) diff --git a/core/application/startup.go b/core/application/startup.go index db9e1caaf..fe8ee1394 100644 --- a/core/application/startup.go +++ b/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) diff --git a/core/services/advisorylock/keys.go b/core/services/advisorylock/keys.go index 277817229..76a1834ed 100644 --- a/core/services/advisorylock/keys.go +++ b/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 ) diff --git a/core/services/failover/schedule.go b/core/services/failover/schedule.go index 739c48b42..c910cb9b4 100644 --- a/core/services/failover/schedule.go +++ b/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() diff --git a/core/services/failover/statesync.go b/core/services/failover/statesync.go index bcab901ef..a14378a04 100644 --- a/core/services/failover/statesync.go +++ b/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) { diff --git a/core/services/failover/statesync_test.go b/core/services/failover/statesync_test.go index 3901d3b95..d379bb674 100644 --- a/core/services/failover/statesync_test.go +++ b/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))) diff --git a/docs/content/features/model-failover.md b/docs/content/features/model-failover.md index 5915ff8b3..24b9ae8dd 100644 --- a/docs/content/features/model-failover.md +++ b/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" %}}).