Skip to content

Commit 848f430

Browse files
imeoerimeoer
committed
feat: support maximum disk usage limit
Signed-off-by: imeoer <yansong.ys@antgroup.com>
1 parent b588994 commit 848f430

File tree

10 files changed

+445
-53
lines changed

10 files changed

+445
-53
lines changed

‎cmd/model-csi-driver/main.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ func main() {
3838
},
3939
},
4040
Action: func(c *cli.Context) error {
41-
cfg, err := config.FromFile(c.String("config"))
41+
cfg, err := config.New(c.String("config"))
4242
if err != nil {
4343
return errors.Wrap(err, "load config")
4444
}

‎go.mod‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,11 @@ module github.com/modelpack/model-csi-driver
33
go 1.24.2
44

55
require (
6+
github.com/agiledragon/gomonkey/v2 v2.13.0
67
github.com/container-storage-interface/spec v1.2.0
78
github.com/containerd/containerd v1.7.27
89
github.com/dustin/go-humanize v1.0.1
10+
github.com/fsnotify/fsnotify v1.8.0
911
github.com/google/uuid v1.6.0
1012
github.com/labstack/echo/v4 v4.13.3
1113
github.com/moby/sys/mountinfo v0.7.2
@@ -27,6 +29,7 @@ require (
2729
go.opentelemetry.io/otel/trace v1.37.0
2830
golang.org/x/net v0.42.0
2931
golang.org/x/sync v0.16.0
32+
golang.org/x/sys v0.35.0
3033
google.golang.org/grpc v1.75.0
3134
gopkg.in/yaml.v2 v2.4.0
3235
k8s.io/api v0.28.4
@@ -128,7 +131,6 @@ require (
128131
go.opentelemetry.io/proto/otlp v1.6.0 // indirect
129132
golang.org/x/crypto v0.41.0 // indirect
130133
golang.org/x/oauth2 v0.30.0 // indirect
131-
golang.org/x/sys v0.35.0 // indirect
132134
golang.org/x/term v0.34.0 // indirect
133135
golang.org/x/text v0.28.0 // indirect
134136
golang.org/x/time v0.8.0 // indirect

‎go.sum‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@ github.com/VividCortex/ewma v1.2.0 h1:f58SaIzcDXrSy3kWaHNvuJgJ3Nmz59Zji6XoJR/q1o
1717
github.com/VividCortex/ewma v1.2.0/go.mod h1:nz4BbCtbLyFDeC9SUHbtcT5644juEuWfUAUnGx7j5l4=
1818
github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d h1:licZJFw2RwpHMqeKTCYkitsPqHNxTmd4SNR5r94FGM8=
1919
github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d/go.mod h1:asat636LX7Bqt5lYEZ27JNDcqxfjdBQuJ/MM4CN/Lzo=
20+
github.com/agiledragon/gomonkey/v2 v2.13.0 h1:B24Jg6wBI1iB8EFR1c+/aoTg7QN/Cum7YffG8KMIyYo=
21+
github.com/agiledragon/gomonkey/v2 v2.13.0/go.mod h1:ap1AmDzcVOAz1YpeJ3TCzIgstoaWLA6jbbgxfB4w2iY=
2022
github.com/akutz/memconn v0.1.0/go.mod h1:Jo8rI7m0NieZyLI5e2CDlRdRqRRB4S7Xp77ukDjH+Fw=
2123
github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
2224
github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
@@ -78,6 +80,8 @@ github.com/emirpasic/gods v1.18.1 h1:FXtiHYKDGKCW2KzwZKx0iC0PQmdlorYgdFG9jPXJ1Bc
7880
github.com/emirpasic/gods v1.18.1/go.mod h1:8tpGGwCnJ5H4r6BWwaV6OrWmMoPhUl5jm/FMNAnJvWQ=
7981
github.com/envoyproxy/protoc-gen-validate v1.2.1 h1:DEo3O99U8j4hBFwbJfrz9VtgcDfUKS7KJ7spH3d86P8=
8082
github.com/envoyproxy/protoc-gen-validate v1.2.1/go.mod h1:d/C80l/jxXLdfEIhX1W2TmLfsJ31lvEjwamM4DxlWXU=
83+
github.com/fsnotify/fsnotify v1.8.0 h1:dAwr6QBTBZIkG8roQaJjGof0pp0EeF+tNV7YBP3F/8M=
84+
github.com/fsnotify/fsnotify v1.8.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0=
8185
github.com/gliderlabs/ssh v0.3.8 h1:a4YXD1V7xMF9g5nTkdfnja3Sxy1PVDCj1Zg4Wb8vY6c=
8286
github.com/gliderlabs/ssh v0.3.8/go.mod h1:xYoytBv1sV0aL3CavoDuJIQNURXkkfPA/wxQ1pL1fAU=
8387
github.com/go-git/gcfg v1.5.1-0.20230307220236-3a3c6141e376 h1:+zs/tPmkDkHx3U66DAb0lQFJrpS6731Oaa12ikc+DiI=
@@ -134,6 +138,7 @@ github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 h1:El6M4kTTCOh6aBiKaU
134138
github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510/go.mod h1:pupxD2MaaD3pAXIBCelhxNneeOaAeabZDe5s4K6zSpQ=
135139
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
136140
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
141+
github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=
137142
github.com/gorilla/mux v1.8.2-0.20240619235004-db9d1d0073d2 h1:oZRjfKe/6Qh676XFYvylkCWd0gu8KVZeZYZwkNw6NAU=
138143
github.com/gorilla/mux v1.8.2-0.20240619235004-db9d1d0073d2/go.mod h1:AKf9I4AEqPTmMytcMc0KkNouC66V3BtZ4qD5fmWSiMQ=
139144
github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.3 h1:5ZPtiqj0JL5oKWmcsq4VMaAW5ukBEgSGXEN89zeH1Jo=
@@ -155,6 +160,7 @@ github.com/json-iterator/go v1.1.6/go.mod h1:+SdeFBvtyEkXs7REEP0seUULqWtbJapLOCV
155160
github.com/json-iterator/go v1.1.7/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4=
156161
github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
157162
github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
163+
github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU=
158164
github.com/julienschmidt/httprouter v1.2.0/go.mod h1:SYymIcj16QtmaHHD7aYtjjsJG7VTCxuUUipMqKk8s4w=
159165
github.com/kevinburke/ssh_config v1.2.0 h1:x584FjTGwHzMwvHx18PXxbBVzfnxogHaAReU4gf13a4=
160166
github.com/kevinburke/ssh_config v1.2.0/go.mod h1:CT57kijsi8u/K/BOFA39wgDQJ9CxiF4nAY/ojJ6r6mM=
@@ -263,6 +269,8 @@ github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ
263269
github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
264270
github.com/skeema/knownhosts v1.3.1 h1:X2osQ+RAjK76shCbvhHHHVl3ZlgDm8apHEHFqRjnBY8=
265271
github.com/skeema/knownhosts v1.3.1/go.mod h1:r7KTdC8l4uxWRyK2TpQZ/1o5HaSzh06ePQNxPwTcfiY=
272+
github.com/smartystreets/assertions v0.0.0-20180927180507-b2de0cb4f26d/go.mod h1:OnSkiWE9lh6wB0YB77sQom3nweQdgAjqCqsofrRNTgc=
273+
github.com/smartystreets/goconvey v1.6.4/go.mod h1:syvi0/a8iFYH4r/RixwvyeAJjdLS9QV7WQ/tjFTllLA=
266274
github.com/spf13/cobra v0.0.1/go.mod h1:1l0Ry5zgKvJasoi3XT1TypsSe7PqH0Sj9dhYf7v3XqQ=
267275
github.com/spf13/pflag v1.0.0/go.mod h1:DYY7MBk1bdzusC3SYhjObp+wFpr4gzcvqqNjLnInEg4=
268276
github.com/spf13/pflag v1.0.6 h1:jFzHGLGAlb3ruxLB8MhbI6A8+AQX/2eW4qeyNZXNp2o=
@@ -368,6 +376,7 @@ golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
368376
golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
369377
golang.org/x/net v0.0.0-20181114220301-adae6a3d119a/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
370378
golang.org/x/net v0.0.0-20181220203305-927f97764cc3/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
379+
golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
371380
golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
372381
golang.org/x/net v0.0.0-20190613194153-d28f0bde5980/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
373382
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
@@ -417,6 +426,7 @@ golang.org/x/time v0.8.0 h1:9i3RxcPv3PZnitoVGMPDKZSq1xW1gK1Xy3ArNOGZfEg=
417426
golang.org/x/time v0.8.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
418427
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
419428
golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
429+
golang.org/x/tools v0.0.0-20190328211700-ab21143f2384/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
420430
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
421431
golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
422432
golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=

‎pkg/config/config.go‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,29 @@ import (
55
"os"
66
"path/filepath"
77

8+
"github.com/dustin/go-humanize"
89
"github.com/pkg/errors"
910
"gopkg.in/yaml.v2"
1011
)
1112

13+
type HumanizeSize uint64
14+
15+
func (s *HumanizeSize) UnmarshalYAML(unmarshal func(interface{}) error) error {
16+
var str string
17+
if err := unmarshal(&str); err != nil {
18+
return err
19+
}
20+
21+
size, err := humanize.ParseBytes(str)
22+
if err != nil {
23+
return err
24+
}
25+
26+
*s = HumanizeSize(size)
27+
28+
return nil
29+
}
30+
1231
type Config struct {
1332
// Pattern:
1433
// static: /var/lib/dragonfly/model-csi/volumes/$volumeName/model
@@ -28,8 +47,10 @@ type Config struct {
2847
NodeID string // From env CSI_NODE_ID
2948
Mode string // From env X_CSI_MODE: "controller" or "node"
3049
}
50+
3151
type Features struct {
32-
CheckDiskQuota bool `yaml:"check_disk_quota"`
52+
CheckDiskQuota bool `yaml:"check_disk_quota"`
53+
DiskQuotaMaximumSize HumanizeSize `yaml:"disk_quota_maximum_size"`
3354
}
3455

3556
type PullConfig struct {
@@ -116,7 +137,7 @@ func (cfg *Config) IsNodeMode() bool {
116137
return cfg.Mode == "node"
117138
}
118139

119-
func FromFile(path string) (*Config, error) {
140+
func parse(path string) (*Config, error) {
120141
data, err := os.ReadFile(path)
121142
if err != nil {
122143
return nil, errors.Wrap(err, "read config file")
@@ -184,3 +205,14 @@ func FromFile(path string) (*Config, error) {
184205

185206
return &cfg, nil
186207
}
208+
209+
func New(path string) (*Config, error) {
210+
cfg, err := parse(path)
211+
if err != nil {
212+
return nil, err
213+
}
214+
215+
go cfg.watch(path)
216+
217+
return cfg, nil
218+
}

‎pkg/config/watcher.go‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
package config
2+
3+
import (
4+
"path/filepath"
5+
"sync"
6+
7+
"github.com/fsnotify/fsnotify"
8+
"github.com/modelpack/model-csi-driver/pkg/logger"
9+
)
10+
11+
var mutex = sync.Mutex{}
12+
13+
func (cfg *Config) watch(path string) {
14+
configDir := filepath.Dir(path)
15+
configFile := filepath.Base(path)
16+
17+
watcher, err := fsnotify.NewWatcher()
18+
if err != nil {
19+
logger.Logger().WithError(err).Error("failed to create fsnotify watcher")
20+
return
21+
}
22+
defer func() { _ = watcher.Close() }()
23+
24+
go func() {
25+
for {
26+
select {
27+
case event, ok := <-watcher.Events:
28+
if !ok {
29+
return
30+
}
31+
relPath := filepath.Base(event.Name)
32+
if relPath == configFile && (event.Op&(fsnotify.Write|fsnotify.Create|fsnotify.Remove|fsnotify.Rename)) != 0 {
33+
logger.Logger().Infof("config file changed: %s, event: %s", event.Name, event.Op)
34+
cfg.reload(filepath.Join(configDir, configFile))
35+
}
36+
case err, ok := <-watcher.Errors:
37+
if !ok {
38+
return
39+
}
40+
logger.Logger().WithError(err).Error("watcher error")
41+
}
42+
}
43+
}()
44+
45+
err = watcher.Add(configDir)
46+
if err != nil {
47+
logger.Logger().WithError(err).Error("failed to add config dir to watcher")
48+
}
49+
}
50+
51+
func (cfg *Config) reload(path string) {
52+
newCfg, err := parse(path)
53+
if err != nil {
54+
logger.Logger().WithError(err).Error("failed to parse config file")
55+
return
56+
}
57+
58+
mutex.Lock()
59+
defer mutex.Unlock()
60+
61+
*cfg = *newCfg
62+
}

‎pkg/server/server_test.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ type mockPuller struct {
4040
hook *service.Hook
4141
}
4242

43-
func (puller *mockPuller) Pull(ctx context.Context, reference, targetDir string, checkDiskQuota bool) error {
43+
func (puller *mockPuller) Pull(ctx context.Context, reference, targetDir string) error {
4444
if err := os.MkdirAll(targetDir, 0755); err != nil {
4545
return err
4646
}
@@ -560,13 +560,13 @@ func TestServer(t *testing.T) {
560560
if configPathFromEnv != "" {
561561
defaultCoofigPath = configPathFromEnv
562562
}
563-
cfg, err := config.FromFile(defaultCoofigPath)
563+
cfg, err := config.New(defaultCoofigPath)
564564
require.NoError(t, err)
565565
cfg.RootDir = rootDir
566566
cfg.PullConfig.ProxyURL = ""
567567
service.CacheSacnInterval = 1 * time.Second
568568

569-
service.NewPuller = func(ctx context.Context, pullCfg *config.PullConfig, hook *service.Hook) service.Puller {
569+
service.NewPuller = func(ctx context.Context, pullCfg *config.PullConfig, hook *service.Hook, diskQuotaChecker *service.DiskQuotaChecker) service.Puller {
570570
return &mockPuller{
571571
pullCfg: pullCfg,
572572
duration: time.Second * 2,

‎pkg/service/puller.go‎

Lines changed: 13 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,9 @@ import (
55
"fmt"
66
"io"
77
"os"
8-
"path/filepath"
98
"sort"
109
"sync"
1110
"sync/atomic"
12-
"syscall"
1311
"time"
1412

1513
"github.com/dustin/go-humanize"
@@ -52,19 +50,21 @@ func NewHook(ctx context.Context, progressCb func(progress status.Progress)) *Ho
5250
}
5351

5452
type Puller interface {
55-
Pull(ctx context.Context, reference, targetDir string, checkDiskQuota bool) error
53+
Pull(ctx context.Context, reference, targetDir string) error
5654
}
5755

58-
var NewPuller = func(ctx context.Context, pullCfg *config.PullConfig, hook *Hook) Puller {
56+
var NewPuller = func(ctx context.Context, pullCfg *config.PullConfig, hook *Hook, diskQuotaChecker *DiskQuotaChecker) Puller {
5957
return &puller{
60-
pullCfg: pullCfg,
61-
hook: hook,
58+
pullCfg: pullCfg,
59+
hook: hook,
60+
diskQuotaChecker: diskQuotaChecker,
6261
}
6362
}
6463

6564
type puller struct {
66-
pullCfg *config.PullConfig
67-
hook *Hook
65+
pullCfg *config.PullConfig
66+
hook *Hook
67+
diskQuotaChecker *DiskQuotaChecker
6868
}
6969

7070
func (h *Hook) getProgressDesc() string {
@@ -201,7 +201,7 @@ func (h *Hook) GetProgress() status.Progress {
201201
return h.getProgress()
202202
}
203203

204-
func (p *puller) Pull(ctx context.Context, reference, targetDir string, checkDiskQuota bool) error {
204+
func (p *puller) Pull(ctx context.Context, reference, targetDir string) error {
205205
keyChain, err := auth.GetKeyChainByRef(reference)
206206
if err != nil {
207207
return errors.Wrapf(err, "get auth for model: %s", reference)
@@ -212,9 +212,10 @@ func (p *puller) Pull(ctx context.Context, reference, targetDir string, checkDis
212212
return errors.Wrap(err, "create modctl backend")
213213
}
214214

215-
if checkDiskQuota {
216-
if err := p.checkDiskQuota(ctx, reference, filepath.Dir(targetDir), keyChain.ServerScheme == "http", b); err != nil {
217-
return err
215+
if p.diskQuotaChecker != nil {
216+
plainHTTP := keyChain.ServerScheme == "http"
217+
if err := p.diskQuotaChecker.Check(ctx, b, reference, plainHTTP); err != nil {
218+
return errors.Wrap(err, "check disk quota")
218219
}
219220
}
220221

@@ -247,34 +248,3 @@ func (p *puller) Pull(ctx context.Context, reference, targetDir string, checkDis
247248

248249
return nil
249250
}
250-
251-
func (p *puller) checkDiskQuota(ctx context.Context, reference, dir string, plainHTTP bool, b backend.Backend) error {
252-
var st syscall.Statfs_t
253-
if err := syscall.Statfs(dir, &st); err != nil {
254-
logger.WithContext(ctx).WithError(err).Errorf("failed to stat dir %s in mounting %s", dir, reference)
255-
} else {
256-
availSpace := int64(st.Bavail) * int64(st.Bsize)
257-
logger.WithContext(ctx).Infof("cache dir available space: %s", humanize.IBytes(uint64(availSpace)))
258-
// get model image size
259-
result, err := b.Inspect(ctx, reference, &modctlConfig.Inspect{Remote: true, Insecure: true, PlainHTTP: plainHTTP})
260-
if err != nil {
261-
logger.WithContext(ctx).WithError(err).Errorf("failed to inspect model image: %s", reference)
262-
return errors.Wrap(err, "inspect model image")
263-
}
264-
265-
modelArtifact, ok := result.(*backend.InspectedModelArtifact)
266-
if !ok {
267-
logger.WithContext(ctx).Errorf("invalid inspected result: %s", result)
268-
return fmt.Errorf("invalid inspected result")
269-
}
270-
271-
totalSize := int64(0)
272-
for _, layer := range modelArtifact.Layers {
273-
totalSize += layer.Size
274-
}
275-
if totalSize > availSpace {
276-
return errors.Wrapf(syscall.ENOSPC, "model image %s is %s, but only %s of disk quota is available", reference, humanize.IBytes(uint64(totalSize)), humanize.IBytes(uint64(availSpace)))
277-
}
278-
}
279-
return nil
280-
}

0 commit comments

Comments
 (0)