From c7b44e9b20075e0bddf98ed069a732687ba914a2 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Sat, 18 Jul 2026 22:57:44 +0000 Subject: [PATCH 1/2] fix(distributed): make the remote LoadModel deadline configurable The router hardcoded a 5 minute gRPC deadline for the remote LoadModel call. Staging finishes before the timer starts, so those five minutes cover only the worker backend's own checkpoint load and pipeline init. A cold load of meituan-longcat/LongCat-Video-Avatar-1.5 (~83 GB) on an ARM64 Thor worker fails at exactly 302s with DeadlineExceeded while the backend process is still making progress (CPU time accumulating, RSS moving as weights are mapped), so the load was cut short rather than wedged. Add LOCALAI_NATS_MODEL_LOAD_TIMEOUT / --model-load-timeout mirroring the existing backend-install timeout knob, defaulting to 5m so unset clusters keep today's behaviour. The cold-load hold ceiling (which bounds how long one load may hold the per-model advisory lock) was derived from the install timeout alone, so raising the load deadline past it would have been silently clipped. Derive it from both budgets via ModelLoadCeilingFor: max(install + load + 5m staging margin, 25m) With the defaults that is 15m + 5m + 5m = 25m, identical to the previous constant, and the 25m floor means shrinking either budget can never tighten the ceiling below what clusters relied on before. Assisted-by: Claude:claude-opus-4-8 golangci-lint Signed-off-by: Ettore Di Giacinto --- core/application/distributed.go | 16 ++- core/cli/run.go | 28 +++- core/cli/run_distributed_timeouts_test.go | 23 ++++ core/config/distributed_config.go | 20 +++ core/config/distributed_config_test.go | 25 ++++ core/services/nodes/router.go | 60 +++++++-- .../nodes/router_load_timeout_test.go | 127 ++++++++++++++++++ docs/content/features/distributed-mode.md | 3 + 8 files changed, 281 insertions(+), 21 deletions(-) create mode 100644 core/cli/run_distributed_timeouts_test.go create mode 100644 core/services/nodes/router_load_timeout_test.go diff --git a/core/application/distributed.go b/core/application/distributed.go index 40cc9fff3dfa..0519d96e787e 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -356,12 +356,16 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade PrefixConfig: prefixCfg, Pressure: pressure, SharedModels: cfg.Distributed.SharedModels, - // Cap how long a cold load may hold the per-model advisory lock: the - // configured backend.install deadline plus a margin for file staging and - // the remote LoadModel. Derived from the install timeout so raising it - // (for slow links pulling multi-GB images) widens the ceiling too, - // instead of letting the static default cut a legitimately slow load. - ModelLoadCeiling: cfg.Distributed.BackendInstallTimeoutOrDefault() + 10*time.Minute, + ModelLoadTimeout: cfg.Distributed.ModelLoadTimeoutOrDefault(), + // Cap how long a cold load may hold the per-model advisory lock. Derived + // from BOTH configured budgets it has to cover, so raising either the + // install timeout (slow links pulling multi-GB images) or the model load + // timeout (very large checkpoints) widens the ceiling too, instead of + // letting a stale bound cut a legitimately slow load short. + ModelLoadCeiling: nodes.ModelLoadCeilingFor( + cfg.Distributed.BackendInstallTimeoutOrDefault(), + cfg.Distributed.ModelLoadTimeoutOrDefault(), + ), }) // Wire staging-progress broadcasting so file-staging shows up on every diff --git a/core/cli/run.go b/core/cli/run.go index 6c4b1c5a16f1..659977cc1dd0 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -173,6 +173,7 @@ type RunCMD struct { DistributedPrefixCacheTTL string `env:"LOCALAI_DISTRIBUTED_PREFIX_CACHE_TTL" help:"Idle-timeout for prefix-cache index entries; also drives the background eviction cadence (every TTL/2). Default 5m." group:"distributed"` BackendInstallTimeout string `env:"LOCALAI_NATS_BACKEND_INSTALL_TIMEOUT" help:"NATS round-trip timeout for backend.install requests sent to worker nodes (default 15m). Increase for slow links pulling multi-GB images." group:"distributed"` BackendUpgradeTimeout string `env:"LOCALAI_NATS_BACKEND_UPGRADE_TIMEOUT" help:"NATS round-trip timeout for backend.upgrade requests (default 15m)." group:"distributed"` + ModelLoadTimeout string `env:"LOCALAI_NATS_MODEL_LOAD_TIMEOUT" help:"gRPC deadline for the remote LoadModel call sent to a worker node once its backend is installed and model files are staged (default 5m). Increase for very large checkpoints (multi-tens-of-GB diffusion/video models) whose load and pipeline init exceed 5 minutes. Raising it also widens the cold-load lock ceiling." group:"distributed"` NatsAccountSeed string `env:"LOCALAI_NATS_ACCOUNT_SEED" help:"NATS account signing seed (SU...) used to mint per-node worker JWTs at registration" group:"distributed"` NatsServiceJWT string `env:"LOCALAI_NATS_SERVICE_JWT" help:"NATS user JWT for the frontend (and agent workers) to publish control-plane messages" group:"distributed"` NatsServiceSeed string `env:"LOCALAI_NATS_SERVICE_SEED" help:"NATS user signing seed (SU...) paired with LOCALAI_NATS_SERVICE_JWT" group:"distributed"` @@ -219,6 +220,18 @@ func DefaultUploadPath() string { return filepath.Join(userScopedTempDir(), "upload") } +// parseDistributedDuration parses an operator-supplied duration for a +// distributed timeout knob. The env var and the rejected value are both named +// in the error: these knobs are usually set from a compose file or a K8s +// manifest, where a bare "invalid duration" gives the operator nothing to grep. +func parseDistributedDuration(envVar, raw string) (time.Duration, error) { + d, err := time.ParseDuration(raw) + if err != nil { + return 0, fmt.Errorf("invalid %s %q: %w", envVar, raw, err) + } + return d, nil +} + func (r *RunCMD) Run(ctx *cliContext.Context) error { warnDeprecatedFlags() @@ -326,19 +339,26 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { opts = append(opts, config.WithStorageSecretKey(r.StorageSecretKey)) } if r.BackendInstallTimeout != "" { - d, err := time.ParseDuration(r.BackendInstallTimeout) + d, err := parseDistributedDuration("LOCALAI_NATS_BACKEND_INSTALL_TIMEOUT", r.BackendInstallTimeout) if err != nil { - return fmt.Errorf("invalid LOCALAI_NATS_BACKEND_INSTALL_TIMEOUT %q: %w", r.BackendInstallTimeout, err) + return err } opts = append(opts, config.WithBackendInstallTimeout(d)) } if r.BackendUpgradeTimeout != "" { - d, err := time.ParseDuration(r.BackendUpgradeTimeout) + d, err := parseDistributedDuration("LOCALAI_NATS_BACKEND_UPGRADE_TIMEOUT", r.BackendUpgradeTimeout) if err != nil { - return fmt.Errorf("invalid LOCALAI_NATS_BACKEND_UPGRADE_TIMEOUT %q: %w", r.BackendUpgradeTimeout, err) + return err } opts = append(opts, config.WithBackendUpgradeTimeout(d)) } + if r.ModelLoadTimeout != "" { + d, err := parseDistributedDuration("LOCALAI_NATS_MODEL_LOAD_TIMEOUT", r.ModelLoadTimeout) + if err != nil { + return err + } + opts = append(opts, config.WithModelLoadTimeout(d)) + } if r.RegistrationToken != "" { opts = append(opts, config.WithRegistrationToken(r.RegistrationToken)) } diff --git a/core/cli/run_distributed_timeouts_test.go b/core/cli/run_distributed_timeouts_test.go new file mode 100644 index 000000000000..5f54c9efc77d --- /dev/null +++ b/core/cli/run_distributed_timeouts_test.go @@ -0,0 +1,23 @@ +package cli + +import ( + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("parseDistributedDuration", func() { + It("parses a valid operator-supplied duration", func() { + d, err := parseDistributedDuration("LOCALAI_NATS_MODEL_LOAD_TIMEOUT", "45m") + Expect(err).ToNot(HaveOccurred()) + Expect(d).To(Equal(45 * time.Minute)) + }) + + It("names the env var and the bad value so a typo is actionable", func() { + _, err := parseDistributedDuration("LOCALAI_NATS_MODEL_LOAD_TIMEOUT", "45mins") + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("LOCALAI_NATS_MODEL_LOAD_TIMEOUT")) + Expect(err.Error()).To(ContainSubstring("45mins")) + }) +}) diff --git a/core/config/distributed_config.go b/core/config/distributed_config.go index 7cb0a7c57ea0..f486ba196d17 100644 --- a/core/config/distributed_config.go +++ b/core/config/distributed_config.go @@ -76,6 +76,12 @@ type DistributedConfig struct { BackendInstallTimeout time.Duration // NATS round-trip timeout for backend.install (default 15m) BackendUpgradeTimeout time.Duration // NATS round-trip timeout for backend.upgrade (default 15m) + // ModelLoadTimeout is the gRPC deadline for the remote LoadModel call the + // router issues once a worker has the backend installed and the model files + // staged. It therefore covers only the backend's own checkpoint load and + // pipeline init, which for a multi-tens-of-GB diffusion/video checkpoint on + // unified memory can far exceed the 5m default. + ModelLoadTimeout time.Duration // gRPC deadline for remote LoadModel (default 5m) MaxUploadSize int64 // Maximum upload body size in bytes (default 50 GB) @@ -142,6 +148,7 @@ func (c DistributedConfig) Validate() error { FlagMCPCIJobTimeout: c.MCPCIJobTimeout, FlagBackendInstallTimeout: c.BackendInstallTimeout, FlagBackendUpgradeTimeout: c.BackendUpgradeTimeout, + FlagModelLoadTimeout: c.ModelLoadTimeout, } { if d < 0 { return fmt.Errorf("%s must not be negative", name) @@ -286,6 +293,12 @@ func WithBackendUpgradeTimeout(d time.Duration) AppOption { } } +func WithModelLoadTimeout(d time.Duration) AppOption { + return func(o *ApplicationConfig) { + o.Distributed.ModelLoadTimeout = d + } +} + var EnableAutoApproveNodes = func(o *ApplicationConfig) { o.Distributed.AutoApproveNodes = true } @@ -342,6 +355,7 @@ const ( FlagMCPCIJobTimeout = "mcp-ci-job-timeout" FlagBackendInstallTimeout = "backend-install-timeout" FlagBackendUpgradeTimeout = "backend-upgrade-timeout" + FlagModelLoadTimeout = "model-load-timeout" ) // Defaults for distributed timeouts. @@ -355,6 +369,7 @@ const ( DefaultMCPCIJobTimeout = 10 * time.Minute DefaultBackendInstallTimeout = 15 * time.Minute DefaultBackendUpgradeTimeout = 15 * time.Minute + DefaultModelLoadTimeout = 5 * time.Minute ) // DefaultMaxUploadSize is the default maximum upload body size (50 GB). @@ -408,6 +423,11 @@ func (c DistributedConfig) BackendUpgradeTimeoutOrDefault() time.Duration { return cmp.Or(c.BackendUpgradeTimeout, DefaultBackendUpgradeTimeout) } +// ModelLoadTimeoutOrDefault returns the configured timeout or the default. +func (c DistributedConfig) ModelLoadTimeoutOrDefault() time.Duration { + return cmp.Or(c.ModelLoadTimeout, DefaultModelLoadTimeout) +} + // MCPToolTimeoutOrDefault returns the configured timeout or the default. func (c DistributedConfig) MCPToolTimeoutOrDefault() time.Duration { return cmp.Or(c.MCPToolTimeout, DefaultMCPToolTimeout) diff --git a/core/config/distributed_config_test.go b/core/config/distributed_config_test.go index 377c397dee95..ec7fbe8dc7e9 100644 --- a/core/config/distributed_config_test.go +++ b/core/config/distributed_config_test.go @@ -33,6 +33,18 @@ var _ = Describe("DistributedConfig backend NATS timeouts", func() { Expect(c.BackendUpgradeTimeoutOrDefault()).To(Equal(30 * time.Minute)) }) }) + + Context("ModelLoadTimeoutOrDefault", func() { + It("returns 5 minutes when unset so existing clusters keep today's behaviour", func() { + c := config.DistributedConfig{} + Expect(c.ModelLoadTimeoutOrDefault()).To(Equal(5 * time.Minute)) + }) + + It("returns the configured value when set", func() { + c := config.DistributedConfig{ModelLoadTimeout: 45 * time.Minute} + Expect(c.ModelLoadTimeoutOrDefault()).To(Equal(45 * time.Minute)) + }) + }) }) var _ = Describe("DistributedConfig flag-name constants", func() { @@ -53,6 +65,7 @@ var _ = Describe("DistributedConfig flag-name constants", func() { Entry("MCP CI job timeout", config.FlagMCPCIJobTimeout, "mcp-ci-job-timeout"), Entry("backend install timeout", config.FlagBackendInstallTimeout, "backend-install-timeout"), Entry("backend upgrade timeout", config.FlagBackendUpgradeTimeout, "backend-upgrade-timeout"), + Entry("model load timeout", config.FlagModelLoadTimeout, "model-load-timeout"), ) }) @@ -80,6 +93,18 @@ var _ = Describe("DistributedConfig.Validate negative-duration errors", func() { Expect(err.Error()).To(ContainSubstring(config.FlagBackendUpgradeTimeout)) }) + It("rejects a negative ModelLoadTimeout with the flag name in the error", func() { + c := config.DistributedConfig{ + Enabled: true, + NatsURL: "nats://localhost:4222", + ModelLoadTimeout: -1 * time.Second, + } + err := c.Validate() + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(config.FlagModelLoadTimeout)) + Expect(err.Error()).To(ContainSubstring("must not be negative")) + }) + It("accepts all-zero durations as valid (defaults apply)", func() { c := config.DistributedConfig{ Enabled: true, diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index 01100123401a..ce3805320132 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -72,9 +72,46 @@ type SmartRouterOptions struct { // attempt (node selection -> backend install -> file staging -> LoadModel) // may run while holding the per-model advisory lock. It backstops every // sub-step's own timeout so a wedged worker can never pin the lock - and - // every other replica's request for that model - indefinitely. Zero selects - // defaultModelLoadCeiling. + // every other replica's request for that model - indefinitely. Zero derives + // it from the default install budget and ModelLoadTimeout via + // ModelLoadCeilingFor. ModelLoadCeiling time.Duration + // ModelLoadTimeout is the gRPC deadline for the remote LoadModel call, which + // runs after the backend install and file staging have already completed and + // so covers only the worker backend's checkpoint load and pipeline init. + // Zero selects config.DefaultModelLoadTimeout. Operators raise it via + // LOCALAI_NATS_MODEL_LOAD_TIMEOUT for very large checkpoints; ModelLoadCeiling + // must be widened in step (see ModelLoadCeilingFor) or the ceiling clips the + // longer deadline before it can be used. + ModelLoadTimeout time.Duration +} + +// modelLoadStagingMargin is the slack ModelLoadCeilingFor adds on top of the +// install + remote-load budgets to cover the steps that have no deadline of +// their own: node selection, replica allocation and model file staging. +const modelLoadStagingMargin = 5 * time.Minute + +// minModelLoadCeiling floors the derived ceiling at the historical 25m constant +// so shrinking the install or load budgets can never tighten the hold ceiling +// below what clusters relied on before it became derived. +const minModelLoadCeiling = 25 * time.Minute + +// ModelLoadCeilingFor derives the cold-load hold ceiling from the budgets it has +// to cover. The ceiling bounds how long one cold load may hold the per-model +// advisory lock; it must comfortably exceed the slowest legitimate load (backend +// install + staging + remote LoadModel) so it only ever fires when a step is +// genuinely wedged - e.g. a worker that died mid-install. Deriving it means +// raising LOCALAI_NATS_MODEL_LOAD_TIMEOUT for an 80 GB video checkpoint actually +// takes effect, instead of being cut short by a constant that silently went +// stale. Non-positive inputs fall back to their package defaults. +func ModelLoadCeilingFor(installTimeout, loadTimeout time.Duration) time.Duration { + if installTimeout <= 0 { + installTimeout = config.DefaultBackendInstallTimeout + } + if loadTimeout <= 0 { + loadTimeout = config.DefaultModelLoadTimeout + } + return max(installTimeout+loadTimeout+modelLoadStagingMargin, minModelLoadCeiling) } // SmartRouter routes inference requests to the best available backend node. @@ -111,15 +148,11 @@ type SmartRouter struct { // modelLoadCeiling bounds how long a cold load may hold the per-model // advisory lock (see SmartRouterOptions.ModelLoadCeiling). modelLoadCeiling time.Duration + // modelLoadTimeout is the deadline for the remote LoadModel gRPC call + // (see SmartRouterOptions.ModelLoadTimeout). + modelLoadTimeout time.Duration } -// defaultModelLoadCeiling is the fallback hold ceiling for a cold model load. -// It must comfortably exceed the slowest legitimate load - a multi-GB backend -// install (DefaultBackendInstallTimeout, 15m) plus staging and the remote -// LoadModel (5m) - so it never cuts a real load short; it only ever fires when -// a step is genuinely wedged (e.g. a worker that died mid-install). -const defaultModelLoadCeiling = 25 * time.Minute - // probeCacheTTL is how long a successful gRPC HealthCheck on a backend is // trusted before the next request re-probes. Matches healthCheckTTL in // pkg/model/model.go so the single-process and distributed paths share a @@ -134,9 +167,13 @@ func NewSmartRouter(registry ModelRouter, opts SmartRouterOptions) *SmartRouter if factory == nil { factory = &tokenClientFactory{token: opts.AuthToken} } + loadTimeout := opts.ModelLoadTimeout + if loadTimeout <= 0 { + loadTimeout = config.DefaultModelLoadTimeout + } ceiling := opts.ModelLoadCeiling if ceiling <= 0 { - ceiling = defaultModelLoadCeiling + ceiling = ModelLoadCeilingFor(config.DefaultBackendInstallTimeout, loadTimeout) } return &SmartRouter{ registry: registry, @@ -153,6 +190,7 @@ func NewSmartRouter(registry ModelRouter, opts SmartRouterOptions) *SmartRouter pressure: opts.Pressure, sharedModels: opts.SharedModels, modelLoadCeiling: ceiling, + modelLoadTimeout: loadTimeout, } } @@ -250,7 +288,7 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking if loadOpts != nil { xlog.Info("Loading model on remote node", "node", node.Name, "model", modelName, "addr", backendAddr) - loadCtx, cancel := context.WithTimeout(ctx, 5*time.Minute) + loadCtx, cancel := context.WithTimeout(ctx, r.modelLoadTimeout) defer cancel() res, err := client.LoadModel(loadCtx, loadOpts) diff --git a/core/services/nodes/router_load_timeout_test.go b/core/services/nodes/router_load_timeout_test.go new file mode 100644 index 000000000000..d6ba09fa0b51 --- /dev/null +++ b/core/services/nodes/router_load_timeout_test.go @@ -0,0 +1,127 @@ +package nodes + +import ( + "context" + "errors" + "sync" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/services/messaging" + grpc "github.com/mudler/LocalAI/pkg/grpc" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + ggrpc "google.golang.org/grpc" +) + +// deadlineBackend records the deadline of the context handed to LoadModel so a +// spec can assert what budget the remote gRPC load actually got. A hardcoded +// deadline in scheduleAndLoad is invisible to every other assertion in this +// package — only the context the client sees proves it. +type deadlineBackend struct { + grpc.Backend // embedded: unused methods panic if called + + mu sync.Mutex + deadline time.Time + hadOne bool +} + +func (b *deadlineBackend) HealthCheck(_ context.Context) (bool, error) { return true, nil } + +func (b *deadlineBackend) IsBusy() bool { return false } + +func (b *deadlineBackend) LoadModel(ctx context.Context, _ *pb.ModelOptions, _ ...ggrpc.CallOption) (*pb.Result, error) { + b.mu.Lock() + defer b.mu.Unlock() + b.deadline, b.hadOne = ctx.Deadline() + return &pb.Result{Success: true}, nil +} + +// budget returns how much time LoadModel was given, rounded to the nearest +// second so scheduling jitter between the router's WithTimeout and the stub's +// clock read doesn't make the assertion flaky. +func (b *deadlineBackend) budget() time.Duration { + b.mu.Lock() + defer b.mu.Unlock() + Expect(b.hadOne).To(BeTrue(), "LoadModel must be called with a deadline") + return time.Until(b.deadline).Round(time.Second) +} + +type deadlineClientFactory struct{ client *deadlineBackend } + +func (f *deadlineClientFactory) NewClient(_ string, _ bool) grpc.Backend { return f.client } + +var _ = Describe("remote LoadModel deadline", func() { + var ( + reg *fakeModelRouter + backend *deadlineBackend + factory *deadlineClientFactory + unloader *fakeUnloader + ) + + BeforeEach(func() { + reg = &fakeModelRouter{findAndLockErr: errors.New("not found")} + reg.findIdleNode = &BackendNode{ID: "n1", Name: "worker", Address: "10.0.0.1:50051"} + backend = &deadlineBackend{} + factory = &deadlineClientFactory{client: backend} + unloader = &fakeUnloader{ + installReply: &messaging.BackendInstallReply{Success: true, Address: "10.0.0.1:9001"}, + } + }) + + route := func(router *SmartRouter) { + _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/big.gguf"}, false) + Expect(err).ToNot(HaveOccurred()) + } + + It("gives LoadModel the configured budget instead of the hardcoded 5 minutes", func() { + router := NewSmartRouter(reg, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + ModelLoadTimeout: 45 * time.Minute, + // Ceiling must not be the thing that clips the deadline here. + ModelLoadCeiling: 2 * time.Hour, + }) + route(router) + Expect(backend.budget()).To(Equal(45 * time.Minute)) + }) + + It("falls back to 5 minutes when no load timeout is configured", func() { + router := NewSmartRouter(reg, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + }) + route(router) + Expect(backend.budget()).To(Equal(5 * time.Minute)) + }) +}) + +var _ = Describe("ModelLoadCeilingFor", func() { + // The ceiling bounds how long a cold load may hold the per-model advisory + // lock. It is derived from the sub-step budgets it must cover, so raising + // the load timeout for a huge checkpoint cannot be silently clipped by a + // stale constant. + It("keeps today's 25m behaviour at the default install + load budgets", func() { + Expect(ModelLoadCeilingFor(config.DefaultBackendInstallTimeout, config.DefaultModelLoadTimeout)). + To(Equal(25 * time.Minute)) + }) + + It("widens when the load timeout is raised", func() { + Expect(ModelLoadCeilingFor(15*time.Minute, 45*time.Minute)).To(Equal(65 * time.Minute)) + }) + + It("widens when the install timeout is raised", func() { + Expect(ModelLoadCeilingFor(40*time.Minute, 5*time.Minute)).To(Equal(50 * time.Minute)) + }) + + It("never drops below the 25m floor when the budgets are shrunk", func() { + Expect(ModelLoadCeilingFor(1*time.Minute, 1*time.Minute)).To(Equal(25 * time.Minute)) + }) + + It("treats non-positive budgets as their defaults", func() { + Expect(ModelLoadCeilingFor(0, 0)).To(Equal(25 * time.Minute)) + }) +}) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 6dfddff2db33..cdeea890e115 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -72,8 +72,11 @@ The frontend is a standard LocalAI instance with distributed mode enabled. These | `--auth-database-url` | `LOCALAI_AUTH_DATABASE_URL` | *(required)* | PostgreSQL connection URL | | `--backend-install-timeout` | `LOCALAI_NATS_BACKEND_INSTALL_TIMEOUT` | `15m` | How long the frontend waits for a worker to acknowledge a backend install before considering the request stalled. Raise it when workers pull large backend images over slow links. If a worker takes longer than this, the operation shows as "still installing in background" in the admin UI and clears once the worker finishes. | | `--backend-upgrade-timeout` | `LOCALAI_NATS_BACKEND_UPGRADE_TIMEOUT` | `15m` | Same as the install timeout, applied to backend upgrades (force-reinstall). | +| `--model-load-timeout` | `LOCALAI_NATS_MODEL_LOAD_TIMEOUT` | `5m` | Deadline for the `LoadModel` gRPC call the frontend issues to a worker. It starts *after* the backend is installed and the model files are staged, so it covers only the worker backend's own checkpoint load and pipeline init. Raise it for very large checkpoints: a multi-tens-of-GB diffusion or video model on unified memory routinely needs more than 5 minutes, and the load then fails with `rpc error: code = DeadlineExceeded` at exactly the timeout. Raising this also widens the cold-load lock ceiling below, so the two knobs never fight. | | `--expose-node-header` | `LOCALAI_EXPOSE_NODE_HEADER` | `false` | When enabled, inference responses carry an `X-LocalAI-Node` header with the ID of the worker node that served the request. Coverage spans the OpenAI-compatible endpoints (chat completions, completions, embeddings, audio transcriptions, audio speech / TTS, image generations, image inpainting), the Jina rerank endpoint (`/v1/rerank`), the VAD endpoints (`/v1/vad`, `/vad`), and the Anthropic Messages (`/v1/messages`) and Ollama (`/api/chat`, `/api/generate`, `/api/embed`) shims. Useful for debugging, observability and load-balancer attribution. Off by default: the node ID reveals internal cluster topology and should not be exposed on a public endpoint. Best-effort: under heavy concurrency for the same model across multiple replicas, the header may reflect a recent routing decision rather than this exact request's. Acceptable for observability and debugging. | +The router also bounds how long a single cold load may hold the per-model advisory lock, so a worker that dies mid-install cannot pin every other replica's request for that model. That ceiling is **derived**, not configured: `max(backend-install-timeout + model-load-timeout + 5m, 25m)`. With the defaults that is `15m + 5m + 5m = 25m`, matching previous releases; raising either timeout widens the ceiling in step, so a longer load deadline is never clipped by it. + ### NATS JWT authentication (recommended for production) By default, NATS connections are anonymous: any client that can reach port `4222` may publish control-plane subjects such as `nodes..backend.install`. Enable JWT auth to scope workers to their own node subjects and give the frontend a dedicated service credential. From 35a330b46044ef3b342c28469d70f0fb8608f2eb Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Sun, 19 Jul 2026 00:37:30 +0000 Subject: [PATCH 2/2] fix(distributed): reap the abandoned replica when a remote load times out The gRPC deadline on the remote LoadModel call only cancels the client side. A backend blocked in a synchronous weight load never observes its cancelled handler context, so when scheduleAndLoad gave up it left the worker loading with nobody waiting for the result. Observed on an ARM64 Thor worker loading LongCat-Video-Avatar-1.5: the client returned DeadlineExceeded at 302s, and the backend process was still alive 30 minutes later having pulled ~57GB from HuggingFace. Every retry stacked another multi-GB loader on the worker; they had to be reaped by hand via POST /api/nodes/:id/models/unload. Send backend.stop for the exact `modelID#replicaIndex` process key we just abandoned. The exact key matters: a bare model ID stops every replica on that node, including healthy ones serving traffic. Only a deadline or cancellation triggers the reap. Any other LoadModel failure is the backend answering, which means its handler returned and the process is idle - stopping it there would discard a warm process and its downloaded weights. The reap is best-effort and never replaces the load error the caller is waiting on. The `modelID#replicaIndex` format was already hand-rolled in two places (the worker's buildProcessKey and pkg/model's log store). Rather than add a third, export model.BackendProcessKey from pkg/model, the lowest common dependency of both sides. Signed-off-by: Ettore Di Giacinto Assisted-by: Claude:claude-opus-4-8 golangci-lint --- core/services/nodes/router.go | 58 +++++++++ core/services/nodes/router_reap_load_test.go | 126 +++++++++++++++++++ pkg/model/backend_log_store.go | 12 ++ 3 files changed, 196 insertions(+) create mode 100644 core/services/nodes/router_reap_load_test.go diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index ce3805320132..145880c7f174 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -18,9 +18,12 @@ import ( "github.com/mudler/LocalAI/pkg/distributedhdr" grpc "github.com/mudler/LocalAI/pkg/grpc" pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/vram" "github.com/mudler/xlog" "golang.org/x/sync/singleflight" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" "google.golang.org/protobuf/proto" "gorm.io/gorm" "gorm.io/gorm/clause" @@ -293,6 +296,16 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking res, err := client.LoadModel(loadCtx, loadOpts) if err != nil { + // A gRPC deadline only cancels the CLIENT side of the call. A + // backend blocked in a synchronous weight load never observes its + // cancelled handler context, so it keeps loading with nobody + // waiting: an 83GB checkpoint was seen still downloading 30 + // minutes past the client timeout, and each retry stacked another + // multi-GB loader process on the worker. Reap the replica we just + // abandoned before handing the failure back. + if loadAbandonedOnWorker(err) { + r.reapAbandonedLoad(node, trackingKey, replicaIndex) + } return nil, fmt.Errorf("loading model %s on node %s: %w", modelName, node.Name, err) } if !res.Success { @@ -323,6 +336,51 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking return &scheduleLoadResult{Node: node, Client: client, BackendAddr: backendAddr, ReplicaIndex: replicaIndex}, nil } +// loadAbandonedOnWorker reports whether a failed remote LoadModel left the +// worker process still running the load. +// +// Only a deadline or a cancellation qualifies: those are OUR side giving up +// while the backend keeps going. Every other failure (an unsupported model, a +// bad option, an OOM) is the backend answering, which means its handler +// returned and the process is idle — stopping it there would throw away a +// warm process and its downloaded weights for the next attempt. +func loadAbandonedOnWorker(err error) bool { + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { + return true + } + // The deadline usually reaches us as a gRPC status rather than a wrapped + // context error. errors.As (not status.Code) so a wrapped status is still + // recognised. + var st interface{ GRPCStatus() *status.Status } + if errors.As(err, &st) { + switch st.GRPCStatus().Code() { + case codes.DeadlineExceeded, codes.Canceled: + return true + } + } + return false +} + +// reapAbandonedLoad tells the worker to stop the one replica whose load we +// just abandoned. The exact `modelID#replicaIndex` process key matters: the +// bare model ID would stop every replica of the model on that node, including +// healthy ones serving traffic. +// +// Best-effort by design — the caller is waiting on the load failure, and a +// failed reap must be visible in the logs, never substituted for it. +func (r *SmartRouter) reapAbandonedLoad(node *BackendNode, trackingKey string, replicaIndex int) { + if r.unloader == nil { + return + } + processKey := model.BackendProcessKey(trackingKey, replicaIndex) + xlog.Warn("Reaping abandoned model load on worker", + "node", node.Name, "model", trackingKey, "replica", replicaIndex) + if err := r.unloader.StopBackend(node.ID, processKey); err != nil { + xlog.Warn("Failed to reap abandoned model load; the worker may still be loading", + "node", node.Name, "backend", processKey, "error", err) + } +} + // ScheduleAndLoadModel implements ModelScheduler for the reconciler. // It retrieves stored model options from an existing replica and performs the // full load sequence (stage files, LoadModel, SetNodeModel) on a new node. diff --git a/core/services/nodes/router_reap_load_test.go b/core/services/nodes/router_reap_load_test.go new file mode 100644 index 000000000000..4c8e7cf8bb45 --- /dev/null +++ b/core/services/nodes/router_reap_load_test.go @@ -0,0 +1,126 @@ +package nodes + +import ( + "context" + "errors" + "sync" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/messaging" + grpc "github.com/mudler/LocalAI/pkg/grpc" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + ggrpc "google.golang.org/grpc" +) + +// failingLoadBackend returns a canned error from LoadModel so a spec can drive +// scheduleAndLoad down its failure path. The gRPC deadline that fires in +// production only cancels the client side of the call, so the error the router +// sees is the only signal that a worker was left loading. +type failingLoadBackend struct { + grpc.Backend // embedded: unused methods panic if called + + mu sync.Mutex + err error + seen int +} + +func (b *failingLoadBackend) HealthCheck(_ context.Context) (bool, error) { return true, nil } + +func (b *failingLoadBackend) IsBusy() bool { return false } + +func (b *failingLoadBackend) LoadModel(_ context.Context, _ *pb.ModelOptions, _ ...ggrpc.CallOption) (*pb.Result, error) { + b.mu.Lock() + defer b.mu.Unlock() + b.seen++ + return nil, b.err +} + +type failingClientFactory struct{ client *failingLoadBackend } + +func (f *failingClientFactory) NewClient(_ string, _ bool) grpc.Backend { return f.client } + +// replicaSlotRouter pins the replica slot scheduleAndLoad allocates so a spec +// can assert the reaped process key carries the real index, not a hardcoded 0. +type replicaSlotRouter struct { + *fakeModelRouter + replica int +} + +func (r *replicaSlotRouter) NextFreeReplicaIndex(_ context.Context, _, _ string, _ int) (int, error) { + return r.replica, nil +} + +var _ = Describe("reaping an abandoned remote load", func() { + var ( + reg *replicaSlotRouter + backend *failingLoadBackend + unloader *fakeUnloader + ) + + BeforeEach(func() { + base := &fakeModelRouter{findAndLockErr: errors.New("not found")} + base.findIdleNode = &BackendNode{ID: "n1", Name: "worker", Address: "10.0.0.1:50051", MaxReplicasPerModel: 4} + reg = &replicaSlotRouter{fakeModelRouter: base, replica: 2} + backend = &failingLoadBackend{} + unloader = &fakeUnloader{ + installReply: &messaging.BackendInstallReply{Success: true, Address: "10.0.0.1:9001"}, + } + }) + + route := func() error { + router := NewSmartRouter(reg, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: &failingClientFactory{client: backend}, + ModelLoadTimeout: time.Minute, + ModelLoadCeiling: time.Hour, + }) + _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/big.gguf"}, false) + return err + } + + It("stops exactly the abandoned replica when LoadModel hits its deadline", func() { + backend.err = status.Error(codes.DeadlineExceeded, "context deadline exceeded") + + err := route() + + Expect(err).To(HaveOccurred()) + Expect(unloader.stopCalls).To(ConsistOf("n1:big-model#2"), + "the abandoned replica must be reaped by its exact worker process key") + }) + + It("still returns the load failure to the caller, not the stop result", func() { + backend.err = status.Error(codes.DeadlineExceeded, "context deadline exceeded") + unloader.stopErr = errors.New("nats publish failed") + + err := route() + + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("context deadline exceeded")) + Expect(err.Error()).ToNot(ContainSubstring("nats publish failed"), + "a failed reap must be logged, never substituted for the load error") + }) + + It("reaps when the deadline surfaces as a wrapped context error", func() { + backend.err = errors.Join(errors.New("loading model"), context.DeadlineExceeded) + + Expect(route()).To(HaveOccurred()) + Expect(unloader.stopCalls).To(ConsistOf("n1:big-model#2")) + }) + + It("does not reap when the backend rejected the model outright", func() { + // A clean application error means the handler returned: the worker is + // idle, not stuck mid-load, and stopping it would throw away a usable + // process (and its downloaded weights) for the next attempt. + backend.err = status.Error(codes.InvalidArgument, "unsupported model architecture") + + Expect(route()).To(HaveOccurred()) + Expect(unloader.stopCalls).To(BeEmpty()) + }) +}) diff --git a/pkg/model/backend_log_store.go b/pkg/model/backend_log_store.go index 808b0d75c831..a0988790ff9e 100644 --- a/pkg/model/backend_log_store.go +++ b/pkg/model/backend_log_store.go @@ -2,6 +2,7 @@ package model import ( "sort" + "strconv" "strings" "sync" "time" @@ -15,6 +16,17 @@ import ( // package free of CLI imports. const replicaSeparator = "#" +// BackendProcessKey builds the worker supervisor's process key for one replica +// of a model (e.g. "qwen3-0.6b#0"). The worker keys its process map by this +// string and resolves an exact key to exactly that replica, so any caller that +// wants to address a single replica over NATS must format it identically — +// a mismatch degrades silently into "no process found" and leaks the process. +// This package is the lowest common dependency of the router and the worker, +// so the format lives here instead of being hand-rolled at each call site. +func BackendProcessKey(modelID string, replicaIndex int) string { + return modelID + replicaSeparator + strconv.Itoa(replicaIndex) +} + // BackendLogLine represents a single line of output from a backend process. type BackendLogLine struct { Timestamp time.Time `json:"timestamp"`