Skip to content

Commit 776d032

Browse files
chlinschlins
authored
feat: support pull the blob from dragonfly unix socket (#193)
Signed-off-by: chlins <[email protected]>
1 parent 7bb9404 commit 776d032

File tree

10 files changed

+403
-48
lines changed

10 files changed

+403
-48
lines changed

‎cmd/pull.go‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,11 @@ import (
2020
"context"
2121
"fmt"
2222

23-
"github.com/CloudNativeAI/modctl/pkg/backend"
24-
"github.com/CloudNativeAI/modctl/pkg/config"
25-
2623
"github.com/spf13/cobra"
2724
"github.com/spf13/viper"
25+
26+
"github.com/CloudNativeAI/modctl/pkg/backend"
27+
"github.com/CloudNativeAI/modctl/pkg/config"
2828
)
2929

3030
var pullConfig = config.NewPull()
@@ -55,6 +55,7 @@ func init() {
5555
flags.StringVar(&pullConfig.Proxy, "proxy", "", "use proxy for the pull operation")
5656
flags.StringVar(&pullConfig.ExtractDir, "extract-dir", "", "specify the extract dir for extracting the model artifact")
5757
flags.BoolVar(&pullConfig.ExtractFromRemote, "extract-from-remote", false, "turning on this flag will pull and extract the data from remote registry and no longer store model artifact locally, so user must specify extract-dir as the output directory")
58+
flags.StringVar(&pullConfig.DragonflyEndpoint, "dragonfly-endpoint", "", "specify the dragonfly endpoint for the pull operation, which will download and hardlink the blob by dragonfly GRPC service, this mode requires extract-from-remote must be true")
5859

5960
if err := viper.BindPFlags(flags); err != nil {
6061
panic(fmt.Errorf("bind cache pull flags to viper: %w", err))
@@ -72,8 +73,14 @@ func runPull(ctx context.Context, target string) error {
7273
return fmt.Errorf("target is required")
7374
}
7475

75-
if err := b.Pull(ctx, target, pullConfig); err != nil {
76-
return err
76+
if pullConfig.DragonflyEndpoint != "" {
77+
if err := b.PullByDragonfly(ctx, target, pullConfig); err != nil {
78+
return err
79+
}
80+
} else {
81+
if err := b.Pull(ctx, target, pullConfig); err != nil {
82+
return err
83+
}
7784
}
7885

7986
fmt.Printf("Successfully pulled model artifact: %s\n", target)

‎go.mod‎

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ module github.com/CloudNativeAI/modctl
33
go 1.24.1
44

55
require (
6+
d7y.io/api/v2 v2.1.39
67
github.com/CloudNativeAI/model-spec v0.0.6
78
github.com/antgroup/hugescm v0.18.0
89
github.com/avast/retry-go/v4 v4.6.1
@@ -23,6 +24,7 @@ require (
2324
github.com/vbauerster/mpb/v8 v8.10.2
2425
golang.org/x/crypto v0.39.0
2526
golang.org/x/sync v0.15.0
27+
google.golang.org/grpc v1.71.1
2628
oras.land/oras-go/v2 v2.6.0
2729
)
2830

@@ -42,6 +44,7 @@ require (
4244
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
4345
github.com/dgraph-io/ristretto/v2 v2.2.0 // indirect
4446
github.com/docker/go-metrics v0.0.1 // indirect
47+
github.com/envoyproxy/protoc-gen-validate v1.2.1 // indirect
4548
github.com/fatih/color v1.7.0 // indirect
4649
github.com/fsnotify/fsnotify v1.8.0 // indirect
4750
github.com/go-git/gcfg v1.5.1-0.20230307220236-3a3c6141e376 // indirect
@@ -84,9 +87,10 @@ require (
8487
github.com/subosito/gotenv v1.6.0 // indirect
8588
github.com/xanzy/ssh-agent v0.3.3 // indirect
8689
github.com/zeebo/blake3 v0.2.4 // indirect
90+
go.opentelemetry.io/auto/sdk v1.1.0 // indirect
8791
go.opentelemetry.io/contrib/bridges/prometheus v0.57.0 // indirect
8892
go.opentelemetry.io/contrib/exporters/autoexport v0.57.0 // indirect
89-
go.opentelemetry.io/otel v1.32.0 // indirect
93+
go.opentelemetry.io/otel v1.34.0 // indirect
9094
go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploggrpc v0.8.0 // indirect
9195
go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploghttp v0.8.0 // indirect
9296
go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetricgrpc v1.32.0 // indirect
@@ -99,21 +103,20 @@ require (
99103
go.opentelemetry.io/otel/exporters/stdout/stdoutmetric v1.32.0 // indirect
100104
go.opentelemetry.io/otel/exporters/stdout/stdouttrace v1.32.0 // indirect
101105
go.opentelemetry.io/otel/log v0.8.0 // indirect
102-
go.opentelemetry.io/otel/metric v1.32.0 // indirect
103-
go.opentelemetry.io/otel/sdk v1.32.0 // indirect
106+
go.opentelemetry.io/otel/metric v1.34.0 // indirect
107+
go.opentelemetry.io/otel/sdk v1.34.0 // indirect
104108
go.opentelemetry.io/otel/sdk/log v0.8.0 // indirect
105-
go.opentelemetry.io/otel/sdk/metric v1.32.0 // indirect
106-
go.opentelemetry.io/otel/trace v1.32.0 // indirect
109+
go.opentelemetry.io/otel/sdk/metric v1.34.0 // indirect
110+
go.opentelemetry.io/otel/trace v1.34.0 // indirect
107111
go.opentelemetry.io/proto/otlp v1.3.1 // indirect
108112
go.uber.org/multierr v1.11.0 // indirect
109113
golang.org/x/net v0.40.0 // indirect
110114
golang.org/x/sys v0.33.0 // indirect
111115
golang.org/x/term v0.32.0 // indirect
112116
golang.org/x/text v0.26.0 // indirect
113-
google.golang.org/genproto/googleapis/api v0.0.0-20241209162323-e6fa225c2576 // indirect
114-
google.golang.org/genproto/googleapis/rpc v0.0.0-20241223144023-3abc09e42ca8 // indirect
115-
google.golang.org/grpc v1.68.0 // indirect
116-
google.golang.org/protobuf v1.36.1 // indirect
117+
google.golang.org/genproto/googleapis/api v0.0.0-20250218202821-56aae31c358a // indirect
118+
google.golang.org/genproto/googleapis/rpc v0.0.0-20250519155744-55703ea1f237 // indirect
119+
google.golang.org/protobuf v1.36.6 // indirect
117120
gopkg.in/warnings.v0 v0.1.2 // indirect
118121
gopkg.in/yaml.v3 v3.0.1 // indirect
119122
)

‎go.sum‎

Lines changed: 24 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
d7y.io/api/v2 v2.1.39 h1:BxbccbEUooXuCr8aqRr5hNQONFdL7DOWo5pz2CnBg/0=
2+
d7y.io/api/v2 v2.1.39/go.mod h1:5DPjID9MNDgw0mOqf49KF6Mho4x2v+C3uKmWWQYCWUs=
13
dario.cat/mergo v1.0.0 h1:AGCNq9Evsj31mOgNPcLyXc+4PNABt905YmuqPYYpBWk=
24
dario.cat/mergo v1.0.0/go.mod h1:uNxQE+84aUszobStD9th8a29P2fMDhsBdgRYvZOxGmk=
35
github.com/BurntSushi/toml v1.5.0 h1:W5quZX/G/csjUnuI8SUYlsHs9M38FC7znL0lIO+DvMg=
@@ -64,6 +66,8 @@ github.com/elazarl/goproxy v1.7.2 h1:Y2o6urb7Eule09PjlhQRGNsqRfPmYI3KKQLFpCAV3+o
6466
github.com/elazarl/goproxy v1.7.2/go.mod h1:82vkLNir0ALaW14Rc399OTTjyNREgmdL2cVoIbS6XaE=
6567
github.com/emirpasic/gods v1.18.1 h1:FXtiHYKDGKCW2KzwZKx0iC0PQmdlorYgdFG9jPXJ1Bc=
6668
github.com/emirpasic/gods v1.18.1/go.mod h1:8tpGGwCnJ5H4r6BWwaV6OrWmMoPhUl5jm/FMNAnJvWQ=
69+
github.com/envoyproxy/protoc-gen-validate v1.2.1 h1:DEo3O99U8j4hBFwbJfrz9VtgcDfUKS7KJ7spH3d86P8=
70+
github.com/envoyproxy/protoc-gen-validate v1.2.1/go.mod h1:d/C80l/jxXLdfEIhX1W2TmLfsJ31lvEjwamM4DxlWXU=
6771
github.com/fatih/color v1.7.0 h1:DkWD4oS2D8LGGgTQ6IvwJJXSL5Vp2ffcQg58nFV38Ys=
6872
github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4=
6973
github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
@@ -247,12 +251,14 @@ github.com/zeebo/blake3 v0.2.4 h1:KYQPkhpRtcqh0ssGYcKLG1JYvddkEA8QwCM/yBqhaZI=
247251
github.com/zeebo/blake3 v0.2.4/go.mod h1:7eeQ6d2iXWRGF6npfaxl2CU+xy2Fjo2gxeyZGCRUjcE=
248252
github.com/zeebo/pcg v1.0.1 h1:lyqfGeWiv4ahac6ttHs+I5hwtH/+1mrhlCtVNQM2kHo=
249253
github.com/zeebo/pcg v1.0.1/go.mod h1:09F0S9iiKrwn9rlI5yjLkmrug154/YRW6KnnXVDM/l4=
254+
go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA=
255+
go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A=
250256
go.opentelemetry.io/contrib/bridges/prometheus v0.57.0 h1:UW0+QyeyBVhn+COBec3nGhfnFe5lwB0ic1JBVjzhk0w=
251257
go.opentelemetry.io/contrib/bridges/prometheus v0.57.0/go.mod h1:ppciCHRLsyCio54qbzQv0E4Jyth/fLWDTJYfvWpcSVk=
252258
go.opentelemetry.io/contrib/exporters/autoexport v0.57.0 h1:jmTVJ86dP60C01K3slFQa2NQ/Aoi7zA+wy7vMOKD9H4=
253259
go.opentelemetry.io/contrib/exporters/autoexport v0.57.0/go.mod h1:EJBheUMttD/lABFyLXhce47Wr6DPWYReCzaZiXadH7g=
254-
go.opentelemetry.io/otel v1.32.0 h1:WnBN+Xjcteh0zdk01SVqV55d/m62NJLJdIyb4y/WO5U=
255-
go.opentelemetry.io/otel v1.32.0/go.mod h1:00DCVSB0RQcnzlwyTfqtxSm+DRr9hpYrHjNGiBHVQIg=
260+
go.opentelemetry.io/otel v1.34.0 h1:zRLXxLCgL1WyKsPVrgbSdMN4c0FMkDAskSTQP+0hdUY=
261+
go.opentelemetry.io/otel v1.34.0/go.mod h1:OWFPOQ+h4G8xpyjgqo4SxJYdDQ/qmRH+wivy7zzx9oI=
256262
go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploggrpc v0.8.0 h1:WzNab7hOOLzdDF/EoWCt4glhrbMPVMOO5JYTmpz36Ls=
257263
go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploggrpc v0.8.0/go.mod h1:hKvJwTzJdp90Vh7p6q/9PAOd55dI6WA6sWj62a/JvSs=
258264
go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploghttp v0.8.0 h1:S+LdBGiQXtJdowoJoQPEtI52syEP/JYBUpjO49EQhV8=
@@ -277,16 +283,16 @@ go.opentelemetry.io/otel/exporters/stdout/stdouttrace v1.32.0 h1:cC2yDI3IQd0Udsu
277283
go.opentelemetry.io/otel/exporters/stdout/stdouttrace v1.32.0/go.mod h1:2PD5Ex6z8CFzDbTdOlwyNIUywRr1DN0ospafJM1wJ+s=
278284
go.opentelemetry.io/otel/log v0.8.0 h1:egZ8vV5atrUWUbnSsHn6vB8R21G2wrKqNiDt3iWertk=
279285
go.opentelemetry.io/otel/log v0.8.0/go.mod h1:M9qvDdUTRCopJcGRKg57+JSQ9LgLBrwwfC32epk5NX8=
280-
go.opentelemetry.io/otel/metric v1.32.0 h1:xV2umtmNcThh2/a/aCP+h64Xx5wsj8qqnkYZktzNa0M=
281-
go.opentelemetry.io/otel/metric v1.32.0/go.mod h1:jH7CIbbK6SH2V2wE16W05BHCtIDzauciCRLoc/SyMv8=
282-
go.opentelemetry.io/otel/sdk v1.32.0 h1:RNxepc9vK59A8XsgZQouW8ue8Gkb4jpWtJm9ge5lEG4=
283-
go.opentelemetry.io/otel/sdk v1.32.0/go.mod h1:LqgegDBjKMmb2GC6/PrTnteJG39I8/vJCAP9LlJXEjU=
286+
go.opentelemetry.io/otel/metric v1.34.0 h1:+eTR3U0MyfWjRDhmFMxe2SsW64QrZ84AOhvqS7Y+PoQ=
287+
go.opentelemetry.io/otel/metric v1.34.0/go.mod h1:CEDrp0fy2D0MvkXE+dPV7cMi8tWZwX3dmaIhwPOaqHE=
288+
go.opentelemetry.io/otel/sdk v1.34.0 h1:95zS4k/2GOy069d321O8jWgYsW3MzVV+KuSPKp7Wr1A=
289+
go.opentelemetry.io/otel/sdk v1.34.0/go.mod h1:0e/pNiaMAqaykJGKbi+tSjWfNNHMTxoC9qANsCzbyxU=
284290
go.opentelemetry.io/otel/sdk/log v0.8.0 h1:zg7GUYXqxk1jnGF/dTdLPrK06xJdrXgqgFLnI4Crxvs=
285291
go.opentelemetry.io/otel/sdk/log v0.8.0/go.mod h1:50iXr0UVwQrYS45KbruFrEt4LvAdCaWWgIrsN3ZQggo=
286-
go.opentelemetry.io/otel/sdk/metric v1.32.0 h1:rZvFnvmvawYb0alrYkjraqJq0Z4ZUJAiyYCU9snn1CU=
287-
go.opentelemetry.io/otel/sdk/metric v1.32.0/go.mod h1:PWeZlq0zt9YkYAp3gjKZ0eicRYvOh1Gd+X99x6GHpCQ=
288-
go.opentelemetry.io/otel/trace v1.32.0 h1:WIC9mYrXf8TmY/EXuULKc8hR17vE+Hjv2cssQDe03fM=
289-
go.opentelemetry.io/otel/trace v1.32.0/go.mod h1:+i4rkvCraA+tG6AzwloGaCtkx53Fa+L+V8e9a7YvhT8=
292+
go.opentelemetry.io/otel/sdk/metric v1.34.0 h1:5CeK9ujjbFVL5c1PhLuStg1wxA7vQv7ce1EK0Gyvahk=
293+
go.opentelemetry.io/otel/sdk/metric v1.34.0/go.mod h1:jQ/r8Ze28zRKoNRdkjCZxfs6YvBTG1+YIqyFVFYec5w=
294+
go.opentelemetry.io/otel/trace v1.34.0 h1:+ouXS2V8Rd4hp4580a8q23bg0azF2nI8cqLYnC8mh/k=
295+
go.opentelemetry.io/otel/trace v1.34.0/go.mod h1:Svm7lSjQD7kG7KJ/MUHPVXSDGz2OX4h0M2jHBhmSfRE=
290296
go.opentelemetry.io/proto/otlp v1.3.1 h1:TrMUixzpM0yuc/znrFTP9MMRh8trP93mkCiDVeXrui0=
291297
go.opentelemetry.io/proto/otlp v1.3.1/go.mod h1:0X1WI4de4ZsLrrJNLAQbFeLCm3T7yBkR0XqQ7niQU+8=
292298
go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
@@ -334,14 +340,14 @@ golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
334340
golang.org/x/text v0.26.0 h1:P42AVeLghgTYr4+xUnTRKDMqpar+PtX7KWuNQL21L8M=
335341
golang.org/x/text v0.26.0/go.mod h1:QK15LZJUUQVJxhz7wXgxSy/CJaTFjd0G+YLonydOVQA=
336342
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
337-
google.golang.org/genproto/googleapis/api v0.0.0-20241209162323-e6fa225c2576 h1:CkkIfIt50+lT6NHAVoRYEyAvQGFM7xEwXUUywFvEb3Q=
338-
google.golang.org/genproto/googleapis/api v0.0.0-20241209162323-e6fa225c2576/go.mod h1:1R3kvZ1dtP3+4p4d3G8uJ8rFk/fWlScl38vanWACI08=
339-
google.golang.org/genproto/googleapis/rpc v0.0.0-20241223144023-3abc09e42ca8 h1:TqExAhdPaB60Ux47Cn0oLV07rGnxZzIsaRhQaqS666A=
340-
google.golang.org/genproto/googleapis/rpc v0.0.0-20241223144023-3abc09e42ca8/go.mod h1:lcTa1sDdWEIHMWlITnIczmw5w60CF9ffkb8Z+DVmmjA=
341-
google.golang.org/grpc v1.68.0 h1:aHQeeJbo8zAkAa3pRzrVjZlbz6uSfeOXlJNQM0RAbz0=
342-
google.golang.org/grpc v1.68.0/go.mod h1:fmSPC5AsjSBCK54MyHRx48kpOti1/jRfOlwEWywNjWA=
343-
google.golang.org/protobuf v1.36.1 h1:yBPeRvTftaleIgM3PZ/WBIZ7XM/eEYAaEyCwvyjq/gk=
344-
google.golang.org/protobuf v1.36.1/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE=
343+
google.golang.org/genproto/googleapis/api v0.0.0-20250218202821-56aae31c358a h1:nwKuGPlUAt+aR+pcrkfFRrTU1BVrSmYyYMxYbUIVHr0=
344+
google.golang.org/genproto/googleapis/api v0.0.0-20250218202821-56aae31c358a/go.mod h1:3kWAYMk1I75K4vykHtKt2ycnOgpA6974V7bREqbsenU=
345+
google.golang.org/genproto/googleapis/rpc v0.0.0-20250519155744-55703ea1f237 h1:cJfm9zPbe1e873mHJzmQ1nwVEeRDU/T1wXDK2kUSU34=
346+
google.golang.org/genproto/googleapis/rpc v0.0.0-20250519155744-55703ea1f237/go.mod h1:qQ0YXyHHx3XkvlzUtpXDkS29lDSafHMZBAZDc03LQ3A=
347+
google.golang.org/grpc v1.71.1 h1:ffsFWr7ygTUscGPI0KKK6TLrGz0476KUvvsbqWK0rPI=
348+
google.golang.org/grpc v1.71.1/go.mod h1:H0GRtasmQOh9LkFoCPDu3ZrwUtD1YGE+b2vYBYd/8Ec=
349+
google.golang.org/protobuf v1.36.6 h1:z1NpPI8ku2WgiWnf+t9wTPsn6eP1L7ksHUlkfLvd9xY=
350+
google.golang.org/protobuf v1.36.6/go.mod h1:jduwjTPXsFjZGTmRluh+L6NjiWu7pchiJ2/5YcXBHnY=
345351
gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
346352
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
347353
gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

‎internal/pb/pb.go‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,9 +59,10 @@ type progressBar struct {
5959
// NewProgressBar creates a new progress bar.
6060
func NewProgressBar(writers ...io.Writer) *ProgressBar {
6161
opts := []mpbv8.ContainerOption{
62+
mpbv8.PopCompletedMode(),
6263
mpbv8.WithAutoRefresh(),
6364
mpbv8.WithWidth(60),
64-
mpbv8.WithRefreshRate(300 * time.Millisecond),
65+
mpbv8.WithRefreshRate(200 * time.Millisecond),
6566
}
6667

6768
// If no writer specified, use stdout.
@@ -118,7 +119,20 @@ func (p *ProgressBar) Add(prompt, name string, size int64, reader io.Reader) io.
118119
p.bars[name] = newBar
119120
p.mu.Unlock()
120121

121-
return newBar.ProxyReader(reader)
122+
if reader != nil {
123+
return newBar.ProxyReader(reader)
124+
}
125+
126+
return reader
127+
}
128+
129+
// Get returns the progress bar.
130+
func (p *ProgressBar) Get(name string) *progressBar {
131+
p.mu.RLock()
132+
bar := p.bars[name]
133+
p.mu.RUnlock()
134+
135+
return bar
122136
}
123137

124138
// Complete completes the progress bar.

‎pkg/backend/backend.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,9 @@ type Backend interface {
4040
// Pull pulls an artifact from a registry.
4141
Pull(ctx context.Context, target string, cfg *config.Pull) error
4242

43+
// PullByDragonfly pulls an artifact from a registry by Dragonfly.
44+
PullByDragonfly(ctx context.Context, target string, cfg *config.Pull) error
45+
4346
// Fetch fetches partial files to the output.
4447
Fetch(ctx context.Context, target string, cfg *config.Fetch) error
4548

0 commit comments

Comments
 (0)