From ead5d37ca40504bd890add882b97ba1cbfaaee16 Mon Sep 17 00:00:00 2001 From: Zhao Chen Date: Mon, 23 Mar 2026 22:09:29 +0800 Subject: [PATCH 1/3] feat: add disk space pre-check for build and pull operations Check available disk space before writing to local storage. When space is insufficient (with 10% safety margin), log a warning but continue the operation without returning an error. Signed-off-by: Zhao Chen --- pkg/backend/backend.go | 6 ++- pkg/backend/build.go | 45 +++++++++++++++++ pkg/backend/pull.go | 16 +++++++ pkg/diskspace/check.go | 96 +++++++++++++++++++++++++++++++++++++ pkg/diskspace/check_test.go | 85 ++++++++++++++++++++++++++++++++ 5 files changed, 246 insertions(+), 2 deletions(-) create mode 100644 pkg/diskspace/check.go create mode 100644 pkg/diskspace/check_test.go diff --git a/pkg/backend/backend.go b/pkg/backend/backend.go index dff0ea9..9515ac6 100644 --- a/pkg/backend/backend.go +++ b/pkg/backend/backend.go @@ -70,7 +70,8 @@ type Backend interface { // backend is the implementation of Backend. type backend struct { - store storage.Storage + store storage.Storage + storageDir string } // New creates a new backend. @@ -81,6 +82,7 @@ func New(storageDir string) (Backend, error) { } return &backend{ - store: store, + store: store, + storageDir: storageDir, }, nil } diff --git a/pkg/backend/build.go b/pkg/backend/build.go index 10f9156..1f0ffec 100644 --- a/pkg/backend/build.go +++ b/pkg/backend/build.go @@ -34,6 +34,7 @@ import ( "github.com/modelpack/modctl/pkg/backend/build/hooks" "github.com/modelpack/modctl/pkg/backend/processor" "github.com/modelpack/modctl/pkg/config" + "github.com/modelpack/modctl/pkg/diskspace" "github.com/modelpack/modctl/pkg/modelfile" "github.com/modelpack/modctl/pkg/source" ) @@ -67,6 +68,14 @@ func (b *backend) Build(ctx context.Context, modelfilePath, workDir, target stri return fmt.Errorf("failed to get source info: %w", err) } + // Check disk space before building (only for local output). + if !cfg.OutputRemote { + totalSize := estimateBuildSize(workDir, modelfile) + if err := diskspace.Check(b.storageDir, totalSize); err != nil { + logrus.Warnf("build: %v", err) + } + } + // using the local output by default. outputType := build.OutputTypeLocal if cfg.OutputRemote { @@ -263,3 +272,39 @@ func getSourceInfo(workspace string, buildConfig *config.Build) (*source.Info, e return info, nil } + +// estimateBuildSize estimates the total size of files that will be built by summing +// the sizes of all files referenced in the modelfile. +func estimateBuildSize(workDir string, mf modelfile.Modelfile) int64 { + var totalSize int64 + + files := []string{} + files = append(files, mf.GetConfigs()...) + files = append(files, mf.GetModels()...) + files = append(files, mf.GetCodes()...) + files = append(files, mf.GetDocs()...) + + for _, file := range files { + path := filepath.Join(workDir, file) + info, err := os.Stat(path) + if err != nil { + logrus.Debugf("build: failed to stat file %s for size estimation: %v", path, err) + continue + } + if info.IsDir() { + _ = filepath.Walk(path, func(_ string, fi os.FileInfo, err error) error { + if err != nil { + return nil + } + if !fi.IsDir() { + totalSize += fi.Size() + } + return nil + }) + } else { + totalSize += info.Size() + } + } + + return totalSize +} diff --git a/pkg/backend/pull.go b/pkg/backend/pull.go index 3cb6a1b..cb4b930 100644 --- a/pkg/backend/pull.go +++ b/pkg/backend/pull.go @@ -33,6 +33,7 @@ import ( "github.com/modelpack/modctl/pkg/backend/remote" "github.com/modelpack/modctl/pkg/codec" "github.com/modelpack/modctl/pkg/config" + "github.com/modelpack/modctl/pkg/diskspace" "github.com/modelpack/modctl/pkg/storage" ) @@ -72,6 +73,21 @@ func (b *backend) Pull(ctx context.Context, target string, cfg *config.Pull) err logrus.Debugf("pull: loaded manifest for target %s [manifest: %+v]", target, manifest) + // Check disk space before pulling layers. + var totalSize int64 + for _, layer := range manifest.Layers { + totalSize += layer.Size + } + totalSize += manifest.Config.Size + + targetDir := b.storageDir + if cfg.ExtractFromRemote && cfg.ExtractDir != "" { + targetDir = cfg.ExtractDir + } + if err := diskspace.Check(targetDir, totalSize); err != nil { + logrus.Warnf("pull: %v", err) + } + // TODO: need refactor as currently use a global flag to control the progress bar render. if cfg.DisableProgress { internalpb.SetDisableProgress(true) diff --git a/pkg/diskspace/check.go b/pkg/diskspace/check.go new file mode 100644 index 0000000..25349af --- /dev/null +++ b/pkg/diskspace/check.go @@ -0,0 +1,96 @@ +/* + * Copyright 2024 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 diskspace + +import ( + "fmt" + "os" + "path/filepath" + + "golang.org/x/sys/unix" +) + +const ( + // safetyMargin is the extra space ratio to account for metadata overhead + // (manifests, temporary files, etc.). 10% extra required. + safetyMargin = 1.1 +) + +// Check checks if the directory has enough disk space for the required bytes. +// It returns a descriptive error if space is insufficient, or nil if space is enough. +// The caller should use the returned error for warning purposes only and not +// treat it as a fatal error. +func Check(dir string, requiredBytes int64) error { + if requiredBytes <= 0 { + return nil + } + + // Ensure the directory exists for statfs; walk up to find an existing parent. + checkDir := dir + for { + if _, err := os.Stat(checkDir); err == nil { + break + } + parent := filepath.Dir(checkDir) + if parent == checkDir { + // Reached filesystem root without finding an existing directory. + return fmt.Errorf("cannot determine disk space: no existing directory found for path %s", dir) + } + checkDir = parent + } + + var stat unix.Statfs_t + if err := unix.Statfs(checkDir, &stat); err != nil { + return fmt.Errorf("failed to check disk space for %s: %w", dir, err) + } + + // Available space for non-root users. + availableBytes := int64(stat.Bavail) * int64(stat.Bsize) + requiredWithMargin := int64(float64(requiredBytes) * safetyMargin) + + if availableBytes < requiredWithMargin { + return fmt.Errorf( + "insufficient disk space in %s: available %s, required %s (with 10%% safety margin)", + dir, formatBytes(availableBytes), formatBytes(requiredWithMargin), + ) + } + + return nil +} + +// formatBytes formats bytes into a human-readable string. +func formatBytes(bytes int64) string { + const ( + kb = 1024 + mb = kb * 1024 + gb = mb * 1024 + tb = gb * 1024 + ) + + switch { + case bytes >= tb: + return fmt.Sprintf("%.2f TB", float64(bytes)/float64(tb)) + case bytes >= gb: + return fmt.Sprintf("%.2f GB", float64(bytes)/float64(gb)) + case bytes >= mb: + return fmt.Sprintf("%.2f MB", float64(bytes)/float64(mb)) + case bytes >= kb: + return fmt.Sprintf("%.2f KB", float64(bytes)/float64(kb)) + default: + return fmt.Sprintf("%d B", bytes) + } +} diff --git a/pkg/diskspace/check_test.go b/pkg/diskspace/check_test.go new file mode 100644 index 0000000..a5020c5 --- /dev/null +++ b/pkg/diskspace/check_test.go @@ -0,0 +1,85 @@ +/* + * Copyright 2024 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 diskspace + +import ( + "strings" + "testing" +) + +func TestCheck_ZeroBytes(t *testing.T) { + err := Check("/tmp", 0) + if err != nil { + t.Errorf("expected nil error for zero bytes, got: %v", err) + } +} + +func TestCheck_NegativeBytes(t *testing.T) { + err := Check("/tmp", -1) + if err != nil { + t.Errorf("expected nil error for negative bytes, got: %v", err) + } +} + +func TestCheck_SmallSize(t *testing.T) { + // 1 byte should always have enough space + err := Check("/tmp", 1) + if err != nil { + t.Errorf("expected nil error for 1 byte, got: %v", err) + } +} + +func TestCheck_ExtremelyLargeSize(t *testing.T) { + // 1 exabyte should always fail + err := Check("/tmp", 1<<60) + if err == nil { + t.Error("expected error for extremely large size, got nil") + } + if !strings.Contains(err.Error(), "insufficient disk space") { + t.Errorf("expected 'insufficient disk space' in error, got: %v", err) + } +} + +func TestCheck_NonExistentDirWalksUp(t *testing.T) { + // Should walk up to find an existing parent directory + err := Check("/tmp/nonexistent-modctl-test-dir-12345/subdir", 1) + if err != nil { + t.Errorf("expected nil error when parent exists, got: %v", err) + } +} + +func TestFormatBytes(t *testing.T) { + tests := []struct { + bytes int64 + expected string + }{ + {0, "0 B"}, + {500, "500 B"}, + {1024, "1.00 KB"}, + {1536, "1.50 KB"}, + {1048576, "1.00 MB"}, + {1073741824, "1.00 GB"}, + {1099511627776, "1.00 TB"}, + } + + for _, tt := range tests { + result := formatBytes(tt.bytes) + if result != tt.expected { + t.Errorf("formatBytes(%d) = %s, want %s", tt.bytes, result, tt.expected) + } + } +} From 007f722a19c4e2b2912d7292d68244303d853ab0 Mon Sep 17 00:00:00 2001 From: Zhao Chen Date: Mon, 23 Mar 2026 22:23:26 +0800 Subject: [PATCH 2/3] fix: log filepath.Walk errors during build size estimation Add debug logging for errors encountered during directory walking in estimateBuildSize, so that inaccessible paths are visible instead of being silently ignored. Signed-off-by: Zhao Chen --- pkg/backend/build.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pkg/backend/build.go b/pkg/backend/build.go index 1f0ffec..01cfd7f 100644 --- a/pkg/backend/build.go +++ b/pkg/backend/build.go @@ -292,8 +292,9 @@ func estimateBuildSize(workDir string, mf modelfile.Modelfile) int64 { continue } if info.IsDir() { - _ = filepath.Walk(path, func(_ string, fi os.FileInfo, err error) error { + _ = filepath.Walk(path, func(walkPath string, fi os.FileInfo, err error) error { if err != nil { + logrus.Debugf("build: failed to access path %s for size estimation: %v", walkPath, err) return nil } if !fi.IsDir() { From 808f64149d4abe7bbd6abcf97a9bc6b8f1d82e82 Mon Sep 17 00:00:00 2001 From: Zhao Chen Date: Tue, 24 Mar 2026 10:49:47 +0800 Subject: [PATCH 3/3] fix: address review feedback for disk space check - Add overflow guard for Bavail * Bsize multiplication (cap at MaxInt64) - Handle negative values in formatBytes (return "0 B") - Add edge case tests for Check (empty path, root dir, deeply nested) - Add unit tests for estimateBuildSize (files, dirs, missing, empty) - Fix copyright year to 2025 in new files Signed-off-by: Zhao Chen --- pkg/backend/build_test.go | 72 +++++++++++++++++++++++++++++++++++-- pkg/diskspace/check.go | 18 ++++++++-- pkg/diskspace/check_test.go | 26 +++++++++++++- 3 files changed, 111 insertions(+), 5 deletions(-) diff --git a/pkg/backend/build_test.go b/pkg/backend/build_test.go index bedaa76..4bf02c2 100644 --- a/pkg/backend/build_test.go +++ b/pkg/backend/build_test.go @@ -17,12 +17,14 @@ package backend import ( + "os" + "path/filepath" "testing" + "github.com/stretchr/testify/assert" + "github.com/modelpack/modctl/pkg/config" "github.com/modelpack/modctl/test/mocks/modelfile" - - "github.com/stretchr/testify/assert" ) func TestGetProcessors(t *testing.T) { @@ -41,3 +43,69 @@ func TestGetProcessors(t *testing.T) { assert.Equal(t, "code", processors[2].Name()) assert.Equal(t, "doc", processors[3].Name()) } + +func TestEstimateBuildSize(t *testing.T) { + t.Run("single files", func(t *testing.T) { + workDir := t.TempDir() + + // Create test files with known sizes. + assert.NoError(t, os.WriteFile(filepath.Join(workDir, "model.bin"), make([]byte, 1024), 0644)) + assert.NoError(t, os.WriteFile(filepath.Join(workDir, "config.json"), make([]byte, 256), 0644)) + + mf := &modelfile.Modelfile{} + mf.On("GetConfigs").Return([]string{"config.json"}) + mf.On("GetModels").Return([]string{"model.bin"}) + mf.On("GetCodes").Return([]string{}) + mf.On("GetDocs").Return([]string{}) + + size := estimateBuildSize(workDir, mf) + assert.Equal(t, int64(1280), size) + }) + + t.Run("directory entry", func(t *testing.T) { + workDir := t.TempDir() + + // Create a subdirectory with files. + subDir := filepath.Join(workDir, "models") + assert.NoError(t, os.MkdirAll(subDir, 0755)) + assert.NoError(t, os.WriteFile(filepath.Join(subDir, "a.bin"), make([]byte, 512), 0644)) + assert.NoError(t, os.WriteFile(filepath.Join(subDir, "b.bin"), make([]byte, 512), 0644)) + + mf := &modelfile.Modelfile{} + mf.On("GetConfigs").Return([]string{}) + mf.On("GetModels").Return([]string{"models"}) + mf.On("GetCodes").Return([]string{}) + mf.On("GetDocs").Return([]string{}) + + size := estimateBuildSize(workDir, mf) + assert.Equal(t, int64(1024), size) + }) + + t.Run("nonexistent file is skipped", func(t *testing.T) { + workDir := t.TempDir() + + assert.NoError(t, os.WriteFile(filepath.Join(workDir, "real.bin"), make([]byte, 100), 0644)) + + mf := &modelfile.Modelfile{} + mf.On("GetConfigs").Return([]string{}) + mf.On("GetModels").Return([]string{"real.bin", "missing.bin"}) + mf.On("GetCodes").Return([]string{}) + mf.On("GetDocs").Return([]string{}) + + size := estimateBuildSize(workDir, mf) + assert.Equal(t, int64(100), size) + }) + + t.Run("empty modelfile", func(t *testing.T) { + workDir := t.TempDir() + + mf := &modelfile.Modelfile{} + mf.On("GetConfigs").Return([]string{}) + mf.On("GetModels").Return([]string{}) + mf.On("GetCodes").Return([]string{}) + mf.On("GetDocs").Return([]string{}) + + size := estimateBuildSize(workDir, mf) + assert.Equal(t, int64(0), size) + }) +} diff --git a/pkg/diskspace/check.go b/pkg/diskspace/check.go index 25349af..b04bfed 100644 --- a/pkg/diskspace/check.go +++ b/pkg/diskspace/check.go @@ -1,5 +1,5 @@ /* - * Copyright 2024 The CNAI Authors + * 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. @@ -18,6 +18,7 @@ package diskspace import ( "fmt" + "math" "os" "path/filepath" @@ -59,7 +60,16 @@ func Check(dir string, requiredBytes int64) error { } // Available space for non-root users. - availableBytes := int64(stat.Bavail) * int64(stat.Bsize) + // Guard against overflow: on Linux Bavail is uint64, and values exceeding + // math.MaxInt64 would wrap negative when cast to int64. Cap at MaxInt64. + bavail := stat.Bavail + bsize := uint64(stat.Bsize) + var availableBytes int64 + if bavail > 0 && bsize > uint64(math.MaxInt64)/bavail { + availableBytes = math.MaxInt64 + } else { + availableBytes = int64(bavail * bsize) + } requiredWithMargin := int64(float64(requiredBytes) * safetyMargin) if availableBytes < requiredWithMargin { @@ -74,6 +84,10 @@ func Check(dir string, requiredBytes int64) error { // formatBytes formats bytes into a human-readable string. func formatBytes(bytes int64) string { + if bytes < 0 { + return "0 B" + } + const ( kb = 1024 mb = kb * 1024 diff --git a/pkg/diskspace/check_test.go b/pkg/diskspace/check_test.go index a5020c5..24bbb62 100644 --- a/pkg/diskspace/check_test.go +++ b/pkg/diskspace/check_test.go @@ -1,5 +1,5 @@ /* - * Copyright 2024 The CNAI Authors + * 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. @@ -62,11 +62,35 @@ func TestCheck_NonExistentDirWalksUp(t *testing.T) { } } +func TestCheck_EmptyPath(t *testing.T) { + // Empty string should walk up to root and succeed for small size + err := Check("", 1) + if err != nil { + t.Errorf("expected nil error for empty path with 1 byte, got: %v", err) + } +} + +func TestCheck_RootDir(t *testing.T) { + err := Check("/", 1) + if err != nil { + t.Errorf("expected nil error for root dir with 1 byte, got: %v", err) + } +} + +func TestCheck_DeeplyNestedNonExistentDir(t *testing.T) { + err := Check("/tmp/a/b/c/d/e/f/g/h", 1) + if err != nil { + t.Errorf("expected nil error for deeply nested non-existent path, got: %v", err) + } +} + func TestFormatBytes(t *testing.T) { tests := []struct { bytes int64 expected string }{ + {-1, "0 B"}, + {-1024, "0 B"}, {0, "0 B"}, {500, "500 B"}, {1024, "1.00 KB"},