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

modal/_grpc_client.py

@@ -0,0 +1,191 @@
+# Copyright Modal Labs 2025
+from typing import TYPE_CHECKING, Any, Collection, Generic, Literal, Mapping, Optional, TypeVar, Union, overload
+
+import grpclib.client
+from google.protobuf.message import Message
+from grpclib import GRPCError, Status
+
+from . import exception
+from ._traceback import suppress_tb_frame
+from ._utils.grpc_utils import Retry, _retry_transient_errors
+from .config import config, logger
+
+if TYPE_CHECKING:
+    from .client import _Client
+
+
+_Value = Union[str, bytes]
+_MetadataLike = Union[Mapping[str, _Value], Collection[tuple[str, _Value]]]
+RequestType = TypeVar("RequestType", bound=Message)
+ResponseType = TypeVar("ResponseType", bound=Message)
+
+
+class WrappedGRPCError(exception.Error, exception._GRPCErrorWrapper): ...
+
+
+_STATUS_TO_EXCEPTION: dict[Status, type[exception._GRPCErrorWrapper]] = {
+    Status.CANCELLED: exception.ServiceError,
+    Status.UNKNOWN: exception.ServiceError,
+    Status.INVALID_ARGUMENT: exception.InvalidError,
+    Status.DEADLINE_EXCEEDED: exception.ServiceError,
+    Status.NOT_FOUND: exception.NotFoundError,
+    Status.ALREADY_EXISTS: exception.AlreadyExistsError,
+    Status.PERMISSION_DENIED: exception.PermissionDeniedError,
+    Status.RESOURCE_EXHAUSTED: exception.ResourceExhaustedError,
+    Status.FAILED_PRECONDITION: exception.ConflictError,
+    Status.ABORTED: exception.ConflictError,
+    Status.OUT_OF_RANGE: exception.InvalidError,
+    Status.UNIMPLEMENTED: exception.UnimplementedError,
+    Status.INTERNAL: exception.InternalError,
+    Status.UNAVAILABLE: exception.ServiceError,
+    Status.DATA_LOSS: exception.DataLossError,
+    Status.UNAUTHENTICATED: exception.AuthError,
+}
+
+
+class grpc_error_converter:
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc, traceback) -> Literal[False]:
+        # skip all internal frames from grpclib
+        use_full_traceback = config.get("traceback")
+        with suppress_tb_frame():
+            if isinstance(exc, GRPCError):
+                modal_exc = _STATUS_TO_EXCEPTION[exc.status](exc.message)
+                modal_exc._grpc_message = exc.message
+                modal_exc._grpc_status = exc.status
+                modal_exc._grpc_details = exc.details
+                if use_full_traceback:
+                    raise modal_exc
+                else:
+                    raise modal_exc from None  # from None to skip the grpc-internal cause
+
+        return False
+
+
+_DEFAULT_RETRY = Retry()
+
+
+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
+
+    @overload
+    async def __call__(
+        self,
+        req: RequestType,
+        *,
+        retry: Retry = _DEFAULT_RETRY,
+        timeout: None = None,
+        metadata: Optional[list[tuple[str, str]]] = None,
+    ) -> ResponseType: ...
+
+    @overload
+    async def __call__(
+        self,
+        req: RequestType,
+        *,
+        retry: None,
+        timeout: Optional[float] = None,
+        metadata: Optional[list[tuple[str, str]]] = None,
+    ) -> ResponseType: ...
+
+    async def __call__(
+        self,
+        req: RequestType,
+        *,
+        retry: Optional[Retry] = _DEFAULT_RETRY,
+        timeout: Optional[float] = None,
+        metadata: Optional[list[tuple[str, str]]] = None,
+    ) -> ResponseType:
+        with suppress_tb_frame():
+            if timeout is not None and retry is not None:
+                raise exception.InvalidError("Retry must be None when timeout is set")
+
+            if retry is None:
+                with grpc_error_converter():
+                    return await self.direct(req, timeout=timeout, metadata=metadata)
+
+            # TODO do we need suppress_error_frames(1) here too?
+            with grpc_error_converter():
+                return await _retry_transient_errors(
+                    self,  # type: ignore
+                    req,
+                    retry=retry,
+                    metadata=metadata,
+                )
+
+    async def direct(
+        self,
+        req: RequestType,
+        *,
+        timeout: Optional[float] = None,
+        metadata: Optional[_MetadataLike] = None,
+    ) -> ResponseType:
+        from .client import _Client
+
+        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)
+        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,
+    ):
+        from .client import _Client
+
+        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)
+        with grpc_error_converter():
+            async for response in self.client._call_stream(self.wrapped_method, request, metadata=metadata):
+                yield response

modal/_ipython.py

@@ -2,10 +2,20 @@
 import sys
 
 
-def is_notebook(stdout=None):
-    ipykernel_iostream = sys.modules.get("ipykernel.iostream")
-    if ipykernel_iostream is None:
+def is_interactive_ipython():
+    """
+    Detect if we're running in an interactive IPython session.
+
+    Returns True for IPython shells (including Jupyter notebooks), False otherwise.
+    """
+    try:
+        # Check if IPython is available and get the current instance
+        ipython = sys.modules.get("IPython")
+        if ipython is None:
+            return False
+
+        # Try to get the active IPython instance
+        shell = ipython.get_ipython()
+        return shell is not None
+    except Exception:
         return False
-    if stdout is None:
-        stdout = sys.stdout
-    return isinstance(stdout, ipykernel_iostream.OutStream)

modal/_load_context.py

@@ -0,0 +1,106 @@
+# Copyright Modal Labs 2025
+from typing import Optional
+
+from .client import _Client
+from .config import config
+
+
+class LoadContext:
+    """Encapsulates optional metadata values used during object loading.
+
+    This metadata is set during object construction and propagated through
+    parent-child relationships (e.g., App -> Function, Cls -> Obj -> bound methods).
+    """
+
+    _client: Optional[_Client] = None
+    _environment_name: Optional[str] = None
+    _app_id: Optional[str] = None
+
+    def __init__(
+        self,
+        *,
+        client: Optional[_Client] = None,
+        environment_name: Optional[str] = None,
+        app_id: Optional[str] = None,
+    ):
+        self._client = client
+        self._environment_name = environment_name
+        self._app_id = app_id
+
+    @property
+    def client(self) -> _Client:
+        assert self._client is not None
+        return self._client
+
+    @property
+    def environment_name(self) -> str:
+        assert self._environment_name is not None
+        return self._environment_name
+
+    @property
+    def app_id(self) -> Optional[str]:
+        return self._app_id
+
+    @classmethod
+    def empty(cls) -> "LoadContext":
+        """Create an empty LoadContext with all fields set to None.
+
+        Used when loading objects that don't have a parent context.
+        """
+        return cls(client=None, environment_name=None, app_id=None)
+
+    def merged_with(self, parent: "LoadContext") -> "LoadContext":
+        """Create a new LoadContext with parent values filling in None fields.
+
+        Returns a new LoadContext without mutating self or parent.
+        Values from self take precedence over values from parent.
+        """
+        return LoadContext(
+            client=self._client if self._client is not None else parent._client,
+            environment_name=self._environment_name if self._environment_name is not None else parent._environment_name,
+            app_id=self._app_id if self._app_id is not None else parent._app_id,
+        )  # TODO (elias):  apply_defaults?
+
+    async def apply_defaults(self) -> "LoadContext":
+        """Infer default client and environment_name if not present
+
+        Returns a new instance (no in place mutation)"""
+
+        is_valid_client = self._client is not None and not self._client._snapshotted
+        return LoadContext(
+            client=self.client if is_valid_client else await _Client.from_env(),
+            environment_name=self._environment_name or config.get("environment") or "",
+            app_id=self._app_id,
+        )
+
+    def reset(self) -> "LoadContext":
+        self._client = None
+        self._environment_name = None
+        self._app_id = None
+        return self
+
+    async def in_place_upgrade(
+        self, client: Optional[_Client] = None, environment_name: Optional[str] = None, app_id: Optional[str] = None
+    ) -> "LoadContext":
+        """In-place set values if they aren't already set, or set default values
+
+        Intended for Function/Cls hydration specifically
+
+        In those cases, it's important to in-place upgrade/apply_defaults since any "sibling" of the function/cls
+        would share the load context with its parent, and the initial load context overrides may not be sufficient
+        since an `app.deploy()` etc could get arguments that set a new client etc.
+
+        E.g.
+        @app.function()
+        def f():
+            ...
+
+        f2 = Function.with_options(...)
+
+        with app.run(client=...): # hydrates f and f2 at this point
+            ...
+        """
+        self._client = self._client or client or await _Client.from_env()
+        self._environment_name = self._environment_name or environment_name or config.get("environment") or ""
+        self._app_id = self._app_id or app_id
+        return self

modal/_object.py

@@ -1,4 +1,5 @@
 # Copyright Modal Labs 2022
+import contextlib
 import typing
 import uuid
 from collections.abc import Awaitable, Hashable, Sequence
@@ -8,8 +9,9 @@ from typing import Callable, ClassVar, Optional
 from google.protobuf.message import Message
 from typing_extensions import Self
 
-from modal._traceback import suppress_tb_frames
+from modal._traceback import suppress_tb_frame
 
+from ._load_context import LoadContext
 from ._resolver import Resolver
 from ._utils.async_utils import aclosing
 from ._utils.deprecation import deprecation_warning
@@ -20,11 +22,19 @@ from .exception import ExecutionError, InvalidError
 EPHEMERAL_OBJECT_HEARTBEAT_SLEEP: int = 300
 
 
-def _get_environment_name(environment_name: Optional[str] = None, resolver: Optional[Resolver] = None) -> Optional[str]:
+def _get_environment_name(
+    environment_name: Optional[str] = None,
+) -> Optional[str]:
+    """Get environment name from various sources.
+
+    Args:
+        environment_name: Explicitly provided environment name (highest priority)
+
+    Returns:
+        Environment name from first available source, or config default
+    """
     if environment_name:
         return environment_name
-    elif resolver and resolver.environment_name:
-        return resolver.environment_name
     else:
         return config.get("environment")
 
@@ -34,13 +44,14 @@ class _Object:
     _prefix_to_type: ClassVar[dict[str, type]] = {}
 
     # For constructors
-    _load: Optional[Callable[[Self, Resolver, Optional[str]], Awaitable[None]]]
-    _preload: Optional[Callable[[Self, Resolver, Optional[str]], Awaitable[None]]]
+    _load: Optional[Callable[[Self, Resolver, LoadContext, Optional[str]], Awaitable[None]]] = None
+    _preload: Optional[Callable[[Self, Resolver, LoadContext, Optional[str]], Awaitable[None]]]
     _rep: str
     _is_another_app: bool
     _hydrate_lazily: bool
     _deps: Optional[Callable[..., Sequence["_Object"]]]
     _deduplication_key: Optional[Callable[[], Awaitable[Hashable]]] = None
+    _load_context_overrides: LoadContext
 
     # For hydrated objects
     _object_id: Optional[str]
@@ -66,13 +77,15 @@ class _Object:
     def _init(
         self,
         rep: str,
-        load: Optional[Callable[[Self, Resolver, Optional[str]], Awaitable[None]]] = None,
+        load: Optional[Callable[[Self, Resolver, LoadContext, Optional[str]], Awaitable[None]]] = None,
         is_another_app: bool = False,
-        preload: Optional[Callable[[Self, Resolver, Optional[str]], Awaitable[None]]] = None,
+        preload: Optional[Callable[[Self, Resolver, LoadContext, Optional[str]], Awaitable[None]]] = None,
         hydrate_lazily: bool = False,
         deps: Optional[Callable[..., Sequence["_Object"]]] = None,
         deduplication_key: Optional[Callable[[], Awaitable[Hashable]]] = None,
         name: Optional[str] = None,
+        *,
+        load_context_overrides: Optional[LoadContext] = None,
     ):
         self._local_uuid = str(uuid.uuid4())
         self._load = load
@@ -82,6 +95,9 @@ class _Object:
         self._hydrate_lazily = hydrate_lazily
         self._deps = deps
         self._deduplication_key = deduplication_key
+        self._load_context_overrides = (
+            load_context_overrides if load_context_overrides is not None else LoadContext.empty()
+        )
 
         self._object_id = None
         self._client = None
@@ -163,18 +179,30 @@ class _Object:
     @classmethod
     def _from_loader(
         cls,
-        load: Callable[[Self, Resolver, Optional[str]], Awaitable[None]],
+        load: Callable[[Self, Resolver, LoadContext, Optional[str]], Awaitable[None]],
         rep: str,
         is_another_app: bool = False,
-        preload: Optional[Callable[[Self, Resolver, Optional[str]], Awaitable[None]]] = None,
+        preload: Optional[Callable[[Self, Resolver, LoadContext, Optional[str]], Awaitable[None]]] = None,
         hydrate_lazily: bool = False,
         deps: Optional[Callable[..., Sequence["_Object"]]] = None,
         deduplication_key: Optional[Callable[[], Awaitable[Hashable]]] = None,
         name: Optional[str] = None,
+        *,
+        load_context_overrides: LoadContext,
     ):
         # TODO(erikbern): flip the order of the two first arguments
         obj = _Object.__new__(cls)
-        obj._init(rep, load, is_another_app, preload, hydrate_lazily, deps, deduplication_key, name)
+        obj._init(
+            rep,
+            load,
+            is_another_app,
+            preload,
+            hydrate_lazily,
+            deps,
+            deduplication_key,
+            name,
+            load_context_overrides=load_context_overrides,
+        )
         return obj
 
     @staticmethod
@@ -275,25 +303,33 @@ class _Object:
 
         *Added in v0.72.39*: This method replaces the deprecated `.resolve()` method.
         """
+        # TODO: add deprecation for the client argument here - should be added in constructors instead
         if self._is_hydrated:
             if self.client._snapshotted and not self._is_rehydrated:
                 # memory snapshots capture references which must be rehydrated
                 # on restore to handle staleness.
                 logger.debug(f"rehydrating {self} after snapshot")
-                self._is_hydrated = False  # un-hydrate and re-resolve
-                c = client if client is not None else await _Client.from_env()
-                resolver = Resolver(c)
-                await resolver.load(typing.cast(_Object, self))
+                if self._hydrate_lazily:
+                    logger.debug(f"reloading lazy {self} from server")
+                    self._is_hydrated = False  # un-hydrate and re-resolve
+                    # we don't set an explicit Client here, relying on the default
+                    # env client to be applied by LoadContext.apply_default
+                    root_load_context = LoadContext.empty()
+                    resolver = Resolver()
+                    await resolver.load(typing.cast(_Object, self), root_load_context)
+                else:
+                    logger.debug(f"reloading non-lazy {self} by replacing client")
+                    self._client = client or await _Client.from_env()
                 self._is_rehydrated = True
-                logger.debug(f"rehydrated {self} with client {id(c)}")
+                logger.debug(f"rehydrated {self} with client {id(self.client)}")
         elif not self._hydrate_lazily:
-            # TODO(michael) can remove _hydrate lazily? I think all objects support it now?
             self._validate_is_hydrated()
         else:
-            c = client if client is not None else await _Client.from_env()
-            resolver = Resolver(c)
-            with suppress_tb_frames(1):  # skip this frame by default
-                await resolver.load(self)
+            # Set the client on LoadContext before loading
+            root_load_context = LoadContext(client=client)
+            resolver = Resolver()
+            with suppress_tb_frame():  # skip this frame by default
+                await resolver.load(self, root_load_context)
         return self
 
 
@@ -315,3 +351,18 @@ def live_method_gen(method):
                 yield item
 
     return wrapped
+
+
+def live_method_contextmanager(method):
+    # make sure a wrapped function returning an async context manager
+    # will not require both an `await func.aio()` and `async with`
+    # which would have been the case if it was wrapped in live_method
+
+    @wraps(method)
+    @contextlib.asynccontextmanager
+    async def wrapped(self, *args, **kwargs):
+        await self.hydrate()
+        async with method(self, *args, **kwargs) as ctx:
+            yield ctx
+
+    return wrapped

modal/_output.py

@@ -12,7 +12,7 @@ from collections.abc import Generator
 from datetime import timedelta
 from typing import Callable, ClassVar
 
-from grpclib.exceptions import GRPCError, StreamTerminatedError
+from grpclib.exceptions import StreamTerminatedError
 from rich.console import Console, Group, RenderableType
 from rich.live import Live
 from rich.panel import Panel
@@ -34,10 +34,11 @@ from rich.text import Text
 from modal._utils.time_utils import timestamp_to_localized_str
 from modal_proto import api_pb2
 
-from ._utils.grpc_utils import RETRYABLE_GRPC_STATUS_CODES, retry_transient_errors
+from ._utils.grpc_utils import Retry
 from ._utils.shell_utils import stream_from_stdin, write_to_fd
 from .client import _Client
 from .config import logger
+from .exception import InternalError, ServiceError
 
 if platform.system() == "Windows":
     default_spinner = "line"
@@ -489,12 +490,11 @@ async def stream_pty_shell_input(client: _Client, exec_id: str, finish_event: as
     """
 
     async def _handle_input(data: bytes, message_index: int):
-        await retry_transient_errors(
-            client.stub.ContainerExecPutInput,
+        await client.stub.ContainerExecPutInput(
             api_pb2.ContainerExecPutInputRequest(
                 exec_id=exec_id, input=api_pb2.RuntimeInputMessage(message=data, message_index=message_index)
             ),
-            total_timeout=10,
+            retry=Retry(total_timeout=10),
         )
 
     async with stream_from_stdin(_handle_input, use_raw_terminal=True):
@@ -557,7 +557,7 @@ async def get_app_logs_loop(
     async def stop_pty_shell():
         nonlocal pty_shell_finish_event, pty_shell_input_task
         if pty_shell_finish_event:
-            print("\r", end="")  # move cursor to beginning of line
+            print("\r", end="")  # move cursor to beginning of line # noqa: T201
             pty_shell_finish_event.set()
             pty_shell_finish_event = None
 
@@ -624,7 +624,7 @@ async def get_app_logs_loop(
                 # This corresponds to the `modal run -i` use case where a breakpoint
                 # triggers and the task drops into an interactive PTY mode
                 if pty_shell_finish_event:
-                    print("ERROR: concurrent PTY shells are not supported.")
+                    print("ERROR: concurrent PTY shells are not supported.")  # noqa: T201
                 else:
                     pty_shell_stdout = output_mgr._stdout
                     pty_shell_finish_event = asyncio.Event()
@@ -645,13 +645,11 @@ async def get_app_logs_loop(
     while True:
         try:
             await _get_logs()
-        except (GRPCError, StreamTerminatedError, socket.gaierror, AttributeError) as exc:
-            if isinstance(exc, GRPCError):
-                if exc.status in RETRYABLE_GRPC_STATUS_CODES:
-                    # Try again if we had a temporary connection drop,
-                    # for example if computer went to sleep.
-                    logger.debug("Log fetching timed out. Retrying ...")
-                    continue
+        except (ServiceError, InternalError, StreamTerminatedError, socket.gaierror, AttributeError) as exc:
+            if isinstance(exc, (ServiceError, InternalError)):
+                # Try again if we had a temporary connection drop, for example if computer went to sleep.
+                logger.debug("Log fetching timed out. Retrying ...")
+                continue
             elif isinstance(exc, StreamTerminatedError):
                 logger.debug("Stream closed. Retrying ...")
                 continue