Fix error handling when parallel pull fails and context is canceled

jozu-ai · Jul 12, 2024 · 9a588e1 · 9a588e1
1 parent 12d3cfc
commit 9a588e1
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/pkg/cmd/pull/cmd.go b/pkg/cmd/pull/cmd.go
@@ -97,7 +97,7 @@ func runCommand(opts *pullOptions) func(*cobra.Command, []string) error {
 		output.Infof("Pulling %s", opts.modelRef.String())
 		desc, err := runPull(cmd.Context(), opts)
 		if err != nil {
-			return output.Fatalf("Failed to pull: %s", err)
+			return output.Fatalln(err)
 		}
 		output.Infof("Pulled %s", desc.Digest)
 		return nil

diff --git a/pkg/lib/repo/local/pull.go b/pkg/lib/repo/local/pull.go
@@ -67,12 +67,15 @@ func (l *localRepo) PullModel(ctx context.Context, src oras.ReadOnlyTarget, ref
 		if err == nil {
 			return nil
 		}
-		return fmt.Errorf("failed to get %s: %w", constants.FormatMediaTypeForUser(desc.MediaType), err)
+		return fmt.Errorf("failed to get %s layer: %w", constants.FormatMediaTypeForUser(desc.MediaType), err)
 	}
+	var semErr error
 	for _, pullDesc := range toPull {
 		pullDesc := pullDesc
 		if err := sem.Acquire(errCtx, 1); err != nil {
-			return ocispec.DescriptorEmptyJSON, fmt.Errorf("failed to acquire lock: %w", err)
+			// Save error and break to get the _actual_ error
+			semErr = err
+			break
 		}
 		errs.Go(func() error {
 			defer sem.Release(1)
@@ -82,6 +85,9 @@ func (l *localRepo) PullModel(ctx context.Context, src oras.ReadOnlyTarget, ref
 	if err := errs.Wait(); err != nil {
 		return ocispec.DescriptorEmptyJSON, err
 	}
+	if semErr != nil {
+		return ocispec.DescriptorEmptyJSON, fmt.Errorf("failed to acquire lock: %w", semErr)
+	}
 
 	// Special handling to make sure local (scoped) repo contains the just-pulled manifest
 	if err := l.localIndex.addManifest(desc); err != nil {