Repository: replicate/cog

Issue:
Handle cancellation in Redis queue worker
Replicate currently kills the whole pod when the user cancels a prediction. Instead, we should be able to handle cancellations inside the running process.

One solution: Create a separate per-prediction cancellation key in Redis, and listen for keyspace notifications inside the Redis queue worker. When a cancellation occurs, cancel the running prediction and send a terminal ~~"cancelled"~~ "canceled" status to the output queue/webhook.
* As a first step we can kill the prediction subprocess which will re-run `setup()`.
* A refinement of that would be to listen for a signal inside the `predict()` function and only kill the current prediction (without re-running `setup()`).


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
@@ -8,6 +8,7 @@ class Status(str, enum.Enum):
     PROCESSING = "processing"
     SUCCEEDED = "succeeded"
     FAILED = "failed"
+    CANCELED = "canceled"
 
 
 def get_response_type(OutputType: Type[BaseModel]) -> Any:
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
@@ -176,6 +176,8 @@ def start(self) -> None:
                         f"Received message {message_id} on {self.input_queue}\n"
                     )
 
+                    should_cancel = self.cancelation_checker(message.get("cancel_key"))
+
                     # use the request message as the basis of our response so
                     # that we echo back any additional fields sent to us
                     response = message
@@ -186,7 +188,9 @@ def start(self) -> None:
                     cleanup_functions: List[Callable] = []
                     try:
                         start_time = time.time()
-                        self.handle_message(send_response, response, cleanup_functions)
+                        self.handle_message(
+                            send_response, response, cleanup_functions, should_cancel
+                        )
                         self.redis.xack(self.input_queue, self.input_queue, message_id)
                         self.redis.xdel(
                             self.input_queue, message_id
@@ -225,6 +229,7 @@ def handle_message(
         send_response: Callable,
         response: Dict[str, Any],
         cleanup_functions: List[Callable],
+        should_cancel: Callable,
     ) -> None:
         span = trace.get_current_span()
 
@@ -254,6 +259,13 @@ def handle_message(
 
         # just send logs until output starts
         while self.runner.is_processing() and not self.runner.has_output_waiting():
+            if should_cancel():
+                self.runner.cancel()
+                response["status"] = Status.CANCELED
+                response["completed_at"] = format_datetime(datetime.datetime.now())
+                send_response(response)
+                return
+
             if self.runner.has_logs_waiting():
                 response["logs"] += self.runner.read_logs()
                 send_response(response)
@@ -274,6 +286,13 @@ def handle_message(
 
             while self.runner.is_processing():
                 # TODO: restructure this to avoid the tight CPU-eating loop
+                if should_cancel():
+                    self.runner.cancel()
+                    response["status"] = Status.CANCELED
+                    response["completed_at"] = format_datetime(datetime.datetime.now())
+                    send_response(response)
+                    return
+
                 if self.runner.has_output_waiting() or self.runner.has_logs_waiting():
                     # Object has already passed through `make_encodeable()` in the Runner, so all we need to do here is upload the files
                     new_output = [
@@ -318,6 +337,13 @@ def handle_message(
         else:
             # just send logs until output ends
             while self.runner.is_processing():
+                if should_cancel():
+                    self.runner.cancel()
+                    response["status"] = Status.CANCELED
+                    response["completed_at"] = format_datetime(datetime.datetime.now())
+                    send_response(response)
+                    return
+
                 if self.runner.has_logs_waiting():
                     response["logs"] += self.runner.read_logs()
                     send_response(response)
@@ -361,6 +387,12 @@ def setter(response: Any) -> None:
 
         return setter
 
+    def cancelation_checker(self, redis_key: str) -> Callable:
+        def checker() -> bool:
+            return redis_key is not None and self.redis.exists(redis_key) > 0
+
+        return checker
+
     def upload_files(self, obj: Any) -> Any:
         def upload_file(fh: io.IOBase) -> str:
             filename = guess_filename(fh)
diff --git a/python/cog/server/runner.py b/python/cog/server/runner.py
--- a/python/cog/server/runner.py
+++ b/python/cog/server/runner.py
@@ -52,6 +52,10 @@ def __exit__(self, type: Any, value: Any, traceback: Any) -> None:
             signal.alarm(0)
 
 
+class CancelPredictionException(Exception):
+    pass
+
+
 class PredictionRunner:
     PROCESSING_DONE = 1
     EXIT_SENTINEL = "exit"
@@ -254,6 +258,11 @@ def _run_prediction(
         drain_pipe(self.error_pipe_reader)
         drain_pipe(self.done_pipe_reader)
 
+        def cancel(_signum: Any, _frame: Any) -> None:
+            raise CancelPredictionException()
+
+        signal.signal(signal.SIGUSR1, cancel)
+
         with capture_log(self.logs_pipe_writer):
             tracer = trace.get_tracer("cog")
             with tracer.start_as_current_span(
@@ -276,6 +285,9 @@ def _run_prediction(
                         else:
                             self.predictor_pipe_writer.send(self.OutputType.SINGLE)
                             self.predictor_pipe_writer.send(make_encodeable(output))
+                except CancelPredictionException:
+                    # we've been canceled, just stop and wait for cleanup
+                    pass
                 except Exception as e:
                     # if it timed out there's no stack trace
                     if type(e) != TimeoutError:
@@ -304,6 +316,12 @@ def close(self) -> None:
         self.prediction_input_pipe_writer.send(PredictionRunner.EXIT_SENTINEL)
         self.predictor_process.join()
 
+    def cancel(self) -> None:
+        """
+        Cancel the active prediction.
+        """
+        os.kill(self.predictor_process.pid, signal.SIGUSR1)  # type: ignore
+
 
 def drain_pipe(pipe_reader: Connection) -> None:
     """


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
@@ -189,7 +189,7 @@ def predict(
                 },
                 "Status": {
                     "title": "Status",
-                    "enum": ["processing", "succeeded", "failed"],
+                    "enum": ["processing", "succeeded", "failed", "canceled"],
                     "description": "An enumeration.",
                     "type": "string",
                 },
@@ -348,7 +348,7 @@ def predict(
                 },
                 "Status": {
                     "title": "Status",
-                    "enum": ["processing", "succeeded", "failed"],
+                    "enum": ["processing", "succeeded", "failed", "canceled"],
                     "description": "An enumeration.",
                     "type": "string",
                 },
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
@@ -1648,6 +1648,98 @@ def test_queue_worker_redis_responses(docker_network, docker_image, redis_client
         }
 
 
+def test_queue_worker_cancel(docker_network, docker_image, redis_client, httpserver):
+    project_dir = Path(__file__).parent / "fixtures/timeout-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",
+            "6379",
+            "predict-queue",
+            "",
+            "test-worker",
+            "model_id",
+            "logs",
+        ],
+    ):
+        predict_id = random_string(10)
+        webhook_url = httpserver.url_for("/webhook").replace(
+            "localhost", "host.docker.internal"
+        )
+
+        httpserver.expect_oneshot_request(
+            "/webhook",
+            json={
+                "id": predict_id,
+                "input": {
+                    "sleep_time": 30,
+                },
+                "webhook": webhook_url,
+                "logs": "",
+                "output": None,
+                "status": "processing",
+                "started_at": mock.ANY,
+                "cancel_key": "cancel-key",
+            },
+            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": {
+                                "sleep_time": 30,
+                            },
+                            "webhook": webhook_url,
+                            "cancel_key": "cancel-key",
+                        }
+                    ),
+                },
+            )
+
+        # check we receive the initial webhook
+        assert waiting.result
+
+        httpserver.expect_oneshot_request(
+            "/webhook",
+            json={
+                "id": predict_id,
+                "input": {
+                    "sleep_time": 30,
+                },
+                "webhook": webhook_url,
+                "logs": "",
+                "output": None,
+                "status": "canceled",
+                "started_at": mock.ANY,
+                "completed_at": mock.ANY,
+                "cancel_key": "cancel-key",
+            },
+            method="POST",
+        )
+
+        with httpserver.wait(timeout=5) as waiting:
+            redis_client.set("cancel-key", 1, ex=5)
+
+        # check we receive the "canceled" webhook
+        assert waiting.result
+
+
 def response_iterator(redis_client, response_queue, timeout=10):
     redis_client.config_set("notify-keyspace-events", "KEA")
     channel = redis_client.pubsub()