modal 1.2.1.dev13__py3-none-any.whl → 1.2.1.dev15__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.dev13",
+        version: str = "1.2.1.dev15",
     ):
         """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.dev15",
     ):
         """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/experimental/flash.py

@@ -321,7 +321,7 @@ class _FlashPrometheusAutoscaler:
 
     async def _compute_target_containers(self, current_replicas: int) -> int:
         """
-        Gets internal metrics from container to autoscale up or down.
+        Gets metrics from container to autoscale up or down.
         """
         containers = await self._get_all_containers()
         if len(containers) > current_replicas:
@@ -334,7 +334,7 @@ class _FlashPrometheusAutoscaler:
         if current_replicas == 0:
             return 1
 
-        # Get metrics based on autoscaler type (prometheus or internal)
+        # Get metrics based on autoscaler type
         sum_metric, n_containers_with_metrics = await self._get_scaling_info(containers)
 
         desired_replicas = self._calculate_desired_replicas(
@@ -406,39 +406,26 @@ class _FlashPrometheusAutoscaler:
         return desired_replicas
 
     async def _get_scaling_info(self, containers) -> tuple[float, int]:
-        """Get metrics using either internal container metrics API or prometheus HTTP endpoints."""
-        if self.metrics_endpoint == "internal":
-            container_metrics_results = await asyncio.gather(
-                *[self._get_container_metrics(container.task_id) for container in containers]
-            )
-            container_metrics_list = []
-            for container_metric in container_metrics_results:
-                if container_metric is None:
-                    continue
-                container_metrics_list.append(getattr(container_metric.metrics, self.target_metric))
-
-            sum_metric = sum(container_metrics_list)
-            n_containers_with_metrics = len(container_metrics_list)
-        else:
-            sum_metric = 0
-            n_containers_with_metrics = 0
-
-            container_metrics_list = await asyncio.gather(
-                *[
-                    self._get_metrics(f"https://{container.host}:{container.port}/{self.metrics_endpoint}")
-                    for container in containers
-                ]
-            )
+        """Get metrics using container exposed metrics endpoints."""
+        sum_metric = 0
+        n_containers_with_metrics = 0
+
+        container_metrics_list = await asyncio.gather(
+            *[
+                self._get_metrics(f"https://{container.host}:{container.port}/{self.metrics_endpoint}")
+                for container in containers
+            ]
+        )
 
-            for container_metrics in container_metrics_list:
-                if (
-                    container_metrics is None
-                    or self.target_metric not in container_metrics
-                    or len(container_metrics[self.target_metric]) == 0
-                ):
-                    continue
-                sum_metric += container_metrics[self.target_metric][0].value
-                n_containers_with_metrics += 1
+        for container_metrics in container_metrics_list:
+            if (
+                container_metrics is None
+                or self.target_metric not in container_metrics
+                or len(container_metrics[self.target_metric]) == 0
+            ):
+                continue
+            sum_metric += container_metrics[self.target_metric][0].value
+            n_containers_with_metrics += 1
 
         return sum_metric, n_containers_with_metrics
 
@@ -474,15 +461,6 @@ class _FlashPrometheusAutoscaler:
 
         return metrics
 
-    async def _get_container_metrics(self, container_id: str) -> Optional[api_pb2.TaskGetAutoscalingMetricsResponse]:
-        req = api_pb2.TaskGetAutoscalingMetricsRequest(task_id=container_id)
-        try:
-            resp = await retry_transient_errors(self.client.stub.TaskGetAutoscalingMetrics, req)
-            return resp
-        except Exception as e:
-            logger.warning(f"[Modal Flash] Error getting metrics for container {container_id}: {e}")
-            return None
-
     async def _get_all_containers(self):
         req = api_pb2.FlashContainerListRequest(function_id=self.fn.object_id)
         resp = await retry_transient_errors(self.client.stub.FlashContainerList, req)
@@ -572,14 +550,10 @@ async def flash_prometheus_autoscaler(
     app_name: str,
     cls_name: str,
     # Endpoint to fetch metrics from. Must be in Prometheus format. Example: "/metrics"
-    # If metrics_endpoint is "internal", we will use containers' internal metrics to autoscale instead.
     metrics_endpoint: str,
     # Target metric to autoscale on. Example: "vllm:num_requests_running"
-    # If metrics_endpoint is "internal", target_metrics options are: [cpu_usage_percent, memory_usage_percent]
     target_metric: str,
     # Target metric value. Example: 25
-    # If metrics_endpoint is "internal", target_metric_value is a percentage value between 0.1 and 1.0 (inclusive),
-    # indicating container's usage of that metric.
     target_metric_value: float,
     min_containers: Optional[int] = None,
     max_containers: Optional[int] = None,