From e0066557d83121bf9ad4d4a015915e4876f1d6c2 Mon Sep 17 00:00:00 2001
From: Will Sackfield <sackfield@replicate.com>
Date: Wed, 15 Jan 2025 12:13:56 -0500
Subject: [PATCH] Fix file not found when evaluating weights

---
 pkg/dockerfile/weights.go                     |  3 +++
 pkg/dockerfile/weights_test.go                | 25 +++++++++++++++++++
 .../fixtures/fast-build/cog.yaml              |  1 +
 3 files changed, 29 insertions(+)

diff --git a/pkg/dockerfile/weights.go b/pkg/dockerfile/weights.go
index 8c09d7b814..a951a3d610 100644
--- 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
index b8173b4905..0204545930 100644
--- 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
index ac4bac61e5..62419bfb88 100644
--- 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"