feat: introduce fetch command to get partial files (#137)

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

Author: chlins

.golangci.yml

@@ -1,8 +1,6 @@
 run:
-  deadline: 3m
+  timeout: 3m
   modules-download-mode: readonly
-  skip-dirs:
-    - test/mocks
 
 linters-settings:
   gocyclo:
@@ -11,13 +9,16 @@ linters-settings:
     sections:
       - standard
       - default
+      - prefix(github.com/CloudNativeAI/modctl)
 
 issues:
   new: true
   exclude-rules:
     - linters:
         - staticcheck
       text: "SA1019:"
+  exclude-dirs:
+    - test/mocks
 
 linters:
   disable-all: true
@@ -34,6 +35,7 @@ linters:
     - errcheck
 
 output:
-  format: colored-line-number
+  formats:
+    - format: colored-line-number
   print-issued-lines: true
   print-linter-name: true

cmd/fetch.go

@@ -0,0 +1,81 @@
+/*
+ *     Copyright 2025 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 fetchConfig = config.NewFetch()
+
+// fetchCmd represents the modctl command for fetch.
+var fetchCmd = &cobra.Command{
+	Use:                "fetch [flags] <target>",
+	Short:              "A command line tool for modctl fetch, please note that this command is designed for remote model fetching only.",
+	Args:               cobra.ExactArgs(1),
+	DisableAutoGenTag:  true,
+	SilenceUsage:       true,
+	FParseErrWhitelist: cobra.FParseErrWhitelist{UnknownFlags: true},
+	RunE: func(cmd *cobra.Command, args []string) error {
+		if err := fetchConfig.Validate(); err != nil {
+			return err
+		}
+
+		return runFetch(context.Background(), args[0])
+	},
+}
+
+// init initializes fetch command.
+func init() {
+	flags := fetchCmd.Flags()
+	flags.IntVar(&fetchConfig.Concurrency, "concurrency", fetchConfig.Concurrency, "specify the number of concurrent fetch operations")
+	flags.BoolVar(&fetchConfig.PlainHTTP, "plain-http", false, "use plain HTTP instead of HTTPS")
+	flags.BoolVar(&fetchConfig.Insecure, "insecure", false, "use insecure connection for the fetch operation and skip TLS verification")
+	flags.StringVar(&fetchConfig.Proxy, "proxy", "", "use proxy for the fetch operation")
+	flags.StringVar(&fetchConfig.Output, "output", "", "specify the directory for fetching the model artifact")
+	flags.StringSliceVar(&fetchConfig.Patterns, "patterns", []string{}, "specify the patterns for fetching the model artifact")
+
+	if err := viper.BindPFlags(flags); err != nil {
+		panic(fmt.Errorf("bind cache pull flags to viper: %w", err))
+	}
+}
+
+// runFetch runs the fetch modctl.
+func runFetch(ctx context.Context, target string) error {
+	b, err := backend.New(rootConfig.StoargeDir)
+	if err != nil {
+		return err
+	}
+
+	if target == "" {
+		return fmt.Errorf("target is required")
+	}
+
+	if err := b.Fetch(ctx, target, fetchConfig); err != nil {
+		return err
+	}
+
+	fmt.Printf("Successfully fetched model artifact: %s\n", target)
+	return nil
+}

cmd/root.go

@@ -85,5 +85,6 @@ func init() {
 	rootCmd.AddCommand(inspectCmd)
 	rootCmd.AddCommand(extractCmd)
 	rootCmd.AddCommand(tagCmd)
+	rootCmd.AddCommand(fetchCmd)
 	rootCmd.AddCommand(modelfile.RootCmd)
 }

pkg/backend/backend.go

@@ -37,6 +37,9 @@ type Backend interface {
 	// Pull pulls an artifact from a registry.
 	Pull(ctx context.Context, target string, cfg *config.Pull) error
 
+	// Fetch fetches partial files to the output.
+	Fetch(ctx context.Context, target string, cfg *config.Fetch) error
+
 	// Push pushes the image to the registry.
 	Push(ctx context.Context, target string, cfg *config.Push) error
 

pkg/backend/fetch.go

@@ -0,0 +1,135 @@
+/*
+ *     Copyright 2025 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"
+	"crypto/tls"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/url"
+	"path/filepath"
+
+	modelspec "github.com/CloudNativeAI/model-spec/specs-go/v1"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+	"golang.org/x/sync/errgroup"
+	"oras.land/oras-go/v2/registry/remote"
+	"oras.land/oras-go/v2/registry/remote/auth"
+	"oras.land/oras-go/v2/registry/remote/credentials"
+	"oras.land/oras-go/v2/registry/remote/retry"
+
+	internalpb "github.com/CloudNativeAI/modctl/internal/pb"
+	"github.com/CloudNativeAI/modctl/pkg/config"
+)
+
+// Fetch fetches partial files to the output.
+func (b *backend) Fetch(ctx context.Context, target string, cfg *config.Fetch) 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)
+	}
+
+	_, tag := ref.Repository(), ref.Tag()
+
+	// create the src storage from the remote repository.
+	src, err := remote.NewRepository(target)
+	if err != nil {
+		return fmt.Errorf("failed to create remote repository: %w", err)
+	}
+
+	// gets the credentials store.
+	credStore, err := credentials.NewStoreFromDocker(credentials.StoreOptions{AllowPlaintextPut: true})
+	if err != nil {
+		return fmt.Errorf("failed to create credential store: %w", err)
+	}
+
+	// create the http client.
+	httpClient := &http.Client{}
+	if cfg.Proxy != "" {
+		proxyURL, err := url.Parse(cfg.Proxy)
+		if err != nil {
+			return fmt.Errorf("failed to parse the proxy URL: %w", err)
+		}
+
+		httpClient.Transport = retry.NewTransport(&http.Transport{
+			Proxy: http.ProxyURL(proxyURL),
+			TLSClientConfig: &tls.Config{
+				InsecureSkipVerify: cfg.Insecure,
+			},
+		})
+	}
+
+	src.Client = &auth.Client{
+		Cache:      auth.NewCache(),
+		Credential: credentials.Credential(credStore),
+		Client:     httpClient,
+	}
+
+	if cfg.PlainHTTP {
+		src.PlainHTTP = true
+	}
+
+	_, manifestReader, err := src.Manifests().FetchReference(ctx, tag)
+	if err != nil {
+		return fmt.Errorf("failed to fetch the manifest: %w", err)
+	}
+
+	defer manifestReader.Close()
+
+	var manifest ocispec.Manifest
+	if err := json.NewDecoder(manifestReader).Decode(&manifest); err != nil {
+		return fmt.Errorf("failed to decode the manifest: %w", err)
+	}
+
+	layers := []ocispec.Descriptor{}
+	// filter the layers by patterns.
+	for _, layer := range manifest.Layers {
+		for _, pattern := range cfg.Patterns {
+			if anno := layer.Annotations; anno != nil {
+				matched, err := filepath.Match(pattern, anno[modelspec.AnnotationFilepath])
+				if err != nil {
+					return fmt.Errorf("failed to match pattern: %w", err)
+				}
+
+				if matched {
+					layers = append(layers, layer)
+				}
+			}
+		}
+	}
+
+	if len(layers) == 0 {
+		return fmt.Errorf("no layers matched the patterns")
+	}
+
+	pb := internalpb.NewProgressBar()
+	pb.Start()
+	defer pb.Stop()
+
+	g := &errgroup.Group{}
+	g.SetLimit(cfg.Concurrency)
+
+	for _, layer := range layers {
+		g.Go(func() error {
+			return pullAndExtractFromRemote(ctx, pb, internalpb.NormalizePrompt("Fetching blob"), src, cfg.Output, layer)
+		})
+	}
+
+	return g.Wait()
+}