optimize: improve pull process and switch the sha256 library to simd (#103)

optimize: improve pull process and switch the sha256 library to sha256-simd

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

Author: chlins

cmd/root.go

@@ -18,9 +18,11 @@ package cmd
 
 import (
 	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/CloudNativeAI/modctl/pkg/config"
-	"github.com/sirupsen/logrus"
+
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 )
@@ -36,15 +38,21 @@ var rootCmd = &cobra.Command{
 	SilenceUsage:       true,
 	FParseErrWhitelist: cobra.FParseErrWhitelist{UnknownFlags: true},
 	RunE: func(cmd *cobra.Command, args []string) error {
-		logrus.Debug("modctl is running")
-
 		return nil
 	},
 }
 
 // Execute adds all child commands to the root command and sets flags appropriately.
 // This is called by main.main(). It only needs to happen once to the rootCmd.
 func Execute() {
+	sig := make(chan os.Signal, 1)
+	signal.Notify(sig, os.Interrupt, syscall.SIGTERM)
+
+	go func() {
+		<-sig
+		os.Exit(1)
+	}()
+
 	if err := rootCmd.Execute(); err != nil {
 		os.Exit(1)
 	}

go.mod

@@ -9,9 +9,9 @@ require (
 	github.com/distribution/reference v0.6.0
 	github.com/dustin/go-humanize v1.0.1
 	github.com/emirpasic/gods v1.18.1
+	github.com/minio/sha256-simd v1.0.1
 	github.com/opencontainers/go-digest v1.0.0
 	github.com/opencontainers/image-spec v1.1.0
-	github.com/sirupsen/logrus v1.9.3
 	github.com/spf13/cobra v1.9.1
 	github.com/spf13/viper v1.19.0
 	github.com/stretchr/testify v1.10.0
@@ -40,6 +40,7 @@ require (
 	github.com/hashicorp/hcl v1.0.1-vault-5 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/klauspost/compress v1.17.11 // indirect
+	github.com/klauspost/cpuid/v2 v2.2.3 // 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
@@ -55,6 +56,7 @@ require (
 	github.com/rivo/uniseg v0.4.7 // indirect
 	github.com/sagikazarmark/locafero v0.6.0 // indirect
 	github.com/sagikazarmark/slog-shim v0.1.0 // indirect
+	github.com/sirupsen/logrus v1.9.3 // indirect
 	github.com/sourcegraph/conc v0.3.0 // indirect
 	github.com/spf13/afero v1.11.0 // indirect
 	github.com/spf13/cast v1.7.0 // indirect

go.sum

@@ -76,6 +76,8 @@ github.com/json-iterator/go v1.1.7/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/u
 github.com/julienschmidt/httprouter v1.2.0/go.mod h1:SYymIcj16QtmaHHD7aYtjjsJG7VTCxuUUipMqKk8s4w=
 github.com/klauspost/compress v1.17.11 h1:In6xLpyWOi1+C7tXUUWv2ot1QvBjxevKAaI6IXrJmUc=
 github.com/klauspost/compress v1.17.11/go.mod h1:pMDklpSncoRMuLFrf1W9Ss9KT+0rH90U12bZKk7uwG0=
+github.com/klauspost/cpuid/v2 v2.2.3 h1:sxCkb+qR91z4vsqw4vGGZlDgPz3G7gjaLyK3V8y70BU=
+github.com/klauspost/cpuid/v2 v2.2.3/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY=
 github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
 github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
@@ -93,6 +95,8 @@ github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D
 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=
+github.com/minio/sha256-simd v1.0.1 h1:6kaan5IFmwTNynnKKpDHe6FWHohJOHhCPchzK49dzMM=
+github.com/minio/sha256-simd v1.0.1/go.mod h1:Pz6AKMiUdngCLpeTL/RJY1M9rUuPMYujV5xJjtbRSN8=
 github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY=
 github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo=
 github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
@@ -238,6 +242,7 @@ golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5h
 golang.org/x/sys v0.0.0-20181116152217-5ac8a444bdc5/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 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-20220704084225-05e143d24a9e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 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=

pkg/backend/build/build.go

@@ -56,7 +56,7 @@ func BuildLayer(ctx context.Context, store storage.Storage, mediaType, workDir,
 		return ocispec.Descriptor{}, fmt.Errorf("failed to tar file: %w", err)
 	}
 
-	digest, size, err := store.PushBlob(ctx, repo, reader)
+	digest, size, err := store.PushBlob(ctx, repo, reader, ocispec.Descriptor{})
 	if err != nil {
 		return ocispec.Descriptor{}, fmt.Errorf("failed to push blob to storage: %w", err)
 	}
@@ -118,7 +118,7 @@ func BuildConfig(ctx context.Context, store storage.Storage, modelfile modelfile
 		return ocispec.Descriptor{}, fmt.Errorf("failed to marshal config: %w", err)
 	}
 
-	digest, size, err := store.PushBlob(ctx, repo, bytes.NewReader(configJSON))
+	digest, size, err := store.PushBlob(ctx, repo, bytes.NewReader(configJSON), ocispec.Descriptor{})
 	if err != nil {
 		return ocispec.Descriptor{}, fmt.Errorf("failed to push config to storage: %w", err)
 	}

pkg/backend/processor/code_test.go

@@ -54,7 +54,7 @@ func (s *codeProcessorSuite) TestName() {
 func (s *codeProcessorSuite) TestProcess() {
 	ctx := context.Background()
 	repo := "test-repo"
-	s.mockStore.On("PushBlob", ctx, repo, mock.Anything).Return("sha256:1234567890abcdef", int64(1024), nil)
+	s.mockStore.On("PushBlob", ctx, repo, mock.Anything, mock.Anything).Return("sha256:1234567890abcdef", int64(1024), nil)
 
 	desc, err := s.processor.Process(ctx, s.workDir, repo)
 	assert.NoError(s.Suite.T(), err)

pkg/backend/processor/doc_test.go

@@ -54,7 +54,7 @@ func (s *docProcessorSuite) TestName() {
 func (s *docProcessorSuite) TestProcess() {
 	ctx := context.Background()
 	repo := "test-repo"
-	s.mockStore.On("PushBlob", ctx, repo, mock.Anything).Return("sha256:1234567890abcdef", int64(1024), nil)
+	s.mockStore.On("PushBlob", ctx, repo, mock.Anything, mock.Anything).Return("sha256:1234567890abcdef", int64(1024), nil)
 
 	desc, err := s.processor.Process(ctx, s.workDir, repo)
 	assert.NoError(s.Suite.T(), err)

pkg/backend/processor/model_config_test.go

@@ -54,7 +54,7 @@ func (s *modelConfigProcessorSuite) TestName() {
 func (s *modelConfigProcessorSuite) TestProcess() {
 	ctx := context.Background()
 	repo := "test-repo"
-	s.mockStore.On("PushBlob", ctx, repo, mock.Anything).Return("sha256:1234567890abcdef", int64(1024), nil)
+	s.mockStore.On("PushBlob", ctx, repo, mock.Anything, mock.Anything).Return("sha256:1234567890abcdef", int64(1024), nil)
 
 	desc, err := s.processor.Process(ctx, s.workDir, repo)
 	assert.NoError(s.Suite.T(), err)

pkg/backend/processor/model_test.go

@@ -54,7 +54,7 @@ func (s *modelProcessorSuite) TestName() {
 func (s *modelProcessorSuite) TestProcess() {
 	ctx := context.Background()
 	repo := "test-repo"
-	s.mockStore.On("PushBlob", ctx, repo, mock.Anything).Return("sha256:1234567890abcdef", int64(1024), nil)
+	s.mockStore.On("PushBlob", ctx, repo, mock.Anything, mock.Anything).Return("sha256:1234567890abcdef", int64(1024), nil)
 
 	desc, err := s.processor.Process(ctx, s.workDir, repo)
 	assert.NoError(s.Suite.T(), err)

pkg/backend/pull.go

@@ -172,7 +172,7 @@ func pullIfNotExist(ctx context.Context, pb *ProgressBar, prompt string, src *re
 			return err
 		}
 	} else {
-		if _, _, err := dst.PushBlob(ctx, repo, pb.Add(prompt, desc, content)); err != nil {
+		if _, _, err := dst.PushBlob(ctx, repo, pb.Add(prompt, desc, content), desc); err != nil {
 			pb.Abort(desc)
 			return err
 		}

pkg/storage/distribution/distribution.go

@@ -18,7 +18,6 @@ package distribution
 
 import (
 	"context"
-	"crypto/sha256"
 	"fmt"
 	"io"
 	"regexp"
@@ -28,6 +27,7 @@ import (
 	"github.com/distribution/distribution/v3/registry/storage/driver"
 	"github.com/distribution/distribution/v3/registry/storage/driver/filesystem"
 	ref "github.com/distribution/reference"
+	sha256 "github.com/minio/sha256-simd"
 	godigest "github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 )
@@ -172,30 +172,35 @@ func (s *storage) PullBlob(ctx context.Context, repo, digest string) (io.ReadClo
 }
 
 // PushBlob pushes the blob to the storage.
-func (s *storage) PushBlob(ctx context.Context, repo string, blobReader io.Reader) (string, int64, error) {
+func (s *storage) PushBlob(ctx context.Context, repo string, blobReader io.Reader, provisional ocispec.Descriptor) (string, int64, error) {
 	repository, err := s.repository(ctx, repo)
 	if err != nil {
 		return "", 0, err
 	}
 
-	// use teeReader to calculate the digest.
 	hash := sha256.New()
-	teeReader := io.TeeReader(blobReader, hash)
+	if provisional.Digest == "" {
+		blobReader = io.TeeReader(blobReader, hash)
+	}
 
 	blob, err := repository.Blobs(ctx).Create(ctx)
 	if err != nil {
 		return "", 0, err
 	}
 
-	size, err := blob.ReadFrom(teeReader)
+	size, err := blob.ReadFrom(blobReader)
 	if err != nil {
 		return "", 0, err
 	}
 
-	desc, err := blob.Commit(ctx, ocispec.Descriptor{
-		Digest: godigest.Digest(fmt.Sprintf("sha256:%x", hash.Sum(nil))),
-		Size:   size,
-	})
+	// if the provided provisional descriptor is not empty, we can just use it to commit,
+	// otherwise we need to calculate the digest.
+	if provisional.Digest == "" {
+		provisional.Digest = godigest.Digest(fmt.Sprintf("sha256:%x", hash.Sum(nil)))
+		provisional.Size = size
+	}
+
+	desc, err := blob.Commit(ctx, provisional)
 	if err != nil {
 		return "", 0, nil
 	}