feat: support concurrent for build operation(5 concurrency by default) (#98)

Author: chlins

cmd/build.go

@@ -49,6 +49,7 @@ var buildCmd = &cobra.Command{
 // init initializes build command.
 func init() {
 	flags := buildCmd.Flags()
+	flags.IntVarP(&buildConfig.Concurrency, "concurrency", "c", buildConfig.Concurrency, "specify the number of concurrent build operations")
 	flags.StringVarP(&buildConfig.Target, "target", "t", "", "target model artifact name")
 	flags.StringVarP(&buildConfig.Modelfile, "modelfile", "f", "Modelfile", "model file path")
 
@@ -64,7 +65,11 @@ func runBuild(ctx context.Context, workDir string) error {
 		return err
 	}
 
-	if err := b.Build(ctx, buildConfig.Modelfile, workDir, buildConfig.Target); err != nil {
+	opts := []backend.Option{
+		backend.WithConcurrency(buildConfig.Concurrency),
+	}
+
+	if err := b.Build(ctx, buildConfig.Modelfile, workDir, buildConfig.Target, opts...); err != nil {
 		return err
 	}
 

go.mod

@@ -4,6 +4,7 @@ go 1.23.3
 
 require (
 	github.com/CloudNativeAI/model-spec v0.0.2
+	github.com/chelnak/ysmrr v0.6.0
 	github.com/distribution/distribution/v3 v3.0.0-rc.3
 	github.com/distribution/reference v0.6.0
 	github.com/dustin/go-humanize v1.0.1
@@ -28,6 +29,7 @@ require (
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	github.com/docker/go-metrics v0.0.1 // indirect
+	github.com/fatih/color v1.18.0 // indirect
 	github.com/fsnotify/fsnotify v1.7.0 // indirect
 	github.com/go-logr/logr v1.4.2 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
@@ -39,6 +41,8 @@ require (
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/klauspost/compress v1.17.11 // indirect
 	github.com/magiconair/properties v1.8.7 // indirect
+	github.com/mattn/go-colorable v0.1.14 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/mattn/go-runewidth v0.0.16 // indirect
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect

go.sum

@@ -14,6 +14,8 @@ github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK3
 github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
 github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
 github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/chelnak/ysmrr v0.6.0 h1:kMhO0oI02tl/9szvxrOE0yeImtrK4KQhER0oXu1K/iM=
+github.com/chelnak/ysmrr v0.6.0/go.mod h1:56JSrmQgb7/7xoMvuD87h3PE/qW6K1+BQcrgWtVLTUo=
 github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -29,6 +31,8 @@ github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkp
 github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
 github.com/emirpasic/gods v1.18.1 h1:FXtiHYKDGKCW2KzwZKx0iC0PQmdlorYgdFG9jPXJ1Bc=
 github.com/emirpasic/gods v1.18.1/go.mod h1:8tpGGwCnJ5H4r6BWwaV6OrWmMoPhUl5jm/FMNAnJvWQ=
+github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM=
+github.com/fatih/color v1.18.0/go.mod h1:4FelSpRwEGDpQ12mAdzqdOukCy4u8WUtOY6lkT/6HfU=
 github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
 github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0=
 github.com/fsnotify/fsnotify v1.7.0 h1:8JEhPFa5W2WU7YfeZzPNqzMP6Lwt7L2715Ggo0nosvA=
@@ -82,6 +86,10 @@ github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0
 github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
 github.com/magiconair/properties v1.8.7 h1:IeQXZAiQcpL9mgcAe1Nu6cX9LLw6ExEHKjN0VQdvPDY=
 github.com/magiconair/properties v1.8.7/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0=
+github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE=
+github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8=
+github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
+github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
 github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc=
 github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
 github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
@@ -231,6 +239,7 @@ golang.org/x/sys v0.0.0-20181116152217-5ac8a444bdc5/go.mod h1:STP8DvDyc/dI5b8T5h
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190801041406-cbf593c0f2f3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc=
 golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/term v0.29.0 h1:L6pJp37ocefwRRtYPKSWOWzOtWSxVajvz2ldH/xi3iU=

pkg/backend/backend.go

@@ -31,7 +31,7 @@ type Backend interface {
 	Logout(ctx context.Context, registry string) error
 
 	// Build builds the user materials into the OCI image which follows the Model Spec.
-	Build(ctx context.Context, modelfilePath, workDir, target string) error
+	Build(ctx context.Context, modelfilePath, workDir, target string, opts ...Option) error
 
 	// Pull pulls an artifact from a registry.
 	Pull(ctx context.Context, target string, opts ...Option) error

pkg/backend/build.go

@@ -30,7 +30,13 @@ import (
 )
 
 // Build builds the user materials into the OCI image which follows the Model Spec.
-func (b *backend) Build(ctx context.Context, modelfilePath, workDir, target string) error {
+func (b *backend) Build(ctx context.Context, modelfilePath, workDir, target string, opts ...Option) error {
+	// apply options.
+	options := &Options{}
+	for _, opt := range opts {
+		opt(options)
+	}
+
 	// parse the repo name and tag name from target.
 	ref, err := ParseReference(target)
 	if err != nil {
@@ -48,7 +54,7 @@ func (b *backend) Build(ctx context.Context, modelfilePath, workDir, target stri
 	}
 
 	layers := []ocispec.Descriptor{}
-	layerDescs, err := b.process(ctx, workDir, repo, b.getProcessors(modelfile)...)
+	layerDescs, err := b.process(ctx, workDir, repo, options, b.getProcessors(modelfile)...)
 	if err != nil {
 		return fmt.Errorf("failed to process files: %w", err)
 	}
@@ -61,15 +67,15 @@ func (b *backend) Build(ctx context.Context, modelfilePath, workDir, target stri
 		return fmt.Errorf("failed to build image config: %w", err)
 	}
 
-	fmt.Printf("%-15s => %s (%s)\n", "Built config", configDesc.Digest, humanize.IBytes(uint64(configDesc.Size)))
+	fmt.Printf("%s => %s (%s)\n", "Built config", configDesc.Digest, humanize.IBytes(uint64(configDesc.Size)))
 
 	// build the image manifest.
 	manifestDesc, err := build.BuildManifest(ctx, b.store, repo, tag, layers, configDesc, manifestAnnotation())
 	if err != nil {
 		return fmt.Errorf("failed to build image manifest: %w", err)
 	}
 
-	fmt.Printf("%-15s => %s (%s)\n", "Built manifest", manifestDesc.Digest, humanize.IBytes(uint64(manifestDesc.Size)))
+	fmt.Printf("%s => %s (%s)\n", "Built manifest", manifestDesc.Digest, humanize.IBytes(uint64(manifestDesc.Size)))
 	return nil
 }
 
@@ -96,10 +102,10 @@ func (b *backend) getProcessors(modelfile modelfile.Modelfile) []processor.Proce
 }
 
 // process walks the user work directory and process the identified files.
-func (b *backend) process(ctx context.Context, workDir string, repo string, processors ...processor.Processor) ([]ocispec.Descriptor, error) {
+func (b *backend) process(ctx context.Context, workDir string, repo string, opts *Options, processors ...processor.Processor) ([]ocispec.Descriptor, error) {
 	descriptors := []ocispec.Descriptor{}
 	for _, p := range processors {
-		descs, err := p.Process(ctx, workDir, repo)
+		descs, err := p.Process(ctx, workDir, repo, processor.WithConcurrency(opts.concurrency))
 		if err != nil {
 			return nil, err
 		}

pkg/backend/build_test.go

@@ -35,7 +35,7 @@ func TestGetProcessors(t *testing.T) {
 	processors := b.getProcessors(modelfile)
 
 	assert.Len(t, processors, 4)
-	assert.Equal(t, "model_config", processors[0].Name())
+	assert.Equal(t, "config", processors[0].Name())
 	assert.Equal(t, "model", processors[1].Name())
 	assert.Equal(t, "code", processors[2].Name())
 	assert.Equal(t, "doc", processors[3].Name())

pkg/backend/processor/base.go

@@ -21,15 +21,21 @@ import (
 	"fmt"
 	"path/filepath"
 	"sort"
+	"sync"
+	"sync/atomic"
 
 	"github.com/CloudNativeAI/modctl/pkg/backend/build"
 	"github.com/CloudNativeAI/modctl/pkg/storage"
 
+	"github.com/chelnak/ysmrr"
 	humanize "github.com/dustin/go-humanize"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+	"golang.org/x/sync/errgroup"
 )
 
 type base struct {
+	// name is the name of the processor.
+	name string
 	// store is the underlying storage backend.
 	store storage.Storage
 	// mediaType is the media type of the processed content.
@@ -39,7 +45,12 @@ type base struct {
 }
 
 // Process implements the Processor interface, which can be reused by other processors.
-func (b *base) Process(ctx context.Context, workDir, repo string) ([]ocispec.Descriptor, error) {
+func (b *base) Process(ctx context.Context, workDir, repo string, opts ...Option) ([]ocispec.Descriptor, error) {
+	baseOpts := &options{}
+	for _, opt := range opts {
+		opt(baseOpts)
+	}
+
 	absWorkDir, err := filepath.Abs(workDir)
 	if err != nil {
 		return nil, err
@@ -57,16 +68,55 @@ func (b *base) Process(ctx context.Context, workDir, repo string) ([]ocispec.Des
 
 	sort.Strings(matchedPaths)
 
-	var descriptors []ocispec.Descriptor
+	var (
+		idx         atomic.Int64
+		mu          sync.Mutex
+		eg          errgroup.Group
+		descriptors []ocispec.Descriptor
+	)
+
+	// Set default concurrency limit to 1 if not specified.
+	if baseOpts.concurrency > 0 {
+		eg.SetLimit(baseOpts.concurrency)
+	} else {
+		eg.SetLimit(1)
+	}
+
+	total := int64(len(matchedPaths))
+	sm := ysmrr.NewSpinnerManager()
+	sm.Start()
+
 	for _, path := range matchedPaths {
-		desc, err := build.BuildLayer(ctx, b.store, b.mediaType, workDir, repo, path)
-		if err != nil {
-			return nil, err
-		}
+		eg.Go(func() error {
+			relPath, err := filepath.Rel(absWorkDir, path)
+			if err != nil {
+				return err
+			}
 
-		fmt.Printf("%-15s => %s (%s)\n", "Built blob", desc.Digest, humanize.IBytes(uint64(desc.Size)))
-		descriptors = append(descriptors, desc)
+			blobMsg := fmt.Sprintf("blob [%s %d/%d]", b.name, idx.Add(1), total)
+			sp := sm.AddSpinner(fmt.Sprintf("Building %s => %s", blobMsg, relPath))
+
+			desc, err := build.BuildLayer(ctx, b.store, b.mediaType, workDir, repo, path)
+			if err != nil {
+				sp.ErrorWithMessagef("Failed to build blob %s: %v", path, relPath)
+				return err
+			}
+
+			sp.CompleteWithMessagef("%s => %s (%s)", fmt.Sprintf("Built %s", blobMsg), desc.Digest, humanize.IBytes(uint64(desc.Size)))
+
+			mu.Lock()
+			descriptors = append(descriptors, desc)
+			mu.Unlock()
+
+			return nil
+		})
 	}
 
+	if err := eg.Wait(); err != nil {
+		return nil, err
+	}
+
+	sm.Stop()
+
 	return descriptors, nil
 }

pkg/backend/processor/code.go

@@ -24,10 +24,15 @@ import (
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 )
 
+const (
+	codeProcessorName = "code"
+)
+
 // NewCodeProcessor creates a new code processor.
 func NewCodeProcessor(store storage.Storage, mediaType string, patterns []string) Processor {
 	return &codeProcessor{
 		base: &base{
+			name:      codeProcessorName,
 			store:     store,
 			mediaType: mediaType,
 			patterns:  patterns,
@@ -41,9 +46,9 @@ type codeProcessor struct {
 }
 
 func (p *codeProcessor) Name() string {
-	return "code"
+	return codeProcessorName
 }
 
-func (p *codeProcessor) Process(ctx context.Context, workDir, repo string) ([]ocispec.Descriptor, error) {
-	return p.base.Process(ctx, workDir, repo)
+func (p *codeProcessor) Process(ctx context.Context, workDir, repo string, opts ...Option) ([]ocispec.Descriptor, error) {
+	return p.base.Process(ctx, workDir, repo, opts...)
 }

pkg/backend/processor/doc.go

@@ -20,13 +20,19 @@ import (
 	"context"
 
 	"github.com/CloudNativeAI/modctl/pkg/storage"
+
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 )
 
+const (
+	docProcessorName = "doc"
+)
+
 // NewDocProcessor creates a new doc processor.
 func NewDocProcessor(store storage.Storage, mediaType string, patterns []string) Processor {
 	return &docProcessor{
 		base: &base{
+			name:      docProcessorName,
 			store:     store,
 			mediaType: mediaType,
 			patterns:  patterns,
@@ -40,9 +46,9 @@ type docProcessor struct {
 }
 
 func (p *docProcessor) Name() string {
-	return "doc"
+	return docProcessorName
 }
 
-func (p *docProcessor) Process(ctx context.Context, workDir, repo string) ([]ocispec.Descriptor, error) {
-	return p.base.Process(ctx, workDir, repo)
+func (p *docProcessor) Process(ctx context.Context, workDir, repo string, opts ...Option) ([]ocispec.Descriptor, error) {
+	return p.base.Process(ctx, workDir, repo, opts...)
 }

pkg/backend/processor/model.go

@@ -24,10 +24,15 @@ import (
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 )
 
+const (
+	modelProcessorName = "model"
+)
+
 // NewModelProcessor creates a new model processor.
 func NewModelProcessor(store storage.Storage, mediaType string, patterns []string) Processor {
 	return &modelProcessor{
 		base: &base{
+			name:      modelProcessorName,
 			store:     store,
 			mediaType: mediaType,
 			patterns:  patterns,
@@ -41,9 +46,9 @@ type modelProcessor struct {
 }
 
 func (p *modelProcessor) Name() string {
-	return "model"
+	return modelProcessorName
 }
 
-func (p *modelProcessor) Process(ctx context.Context, workDir, repo string) ([]ocispec.Descriptor, error) {
-	return p.base.Process(ctx, workDir, repo)
+func (p *modelProcessor) Process(ctx context.Context, workDir, repo string, opts ...Option) ([]ocispec.Descriptor, error) {
+	return p.base.Process(ctx, workDir, repo, opts...)
 }