modAL 1.2.1.dev22__tar.gz → 1.2.2__tar.gz

{modal-1.2.1.dev22 → modal-1.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modal
-Version: 1.2.1.dev22
+Version: 1.2.2
 Summary: Python client library for Modal
 Author-email: Modal Labs <support@modal.com>
 License: Apache-2.0
@@ -13,7 +13,7 @@ Classifier: Topic :: System :: Distributed Computing
 Classifier: Operating System :: OS Independent
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
-Requires-Python: >=3.9
+Requires-Python: <3.14,>=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: aiohttp

{modal-1.2.1.dev22 → modal-1.2.2}/modal/_clustered_functions.py

@@ -5,7 +5,6 @@ from dataclasses import dataclass
 from typing import Optional
 
 from modal._utils.async_utils import synchronize_api
-from modal._utils.grpc_utils import retry_transient_errors
 from modal.client import _Client
 from modal.exception import InvalidError
 from modal_proto import api_pb2
@@ -61,8 +60,7 @@ async def _initialize_clustered_function(client: _Client, task_id: str, world_si
         os.environ["NCCL_NSOCKS_PERTHREAD"] = "1"
 
     if world_size > 1:
-        resp: api_pb2.TaskClusterHelloResponse = await retry_transient_errors(
-            client.stub.TaskClusterHello,
+        resp = await client.stub.TaskClusterHello(
             api_pb2.TaskClusterHelloRequest(
                 task_id=task_id,
                 container_ip=container_ip,

{modal-1.2.1.dev22 → modal-1.2.2}/modal/_functions.py

@@ -19,7 +19,8 @@ from synchronicity.combined_types import MethodWithAio
 from modal_proto import api_pb2
 from modal_proto.modal_api_grpc import ModalClientModal
 
-from ._object import _get_environment_name, _Object, live_method, live_method_gen
+from ._load_context import LoadContext
+from ._object import _Object, live_method, live_method_gen
 from ._pty import get_pty_info
 from ._resolver import Resolver
 from ._resources import convert_fn_config_to_resources_config
@@ -53,7 +54,7 @@ from ._utils.function_utils import (
     get_function_type,
     is_async,
 )
-from ._utils.grpc_utils import RetryWarningMessage, retry_transient_errors
+from ._utils.grpc_utils import Retry, RetryWarningMessage
 from ._utils.mount_utils import validate_network_file_systems, validate_volumes
 from .call_graph import InputInfo, _reconstruct_call_graph
 from .client import _Client
@@ -164,21 +165,22 @@ class _Invocation:
 
         if from_spawn_map:
             request.from_spawn_map = True
-            response = await retry_transient_errors(
-                client.stub.FunctionMap,
+            response = await client.stub.FunctionMap(
                 request,
-                max_retries=None,
-                max_delay=30.0,
-                retry_warning_message=RetryWarningMessage(
-                    message="Warning: `.spawn_map(...)` for function `{self._function_name}` is waiting to create"
-                    "more function calls. This may be due to hitting rate limits or function backlog limits.",
-                    warning_interval=10,
-                    errors_to_warn_for=[Status.RESOURCE_EXHAUSTED],
+                retry=Retry(
+                    max_retries=None,
+                    max_delay=30.0,
+                    warning_message=RetryWarningMessage(
+                        message="Warning: `.spawn_map(...)` for function `{self._function_name}` is waiting to create"
+                        "more function calls. This may be due to hitting rate limits or function backlog limits.",
+                        warning_interval=10,
+                        errors_to_warn_for=[Status.RESOURCE_EXHAUSTED],
+                    ),
+                    additional_status_codes=[Status.RESOURCE_EXHAUSTED],
                 ),
-                additional_status_codes=[Status.RESOURCE_EXHAUSTED],
             )
         else:
-            response = await retry_transient_errors(client.stub.FunctionMap, request)
+            response = await client.stub.FunctionMap(request)
 
         function_call_id = response.function_call_id
         if response.pipelined_inputs:
@@ -198,10 +200,7 @@ class _Invocation:
         request_put = api_pb2.FunctionPutInputsRequest(
             function_id=function_id, inputs=[item], function_call_id=function_call_id
         )
-        inputs_response: api_pb2.FunctionPutInputsResponse = await retry_transient_errors(
-            client.stub.FunctionPutInputs,
-            request_put,
-        )
+        inputs_response: api_pb2.FunctionPutInputsResponse = await client.stub.FunctionPutInputs(request_put)
         processed_inputs = inputs_response.inputs
         if not processed_inputs:
             raise Exception("Could not create function call - the input queue seems to be full")
@@ -243,10 +242,9 @@ class _Invocation:
                 start_idx=index,
                 end_idx=index,
             )
-            response: api_pb2.FunctionGetOutputsResponse = await retry_transient_errors(
-                self.stub.FunctionGetOutputs,
+            response: api_pb2.FunctionGetOutputsResponse = await self.stub.FunctionGetOutputs(
                 request,
-                attempt_timeout=backend_timeout + ATTEMPT_TIMEOUT_GRACE_PERIOD,
+                retry=Retry(attempt_timeout=backend_timeout + ATTEMPT_TIMEOUT_GRACE_PERIOD),
             )
 
             if len(response.outputs) > 0:
@@ -266,10 +264,7 @@ class _Invocation:
 
         item = api_pb2.FunctionRetryInputsItem(input_jwt=ctx.input_jwt, input=ctx.item.input)
         request = api_pb2.FunctionRetryInputsRequest(function_call_jwt=ctx.function_call_jwt, inputs=[item])
-        await retry_transient_errors(
-            self.stub.FunctionRetryInputs,
-            request,
-        )
+        await self.stub.FunctionRetryInputs(request)
 
     async def _get_single_output(self, expected_jwt: Optional[str] = None) -> api_pb2.FunctionGetOutputsItem:
         # waits indefinitely for a single result for the function, and clear the outputs buffer after
@@ -373,10 +368,8 @@ class _Invocation:
                 start_idx=current_index,
                 end_idx=batch_end_index,
             )
-            response: api_pb2.FunctionGetOutputsResponse = await retry_transient_errors(
-                self.stub.FunctionGetOutputs,
-                request,
-                attempt_timeout=ATTEMPT_TIMEOUT_GRACE_PERIOD,
+            response: api_pb2.FunctionGetOutputsResponse = await self.stub.FunctionGetOutputs(
+                request, retry=Retry(attempt_timeout=ATTEMPT_TIMEOUT_GRACE_PERIOD)
             )
 
             outputs = list(response.outputs)
@@ -448,7 +441,7 @@ class _InputPlaneInvocation:
         )
 
         metadata = await client.get_input_plane_metadata(input_plane_region)
-        response = await retry_transient_errors(stub.AttemptStart, request, metadata=metadata)
+        response = await stub.AttemptStart(request, metadata=metadata)
         attempt_token = response.attempt_token
 
         return _InputPlaneInvocation(
@@ -468,10 +461,9 @@ class _InputPlaneInvocation:
                 requested_at=time.time(),
             )
             metadata = await self.client.get_input_plane_metadata(self.input_plane_region)
-            await_response: api_pb2.AttemptAwaitResponse = await retry_transient_errors(
-                self.stub.AttemptAwait,
+            await_response: api_pb2.AttemptAwaitResponse = await self.stub.AttemptAwait(
                 await_request,
-                attempt_timeout=OUTPUTS_TIMEOUT + ATTEMPT_TIMEOUT_GRACE_PERIOD,
+                retry=Retry(attempt_timeout=OUTPUTS_TIMEOUT + ATTEMPT_TIMEOUT_GRACE_PERIOD),
                 metadata=metadata,
             )
 
@@ -511,11 +503,7 @@ class _InputPlaneInvocation:
             input=self.input_item,
             attempt_token=self.attempt_token,
         )
-        retry_response = await retry_transient_errors(
-            self.stub.AttemptRetry,
-            retry_request,
-            metadata=metadata,
-        )
+        retry_response = await self.stub.AttemptRetry(retry_request, metadata=metadata)
         return retry_response.attempt_token
 
     async def run_generator(self):
@@ -669,7 +657,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
     @staticmethod
     def from_local(
         info: FunctionInfo,
-        app,
+        app: Optional["modal.app._App"],  # App here should only be None in case of Image.run_function
         image: _Image,
         env: Optional[dict[str, Optional[str]]] = None,
         secrets: Optional[Collection[_Secret]] = None,
@@ -895,12 +883,12 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                 is_web_endpoint, is_generator, restrict_output
             )
 
-        async def _preload(self: _Function, resolver: Resolver, existing_object_id: Optional[str]):
-            assert resolver.client and resolver.client.stub
-
-            assert resolver.app_id
+        async def _preload(
+            self: _Function, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]
+        ):
+            assert load_context.app_id
             req = api_pb2.FunctionPrecreateRequest(
-                app_id=resolver.app_id,
+                app_id=load_context.app_id,
                 function_name=info.function_name,
                 function_type=function_type,
                 existing_function_id=existing_object_id or "",
@@ -916,11 +904,12 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
             elif webhook_config:
                 req.webhook_config.CopyFrom(webhook_config)
 
-            response = await retry_transient_errors(resolver.client.stub.FunctionPrecreate, req)
-            self._hydrate(response.function_id, resolver.client, response.handle_metadata)
+            response = await load_context.client.stub.FunctionPrecreate(req)
+            self._hydrate(response.function_id, load_context.client, response.handle_metadata)
 
-        async def _load(self: _Function, resolver: Resolver, existing_object_id: Optional[str]):
-            assert resolver.client and resolver.client.stub
+        async def _load(
+            self: _Function, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]
+        ):
             with FunctionCreationStatus(resolver, tag) as function_creation_status:
                 timeout_secs = timeout
 
@@ -1116,18 +1105,16 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                         ),
                     )
 
-                assert resolver.app_id
+                assert load_context.app_id
                 assert (function_definition is None) != (function_data is None)  # xor
                 request = api_pb2.FunctionCreateRequest(
-                    app_id=resolver.app_id,
+                    app_id=load_context.app_id,
                     function=function_definition,
                     function_data=function_data,
                     existing_function_id=existing_object_id or "",
                 )
                 try:
-                    response: api_pb2.FunctionCreateResponse = await retry_transient_errors(
-                        resolver.client.stub.FunctionCreate, request
-                    )
+                    response: api_pb2.FunctionCreateResponse = await load_context.client.stub.FunctionCreate(request)
                 except GRPCError as exc:
                     if exc.status == Status.INVALID_ARGUMENT:
                         raise InvalidError(exc.message)
@@ -1142,10 +1129,14 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
             serve_mounts = {m for m in all_mounts if m.is_local()}
             serve_mounts |= image._serve_mounts
             obj._serve_mounts = frozenset(serve_mounts)
-            self._hydrate(response.function_id, resolver.client, response.handle_metadata)
+            self._hydrate(response.function_id, load_context.client, response.handle_metadata)
 
         rep = f"Function({tag})"
-        obj = _Function._from_loader(_load, rep, preload=_preload, deps=_deps)
+        # Pass a *reference* to the App's LoadContext - this is important since the App is
+        # the only way to infer a LoadContext for an `@app.function`, and the App doesn't
+        # get its client until *after* the Function is created.
+        load_context = app._root_load_context if app else LoadContext.empty()
+        obj = _Function._from_loader(_load, rep, preload=_preload, deps=_deps, load_context_overrides=load_context)
 
         obj._raw_f = info.raw_f
         obj._info = info
@@ -1187,7 +1178,12 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
 
         parent = self
 
-        async def _load(param_bound_func: _Function, resolver: Resolver, existing_object_id: Optional[str]):
+        async def _load(
+            param_bound_func: _Function,
+            resolver: Resolver,
+            load_context: LoadContext,
+            existing_object_id: Optional[str],
+        ):
             if not parent.is_hydrated:
                 # While the base Object.hydrate() method appears to be idempotent, it's not always safe
                 await parent.hydrate()
@@ -1220,7 +1216,6 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                 param_bound_func._hydrate_from_other(parent)
                 return
 
-            environment_name = _get_environment_name(None, resolver)
             assert parent is not None and parent.is_hydrated
 
             if options:
@@ -1260,11 +1255,11 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                 function_id=parent.object_id,
                 serialized_params=serialized_params,
                 function_options=options_pb,
-                environment_name=environment_name
+                environment_name=load_context.environment_name
                 or "",  # TODO: investigate shouldn't environment name always be specified here?
             )
 
-            response = await retry_transient_errors(parent._client.stub.FunctionBindParams, req)
+            response = await parent._client.stub.FunctionBindParams(req)
             param_bound_func._hydrate(response.bound_function_id, parent._client, response.handle_metadata)
 
         def _deps():
@@ -1277,7 +1272,13 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                 return [dep for dep in all_deps if not dep.is_hydrated]
             return []
 
-        fun: _Function = _Function._from_loader(_load, "Function(parametrized)", hydrate_lazily=True, deps=_deps)
+        fun: _Function = _Function._from_loader(
+            _load,
+            "Function(parametrized)",
+            hydrate_lazily=True,
+            deps=_deps,
+            load_context_overrides=self._load_context_overrides,
+        )
 
         fun._info = self._info
         fun._obj = obj
@@ -1328,7 +1329,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
             scaledown_window=scaledown_window,
         )
         request = api_pb2.FunctionUpdateSchedulingParamsRequest(function_id=self.object_id, settings=settings)
-        await retry_transient_errors(self.client.stub.FunctionUpdateSchedulingParams, request)
+        await self.client.stub.FunctionUpdateSchedulingParams(request)
 
         # One idea would be for FunctionUpdateScheduleParams to return the current (coalesced) settings
         # and then we could return them here (would need some ad hoc dataclass, which I don't love)
@@ -1375,34 +1376,43 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         cls,
         app_name: str,
         name: str,
-        namespace=None,  # mdmd:line-hidden
-        environment_name: Optional[str] = None,
+        *,
+        load_context_overrides: LoadContext,
     ):
         # internal function lookup implementation that allows lookup of class "service functions"
         # in addition to non-class functions
-        async def _load_remote(self: _Function, resolver: Resolver, existing_object_id: Optional[str]):
-            assert resolver.client and resolver.client.stub
+        async def _load_remote(
+            self: _Function, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]
+        ):
             request = api_pb2.FunctionGetRequest(
                 app_name=app_name,
                 object_tag=name,
-                environment_name=_get_environment_name(environment_name, resolver) or "",
+                environment_name=load_context.environment_name,
             )
             try:
-                response = await retry_transient_errors(resolver.client.stub.FunctionGet, request)
+                response = await load_context.client.stub.FunctionGet(request)
             except NotFoundError as exc:
                 # refine the error message
-                env_context = f" (in the '{environment_name}' environment)" if environment_name else ""
+                env_context = (
+                    f" (in the '{load_context.environment_name}' environment)" if load_context.environment_name else ""
+                )
                 raise NotFoundError(
                     f"Lookup failed for Function '{name}' from the '{app_name}' app{env_context}: {exc}."
                 ) from None
 
             print_server_warnings(response.server_warnings)
 
-            self._hydrate(response.function_id, resolver.client, response.handle_metadata)
+            self._hydrate(response.function_id, load_context.client, response.handle_metadata)
 
-        environment_rep = f", environment_name={environment_name!r}" if environment_name else ""
+        environment_rep = (
+            f", environment_name={load_context_overrides.environment_name!r}"
+            if load_context_overrides._environment_name  # slightly ugly - checking if _environment_name is overridden
+            else ""
+        )
         rep = f"modal.Function.from_name('{app_name}', '{name}'{environment_rep})"
-        return cls._from_loader(_load_remote, rep, is_another_app=True, hydrate_lazily=True)
+        return cls._from_loader(
+            _load_remote, rep, is_another_app=True, hydrate_lazily=True, load_context_overrides=load_context_overrides
+        )
 
     @classmethod
     def from_name(
@@ -1412,6 +1422,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         *,
         namespace=None,  # mdmd:line-hidden
         environment_name: Optional[str] = None,
+        client: Optional[_Client] = None,
     ) -> "_Function":
         """Reference a Function from a deployed App by its name.
 
@@ -1435,7 +1446,9 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
             )
 
         warn_if_passing_namespace(namespace, "modal.Function.from_name")
-        return cls._from_name(app_name, name, environment_name=environment_name)
+        return cls._from_name(
+            app_name, name, load_context_overrides=LoadContext(environment_name=environment_name, client=client)
+        )
 
     @property
     def tag(self) -> str:
@@ -1658,8 +1671,8 @@ Use the `Function.get_web_url()` method instead.
             input_queue,
             self.client,
         )
-        metadata = api_pb2.FunctionCallFromIdResponse(function_call_id=function_call_id, num_inputs=num_inputs)
-        fc: _FunctionCall[ReturnType] = _FunctionCall._new_hydrated(function_call_id, self.client, metadata)
+        fc: _FunctionCall[ReturnType] = _FunctionCall._new_hydrated(function_call_id, self.client, None)
+        fc._num_inputs = num_inputs  # set the cached value of num_inputs
         return fc
 
     async def _call_function(self, args, kwargs) -> ReturnType:
@@ -1888,10 +1901,9 @@ Use the `Function.get_web_url()` method instead.
     @live_method
     async def get_current_stats(self) -> FunctionStats:
         """Return a `FunctionStats` object describing the current function's queue and runner counts."""
-        resp = await retry_transient_errors(
-            self.client.stub.FunctionGetCurrentStats,
+        resp = await self.client.stub.FunctionGetCurrentStats(
             api_pb2.FunctionGetCurrentStatsRequest(function_id=self.object_id),
-            total_timeout=10.0,
+            retry=Retry(total_timeout=10.0),
         )
         return FunctionStats(backlog=resp.backlog, num_total_runners=resp.num_total_tasks)
 
@@ -1929,19 +1941,16 @@ class _FunctionCall(typing.Generic[ReturnType], _Object, type_prefix="fc"):
     def _invocation(self):
         return _Invocation(self.client.stub, self.object_id, self.client)
 
-    def _hydrate_metadata(self, metadata: Optional[Message]):
-        if not metadata:
-            return
-        assert isinstance(metadata, api_pb2.FunctionCallFromIdResponse)
-        self._num_inputs = metadata.num_inputs
-
     @live_method
     async def num_inputs(self) -> int:
         """Get the number of inputs in the function call."""
-        # Should have been hydrated.
-        assert self._num_inputs is not None
+        if self._num_inputs is None:
+            request = api_pb2.FunctionCallFromIdRequest(function_call_id=self.object_id)
+            resp = await self.client.stub.FunctionCallFromId(request)
+            self._num_inputs = resp.num_inputs  # cached
         return self._num_inputs
 
+    @live_method
     async def get(self, timeout: Optional[float] = None, *, index: int = 0) -> ReturnType:
         """Get the result of the index-th input of the function call.
         `.spawn()` calls have a single output, so only specifying `index=0` is valid.
@@ -1985,6 +1994,7 @@ class _FunctionCall(typing.Generic[ReturnType], _Object, type_prefix="fc"):
         async for _, item in self._invocation().enumerate(start_index=start, end_index=end):
             yield item
 
+    @live_method
     async def get_call_graph(self) -> list[InputInfo]:
         """Returns a structure representing the call graph from a given root
         call ID, along with the status of execution for each node.
@@ -1994,9 +2004,10 @@ class _FunctionCall(typing.Generic[ReturnType], _Object, type_prefix="fc"):
         """
         assert self._client and self._client.stub
         request = api_pb2.FunctionGetCallGraphRequest(function_call_id=self.object_id)
-        response = await retry_transient_errors(self._client.stub.FunctionGetCallGraph, request)
+        response = await self._client.stub.FunctionGetCallGraph(request)
         return _reconstruct_call_graph(response)
 
+    @live_method
     async def cancel(
         self,
         # if true, containers running the inputs are forcibly terminated
@@ -2012,7 +2023,7 @@ class _FunctionCall(typing.Generic[ReturnType], _Object, type_prefix="fc"):
             function_call_id=self.object_id, terminate_containers=terminate_containers
         )
         assert self._client and self._client.stub
-        await retry_transient_errors(self._client.stub.FunctionCallCancel, request)
+        await self._client.stub.FunctionCallCancel(request)
 
     @staticmethod
     async def from_id(function_call_id: str, client: Optional[_Client] = None) -> "_FunctionCall[Any]":
@@ -2034,20 +2045,18 @@ class _FunctionCall(typing.Generic[ReturnType], _Object, type_prefix="fc"):
         if you no longer have access to the original object returned from `Function.spawn`.
 
         """
-        if client is None:
-            client = await _Client.from_env()
 
-        async def _load(self: _FunctionCall, resolver: Resolver, existing_object_id: Optional[str]):
-            request = api_pb2.FunctionCallFromIdRequest(function_call_id=function_call_id)
-            resp = await retry_transient_errors(resolver.client.stub.FunctionCallFromId, request)
-            self._hydrate(function_call_id, resolver.client, resp)
+        async def _load(
+            self: _FunctionCall, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]
+        ):
+            # this loader doesn't do anything in practice, but it will get the client from the load_context
+            self._hydrate(function_call_id, load_context.client, None)
 
         rep = f"FunctionCall.from_id({function_call_id!r})"
-        fc: _FunctionCall[Any] = _FunctionCall._from_loader(_load, rep, hydrate_lazily=True)
-        # We already know the object ID, so we can set it directly
-        fc._object_id = function_call_id
-        fc._client = client
-        return fc
+
+        return _FunctionCall._from_loader(
+            _load, rep, hydrate_lazily=True, load_context_overrides=LoadContext(client=client)
+        )
 
     @staticmethod
     async def gather(*function_calls: "_FunctionCall[T]") -> typing.Sequence[T]:

modal-1.2.2/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