Fix file not found when evaluating weights

replicate · Jan 15, 2025 · e006655 · e006655
1 parent 9dbbe5a
commit e006655
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 0 deletions.
diff --git a/pkg/dockerfile/weights.go b/pkg/dockerfile/weights.go
@@ -139,6 +139,9 @@ func checkWeights(folder string, weightFile string) ([]Weight, error) {
 	for _, weight := range weights {
 		info, err := os.Stat(weight.Path)
 		if err != nil {
+			if errors.Is(err, os.ErrNotExist) {
+				continue
+			}
 			return nil, err
 		}
 

diff --git a/pkg/dockerfile/weights_test.go b/pkg/dockerfile/weights_test.go
@@ -1,7 +1,11 @@
 package dockerfile
 
 import (
+	"encoding/json"
+	"os"
+	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 )
@@ -13,3 +17,24 @@ func TestFindWeights(t *testing.T) {
 	require.NoError(t, err)
 	require.Empty(t, weights)
 }
+
+func TestFindWeightsWithRemovedWeight(t *testing.T) {
+	folder := t.TempDir()
+	tmpDir := t.TempDir()
+	weightFile := filepath.Join(tmpDir, WEIGHT_FILE)
+	weights := []Weight{
+		{
+			Path:      "nonexistant_weight.h5",
+			Digest:    "1",
+			Timestamp: time.Now(),
+			Size:      10,
+		},
+	}
+	jsonData, err := json.MarshalIndent(weights, "", "  ")
+	require.NoError(t, err)
+	err = os.WriteFile(weightFile, jsonData, 0o644)
+	require.NoError(t, err)
+	weights, err = FindWeights(folder, tmpDir)
+	require.NoError(t, err)
+	require.Empty(t, weights)
+}
diff --git a/test-integration/test_integration/fixtures/fast-build/cog.yaml b/test-integration/test_integration/fixtures/fast-build/cog.yaml
@@ -6,3 +6,4 @@ build:
   system_packages:
     - "git"
 predict: "predict.py:Predictor"
+image: "test"