patch cancel bug: immediately mark predictions as cancelled

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