diff --git a/cmd/cli/desktop/progress.go b/cmd/cli/desktop/progress.go index c771b6c5..3dc2b714 100644 --- a/cmd/cli/desktop/progress.go +++ b/cmd/cli/desktop/progress.go @@ -12,6 +12,7 @@ import ( "github.com/docker/docker/pkg/jsonmessage" "github.com/docker/go-units" "github.com/docker/model-runner/cmd/cli/pkg/standalone" + "github.com/docker/model-runner/pkg/distribution/oci" ) // DisplayProgress displays progress messages from a model pull/push operation @@ -155,6 +156,9 @@ func writeDockerProgress(w io.Writer, msg *ProgressMessage, layerStatus map[stri return nil } + // Detect if this is a push operation based on the sentinel layer ID + isPush := layerID == oci.UploadingLayerID + // Determine status based on progress var status string var progressDetail *jsonmessage.JSONProgress @@ -162,13 +166,21 @@ func writeDockerProgress(w io.Writer, msg *ProgressMessage, layerStatus map[stri if msg.Layer.Current == 0 { status = "Waiting" } else if msg.Layer.Current < msg.Layer.Size { - status = "Downloading" + if isPush { + status = "Uploading" + } else { + status = "Downloading" + } progressDetail = &jsonmessage.JSONProgress{ Current: int64(msg.Layer.Current), Total: int64(msg.Layer.Size), } } else if msg.Layer.Current >= msg.Layer.Size && msg.Layer.Size > 0 { - status = "Pull complete" + if isPush { + status = "Push complete" + } else { + status = "Pull complete" + } progressDetail = &jsonmessage.JSONProgress{ Current: int64(msg.Layer.Current), Total: int64(msg.Layer.Size), diff --git a/pkg/distribution/internal/progress/reporter.go b/pkg/distribution/internal/progress/reporter.go index acfcf15d..b5e0cf74 100644 --- a/pkg/distribution/internal/progress/reporter.go +++ b/pkg/distribution/internal/progress/reporter.go @@ -84,7 +84,7 @@ func (r *Reporter) Updates() chan<- oci.Update { now := time.Now() var layerSize uint64 var layerID string - if r.layer != nil { // In case of Push there is no layer yet + if r.layer != nil { // In case of Pull id, err := r.layer.DiffID() if err != nil { r.err = err @@ -97,8 +97,10 @@ func (r *Reporter) Updates() chan<- oci.Update { continue } layerSize = safeUint64(size) - } else { - layerSize = safeUint64(p.Total) + } else { // In case of Push there is no layer yet + // Use imageSize as layer is not known at this point + layerSize = r.imageSize + layerID = oci.UploadingLayerID // Fake ID for push operations to enable progress display } incrementalBytes := p.Complete - lastComplete diff --git a/pkg/distribution/oci/progress.go b/pkg/distribution/oci/progress.go index 9e2026bf..d65ae237 100644 --- a/pkg/distribution/oci/progress.go +++ b/pkg/distribution/oci/progress.go @@ -1,5 +1,10 @@ package oci +// UploadingLayerID is a sentinel layer ID used to identify push operations. +// During push, there is no real layer available yet, so this fake ID signals +// that the operation is an upload rather than a download. +const UploadingLayerID = "uploading" + // Update represents a progress update during image operations. type Update struct { Complete int64 diff --git a/pkg/distribution/oci/remote/remote.go b/pkg/distribution/oci/remote/remote.go index d7344bfe..ce9d9888 100644 --- a/pkg/distribution/oci/remote/remote.go +++ b/pkg/distribution/oci/remote/remote.go @@ -18,6 +18,7 @@ import ( "github.com/containerd/containerd/v2/core/remotes/docker" "github.com/containerd/containerd/v2/plugins/content/local" "github.com/containerd/errdefs" + "github.com/docker/model-runner/pkg/distribution/internal/progress" "github.com/docker/model-runner/pkg/distribution/oci" "github.com/docker/model-runner/pkg/distribution/oci/authn" "github.com/docker/model-runner/pkg/distribution/oci/reference" @@ -771,7 +772,14 @@ func Write(ref reference.Reference, img oci.Image, opts ...Option) error { return fmt.Errorf("pushing layer: %w", err) } - if _, err := io.Copy(cw, rc); err != nil { + // Wrap the reader with progress tracking to report incremental upload progress + // Uses the shared progress.Reader from internal/progress package + var reader io.Reader = rc + if o.progress != nil { + reader = progress.NewReaderWithOffset(rc, o.progress, completed) + } + + if _, err := io.Copy(cw, reader); err != nil { cw.Close() rc.Close() closeProgress(o.progress)