feat: support export to output path for pull command

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

Author: chlins

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.OutputPath, "output", "", "specify the output path for exporting 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.OutputPath != "" {
+		opts = append(opts, backend.WithOutputPath(pullConfig.OutputPath))
+	}
+
 	if err := b.Pull(ctx, target, opts...); err != nil {
 		return err
 	}

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/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/options.go

@@ -19,8 +19,9 @@ package backend
 type Option func(*Options)
 
 type Options struct {
-	plainHTTP bool
-	proxy     string
+	plainHTTP  bool
+	proxy      string
+	outputPath string
 }
 
 // WithPlainHTTP sets the plain HTTP option.
@@ -36,3 +37,10 @@ func WithProxy(proxy string) Option {
 		opts.proxy = proxy
 	}
 }
+
+// WithOutputPath sets the output path option.
+func WithOutputPath(outputPath string) Option {
+	return func(opts *Options) {
+		opts.outputPath = outputPath
+	}
+}

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())),

pkg/backend/pull.go

@@ -23,9 +23,12 @@ import (
 	"io"
 	"net/http"
 	"net/url"
+	"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"
 	"oras.land/oras-go/v2/registry/remote"
 	"oras.land/oras-go/v2/registry/remote/auth"
@@ -124,6 +127,13 @@ func (b *backend) Pull(ctx context.Context, target string, opts ...Option) error
 		return fmt.Errorf("failed to pull manifest to local: %w", err)
 	}
 
+	// export the target model artifact to the output directory if needed.
+	if options.outputPath != "" {
+		if err := exportModelArtifact(ctx, dst, manifest, repo, options.outputPath); err != nil {
+			return fmt.Errorf("failed to export the artifact to the output directory: %w", err)
+		}
+	}
+
 	return nil
 }
 
@@ -162,3 +172,30 @@ func pullIfNotExist(ctx context.Context, pb *ProgressBar, prompt string, src *re
 
 	return nil
 }
+
+// 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, outputPath 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()
+
+		targetPath := outputPath
+		// get the original filepath in order to restore the original repo structure.
+		originalFilePath := layer.Annotations[modelspec.AnnotationFilepath]
+		if dir := filepath.Dir(originalFilePath); dir != "" {
+			targetPath = filepath.Join(targetPath, dir)
+		}
+
+		// untar the blob to the output directory.
+		if err := archiver.Untar(reader, targetPath); err != nil {
+			return fmt.Errorf("failed to untar the blob to output directory: %w", err)
+		}
+	}
+
+	return nil
+}

pkg/backend/push.go

@@ -17,6 +17,7 @@
 package backend
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -127,17 +128,16 @@ func pushIfNotExist(ctx context.Context, pb *ProgressBar, prompt string, src sto
 		return nil
 	}
 
-	// fetch the content from the source storage.
-	content, err := src.PullBlob(ctx, repo, desc.Digest.String())
-	if err != nil {
-		return fmt.Errorf("failed to fetch the content from source: %w", err)
-	}
-
-	defer content.Close()
 	// push the content to the destination, and wrap the content reader for progress bar,
 	// manifest should use dst.Manifests().Push, others should use dst.Blobs().Push.
 	if desc.MediaType == ocispec.MediaTypeImageManifest {
-		if err := dst.Manifests().Push(ctx, desc, pb.Add(prompt, desc, content)); err != nil {
+		// fetch the manifest from the source storage.
+		manifestRaw, _, err := src.PullManifest(ctx, repo, tag)
+		if err != nil {
+			return fmt.Errorf("failed to fetch the manifest from source: %w", err)
+		}
+
+		if err := dst.Manifests().Push(ctx, desc, pb.Add(prompt, desc, bytes.NewReader(manifestRaw))); err != nil {
 			return err
 		}
 
@@ -146,6 +146,14 @@ func pushIfNotExist(ctx context.Context, pb *ProgressBar, prompt string, src sto
 			return err
 		}
 	} else {
+		// fetch the content from the source storage.
+		content, err := src.PullBlob(ctx, repo, desc.Digest.String())
+		if err != nil {
+			return fmt.Errorf("failed to fetch the content from source: %w", err)
+		}
+
+		defer content.Close()
+
 		if err := dst.Blobs().Push(ctx, desc, pb.Add(prompt, desc, content)); err != nil {
 			return err
 		}

pkg/config/pull.go

@@ -17,13 +17,15 @@
 package config
 
 type Pull struct {
-	PlainHTTP bool
-	Proxy     string
+	PlainHTTP  bool
+	Proxy      string
+	OutputPath string
 }
 
 func NewPull() *Pull {
 	return &Pull{
-		PlainHTTP: false,
-		Proxy:     "",
+		PlainHTTP:  false,
+		Proxy:      "",
+		OutputPath: "",
 	}
 }