modAL 1.2.2__tar.gz → 1.2.2.dev2__tar.gz

{modal-1.2.2 → modal-1.2.2.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modal
-Version: 1.2.2
+Version: 1.2.2.dev2
 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.14,>=3.9
+Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: aiohttp

{modal-1.2.2 → modal-1.2.2.dev2}/modal/_clustered_functions.py

@@ -5,6 +5,7 @@ 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
@@ -60,7 +61,8 @@ async def _initialize_clustered_function(client: _Client, task_id: str, world_si
         os.environ["NCCL_NSOCKS_PERTHREAD"] = "1"
 
     if world_size > 1:
-        resp = await client.stub.TaskClusterHello(
+        resp: api_pb2.TaskClusterHelloResponse = await retry_transient_errors(
+            client.stub.TaskClusterHello,
             api_pb2.TaskClusterHelloRequest(
                 task_id=task_id,
                 container_ip=container_ip,

{modal-1.2.2 → modal-1.2.2.dev2}/modal/_functions.py

@@ -19,8 +19,7 @@ from synchronicity.combined_types import MethodWithAio
 from modal_proto import api_pb2
 from modal_proto.modal_api_grpc import ModalClientModal
 
-from ._load_context import LoadContext
-from ._object import _Object, live_method, live_method_gen
+from ._object import _get_environment_name, _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
@@ -54,7 +53,7 @@ from ._utils.function_utils import (
     get_function_type,
     is_async,
 )
-from ._utils.grpc_utils import Retry, RetryWarningMessage
+from ._utils.grpc_utils import RetryWarningMessage, retry_transient_errors
 from ._utils.mount_utils import validate_network_file_systems, validate_volumes
 from .call_graph import InputInfo, _reconstruct_call_graph
 from .client import _Client
@@ -165,22 +164,21 @@ class _Invocation:
 
         if from_spawn_map:
             request.from_spawn_map = True
-            response = await client.stub.FunctionMap(
+            response = await retry_transient_errors(
+                client.stub.FunctionMap,
                 request,
-                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],
+                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],
                 ),
+                additional_status_codes=[Status.RESOURCE_EXHAUSTED],
             )
         else:
-            response = await client.stub.FunctionMap(request)
+            response = await retry_transient_errors(client.stub.FunctionMap, request)
 
         function_call_id = response.function_call_id
         if response.pipelined_inputs:
@@ -200,7 +198,10 @@ 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 client.stub.FunctionPutInputs(request_put)
+        inputs_response: api_pb2.FunctionPutInputsResponse = await retry_transient_errors(
+            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")
@@ -242,9 +243,10 @@ class _Invocation:
                 start_idx=index,
                 end_idx=index,
             )
-            response: api_pb2.FunctionGetOutputsResponse = await self.stub.FunctionGetOutputs(
+            response: api_pb2.FunctionGetOutputsResponse = await retry_transient_errors(
+                self.stub.FunctionGetOutputs,
                 request,
-                retry=Retry(attempt_timeout=backend_timeout + ATTEMPT_TIMEOUT_GRACE_PERIOD),
+                attempt_timeout=backend_timeout + ATTEMPT_TIMEOUT_GRACE_PERIOD,
             )
 
             if len(response.outputs) > 0:
@@ -264,7 +266,10 @@ 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 self.stub.FunctionRetryInputs(request)
+        await retry_transient_errors(
+            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
@@ -368,8 +373,10 @@ class _Invocation:
                 start_idx=current_index,
                 end_idx=batch_end_index,
             )
-            response: api_pb2.FunctionGetOutputsResponse = await self.stub.FunctionGetOutputs(
-                request, retry=Retry(attempt_timeout=ATTEMPT_TIMEOUT_GRACE_PERIOD)
+            response: api_pb2.FunctionGetOutputsResponse = await retry_transient_errors(
+                self.stub.FunctionGetOutputs,
+                request,
+                attempt_timeout=ATTEMPT_TIMEOUT_GRACE_PERIOD,
             )
 
             outputs = list(response.outputs)
@@ -441,7 +448,7 @@ class _InputPlaneInvocation:
         )
 
         metadata = await client.get_input_plane_metadata(input_plane_region)
-        response = await stub.AttemptStart(request, metadata=metadata)
+        response = await retry_transient_errors(stub.AttemptStart, request, metadata=metadata)
         attempt_token = response.attempt_token
 
         return _InputPlaneInvocation(
@@ -461,9 +468,10 @@ class _InputPlaneInvocation:
                 requested_at=time.time(),
             )
             metadata = await self.client.get_input_plane_metadata(self.input_plane_region)
-            await_response: api_pb2.AttemptAwaitResponse = await self.stub.AttemptAwait(
+            await_response: api_pb2.AttemptAwaitResponse = await retry_transient_errors(
+                self.stub.AttemptAwait,
                 await_request,
-                retry=Retry(attempt_timeout=OUTPUTS_TIMEOUT + ATTEMPT_TIMEOUT_GRACE_PERIOD),
+                attempt_timeout=OUTPUTS_TIMEOUT + ATTEMPT_TIMEOUT_GRACE_PERIOD,
                 metadata=metadata,
             )
 
@@ -503,7 +511,11 @@ class _InputPlaneInvocation:
             input=self.input_item,
             attempt_token=self.attempt_token,
         )
-        retry_response = await self.stub.AttemptRetry(retry_request, metadata=metadata)
+        retry_response = await retry_transient_errors(
+            self.stub.AttemptRetry,
+            retry_request,
+            metadata=metadata,
+        )
         return retry_response.attempt_token
 
     async def run_generator(self):
@@ -657,7 +669,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
     @staticmethod
     def from_local(
         info: FunctionInfo,
-        app: Optional["modal.app._App"],  # App here should only be None in case of Image.run_function
+        app,
         image: _Image,
         env: Optional[dict[str, Optional[str]]] = None,
         secrets: Optional[Collection[_Secret]] = None,
@@ -883,12 +895,12 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                 is_web_endpoint, is_generator, restrict_output
             )
 
-        async def _preload(
-            self: _Function, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]
-        ):
-            assert load_context.app_id
+        async def _preload(self: _Function, resolver: Resolver, existing_object_id: Optional[str]):
+            assert resolver.client and resolver.client.stub
+
+            assert resolver.app_id
             req = api_pb2.FunctionPrecreateRequest(
-                app_id=load_context.app_id,
+                app_id=resolver.app_id,
                 function_name=info.function_name,
                 function_type=function_type,
                 existing_function_id=existing_object_id or "",
@@ -904,12 +916,11 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
             elif webhook_config:
                 req.webhook_config.CopyFrom(webhook_config)
 
-            response = await load_context.client.stub.FunctionPrecreate(req)
-            self._hydrate(response.function_id, load_context.client, response.handle_metadata)
+            response = await retry_transient_errors(resolver.client.stub.FunctionPrecreate, req)
+            self._hydrate(response.function_id, resolver.client, response.handle_metadata)
 
-        async def _load(
-            self: _Function, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]
-        ):
+        async def _load(self: _Function, resolver: Resolver, existing_object_id: Optional[str]):
+            assert resolver.client and resolver.client.stub
             with FunctionCreationStatus(resolver, tag) as function_creation_status:
                 timeout_secs = timeout
 
@@ -1105,16 +1116,18 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                         ),
                     )
 
-                assert load_context.app_id
+                assert resolver.app_id
                 assert (function_definition is None) != (function_data is None)  # xor
                 request = api_pb2.FunctionCreateRequest(
-                    app_id=load_context.app_id,
+                    app_id=resolver.app_id,
                     function=function_definition,
                     function_data=function_data,
                     existing_function_id=existing_object_id or "",
                 )
                 try:
-                    response: api_pb2.FunctionCreateResponse = await load_context.client.stub.FunctionCreate(request)
+                    response: api_pb2.FunctionCreateResponse = await retry_transient_errors(
+                        resolver.client.stub.FunctionCreate, request
+                    )
                 except GRPCError as exc:
                     if exc.status == Status.INVALID_ARGUMENT:
                         raise InvalidError(exc.message)
@@ -1129,14 +1142,10 @@ 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, load_context.client, response.handle_metadata)
+            self._hydrate(response.function_id, resolver.client, response.handle_metadata)
 
         rep = f"Function({tag})"
-        # 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 = _Function._from_loader(_load, rep, preload=_preload, deps=_deps)
 
         obj._raw_f = info.raw_f
         obj._info = info
@@ -1178,12 +1187,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
 
         parent = self
 
-        async def _load(
-            param_bound_func: _Function,
-            resolver: Resolver,
-            load_context: LoadContext,
-            existing_object_id: Optional[str],
-        ):
+        async def _load(param_bound_func: _Function, resolver: Resolver, 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()
@@ -1216,6 +1220,7 @@ 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:
@@ -1255,11 +1260,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=load_context.environment_name
+                environment_name=environment_name
                 or "",  # TODO: investigate shouldn't environment name always be specified here?
             )
 
-            response = await parent._client.stub.FunctionBindParams(req)
+            response = await retry_transient_errors(parent._client.stub.FunctionBindParams, req)
             param_bound_func._hydrate(response.bound_function_id, parent._client, response.handle_metadata)
 
         def _deps():
@@ -1272,13 +1277,7 @@ 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,
-            load_context_overrides=self._load_context_overrides,
-        )
+        fun: _Function = _Function._from_loader(_load, "Function(parametrized)", hydrate_lazily=True, deps=_deps)
 
         fun._info = self._info
         fun._obj = obj
@@ -1329,7 +1328,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 self.client.stub.FunctionUpdateSchedulingParams(request)
+        await retry_transient_errors(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)
@@ -1376,43 +1375,34 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         cls,
         app_name: str,
         name: str,
-        *,
-        load_context_overrides: LoadContext,
+        namespace=None,  # mdmd:line-hidden
+        environment_name: Optional[str] = None,
     ):
         # 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, load_context: LoadContext, existing_object_id: Optional[str]
-        ):
+        async def _load_remote(self: _Function, resolver: Resolver, existing_object_id: Optional[str]):
+            assert resolver.client and resolver.client.stub
             request = api_pb2.FunctionGetRequest(
                 app_name=app_name,
                 object_tag=name,
-                environment_name=load_context.environment_name,
+                environment_name=_get_environment_name(environment_name, resolver) or "",
             )
             try:
-                response = await load_context.client.stub.FunctionGet(request)
+                response = await retry_transient_errors(resolver.client.stub.FunctionGet, request)
             except NotFoundError as exc:
                 # refine the error message
-                env_context = (
-                    f" (in the '{load_context.environment_name}' environment)" if load_context.environment_name else ""
-                )
+                env_context = f" (in the '{environment_name}' environment)" if 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, load_context.client, response.handle_metadata)
+            self._hydrate(response.function_id, resolver.client, response.handle_metadata)
 
-        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 ""
-        )
+        environment_rep = f", environment_name={environment_name!r}" if environment_name 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, load_context_overrides=load_context_overrides
-        )
+        return cls._from_loader(_load_remote, rep, is_another_app=True, hydrate_lazily=True)
 
     @classmethod
     def from_name(
@@ -1422,7 +1412,6 @@ 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.
 
@@ -1446,9 +1435,7 @@ 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, load_context_overrides=LoadContext(environment_name=environment_name, client=client)
-        )
+        return cls._from_name(app_name, name, environment_name=environment_name)
 
     @property
     def tag(self) -> str:
@@ -1671,8 +1658,8 @@ Use the `Function.get_web_url()` method instead.
             input_queue,
             self.client,
         )
-        fc: _FunctionCall[ReturnType] = _FunctionCall._new_hydrated(function_call_id, self.client, None)
-        fc._num_inputs = num_inputs  # set the cached value of num_inputs
+        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)
         return fc
 
     async def _call_function(self, args, kwargs) -> ReturnType:
@@ -1901,9 +1888,10 @@ 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 self.client.stub.FunctionGetCurrentStats(
+        resp = await retry_transient_errors(
+            self.client.stub.FunctionGetCurrentStats,
             api_pb2.FunctionGetCurrentStatsRequest(function_id=self.object_id),
-            retry=Retry(total_timeout=10.0),
+            total_timeout=10.0,
         )
         return FunctionStats(backlog=resp.backlog, num_total_runners=resp.num_total_tasks)
 
@@ -1941,16 +1929,19 @@ 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."""
-        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
+        # Should have been hydrated.
+        assert self._num_inputs is not None
         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.
@@ -1994,7 +1985,6 @@ 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.
@@ -2004,10 +1994,9 @@ 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 self._client.stub.FunctionGetCallGraph(request)
+        response = await retry_transient_errors(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
@@ -2023,7 +2012,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 self._client.stub.FunctionCallCancel(request)
+        await retry_transient_errors(self._client.stub.FunctionCallCancel, request)
 
     @staticmethod
     async def from_id(function_call_id: str, client: Optional[_Client] = None) -> "_FunctionCall[Any]":
@@ -2045,18 +2034,20 @@ 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, 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)
+        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)
 
         rep = f"FunctionCall.from_id({function_call_id!r})"
-
-        return _FunctionCall._from_loader(
-            _load, rep, hydrate_lazily=True, load_context_overrides=LoadContext(client=client)
-        )
+        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
 
     @staticmethod
     async def gather(*function_calls: "_FunctionCall[T]") -> typing.Sequence[T]:

{modal-1.2.2 → modal-1.2.2.dev2}/modal/_object.py

@@ -10,7 +10,6 @@ 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
@@ -21,19 +20,11 @@ from .exception import ExecutionError, InvalidError
 EPHEMERAL_OBJECT_HEARTBEAT_SLEEP: int = 300
 
 
-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
-    """
+def _get_environment_name(environment_name: Optional[str] = None, resolver: Optional[Resolver] = None) -> Optional[str]:
     if environment_name:
         return environment_name
+    elif resolver and resolver.environment_name:
+        return resolver.environment_name
     else:
         return config.get("environment")
 
@@ -43,14 +34,13 @@ class _Object:
     _prefix_to_type: ClassVar[dict[str, type]] = {}
 
     # For constructors
-    _load: Optional[Callable[[Self, Resolver, LoadContext, Optional[str]], Awaitable[None]]]
-    _preload: Optional[Callable[[Self, Resolver, LoadContext, Optional[str]], Awaitable[None]]]
+    _load: Optional[Callable[[Self, Resolver, Optional[str]], Awaitable[None]]]
+    _preload: Optional[Callable[[Self, Resolver, 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]
@@ -76,15 +66,13 @@ class _Object:
     def _init(
         self,
         rep: str,
-        load: Optional[Callable[[Self, Resolver, LoadContext, Optional[str]], Awaitable[None]]] = None,
+        load: Optional[Callable[[Self, Resolver, Optional[str]], Awaitable[None]]] = None,
         is_another_app: bool = False,
-        preload: Optional[Callable[[Self, Resolver, LoadContext, Optional[str]], Awaitable[None]]] = None,
+        preload: Optional[Callable[[Self, Resolver, 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
@@ -94,9 +82,6 @@ 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
@@ -178,30 +163,18 @@ class _Object:
     @classmethod
     def _from_loader(
         cls,
-        load: Callable[[Self, Resolver, LoadContext, Optional[str]], Awaitable[None]],
+        load: Callable[[Self, Resolver, Optional[str]], Awaitable[None]],
         rep: str,
         is_another_app: bool = False,
-        preload: Optional[Callable[[Self, Resolver, LoadContext, Optional[str]], Awaitable[None]]] = None,
+        preload: Optional[Callable[[Self, Resolver, 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,
-            load_context_overrides=load_context_overrides,
-        )
+        obj._init(rep, load, is_another_app, preload, hydrate_lazily, deps, deduplication_key, name)
         return obj
 
     @staticmethod
@@ -302,27 +275,25 @@ 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
-                # 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)
+                c = client if client is not None else await _Client.from_env()
+                resolver = Resolver(c)
+                await resolver.load(typing.cast(_Object, self))
                 self._is_rehydrated = True
-                logger.debug(f"rehydrated {self} with client {id(self.client)}")
+                logger.debug(f"rehydrated {self} with client {id(c)}")
         elif not self._hydrate_lazily:
+            # TODO(michael) can remove _hydrate lazily? I think all objects support it now?
             self._validate_is_hydrated()
         else:
-            # Set the client on LoadContext before loading
-            root_load_context = LoadContext(client=client)
-            resolver = Resolver()
+            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, root_load_context)
+                await resolver.load(self)
         return self
 
 

{modal-1.2.2 → modal-1.2.2.dev2}/modal/_output.py

@@ -34,7 +34,7 @@ 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
+from ._utils.grpc_utils import RETRYABLE_GRPC_STATUS_CODES, retry_transient_errors
 from ._utils.shell_utils import stream_from_stdin, write_to_fd
 from .client import _Client
 from .config import logger
@@ -489,11 +489,12 @@ async def stream_pty_shell_input(client: _Client, exec_id: str, finish_event: as
     """
 
     async def _handle_input(data: bytes, message_index: int):
-        await client.stub.ContainerExecPutInput(
+        await retry_transient_errors(
+            client.stub.ContainerExecPutInput,
             api_pb2.ContainerExecPutInputRequest(
                 exec_id=exec_id, input=api_pb2.RuntimeInputMessage(message=data, message_index=message_index)
             ),
-            retry=Retry(total_timeout=10),
+            total_timeout=10,
         )
 
     async with stream_from_stdin(_handle_input, use_raw_terminal=True):