fix(distributed): count staging verification as progress, not as a stall (#11026)

Testing the progress-based cold-load deadline on the live cluster surfaced a
false positive. The stall window observed UPLOAD bytes only, but the staging
path has a phase that does real work while moving zero upload bytes: the
resumable-upload verify phase.

When a shard is already present on the worker from an earlier attempt, the
frontend HEADs it, hashes the local copy to confirm it matches, and skips the
transfer. Staging a 70 GB model with 56 GB already staged:

  17:27:34 INFO Upload skipped (file already exists with matching hash) ...
  17:28:20 INFO Upload skipped (file already exists with matching hash) ...
  17:29:07 INFO Upload skipped (file already exists with matching hash) ...
  ... six-plus consecutive minutes, no bytes uploaded at all

~45s per skipped ~4 GB shard. That is correct and desirable - it is what makes
resume work - but it was indistinguishable from a stall. At 45s per shard it
sits inside the 5m window, so the run in flight was fine; the problem is the
600 GB scale this machinery exists to enable, where one shard can plausibly hash
for longer than the window. The guard would then fire during verification of a
transfer that is working perfectly.

Verified mechanism: probeExisting() HEADs the worker and then calls
downloader.CalculateSHA(). The staging progress callback is only consulted
inside doUpload(), which the skip path never reaches, so observeLoadProgress was
called zero times for the whole verify phase.

Verification exposed a second, worse bug in the same path: CalculateSHA consults
no context at all. An expired cold load kept hashing to completion, compared the
hashes, and returned success - reporting a file as staged on a dead load. The
failure only surfaced on the NEXT file, whose HEAD died immediately. That is
exactly the shape of the red test here, which fails on shard 3.

Fix: hash in 1 MiB chunks via hashFileWithActivity(), ticking the cold-load
deadline per chunk and checking ctx per chunk. A successful HEAD also counts,
since a 200 with a content hash proves the worker is serving right now.

Counting hash progress does not make a dead transfer look alive: hashing is
bounded, terminating work proportional to file size, in probeExisting it runs
only after a HEAD proved the worker was up, and the 24h absolute cap still
bounds the whole hold. The alternative of simply widening the window was
rejected - it would reintroduce the size cliff this work removes.


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

@@ -2,6 +2,8 @@ package nodes
 
 import (
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -19,7 +21,6 @@ import (
 	"github.com/mudler/xlog"
 
 	"github.com/mudler/LocalAI/core/services/storage"
-	"github.com/mudler/LocalAI/pkg/downloader"
 	"github.com/mudler/LocalAI/pkg/httpclient"
 )
 
@@ -106,8 +107,14 @@ func (h *HTTPFileStager) EnsureRemote(ctx context.Context, nodeID, localPath, ke
 	// attempt — the server uses it to detect mid-flight content drift and
 	// reject (409) if a partial upload claims a new identity, forcing a clean
 	// restart.
-	localHash, err := downloader.CalculateSHA(localPath)
+	localHash, err := hashFileWithActivity(ctx, localPath)
 	if err != nil {
+		if ctx.Err() != nil {
+			// The cold load was cancelled or expired while hashing. Uploading
+			// on a dead context can only fail, so surface it rather than
+			// pressing on without resume-safety.
+			return "", fmt.Errorf("hashing %s for upload to node %s: %w", localPath, nodeID, err)
+		}
 		// Hash failure isn't fatal — we can still upload; we just lose
 		// resume-safety and end-of-transfer integrity checks.
 		xlog.Warn("Failed to hash local file for upload integrity check", "localPath", localPath, "error", err)
@@ -461,7 +468,12 @@ func (h *HTTPFileStager) probeExisting(ctx context.Context, addr, localPath, key
 		return "", false
 	}
 
-	localHash, err := downloader.CalculateSHA(localPath)
+	// A 200 with a content hash is proof the worker is alive and serving right
+	// now, so it counts as progress in its own right — otherwise a long run of
+	// legitimately skipped shards accrues no activity at all.
+	observeLoadProgress(ctx)
+
+	localHash, err := hashFileWithActivity(ctx, localPath)
 	if err != nil {
 		return "", false
 	}
@@ -473,6 +485,59 @@ func (h *HTTPFileStager) probeExisting(ctx context.Context, addr, localPath, key
 	return remotePath, true
 }
 
+// hashChunkSize is how much of a file is hashed between activity ticks and
+// cancellation checks. 1 MiB is small enough that a cancelled cold load aborts
+// promptly even on a huge shard, and large enough that the per-chunk bookkeeping
+// is irrelevant next to the hashing itself.
+const hashChunkSize = 1 << 20
+
+// hashFileWithActivity computes a file's SHA-256 while reporting activity to the
+// cold-load deadline, and aborts if that deadline expires.
+//
+// Hashing is the long pole of the resumable-upload verify phase, which moves
+// zero upload bytes: the client HEADs the worker, hashes the local file to
+// confirm it matches, and skips the transfer. On the live cluster that measured
+// ~45s per ~4 GB shard, with six-plus consecutive minutes of no uploaded bytes.
+// A stall window watching only upload bytes would eventually mistake that for a
+// wedged worker — and precisely in the 600 GB case this machinery exists to
+// enable, where one shard can hash for longer than the window.
+//
+// Counting hash progress as progress is safe because it is bounded, terminating
+// work proportional to file size: it cannot make a dead transfer look alive
+// indefinitely. In probeExisting it also runs only after a successful HEAD
+// proved the worker was serving, and the absolute cap still bounds the whole
+// hold regardless.
+func hashFileWithActivity(ctx context.Context, path string) (string, error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return "", err
+	}
+	defer func() { _ = f.Close() }()
+
+	h := sha256.New()
+	buf := make([]byte, hashChunkSize)
+	for {
+		// Checked per chunk: CalculateSHA consulted no context at all, so an
+		// expired cold load used to hash to completion and report the file as
+		// staged, turning a dead load into a silent success.
+		if err := ctx.Err(); err != nil {
+			return "", err
+		}
+		n, readErr := f.Read(buf)
+		if n > 0 {
+			h.Write(buf[:n])
+			observeLoadProgress(ctx)
+		}
+		if errors.Is(readErr, io.EOF) {
+			break
+		}
+		if readErr != nil {
+			return "", readErr
+		}
+	}
+	return hex.EncodeToString(h.Sum(nil)), nil
+}
+
 // progressReader wraps an io.Reader and logs upload progress periodically.
 // If a StagingProgressCallback is present in the context, it also calls it
 // for UI-visible progress updates.

core/services/nodes/file_stager_verify_deadline_test.go

@@ -0,0 +1,165 @@
+package nodes
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"os"
+	"path/filepath"
+	"time"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+// The resumable-upload fast path (HEAD the worker, hash the local file, skip
+// when the hashes match) is what makes resume work after a partial transfer.
+// It reports ZERO upload bytes for its entire duration, so a stall window that
+// watches only upload bytes cannot tell it apart from a wedged worker.
+//
+// Measured on the live cluster staging a 70 GB model with 56 GB already
+// present: ~45s per skipped ~4 GB shard, six-plus consecutive minutes with no
+// bytes uploaded at all. At the 600 GB scale this work exists to support, a
+// single shard can plausibly hash for longer than the 5m stall window.
+
+// verifyShardSize keeps each shard small enough to stay friendly to /tmp while
+// still costing real, measurable hashing time.
+const verifyShardSize = 16 << 20
+
+// verifyShardCount is chosen so the cumulative hashing time across consecutive
+// skipped shards comfortably outlasts the scaled-down stall window, mirroring
+// the run of consecutive skips seen on the cluster.
+const verifyShardCount = 6
+
+func writeShard(path string, size int) string {
+	f, err := os.Create(path)
+	Expect(err).ToNot(HaveOccurred())
+	defer func() { _ = f.Close() }()
+
+	h := sha256.New()
+	chunk := make([]byte, 1<<20)
+	for i := range chunk {
+		chunk[i] = byte(i * 7)
+	}
+	for written := 0; written < size; written += len(chunk) {
+		n := len(chunk)
+		if remaining := size - written; remaining < n {
+			n = remaining
+		}
+		_, err := f.Write(chunk[:n])
+		Expect(err).ToNot(HaveOccurred())
+		h.Write(chunk[:n])
+	}
+	return hex.EncodeToString(h.Sum(nil))
+}
+
+var _ = Describe("staging verify phase and the cold-load stall window", func() {
+	newStagerFor := func(srv *httptest.Server) *HTTPFileStager {
+		return NewHTTPFileStager(func(string) (string, error) {
+			u, err := url.Parse(srv.URL)
+			if err != nil {
+				return "", err
+			}
+			return u.Host, nil
+		}, "")
+	}
+
+	It("survives a run of verified-and-skipped shards that upload no bytes at all", func() {
+		tmp := GinkgoT().TempDir()
+		paths := make([]string, verifyShardCount)
+		hashes := make(map[string]string, verifyShardCount)
+		for i := range paths {
+			key := fmt.Sprintf("shard-%d", i)
+			paths[i] = filepath.Join(tmp, key+".safetensors")
+			hashes[key] = writeShard(paths[i], verifyShardSize)
+		}
+
+		// The worker already holds every shard from a previous attempt, so each
+		// one is HEADed, hashed locally, and skipped.
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			Expect(r.Method).To(Equal(http.MethodHead),
+				"every shard must take the verify/skip path, never upload")
+			key := filepath.Base(r.URL.Path)
+			w.Header().Set(HeaderLocalPath, "/worker/models/"+key)
+			w.Header().Set(HeaderContentSHA256, hashes[key])
+			w.WriteHeader(http.StatusOK)
+		}))
+		defer srv.Close()
+
+		// Scaled-down budgets: hashing these shards takes far longer than the
+		// window, standing in for 600 GB shards against a 5m window.
+		ctx, cancel := newLoadDeadlineContext(context.Background(),
+			20*time.Millisecond, 20*time.Millisecond, time.Minute)
+		defer cancel()
+
+		stager := newStagerFor(srv)
+		start := time.Now()
+		for i, p := range paths {
+			key := fmt.Sprintf("shard-%d", i)
+			remotePath, err := stager.EnsureRemote(ctx, "nvidia-thor", p, key)
+			Expect(err).ToNot(HaveOccurred(),
+				"verifying shard %d does real work and must not be mistaken for a stall", i)
+			Expect(remotePath).To(Equal("/worker/models/" + key))
+		}
+		Expect(time.Since(start)).To(BeNumerically(">", 20*time.Millisecond),
+			"the verify run must actually have outlasted the stall window")
+	})
+
+	It("does not silently succeed on the verify path after the deadline has expired", func() {
+		// The verify path consulted no context at all: it hashed to completion
+		// and returned success even on a dead context, so an expired cold load
+		// looked like a staged file.
+		tmp := GinkgoT().TempDir()
+		localPath := filepath.Join(tmp, "shard.safetensors")
+		hash := writeShard(localPath, 1<<20)
+
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set(HeaderLocalPath, "/worker/models/shard")
+			w.Header().Set(HeaderContentSHA256, hash)
+			w.WriteHeader(http.StatusOK)
+		}))
+		defer srv.Close()
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		_, err := newStagerFor(srv).EnsureRemote(ctx, "nvidia-thor", localPath, "shard")
+		Expect(err).To(HaveOccurred(),
+			"a cancelled load must not report a file as staged")
+	})
+
+	It("still kills a worker that goes silent during the verify phase", func() {
+		tmp := GinkgoT().TempDir()
+		localPath := filepath.Join(tmp, "small.safetensors")
+		writeShard(localPath, 1<<20)
+
+		// The worker accepted the connection and then died: no HEAD response,
+		// no bytes, no hashing. Nothing here is real work.
+		//
+		// Ordering matters: srv.Close() blocks until in-flight handlers return,
+		// so the handler must be released BEFORE Close runs. Defers are LIFO.
+		release := make(chan struct{})
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			<-release
+		}))
+		defer srv.Close()
+		defer close(release)
+
+		ctx, cancel := newLoadDeadlineContext(context.Background(),
+			100*time.Millisecond, 100*time.Millisecond, time.Minute)
+		defer cancel()
+
+		start := time.Now()
+		_, err := newStagerFor(srv).EnsureRemote(ctx, "dead-node", localPath, "shard")
+		elapsed := time.Since(start)
+
+		Expect(err).To(HaveOccurred(),
+			"a silent worker must still be caught by the stall window")
+		Expect(elapsed).To(BeNumerically("<", 10*time.Second),
+			"the advisory lock must be released promptly")
+	})
+})

docs/content/features/distributed-mode.md

@@ -79,7 +79,7 @@ The router also bounds how long a single cold load may hold the per-model adviso
 
 The load starts with a base budget of `max(backend-install-timeout + model-load-timeout + 5m, 25m)` — with the defaults, `15m + 5m + 5m = 25m`. That budget covers the steps that report no progress: node selection, backend install, and the remote `LoadModel` call. Raising either timeout widens it in step, so a longer load deadline is never clipped.
 
-While **model files are staging**, however, the deadline extends every time bytes actually move, and expires only once the transfer has been silent for a 5-minute stall window. Staging time is a function of checkpoint size and available bandwidth, not a constant: a 70 GB model at 26 MB/s needs about 45 minutes, and a 600 GB checkpoint needs hours. A fixed ceiling would therefore be a model-size cliff — every increase just moves the cliff to the next larger model. Extending on progress means a large model transfers for as long as it legitimately needs, while a worker that dies mid-transfer still releases the lock within the stall window.
+While **model files are staging**, however, the deadline extends every time staging does real work, and expires only once staging has been silent for a 5-minute stall window. Real work means uploaded bytes, and also the resumable-upload verify phase: when a shard is already present on the worker from an earlier attempt, the frontend HEADs it and hashes the local copy to confirm it matches, then skips the transfer. That phase uploads nothing at all — on a 70 GB model resuming with 56 GB already staged it ran for six-plus consecutive minutes at ~45s per shard — so hashing counts as progress too. Otherwise a resumed transfer would be mistaken for a wedged one. Staging time is a function of checkpoint size and available bandwidth, not a constant: a 70 GB model at 26 MB/s needs about 45 minutes, and a 600 GB checkpoint needs hours. A fixed ceiling would therefore be a model-size cliff — every increase just moves the cliff to the next larger model. Extending on progress means a large model transfers for as long as it legitimately needs, while a worker that dies mid-transfer still releases the lock within the stall window.
 
 An absolute cap of 24h ends the hold even if progress keeps arriving, so a degenerate peer trickling a few bytes at a time cannot pin the lock forever. No configuration is needed for either value; both are sized well above any legitimate transfer.