modal 0.73.116__py3-none-any.whl → 0.73.126__py3-none-any.whl

modal/cls.py

@@ -21,7 +21,7 @@ from ._partial_function import (
 )
 from ._resolver import Resolver
 from ._resources import convert_fn_config_to_resources_config
-from ._serialization import PYTHON_TO_PROTO_TYPE, check_valid_cls_constructor_arg
+from ._serialization import check_valid_cls_constructor_arg, get_proto_parameter_type
 from ._traceback import print_server_warnings
 from ._utils.async_utils import synchronize_api, synchronizer
 from ._utils.deprecation import deprecation_warning, renamed_parameter, warn_on_renamed_autoscaler_settings
@@ -362,15 +362,6 @@ class _Obj:
 Obj = synchronize_api(_Obj)
 
 
-def _validate_parameter_type(cls_name: str, parameter_name: str, parameter_type: type):
-    if parameter_type not in PYTHON_TO_PROTO_TYPE:
-        type_name = getattr(parameter_type, "__name__", repr(parameter_type))
-        supported = ", ".join(parameter_type.__name__ for parameter_type in PYTHON_TO_PROTO_TYPE.keys())
-        raise InvalidError(
-            f"{cls_name}.{parameter_name}: {type_name} is not a supported parameter type. Use one of: {supported}"
-        )
-
-
 class _Cls(_Object, type_prefix="cs"):
     """
     Cls adds method pooling and [lifecycle hook](/docs/guide/lifecycle-functions) behavior
@@ -467,12 +458,11 @@ class _Cls(_Object, type_prefix="cs"):
         annotations = user_cls.__dict__.get("__annotations__", {})  # compatible with older pythons
         missing_annotations = params.keys() - annotations.keys()
         if missing_annotations:
-            raise InvalidError("All modal.parameter() specifications need to be type annotated")
+            raise InvalidError("All modal.parameter() specifications need to be type-annotated")
 
         annotated_params = {k: t for k, t in annotations.items() if k in params}
         for k, t in annotated_params.items():
-            if t not in PYTHON_TO_PROTO_TYPE:
-                _validate_parameter_type(user_cls.__name__, k, t)
+            get_proto_parameter_type(t)
 
     @staticmethod
     def from_local(user_cls, app: "modal.app._App", class_service_function: _Function) -> "_Cls":

modal/cls.pyi

@@ -109,8 +109,6 @@ class Obj:
     async def _aenter(self): ...
     def __getattr__(self, k): ...
 
-def _validate_parameter_type(cls_name: str, parameter_name: str, parameter_type: type): ...
-
 class _Cls(modal._object._Object):
     _class_service_function: typing.Optional[modal._functions._Function]
     _options: typing.Optional[_ServiceOptions]

modal/functions.pyi

@@ -198,11 +198,11 @@ class Function(
 
     _call_generator_nowait: ___call_generator_nowait_spec[typing_extensions.Self]
 
-    class __remote_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
+    class __remote_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
         def __call__(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> ReturnType_INNER: ...
         async def aio(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> ReturnType_INNER: ...
 
-    remote: __remote_spec[modal._functions.ReturnType, modal._functions.P, typing_extensions.Self]
+    remote: __remote_spec[modal._functions.P, modal._functions.ReturnType, typing_extensions.Self]
 
     class __remote_gen_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, *args, **kwargs) -> typing.Generator[typing.Any, None, None]: ...
@@ -217,19 +217,19 @@ class Function(
         self, *args: modal._functions.P.args, **kwargs: modal._functions.P.kwargs
     ) -> modal._functions.OriginalReturnType: ...
 
-    class ___experimental_spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
+    class ___experimental_spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
         def __call__(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
         async def aio(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
 
     _experimental_spawn: ___experimental_spawn_spec[
-        modal._functions.ReturnType, modal._functions.P, typing_extensions.Self
+        modal._functions.P, modal._functions.ReturnType, typing_extensions.Self
     ]
 
-    class __spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
+    class __spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
         def __call__(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
         async def aio(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
 
-    spawn: __spawn_spec[modal._functions.ReturnType, modal._functions.P, typing_extensions.Self]
+    spawn: __spawn_spec[modal._functions.P, modal._functions.ReturnType, typing_extensions.Self]
 
     def get_raw_f(self) -> collections.abc.Callable[..., typing.Any]: ...
 

modal/parallel_map.py

@@ -1,5 +1,6 @@
 # Copyright Modal Labs 2024
 import asyncio
+import enum
 import time
 import typing
 from dataclasses import dataclass
@@ -10,6 +11,7 @@ from grpclib import Status
 from modal._runtime.execution_context import current_input_id
 from modal._utils.async_utils import (
     AsyncOrSyncIterable,
+    TimestampPriorityQueue,
     aclosing,
     async_map_ordered,
     async_merge,
@@ -28,7 +30,9 @@ from modal._utils.function_utils import (
     _process_result,
 )
 from modal._utils.grpc_utils import RETRYABLE_GRPC_STATUS_CODES, RetryWarningMessage, retry_transient_errors
+from modal._utils.jwt_utils import DecodedJwt
 from modal.config import logger
+from modal.retries import RetryManager
 from modal_proto import api_pb2
 
 if typing.TYPE_CHECKING:
@@ -66,6 +70,12 @@ class _OutputValue:
     value: Any
 
 
+# maximum number of inputs that can be in progress (either queued to be sent,
+# or waiting for completion). if this limit is reached, we will block sending
+# more inputs to the server until some of the existing inputs are completed.
+MAP_MAX_INPUTS_OUTSTANDING = 1000
+
+# maximum number of inputs to send to the server in a single request
 MAP_INVOCATION_CHUNK_SIZE = 49
 
 if typing.TYPE_CHECKING:
@@ -79,6 +89,7 @@ async def _map_invocation(
     order_outputs: bool,
     return_exceptions: bool,
     count_update_callback: Optional[Callable[[int, int], None]],
+    function_call_invocation_type: "api_pb2.FunctionCallInvocationType.ValueType",
 ):
     assert client.stub
     request = api_pb2.FunctionMapRequest(
@@ -86,28 +97,43 @@ async def _map_invocation(
         parent_input_id=current_input_id() or "",
         function_call_type=api_pb2.FUNCTION_CALL_TYPE_MAP,
         return_exceptions=return_exceptions,
+        function_call_invocation_type=function_call_invocation_type,
     )
-    response = await retry_transient_errors(client.stub.FunctionMap, request)
+    response: api_pb2.FunctionMapResponse = await retry_transient_errors(client.stub.FunctionMap, request)
 
     function_call_id = response.function_call_id
+    function_call_jwt = response.function_call_jwt
+    retry_policy = response.retry_policy
+    sync_client_retries_enabled = response.sync_client_retries_enabled
 
     have_all_inputs = False
-    num_inputs = 0
-    num_outputs = 0
+    inputs_created = 0
+    inputs_sent = 0
+    inputs_retried = 0
+    outputs_completed = 0
+    outputs_received = 0
+    retried_outputs = 0
+    successful_completions = 0
+    failed_completions = 0
+    already_complete_duplicates = 0
+    stale_retry_duplicates = 0
+    no_context_duplicates = 0
 
     def count_update():
         if count_update_callback is not None:
-            count_update_callback(num_outputs, num_inputs)
+            count_update_callback(outputs_completed, inputs_created)
 
-    pending_outputs: dict[str, int] = {}  # Map input_id -> next expected gen_index value
+    retry_queue = TimestampPriorityQueue()
     completed_outputs: set[str] = set()  # Set of input_ids whose outputs are complete (expecting no more values)
-
-    input_queue: asyncio.Queue = asyncio.Queue()
+    input_queue: asyncio.Queue[api_pb2.FunctionPutInputsItem | None] = asyncio.Queue()
+    map_items_manager = _MapItemsManager(
+        retry_policy, function_call_invocation_type, retry_queue, sync_client_retries_enabled
+    )
 
     async def create_input(argskwargs):
-        nonlocal num_inputs
-        idx = num_inputs
-        num_inputs += 1
+        nonlocal inputs_created
+        idx = inputs_created
+        inputs_created += 1
         (args, kwargs) = argskwargs
         return await _create_input(args, kwargs, client, idx=idx, method_name=function._use_method_name)
 
@@ -119,6 +145,8 @@ async def _map_invocation(
             yield raw_input  # args, kwargs
 
     async def drain_input_generator():
+        nonlocal have_all_inputs
+
         # Parallelize uploading blobs
         async with aclosing(
             async_map_ordered(input_iter(), create_input, concurrency=BLOB_MAX_PARALLELISM)
@@ -132,49 +160,100 @@ async def _map_invocation(
 
     async def pump_inputs():
         assert client.stub
-        nonlocal have_all_inputs, num_inputs
-        async for items in queue_batch_iterator(input_queue, MAP_INVOCATION_CHUNK_SIZE):
+        nonlocal have_all_inputs, inputs_created, inputs_sent
+        async for items in queue_batch_iterator(input_queue, max_batch_size=MAP_INVOCATION_CHUNK_SIZE):
+            # Add items to the manager. Their state will be SENDING.
+            await map_items_manager.add_items(items)
             request = api_pb2.FunctionPutInputsRequest(
-                function_id=function.object_id, inputs=items, function_call_id=function_call_id
+                function_id=function.object_id,
+                inputs=items,
+                function_call_id=function_call_id,
             )
             logger.debug(
                 f"Pushing {len(items)} inputs to server. Num queued inputs awaiting push is {input_queue.qsize()}."
             )
-            # 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 {function._function_name} is limited."
-                " Common bottlenecks include slow iteration over results, or function backlogs.",
-                warning_interval=8,
-                errors_to_warn_for=[Status.RESOURCE_EXHAUSTED])
-            resp = await retry_transient_errors(
-                client.stub.FunctionPutInputs,
-                request,
-                max_retries=None,
-                max_delay=PUMP_INPUTS_MAX_RETRY_DELAY,
-                additional_status_codes=[Status.RESOURCE_EXHAUSTED],
-                retry_warning_message=retry_warning_message)
+
+            resp = await send_inputs(client.stub.FunctionPutInputs, request)
             count_update()
-            for item in resp.inputs:
-                pending_outputs.setdefault(item.input_id, 0)
+            inputs_sent += len(items)
+            # Change item state to WAITING_FOR_OUTPUT, and set the input_id and input_jwt which are in the response.
+            map_items_manager.handle_put_inputs_response(resp.inputs)
             logger.debug(
                 f"Successfully pushed {len(items)} inputs to server. "
                 f"Num queued inputs awaiting push is {input_queue.qsize()}."
             )
-
         have_all_inputs = True
         yield
 
+    async def retry_inputs():
+        nonlocal inputs_retried
+        async for retriable_idxs in queue_batch_iterator(retry_queue, max_batch_size=MAP_INVOCATION_CHUNK_SIZE):
+            # For each index, use the context in the manager to create a FunctionRetryInputsItem.
+            # This will also update the context state to RETRYING.
+            inputs: list[api_pb2.FunctionRetryInputsItem] = await map_items_manager.prepare_items_for_retry(
+                retriable_idxs
+            )
+            request = api_pb2.FunctionRetryInputsRequest(
+                function_call_jwt=function_call_jwt,
+                inputs=inputs,
+            )
+            resp = await send_inputs(client.stub.FunctionRetryInputs, request)
+            # Update the state to WAITING_FOR_OUTPUT, and update the input_jwt in the context
+            # to the new value in the response.
+            map_items_manager.handle_retry_response(resp.input_jwts)
+            logger.debug(f"Successfully pushed retry for {len(inputs)} to server.")
+            inputs_retried += len(inputs)
+        yield
+
+    async def send_inputs(
+        fn: "modal.client.UnaryUnaryWrapper",
+        request: typing.Union[api_pb2.FunctionPutInputsRequest, api_pb2.FunctionRetryInputsRequest],
+    ) -> typing.Union[api_pb2.FunctionPutInputsResponse, api_pb2.FunctionRetryInputsResponse]:
+        # 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 {function._function_name} is limited."
+            " Common bottlenecks include slow iteration over results, or function backlogs.",
+            warning_interval=8,
+            errors_to_warn_for=[Status.RESOURCE_EXHAUSTED],
+        )
+        return await retry_transient_errors(
+            fn,
+            request,
+            max_retries=None,
+            max_delay=PUMP_INPUTS_MAX_RETRY_DELAY,
+            additional_status_codes=[Status.RESOURCE_EXHAUSTED],
+            retry_warning_message=retry_warning_message,
+        )
+
     async def get_all_outputs():
         assert client.stub
-        nonlocal num_inputs, num_outputs, have_all_inputs
+        nonlocal \
+            inputs_created, \
+            successful_completions, \
+            failed_completions, \
+            outputs_completed, \
+            have_all_inputs, \
+            outputs_received, \
+            already_complete_duplicates, \
+            no_context_duplicates, \
+            stale_retry_duplicates, \
+            retried_outputs
+
         last_entry_id = "0-0"
-        while not have_all_inputs or len(pending_outputs) > len(completed_outputs):
+
+        while not have_all_inputs or outputs_completed < inputs_created:
+            logger.debug(f"Requesting outputs. Have {outputs_completed} outputs, {inputs_created} inputs.")
+            # Get input_jwts of all items in the WAITING_FOR_OUTPUT state.
+            # The server uses these to track for lost inputs.
+            input_jwts = [input_jwt for input_jwt in map_items_manager.get_input_jwts_waiting_for_output()]
+
             request = api_pb2.FunctionGetOutputsRequest(
                 function_call_id=function_call_id,
                 timeout=OUTPUTS_TIMEOUT,
                 last_entry_id=last_entry_id,
                 clear_on_success=False,
                 requested_at=time.time(),
+                input_jwts=input_jwts,
             )
             response = await retry_transient_errors(
                 client.stub.FunctionGetOutputs,
@@ -183,19 +262,31 @@ async def _map_invocation(
                 attempt_timeout=OUTPUTS_TIMEOUT + ATTEMPT_TIMEOUT_GRACE_PERIOD,
             )
 
-            if len(response.outputs) == 0:
-                continue
-
             last_entry_id = response.last_entry_id
+            now_seconds = int(time.time())
             for item in response.outputs:
-                pending_outputs.setdefault(item.input_id, 0)
-                if item.input_id in completed_outputs:
-                    # If this input is already completed, it means the output has already been
-                    # processed and was received again due to a duplicate.
-                    continue
-                completed_outputs.add(item.input_id)
-                num_outputs += 1
-                yield item
+                outputs_received += 1
+                # If the output failed, and there are retries remaining, the input will be placed on the
+                # retry queue, and state updated to WAITING_FOR_RETRY. Otherwise, the output is considered
+                # complete and the item is removed from the manager.
+                output_type = await map_items_manager.handle_get_outputs_response(item, now_seconds)
+                if output_type == _OutputType.SUCCESSFUL_COMPLETION:
+                    successful_completions += 1
+                elif output_type == _OutputType.FAILED_COMPLETION:
+                    failed_completions += 1
+                elif output_type == _OutputType.NO_CONTEXT_DUPLICATE:
+                    no_context_duplicates += 1
+                elif output_type == _OutputType.STALE_RETRY_DUPLICATE:
+                    stale_retry_duplicates += 1
+                elif output_type == _OutputType.ALREADY_COMPLETE_DUPLICATE:
+                    already_complete_duplicates += 1
+                elif output_type == _OutputType.RETRYING:
+                    retried_outputs += 1
+
+                if output_type == _OutputType.SUCCESSFUL_COMPLETION or output_type == _OutputType.FAILED_COMPLETION:
+                    completed_outputs.add(item.input_id)
+                    outputs_completed += 1
+                    yield item
 
     async def get_all_outputs_and_clean_up():
         assert client.stub
@@ -213,6 +304,7 @@ async def _map_invocation(
                 requested_at=time.time(),
             )
             await retry_transient_errors(client.stub.FunctionGetOutputs, request)
+            await retry_queue.close()
 
     async def fetch_output(item: api_pb2.FunctionGetOutputsItem) -> tuple[int, Any]:
         try:
@@ -239,17 +331,50 @@ async def _map_invocation(
                 else:
                     # hold on to outputs for function maps, so we can reorder them correctly.
                     received_outputs[idx] = output
-                    while output_idx in received_outputs:
+
+                    while True:
+                        if output_idx not in received_outputs:
+                            # we haven't received the output for the current index yet.
+                            # stop returning outputs to the caller and instead wait for
+                            # the next output to arrive from the server.
+                            break
+
                         output = received_outputs.pop(output_idx)
                         yield _OutputValue(output)
                         output_idx += 1
 
         assert len(received_outputs) == 0
 
-    async with aclosing(async_merge(drain_input_generator(), pump_inputs(), poll_outputs())) as streamer:
+    async def log_debug_stats():
+        def log_stats():
+            logger.debug(
+                f"Map stats: sync_client_retries_enabled={sync_client_retries_enabled} "
+                f"have_all_inputs={have_all_inputs} inputs_created={inputs_created} input_sent={inputs_sent} "
+                f"inputs_retried={inputs_retried} outputs_received={outputs_received} "
+                f"successful_completions={successful_completions} failed_completions={failed_completions} "
+                f"no_context_duplicates={no_context_duplicates} old_retry_duplicates={stale_retry_duplicates} "
+                f"already_complete_duplicates={already_complete_duplicates} "
+                f"retried_outputs={retried_outputs} input_queue_size={input_queue.qsize()} "
+                f"retry_queue_size={retry_queue.qsize()} map_items_manager={len(map_items_manager)}"
+            )
+        while True:
+            log_stats()
+            try:
+                await asyncio.sleep(10)
+            except asyncio.CancelledError:
+                # Log final stats before exiting
+                log_stats()
+                break
+
+    log_debug_stats_task = asyncio.create_task(log_debug_stats())
+    async with aclosing(
+        async_merge(drain_input_generator(), pump_inputs(), poll_outputs(), retry_inputs())
+    ) as streamer:
         async for response in streamer:
             if response is not None:
                 yield response.value
+    log_debug_stats_task.cancel()
+    await log_debug_stats_task
 
 
 @warn_if_generator_is_not_consumed(function_name="Function.map")
@@ -431,3 +556,227 @@ def _starmap_sync(
             "Use Function.map.aio()/Function.for_each.aio() instead."
         ),
     )
+
+
+class _MapItemState(enum.Enum):
+    # The input is being sent the server with a PutInputs request, but the response has not been received yet.
+    SENDING = 1
+    # A call to either PutInputs or FunctionRetry has completed, and we are waiting to receive the output.
+    WAITING_FOR_OUTPUT = 2
+    # The input is on the retry queue, and waiting for its delay to expire.
+    WAITING_TO_RETRY = 3
+    # The input is being sent to the server with a FunctionRetry request, but the response has not been received yet.
+    RETRYING = 4
+    # The output has been received and was either successful, or failed with no more retries remaining.
+    COMPLETE = 5
+
+class _OutputType(enum.Enum):
+    SUCCESSFUL_COMPLETION = 1
+    FAILED_COMPLETION = 2
+    RETRYING = 3
+    ALREADY_COMPLETE_DUPLICATE = 4
+    STALE_RETRY_DUPLICATE = 5
+    NO_CONTEXT_DUPLICATE = 6
+
+class _MapItemContext:
+    state: _MapItemState
+    input: api_pb2.FunctionInput
+    retry_manager: RetryManager
+    sync_client_retries_enabled:bool
+    # Both these futures are strings. Omitting generic type because
+    # it causes an error when running `inv protoc type-stubs`.
+    input_id: asyncio.Future
+    input_jwt: asyncio.Future
+    previous_input_jwt: Optional[str]
+    _event_loop: asyncio.AbstractEventLoop
+
+    def __init__(self, input: api_pb2.FunctionInput, retry_manager: RetryManager, sync_client_retries_enabled: bool):
+        self.state = _MapItemState.SENDING
+        self.input = input
+        self.retry_manager = retry_manager
+        self.sync_client_retries_enabled = sync_client_retries_enabled
+        self._event_loop = asyncio.get_event_loop()
+        # create a future for each input, to be resolved when we have
+        # received the input ID and JWT from the server. this addresses
+        # a race condition where we could receive outputs before we have
+        # recorded the input ID and JWT in `pending_outputs`.
+        self.input_jwt = self._event_loop.create_future()
+        self.input_id = self._event_loop.create_future()
+
+    def handle_put_inputs_response(self, item: api_pb2.FunctionPutInputsResponseItem):
+        self.input_jwt.set_result(item.input_jwt)
+        self.input_id.set_result(item.input_id)
+        # Set state to WAITING_FOR_OUTPUT only if current state is SENDING. If state is
+        # RETRYING, WAITING_TO_RETRY, or COMPLETE, then we already got the output.
+        if self.state == _MapItemState.SENDING:
+            self.state = _MapItemState.WAITING_FOR_OUTPUT
+
+    async def handle_get_outputs_response(
+        self,
+        item: api_pb2.FunctionGetOutputsItem,
+        now_seconds: int,
+        function_call_invocation_type: "api_pb2.FunctionCallInvocationType.ValueType",
+        retry_queue: TimestampPriorityQueue,
+    ) -> _OutputType:
+        """
+        Processes the output, and determines if it is complete or needs to be retried.
+
+        Return True if input state was changed to COMPLETE, otherwise False.
+        """
+        # If the item is already complete, this is a duplicate output and can be ignored.
+        if self.state == _MapItemState.COMPLETE:
+            logger.debug(
+                f"Received output for input marked as complete. Must be duplicate, so ignoring. "
+                f"idx={item.idx} input_id={item.input_id}, retry_count={item.retry_count}"
+            )
+            return _OutputType.ALREADY_COMPLETE_DUPLICATE
+        # If the item's retry count doesn't match our retry count, this is probably a duplicate of an old output.
+        if item.retry_count != self.retry_manager.retry_count:
+            logger.debug(
+                f"Received output with stale retry_count, so ignoring. "
+                f"idx={item.idx} input_id={item.input_id} retry_count={item.retry_count} "
+                f"expected_retry_count={self.retry_manager.retry_count}"
+            )
+            return _OutputType.STALE_RETRY_DUPLICATE
+
+        # retry failed inputs when the function call invocation type is SYNC
+        if (
+            item.result.status == api_pb2.GenericResult.GENERIC_STATUS_SUCCESS
+            or function_call_invocation_type != api_pb2.FUNCTION_CALL_INVOCATION_TYPE_SYNC
+            or not self.sync_client_retries_enabled
+        ):
+            self.state = _MapItemState.COMPLETE
+            if item.result.status == api_pb2.GenericResult.GENERIC_STATUS_SUCCESS:
+                return _OutputType.SUCCESSFUL_COMPLETION
+            else:
+                return _OutputType.FAILED_COMPLETION
+
+        # Get the retry delay and increment the retry count.
+        # TODO(ryan): We must call this for lost inputs - even though we will set the retry delay to 0 later -
+        # because we must increment the retry count. That's awkward, let's come up with something better.
+        # TODO(ryan):To maintain parity with server-side retries, retrying lost inputs should not count towards
+        # the retry policy. However we use the retry_count number as a unique identifier on each attempt to:
+        #  1) ignore duplicate outputs
+        #  2) ignore late outputs received from previous attempts
+        #  3) avoid a server race condition between FunctionRetry and GetOutputs that results in deleted input metadata
+        # For now, lost inputs will count towards the retry policy. But let's address this in another PR, perhaps by
+        # tracking total attempts and attempts which count towards the retry policy separately.
+        delay_ms = self.retry_manager.get_delay_ms()
+
+        # For system failures on the server, we retry immediately.
+        # and the failure does not count towards the retry policy.
+        if item.result.status == api_pb2.GenericResult.GENERIC_STATUS_INTERNAL_FAILURE:
+            delay_ms = 0
+
+        # None means the maximum number of retries has been reached, so output the error
+        if delay_ms is None:
+            self.state = _MapItemState.COMPLETE
+            return _OutputType.FAILED_COMPLETION
+
+        self.state = _MapItemState.WAITING_TO_RETRY
+        await retry_queue.put(now_seconds + (delay_ms / 1000), item.idx)
+
+        return _OutputType.RETRYING
+
+    async def prepare_item_for_retry(self) -> api_pb2.FunctionRetryInputsItem:
+        self.state = _MapItemState.RETRYING
+        # If the input_jwt is not set, then put_inputs hasn't returned yet. Block until we have it.
+        input_jwt = await self.input_jwt
+        self.input_jwt = self._event_loop.create_future()
+        return api_pb2.FunctionRetryInputsItem(
+            input_jwt=input_jwt,
+            input=self.input,
+            retry_count=self.retry_manager.retry_count,
+        )
+
+    def handle_retry_response(self, input_jwt: str):
+        self.input_jwt.set_result(input_jwt)
+        self.state = _MapItemState.WAITING_FOR_OUTPUT
+
+
+class _MapItemsManager:
+    def __init__(
+        self,
+        retry_policy: api_pb2.FunctionRetryPolicy,
+        function_call_invocation_type: "api_pb2.FunctionCallInvocationType.ValueType",
+        retry_queue: TimestampPriorityQueue,
+        sync_client_retries_enabled: bool
+    ):
+        self._retry_policy = retry_policy
+        self.function_call_invocation_type = function_call_invocation_type
+        self._retry_queue = retry_queue
+        # semaphore to limit the number of inputs that can be in progress at once
+        self._inputs_outstanding = asyncio.BoundedSemaphore(MAP_MAX_INPUTS_OUTSTANDING)
+        self._item_context: dict[int, _MapItemContext] = {}
+        self._sync_client_retries_enabled = sync_client_retries_enabled
+
+    async def add_items(self, items: list[api_pb2.FunctionPutInputsItem]):
+        for item in items:
+            # acquire semaphore to limit the number of inputs in progress
+            # (either queued to be sent, waiting for completion, or retrying)
+            await self._inputs_outstanding.acquire()
+            self._item_context[item.idx] = _MapItemContext(
+                input=item.input,
+                retry_manager=RetryManager(self._retry_policy),
+                sync_client_retries_enabled=self._sync_client_retries_enabled,
+            )
+
+    async def prepare_items_for_retry(self, retriable_idxs: list[int]) -> list[api_pb2.FunctionRetryInputsItem]:
+        return [await self._item_context[idx].prepare_item_for_retry() for idx in retriable_idxs]
+
+    def get_input_jwts_waiting_for_output(self) -> list[str]:
+        """
+        Returns a list of input_jwts for inputs that are waiting for output.
+        """
+        # If input_jwt is not done, the call to PutInputs has not completed, so omit it from results.
+        return [
+            ctx.input_jwt.result()
+            for ctx in self._item_context.values()
+            if ctx.state == _MapItemState.WAITING_FOR_OUTPUT and ctx.input_jwt.done()
+        ]
+
+    def _remove_item(self, item_idx: int):
+        del self._item_context[item_idx]
+        self._inputs_outstanding.release()
+
+    def get_item_context(self, item_idx: int) -> _MapItemContext:
+        return self._item_context.get(item_idx)
+
+    def handle_put_inputs_response(self, items: list[api_pb2.FunctionPutInputsResponseItem]):
+        for item in items:
+            ctx = self._item_context.get(item.idx, None)
+            # If the context is None, then get_all_outputs() has already received a successful
+            # output, and deleted the context. This happens if FunctionGetOutputs completes
+            # before FunctionPutInputsResponse is received.
+            if ctx is not None:
+                ctx.handle_put_inputs_response(item)
+
+    def handle_retry_response(self, input_jwts: list[str]):
+        for input_jwt in input_jwts:
+            decoded_jwt = DecodedJwt.decode_without_verification(input_jwt)
+            ctx = self._item_context.get(decoded_jwt.payload["idx"], None)
+            # If the context is None, then get_all_outputs() has already received a successful
+            # output, and deleted the context. This happens if FunctionGetOutputs completes
+            # before FunctionRetryInputsResponse is received.
+            if ctx is not None:
+                ctx.handle_retry_response(input_jwt)
+
+    async def handle_get_outputs_response(self, item: api_pb2.FunctionGetOutputsItem, now_seconds: int) -> _OutputType:
+        ctx = self._item_context.get(item.idx, None)
+        if ctx is None:
+            # We've already processed this output, so we can skip it.
+            # This can happen because the worker can sometimes send duplicate outputs.
+            logger.debug(
+                f"Received output that does not have entry in item_context map, so ignoring. "
+                f"idx={item.idx} input_id={item.input_id} retry_count={item.retry_count} "
+            )
+            return _OutputType.NO_CONTEXT_DUPLICATE
+        output_type = await ctx.handle_get_outputs_response(
+            item, now_seconds, self.function_call_invocation_type, self._retry_queue
+        )
+        if output_type == _OutputType.SUCCESSFUL_COMPLETION or output_type == _OutputType.FAILED_COMPLETION:
+            self._remove_item(item.idx)
+        return output_type
+
+    def __len__(self):
+        return len(self._item_context)