patch cancel bug: immediately mark predictions as cancelled (#1798)

* 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>

* make Mux.write a sync method

this makes the cancel patch cleaner, but might be a small speedup for high-throughput outputs

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

* gate this behind a feature flag

---------

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

Author: technillogue

python/cog/server/runner.py

@@ -404,6 +404,16 @@ 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
+            if os.getenv("COG_FAKE_CANCEL"):
+                self._mux.write(prediction_id, Done(canceled=True))
+
     _read_events_task: "Optional[asyncio.Task[None]]" = None
 
     def _start_event_reader(self) -> None:
@@ -427,7 +437,7 @@ async def _read_events(self) -> None:
                 id = "SETUP"
             if id == "LOG" and len(self._predictions_in_flight) == 1:
                 id = list(self._predictions_in_flight)[0]
-            await self._mux.write(id, event)
+            self._mux.write(id, event)
         # If we dropped off the end off the end of the loop, check if it's
         # because the child process died.
         if not self._child.is_alive() and not self._terminating.is_set():

python/cog/server/worker.py

@@ -60,8 +60,8 @@ def __init__(self, terminating: asyncio.Event) -> None:
         self.terminating = terminating
         self.fatal: Optional[FatalWorkerException] = None
 
-    async def write(self, id: str, item: PublicEventType) -> None:
-        await self.outs[id].put(item)
+    def write(self, id: str, item: PublicEventType) -> None:
+        self.outs[id].put_nowait(item)
 
     async def read(
         self, id: str, poll: Optional[float] = None