modal 1.4.3.dev23__tar.gz → 1.4.3.dev25__tar.gz

{modal-1.4.3.dev23 → modal-1.4.3.dev25}/PKG-INFO

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

{modal-1.4.3.dev23 → modal-1.4.3.dev25}/modal/_functions.py

@@ -549,6 +549,7 @@ class FunctionStats:
     backlog: int
     num_total_runners: int
     num_running_inputs: int
+    input_headroom: int
 
 
 def _parse_retries(
@@ -1851,7 +1852,10 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
             retry=Retry(total_timeout=10.0),
         )
         return FunctionStats(
-            backlog=resp.backlog, num_total_runners=resp.num_total_tasks, num_running_inputs=resp.num_running_inputs
+            backlog=resp.backlog,
+            num_total_runners=resp.num_total_tasks,
+            num_running_inputs=resp.num_running_inputs,
+            input_headroom=resp.input_headroom,
         )
 
     @live_method

{modal-1.4.3.dev23 → modal-1.4.3.dev25}/modal/_utils/grpc_utils.py

@@ -403,7 +403,7 @@ async def _retry_transient_errors(
                 and isinstance(exc, GRPCError)
                 and (server_retry_policy := get_server_retry_policy(exc))
             ):
-                server_delay = server_retry_policy.retry_after_secs
+                server_delay = max(server_retry_policy.retry_after_secs, 0.1)
 
                 now = time.time()
 
@@ -417,6 +417,7 @@ async def _retry_transient_errors(
                 )
                 final_attempt = total_timeout_will_be_reached or max_throttle_will_be_reached
 
+                elapsed_time = time.monotonic() - attempt_started_at
                 with suppress_tb_frame():
                     process_exception_before_retry(
                         exc,
@@ -425,20 +426,23 @@ async def _retry_transient_errors(
                         n_retries,
                         server_delay,
                         idempotency_key,
-                        time.monotonic() - attempt_started_at,
+                        elapsed_time,
                     )
 
-                now = time.time()
                 if last_server_retry_warning_time is None or (
                     now - last_server_retry_warning_time >= SERVER_RETRY_WARNING_TIME_INTERVAL
                 ):
                     last_server_retry_warning_time = now
                     logger.warning(
-                        f"Warning: Received {exc.status}{os.linesep}"
-                        f"{exc.message}{os.linesep}"
-                        f"Will retry in {server_delay:0.2f} seconds."
+                        f"Warning: Received {exc.status}{os.linesep}{exc.message}{os.linesep}"
+                        f"for {fn.name} ({idempotency_key[:8]}). Retrying..."
                     )
 
+                logger.debug(
+                    f"Server-side retry triggered with {repr(exc)} {n_throttled_retries=} {server_delay=} "
+                    f"{elapsed_time=} for {fn.name} ({idempotency_key[:8]})."
+                )
+
                 n_throttled_retries += 1
                 await asyncio.sleep(server_delay)
                 continue

{modal-1.4.3.dev23 → modal-1.4.3.dev25}/modal/_utils/task_command_router_client.py

@@ -16,7 +16,7 @@ from grpclib import GRPCError, Status
 from grpclib.exceptions import StreamTerminatedError
 
 from modal.config import logger
-from modal.exception import ConflictError, ExecTimeoutError
+from modal.exception import ExecTimeoutError
 from modal_proto import api_pb2, task_command_router_pb2 as sr_pb2
 from modal_proto.task_command_router_grpc import TaskCommandRouterStub
 
@@ -196,21 +196,12 @@ class TaskCommandRouterClient:
         return cls(server_client, task_id, url, jwt, channel, loop, jwt_refresh_lock, sandbox_id=sandbox_id)
 
     @classmethod
-    async def try_init(
+    async def init(
         cls,
         server_client,
         task_id: str,
-    ) -> Optional["TaskCommandRouterClient"]:
-        """Attempt to initialize a TaskCommandRouterClient by fetching direct access.
-
-        Returns None if command router access is not enabled (FAILED_PRECONDITION).
-        """
-        try:
-            resp = await fetch_command_router_access(server_client, task_id)
-        except ConflictError:
-            logger.debug(f"Command router access is not enabled for task {task_id}")
-            return None
-
+    ) -> "TaskCommandRouterClient":
+        resp = await fetch_command_router_access(server_client, task_id)
         logger.debug(f"Using command router access for task {task_id}")
         return await cls._connect(server_client, task_id, resp.url, resp.jwt)
 
@@ -241,7 +232,7 @@ class TaskCommandRouterClient:
         stream_stdio_retry_delay_factor: float = 2,
         stream_stdio_max_retries: int = 10,
     ) -> None:
-        """Callers should not use this directly. Use TaskCommandRouterClient.try_init() instead."""
+        """Callers should not use this directly. Use TaskCommandRouterClient.init() instead."""
         # Record the loop this instance is bound to so __del__ can safely schedule cleanup
         # even if finalization happens from a different thread (e.g. via synchronicity).
         self._loop = loop

{modal-1.4.3.dev23 → modal-1.4.3.dev25}/modal/cli/cluster.py

@@ -16,7 +16,6 @@ from modal.cli.utils import display_table, env_option, is_tty
 from modal.client import _Client
 from modal.config import config
 from modal.container_process import _ContainerProcess
-from modal.exception import InvalidError
 from modal.output import OutputManager
 from modal.stream_type import StreamType
 from modal_proto import api_pb2, task_command_router_pb2 as sr_pb2
@@ -79,9 +78,7 @@ async def shell(cluster_id: str, rank: int = 0):
 
     pty = is_tty()
 
-    command_router_client = await TaskCommandRouterClient.try_init(client, task_id)
-    if command_router_client is None:
-        raise InvalidError(f"Command router access is not available for container {task_id}")
+    command_router_client = await TaskCommandRouterClient.init(client, task_id)
 
     process_id = str(uuid.uuid4())
 

{modal-1.4.3.dev23 → modal-1.4.3.dev25}/modal/cli/container.py

@@ -278,9 +278,7 @@ async def _exec_impl(
 
     client = await _Client.from_env()
 
-    command_router_client = await TaskCommandRouterClient.try_init(client, container_id)
-    if command_router_client is None:
-        raise InvalidError(f"Command router access is not available for container {container_id}")
+    command_router_client = await TaskCommandRouterClient.init(client, container_id)
 
     process_id = str(uuid.uuid4())
 

{modal-1.4.3.dev23 → modal-1.4.3.dev25}/modal/client.pyi

@@ -35,7 +35,7 @@ class _Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.4.3.dev23",
+        version: str = "1.4.3.dev25",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.
@@ -175,7 +175,7 @@ class Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.4.3.dev23",
+        version: str = "1.4.3.dev25",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.

modal-1.4.3.dev25/modal/container_process.py

@@ -0,0 +1,224 @@
+# Copyright Modal Labs 2024
+import asyncio
+import platform
+from typing import Generic, Optional, TypeVar
+
+from modal_proto import api_pb2
+
+from ._utils.async_utils import TaskContext, synchronize_api
+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,
+    _StreamReaderThroughCommandRouterParams,
+    _StreamWriter,
+    _StreamWriterThroughCommandRouterParams,
+)
+from .stream_type import StreamType
+
+T = TypeVar("T", str, bytes)
+
+
+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 _ContainerProcess(Generic[T]):
+    """Represents a running process in a container.
+
+    Container processes communicate via direct communication with
+    the Modal worker where the container is running.
+    """
+
+    def __init__(
+        self,
+        process_id: str,
+        task_id: str,
+        client: _Client,
+        command_router_client: TaskCommandRouterClient,
+        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](
+            _StreamReaderThroughCommandRouterParams(
+                file_descriptor=api_pb2.FILE_DESCRIPTOR_STDOUT,
+                task_id=self._task_id,
+                object_id=process_id,
+                command_router_client=self._command_router_client,
+                deadline=exec_deadline,
+            ),
+            stream_type=stdout,
+            text=text,
+            by_line=by_line,
+        )
+        self._stderr = _StreamReader[T](
+            _StreamReaderThroughCommandRouterParams(
+                file_descriptor=api_pb2.FILE_DESCRIPTOR_STDERR,
+                task_id=self._task_id,
+                object_id=process_id,
+                command_router_client=self._command_router_client,
+                deadline=exec_deadline,
+            ),
+            stream_type=stderr,
+            text=text,
+            by_line=by_line,
+        )
+        self._stdin = _StreamWriter(
+            _StreamWriterThroughCommandRouterParams(
+                task_id=self._task_id,
+                object_id=process_id,
+                command_router_client=self._command_router_client,
+            )
+        )
+        self._returncode = None
+
+    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
+        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:
+        """Wait for the container process to finish running. Returns the exit code."""
+        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):
+        """mdmd:hidden"""
+        if platform.system() == "Windows":
+            print("interactive exec is not currently supported on Windows.")  # noqa: T201
+            return
+
+        from .output import OutputManager
+
+        output = OutputManager.get()
+        connecting_status = output.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.")
+
+
+ContainerProcess = synchronize_api(_ContainerProcess)

{modal-1.4.3.dev23 → modal-1.4.3.dev25}/modal/container_process.pyi

@@ -7,140 +7,27 @@ import typing_extensions
 
 T = typing.TypeVar("T")
 
-class _ContainerProcessThroughServer(typing.Generic[T]):
-    """Abstract base class for generic types.
-
-    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::
-
-      class Mapping(Generic[KT, VT]):
-          def __getitem__(self, key: KT) -> VT:
-              ...
-          # Etc.
-
-    This class can then be used as follows::
-
-      def lookup_name(mapping: Mapping[KT, VT], key: KT, default: VT) -> VT:
-          try:
-              return mapping[key]
-          except KeyError:
-              return default
-    """
-
-    _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]
-
-    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,
-    ) -> None:
-        """Initialize self.  See help(type(self)) for accurate signature."""
-        ...
-
-    def __repr__(self) -> str:
-        """Return repr(self)."""
-        ...
-
-    @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_for_completion(self) -> int: ...
-    async def wait(self) -> int:
-        """Wait for the container process to finish running. Returns the exit code."""
-        ...
-
-    async def attach(self):
-        """mdmd:hidden"""
-        ...
-
 def _iter_stream_as_bytes(stream: modal.io_streams._StreamReader[T]):
     """Yield raw bytes from a StreamReader regardless of text mode/backend."""
     ...
 
-class _ContainerProcessThroughCommandRouter(typing.Generic[T]):
-    """Container process implementation that works via direct communication with
+class _ContainerProcess(typing.Generic[T]):
+    """Represents a running process in a container.
+
+    Container processes communicate 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."""
-        ...
-
-    def __repr__(self) -> str:
-        """Return repr(self)."""
-        ...
-
-    @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 _ContainerProcess(typing.Generic[T]):
-    """Represents a running process in a container."""
     def __init__(
         self,
         process_id: str,
         task_id: str,
         client: modal.client._Client,
+        command_router_client: modal._utils.task_command_router_client.TaskCommandRouterClient,
         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."""
         ...
@@ -182,18 +69,22 @@ class _ContainerProcess(typing.Generic[T]):
         ...
 
 class ContainerProcess(typing.Generic[T]):
-    """Represents a running process in a container."""
+    """Represents a running process in a container.
+
+    Container processes communicate via direct communication with
+    the Modal worker where the container is running.
+    """
     def __init__(
         self,
         process_id: str,
         task_id: str,
         client: modal.client.Client,
+        command_router_client: modal._utils.task_command_router_client.TaskCommandRouterClient,
         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