modal 1.0.6.dev58__py3-none-any.whl → 1.2.3.dev7__py3-none-any.whl

modal/_grpc_client.py

@@ -0,0 +1,171 @@
+# 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 ._traceback import suppress_tb_frames
+from ._utils.grpc_utils import Retry, _retry_transient_errors
+from .config import config, logger
+from .exception import InvalidError, NotFoundError
+
+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 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_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
+
+
+_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_frames(1):
+            if timeout is not None and retry is not None:
+                raise InvalidError("Retry must be None when timeout is set")
+
+            if retry is None:
+                return await self.direct(req, timeout=timeout, metadata=metadata)
+
+            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)
+        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,
+    ):
+        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)
+        async for response in self.client._call_stream(self.wrapped_method, request, metadata=metadata):
+            yield response

modal/_load_context.py

@@ -0,0 +1,105 @@
+# 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)"""
+
+        return LoadContext(
+            client=await _Client.from_env() if self._client is None else self.client,
+            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

@@ -10,6 +10,7 @@ from typing_extensions import Self
 
 from modal._traceback import suppress_tb_frames
 
+from ._load_context import LoadContext
 from ._resolver import Resolver
 from ._utils.async_utils import aclosing
 from ._utils.deprecation import deprecation_warning
@@ -20,11 +21,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 +43,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]]]
+    _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]
@@ -48,6 +58,10 @@ class _Object:
     _is_hydrated: bool
     _is_rehydrated: bool
 
+    # Not all object subclasses have a meaningful "name" concept
+    # So whether they expose this is a matter of having a name property
+    _name: Optional[str]
+
     @classmethod
     def __init_subclass__(cls, type_prefix: Optional[str] = None):
         super().__init_subclass__()
@@ -62,12 +76,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
@@ -77,12 +94,17 @@ 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
         self._is_hydrated = False
         self._is_rehydrated = False
 
+        self._name = name
+
         self._initialize_from_empty()
 
     def _unhydrate(self):
@@ -156,17 +178,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)
+        obj._init(
+            rep,
+            load,
+            is_another_app,
+            preload,
+            hydrate_lazily,
+            deps,
+            deduplication_key,
+            name,
+            load_context_overrides=load_context_overrides,
+        )
         return obj
 
     @staticmethod
@@ -183,9 +218,20 @@ class _Object:
     def _is_id_type(cls, object_id) -> bool:
         return cls._get_type_from_id(object_id) == cls
 
+    @classmethod
+    def _repr(cls, name: str, environment_name: Optional[str] = None) -> str:
+        public_cls = cls.__name__.strip("_")
+        environment_repr = f", environment_name={environment_name!r}" if environment_name else ""
+        return f"modal.{public_cls}.from_name({name!r}{environment_repr})"
+
     @classmethod
     def _new_hydrated(
-        cls, object_id: str, client: _Client, handle_metadata: Optional[Message], is_another_app: bool = False
+        cls,
+        object_id: str,
+        client: _Client,
+        handle_metadata: Optional[Message],
+        is_another_app: bool = False,
+        rep: Optional[str] = None,
     ) -> Self:
         obj_cls: type[Self]
         if cls._type_prefix is not None:
@@ -202,7 +248,7 @@ class _Object:
 
         # Instantiate provider
         obj = _Object.__new__(obj_cls)
-        rep = f"Object({object_id})"  # TODO(erikbern): dumb
+        rep = rep or f"modal.{obj_cls.__name__.strip('_')}.from_id({object_id!r})"
         obj._init(rep, is_another_app=is_another_app)
         obj._hydrate(object_id, client, handle_metadata)
 
@@ -256,25 +302,27 @@ 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))
+                # Set the client on LoadContext before loading
+                root_load_context = LoadContext(client=client)
+                resolver = Resolver()
+                await resolver.load(typing.cast(_Object, self), root_load_context)
                 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)
+            # Set the client on LoadContext before loading
+            root_load_context = LoadContext(client=client)
+            resolver = Resolver()
             with suppress_tb_frames(1):  # skip this frame by default
-                await resolver.load(self)
+                await resolver.load(self, root_load_context)
         return self
 
 

modal/_output.py

@@ -4,7 +4,6 @@ from __future__ import annotations
 import asyncio
 import contextlib
 import functools
-import io
 import platform
 import re
 import socket
@@ -32,11 +31,11 @@ from rich.progress import (
 from rich.spinner import Spinner
 from rich.text import Text
 
-from modal._utils.time_utils import timestamp_to_local
+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.shell_utils import stream_from_stdin
+from ._utils.grpc_utils import RETRYABLE_GRPC_STATUS_CODES, Retry
+from ._utils.shell_utils import stream_from_stdin, write_to_fd
 from .client import _Client
 from .config import logger
 
@@ -46,6 +45,16 @@ else:
     default_spinner = "dots"
 
 
+def make_console(*, stderr: bool = False, highlight: bool = True) -> Console:
+    """Create a rich Console tuned for Modal CLI output."""
+    return Console(
+        stderr=stderr,
+        highlight=highlight,
+        # CLI does not work with auto-detected Jupyter HTML display_data.
+        force_jupyter=False,
+    )
+
+
 class FunctionQueuingColumn(ProgressColumn):
     """Renders time elapsed, including task.completed as additional elapsed time."""
 
@@ -63,25 +72,6 @@ class FunctionQueuingColumn(ProgressColumn):
         return Text(str(delta), style="progress.elapsed")
 
 
-def download_progress_bar() -> Progress:
-    """
-    Returns a progress bar suitable for showing file download progress.
-    Requires passing a `path: str` data field for rendering.
-    """
-    return Progress(
-        TextColumn("[bold white]{task.fields[path]}", justify="right"),
-        BarColumn(bar_width=None),
-        "[progress.percentage]{task.percentage:>3.1f}%",
-        "•",
-        DownloadColumn(),
-        "•",
-        TransferSpeedColumn(),
-        "•",
-        TimeRemainingColumn(),
-        transient=True,
-    )
-
-
 class LineBufferedOutput:
     """Output stream that buffers lines and passes them to a callback."""
 
@@ -101,7 +91,7 @@ class LineBufferedOutput:
 
         if self._show_timestamps:
             for i in range(0, len(chunks) - 1, 2):
-                chunks[i] = f"{timestamp_to_local(log.timestamp)} {chunks[i]}"
+                chunks[i] = f"{timestamp_to_localized_str(log.timestamp)} {chunks[i]}"
 
         completed_lines = "".join(chunks[:-1])
         remainder = chunks[-1]
@@ -147,12 +137,11 @@ class OutputManager:
     def __init__(
         self,
         *,
-        stdout: io.TextIOWrapper | None = None,
         status_spinner_text: str = "Running app...",
         show_timestamps: bool = False,
     ):
-        self._stdout = stdout or sys.stdout
-        self._console = Console(file=stdout, highlight=False)
+        self._stdout = sys.stdout
+        self._console = make_console(highlight=False)
         self._task_states = {}
         self._task_progress_items = {}
         self._current_render_group = None
@@ -500,12 +489,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):
@@ -518,17 +506,32 @@ async def put_pty_content(log: api_pb2.TaskLogs, stdout):
         # because the progress spinner can't interfere with output.
 
         data = log.data.encode("utf-8")
-        written = 0
-        n_retries = 0
-        while written < len(data):
-            try:
-                written += stdout.buffer.write(data[written:])
-                stdout.flush()
-            except BlockingIOError:
-                if n_retries >= 5:
-                    raise
-                n_retries += 1
-                await asyncio.sleep(0.1)
+        # Non-blocking terminals can fill the kernel buffer on output bursts, making flush() raise
+        # BlockingIOError (EAGAIN) and appear frozen until a key is pressed (this happened e.g. when
+        # printing large data from a pdb breakpoint). If stdout has a real fd, we await a
+        # non-blocking fd write (write_to_fd) instead.
+        fd = None
+        try:
+            if hasattr(stdout, "fileno"):
+                fd = stdout.fileno()
+        except Exception:
+            fd = None
+
+        if fd is not None:
+            await write_to_fd(fd, data)
+        else:
+            # For streams without fileno(), use the normal write/flush path.
+            written = 0
+            n_retries = 0
+            while written < len(data):
+                try:
+                    written += stdout.buffer.write(data[written:])
+                    stdout.flush()
+                except BlockingIOError:
+                    if n_retries >= 5:
+                        raise
+                    n_retries += 1
+                    await asyncio.sleep(0.1)
     else:
         # `stdout` isn't always buffered (e.g. %%capture in Jupyter notebooks redirects it to
         # io.StringIO).
@@ -548,14 +551,22 @@ async def get_app_logs_loop(
     pty_shell_stdout = None
     pty_shell_finish_event: asyncio.Event | None = None
     pty_shell_task_id: str | None = None
+    pty_shell_input_task: asyncio.Task | None = None
 
     async def stop_pty_shell():
-        nonlocal pty_shell_finish_event
+        nonlocal pty_shell_finish_event, pty_shell_input_task
         if pty_shell_finish_event:
             print("\r", end="")  # move cursor to beginning of line
             pty_shell_finish_event.set()
             pty_shell_finish_event = None
-            await asyncio.sleep(0)  # yield to handle_exec_input() so it can disable raw terminal
+
+            if pty_shell_input_task:
+                try:
+                    await pty_shell_input_task
+                except Exception as exc:
+                    logger.exception(f"Exception in PTY shell input task: {exc}")
+                finally:
+                    pty_shell_input_task = None
 
     async def _put_log(log_batch: api_pb2.TaskLogsBatch, log: api_pb2.TaskLogs):
         if log.task_state:
@@ -583,7 +594,7 @@ async def get_app_logs_loop(
 
     async def _get_logs():
         nonlocal last_log_batch_entry_id
-        nonlocal pty_shell_stdout, pty_shell_finish_event, pty_shell_task_id
+        nonlocal pty_shell_stdout, pty_shell_finish_event, pty_shell_task_id, pty_shell_input_task
 
         request = api_pb2.AppGetLogsRequest(
             app_id=app_id or "",
@@ -618,7 +629,9 @@ async def get_app_logs_loop(
                     pty_shell_finish_event = asyncio.Event()
                     pty_shell_task_id = log_batch.task_id
                     output_mgr.disable()
-                    asyncio.create_task(stream_pty_shell_input(client, log_batch.pty_exec_id, pty_shell_finish_event))
+                    pty_shell_input_task = asyncio.create_task(
+                        stream_pty_shell_input(client, log_batch.pty_exec_id, pty_shell_finish_event)
+                    )
             else:
                 for log in log_batch.items:
                     await _put_log(log_batch, log)