feat: support maximum disk usage limit

Signed-off-by: imeoer <yansong.ys@antgroup.com>

Author: imeoer

cmd/model-csi-driver/main.go

@@ -38,7 +38,7 @@ func main() {
 			},
 		},
 		Action: func(c *cli.Context) error {
-			cfg, err := config.FromFile(c.String("config"))
+			cfg, err := config.New(c.String("config"))
 			if err != nil {
 				return errors.Wrap(err, "load config")
 			}

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/container-storage-interface/spec v1.2.0
 	github.com/containerd/containerd v1.7.27
 	github.com/dustin/go-humanize v1.0.1
+	github.com/fsnotify/fsnotify v1.8.0
 	github.com/google/uuid v1.6.0
 	github.com/labstack/echo/v4 v4.13.3
 	github.com/moby/sys/mountinfo v0.7.2
@@ -27,6 +28,7 @@ require (
 	go.opentelemetry.io/otel/trace v1.37.0
 	golang.org/x/net v0.42.0
 	golang.org/x/sync v0.16.0
+	golang.org/x/sys v0.35.0
 	google.golang.org/grpc v1.75.0
 	gopkg.in/yaml.v2 v2.4.0
 	k8s.io/api v0.28.4
@@ -128,7 +130,6 @@ require (
 	go.opentelemetry.io/proto/otlp v1.6.0 // indirect
 	golang.org/x/crypto v0.41.0 // indirect
 	golang.org/x/oauth2 v0.30.0 // indirect
-	golang.org/x/sys v0.35.0 // indirect
 	golang.org/x/term v0.34.0 // indirect
 	golang.org/x/text v0.28.0 // indirect
 	golang.org/x/time v0.8.0 // indirect

go.sum

@@ -78,6 +78,8 @@ github.com/emirpasic/gods v1.18.1 h1:FXtiHYKDGKCW2KzwZKx0iC0PQmdlorYgdFG9jPXJ1Bc
 github.com/emirpasic/gods v1.18.1/go.mod h1:8tpGGwCnJ5H4r6BWwaV6OrWmMoPhUl5jm/FMNAnJvWQ=
 github.com/envoyproxy/protoc-gen-validate v1.2.1 h1:DEo3O99U8j4hBFwbJfrz9VtgcDfUKS7KJ7spH3d86P8=
 github.com/envoyproxy/protoc-gen-validate v1.2.1/go.mod h1:d/C80l/jxXLdfEIhX1W2TmLfsJ31lvEjwamM4DxlWXU=
+github.com/fsnotify/fsnotify v1.8.0 h1:dAwr6QBTBZIkG8roQaJjGof0pp0EeF+tNV7YBP3F/8M=
+github.com/fsnotify/fsnotify v1.8.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0=
 github.com/gliderlabs/ssh v0.3.8 h1:a4YXD1V7xMF9g5nTkdfnja3Sxy1PVDCj1Zg4Wb8vY6c=
 github.com/gliderlabs/ssh v0.3.8/go.mod h1:xYoytBv1sV0aL3CavoDuJIQNURXkkfPA/wxQ1pL1fAU=
 github.com/go-git/gcfg v1.5.1-0.20230307220236-3a3c6141e376 h1:+zs/tPmkDkHx3U66DAb0lQFJrpS6731Oaa12ikc+DiI=

pkg/config/config.go

@@ -5,10 +5,29 @@ import (
 	"os"
 	"path/filepath"
 
+	"github.com/dustin/go-humanize"
 	"github.com/pkg/errors"
 	"gopkg.in/yaml.v2"
 )
 
+type HumanizeSize uint64
+
+func (s *HumanizeSize) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var str string
+	if err := unmarshal(&str); err != nil {
+		return err
+	}
+
+	size, err := humanize.ParseBytes(str)
+	if err != nil {
+		return err
+	}
+
+	*s = HumanizeSize(size)
+
+	return nil
+}
+
 type Config struct {
 	// Pattern:
 	// 	static: /var/lib/dragonfly/model-csi/volumes/$volumeName/model
@@ -28,8 +47,10 @@ type Config struct {
 	NodeID                   string     // From env CSI_NODE_ID
 	Mode                     string     // From env X_CSI_MODE: "controller" or "node"
 }
+
 type Features struct {
-	CheckDiskQuota bool `yaml:"check_disk_quota"`
+	CheckDiskQuota       bool         `yaml:"check_disk_quota"`
+	DiskQuotaMaximumSize HumanizeSize `yaml:"disk_quota_maximum_size"`
 }
 
 type PullConfig struct {
@@ -116,7 +137,7 @@ func (cfg *Config) IsNodeMode() bool {
 	return cfg.Mode == "node"
 }
 
-func FromFile(path string) (*Config, error) {
+func parse(path string) (*Config, error) {
 	data, err := os.ReadFile(path)
 	if err != nil {
 		return nil, errors.Wrap(err, "read config file")
@@ -184,3 +205,14 @@ func FromFile(path string) (*Config, error) {
 
 	return &cfg, nil
 }
+
+func New(path string) (*Config, error) {
+	cfg, err := parse(path)
+	if err != nil {
+		return nil, err
+	}
+
+	go cfg.watch(path)
+
+	return cfg, nil
+}

pkg/config/watcher.go

@@ -0,0 +1,62 @@
+package config
+
+import (
+	"path/filepath"
+	"sync"
+
+	"github.com/fsnotify/fsnotify"
+	"github.com/modelpack/model-csi-driver/pkg/logger"
+)
+
+var mutex = sync.Mutex{}
+
+func (cfg *Config) watch(path string) {
+	configDir := filepath.Dir(path)
+	configFile := filepath.Base(path)
+
+	watcher, err := fsnotify.NewWatcher()
+	if err != nil {
+		logger.Logger().WithError(err).Error("failed to create fsnotify watcher")
+		return
+	}
+	defer func() { _ = watcher.Close() }()
+
+	go func() {
+		for {
+			select {
+			case event, ok := <-watcher.Events:
+				if !ok {
+					return
+				}
+				relPath := filepath.Base(event.Name)
+				if relPath == configFile && (event.Op&(fsnotify.Write|fsnotify.Create|fsnotify.Remove|fsnotify.Rename)) != 0 {
+					logger.Logger().Infof("config file changed: %s, event: %s", event.Name, event.Op)
+					cfg.reload(filepath.Join(configDir, configFile))
+				}
+			case err, ok := <-watcher.Errors:
+				if !ok {
+					return
+				}
+				logger.Logger().WithError(err).Error("watcher error")
+			}
+		}
+	}()
+
+	err = watcher.Add(configDir)
+	if err != nil {
+		logger.Logger().WithError(err).Error("failed to add config dir to watcher")
+	}
+}
+
+func (cfg *Config) reload(path string) {
+	newCfg, err := parse(path)
+	if err != nil {
+		logger.Logger().WithError(err).Error("failed to parse config file")
+		return
+	}
+
+	mutex.Lock()
+	defer mutex.Unlock()
+
+	*cfg = *newCfg
+}

pkg/server/server_test.go

@@ -40,7 +40,7 @@ type mockPuller struct {
 	hook     *service.Hook
 }
 
-func (puller *mockPuller) Pull(ctx context.Context, reference, targetDir string, checkDiskQuota bool) error {
+func (puller *mockPuller) Pull(ctx context.Context, reference, targetDir string) error {
 	if err := os.MkdirAll(targetDir, 0755); err != nil {
 		return err
 	}
@@ -560,13 +560,13 @@ func TestServer(t *testing.T) {
 	if configPathFromEnv != "" {
 		defaultCoofigPath = configPathFromEnv
 	}
-	cfg, err := config.FromFile(defaultCoofigPath)
+	cfg, err := config.New(defaultCoofigPath)
 	require.NoError(t, err)
 	cfg.RootDir = rootDir
 	cfg.PullConfig.ProxyURL = ""
 	service.CacheSacnInterval = 1 * time.Second
 
-	service.NewPuller = func(ctx context.Context, pullCfg *config.PullConfig, hook *service.Hook) service.Puller {
+	service.NewPuller = func(ctx context.Context, pullCfg *config.PullConfig, hook *service.Hook, diskQuotaChecker *service.DiskQuotaChecker) service.Puller {
 		return &mockPuller{
 			pullCfg:  pullCfg,
 			duration: time.Second * 2,

pkg/service/puller.go

@@ -5,11 +5,9 @@ import (
 	"fmt"
 	"io"
 	"os"
-	"path/filepath"
 	"sort"
 	"sync"
 	"sync/atomic"
-	"syscall"
 	"time"
 
 	"github.com/dustin/go-humanize"
@@ -52,19 +50,21 @@ func NewHook(ctx context.Context, progressCb func(progress status.Progress)) *Ho
 }
 
 type Puller interface {
-	Pull(ctx context.Context, reference, targetDir string, checkDiskQuota bool) error
+	Pull(ctx context.Context, reference, targetDir string) error
 }
 
-var NewPuller = func(ctx context.Context, pullCfg *config.PullConfig, hook *Hook) Puller {
+var NewPuller = func(ctx context.Context, pullCfg *config.PullConfig, hook *Hook, diskQuotaChecker *DiskQuotaChecker) Puller {
 	return &puller{
-		pullCfg: pullCfg,
-		hook:    hook,
+		pullCfg:          pullCfg,
+		hook:             hook,
+		diskQuotaChecker: diskQuotaChecker,
 	}
 }
 
 type puller struct {
-	pullCfg *config.PullConfig
-	hook    *Hook
+	pullCfg          *config.PullConfig
+	hook             *Hook
+	diskQuotaChecker *DiskQuotaChecker
 }
 
 func (h *Hook) getProgressDesc() string {
@@ -201,7 +201,7 @@ func (h *Hook) GetProgress() status.Progress {
 	return h.getProgress()
 }
 
-func (p *puller) Pull(ctx context.Context, reference, targetDir string, checkDiskQuota bool) error {
+func (p *puller) Pull(ctx context.Context, reference, targetDir string) error {
 	keyChain, err := auth.GetKeyChainByRef(reference)
 	if err != nil {
 		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
 		return errors.Wrap(err, "create modctl backend")
 	}
 
-	if checkDiskQuota {
-		if err := p.checkDiskQuota(ctx, reference, filepath.Dir(targetDir), keyChain.ServerScheme == "http", b); err != nil {
-			return err
+	if p.diskQuotaChecker != nil {
+		plainHTTP := keyChain.ServerScheme == "http"
+		if err := p.diskQuotaChecker.Check(ctx, b, reference, plainHTTP); err != nil {
+			return errors.Wrap(err, "check disk quota")
 		}
 	}
 
@@ -247,34 +248,3 @@ func (p *puller) Pull(ctx context.Context, reference, targetDir string, checkDis
 
 	return nil
 }
-
-func (p *puller) checkDiskQuota(ctx context.Context, reference, dir string, plainHTTP bool, b backend.Backend) error {
-	var st syscall.Statfs_t
-	if err := syscall.Statfs(dir, &st); err != nil {
-		logger.WithContext(ctx).WithError(err).Errorf("failed to stat dir %s in mounting %s", dir, reference)
-	} else {
-		availSpace := int64(st.Bavail) * int64(st.Bsize)
-		logger.WithContext(ctx).Infof("cache dir available space: %s", humanize.IBytes(uint64(availSpace)))
-		// get model image size
-		result, err := b.Inspect(ctx, reference, &modctlConfig.Inspect{Remote: true, Insecure: true, PlainHTTP: plainHTTP})
-		if err != nil {
-			logger.WithContext(ctx).WithError(err).Errorf("failed to inspect model image: %s", reference)
-			return errors.Wrap(err, "inspect model image")
-		}
-
-		modelArtifact, ok := result.(*backend.InspectedModelArtifact)
-		if !ok {
-			logger.WithContext(ctx).Errorf("invalid inspected result: %s", result)
-			return fmt.Errorf("invalid inspected result")
-		}
-
-		totalSize := int64(0)
-		for _, layer := range modelArtifact.Layers {
-			totalSize += layer.Size
-		}
-		if totalSize > availSpace {
-			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)))
-		}
-	}
-	return nil
-}