feat: support the concurrency for extract (#136)

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

Author: chlins

Makefile

@@ -98,7 +98,7 @@ LOCALBIN ?= $(shell pwd)/bin
 $(LOCALBIN):
 	mkdir -p $(LOCALBIN)
 
-MOCKERY_VERSION=v2.52.1
+MOCKERY_VERSION=v2.53.3
 
 .PHONY: gen
 gen: gen-mockery## Generate all we need!

cmd/extract.go

@@ -38,6 +38,10 @@ var extractCmd = &cobra.Command{
 	SilenceUsage:       true,
 	FParseErrWhitelist: cobra.FParseErrWhitelist{UnknownFlags: true},
 	RunE: func(cmd *cobra.Command, args []string) error {
+		if err := extractConfig.Validate(); err != nil {
+			return err
+		}
+
 		return runExtract(context.Background(), args[0])
 	},
 }
@@ -46,6 +50,7 @@ var extractCmd = &cobra.Command{
 func init() {
 	flags := extractCmd.Flags()
 	flags.StringVar(&extractConfig.Output, "output", "", "specify the output for extracting the model artifact")
+	flags.IntVar(&extractConfig.Concurrency, "concurrency", extractConfig.Concurrency, "specify the concurrency for extracting the model artifact")
 
 	if err := viper.BindPFlags(flags); err != nil {
 		panic(fmt.Errorf("bind cache extract flags to viper: %w", err))
@@ -63,11 +68,7 @@ func runExtract(ctx context.Context, target string) error {
 		return fmt.Errorf("target is required")
 	}
 
-	if extractConfig.Output == "" {
-		return fmt.Errorf("output is required")
-	}
-
-	if err := b.Extract(ctx, target, extractConfig.Output); err != nil {
+	if err := b.Extract(ctx, target, extractConfig); err != nil {
 		return err
 	}
 

go.mod

@@ -1,7 +1,6 @@
 module github.com/CloudNativeAI/modctl
 
-go 1.23.3
-toolchain go1.24.1
+go 1.24.1
 
 require (
 	github.com/CloudNativeAI/model-spec v0.0.3

pkg/backend/backend.go

@@ -53,7 +53,7 @@ type Backend interface {
 	Inspect(ctx context.Context, target string) (*InspectedModelArtifact, error)
 
 	// Extract extracts the model artifact.
-	Extract(ctx context.Context, target string, output string) error
+	Extract(ctx context.Context, target string, cfg *config.Extract) error
 
 	// Tag creates a new tag that refers to the source model artifact.
 	Tag(ctx context.Context, source, target string) error

pkg/backend/extract.go

@@ -24,8 +24,10 @@ import (
 	"io"
 
 	"github.com/CloudNativeAI/modctl/pkg/codec"
+	"github.com/CloudNativeAI/modctl/pkg/config"
 	"github.com/CloudNativeAI/modctl/pkg/storage"
 	modelspec "github.com/CloudNativeAI/model-spec/specs-go/v1"
+	"golang.org/x/sync/errgroup"
 
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 )
@@ -36,7 +38,7 @@ const (
 )
 
 // Extract extracts the model artifact.
-func (b *backend) Extract(ctx context.Context, target string, output string) error {
+func (b *backend) Extract(ctx context.Context, target string, cfg *config.Extract) error {
 	// parse the repository and tag from the target.
 	ref, err := ParseReference(target)
 	if err != nil {
@@ -55,26 +57,33 @@ func (b *backend) Extract(ctx context.Context, target string, output string) err
 		return fmt.Errorf("failed to unmarshal the manifest: %w", err)
 	}
 
-	return exportModelArtifact(ctx, b.store, manifest, repo, output)
+	return exportModelArtifact(ctx, b.store, manifest, repo, cfg)
 }
 
 // exportModelArtifact exports the target model artifact to the output directory, which will open the artifact and extract to restore the original repo structure.
-func exportModelArtifact(ctx context.Context, store storage.Storage, manifest ocispec.Manifest, repo, outputDir string) error {
+func exportModelArtifact(ctx context.Context, store storage.Storage, manifest ocispec.Manifest, repo string, cfg *config.Extract) error {
+	g := &errgroup.Group{}
+	g.SetLimit(cfg.Concurrency)
+
 	for _, layer := range manifest.Layers {
-		// pull the blob from the storage.
-		reader, err := store.PullBlob(ctx, repo, layer.Digest.String())
-		if err != nil {
-			return fmt.Errorf("failed to pull the blob from storage: %w", err)
-		}
-		defer reader.Close()
-
-		bufferedReader := bufio.NewReaderSize(reader, defaultBufferSize)
-		if err := extractLayer(layer, outputDir, bufferedReader); err != nil {
-			return fmt.Errorf("failed to extract layer %s: %w", layer.Digest.String(), err)
-		}
+		g.Go(func() error {
+			// pull the blob from the storage.
+			reader, err := store.PullBlob(ctx, repo, layer.Digest.String())
+			if err != nil {
+				return fmt.Errorf("failed to pull the blob from storage: %w", err)
+			}
+			defer reader.Close()
+
+			bufferedReader := bufio.NewReaderSize(reader, defaultBufferSize)
+			if err := extractLayer(layer, cfg.Output, bufferedReader); err != nil {
+				return fmt.Errorf("failed to extract layer %s: %w", layer.Digest.String(), err)
+			}
+
+			return nil
+		})
 	}
 
-	return nil
+	return g.Wait()
 }
 
 // extractLayer extracts the layer to the output directory.

pkg/backend/pull.go

@@ -150,7 +150,9 @@ func (b *backend) Pull(ctx context.Context, target string, cfg *config.Pull) err
 
 	// export the target model artifact to the output directory if needed.
 	if cfg.ExtractDir != "" {
-		if err := exportModelArtifact(ctx, dst, manifest, repo, cfg.ExtractDir); err != nil {
+		// set the concurrency to 1 because the pull already has concurrency control.
+		extractCfg := &config.Extract{Concurrency: 1, Output: cfg.ExtractDir}
+		if err := exportModelArtifact(ctx, dst, manifest, repo, extractCfg); err != nil {
 			return fmt.Errorf("failed to export the artifact to the output directory: %w", err)
 		}
 	}

pkg/config/extract.go

@@ -16,12 +16,33 @@
 
 package config
 
+import "fmt"
+
+const (
+	// defaultExtractConcurrency is the default number of concurrent extracts.
+	defaultExtractConcurrency = 5
+)
+
 type Extract struct {
-	Output string
+	Output      string
+	Concurrency int
 }
 
 func NewExtract() *Extract {
 	return &Extract{
-		Output: "",
+		Output:      "",
+		Concurrency: defaultExtractConcurrency,
+	}
+}
+
+func (e *Extract) Validate() error {
+	if e.Concurrency <= 0 {
+		return fmt.Errorf("concurrency must be greater than 0")
+	}
+
+	if e.Output == "" {
+		return fmt.Errorf("output is required")
 	}
+
+	return nil
 }