Repository: replicate/cog

Issue:
Design proposal: Progressive output as a list
## Background

In-progress output can be returned by using `yield` in a `predict()` function. When the function yields, the output is put on the Redis queue to return to the caller while the `predict()` function is still running.

This is useful for generative models like [pixray](https://replicate.com/pixray/text2image) which take a while to process, and users want to see the output as it is running.

The way this currently works is with a single output value that changes over time.

## Problem

Previous output gets discarded, and it is hard to track when output changes.

Due to this, displaying progressive output in Replicate is a big 'ol hack. If the output changes while the model is still running, we save that to the database and loosely keep track of what outputs get produced (with the `File.output_for_prediction` field). It’s hard to dedupe files (we have another hacky field `File.serving_url` for that) and we might miss output if we don’t poll in time (though this will be less of of a problem if we push with webhooks).

## Design

We could model progressive output as a list. Instead of each output being considered a point-in-time value, it could change over time. For example, consider these outputs changing over time in the previous design:

- `1.jpg`
- `2.jpg`
- `3.jpg`

We’d instead model it like this:

- `["1.jpg"]`
- `["1.jpg", "2.jpg"]`
- `["1.jpg", "2.jpg", "3.jpg"]`

In other words, the progressive output is considered part of the primary output, and part of the primary data model, rather than it being an effect of that data changing over time.

Some neat properties of this:

- It’s always consistent, in case data gets missed if we don’t poll fast enough or webhooks get missed
- It’s saved alongside the prediction, so wherever we display predictions the progressive output data is right there
- It kinda makes semantic sense. From a Python programmer’s point of view, a function that returns an iterator is returning a list that changes.
- For a backwards compatibility shim we're making for Replicate, we can keep track of the output changing. So instead of having to have the `serving_url` hack we can just be like `if len(new_output) > len(existing_output)`

[This is reminiscent of whether to send snapshots or incremental updates for logs](https://github.com/replicate/cog/issues/438). This would be inefficient if there are large numbers of outputs and we don’t care about the previous outputs.

Perhaps a difference with output is that progressive output is likely to be small amounts of data. And, it is primary output we care about, not just metadata.


## Implementation

- In the schema for Pydantic models, generators would be considered an `array`
- The schema for old models for backwards compatibility would *always* be `[{text: ..., file: ...}]`.
    - Most models will just return list of length 1.
    - We could be smart and figure out whether the function is a generator, but we’d then have to generate a different output schema for each model somehow. Currently, all old models have the same schema which makes the data migration really easy.
- Pydantic Cog models and the backwards compatibility shim would keep track of values that have been yielded and send a progressively growing list of outputs over the Redis API
- On Replicate, we’d just blindly store the output in `Prediction.output`. We would display outputs of list length 1 as if it was just a normal output. If the length of the list is >1, then we display a slider to skip through the outputs.


Discussion:

Resolving patch:
diff --git a/python/cog/predictor.py b/python/cog/predictor.py
--- a/python/cog/predictor.py
+++ b/python/cog/predictor.py
@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from collections.abc import Generator
+from collections.abc import Iterator
 import enum
 import importlib
 import inspect
@@ -7,9 +7,10 @@
 from pathlib import Path
 from pydantic import create_model, BaseModel
 from pydantic.fields import FieldInfo
+from typing import List
 
 # Added in Python 3.8. Can be from typing if we drop support for <3.8.
-from typing_extensions import Literal, get_origin, get_args
+from typing_extensions import get_origin, get_args
 import yaml
 
 from .errors import ConfigDoesNotExist, PredictorNotSet
@@ -174,11 +175,11 @@ def predict(
     else:
         OutputType = signature.return_annotation
 
-    # The type that goes in the response is the type that is yielded
-    if get_origin(OutputType) is Generator:
-        OutputType = get_args(OutputType)[0]
+    # The type that goes in the response is a list of the yielded type
+    if get_origin(OutputType) is Iterator:
+        OutputType = List[get_args(OutputType)[0]]
 
-    if OutputType.__name__ != "Output":
+    if not hasattr(OutputType, "__name__") or OutputType.__name__ != "Output":
         # Wrap the type in a model called "Output" so it is a consistent name in the OpenAPI schema
         class Output(BaseModel):
             __root__: OutputType
diff --git a/python/cog/server/http.py b/python/cog/server/http.py
--- a/python/cog/server/http.py
+++ b/python/cog/server/http.py
@@ -65,15 +65,6 @@ def predict(request: Request = Body(default=None)):
         if request:
             output_file_prefix = request.output_file_prefix
 
-        # loop over generator function to get the last result
-        if isinstance(output, types.GeneratorType):
-            last_result = None
-            for iteration in enumerate(output):
-                last_result = iteration
-                # TODO: clean up output files
-            # last result is a tuple with (index, value)
-            output = last_result[1]
-
         OutputType = get_output_type(predictor)
         Response = get_response_type(OutputType)
 
diff --git a/python/cog/server/redis_queue.py b/python/cog/server/redis_queue.py
--- a/python/cog/server/redis_queue.py
+++ b/python/cog/server/redis_queue.py
@@ -220,7 +220,7 @@ def handle_message(self, response_queue, message, cleanup_functions):
             finally:
                 input_obj.cleanup()
         if isinstance(return_value, types.GeneratorType):
-            last_result = None
+            output = []
 
             while True:
                 # we consume iterator manually to capture log
@@ -231,18 +231,20 @@ def handle_message(self, response_queue, message, cleanup_functions):
                         result = next(return_value)
                 except StopIteration:
                     break
+                # Encode as JSON here instead of each time we push a result so the file URLs remain consistent
+                result = self.encode_json(result)
+                output.append(result)
                 # push the previous result, so we can eventually detect the last iteration
-                if last_result is not None:
-                    self.push_result(response_queue, last_result, status="processing")
+                self.push_result(response_queue, output, status="processing")
                 if isinstance(result, Path):
                     cleanup_functions.append(result.unlink)
-                last_result = result
 
             # push the last result
-            self.push_result(response_queue, last_result, status="success")
+            self.push_result(response_queue, output, status="success")
         else:
             if isinstance(return_value, Path):
                 cleanup_functions.append(return_value.unlink)
+            return_value = self.encode_json(return_value)
             self.push_result(response_queue, return_value, status="success")
 
     def download(self, url):
@@ -262,7 +264,7 @@ def push_error(self, response_queue, error):
 
     def push_result(self, response_queue, result, status):
         message = {
-            "value": self.encode_json(result),
+            "value": result,
         }
 
         message["status"] = status


Test patch:
diff --git a/python/tests/server/test_http.py b/python/tests/server/test_http.py
--- a/python/tests/server/test_http.py
+++ b/python/tests/server/test_http.py
@@ -2,7 +2,7 @@
 import io
 import os
 import tempfile
-from typing import Generator
+from typing import Iterator
 from unittest import mock
 
 from fastapi.testclient import TestClient
@@ -382,9 +382,29 @@ def predict(
     }
 
 
+def test_openapi_specification_with_yield():
+    class Predictor(BasePredictor):
+        def predict(
+            self,
+        ) -> Iterator[str]:
+            pass
+
+    client = make_client(Predictor())
+    resp = client.get("/openapi.json")
+    assert resp.status_code == 200
+
+    assert resp.json()["components"]["schemas"]["Output"] == {
+        "title": "Output",
+        "type": "array",
+        "items": {
+            "type": "string",
+        },
+    }
+
+
 def test_yielding_strings_from_generator_predictors():
     class Predictor(BasePredictor):
-        def predict(self) -> Generator[str, None, None]:
+        def predict(self) -> Iterator[str]:
             predictions = ["foo", "bar", "baz"]
             for prediction in predictions:
                 yield prediction
@@ -392,12 +412,12 @@ def predict(self) -> Generator[str, None, None]:
     client = make_client(Predictor())
     resp = client.post("/predictions")
     assert resp.status_code == 200
-    assert resp.json() == {"status": "success", "output": "baz"}
+    assert resp.json() == {"status": "success", "output": ["foo", "bar", "baz"]}
 
 
 def test_yielding_files_from_generator_predictors():
     class Predictor(BasePredictor):
-        def predict(self) -> Generator[Path, None, None]:
+        def predict(self) -> Iterator[Path]:
             colors = ["red", "blue", "yellow"]
             for i, color in enumerate(colors):
                 temp_dir = tempfile.mkdtemp()
@@ -410,10 +430,16 @@ def predict(self) -> Generator[Path, None, None]:
     resp = client.post("/predictions")
 
     assert resp.status_code == 200
-    header, b64data = resp.json()["output"].split(",", 1)
-    image = Image.open(io.BytesIO(base64.b64decode(b64data)))
-    image_color = Image.Image.getcolors(image)[0][1]
-    assert image_color == (255, 255, 0)  # yellow
+    output = resp.json()["output"]
+
+    def image_color(data_url):
+        header, b64data = data_url.split(",", 1)
+        image = Image.open(io.BytesIO(base64.b64decode(b64data)))
+        return Image.Image.getcolors(image)[0][1]
+
+    assert image_color(output[0]) == (255, 0, 0)  # red
+    assert image_color(output[1]) == (0, 0, 255)  # blue
+    assert image_color(output[2]) == (255, 255, 0)  # yellow
 
 
 # TODO: timing
diff --git a/test-integration/test_integration/fixtures/yielding-project/predict.py b/test-integration/test_integration/fixtures/yielding-project/predict.py
--- a/test-integration/test_integration/fixtures/yielding-project/predict.py
+++ b/test-integration/test_integration/fixtures/yielding-project/predict.py
@@ -1,10 +1,10 @@
-from typing import Generator
+from typing import Iterator
 
 from cog import BasePredictor
 
 
 class Predictor(BasePredictor):
-    def predict(self, text: str) -> Generator[str, None, None]:
+    def predict(self, text: str) -> Iterator[str]:
         predictions = ["foo", "bar", "baz"]
         for prediction in predictions:
             yield prediction
diff --git a/test-integration/test_integration/fixtures/yielding-timeout-project/predict.py b/test-integration/test_integration/fixtures/yielding-timeout-project/predict.py
--- a/test-integration/test_integration/fixtures/yielding-timeout-project/predict.py
+++ b/test-integration/test_integration/fixtures/yielding-timeout-project/predict.py
@@ -1,13 +1,11 @@
 import time
-from typing import Generator
+from typing import Iterator
 
 from cog import BasePredictor
 
 
 class Predictor(BasePredictor):
-    def predict(
-        self, sleep_time: float, n_iterations: int
-    ) -> Generator[str, None, None]:
+    def predict(self, sleep_time: float, n_iterations: int) -> Iterator[str]:
         for i in range(n_iterations):
             time.sleep(sleep_time)
             yield f"yield {i}"
diff --git a/test-integration/test_integration/test_redis_queue.py b/test-integration/test_integration/test_redis_queue.py
--- a/test-integration/test_integration/test_redis_queue.py
+++ b/test-integration/test_integration/test_redis_queue.py
@@ -106,13 +106,16 @@ def test_queue_worker_yielding(docker_network, docker_image, redis_client):
         )
 
         response = json.loads(redis_client.brpop("response-queue", timeout=10)[1])
-        assert response == {"value": "foo", "status": "processing"}
+        assert response == {"value": ["foo"], "status": "processing"}
 
         response = json.loads(redis_client.brpop("response-queue", timeout=10)[1])
-        assert response == {"value": "bar", "status": "processing"}
+        assert response == {"value": ["foo", "bar"], "status": "processing"}
 
         response = json.loads(redis_client.brpop("response-queue", timeout=10)[1])
-        assert response == {"value": "baz", "status": "success"}
+        assert response == {"value": ["foo", "bar", "baz"], "status": "processing"}
+
+        response = json.loads(redis_client.brpop("response-queue", timeout=10)[1])
+        assert response == {"value": ["foo", "bar", "baz"], "status": "success"}
 
         response = redis_client.rpop("response-queue")
         assert response == None
@@ -388,7 +391,10 @@ def test_queue_worker_yielding_timeout(docker_image, docker_network, redis_clien
         )
 
         response = json.loads(redis_client.brpop("response-queue", timeout=10)[1])
-        assert response == {"status": "success", "value": "yield 0"}
+        assert response == {"status": "processing", "value": ["yield 0"]}
+
+        response = json.loads(redis_client.brpop("response-queue", timeout=10)[1])
+        assert response == {"status": "success", "value": ["yield 0"]}
 
         predict_id = random_string(10)
         redis_client.xadd(
@@ -409,7 +415,10 @@ def test_queue_worker_yielding_timeout(docker_image, docker_network, redis_clien
 
         # TODO(andreas): revisit this test design if it starts being flakey
         response = json.loads(redis_client.brpop("response-queue", timeout=10)[1])
-        assert response == {"value": "yield 0", "status": "processing"}
+        assert response == {"value": ["yield 0"], "status": "processing"}
+
+        response = json.loads(redis_client.brpop("response-queue", timeout=10)[1])
+        assert response == {"value": ["yield 0", "yield 1"], "status": "processing"}
 
         response = json.loads(redis_client.brpop("response-queue", timeout=10)[1])
         assert response == {"status": "failed", "error": "Prediction timed out"}