From 4d7ced0ef602b22477639a2817696d334980c6de Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Sat, 26 Sep 2026 15:48:05 +0000 Subject: [PATCH] feat(failover): schedule liveness and recovery probes Idle targets get a liveness probe each interval, recovering targets an inference probe. Cold local targets are never loaded to be probed. Assisted-by: Claude:claude-opus-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/failover/manager.go | 1 - core/services/failover/schedule.go | 148 ++++++++++++++++++++++ core/services/failover/schedule_test.go | 160 ++++++++++++++++++++++++ 3 files changed, 308 insertions(+), 1 deletion(-) create mode 100644 core/services/failover/schedule.go create mode 100644 core/services/failover/schedule_test.go diff --git a/core/services/failover/manager.go b/core/services/failover/manager.go index 1169061bf..cfb2d8b66 100644 --- a/core/services/failover/manager.go +++ b/core/services/failover/manager.go @@ -581,7 +581,6 @@ func (m *Manager) emitLocked(ev Event) { } } -//nolint:unused // wired by the Task 5 probe scheduler's Stop, which owns the manager's lifecycle func (m *Manager) close() { m.mu.Lock() defer m.mu.Unlock() diff --git a/core/services/failover/schedule.go b/core/services/failover/schedule.go new file mode 100644 index 000000000..170de3cfa --- /dev/null +++ b/core/services/failover/schedule.go @@ -0,0 +1,148 @@ +package failover + +import ( + "context" + "sync" + "time" + + "github.com/mudler/LocalAI/core/config" +) + +// Run drives probes and dwell-based fail-back until ctx ends. Only this +// scheduler goroutine calls Sync: onWarm callbacks run after the manager's +// lock is released, and a concurrent Sync from elsewhere could reorder them. +func (m *Manager) Run(ctx context.Context) { + ticker := time.NewTicker(time.Second) + defer ticker.Stop() + m.Tick(ctx) + for { + select { + case <-ctx.Done(): + m.close() + return + case <-ticker.C: + m.Tick(ctx) + } + } +} + +// Tick runs one pass: sync configs, run due probes, recompute chains. It is +// exported so tests can drive the manager without a real ticker. Like Run, it +// must only be called from the scheduler goroutine (see Run's comment on Sync). +func (m *Manager) Tick(ctx context.Context) { + m.Sync() + var wg sync.WaitGroup + for _, j := range m.dueProbes() { + wg.Add(1) + go func(j probeJob) { + defer wg.Done() + m.runProbe(ctx, j) + }(j) + } + wg.Wait() + m.Reevaluate() +} + +type probeJob struct { + target string + cfg config.ModelConfig + kind Kind + warm bool + inference bool + timeout time.Duration +} + +func (m *Manager) dueProbes() []probeJob { + m.mu.Lock() + defer m.mu.Unlock() + now := m.clock.Now() + var jobs []probeJob + for _, ts := range m.targets { + interval := ts.params.ProbeInterval() + inference := false + switch ts.state { + case StateMissing: + continue + case StateHealthy: + // A served request is as good as a liveness probe. + if now.Sub(ts.lastActivity) < interval || now.Sub(ts.lastProbe) < interval { + continue + } + case StateDown: + if ts.cold() { + // Loading a cold model only to probe it could evict others. + if now.Sub(ts.downSince) >= ts.params.MinDwell() { + m.setTargetLocked(ts, StateHealthy, ReasonRecovery, "") + m.recomputeForLocked(ts.name) + } + continue + } + if now.Sub(ts.lastProbe) < interval { + continue + } + case StateRecovering: + if ts.cold() || now.Sub(ts.lastProbe) < interval { + continue + } + inference = true + } + if m.prober == nil { + continue + } + cfg, ok := m.lookupTarget(ts.name) + if !ok { + continue + } + ts.lastProbe = now + jobs = append(jobs, probeJob{ + target: ts.name, cfg: cfg, kind: ts.kind, warm: ts.warm, + inference: inference, timeout: ts.params.ProbeTimeout(), + }) + } + return jobs +} + +func (m *Manager) runProbe(ctx context.Context, j probeJob) { + pctx, cancel := context.WithTimeout(ctx, j.timeout) + defer cancel() + var err error + if j.inference { + err = m.prober.Inference(pctx, j.cfg, j.kind, j.warm) + } else { + err = m.prober.Liveness(pctx, j.cfg, j.kind, j.warm) + } + if ctx.Err() != nil { + return // shutting down: a cancelled probe says nothing about the target + } + m.applyProbe(j, err) +} + +func (m *Manager) applyProbe(j probeJob, err error) { + m.mu.Lock() + defer m.mu.Unlock() + ts := m.targets[j.target] + if ts == nil || ts.state == StateMissing { + return + } + if err != nil { + if ts.state == StateDown { + ts.lastError = err.Error() + } else { + m.recordFailureLocked(ts, err.Error()) + } + m.recomputeForLocked(ts.name) + return + } + switch ts.state { + case StateHealthy: + ts.lastActivity = m.clock.Now() + ts.failures = nil + case StateDown: + m.setTargetLocked(ts, StateRecovering, ReasonRecovery, "") + case StateRecovering: + if j.inference { + m.recordPassLocked(ts) + } + } + m.recomputeForLocked(ts.name) +} diff --git a/core/services/failover/schedule_test.go b/core/services/failover/schedule_test.go new file mode 100644 index 000000000..931f1a151 --- /dev/null +++ b/core/services/failover/schedule_test.go @@ -0,0 +1,160 @@ +package failover + +import ( + "context" + "sync" + "time" + + "github.com/mudler/LocalAI/core/config" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +type probeCall struct { + target string + inference bool +} + +type fakeProber struct { + mu sync.Mutex + calls []probeCall + fail map[string]error // target -> error returned by every probe +} + +func (p *fakeProber) record(target string, inference bool) error { + p.mu.Lock() + defer p.mu.Unlock() + p.calls = append(p.calls, probeCall{target, inference}) + return p.fail[target] +} +func (p *fakeProber) Liveness(_ context.Context, c config.ModelConfig, _ Kind, _ bool) error { + return p.record(c.Name, false) +} +func (p *fakeProber) Inference(_ context.Context, c config.ModelConfig, _ Kind, _ bool) error { + return p.record(c.Name, true) +} +func (p *fakeProber) take() []probeCall { + p.mu.Lock() + defer p.mu.Unlock() + out := p.calls + p.calls = nil + return out +} + +var _ = Describe("Manager probes", func() { + var ( + clock *fakeClock + src *fakeSource + prober *fakeProber + m *Manager + ctx = context.Background() + ) + + BeforeEach(func() { + clock = newFakeClock() + prober = &fakeProber{fail: map[string]error{}} + src = newFakeSource(remote("a"), local("b"), local("cold"), + chainCfg("chain", nil, t("a"), warmT("b"))) + m = New(src, WithClock(clock), WithProber(prober)) + }) + + It("probes idle targets on the first tick and not again before the interval", func() { + m.Tick(ctx) + Expect(prober.take()).To(ConsistOf(probeCall{"a", false}, probeCall{"b", false})) + clock.Advance(5 * time.Second) + m.Tick(ctx) + Expect(prober.take()).To(BeEmpty()) + }) + + It("skips the liveness probe for a target with recent traffic", func() { + m.Tick(ctx) + prober.take() + clock.Advance(14 * time.Second) + m.ReportSuccess("a") + clock.Advance(2 * time.Second) + m.Tick(ctx) + Expect(prober.take()).To(ConsistOf(probeCall{"b", false})) + }) + + It("trips a target whose liveness probe fails", func() { + prober.fail["a"] = errBoom + m.Tick(ctx) + st, _ := m.ChainStatus("chain") + Expect(st.Targets[0].State).To(Equal(StateDown)) + Expect(st.Active).To(Equal("b")) + }) + + It("recovers through liveness, then inference probes, then fails back after dwell", func() { + prober.fail["a"] = errBoom + m.Tick(ctx) + delete(prober.fail, "a") + prober.take() + + clock.Advance(15 * time.Second) + m.Tick(ctx) // liveness passes: recovering + st, _ := m.ChainStatus("chain") + Expect(st.Targets[0].State).To(Equal(StateRecovering)) + + for i := 0; i < 3; i++ { + clock.Advance(15 * time.Second) + m.Tick(ctx) + } + calls := prober.take() + Expect(calls).To(ContainElement(probeCall{"a", true})) + st, _ = m.ChainStatus("chain") + Expect(st.Targets[0].State).To(Equal(StateHealthy)) + Expect(st.Active).To(Equal("a"), "60s min_dwell passed during the 4 ticks") + }) + + It("sends a recovering target back down when an inference probe fails", func() { + m.ReportFailure("a", errBoom) + clock.Advance(15 * time.Second) + m.Tick(ctx) // liveness passes: recovering + prober.fail["a"] = errBoom + clock.Advance(15 * time.Second) + m.Tick(ctx) + st, _ := m.ChainStatus("chain") + Expect(st.Targets[0].State).To(Equal(StateDown)) + }) + + It("never probes a down cold target and restores it after min_dwell", func() { + src.Put(chainCfg("chain", nil, t("cold"), warmT("b"))) + m.Sync() + m.ReportFailure("cold", errBoom) + prober.take() + clock.Advance(30 * time.Second) + m.Tick(ctx) + for _, c := range prober.take() { + Expect(c.target).ToNot(Equal("cold")) + } + st, _ := m.ChainStatus("chain") + Expect(st.Targets[0].State).To(Equal(StateDown)) + clock.Advance(31 * time.Second) + m.Tick(ctx) + st, _ = m.ChainStatus("chain") + Expect(st.Targets[0].State).To(Equal(StateHealthy)) + }) + + It("probes a target shared by two chains once per tick", func() { + src.Put(chainCfg("chain2", nil, t("a"), warmT("b"))) + m.Tick(ctx) + calls := prober.take() + n := 0 + for _, c := range calls { + if c.target == "a" { + n++ + } + } + Expect(n).To(Equal(1)) + }) + + It("closes subscriptions when Run stops", func() { + events, _ := m.Subscribe(1) + rctx, cancel := context.WithCancel(ctx) + done := make(chan struct{}) + go func() { m.Run(rctx); close(done) }() + cancel() + Eventually(done).Should(BeClosed()) + Eventually(events).Should(BeClosed()) + }) +})