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 <winters.zc@antgroup.com>

Author: aftersnow

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
 }

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
+}

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)

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)
+	}
+}

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)
+		}
+	}
+}