diff --git a/core/application/application.go b/core/application/application.go index 83057c9cd..05a4ed619 100644 --- a/core/application/application.go +++ b/core/application/application.go @@ -98,6 +98,11 @@ type Application struct { func newApplication(appConfig *config.ApplicationConfig) *Application { ml := model.NewModelLoader(appConfig.SystemState) + // Apply the per-model load-failure cooldown (0 disables). Set here rather + // than in the watchdog block so it takes effect regardless of whether the + // watchdog/LRU limiter is enabled. + ml.SetLoadFailureCooldown(appConfig.ModelLoadFailureCooldown, 0) + // Close MCP sessions when a model is unloaded (watchdog eviction, manual shutdown, etc.) ml.OnModelUnload(func(modelName string) { mcpTools.CloseMCPSessions(modelName) diff --git a/core/cli/run.go b/core/cli/run.go index 85f82e5d5..542d441c1 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -96,6 +96,7 @@ type RunCMD struct { SizeAwareEviction bool `env:"LOCALAI_SIZE_AWARE_EVICTION,SIZE_AWARE_EVICTION" default:"false" help:"Evict the largest loaded model first rather than the least-recently-used one, keeping small utility models resident and maximizing freed memory per eviction" group:"backends"` LRUEvictionMaxRetries int `env:"LOCALAI_LRU_EVICTION_MAX_RETRIES,LRU_EVICTION_MAX_RETRIES" default:"30" help:"Maximum number of retries when waiting for busy models to become idle before eviction (default: 30)" group:"backends"` LRUEvictionRetryInterval string `env:"LOCALAI_LRU_EVICTION_RETRY_INTERVAL,LRU_EVICTION_RETRY_INTERVAL" default:"1s" help:"Interval between retries when waiting for busy models to become idle (e.g., 1s, 2s) (default: 1s)" group:"backends"` + ModelLoadFailureCooldown string `env:"LOCALAI_MODEL_LOAD_FAILURE_COOLDOWN,MODEL_LOAD_FAILURE_COOLDOWN" default:"10s" help:"After a model load fails, refuse new load attempts for that model for this long (returned as HTTP 503 + Retry-After) so a client polling a broken model doesn't respawn a crashing backend every request. Doubles per consecutive failure up to 5m; reset on success. Set to 0 to disable (e.g., 10s, 30s)" group:"backends"` Federated bool `env:"LOCALAI_FEDERATED,FEDERATED" help:"Enable federated instance" group:"federated"` DisableGalleryEndpoint bool `env:"LOCALAI_DISABLE_GALLERY_ENDPOINT,DISABLE_GALLERY_ENDPOINT" help:"Disable the gallery endpoints" group:"api"` DisableMCP bool `env:"LOCALAI_DISABLE_MCP,DISABLE_MCP" help:"Disable MCP (Model Context Protocol) support" group:"api" default:"false"` @@ -614,6 +615,13 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { } opts = append(opts, config.WithLRUEvictionRetryInterval(dur)) } + if r.ModelLoadFailureCooldown != "" { + dur, err := time.ParseDuration(r.ModelLoadFailureCooldown) + if err != nil { + return fmt.Errorf("invalid model load failure cooldown: %w", err) + } + opts = append(opts, config.WithModelLoadFailureCooldown(dur)) + } // Handle Open Responses store TTL if r.OpenResponsesStoreTTL != "" && r.OpenResponsesStoreTTL != "0" { diff --git a/core/config/application_config.go b/core/config/application_config.go index 8e4bb08f5..042c9d826 100644 --- a/core/config/application_config.go +++ b/core/config/application_config.go @@ -131,6 +131,13 @@ type ApplicationConfig struct { LRUEvictionMaxRetries int // Maximum number of retries when waiting for busy models to become idle (default: 30) LRUEvictionRetryInterval time.Duration // Interval between retries when waiting for busy models (default: 1s) + // ModelLoadFailureCooldown is the base cooldown applied after a model load + // fails: new load attempts for that model are refused (HTTP 503 + Retry-After) + // until it elapses, doubling per consecutive failure up to a 5m cap and reset + // on success. Prevents a client polling a broken model from respawning a + // crashing backend on every request. 0 disables it. Default: 10s. + ModelLoadFailureCooldown time.Duration + ModelsURL []string WatchDogBusyTimeout, WatchDogIdleTimeout time.Duration @@ -245,6 +252,7 @@ func NewApplicationConfig(o ...AppOption) *ApplicationConfig { AgentJobRetentionDays: 30, // Default: 30 days LRUEvictionMaxRetries: 30, // Default: 30 retries LRUEvictionRetryInterval: 1 * time.Second, // Default: 1 second + ModelLoadFailureCooldown: 10 * time.Second, // Default: 10s base cooldown after a failed load // WatchDogInterval is intentionally left at the zero value here. // The startup loader applies a persisted runtime_settings.json value // only when the interval is still 0 (its "not set by env var" @@ -531,6 +539,17 @@ func WithLRUEvictionRetryInterval(interval time.Duration) AppOption { } } +// WithModelLoadFailureCooldown sets the base cooldown applied after a failed +// model load. 0 disables the cooldown, so unlike most options it accepts any +// non-negative value. +func WithModelLoadFailureCooldown(cooldown time.Duration) AppOption { + return func(o *ApplicationConfig) { + if cooldown >= 0 { + o.ModelLoadFailureCooldown = cooldown + } + } +} + var EnableGalleriesAutoload = func(o *ApplicationConfig) { o.AutoloadGalleries = true } diff --git a/core/http/app.go b/core/http/app.go index fff8a3468..11661ea07 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -5,16 +5,20 @@ import ( "errors" "fmt" "io/fs" + "math" "mime" "net/http" "os" "path/filepath" + "strconv" "strings" "time" "github.com/labstack/echo/v4" "github.com/labstack/echo/v4/middleware" + "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/core/http/auth" "github.com/mudler/LocalAI/core/http/endpoints/localai" @@ -45,6 +49,23 @@ var reactUI embed.FS var quietPaths = []string{"/api/operations", "/api/resources", "/healthz", "/readyz"} +// applyModelLoadCooldown maps a ModelLoadCooldownError anywhere in err's chain +// to HTTP 503 with a Retry-After header (whole seconds, floor 1), so a client +// polling a model whose load recently failed backs off instead of triggering a +// fresh backend start. If err is not a cooldown error, code is returned as-is. +func applyModelLoadCooldown(err error, code int, c echo.Context) int { + var coolErr *model.ModelLoadCooldownError + if !errors.As(err, &coolErr) { + return code + } + secs := int(math.Ceil(coolErr.RetryAfter.Seconds())) + if secs < 1 { + secs = 1 + } + c.Response().Header().Set("Retry-After", strconv.Itoa(secs)) + return http.StatusServiceUnavailable +} + // @title LocalAI API // @version 2.0.0 // @description The LocalAI Rest API. @@ -110,6 +131,7 @@ func API(application *application.Application) (*echo.Echo, error) { if errors.As(err, &he) { code = he.Code } + code = applyModelLoadCooldown(err, code, c) // Handle 404 errors: serve React SPA for HTML requests, JSON otherwise if code == http.StatusNotFound { @@ -137,6 +159,7 @@ func API(application *application.Application) (*echo.Echo, error) { if errors.As(err, &he) { code = he.Code } + code = applyModelLoadCooldown(err, code, c) c.NoContent(code) } } diff --git a/docs/content/reference/cli-reference.md b/docs/content/reference/cli-reference.md index 3e919efb8..50ba18705 100644 --- a/docs/content/reference/cli-reference.md +++ b/docs/content/reference/cli-reference.md @@ -50,6 +50,7 @@ Complete reference for all LocalAI command-line interface (CLI) parameters and e | `--force-eviction-when-busy` | `false` | Force eviction even when models have active API calls (default: false for safety). **Warning:** Enabling this can interrupt active requests | `$LOCALAI_FORCE_EVICTION_WHEN_BUSY`, `$FORCE_EVICTION_WHEN_BUSY` | | `--lru-eviction-max-retries` | `30` | Maximum number of retries when waiting for busy models to become idle before eviction | `$LOCALAI_LRU_EVICTION_MAX_RETRIES`, `$LRU_EVICTION_MAX_RETRIES` | | `--lru-eviction-retry-interval` | `1s` | Interval between retries when waiting for busy models to become idle (e.g., `1s`, `2s`) | `$LOCALAI_LRU_EVICTION_RETRY_INTERVAL`, `$LRU_EVICTION_RETRY_INTERVAL` | +| `--model-load-failure-cooldown` | `10s` | After a model load fails, refuse new load attempts for that model for this long (HTTP 503 + `Retry-After`) so a client polling a broken model doesn't respawn a crashing backend every request. Doubles per consecutive failure up to 5m; reset on success. `0` disables | `$LOCALAI_MODEL_LOAD_FAILURE_COOLDOWN`, `$MODEL_LOAD_FAILURE_COOLDOWN` | For more information on VRAM management, see [VRAM and Memory Management]({{%relref "advanced/vram-management" %}}). diff --git a/pkg/model/loader.go b/pkg/model/loader.go index b4018a979..17420f3e2 100644 --- a/pkg/model/loader.go +++ b/pkg/model/loader.go @@ -82,6 +82,33 @@ type ModelLoader struct { // the exit code can't, since a child killed by our own SIGTERM/SIGKILL // reports -1, indistinguishable from a signal-induced crash. stoppingProcs sync.Map + // loadFailures records, per modelID, the cooldown window applied after a + // failed load so that a client repeatedly polling a broken model does not + // spawn (and leak) a fresh backend process on every request. Guarded by mu. + loadFailures map[string]*loadFailureState + loadFailureBaseCooldown time.Duration // first cooldown after a failure + loadFailureMaxCooldown time.Duration // cap for the exponential backoff +} + +// loadFailureState tracks consecutive load failures for a single modelID and +// the instant at which its cooldown window expires. +type loadFailureState struct { + consecutive int + cooldownUntil time.Time +} + +// ModelLoadCooldownError is returned when a model load is skipped because a +// recent attempt failed and the per-model cooldown window has not yet elapsed. +// The HTTP layer maps it to 503 with a Retry-After header so a polling client +// backs off instead of triggering a fresh backend start on every request. +type ModelLoadCooldownError struct { + ModelID string + RetryAfter time.Duration +} + +func (e *ModelLoadCooldownError) Error() string { + return fmt.Sprintf("model %q load is in cooldown after a recent failure; retry after %s", + e.ModelID, e.RetryAfter.Round(time.Second)) } // NewModelLoader creates a new ModelLoader instance. @@ -95,11 +122,79 @@ func NewModelLoader(system *system.SystemState) *ModelLoader { lruEvictionMaxRetries: 30, // Default: 30 retries lruEvictionRetryInterval: 1 * time.Second, // Default: 1 second backendLogs: NewBackendLogStore(1000), + loadFailures: make(map[string]*loadFailureState), + loadFailureBaseCooldown: 10 * time.Second, // Default: 10s after the first failure + loadFailureMaxCooldown: 5 * time.Minute, // Default cap for the backoff } return nml } +// SetLoadFailureCooldown configures the per-model load-failure backoff. base is +// the cooldown applied after the first failure; it doubles on each consecutive +// failure up to max. base is authoritative (base <= 0 disables the cooldown +// entirely); max only overrides the current cap when > 0. +func (ml *ModelLoader) SetLoadFailureCooldown(base, max time.Duration) { + ml.mu.Lock() + defer ml.mu.Unlock() + if base < 0 { + base = 0 + } + ml.loadFailureBaseCooldown = base + if max > 0 { + ml.loadFailureMaxCooldown = max + } + if ml.loadFailureMaxCooldown < ml.loadFailureBaseCooldown { + ml.loadFailureMaxCooldown = ml.loadFailureBaseCooldown + } +} + +// cooldownRemaining returns how long the modelID's load cooldown still has to +// run, or 0 if there is none. Callers must hold ml.mu. +func (ml *ModelLoader) cooldownRemaining(modelID string) time.Duration { + st := ml.loadFailures[modelID] + if st == nil { + return 0 + } + if remaining := time.Until(st.cooldownUntil); remaining > 0 { + return remaining + } + return 0 +} + +// recordLoadFailure grows the modelID's consecutive-failure count and arms the +// next cooldown window using exponential backoff capped at loadFailureMaxCooldown. +func (ml *ModelLoader) recordLoadFailure(modelID string) { + ml.mu.Lock() + defer ml.mu.Unlock() + if ml.loadFailureBaseCooldown <= 0 { + return // cooldown disabled + } + st := ml.loadFailures[modelID] + if st == nil { + st = &loadFailureState{} + ml.loadFailures[modelID] = st + } + st.consecutive++ + // base * 2^(consecutive-1), clamped. Cap the shift to avoid overflowing + // the Duration; anything past the cap collapses to loadFailureMaxCooldown. + shift := st.consecutive - 1 + if shift > 20 { + shift = 20 + } + backoff := ml.loadFailureBaseCooldown * (1 << shift) + if backoff <= 0 || backoff > ml.loadFailureMaxCooldown { + backoff = ml.loadFailureMaxCooldown + } + st.cooldownUntil = time.Now().Add(backoff) +} + +// clearLoadFailure resets the modelID's failure state after a successful load. +// Callers must hold ml.mu. +func (ml *ModelLoader) clearLoadFailure(modelID string) { + delete(ml.loadFailures, modelID) +} + // GetLoadingCount returns the number of models currently being loaded func (ml *ModelLoader) GetLoadingCount() int { ml.mu.Lock() @@ -302,6 +397,14 @@ func (ml *ModelLoader) ListLoadedModels() []*Model { } func (ml *ModelLoader) LoadModel(modelID, modelName string, loader func(string, string, string) (*Model, error)) (*Model, error) { + return ml.loadModel(modelID, modelName, loader, true) +} + +// loadModel is the implementation behind LoadModel. checkCooldown gates fresh, +// independent load triggers behind the per-model failure cooldown; it is set to +// false for the coalesced retry of an in-flight burst (a follower whose leader +// just failed), which is not a new trigger and should still get its one retry. +func (ml *ModelLoader) loadModel(modelID, modelName string, loader func(string, string, string) (*Model, error), checkCooldown bool) (*Model, error) { ml.mu.Lock() distributed := ml.modelRouter != nil ml.mu.Unlock() @@ -353,6 +456,19 @@ func (ml *ModelLoader) LoadModel(modelID, modelName string, loader func(string, return model, nil } + // If a recent load attempt for this model failed, short-circuit fresh load + // triggers until the cooldown elapses. This stops a client that keeps + // polling a broken model from spawning (and leaking) a new backend process + // on every request. The coalesced follower-retry below passes + // checkCooldown=false so an in-flight burst still gets its one retry. + if checkCooldown { + if retryAfter := ml.cooldownRemaining(modelID); retryAfter > 0 { + ml.mu.Unlock() + xlog.Debug("Model load in cooldown after a recent failure", "modelID", modelID, "retryAfter", retryAfter) + return nil, &ModelLoadCooldownError{ModelID: modelID, RetryAfter: retryAfter} + } + } + // Check if another goroutine is already loading this model if loadingChan, isLoading := ml.loading[modelID]; isLoading { ml.mu.Unlock() @@ -366,8 +482,9 @@ func (ml *ModelLoader) LoadModel(modelID, modelName string, loader func(string, if model != nil { return model, nil } - // If still not loaded, the other goroutine failed - we'll try again - return ml.LoadModel(modelID, modelName, loader) + // If still not loaded, the other goroutine failed. Retry once as part of + // this burst, bypassing the cooldown gate (we are not a new trigger). + return ml.loadModel(modelID, modelName, loader, false) } // Mark this model as loading (create a channel that will be closed when done) @@ -389,15 +506,18 @@ func (ml *ModelLoader) LoadModel(modelID, modelName string, loader func(string, model, err := loader(modelID, modelName, modelFile) if err != nil { + ml.recordLoadFailure(modelID) return nil, fmt.Errorf("failed to load model with internal loader: %s", err) } if model == nil { + ml.recordLoadFailure(modelID) return nil, fmt.Errorf("loader didn't return a model") } - // Add to models map + // Add to models map and reset any prior failure cooldown for this model. ml.mu.Lock() + ml.clearLoadFailure(modelID) ml.store.Set(modelID, model) ml.mu.Unlock() diff --git a/pkg/model/loader_test.go b/pkg/model/loader_test.go index bb03c8942..e6a6148c3 100644 --- a/pkg/model/loader_test.go +++ b/pkg/model/loader_test.go @@ -255,6 +255,11 @@ var _ = Describe("ModelLoader", func() { mockLoader := func(modelID, modelName, modelFile string) (*model.Model, error) { count := atomic.AddInt32(&attemptCount, 1) if count == 1 { + // Hold the loading slot so the second request coalesces as a + // follower before this leader fails. That follower then gets + // the one in-burst retry, which bypasses the failure cooldown + // (the cooldown only gates fresh, independent load triggers). + time.Sleep(50 * time.Millisecond) return nil, errors.New("first attempt fails") } return model.NewModel(modelID, modelName, nil), nil @@ -270,7 +275,7 @@ var _ = Describe("ModelLoader", func() { m1, err1 = modelLoader.LoadModel("retry-model", "test.model", mockLoader) }) - // Give first goroutine a head start + // Give first goroutine a head start so it owns the loading slot. time.Sleep(10 * time.Millisecond) wg.Go(func() { @@ -305,4 +310,90 @@ var _ = Describe("ModelLoader", func() { Expect(modelLoader).ToNot(BeNil()) }) }) + + Context("Load failure cooldown", func() { + It("refuses a fresh load within the cooldown window without re-invoking the loader", func() { + modelLoader.SetLoadFailureCooldown(60*time.Millisecond, 240*time.Millisecond) + + var loadCount int32 + failing := func(modelID, modelName, modelFile string) (*model.Model, error) { + atomic.AddInt32(&loadCount, 1) + return nil, errors.New("boom") + } + + // First attempt runs the loader and fails (not a cooldown error). + _, err := modelLoader.LoadModel("broken", "test.model", failing) + Expect(err).To(HaveOccurred()) + var coolErr *model.ModelLoadCooldownError + Expect(errors.As(err, &coolErr)).To(BeFalse()) + Expect(atomic.LoadInt32(&loadCount)).To(Equal(int32(1))) + + // An immediate retry is short-circuited: cooldown error, loader untouched. + _, err = modelLoader.LoadModel("broken", "test.model", failing) + Expect(errors.As(err, &coolErr)).To(BeTrue()) + Expect(coolErr.ModelID).To(Equal("broken")) + Expect(coolErr.RetryAfter).To(BeNumerically(">", time.Duration(0))) + Expect(atomic.LoadInt32(&loadCount)).To(Equal(int32(1))) + + // Once the window elapses the loader is attempted again. + Eventually(func() int32 { + _, _ = modelLoader.LoadModel("broken", "test.model", failing) + return atomic.LoadInt32(&loadCount) + }, "1s", "20ms").Should(BeNumerically(">=", 2)) + }) + + It("clears the cooldown after a successful load", func() { + modelLoader.SetLoadFailureCooldown(60*time.Millisecond, 240*time.Millisecond) + + var attempts int32 + loader := func(modelID, modelName, modelFile string) (*model.Model, error) { + if atomic.AddInt32(&attempts, 1) == 1 { + return nil, errors.New("boom") + } + m := model.NewModel(modelID, modelName, nil) + m.MarkHealthy() + return m, nil + } + + _, err := modelLoader.LoadModel("flaky", "test.model", loader) + Expect(err).To(HaveOccurred()) + + // After the window, the retry succeeds and resets the failure state. + var m *model.Model + Eventually(func() error { + m, err = modelLoader.LoadModel("flaky", "test.model", loader) + return err + }, "1s", "20ms").Should(Succeed()) + Expect(m).ToNot(BeNil()) + + // A subsequent load returns the cached model, never a cooldown error. + m2, err := modelLoader.LoadModel("flaky", "test.model", loader) + Expect(err).ToNot(HaveOccurred()) + Expect(m2).To(Equal(m)) + }) + + It("grows the cooldown on consecutive failures", func() { + modelLoader.SetLoadFailureCooldown(50*time.Millisecond, 10*time.Second) + failing := func(modelID, modelName, modelFile string) (*model.Model, error) { + return nil, errors.New("boom") + } + + // Failure 1, then read its cooldown. + _, err := modelLoader.LoadModel("bad", "test.model", failing) + Expect(err).To(HaveOccurred()) + _, err = modelLoader.LoadModel("bad", "test.model", failing) + var c1 *model.ModelLoadCooldownError + Expect(errors.As(err, &c1)).To(BeTrue()) + + // Wait out the first window, trigger failure 2, read its (larger) cooldown. + time.Sleep(70 * time.Millisecond) + _, err = modelLoader.LoadModel("bad", "test.model", failing) + Expect(err).To(HaveOccurred()) + _, err = modelLoader.LoadModel("bad", "test.model", failing) + var c2 *model.ModelLoadCooldownError + Expect(errors.As(err, &c2)).To(BeTrue()) + + Expect(c2.RetryAfter).To(BeNumerically(">", c1.RetryAfter)) + }) + }) })