modal 1.2.1.dev13__tar.gz → 1.2.1.dev14__tar.gz

{modal-1.2.1.dev13 → modal-1.2.1.dev14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modal
-Version: 1.2.1.dev13
+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.dev13 → modal-1.2.1.dev14}/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-1.2.1.dev13 → modal-1.2.1.dev14}/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-1.2.1.dev13 → modal-1.2.1.dev14}/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-1.2.1.dev13 → modal-1.2.1.dev14}/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.dev13",
+        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.dev13",
+        version: str = "1.2.1.dev14",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.

modal-1.2.1.dev14/modal/container_process.py

@@ -0,0 +1,470 @@
+# Copyright Modal Labs 2024
+import asyncio
+import platform
+import time
+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 .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 _ContainerProcessThroughServer(Generic[T]):
+    _process_id: Optional[str] = None
+    _stdout: _StreamReader[T]
+    _stderr: _StreamReader[T]
+    _stdin: _StreamWriter
+    _exec_deadline: Optional[float] = None
+    _text: bool
+    _by_line: bool
+    _returncode: Optional[int] = None
+
+    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,
+    ) -> None:
+        self._process_id = process_id
+        self._client = client
+        self._exec_deadline = exec_deadline
+        self._text = text
+        self._by_line = by_line
+        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,
+            task_id=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,
+            task_id=task_id,
+        )
+        self._stdin = _StreamWriter(process_id, "container_process", self._client)
+
+    def __repr__(self) -> str:
+        return f"ContainerProcess(process_id={self._process_id!r})"
+
+    @property
+    def stdout(self) -> _StreamReader[T]:
+        """StreamReader for the container process's stdout stream."""
+        return self._stdout
+
+    @property
+    def stderr(self) -> _StreamReader[T]:
+        """StreamReader for the container process's stderr stream."""
+        return self._stderr
+
+    @property
+    def stdin(self) -> _StreamWriter:
+        """StreamWriter for the container process's stdin stream."""
+        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]:
+        """Check if the container process has finished running.
+
+        Returns `None` if the process is still running, else returns the exit code.
+        """
+        if self._returncode is not None:
+            return self._returncode
+        if self._exec_deadline and time.monotonic() >= self._exec_deadline:
+            # TODO(matt): In the future, it would be nice to raise a ContainerExecTimeoutError to make it
+            # clear to the user that their sandbox terminated due to a timeout
+            self._returncode = -1
+            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)
+
+        if resp.completed:
+            self._returncode = resp.exit_code
+            return self._returncode
+
+        return None
+
+    async def _wait_for_completion(self) -> int:
+        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
+            )
+            if resp.completed:
+                return resp.exit_code
+
+    async def wait(self) -> int:
+        """Wait for the container process to finish running. Returns the exit code."""
+        if self._returncode is not None:
+            return self._returncode
+
+        try:
+            timeout = None
+            if self._exec_deadline:
+                timeout = self._exec_deadline - time.monotonic()
+                if timeout <= 0:
+                    raise TimeoutError()
+            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):
+        """mdmd:hidden"""
+        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):
+            # 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_impl._get_logs(skip_empty_messages=False):
+                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 to the server 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
+
+                # TODO: this doesn't work right now.
+                # if exit_status != 0:
+                #     raise ExecutionError(f"Process exited with status code {exit_status}")
+
+            except (asyncio.TimeoutError, TimeoutError):
+                connecting_status.stop()
+                stdout_task.cancel()
+                stderr_task.cancel()
+                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)