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

modal/cli/shell.py

@@ -0,0 +1,375 @@
+# Copyright Modal Labs 2022
+import inspect
+import platform
+import shlex
+from pathlib import Path, PurePosixPath
+from typing import Any, Callable, Optional
+
+import typer
+from click import ClickException
+
+from .._functions import _FunctionSpec
+from ..app import App
+from ..environments import ensure_env
+from ..exception import InvalidError, NotFoundError
+from ..functions import Function
+from ..image import Image
+from ..mount import _Mount
+from ..runner import interactive_shell
+from ..sandbox import Sandbox
+from ..secret import Secret
+from ..volume import Volume
+from .container import exec
+from .import_refs import (
+    MethodReference,
+    import_and_filter,
+    parse_import_ref,
+)
+from .run import _get_runnable_list
+from .utils import ENV_OPTION, is_tty
+
+
+def _params_from_signature(
+    func: Callable[..., Any],
+) -> dict[str, typer.models.ParameterInfo]:
+    sig = inspect.signature(func)
+    params = {param_name: param.default for param_name, param in sig.parameters.items()}
+    assert all(isinstance(param, typer.models.ParameterInfo) for param in params.values()), (
+        f"All params to {func.__name__} must be of type typer.models.ParameterInfo."
+    )
+    return params
+
+
+def _passed_forbidden_args(
+    param_objs: dict[str, typer.models.ParameterInfo],
+    passed_args: dict[str, Any],
+    allowed: Callable[[str], bool],
+) -> list[str]:
+    """Check which forbidden arguments were passed with non-default values."""
+    passed_forbidden: list[str] = []
+    for param_name, param_obj in param_objs.items():
+        if allowed(param_name):
+            continue
+
+        assert param_obj.param_decls is not None, "All params must be typer.models.ParameterInfo, and have param_decls."
+
+        if passed_args.get(param_name) != param_obj.default:
+            passed_forbidden.append("/".join(param_obj.param_decls))
+
+    return passed_forbidden
+
+
+def _is_valid_modal_id(ref: str, prefix: str) -> bool:
+    assert prefix.endswith("-")
+    return ref.startswith(prefix) and len(ref[len(prefix) :]) > 0 and ref[len(prefix) :].isalnum()
+
+
+def _is_running_container_ref(ref: Optional[str]) -> bool:
+    if ref is None:
+        return False
+    return _is_valid_modal_id(ref, "sb-") or _is_valid_modal_id(ref, "ta-")
+
+
+def _start_shell_in_running_container(ref: str, cmd: str, pty: bool) -> None:
+    if _is_valid_modal_id(ref, "sb-"):
+        try:
+            sandbox = Sandbox.from_id(ref)
+            ref = sandbox._get_task_id()
+        except NotFoundError as e:
+            raise ClickException(f"Sandbox '{ref}' not found (is it still running?)")
+        except Exception as e:
+            raise ClickException(f"Error connecting to Sandbox '{ref}': {str(e)}")
+
+    assert _is_valid_modal_id(ref, "ta-")
+    try:
+        exec(container_id=ref, command=shlex.split(cmd), pty=pty)
+    except NotFoundError as e:
+        raise ClickException(f"Container '{ref}' not found (is it still running?)")
+    except Exception as e:
+        raise ClickException(f"Error connecting to container '{ref}': {str(e)}")
+
+
+def _function_spec_from_ref(ref: str, use_module_mode: bool) -> _FunctionSpec:
+    import_ref = parse_import_ref(ref, use_module_mode=use_module_mode)
+    runnable, all_usable_commands = import_and_filter(
+        import_ref, base_cmd="modal shell", accept_local_entrypoint=False, accept_webhook=True
+    )
+    if not runnable:
+        help_header = (
+            "Specify a Modal function to start a shell session for. E.g.\n"
+            f"> modal shell {import_ref.file_or_module}::my_function"
+        )
+
+        if all_usable_commands:
+            help_footer = f"The selected module '{import_ref.file_or_module}' has the following choices:\n\n"
+            help_footer += _get_runnable_list(all_usable_commands)
+        else:
+            help_footer = f"The selected module '{import_ref.file_or_module}' has no Modal functions or classes."
+
+        raise ClickException(f"{help_header}\n\n{help_footer}")
+
+    if isinstance(runnable, MethodReference):
+        # TODO: let users specify a class instead of a method, since they use the same environment
+        class_service_function = runnable.cls._get_class_service_function()
+        return class_service_function.spec
+    elif isinstance(runnable, Function):
+        return runnable.spec
+
+    raise ValueError("Referenced entity is not a Modal Function or Cls")
+
+
+def _start_shell_from_function_spec(
+    app: App,
+    cmds: list[str],
+    env: str,
+    timeout: int,
+    function_spec: _FunctionSpec,
+    pty: bool,
+) -> None:
+    interactive_shell(
+        app,
+        cmds=cmds,
+        environment_name=env,
+        timeout=timeout,
+        image=function_spec.image,
+        mounts=function_spec.mounts,
+        secrets=function_spec.secrets,
+        network_file_systems=function_spec.network_file_systems,
+        gpu=function_spec.gpus,
+        cloud=function_spec.cloud,
+        cpu=function_spec.cpu,
+        memory=function_spec.memory,
+        volumes=function_spec.volumes,
+        region=function_spec.scheduler_placement.regions if function_spec.scheduler_placement else None,
+        pty=pty,
+        proxy=function_spec.proxy,
+    )
+
+
+def _start_shell_from_image(
+    app: App,
+    cmds: list[str],
+    env: str,
+    timeout: int,
+    modal_image: Optional[Image],
+    volume: list[str],
+    secret: list[str],
+    add_local: list[str],
+    cpu: Optional[int],
+    memory: Optional[int],
+    gpu: Optional[str],
+    cloud: Optional[str],
+    region: Optional[str],
+    pty: bool,
+) -> None:
+    volumes = {f"/mnt/{vol}": Volume.from_name(vol) for vol in volume}
+    secrets = [Secret.from_name(s) for s in secret]
+
+    mounts = []
+    for local_path_str in add_local:
+        local_path = Path(local_path_str).expanduser().resolve()
+        remote_path = PurePosixPath(f"/mnt/{local_path.name}")
+
+        if local_path.is_dir():
+            m = _Mount._from_local_dir(local_path, remote_path=remote_path)
+        else:
+            m = _Mount._from_local_file(local_path, remote_path=remote_path)
+        mounts.append(m)
+
+    interactive_shell(
+        app,
+        cmds=cmds,
+        environment_name=env,
+        timeout=timeout,
+        image=modal_image,
+        mounts=mounts,
+        cpu=cpu,
+        memory=memory,
+        gpu=gpu,
+        cloud=cloud,
+        volumes=volumes,
+        secrets=secrets,
+        region=region.split(",") if region else [],
+        pty=pty,
+    )
+
+
+def shell(
+    ref: Optional[str] = typer.Argument(
+        default=None,
+        help=(
+            "ID of running container or Sandbox, or path to a Python file containing an App."
+            " Can also include a Function specifier, like `module.py::func`, if the file defines multiple Functions."
+        ),
+    ),
+    cmd: str = typer.Option("/bin/bash", "-c", "--cmd", help="Command to run inside the Modal image."),
+    env: Optional[str] = ENV_OPTION,
+    image: Optional[str] = typer.Option(
+        None, "--image", help="Container image tag for inside the shell (if not using REF)."
+    ),
+    add_python: Optional[str] = typer.Option(None, "--add-python", help="Add Python to the image (if not using REF)."),
+    volume: Optional[list[str]] = typer.Option(
+        None,
+        "--volume",
+        help=(
+            "Name of a `modal.Volume` to mount inside the shell at `/mnt/{name}` (if not using REF)."
+            " Can be used multiple times."
+        ),
+    ),
+    add_local: Optional[list[str]] = typer.Option(
+        None,
+        "--add-local",
+        help=(
+            "Local file or directory to mount inside the shell at `/mnt/{basename}` (if not using REF)."
+            " Can be used multiple times."
+        ),
+    ),
+    secret: Optional[list[str]] = typer.Option(
+        None,
+        "--secret",
+        help=("Name of a `modal.Secret` to mount inside the shell (if not using REF). Can be used multiple times."),
+    ),
+    cpu: Optional[int] = typer.Option(
+        None, "--cpu", help="Number of CPUs to allocate to the shell (if not using REF)."
+    ),
+    memory: Optional[int] = typer.Option(
+        None, "--memory", help="Memory to allocate for the shell, in MiB (if not using REF)."
+    ),
+    gpu: Optional[str] = typer.Option(
+        None,
+        "--gpu",
+        help="GPUs to request for the shell, if any. Examples are `any`, `a10g`, `a100:4` (if not using REF).",
+    ),
+    cloud: Optional[str] = typer.Option(
+        None,
+        "--cloud",
+        help=(
+            "Cloud provider to run the shell on. Possible values are `aws`, `gcp`, `oci`, `auto` (if not using REF)."
+        ),
+    ),
+    region: Optional[str] = typer.Option(
+        None,
+        "--region",
+        help=(
+            "Region(s) to run the container on. "
+            "Can be a single region or a comma-separated list to choose from (if not using REF)."
+        ),
+    ),
+    pty: Optional[bool] = typer.Option(None, "--pty", help="Run the command using a PTY."),
+    use_module_mode: bool = typer.Option(
+        False, "-m", help="Interpret argument as a Python module path instead of a file/script path"
+    ),
+):
+    """Run a command or interactive shell inside a Modal container.
+
+    **Examples:**
+
+    Start an interactive shell inside the default Debian-based image:
+
+    ```
+    modal shell
+    ```
+
+    Start an interactive shell with the spec for `my_function` in your App
+    (uses the same image, volumes, mounts, etc.):
+
+    ```
+    modal shell hello_world.py::my_function
+    ```
+
+    Or, if you're using a [modal.Cls](https://modal.com/docs/reference/modal.Cls)
+    you can refer to a `@modal.method` directly:
+
+    ```
+    modal shell hello_world.py::MyClass.my_method
+    ```
+
+    Start a `python` shell:
+
+    ```
+    modal shell hello_world.py --cmd=python
+    ```
+
+    Run a command with your function's spec and pipe the output to a file:
+
+    ```
+    modal shell hello_world.py -c 'uv pip list' > env.txt
+    ```
+
+    Connect to a running Sandbox by ID:
+
+    ```
+    modal shell sb-abc123xyz
+    ```
+    """
+    if pty is None:
+        pty = is_tty()
+
+    if platform.system() == "Windows":
+        raise InvalidError("`modal shell` is currently not supported on Windows")
+
+    param_objs = _params_from_signature(shell)
+
+    if ref is not None and _is_running_container_ref(ref):
+        # We're attaching to an already running container or Sandbox.
+        if passed_forbidden := _passed_forbidden_args(
+            param_objs, locals(), allowed=lambda p: p in {"cmd", "pty", "ref"}
+        ):
+            raise ClickException(
+                f"Cannot specify container configuration arguments ({', '.join(passed_forbidden)}) "
+                f"when attaching to an already running container or Sandbox ('{ref}')."
+            )
+
+        _start_shell_in_running_container(ref, cmd, pty)
+        return
+
+    # We're not attaching to an existing container, so we need to create a new one.
+    env = ensure_env(env)
+    app = App("modal shell")
+
+    # NB: invoking under bash makes --cmd a lot more flexible.
+    cmds = shlex.split(f'/bin/bash -c "{cmd}"')
+    timeout = 3600
+
+    if ref is not None and not _is_valid_modal_id(ref, "im-"):
+        # If ref it not a Modal Image ID, then it's a function reference, and we'll start a new container from its spec.
+        if passed_forbidden := _passed_forbidden_args(
+            param_objs, locals(), allowed=lambda p: p in {"cmd", "env", "pty", "ref", "use_module_mode"}
+        ):
+            raise ClickException(
+                f"Cannot specify container configuration arguments ({', '.join(passed_forbidden)}) "
+                f"when starting a new container from a function reference ('{ref}')."
+            )
+
+        function_spec = _function_spec_from_ref(ref, use_module_mode)
+        _start_shell_from_function_spec(app, cmds, env, timeout, function_spec, pty)
+        return
+
+    if ref is not None and _is_valid_modal_id(ref, "im-"):
+        if passed_forbidden := _passed_forbidden_args(
+            param_objs, locals(), allowed=lambda p: p not in {"add_python", "image"}
+        ):
+            raise ClickException(
+                f"Cannot specify {', '.join(passed_forbidden)} argument(s) "
+                f"when starting a new container from a Modal Image ID ('{ref}')."
+            )
+        modal_image = Image.from_id(ref)
+    else:
+        modal_image = Image.from_registry(image, add_python=add_python) if image else None
+
+    _start_shell_from_image(
+        app,
+        cmds,
+        env,
+        timeout,
+        modal_image,
+        volume or [],
+        secret or [],
+        add_local or [],
+        cpu,
+        memory,
+        gpu,
+        cloud,
+        region,
+        pty,
+    )

modal/cli/utils.py

@@ -5,8 +5,6 @@ from json import dumps
 from typing import Optional, Union
 
 import typer
-from click import UsageError
-from grpclib import GRPCError, Status
 from rich.table import Column, Table
 from rich.text import Text
 
@@ -33,11 +31,6 @@ async def stream_app_logs(
             await get_app_logs_loop(client, output_mgr, app_id=app_id, task_id=task_id, app_logs_url=app_logs_url)
     except asyncio.CancelledError:
         pass
-    except GRPCError as exc:
-        if exc.status in (Status.INVALID_ARGUMENT, Status.NOT_FOUND):
-            raise UsageError(exc.message)
-        else:
-            raise
     except KeyboardInterrupt:
         pass
 
@@ -48,12 +41,7 @@ async def get_app_id_from_name(name: str, env: Optional[str], client: Optional[_
         client = await _Client.from_env()
     env_name = ensure_env(env)
     request = api_pb2.AppGetByDeploymentNameRequest(name=name, environment_name=env_name)
-    try:
-        resp = await client.stub.AppGetByDeploymentName(request)
-    except GRPCError as exc:
-        if exc.status in (Status.INVALID_ARGUMENT, Status.NOT_FOUND):
-            raise UsageError(exc.message or "")
-        raise
+    resp = await client.stub.AppGetByDeploymentName(request)
     if not resp.app_id:
         env_comment = f" in the '{env_name}' environment" if env_name else ""
         raise NotFoundError(f"Could not find a deployed app named '{name}'{env_comment}.")

modal/cli/volume.py

@@ -6,7 +6,6 @@ from typing import Optional
 
 import typer
 from click import UsageError
-from grpclib import GRPCError, Status
 from rich.syntax import Syntax
 from typer import Argument, Option, Typer
 
@@ -96,7 +95,13 @@ async def get(
     console = make_console()
     progress_handler = ProgressHandler(type="download", console=console)
     with progress_handler.live:
-        await _volume_download(volume, remote_path, destination, force, progress_cb=progress_handler.progress)
+        await _volume_download(
+            volume=volume,
+            remote_path=remote_path,
+            local_destination=destination,
+            overwrite=force,
+            progress_cb=progress_handler.progress,
+        )
     console.print(OutputManager.step_completed("Finished downloading files to local!"))
 
 
@@ -131,18 +136,12 @@ async def ls(
 ):
     ensure_env(env)
     vol = _Volume.from_name(volume_name, environment_name=env)
-
-    try:
-        entries = await vol.listdir(path)
-    except GRPCError as exc:
-        if exc.status in (Status.INVALID_ARGUMENT, Status.NOT_FOUND):
-            raise UsageError(exc.message)
-        raise
+    entries = await vol.listdir(path)
 
     if not json and not sys.stdout.isatty():
         # Legacy behavior -- I am not sure why exactly we did this originally but I don't want to break it
         for entry in entries:
-            print(entry.path)
+            print(entry.path)  # noqa: T201
     else:
         rows = []
         for entry in entries:
@@ -241,14 +240,9 @@ async def rm(
 ):
     ensure_env(env)
     volume = _Volume.from_name(volume_name, environment_name=env)
+    await volume.remove_file(remote_path, recursive=recursive)
     console = make_console()
-    try:
-        await volume.remove_file(remote_path, recursive=recursive)
-        console.print(OutputManager.step_completed(f"{remote_path} was deleted successfully!"))
-    except GRPCError as exc:
-        if exc.status in (Status.NOT_FOUND, Status.INVALID_ARGUMENT):
-            raise UsageError(exc.message)
-        raise
+    console.print(OutputManager.step_completed(f"{remote_path} was deleted successfully!"))
 
 
 @volume_cli.command(

modal/client.py

@@ -5,33 +5,25 @@ import platform
 import sys
 import urllib.parse
 import warnings
-from collections.abc import AsyncGenerator, AsyncIterator, Collection, Mapping
-from typing import (
-    Any,
-    ClassVar,
-    Generic,
-    Optional,
-    TypeVar,
-    Union,
-)
+from collections.abc import AsyncGenerator, Collection, Mapping
+from typing import Any, ClassVar, Optional, TypeVar, Union
 
 import grpclib.client
 from google.protobuf import empty_pb2
 from google.protobuf.message import Message
-from grpclib import GRPCError, Status
 from synchronicity.async_wrap import asynccontextmanager
 
 from modal._utils.async_utils import synchronizer
-from modal_proto import api_grpc, api_pb2, modal_api_grpc
+from modal_proto import api_pb2, modal_api_grpc
 from modal_version import __version__
 
-from ._traceback import print_server_warnings, suppress_tb_frames
+from ._traceback import print_server_warnings
 from ._utils import async_utils
 from ._utils.async_utils import TaskContext, synchronize_api
 from ._utils.auth_token_manager import _AuthTokenManager
-from ._utils.grpc_utils import ConnectionManager, retry_transient_errors
+from ._utils.grpc_utils import ConnectionManager
 from .config import _check_config, _is_remote, config, logger
-from .exception import AuthError, ClientClosed, NotFoundError
+from .exception import AuthError, ClientClosed
 
 HEARTBEAT_INTERVAL: float = config.get("heartbeat_interval")
 HEARTBEAT_TIMEOUT: float = HEARTBEAT_INTERVAL + 0.1
@@ -78,15 +70,16 @@ class _Client:
     _client_from_env: ClassVar[Optional["_Client"]] = None
     _client_from_env_lock: ClassVar[Optional[asyncio.Lock]] = None
     _cancellation_context: TaskContext
-    _cancellation_context_event_loop: asyncio.AbstractEventLoop = None
-    _stub: Optional[api_grpc.ModalClientStub]
-    _auth_token_manager: _AuthTokenManager = None
-    _snapshotted: bool
+    _cancellation_context_event_loop: Optional[asyncio.AbstractEventLoop] = None
+    _stub: Optional[modal_api_grpc.ModalClientModal] = None
+    _auth_token_manager: Optional[_AuthTokenManager] = None
+    _snapshotted: bool = False
+    client_type: "api_pb2.ClientType.ValueType"
 
     def __init__(
         self,
         server_url: str,
-        client_type: int,
+        client_type: "api_pb2.ClientType.ValueType",
         credentials: Optional[tuple[str, str]],
         version: str = __version__,
     ):
@@ -98,8 +91,8 @@ class _Client:
         self._credentials = credentials
         self.version = version
         self._closed = False
-        self._stub: Optional[modal_api_grpc.ModalClientModal] = None
-        self._auth_token_manager: Optional[_AuthTokenManager] = None
+        self._stub = None
+        self._auth_token_manager = None
         self._snapshotted = False
         self._owner_pid = None
 
@@ -159,7 +152,7 @@ class _Client:
     async def hello(self):
         """Connect to server and retrieve version information; raise appropriate error for various failures."""
         logger.debug(f"Client ({id(self)}): Starting")
-        resp = await retry_transient_errors(self.stub.ClientHello, empty_pb2.Empty())
+        resp = await self.stub.ClientHello(empty_pb2.Empty())
         print_server_warnings(resp.server_warnings)
 
     async def __aenter__(self):
@@ -171,7 +164,7 @@ class _Client:
 
     @classmethod
     @asynccontextmanager
-    async def anonymous(cls, server_url: str) -> AsyncIterator["_Client"]:
+    async def anonymous(cls, server_url: str) -> AsyncGenerator["_Client", None]:
         """mdmd:hidden
         Create a connection with no credentials; to be used for token creation.
         """
@@ -362,105 +355,3 @@ class _Client:
 
 
 Client = synchronize_api(_Client)
-
-
-class grpc_error_converter:
-    def __enter__(self):
-        pass
-
-    def __exit__(self, exc_type, exc, traceback) -> bool:
-        # skip all internal frames from grpclib
-        use_full_traceback = config.get("traceback")
-        with suppress_tb_frames(1):
-            if isinstance(exc, GRPCError):
-                if exc.status == Status.NOT_FOUND:
-                    if use_full_traceback:
-                        raise NotFoundError(exc.message)
-                    else:
-                        raise NotFoundError(exc.message) from None  # from None to skip the grpc-internal cause
-
-                if not use_full_traceback:
-                    # just include the frame in grpclib that actually raises the GRPCError
-                    tb = exc.__traceback__
-                    while tb.tb_next:
-                        tb = tb.tb_next
-                    exc.with_traceback(tb)
-                    raise exc from None  # from None to skip the grpc-internal cause
-                raise exc
-
-        return False
-
-
-class UnaryUnaryWrapper(Generic[RequestType, ResponseType]):
-    # Calls a grpclib.UnaryUnaryMethod using a specific Client instance, respecting
-    # if that client is closed etc. and possibly introducing Modal-specific retry logic
-    wrapped_method: grpclib.client.UnaryUnaryMethod[RequestType, ResponseType]
-    client: _Client
-
-    def __init__(
-        self,
-        wrapped_method: grpclib.client.UnaryUnaryMethod[RequestType, ResponseType],
-        client: _Client,
-        server_url: str,
-    ):
-        self.wrapped_method = wrapped_method
-        self.client = client
-        self.server_url = server_url
-
-    @property
-    def name(self) -> str:
-        return self.wrapped_method.name
-
-    async def __call__(
-        self,
-        req: RequestType,
-        *,
-        timeout: Optional[float] = None,
-        metadata: Optional[_MetadataLike] = None,
-    ) -> ResponseType:
-        if self.client._snapshotted:
-            logger.debug(f"refreshing client after snapshot for {self.name.rsplit('/', 1)[1]}")
-            self.client = await _Client.from_env()
-
-        # Note: We override the grpclib method's channel (see grpclib's code [1]). I think this is fine
-        # since grpclib's code doesn't seem to change very much, but we could also recreate the
-        # grpclib stub if we aren't comfortable with this. The downside is then we need to cache
-        # the grpclib stub so the rest of our code becomes a bit more complicated.
-        #
-        # We need to override the channel because after the process is forked or the client is
-        # snapshotted, the existing channel may be stale / unusable.
-        #
-        # [1]: https://github.com/vmagamedov/grpclib/blob/62f968a4c84e3f64e6966097574ff0a59969ea9b/grpclib/client.py#L844
-        self.wrapped_method.channel = await self.client._get_channel(self.server_url)
-        with suppress_tb_frames(1), grpc_error_converter():
-            return await self.client._call_unary(self.wrapped_method, req, timeout=timeout, metadata=metadata)
-
-
-class UnaryStreamWrapper(Generic[RequestType, ResponseType]):
-    wrapped_method: grpclib.client.UnaryStreamMethod[RequestType, ResponseType]
-
-    def __init__(
-        self,
-        wrapped_method: grpclib.client.UnaryStreamMethod[RequestType, ResponseType],
-        client: _Client,
-        server_url: str,
-    ):
-        self.wrapped_method = wrapped_method
-        self.client = client
-        self.server_url = server_url
-
-    @property
-    def name(self) -> str:
-        return self.wrapped_method.name
-
-    async def unary_stream(
-        self,
-        request,
-        metadata: Optional[Any] = None,
-    ):
-        if self.client._snapshotted:
-            logger.debug(f"refreshing client after snapshot for {self.name.rsplit('/', 1)[1]}")
-            self.client = await _Client.from_env()
-        self.wrapped_method.channel = await self.client._get_channel(self.server_url)
-        async for response in self.client._call_stream(self.wrapped_method, request, metadata=metadata):
-            yield response