Implement resumable downloads

So it a download gets interrupted, we do not have to start again from scratch

Signed-off-by: Eric Curtin <eric.curtin@docker.com>

Author: ericcurtin

pkg/distribution/distribution/client.go

@@ -17,6 +17,7 @@ import (
 	"github.com/docker/model-runner/pkg/distribution/registry"
 	"github.com/docker/model-runner/pkg/distribution/tarball"
 	"github.com/docker/model-runner/pkg/distribution/types"
+	"github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1/remote"
 	"github.com/docker/model-runner/pkg/inference/platform"
 )
 
@@ -140,11 +141,57 @@ func NewClient(opts ...Option) (*Client, error) {
 func (c *Client) PullModel(ctx context.Context, reference string, progressWriter io.Writer) error {
 	c.log.Infoln("Starting model pull:", utils.SanitizeForLog(reference))
 
+	// First, fetch the remote model to get the manifest
 	remoteModel, err := c.registry.Model(ctx, reference)
 	if err != nil {
 		return fmt.Errorf("reading model from registry: %w", err)
 	}
 
+	// Check for incomplete downloads and prepare resume offsets
+	layers, err := remoteModel.Layers()
+	if err != nil {
+		return fmt.Errorf("getting layers: %w", err)
+	}
+
+	// Build a map of digest -> resume offset for layers with incomplete downloads
+	resumeOffsets := make(map[string]int64)
+	for _, layer := range layers {
+		digest, err := layer.Digest()
+		if err != nil {
+			c.log.Warnf("Failed to get layer digest: %v", err)
+			continue
+		}
+
+		// Check if there's an incomplete download for this layer (use DiffID for uncompressed models)
+		diffID, err := layer.DiffID()
+		if err != nil {
+			c.log.Warnf("Failed to get layer diffID: %v", err)
+			continue
+		}
+
+		incompleteSize, err := c.store.GetIncompleteSize(diffID)
+		if err != nil {
+			c.log.Warnf("Failed to check incomplete size for layer %s: %v", digest, err)
+			continue
+		}
+
+		if incompleteSize > 0 {
+			c.log.Infof("Found incomplete download for layer %s: %d bytes", digest, incompleteSize)
+			resumeOffsets[digest.String()] = incompleteSize
+		}
+	}
+
+	// If we have any incomplete downloads, create a new context with resume offsets
+	if len(resumeOffsets) > 0 {
+		c.log.Infof("Resuming %d interrupted layer download(s)", len(resumeOffsets))
+		ctx = remote.WithResumeOffsets(ctx, resumeOffsets)
+		// Re-fetch the model with the updated context
+		remoteModel, err = c.registry.Model(ctx, reference)
+		if err != nil {
+			return fmt.Errorf("reading model from registry with resume context: %w", err)
+		}
+	}
+
 	// Check for supported type
 	if err := checkCompat(remoteModel, c.log, reference); err != nil {
 		return err

pkg/distribution/internal/progress/reader.go

@@ -24,6 +24,19 @@ func NewReader(r io.Reader, updates chan<- v1.Update) io.Reader {
 	}
 }
 
+// NewReaderWithOffset returns a reader that reports progress starting from an initial offset.
+// This is useful for resuming interrupted downloads.
+func NewReaderWithOffset(r io.Reader, updates chan<- v1.Update, initialOffset int64) io.Reader {
+	if updates == nil {
+		return r
+	}
+	return &Reader{
+		Reader:       r,
+		ProgressChan: updates,
+		Total:        initialOffset,
+	}
+}
+
 func (pr *Reader) Read(p []byte) (int, error) {
 	n, err := pr.Reader.Read(p)
 	pr.Total += int64(n)

pkg/distribution/internal/store/blobs.go

@@ -78,6 +78,27 @@ type blob interface {
 	Uncompressed() (io.ReadCloser, error)
 }
 
+// layerWithDigest extends blob to include the Digest method
+type layerWithDigest interface {
+	blob
+	Digest() (v1.Hash, error)
+}
+
+// resumableLayer wraps a layer to add resume support
+type resumableLayer struct {
+	v1.Layer
+	store *LocalStore
+}
+
+func (rl *resumableLayer) Uncompressed() (io.ReadCloser, error) {
+	// For resumable downloads, we need to check for incomplete downloads
+	// and wrap the layer to inject the resume offset into the context.
+	// However, since the HTTP request happens inside the layer's Compressed() call,
+	// we can't easily intercept it here without modifying the layer interface.
+	// For now, we'll rely on the WriteBlob append logic and fetcher Range support.
+	return rl.Layer.Uncompressed()
+}
+
 // writeLayer writes the layer blob to the store.
 // It returns true when a new blob was created and the blob's DiffID.
 func (s *LocalStore) writeLayer(layer blob, updates chan<- v1.Update) (bool, v1.Hash, error) {
@@ -94,13 +115,28 @@ func (s *LocalStore) writeLayer(layer blob, updates chan<- v1.Update) (bool, v1.
 		return false, hash, nil
 	}
 
+	// Check if we're resuming an incomplete download
+	incompleteSize, err := s.GetIncompleteSize(hash)
+	if err != nil {
+		return false, v1.Hash{}, fmt.Errorf("check incomplete size: %w", err)
+	}
+
 	lr, err := layer.Uncompressed()
 	if err != nil {
 		return false, v1.Hash{}, fmt.Errorf("get blob contents: %w", err)
 	}
 	defer lr.Close()
-	r := progress.NewReader(lr, updates)
 
+	// Wrap the reader with progress reporting, accounting for already downloaded bytes
+	var r io.Reader
+	if incompleteSize > 0 {
+		r = progress.NewReaderWithOffset(lr, updates, incompleteSize)
+	} else {
+		r = progress.NewReader(lr, updates)
+	}
+
+	// WriteBlob will handle appending to incomplete files
+	// The HTTP layer will handle resuming via Range headers
 	if err := s.WriteBlob(hash, r); err != nil {
 		return false, hash, err
 	}
@@ -109,6 +145,7 @@ func (s *LocalStore) writeLayer(layer blob, updates chan<- v1.Update) (bool, v1.
 
 // WriteBlob writes the blob to the store, reporting progress to the given channel.
 // If the blob is already in the store, it is a no-op and the blob is not consumed from the reader.
+// If an incomplete download exists, it will be resumed by appending to the existing file.
 func (s *LocalStore) WriteBlob(diffID v1.Hash, r io.Reader) error {
 	hasBlob, err := s.hasBlob(diffID)
 	if err != nil {
@@ -122,21 +159,61 @@ func (s *LocalStore) WriteBlob(diffID v1.Hash, r io.Reader) error {
 	if err != nil {
 		return fmt.Errorf("get blob path: %w", err)
 	}
-	f, err := createFile(incompletePath(path))
-	if err != nil {
-		return fmt.Errorf("create blob file: %w", err)
+
+	incompletePath := incompletePath(path)
+
+	// Check if we're resuming a partial download
+	var f *os.File
+	var isResume bool
+	if _, err := os.Stat(incompletePath); err == nil {
+		// Resume: open file in append mode
+		isResume = true
+		f, err = os.OpenFile(incompletePath, os.O_WRONLY|os.O_APPEND, 0666)
+		if err != nil {
+			return fmt.Errorf("open incomplete blob file for resume: %w", err)
+		}
+	} else {
+		// New download: create file
+		f, err = createFile(incompletePath)
+		if err != nil {
+			return fmt.Errorf("create blob file: %w", err)
+		}
 	}
-	defer os.Remove(incompletePath(path))
 	defer f.Close()
 
 	if _, err := io.Copy(f, r); err != nil {
+		// Don't delete the incomplete file on error - we want to resume later
 		return fmt.Errorf("copy blob %q to store: %w", diffID.String(), err)
 	}
 
 	f.Close() // Rename will fail on Windows if the file is still open.
-	if err := os.Rename(incompletePath(path), path); err != nil {
+
+	// For resumed downloads, verify the complete file's hash before finalizing
+	// (For new downloads, the stream was already verified during download)
+	if isResume {
+		completeFile, err := os.Open(incompletePath)
+		if err != nil {
+			return fmt.Errorf("open completed file for verification: %w", err)
+		}
+		defer completeFile.Close()
+
+		computedHash, _, err := v1.SHA256(completeFile)
+		if err != nil {
+			return fmt.Errorf("compute hash of completed file: %w", err)
+		}
+
+		if computedHash.String() != diffID.String() {
+			return fmt.Errorf("hash mismatch after download: got %s, want %s", computedHash, diffID)
+		}
+	}
+
+	if err := os.Rename(incompletePath, path); err != nil {
 		return fmt.Errorf("rename blob file: %w", err)
 	}
+
+	// Only remove incomplete file if rename succeeded (though rename should have moved it)
+	// This is a safety cleanup in case rename didn't remove the source
+	os.Remove(incompletePath)
 	return nil
 }
 
@@ -160,6 +237,25 @@ func (s *LocalStore) hasBlob(hash v1.Hash) (bool, error) {
 	return false, nil
 }
 
+// GetIncompleteSize returns the size of an incomplete blob if it exists, or 0 if it doesn't.
+func (s *LocalStore) GetIncompleteSize(hash v1.Hash) (int64, error) {
+	path, err := s.blobPath(hash)
+	if err != nil {
+		return 0, fmt.Errorf("get blob path: %w", err)
+	}
+
+	incompletePath := incompletePath(path)
+	stat, err := os.Stat(incompletePath)
+	if err != nil {
+		if os.IsNotExist(err) {
+			return 0, nil
+		}
+		return 0, fmt.Errorf("stat incomplete file: %w", err)
+	}
+
+	return stat.Size(), nil
+}
+
 // createFile is a wrapper around os.Create that creates any parent directories as needed.
 func createFile(path string) (*os.File, error) {
 	if err := os.MkdirAll(filepath.Dir(path), 0777); err != nil {

pkg/go-containerregistry/pkg/v1/remote/fetcher.go

@@ -245,30 +245,96 @@ func (f *fetcher) headManifest(ctx context.Context, ref name.Reference, acceptab
 	}, nil
 }
 
+// contextKey is a type for context keys used in this package
+type contextKey string
+
+const resumeOffsetKey contextKey = "resumeOffset"
+const resumeOffsetsKey contextKey = "resumeOffsets"
+
+// WithResumeOffset returns a context with the resume offset set for a single blob
+func WithResumeOffset(ctx context.Context, offset int64) context.Context {
+	return context.WithValue(ctx, resumeOffsetKey, offset)
+}
+
+// WithResumeOffsets returns a context with resume offsets for multiple blobs (keyed by digest)
+func WithResumeOffsets(ctx context.Context, offsets map[string]int64) context.Context {
+	return context.WithValue(ctx, resumeOffsetsKey, offsets)
+}
+
+// getResumeOffset retrieves the resume offset from context for a given digest
+func getResumeOffset(ctx context.Context, digest string) int64 {
+	// First check if there's a specific offset for this digest
+	if offsets, ok := ctx.Value(resumeOffsetsKey).(map[string]int64); ok {
+		if offset, found := offsets[digest]; found && offset > 0 {
+			return offset
+		}
+	}
+	// Fall back to single offset (for fetchBlob)
+	if offset, ok := ctx.Value(resumeOffsetKey).(int64); ok {
+		return offset
+	}
+	return 0
+}
+
 func (f *fetcher) fetchBlob(ctx context.Context, size int64, h v1.Hash) (io.ReadCloser, error) {
 	u := f.url("blobs", h.String())
 	req, err := http.NewRequest(http.MethodGet, u.String(), nil)
 	if err != nil {
 		return nil, err
 	}
 
+	// Check if we should resume from a specific offset
+	resumeOffset := getResumeOffset(ctx, h.String())
+	if resumeOffset > 0 {
+		// Add Range header to resume download
+		req.Header.Set("Range", fmt.Sprintf("bytes=%d-", resumeOffset))
+	}
+
 	resp, err := f.client.Do(req.WithContext(ctx))
 	if err != nil {
 		return nil, redact.Error(err)
 	}
 
-	if err := transport.CheckError(resp, http.StatusOK); err != nil {
+	// Accept both 200 OK (full content) and 206 Partial Content (resumed)
+	if resumeOffset > 0 {
+		// If we requested a Range but got 200, the server doesn't support ranges
+		// We'll have to download from scratch
+		if resp.StatusCode == http.StatusOK {
+			// Server doesn't support range requests, will download full content
+			resumeOffset = 0
+		}
+	}
+
+	if err := transport.CheckError(resp, http.StatusOK, http.StatusPartialContent); err != nil {
 		resp.Body.Close()
 		return nil, err
 	}
 
-	// Do whatever we can.
-	// If we have an expected size and Content-Length doesn't match, return an error.
-	// If we don't have an expected size and we do have a Content-Length, use Content-Length.
+	// For partial content (resumed downloads), we can't verify the hash on the stream
+	// since we're only getting part of the file. The complete file will be verified
+	// after all bytes are written to disk.
+	if resumeOffset > 0 && resp.StatusCode == http.StatusPartialContent {
+		// Verify Content-Length matches expected remaining size
+		if hsize := resp.ContentLength; hsize != -1 {
+			if size != verify.SizeUnknown {
+				expectedRemaining := size - resumeOffset
+				if hsize != expectedRemaining {
+					resp.Body.Close()
+					return nil, fmt.Errorf("GET %s: Content-Length header %d does not match expected remaining size %d", u.String(), hsize, expectedRemaining)
+				}
+			}
+		}
+		// Return the body without verification - we'll verify the complete file later
+		return io.NopCloser(resp.Body), nil
+	}
+
+	// For full downloads, verify the stream
+	// Do whatever we can with size validation
 	if hsize := resp.ContentLength; hsize != -1 {
 		if size == verify.SizeUnknown {
 			size = hsize
 		} else if hsize != size {
+			resp.Body.Close()
 			return nil, fmt.Errorf("GET %s: Content-Length header %d does not match expected size %d", u.String(), hsize, size)
 		}
 	}