modal 1.1.5.dev66__py3-none-any.whl → 1.3.1.dev8__py3-none-any.whl

modal/sandbox.py

@@ -1,10 +1,13 @@
 # Copyright Modal Labs 2022
 import asyncio
+import builtins
 import json
 import os
 import time
+import uuid
 from collections.abc import AsyncGenerator, Collection, Sequence
 from dataclasses import dataclass
+from pathlib import PurePosixPath
 from typing import TYPE_CHECKING, Any, AsyncIterator, Literal, Optional, Union, overload
 
 from ._pty import get_pty_info
@@ -14,32 +17,31 @@ if TYPE_CHECKING:
     import _typeshed
 
 from google.protobuf.message import Message
-from grpclib import GRPCError, Status
 
 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 ._load_context import LoadContext
 from ._object import _get_environment_name, _Object
 from ._resolver import Resolver
 from ._resources import convert_fn_config_to_resources_config
 from ._utils.async_utils import TaskContext, synchronize_api
 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.name_utils import check_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
+from .exception import ExecutionError, InvalidError, SandboxTerminatedError, SandboxTimeoutError
 from .file_io import FileWatchEvent, FileWatchEventType, _FileIO
 from .gpu import GPU_T
 from .image import _Image
 from .io_streams import StreamReader, StreamWriter, _StreamReader, _StreamWriter
 from .network_file_system import _NetworkFileSystem, network_file_system_mount_protos
 from .proxy import _Proxy
-from .scheduler_placement import SchedulerPlacement
 from .secret import _Secret
 from .snapshot import _SandboxSnapshot
 from .stream_type import StreamType
@@ -78,16 +80,6 @@ def _validate_exec_args(args: Sequence[str]) -> None:
         )
 
 
-def _warn_if_invalid_name(name: str) -> None:
-    if not is_valid_object_name(name):
-        deprecation_warning(
-            (2025, 9, 3),
-            f"Sandbox name '{name}' will be considered invalid in a future release."
-            "\n\nNames may contain only alphanumeric characters, dashes, periods, and underscores,"
-            " must be shorter than 64 characters, and cannot conflict with App ID strings.",
-        )
-
-
 class DefaultSandboxNameOverride(str):
     """A singleton class that represents the default sandbox name override.
 
@@ -121,9 +113,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:
@@ -155,7 +148,6 @@ class _Sandbox(_Object, type_prefix="sb"):
         unencrypted_ports: Sequence[int] = [],
         proxy: Optional[_Proxy] = None,
         experimental_options: Optional[dict[str, bool]] = None,
-        _experimental_scheduler_placement: Optional[SchedulerPlacement] = None,
         enable_snapshot: bool = False,
         verbose: bool = False,
     ) -> "_Sandbox":
@@ -163,12 +155,6 @@ class _Sandbox(_Object, type_prefix="sb"):
 
         validated_network_file_systems = validate_network_file_systems(network_file_systems)
 
-        scheduler_placement: Optional[SchedulerPlacement] = _experimental_scheduler_placement
-        if region:
-            if scheduler_placement:
-                raise InvalidError("`region` and `_experimental_scheduler_placement` cannot be used together")
-            scheduler_placement = SchedulerPlacement(region=region)
-
         if isinstance(gpu, list):
             raise InvalidError(
                 "Sandboxes do not support configuring a list of GPUs. "
@@ -183,6 +169,11 @@ class _Sandbox(_Object, type_prefix="sb"):
         cloud_bucket_mounts = [(k, v) for k, v in validated_volumes if isinstance(v, _CloudBucketMount)]
         validated_volumes = [(k, v) for k, v in validated_volumes if isinstance(v, _Volume)]
 
+        scheduler_placement: Optional[api_pb2.SchedulerPlacement] = None
+        if region:
+            regions = [region] if isinstance(region, str) else (list(region) if region else None)
+            scheduler_placement = api_pb2.SchedulerPlacement(regions=regions)
+
         if pty:
             pty_info = _Sandbox._default_pty_info()
 
@@ -199,7 +190,9 @@ class _Sandbox(_Object, type_prefix="sb"):
                 deps.append(proxy)
             return deps
 
-        async def _load(self: _Sandbox, resolver: Resolver, _existing_object_id: Optional[str]):
+        async def _load(
+            self: _Sandbox, resolver: Resolver, load_context: LoadContext, _existing_object_id: Optional[str]
+        ):
             # Relies on dicts being ordered (true as of Python 3.6).
             volume_mounts = [
                 api_pb2.VolumeMount(
@@ -257,7 +250,7 @@ class _Sandbox(_Object, type_prefix="sb"):
                 cloud_bucket_mounts=cloud_bucket_mounts_to_proto(cloud_bucket_mounts),
                 volume_mounts=volume_mounts,
                 pty_info=pty_info,
-                scheduler_placement=scheduler_placement.proto if scheduler_placement else None,
+                scheduler_placement=scheduler_placement,
                 worker_id=config.get("worker_id"),
                 open_ports=api_pb2.PortSpecs(ports=open_ports),
                 network_access=network_access,
@@ -268,18 +261,12 @@ class _Sandbox(_Object, type_prefix="sb"):
                 experimental_options=experimental_options,
             )
 
-            create_req = api_pb2.SandboxCreateRequest(app_id=resolver.app_id, definition=definition)
-            try:
-                create_resp = await retry_transient_errors(resolver.client.stub.SandboxCreate, create_req)
-            except GRPCError as exc:
-                if exc.status == Status.ALREADY_EXISTS:
-                    raise AlreadyExistsError(exc.message)
-                raise exc
-
+            create_req = api_pb2.SandboxCreateRequest(app_id=load_context.app_id, definition=definition)
+            create_resp = await load_context.client.stub.SandboxCreate(create_req)
             sandbox_id = create_resp.sandbox_id
-            self._hydrate(sandbox_id, resolver.client, None)
+            self._hydrate(sandbox_id, load_context.client, None)
 
-        return _Sandbox._from_loader(_load, "Sandbox()", deps=_deps)
+        return _Sandbox._from_loader(_load, "Sandbox()", deps=_deps, load_context_overrides=LoadContext.empty())
 
     @staticmethod
     async def create(
@@ -325,9 +312,6 @@ class _Sandbox(_Object, type_prefix="sb"):
         experimental_options: Optional[dict[str, bool]] = None,
         # Enable memory snapshots.
         _experimental_enable_snapshot: bool = False,
-        _experimental_scheduler_placement: Optional[
-            SchedulerPlacement
-        ] = None,  # Experimental controls over fine-grained scheduling (alpha).
         client: Optional[_Client] = None,
         environment_name: Optional[str] = None,  # *DEPRECATED* Optionally override the default environment
         pty_info: Optional[api_pb2.PTYInfo] = None,  # *DEPRECATED* Use `pty` instead. `pty` will override `pty_info`.
@@ -389,7 +373,6 @@ class _Sandbox(_Object, type_prefix="sb"):
             proxy=proxy,
             experimental_options=experimental_options,
             _experimental_enable_snapshot=_experimental_enable_snapshot,
-            _experimental_scheduler_placement=_experimental_scheduler_placement,
             client=client,
             verbose=verbose,
             pty_info=pty_info,
@@ -423,7 +406,6 @@ class _Sandbox(_Object, type_prefix="sb"):
         proxy: Optional[_Proxy] = None,
         experimental_options: Optional[dict[str, bool]] = None,
         _experimental_enable_snapshot: bool = False,
-        _experimental_scheduler_placement: Optional[SchedulerPlacement] = None,
         client: Optional[_Client] = None,
         verbose: bool = False,
         pty_info: Optional[api_pb2.PTYInfo] = None,
@@ -438,7 +420,7 @@ class _Sandbox(_Object, type_prefix="sb"):
 
         _validate_exec_args(args)
         if name is not None:
-            _warn_if_invalid_name(name)
+            check_object_name(name, "Sandbox")
 
         if block_network and (encrypted_ports or h2_ports or unencrypted_ports):
             raise InvalidError("Cannot specify open ports when `block_network` is enabled")
@@ -473,7 +455,6 @@ class _Sandbox(_Object, type_prefix="sb"):
             unencrypted_ports=unencrypted_ports,
             proxy=proxy,
             experimental_options=experimental_options,
-            _experimental_scheduler_placement=_experimental_scheduler_placement,
             enable_snapshot=_experimental_enable_snapshot,
             verbose=verbose,
         )
@@ -494,6 +475,7 @@ class _Sandbox(_Object, type_prefix="sb"):
             app_id = app.app_id
             app_client = app._client
         elif (container_app := _App._get_container_app()) is not None:
+            # implicit app/client provided by running in a modal Function
             app_id = container_app.app_id
             app_client = container_app._client
         else:
@@ -506,21 +488,26 @@ class _Sandbox(_Object, type_prefix="sb"):
                 "```",
             )
 
-        client = client or app_client or await _Client.from_env()
+        client = client or app_client
 
-        resolver = Resolver(client, app_id=app_id)
-        await resolver.load(obj)
+        resolver = Resolver()
+        load_context = LoadContext(client=client, app_id=app_id)
+        await resolver.load(obj, load_context)
         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(
@@ -540,7 +527,7 @@ class _Sandbox(_Object, type_prefix="sb"):
         env_name = _get_environment_name(environment_name)
 
         req = api_pb2.SandboxGetFromNameRequest(sandbox_name=name, app_name=app_name, environment_name=env_name)
-        resp = await retry_transient_errors(client.stub.SandboxGetFromName, req)
+        resp = await client.stub.SandboxGetFromName(req)
         return _Sandbox._new_hydrated(resp.sandbox_id, client, None)
 
     @staticmethod
@@ -553,7 +540,7 @@ class _Sandbox(_Object, type_prefix="sb"):
             client = await _Client.from_env()
 
         req = api_pb2.SandboxWaitRequest(sandbox_id=sandbox_id, timeout=0)
-        resp = await retry_transient_errors(client.stub.SandboxWait, req)
+        resp = await client.stub.SandboxWait(req)
 
         obj = _Sandbox._new_hydrated(sandbox_id, client, None)
 
@@ -565,10 +552,7 @@ class _Sandbox(_Object, type_prefix="sb"):
     async def get_tags(self) -> dict[str, str]:
         """Fetches any tags (key-value pairs) currently attached to this Sandbox from the server."""
         req = api_pb2.SandboxTagsGetRequest(sandbox_id=self.object_id)
-        try:
-            resp = await retry_transient_errors(self._client.stub.SandboxTagsGet, req)
-        except GRPCError as exc:
-            raise InvalidError(exc.message) if exc.status == Status.INVALID_ARGUMENT else exc
+        resp = await self._client.stub.SandboxTagsGet(req)
 
         return {tag.tag_name: tag.tag_value for tag in resp.tags}
 
@@ -589,10 +573,7 @@ class _Sandbox(_Object, type_prefix="sb"):
             sandbox_id=self.object_id,
             tags=tags_list,
         )
-        try:
-            await retry_transient_errors(self._client.stub.SandboxTagsSet, req)
-        except GRPCError as exc:
-            raise InvalidError(exc.message) if exc.status == Status.INVALID_ARGUMENT else exc
+        await self._client.stub.SandboxTagsSet(req)
 
     async def snapshot_filesystem(self, timeout: int = 55) -> _Image:
         """Snapshot the filesystem of the Sandbox.
@@ -602,7 +583,7 @@ class _Sandbox(_Object, type_prefix="sb"):
         """
         await self._get_task_id()  # Ensure the sandbox has started
         req = api_pb2.SandboxSnapshotFsRequest(sandbox_id=self.object_id, timeout=timeout)
-        resp = await retry_transient_errors(self._client.stub.SandboxSnapshotFs, req)
+        resp = await self._client.stub.SandboxSnapshotFs(req)
 
         if resp.result.status != api_pb2.GenericResult.GENERIC_STATUS_SUCCESS:
             raise ExecutionError(resp.result.exception)
@@ -610,16 +591,60 @@ class _Sandbox(_Object, type_prefix="sb"):
         image_id = resp.image_id
         metadata = resp.image_metadata
 
-        async def _load(self: _Image, resolver: Resolver, existing_object_id: Optional[str]):
+        async def _load(self: _Image, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]):
             # no need to hydrate again since we do it eagerly below
             pass
 
         rep = "Image()"
-        image = _Image._from_loader(_load, rep, hydrate_lazily=True)
+        # TODO: use ._new_hydrated instead
+        image = _Image._from_loader(_load, rep, hydrate_lazily=True, load_context_overrides=LoadContext.empty())
         image._hydrate(image_id, self._client, metadata)  # hydrating eagerly since we have all of the data
 
         return image
 
+    async def _experimental_mount_image(self, path: Union[PurePosixPath, str], image: Optional[_Image]):
+        """Mount an Image at a path in the Sandbox filesystem."""
+
+        image_id = None
+
+        if image:
+            if not image._object_id:
+                # FIXME
+                raise InvalidError("Image has not been built.")
+            image_id = image._object_id
+        else:
+            image_id = ""  # empty string indicates mount an empty dir
+
+        task_id = await self._get_task_id()
+        if (command_router_client := await self._get_command_router_client(task_id)) is None:
+            raise InvalidError("Mounting directories requires direct Sandbox control - please contact Modal support.")
+
+        posix_path = PurePosixPath(path)
+        if not posix_path.is_absolute():
+            raise InvalidError(f"Mount path must be absolute; got: {posix_path}")
+        path_bytes = posix_path.as_posix().encode("utf8")
+
+        req = sr_pb2.TaskMountDirectoryRequest(task_id=task_id, path=path_bytes, image_id=image_id)
+        await command_router_client.mount_image(req)
+
+    async def _experimental_snapshot_directory(self, path: Union[PurePosixPath, str]) -> _Image:
+        """Snapshot local changes to a previously mounted Image, creating a new Image."""
+
+        task_id = await self._get_task_id()
+        if (command_router_client := await self._get_command_router_client(task_id)) is None:
+            raise InvalidError(
+                "Snapshotting directories requires direct Sandbox control - please contact Modal support."
+            )
+
+        posix_path = PurePosixPath(path)
+        if not posix_path.is_absolute():
+            raise InvalidError(f"Snapshot path must be absolute; got: {posix_path}")
+        path_bytes = posix_path.as_posix().encode("utf8")
+
+        req = sr_pb2.TaskSnapshotDirectoryRequest(task_id=task_id, path=path_bytes)
+        res = await command_router_client.snapshot_directory(req)
+        return _Image._new_hydrated(res.image_id, self._client, None)
+
     # Live handle methods
 
     async def wait(self, raise_on_termination: bool = True):
@@ -627,7 +652,7 @@ class _Sandbox(_Object, type_prefix="sb"):
 
         while True:
             req = api_pb2.SandboxWaitRequest(sandbox_id=self.object_id, timeout=10)
-            resp = await retry_transient_errors(self._client.stub.SandboxWait, req)
+            resp = await self._client.stub.SandboxWait(req)
             if resp.result.status:
                 logger.debug(f"Sandbox {self.object_id} wait completed with status {resp.result.status}")
                 self._result = resp.result
@@ -653,7 +678,7 @@ class _Sandbox(_Object, type_prefix="sb"):
             return self._tunnels
 
         req = api_pb2.SandboxGetTunnelsRequest(sandbox_id=self.object_id, timeout=timeout)
-        resp = await retry_transient_errors(self._client.stub.SandboxGetTunnels, req)
+        resp = await self._client.stub.SandboxGetTunnels(req)
 
         # If we couldn't get the tunnels in time, report the timeout.
         if resp.result.status == api_pb2.GenericResult.GENERIC_STATUS_TIMEOUT:
@@ -669,10 +694,11 @@ class _Sandbox(_Object, type_prefix="sb"):
     async def create_connect_token(
         self, user_metadata: Optional[Union[str, dict[str, Any]]] = None
     ) -> SandboxConnectCredentials:
-        """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)
@@ -680,7 +706,7 @@ class _Sandbox(_Object, type_prefix="sb"):
                 raise InvalidError(f"Failed to serialize user_metadata: {e}")
 
         req = api_pb2.SandboxCreateConnectTokenRequest(sandbox_id=self.object_id, user_metadata=user_metadata)
-        resp = await retry_transient_errors(self._client.stub.SandboxCreateConnectToken, req)
+        resp = await self._client.stub.SandboxCreateConnectToken(req)
         return SandboxConnectCredentials(resp.url, resp.token)
 
     async def reload_volumes(self) -> None:
@@ -689,8 +715,7 @@ class _Sandbox(_Object, type_prefix="sb"):
         Added in v1.1.0.
         """
         task_id = await self._get_task_id()
-        await retry_transient_errors(
-            self._client.stub.ContainerReloadVolumes,
+        await self._client.stub.ContainerReloadVolumes(
             api_pb2.ContainerReloadVolumesRequest(
                 task_id=task_id,
             ),
@@ -701,9 +726,7 @@ class _Sandbox(_Object, type_prefix="sb"):
 
         This is a no-op if the Sandbox has already finished running."""
 
-        await retry_transient_errors(
-            self._client.stub.SandboxTerminate, api_pb2.SandboxTerminateRequest(sandbox_id=self.object_id)
-        )
+        await self._client.stub.SandboxTerminate(api_pb2.SandboxTerminateRequest(sandbox_id=self.object_id))
 
     async def poll(self) -> Optional[int]:
         """Check if the Sandbox has finished running.
@@ -712,7 +735,7 @@ class _Sandbox(_Object, type_prefix="sb"):
         """
 
         req = api_pb2.SandboxWaitRequest(sandbox_id=self.object_id, timeout=0)
-        resp = await retry_transient_errors(self._client.stub.SandboxWait, req)
+        resp = await self._client.stub.SandboxWait(req)
 
         if resp.result.status:
             self._result = resp.result
@@ -721,14 +744,19 @@ class _Sandbox(_Object, type_prefix="sb"):
 
     async def _get_task_id(self) -> str:
         while not self._task_id:
-            resp = await retry_transient_errors(
-                self._client.stub.SandboxGetTaskId, api_pb2.SandboxGetTaskIdRequest(sandbox_id=self.object_id)
-            )
+            resp = await self._client.stub.SandboxGetTaskId(api_pb2.SandboxGetTaskIdRequest(sandbox_id=self.object_id))
             self._task_id = resp.task_id
             if not self._task_id:
                 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,
@@ -790,13 +818,8 @@ class _Sandbox(_Object, type_prefix="sb"):
 
         **Usage**
 
-        ```python
-        app = modal.App.lookup("my-app", create_if_missing=True)
-
-        sandbox = modal.Sandbox.create("sleep", "infinity", app=app)
-
-        process = sandbox.exec("bash", "-c", "for i in $(seq 1 10); do echo foo $i; sleep 0.5; done")
-
+        ```python fixture:sandbox
+        process = sandbox.exec("bash", "-c", "for i in $(seq 1 3); do echo foo $i; sleep 0.1; done")
         for line in process.stdout:
             print(line)
         ```
@@ -854,21 +877,57 @@ 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)
+        resp = await self._client.stub.ContainerExec(req)
         by_line = bufsize == 1
         exec_deadline = time.monotonic() + int(timeout) + CONTAINER_EXEC_TIMEOUT_BUFFER if timeout else None
         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,
@@ -877,26 +936,95 @@ 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:
+            # Stream stdout to the client so that it can be printed locally in the reader.
+            stdout_config = sr_pb2.TaskExecStdoutConfig.TASK_EXEC_STDOUT_CONFIG_PIPE
+        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:
+            # Stream stderr to the client so that it can be printed locally in the reader.
+            stderr_config = sr_pb2.TaskExecStderrConfig.TASK_EXEC_STDERR_CONFIG_PIPE
+        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)
-        snap_resp = await retry_transient_errors(self._client.stub.SandboxSnapshot, snap_req)
+        snap_resp = await self._client.stub.SandboxSnapshot(snap_req)
 
         snapshot_id = snap_resp.snapshot_id
 
         # wait for the snapshot to succeed. this is implemented as a second idempotent rpc
         # because the snapshot itself may take a while to complete.
         wait_req = api_pb2.SandboxSnapshotWaitRequest(snapshot_id=snapshot_id, timeout=55.0)
-        wait_resp = await retry_transient_errors(self._client.stub.SandboxSnapshotWait, wait_req)
+        wait_resp = await self._client.stub.SandboxSnapshotWait(wait_req)
         if wait_resp.result.status != api_pb2.GenericResult.GENERIC_STATUS_SUCCESS:
             raise ExecutionError(wait_resp.result.exception)
 
-        async def _load(self: _SandboxSnapshot, resolver: Resolver, existing_object_id: Optional[str]):
+        async def _load(
+            self: _SandboxSnapshot, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]
+        ):
             # we eagerly hydrate the sandbox snapshot below
             pass
 
         rep = "SandboxSnapshot()"
-        obj = _SandboxSnapshot._from_loader(_load, rep, hydrate_lazily=True)
+        # TODO: use ._new_hydrated instead
+        obj = _SandboxSnapshot._from_loader(_load, rep, hydrate_lazily=True, load_context_overrides=LoadContext.empty())
         obj._hydrate(snapshot_id, self._client, None)
 
         return obj
@@ -911,7 +1039,7 @@ class _Sandbox(_Object, type_prefix="sb"):
         client = client or await _Client.from_env()
 
         if name is not None and name != _DEFAULT_SANDBOX_NAME_OVERRIDE:
-            _warn_if_invalid_name(name)
+            check_object_name(name, "Sandbox")
 
         if name is _DEFAULT_SANDBOX_NAME_OVERRIDE:
             restore_req = api_pb2.SandboxRestoreRequest(
@@ -929,21 +1057,14 @@ class _Sandbox(_Object, type_prefix="sb"):
                 sandbox_name_override=name,
                 sandbox_name_override_type=api_pb2.SandboxRestoreRequest.SANDBOX_NAME_OVERRIDE_TYPE_STRING,
             )
-        try:
-            restore_resp: api_pb2.SandboxRestoreResponse = await retry_transient_errors(
-                client.stub.SandboxRestore, restore_req
-            )
-        except GRPCError as exc:
-            if exc.status == Status.ALREADY_EXISTS:
-                raise AlreadyExistsError(exc.message)
-            raise exc
+        restore_resp: api_pb2.SandboxRestoreResponse = await client.stub.SandboxRestore(restore_req)
 
         sandbox = await _Sandbox.from_id(restore_resp.sandbox_id, client)
 
         task_id_req = api_pb2.SandboxGetTaskIdRequest(
             sandbox_id=restore_resp.sandbox_id, wait_until_ready=True, timeout=55.0
         )
-        resp = await retry_transient_errors(client.stub.SandboxGetTaskId, task_id_req)
+        resp = await client.stub.SandboxGetTaskId(task_id_req)
         if resp.task_result.status not in [
             api_pb2.GenericResult.GENERIC_STATUS_UNSPECIFIED,
             api_pb2.GenericResult.GENERIC_STATUS_SUCCESS,
@@ -986,7 +1107,7 @@ class _Sandbox(_Object, type_prefix="sb"):
         task_id = await self._get_task_id()
         return await _FileIO.create(path, mode, self._client, task_id)
 
-    async def ls(self, path: str) -> list[str]:
+    async def ls(self, path: str) -> builtins.list[str]:
         """[Alpha] List the contents of a directory in the Sandbox."""
         task_id = await self._get_task_id()
         return await _FileIO.ls(path, self._client, task_id)
@@ -1004,7 +1125,7 @@ class _Sandbox(_Object, type_prefix="sb"):
     async def watch(
         self,
         path: str,
-        filter: Optional[list[FileWatchEventType]] = None,
+        filter: Optional[builtins.list[FileWatchEventType]] = None,
         recursive: Optional[bool] = None,
         timeout: Optional[int] = None,
     ) -> AsyncIterator[FileWatchEvent]:
@@ -1077,10 +1198,7 @@ class _Sandbox(_Object, type_prefix="sb"):
             )
 
             # Fetches a batch of sandboxes.
-            try:
-                resp = await retry_transient_errors(client.stub.SandboxList, req)
-            except GRPCError as exc:
-                raise InvalidError(exc.message) if exc.status == Status.INVALID_ARGUMENT else exc
+            resp = await client.stub.SandboxList(req)
 
             if not resp.sandboxes:
                 return