modal 1.2.0__py3-none-any.whl → 1.2.1__py3-none-any.whl

modal/runner.pyi

@@ -1,6 +1,5 @@
-import modal._functions
+import modal.app
 import modal.client
-import modal.cls
 import modal.running_app
 import modal_proto.api_pb2
 import multiprocessing.synchronize
@@ -8,8 +7,6 @@ import synchronicity.combined_types
 import typing
 import typing_extensions
 
-_App = typing.TypeVar("_App")
-
 V = typing.TypeVar("V")
 
 async def _heartbeat(client: modal.client._Client, app_id: str) -> None: ...
@@ -29,8 +26,7 @@ async def _init_local_app_from_name(
 async def _create_all_objects(
     client: modal.client._Client,
     running_app: modal.running_app.RunningApp,
-    functions: dict[str, modal._functions._Function],
-    classes: dict[str, modal.cls._Cls],
+    local_app_state: modal.app._LocalAppState,
     environment_name: str,
 ) -> None:
     """Create objects that have been defined but not created on the server."""
@@ -40,10 +36,8 @@ async def _publish_app(
     client: modal.client._Client,
     running_app: modal.running_app.RunningApp,
     app_state: int,
-    functions: dict[str, modal._functions._Function],
-    classes: dict[str, modal.cls._Cls],
+    app_local_state: modal.app._LocalAppState,
     name: str = "",
-    tags: dict[str, str] = {},
     deployment_tag: str = "",
     commit_info: typing.Optional[modal_proto.api_pb2.CommitInfo] = None,
 ) -> tuple[str, list[modal_proto.api_pb2.Warning]]:
@@ -66,18 +60,18 @@ async def _status_based_disconnect(
     ...
 
 def _run_app(
-    app: _App,
+    app: modal.app._App,
     *,
     client: typing.Optional[modal.client._Client] = None,
     detach: bool = False,
     environment_name: typing.Optional[str] = None,
     interactive: bool = False,
-) -> typing.AsyncContextManager[_App]:
+) -> typing.AsyncContextManager[modal.app._App]:
     """mdmd:hidden"""
     ...
 
 async def _serve_update(
-    app: _App, existing_app_id: str, is_ready: multiprocessing.synchronize.Event, environment_name: str
+    app: modal.app._App, existing_app_id: str, is_ready: multiprocessing.synchronize.Event, environment_name: str
 ) -> None:
     """mdmd:hidden"""
     ...
@@ -115,7 +109,7 @@ class DeployResult:
         ...
 
 async def _deploy_app(
-    app: _App,
+    app: modal.app._App,
     name: typing.Optional[str] = None,
     namespace: typing.Any = None,
     client: typing.Optional[modal.client._Client] = None,
@@ -129,7 +123,7 @@ async def _deploy_app(
     ...
 
 async def _interactive_shell(
-    _app: _App, cmds: list[str], environment_name: str = "", pty: bool = True, **kwargs: typing.Any
+    _app: modal.app._App, cmds: list[str], environment_name: str = "", pty: bool = True, **kwargs: typing.Any
 ) -> None:
     """Run an interactive shell (like `bash`) within the image for this app.
 
@@ -159,26 +153,26 @@ class __run_app_spec(typing_extensions.Protocol):
     def __call__(
         self,
         /,
-        app: _App,
+        app: modal.app.App,
         *,
         client: typing.Optional[modal.client.Client] = None,
         detach: bool = False,
         environment_name: typing.Optional[str] = None,
         interactive: bool = False,
-    ) -> synchronicity.combined_types.AsyncAndBlockingContextManager[_App]:
+    ) -> synchronicity.combined_types.AsyncAndBlockingContextManager[modal.app.App]:
         """mdmd:hidden"""
         ...
 
     def aio(
         self,
         /,
-        app: _App,
+        app: modal.app.App,
         *,
         client: typing.Optional[modal.client.Client] = None,
         detach: bool = False,
         environment_name: typing.Optional[str] = None,
         interactive: bool = False,
-    ) -> typing.AsyncContextManager[_App]:
+    ) -> typing.AsyncContextManager[modal.app.App]:
         """mdmd:hidden"""
         ...
 
@@ -186,13 +180,23 @@ run_app: __run_app_spec
 
 class __serve_update_spec(typing_extensions.Protocol):
     def __call__(
-        self, /, app: _App, existing_app_id: str, is_ready: multiprocessing.synchronize.Event, environment_name: str
+        self,
+        /,
+        app: modal.app.App,
+        existing_app_id: str,
+        is_ready: multiprocessing.synchronize.Event,
+        environment_name: str,
     ) -> None:
         """mdmd:hidden"""
         ...
 
     async def aio(
-        self, /, app: _App, existing_app_id: str, is_ready: multiprocessing.synchronize.Event, environment_name: str
+        self,
+        /,
+        app: modal.app.App,
+        existing_app_id: str,
+        is_ready: multiprocessing.synchronize.Event,
+        environment_name: str,
     ) -> None:
         """mdmd:hidden"""
         ...
@@ -203,7 +207,7 @@ class __deploy_app_spec(typing_extensions.Protocol):
     def __call__(
         self,
         /,
-        app: _App,
+        app: modal.app.App,
         name: typing.Optional[str] = None,
         namespace: typing.Any = None,
         client: typing.Optional[modal.client.Client] = None,
@@ -219,7 +223,7 @@ class __deploy_app_spec(typing_extensions.Protocol):
     async def aio(
         self,
         /,
-        app: _App,
+        app: modal.app.App,
         name: typing.Optional[str] = None,
         namespace: typing.Any = None,
         client: typing.Optional[modal.client.Client] = None,
@@ -236,7 +240,13 @@ deploy_app: __deploy_app_spec
 
 class __interactive_shell_spec(typing_extensions.Protocol):
     def __call__(
-        self, /, _app: _App, cmds: list[str], environment_name: str = "", pty: bool = True, **kwargs: typing.Any
+        self,
+        /,
+        _app: modal.app.App,
+        cmds: list[str],
+        environment_name: str = "",
+        pty: bool = True,
+        **kwargs: typing.Any,
     ) -> None:
         """Run an interactive shell (like `bash`) within the image for this app.
 
@@ -263,7 +273,13 @@ class __interactive_shell_spec(typing_extensions.Protocol):
         ...
 
     async def aio(
-        self, /, _app: _App, cmds: list[str], environment_name: str = "", pty: bool = True, **kwargs: typing.Any
+        self,
+        /,
+        _app: modal.app.App,
+        cmds: list[str],
+        environment_name: str = "",
+        pty: bool = True,
+        **kwargs: typing.Any,
     ) -> None:
         """Run an interactive shell (like `bash`) within the image for this app.
 

modal/sandbox.py

@@ -3,6 +3,7 @@ import asyncio
 import json
 import os
 import time
+import uuid
 from collections.abc import AsyncGenerator, Collection, Sequence
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, AsyncIterator, Literal, Optional, Union, overload
@@ -20,7 +21,7 @@ from modal._tunnel import Tunnel
 from modal.cloud_bucket_mount import _CloudBucketMount, cloud_bucket_mounts_to_proto
 from modal.mount import _Mount
 from modal.volume import _Volume
-from modal_proto import api_pb2
+from modal_proto import api_pb2, task_command_router_pb2 as sr_pb2
 
 from ._object import _get_environment_name, _Object
 from ._resolver import Resolver
@@ -30,6 +31,7 @@ from ._utils.deprecation import deprecation_warning
 from ._utils.grpc_utils import retry_transient_errors
 from ._utils.mount_utils import validate_network_file_systems, validate_volumes
 from ._utils.name_utils import is_valid_object_name
+from ._utils.task_command_router_client import TaskCommandRouterClient
 from .client import _Client
 from .container_process import _ContainerProcess
 from .exception import AlreadyExistsError, ExecutionError, InvalidError, SandboxTerminatedError, SandboxTimeoutError
@@ -121,9 +123,10 @@ class _Sandbox(_Object, type_prefix="sb"):
     _stdout: _StreamReader[str]
     _stderr: _StreamReader[str]
     _stdin: _StreamWriter
-    _task_id: Optional[str] = None
-    _tunnels: Optional[dict[int, Tunnel]] = None
-    _enable_snapshot: bool = False
+    _task_id: Optional[str]
+    _tunnels: Optional[dict[int, Tunnel]]
+    _enable_snapshot: bool
+    _command_router_client: Optional[TaskCommandRouterClient]
 
     @staticmethod
     def _default_pty_info() -> api_pb2.PTYInfo:
@@ -513,14 +516,18 @@ class _Sandbox(_Object, type_prefix="sb"):
         return obj
 
     def _hydrate_metadata(self, handle_metadata: Optional[Message]):
-        self._stdout: _StreamReader[str] = StreamReader[str](
+        self._stdout = StreamReader(
             api_pb2.FILE_DESCRIPTOR_STDOUT, self.object_id, "sandbox", self._client, by_line=True
         )
-        self._stderr: _StreamReader[str] = StreamReader[str](
+        self._stderr = StreamReader(
             api_pb2.FILE_DESCRIPTOR_STDERR, self.object_id, "sandbox", self._client, by_line=True
         )
         self._stdin = StreamWriter(self.object_id, "sandbox", self._client)
         self._result = None
+        self._task_id = None
+        self._tunnels = None
+        self._enable_snapshot = False
+        self._command_router_client = None
 
     @staticmethod
     async def from_name(
@@ -669,11 +676,11 @@ class _Sandbox(_Object, type_prefix="sb"):
     async def create_connect_token(
         self, user_metadata: Optional[Union[str, dict[str, Any]]] = None
     ) -> SandboxConnectCredentials:
-        """mdmd:hidden
-        [Alpha] Create a token for making HTTP connections to the sandbox.
+        """
+        [Alpha] Create a token for making HTTP connections to the Sandbox.
 
         Also accepts an optional user_metadata string or dict to associate with the token. This metadata
-        will be added to the headers by the proxy when forwarding requests to the sandbox."""
+        will be added to the headers by the proxy when forwarding requests to the Sandbox."""
         if user_metadata is not None and isinstance(user_metadata, dict):
             try:
                 user_metadata = json.dumps(user_metadata)
@@ -730,6 +737,13 @@ class _Sandbox(_Object, type_prefix="sb"):
                 await asyncio.sleep(0.5)
         return self._task_id
 
+    async def _get_command_router_client(self, task_id: str) -> Optional[TaskCommandRouterClient]:
+        if self._command_router_client is None:
+            # Attempt to initialize a router client. Returns None if the new exec path not enabled
+            # for this sandbox.
+            self._command_router_client = await TaskCommandRouterClient.try_init(self._client, task_id)
+        return self._command_router_client
+
     @overload
     async def exec(
         self,
@@ -855,14 +869,49 @@ class _Sandbox(_Object, type_prefix="sb"):
         await TaskContext.gather(*secret_coros)
 
         task_id = await self._get_task_id()
+        kwargs = {
+            "task_id": task_id,
+            "pty_info": pty_info,
+            "stdout": stdout,
+            "stderr": stderr,
+            "timeout": timeout,
+            "workdir": workdir,
+            "secret_ids": [secret.object_id for secret in secrets],
+            "text": text,
+            "bufsize": bufsize,
+            "runtime_debug": config.get("function_runtime_debug"),
+        }
+        # NB: This must come after the task ID is set, since the sandbox must be
+        # scheduled before we can create a router client.
+        if (command_router_client := await self._get_command_router_client(task_id)) is not None:
+            kwargs["command_router_client"] = command_router_client
+            return await self._exec_through_command_router(*args, **kwargs)
+        else:
+            return await self._exec_through_server(*args, **kwargs)
+
+    async def _exec_through_server(
+        self,
+        *args: str,
+        task_id: str,
+        pty_info: Optional[api_pb2.PTYInfo] = None,
+        stdout: StreamType = StreamType.PIPE,
+        stderr: StreamType = StreamType.PIPE,
+        timeout: Optional[int] = None,
+        workdir: Optional[str] = None,
+        secret_ids: Optional[Collection[str]] = None,
+        text: bool = True,
+        bufsize: Literal[-1, 1] = -1,
+        runtime_debug: bool = False,
+    ) -> Union[_ContainerProcess[bytes], _ContainerProcess[str]]:
+        """Execute a command through the Modal server."""
         req = api_pb2.ContainerExecRequest(
             task_id=task_id,
             command=args,
             pty_info=pty_info,
-            runtime_debug=config.get("function_runtime_debug"),
+            runtime_debug=runtime_debug,
             timeout_secs=timeout or 0,
             workdir=workdir,
-            secret_ids=[secret.object_id for secret in secrets],
+            secret_ids=secret_ids,
         )
         resp = await retry_transient_errors(self._client.stub.ContainerExec, req)
         by_line = bufsize == 1
@@ -870,6 +919,7 @@ class _Sandbox(_Object, type_prefix="sb"):
         logger.debug(f"Created ContainerProcess for exec_id {resp.exec_id} on Sandbox {self.object_id}")
         return _ContainerProcess(
             resp.exec_id,
+            task_id,
             self._client,
             stdout=stdout,
             stderr=stderr,
@@ -878,6 +928,75 @@ class _Sandbox(_Object, type_prefix="sb"):
             by_line=by_line,
         )
 
+    async def _exec_through_command_router(
+        self,
+        *args: str,
+        task_id: str,
+        command_router_client: TaskCommandRouterClient,
+        pty_info: Optional[api_pb2.PTYInfo] = None,
+        stdout: StreamType = StreamType.PIPE,
+        stderr: StreamType = StreamType.PIPE,
+        timeout: Optional[int] = None,
+        workdir: Optional[str] = None,
+        secret_ids: Optional[Collection[str]] = None,
+        text: bool = True,
+        bufsize: Literal[-1, 1] = -1,
+        runtime_debug: bool = False,
+    ) -> Union[_ContainerProcess[bytes], _ContainerProcess[str]]:
+        """Execute a command through a task command router running on the Modal worker."""
+
+        # Generate a random process ID to use as a combination of idempotency key/process identifier.
+        process_id = str(uuid.uuid4())
+        if stdout == StreamType.PIPE:
+            stdout_config = sr_pb2.TaskExecStdoutConfig.TASK_EXEC_STDOUT_CONFIG_PIPE
+        elif stdout == StreamType.DEVNULL:
+            stdout_config = sr_pb2.TaskExecStdoutConfig.TASK_EXEC_STDOUT_CONFIG_DEVNULL
+        elif stdout == StreamType.STDOUT:
+            # TODO(saltzm): This is a behavior change from the old implementation. We should
+            # probably implement the old behavior of printing to stdout before moving out of beta.
+            raise NotImplementedError(
+                "Currently the STDOUT stream type is not supported when using exec "
+                "through a task command router, which is currently in beta."
+            )
+        else:
+            raise ValueError("Unsupported StreamType for stdout")
+
+        if stderr == StreamType.PIPE:
+            stderr_config = sr_pb2.TaskExecStderrConfig.TASK_EXEC_STDERR_CONFIG_PIPE
+        elif stderr == StreamType.DEVNULL:
+            stderr_config = sr_pb2.TaskExecStderrConfig.TASK_EXEC_STDERR_CONFIG_DEVNULL
+        elif stderr == StreamType.STDOUT:
+            stderr_config = sr_pb2.TaskExecStderrConfig.TASK_EXEC_STDERR_CONFIG_STDOUT
+        else:
+            raise ValueError("Unsupported StreamType for stderr")
+
+        # Start the process.
+        start_req = sr_pb2.TaskExecStartRequest(
+            task_id=task_id,
+            exec_id=process_id,
+            command_args=args,
+            stdout_config=stdout_config,
+            stderr_config=stderr_config,
+            timeout_secs=timeout,
+            workdir=workdir,
+            secret_ids=secret_ids,
+            pty_info=pty_info,
+            runtime_debug=runtime_debug,
+        )
+        _ = await command_router_client.exec_start(start_req)
+
+        return _ContainerProcess(
+            process_id,
+            task_id,
+            self._client,
+            command_router_client=command_router_client,
+            stdout=stdout,
+            stderr=stderr,
+            text=text,
+            by_line=bufsize == 1,
+            exec_deadline=time.monotonic() + int(timeout) if timeout else None,
+        )
+
     async def _experimental_snapshot(self) -> _SandboxSnapshot:
         await self._get_task_id()
         snap_req = api_pb2.SandboxSnapshotRequest(sandbox_id=self.object_id)

modal/sandbox.pyi

@@ -3,6 +3,7 @@ import collections.abc
 import google.protobuf.message
 import modal._object
 import modal._tunnel
+import modal._utils.task_command_router_client
 import modal.app
 import modal.client
 import modal.cloud_bucket_mount
@@ -83,6 +84,7 @@ class _Sandbox(modal._object._Object):
     _task_id: typing.Optional[str]
     _tunnels: typing.Optional[dict[int, modal._tunnel.Tunnel]]
     _enable_snapshot: bool
+    _command_router_client: typing.Optional[modal._utils.task_command_router_client.TaskCommandRouterClient]
 
     @staticmethod
     def _default_pty_info() -> modal_proto.api_pb2.PTYInfo: ...
@@ -276,11 +278,10 @@ class _Sandbox(modal._object._Object):
     async def create_connect_token(
         self, user_metadata: typing.Union[str, dict[str, typing.Any], None] = None
     ) -> SandboxConnectCredentials:
-        """mdmd:hidden
-        [Alpha] Create a token for making HTTP connections to the sandbox.
+        """[Alpha] Create a token for making HTTP connections to the Sandbox.
 
         Also accepts an optional user_metadata string or dict to associate with the token. This metadata
-        will be added to the headers by the proxy when forwarding requests to the sandbox.
+        will be added to the headers by the proxy when forwarding requests to the Sandbox.
         """
         ...
 
@@ -306,6 +307,9 @@ class _Sandbox(modal._object._Object):
         ...
 
     async def _get_task_id(self) -> str: ...
+    async def _get_command_router_client(
+        self, task_id: str
+    ) -> typing.Optional[modal._utils.task_command_router_client.TaskCommandRouterClient]: ...
     @typing.overload
     async def exec(
         self,
@@ -357,6 +361,41 @@ class _Sandbox(modal._object._Object):
         """
         ...
 
+    async def _exec_through_server(
+        self,
+        *args: str,
+        task_id: str,
+        pty_info: typing.Optional[modal_proto.api_pb2.PTYInfo] = None,
+        stdout: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+        stderr: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+        timeout: typing.Optional[int] = None,
+        workdir: typing.Optional[str] = None,
+        secret_ids: typing.Optional[collections.abc.Collection[str]] = None,
+        text: bool = True,
+        bufsize: typing.Literal[-1, 1] = -1,
+        runtime_debug: bool = False,
+    ) -> typing.Union[modal.container_process._ContainerProcess[bytes], modal.container_process._ContainerProcess[str]]:
+        """Execute a command through the Modal server."""
+        ...
+
+    async def _exec_through_command_router(
+        self,
+        *args: str,
+        task_id: str,
+        command_router_client: modal._utils.task_command_router_client.TaskCommandRouterClient,
+        pty_info: typing.Optional[modal_proto.api_pb2.PTYInfo] = None,
+        stdout: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+        stderr: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+        timeout: typing.Optional[int] = None,
+        workdir: typing.Optional[str] = None,
+        secret_ids: typing.Optional[collections.abc.Collection[str]] = None,
+        text: bool = True,
+        bufsize: typing.Literal[-1, 1] = -1,
+        runtime_debug: bool = False,
+    ) -> typing.Union[modal.container_process._ContainerProcess[bytes], modal.container_process._ContainerProcess[str]]:
+        """Execute a command through a task command router running on the Modal worker."""
+        ...
+
     async def _experimental_snapshot(self) -> modal.snapshot._SandboxSnapshot: ...
     @staticmethod
     async def _experimental_from_snapshot(
@@ -445,6 +484,7 @@ class Sandbox(modal.object.Object):
     _task_id: typing.Optional[str]
     _tunnels: typing.Optional[dict[int, modal._tunnel.Tunnel]]
     _enable_snapshot: bool
+    _command_router_client: typing.Optional[modal._utils.task_command_router_client.TaskCommandRouterClient]
 
     def __init__(self, *args, **kwargs):
         """mdmd:hidden"""
@@ -832,22 +872,20 @@ class Sandbox(modal.object.Object):
         def __call__(
             self, /, user_metadata: typing.Union[str, dict[str, typing.Any], None] = None
         ) -> SandboxConnectCredentials:
-            """mdmd:hidden
-            [Alpha] Create a token for making HTTP connections to the sandbox.
+            """[Alpha] Create a token for making HTTP connections to the Sandbox.
 
             Also accepts an optional user_metadata string or dict to associate with the token. This metadata
-            will be added to the headers by the proxy when forwarding requests to the sandbox.
+            will be added to the headers by the proxy when forwarding requests to the Sandbox.
             """
             ...
 
         async def aio(
             self, /, user_metadata: typing.Union[str, dict[str, typing.Any], None] = None
         ) -> SandboxConnectCredentials:
-            """mdmd:hidden
-            [Alpha] Create a token for making HTTP connections to the sandbox.
+            """[Alpha] Create a token for making HTTP connections to the Sandbox.
 
             Also accepts an optional user_metadata string or dict to associate with the token. This metadata
-            will be added to the headers by the proxy when forwarding requests to the sandbox.
+            will be added to the headers by the proxy when forwarding requests to the Sandbox.
             """
             ...
 
@@ -910,6 +948,16 @@ class Sandbox(modal.object.Object):
 
     _get_task_id: ___get_task_id_spec[typing_extensions.Self]
 
+    class ___get_command_router_client_spec(typing_extensions.Protocol[SUPERSELF]):
+        def __call__(
+            self, /, task_id: str
+        ) -> typing.Optional[modal._utils.task_command_router_client.TaskCommandRouterClient]: ...
+        async def aio(
+            self, /, task_id: str
+        ) -> typing.Optional[modal._utils.task_command_router_client.TaskCommandRouterClient]: ...
+
+    _get_command_router_client: ___get_command_router_client_spec[typing_extensions.Self]
+
     class __exec_spec(typing_extensions.Protocol[SUPERSELF]):
         @typing.overload
         def __call__(
@@ -1029,6 +1077,94 @@ class Sandbox(modal.object.Object):
 
     _exec: ___exec_spec[typing_extensions.Self]
 
+    class ___exec_through_server_spec(typing_extensions.Protocol[SUPERSELF]):
+        def __call__(
+            self,
+            /,
+            *args: str,
+            task_id: str,
+            pty_info: typing.Optional[modal_proto.api_pb2.PTYInfo] = None,
+            stdout: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+            stderr: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+            timeout: typing.Optional[int] = None,
+            workdir: typing.Optional[str] = None,
+            secret_ids: typing.Optional[collections.abc.Collection[str]] = None,
+            text: bool = True,
+            bufsize: typing.Literal[-1, 1] = -1,
+            runtime_debug: bool = False,
+        ) -> typing.Union[
+            modal.container_process.ContainerProcess[bytes], modal.container_process.ContainerProcess[str]
+        ]:
+            """Execute a command through the Modal server."""
+            ...
+
+        async def aio(
+            self,
+            /,
+            *args: str,
+            task_id: str,
+            pty_info: typing.Optional[modal_proto.api_pb2.PTYInfo] = None,
+            stdout: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+            stderr: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+            timeout: typing.Optional[int] = None,
+            workdir: typing.Optional[str] = None,
+            secret_ids: typing.Optional[collections.abc.Collection[str]] = None,
+            text: bool = True,
+            bufsize: typing.Literal[-1, 1] = -1,
+            runtime_debug: bool = False,
+        ) -> typing.Union[
+            modal.container_process.ContainerProcess[bytes], modal.container_process.ContainerProcess[str]
+        ]:
+            """Execute a command through the Modal server."""
+            ...
+
+    _exec_through_server: ___exec_through_server_spec[typing_extensions.Self]
+
+    class ___exec_through_command_router_spec(typing_extensions.Protocol[SUPERSELF]):
+        def __call__(
+            self,
+            /,
+            *args: str,
+            task_id: str,
+            command_router_client: modal._utils.task_command_router_client.TaskCommandRouterClient,
+            pty_info: typing.Optional[modal_proto.api_pb2.PTYInfo] = None,
+            stdout: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+            stderr: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+            timeout: typing.Optional[int] = None,
+            workdir: typing.Optional[str] = None,
+            secret_ids: typing.Optional[collections.abc.Collection[str]] = None,
+            text: bool = True,
+            bufsize: typing.Literal[-1, 1] = -1,
+            runtime_debug: bool = False,
+        ) -> typing.Union[
+            modal.container_process.ContainerProcess[bytes], modal.container_process.ContainerProcess[str]
+        ]:
+            """Execute a command through a task command router running on the Modal worker."""
+            ...
+
+        async def aio(
+            self,
+            /,
+            *args: str,
+            task_id: str,
+            command_router_client: modal._utils.task_command_router_client.TaskCommandRouterClient,
+            pty_info: typing.Optional[modal_proto.api_pb2.PTYInfo] = None,
+            stdout: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+            stderr: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+            timeout: typing.Optional[int] = None,
+            workdir: typing.Optional[str] = None,
+            secret_ids: typing.Optional[collections.abc.Collection[str]] = None,
+            text: bool = True,
+            bufsize: typing.Literal[-1, 1] = -1,
+            runtime_debug: bool = False,
+        ) -> typing.Union[
+            modal.container_process.ContainerProcess[bytes], modal.container_process.ContainerProcess[str]
+        ]:
+            """Execute a command through a task command router running on the Modal worker."""
+            ...
+
+    _exec_through_command_router: ___exec_through_command_router_spec[typing_extensions.Self]
+
     class ___experimental_snapshot_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, /) -> modal.snapshot.SandboxSnapshot: ...
         async def aio(self, /) -> modal.snapshot.SandboxSnapshot: ...

modal/volume.py

@@ -655,6 +655,7 @@ class _Volume(_Object, type_prefix="vo"):
         @retry(n_attempts=5, base_delay=0.1, timeout=None)
         async def read_block(block_url: str) -> bytes:
             async with ClientSessionRegistry.get_session().get(block_url) as get_response:
+                get_response.raise_for_status()
                 return await get_response.content.read()
 
         async def iter_urls() -> AsyncGenerator[str]:
@@ -670,16 +671,33 @@ class _Volume(_Object, type_prefix="vo"):
 
     @live_method
     async def read_file_into_fileobj(
-        self, path: str, fileobj: typing.IO[bytes], progress_cb: Optional[Callable[..., Any]] = None
+        self,
+        path: str,
+        fileobj: typing.IO[bytes],
+        progress_cb: Optional[Callable[..., Any]] = None,
     ) -> int:
         """mdmd:hidden
         Read volume file into file-like IO object.
         """
+        return await self._read_file_into_fileobj(path, fileobj, progress_cb=progress_cb)
+
+    @live_method
+    async def _read_file_into_fileobj(
+        self,
+        path: str,
+        fileobj: typing.IO[bytes],
+        concurrency: Optional[int] = None,
+        download_semaphore: Optional[asyncio.Semaphore] = None,
+        progress_cb: Optional[Callable[..., Any]] = None,
+    ) -> int:
         if progress_cb is None:
 
             def progress_cb(*_, **__):
                 pass
 
+        if concurrency is None:
+            concurrency = multiprocessing.cpu_count()
+
         req = api_pb2.VolumeGetFile2Request(volume_id=self.object_id, path=path)
 
         try:
@@ -687,8 +705,9 @@ class _Volume(_Object, type_prefix="vo"):
         except modal.exception.NotFoundError as exc:
             raise FileNotFoundError(exc.args[0])
 
-        # TODO(dflemstr): Sane default limit? Make configurable?
-        download_semaphore = asyncio.Semaphore(multiprocessing.cpu_count())
+        if download_semaphore is None:
+            download_semaphore = asyncio.Semaphore(concurrency)
+
         write_lock = asyncio.Lock()
         start_pos = fileobj.tell()
 
@@ -698,6 +717,7 @@ class _Volume(_Object, type_prefix="vo"):
             num_bytes_written = 0
 
             async with download_semaphore, ClientSessionRegistry.get_session().get(url) as get_response:
+                get_response.raise_for_status()
                 async for chunk in get_response.content.iter_any():
                     num_chunk_bytes_written = 0