modal 1.2.1.dev8__py3-none-any.whl → 1.2.2.dev19__py3-none-any.whl

modal/parallel_map.py

@@ -35,7 +35,7 @@ from modal._utils.function_utils import (
     _create_input,
     _process_result,
 )
-from modal._utils.grpc_utils import RETRYABLE_GRPC_STATUS_CODES, RetryWarningMessage, retry_transient_errors
+from modal._utils.grpc_utils import RETRYABLE_GRPC_STATUS_CODES, Retry, RetryWarningMessage
 from modal._utils.jwt_utils import DecodedJwt
 from modal.config import logger
 from modal.retries import RetryManager
@@ -187,7 +187,7 @@ class InputPumper:
                 f" push is {self.input_queue.qsize()}. "
             )
 
-            resp = await self._send_inputs(self.client.stub.FunctionPutInputs, request)
+            resp = await self.client.stub.FunctionPutInputs(request, retry=self._function_inputs_retry)
             self.inputs_sent += len(items)
             # Change item state to WAITING_FOR_OUTPUT, and set the input_id and input_jwt which are in the response.
             if self.map_items_manager is not None:
@@ -198,11 +198,8 @@ class InputPumper:
             )
         yield
 
-    async def _send_inputs(
-        self,
-        fn: "modal.client.UnaryUnaryWrapper",
-        request: typing.Union[api_pb2.FunctionPutInputsRequest, api_pb2.FunctionRetryInputsRequest],
-    ) -> typing.Union[api_pb2.FunctionPutInputsResponse, api_pb2.FunctionRetryInputsResponse]:
+    @property
+    def _function_inputs_retry(self) -> Retry:
         # with 8 retries we log the warning below about every 30 seconds which isn't too spammy.
         retry_warning_message = RetryWarningMessage(
             message=f"Warning: map progress for function {self.function._function_name} is limited."
@@ -210,13 +207,11 @@ class InputPumper:
             warning_interval=8,
             errors_to_warn_for=[Status.RESOURCE_EXHAUSTED],
         )
-        return await retry_transient_errors(
-            fn,
-            request,
+        return Retry(
             max_retries=None,
             max_delay=PUMP_INPUTS_MAX_RETRY_DELAY,
             additional_status_codes=[Status.RESOURCE_EXHAUSTED],
-            retry_warning_message=retry_warning_message,
+            warning_message=retry_warning_message,
         )
 
 
@@ -255,7 +250,7 @@ class SyncInputPumper(InputPumper):
                 function_call_jwt=self.function_call_jwt,
                 inputs=inputs,
             )
-            resp = await self._send_inputs(self.client.stub.FunctionRetryInputs, request)
+            resp = await self.client.stub.FunctionRetryInputs(request, retry=self._function_inputs_retry)
             # Update the state to WAITING_FOR_OUTPUT, and update the input_jwt in the context
             # to the new value in the response.
             self.map_items_manager.handle_retry_response(resp.input_jwts)
@@ -289,7 +284,7 @@ class AsyncInputPumper(InputPumper):
             function_call_id=self.function_call_id,
             num_inputs=self.inputs_sent,
         )
-        await retry_transient_errors(self.client.stub.FunctionFinishInputs, request, max_retries=None)
+        await self.client.stub.FunctionFinishInputs(request, retry=Retry(max_retries=None))
         yield
 
 
@@ -303,7 +298,7 @@ async def _spawn_map_invocation(
         function_call_type=api_pb2.FUNCTION_CALL_TYPE_MAP,
         function_call_invocation_type=api_pb2.FUNCTION_CALL_INVOCATION_TYPE_ASYNC,
     )
-    response: api_pb2.FunctionMapResponse = await retry_transient_errors(client.stub.FunctionMap, request)
+    response: api_pb2.FunctionMapResponse = await client.stub.FunctionMap(request)
     function_call_id = response.function_call_id
 
     have_all_inputs = False
@@ -382,7 +377,7 @@ async def _map_invocation(
         return_exceptions=return_exceptions,
         function_call_invocation_type=function_call_invocation_type,
     )
-    response: api_pb2.FunctionMapResponse = await retry_transient_errors(client.stub.FunctionMap, request)
+    response: api_pb2.FunctionMapResponse = await client.stub.FunctionMap(request)
 
     function_call_id = response.function_call_id
     function_call_jwt = response.function_call_jwt
@@ -478,11 +473,12 @@ async def _map_invocation(
                 input_jwts=input_jwts,
             )
             get_response_task = asyncio.create_task(
-                retry_transient_errors(
-                    client.stub.FunctionGetOutputs,
+                client.stub.FunctionGetOutputs(
                     request,
-                    max_retries=20,
-                    attempt_timeout=OUTPUTS_TIMEOUT + ATTEMPT_TIMEOUT_GRACE_PERIOD,
+                    retry=Retry(
+                        max_retries=20,
+                        attempt_timeout=OUTPUTS_TIMEOUT + ATTEMPT_TIMEOUT_GRACE_PERIOD,
+                    ),
                 )
             )
             map_done_task = asyncio.create_task(map_done_event.wait())
@@ -541,7 +537,7 @@ async def _map_invocation(
                 clear_on_success=True,
                 requested_at=time.time(),
             )
-            await retry_transient_errors(client.stub.FunctionGetOutputs, request)
+            await client.stub.FunctionGetOutputs(request)
             await retry_queue.close()
 
     async def fetch_output(item: api_pb2.FunctionGetOutputsItem) -> tuple[int, Any]:
@@ -770,13 +766,14 @@ async def _map_invocation_inputplane(
 
             metadata = await client.get_input_plane_metadata(function._input_plane_region)
 
-            response: api_pb2.MapStartOrContinueResponse = await retry_transient_errors(
-                input_plane_stub.MapStartOrContinue,
+            response: api_pb2.MapStartOrContinueResponse = await input_plane_stub.MapStartOrContinue(
                 request,
+                retry=Retry(
+                    additional_status_codes=[Status.RESOURCE_EXHAUSTED],
+                    max_delay=PUMP_INPUTS_MAX_RETRY_DELAY,
+                    max_retries=None,
+                ),
                 metadata=metadata,
-                additional_status_codes=[Status.RESOURCE_EXHAUSTED],
-                max_delay=PUMP_INPUTS_MAX_RETRY_DELAY,
-                max_retries=None,
             )
 
             # match response items to the corresponding request item index
@@ -824,8 +821,8 @@ async def _map_invocation_inputplane(
                 )
 
                 metadata = await client.get_input_plane_metadata(function._input_plane_region)
-                response: api_pb2.MapCheckInputsResponse = await retry_transient_errors(
-                    input_plane_stub.MapCheckInputs, request, metadata=metadata
+                response: api_pb2.MapCheckInputsResponse = await input_plane_stub.MapCheckInputs(
+                    request, metadata=metadata
                 )
                 check_inputs_response = [
                     (check_inputs[resp_idx][0], response.lost[resp_idx]) for resp_idx, _ in enumerate(response.lost)
@@ -859,11 +856,12 @@ async def _map_invocation_inputplane(
             )
             metadata = await client.get_input_plane_metadata(function._input_plane_region)
             get_response_task = asyncio.create_task(
-                retry_transient_errors(
-                    input_plane_stub.MapAwait,
+                input_plane_stub.MapAwait(
                     request,
-                    max_retries=20,
-                    attempt_timeout=OUTPUTS_TIMEOUT + ATTEMPT_TIMEOUT_GRACE_PERIOD,
+                    retry=Retry(
+                        max_retries=20,
+                        attempt_timeout=OUTPUTS_TIMEOUT + ATTEMPT_TIMEOUT_GRACE_PERIOD,
+                    ),
                     metadata=metadata,
                 )
             )

modal/parallel_map.pyi

@@ -5,6 +5,7 @@ import collections.abc
 import enum
 import modal._functions
 import modal._utils.async_utils
+import modal._utils.grpc_utils
 import modal.client
 import modal.functions
 import modal.retries
@@ -96,15 +97,8 @@ class InputPumper:
         ...
 
     def pump_inputs(self): ...
-    async def _send_inputs(
-        self,
-        fn: modal.client.UnaryUnaryWrapper,
-        request: typing.Union[
-            modal_proto.api_pb2.FunctionPutInputsRequest, modal_proto.api_pb2.FunctionRetryInputsRequest
-        ],
-    ) -> typing.Union[
-        modal_proto.api_pb2.FunctionPutInputsResponse, modal_proto.api_pb2.FunctionRetryInputsResponse
-    ]: ...
+    @property
+    def _function_inputs_retry(self) -> modal._utils.grpc_utils.Retry: ...
 
 class SyncInputPumper(InputPumper):
     """Reads inputs from a queue of FunctionPutInputsItems, and sends them to the server."""

modal/partial_function.pyi

@@ -329,7 +329,10 @@ def batched(
     ...
 
 def concurrent(
-    _warn_parentheses_missing=None, *, max_inputs: int, target_inputs: typing.Optional[int] = None
+    _warn_parentheses_missing=None,
+    *,
+    max_inputs: typing.Optional[int] = None,
+    target_inputs: typing.Optional[int] = None,
 ) -> collections.abc.Callable[
     [
         typing.Union[

modal/queue.py

@@ -25,7 +25,7 @@ from ._resolver import Resolver
 from ._serialization import deserialize, serialize
 from ._utils.async_utils import TaskContext, synchronize_api, warn_if_generator_is_not_consumed
 from ._utils.deprecation import deprecation_warning, warn_if_passing_namespace
-from ._utils.grpc_utils import retry_transient_errors
+from ._utils.grpc_utils import Retry
 from ._utils.name_utils import check_object_name
 from ._utils.time_utils import as_timestamp, timestamp_to_localized_dt
 from .client import _Client
@@ -95,7 +95,7 @@ class _QueueManager:
             object_creation_type=object_creation_type,
         )
         try:
-            await retry_transient_errors(client.stub.QueueGetOrCreate, req)
+            await client.stub.QueueGetOrCreate(req)
         except GRPCError as exc:
             if exc.status == Status.ALREADY_EXISTS and not allow_existing:
                 raise AlreadyExistsError(exc.message)
@@ -147,7 +147,7 @@ class _QueueManager:
             req = api_pb2.QueueListRequest(
                 environment_name=_get_environment_name(environment_name), pagination=pagination
             )
-            resp = await retry_transient_errors(client.stub.QueueList, req)
+            resp = await client.stub.QueueList(req)
             items.extend(resp.queues)
             finished = (len(resp.queues) < max_page_size) or (max_objects is not None and len(items) >= max_objects)
             return finished
@@ -205,7 +205,7 @@ class _QueueManager:
                 raise
         else:
             req = api_pb2.QueueDeleteRequest(queue_id=obj.object_id)
-            await retry_transient_errors(obj._client.stub.QueueDelete, req)
+            await obj._client.stub.QueueDelete(req)
 
 
 QueueManager = synchronize_api(_QueueManager)
@@ -424,7 +424,7 @@ class _Queue(_Object, type_prefix="qu"):
             n_values=n_values,
         )
 
-        response = await retry_transient_errors(self._client.stub.QueueGet, request)
+        response = await self._client.stub.QueueGet(request)
         if response.values:
             return [deserialize(value, self._client) for value in response.values]
         else:
@@ -449,7 +449,7 @@ class _Queue(_Object, type_prefix="qu"):
                 n_values=n_values,
             )
 
-            response = await retry_transient_errors(self._client.stub.QueueGet, request)
+            response = await self._client.stub.QueueGet(request)
 
             if response.values:
                 return [deserialize(value, self._client) for value in response.values]
@@ -469,7 +469,7 @@ class _Queue(_Object, type_prefix="qu"):
             partition_key=self.validate_partition_key(partition),
             all_partitions=all,
         )
-        await retry_transient_errors(self._client.stub.QueueClear, request)
+        await self._client.stub.QueueClear(request)
 
     @live_method
     async def get(
@@ -578,14 +578,15 @@ class _Queue(_Object, type_prefix="qu"):
             partition_ttl_seconds=partition_ttl,
         )
         try:
-            await retry_transient_errors(
-                self._client.stub.QueuePut,
+            await self._client.stub.QueuePut(
                 request,
                 # A full queue will return this status.
-                additional_status_codes=[Status.RESOURCE_EXHAUSTED],
-                max_delay=30.0,
-                max_retries=None,
-                total_timeout=timeout,
+                retry=Retry(
+                    additional_status_codes=[Status.RESOURCE_EXHAUSTED],
+                    max_delay=30.0,
+                    max_retries=None,
+                    total_timeout=timeout,
+                ),
             )
         except GRPCError as exc:
             if exc.status == Status.RESOURCE_EXHAUSTED:
@@ -605,7 +606,7 @@ class _Queue(_Object, type_prefix="qu"):
             partition_ttl_seconds=partition_ttl,
         )
         try:
-            await retry_transient_errors(self._client.stub.QueuePut, request)
+            await self._client.stub.QueuePut(request)
         except GRPCError as exc:
             if exc.status == Status.RESOURCE_EXHAUSTED:
                 raise queue.Full(exc.message)
@@ -625,7 +626,7 @@ class _Queue(_Object, type_prefix="qu"):
             partition_key=self.validate_partition_key(partition),
             total=total,
         )
-        response = await retry_transient_errors(self._client.stub.QueueLen, request)
+        response = await self._client.stub.QueueLen(request)
         return response.len
 
     @warn_if_generator_is_not_consumed()
@@ -651,9 +652,7 @@ class _Queue(_Object, type_prefix="qu"):
                 item_poll_timeout=poll_duration,
             )
 
-            response: api_pb2.QueueNextItemsResponse = await retry_transient_errors(
-                self._client.stub.QueueNextItems, request
-            )
+            response: api_pb2.QueueNextItemsResponse = await self._client.stub.QueueNextItems(request)
             if response.items:
                 for item in response.items:
                     yield deserialize(item.value, self._client)

modal/runner.py

@@ -19,6 +19,7 @@ from synchronicity.async_wrap import asynccontextmanager
 
 import modal._runtime.execution_context
 import modal_proto.api_pb2
+from modal._utils.grpc_utils import Retry
 from modal_proto import api_pb2
 
 from ._functions import _Function
@@ -29,7 +30,6 @@ from ._traceback import print_server_warnings, traceback_contains_remote_call
 from ._utils.async_utils import TaskContext, gather_cancel_on_exc, synchronize_api
 from ._utils.deprecation import warn_if_passing_namespace
 from ._utils.git_utils import get_git_commit_info
-from ._utils.grpc_utils import retry_transient_errors
 from ._utils.name_utils import check_object_name, is_valid_tag
 from .client import HEARTBEAT_INTERVAL, HEARTBEAT_TIMEOUT, _Client
 from .cls import _Cls
@@ -54,14 +54,14 @@ async def _heartbeat(client: _Client, app_id: str) -> None:
     # TODO(erikbern): we should capture exceptions here
     # * if request fails: destroy the client
     # * if server says the app is gone: print a helpful warning about detaching
-    await retry_transient_errors(client.stub.AppHeartbeat, request, attempt_timeout=HEARTBEAT_TIMEOUT)
+    await client.stub.AppHeartbeat(request, retry=Retry(attempt_timeout=HEARTBEAT_TIMEOUT))
 
 
 async def _init_local_app_existing(client: _Client, existing_app_id: str, environment_name: str) -> RunningApp:
     # Get all the objects first
     obj_req = api_pb2.AppGetLayoutRequest(app_id=existing_app_id)
     obj_resp, _ = await gather_cancel_on_exc(
-        retry_transient_errors(client.stub.AppGetLayout, obj_req),
+        client.stub.AppGetLayout(obj_req),
         # Cache the environment associated with the app now as we will use it later
         _get_environment_cached(environment_name, client),
     )
@@ -86,7 +86,7 @@ async def _init_local_app_new(
         app_state=app_state,  # type: ignore
     )
     app_resp, _ = await gather_cancel_on_exc(  # TODO: use TaskGroup?
-        retry_transient_errors(client.stub.AppCreate, app_req),
+        client.stub.AppCreate(app_req),
         # Cache the environment associated with the app now as we will use it later
         _get_environment_cached(environment_name, client),
     )
@@ -109,7 +109,7 @@ async def _init_local_app_from_name(
         name=name,
         environment_name=environment_name,
     )
-    app_resp = await retry_transient_errors(client.stub.AppGetByDeploymentName, app_req)
+    app_resp = await client.stub.AppGetByDeploymentName(app_req)
     existing_app_id = app_resp.app_id or None
 
     # Grab the app
@@ -201,7 +201,7 @@ async def _publish_app(
     )
 
     try:
-        response = await retry_transient_errors(client.stub.AppPublish, request)
+        response = await client.stub.AppPublish(request)
     except GRPCError as exc:
         if exc.status == Status.INVALID_ARGUMENT or exc.status == Status.FAILED_PRECONDITION:
             raise InvalidError(exc.message)
@@ -225,7 +225,7 @@ async def _disconnect(
 
     logger.debug("Sending app disconnect/stop request")
     req_disconnect = api_pb2.AppClientDisconnectRequest(app_id=app_id, reason=reason, exception=exc_str)
-    await retry_transient_errors(client.stub.AppClientDisconnect, req_disconnect)
+    await client.stub.AppClientDisconnect(req_disconnect)
     logger.debug("App disconnected")
 
 
@@ -354,9 +354,10 @@ async def _run_app(
         detached_disconnect_msg = (
             "The detached App will keep running. You can track its progress on the Dashboard: "
             f"[magenta underline]{running_app.app_page_url}[/magenta underline]"
-            "\n"
-            f"\nStream logs: [green]modal app logs {running_app.app_id}[/green]"
-            f"\nStop the App: [green]modal app stop {running_app.app_id}[/green]"
+            "\n\nStream App logs:\n"
+            f"[green]modal app logs {running_app.app_id}[/green]"
+            "\n\nStop the App:\n"
+            f"[green]modal app stop {running_app.app_id}[/green]"
         )
 
         try:
@@ -634,7 +635,7 @@ async def _interactive_shell(
         except InteractiveTimeoutError:
             # Check on status of Sandbox. It may have crashed, causing connection failure.
             req = api_pb2.SandboxWaitRequest(sandbox_id=sandbox._object_id, timeout=0)
-            resp = await retry_transient_errors(sandbox._client.stub.SandboxWait, req)
+            resp = await sandbox._client.stub.SandboxWait(req)
             if resp.result.exception:
                 raise RemoteError(resp.result.exception)
             else: