modal 1.2.2.dev8__tar.gz → 1.2.2.dev24__tar.gz

{modal-1.2.2.dev8 → modal-1.2.2.dev24}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modal
-Version: 1.2.2.dev8
+Version: 1.2.2.dev24
 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.2.dev8 → modal-1.2.2.dev24}/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.2.dev8 → modal-1.2.2.dev24}/modal/_functions.py

@@ -53,7 +53,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 +164,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 +199,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 +241,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 +263,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 +367,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 +440,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 +460,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 +502,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):
@@ -916,7 +903,7 @@ 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)
+            response = await resolver.client.stub.FunctionPrecreate(req)
             self._hydrate(response.function_id, resolver.client, response.handle_metadata)
 
         async def _load(self: _Function, resolver: Resolver, existing_object_id: Optional[str]):
@@ -1125,9 +1112,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                     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 resolver.client.stub.FunctionCreate(request)
                 except GRPCError as exc:
                     if exc.status == Status.INVALID_ARGUMENT:
                         raise InvalidError(exc.message)
@@ -1264,7 +1249,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                 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():
@@ -1328,7 +1313,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)
@@ -1388,7 +1373,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                 environment_name=_get_environment_name(environment_name, resolver) or "",
             )
             try:
-                response = await retry_transient_errors(resolver.client.stub.FunctionGet, request)
+                response = await resolver.client.stub.FunctionGet(request)
             except NotFoundError as exc:
                 # refine the error message
                 env_context = f" (in the '{environment_name}' environment)" if environment_name else ""
@@ -1888,10 +1873,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)
 
@@ -1994,7 +1978,7 @@ 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)
 
     async def cancel(
@@ -2012,7 +1996,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]":
@@ -2039,7 +2023,7 @@ class _FunctionCall(typing.Generic[ReturnType], _Object, type_prefix="fc"):
 
         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)
+            resp = await resolver.client.stub.FunctionCallFromId(request)
             self._hydrate(function_call_id, resolver.client, resp)
 
         rep = f"FunctionCall.from_id({function_call_id!r})"

modal-1.2.2.dev24/modal/_grpc_client.py

@@ -0,0 +1,148 @@
+# Copyright Modal Labs 2025
+from typing import TYPE_CHECKING, Any, Collection, Generic, Literal, Mapping, Optional, TypeVar, Union
+
+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
+
+
+class UnaryUnaryWrapper(Generic[RequestType, ResponseType]):
+    # Calls a grpclib.UnaryUnaryMethod using a specific Client instance, respecting
+    # if that client is closed etc. and possibly introducing Modal-specific retry logic
+    wrapped_method: grpclib.client.UnaryUnaryMethod[RequestType, ResponseType]
+    client: "_Client"
+
+    def __init__(
+        self,
+        wrapped_method: grpclib.client.UnaryUnaryMethod[RequestType, ResponseType],
+        client: "_Client",
+        server_url: str,
+    ):
+        self.wrapped_method = wrapped_method
+        self.client = client
+        self.server_url = server_url
+
+    @property
+    def name(self) -> str:
+        return self.wrapped_method.name
+
+    async def __call__(
+        self,
+        req: RequestType,
+        *,
+        retry: Optional[Retry] = 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-1.2.2.dev8 → modal-1.2.2.dev24}/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_transient_errors
+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
@@ -489,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):

{modal-1.2.2.dev8 → modal-1.2.2.dev24}/modal/_runtime/container_io_manager.py

@@ -36,7 +36,7 @@ from modal._traceback import print_exception
 from modal._utils.async_utils import TaskContext, aclosing, asyncify, synchronize_api, synchronizer
 from modal._utils.blob_utils import MAX_OBJECT_SIZE_BYTES, blob_download, blob_upload, format_blob_data
 from modal._utils.function_utils import _stream_function_call_data
-from modal._utils.grpc_utils import retry_transient_errors
+from modal._utils.grpc_utils import Retry
 from modal._utils.package_utils import parse_major_minor_version
 from modal.client import HEARTBEAT_INTERVAL, HEARTBEAT_TIMEOUT, _Client
 from modal.config import config, logger
@@ -623,8 +623,8 @@ class _ContainerIOManager:
                 await self.heartbeat_condition.wait()
 
             request = api_pb2.ContainerHeartbeatRequest(canceled_inputs_return_outputs_v2=True)
-            response = await retry_transient_errors(
-                self._client.stub.ContainerHeartbeat, request, attempt_timeout=HEARTBEAT_TIMEOUT
+            response = await self._client.stub.ContainerHeartbeat(
+                request, retry=Retry(attempt_timeout=HEARTBEAT_TIMEOUT)
             )
 
         if response.HasField("cancel_input_event"):
@@ -671,10 +671,9 @@ class _ContainerIOManager:
                     target_concurrency=self._target_concurrency,
                     max_concurrency=self._max_concurrency,
                 )
-                resp = await retry_transient_errors(
-                    self._client.stub.FunctionGetDynamicConcurrency,
+                resp = await self._client.stub.FunctionGetDynamicConcurrency(
                     request,
-                    attempt_timeout=DYNAMIC_CONCURRENCY_TIMEOUT_SECS,
+                    retry=Retry(attempt_timeout=DYNAMIC_CONCURRENCY_TIMEOUT_SECS),
                 )
                 if resp.concurrency != self._input_slots.value and not self._stop_concurrency_loop:
                     logger.debug(f"Dynamic concurrency set from {self._input_slots.value} to {resp.concurrency}")
@@ -725,9 +724,9 @@ class _ContainerIOManager:
 
         if self.input_plane_server_url:
             stub = await self._client.get_stub(self.input_plane_server_url)
-            await retry_transient_errors(stub.FunctionCallPutDataOut, req)
+            await stub.FunctionCallPutDataOut(req)
         else:
-            await retry_transient_errors(self._client.stub.FunctionCallPutDataOut, req)
+            await self._client.stub.FunctionCallPutDataOut(req)
 
     @asynccontextmanager
     async def generator_output_sender(
@@ -815,9 +814,7 @@ class _ContainerIOManager:
             try:
                 # If number of active inputs is at max queue size, this will block.
                 iteration += 1
-                response: api_pb2.FunctionGetInputsResponse = await retry_transient_errors(
-                    self._client.stub.FunctionGetInputs, request
-                )
+                response: api_pb2.FunctionGetInputsResponse = await self._client.stub.FunctionGetInputs(request)
 
                 if response.rate_limit_sleep_duration:
                     logger.info(
@@ -887,11 +884,12 @@ class _ContainerIOManager:
         # Limit the batch size to 20 to stay within message size limits and buffer size limits.
         output_batch_size = 20
         for i in range(0, len(outputs), output_batch_size):
-            await retry_transient_errors(
-                self._client.stub.FunctionPutOutputs,
+            await self._client.stub.FunctionPutOutputs(
                 api_pb2.FunctionPutOutputsRequest(outputs=outputs[i : i + output_batch_size]),
-                additional_status_codes=[Status.RESOURCE_EXHAUSTED],
-                max_retries=None,  # Retry indefinitely, trying every 1s.
+                retry=Retry(
+                    additional_status_codes=[Status.RESOURCE_EXHAUSTED],
+                    max_retries=None,  # Retry indefinitely, trying every 1s.
+                ),
             )
         input_ids = [output.input_id for output in outputs]
         self.exit_context(started_at, input_ids)
@@ -932,7 +930,7 @@ class _ContainerIOManager:
             )
 
             req = api_pb2.TaskResultRequest(result=result)
-            await retry_transient_errors(self._client.stub.TaskResult, req)
+            await self._client.stub.TaskResult(req)
 
             # Shut down the task gracefully
             raise UserException()
@@ -1082,13 +1080,14 @@ class _ContainerIOManager:
         await asyncify(os.sync)()
         results = await asyncio.gather(
             *[
-                retry_transient_errors(
-                    self._client.stub.VolumeCommit,
+                self._client.stub.VolumeCommit(
                     api_pb2.VolumeCommitRequest(volume_id=v_id),
-                    max_retries=9,
-                    base_delay=0.25,
-                    max_delay=256,
-                    delay_factor=2,
+                    retry=Retry(
+                        max_retries=9,
+                        base_delay=0.25,
+                        max_delay=256,
+                        delay_factor=2,
+                    ),
                 )
                 for v_id in volume_ids
             ],

{modal-1.2.2.dev8 → modal-1.2.2.dev24}/modal/_utils/async_utils.py

@@ -51,6 +51,10 @@ def synchronize_api(obj, target_module=None):
     return synchronizer.create_blocking(obj, blocking_name, target_module=target_module)
 
 
+# Used for testing to configure the `n_attempts` that `retry` will use.
+RETRY_N_ATTEMPTS_OVERRIDE: Optional[int] = None
+
+
 def retry(direct_fn=None, *, n_attempts=3, base_delay=0, delay_factor=2, timeout=90):
     """Decorator that calls an async function multiple times, with a given timeout.
 
@@ -75,8 +79,13 @@ def retry(direct_fn=None, *, n_attempts=3, base_delay=0, delay_factor=2, timeout
     def decorator(fn):
         @functools.wraps(fn)
         async def f_wrapped(*args, **kwargs):
+            if RETRY_N_ATTEMPTS_OVERRIDE is not None:
+                local_n_attempts = RETRY_N_ATTEMPTS_OVERRIDE
+            else:
+                local_n_attempts = n_attempts
+
             delay = base_delay
-            for i in range(n_attempts):
+            for i in range(local_n_attempts):
                 t0 = time.time()
                 try:
                     return await asyncio.wait_for(fn(*args, **kwargs), timeout=timeout)
@@ -84,12 +93,12 @@ def retry(direct_fn=None, *, n_attempts=3, base_delay=0, delay_factor=2, timeout
                     logger.debug(f"Function {fn} was cancelled")
                     raise
                 except Exception as e:
-                    if i >= n_attempts - 1:
+                    if i >= local_n_attempts - 1:
                         raise
                     logger.debug(
                         f"Failed invoking function {fn}: {e}"
                         f" (took {time.time() - t0}s, sleeping {delay}s"
-                        f" and trying {n_attempts - i - 1} more times)"
+                        f" and trying {local_n_attempts - i - 1} more times)"
                     )
                 await asyncio.sleep(delay)
                 delay *= delay_factor

{modal-1.2.2.dev8 → modal-1.2.2.dev24}/modal/_utils/auth_token_manager.py

@@ -9,7 +9,6 @@ from typing import Any
 from modal.exception import ExecutionError
 from modal_proto import api_pb2, modal_api_grpc
 
-from .grpc_utils import retry_transient_errors
 from .logger import logger
 
 
@@ -66,9 +65,7 @@ class _AuthTokenManager:
             # new token. Once we have a new token, the other coroutines will unblock and return from here.
             if self._token and not self._needs_refresh():
                 return
-            resp: api_pb2.AuthTokenGetResponse = await retry_transient_errors(
-                self._stub.AuthTokenGet, api_pb2.AuthTokenGetRequest()
-            )
+            resp: api_pb2.AuthTokenGetResponse = await self._stub.AuthTokenGet(api_pb2.AuthTokenGetRequest())
             if not resp.token:
                 # Not expected
                 raise ExecutionError(

{modal-1.2.2.dev8 → modal-1.2.2.dev24}/modal/_utils/blob_utils.py

@@ -27,7 +27,6 @@ from modal_proto.modal_api_grpc import ModalClientModal
 
 from ..exception import ExecutionError
 from .async_utils import TaskContext, retry
-from .grpc_utils import retry_transient_errors
 from .hash_utils import UploadHashes, get_upload_hashes
 from .http_utils import ClientSessionRegistry
 from .logger import logger
@@ -229,7 +228,7 @@ async def _blob_upload(
         content_sha256_base64=upload_hashes.sha256_base64,
         content_length=content_length,
     )
-    resp = await retry_transient_errors(stub.BlobCreate, req)
+    resp = await stub.BlobCreate(req)
 
     if resp.WhichOneof("upload_types_oneof") == "multiparts":
 
@@ -335,7 +334,7 @@ async def blob_download(blob_id: str, stub: ModalClientModal) -> bytes:
     logger.debug(f"Downloading large blob {blob_id}")
     t0 = time.time()
     req = api_pb2.BlobGetRequest(blob_id=blob_id)
-    resp = await retry_transient_errors(stub.BlobGet, req)
+    resp = await stub.BlobGet(req)
     data = await _download_from_url(resp.download_url)
     size_mib = len(data) / 1024 / 1024
     dur_s = max(time.time() - t0, 0.001)  # avoid division by zero
@@ -348,7 +347,7 @@ async def blob_download(blob_id: str, stub: ModalClientModal) -> bytes:
 
 async def blob_iter(blob_id: str, stub: ModalClientModal) -> AsyncIterator[bytes]:
     req = api_pb2.BlobGetRequest(blob_id=blob_id)
-    resp = await retry_transient_errors(stub.BlobGet, req)
+    resp = await stub.BlobGet(req)
     download_url = resp.download_url
     async with ClientSessionRegistry.get_session().get(download_url) as s3_resp:
         # S3 signal to slow down request rate.