Support async predictors (#2010)

* Permit async predictors

First steps towards allowing `async def predict` method signatures. This
commit adds support to Worker for starting an `asyncio` event loop if
the `predict` function returns an awaitable or an async generator.

For now, we drop support for output capture as well as cancelation.

* Async-compatible stream interception

In an async context, attempting to intercept stream writes at the file
descriptor layer is futile. We can do it, but we will have no way of
associating a write made from native code with a specific prediction --
and the only reason to intercept/swap out the STDOUT/STDERR file
descriptors is so that we can catch writes from native code.

This commit adds an altogether simpler implementation which can work for
async code with that restriction. All it does is patch `sys.stdout` and
`sys.stderr` with objects that can redirect (or tee) the output to a
callback function.

* Implement single-task cancelation for async predictors

This implements basic cancelation for async predictors. Whereas regular
predictors implement cancelation using a custom CancelationException,
asyncio already has a concept of task cancelation, so we use that.

When cancelation is requested, we send a `Cancel()` event down the
events pipe to the child. Regular predictors ignore these, but async
predictors cancel the currently-running task when they receive one.

In future, these `Cancel()` events will specify which running prediction
they are intended to cancel.

* Ensure that graceful shutdown works as expected for async predictors

When a `Shutdown()` event is sent, any running prediction should be
allowed to completed.

For now, we implement this by awaiting any task that is tracked when we
break out of the child worker's event loop.

* Update support-async-predictors branch for Pydantic v2

- Use renamed _ChildWorker type
- Set initial `__url__` to `None` prior to URL parsing potentially
  throwing `ValueError`
- Declare `pid` field in `FakeChildWorker`
- Do not use nested redirectors

---------

Co-authored-by: Nick Stenning <[email protected]>
Co-authored-by: Dominic Baggott <[email protected]>

python/cog/server/connection.py

@@ -0,0 +1,91 @@
+import asyncio
+import multiprocessing
+from multiprocessing.connection import Connection
+from typing import Any, Optional
+
+from typing_extensions import Buffer
+
+_spawn = multiprocessing.get_context("spawn")
+
+
+class AsyncConnection:
+    def __init__(self, connection: Connection) -> None:
+        self._connection = connection
+        self._event = asyncio.Event()
+        loop = asyncio.get_event_loop()
+        loop.add_reader(self._connection.fileno(), self._event.set)
+
+    def send(self, obj: Any) -> None:
+        """Send a (picklable) object"""
+
+        self._connection.send(obj)
+
+    async def _wait_for_input(self) -> None:
+        """Wait until there is an input available to be read"""
+
+        while not self._connection.poll():
+            await self._event.wait()
+            self._event.clear()
+
+    async def recv(self) -> Any:
+        """Receive a (picklable) object"""
+
+        await self._wait_for_input()
+        return self._connection.recv()
+
+    def fileno(self) -> int:
+        """File descriptor or handle of the connection"""
+        return self._connection.fileno()
+
+    def close(self) -> None:
+        """Close the connection"""
+        self._connection.close()
+
+    async def poll(self, timeout: float = 0.0) -> bool:
+        """Whether there is an input available to be read"""
+
+        if self._connection.poll():
+            return True
+
+        try:
+            await asyncio.wait_for(self._wait_for_input(), timeout=timeout)
+        except asyncio.TimeoutError:
+            return False
+        return self._connection.poll()
+
+    def send_bytes(
+        self, buf: Buffer, offset: int = 0, size: Optional[int] = None
+    ) -> None:
+        """Send the bytes data from a bytes-like object"""
+
+        self._connection.send_bytes(buf, offset, size)
+
+    async def recv_bytes(self, maxlength: Optional[int] = None) -> bytes:
+        """
+        Receive bytes data as a bytes object.
+        """
+
+        await self._wait_for_input()
+        return self._connection.recv_bytes(maxlength)
+
+    async def recv_bytes_into(self, buf: Buffer, offset: int = 0) -> int:
+        """
+        Receive bytes data into a writeable bytes-like object.
+        Return the number of bytes read.
+        """
+
+        await self._wait_for_input()
+        return self._connection.recv_bytes_into(buf, offset)
+
+
+class LockedConnection:
+    def __init__(self, connection: Connection) -> None:
+        self.connection = connection
+        self._lock = _spawn.Lock()
+
+    def send(self, obj: Any) -> None:
+        with self._lock:
+            self.connection.send(obj)
+
+    def recv(self) -> Any:
+        return self.connection.recv()

python/cog/server/eventtypes.py

@@ -5,6 +5,12 @@
 
 # From worker parent process
 #
+@define
+class Cancel:
+    # TODO: identify which prediction!
+    pass
+
+
 @define
 class PredictionInput:
     payload: Dict[str, Any]

python/cog/server/helpers.py

@@ -10,7 +10,7 @@
 import threading
 import uuid
 from types import TracebackType
-from typing import Any, Callable, Dict, List, Sequence, TextIO, Union
+from typing import Any, BinaryIO, Callable, Dict, List, Sequence, TextIO, Union
 
 import pydantic
 from typing_extensions import Self
@@ -19,6 +19,45 @@
 from .errors import CogRuntimeError, CogTimeoutError
 
 
+class _SimpleStreamWrapper(io.TextIOWrapper):
+    """
+    _SimpleStreamWrapper wraps a binary I/O buffer and provides a TextIOWrapper
+    interface (primarily write and flush methods) which call a provided
+    callback function instead of (or, if `tee` is True, in addition to) writing
+    to the underlying buffer.
+    """
+
+    def __init__(
+        self,
+        buffer: BinaryIO,
+        callback: Callable[[str, str], None],
+        tee: bool = False,
+    ) -> None:
+        super().__init__(buffer, line_buffering=True)
+
+        self._callback = callback
+        self._tee = tee
+        self._buffer = []
+
+    def write(self, s: str) -> int:
+        length = len(s)
+        self._buffer.append(s)
+        if self._tee:
+            super().write(s)
+        else:
+            # If we're not teeing, we have to handle automatic flush on
+            # newline. When `tee` is true, this is handled by the write method.
+            if "\n" in s or "\r" in s:
+                self.flush()
+        return length
+
+    def flush(self) -> None:
+        self._callback(self.name, "".join(self._buffer))
+        self._buffer.clear()
+        if self._tee:
+            super().flush()
+
+
 class _StreamWrapper:
     def __init__(self, name: str, stream: TextIO) -> None:
         self.name = name
@@ -86,6 +125,66 @@ def original(self) -> TextIO:
         return self._original_fp
 
 
+if sys.version_info < (3, 9):
+
+    class _AsyncStreamRedirectorBase(contextlib.AbstractContextManager):
+        pass
+else:
+
+    class _AsyncStreamRedirectorBase(
+        contextlib.AbstractContextManager["AsyncStreamRedirector"]
+    ):
+        pass
+
+
+class AsyncStreamRedirector(_AsyncStreamRedirectorBase):
+    """
+    AsyncStreamRedirector is a context manager that redirects I/O streams to a
+    callback function. If `tee` is True, it also writes output to the original
+    streams.
+
+    Unlike StreamRedirector, the underlying stream file descriptors are not
+    modified, which means that only stream writes from Python code will be
+    captured. Writes from native code will not be captured.
+
+    Unlike StreamRedirector, the streams redirected cannot be configured. The
+    context manager is only able to redirect STDOUT and STDERR.
+    """
+
+    def __init__(
+        self,
+        callback: Callable[[str, str], None],
+        tee: bool = False,
+    ) -> None:
+        self._callback = callback
+        self._tee = tee
+
+        stdout_wrapper = _SimpleStreamWrapper(sys.stdout.buffer, callback, tee)
+        stderr_wrapper = _SimpleStreamWrapper(sys.stderr.buffer, callback, tee)
+        self._stdout_ctx = contextlib.redirect_stdout(stdout_wrapper)
+        self._stderr_ctx = contextlib.redirect_stderr(stderr_wrapper)
+
+    def __enter__(self) -> Self:
+        self._stdout_ctx.__enter__()
+        self._stderr_ctx.__enter__()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self._stdout_ctx.__exit__(exc_type, exc_value, traceback)
+        self._stderr_ctx.__exit__(exc_type, exc_value, traceback)
+
+    def drain(self, timeout: float = 0.0) -> None:
+        # Draining isn't complicated for AsyncStreamRedirector, since we're not
+        # moving data between threads. We just need to flush the streams.
+        sys.stdout.flush()
+        sys.stderr.flush()
+
+
 if sys.version_info < (3, 9):
 
     class _StreamRedirectorBase(contextlib.AbstractContextManager):