maybe fix the race condition

Author: technillogue

python/cog/server/connection.py

@@ -25,11 +25,12 @@ async def async_init(self) -> None:
         if self.started:
             return
         fd = self.wrapped_conn.fileno()
-        # # mp may have handled something already but let's dup so exit is clean
+        # mp may have handled something already but let's dup so exit is clean
         dup_fd = os.dup(fd)
         sock = socket.socket(fileno=dup_fd)
         # sock = socket.socket(fileno=fd)
         # we don't want to see EAGAIN, we'd rather wait
+        # however, perhaps this is wrong and in some cases this could still block terribly
         # sock.setblocking(False)
         # TODO: use /proc/sys/net/core/rmem_max, but special-case language models
         sz = 65536

python/cog/server/helpers.py

@@ -114,6 +114,7 @@ def drain(self) -> None:
             stream.flush()
         debug("wait drain")
         if not self.drain_event.wait(timeout=1):
+            debug("drain timed out")
             raise RuntimeError("output streams failed to drain")
         debug("drain done")
 
@@ -154,12 +155,12 @@ def run(self) -> None:
                 stream = key.data
 
                 for line in stream.wrapped:
+                    debug("redirector saw", line)
                     if not line.endswith("\n"):
                         # TODO: limit how much we're prepared to buffer on a
                         # single line
                         buffers[stream.name].write(line)
                         continue
-                    debug("redirector saw", line)
 
                     full_line = buffers[stream.name].getvalue() + line.strip()
 
@@ -179,7 +180,9 @@ def run(self) -> None:
                     # thing in the line was a drain token (or a terminate
                     # token).
                     if full_line:
+                        debug("write hook")
                         self._write_hook(stream.name, stream.original, full_line + "\n")
+                        debug("write hook done")
 
                     if drain_tokens_seen >= drain_tokens_needed:
                         debug("drain event set")
@@ -202,14 +205,20 @@ async def switch_to_async(self) -> None:
         """
         debug("switch async, drain")
         # Drain the streams to ensure all buffered data is processed
-        self.drain()
+        try:
+            self.drain()
+        except RuntimeError:
+            debug("drain failed")
+            raise
+        debug("drain done, shutdown")
 
         # Shut down the thread
         # we do this before starting a coroutine that will also read from the same fd
         # so that shutdown can find the terminate tokens correctly
         self.shutdown()
         self.stream_tasks = []
         self.is_async = True
+        debug("set is async")
 
         for stream in self._streams:
             # Open each stream as a StreamReader
@@ -220,6 +229,9 @@ async def switch_to_async(self) -> None:
             task = asyncio.create_task(self.process_stream(stream, reader))
             self.stream_tasks.append(task)
 
+        # give the tasks a chance to start
+        await asyncio.sleep(0)
+
     async def process_stream(
         self, stream: WrappedStream, reader: asyncio.StreamReader
     ) -> None:

python/cog/server/worker.py

@@ -158,9 +158,12 @@ async def _async_init(self) -> None:
         if self._events_async:
             debug("async_init finished")
             return
+        # if AsyncConnection is created before switch_to_async, a race condition can cause drain to fail
+        # and write, seemingly, to block
+        # maybe because we're trying to call StreamWriter.write when no event loop is running?
+        await self._stream_redirector.switch_to_async()
         self._events_async = AsyncConnection(self._events)
         await self._events_async.async_init()
-        await self._stream_redirector.switch_to_async()
         debug("async_init done")
 
     def _setup(self) -> None:
@@ -189,15 +192,19 @@ def _setup(self) -> None:
                 debug("inspect")
                 if inspect.iscoroutinefunction(self._predictor.setup):
 
-                    async def setup_async():
+                    async def setup_async() -> None:
                         # we prefer to not stop-start the event loop between these calls
-                        await sef._async_init()  # this creates tasks
+                        await self._async_init()  # this creates tasks
+                        debug("async_init done")
                         await run_setup_async(self._predictor)
+                        debug("ran setup async")
 
                     self.loop.run_until_complete(setup_async())
+                    debug("setup_async loop done")
                 else:
                     debug("sync setup")
                     run_setup(self._predictor)
+            debug("_setup done inside ctx mgr")
         debug("_setup done")
 
     @contextlib.contextmanager
@@ -220,8 +227,15 @@ def _handle_setup_error(self) -> Iterator[None]:
             done.error_detail = str(e)
             raise
         finally:
+            # we can arrive here if there was an error setting up stream_redirector
+            # for example, because drain failed
+            # in this case this drain could block or fail
             debug("setup done, calling drain")
-            self._stream_redirector.drain()
+            try:
+                self._stream_redirector.drain()
+            except Exception as e:
+                debug("exc", e)
+                raise
             debug("sending setup done")
             self.send(("SETUP", done))
             debug("sent setup done")
@@ -242,6 +256,7 @@ def _loop_sync(self) -> None:
         self._stream_redirector.shutdown()
 
     async def _loop_async(self) -> None:
+        debug("loop async")
         await self._async_init()
         assert self._events_async
         tasks: dict[str, asyncio.Task[None]] = {}
@@ -307,10 +322,14 @@ def _emit_metric(self, name: str, value: "int | float") -> None:
 
     def send(self, obj: Any) -> None:
         if self._events_async:
+            debug("sending on async")
             self._events_async.send(obj)
+            debug("sent on async")
         else:
+            debug("send lock")
             with self._sync_events_lock:
                 self._events.send(obj)
+                debug("finished sync send")
 
     def _mk_send(self, id: str) -> Callable[[PublicEventType], None]:
         def send(event: PublicEventType) -> None: