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

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 {