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

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()

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)
+}

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())
+	})
+})