feat: support export to output path for pull command (#36)

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

Author: chlins

cmd/extract.go

@@ -0,0 +1,76 @@
+/*
+ *     Copyright 2024 The CNAI Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package cmd
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/CloudNativeAI/modctl/pkg/backend"
+	"github.com/CloudNativeAI/modctl/pkg/config"
+
+	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+)
+
+var extractConfig = config.NewExtract()
+
+// extractCmd represents the modctl command for extract.
+var extractCmd = &cobra.Command{
+	Use:                "extract <target> --output <output>",
+	Short:              "A command line tool for modctl extract",
+	Args:               cobra.ExactArgs(1),
+	DisableAutoGenTag:  true,
+	SilenceUsage:       true,
+	FParseErrWhitelist: cobra.FParseErrWhitelist{UnknownFlags: true},
+	RunE: func(cmd *cobra.Command, args []string) error {
+		return runExtract(context.Background(), args[0])
+	},
+}
+
+// init initializes extract command.
+func init() {
+	flags := extractCmd.Flags()
+	flags.StringVar(&extractConfig.Output, "output", "", "specify the output for extracting the model artifact")
+
+	if err := viper.BindPFlags(flags); err != nil {
+		panic(fmt.Errorf("bind cache extract flags to viper: %w", err))
+	}
+}
+
+// runExtract runs the extract modctl.
+func runExtract(ctx context.Context, target string) error {
+	b, err := backend.New()
+	if err != nil {
+		return err
+	}
+
+	if target == "" {
+		return fmt.Errorf("target is required")
+	}
+
+	if extractConfig.Output == "" {
+		return fmt.Errorf("output is required")
+	}
+
+	if err := b.Extract(ctx, target, extractConfig.Output); err != nil {
+		return err
+	}
+
+	fmt.Printf("Successfully extracted model artifact %s to %s\n", target, extractConfig.Output)
+	return nil
+}

cmd/pull.go

@@ -47,6 +47,7 @@ func init() {
 	flags := pullCmd.Flags()
 	flags.BoolVar(&pullConfig.PlainHTTP, "plain-http", false, "use plain HTTP instead of HTTPS")
 	flags.StringVar(&pullConfig.Proxy, "proxy", "", "use proxy for the pull operation")
+	flags.StringVar(&pullConfig.ExtractDir, "extract-dir", "", "specify the extract dir for extracting the model artifact")
 
 	if err := viper.BindPFlags(flags); err != nil {
 		panic(fmt.Errorf("bind cache pull flags to viper: %w", err))
@@ -73,6 +74,10 @@ func runPull(ctx context.Context, target string) error {
 		opts = append(opts, backend.WithProxy(pullConfig.Proxy))
 	}
 
+	if pullConfig.ExtractDir != "" {
+		opts = append(opts, backend.WithOutput(pullConfig.ExtractDir))
+	}
+
 	if err := b.Pull(ctx, target, opts...); err != nil {
 		return err
 	}

cmd/root.go

@@ -66,4 +66,5 @@ func init() {
 	rootCmd.AddCommand(rmCmd)
 	rootCmd.AddCommand(pruneCmd)
 	rootCmd.AddCommand(inspectCmd)
+	rootCmd.AddCommand(extractCmd)
 }

docs/getting-started.md

@@ -82,6 +82,14 @@ Push the model artifact to the registry:
 $ modctl push registry.com/models/llama3:v1.0.0
 ```
 
+### Extract
+
+Extract the model artifact to the specified directory:
+
+```shell
+$ modctl extract registry.com/models/llama3:v1.0.0 --output /path/to/extract
+```
+
 ### List
 
 List the model artifacts in the local storage:

pkg/backend/build/tar.go pkg/archiver/archiver.gopkg/backend/build/tar.go renamed to pkg/archiver/archiver.go

@@ -14,17 +14,19 @@
  * limitations under the License.
  */
 
-package build
+package archiver
 
 import (
 	"archive/tar"
 	"fmt"
 	"io"
 	"os"
+	"path/filepath"
+	"strings"
 )
 
-// TarFileToStream tars the target file and return the content by stream.
-func TarFileToStream(path string) (io.Reader, error) {
+// Tar tars the target file and return the content by stream.
+func Tar(path string) (io.Reader, error) {
 	pr, pw := io.Pipe()
 	go func() {
 		defer pw.Close()
@@ -65,3 +67,62 @@ func TarFileToStream(path string) (io.Reader, error) {
 
 	return pr, nil
 }
+
+// Untar untars the target stream to the destination path.
+func Untar(reader io.Reader, destPath string) error {
+	// uncompress gzip if it is a .tar.gz file
+	// gzipReader, err := gzip.NewReader(reader)
+	// if err != nil {
+	//     return err
+	// }
+	// defer gzipReader.Close()
+	// tarReader := tar.NewReader(gzipReader)
+
+	tarReader := tar.NewReader(reader)
+
+	if err := os.MkdirAll(destPath, 0755); err != nil {
+		return err
+	}
+
+	for {
+		header, err := tarReader.Next()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return err
+		}
+
+		// sanitize filepaths to prevent directory traversal.
+		cleanPath := filepath.Clean(header.Name)
+		if strings.Contains(cleanPath, "..") {
+			return fmt.Errorf("tar file contains invalid path: %s", cleanPath)
+		}
+
+		path := filepath.Join(destPath, cleanPath)
+		// check the file type.
+		switch header.Typeflag {
+		case tar.TypeDir:
+			if err := os.MkdirAll(path, 0755); err != nil {
+				return err
+			}
+		case tar.TypeReg:
+			file, err := os.Create(path)
+			if err != nil {
+				return err
+			}
+
+			if _, err := io.Copy(file, tarReader); err != nil {
+				file.Close()
+				return err
+			}
+			file.Close()
+
+			if err := os.Chmod(path, os.FileMode(header.Mode)); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}

pkg/backend/backend.go

@@ -50,6 +50,9 @@ type Backend interface {
 
 	// Inspect inspects the model artifact.
 	Inspect(ctx context.Context, target string) (*InspectedModelArtifact, error)
+
+	// Extract extracts the model artifact.
+	Extract(ctx context.Context, target string, output string) error
 }
 
 // backend is the implementation of Backend.

pkg/backend/build/build.go

@@ -24,8 +24,9 @@ import (
 	"path/filepath"
 	"time"
 
-	modelspec "github.com/CloudNativeAI/model-spec/specs-go/v1"
+	"github.com/CloudNativeAI/modctl/pkg/archiver"
 	"github.com/CloudNativeAI/modctl/pkg/storage"
+	modelspec "github.com/CloudNativeAI/model-spec/specs-go/v1"
 
 	godigest "github.com/opencontainers/go-digest"
 	spec "github.com/opencontainers/image-spec/specs-go"
@@ -40,7 +41,7 @@ type ModelConfig struct {
 
 // BuildLayer converts the file to the image blob and push it to the storage.
 func BuildLayer(ctx context.Context, store storage.Storage, repo, path, workDir string) (ocispec.Descriptor, error) {
-	reader, err := TarFileToStream(path)
+	reader, err := archiver.Tar(path)
 	if err != nil {
 		return ocispec.Descriptor{}, fmt.Errorf("failed to tar file: %w", err)
 	}

pkg/backend/extract.go

@@ -0,0 +1,80 @@
+/*
+ *     Copyright 2024 The CNAI Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package backend
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"path/filepath"
+
+	"github.com/CloudNativeAI/modctl/pkg/archiver"
+	"github.com/CloudNativeAI/modctl/pkg/storage"
+	modelspec "github.com/CloudNativeAI/model-spec/specs-go/v1"
+
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+)
+
+// Extract extracts the model artifact.
+func (b *backend) Extract(ctx context.Context, target string, output string) error {
+	// parse the repository and tag from the target.
+	ref, err := ParseReference(target)
+	if err != nil {
+		return fmt.Errorf("failed to parse the target: %w", err)
+	}
+
+	repo, tag := ref.Repository(), ref.Tag()
+	// pull the manifest from the storage.
+	manifestRaw, _, err := b.store.PullManifest(ctx, repo, tag)
+	if err != nil {
+		return fmt.Errorf("failed to pull the manifest from storage: %w", err)
+	}
+	// unmarshal the manifest.
+	var manifest ocispec.Manifest
+	if err := json.Unmarshal(manifestRaw, &manifest); err != nil {
+		return fmt.Errorf("failed to unmarshal the manifest: %w", err)
+	}
+
+	return exportModelArtifact(ctx, b.store, manifest, repo, output)
+}
+
+// exportModelArtifact exports the target model artifact to the output directory, which will open the artifact and extract to restore the orginal repo structure.
+func exportModelArtifact(ctx context.Context, store storage.Storage, manifest ocispec.Manifest, repo, output string) error {
+	for _, layer := range manifest.Layers {
+		// pull the blob from the storage.
+		reader, err := store.PullBlob(ctx, repo, layer.Digest.String())
+		if err != nil {
+			return fmt.Errorf("failed to pull the blob from storage: %w", err)
+		}
+
+		defer reader.Close()
+
+		targetDir := output
+		// get the original filepath in order to restore the original repo structure.
+		originalFilePath := layer.Annotations[modelspec.AnnotationFilepath]
+		if dir := filepath.Dir(originalFilePath); dir != "" {
+			targetDir = filepath.Join(targetDir, dir)
+		}
+
+		// untar the blob to the target directory.
+		if err := archiver.Untar(reader, targetDir); err != nil {
+			return fmt.Errorf("failed to untar the blob to output directory: %w", err)
+		}
+	}
+
+	return nil
+}

pkg/backend/options.go

@@ -21,6 +21,7 @@ type Option func(*Options)
 type Options struct {
 	plainHTTP bool
 	proxy     string
+	output    string
 }
 
 // WithPlainHTTP sets the plain HTTP option.
@@ -36,3 +37,10 @@ func WithProxy(proxy string) Option {
 		opts.proxy = proxy
 	}
 }
+
+// WithOutput sets the output option.
+func WithOutput(output string) Option {
+	return func(opts *Options) {
+		opts.output = output
+	}
+}

pkg/backend/progress.go

@@ -59,7 +59,7 @@ func (p *ProgressBar) Add(prompt string, desc ocispec.Descriptor, reader io.Read
 
 	// create a new bar if it does not exist.
 	bar := p.mpb.New(desc.Size,
-		mpbv8.BarStyle().Rbound("|"),
+		mpbv8.BarStyle(),
 		mpbv8.BarFillerOnComplete("|"),
 		mpbv8.PrependDecorators(
 			decor.Name(fmt.Sprintf("%s%s", prompt, desc.Digest.String())),