modAL 1.4.3.dev23__py3-none-any.whl → 1.4.3.dev24__py3-none-any.whl

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/_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/_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/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/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/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.dev24",
     ):
         """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.dev24",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.

modal/container_process.py

@@ -1,7 +1,6 @@
 # Copyright Modal Labs 2024
 import asyncio
 import platform
-import time
 from typing import Generic, Optional, TypeVar
 
 from modal_proto import api_pb2
@@ -12,194 +11,17 @@ 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 .io_streams import (
+    _StreamReader,
+    _StreamReaderThroughCommandRouterParams,
+    _StreamWriter,
+    _StreamWriterThroughCommandRouterParams,
+)
 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.
-        """
-        assert self._process_id
-        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 = await 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:
-        assert self._process_id
-        while True:
-            req = api_pb2.ContainerExecWaitRequest(exec_id=self._process_id, timeout=10)
-            resp = await 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.")  # 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):
-            # 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 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:
@@ -209,19 +31,19 @@ async def _iter_stream_as_bytes(stream: _StreamReader[T]):
             yield part
 
 
-class _ContainerProcessThroughCommandRouter(Generic[T]):
-    """
-    Container process implementation that works via direct communication with
+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,
-        task_id: str,
-        *,
         stdout: StreamType = StreamType.PIPE,
         stderr: StreamType = StreamType.PIPE,
         exec_deadline: Optional[float] = None,
@@ -236,35 +58,35 @@ class _ContainerProcessThroughCommandRouter(Generic[T]):
         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,
+            _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,
-            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,
+            _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,
-            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,
+            _StreamWriterThroughCommandRouterParams(
+                task_id=self._task_id,
+                object_id=process_id,
+                command_router_client=self._command_router_client,
+            )
         )
         self._returncode = None
 
@@ -273,14 +95,17 @@ class _ContainerProcessThroughCommandRouter(Generic[T]):
 
     @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
@@ -293,6 +118,10 @@ class _ContainerProcessThroughCommandRouter(Generic[T]):
         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:
@@ -322,6 +151,7 @@ class _ContainerProcessThroughCommandRouter(Generic[T]):
             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
 
@@ -345,6 +175,7 @@ class _ContainerProcessThroughCommandRouter(Generic[T]):
         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
@@ -390,81 +221,4 @@ class _ContainerProcessThroughCommandRouter(Generic[T]):
                 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)