modal 1.1.5.dev66__py3-none-any.whl → 1.3.1.dev8__py3-none-any.whl

modal/io_streams.py

@@ -1,27 +1,33 @@
 # Copyright Modal Labs 2022
 import asyncio
+import codecs
+import contextlib
+import io
+import sys
 import time
 from collections.abc import AsyncGenerator, AsyncIterator
+from dataclasses import dataclass
 from typing import (
     TYPE_CHECKING,
     Generic,
     Literal,
     Optional,
+    TextIO,
     TypeVar,
     Union,
     cast,
 )
 
-from grpclib import Status
-from grpclib.exceptions import GRPCError, StreamTerminatedError
+from grpclib.exceptions import StreamTerminatedError
 
-from modal.exception import ClientClosed, InvalidError
+from modal.exception import ClientClosed, ExecTimeoutError, 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 ._utils.async_utils import aclosing, synchronize_api, synchronizer
+from ._utils.task_command_router_client import TaskCommandRouterClient
 from .client import _Client
 from .config import logger
+from .exception import ConflictError, InternalError, ServiceError
 from .stream_type import StreamType
 
 if TYPE_CHECKING:
@@ -61,7 +67,6 @@ async def _container_process_logs_iterator(
         get_raw_bytes=True,
         last_batch_index=last_index,
     )
-
     stream = client.stub.ContainerExecGetOutput.unary_stream(req)
     while True:
         # Check deadline before attempting to receive the next batch
@@ -73,39 +78,22 @@ async def _container_process_logs_iterator(
             break
         except StopAsyncIteration:
             break
+
+        for item in batch.items:
+            yield item.message_bytes, batch.batch_index
+
         if batch.HasField("exit_code"):
             yield None, batch.batch_index
             break
-        for item in batch.items:
-            yield item.message_bytes, batch.batch_index
 
 
 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**
+class _StreamReaderThroughServer(Generic[T]):
+    """A StreamReader implementation that reads from the server."""
 
-    ```python fixture:running_app
-    from modal import Sandbox
-
-    sandbox = Sandbox.create(
-        "bash",
-        "-c",
-        "for i in $(seq 1 10); do echo foo; sleep 0.1; done",
-        app=running_app,
-    )
-    for message in sandbox.stdout:
-        print(f"Message: {message}")
-    ```
-    """
-
-    _stream: Optional[AsyncGenerator[Optional[bytes], None]]
+    _stream: Optional[AsyncGenerator[T, None]]
 
     def __init__(
         self,
@@ -133,10 +121,6 @@ class _StreamReader(Generic[T]):
         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
 
@@ -154,10 +138,9 @@ class _StreamReader(Generic[T]):
         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]] = []
+            # TODO: we should not have this async code in constructors!
+            #  it only works as long as all the construction happens inside of synchronicity code
+            self._container_process_buffer: list[Optional[bytes]] = []  # TODO: change this to an asyncio.Queue
             self._consume_container_process_task = asyncio.create_task(self._consume_container_process_stream())
 
     @property
@@ -166,35 +149,20 @@ class _StreamReader(Generic[T]):
         return self._file_descriptor
 
     async def read(self) -> T:
-        """Fetch the entire contents of the stream until EOF.
-
-        **Usage**
-
-        ```python fixture:running_app
-        from modal import Sandbox
-
-        sandbox = Sandbox.create("echo", "hello", app=running_app)
-        sandbox.wait()
-
-        print(sandbox.stdout.read())
-        ```
-        """
-        data_str = ""
-        data_bytes = b""
+        """Fetch the entire contents of the stream until EOF."""
         logger.debug(f"{self._object_id} StreamReader fd={self._file_descriptor} read starting")
-        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
-
-        logger.debug(f"{self._object_id} StreamReader fd={self._file_descriptor} read completed after EOF")
         if self._text:
-            return cast(T, data_str)
+            buffer = io.StringIO()
+            async for message in _decode_bytes_stream_to_str(self._get_logs()):
+                buffer.write(message)
+            logger.debug(f"{self._object_id} StreamReader fd={self._file_descriptor} read completed after EOF")
+            return cast(T, buffer.getvalue())
         else:
-            return cast(T, data_bytes)
+            buffer = io.BytesIO()
+            async for message in self._get_logs():
+                buffer.write(message)
+            logger.debug(f"{self._object_id} StreamReader fd={self._file_descriptor} read completed after EOF")
+            return cast(T, buffer.getvalue())
 
     async def _consume_container_process_stream(self):
         """Consume the container process stream and store messages in the buffer."""
@@ -213,7 +181,8 @@ class _StreamReader(Generic[T]):
                 )
                 async for message, batch_index in iterator:
                     if self._stream_type == StreamType.STDOUT and message:
-                        print(message.decode("utf-8"), end="")
+                        # TODO: rearchitect this, since these bytes aren't necessarily decodable
+                        print(message.decode("utf-8"), end="")  # noqa: T201
                     elif self._stream_type == StreamType.PIPE:
                         self._container_process_buffer.append(message)
 
@@ -223,13 +192,12 @@ class _StreamReader(Generic[T]):
                     else:
                         last_index = batch_index
 
-            except (GRPCError, StreamTerminatedError, ClientClosed) as exc:
+            except (ServiceError, InternalError, 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
+                    if isinstance(exc, (ServiceError, InternalError)):
+                        await asyncio.sleep(1.0)
+                        continue
                     elif isinstance(exc, StreamTerminatedError):
                         continue
                     elif isinstance(exc, ClientClosed):
@@ -240,6 +208,9 @@ class _StreamReader(Generic[T]):
 
     async def _stream_container_process(self) -> AsyncGenerator[tuple[Optional[bytes], str], None]:
         """Streams the container process buffer to the reader."""
+        # 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.
         entry_id = 0
         if self._last_entry_id:
             entry_id = int(self._last_entry_id) + 1
@@ -257,7 +228,7 @@ class _StreamReader(Generic[T]):
 
             entry_id += 1
 
-    async def _get_logs(self, skip_empty_messages: bool = True) -> AsyncGenerator[Optional[bytes], None]:
+    async def _get_logs(self, skip_empty_messages: bool = True) -> AsyncGenerator[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.
@@ -269,7 +240,6 @@ class _StreamReader(Generic[T]):
             raise InvalidError("Logs can only be retrieved using the PIPE stream type.")
 
         if self.eof:
-            yield None
             return
 
         completed = False
@@ -294,72 +264,373 @@ class _StreamReader(Generic[T]):
                     if message is None:
                         completed = True
                         self.eof = True
+                        return
+
                     yield message
 
-            except (GRPCError, StreamTerminatedError) as exc:
+            except (ServiceError, InternalError, 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
+                    if isinstance(exc, (ServiceError, InternalError)):
+                        await asyncio.sleep(1.0)
+                        continue
                     elif isinstance(exc, StreamTerminatedError):
                         continue
                 raise
 
-    async def _get_logs_by_line(self) -> AsyncGenerator[Optional[bytes], None]:
+    async def _get_logs_by_line(self) -> AsyncGenerator[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"
+            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"
+
+        if self._line_buffer:
+            yield self._line_buffer
+            self._line_buffer = b""
 
-    def _ensure_stream(self) -> AsyncGenerator[Optional[bytes], None]:
+    def __aiter__(self) -> AsyncGenerator[T, None]:
         if not self._stream:
             if self._by_line:
-                self._stream = self._get_logs_by_line()
+                # TODO: This is quite odd - it does line buffering in binary mode
+                # but we then always add the buffered text decoding on top of that.
+                # feels a bit upside down...
+                stream = self._get_logs_by_line()
             else:
-                self._stream = self._get_logs()
+                stream = self._get_logs()
+            if self._text:
+                stream = _decode_bytes_stream_to_str(stream)
+            self._stream = cast(AsyncGenerator[T, None], stream)
         return self._stream
 
-    def __aiter__(self) -> AsyncIterator[T]:
+    async def aclose(self):
         """mdmd:hidden"""
-        self._ensure_stream()
-        return self
+        if self._stream:
+            await self._stream.aclose()
+
+
+async def _decode_bytes_stream_to_str(stream: AsyncGenerator[bytes, None]) -> AsyncGenerator[str, None]:
+    """Incrementally decode a bytes async generator as UTF-8 without breaking on chunk boundaries.
+
+    This function uses a streaming UTF-8 decoder so that multi-byte characters split across
+    chunks are handled correctly instead of raising ``UnicodeDecodeError``.
+    """
+    decoder = codecs.getincrementaldecoder("utf-8")(errors="strict")
+    async for item in stream:
+        text = decoder.decode(item, final=False)
+        if text:
+            yield text
+
+    # Flush any buffered partial character at end-of-stream
+    tail = decoder.decode(b"", final=True)
+    if tail:
+        yield tail
+
+
+async def _stream_by_line(stream: AsyncGenerator[bytes, None]) -> AsyncGenerator[bytes, None]:
+    """Yield complete lines only (ending with \n), buffering partial lines until complete.
+
+    When this generator returns, the underlying generator is closed.
+    """
+    line_buffer = b""
+    try:
+        async for message in stream:
+            assert isinstance(message, bytes)
+            line_buffer += message
+            while b"\n" in line_buffer:
+                line, line_buffer = line_buffer.split(b"\n", 1)
+                yield line + b"\n"
+
+        if line_buffer:
+            yield line_buffer
+    finally:
+        await stream.aclose()
+
+
+@dataclass
+class _StreamReaderThroughCommandRouterParams:
+    file_descriptor: "api_pb2.FileDescriptor.ValueType"
+    task_id: str
+    object_id: str
+    command_router_client: TaskCommandRouterClient
+    deadline: Optional[float]
+
+
+async def _stdio_stream_from_command_router(
+    params: _StreamReaderThroughCommandRouterParams,
+) -> AsyncGenerator[bytes, None]:
+    """Stream raw bytes from the router client."""
+    async with aclosing(
+        params.command_router_client.exec_stdio_read(
+            params.task_id, params.object_id, params.file_descriptor, params.deadline
+        )
+    ) as stream:
+        try:
+            async for item in stream:
+                if len(item.data) == 0:
+                    # This is an error.
+                    raise ValueError("Received empty message streaming stdio from sandbox.")
+
+                yield item.data
+        except ExecTimeoutError:
+            logger.debug(f"Deadline exceeded while streaming stdio for exec {params.object_id}")
+            # TODO(saltzm): This is a weird API, but customers currently may rely on it. We
+            # should probably raise this error rather than just ending the stream.
+            return
+
+
+class _BytesStreamReaderThroughCommandRouter:
+    """
+    StreamReader implementation that will read directly from the worker that
+    hosts the sandbox.
+
+    This implementation is used for non-text streams.
+    """
+
+    def __init__(
+        self,
+        params: _StreamReaderThroughCommandRouterParams,
+    ) -> None:
+        self._params = params
+        self._stream = None
+
+    @property
+    def file_descriptor(self) -> int:
+        return self._params.file_descriptor
+
+    async def read(self) -> bytes:
+        buffer = io.BytesIO()
+        async for part in self:
+            buffer.write(part)
+        return buffer.getvalue()
+
+    def __aiter__(self) -> AsyncGenerator[bytes, None]:
+        return _stdio_stream_from_command_router(self._params)
+
+    async def _print_all(self, output_stream: TextIO) -> None:
+        async for part in self:
+            output_stream.buffer.write(part)
+            output_stream.buffer.flush()
+
+
+class _TextStreamReaderThroughCommandRouter:
+    """
+    StreamReader implementation that will read directly from the worker
+    that hosts the sandbox.
+
+    This implementation is used for text streams.
+    """
+
+    def __init__(
+        self,
+        params: _StreamReaderThroughCommandRouterParams,
+        by_line: bool,
+    ) -> None:
+        self._params = params
+        self._by_line = by_line
+
+    @property
+    def file_descriptor(self) -> int:
+        return self._params.file_descriptor
+
+    async def read(self) -> str:
+        buffer = io.StringIO()
+        async for part in self:
+            buffer.write(part)
+        return buffer.getvalue()
+
+    async def __aiter__(self) -> AsyncGenerator[str, None]:
+        async with aclosing(_stdio_stream_from_command_router(self._params)) as bytes_stream:
+            if self._by_line:
+                stream = _decode_bytes_stream_to_str(_stream_by_line(bytes_stream))
+            else:
+                stream = _decode_bytes_stream_to_str(bytes_stream)
+
+            async with aclosing(stream):
+                async for part in stream:
+                    yield part
+
+    async def _print_all(self, output_stream: TextIO) -> None:
+        async with aclosing(self.__aiter__()) as stream:
+            async for part in stream:
+                output_stream.write(part)
+
+
+class _StdoutPrintingStreamReaderThroughCommandRouter(Generic[T]):
+    """
+    StreamReader implementation for StreamType.STDOUT when using the task command router.
+
+    This mirrors the behavior from the server-backed implementation: the stream is printed to
+    the local stdout immediately and is not readable via StreamReader methods.
+    """
+
+    _reader: Union[_TextStreamReaderThroughCommandRouter, _BytesStreamReaderThroughCommandRouter]
+
+    def __init__(
+        self,
+        reader: Union[_TextStreamReaderThroughCommandRouter, _BytesStreamReaderThroughCommandRouter],
+    ) -> None:
+        self._reader = reader
+        self._task: Optional[asyncio.Task[None]] = None
+        # Kick off a background task that reads from the underlying text stream and prints to stdout.
+        self._start_printing_task()
+
+    @property
+    def file_descriptor(self) -> int:
+        return self._reader.file_descriptor
+
+    def _start_printing_task(self) -> None:
+        async def _run():
+            try:
+                await self._reader._print_all(sys.stdout)
+            except Exception as e:
+                logger.exception(f"Error printing stream: {e}")
+
+        self._task = asyncio.create_task(_run())
+
+    async def read(self) -> T:
+        raise InvalidError("Output can only be retrieved using the PIPE stream type.")
+
+    def __aiter__(self) -> AsyncIterator[T]:
+        raise InvalidError("Output can only be retrieved using the PIPE stream type.")
 
     async def __anext__(self) -> T:
-        """mdmd:hidden"""
-        stream = self._ensure_stream()
+        raise InvalidError("Output can only be retrieved using the PIPE stream type.")
 
-        value = await stream.__anext__()
+    async def aclose(self):
+        if self._task is not None:
+            self._task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._task
+            self._task = None
 
-        # 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"))
+class _DevnullStreamReader(Generic[T]):
+    """StreamReader implementation for a stream configured with
+    StreamType.DEVNULL. Throws an error if read or any other method is
+    called.
+    """
+
+    def __init__(self, file_descriptor: "api_pb2.FileDescriptor.ValueType") -> None:
+        self._file_descriptor = file_descriptor
+
+    @property
+    def file_descriptor(self) -> int:
+        return self._file_descriptor
+
+    async def read(self) -> T:
+        raise ValueError("read is not supported for a stream configured with StreamType.DEVNULL")
+
+    def __aiter__(self) -> AsyncIterator[T]:
+        raise ValueError("__aiter__ is not supported for a stream configured with StreamType.DEVNULL")
+
+    async def __anext__(self) -> T:
+        raise ValueError("__anext__ is not supported for a stream configured with StreamType.DEVNULL")
+
+    async def aclose(self):
+        raise ValueError("aclose is not supported for a stream configured with StreamType.DEVNULL")
+
+
+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.
+    """
+
+    _impl: Union[
+        _StreamReaderThroughServer,
+        _DevnullStreamReader,
+        _TextStreamReaderThroughCommandRouter,
+        _BytesStreamReaderThroughCommandRouter,
+        _StdoutPrintingStreamReaderThroughCommandRouter,
+    ]
+    _read_gen: Optional[AsyncGenerator[T, None]] = 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,
+        deadline: Optional[float] = None,
+        command_router_client: Optional[TaskCommandRouterClient] = None,
+        task_id: Optional[str] = None,
+    ) -> None:
+        """mdmd:hidden"""
+        # we can remove this once we ensure no constructors use async code
+        assert asyncio.get_running_loop() == synchronizer._get_loop(start=False)
+
+        if by_line and not text:
+            raise ValueError("line-buffering is only supported when text=True")
+
+        if command_router_client is None:
+            self._impl = _StreamReaderThroughServer(
+                file_descriptor, object_id, object_type, client, stream_type, text, by_line, deadline
+            )
         else:
-            return cast(T, value)
+            # The only reason task_id is optional is because StreamReader is also used for sandbox
+            # logs, which don't have a task ID available when the StreamReader is created.
+            assert task_id is not None
+            assert object_type == "container_process"
+            if stream_type == StreamType.DEVNULL:
+                self._impl = _DevnullStreamReader(file_descriptor)
+            else:
+                assert stream_type == StreamType.PIPE or stream_type == StreamType.STDOUT
+                params = _StreamReaderThroughCommandRouterParams(
+                    file_descriptor, task_id, object_id, command_router_client, deadline
+                )
+                if text:
+                    reader = _TextStreamReaderThroughCommandRouter(params, by_line)
+                else:
+                    reader = _BytesStreamReaderThroughCommandRouter(params)
+
+                if stream_type == StreamType.STDOUT:
+                    self._impl = _StdoutPrintingStreamReaderThroughCommandRouter(reader)
+                else:
+                    self._impl = reader
+
+    @property
+    def file_descriptor(self) -> int:
+        """Possible values are `1` for stdout and `2` for stderr."""
+        return self._impl.file_descriptor
+
+    async def read(self) -> T:
+        """Fetch the entire contents of the stream until EOF."""
+        return cast(T, await self._impl.read())
+
+    def __aiter__(self) -> AsyncGenerator[T, None]:
+        if not self._read_gen:
+            self._read_gen = cast(AsyncGenerator[T, None], self._impl.__aiter__())
+        return self._read_gen
+
+    async def __anext__(self) -> T:
+        """Deprecated: This exists for backwards compatibility and will be removed in a future version of Modal
+
+        Only use next/anext on the return value of iter/aiter on the StreamReader object (treat streamreader as
+        an iterable, not an iterator).
+        """
+        if not self._read_gen:
+            self.__aiter__()  # initialize the read generator
+        assert self._read_gen
+        return await self._read_gen.__anext__()
 
     async def aclose(self):
         """mdmd:hidden"""
-        if self._stream:
-            await self._stream.aclose()
+        if self._read_gen:
+            await self._read_gen.aclose()
+            self._read_gen = None
 
 
 MAX_BUFFER_SIZE = 2 * 1024 * 1024
 
 
-class _StreamWriter:
+class _StreamWriterThroughServer:
     """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:
@@ -381,25 +652,6 @@ class _StreamWriter:
 
         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 fixture:running_app
-        from modal import Sandbox
-
-        sandbox = Sandbox.create(
-            "bash",
-            "-c",
-            "while read line; do echo $line; done",
-            app=running_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.")
@@ -407,7 +659,7 @@ class _StreamWriter:
             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.")
+                raise BufferError("Buffer size exceed limit. Call drain to flush the buffer.")
             self._buffer.extend(data)
         else:
             raise TypeError(f"data argument must be a bytes-like object, not {type(data).__name__}")
@@ -426,19 +678,6 @@ class _StreamWriter:
 
         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()
@@ -446,25 +685,137 @@ class _StreamWriter:
 
         try:
             if self._object_type == "sandbox":
-                await retry_transient_errors(
-                    self._client.stub.SandboxStdinWrite,
+                await 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,
+                await 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
+        except ConflictError as exc:
+            raise ValueError(str(exc))
+
+
+class _StreamWriterThroughCommandRouter:
+    def __init__(
+        self,
+        object_id: str,
+        command_router_client: TaskCommandRouterClient,
+        task_id: str,
+    ) -> None:
+        self._object_id = object_id
+        self._command_router_client = command_router_client
+        self._task_id = task_id
+        self._is_closed = False
+        self._buffer = bytearray()
+        self._offset = 0
+
+    def write(self, data: Union[bytes, bytearray, memoryview, str]) -> None:
+        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 flush 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:
+        self._is_closed = True
+
+    async def drain(self) -> None:
+        eof = self._is_closed
+        # NB: There's no need to prevent writing eof twice, because the command router will ignore the second EOF.
+        if self._buffer or eof:
+            data = bytes(self._buffer)
+            await self._command_router_client.exec_stdin_write(
+                task_id=self._task_id, exec_id=self._object_id, offset=self._offset, data=data, eof=eof
+            )
+            # Only clear the buffer after writing the data to the command router is successful.
+            # This allows the client to retry drain() in the event of an exception (though
+            # exec_stdin_write already retries on transient errors, so most users will probably
+            # not do this).
+            self._buffer.clear()
+            self._offset += len(data)
+
+
+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,
+        command_router_client: Optional[TaskCommandRouterClient] = None,
+        task_id: Optional[str] = None,
+    ) -> None:
+        """mdmd:hidden"""
+        if command_router_client is None:
+            self._impl = _StreamWriterThroughServer(object_id, object_type, client)
+        else:
+            assert task_id is not None
+            assert object_type == "container_process"
+            self._impl = _StreamWriterThroughCommandRouter(object_id, command_router_client, task_id=task_id)
+
+    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 fixture:sandbox
+        proc = sandbox.exec(
+            "bash",
+            "-c",
+            "while read line; do echo $line; done",
+        )
+        proc.stdin.write(b"foo\\n")
+        proc.stdin.write(b"bar\\n")
+        proc.stdin.write_eof()
+        proc.stdin.drain()
+        ```
+        """
+        self._impl.write(data)
+
+    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._impl.write_eof()
+
+    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()
+        ```
+        """
+        await self._impl.drain()
 
 
 StreamReader = synchronize_api(_StreamReader)