Implement resumable downloads

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:

cmd/cli/commands/package.go

@@ -17,7 +17,7 @@ import (
 	"github.com/docker/model-runner/pkg/distribution/tarball"
 	"github.com/docker/model-runner/pkg/distribution/types"
 
-	"github.com/google/go-containerregistry/pkg/name"
+	"github.com/docker/model-runner/pkg/go-containerregistry/pkg/name"
 	"github.com/spf13/cobra"
 
 	"github.com/docker/model-runner/cmd/cli/commands/completion"

cmd/cli/commands/tag.go

@@ -7,7 +7,7 @@ import (
 	"github.com/docker/model-runner/cmd/cli/commands/completion"
 	"github.com/docker/model-runner/cmd/cli/desktop"
 	"github.com/docker/model-runner/pkg/distribution/registry"
-	"github.com/google/go-containerregistry/pkg/name"
+	"github.com/docker/model-runner/pkg/go-containerregistry/pkg/name"
 	"github.com/spf13/cobra"
 )
 

cmd/cli/commands/utils.go

@@ -11,7 +11,7 @@ import (
 	"github.com/docker/model-runner/cmd/cli/desktop"
 	"github.com/docker/model-runner/cmd/cli/pkg/standalone"
 	"github.com/docker/model-runner/pkg/inference/backends/vllm"
-	"github.com/google/go-containerregistry/pkg/name"
+	"github.com/docker/model-runner/pkg/go-containerregistry/pkg/name"
 	"github.com/moby/term"
 	"github.com/spf13/cobra"
 )

cmd/cli/go.mod

@@ -11,9 +11,9 @@ require (
 	github.com/docker/go-connections v0.6.0
 	github.com/docker/go-units v0.5.0
 	github.com/docker/model-runner v0.0.0
+	github.com/docker/model-runner/pkg/go-containerregistry v0.0.0-00010101000000-000000000000
 	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-runewidth v0.0.16
 	github.com/moby/term v0.5.2
 	github.com/muesli/termenv v0.16.0
@@ -154,3 +154,5 @@ require (
 replace github.com/kolesnikovae/go-winjob => github.com/docker/go-winjob v0.0.0-20250829235554-57b487ebcbc5
 
 replace github.com/docker/model-runner => ../..
+
+replace github.com/docker/model-runner/pkg/go-containerregistry => ../../pkg/go-containerregistry

cmd/cli/go.sum

@@ -155,8 +155,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=

go.mod

@@ -6,8 +6,8 @@ require (
 	github.com/containerd/containerd/v2 v2.1.5
 	github.com/containerd/platforms v1.0.0-rc.1
 	github.com/docker/go-units v0.5.0
+	github.com/docker/model-runner/pkg/go-containerregistry v0.0.0-00010101000000-000000000000
 	github.com/elastic/go-sysinfo v1.15.4
-	github.com/google/go-containerregistry v0.20.6
 	github.com/gpustack/gguf-parser-go v0.22.1
 	github.com/jaypipes/ghw v0.19.1
 	github.com/kolesnikovae/go-winjob v1.0.0
@@ -72,3 +72,5 @@ require (
 )
 
 replace github.com/kolesnikovae/go-winjob => github.com/docker/go-winjob v0.0.0-20250829235554-57b487ebcbc5
+
+replace github.com/docker/model-runner/pkg/go-containerregistry => ./pkg/go-containerregistry

go.sum

@@ -60,8 +60,6 @@ github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da h1:oI5xCqsCo564l
 github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
 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=

pkg/distribution/builder/builder.go

@@ -5,7 +5,7 @@ import (
 	"fmt"
 	"io"
 
-	v1 "github.com/google/go-containerregistry/pkg/v1"
+	v1 "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1"
 
 	"github.com/docker/model-runner/pkg/distribution/internal/gguf"
 	"github.com/docker/model-runner/pkg/distribution/internal/mutate"

pkg/distribution/builder/builder_test.go

@@ -10,7 +10,7 @@ import (
 
 	"github.com/docker/model-runner/pkg/distribution/builder"
 	"github.com/docker/model-runner/pkg/distribution/types"
-	v1 "github.com/google/go-containerregistry/pkg/v1"
+	v1 "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1"
 )
 
 func TestBuilder(t *testing.T) {

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/distribution/client_test.go

@@ -16,9 +16,9 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/google/go-containerregistry/pkg/name"
-	"github.com/google/go-containerregistry/pkg/registry"
-	"github.com/google/go-containerregistry/pkg/v1/remote"
+	"github.com/docker/model-runner/pkg/go-containerregistry/pkg/name"
+	"github.com/docker/model-runner/pkg/go-containerregistry/pkg/registry"
+	"github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1/remote"
 	"github.com/sirupsen/logrus"
 
 	"github.com/docker/model-runner/pkg/distribution/internal/gguf"

pkg/distribution/internal/bundle/unpack.go

@@ -9,7 +9,7 @@ import (
 	"path/filepath"
 
 	"github.com/docker/model-runner/pkg/distribution/types"
-	ggcrtypes "github.com/google/go-containerregistry/pkg/v1/types"
+	ggcrtypes "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1/types"
 )
 
 // Unpack creates and return a Bundle by unpacking files and config from model into dir.

pkg/distribution/internal/gguf/create.go

@@ -6,7 +6,7 @@ import (
 	"strings"
 	"time"
 
-	v1 "github.com/google/go-containerregistry/pkg/v1"
+	v1 "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1"
 	parser "github.com/gpustack/gguf-parser-go"
 
 	"github.com/docker/model-runner/pkg/distribution/internal/partial"

pkg/distribution/internal/mutate/model.go

@@ -4,9 +4,9 @@ import (
 	"encoding/json"
 	"fmt"
 
-	v1 "github.com/google/go-containerregistry/pkg/v1"
-	ggcrpartial "github.com/google/go-containerregistry/pkg/v1/partial"
-	ggcr "github.com/google/go-containerregistry/pkg/v1/types"
+	v1 "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1"
+	ggcrpartial "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1/partial"
+	ggcr "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1/types"
 
 	"github.com/docker/model-runner/pkg/distribution/internal/partial"
 	"github.com/docker/model-runner/pkg/distribution/types"

pkg/distribution/internal/mutate/mutate.go

@@ -1,8 +1,8 @@
 package mutate
 
 import (
-	v1 "github.com/google/go-containerregistry/pkg/v1"
-	ggcr "github.com/google/go-containerregistry/pkg/v1/types"
+	v1 "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1"
+	ggcr "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1/types"
 
 	"github.com/docker/model-runner/pkg/distribution/types"
 )

pkg/distribution/internal/mutate/mutate_test.go

@@ -5,8 +5,8 @@ import (
 	"path/filepath"
 	"testing"
 
-	"github.com/google/go-containerregistry/pkg/v1/static"
-	ggcr "github.com/google/go-containerregistry/pkg/v1/types"
+	"github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1/static"
+	ggcr "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1/types"
 
 	"github.com/docker/model-runner/pkg/distribution/internal/gguf"
 	"github.com/docker/model-runner/pkg/distribution/internal/mutate"

pkg/distribution/internal/partial/layer.go

@@ -6,8 +6,8 @@ import (
 	"os"
 	"path/filepath"
 
-	"github.com/google/go-containerregistry/pkg/v1"
-	ggcrtypes "github.com/google/go-containerregistry/pkg/v1/types"
+	"github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1"
+	ggcrtypes "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1/types"
 
 	"github.com/docker/model-runner/pkg/distribution/types"
 )

pkg/distribution/internal/partial/model.go

@@ -4,9 +4,9 @@ import (
 	"encoding/json"
 	"fmt"
 
-	v1 "github.com/google/go-containerregistry/pkg/v1"
-	"github.com/google/go-containerregistry/pkg/v1/partial"
-	ggcr "github.com/google/go-containerregistry/pkg/v1/types"
+	v1 "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1"
+	"github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1/partial"
+	ggcr "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1/types"
 
 	"github.com/docker/model-runner/pkg/distribution/types"
 )

pkg/distribution/internal/partial/partial.go

@@ -4,9 +4,9 @@ import (
 	"encoding/json"
 	"fmt"
 
-	v1 "github.com/google/go-containerregistry/pkg/v1"
-	"github.com/google/go-containerregistry/pkg/v1/partial"
-	ggcr "github.com/google/go-containerregistry/pkg/v1/types"
+	v1 "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1"
+	"github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1/partial"
+	ggcr "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1/types"
 
 	"github.com/docker/model-runner/pkg/distribution/types"
 )

pkg/distribution/internal/progress/reader.go

@@ -3,7 +3,7 @@ package progress
 import (
 	"io"
 
-	"github.com/google/go-containerregistry/pkg/v1"
+	"github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1"
 )
 
 // Reader wraps an io.Reader to track reading progress
@@ -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/progress/reporter.go

@@ -6,7 +6,7 @@ import (
 	"io"
 	"time"
 
-	v1 "github.com/google/go-containerregistry/pkg/v1"
+	v1 "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1"
 )
 
 // UpdateInterval defines how often progress updates should be sent

pkg/distribution/internal/progress/reporter_test.go

@@ -7,8 +7,8 @@ import (
 	"testing"
 	"time"
 
-	v1 "github.com/google/go-containerregistry/pkg/v1"
-	v1types "github.com/google/go-containerregistry/pkg/v1/types"
+	v1 "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1"
+	v1types "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1/types"
 
 	"github.com/docker/model-runner/pkg/distribution/types"
 )

pkg/distribution/internal/safetensors/create.go

@@ -8,7 +8,7 @@ import (
 	"strconv"
 	"time"
 
-	v1 "github.com/google/go-containerregistry/pkg/v1"
+	v1 "github.com/docker/model-runner/pkg/go-containerregistry/pkg/v1"
 
 	"github.com/docker/model-runner/pkg/distribution/internal/partial"
 	"github.com/docker/model-runner/pkg/distribution/types"