modal 1.0.6.dev58__py3-none-any.whl → 1.2.3.dev7__py3-none-any.whl

modal/container_process.py

@@ -7,17 +7,18 @@ from typing import Generic, Optional, TypeVar
 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 .exception import InteractiveTimeoutError, InvalidError
+from .config import logger
+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]
@@ -30,6 +31,7 @@ class _ContainerProcess(Generic[T]):
     def __init__(
         self,
         process_id: str,
+        task_id: str,
         client: _Client,
         stdout: StreamType = StreamType.PIPE,
         stderr: StreamType = StreamType.PIPE,
@@ -51,6 +53,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,
@@ -61,6 +64,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)
 
@@ -96,6 +100,7 @@ class _ContainerProcess(Generic[T]):
 
         Returns `None` if the process is still running, else returns the exit code.
         """
+        assert self._process_id
         if self._returncode is not None:
             return self._returncode
         if self._exec_deadline and time.monotonic() >= self._exec_deadline:
@@ -105,7 +110,7 @@ class _ContainerProcess(Generic[T]):
             return self._returncode
 
         req = api_pb2.ContainerExecWaitRequest(exec_id=self._process_id, timeout=0)
-        resp: api_pb2.ContainerExecWaitResponse = await retry_transient_errors(self._client.stub.ContainerExecWait, req)
+        resp = await self._client.stub.ContainerExecWait(req)
 
         if resp.completed:
             self._returncode = resp.exit_code
@@ -114,11 +119,10 @@ class _ContainerProcess(Generic[T]):
         return None
 
     async def _wait_for_completion(self) -> int:
+        assert self._process_id
         while True:
             req = api_pb2.ContainerExecWaitRequest(exec_id=self._process_id, timeout=10)
-            resp: api_pb2.ContainerExecWaitResponse = await retry_transient_errors(
-                self._client.stub.ContainerExecWait, req
-            )
+            resp = await self._client.stub.ContainerExecWait(req)
             if resp.completed:
                 return resp.exit_code
 
@@ -136,6 +140,7 @@ class _ContainerProcess(Generic[T]):
             self._returncode = await asyncio.wait_for(self._wait_for_completion(), timeout=timeout)
         except (asyncio.TimeoutError, TimeoutError):
             self._returncode = -1
+        logger.debug(f"ContainerProcess {self._process_id} wait completed with returncode {self._returncode}")
         return self._returncode
 
     async def attach(self):
@@ -144,20 +149,25 @@ class _ContainerProcess(Generic[T]):
             print("interactive exec is not currently supported on Windows.")
             return
 
-        from rich.console import Console
+        from ._output import make_console
 
-        console = Console()
+        console = make_console()
 
         connecting_status = console.status("Connecting...")
         connecting_status.start()
         on_connect = asyncio.Event()
 
         async def _write_to_fd_loop(stream: _StreamReader):
+            # This is required to make modal shell to an existing task work,
+            # since that uses ContainerExec RPCs directly, but this is hacky.
+            #
+            # TODO(saltzm): Once we use the new exec path for that use case, this code can all be removed.
+            from .io_streams import _StreamReaderThroughServer
+
+            assert isinstance(stream._impl, _StreamReaderThroughServer)
+            stream_impl = stream._impl
             # Don't skip empty messages so we can detect when the process has booted.
-            async for chunk in stream._get_logs(skip_empty_messages=False):
-                if chunk is None:
-                    break
-
+            async for chunk in stream_impl._get_logs(skip_empty_messages=False):
                 if not on_connect.is_set():
                     connecting_status.stop()
                     on_connect.set()
@@ -191,4 +201,272 @@ 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
+
+    def __repr__(self) -> str:
+        return f"ContainerProcess(process_id={self._process_id!r})"
+
+    @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")
+            # TODO(saltzm): This is a weird API, but customers currently may rely on it. This
+            # should probably raise an ExecTimeoutError instead.
+            self._returncode = -1
+            return self._returncode
+        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,114 @@ 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::
+    def __repr__(self) -> str:
+        """Return repr(self)."""
+        ...
 
-      class Mapping(Generic[KT, VT]):
-          def __getitem__(self, key: KT) -> VT:
-              ...
-          # Etc.
+    @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): ...
 
-    This class can then be used as follows::
+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."""
+        ...
 
-      def lookup_name(mapping: Mapping[KT, VT], key: KT, default: VT) -> VT:
-          try:
-              return mapping[key]
-          except KeyError:
-              return default
-    """
+    def __repr__(self) -> str:
+        """Return repr(self)."""
+        ...
 
-    _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 stdout(self) -> modal.io_streams._StreamReader[T]:
+        """StreamReader for the container process's stdout stream."""
+        ...
+
+    @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 +233,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."""