Skip to content

Commit 701ca73

Browse files
imeoerimeoer
committed
feat: support maximum disk usage limit
Add features.disk_usage_limit option to limit the disk usage of model storage: - disk_usage_limit == 0: reject if available disk space < model size; - disk_usage_limit > 0: reject if (disk_usage_limit - used space) < model size; Signed-off-by: imeoer <yansong.ys@antgroup.com>
1 parent b588994 commit 701ca73

File tree

12 files changed

+508
-53
lines changed

12 files changed

+508
-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+
DiskUsageLimit HumanizeSize `yaml:"disk_usage_limit"`
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/config_test.go‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
package config
2+
3+
import (
4+
"os"
5+
"path/filepath"
6+
"strings"
7+
"testing"
8+
"time"
9+
10+
"github.com/stretchr/testify/require"
11+
)
12+
13+
func copyFile(t *testing.T, src, dst string) {
14+
data, err := os.ReadFile(src)
15+
require.NoError(t, err)
16+
17+
err = os.WriteFile(dst, data, 0644)
18+
require.NoError(t, err)
19+
}
20+
21+
func TestConfig(t *testing.T) {
22+
// Create a temporary directory for testing
23+
tmpDir, err := os.MkdirTemp("", "config-test-")
24+
require.NoError(t, err)
25+
defer func() { _ = os.RemoveAll(tmpDir) }()
26+
27+
require.NoError(t, os.Setenv("X_CSI_MODE", "node"))
28+
require.NoError(t, os.Setenv("CSI_NODE_ID", "test-node"))
29+
30+
// Prepare the origin config file
31+
testConfigPath := "../../test/testdata/config.test.yaml"
32+
configPath := filepath.Join(tmpDir, "config.yaml")
33+
copyFile(t, testConfigPath, configPath)
34+
cfg, err := New(configPath)
35+
require.NoError(t, err)
36+
37+
// Wait watcher to start
38+
time.Sleep(time.Second)
39+
40+
// Update the origin config file (k8s configmap's atomic update is rename)
41+
tmpConfigPath := filepath.Join(tmpDir, "config.tmp.yaml")
42+
copyFile(t, testConfigPath, tmpConfigPath)
43+
data, err := os.ReadFile(tmpConfigPath)
44+
require.NoError(t, err)
45+
updatedData := strings.Replace(string(data), "disk_usage_limit: 10TiB", "disk_usage_limit: 5TiB", 1)
46+
require.NoError(t, os.WriteFile(tmpConfigPath, []byte(updatedData), 0644))
47+
require.NoError(t, os.Rename(tmpConfigPath, configPath))
48+
49+
// Wait watcher to reload the config
50+
time.Sleep(time.Second * 1)
51+
52+
// Verify the config is reloaded
53+
require.Equal(t, uint64(0x50000000000), uint64(cfg.Features.DiskUsageLimit))
54+
}

‎pkg/config/watcher.go‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
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+
defer logger.Logger().Warn("fsnotify watcher goroutine exited")
26+
for {
27+
select {
28+
case event, ok := <-watcher.Events:
29+
if !ok {
30+
return
31+
}
32+
relPath := filepath.Base(event.Name)
33+
if relPath == configFile && (event.Op&(fsnotify.Write|fsnotify.Create|fsnotify.Remove|fsnotify.Rename)) != 0 {
34+
logger.Logger().Infof("config file changed: %s, event: %s", event.Name, event.Op)
35+
cfg.reload(filepath.Join(configDir, configFile))
36+
}
37+
case err, ok := <-watcher.Errors:
38+
if !ok {
39+
return
40+
}
41+
logger.Logger().WithError(err).Error("watcher error")
42+
}
43+
}
44+
}()
45+
46+
err = watcher.Add(configDir)
47+
if err != nil {
48+
logger.Logger().WithError(err).Error("failed to add config dir to watcher")
49+
}
50+
51+
select {}
52+
}
53+
54+
func (cfg *Config) reload(path string) {
55+
newCfg, err := parse(path)
56+
if err != nil {
57+
logger.Logger().WithError(err).Error("failed to parse config file")
58+
return
59+
}
60+
61+
mutex.Lock()
62+
defer mutex.Unlock()
63+
64+
*cfg = *newCfg
65+
66+
logger.Logger().Infof("config reloaded: %s", path)
67+
}

‎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,

0 commit comments

Comments
 (0)