Merge pull request #873 from docker/fix-pull-resume-and-progress-bars

fix: improve pull resume reliability and progress bar display on retry

Author: ericcurtin

cmd/cli/desktop/desktop.go

@@ -117,7 +117,7 @@ func (c *Client) Pull(model string, printer standalone.StatusPrinter) (string, b
 		hfToken = os.Getenv("HF_TOKEN")
 	}
 
-	return c.withRetries("download", 3, printer, func(attempt int) (string, bool, error, bool) {
+	return c.withRetries("download", 4, printer, func(attempt int) (string, bool, error, bool) {
 		jsonData, err := json.Marshal(dmrm.ModelCreateRequest{
 			From:        model,
 			BearerToken: hfToken,
@@ -237,6 +237,11 @@ func (c *Client) withRetries(
 		if attempt > 0 {
 			// Calculate exponential backoff: 2^(attempt-1) seconds (1s, 2s, 4s)
 			backoffDuration := time.Duration(1<<uint(attempt-1)) * time.Second
+			// Print a blank line to stdout so that any progress bars drawn during
+			// the previous attempt are visually separated from the retry attempt.
+			// This prevents the new progress bars from overwriting the old ones
+			// when the terminal display is reset on each retry.
+			printer.Println("")
 			printer.PrintErrf("Retrying %s (attempt %d/%d) in %v...\n", operationName, attempt, maxRetries, backoffDuration)
 			time.Sleep(backoffDuration)
 		}
@@ -263,7 +268,7 @@ func (c *Client) Push(model string, printer standalone.StatusPrinter) (string, b
 		hfToken = os.Getenv("HF_TOKEN")
 	}
 
-	return c.withRetries("push", 3, printer, func(attempt int) (string, bool, error, bool) {
+	return c.withRetries("push", 4, printer, func(attempt int) (string, bool, error, bool) {
 		pushPath := inference.ModelsPrefix + "/" + model + "/push"
 		var body io.Reader
 		if hfToken != "" {

cmd/cli/desktop/desktop_test.go

@@ -191,13 +191,13 @@ func TestPullMaxRetriesExhausted(t *testing.T) {
 	mockContext := NewContextForMock(mockClient)
 	client := New(mockContext)
 
-	// All 4 attempts (1 initial + 3 retries) fail with network error
-	mockClient.EXPECT().Do(gomock.Any()).Return(nil, io.EOF).Times(4)
+	// All 5 attempts (1 initial + 4 retries) fail with network error
+	mockClient.EXPECT().Do(gomock.Any()).Return(nil, io.EOF).Times(5)
 
 	printer := NewSimplePrinter(func(s string) {})
 	_, _, err := client.Pull(modelName, printer)
 	assert.Error(t, err)
-	assert.Contains(t, err.Error(), "download failed after 3 retries")
+	assert.Contains(t, err.Error(), "download failed after 4 retries")
 }
 
 func TestPushRetryOnNetworkError(t *testing.T) {

pkg/distribution/internal/store/blobs.go

@@ -1,10 +1,8 @@
 package store
 
 import (
-	"context"
 	"crypto/sha256"
 	"encoding/hex"
-	"errors"
 	"fmt"
 	"io"
 	"os"
@@ -199,11 +197,10 @@ func (s *LocalStore) WriteBlobWithResume(diffID oci.Hash, r io.Reader, digestStr
 		buf := make([]byte, 1)
 		n, readErr := r.Read(buf)
 		if readErr != nil && readErr != io.EOF {
-			// Clean up the incomplete file on read error (unless it's a context cancellation
-			// which should preserve the file for future resume attempts)
-			if !errors.Is(readErr, context.Canceled) && !errors.Is(readErr, context.DeadlineExceeded) {
-				_ = os.Remove(incompletePath)
-			}
+			// Preserve the incomplete file on all errors so that the next
+			// attempt can resume from where this one left off. Stale
+			// incomplete files are cleaned up by CleanupStaleIncompleteFiles
+			// during store initialisation (default: files older than 7 days).
 			return fmt.Errorf("read first byte: %w", readErr)
 		}
 

pkg/distribution/internal/store/store.go

@@ -374,29 +374,16 @@ func (s *LocalStore) Write(mdl oci.Image, tags []string, w io.Writer, opts ...Wr
 		return err
 	}
 
-	// Collect new layer digests
-	var newLayerDigests []oci.Hash
-	for _, result := range results {
-		if result.created {
-			newLayerDigests = append(newLayerDigests, result.diffID)
-		}
-	}
-
-	if len(newLayerDigests) > 0 {
-		digests := append([]oci.Hash(nil), newLayerDigests...)
-		cleanups = append(cleanups, func() error {
-			var errs []error
-			for _, dg := range digests {
-				if err := s.removeBlob(dg); err != nil && !errors.Is(err, os.ErrNotExist) {
-					errs = append(errs, fmt.Errorf("remove blob %s: %w", dg, err))
-				}
-			}
-			if len(errs) > 0 {
-				return errors.Join(errs...)
-			}
-			return nil
-		})
-	}
+	// Do not register completed layer blobs in the rollback list.
+	//
+	// Layer blobs are content-addressed: a blob identified by its diffID is
+	// immutable and may be shared across multiple models. Rolling them back
+	// when a later step fails (e.g. writing the manifest, tagging) would
+	// discard already-downloaded data that is still valid, forcing a full
+	// re-download on the next attempt instead of resuming only the layer
+	// that was actually in progress. The manifest/config/index cleanup below
+	// is sufficient to leave the store in a consistent state.
+	_ = results // results used above for error checking; layer blobs are retained
 
 	// Write the manifest
 	digest, err := mdl.Digest()

pkg/distribution/internal/store/store_test.go

@@ -424,14 +424,19 @@ func TestWriteRollsBackOnTagFailure(t *testing.T) {
 		t.Fatalf("expected config blob to be cleaned up, stat error: %v", err)
 	}
 
+	// Layer blobs are content-addressed and are intentionally retained even
+	// after a failed write. They may be reused by a subsequent pull of the
+	// same or another model, and they allow the download to resume rather
+	// than restart from byte 0. Only the manifest, config, and index are
+	// rolled back to leave the store in a consistent (non-indexed) state.
 	for _, digestStr := range diffIDs {
 		parts := strings.SplitN(digestStr, ":", 2)
 		if len(parts) != 2 {
 			t.Fatalf("unexpected diffID format: %q", digestStr)
 		}
 		layerPath := filepath.Join(storePath, "blobs", parts[0], parts[1])
-		if _, err := os.Stat(layerPath); !errors.Is(err, os.ErrNotExist) {
-			t.Fatalf("expected layer blob %q to be cleaned up, stat error: %v", layerPath, err)
+		if _, err := os.Stat(layerPath); err != nil {
+			t.Fatalf("expected layer blob %q to be retained for future resume, stat error: %v", layerPath, err)
 		}
 	}
 

pkg/distribution/oci/remote/remote.go

@@ -189,24 +189,48 @@ func (t *rangeTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 		return resp, err
 	}
 
-	// If we requested a Range, record success only if the server accepted the range request
-	// Servers should return 206 (Partial Content) for successful range requests,
-	// but some may return 200 with the partial content, so we record success for both
-	if requestedOffset > 0 {
-		if resp.StatusCode == http.StatusPartialContent || resp.StatusCode == http.StatusOK {
-			// Record in RangeSuccess tracker so WriteBlob can check it
+	// If we requested a Range, record success only when the server honoured it
+	// with 206 Partial Content and a matching Content-Range start offset. A 200
+	// response means the server ignored the Range header and is sending the full
+	// file from byte 0; appending that stream to the existing partial file would
+	// produce a corrupt blob. We also validate the Content-Range start offset to
+	// guard against a misbehaving server that returns 206 with a different range.
+	if requestedOffset > 0 && resp.StatusCode == http.StatusPartialContent {
+		if rangeStartMatchesOffset(resp.Header.Get("Content-Range"), requestedOffset) {
 			if rs := GetRangeSuccess(req.Context()); rs != nil {
 				rs.Add(digest, requestedOffset)
 			}
 		}
-		// If range request was not successful (e.g., 416 Range Not Satisfiable),
-		// don't record in RangeSuccess, which will cause WriteBlob to start fresh
-		// (no explicit action needed in the else case)
 	}
 
 	return resp, nil
 }
 
+// rangeStartMatchesOffset parses the Content-Range response header and reports
+// whether its start byte equals the given offset. The format is defined by
+// RFC 9110: "bytes START-END/TOTAL" (TOTAL may be "*"). We fail closed: if the
+// header is absent or cannot be parsed we return false so that the caller does
+// not treat an ambiguous response as a successful range request.
+func rangeStartMatchesOffset(contentRange string, offset int64) bool {
+	if contentRange == "" {
+		return false
+	}
+	// Trim the unit prefix "bytes " and split on "-"
+	after, ok := strings.CutPrefix(contentRange, "bytes ")
+	if !ok {
+		return false
+	}
+	dashIdx := strings.Index(after, "-")
+	if dashIdx < 0 {
+		return false
+	}
+	var start int64
+	if _, err := fmt.Sscanf(after[:dashIdx], "%d", &start); err != nil {
+		return false
+	}
+	return start == offset
+}
+
 // extractDigestAndOffset extracts the blob digest from the request URL and returns
 // the corresponding resume offset if one exists.
 func (t *rangeTransport) extractDigestAndOffset(req *http.Request, offsets map[string]int64) (string, int64) {