modal 1.1.5.dev66__py3-none-any.whl → 1.3.1.dev8__py3-none-any.whl

modal/network_file_system.pyi

@@ -54,7 +54,12 @@ class _NetworkFileSystem(modal._object._Object):
     """
     @staticmethod
     def from_name(
-        name: str, *, namespace=None, environment_name: typing.Optional[str] = None, create_if_missing: bool = False
+        name: str,
+        *,
+        namespace=None,
+        environment_name: typing.Optional[str] = None,
+        create_if_missing: bool = False,
+        client: typing.Optional[modal.client._Client] = None,
     ) -> _NetworkFileSystem:
         """Reference a NetworkFileSystem by its name, creating if necessary.
 
@@ -163,8 +168,6 @@ class _NetworkFileSystem(modal._object._Object):
         """Remove a file in a network file system."""
         ...
 
-SUPERSELF = typing.TypeVar("SUPERSELF", covariant=True)
-
 class NetworkFileSystem(modal.object.Object):
     """A shared, writable file system accessible by one or more Modal functions.
 
@@ -210,7 +213,12 @@ class NetworkFileSystem(modal.object.Object):
 
     @staticmethod
     def from_name(
-        name: str, *, namespace=None, environment_name: typing.Optional[str] = None, create_if_missing: bool = False
+        name: str,
+        *,
+        namespace=None,
+        environment_name: typing.Optional[str] = None,
+        create_if_missing: bool = False,
+        client: typing.Optional[modal.client.Client] = None,
     ) -> NetworkFileSystem:
         """Reference a NetworkFileSystem by its name, creating if necessary.
 
@@ -228,27 +236,52 @@ class NetworkFileSystem(modal.object.Object):
         """
         ...
 
-    @classmethod
-    def ephemeral(
-        cls: type[NetworkFileSystem],
-        client: typing.Optional[modal.client.Client] = None,
-        environment_name: typing.Optional[str] = None,
-        _heartbeat_sleep: float = 300,
-    ) -> synchronicity.combined_types.AsyncAndBlockingContextManager[NetworkFileSystem]:
-        """Creates a new ephemeral network filesystem within a context manager:
+    class __ephemeral_spec(typing_extensions.Protocol):
+        def __call__(
+            self,
+            /,
+            client: typing.Optional[modal.client.Client] = None,
+            environment_name: typing.Optional[str] = None,
+            _heartbeat_sleep: float = 300,
+        ) -> synchronicity.combined_types.AsyncAndBlockingContextManager[NetworkFileSystem]:
+            """Creates a new ephemeral network filesystem within a context manager:
+
+            Usage:
+            ```python
+            with modal.NetworkFileSystem.ephemeral() as nfs:
+                assert nfs.listdir("/") == []
+            ```
+
+            ```python notest
+            async with modal.NetworkFileSystem.ephemeral() as nfs:
+                assert await nfs.listdir("/") == []
+            ```
+            """
+            ...
 
-        Usage:
-        ```python
-        with modal.NetworkFileSystem.ephemeral() as nfs:
-            assert nfs.listdir("/") == []
-        ```
+        def aio(
+            self,
+            /,
+            client: typing.Optional[modal.client.Client] = None,
+            environment_name: typing.Optional[str] = None,
+            _heartbeat_sleep: float = 300,
+        ) -> typing.AsyncContextManager[NetworkFileSystem]:
+            """Creates a new ephemeral network filesystem within a context manager:
+
+            Usage:
+            ```python
+            with modal.NetworkFileSystem.ephemeral() as nfs:
+                assert nfs.listdir("/") == []
+            ```
+
+            ```python notest
+            async with modal.NetworkFileSystem.ephemeral() as nfs:
+                assert await nfs.listdir("/") == []
+            ```
+            """
+            ...
 
-        ```python notest
-        async with modal.NetworkFileSystem.ephemeral() as nfs:
-            assert await nfs.listdir("/") == []
-        ```
-        """
-        ...
+    ephemeral: typing.ClassVar[__ephemeral_spec]
 
     class __create_deployed_spec(typing_extensions.Protocol):
         def __call__(
@@ -273,7 +306,7 @@ class NetworkFileSystem(modal.object.Object):
             """mdmd:hidden"""
             ...
 
-    create_deployed: __create_deployed_spec
+    create_deployed: typing.ClassVar[__create_deployed_spec]
 
     class __delete_spec(typing_extensions.Protocol):
         def __call__(
@@ -291,9 +324,9 @@ class NetworkFileSystem(modal.object.Object):
             environment_name: typing.Optional[str] = None,
         ): ...
 
-    delete: __delete_spec
+    delete: typing.ClassVar[__delete_spec]
 
-    class __write_file_spec(typing_extensions.Protocol[SUPERSELF]):
+    class __write_file_spec(typing_extensions.Protocol):
         def __call__(
             self,
             /,
@@ -326,9 +359,9 @@ class NetworkFileSystem(modal.object.Object):
             """
             ...
 
-    write_file: __write_file_spec[typing_extensions.Self]
+    write_file: __write_file_spec
 
-    class __read_file_spec(typing_extensions.Protocol[SUPERSELF]):
+    class __read_file_spec(typing_extensions.Protocol):
         def __call__(self, /, path: str) -> typing.Iterator[bytes]:
             """Read a file from the network file system"""
             ...
@@ -337,9 +370,9 @@ class NetworkFileSystem(modal.object.Object):
             """Read a file from the network file system"""
             ...
 
-    read_file: __read_file_spec[typing_extensions.Self]
+    read_file: __read_file_spec
 
-    class __iterdir_spec(typing_extensions.Protocol[SUPERSELF]):
+    class __iterdir_spec(typing_extensions.Protocol):
         def __call__(self, /, path: str) -> typing.Iterator[modal.volume.FileEntry]:
             """Iterate over all files in a directory in the network file system.
 
@@ -360,9 +393,9 @@ class NetworkFileSystem(modal.object.Object):
             """
             ...
 
-    iterdir: __iterdir_spec[typing_extensions.Self]
+    iterdir: __iterdir_spec
 
-    class __add_local_file_spec(typing_extensions.Protocol[SUPERSELF]):
+    class __add_local_file_spec(typing_extensions.Protocol):
         def __call__(
             self,
             /,
@@ -378,9 +411,9 @@ class NetworkFileSystem(modal.object.Object):
             progress_cb: typing.Optional[collections.abc.Callable[..., typing.Any]] = None,
         ): ...
 
-    add_local_file: __add_local_file_spec[typing_extensions.Self]
+    add_local_file: __add_local_file_spec
 
-    class __add_local_dir_spec(typing_extensions.Protocol[SUPERSELF]):
+    class __add_local_dir_spec(typing_extensions.Protocol):
         def __call__(
             self,
             /,
@@ -396,9 +429,9 @@ class NetworkFileSystem(modal.object.Object):
             progress_cb: typing.Optional[collections.abc.Callable[..., typing.Any]] = None,
         ): ...
 
-    add_local_dir: __add_local_dir_spec[typing_extensions.Self]
+    add_local_dir: __add_local_dir_spec
 
-    class __listdir_spec(typing_extensions.Protocol[SUPERSELF]):
+    class __listdir_spec(typing_extensions.Protocol):
         def __call__(self, /, path: str) -> list[modal.volume.FileEntry]:
             """List all files in a directory in the network file system.
 
@@ -419,9 +452,9 @@ class NetworkFileSystem(modal.object.Object):
             """
             ...
 
-    listdir: __listdir_spec[typing_extensions.Self]
+    listdir: __listdir_spec
 
-    class __remove_file_spec(typing_extensions.Protocol[SUPERSELF]):
+    class __remove_file_spec(typing_extensions.Protocol):
         def __call__(self, /, path: str, recursive=False):
             """Remove a file in a network file system."""
             ...
@@ -430,4 +463,4 @@ class NetworkFileSystem(modal.object.Object):
             """Remove a file in a network file system."""
             ...
 
-    remove_file: __remove_file_spec[typing_extensions.Self]
+    remove_file: __remove_file_spec

modal/object.pyi

@@ -1,5 +1,6 @@
 import collections.abc
 import google.protobuf.message
+import modal._load_context
 import modal._resolver
 import modal.client
 import typing
@@ -12,12 +13,14 @@ class Object:
     _prefix_to_type: typing.ClassVar[dict[str, type]]
     _load: typing.Optional[
         collections.abc.Callable[
-            [typing_extensions.Self, modal._resolver.Resolver, typing.Optional[str]], collections.abc.Awaitable[None]
+            [typing_extensions.Self, modal._resolver.Resolver, modal._load_context.LoadContext, typing.Optional[str]],
+            collections.abc.Awaitable[None],
         ]
     ]
     _preload: typing.Optional[
         collections.abc.Callable[
-            [typing_extensions.Self, modal._resolver.Resolver, typing.Optional[str]], collections.abc.Awaitable[None]
+            [typing_extensions.Self, modal._resolver.Resolver, modal._load_context.LoadContext, typing.Optional[str]],
+            collections.abc.Awaitable[None],
         ]
     ]
     _rep: str
@@ -27,6 +30,7 @@ class Object:
     _deduplication_key: typing.Optional[
         collections.abc.Callable[[], collections.abc.Awaitable[collections.abc.Hashable]]
     ]
+    _load_context_overrides: modal._load_context.LoadContext
     _object_id: typing.Optional[str]
     _client: typing.Optional[modal.client.Client]
     _is_hydrated: bool
@@ -40,22 +44,40 @@ class Object:
     @classmethod
     def __init_subclass__(cls, type_prefix: typing.Optional[str] = None): ...
 
-    class ___init_spec(typing_extensions.Protocol[SUPERSELF]):
+    class ___init_spec(typing_extensions.Protocol):
         def __call__(
             self,
             /,
             rep: str,
             load: typing.Optional[
-                collections.abc.Callable[[SUPERSELF, modal._resolver.Resolver, typing.Optional[str]], None]
+                collections.abc.Callable[
+                    [
+                        typing_extensions.Self,
+                        modal._resolver.Resolver,
+                        modal._load_context.LoadContext,
+                        typing.Optional[str],
+                    ],
+                    None,
+                ]
             ] = None,
             is_another_app: bool = False,
             preload: typing.Optional[
-                collections.abc.Callable[[SUPERSELF, modal._resolver.Resolver, typing.Optional[str]], None]
+                collections.abc.Callable[
+                    [
+                        typing_extensions.Self,
+                        modal._resolver.Resolver,
+                        modal._load_context.LoadContext,
+                        typing.Optional[str],
+                    ],
+                    None,
+                ]
             ] = None,
             hydrate_lazily: bool = False,
             deps: typing.Optional[collections.abc.Callable[..., collections.abc.Sequence[Object]]] = None,
             deduplication_key: typing.Optional[collections.abc.Callable[[], collections.abc.Hashable]] = None,
             name: typing.Optional[str] = None,
+            *,
+            load_context_overrides: typing.Optional[modal._load_context.LoadContext] = None,
         ): ...
         def aio(
             self,
@@ -63,13 +85,25 @@ class Object:
             rep: str,
             load: typing.Optional[
                 collections.abc.Callable[
-                    [SUPERSELF, modal._resolver.Resolver, typing.Optional[str]], collections.abc.Awaitable[None]
+                    [
+                        typing_extensions.Self,
+                        modal._resolver.Resolver,
+                        modal._load_context.LoadContext,
+                        typing.Optional[str],
+                    ],
+                    collections.abc.Awaitable[None],
                 ]
             ] = None,
             is_another_app: bool = False,
             preload: typing.Optional[
                 collections.abc.Callable[
-                    [SUPERSELF, modal._resolver.Resolver, typing.Optional[str]], collections.abc.Awaitable[None]
+                    [
+                        typing_extensions.Self,
+                        modal._resolver.Resolver,
+                        modal._load_context.LoadContext,
+                        typing.Optional[str],
+                    ],
+                    collections.abc.Awaitable[None],
                 ]
             ] = None,
             hydrate_lazily: bool = False,
@@ -78,9 +112,11 @@ class Object:
                 collections.abc.Callable[[], collections.abc.Awaitable[collections.abc.Hashable]]
             ] = None,
             name: typing.Optional[str] = None,
+            *,
+            load_context_overrides: typing.Optional[modal._load_context.LoadContext] = None,
         ): ...
 
-    _init: ___init_spec[typing_extensions.Self]
+    _init: ___init_spec
 
     def _unhydrate(self): ...
     def _initialize_from_empty(self): ...
@@ -98,20 +134,76 @@ class Object:
         """
         ...
 
-    @classmethod
-    def _from_loader(
-        cls,
-        load: collections.abc.Callable[[typing_extensions.Self, modal._resolver.Resolver, typing.Optional[str]], None],
-        rep: str,
-        is_another_app: bool = False,
-        preload: typing.Optional[
-            collections.abc.Callable[[typing_extensions.Self, modal._resolver.Resolver, typing.Optional[str]], None]
-        ] = None,
-        hydrate_lazily: bool = False,
-        deps: typing.Optional[collections.abc.Callable[..., collections.abc.Sequence[Object]]] = None,
-        deduplication_key: typing.Optional[collections.abc.Callable[[], collections.abc.Hashable]] = None,
-        name: typing.Optional[str] = None,
-    ): ...
+    class ___from_loader_spec(typing_extensions.Protocol):
+        def __call__(
+            self,
+            /,
+            load: collections.abc.Callable[
+                [
+                    typing_extensions.Self,
+                    modal._resolver.Resolver,
+                    modal._load_context.LoadContext,
+                    typing.Optional[str],
+                ],
+                None,
+            ],
+            rep: str,
+            is_another_app: bool = False,
+            preload: typing.Optional[
+                collections.abc.Callable[
+                    [
+                        typing_extensions.Self,
+                        modal._resolver.Resolver,
+                        modal._load_context.LoadContext,
+                        typing.Optional[str],
+                    ],
+                    None,
+                ]
+            ] = None,
+            hydrate_lazily: bool = False,
+            deps: typing.Optional[collections.abc.Callable[..., collections.abc.Sequence[Object]]] = None,
+            deduplication_key: typing.Optional[collections.abc.Callable[[], collections.abc.Hashable]] = None,
+            name: typing.Optional[str] = None,
+            *,
+            load_context_overrides: modal._load_context.LoadContext,
+        ): ...
+        def aio(
+            self,
+            /,
+            load: collections.abc.Callable[
+                [
+                    typing_extensions.Self,
+                    modal._resolver.Resolver,
+                    modal._load_context.LoadContext,
+                    typing.Optional[str],
+                ],
+                collections.abc.Awaitable[None],
+            ],
+            rep: str,
+            is_another_app: bool = False,
+            preload: typing.Optional[
+                collections.abc.Callable[
+                    [
+                        typing_extensions.Self,
+                        modal._resolver.Resolver,
+                        modal._load_context.LoadContext,
+                        typing.Optional[str],
+                    ],
+                    collections.abc.Awaitable[None],
+                ]
+            ] = None,
+            hydrate_lazily: bool = False,
+            deps: typing.Optional[collections.abc.Callable[..., collections.abc.Sequence[Object]]] = None,
+            deduplication_key: typing.Optional[
+                collections.abc.Callable[[], collections.abc.Awaitable[collections.abc.Hashable]]
+            ] = None,
+            name: typing.Optional[str] = None,
+            *,
+            load_context_overrides: modal._load_context.LoadContext,
+        ): ...
+
+    _from_loader: typing.ClassVar[___from_loader_spec]
+
     @staticmethod
     def _get_type_from_id(object_id: str) -> type[Object]: ...
     @classmethod

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]:
@@ -671,8 +667,8 @@ async def _map_invocation_inputplane(
     last_entry_id = ""
 
     # The input-plane server returns this after the first request.
-    function_call_id = None
-    function_call_id_received = asyncio.Event()
+    map_token = None
+    map_token_received = asyncio.Event()
 
     # Single priority queue that holds *both* fresh inputs (timestamp == now)
     # and future retries (timestamp > now).
@@ -751,7 +747,7 @@ async def _map_invocation_inputplane(
         yield
 
     async def pump_inputs():
-        nonlocal function_call_id, max_inputs_outstanding
+        nonlocal map_token, max_inputs_outstanding
         async for batch in queue_batch_iterator(queue, max_batch_size=MAP_INVOCATION_CHUNK_SIZE):
             # Convert the queued items into the proto format expected by the RPC.
             request_items: list[api_pb2.MapStartOrContinueItem] = [
@@ -763,20 +759,21 @@ async def _map_invocation_inputplane(
             # Build request
             request = api_pb2.MapStartOrContinueRequest(
                 function_id=function.object_id,
-                function_call_id=function_call_id,
+                map_token=map_token,
                 parent_input_id=current_input_id() or "",
                 items=request_items,
             )
 
             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
@@ -789,9 +786,9 @@ async def _map_invocation_inputplane(
 
             # Set the function call id and actual retry policy with the data from the first response.
             # This conditional is skipped for subsequent iterations of this for-loop.
-            if function_call_id is None:
-                function_call_id = response.function_call_id
-                function_call_id_received.set()
+            if map_token is None:
+                map_token = response.map_token
+                map_token_received.set()
                 max_inputs_outstanding = response.max_inputs_outstanding or MAX_INPUTS_OUTSTANDING_DEFAULT
                 map_items_manager.set_retry_policy(response.retry_policy)
                 # Update the retry policy for the first batch of inputs.
@@ -804,8 +801,8 @@ async def _map_invocation_inputplane(
         nonlocal last_entry_id  # shared with get_all_outputs
         try:
             while not map_done_event.is_set():
-                if function_call_id is None:
-                    await function_call_id_received.wait()
+                if map_token is None:
+                    await map_token_received.wait()
                     continue
 
                 sleep_task = asyncio.create_task(asyncio.sleep(1))
@@ -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)
@@ -847,23 +844,24 @@ async def _map_invocation_inputplane(
             last_entry_id
 
         while not map_done_event.is_set():
-            if function_call_id is None:
-                await function_call_id_received.wait()
+            if map_token is None:
+                await map_token_received.wait()
                 continue
 
             request = api_pb2.MapAwaitRequest(
-                function_call_id=function_call_id,
+                map_token=map_token,
                 last_entry_id=last_entry_id,
                 requested_at=time.time(),
                 timeout=OUTPUTS_TIMEOUT,
             )
             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,
                 )
             )
@@ -963,7 +961,7 @@ async def _map_invocation_inputplane(
                 f"Map stats:\nsuccessful_completions={successful_completions} failed_completions={failed_completions} "
                 f"no_context_duplicates={no_context_duplicates} stale_retry_duplicates={stale_retry_duplicates} "
                 f"already_complete_duplicates={already_complete_duplicates} retried_outputs={retried_outputs} "
-                f"function_call_id={function_call_id} max_inputs_outstanding={max_inputs_outstanding} "
+                f"map_token={map_token} max_inputs_outstanding={max_inputs_outstanding} "
                 f"map_items_manager_size={len(map_items_manager)} input_queue_size={input_queue_size}"
             )