fix(distributed): make the cold-load hold scale with progress, not wall-clock (#11019)

A 70 GB video checkpoint (longcat-video-avatar-1.5) could not be loaded on a
distributed cluster. The request failed with HTTP 500 after 1499.98s - exactly
the 25m00s cold-load ceiling - while staging was demonstrably healthy: 26 of 57
files and 39 GB transferred at a sustained ~26 MB/s, zero errors, no stalls. It
was not wedged, it was killed by a timer.

ModelLoadCeilingFor covers node selection, backend install, file staging and the
remote LoadModel. Install and load carry their own budgets; staging was covered
only by a FIXED 5-minute margin. But staging time is bytes over bandwidth, not a
constant: 70 GB at 26 MB/s needs ~45m against a 25m ceiling, so the failure is
deterministic for any sufficiently large model rather than a flake. Simply
raising the constant moves the cliff to the next model size - the deployment
target here is checkpoints of 600 GB and beyond.

The ceiling's real purpose is that "a wedged worker can never pin the lock
indefinitely". Progress, not elapsed time, is what distinguishes a wedged worker
from a large one. The hold is now a deadline that extends whenever the transfer
reports bytes and expires a 5-minute stall window after they stop:

- A large model transferring fine continues, for hours if needed.
- A worker that died mid-transfer still fails within the stall window.

Progress is observed at byte level on the transfer itself, via the existing
staging progress callback. Per-file completion would be too coarse - a single
600 GB shard would be indistinguishable from a stall for hours. The observation
point is back-pressured by the socket, so it reflects the network rather than
local disk reads. Observation is coarsened to one timer touch per stall/20 so
the per-read callback stays cheap.

The base budget (unchanged, and still derived from the install and load
timeouts) continues to cover the steps that report no progress, so
LOCALAI_NATS_MODEL_LOAD_TIMEOUT keeps working exactly as before. An absolute
cap of 24h bounds the hold even while progress keeps arriving, so a peer
trickling bytes forever cannot pin the advisory lock; 600 GB at the measured
26 MB/s is ~6.5h, so the cap sits far above any legitimate transfer.

Also fixes the incoherent layering the same error exposed: the resumable upload
carried a 1h retry budget nested inside the 25m ceiling, so the inner budget was
unreachable and the message still blamed it ("failed after 1 attempts within
1h0m0s budget") while the 25m parent was the actual killer. The upload now
adopts the caller's deadline when there is one, and applies its fixed budget
only when nothing above bounded it - which also stops a fixed 1h from
reintroducing the size cliff under the now-extendable parent.

This is the successor to #10968, where a hardcoded 5-minute LoadModel gRPC
timeout was replaced by this derived ceiling. Fixing the inner timeout exposed
the outer ceiling as the new binding constraint.


Assisted-by: Claude Code:claude-opus-4-8[1m] [Read] [Edit] [Bash]

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

Author: localai-bot

core/services/nodes/file_stager_http.go

@@ -117,11 +117,8 @@ func (h *HTTPFileStager) EnsureRemote(ctx context.Context, nodeID, localPath, ke
 	xlog.Info("Uploading file to remote node", "node", nodeID, "file", filepath.Base(localPath), "size", humanFileSize(fileSize), "url", url)
 
 	// Outer time budget: bound the total resumable-upload duration so a
-	// permanently-unreachable worker doesn't hold the request forever. Default
-	// matches the existing per-response timeout.
-	outerBudget := h.resumeBudget()
-
-	resumeCtx, cancel := context.WithTimeout(ctx, outerBudget)
+	// permanently-unreachable worker doesn't hold the request forever.
+	resumeCtx, cancel, outerBudget := h.resumeContext(ctx)
 	defer cancel()
 
 	var lastErr error
@@ -178,14 +175,34 @@ func (h *HTTPFileStager) EnsureRemote(ctx context.Context, nodeID, localPath, ke
 	}
 }
 
-// resumeBudget returns the maximum total time the resumable upload loop will
-// spend retrying transient failures end-to-end. Past this budget the upload
-// fails rather than spinning forever — 1h covers multi-GB transfers on
-// pathological links without letting a wedged server jam the master.
-func (h *HTTPFileStager) resumeBudget() time.Duration {
-	return 1 * time.Hour
+// resumeContext bounds the resumable upload loop so it can't spin forever, and
+// returns the budget it ended up with for error reporting.
+//
+// When the caller already imposed a deadline, that deadline IS the budget:
+// nesting a second one under it was actively misleading. In production a 1h
+// resume budget sat inside a 25m cold-load ceiling, so the inner budget was
+// unreachable and the failure still reported "failed after 1 attempts within
+// 1h0m0s budget" while the real killer was the 25m parent. Worse, the parent's
+// cold-load hold is now progress-extended (see load_deadline.go) precisely so a
+// 600 GB transfer can run for hours — a fixed 1h nested inside it would
+// reintroduce the very size cliff that change removes.
+//
+// The fixed fallback only applies when nobody above bounded the transfer.
+func (h *HTTPFileStager) resumeContext(ctx context.Context) (context.Context, context.CancelFunc, time.Duration) {
+	if deadline, ok := ctx.Deadline(); ok {
+		resumeCtx, cancel := context.WithCancel(ctx)
+		return resumeCtx, cancel, time.Until(deadline)
+	}
+	resumeCtx, cancel := context.WithTimeout(ctx, defaultResumeBudget)
+	return resumeCtx, cancel, defaultResumeBudget
 }
 
+// defaultResumeBudget bounds an unparented resumable upload. 1h covers multi-GB
+// transfers on pathological links without letting a wedged server jam the
+// master. Callers that need longer (a cold load staging a 600 GB checkpoint)
+// impose their own, progress-extended deadline instead.
+const defaultResumeBudget = 1 * time.Hour
+
 // nextBackoff returns the sleep before retry #attempt: 1s, 2s, 4s, ..., capped
 // at 30s, with the first sleep (attempt=2) being 1s.
 func nextBackoff(attempt int) time.Duration {

core/services/nodes/load_deadline.go

@@ -0,0 +1,256 @@
+package nodes
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"time"
+
+	"github.com/mudler/xlog"
+)
+
+// The cold-load hold has to satisfy two requirements that a single wall-clock
+// timeout cannot express at once:
+//
+//  1. A legitimately huge checkpoint must be allowed to finish. Staging time is
+//     bytes/bandwidth, not a constant: a 70 GB model at 26 MB/s needs ~45m and a
+//     600 GB checkpoint needs hours. Any fixed ceiling is therefore a model-size
+//     cliff — raise it and the cliff simply moves to the next larger model.
+//  2. A worker that died mid-transfer must still release the per-model advisory
+//     lock promptly, or it pins every other replica's request for that model.
+//
+// The discriminator between the two is progress: bytes moving means healthy,
+// silence means wedged. So the hold is a deadline that is *pushed forward* every
+// time the transfer reports bytes, rather than a countdown from when the load
+// began. Duration then scales with the work instead of with the clock.
+
+const (
+	// stagingStallWindow is how long the transfer may report zero bytes before
+	// the load is declared wedged. It replaces the old fixed
+	// modelLoadStagingMargin: the same 5 minutes, but rolling rather than
+	// one-shot, so it bounds silence instead of bounding total transfer size.
+	//
+	// 5m is well beyond any legitimate gap in a healthy stream (the receiving
+	// worker fsyncing a multi-GB shard, a transient LAN blip inside the
+	// resumable-upload retry loop, a GC pause) while still releasing the
+	// advisory lock quickly enough that a dead worker is not felt as an outage.
+	stagingStallWindow = 5 * time.Minute
+
+	// modelLoadAbsoluteMax bounds the hold even while progress keeps arriving,
+	// so a degenerate peer that trickles a few bytes per stall window forever
+	// cannot pin the advisory lock for all time. It is deliberately far above
+	// any legitimate transfer: 600 GB at the 26 MB/s measured in production is
+	// ~6.5h, and at a pessimistic 10 MB/s ~17h, so 24h leaves real headroom
+	// while still guaranteeing the lock is eventually released.
+	modelLoadAbsoluteMax = 24 * time.Hour
+
+	// loadProgressQuantumDivisor coarsens progress observation. The byte-level
+	// callback fires on every read of the upload body (~32 KB), which is far too
+	// often to touch a timer; extending at most once per stall/20 keeps the hot
+	// path cheap while costing at most 5% of the stall window in precision.
+	loadProgressQuantumDivisor = 20
+
+	// loadProgressQuantumMax caps the coarsening so a long stall window doesn't
+	// make progress observation itself laggy.
+	loadProgressQuantumMax = time.Second
+)
+
+// errLoadDeadlineExpired is the cancellation cause recorded when the cold-load
+// hold expires, so loadDeadlineContext can report context.DeadlineExceeded and
+// stay distinguishable from a caller-driven cancel.
+var errLoadDeadlineExpired = errors.New("cold-load deadline expired")
+
+// loadDeadline is the mutable expiry behind a loadDeadlineContext. Observe()
+// pushes the expiry forward; the absolute cap and the parent context still
+// bound it.
+type loadDeadline struct {
+	mu          sync.Mutex
+	timer       *time.Timer
+	stall       time.Duration
+	quantum     time.Duration
+	expiry      time.Time
+	hardExpiry  time.Time
+	lastObserve time.Time
+	extensions  int
+	stopped     bool
+	cancel      context.CancelCauseFunc
+}
+
+// loadDeadlineKey retrieves the active loadDeadline from a context chain.
+type loadDeadlineKey struct{}
+
+// loadDeadlineContext is a context whose expiry moves forward while the work it
+// covers reports progress.
+//
+// It reports context.DeadlineExceeded (not context.Canceled) on expiry, because
+// downstream code — notably the resumable upload loop in file_stager_http.go —
+// branches on that distinction to tell "we ran out of budget" apart from "the
+// caller gave up", and the two lead to different retry decisions.
+type loadDeadlineContext struct {
+	context.Context
+	d *loadDeadline
+}
+
+// Deadline reports the ABSOLUTE cap rather than the current rolling expiry.
+// Children derive their own budgets from the parent deadline (a gRPC
+// context.WithTimeout takes the earlier of the two), and the rolling expiry is
+// a floor that moves, not a promise that the work stops then. Reporting the
+// rolling value would let a child silently inherit a few seconds of budget.
+func (c *loadDeadlineContext) Deadline() (time.Time, bool) {
+	c.d.mu.Lock()
+	defer c.d.mu.Unlock()
+	return c.d.hardExpiry, true
+}
+
+func (c *loadDeadlineContext) Err() error {
+	err := c.Context.Err()
+	if err == nil {
+		return nil
+	}
+	if errors.Is(context.Cause(c.Context), errLoadDeadlineExpired) {
+		return context.DeadlineExceeded
+	}
+	return err
+}
+
+func (c *loadDeadlineContext) Value(key any) any {
+	if _, ok := key.(loadDeadlineKey); ok {
+		return c.d
+	}
+	return c.Context.Value(key)
+}
+
+// newLoadDeadlineContext builds a cold-load hold context. base is the initial
+// budget granted before any progress is seen (it covers the steps that report
+// none: node selection, backend install, the remote LoadModel). stall is how
+// long zero progress is tolerated once the transfer has started, and absoluteMax
+// bounds the whole hold regardless of progress.
+//
+// Non-positive base/stall/absoluteMax fall back to their package defaults, and
+// stall is clamped to base so a caller with a deliberately tight ceiling (tests,
+// or an operator who wants fast failure) doesn't get a stall window wider than
+// the ceiling it was derived from.
+func newLoadDeadlineContext(parent context.Context, base, stall, absoluteMax time.Duration) (context.Context, context.CancelFunc) {
+	if base <= 0 {
+		base = minModelLoadCeiling
+	}
+	if stall <= 0 {
+		stall = stagingStallWindow
+	}
+	if stall > base {
+		stall = base
+	}
+	if absoluteMax <= 0 {
+		absoluteMax = modelLoadAbsoluteMax
+	}
+	if absoluteMax < base {
+		absoluteMax = base
+	}
+
+	ctx, cancel := context.WithCancelCause(parent)
+	now := time.Now()
+
+	quantum := stall / loadProgressQuantumDivisor
+	if quantum > loadProgressQuantumMax {
+		quantum = loadProgressQuantumMax
+	}
+
+	d := &loadDeadline{
+		stall:      stall,
+		quantum:    quantum,
+		expiry:     now.Add(base),
+		hardExpiry: now.Add(absoluteMax),
+		cancel:     cancel,
+	}
+	d.timer = time.AfterFunc(base, d.fire)
+
+	// The absolute cap is a plain one-shot: progress can never move it.
+	hardTimer := time.AfterFunc(absoluteMax, func() {
+		xlog.Warn("Cold-load hold hit its absolute cap while still reporting progress; releasing the model lock",
+			"cap", absoluteMax)
+		d.stop()
+		cancel(errLoadDeadlineExpired)
+	})
+
+	stopAll := func() {
+		hardTimer.Stop()
+		d.stop()
+		cancel(context.Canceled)
+	}
+	return &loadDeadlineContext{Context: ctx, d: d}, stopAll
+}
+
+func (d *loadDeadline) fire() {
+	d.mu.Lock()
+	if d.stopped {
+		d.mu.Unlock()
+		return
+	}
+	// A late-arriving Observe may have pushed the expiry out after this timer
+	// was already scheduled to fire; re-arm instead of cancelling.
+	if remaining := time.Until(d.expiry); remaining > 0 {
+		d.timer.Reset(remaining)
+		d.mu.Unlock()
+		return
+	}
+	d.stopped = true
+	extensions := d.extensions
+	stall := d.stall
+	d.mu.Unlock()
+
+	if extensions > 0 {
+		xlog.Warn("Cold load stalled: no staging progress within the stall window, releasing the model lock",
+			"stallWindow", stall, "progressExtensions", extensions)
+	}
+	d.cancel(errLoadDeadlineExpired)
+}
+
+func (d *loadDeadline) stop() {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	d.stopped = true
+	if d.timer != nil {
+		d.timer.Stop()
+	}
+}
+
+// Observe records that real work happened, pushing the expiry out by a full
+// stall window. It never pulls the expiry in, so it cannot shorten the base
+// budget, and it never pushes past the absolute cap.
+func (d *loadDeadline) Observe() {
+	if d == nil {
+		return
+	}
+	now := time.Now()
+
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	if d.stopped {
+		return
+	}
+	// Coarsen: the byte callback fires per read, we only need per-quantum.
+	if !d.lastObserve.IsZero() && now.Sub(d.lastObserve) < d.quantum {
+		return
+	}
+	d.lastObserve = now
+
+	next := now.Add(d.stall)
+	if next.After(d.hardExpiry) {
+		next = d.hardExpiry
+	}
+	if !next.After(d.expiry) {
+		return
+	}
+	d.expiry = next
+	d.extensions++
+	d.timer.Reset(time.Until(next))
+}
+
+// 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.
+func observeLoadProgress(ctx context.Context) {
+	if d, ok := ctx.Value(loadDeadlineKey{}).(*loadDeadline); ok {
+		d.Observe()
+	}
+}

core/services/nodes/router.go

@@ -87,6 +87,15 @@ type SmartRouterOptions struct {
 	// must be widened in step (see ModelLoadCeilingFor) or the ceiling clips the
 	// longer deadline before it can be used.
 	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
+	// bytes keep moving the hold extends, so transfer size no longer bounds the
+	// load. Zero selects stagingStallWindow. It is clamped to ModelLoadCeiling.
+	StagingStallWindow time.Duration
+	// ModelLoadAbsoluteMax bounds the cold-load hold even while progress keeps
+	// arriving, so a peer trickling bytes forever cannot pin the advisory lock
+	// indefinitely. Zero selects modelLoadAbsoluteMax (24h).
+	ModelLoadAbsoluteMax time.Duration
 }
 
 // modelLoadStagingMargin is the slack ModelLoadCeilingFor adds on top of the
@@ -154,6 +163,10 @@ type SmartRouter struct {
 	// modelLoadTimeout is the deadline for the remote LoadModel gRPC call
 	// (see SmartRouterOptions.ModelLoadTimeout).
 	modelLoadTimeout time.Duration
+	// stagingStallWindow and modelLoadAbsoluteMax turn modelLoadCeiling from a
+	// hard countdown into a progress-extended hold (see load_deadline.go).
+	stagingStallWindow   time.Duration
+	modelLoadAbsoluteMax time.Duration
 }
 
 // probeCacheTTL is how long a successful gRPC HealthCheck on a backend is
@@ -194,6 +207,11 @@ func NewSmartRouter(registry ModelRouter, opts SmartRouterOptions) *SmartRouter
 		sharedModels:     opts.SharedModels,
 		modelLoadCeiling: ceiling,
 		modelLoadTimeout: loadTimeout,
+		// Zero values are resolved to their defaults inside
+		// newLoadDeadlineContext, which also clamps the stall window to the
+		// ceiling, so nothing to normalize here.
+		stagingStallWindow:   opts.StagingStallWindow,
+		modelLoadAbsoluteMax: opts.ModelLoadAbsoluteMax,
 	}
 }
 
@@ -526,7 +544,13 @@ func (r *SmartRouter) Route(ctx context.Context, modelID, modelName, backendType
 	// even if a sub-step wedges. Each long step still has its own (tighter)
 	// bound; this only backstops them. The per-model advisory lock below
 	// de-dupes concurrent loaders across replicas.
-	loadCtx, cancelLoad := context.WithTimeout(context.WithoutCancel(ctx), r.modelLoadCeiling)
+	// The backstop is progress-based, not wall-clock: staging time is bytes over
+	// bandwidth, so a fixed ceiling is a model-size cliff (a 70 GB checkpoint
+	// transferring healthily at 26 MB/s needs ~45m and was killed at exactly
+	// 25m00s). The hold instead extends while the transfer reports bytes and
+	// expires a stall window after they stop. See load_deadline.go.
+	loadCtx, cancelLoad := newLoadDeadlineContext(context.WithoutCancel(ctx),
+		r.modelLoadCeiling, r.stagingStallWindow, r.modelLoadAbsoluteMax)
 	defer cancelLoad()
 	loadModel := func(ctx context.Context) (*RouteResult, error) {
 		// Re-check after acquiring lock — another request may have loaded it
@@ -1328,6 +1352,12 @@ func (r *SmartRouter) stageModelFiles(ctx context.Context, node *BackendNode, op
 func (r *SmartRouter) withStagingCallback(ctx context.Context, trackingKey, fileName string, fileIdx, totalFiles int) context.Context {
 	start := time.Now()
 	return WithStagingProgress(ctx, func(fn string, bytesSent, totalBytes int64) {
+		// Byte-level movement is what keeps the cold-load hold alive. Observing
+		// here (rather than at per-file completion) is what makes a single
+		// 600 GB shard distinguishable from a wedged worker: file-granular
+		// progress would look identical to a stall for hours.
+		observeLoadProgress(ctx)
+
 		var speed string
 		elapsed := time.Since(start)
 		if elapsed > 0 {