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

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.