Repository: replicate/cog

Issue:
Make it possible for queue worker to only send webhooks for new output
Currently the queue worker sends a webhook whenever there's new output **or** new logs. For some use cases, getting those logs immediately is great. For others, it's just noise. It'd be great to have a way to pick between those behaviours by passing a parameter with the queued request.


Discussion:

Resolving patch:
diff --git a/python/cog/response.py b/python/cog/response.py
--- a/python/cog/response.py
+++ b/python/cog/response.py
@@ -1,9 +1,41 @@
 import enum
-from typing import Optional, Type, Any
+from typing import Any, List, Optional, Type
 
 from pydantic import BaseModel, Field
 
 
+class Event(str, enum.Enum):
+    START = "start"
+    OUTPUT = "output"
+    LOGS = "logs"
+    COMPLETED = "completed"
+
+    @staticmethod
+    def validate(events: List[str]) -> List[str]:
+        for event in events:
+            assert event in {
+                Event.START,
+                Event.OUTPUT,
+                Event.LOGS,
+                Event.COMPLETED,
+            }, f"Unexpected event {event}. Must be from {Event.START}, {Event.OUTPUT}, {Event.LOGS}, {Event.COMPLETED}"
+
+        # Always include COMPLETED events
+        if Event.COMPLETED not in events:
+            events.append(Event.COMPLETED)
+
+        return events
+
+    @staticmethod
+    def default_events() -> List[str]:
+        return [
+            Event.START,
+            Event.OUTPUT,
+            Event.LOGS,
+            Event.COMPLETED,
+        ]
+
+
 class Status(str, enum.Enum):
     PROCESSING = "processing"
     SUCCEEDED = "succeeded"
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
@@ -33,7 +33,7 @@
     load_config,
     load_predictor_from_ref,
 )
-from ..response import Status
+from ..response import Event, Status
 from .eventtypes import Done, Heartbeat, Log, PredictionOutput, PredictionOutputType
 from .probes import ProbeHelper
 from .webhook import requests_session, webhook_caller
@@ -226,8 +226,17 @@ def start(self) -> None:
 
                     should_cancel = self.cancelation_checker(message.get("cancel_key"))
 
-                    for response in self.run_prediction(message, should_cancel):
-                        send_response(response)
+                    if "webhook_events_filter" in message:
+                        # this will raise an exception if the list isn't valid
+                        events_filter = Event.validate(message["webhook_events_filter"])
+                    else:
+                        events_filter = Event.default_events()
+
+                    for response_event, response in self.run_prediction(
+                        message, should_cancel
+                    ):
+                        if response_event in events_filter:
+                            send_response(response)
 
                     if self.max_failure_count is not None:
                         # Keep track of runs of failures to catch the situation
@@ -256,7 +265,7 @@ def start(self) -> None:
 
     def run_prediction(
         self, message: Dict[str, Any], should_cancel: Callable
-    ) -> Iterable[Dict[str, Any]]:
+    ) -> Iterable[Tuple[Event, Dict[str, Any]]]:
         # use the request message as the basis of our response so
         # that we echo back any additional fields sent to us
         response = message
@@ -271,7 +280,7 @@ def run_prediction(
         except Exception as e:
             response["status"] = Status.FAILED
             response["error"] = str(e)
-            yield response
+            yield (Event.COMPLETED, response)
 
             try:
                 input_obj.cleanup()
@@ -283,7 +292,7 @@ def run_prediction(
         response["started_at"] = format_datetime(started_at)
         response["logs"] = ""
 
-        yield response
+        yield (Event.START, response)
 
         timed_out = False
         was_canceled = False
@@ -311,7 +320,7 @@ def run_prediction(
                     pass
                 elif isinstance(event, Log):
                     response["logs"] += event.message
-                    yield response
+                    yield (Event.LOGS, response)
                 elif isinstance(event, PredictionOutputType):
                     # Note: this error message will be seen by users so it is
                     # intentionally vague about what has gone wrong.
@@ -330,7 +339,7 @@ def run_prediction(
 
                     if output_type.multi:
                         response["output"].append(output)
-                        yield response
+                        yield (Event.OUTPUT, response)
                     else:
                         assert (
                             response["output"] is None
@@ -369,7 +378,7 @@ def run_prediction(
             response["status"] = Status.FAILED
             response["error"] = str(e)
         finally:
-            yield response
+            yield (Event.COMPLETED, response)
 
             try:
                 input_obj.cleanup()


Test patch:
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
@@ -2084,6 +2084,80 @@ def capture_response(request):
         )
 
 
+def test_queue_worker_webhook_events_filter(
+    docker_network, docker_image, redis_client, httpserver
+):
+    project_dir = Path(__file__).parent / "fixtures/logging-project"
+    subprocess.run(["cog", "build", "-t", docker_image], check=True, cwd=project_dir)
+
+    with docker_run(
+        image=docker_image,
+        interactive=True,
+        network=docker_network,
+        command=[
+            "python",
+            "-m",
+            "cog.server.redis_queue",
+            "--redis-host=redis",
+            "--redis-port=6379",
+            "--input-queue=predict-queue",
+            "--consumer-id=test-worker",
+        ],
+        env=DEFAULT_ENV,
+    ):
+        predict_id = random_string(10)
+        webhook_url = httpserver.url_for("/webhook").replace(
+            "localhost", "host.docker.internal"
+        )
+
+        # We're only expecting a single webhook: after the prediction is done
+        httpserver.expect_oneshot_request(
+            "/webhook",
+            json={
+                "id": predict_id,
+                "input": {},
+                "webhook": webhook_url,
+                "webhook_events_filter": ["completed"],
+                "logs": (
+                    "WARNING:root:writing log message\n"
+                    + "writing from C\n"
+                    + "writing to stderr\n"
+                    + "writing with print\n"
+                ),
+                "output": "output",
+                "status": "succeeded",
+                "started_at": mock.ANY,
+                "completed_at": mock.ANY,
+                "metrics": {
+                    "predict_time": mock.ANY,
+                },
+            },
+            method="POST",
+        )
+
+        redis_client.xgroup_create(
+            mkstream=True, groupname="predict-queue", name="predict-queue", id="$"
+        )
+
+        with httpserver.wait(timeout=15) as waiting:
+            redis_client.xadd(
+                name="predict-queue",
+                fields={
+                    "value": json.dumps(
+                        {
+                            "id": predict_id,
+                            "input": {},
+                            "webhook": webhook_url,
+                            "webhook_events_filter": ["completed"],
+                        }
+                    ),
+                },
+            )
+
+        # check we received all the webhooks
+        assert waiting.result
+
+
 def response_iterator(redis_client, response_queue, timeout=10):
     redis_client.config_set("notify-keyspace-events", "KEA")
     channel = redis_client.pubsub()