fix(distributed): scale the remote model-load deadline with checkpoint size

core/application/distributed.go

@@ -356,7 +356,12 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
 		PrefixConfig:     prefixCfg,
 		Pressure:         pressure,
 		SharedModels:     cfg.Distributed.SharedModels,
-		ModelLoadTimeout: cfg.Distributed.ModelLoadTimeoutOrDefault(),
+		// RAW, not OrDefault: zero means "derive the budget per model from the
+		// checkpoint size" (config.ModelLoadTimeoutForSize), which is what makes
+		// a 70 GB video checkpoint work without the operator first hitting a
+		// DeadlineExceeded and going looking for a knob. A non-zero value here is
+		// an explicit override and is used verbatim.
+		ModelLoadTimeout: cfg.Distributed.ModelLoadTimeout,
 		// 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

core/cli/run.go

@@ -173,7 +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"`
+	ModelLoadTimeout          string `env:"LOCALAI_NATS_MODEL_LOAD_TIMEOUT" help:"Fixed gRPC deadline for the remote LoadModel call sent to a worker node once its backend is installed and model files are staged. Unset (the default), the deadline is derived from the checkpoint size instead: 5m plus 20s per GiB, capped at 6h, so multi-tens-of-GB diffusion/video checkpoints get the minutes they need without a fixed cliff. Set this only to pin a specific budget; the value is used verbatim, including when it is shorter than the derived one." 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"`

core/config/model_load_budget.go

@@ -0,0 +1,78 @@
+package config
+
+import "time"
+
+// The remote LoadModel deadline used to be the fixed DefaultModelLoadTimeout.
+// That is a model-size cliff, not a timeout: the deadline starts only after the
+// backend install and file staging have finished, so it covers the worker's
+// checkpoint read and pipeline init alone — work whose duration is proportional
+// to the bytes on disk. A 70 GB video checkpoint on a Jetson Thor worker
+// therefore failed reproducibly (953.5s wall clock, ~11m of it install and
+// staging, then DeadlineExceeded on a load that never had a chance).
+//
+// Raising the constant does not fix that: it moves the cliff to the next larger
+// model, and it makes a genuinely wedged SMALL model hang for the whole inflated
+// duration before anyone notices. So the budget is derived from the size instead.
+const (
+	// ModelLoadTimeoutPerGiB is the budget granted per GiB of checkpoint on top
+	// of DefaultModelLoadTimeout.
+	//
+	// It is deliberately generous. This is a timeout: erring long costs only
+	// failure LATENCY on a load that was going to fail anyway, while erring
+	// short costs a guaranteed FALSE failure on a load that was healthy. 20s/GiB
+	// corresponds to reading weights at ~54 MB/s, which is below what any
+	// supported medium sustains (NVMe is orders of magnitude faster; the slow end
+	// is an eMMC or SD-backed Jetson, or a checkpoint faulted in over a network
+	// filesystem) and so leaves headroom for the dequantisation and pipeline init
+	// that follow the read. Hardware and quantisation both move the real figure
+	// by an order of magnitude, which is exactly why the constant sits at the
+	// pessimistic end rather than at a measured average.
+	//
+	// Worked examples: 2 GiB -> 5m40s (a wedged small model still fails fast),
+	// 70 GiB -> 28m20s (the production checkpoint that failed at 5m),
+	// 600 GiB -> 3h25m (the size the cluster has to support).
+	ModelLoadTimeoutPerGiB = 20 * time.Second
+
+	// MaxModelLoadTimeout caps the derived budget so a nonsense size (a corrupted
+	// stat, a future 10 TB artifact) cannot hand out an effectively infinite
+	// deadline. At ModelLoadTimeoutPerGiB the cap binds only above ~1 TiB,
+	// comfortably past the 600 GB requirement.
+	MaxModelLoadTimeout = 6 * time.Hour
+)
+
+// bytesPerGiB is the divisor for the per-GiB rate above.
+const bytesPerGiB int64 = 1 << 30
+
+// ModelLoadTimeoutForSize derives the gRPC deadline for the remote LoadModel
+// call from the checkpoint's on-disk size.
+//
+// A non-positive size means the frontend could not measure the payload — a
+// backend given a bare HuggingFace repo id fetches its own weights on the
+// worker, so there is nothing local to stat. Rather than guess, that case keeps
+// the historical DefaultModelLoadTimeout, which is also the floor of the derived
+// range: size only ever adds budget.
+//
+// An explicit LOCALAI_NATS_MODEL_LOAD_TIMEOUT always wins over this; see
+// SmartRouterOptions.ModelLoadTimeout.
+func ModelLoadTimeoutForSize(bytes int64) time.Duration {
+	if bytes <= 0 {
+		return DefaultModelLoadTimeout
+	}
+
+	// Split into whole GiB plus remainder before scaling: multiplying a 600 GB
+	// byte count by a 20e9-nanosecond rate overflows int64 long before the cap
+	// could clamp it.
+	gib := bytes / bytesPerGiB
+	remainder := bytes % bytesPerGiB
+
+	extra := time.Duration(gib) * ModelLoadTimeoutPerGiB
+	extra += time.Duration(int64(ModelLoadTimeoutPerGiB) * remainder / bytesPerGiB)
+
+	// A large enough gib still overflows the multiply above into a negative
+	// duration; treat any non-positive extra as "past the cap".
+	if extra <= 0 {
+		return MaxModelLoadTimeout
+	}
+
+	return min(DefaultModelLoadTimeout+extra, MaxModelLoadTimeout)
+}

core/config/model_load_budget_test.go

@@ -0,0 +1,53 @@
+package config_test
+
+import (
+	"time"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	"github.com/mudler/LocalAI/core/config"
+)
+
+const gib int64 = 1 << 30
+
+var _ = Describe("ModelLoadTimeoutForSize", func() {
+	// The remote LoadModel deadline used to be a fixed 5m. That is a model-size
+	// cliff: a 70 GB video checkpoint on a Jetson Thor worker failed
+	// reproducibly with DeadlineExceeded because its weight load and pipeline
+	// init alone exceed 5 minutes. Raising the constant only moves the cliff, so
+	// the budget is derived from the bytes the worker has to read.
+
+	It("keeps a small checkpoint close to the historical 5m default", func() {
+		// A wedged 2 GB model must still fail fast: inflating every load's
+		// budget is a real regression in failure latency.
+		Expect(config.ModelLoadTimeoutForSize(2 * gib)).To(BeNumerically("<", 10*time.Minute))
+	})
+
+	It("gives a 70 GB checkpoint materially more budget than a 2 GB one", func() {
+		small := config.ModelLoadTimeoutForSize(2 * gib)
+		big := config.ModelLoadTimeoutForSize(70 * gib)
+		Expect(big).To(BeNumerically(">", small*3))
+		// The measured production failure had ~5m of load budget and needed
+		// more; anything under 20m would still be a cliff for this exact model.
+		Expect(big).To(BeNumerically(">=", 20*time.Minute))
+	})
+
+	It("scales monotonically with size", func() {
+		Expect(config.ModelLoadTimeoutForSize(600 * gib)).
+			To(BeNumerically(">", config.ModelLoadTimeoutForSize(70*gib)))
+	})
+
+	It("still gives a 600 GB checkpoint hours, not minutes", func() {
+		Expect(config.ModelLoadTimeoutForSize(600 * gib)).To(BeNumerically(">=", 3*time.Hour))
+	})
+
+	It("falls back to the plain default when the size is unknown", func() {
+		Expect(config.ModelLoadTimeoutForSize(0)).To(Equal(config.DefaultModelLoadTimeout))
+		Expect(config.ModelLoadTimeoutForSize(-1)).To(Equal(config.DefaultModelLoadTimeout))
+	})
+
+	It("never exceeds the absolute maximum, however absurd the size", func() {
+		Expect(config.ModelLoadTimeoutForSize(100_000 * gib)).To(Equal(config.MaxModelLoadTimeout))
+	})
+})

core/services/nodes/load_deadline.go

@@ -246,6 +246,44 @@ func (d *loadDeadline) Observe() {
 	d.timer.Reset(time.Until(next))
 }
 
+// Extend pushes the expiry out to at least now+d, for a phase whose duration is
+// known up front but which reports no progress of its own.
+//
+// Observe() cannot serve that case: it grants exactly one stall window, so a
+// step that legitimately runs for half an hour in silence — the remote
+// LoadModel reading a 70 GB checkpoint — is cancelled long before it finishes.
+// Extend states the budget once instead of inferring it from a heartbeat that
+// will never arrive. Like Observe it only ever moves the expiry forward and
+// never past the absolute cap, so it cannot be used to escape the hard bound.
+func (d *loadDeadline) Extend(dur time.Duration) {
+	if d == nil || dur <= 0 {
+		return
+	}
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	if d.stopped {
+		return
+	}
+	next := time.Now().Add(dur)
+	if next.After(d.hardExpiry) {
+		next = d.hardExpiry
+	}
+	if !next.After(d.expiry) {
+		return
+	}
+	d.expiry = next
+	d.timer.Reset(time.Until(next))
+}
+
+// extendLoadDeadline widens the cold-load hold attached to ctx to cover a phase
+// of known duration. A context without a load deadline (single-host paths, tests
+// constructing routers directly) is a no-op.
+func extendLoadDeadline(ctx context.Context, d time.Duration) {
+	if ld, ok := ctx.Value(loadDeadlineKey{}).(*loadDeadline); ok {
+		ld.Extend(d)
+	}
+}
+
 // observeLoadProgress pushes out the cold-load deadline attached to ctx, if any.
 // Callers report byte-level movement here; a context without a load deadline
 // (single-host paths, tests constructing stagers directly) is a no-op.

core/services/nodes/router.go

@@ -79,13 +79,17 @@ type SmartRouterOptions struct {
 	// 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 pins 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 means "derive it per model from the checkpoint size" via
+	// config.ModelLoadTimeoutForSize — the default, because load duration is
+	// proportional to the bytes the worker reads and any fixed value is a
+	// model-size cliff. A non-zero value is an explicit operator override
+	// (LOCALAI_NATS_MODEL_LOAD_TIMEOUT) and always wins over the derived budget,
+	// in BOTH directions: an operator who wants faster failure gets it.
 	ModelLoadTimeout time.Duration
 	// StagingStallWindow is how long file staging may report zero bytes before
 	// the cold load is declared wedged and the advisory lock released. While
@@ -116,6 +120,12 @@ const minModelLoadCeiling = 25 * time.Minute
 // 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.
+//
+// This is the hold's STARTING budget, not its maximum. A per-model budget
+// derived from checkpoint size (config.ModelLoadTimeoutForSize) can exceed any
+// ceiling computed here, so scheduleAndLoad widens the hold as it enters the
+// load phase — see extendLoadDeadline. Without that, a 70 GB checkpoint's ~28m
+// load budget would be cancelled by a 25m ceiling that knew nothing about it.
 func ModelLoadCeilingFor(installTimeout, loadTimeout time.Duration) time.Duration {
 	if installTimeout <= 0 {
 		installTimeout = config.DefaultBackendInstallTimeout
@@ -160,8 +170,9 @@ 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 is the operator's explicit override for the remote
+	// LoadModel deadline, or zero to derive it per model from the checkpoint
+	// size (see SmartRouterOptions.ModelLoadTimeout and loadTimeoutFor).
 	modelLoadTimeout time.Duration
 	// stagingStallWindow and modelLoadAbsoluteMax turn modelLoadCeiling from a
 	// hard countdown into a progress-extended hold (see load_deadline.go).
@@ -183,10 +194,10 @@ func NewSmartRouter(registry ModelRouter, opts SmartRouterOptions) *SmartRouter
 	if factory == nil {
 		factory = &tokenClientFactory{token: opts.AuthToken}
 	}
-	loadTimeout := opts.ModelLoadTimeout
-	if loadTimeout <= 0 {
-		loadTimeout = config.DefaultModelLoadTimeout
-	}
+	// Keep the override RAW: zero has to stay distinguishable from an explicit
+	// 5m, because zero now means "derive per model from the checkpoint size"
+	// while an explicit 5m means "hold every load to 5m".
+	loadTimeout := max(opts.ModelLoadTimeout, 0)
 	ceiling := opts.ModelLoadCeiling
 	if ceiling <= 0 {
 		ceiling = ModelLoadCeilingFor(config.DefaultBackendInstallTimeout, loadTimeout)
@@ -293,6 +304,12 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking
 	// that built modelOpts may have no GPU at all in distributed mode.
 	applyNodeHardwareDefaults(modelOpts, node, backendType)
 
+	// Size the remote load budget BEFORE staging: stageModelFiles rewrites the
+	// path fields to their remote equivalents on a clone, and only the local
+	// paths can be stat'ed here.
+	payloadBytes := modelPayloadBytes(modelOpts)
+	loadTimeout := r.loadTimeoutFor(payloadBytes)
+
 	// Pre-stage model files via FileStager before loading
 	loadOpts := modelOpts
 	if r.fileStager != nil && modelOpts != nil {
@@ -307,13 +324,30 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking
 
 	// Load the model on the remote node
 	if loadOpts != nil {
-		xlog.Info("Loading model on remote node", "node", node.Name, "model", modelName, "addr", backendAddr)
+		xlog.Info("Loading model on remote node", "node", node.Name, "model", modelName, "addr", backendAddr,
+			"payloadBytes", payloadBytes, "loadBudget", loadTimeout)
+
+		// The cold-load hold above this call extends on STAGING progress, and
+		// the remote LoadModel reports none — so once the last byte lands the
+		// hold expires a stall window later and would cancel a load that is
+		// still well inside its own budget (#11026 in miniature). Widen the hold
+		// to cover the load phase before entering it.
+		extendLoadDeadline(ctx, loadTimeout+modelLoadStagingMargin)
 
-		loadCtx, cancel := context.WithTimeout(ctx, r.modelLoadTimeout)
+		loadCtx, cancel := context.WithTimeout(ctx, loadTimeout)
 		defer cancel()
 
 		res, err := client.LoadModel(loadCtx, loadOpts)
 		if err != nil {
+			// A bare "context deadline exceeded" tells the operator nothing
+			// about which of the several budgets in a cold load ran out, and
+			// cost real debugging time in production. Name the budget, the
+			// payload it was derived from, and the knob that overrides it.
+			if errors.Is(err, context.DeadlineExceeded) {
+				err = fmt.Errorf("model load budget of %s for a %s checkpoint exceeded; "+
+					"raise it with LOCALAI_NATS_MODEL_LOAD_TIMEOUT (%s): %w",
+					loadTimeout, humanFileSize(payloadBytes), config.FlagModelLoadTimeout, err)
+			}
 			// 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
@@ -1368,6 +1402,80 @@ func (r *SmartRouter) withStagingCallback(ctx context.Context, trackingKey, file
 	})
 }
 
+// loadTimeoutFor resolves the gRPC deadline for one remote LoadModel call.
+//
+// An explicit LOCALAI_NATS_MODEL_LOAD_TIMEOUT wins outright — including when it
+// is SHORTER than the derived value, because an operator who deliberately wants
+// fast failure must not have a size heuristic silently extend their loads.
+// Otherwise the budget scales with the checkpoint, which is what the worker
+// actually spends its time reading.
+func (r *SmartRouter) loadTimeoutFor(payloadBytes int64) time.Duration {
+	if r.modelLoadTimeout > 0 {
+		return r.modelLoadTimeout
+	}
+	return config.ModelLoadTimeoutForSize(payloadBytes)
+}
+
+// modelPayloadBytes totals the on-disk size of everything the worker will have
+// to read for this model, over the same field set stageModelFiles uploads.
+//
+// Paths that do not exist locally contribute nothing: a backend handed a bare
+// HuggingFace repo id gets an optimistically constructed path that was never
+// materialized and fetches its own weights on the worker. There is no way to
+// size that from here, so it falls back to the plain default budget rather than
+// to a guess (see config.ModelLoadTimeoutForSize).
+func modelPayloadBytes(opts *pb.ModelOptions) int64 {
+	if opts == nil {
+		return 0
+	}
+	paths := []string{
+		opts.ModelFile, opts.MMProj, opts.LoraAdapter, opts.DraftModel,
+		opts.CLIPModel, opts.Tokenizer, opts.AudioPath, opts.LoraBase,
+	}
+	paths = append(paths, opts.LoraAdapters...)
+
+	// The same file can legitimately appear in two fields (a GGUF that is both
+	// ModelFile and Tokenizer, say); counting it twice would inflate the budget.
+	seen := make(map[string]struct{}, len(paths))
+	var total int64
+	for _, p := range paths {
+		if p == "" {
+			continue
+		}
+		if _, dup := seen[p]; dup {
+			continue
+		}
+		seen[p] = struct{}{}
+		total += pathBytes(p)
+	}
+	return total
+}
+
+// pathBytes returns the size of a regular file, the total size of a directory's
+// contents, or 0 if the path cannot be stat'ed.
+func pathBytes(path string) int64 {
+	fi, err := os.Stat(path)
+	if err != nil {
+		return 0
+	}
+	if !fi.IsDir() {
+		return fi.Size()
+	}
+	var total int64
+	_ = filepath.WalkDir(path, func(_ string, d fs.DirEntry, walkErr error) error {
+		if walkErr != nil || d.IsDir() {
+			return nil
+		}
+		info, infoErr := d.Info()
+		if infoErr != nil {
+			return nil
+		}
+		total += info.Size()
+		return nil
+	})
+	return total
+}
+
 // countStageableFiles returns the number of regular files a model path expands
 // to for staging: 1 for a regular file, the contained file count for a
 // directory, and 0 if the path does not exist.