fix(distributed): configurable remote model-load timeout, and reap the load when it times out (#10948)

* 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 <mudler@localai.io>

* 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 <mudler@localai.io>
Assisted-by: Claude:claude-opus-4-8 golangci-lint

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Co-authored-by: Ettore Di Giacinto <mudler@localai.io>

Author: localai-bot

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

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

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

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)

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,

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"
@@ -72,9 +75,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 +151,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 +170,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 +193,7 @@ func NewSmartRouter(registry ModelRouter, opts SmartRouterOptions) *SmartRouter
 		pressure:         opts.Pressure,
 		sharedModels:     opts.SharedModels,
 		modelLoadCeiling: ceiling,
+		modelLoadTimeout: loadTimeout,
 	}
 }
 
@@ -250,11 +291,21 @@ 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)
 		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 {
@@ -285,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.