Revert "Revert "Support async predictors (#2010)" (#2022)"

This reverts commit 8333a83.

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):