Implement resumable downloads

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

Signed-off-by: Eric Curtin <[email protected]>

Author: ericcurtin

Dockerfile

@@ -19,6 +19,9 @@ WORKDIR /app
 # Copy go mod/sum first for better caching
 COPY --link go.mod go.sum ./
 
+# Copy pkg/go-containerregistry for the replace directive in go.mod
+COPY --link pkg/go-containerregistry ./pkg/go-containerregistry
+
 # Download dependencies (with cache mounts)
 RUN --mount=type=cache,target=/go/pkg/mod \
     --mount=type=cache,target=/root/.cache/go-build \

Makefile

@@ -77,7 +77,7 @@ integration-tests:
 	@echo "Integration tests completed!"
 
 validate:
-	find . -type f -name "*.sh" | xargs shellcheck
+	find . -type f -name "*.sh" | grep -v pkg/go-containerregistry | xargs shellcheck
 
 # Build Docker image
 docker-build:

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,24 +141,90 @@ 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 supported type
-	if err := checkCompat(remoteModel, c.log, reference); err != nil {
-		return err
-	}
-
-	// Get the remote image digest
+	// Get the remote image digest immediately to ensure we work with a consistent manifest
+	// This prevents race conditions if the tag is updated during the pull
 	remoteDigest, err := remoteModel.Digest()
 	if err != nil {
 		c.log.Errorln("Failed to get remote image digest:", err)
 		return fmt.Errorf("getting remote image digest: %w", err)
 	}
 	c.log.Infoln("Remote model digest:", remoteDigest.String())
 
+	// 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
+	// and re-fetch using the digest to ensure we're resuming the same manifest
+	if len(resumeOffsets) > 0 {
+		c.log.Infof("Resuming %d interrupted layer download(s)", len(resumeOffsets))
+		ctx = remote.WithResumeOffsets(ctx, resumeOffsets)
+		// Re-fetch the model using the digest reference to prevent race conditions
+		// Extract repository name from the original reference and construct digest reference
+		repository := reference
+		// Find the last occurrence of : or @ (tag or digest separator)
+		// We need to search after the last / to avoid matching port separators
+		if lastSlash := strings.LastIndex(reference, "/"); lastSlash != -1 {
+			// Search for : or @ after the last slash
+			suffix := reference[lastSlash:]
+			if idx := strings.IndexAny(suffix, ":@"); idx != -1 {
+				repository = reference[:lastSlash+idx]
+			}
+		} else {
+			// No slash found, search from beginning (e.g., "library/image:tag" or "image:tag")
+			if idx := strings.IndexAny(reference, ":@"); idx != -1 {
+				repository = reference[:idx]
+			}
+		}
+		digestReference := repository + "@" + remoteDigest.String()
+		c.log.Infof("Re-fetching model with digest reference: %s", utils.SanitizeForLog(digestReference))
+		remoteModel, err = c.registry.Model(ctx, digestReference)
+		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
+	}
+
 	// Check if model exists in local store
 	localModel, err := c.store.Read(remoteDigest.String())
 	if err == nil {

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

@@ -6,6 +6,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"time"
 
 	"github.com/docker/model-runner/pkg/distribution/internal/progress"
 
@@ -78,6 +79,12 @@ type blob interface {
 	Uncompressed() (io.ReadCloser, error)
 }
 
+// layerWithDigest extends blob to include the Digest method
+type layerWithDigest interface {
+	blob
+	Digest() (v1.Hash, error)
+}
+
 // 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 +101,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 +131,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 +145,83 @@ 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 {
+		// Before resuming, verify that the incomplete file isn't already complete
+		existingFile, err := os.Open(incompletePath)
+		if err != nil {
+			return fmt.Errorf("open incomplete file for verification: %w", err)
+		}
+
+		computedHash, _, err := v1.SHA256(existingFile)
+		existingFile.Close()
+
+		if err == nil && computedHash.String() == diffID.String() {
+			// File is already complete, just rename it
+			if err := os.Rename(incompletePath, path); err != nil {
+				return fmt.Errorf("rename completed blob file: %w", err)
+			}
+			return nil
+		}
+
+		// File is incomplete or corrupt, try to resume
+		isResume = true
+		f, err = os.OpenFile(incompletePath, os.O_WRONLY|os.O_APPEND, 0644)
+		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 {
+		// If we were resuming and copy failed, the incomplete file might be corrupt
+		if isResume {
+			_ = os.Remove(incompletePath)
+		}
 		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() {
+			// The resumed download is corrupt, remove it so we can start fresh next time
+			_ = os.Remove(incompletePath)
+			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 +245,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 {
@@ -201,3 +305,59 @@ func (s *LocalStore) writeConfigFile(mdl v1.Image) (bool, error) {
 	}
 	return true, nil
 }
+
+// CleanupStaleIncompleteFiles removes incomplete download files that haven't been modified
+// for more than the specified duration. This prevents disk space leaks from abandoned downloads.
+func (s *LocalStore) CleanupStaleIncompleteFiles(maxAge time.Duration) error {
+	blobsPath := s.blobsDir()
+	if _, err := os.Stat(blobsPath); os.IsNotExist(err) {
+		// Blobs directory doesn't exist yet, nothing to clean up
+		return nil
+	}
+
+	var cleanedCount int
+	var cleanupErrors []error
+
+	// Walk through the blobs directory looking for .incomplete files
+	err := filepath.Walk(blobsPath, func(path string, info os.FileInfo, err error) error {
+		if err != nil {
+			// Continue walking even if we encounter errors on individual files
+			return nil
+		}
+
+		// Skip directories
+		if info.IsDir() {
+			return nil
+		}
+
+		// Only process .incomplete files
+		if !strings.HasSuffix(path, ".incomplete") {
+			return nil
+		}
+
+		// Check if file is older than maxAge
+		if time.Since(info.ModTime()) > maxAge {
+			if removeErr := os.Remove(path); removeErr != nil {
+				cleanupErrors = append(cleanupErrors, fmt.Errorf("failed to remove stale incomplete file %s: %w", path, removeErr))
+			} else {
+				cleanedCount++
+			}
+		}
+
+		return nil
+	})
+
+	if err != nil {
+		return fmt.Errorf("walking blobs directory: %w", err)
+	}
+
+	if len(cleanupErrors) > 0 {
+		return fmt.Errorf("encountered %d errors during cleanup (cleaned %d files): %v", len(cleanupErrors), cleanedCount, cleanupErrors[0])
+	}
+
+	if cleanedCount > 0 {
+		fmt.Printf("Cleaned up %d stale incomplete download file(s)\n", cleanedCount)
+	}
+
+	return nil
+}

pkg/distribution/internal/store/store.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"time"
 
 	v1 "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1"
 
@@ -81,6 +82,13 @@ func (s *LocalStore) initialize() error {
 		}
 	}
 
+	// Clean up stale incomplete files (older than 7 days)
+	// This prevents disk space leaks from abandoned downloads
+	if err := s.CleanupStaleIncompleteFiles(7 * 24 * time.Hour); err != nil {
+		// Log the error but don't fail initialization
+		fmt.Printf("Warning: failed to clean up stale incomplete files: %v\n", err)
+	}
+
 	return nil
 }