modal 1.2.1.dev12__py3-none-any.whl → 1.2.1.dev14__py3-none-any.whl

modal/_utils/task_command_router_client.py

@@ -292,7 +292,9 @@ class TaskCommandRouterClient:
             lambda: self._call_with_auth_retry(self._stub.TaskExecStdinWrite, request)
         )
 
-    async def exec_poll(self, task_id: str, exec_id: str) -> sr_pb2.TaskExecPollResponse:
+    async def exec_poll(
+        self, task_id: str, exec_id: str, deadline: Optional[float] = None
+    ) -> sr_pb2.TaskExecPollResponse:
         """Poll for the exit status of an exec'd command, properly retrying on transient errors.
 
         Args:
@@ -302,13 +304,25 @@ class TaskCommandRouterClient:
             sr_pb2.TaskExecPollResponse: The exit status of the command if it has completed.
 
         Raises:
+            ExecTimeoutError: If the deadline is exceeded.
             Other errors: If retries are exhausted on transient errors or if there's an error
               from the RPC itself.
         """
         request = sr_pb2.TaskExecPollRequest(task_id=task_id, exec_id=exec_id)
-        return await call_with_retries_on_transient_errors(
-            lambda: self._call_with_auth_retry(self._stub.TaskExecPoll, request)
-        )
+        # The timeout here is really a backstop in the event of a hang contacting
+        # the command router. Poll should usually be instantaneous.
+        timeout = deadline - time.monotonic() if deadline is not None else None
+        if timeout is not None and timeout <= 0:
+            raise ExecTimeoutError(f"Deadline exceeded while polling for exec {exec_id}")
+        try:
+            return await asyncio.wait_for(
+                call_with_retries_on_transient_errors(
+                    lambda: self._call_with_auth_retry(self._stub.TaskExecPoll, request)
+                ),
+                timeout=timeout,
+            )
+        except asyncio.TimeoutError:
+            raise ExecTimeoutError(f"Deadline exceeded while polling for exec {exec_id}")
 
     async def exec_wait(
         self,

modal/cli/cluster.py

@@ -83,7 +83,9 @@ async def shell(
     )
     exec_res: api_pb2.ContainerExecResponse = await client.stub.ContainerExec(req)
     if pty:
-        await _ContainerProcess(exec_res.exec_id, client).attach()
+        await _ContainerProcess(exec_res.exec_id, task_id, client).attach()
     else:
         # TODO: redirect stderr to its own stream?
-        await _ContainerProcess(exec_res.exec_id, client, stdout=StreamType.STDOUT, stderr=StreamType.STDOUT).wait()
+        await _ContainerProcess(
+            exec_res.exec_id, task_id, client, stdout=StreamType.STDOUT, stderr=StreamType.STDOUT
+        ).wait()

modal/cli/container.py

@@ -80,10 +80,12 @@ async def exec(
     res: api_pb2.ContainerExecResponse = await client.stub.ContainerExec(req)
 
     if pty:
-        await _ContainerProcess(res.exec_id, client).attach()
+        await _ContainerProcess(res.exec_id, container_id, client).attach()
     else:
         # TODO: redirect stderr to its own stream?
-        await _ContainerProcess(res.exec_id, client, stdout=StreamType.STDOUT, stderr=StreamType.STDOUT).wait()
+        await _ContainerProcess(
+            res.exec_id, container_id, client, stdout=StreamType.STDOUT, stderr=StreamType.STDOUT
+        ).wait()
 
 
 @container_cli.command("stop")

modal/client.pyi

@@ -33,7 +33,7 @@ class _Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.2.1.dev12",
+        version: str = "1.2.1.dev14",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.
@@ -164,7 +164,7 @@ class Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.2.1.dev12",
+        version: str = "1.2.1.dev14",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.

modal/container_process.py

@@ -9,16 +9,17 @@ from modal_proto import api_pb2
 from ._utils.async_utils import TaskContext, synchronize_api
 from ._utils.grpc_utils import retry_transient_errors
 from ._utils.shell_utils import stream_from_stdin, write_to_fd
+from ._utils.task_command_router_client import TaskCommandRouterClient
 from .client import _Client
 from .config import logger
-from .exception import InteractiveTimeoutError, InvalidError
+from .exception import ExecTimeoutError, InteractiveTimeoutError, InvalidError
 from .io_streams import _StreamReader, _StreamWriter
 from .stream_type import StreamType
 
 T = TypeVar("T", str, bytes)
 
 
-class _ContainerProcess(Generic[T]):
+class _ContainerProcessThroughServer(Generic[T]):
     _process_id: Optional[str] = None
     _stdout: _StreamReader[T]
     _stderr: _StreamReader[T]
@@ -31,6 +32,7 @@ class _ContainerProcess(Generic[T]):
     def __init__(
         self,
         process_id: str,
+        task_id: str,
         client: _Client,
         stdout: StreamType = StreamType.PIPE,
         stderr: StreamType = StreamType.PIPE,
@@ -52,6 +54,7 @@ class _ContainerProcess(Generic[T]):
             text=text,
             by_line=by_line,
             deadline=exec_deadline,
+            task_id=task_id,
         )
         self._stderr = _StreamReader[T](
             api_pb2.FILE_DESCRIPTOR_STDERR,
@@ -62,6 +65,7 @@ class _ContainerProcess(Generic[T]):
             text=text,
             by_line=by_line,
             deadline=exec_deadline,
+            task_id=task_id,
         )
         self._stdin = _StreamWriter(process_id, "container_process", self._client)
 
@@ -201,4 +205,266 @@ class _ContainerProcess(Generic[T]):
                 raise InteractiveTimeoutError("Failed to establish connection to container. Please try again.")
 
 
+async def _iter_stream_as_bytes(stream: _StreamReader[T]):
+    """Yield raw bytes from a StreamReader regardless of text mode/backend."""
+    async for part in stream:
+        if isinstance(part, str):
+            yield part.encode("utf-8")
+        else:
+            yield part
+
+
+class _ContainerProcessThroughCommandRouter(Generic[T]):
+    """
+    Container process implementation that works via direct communication with
+    the Modal worker where the container is running.
+    """
+
+    def __init__(
+        self,
+        process_id: str,
+        client: _Client,
+        command_router_client: TaskCommandRouterClient,
+        task_id: str,
+        *,
+        stdout: StreamType = StreamType.PIPE,
+        stderr: StreamType = StreamType.PIPE,
+        exec_deadline: Optional[float] = None,
+        text: bool = True,
+        by_line: bool = False,
+    ) -> None:
+        self._client = client
+        self._command_router_client = command_router_client
+        self._process_id = process_id
+        self._exec_deadline = exec_deadline
+        self._text = text
+        self._by_line = by_line
+        self._task_id = task_id
+        self._stdout = _StreamReader[T](
+            api_pb2.FILE_DESCRIPTOR_STDOUT,
+            process_id,
+            "container_process",
+            self._client,
+            stream_type=stdout,
+            text=text,
+            by_line=by_line,
+            deadline=exec_deadline,
+            command_router_client=self._command_router_client,
+            task_id=self._task_id,
+        )
+        self._stderr = _StreamReader[T](
+            api_pb2.FILE_DESCRIPTOR_STDERR,
+            process_id,
+            "container_process",
+            self._client,
+            stream_type=stderr,
+            text=text,
+            by_line=by_line,
+            deadline=exec_deadline,
+            command_router_client=self._command_router_client,
+            task_id=self._task_id,
+        )
+        self._stdin = _StreamWriter(
+            process_id,
+            "container_process",
+            self._client,
+            command_router_client=self._command_router_client,
+            task_id=self._task_id,
+        )
+        self._returncode = None
+
+    @property
+    def stdout(self) -> _StreamReader[T]:
+        return self._stdout
+
+    @property
+    def stderr(self) -> _StreamReader[T]:
+        return self._stderr
+
+    @property
+    def stdin(self) -> _StreamWriter:
+        return self._stdin
+
+    @property
+    def returncode(self) -> int:
+        if self._returncode is None:
+            raise InvalidError(
+                "You must call wait() before accessing the returncode. "
+                "To poll for the status of a running process, use poll() instead."
+            )
+        return self._returncode
+
+    async def poll(self) -> Optional[int]:
+        if self._returncode is not None:
+            return self._returncode
+        try:
+            resp = await self._command_router_client.exec_poll(self._task_id, self._process_id, self._exec_deadline)
+            which = resp.WhichOneof("exit_status")
+            if which is None:
+                return None
+
+            if which == "code":
+                self._returncode = int(resp.code)
+                return self._returncode
+            elif which == "signal":
+                self._returncode = 128 + int(resp.signal)
+                return self._returncode
+            else:
+                logger.debug(f"ContainerProcess {self._process_id} exited with unexpected status: {which}")
+                raise InvalidError("Unexpected exit status")
+        except ExecTimeoutError:
+            logger.debug(f"ContainerProcess poll for {self._process_id} did not complete within deadline")
+            return None
+        except Exception as e:
+            # Re-raise non-transient errors or errors resulting from exceeding retries on transient errors.
+            logger.warning(f"ContainerProcess poll for {self._process_id} failed: {e}")
+            raise
+
+    async def wait(self) -> int:
+        if self._returncode is not None:
+            return self._returncode
+
+        try:
+            resp = await self._command_router_client.exec_wait(self._task_id, self._process_id, self._exec_deadline)
+            which = resp.WhichOneof("exit_status")
+            if which == "code":
+                self._returncode = int(resp.code)
+            elif which == "signal":
+                self._returncode = 128 + int(resp.signal)
+            else:
+                logger.debug(f"ContainerProcess {self._process_id} exited with unexpected status: {which}")
+                self._returncode = -1
+                raise InvalidError("Unexpected exit status")
+        except ExecTimeoutError:
+            logger.debug(f"ContainerProcess {self._process_id} did not complete within deadline")
+            # TODO(saltzm): This is a weird API, but customers currently may rely on it. This
+            # should be a ExecTimeoutError.
+            self._returncode = -1
+
+        return self._returncode
+
+    async def attach(self):
+        if platform.system() == "Windows":
+            print("interactive exec is not currently supported on Windows.")
+            return
+
+        from ._output import make_console
+
+        console = make_console()
+
+        connecting_status = console.status("Connecting...")
+        connecting_status.start()
+        on_connect = asyncio.Event()
+
+        async def _write_to_fd_loop(stream: _StreamReader[T]):
+            async for chunk in _iter_stream_as_bytes(stream):
+                if chunk is None:
+                    break
+
+                if not on_connect.is_set():
+                    connecting_status.stop()
+                    on_connect.set()
+
+                await write_to_fd(stream.file_descriptor, chunk)
+
+        async def _handle_input(data: bytes, message_index: int):
+            self.stdin.write(data)
+            await self.stdin.drain()
+
+        async with TaskContext() as tc:
+            stdout_task = tc.create_task(_write_to_fd_loop(self.stdout))
+            stderr_task = tc.create_task(_write_to_fd_loop(self.stderr))
+
+            try:
+                # Time out if we can't connect fast enough.
+                await asyncio.wait_for(on_connect.wait(), timeout=60)
+
+                async with stream_from_stdin(_handle_input, use_raw_terminal=True):
+                    await stdout_task
+                    await stderr_task
+
+            except (asyncio.TimeoutError, TimeoutError):
+                connecting_status.stop()
+                stdout_task.cancel()
+                stderr_task.cancel()
+                raise InteractiveTimeoutError("Failed to establish connection to container. Please try again.")
+
+
+class _ContainerProcess(Generic[T]):
+    """Represents a running process in a container."""
+
+    def __init__(
+        self,
+        process_id: str,
+        task_id: str,
+        client: _Client,
+        stdout: StreamType = StreamType.PIPE,
+        stderr: StreamType = StreamType.PIPE,
+        exec_deadline: Optional[float] = None,
+        text: bool = True,
+        by_line: bool = False,
+        command_router_client: Optional[TaskCommandRouterClient] = None,
+    ) -> None:
+        if command_router_client is None:
+            self._impl = _ContainerProcessThroughServer(
+                process_id,
+                task_id,
+                client,
+                stdout=stdout,
+                stderr=stderr,
+                exec_deadline=exec_deadline,
+                text=text,
+                by_line=by_line,
+            )
+        else:
+            self._impl = _ContainerProcessThroughCommandRouter(
+                process_id,
+                client,
+                command_router_client,
+                task_id,
+                stdout=stdout,
+                stderr=stderr,
+                exec_deadline=exec_deadline,
+                text=text,
+                by_line=by_line,
+            )
+
+    def __repr__(self) -> str:
+        return self._impl.__repr__()
+
+    @property
+    def stdout(self) -> _StreamReader[T]:
+        """StreamReader for the container process's stdout stream."""
+        return self._impl.stdout
+
+    @property
+    def stderr(self) -> _StreamReader[T]:
+        """StreamReader for the container process's stderr stream."""
+        return self._impl.stderr
+
+    @property
+    def stdin(self) -> _StreamWriter:
+        """StreamWriter for the container process's stdin stream."""
+        return self._impl.stdin
+
+    @property
+    def returncode(self) -> int:
+        return self._impl.returncode
+
+    async def poll(self) -> Optional[int]:
+        """Check if the container process has finished running.
+
+        Returns `None` if the process is still running, else returns the exit code.
+        """
+        return await self._impl.poll()
+
+    async def wait(self) -> int:
+        """Wait for the container process to finish running. Returns the exit code."""
+        return await self._impl.wait()
+
+    async def attach(self):
+        """mdmd:hidden"""
+        await self._impl.attach()
+
+
 ContainerProcess = synchronize_api(_ContainerProcess)

modal/container_process.pyi

@@ -1,3 +1,4 @@
+import modal._utils.task_command_router_client
 import modal.client
 import modal.io_streams
 import modal.stream_type
@@ -6,7 +7,7 @@ import typing_extensions
 
 T = typing.TypeVar("T")
 
-class _ContainerProcess(typing.Generic[T]):
+class _ContainerProcessThroughServer(typing.Generic[T]):
     """Abstract base class for generic types.
 
     A generic type is typically declared by inheriting from
@@ -39,6 +40,7 @@ class _ContainerProcess(typing.Generic[T]):
     def __init__(
         self,
         process_id: str,
+        task_id: str,
         client: modal.client._Client,
         stdout: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
         stderr: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
@@ -86,47 +88,110 @@ class _ContainerProcess(typing.Generic[T]):
         """mdmd:hidden"""
         ...
 
-SUPERSELF = typing.TypeVar("SUPERSELF", covariant=True)
+def _iter_stream_as_bytes(stream: modal.io_streams._StreamReader[T]):
+    """Yield raw bytes from a StreamReader regardless of text mode/backend."""
+    ...
 
-class ContainerProcess(typing.Generic[T]):
-    """Abstract base class for generic types.
+class _ContainerProcessThroughCommandRouter(typing.Generic[T]):
+    """Container process implementation that works via direct communication with
+    the Modal worker where the container is running.
+    """
+    def __init__(
+        self,
+        process_id: str,
+        client: modal.client._Client,
+        command_router_client: modal._utils.task_command_router_client.TaskCommandRouterClient,
+        task_id: str,
+        *,
+        stdout: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+        stderr: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+        exec_deadline: typing.Optional[float] = None,
+        text: bool = True,
+        by_line: bool = False,
+    ) -> None:
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
 
-    A generic type is typically declared by inheriting from
-    this class parameterized with one or more type variables.
-    For example, a generic mapping type might be defined as::
+    @property
+    def stdout(self) -> modal.io_streams._StreamReader[T]: ...
+    @property
+    def stderr(self) -> modal.io_streams._StreamReader[T]: ...
+    @property
+    def stdin(self) -> modal.io_streams._StreamWriter: ...
+    @property
+    def returncode(self) -> int: ...
+    async def poll(self) -> typing.Optional[int]: ...
+    async def wait(self) -> int: ...
+    async def attach(self): ...
 
-      class Mapping(Generic[KT, VT]):
-          def __getitem__(self, key: KT) -> VT:
-              ...
-          # Etc.
+class _ContainerProcess(typing.Generic[T]):
+    """Represents a running process in a container."""
+    def __init__(
+        self,
+        process_id: str,
+        task_id: str,
+        client: modal.client._Client,
+        stdout: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+        stderr: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+        exec_deadline: typing.Optional[float] = None,
+        text: bool = True,
+        by_line: bool = False,
+        command_router_client: typing.Optional[modal._utils.task_command_router_client.TaskCommandRouterClient] = None,
+    ) -> None:
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
 
-    This class can then be used as follows::
+    def __repr__(self) -> str:
+        """Return repr(self)."""
+        ...
 
-      def lookup_name(mapping: Mapping[KT, VT], key: KT, default: VT) -> VT:
-          try:
-              return mapping[key]
-          except KeyError:
-              return default
-    """
+    @property
+    def stdout(self) -> modal.io_streams._StreamReader[T]:
+        """StreamReader for the container process's stdout stream."""
+        ...
 
-    _process_id: typing.Optional[str]
-    _stdout: modal.io_streams.StreamReader[T]
-    _stderr: modal.io_streams.StreamReader[T]
-    _stdin: modal.io_streams.StreamWriter
-    _exec_deadline: typing.Optional[float]
-    _text: bool
-    _by_line: bool
-    _returncode: typing.Optional[int]
+    @property
+    def stderr(self) -> modal.io_streams._StreamReader[T]:
+        """StreamReader for the container process's stderr stream."""
+        ...
+
+    @property
+    def stdin(self) -> modal.io_streams._StreamWriter:
+        """StreamWriter for the container process's stdin stream."""
+        ...
+
+    @property
+    def returncode(self) -> int: ...
+    async def poll(self) -> typing.Optional[int]:
+        """Check if the container process has finished running.
 
+        Returns `None` if the process is still running, else returns the exit code.
+        """
+        ...
+
+    async def wait(self) -> int:
+        """Wait for the container process to finish running. Returns the exit code."""
+        ...
+
+    async def attach(self):
+        """mdmd:hidden"""
+        ...
+
+SUPERSELF = typing.TypeVar("SUPERSELF", covariant=True)
+
+class ContainerProcess(typing.Generic[T]):
+    """Represents a running process in a container."""
     def __init__(
         self,
         process_id: str,
+        task_id: str,
         client: modal.client.Client,
         stdout: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
         stderr: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
         exec_deadline: typing.Optional[float] = None,
         text: bool = True,
         by_line: bool = False,
+        command_router_client: typing.Optional[modal._utils.task_command_router_client.TaskCommandRouterClient] = None,
     ) -> None: ...
     def __repr__(self) -> str: ...
     @property
@@ -164,12 +229,6 @@ class ContainerProcess(typing.Generic[T]):
 
     poll: __poll_spec[typing_extensions.Self]
 
-    class ___wait_for_completion_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(self, /) -> int: ...
-        async def aio(self, /) -> int: ...
-
-    _wait_for_completion: ___wait_for_completion_spec[typing_extensions.Self]
-
     class __wait_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, /) -> int:
             """Wait for the container process to finish running. Returns the exit code."""

modal/functions.pyi

@@ -401,7 +401,7 @@ class Function(
 
     _call_generator: ___call_generator_spec[typing_extensions.Self]
 
-    class __remote_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
+    class __remote_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
         def __call__(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> ReturnType_INNER:
             """Calls the function remotely, executing it with the given arguments and returning the execution's result."""
             ...
@@ -410,7 +410,7 @@ class Function(
             """Calls the function remotely, executing it with the given arguments and returning the execution's result."""
             ...
 
-    remote: __remote_spec[modal._functions.ReturnType, modal._functions.P, typing_extensions.Self]
+    remote: __remote_spec[modal._functions.P, modal._functions.ReturnType, typing_extensions.Self]
 
     class __remote_gen_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, /, *args, **kwargs) -> typing.Generator[typing.Any, None, None]:
@@ -437,7 +437,7 @@ class Function(
         """
         ...
 
-    class ___experimental_spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
+    class ___experimental_spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
         def __call__(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]:
             """[Experimental] Calls the function with the given arguments, without waiting for the results.
 
@@ -461,7 +461,7 @@ class Function(
             ...
 
     _experimental_spawn: ___experimental_spawn_spec[
-        modal._functions.ReturnType, modal._functions.P, typing_extensions.Self
+        modal._functions.P, modal._functions.ReturnType, typing_extensions.Self
     ]
 
     class ___spawn_map_inner_spec(typing_extensions.Protocol[P_INNER, SUPERSELF]):
@@ -470,7 +470,7 @@ class Function(
 
     _spawn_map_inner: ___spawn_map_inner_spec[modal._functions.P, typing_extensions.Self]
 
-    class __spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
+    class __spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
         def __call__(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]:
             """Calls the function with the given arguments, without waiting for the results.
 
@@ -491,7 +491,7 @@ class Function(
             """
             ...
 
-    spawn: __spawn_spec[modal._functions.ReturnType, modal._functions.P, typing_extensions.Self]
+    spawn: __spawn_spec[modal._functions.P, modal._functions.ReturnType, typing_extensions.Self]
 
     def get_raw_f(self) -> collections.abc.Callable[..., typing.Any]:
         """Return the inner Python object wrapped by this Modal Function."""

modal/io_streams.py

@@ -605,7 +605,7 @@ class _StreamReader(Generic[T]):
 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:
@@ -627,25 +627,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.")
@@ -653,7 +634,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__}")
@@ -672,19 +653,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()
@@ -713,5 +681,127 @@ class _StreamWriter:
                 raise 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: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()
+        ```
+        """
+        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)
 StreamWriter = synchronize_api(_StreamWriter)

modal/io_streams.pyi

@@ -249,7 +249,7 @@ class _StreamReader(typing.Generic[T]):
         """mdmd:hidden"""
         ...
 
-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: typing.Literal["sandbox", "container_process"], client: modal.client._Client
@@ -258,6 +258,58 @@ class _StreamWriter:
         ...
 
     def _get_next_index(self) -> int: ...
+    def write(self, data: typing.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.
+        """
+        ...
+
+    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.
+        """
+        ...
+
+    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.
+        """
+        ...
+
+class _StreamWriterThroughCommandRouter:
+    def __init__(
+        self,
+        object_id: str,
+        command_router_client: modal._utils.task_command_router_client.TaskCommandRouterClient,
+        task_id: str,
+    ) -> None:
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
+
+    def write(self, data: typing.Union[bytes, bytearray, memoryview, str]) -> None: ...
+    def write_eof(self) -> None: ...
+    async def drain(self) -> None: ...
+
+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: typing.Literal["sandbox", "container_process"],
+        client: modal.client._Client,
+        command_router_client: typing.Optional[modal._utils.task_command_router_client.TaskCommandRouterClient] = None,
+        task_id: typing.Optional[str] = None,
+    ) -> None:
+        """mdmd:hidden"""
+        ...
+
     def write(self, data: typing.Union[bytes, bytearray, memoryview, str]) -> None:
         """Write data to the stream but does not send it immediately.
 
@@ -423,12 +475,16 @@ class StreamReader(typing.Generic[T]):
 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: typing.Literal["sandbox", "container_process"], client: modal.client.Client
+        self,
+        object_id: str,
+        object_type: typing.Literal["sandbox", "container_process"],
+        client: modal.client.Client,
+        command_router_client: typing.Optional[modal._utils.task_command_router_client.TaskCommandRouterClient] = None,
+        task_id: typing.Optional[str] = None,
     ) -> None:
         """mdmd:hidden"""
         ...
 
-    def _get_next_index(self) -> int: ...
     def write(self, data: typing.Union[bytes, bytearray, memoryview, str]) -> None:
         """Write data to the stream but does not send it immediately.
 

modal/sandbox.py

@@ -870,6 +870,7 @@ class _Sandbox(_Object, type_prefix="sb"):
         logger.debug(f"Created ContainerProcess for exec_id {resp.exec_id} on Sandbox {self.object_id}")
         return _ContainerProcess(
             resp.exec_id,
+            task_id,
             self._client,
             stdout=stdout,
             stderr=stderr,

{modal-1.2.1.dev12.dist-info → modal-1.2.1.dev14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modal
-Version: 1.2.1.dev12
+Version: 1.2.1.dev14
 Summary: Python client library for Modal
 Author-email: Modal Labs <support@modal.com>
 License: Apache-2.0

{modal-1.2.1.dev12.dist-info → modal-1.2.1.dev14.dist-info}/RECORD

@@ -24,14 +24,14 @@ modal/app.pyi,sha256=AUV5Rp8qQrZJTP2waoKHFY7rYgsXNMYibMcCAQKuSeo,50544
 modal/billing.py,sha256=zmQ3bcCJlwa4KD1IA_QgdWpm1pn13c-7qfy79iEauYI,195
 modal/call_graph.py,sha256=1g2DGcMIJvRy-xKicuf63IVE98gJSnQsr8R_NVMptNc,2581
 modal/client.py,sha256=kyAIVB3Ay-XKJizQ_1ufUFB__EagV0MLmHJpyYyJ7J0,18636
-modal/client.pyi,sha256=46-sOq2ZFWLaRP_M6tYYyxhIpAggNOg8AD2JAAzevN4,15831
+modal/client.pyi,sha256=2_DiKPDQsg7NASk_5NqJmNb9HYieLChxPmte-GpCPIY,15831
 modal/cloud_bucket_mount.py,sha256=I2GRXYhOWLIz2kJZjXu75jAm9EJkBNcutGc6jR2ReUw,5928
 modal/cloud_bucket_mount.pyi,sha256=VuUOipMIHqFXMkD-3g2bsoqpSxV5qswlFHDOqPQzYAo,7405
 modal/cls.py,sha256=ZxzivE3fNci4-A5uyBYNAzXMXtdqDg3gnYvgbdy5fhg,40384
 modal/cls.pyi,sha256=jJsDPFoqzM4ht-V-e-xEJKJ5TINLF0fYtoBm_UeAW5Y,27281
 modal/config.py,sha256=hpgkgQKbjzo6gVbRzXQrky72_KpdSEm65RNi1M2iNjc,13038
-modal/container_process.py,sha256=DnqlgHiM-7rgVdJNcaXyZlXFD6DFLxEgMSFkieQ6Oj0,7452
-modal/container_process.pyi,sha256=9m-st3hCUlNN1GOTctfPPvIvoLtEl7FbuGWwif5-7YU,6037
+modal/container_process.py,sha256=X5gkIpXeoeaRemrOFHIpI2XYSVq5NmP7iP54cTqvMFE,16643
+modal/container_process.pyi,sha256=raXBD-ImyYUum4UtgEXCgTGuZoWQ7SriDkRuD91Nfuc,8296
 modal/dict.py,sha256=XkaxuojMVtcc4bZvCjJcd6DedU5xxfF8H4w-mDzFPCo,21580
 modal/dict.pyi,sha256=deOiwuwZtwXqedC3h19SwoQIWc4mUnDTBM5XkONt48Y,31712
 modal/environments.py,sha256=xXYDfgzd20CuFdww_zQ53OB0qANQG-j_ls_fT7mGdoQ,6028
@@ -41,12 +41,12 @@ modal/file_io.py,sha256=OSKr77TujcXGJW1iikzYiHckLSmv07QBgBHcxxYEkoI,21456
 modal/file_io.pyi,sha256=xtO6Glf_BFwDE7QiQQo24QqcMf_Vv-iz7WojcGVlLBU,15932
 modal/file_pattern_matcher.py,sha256=A_Kdkej6q7YQyhM_2-BvpFmPqJ0oHb54B6yf9VqvPVE,8116
 modal/functions.py,sha256=kcNHvqeGBxPI7Cgd57NIBBghkfbeFJzXO44WW0jSmao,325
-modal/functions.pyi,sha256=Z6VuukLrjASAgf0kV9I6c09WvP_b2gCujX6f9j2bBaw,37988
+modal/functions.pyi,sha256=CMwApS396tdElFrjnV6RuL2DTCz4C3jYzYoq1y_LPUQ,37988
 modal/gpu.py,sha256=Fe5ORvVPDIstSq1xjmM6OoNgLYFWvogP9r5BgmD3hYg,6769
 modal/image.py,sha256=HDkOnhIAN8g63a8LTN4J5SjC9ciReFQQJIxTS2z5KFM,107216
 modal/image.pyi,sha256=dMvMwAuvWkNN2BRYJFijkEy2m_xtEXgCKK0T7FVldsc,77514
-modal/io_streams.py,sha256=APkECBiP1-D7xeBdFLuWKYOZq6MZFNAhVwuiO1ymcNE,25835
-modal/io_streams.pyi,sha256=dnuWjJrkQsMla_twPTe05NnixsYtT2gjWYjVFnyff5I,15960
+modal/io_streams.py,sha256=Kv-No6WcNBouwdoogwHafOsmPOKqxTpvVGLU0mM6xMc,29564
+modal/io_streams.pyi,sha256=h7qtAbj8LsN-eJKAGjBhnMBegvWprc_0AmwVFi6rj2Y,18084
 modal/mount.py,sha256=G7_xhQMZqokgfsaFLMch0YR3fs-OUNqYUm3f4jHTSMQ,33161
 modal/mount.pyi,sha256=MD_zV2M7eCWxbOpQRjU60aHevN-bmbiywaCX82QoFlw,15380
 modal/network_file_system.py,sha256=ZdEIRgdcR-p_ILyw_AecEtPOhhrSWJeADYCtFnhtaHM,13509
@@ -67,7 +67,7 @@ modal/retries.py,sha256=IvNLDM0f_GLUDD5VgEDoN09C88yoxSrCquinAuxT1Sc,5205
 modal/runner.py,sha256=Ni54hwa42SEBxLPpqFwKMsUPYY8Dv-I-Kz3_jL1StCI,25220
 modal/runner.pyi,sha256=DV3Z7h0owgRyOu9W5KU5O3UbRftX99KGrZQId91fpsU,8671
 modal/running_app.py,sha256=v61mapYNV1-O-Uaho5EfJlryMLvIT9We0amUOSvSGx8,1188
-modal/sandbox.py,sha256=R1WOFDTmg_SiG7NIEbzQMeTYHwLvX9r3j9lpINhIufc,45970
+modal/sandbox.py,sha256=TVENhuPVJXu3RiOKkqSKVWMnHTP7zMUO61StZ2-Ldb8,45991
 modal/sandbox.pyi,sha256=elVE1xEy_ZhD009oNPnCwlZi4tK-RUb1qAoxkVteG9E,50713
 modal/schedule.py,sha256=ng0g0AqNY5GQI9KhkXZQ5Wam5G42glbkqVQsNpBtbDE,3078
 modal/scheduler_placement.py,sha256=BAREdOY5HzHpzSBqt6jDVR6YC_jYfHMVqOzkyqQfngU,1235
@@ -113,7 +113,7 @@ modal/_utils/package_utils.py,sha256=LcL2olGN4xaUzu2Tbv-C-Ft9Qp6bsLxEfETOAVd-mjU
 modal/_utils/pattern_utils.py,sha256=ZUffaECfe2iYBhH6cvCB-0-UWhmEBTZEl_TwG_So3ag,6714
 modal/_utils/rand_pb_testing.py,sha256=mmVPk1rZldHwHZx0DnHTuHQlRLAiiAYdxjwEJpxvT9c,3900
 modal/_utils/shell_utils.py,sha256=hWHzv730Br2Xyj6cGPiMZ-198Z3RZuOu3pDXhFSZ22c,2157
-modal/_utils/task_command_router_client.py,sha256=EABQz9A-WX_O2OnvPWepxQJ1Nu9U8wd8yfj3f6br1hc,22859
+modal/_utils/task_command_router_client.py,sha256=c8WD8DHkQ91lcMb9M_8eBj-NFm6cxeZSv_hgJjRfpJo,23577
 modal/_utils/time_utils.py,sha256=43tpFVwT7ykOjlETIFLVt9auMsRZqYYRYBEKxGCrRSA,1212
 modal/_vendor/__init__.py,sha256=MIEP8jhXUeGq_eCjYFcqN5b1bxBM4fdk0VESpjWR0fc,28
 modal/_vendor/a2wsgi_wsgi.py,sha256=Q1AsjpV_Q_vzQsz_cSqmP9jWzsGsB-ARFU6vpQYml8k,21878
@@ -131,9 +131,9 @@ modal/cli/__init__.py,sha256=6FRleWQxBDT19y7OayO4lBOzuL6Bs9r0rLINYYYbHwQ,769
 modal/cli/_download.py,sha256=_Q_CG1Qb4cjVSAwHGGPOPoNTsmK9gHvF-HNWCNdFjaY,4900
 modal/cli/_traceback.py,sha256=IKj9xtc6LjAxyhGJWolNIXEX3MhAIulnRqywZNOFmkU,7324
 modal/cli/app.py,sha256=rbuAG92my-1eZN0olk6p2eD4oBnyBliUsrCOUW-U-9k,7832
-modal/cli/cluster.py,sha256=MMdFXqyOOeNO7P2VIUN78xKJLvesLIy_28Dt56XEXhM,3307
+modal/cli/cluster.py,sha256=BLcKDpwpDmlqE2UC4V0qNpJKiQ-ZXfI9g_SE7u6vnIU,3347
 modal/cli/config.py,sha256=lhp2Pq4RbTDhaZJ-ZJvhrMqJj8c-WjuRX6gjE3TrvXc,1691
-modal/cli/container.py,sha256=9Ti-TIZ6vjDSmn9mk9h6SRwyhkQjtwirBN18LjpLyvE,3719
+modal/cli/container.py,sha256=IdJ07p7hwFGXz3oYB__NYilucWpVv-K1JL3QKduu7ok,3769
 modal/cli/dict.py,sha256=YAJtiv41YcCd5Fqam3hXCNTs4Y0yOgGR_i6RfQNSAFM,4572
 modal/cli/entry_point.py,sha256=F06p54rPOs1xAUeYW76RaimFOgLW_I17RCvNwfZRqPc,4747
 modal/cli/environment.py,sha256=LGBq8RVQjfBH3EWz8QgmYe19UO66JKSDNxOXMUjw7JM,4285
@@ -156,7 +156,7 @@ modal/experimental/__init__.py,sha256=9gkVuDmu3m4TlKoU3MzEtTOemUSs8EEOWba40s7Aa0
 modal/experimental/flash.py,sha256=C4sef08rARYFllsgtqukFmYL18SZW0_JpMS0BejDcUs,28552
 modal/experimental/flash.pyi,sha256=vV_OQhtdrPn8SW0XrBK-aLLHHIvxAzLzwFbWrke-m74,15463
 modal/experimental/ipython.py,sha256=TrCfmol9LGsRZMeDoeMPx3Hv3BFqQhYnmD_iH0pqdhk,2904
-modal-1.2.1.dev12.dist-info/licenses/LICENSE,sha256=psuoW8kuDP96RQsdhzwOqi6fyWv0ct8CR6Jr7He_P_k,10173
+modal-1.2.1.dev14.dist-info/licenses/LICENSE,sha256=psuoW8kuDP96RQsdhzwOqi6fyWv0ct8CR6Jr7He_P_k,10173
 modal_docs/__init__.py,sha256=svYKtV8HDwDCN86zbdWqyq5T8sMdGDj0PVlzc2tIxDM,28
 modal_docs/gen_cli_docs.py,sha256=c1yfBS_x--gL5bs0N4ihMwqwX8l3IBWSkBAKNNIi6bQ,3801
 modal_docs/gen_reference_docs.py,sha256=d_CQUGQ0rfw28u75I2mov9AlS773z9rG40-yq5o7g2U,6359
@@ -184,10 +184,10 @@ modal_proto/task_command_router_pb2.py,sha256=_pD2ZpU0bNzhwBdzmLoLyLtAtftI_Agxwn
 modal_proto/task_command_router_pb2.pyi,sha256=EyDgXPLr7alqjXYERV8w_MPuO404x0uCppmSkrfE9IE,14589
 modal_proto/task_command_router_pb2_grpc.py,sha256=uEQ0HdrCp8v-9bB5yIic9muA8spCShLHY6Bz9cCgOUE,10114
 modal_proto/task_command_router_pb2_grpc.pyi,sha256=s3Yxsrawdj4nr8vqQqsAxyX6ilWaGbdECy425KKbLIA,3301
-modal_version/__init__.py,sha256=OYOaillNG-YG3dph7s36QFNeIzc-mCx6s7KPhjISWzY,121
+modal_version/__init__.py,sha256=743KmmOlJb-M5EkswhYesVjee8eiLrK4PizfYr2lZCM,121
 modal_version/__main__.py,sha256=2FO0yYQQwDTh6udt1h-cBnGd1c4ZyHnHSI4BksxzVac,105
-modal-1.2.1.dev12.dist-info/METADATA,sha256=aGlYkorlrlPY8h1THL2Ece76Z2PspqkANMRmy_F8xQk,2484
-modal-1.2.1.dev12.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
-modal-1.2.1.dev12.dist-info/entry_points.txt,sha256=An-wYgeEUnm6xzrAP9_NTSTSciYvvEWsMZILtYrvpAI,46
-modal-1.2.1.dev12.dist-info/top_level.txt,sha256=4BWzoKYREKUZ5iyPzZpjqx4G8uB5TWxXPDwibLcVa7k,43
-modal-1.2.1.dev12.dist-info/RECORD,,
+modal-1.2.1.dev14.dist-info/METADATA,sha256=kp21p_YTNC2q_oQRDT9xdagRNkhSrnlIfh-uQayaTqo,2484
+modal-1.2.1.dev14.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
+modal-1.2.1.dev14.dist-info/entry_points.txt,sha256=An-wYgeEUnm6xzrAP9_NTSTSciYvvEWsMZILtYrvpAI,46
+modal-1.2.1.dev14.dist-info/top_level.txt,sha256=4BWzoKYREKUZ5iyPzZpjqx4G8uB5TWxXPDwibLcVa7k,43
+modal-1.2.1.dev14.dist-info/RECORD,,

modal_version/__init__.py

@@ -1,4 +1,4 @@
 # Copyright Modal Labs 2025
 """Supplies the current version of the modal client library."""
 
-__version__ = "1.2.1.dev12"
+__version__ = "1.2.1.dev14"