modal 0.62.115__py3-none-any.whl → 0.72.13__py3-none-any.whl

modal/io_streams.py

@@ -0,0 +1,434 @@
+# Copyright Modal Labs 2022
+import asyncio
+from collections.abc import AsyncGenerator, AsyncIterator
+from typing import (
+    TYPE_CHECKING,
+    Generic,
+    Literal,
+    Optional,
+    TypeVar,
+    Union,
+    cast,
+)
+
+from grpclib import Status
+from grpclib.exceptions import GRPCError, StreamTerminatedError
+
+from modal.exception import ClientClosed, InvalidError
+from modal_proto import api_pb2
+
+from ._utils.async_utils import synchronize_api
+from ._utils.grpc_utils import RETRYABLE_GRPC_STATUS_CODES, retry_transient_errors
+from .client import _Client
+from .stream_type import StreamType
+
+if TYPE_CHECKING:
+    pass
+
+
+async def _sandbox_logs_iterator(
+    sandbox_id: str, file_descriptor: "api_pb2.FileDescriptor.ValueType", last_entry_id: str, client: _Client
+) -> AsyncGenerator[tuple[Optional[bytes], str], None]:
+    req = api_pb2.SandboxGetLogsRequest(
+        sandbox_id=sandbox_id,
+        file_descriptor=file_descriptor,
+        timeout=55,
+        last_entry_id=last_entry_id,
+    )
+    async for log_batch in client.stub.SandboxGetLogs.unary_stream(req):
+        last_entry_id = log_batch.entry_id
+
+        for message in log_batch.items:
+            yield (message.data.encode("utf-8"), last_entry_id)
+        if log_batch.eof:
+            yield (None, last_entry_id)
+            break
+
+
+async def _container_process_logs_iterator(
+    process_id: str, file_descriptor: "api_pb2.FileDescriptor.ValueType", client: _Client
+) -> AsyncGenerator[Optional[bytes], None]:
+    req = api_pb2.ContainerExecGetOutputRequest(
+        exec_id=process_id, timeout=55, file_descriptor=file_descriptor, get_raw_bytes=True
+    )
+    async for batch in client.stub.ContainerExecGetOutput.unary_stream(req):
+        if batch.HasField("exit_code"):
+            yield None
+            break
+        for item in batch.items:
+            yield item.message_bytes
+
+
+T = TypeVar("T", str, bytes)
+
+
+class _StreamReader(Generic[T]):
+    """Retrieve logs from a stream (`stdout` or `stderr`).
+
+    As an asynchronous iterable, the object supports the `for` and `async for`
+    statements. Just loop over the object to read in chunks.
+
+    **Usage**
+
+    ```python
+    from modal import Sandbox
+
+    sandbox = Sandbox.create(
+        "bash",
+        "-c",
+        "for i in $(seq 1 10); do echo foo; sleep 0.1; done",
+        app=app,
+    )
+    for message in sandbox.stdout:
+        print(f"Message: {message}")
+    ```
+    """
+
+    _stream: Optional[AsyncGenerator[Optional[bytes], None]]
+
+    def __init__(
+        self,
+        file_descriptor: "api_pb2.FileDescriptor.ValueType",
+        object_id: str,
+        object_type: Literal["sandbox", "container_process"],
+        client: _Client,
+        stream_type: StreamType = StreamType.PIPE,
+        text: bool = True,
+        by_line: bool = False,
+    ) -> None:
+        """mdmd:hidden"""
+        self._file_descriptor = file_descriptor
+        self._object_type = object_type
+        self._object_id = object_id
+        self._client = client
+        self._stream = None
+        self._last_entry_id: str = ""
+        self._line_buffer = b""
+
+        # Sandbox logs are streamed to the client as strings, so StreamReaders reading
+        # them must have text mode enabled.
+        if object_type == "sandbox" and not text:
+            raise ValueError("Sandbox streams must have text mode enabled.")
+
+        # line-buffering is only supported when text=True
+        if by_line and not text:
+            raise ValueError("line-buffering is only supported when text=True")
+
+        self._text = text
+        self._by_line = by_line
+
+        # Whether the reader received an EOF. Once EOF is True, it returns
+        # an empty string for any subsequent reads (including async for)
+        self.eof = False
+
+        if not isinstance(stream_type, StreamType):
+            raise TypeError(f"stream_type must be of type StreamType, got {type(stream_type)}")
+
+        # We only support piping sandbox logs because they're meant to be durable logs stored
+        # on the user's application.
+        if object_type == "sandbox" and stream_type != StreamType.PIPE:
+            raise ValueError("Sandbox streams must be piped.")
+        self._stream_type = stream_type
+
+        if self._object_type == "container_process":
+            # Container process streams need to be consumed as they are produced,
+            # otherwise the process will block. Use a buffer to store the stream
+            # until the client consumes it.
+            self._container_process_buffer: list[Optional[bytes]] = []
+            self._consume_container_process_task = asyncio.create_task(self._consume_container_process_stream())
+
+    @property
+    def file_descriptor(self) -> int:
+        """Possible values are `1` for stdout and `2` for stderr."""
+        return self._file_descriptor
+
+    async def read(self) -> T:
+        """Fetch the entire contents of the stream until EOF.
+
+        **Usage**
+
+        ```python
+        from modal import Sandbox
+
+        sandbox = Sandbox.create("echo", "hello", app=app)
+        sandbox.wait()
+
+        print(sandbox.stdout.read())
+        ```
+        """
+        data_str = ""
+        data_bytes = b""
+        async for message in self._get_logs():
+            if message is None:
+                break
+            if self._text:
+                data_str += message.decode("utf-8")
+            else:
+                data_bytes += message
+
+        if self._text:
+            return cast(T, data_str)
+        else:
+            return cast(T, data_bytes)
+
+    async def _consume_container_process_stream(self):
+        """Consume the container process stream and store messages in the buffer."""
+        if self._stream_type == StreamType.DEVNULL:
+            return
+
+        completed = False
+        retries_remaining = 10
+        while not completed:
+            try:
+                iterator = _container_process_logs_iterator(self._object_id, self._file_descriptor, self._client)
+
+                async for message in iterator:
+                    if self._stream_type == StreamType.STDOUT and message:
+                        print(message.decode("utf-8"), end="")
+                    elif self._stream_type == StreamType.PIPE:
+                        self._container_process_buffer.append(message)
+                    if message is None:
+                        completed = True
+                        break
+
+            except (GRPCError, StreamTerminatedError, ClientClosed) as exc:
+                if retries_remaining > 0:
+                    retries_remaining -= 1
+                    if isinstance(exc, GRPCError):
+                        if exc.status in RETRYABLE_GRPC_STATUS_CODES:
+                            await asyncio.sleep(1.0)
+                            continue
+                    elif isinstance(exc, StreamTerminatedError):
+                        continue
+                    elif isinstance(exc, ClientClosed):
+                        # If the client was closed, the user has triggered a cleanup.
+                        break
+                raise exc
+
+    async def _stream_container_process(self) -> AsyncGenerator[tuple[Optional[bytes], str], None]:
+        """Streams the container process buffer to the reader."""
+        entry_id = 0
+        if self._last_entry_id:
+            entry_id = int(self._last_entry_id) + 1
+
+        while True:
+            if entry_id >= len(self._container_process_buffer):
+                await asyncio.sleep(0.1)
+                continue
+
+            item = self._container_process_buffer[entry_id]
+
+            yield (item, str(entry_id))
+            if item is None:
+                break
+
+            entry_id += 1
+
+    async def _get_logs(self, skip_empty_messages: bool = True) -> AsyncGenerator[Optional[bytes], None]:
+        """Streams sandbox or process logs from the server to the reader.
+
+        Logs returned by this method may contain partial or multiple lines at a time.
+
+        When the stream receives an EOF, it yields None. Once an EOF is received,
+        subsequent invocations will not yield logs.
+        """
+        if self._stream_type != StreamType.PIPE:
+            raise InvalidError("Logs can only be retrieved using the PIPE stream type.")
+
+        if self.eof:
+            yield None
+            return
+
+        completed = False
+
+        retries_remaining = 10
+        while not completed:
+            try:
+                if self._object_type == "sandbox":
+                    iterator = _sandbox_logs_iterator(
+                        self._object_id, self._file_descriptor, self._last_entry_id, self._client
+                    )
+                else:
+                    iterator = self._stream_container_process()
+
+                async for message, entry_id in iterator:
+                    self._last_entry_id = entry_id
+                    # Empty messages are sent when the process boots up. Don't yield them unless
+                    # we're using the empty message to signal process liveness.
+                    if skip_empty_messages and message == b"":
+                        continue
+
+                    yield message
+                    if message is None:
+                        completed = True
+                        self.eof = True
+
+            except (GRPCError, StreamTerminatedError) as exc:
+                if retries_remaining > 0:
+                    retries_remaining -= 1
+                    if isinstance(exc, GRPCError):
+                        if exc.status in RETRYABLE_GRPC_STATUS_CODES:
+                            await asyncio.sleep(1.0)
+                            continue
+                    elif isinstance(exc, StreamTerminatedError):
+                        continue
+                raise
+
+    async def _get_logs_by_line(self) -> AsyncGenerator[Optional[bytes], None]:
+        """Process logs from the server and yield complete lines only."""
+        async for message in self._get_logs():
+            if message is None:
+                if self._line_buffer:
+                    yield self._line_buffer
+                    self._line_buffer = b""
+                yield None
+            else:
+                assert isinstance(message, bytes)
+                self._line_buffer += message
+                while b"\n" in self._line_buffer:
+                    line, self._line_buffer = self._line_buffer.split(b"\n", 1)
+                    yield line + b"\n"
+
+    def __aiter__(self) -> AsyncIterator[T]:
+        """mdmd:hidden"""
+        if not self._stream:
+            if self._by_line:
+                self._stream = self._get_logs_by_line()
+            else:
+                self._stream = self._get_logs()
+        return self
+
+    async def __anext__(self) -> T:
+        """mdmd:hidden"""
+        assert self._stream is not None
+
+        value = await self._stream.__anext__()
+
+        # The stream yields None if it receives an EOF batch.
+        if value is None:
+            raise StopAsyncIteration
+
+        if self._text:
+            return cast(T, value.decode("utf-8"))
+        else:
+            return cast(T, value)
+
+    async def aclose(self):
+        """mdmd:hidden"""
+        if self._stream:
+            await self._stream.aclose()
+
+
+MAX_BUFFER_SIZE = 2 * 1024 * 1024
+
+
+class _StreamWriter:
+    """Provides an interface to buffer and write logs to a sandbox or container process stream (`stdin`)."""
+
+    def __init__(self, object_id: str, object_type: Literal["sandbox", "container_process"], client: _Client) -> None:
+        """mdmd:hidden"""
+        self._index = 1
+        self._object_id = object_id
+        self._object_type = object_type
+        self._client = client
+        self._is_closed = False
+        self._buffer = bytearray()
+
+    def _get_next_index(self) -> int:
+        index = self._index
+        self._index += 1
+        return index
+
+    def write(self, data: Union[bytes, bytearray, memoryview, str]) -> None:
+        """Write data to the stream but does not send it immediately.
+
+        This is non-blocking and queues the data to an internal buffer. Must be
+        used along with the `drain()` method, which flushes the buffer.
+
+        **Usage**
+
+        ```python
+        from modal import Sandbox
+
+        sandbox = Sandbox.create(
+            "bash",
+            "-c",
+            "while read line; do echo $line; done",
+            app=app,
+        )
+        sandbox.stdin.write(b"foo\\n")
+        sandbox.stdin.write(b"bar\\n")
+        sandbox.stdin.write_eof()
+
+        sandbox.stdin.drain()
+        sandbox.wait()
+        ```
+        """
+        if self._is_closed:
+            raise ValueError("Stdin is closed. Cannot write to it.")
+        if isinstance(data, (bytes, bytearray, memoryview, str)):
+            if isinstance(data, str):
+                data = data.encode("utf-8")
+            if len(self._buffer) + len(data) > MAX_BUFFER_SIZE:
+                raise BufferError("Buffer size exceed limit. Call drain to clear the buffer.")
+            self._buffer.extend(data)
+        else:
+            raise TypeError(f"data argument must be a bytes-like object, not {type(data).__name__}")
+
+    def write_eof(self) -> None:
+        """Close the write end of the stream after the buffered data is drained.
+
+        If the process was blocked on input, it will become unblocked after
+        `write_eof()`. This method needs to be used along with the `drain()`
+        method, which flushes the EOF to the process.
+        """
+        self._is_closed = True
+
+    async def drain(self) -> None:
+        """Flush the write buffer and send data to the running process.
+
+        This is a flow control method that blocks until data is sent. It returns
+        when it is appropriate to continue writing data to the stream.
+
+        **Usage**
+
+        ```python notest
+        writer.write(data)
+        writer.drain()
+        ```
+
+        Async usage:
+        ```python notest
+        writer.write(data)  # not a blocking operation
+        await writer.drain.aio()
+        ```
+        """
+        data = bytes(self._buffer)
+        self._buffer.clear()
+        index = self._get_next_index()
+
+        try:
+            if self._object_type == "sandbox":
+                await retry_transient_errors(
+                    self._client.stub.SandboxStdinWrite,
+                    api_pb2.SandboxStdinWriteRequest(
+                        sandbox_id=self._object_id, index=index, eof=self._is_closed, input=data
+                    ),
+                )
+            else:
+                await retry_transient_errors(
+                    self._client.stub.ContainerExecPutInput,
+                    api_pb2.ContainerExecPutInputRequest(
+                        exec_id=self._object_id,
+                        input=api_pb2.RuntimeInputMessage(message=data, message_index=index, eof=self._is_closed),
+                    ),
+                )
+        except GRPCError as exc:
+            if exc.status == Status.FAILED_PRECONDITION:
+                raise ValueError(exc.message)
+            else:
+                raise exc
+
+
+StreamReader = synchronize_api(_StreamReader)
+StreamWriter = synchronize_api(_StreamWriter)

modal/io_streams.pyi

@@ -0,0 +1,122 @@
+import collections.abc
+import modal.client
+import modal.stream_type
+import typing
+import typing_extensions
+
+def _sandbox_logs_iterator(
+    sandbox_id: str, file_descriptor: int, last_entry_id: str, client: modal.client._Client
+) -> collections.abc.AsyncGenerator[tuple[typing.Optional[bytes], str], None]: ...
+def _container_process_logs_iterator(
+    process_id: str, file_descriptor: int, client: modal.client._Client
+) -> collections.abc.AsyncGenerator[typing.Optional[bytes], None]: ...
+
+T = typing.TypeVar("T")
+
+class _StreamReader(typing.Generic[T]):
+    _stream: typing.Optional[collections.abc.AsyncGenerator[typing.Optional[bytes], None]]
+
+    def __init__(
+        self,
+        file_descriptor: int,
+        object_id: str,
+        object_type: typing.Literal["sandbox", "container_process"],
+        client: modal.client._Client,
+        stream_type: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+        text: bool = True,
+        by_line: bool = False,
+    ) -> None: ...
+    @property
+    def file_descriptor(self) -> int: ...
+    async def read(self) -> T: ...
+    async def _consume_container_process_stream(self): ...
+    def _stream_container_process(self) -> collections.abc.AsyncGenerator[tuple[typing.Optional[bytes], str], None]: ...
+    def _get_logs(
+        self, skip_empty_messages: bool = True
+    ) -> collections.abc.AsyncGenerator[typing.Optional[bytes], None]: ...
+    def _get_logs_by_line(self) -> collections.abc.AsyncGenerator[typing.Optional[bytes], None]: ...
+    def __aiter__(self) -> collections.abc.AsyncIterator[T]: ...
+    async def __anext__(self) -> T: ...
+    async def aclose(self): ...
+
+class _StreamWriter:
+    def __init__(
+        self, object_id: str, object_type: typing.Literal["sandbox", "container_process"], client: modal.client._Client
+    ) -> None: ...
+    def _get_next_index(self) -> int: ...
+    def write(self, data: typing.Union[bytes, bytearray, memoryview, str]) -> None: ...
+    def write_eof(self) -> None: ...
+    async def drain(self) -> None: ...
+
+T_INNER = typing.TypeVar("T_INNER", covariant=True)
+
+class StreamReader(typing.Generic[T]):
+    _stream: typing.Optional[collections.abc.AsyncGenerator[typing.Optional[bytes], None]]
+
+    def __init__(
+        self,
+        file_descriptor: int,
+        object_id: str,
+        object_type: typing.Literal["sandbox", "container_process"],
+        client: modal.client.Client,
+        stream_type: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+        text: bool = True,
+        by_line: bool = False,
+    ) -> None: ...
+    @property
+    def file_descriptor(self) -> int: ...
+
+    class __read_spec(typing_extensions.Protocol[T_INNER]):
+        def __call__(self) -> T_INNER: ...
+        async def aio(self) -> T_INNER: ...
+
+    read: __read_spec[T]
+
+    class ___consume_container_process_stream_spec(typing_extensions.Protocol):
+        def __call__(self): ...
+        async def aio(self): ...
+
+    _consume_container_process_stream: ___consume_container_process_stream_spec
+
+    class ___stream_container_process_spec(typing_extensions.Protocol):
+        def __call__(self) -> typing.Generator[tuple[typing.Optional[bytes], str], None, None]: ...
+        def aio(self) -> collections.abc.AsyncGenerator[tuple[typing.Optional[bytes], str], None]: ...
+
+    _stream_container_process: ___stream_container_process_spec
+
+    class ___get_logs_spec(typing_extensions.Protocol):
+        def __call__(
+            self, skip_empty_messages: bool = True
+        ) -> typing.Generator[typing.Optional[bytes], None, None]: ...
+        def aio(
+            self, skip_empty_messages: bool = True
+        ) -> collections.abc.AsyncGenerator[typing.Optional[bytes], None]: ...
+
+    _get_logs: ___get_logs_spec
+
+    class ___get_logs_by_line_spec(typing_extensions.Protocol):
+        def __call__(self) -> typing.Generator[typing.Optional[bytes], None, None]: ...
+        def aio(self) -> collections.abc.AsyncGenerator[typing.Optional[bytes], None]: ...
+
+    _get_logs_by_line: ___get_logs_by_line_spec
+
+    def __iter__(self) -> typing.Iterator[T]: ...
+    def __aiter__(self) -> collections.abc.AsyncIterator[T]: ...
+    def __next__(self) -> T: ...
+    async def __anext__(self) -> T: ...
+    def close(self): ...
+    async def aclose(self): ...
+
+class StreamWriter:
+    def __init__(
+        self, object_id: str, object_type: typing.Literal["sandbox", "container_process"], client: modal.client.Client
+    ) -> None: ...
+    def _get_next_index(self) -> int: ...
+    def write(self, data: typing.Union[bytes, bytearray, memoryview, str]) -> None: ...
+    def write_eof(self) -> None: ...
+
+    class __drain_spec(typing_extensions.Protocol):
+        def __call__(self) -> None: ...
+        async def aio(self) -> None: ...
+
+    drain: __drain_spec