feat: add implemention of nydus intercept (#155)

Signed-off-by: Yang Kaiyong <[email protected]>
Co-authored-by: Yang Kaiyong <[email protected]>

Author: BraveY

pkg/backend/build/builder_test.go

@@ -19,15 +19,18 @@ package build
 import (
 	"context"
 	"errors"
+	"io"
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
 	modelspec "github.com/CloudNativeAI/model-spec/specs-go/v1"
 	godigest "github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
 
@@ -268,3 +271,18 @@ func (s *BuilderTestSuite) TestBuildModelConfig() {
 func TestBuilderSuite(t *testing.T) {
 	suite.Run(t, new(BuilderTestSuite))
 }
+
+func TestPipeReader(t *testing.T) {
+	r := strings.NewReader("some io.Reader stream to be read\n")
+	r1, r2 := splitReader(r)
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		_, err := io.Copy(os.Stdout, r2)
+		assert.NoError(t, err)
+	}()
+	_, err := io.Copy(os.Stdout, r1)
+	assert.NoError(t, err)
+	wg.Wait()
+}

pkg/backend/build/interceptor/nydus.go

@@ -17,17 +17,159 @@
 package interceptor
 
 import (
+	"archive/tar"
 	"context"
+	"encoding/json"
+	"fmt"
+	"hash/crc32"
 	"io"
+	"strings"
+
+	"github.com/CloudNativeAI/modctl/pkg/codec"
+	modelspec "github.com/CloudNativeAI/model-spec/specs-go/v1"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+)
+
+const (
+	CrcsKey              = "org.cnai.nydus.crcs"
+	DefaultFileChunkSize = 4 * 1024 * 1024
 )
 
+var mediaTypeChunkSizeMap = map[string]int{
+	modelspec.MediaTypeModelWeight:  64 * 1024 * 1024,
+	modelspec.MediaTypeModelDataset: 64 * 1024 * 1024,
+}
+
+var table = crc32.MakeTable(crc32.Castagnoli)
+
 type nydus struct{}
 
+type FileCrcList struct {
+	Files []FileCrcInfo `json:"files"`
+}
+
+type FileCrcInfo struct {
+	FilePath  string `json:"file_path"`
+	ChunkCrcs string `json:"chunk_crcs"`
+}
+
 func NewNydus() *nydus {
 	return &nydus{}
 }
 
 func (n *nydus) Intercept(ctx context.Context, mediaType string, filepath string, readerType string, reader io.Reader) (ApplyDescriptorFn, error) {
-	// TODO: Implement nydus interceptor
-	return nil, nil
+	crcsStr := ""
+	chunkSize := int64(DefaultFileChunkSize)
+	if c, ok := mediaTypeChunkSizeMap[mediaType]; ok {
+		chunkSize = int64(c)
+	}
+
+	switch readerType {
+	case codec.Tar:
+		fileCrcs, err := calcCrc32inTar(ctx, reader, chunkSize)
+		if err != nil {
+			return nil, fmt.Errorf("failed to calculate crc32 in tar: %w", err)
+		}
+		b, err := json.Marshal(fileCrcs)
+		if err != nil {
+			return nil, fmt.Errorf("failed to marshal crcs: %w", err)
+		}
+		return func(desc *ocispec.Descriptor) {
+			if desc.Annotations == nil {
+				desc.Annotations = make(map[string]string)
+			}
+			desc.Annotations[CrcsKey] = string(b)
+		}, nil
+	case codec.Raw:
+		crc32Results, err := calcCrc32(ctx, reader, chunkSize)
+		if err != nil {
+			return nil, fmt.Errorf("failed to calculate crc32: %w", err)
+		}
+		crcsStr = buildCrc32Str(crc32Results)
+		crcInfo := FileCrcInfo{
+			FilePath:  filepath,
+			ChunkCrcs: crcsStr,
+		}
+		crcs := FileCrcList{
+			Files: []FileCrcInfo{crcInfo},
+		}
+		b, err := json.Marshal(crcs)
+		if err != nil {
+			return nil, fmt.Errorf("failed to marshal crcs: %w", err)
+		}
+		return func(desc *ocispec.Descriptor) {
+			if desc.Annotations == nil {
+				desc.Annotations = make(map[string]string)
+			}
+			desc.Annotations[CrcsKey] = string(b)
+		}, nil
+	default:
+		return nil, fmt.Errorf("unsupported reader type: %s", readerType)
+	}
+}
+
+func calcCrc32inTar(ctx context.Context, r io.Reader, chunkSize int64) (*FileCrcList, error) {
+	fileCrcList := FileCrcList{
+		Files: make([]FileCrcInfo, 0),
+	}
+	tarReader := tar.NewReader(r)
+	for {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		default:
+			header, err := tarReader.Next()
+			if err == io.EOF {
+				return &fileCrcList, nil
+			}
+			if err != nil {
+				return nil, fmt.Errorf("error reading tar: %w", err)
+			}
+			if header.Typeflag == tar.TypeReg {
+				crc32Results, err := calcCrc32(ctx, tarReader, chunkSize)
+				if err != nil {
+					return nil, fmt.Errorf("failed to calculate crc32: %w", err)
+				}
+				crcsStr := buildCrc32Str(crc32Results)
+				crcInfo := FileCrcInfo{
+					FilePath:  header.Name,
+					ChunkCrcs: crcsStr,
+				}
+				fileCrcList.Files = append(fileCrcList.Files, crcInfo)
+			}
+		}
+	}
+}
+
+func calcCrc32(ctx context.Context, r io.Reader, chunkSize int64) ([]uint32, error) {
+	var crc32Results []uint32
+	for {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		default:
+			limitedReader := io.LimitReader(r, chunkSize)
+			hash := crc32.New(table)
+			n, err := io.Copy(hash, limitedReader)
+			if n == 0 || err == io.EOF {
+				return crc32Results, nil
+			}
+
+			if err != nil {
+				return nil, fmt.Errorf("failed to read data: %w", err)
+			}
+
+			if n > 0 {
+				crc32Results = append(crc32Results, hash.Sum32())
+			}
+		}
+	}
+}
+
+func buildCrc32Str(crc32Results []uint32) string {
+	hexCrcs := make([]string, len(crc32Results))
+	for i, crc := range crc32Results {
+		hexCrcs[i] = fmt.Sprintf("0x%x", crc)
+	}
+	return strings.Join(hexCrcs, ",")
 }

pkg/backend/build/remote.go

@@ -69,6 +69,8 @@ func (ro *remoteOutput) OutputLayer(ctx context.Context, mediaType, relPath, dig
 	}
 
 	if exist {
+		// In case the reader is from PipeReader, we need to read the whole reader to avoid the pipe being blocked.
+		io.Copy(io.Discard, reader)
 		hooks.OnComplete(relPath, desc)
 		return desc, nil
 	}

pkg/backend/nydusify.go

@@ -18,6 +18,7 @@ package backend
 
 import (
 	"context"
+	"os"
 	"os/exec"
 )
 
@@ -44,7 +45,8 @@ func (b *backend) Nydusify(ctx context.Context, source string) (string, error) {
 		"--target",
 		target,
 	)
-
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
 	if err := cmd.Run(); err != nil {
 		return "", err
 	}