feat: support raw format for attach command (#209)

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

Author: chlins

cmd/attach.go

@@ -59,6 +59,7 @@ func init() {
 	flags.BoolVarP(&attachConfig.Force, "force", "f", false, "turning on this flag will force the attach, which will overwrite the layer if it already exists with same filepath")
 	flags.BoolVar(&attachConfig.Nydusify, "nydusify", false, "[EXPERIMENTAL] nydusify the model artifact")
 	flags.MarkHidden("nydusify")
+	flags.BoolVar(&attachConfig.Raw, "raw", false, "turning on this flag will attach model artifact layer in raw format")
 
 	if err := viper.BindPFlags(flags); err != nil {
 		panic(fmt.Errorf("bind cache list flags to viper: %w", err))

internal/pb/pb.go

@@ -62,7 +62,7 @@ func NewProgressBar(writers ...io.Writer) *ProgressBar {
 		mpbv8.PopCompletedMode(),
 		mpbv8.WithAutoRefresh(),
 		mpbv8.WithWidth(60),
-		mpbv8.WithRefreshRate(200 * time.Millisecond),
+		mpbv8.WithRefreshRate(300 * time.Millisecond),
 	}
 
 	// If no writer specified, use stdout.

pkg/backend/attach.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"io"
 	"reflect"
+	"slices"
 	"sort"
 
 	modelspec "github.com/CloudNativeAI/model-spec/specs-go/v1"
@@ -87,13 +88,13 @@ func (b *backend) Attach(ctx context.Context, filepath string, cfg *config.Attac
 		// Remove the found layer from the layers slice as we need to replace it with the new layer.
 		for i, layer := range layers {
 			if layer.Digest == foundLayer.Digest && layer.MediaType == foundLayer.MediaType {
-				layers = append(layers[:i], layers[i+1:]...)
+				layers = slices.Delete(layers, i, i+1)
 				break
 			}
 		}
 	}
 
-	proc := b.getProcessor(filepath)
+	proc := b.getProcessor(filepath, cfg)
 	if proc == nil {
 		return fmt.Errorf("failed to get processor for file %s", filepath)
 	}
@@ -260,21 +261,37 @@ func (b *backend) getModelConfig(ctx context.Context, reference string, desc oci
 	return &model, nil
 }
 
-func (b *backend) getProcessor(filepath string) processor.Processor {
+func (b *backend) getProcessor(filepath string, cfg *config.Attach) processor.Processor {
 	if modelfile.IsFileType(filepath, modelfile.ConfigFilePatterns) {
-		return processor.NewModelConfigProcessor(b.store, modelspec.MediaTypeModelWeightConfig, []string{filepath})
+		mediaType := modelspec.MediaTypeModelWeightConfig
+		if cfg.Raw {
+			mediaType = modelspec.MediaTypeModelWeightConfigRaw
+		}
+		return processor.NewModelConfigProcessor(b.store, mediaType, []string{filepath})
 	}
 
 	if modelfile.IsFileType(filepath, modelfile.ModelFilePatterns) {
-		return processor.NewModelProcessor(b.store, modelspec.MediaTypeModelWeight, []string{filepath})
+		mediaType := modelspec.MediaTypeModelWeight
+		if cfg.Raw {
+			mediaType = modelspec.MediaTypeModelWeightRaw
+		}
+		return processor.NewModelProcessor(b.store, mediaType, []string{filepath})
 	}
 
 	if modelfile.IsFileType(filepath, modelfile.CodeFilePatterns) {
-		return processor.NewCodeProcessor(b.store, modelspec.MediaTypeModelCode, []string{filepath})
+		mediaType := modelspec.MediaTypeModelCode
+		if cfg.Raw {
+			mediaType = modelspec.MediaTypeModelCodeRaw
+		}
+		return processor.NewCodeProcessor(b.store, mediaType, []string{filepath})
 	}
 
 	if modelfile.IsFileType(filepath, modelfile.DocFilePatterns) {
-		return processor.NewDocProcessor(b.store, modelspec.MediaTypeModelDoc, []string{filepath})
+		mediaType := modelspec.MediaTypeModelDoc
+		if cfg.Raw {
+			mediaType = modelspec.MediaTypeModelDocRaw
+		}
+		return processor.NewDocProcessor(b.store, mediaType, []string{filepath})
 	}
 
 	return nil

pkg/backend/attach_test.go

@@ -73,7 +73,7 @@ func TestGetProcessor(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.filepath, func(t *testing.T) {
-			proc := b.getProcessor(tt.filepath)
+			proc := b.getProcessor(tt.filepath, &config.Attach{})
 			if tt.wantType == "" {
 				assert.Nil(t, proc)
 			} else {

pkg/config/attach.go

@@ -26,6 +26,7 @@ type Attach struct {
 	Insecure     bool
 	Nydusify     bool
 	Force        bool
+	Raw          bool
 }
 
 func NewAttach() *Attach {
@@ -37,6 +38,7 @@ func NewAttach() *Attach {
 		Insecure:     false,
 		Nydusify:     false,
 		Force:        false,
+		Raw:          false,
 	}
 }