feat(backend): use retrypolicy.Do with independent retry, remove cascading cancellation

Signed-off-by: Zhao Chen <winters.zc@antgroup.com>

Author: aftersnow

internal/pb/pb.go

@@ -136,6 +136,28 @@ 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.
+func (p *ProgressBar) Placeholder(name string, prompt string, size int64) {
+	if disableProgress {
+		return
+	}
+
+	p.mu.RLock()
+	existing := p.bars[name]
+	p.mu.RUnlock()
+
+	// If the bar already exists, just reset its message.
+	if existing != nil {
+		existing.msg = fmt.Sprintf("%s %s", prompt, name)
+		existing.Bar.SetCurrent(0)
+		return
+	}
+
+	// Create a new placeholder bar.
+	p.Add(prompt, name, size, nil)
+}
+
 // Get returns the progress bar.
 func (p *ProgressBar) Get(name string) *progressBar {
 	p.mu.RLock()

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"
 )
 
@@ -123,8 +123,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 +136,17 @@ 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",
+		Config:   &cfg.RetryConfig,
+	}); 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,7 +158,11 @@ 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",
+		Config:   &cfg.RetryConfig,
+	}); err != nil {
 		return fmt.Errorf("failed to build model manifest: %w", err)
 	}
 
@@ -204,7 +212,7 @@ func (b *backend) getProcessors(modelfile modelfile.Modelfile, cfg *config.Build
 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), processor.WithRetryConfig(cfg.RetryConfig))
 		if err != nil {
 			return nil, err
 		}

pkg/backend/fetch.go

@@ -19,7 +19,10 @@ package backend
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
+	"sync"
+	"time"
 
 	"github.com/bmatcuk/doublestar/v4"
 	legacymodelspec "github.com/dragonflyoss/model-spec/specs-go/v1"
@@ -31,6 +34,7 @@ import (
 	internalpb "github.com/modelpack/modctl/internal/pb"
 	"github.com/modelpack/modctl/pkg/backend/remote"
 	"github.com/modelpack/modctl/pkg/config"
+	"github.com/modelpack/modctl/pkg/retrypolicy"
 )
 
 // Fetch fetches partial files to the output.
@@ -101,9 +105,12 @@ func (b *backend) Fetch(ctx context.Context, target string, cfg *config.Fetch) e
 	pb.Start()
 	defer pb.Stop()
 
-	g, ctx := errgroup.WithContext(ctx)
+	g := new(errgroup.Group)
 	g.SetLimit(cfg.Concurrency)
 
+	var mu sync.Mutex
+	var errs []error
+
 	logrus.Infof("fetch: fetching %d matched layers", len(layers))
 	for _, layer := range layers {
 		g.Go(func() error {
@@ -113,17 +120,42 @@ func (b *backend) Fetch(ctx context.Context, target string, cfg *config.Fetch) e
 			default:
 			}
 
-			logrus.Debugf("fetch: processing layer %s", layer.Digest)
-			if err := pullAndExtractFromRemote(ctx, pb, internalpb.NormalizePrompt("Fetching blob"), client, cfg.Output, layer); err != nil {
-				return err
+			var annoFilepath string
+			if layer.Annotations != nil {
+				if layer.Annotations[modelspec.AnnotationFilepath] != "" {
+					annoFilepath = layer.Annotations[modelspec.AnnotationFilepath]
+				} else {
+					annoFilepath = layer.Annotations[legacymodelspec.AnnotationFilepath]
+				}
 			}
 
-			logrus.Debugf("fetch: successfully processed layer %s", layer.Digest)
+			logrus.Debugf("fetch: processing layer %s", layer.Digest)
+			if err := retrypolicy.Do(ctx, func(ctx context.Context) error {
+				return pullAndExtractFromRemote(ctx, pb, internalpb.NormalizePrompt("Fetching blob"), client, cfg.Output, layer)
+			}, retrypolicy.DoOpts{
+				FileSize: layer.Size,
+				FileName: annoFilepath,
+				Config:   &cfg.RetryConfig,
+				OnRetry: func(attempt uint, reason string, backoff time.Duration) {
+					if bar := pb.Get(layer.Digest.String()); bar != nil {
+						bar.SetRefill(bar.Current())
+						bar.SetCurrent(0)
+						bar.EwmaSetCurrent(0, time.Second)
+					}
+				},
+			}); err != nil {
+				mu.Lock()
+				errs = append(errs, err)
+				mu.Unlock()
+			} else {
+				logrus.Debugf("fetch: successfully processed layer %s", layer.Digest)
+			}
 			return nil
 		})
 	}
 
-	if err := g.Wait(); err != nil {
+	_ = g.Wait()
+	if err := errors.Join(errs...); err != nil {
 		return err
 	}
 

pkg/backend/fetch_by_d7y.go

@@ -19,15 +19,17 @@ package backend
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
+	"time"
 
 	common "d7y.io/api/v2/pkg/apis/common/v2"
 	dfdaemon "d7y.io/api/v2/pkg/apis/dfdaemon/v2"
-	"github.com/avast/retry-go/v4"
 	"github.com/bmatcuk/doublestar/v4"
 	legacymodelspec "github.com/dragonflyoss/model-spec/specs-go/v1"
 	modelspec "github.com/modelpack/model-spec/specs-go/v1"
@@ -41,6 +43,7 @@ import (
 	"github.com/modelpack/modctl/pkg/archiver"
 	"github.com/modelpack/modctl/pkg/backend/remote"
 	"github.com/modelpack/modctl/pkg/config"
+	"github.com/modelpack/modctl/pkg/retrypolicy"
 )
 
 // fetchByDragonfly fetches partial files via Dragonfly gRPC service based on pattern matching.
@@ -124,9 +127,12 @@ func (b *backend) fetchByDragonfly(ctx context.Context, target string, cfg *conf
 	defer pb.Stop()
 
 	// Process layers concurrently.
-	g, ctx := errgroup.WithContext(ctx)
+	g := new(errgroup.Group)
 	g.SetLimit(cfg.Concurrency)
 
+	var mu sync.Mutex
+	var errs []error
+
 	logrus.Infof("fetch: fetching %d matched layers via dragonfly", len(layers))
 	for _, layer := range layers {
 		g.Go(func() error {
@@ -138,14 +144,18 @@ func (b *backend) fetchByDragonfly(ctx context.Context, target string, cfg *conf
 
 			logrus.Debugf("fetch: processing layer %s via dragonfly", layer.Digest)
 			if err := fetchLayerByDragonfly(ctx, pb, dfdaemon.NewDfdaemonDownloadClient(conn), ref, manifest, layer, authToken, cfg); err != nil {
-				return err
+				mu.Lock()
+				errs = append(errs, err)
+				mu.Unlock()
+			} else {
+				logrus.Debugf("fetch: successfully processed layer %s via dragonfly", layer.Digest)
 			}
-			logrus.Debugf("fetch: successfully processed layer %s via dragonfly", layer.Digest)
 			return nil
 		})
 	}
 
-	if err := g.Wait(); err != nil {
+	_ = g.Wait()
+	if err := errors.Join(errs...); err != nil {
 		return err
 	}
 
@@ -155,7 +165,16 @@ func (b *backend) fetchByDragonfly(ctx context.Context, target string, cfg *conf
 
 // fetchLayerByDragonfly handles downloading and extracting a single layer via Dragonfly.
 func fetchLayerByDragonfly(ctx context.Context, pb *internalpb.ProgressBar, client dfdaemon.DfdaemonDownloadClient, ref Referencer, manifest ocispec.Manifest, desc ocispec.Descriptor, authToken string, cfg *config.Fetch) error {
-	err := retry.Do(func() error {
+	var annoFilepath string
+	if desc.Annotations != nil {
+		if desc.Annotations[modelspec.AnnotationFilepath] != "" {
+			annoFilepath = desc.Annotations[modelspec.AnnotationFilepath]
+		} else {
+			annoFilepath = desc.Annotations[legacymodelspec.AnnotationFilepath]
+		}
+	}
+
+	err := retrypolicy.Do(ctx, func(ctx context.Context) error {
 		logrus.Debugf("fetch: processing layer %s", desc.Digest)
 		cfg.Hooks.BeforePullLayer(desc, manifest) // Call before hook
 		err := downloadAndExtractFetchLayer(ctx, pb, client, ref, desc, authToken, cfg)
@@ -166,7 +185,18 @@ func fetchLayerByDragonfly(ctx context.Context, pb *internalpb.ProgressBar, clie
 		}
 
 		return err
-	}, append(defaultRetryOpts, retry.Context(ctx))...)
+	}, retrypolicy.DoOpts{
+		FileSize: desc.Size,
+		FileName: annoFilepath,
+		Config:   &cfg.RetryConfig,
+		OnRetry: func(attempt uint, reason string, backoff time.Duration) {
+			if bar := pb.Get(desc.Digest.String()); bar != nil {
+				bar.SetRefill(bar.Current())
+				bar.SetCurrent(0)
+				bar.EwmaSetCurrent(0, time.Second)
+			}
+		},
+	})
 
 	if err != nil {
 		err = fmt.Errorf("fetch: failed to download and extract layer %s: %w", desc.Digest, err)

pkg/backend/processor/base.go

@@ -18,6 +18,7 @@ package processor
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"os"
@@ -26,7 +27,6 @@ import (
 	"strings"
 	"sync"
 
-	"github.com/avast/retry-go/v4"
 	legacymodelspec "github.com/dragonflyoss/model-spec/specs-go/v1"
 	modelspec "github.com/modelpack/model-spec/specs-go/v1"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
@@ -36,6 +36,7 @@ import (
 	internalpb "github.com/modelpack/modctl/internal/pb"
 	"github.com/modelpack/modctl/pkg/backend/build"
 	"github.com/modelpack/modctl/pkg/backend/build/hooks"
+	"github.com/modelpack/modctl/pkg/retrypolicy"
 	"github.com/modelpack/modctl/pkg/storage"
 )
 
@@ -105,14 +106,11 @@ func (b *base) Process(ctx context.Context, builder build.Builder, workDir strin
 
 	var (
 		mu          sync.Mutex
-		eg          *errgroup.Group
 		descriptors []ocispec.Descriptor
+		errs        []error
 	)
 
-	// Initialize errgroup with a context can be canceled.
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
-	eg, ctx = errgroup.WithContext(ctx)
+	eg := new(errgroup.Group)
 
 	// Set default concurrency limit to 1 if not specified.
 	if processOpts.concurrency > 0 {
@@ -137,19 +135,25 @@ func (b *base) Process(ctx context.Context, builder build.Builder, workDir strin
 		eg.Go(func() error {
 			select {
 			case <-ctx.Done():
-				return ctx.Err()
+				return nil
 			default:
 			}
 
-			if err := retry.Do(func() error {
+			// Get file size for dynamic retry parameters.
+			var fileSize int64
+			if fi, err := os.Stat(path); err == nil {
+				fileSize = fi.Size()
+			}
+
+			if err := retrypolicy.Do(ctx, func(rctx context.Context) error {
 				logrus.Debugf("processor: processing %s file %s", b.name, path)
 
 				var destPath string
 				if b.destDir != "" {
 					destPath = filepath.Join(b.destDir, filepath.Base(path))
 				}
 
-				desc, err := builder.BuildLayer(ctx, b.mediaType, workDir, path, destPath, hooks.NewHooks(
+				desc, err := builder.BuildLayer(rctx, b.mediaType, workDir, path, destPath, hooks.NewHooks(
 					hooks.WithOnStart(func(name string, size int64, reader io.Reader) io.Reader {
 						return tracker.Add(internalpb.NormalizePrompt("Building layer"), name, size, reader)
 					}),
@@ -170,19 +174,24 @@ func (b *base) Process(ctx context.Context, builder build.Builder, workDir strin
 				mu.Unlock()
 
 				return nil
-			}, append(defaultRetryOpts, retry.Context(ctx))...); err != nil {
+			}, retrypolicy.DoOpts{
+				FileSize: fileSize,
+				FileName: filepath.Base(path),
+				Config:   processOpts.retryConfig,
+			}); err != nil {
 				logrus.Error(err)
-				// Cancel manually to abort other tasks because if one fails,
-				// we should abort all to avoid useless waiting.
-				cancel()
-				return err
+				mu.Lock()
+				errs = append(errs, err)
+				mu.Unlock()
 			}
 
 			return nil
 		})
 	}
 
-	if err := eg.Wait(); err != nil {
+	_ = eg.Wait()
+
+	if err := errors.Join(errs...); err != nil {
 		return nil, err
 	}