Repository: replicate/cog

Issue:
exposing timing metrics
It's useful to know how long a prediction takes to run. Let's update to Cog to expose metrics about prediction time.


Discussion:

Resolving patch:
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
@@ -1,3 +1,4 @@
+import datetime
 import io
 import json
 import os
@@ -191,10 +192,18 @@ def start(self) -> None:
                     sys.stderr.write(
                         f"Received message {message_id} on {self.input_queue}\n"
                     )
+                    # create this here so it's available during exception handling
+                    response: Dict[str, Any] = {
+                        "status": Status.PROCESSING,
+                        "output": None,
+                        "logs": [],
+                    }
                     cleanup_functions: List[Callable] = []
                     try:
                         start_time = time.time()
-                        self.handle_message(response_queue, message, cleanup_functions)
+                        self.handle_message(
+                            response_queue, response, message, cleanup_functions
+                        )
                         self.redis.xack(self.input_queue, self.input_queue, message_id)
                         self.redis.xdel(
                             self.input_queue, message_id
@@ -207,7 +216,12 @@ def start(self) -> None:
                         )
                         sys.stderr.write(f"Run time for {message_id}: {run_time:.2f}\n")
                     except Exception as e:
-                        self.push_error(response_queue, e)
+                        response["status"] = Status.FAILED
+                        response["error"] = str(e)
+                        response["x-experimental-timestamps"][
+                            "completed_at"
+                        ] = datetime.datetime.now().isoformat()
+                        self.push_message(response_queue, response)
                         self.redis.xack(self.input_queue, self.input_queue, message_id)
                         self.redis.xdel(self.input_queue, message_id)
                     finally:
@@ -226,6 +240,7 @@ def start(self) -> None:
     def handle_message(
         self,
         response_queue: str,
+        response: Dict[str, Any],
         message: Dict[str, Any],
         cleanup_functions: List[Callable],
     ) -> None:
@@ -236,7 +251,9 @@ def handle_message(
         except ValidationError as e:
             tb = traceback.format_exc()
             sys.stderr.write(tb)
-            self.push_error(response_queue, e)
+            response["status"] = Status.FAILED
+            response["error"] = str(e)
+            self.push_message(response_queue, response)
             span.record_exception(e)
             span.set_status(TraceStatus(status_code=StatusCode.ERROR))
             return
@@ -246,25 +263,28 @@ def handle_message(
         with timeout(seconds=self.predict_timeout):
             self.runner.run(**input_obj.dict())
 
-            logs: List[str] = []
-            response = {
-                "status": Status.PROCESSING,
-                "output": None,
-                "logs": logs,
+            response["x-experimental-timestamps"] = {
+                "started_at": datetime.datetime.now().isoformat()
             }
 
-            self.redis.rpush(response_queue, json.dumps(response))
+            logs: List[str] = []
+            response["logs"] = logs
+
+            self.push_message(response_queue, response)
 
             # just send logs until output starts
             while self.runner.is_processing() and not self.runner.has_output_waiting():
                 if self.runner.has_logs_waiting():
                     logs.extend(self.runner.read_logs())
-                    self.redis.rpush(response_queue, json.dumps(response))
+                    self.push_message(response_queue, response)
 
             if self.runner.error() is not None:
                 response["status"] = Status.FAILED
                 response["error"] = str(self.runner.error())  # type: ignore
-                self.redis.rpush(response_queue, json.dumps(response))
+                response["x-experimental-timestamps"][
+                    "completed_at"
+                ] = datetime.datetime.now().isoformat()
+                self.push_message(response_queue, response)
                 span.record_exception(self.runner.error())
                 span.set_status(TraceStatus(status_code=StatusCode.ERROR))
                 return
@@ -297,12 +317,15 @@ def handle_message(
                         # here to give the predictor subprocess a chance to exit
                         # so we don't send a double message for final output, at
                         # the cost of extra latency
-                        self.redis.rpush(response_queue, json.dumps(response))
+                        self.push_message(response_queue, response)
 
                 if self.runner.error() is not None:
                     response["status"] = Status.FAILED
                     response["error"] = str(self.runner.error())  # type: ignore
-                    self.redis.rpush(response_queue, json.dumps(response))
+                    response["x-experimental-timestamps"][
+                        "completed_at"
+                    ] = datetime.datetime.now().isoformat()
+                    self.push_message(response_queue, response)
                     span.record_exception(self.runner.error())
                     span.set_status(TraceStatus(status_code=StatusCode.ERROR))
                     return
@@ -310,21 +333,27 @@ def handle_message(
                 span.add_event("received final output")
 
                 response["status"] = Status.SUCCEEDED
+                response["x-experimental-timestamps"][
+                    "completed_at"
+                ] = datetime.datetime.now().isoformat()
                 output.extend(self.upload_files(o) for o in self.runner.read_output())
                 logs.extend(self.runner.read_logs())
-                self.redis.rpush(response_queue, json.dumps(response))
+                self.push_message(response_queue, response)
 
             else:
                 # just send logs until output ends
                 while self.runner.is_processing():
                     if self.runner.has_logs_waiting():
                         logs.extend(self.runner.read_logs())
-                        self.redis.rpush(response_queue, json.dumps(response))
+                        self.push_message(response_queue, response)
 
                 if self.runner.error() is not None:
                     response["status"] = Status.FAILED
                     response["error"] = str(self.runner.error())  # type: ignore
-                    self.redis.rpush(response_queue, json.dumps(response))
+                    response["x-experimental-timestamps"][
+                        "completed_at"
+                    ] = datetime.datetime.now().isoformat()
+                    self.push_message(response_queue, response)
                     span.record_exception(self.runner.error())
                     span.set_status(TraceStatus(status_code=StatusCode.ERROR))
                     return
@@ -333,24 +362,20 @@ def handle_message(
                 assert len(output) == 1
 
                 response["status"] = Status.SUCCEEDED
+                response["x-experimental-timestamps"][
+                    "completed_at"
+                ] = datetime.datetime.now().isoformat()
                 response["output"] = self.upload_files(output[0])
                 logs.extend(self.runner.read_logs())
-                self.redis.rpush(response_queue, json.dumps(response))
+                self.push_message(response_queue, response)
 
     def download(self, url: str) -> bytes:
         resp = requests.get(url)
         resp.raise_for_status()
         return resp.content
 
-    def push_error(self, response_queue: str, error: Any) -> None:
-        message = json.dumps(
-            {
-                "status": "failed",
-                "error": str(error),
-            }
-        )
-        sys.stderr.write(f"Pushing error to {response_queue}\n")
-        self.redis.rpush(response_queue, message)
+    def push_message(self, response_queue: str, response: Any) -> None:
+        self.redis.rpush(response_queue, json.dumps(response))
 
     def upload_files(self, obj: Any) -> Any:
         def upload_file(fh: io.IOBase) -> str:


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
@@ -1,6 +1,8 @@
 import json
 from pathlib import Path
+import re
 import subprocess
+import unittest.mock as mock
 
 from .util import docker_run, random_string
 
@@ -52,6 +54,9 @@ def test_queue_worker_files(docker_image, docker_network, redis_client, upload_s
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": None,
             "logs": [],
@@ -59,11 +64,24 @@ def test_queue_worker_files(docker_image, docker_network, redis_client, upload_s
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+                "completed_at": mock.ANY,
+            },
             "status": "succeeded",
             "output": "http://upload-server:5000/download/output.txt",
             "logs": [],
         }
 
+        assert re.match(
+            r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}.\d{6}",
+            response["x-experimental-timestamps"]["started_at"],
+        )
+        assert re.match(
+            r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}.\d{6}",
+            response["x-experimental-timestamps"]["completed_at"],
+        )
+
         with open(upload_server / "output.txt") as f:
             assert f.read() == "foobaztest"
 
@@ -116,6 +134,9 @@ def test_queue_worker_yielding_file(
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": None,
             "logs": [],
@@ -123,6 +144,9 @@ def test_queue_worker_yielding_file(
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": ["http://upload-server:5000/download/out-0.txt"],
             "logs": [],
@@ -133,6 +157,9 @@ def test_queue_worker_yielding_file(
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": [
                 "http://upload-server:5000/download/out-0.txt",
@@ -146,6 +173,9 @@ def test_queue_worker_yielding_file(
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": [
                 "http://upload-server:5000/download/out-0.txt",
@@ -160,6 +190,10 @@ def test_queue_worker_yielding_file(
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+                "completed_at": mock.ANY,
+            },
             "status": "succeeded",
             "output": [
                 "http://upload-server:5000/download/out-0.txt",
@@ -216,6 +250,9 @@ def test_queue_worker_yielding(docker_network, docker_image, redis_client):
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": None,
             "logs": [],
@@ -223,6 +260,9 @@ def test_queue_worker_yielding(docker_network, docker_image, redis_client):
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": ["foo", "bar", "baz"],
             "logs": [],
@@ -230,6 +270,10 @@ def test_queue_worker_yielding(docker_network, docker_image, redis_client):
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+                "completed_at": mock.ANY,
+            },
             "status": "succeeded",
             "output": ["foo", "bar", "baz"],
             "logs": [],
@@ -282,6 +326,9 @@ def test_queue_worker_error(docker_network, docker_image, redis_client):
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": None,
             "logs": [],
@@ -289,6 +336,10 @@ def test_queue_worker_error(docker_network, docker_image, redis_client):
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+                "completed_at": mock.ANY,
+            },
             "status": "failed",
             "output": None,
             "logs": [],
@@ -342,6 +393,9 @@ def test_queue_worker_error_after_output(docker_network, docker_image, redis_cli
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": None,
             "logs": [],
@@ -349,6 +403,9 @@ def test_queue_worker_error_after_output(docker_network, docker_image, redis_cli
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": ["hello bar"],
             "logs": [],
@@ -356,6 +413,9 @@ def test_queue_worker_error_after_output(docker_network, docker_image, redis_cli
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": ["hello bar"],
             "logs": ["a printed log message"],
@@ -363,6 +423,10 @@ def test_queue_worker_error_after_output(docker_network, docker_image, redis_cli
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+                "completed_at": mock.ANY,
+            },
             "status": "failed",
             "output": ["hello bar"],
             "logs": ["a printed log message"],
@@ -463,6 +527,9 @@ def test_queue_worker_logging(docker_network, docker_image, redis_client):
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": None,
             "logs": [],
@@ -470,6 +537,9 @@ def test_queue_worker_logging(docker_network, docker_image, redis_client):
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": None,
             "logs": [
@@ -479,6 +549,9 @@ def test_queue_worker_logging(docker_network, docker_image, redis_client):
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": None,
             "logs": [
@@ -489,6 +562,9 @@ def test_queue_worker_logging(docker_network, docker_image, redis_client):
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": None,
             "logs": [
@@ -500,6 +576,9 @@ def test_queue_worker_logging(docker_network, docker_image, redis_client):
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": None,
             "logs": [
@@ -512,6 +591,10 @@ def test_queue_worker_logging(docker_network, docker_image, redis_client):
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+                "completed_at": mock.ANY,
+            },
             "status": "succeeded",
             "output": "output",
             "logs": [
@@ -570,13 +653,24 @@ def test_queue_worker_timeout(docker_network, docker_image, redis_client):
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": None,
             "logs": [],
         }
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
-        assert response == {"status": "succeeded", "output": "it worked!", "logs": []}
+        assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+                "completed_at": mock.ANY,
+            },
+            "status": "succeeded",
+            "output": "it worked!",
+            "logs": [],
+        }
 
         predict_id = random_string(10)
         redis_client.xadd(
@@ -596,13 +690,25 @@ def test_queue_worker_timeout(docker_network, docker_image, redis_client):
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": None,
             "logs": [],
         }
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
-        assert response == {"status": "failed", "error": "Prediction timed out"}
+        assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+                "completed_at": mock.ANY,
+            },
+            "status": "failed",
+            "output": None,
+            "logs": [],
+            "error": "Prediction timed out",
+        }
 
 
 def test_queue_worker_yielding_timeout(docker_image, docker_network, redis_client):
@@ -650,16 +756,34 @@ def test_queue_worker_yielding_timeout(docker_image, docker_network, redis_clien
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": None,
             "logs": [],
         }
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
-        assert response == {"status": "processing", "output": ["yield 0"], "logs": []}
+        assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
+            "status": "processing",
+            "output": ["yield 0"],
+            "logs": [],
+        }
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
-        assert response == {"status": "succeeded", "output": ["yield 0"], "logs": []}
+        assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+                "completed_at": mock.ANY,
+            },
+            "status": "succeeded",
+            "output": ["yield 0"],
+            "logs": [],
+        }
 
         predict_id = random_string(10)
         redis_client.xadd(
@@ -680,6 +804,9 @@ def test_queue_worker_yielding_timeout(docker_image, docker_network, redis_clien
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": None,
             "logs": [],
@@ -687,17 +814,36 @@ 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.blpop("response-queue", timeout=10)[1])
-        assert response == {"status": "processing", "output": ["yield 0"], "logs": []}
+        assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
+            "status": "processing",
+            "output": ["yield 0"],
+            "logs": [],
+        }
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": ["yield 0", "yield 1"],
             "logs": [],
         }
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
-        assert response == {"status": "failed", "error": "Prediction timed out"}
+        assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+                "completed_at": mock.ANY,
+            },
+            "status": "failed",
+            "output": ["yield 0", "yield 1"],
+            "logs": [],
+            "error": "Prediction timed out",
+        }
 
 
 def test_queue_worker_complex_output(docker_network, docker_image, redis_client):
@@ -743,6 +889,9 @@ def test_queue_worker_complex_output(docker_network, docker_image, redis_client)
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": None,
             "logs": [],
@@ -750,6 +899,10 @@ def test_queue_worker_complex_output(docker_network, docker_image, redis_client)
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+                "completed_at": mock.ANY,
+            },
             "status": "succeeded",
             "output": {
                 "hello": "hello world",
@@ -811,6 +964,9 @@ def test_queue_worker_yielding_list_of_complex_output(
 
         response = json.loads(redis_client.brpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": None,
             "logs": [],
@@ -818,6 +974,9 @@ def test_queue_worker_yielding_list_of_complex_output(
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+            },
             "status": "processing",
             "output": [
                 [{"file": "http://upload-server:5000/download/file", "text": "hello"}]
@@ -827,6 +986,10 @@ def test_queue_worker_yielding_list_of_complex_output(
 
         response = json.loads(redis_client.blpop("response-queue", timeout=10)[1])
         assert response == {
+            "x-experimental-timestamps": {
+                "started_at": mock.ANY,
+                "completed_at": mock.ANY,
+            },
             "status": "succeeded",
             "output": [
                 [{"file": "http://upload-server:5000/download/file", "text": "hello"}]