unconditionally mark predictions as cancelled without waiting for cancellation to succeed

while I struggle with #1786, we see queue spikes because not responding to cancellation promptly causes the pod to get restarted. this is a dirty hack to pretend like cancellation works immediately. as soon as we fix the race condition (and possibly any issues with task.cancel() behaving differently from signal handlers), we can drop this.

Signed-off-by: technillogue <technillogue@gmail.com>

Author: technillogue

python/cog/server/runner.py

@@ -122,7 +122,9 @@ def __init__(
         self._shutdown_event = shutdown_event  # __main__ waits for this event
 
         self._upload_url = upload_url
-        self._predictions: dict[str, tuple[schema.PredictionResponse, PredictionTask]] = {}
+        self._predictions: dict[
+            str, tuple[schema.PredictionResponse, PredictionTask]
+        ] = {}
         self._predictions_in_flight: set[str] = set()
         # it would be lovely to merge these but it's not fully clear how best to handle it
         # since idempotent requests can kinda come whenever?
@@ -402,6 +404,15 @@ def cancel(self, prediction_id: str) -> None:
             self._events.send(Cancel(prediction_id))
             # maybe this should probably check self._semaphore._value == self._concurrent
 
+            # HACK: sometimes, Done events may be dropped due to a race condition with logging
+            # (or possibly because asyncio cancellation is less strict than signal handlers)
+            # while we're fixing that, here's a bodge so that we always promptly respond
+            # to cancellation requests so that director doesn't kill us.
+            #
+            # if the real Done event comes through, it shouldn't cause any problems
+            # it will just stay in an ignored queue in the mux
+            asyncio.create_task(self._mux.write(prediction_id, Done(canceled=True)))
+
     _read_events_task: "Optional[asyncio.Task[None]]" = None
 
     def _start_event_reader(self) -> None: