Implement resumable downloads with HTTP Range requests for model-cli pull

cmd/cli/go.mod

@@ -13,7 +13,6 @@ require (
 	github.com/docker/model-runner v0.0.0
 	github.com/emirpasic/gods/v2 v2.0.0-alpha
 	github.com/fatih/color v1.18.0
-	github.com/google/go-containerregistry v0.20.6
 	github.com/mattn/go-isatty v0.0.20
 	github.com/mattn/go-runewidth v0.0.16
 	github.com/moby/term v0.5.2
@@ -51,15 +50,13 @@ require (
 	github.com/containerd/errdefs/pkg v0.3.0 // indirect
 	github.com/containerd/log v0.1.0 // indirect
 	github.com/containerd/platforms v1.0.0-rc.1 // indirect
-	github.com/containerd/stargz-snapshotter/estargz v0.16.3 // indirect
 	github.com/containerd/typeurl/v2 v2.2.3 // indirect
 	github.com/cpuguy83/dockercfg v0.3.2 // indirect
 	github.com/cpuguy83/go-md2man/v2 v2.0.7 // indirect
 	github.com/creack/pty v1.1.24 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	github.com/distribution/reference v0.6.0 // indirect
 	github.com/dlclark/regexp2 v1.11.0 // indirect
-	github.com/docker/distribution v2.8.3+incompatible // indirect
 	github.com/docker/docker-credential-helpers v0.9.3 // indirect
 	github.com/ebitengine/purego v0.8.4 // indirect
 	github.com/elastic/go-sysinfo v1.15.4 // indirect
@@ -89,7 +86,6 @@ require (
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-shellwords v1.0.12 // indirect
 	github.com/microcosm-cc/bluemonday v1.0.27 // indirect
-	github.com/mitchellh/go-homedir v1.1.0 // indirect
 	github.com/moby/docker-image-spec v1.3.1 // indirect
 	github.com/moby/go-archive v0.1.0 // indirect
 	github.com/moby/locker v1.0.1 // indirect
@@ -118,7 +114,6 @@ require (
 	github.com/theupdateframework/notary v0.7.1-0.20210315103452-bf96a202a09a // indirect
 	github.com/tklauser/go-sysconf v0.3.12 // indirect
 	github.com/tklauser/numcpus v0.6.1 // indirect
-	github.com/vbatts/tar-split v0.12.1 // indirect
 	github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
 	github.com/yuin/goldmark v1.7.8 // indirect
 	github.com/yuin/goldmark-emoji v1.0.5 // indirect

cmd/cli/go.sum

@@ -67,8 +67,6 @@ github.com/containerd/log v0.1.0 h1:TCJt7ioM2cr/tfR8GPbGf9/VRAX8D2B4PjzCpfX540I=
 github.com/containerd/log v0.1.0/go.mod h1:VRRf09a7mHDIRezVKTRCrOq78v577GXq3bSa3EhrzVo=
 github.com/containerd/platforms v1.0.0-rc.1 h1:83KIq4yy1erSRgOVHNk1HYdPvzdJ5CnsWaRoJX4C41E=
 github.com/containerd/platforms v1.0.0-rc.1/go.mod h1:J71L7B+aiM5SdIEqmd9wp6THLVRzJGXfNuWCZCllLA4=
-github.com/containerd/stargz-snapshotter/estargz v0.16.3 h1:7evrXtoh1mSbGj/pfRccTampEyKpjpOnS3CyiV1Ebr8=
-github.com/containerd/stargz-snapshotter/estargz v0.16.3/go.mod h1:uyr4BfYfOj3G9WBVE8cOlQmXAbPN9VEQpBBeJIuOipU=
 github.com/containerd/typeurl/v2 v2.2.3 h1:yNA/94zxWdvYACdYO8zofhrTVuQY73fFU1y++dYSw40=
 github.com/containerd/typeurl/v2 v2.2.3/go.mod h1:95ljDnPfD3bAbDJRugOiShd/DlAAsxGtUBhJxIn7SCk=
 github.com/cpuguy83/dockercfg v0.3.2 h1:DlJTyZGBDlXqUZ2Dk2Q3xHs/FtnooJJVaad2S9GKorA=
@@ -155,8 +153,6 @@ github.com/google/certificate-transparency-go v1.0.10-0.20180222191210-5ab67e519
 github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
-github.com/google/go-containerregistry v0.20.6 h1:cvWX87UxxLgaH76b4hIvya6Dzz9qHB31qAwjAohdSTU=
-github.com/google/go-containerregistry v0.20.6/go.mod h1:T0x8MuoAoKX/873bkeSfLD2FAkwCDf9/HZgsFJ02E2Y=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
@@ -226,8 +222,6 @@ github.com/microcosm-cc/bluemonday v1.0.27/go.mod h1:jFi9vgW+H7c3V0lb6nR74Ib/DIB
 github.com/miekg/pkcs11 v1.0.3/go.mod h1:XsNlhZGX73bx86s2hdc/FuaLm2CPZJemRLMA+WTFxgs=
 github.com/miekg/pkcs11 v1.1.1 h1:Ugu9pdy6vAYku5DEpVWVFPYnzV+bxB+iRdbuFSu7TvU=
 github.com/miekg/pkcs11 v1.1.1/go.mod h1:XsNlhZGX73bx86s2hdc/FuaLm2CPZJemRLMA+WTFxgs=
-github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y=
-github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
 github.com/mitchellh/mapstructure v1.0.0/go.mod h1:FVVH3fgwuzCH5S8UJGiWEs2h04kUh9fWfEaFds41c1Y=
 github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0=
 github.com/moby/docker-image-spec v1.3.1/go.mod h1:eKmb5VW8vQEh/BAr2yvVNvuiJuY6UIocYsFu/DxxRpo=
@@ -342,8 +336,6 @@ github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFA
 github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI=
 github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk=
 github.com/tklauser/numcpus v0.6.1/go.mod h1:1XfjsgE2zo8GVw7POkMbHENHzVg3GzmoZ9fESEdAacY=
-github.com/vbatts/tar-split v0.12.1 h1:CqKoORW7BUWBe7UL/iqTVvkTBOF8UvOMKOIZykxnnbo=
-github.com/vbatts/tar-split v0.12.1/go.mod h1:eF6B6i6ftWQcDqEn3/iGFRFRo8cBIMSJVOpnNdfTMFA=
 github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no=
 github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=

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,34 @@ type blob interface {
 	Uncompressed() (io.ReadCloser, error)
 }
 
+// ResumableBlob extends the blob interface to support resumable downloads.
+type ResumableBlob interface {
+	blob
+	// CompressedWithOffset returns a reader starting from the given offset.
+	CompressedWithOffset(offset int64) (io.ReadCloser, error)
+	// SupportsRangeRequests checks if the remote server supports HTTP Range requests.
+	SupportsRangeRequests() (bool, error)
+}
+
+// tryResumable attempts to extract a ResumableBlob from a layer.
+// It returns nil if the layer doesn't support resumable downloads.
+func tryResumable(layer blob) ResumableBlob {
+	// Try direct type assertion
+	if r, ok := layer.(ResumableBlob); ok {
+		return r
+	}
+
+	// Try to unwrap compressedLayerExtender or other wrappers
+	type unwrapper interface {
+		Unwrap() blob
+	}
+	if u, ok := layer.(unwrapper); ok {
+		return tryResumable(u.Unwrap())
+	}
+
+	return nil
+}
+
 // 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,6 +123,13 @@ func (s *LocalStore) writeLayer(layer blob, updates chan<- v1.Update) (bool, v1.
 		return false, hash, nil
 	}
 
+	// Check if layer supports resumable downloads
+	if resumableLayer := tryResumable(layer); resumableLayer != nil {
+		created, err := s.writeBlobResumable(hash, resumableLayer, updates)
+		return created, hash, err
+	}
+
+	// Fall back to regular download
 	lr, err := layer.Uncompressed()
 	if err != nil {
 		return false, v1.Hash{}, fmt.Errorf("get blob contents: %w", err)
@@ -140,6 +176,114 @@ func (s *LocalStore) WriteBlob(diffID v1.Hash, r io.Reader) error {
 	return nil
 }
 
+// writeBlobResumable writes a blob to the store with resume support and retry logic.
+// It returns true when a new blob was created.
+func (s *LocalStore) writeBlobResumable(diffID v1.Hash, layer ResumableBlob, updates chan<- v1.Update) (bool, error) {
+	const maxRetries = 3
+	const initialBackoffSec = 1
+	const backoffMultiplier = 2
+
+	path, err := s.blobPath(diffID)
+	if err != nil {
+		return false, fmt.Errorf("get blob path: %w", err)
+	}
+
+	incompletePath := incompletePath(path)
+
+	// Check for existing incomplete file
+	var offset int64 = 0
+	var supportsRange bool
+	if info, statErr := os.Stat(incompletePath); statErr == nil {
+		offset = info.Size()
+
+		// Check if server supports range requests
+		supportsRange, err = layer.SupportsRangeRequests()
+		if err != nil {
+			// If we can't check range support, log and try anyway
+			fmt.Printf("Warning: failed to check range request support: %v\n", err)
+			supportsRange = false
+		}
+
+		if !supportsRange {
+			// Server doesn't support range requests, remove incomplete file and start fresh
+			fmt.Printf("Server doesn't support range requests, starting download from scratch\n")
+			if removeErr := os.Remove(incompletePath); removeErr != nil {
+				fmt.Printf("Warning: failed to remove incomplete file: %v\n", removeErr)
+			}
+			offset = 0
+		} else {
+			fmt.Printf("Resuming download from offset %d bytes\n", offset)
+		}
+	}
+
+	// Retry loop with exponential backoff
+	var lastErr error
+	for attempt := 0; attempt <= maxRetries; attempt++ {
+		if attempt > 0 {
+			// Calculate backoff
+			backoffDuration := time.Duration(initialBackoffSec * (1 << uint(attempt-1)) * backoffMultiplier) * time.Second
+			fmt.Printf("Retry attempt %d/%d after %v\n", attempt, maxRetries, backoffDuration)
+			time.Sleep(backoffDuration)
+
+			// Re-check incomplete file size in case it changed
+			if info, statErr := os.Stat(incompletePath); statErr == nil {
+				offset = info.Size()
+			}
+		}
+
+		// Open or create the incomplete file
+		var f *os.File
+		if offset > 0 {
+			// Open for appending
+			f, err = os.OpenFile(incompletePath, os.O_WRONLY|os.O_APPEND, 0666)
+		} else {
+			// Create new file
+			f, err = createFile(incompletePath)
+		}
+		if err != nil {
+			lastErr = fmt.Errorf("open/create incomplete file: %w", err)
+			continue
+		}
+
+		// Get the reader (with offset if resuming)
+		var lr io.ReadCloser
+		if offset > 0 && supportsRange {
+			lr, err = layer.CompressedWithOffset(offset)
+		} else {
+			lr, err = layer.Uncompressed()
+		}
+		if err != nil {
+			f.Close()
+			lastErr = fmt.Errorf("get blob contents: %w", err)
+			continue
+		}
+
+		// Wrap with progress reporter
+		r := progress.NewReader(lr, updates)
+
+		// Copy data
+		_, copyErr := io.Copy(f, r)
+		lr.Close()
+		f.Close()
+
+		if copyErr != nil {
+			lastErr = fmt.Errorf("copy blob data: %w", copyErr)
+			// Don't remove incomplete file - keep it for resume
+			continue
+		}
+
+		// Success! Rename the file
+		if err := os.Rename(incompletePath, path); err != nil {
+			return false, fmt.Errorf("rename blob file: %w", err)
+		}
+
+		return true, nil
+	}
+
+	// All retries failed
+	return false, fmt.Errorf("failed after %d retries: %w", maxRetries, lastErr)
+}
+
 // removeBlob removes the blob with the given hash from the store.
 func (s *LocalStore) removeBlob(hash v1.Hash) error {
 	path, err := s.blobPath(hash)

pkg/go-containerregistry/cmd/krane/go.mod

@@ -7,11 +7,9 @@ replace github.com/google/go-containerregistry => ../../
 require (
 	github.com/awslabs/amazon-ecr-credential-helper/ecr-login v0.9.1
 	github.com/chrismellard/docker-credential-acr-env v0.0.0-20230304212654-82a0ddb27589
-	github.com/google/go-containerregistry v0.20.3
 )
 
 require (
-	cloud.google.com/go/compute/metadata v0.7.0 // indirect
 	github.com/Azure/azure-sdk-for-go v68.0.0+incompatible // indirect
 	github.com/Azure/go-autorest v14.2.0+incompatible // indirect
 	github.com/Azure/go-autorest/autorest v0.11.30 // indirect
@@ -36,26 +34,11 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.30.1 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sts v1.33.19 // indirect
 	github.com/aws/smithy-go v1.22.3 // indirect
-	github.com/containerd/stargz-snapshotter/estargz v0.16.3 // indirect
 	github.com/dimchansky/utfbom v1.1.1 // indirect
-	github.com/docker/cli v28.2.2+incompatible // indirect
-	github.com/docker/distribution v2.8.3+incompatible // indirect
 	github.com/docker/docker-credential-helpers v0.9.3 // indirect
 	github.com/golang-jwt/jwt/v4 v4.5.2 // indirect
-	github.com/google/go-cmp v0.7.0 // indirect
-	github.com/inconshreveable/mousetrap v1.1.0 // indirect
-	github.com/klauspost/compress v1.18.0 // indirect
 	github.com/mitchellh/go-homedir v1.1.0 // indirect
-	github.com/opencontainers/go-digest v1.0.0 // indirect
-	github.com/opencontainers/image-spec v1.1.1 // indirect
-	github.com/pkg/errors v0.9.1 // indirect
 	github.com/sirupsen/logrus v1.9.3 // indirect
-	github.com/spf13/cobra v1.9.1 // indirect
-	github.com/spf13/pflag v1.0.6 // indirect
-	github.com/vbatts/tar-split v0.12.1 // indirect
 	golang.org/x/crypto v0.38.0 // indirect
-	golang.org/x/oauth2 v0.30.0 // indirect
-	golang.org/x/sync v0.15.0 // indirect
 	golang.org/x/sys v0.33.0 // indirect
-	gotest.tools/v3 v3.1.0 // indirect
 )