fix(failover): never load a warm target inside a probe

A warm target's liveness probe called ModelLoader.Load, which blocked
until the model finished loading (while the warm preload loaded it
too). Tick waited for every probe, so all probing froze, and the probe
then ran HealthCheck on an expired context and tripped the target at
every startup.

The prober now takes a function that returns the running backend
without loading it. A target that is not loaded passes liveness; its
recovery is neither confirmed nor failed and it returns to healthy
after min_dwell, like a cold target. Tick no longer waits for probes:
each probe applies its own result and a target whose probe is running
is skipped.

Assisted-by: Claude:claude-opus-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto committed 2026-09-27 07:42:20 +00:00
1 parent b67b75dcdd
commit 5d2b90b848
10 files changed
+242 -57

No files matched your search

+20
View File
@@ -2,6 +2,10 @@ package application
import (
"github.com/mudler/LocalAI/core/backend"
"github.com/mudler/LocalAI/core/config"
"github.com/mudler/LocalAI/core/services/failover"
"github.com/mudler/LocalAI/pkg/grpc"
"github.com/mudler/LocalAI/pkg/model"
"github.com/mudler/xlog"
)
@@ -30,3 +34,19 @@ func (a *Application) applyFailoverWarmTargets(warm []string) {
}
}()
}
// failoverLoadedBackend gives the failover prober the running backend of a
// local target without ever loading it. CheckIsLoaded may run the loader's
// own health check and drop a dead process; the target is then "not loaded"
// and the next real request loads and judges it.
func failoverLoadedBackend(ml *model.ModelLoader) failover.LoadedFunc {
return func(cfg config.ModelConfig) grpc.Backend {
m := ml.CheckIsLoaded(cfg.ModelID())
if m == nil {
return nil
}
// Load always enables parallel requests; match it in case this is
// the first client built for the model.
return m.GRPC(true, ml.GetWatchDog())
}
}
+22
View File
@@ -5,7 +5,9 @@ import (
"time"
"github.com/mudler/LocalAI/core/config"
"github.com/mudler/LocalAI/pkg/grpc"
"github.com/mudler/LocalAI/pkg/model"
"github.com/mudler/LocalAI/pkg/system"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
@@ -39,3 +41,23 @@ var _ = Describe("applyFailoverWarmTargets", func() {
Eventually(started, time.Second).Should(BeClosed(), "the preload goroutine should still run in the background")
})
})
type healthyBackend struct{ grpc.Backend }
func (healthyBackend) HealthCheck(context.Context) (bool, error) { return true, nil }
var _ = Describe("failoverLoadedBackend", func() {
It("returns the running backend and never loads a model that is not loaded", func() {
ml := model.NewModelLoader(&system.SystemState{Model: system.Model{ModelsPath: GinkgoT().TempDir()}})
store := model.NewInMemoryModelStore()
ml.SetModelStore(store)
loaded := failoverLoadedBackend(ml)
Expect(loaded(config.ModelConfig{Name: "gemma", Backend: "llama-cpp"})).To(BeNil())
Expect(ml.ListLoadedModels()).To(BeEmpty(), "the lookup must not start a load")
client := healthyBackend{}
store.Set("gemma", model.NewModelWithClient("gemma", "127.0.0.1:0", client))
Expect(loaded(config.ModelConfig{Name: "gemma", Backend: "llama-cpp"})).To(Equal(client))
})
})
+1 -5
View File
@@ -1,7 +1,6 @@
package application
import (
"context"
"crypto/rand"
"encoding/hex"
"fmt"
@@ -29,7 +28,6 @@ import (
"github.com/mudler/LocalAI/core/trace"
"github.com/mudler/LocalAI/internal"
"github.com/mudler/LocalAI/pkg/downloader"
"github.com/mudler/LocalAI/pkg/grpc"
"github.com/mudler/LocalAI/pkg/modelartifacts"
"github.com/mudler/LocalAI/pkg/signals"
"github.com/mudler/LocalAI/pkg/vram"
@@ -257,9 +255,7 @@ func New(opts ...config.AppOption) (*Application, error) {
// chain. WithOnWarmChanged pins and preloads warm local targets so a
// switch to them does not wait for a cold load.
application.failoverManager = failover.New(application.ModelConfigLoader(),
failover.WithProber(failover.NewProber(func(ctx context.Context, cfg config.ModelConfig) (grpc.Backend, error) {
return application.ModelLoader().Load(backend.ModelOptions(cfg, options)...)
})),
failover.WithProber(failover.NewProber(failoverLoadedBackend(application.ModelLoader()))),
failover.WithOnWarmChanged(application.applyFailoverWarmTargets),
)
// The assistant client was built in start() (above), before this
+5
View File
@@ -58,6 +58,8 @@ type Manager struct {
// hasChains mirrors len(chains) > 0 as of the last sync, so the request
// path can check it without the lock or a config-source scan.
hasChains atomic.Bool
// probes counts running probes; only tests wait on it.
probes sync.WaitGroup
}
type targetState struct {
@@ -71,6 +73,9 @@ type targetState struct {
lastProbe time.Time
lastActivity time.Time
lastError string
// probing is set while a probe runs, so the scheduler does not start a
// second one for the same target.
probing bool
// params come from the first chain, in name order, that lists the target.
params config.FailoverConfig
}
+36 -20
View File
@@ -18,18 +18,26 @@ import (
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
)
// LoadFunc returns the backend for a local target, loading it if needed.
type LoadFunc func(ctx context.Context, cfg config.ModelConfig) (grpc.Backend, error)
// ErrNotLoaded is what Inference returns for a local target whose backend is
// not running. It neither confirms nor fails recovery: the manager judges the
// target like a cold one, by real requests once min_dwell has passed.
var ErrNotLoaded = errors.New("failover: target is not loaded")
// LoadedFunc returns the running backend of a local target, or nil when it is
// not loaded. It must never load the model: a probe that loads blocks until
// the load ends (while the warm preload loads the same model) and then judges
// the target on an expired context.
type LoadedFunc func(cfg config.ModelConfig) grpc.Backend
// DefaultProber probes remote targets over the upstream's OpenAI-compatible
// API and local targets through their gRPC backend.
type DefaultProber struct {
HTTP *http.Client
Load LoadFunc
HTTP *http.Client
Loaded LoadedFunc
}
func NewProber(load LoadFunc) *DefaultProber {
return &DefaultProber{HTTP: &http.Client{}, Load: load}
func NewProber(loaded LoadedFunc) *DefaultProber {
return &DefaultProber{HTTP: &http.Client{}, Loaded: loaded}
}
func (p *DefaultProber) Liveness(ctx context.Context, cfg config.ModelConfig, kind Kind, warm bool) error {
@@ -225,15 +233,25 @@ func silenceWAV() []byte {
return b
}
func (p *DefaultProber) loaded(cfg config.ModelConfig) grpc.Backend {
if p.Loaded == nil {
return nil
}
return p.Loaded(cfg)
}
// localHealth checks a warm target's running backend. A target that is not
// loaded passes: the warm preload is loading it, or a crash removed it and
// the next real request loads it again and judges it.
func (p *DefaultProber) localHealth(ctx context.Context, cfg config.ModelConfig) error {
if p.Load == nil {
return errors.New("failover: no backend loader configured")
}
// Load returns the running backend, or starts it again after a crash.
b, err := p.Load(ctx, cfg)
if err != nil {
return err
b := p.loaded(cfg)
if b == nil {
return nil
}
return healthCheck(ctx, b)
}
func healthCheck(ctx context.Context, b grpc.Backend) error {
ok, err := b.HealthCheck(ctx)
if err != nil {
return err
@@ -245,13 +263,11 @@ func (p *DefaultProber) localHealth(ctx context.Context, cfg config.ModelConfig)
}
func (p *DefaultProber) localInference(ctx context.Context, cfg config.ModelConfig) error {
if p.Load == nil {
return errors.New("failover: no backend loader configured")
}
b, err := p.Load(ctx, cfg)
if err != nil {
return err
b := p.loaded(cfg)
if b == nil {
return ErrNotLoaded
}
var err error
switch {
case cfg.HasUsecases(config.FLAG_CHAT) || cfg.HasUsecases(config.FLAG_COMPLETION):
_, err = b.Predict(ctx, &pb.PredictOptions{Prompt: "ping", Tokens: 1})
@@ -262,5 +278,5 @@ func (p *DefaultProber) localInference(ctx context.Context, cfg config.ModelConf
}
// A backend process that answers HealthCheck rarely fails only for TTS or
// transcription, so a real request adds little here.
return p.localHealth(ctx, cfg)
return healthCheck(ctx, b)
}
+18 -4
View File
@@ -138,7 +138,7 @@ var _ = Describe("DefaultProber", func() {
It("uses HealthCheck for warm local liveness and Predict for local chat inference", func() {
b := &fakeBackend{healthy: true}
p = NewProber(func(context.Context, config.ModelConfig) (grpc.Backend, error) { return b, nil })
p = NewProber(func(config.ModelConfig) grpc.Backend { return b })
c := config.ModelConfig{Name: "gemma", Backend: "llama-cpp", KnownUsecaseStrings: []string{"chat"}}
c.KnownUsecases = config.GetUsecasesFromYAML(c.KnownUsecaseStrings)
Expect(p.Liveness(ctx, c, KindLocal, true)).To(Succeed())
@@ -150,10 +150,24 @@ var _ = Describe("DefaultProber", func() {
Expect(p.Inference(ctx, c, KindLocal, true)).To(HaveOccurred())
})
It("passes warm liveness for a target that is not loaded, and leaves recovery unconfirmed", func() {
// The warm preload loads the model; a probe that loaded it too would
// block until the load finished and then judge it on an expired ctx.
asked := 0
p = NewProber(func(config.ModelConfig) grpc.Backend { asked++; return nil })
for _, uc := range []string{"chat", "tts"} {
c := config.ModelConfig{Name: "gemma", Backend: "llama-cpp", KnownUsecaseStrings: []string{uc}}
c.KnownUsecases = config.GetUsecasesFromYAML(c.KnownUsecaseStrings)
Expect(p.Liveness(ctx, c, KindLocal, true)).To(Succeed())
Expect(p.Inference(ctx, c, KindLocal, true)).To(MatchError(ErrNotLoaded), uc)
}
Expect(asked).To(Equal(4))
})
It("passes cold local liveness without a model file and without loading", func() {
p = NewProber(func(context.Context, config.ModelConfig) (grpc.Backend, error) {
Fail("cold liveness must not load the model")
return nil, nil
p = NewProber(func(config.ModelConfig) grpc.Backend {
Fail("cold liveness must not look up the backend")
return nil
})
// None of these files exist: a missing file says nothing about whether
// the target can serve (download on first use, dotted names, backends
+40 -7
View File
@@ -2,7 +2,7 @@ package failover
import (
"context"
"sync"
"errors"
"time"
"github.com/mudler/LocalAI/core/config"
@@ -26,23 +26,30 @@ func (m *Manager) Run(ctx context.Context) {
}
}
// Tick runs one pass: sync configs, run due probes, recompute chains. It is
// Tick runs one pass: sync configs, start 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).
//
// Tick does not wait for the probes it starts: one slow target (a probe that
// hangs until its timeout) must not delay probing and fail-back of every
// other chain. Each probe applies its own result, and a target whose probe is
// still running is skipped until it ends.
func (m *Manager) Tick(ctx context.Context) {
m.Sync()
var wg sync.WaitGroup
for _, j := range m.dueProbes() {
wg.Add(1)
m.probes.Add(1)
go func(j probeJob) {
defer wg.Done()
defer m.probes.Done()
m.runProbe(ctx, j)
}(j)
}
wg.Wait()
m.Reevaluate()
}
// waitProbes waits for the probes started so far. Tests use it to see a
// tick's results; the scheduler never waits.
func (m *Manager) waitProbes() { m.probes.Wait() }
type probeJob struct {
target string
cfg config.ModelConfig
@@ -58,6 +65,9 @@ func (m *Manager) dueProbes() []probeJob {
now := m.clock.Now()
var jobs []probeJob
for _, ts := range m.targets {
if ts.probing {
continue
}
interval := ts.params.ProbeInterval()
inference := false
switch ts.state {
@@ -94,6 +104,7 @@ func (m *Manager) dueProbes() []probeJob {
continue
}
ts.lastProbe = now
ts.probing = true
jobs = append(jobs, probeJob{
target: ts.name, cfg: cfg, kind: ts.kind, warm: ts.warm,
inference: inference, timeout: ts.params.ProbeTimeout(),
@@ -112,16 +123,38 @@ func (m *Manager) runProbe(ctx context.Context, j probeJob) {
err = m.prober.Liveness(pctx, j.cfg, j.kind, j.warm)
}
if ctx.Err() != nil {
m.endProbe(j.target)
return // shutting down: a cancelled probe says nothing about the target
}
m.applyProbe(j, err)
}
func (m *Manager) endProbe(target string) {
m.mu.Lock()
defer m.mu.Unlock()
if ts := m.targets[target]; ts != nil {
ts.probing = false
}
}
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 {
if ts == nil {
return
}
ts.probing = false
if ts.state == StateMissing {
return
}
if errors.Is(err, ErrNotLoaded) {
// Nothing running to confirm recovery against: judge the target like
// a cold one, by real requests once min_dwell has passed.
if ts.state == StateRecovering && m.clock.Now().Sub(ts.downSince) >= ts.params.MinDwell() {
m.setTargetLocked(ts, StateHealthy, ReasonRecovery, "")
m.recomputeForLocked(ts.name)
}
return
}
if err != nil {
+94 -17
View File
@@ -2,6 +2,7 @@ package failover
import (
"context"
"slices"
"sync"
"time"
@@ -18,14 +19,19 @@ type probeCall struct {
type fakeProber struct {
mu sync.Mutex
calls []probeCall
fail map[string]error // target -> error returned by every probe
fail map[string]error // target -> error returned by every probe
block map[string]chan struct{} // target -> probes wait until it is closed
}
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]
err, block := p.fail[target], p.block[target]
p.mu.Unlock()
if block != nil {
<-block
}
return err
}
func (p *fakeProber) Liveness(_ context.Context, c config.ModelConfig, _ Kind, _ bool) error {
return p.record(c.Name, false)
@@ -52,33 +58,40 @@ var _ = Describe("Manager probes", func() {
BeforeEach(func() {
clock = newFakeClock()
prober = &fakeProber{fail: map[string]error{}}
prober = &fakeProber{fail: map[string]error{}, block: map[string]chan struct{}{}}
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() {
// tick runs one scheduler pass and waits for the probes it started, so
// each spec sees their results.
tick := func() {
m.Tick(ctx)
m.waitProbes()
}
It("probes idle targets on the first tick and not again before the interval", func() {
tick()
Expect(prober.take()).To(ConsistOf(probeCall{"a", false}, probeCall{"b", false}))
clock.Advance(5 * time.Second)
m.Tick(ctx)
tick()
Expect(prober.take()).To(BeEmpty())
})
It("skips the liveness probe for a target with recent traffic", func() {
m.Tick(ctx)
tick()
prober.take()
clock.Advance(14 * time.Second)
m.ReportSuccess("a")
clock.Advance(2 * time.Second)
m.Tick(ctx)
tick()
Expect(prober.take()).To(ConsistOf(probeCall{"b", false}))
})
It("trips a target whose liveness probe fails", func() {
prober.fail["a"] = errBoom
m.Tick(ctx)
tick()
st, _ := m.ChainStatus("chain")
Expect(st.Targets[0].State).To(Equal(StateDown))
Expect(st.Active).To(Equal("b"))
@@ -86,18 +99,18 @@ var _ = Describe("Manager probes", func() {
It("recovers through liveness, then inference probes, then fails back after dwell", func() {
prober.fail["a"] = errBoom
m.Tick(ctx)
tick()
delete(prober.fail, "a")
prober.take()
clock.Advance(15 * time.Second)
m.Tick(ctx) // liveness passes: recovering
tick() // 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)
tick()
}
calls := prober.take()
Expect(calls).To(ContainElement(probeCall{"a", true}))
@@ -109,10 +122,10 @@ var _ = Describe("Manager probes", func() {
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
tick() // liveness passes: recovering
prober.fail["a"] = errBoom
clock.Advance(15 * time.Second)
m.Tick(ctx)
tick()
st, _ := m.ChainStatus("chain")
Expect(st.Targets[0].State).To(Equal(StateDown))
})
@@ -123,21 +136,21 @@ var _ = Describe("Manager probes", func() {
m.ReportFailure("cold", errBoom)
prober.take()
clock.Advance(30 * time.Second)
m.Tick(ctx)
tick()
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)
tick()
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)
tick()
calls := prober.take()
n := 0
for _, c := range calls {
@@ -148,6 +161,70 @@ var _ = Describe("Manager probes", func() {
Expect(n).To(Equal(1))
})
It("does not hold other targets' probes behind a slow one", func() {
src.Put(remote("c"))
src.Put(chainCfg("chain2", nil, t("c"), warmT("b")))
release := make(chan struct{})
DeferCleanup(func() { close(release); m.waitProbes() })
prober.block["a"] = release
m.ReportFailure("c", errBoom)
prober.take()
m.Tick(ctx) // a hangs; c's liveness still runs and starts its recovery
Eventually(func() TargetState {
st, _ := m.ChainStatus("chain2")
return st.Targets[0].State
}).Should(Equal(StateRecovering))
clock.Advance(15 * time.Second)
m.Tick(ctx) // a is still in flight: no second probe for it
Eventually(func() []probeCall { return prober.take() }).Should(ContainElement(probeCall{"c", true}))
})
It("does not probe a target again while its probe is in flight", func() {
release := make(chan struct{})
prober.block["a"] = release
m.Tick(ctx)
Eventually(func() []probeCall {
prober.mu.Lock()
defer prober.mu.Unlock()
return slices.Clone(prober.calls)
}).Should(ContainElement(probeCall{"a", false}))
for range 3 {
clock.Advance(15 * time.Second)
m.Tick(ctx)
}
close(release)
m.waitProbes()
n := 0
for _, c := range prober.take() {
if c.target == "a" {
n++
}
}
Expect(n).To(Equal(1))
})
It("restores a warm target that is not loaded like a cold one, after min_dwell", func() {
m.ReportFailure("b", errBoom)
prober.fail["b"] = ErrNotLoaded
prober.take()
clock.Advance(15 * time.Second)
prober.fail["b"] = nil
tick() // liveness passes: recovering
st, _ := m.ChainStatus("chain")
Expect(st.Targets[1].State).To(Equal(StateRecovering))
prober.fail["b"] = ErrNotLoaded
clock.Advance(15 * time.Second)
tick() // nothing to confirm against yet, and no trip
st, _ = m.ChainStatus("chain")
Expect(st.Targets[1].State).To(Equal(StateRecovering))
clock.Advance(31 * time.Second)
tick()
st, _ = m.ChainStatus("chain")
Expect(st.Targets[1].State).To(Equal(StateHealthy))
})
It("closes subscriptions when Run stops", func() {
events, _ := m.Subscribe(1)
rctx, cancel := context.WithCancel(ctx)
+1 -1
View File
@@ -87,7 +87,7 @@ targets were down.
| Target | Regular check | Check before moving back |
|---|---|---|
| Remote (`cloud-proxy`) | `GET /v1/models` on the upstream lists the model | one small real request, for example a 1-token completion |
| Local, `warm: true` | the backend answers a health check | one small real request |
| Local, `warm: true` | the backend answers a health check. A check never loads the model: while it is not loaded, the check passes and real requests judge it | one small real request. While the model is not loaded, the target is used again after `min_dwell` |
| Local, not warm | none: judged only by real requests; it is never loaded only to check it | none: the target is used again after `min_dwell` |
A request that succeeds counts as a check, so a busy target is almost never
@@ -151,7 +151,7 @@ Chain states:
| Target | Liveness (steady state) | Recovery confirmation |
|---|---|---|
| remote | `GET <base>/v1/models` returns 2xx and lists the upstream model. `<base>` is the scheme and host of `proxy.upstream_url` plus any path prefix before `/v1`. The upstream model is `proxy.upstream_model`, or the target name when it is empty. `/v1/models` works on any OpenAI-compatible upstream, and `/readyz` exists only on LocalAI. | one minimal real request, chosen by usecase |
| local, `warm: true` | gRPC `HealthCheck` on the loaded backend. If the backend is not loaded (it crashed), a reload is the recovery attempt. | chat and completion: `Predict` with 1 token; embeddings: `Embedding` of `"ping"`; other usecases: `HealthCheck`. A local backend process that answers `HealthCheck` rarely fails only for TTS or transcription. |
| local, `warm: true` | gRPC `HealthCheck` on the loaded backend, with the probe timeout. A probe never loads the model: when the backend is not loaded (the warm preload is still loading it, or a crash removed it), liveness passes and the next real request loads and judges it. | chat and completion: `Predict` with 1 token; embeddings: `Embedding` of `"ping"`; other usecases: `HealthCheck`. A local backend process that answers `HealthCheck` rarely fails only for TTS or transcription. When the backend is not loaded there is nothing to confirm against: the probe neither passes nor trips, and the target returns to `healthy` like a cold one, when `min_dwell` has passed since the trip. |
| local, cold | none: a cold target is judged only by real requests; it is never loaded only to probe it. | none. After a trip, the target returns to `healthy` when `min_dwell` has passed. The next real request is the test. |
Minimal requests by usecase:
@@ -165,8 +165,10 @@ Minimal requests by usecase:
Probe load rules:
- Each chain has one ticker with jitter. A target shared by chains is probed
once.
- One global scheduler ticks every second, without jitter, and starts the
probes that are due. A target shared by chains is probed once. The
scheduler does not wait for a probe: a target whose probe is still running
is skipped, so one slow target does not delay the others.
- A successful real request counts as a liveness pass, so a busy target is
almost never probed.
- Inference probes run only while a target is `recovering`.