fix: pass the context for retry (#180)

Signed-off-by: chlins <[email protected]>

Author: chlins

pkg/backend/build.go

@@ -129,7 +129,7 @@ func (b *backend) Build(ctx context.Context, modelfilePath, workDir, target stri
 			}),
 		))
 		return err
-	}, retryOpts...); err != nil {
+	}, append(defaultRetryOpts, retry.Context(ctx))...); err != nil {
 		return fmt.Errorf("failed to build model config: %w", err)
 	}
 
@@ -147,7 +147,7 @@ func (b *backend) Build(ctx context.Context, modelfilePath, workDir, target stri
 			}),
 		))
 		return err
-	}, retryOpts...); err != nil {
+	}, append(defaultRetryOpts, retry.Context(ctx))...); err != nil {
 		return fmt.Errorf("failed to build model manifest: %w", err)
 	}
 

pkg/backend/processor/base.go

@@ -124,7 +124,7 @@ func (b *base) Process(ctx context.Context, builder build.Builder, workDir strin
 				mu.Unlock()
 
 				return nil
-			}, retryOpts...)
+			}, append(defaultRetryOpts, retry.Context(ctx))...)
 		})
 	}
 

pkg/backend/processor/options.go

@@ -45,7 +45,7 @@ func WithProgressTracker(tracker *pb.ProgressBar) ProcessOption {
 	}
 }
 
-var retryOpts = []retry.Option{
+var defaultRetryOpts = []retry.Option{
 	retry.Attempts(3),
 	retry.DelayType(retry.BackOffDelay),
 	retry.Delay(5 * time.Second),

pkg/backend/pull.go

@@ -95,7 +95,7 @@ func (b *backend) Pull(ctx context.Context, target string, cfg *config.Pull) err
 				// call the after hook.
 				cfg.Hooks.AfterPullLayer(layer, err)
 				return err
-			}, retryOpts...)
+			}, append(defaultRetryOpts, retry.Context(ctx))...)
 		})
 	}
 
@@ -112,14 +112,14 @@ func (b *backend) Pull(ctx context.Context, target string, cfg *config.Pull) err
 	// copy the config.
 	if err := retry.Do(func() error {
 		return pullIfNotExist(ctx, pb, internalpb.NormalizePrompt("Pulling config"), src, dst, manifest.Config, repo, tag)
-	}, retryOpts...); err != nil {
+	}, append(defaultRetryOpts, retry.Context(ctx))...); err != nil {
 		return fmt.Errorf("failed to pull config to local: %w", err)
 	}
 
 	// copy the manifest.
 	if err := retry.Do(func() error {
 		return pullIfNotExist(ctx, pb, internalpb.NormalizePrompt("Pulling manifest"), src, dst, manifestDesc, repo, tag)
-	}, retryOpts...); err != nil {
+	}, append(defaultRetryOpts, retry.Context(ctx))...); err != nil {
 		return fmt.Errorf("failed to pull manifest to local: %w", err)
 	}
 

pkg/backend/push.go

@@ -79,7 +79,7 @@ func (b *backend) Push(ctx context.Context, target string, cfg *config.Push) err
 		g.Go(func() error {
 			return retry.Do(func() error {
 				return pushIfNotExist(ctx, pb, internalpb.NormalizePrompt("Copying blob"), src, dst, layer, repo, tag)
-			}, retryOpts...)
+			}, append(defaultRetryOpts, retry.Context(ctx))...)
 		})
 	}
 
@@ -90,7 +90,7 @@ func (b *backend) Push(ctx context.Context, target string, cfg *config.Push) err
 	// copy the config.
 	if err := retry.Do(func() error {
 		return pushIfNotExist(ctx, pb, internalpb.NormalizePrompt("Copying config"), src, dst, manifest.Config, repo, tag)
-	}, retryOpts...); err != nil {
+	}, append(defaultRetryOpts, retry.Context(ctx))...); err != nil {
 		return fmt.Errorf("failed to push config to remote: %w", err)
 	}
 
@@ -102,7 +102,7 @@ func (b *backend) Push(ctx context.Context, target string, cfg *config.Push) err
 			Digest:    godigest.FromBytes(manifestRaw),
 			Data:      manifestRaw,
 		}, repo, tag)
-	}, retryOpts...); err != nil {
+	}, append(defaultRetryOpts, retry.Context(ctx))...); err != nil {
 		return fmt.Errorf("failed to push manifest to remote: %w", err)
 	}
 

pkg/backend/retry.go

@@ -22,7 +22,7 @@ import (
 	retry "github.com/avast/retry-go/v4"
 )
 
-var retryOpts = []retry.Option{
+var defaultRetryOpts = []retry.Option{
 	retry.Attempts(3),
 	retry.DelayType(retry.BackOffDelay),
 	retry.Delay(5 * time.Second),