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

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.