Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions cmd/attach.go
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
2 changes: 1 addition & 1 deletion internal/pb/pb.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
31 changes: 24 additions & 7 deletions pkg/backend/attach.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import (
"fmt"
"io"
"reflect"
"slices"
"sort"

modelspec "github.com/CloudNativeAI/model-spec/specs-go/v1"
Expand Down Expand Up @@ -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)
}
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion pkg/backend/attach_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
2 changes: 2 additions & 0 deletions pkg/config/attach.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ type Attach struct {
Insecure bool
Nydusify bool
Force bool
Raw bool
}

func NewAttach() *Attach {
Expand All @@ -37,6 +38,7 @@ func NewAttach() *Attach {
Insecure: false,
Nydusify: false,
Force: false,
Raw: false,
}
}

Expand Down