Fix output Iterator[Path] case

Author: nevillelyh

internal/server/runner.go

@@ -33,6 +33,7 @@ type PendingPrediction struct {
 	response    PredictionResponse
 	lastUpdated time.Time
 	inputPaths  []string
+	outputCache map[string]string
 	mu          sync.Mutex
 	c           chan PredictionResponse
 }
@@ -250,9 +251,10 @@ func (r *Runner) Predict(req PredictionRequest) (chan PredictionResponse, error)
 		StartedAt: req.StartedAt,
 	}
 	pr := PendingPrediction{
-		request:    req,
-		response:   resp,
-		inputPaths: inputPaths,
+		request:     req,
+		response:    resp,
+		inputPaths:  inputPaths,
+		outputCache: make(map[string]string),
 	}
 	if req.Webhook == "" {
 		pr.c = make(chan PredictionResponse, 1)
@@ -515,7 +517,23 @@ func (r *Runner) handleResponses() {
 		} else if r.uploadUrl != "" {
 			outputFn = outputToUpload(r.uploadUrl, pr.response.Id)
 		}
-		if output, err := handlePath(pr.response.Output, &paths, outputFn); err != nil {
+		cachedOutputFn := func(s string, paths *[]string) (string, error) {
+			// Cache already handled output files to avoid duplicates or deleted files in Iterator[Path]
+			if cache, ok := pr.outputCache[s]; ok {
+				return cache, nil
+			}
+			o, err := outputFn(s, paths)
+			if err != nil {
+				return "", err
+			}
+			if o != s {
+				// Output path converted to base64 or upload URL, cache it
+				pr.outputCache[s] = o
+			}
+			return o, nil
+		}
+
+		if output, err := handlePath(pr.response.Output, &paths, cachedOutputFn); err != nil {
 			log.Errorw("failed to handle output path", "id", pid, "error", err)
 			pr.response.Status = PredictionFailed
 			pr.response.Error = fmt.Sprintf("failed to handle output path: %s", err)

internal/tests/path_test.go

@@ -144,6 +144,41 @@ func TestPredictionPathUploadUrlSucceeded(t *testing.T) {
 	assert.NoError(t, ct.Cleanup())
 }
 
+func TestPredictionPathUploadIterator(t *testing.T) {
+	ct := NewCogTest(t, "path_out_iter")
+	ct.StartWebhook()
+	ct.AppendArgs(fmt.Sprintf("--upload-url=http://localhost:%d/upload/", ct.webhookPort))
+	assert.NoError(t, ct.Start())
+
+	hc := ct.WaitForSetup()
+	assert.Equal(t, server.StatusReady.String(), hc.Status)
+	assert.Equal(t, server.SetupSucceeded, hc.Setup.Status)
+
+	ct.AsyncPrediction(map[string]any{"n": 3})
+	wr := ct.WaitForWebhookCompletion()
+	ul := ct.GetUploads()
+
+	assert.Len(t, wr, 5)
+	assert.Equal(t, server.PredictionProcessing, wr[0].Status)
+	assert.Nil(t, wr[0].Output)
+	assert.Equal(t, server.PredictionProcessing, wr[1].Status)
+	assert.Len(t, wr[1].Output.([]any), 1)
+	assert.Equal(t, server.PredictionProcessing, wr[2].Status)
+	assert.Len(t, wr[2].Output.([]any), 2)
+	assert.Equal(t, server.PredictionProcessing, wr[3].Status)
+	assert.Len(t, wr[3].Output.([]any), 3)
+	assert.Equal(t, server.PredictionSucceeded, wr[4].Status)
+	assert.Len(t, wr[4].Output.([]any), 3)
+
+	assert.Len(t, ul, 3)
+	assert.Equal(t, "out0", string(ul[0].Body))
+	assert.Equal(t, "out1", string(ul[1].Body))
+	assert.Equal(t, "out2", string(ul[2].Body))
+
+	ct.Shutdown()
+	assert.NoError(t, ct.Cleanup())
+}
+
 const TestDataPrefix = "https://raw.githubusercontent.com/gabriel-vasile/mimetype/refs/heads/master/testdata/"
 
 func TestPredictionPathMimeTypes(t *testing.T) {

python/tests/runners/path_out_iter.py

@@ -0,0 +1,18 @@
+import tempfile
+import time
+from typing import Iterator
+
+from cog import BasePredictor, Path
+
+
+class Predictor(BasePredictor):
+    test_inputs = {'n': 2}
+
+    def predict(self, n: int) -> Iterator[Path]:
+        for i in range(n):
+            time.sleep(1)
+            with tempfile.NamedTemporaryFile(
+                mode='w', suffix='.txt', delete=False
+            ) as f:
+                f.write(f'out{i}')
+            yield Path(f.name)