feat: add dynamic timeout retry with file-size-based parameters

cmd/push.go

@@ -20,11 +20,11 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/modelpack/modctl/pkg/backend"
-	"github.com/modelpack/modctl/pkg/config"
-
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
+
+	"github.com/modelpack/modctl/pkg/backend"
+	"github.com/modelpack/modctl/pkg/config"
 )
 
 var pushConfig = config.NewPush()
@@ -48,10 +48,25 @@ var pushCmd = &cobra.Command{
 // init initializes push command.
 func init() {
 	flags := pushCmd.Flags()
-	flags.IntVar(&pushConfig.Concurrency, "concurrency", pushConfig.Concurrency, "specify the number of concurrent push operations")
+	flags.IntVar(
+		&pushConfig.Concurrency,
+		"concurrency",
+		pushConfig.Concurrency,
+		"specify the number of concurrent push operations",
+	)
 	flags.BoolVar(&pushConfig.PlainHTTP, "plain-http", false, "use plain HTTP instead of HTTPS")
-	flags.BoolVar(&pushConfig.Insecure, "insecure", false, "turning on this flag will disable TLS verification")
-	flags.BoolVar(&pushConfig.Nydusify, "nydusify", false, "[EXPERIMENTAL] nydusify the model artifact")
+	flags.BoolVar(
+		&pushConfig.Insecure,
+		"insecure",
+		false,
+		"turning on this flag will disable TLS verification",
+	)
+	flags.BoolVar(
+		&pushConfig.Nydusify,
+		"nydusify",
+		false,
+		"[EXPERIMENTAL] nydusify the model artifact",
+	)
 	flags.MarkHidden("nydusify")
 
 	if err := viper.BindPFlags(flags); err != nil {

internal/pb/pb.go

@@ -57,11 +57,30 @@ type ProgressBar struct {
 
 type progressBar struct {
 	*mpbv8.Bar
-	size      int64
+	size int64
+	// msgMu guards msg, which is read by mpb's render goroutine (via the
+	// prepend decorator) while being written by Placeholder/Complete from
+	// transfer goroutines.
+	msgMu     sync.RWMutex
 	msg       string
 	startTime time.Time
 }
 
+// setMsg updates the bar's prepended message under lock.
+func (b *progressBar) setMsg(msg string) {
+	b.msgMu.Lock()
+	b.msg = msg
+	b.msgMu.Unlock()
+}
+
+// msgText returns the bar's prepended message under lock. It is called from
+// mpb's render goroutine, so it must never block on anything but msgMu.
+func (b *progressBar) msgText() string {
+	b.msgMu.RLock()
+	defer b.msgMu.RUnlock()
+	return b.msg
+}
+
 // NewProgressBar creates a new progress bar.
 func NewProgressBar(writers ...io.Writer) *ProgressBar {
 	opts := []mpbv8.ContainerOption{
@@ -113,7 +132,7 @@ func (p *ProgressBar) Add(prompt, name string, size int64, reader io.Reader) io.
 		mpbv8.BarFillerOnComplete("|"),
 		mpbv8.PrependDecorators(
 			decor.Any(func(s decor.Statistics) string {
-				return newBar.msg
+				return newBar.msgText()
 			}, decor.WCSyncSpaceR),
 		),
 		mpbv8.AppendDecorators(
@@ -140,6 +159,36 @@ func (p *ProgressBar) Add(prompt, name string, size int64, reader io.Reader) io.
 	return reader
 }
 
+// Placeholder creates or resets a progress bar entry without a reader. It is
+// used during retry backoff to keep a visible bar for the item.
+//
+// A prior failed attempt typically calls Abort, which drops the bar from mpb's
+// render loop; an aborted (or completed) bar cannot be revived, so in that case
+// the bar is recreated via Add. Otherwise the live bar is reset in place: the
+// bytes transferred so far are stashed as a refill mark, progress is rewound to
+// zero, and the speed estimate is reset so the ETA does not count the backoff
+// wait as stalled transfer time.
+func (p *ProgressBar) Placeholder(name string, prompt string, size int64) {
+	if disableProgress.Load() {
+		return
+	}
+
+	p.mu.RLock()
+	existing := p.bars[name]
+	p.mu.RUnlock()
+
+	if existing == nil || existing.Bar.AbortedOrCompleted() {
+		// A nil or aborted/completed bar would render nothing, so recreate it.
+		p.Add(prompt, name, size, nil)
+		return
+	}
+
+	existing.setMsg(fmt.Sprintf("%s %s", prompt, name))
+	existing.Bar.SetRefill(existing.Bar.Current())
+	existing.Bar.SetCurrent(0)
+	existing.Bar.EwmaSetCurrent(0, time.Second)
+}
+
 // Get returns the progress bar.
 func (p *ProgressBar) Get(name string) *progressBar {
 	p.mu.RLock()
@@ -156,7 +205,7 @@ func (p *ProgressBar) Complete(name string, msg string) {
 	p.mu.RUnlock()
 
 	if ok {
-		bar.msg = msg
+		bar.setMsg(msg)
 		bar.Bar.SetCurrent(bar.size)
 	}
 }

internal/pb/pb_test.go

@@ -0,0 +1,85 @@
+/*
+ *     Copyright 2024 The ModelPack Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package pb
+
+import (
+	"io"
+	"sync"
+	"testing"
+	"time"
+)
+
+// TestProgressBarMsgConcurrency exercises concurrent message updates against
+// the read path used by mpb's render goroutine. Placeholder and Complete are
+// fired from transfer goroutines (e.g. on retry backoff) while the bar is being
+// rendered, so progressBar.msg is read and written concurrently. Run with
+// -race; before msg was guarded this reproduces the reported data race.
+func TestProgressBarMsgConcurrency(t *testing.T) {
+	// Render to a discarded writer so the auto-refresh goroutine still runs the
+	// prepend decorator (which reads msg) without touching the test output.
+	pb := NewProgressBar(io.Discard)
+	pb.Start()
+
+	const name = "sha256:deadbeef"
+	pb.Add("Copying blob", name, 1024, nil)
+
+	var wg sync.WaitGroup
+	stop := make(chan struct{})
+
+	// Writers: simulate retry backoff resets (Placeholder) and completion.
+	for i := 0; i < 4; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for {
+				select {
+				case <-stop:
+					return
+				default:
+				}
+				pb.Placeholder(name, "Copying blob (retry)", 1024)
+				pb.Complete(name, "done")
+			}
+		}()
+	}
+
+	// Readers: mirror the render goroutine's read of the message.
+	for i := 0; i < 4; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for {
+				select {
+				case <-stop:
+					return
+				default:
+				}
+				if b := pb.Get(name); b != nil {
+					_ = b.msgText()
+				}
+			}
+		}()
+	}
+
+	time.Sleep(150 * time.Millisecond)
+	close(stop)
+	wg.Wait()
+
+	// Leave the bar completed so Stop()'s mpb.Shutdown() does not block.
+	pb.Complete(name, "done")
+	pb.Stop()
+}

pkg/backend/retry.go → pkg/backend/annotation.go

@@ -17,14 +17,20 @@
 package backend
 
 import (
-	"time"
-
-	retry "github.com/avast/retry-go/v4"
+	legacymodelspec "github.com/dragonflyoss/model-spec/specs-go/v1"
+	modelspec "github.com/modelpack/model-spec/specs-go/v1"
 )
 
-var defaultRetryOpts = []retry.Option{
-	retry.Attempts(6),
-	retry.DelayType(retry.BackOffDelay),
-	retry.Delay(5 * time.Second),
-	retry.MaxDelay(60 * time.Second),
+// getAnnotationFilepath returns the filepath stored on a descriptor's
+// annotations, preferring the modelpack key and falling back to the legacy
+// dragonflyoss key so older artifacts remain readable. Returns empty string
+// when neither key is present.
+func getAnnotationFilepath(annotations map[string]string) string {
+	if annotations == nil {
+		return ""
+	}
+	if path := annotations[modelspec.AnnotationFilepath]; path != "" {
+		return path
+	}
+	return annotations[legacymodelspec.AnnotationFilepath]
 }

pkg/backend/build.go

@@ -23,7 +23,6 @@ import (
 	"os"
 	"path/filepath"
 
-	retry "github.com/avast/retry-go/v4"
 	modelspec "github.com/modelpack/model-spec/specs-go/v1"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/sirupsen/logrus"
@@ -35,6 +34,7 @@ import (
 	"github.com/modelpack/modctl/pkg/backend/processor"
 	"github.com/modelpack/modctl/pkg/config"
 	"github.com/modelpack/modctl/pkg/modelfile"
+	"github.com/modelpack/modctl/pkg/retrypolicy"
 	"github.com/modelpack/modctl/pkg/source"
 )
 
@@ -44,7 +44,11 @@ const (
 )
 
 // Build builds the user materials into the model artifact which follows the Model Spec.
-func (b *backend) Build(ctx context.Context, modelfilePath, workDir, target string, cfg *config.Build) error {
+func (b *backend) Build(
+	ctx context.Context,
+	modelfilePath, workDir, target string,
+	cfg *config.Build,
+) error {
 	logrus.Infof("build: building artifact %s", target)
 	// parse the repo name and tag name from target.
 	ref, err := ParseReference(target)
@@ -123,8 +127,8 @@ func (b *backend) Build(ctx context.Context, modelfilePath, workDir, target stri
 
 	var configDesc ocispec.Descriptor
 	// Build the model config.
-	if err := retry.Do(func() error {
-		configDesc, err = builder.BuildConfig(ctx, config, hooks.NewHooks(
+	if err := retrypolicy.Do(ctx, func(rctx context.Context) error {
+		configDesc, err = builder.BuildConfig(rctx, config, hooks.NewHooks(
 			hooks.WithOnStart(func(name string, size int64, reader io.Reader) io.Reader {
 				return pb.Add(internalpb.NormalizePrompt("Building config"), name, size, reader)
 			}),
@@ -136,13 +140,16 @@ func (b *backend) Build(ctx context.Context, modelfilePath, workDir, target stri
 			}),
 		))
 		return err
-	}, append(defaultRetryOpts, retry.Context(ctx))...); err != nil {
+	}, retrypolicy.DoOpts{
+		FileSize: 0, // config is small
+		FileName: "config",
+	}); err != nil {
 		return fmt.Errorf("failed to build model config: %w", err)
 	}
 
 	// Build the model manifest.
-	if err := retry.Do(func() error {
-		_, err = builder.BuildManifest(ctx, layers, configDesc, manifestAnnotation(modelfile), hooks.NewHooks(
+	if err := retrypolicy.Do(ctx, func(rctx context.Context) error {
+		_, err = builder.BuildManifest(rctx, layers, configDesc, manifestAnnotation(modelfile), hooks.NewHooks(
 			hooks.WithOnStart(func(name string, size int64, reader io.Reader) io.Reader {
 				return pb.Add(internalpb.NormalizePrompt("Building manifest"), name, size, reader)
 			}),
@@ -154,23 +161,32 @@ func (b *backend) Build(ctx context.Context, modelfilePath, workDir, target stri
 			}),
 		))
 		return err
-	}, append(defaultRetryOpts, retry.Context(ctx))...); err != nil {
+	}, retrypolicy.DoOpts{
+		FileSize: 0, // manifest is small
+		FileName: "manifest",
+	}); err != nil {
 		return fmt.Errorf("failed to build model manifest: %w", err)
 	}
 
 	logrus.Infof("build: built artifact %s", target)
 	return nil
 }
 
-func (b *backend) getProcessors(modelfile modelfile.Modelfile, cfg *config.Build) []processor.Processor {
+func (b *backend) getProcessors(
+	modelfile modelfile.Modelfile,
+	cfg *config.Build,
+) []processor.Processor {
 	processors := []processor.Processor{}
 
 	if configs := modelfile.GetConfigs(); len(configs) > 0 {
 		mediaType := modelspec.MediaTypeModelWeightConfig
 		if cfg.Raw {
 			mediaType = modelspec.MediaTypeModelWeightConfigRaw
 		}
-		processors = append(processors, processor.NewModelConfigProcessor(b.store, mediaType, configs, ""))
+		processors = append(
+			processors,
+			processor.NewModelConfigProcessor(b.store, mediaType, configs, ""),
+		)
 	}
 
 	if models := modelfile.GetModels(); len(models) > 0 {
@@ -201,10 +217,23 @@ func (b *backend) getProcessors(modelfile modelfile.Modelfile, cfg *config.Build
 }
 
 // process walks the user work directory and process the identified files.
-func (b *backend) process(ctx context.Context, builder build.Builder, workDir string, pb *internalpb.ProgressBar, cfg *config.Build, processors ...processor.Processor) ([]ocispec.Descriptor, error) {
+func (b *backend) process(
+	ctx context.Context,
+	builder build.Builder,
+	workDir string,
+	pb *internalpb.ProgressBar,
+	cfg *config.Build,
+	processors ...processor.Processor,
+) ([]ocispec.Descriptor, error) {
 	descriptors := []ocispec.Descriptor{}
 	for _, p := range processors {
-		descs, err := p.Process(ctx, builder, workDir, processor.WithConcurrency(cfg.Concurrency), processor.WithProgressTracker(pb))
+		descs, err := p.Process(
+			ctx,
+			builder,
+			workDir,
+			processor.WithConcurrency(cfg.Concurrency),
+			processor.WithProgressTracker(pb),
+		)
 		if err != nil {
 			return nil, err
 		}