modAL 1.5.2.dev22__py3-none-any.whl → 1.5.3.dev0__py3-none-any.whl

modal/_utils/grpc_utils.py

@@ -164,6 +164,9 @@ RETRYABLE_GRPC_STATUS_CODES = [
 ]
 SERVER_RETRY_WARNING_TIME_INTERVAL = 30.0
 DEFAULT_MAX_RETRIES = 3
+GRPC_KEEPALIVE_TIME_SECS = 30.0
+GRPC_KEEPALIVE_TIMEOUT_SECS = 10.0
+GRPC_WINDOW_SIZE = 64 * 1024 * 1024  # 64 MiB
 
 
 @dataclass
@@ -251,6 +254,21 @@ class CustomProtoStatusDetailsCodec(StatusDetailsCodecBase):
 custom_detail_codec = CustomProtoStatusDetailsCodec()
 
 
+def create_channel_config() -> grpclib.config.Configuration:
+    """Shared grpclib channel settings for Modal client connections.
+
+    HTTP/2 keepalive probes keep idle transports warm through stateful
+    middleboxes and surface dead connections before retries reuse them.
+    """
+    return grpclib.config.Configuration(
+        _keepalive_time=GRPC_KEEPALIVE_TIME_SECS,
+        _keepalive_timeout=GRPC_KEEPALIVE_TIMEOUT_SECS,
+        _keepalive_permit_without_calls=True,
+        http2_connection_window_size=GRPC_WINDOW_SIZE,
+        http2_stream_window_size=GRPC_WINDOW_SIZE,
+    )
+
+
 def create_channel(
     server_url: str,
     metadata: dict[str, str] = {},
@@ -262,10 +280,7 @@ def create_channel(
     o = urllib.parse.urlparse(server_url)
 
     channel: grpclib.client.Channel
-    config = grpclib.config.Configuration(
-        http2_connection_window_size=64 * 1024 * 1024,  # 64 MiB
-        http2_stream_window_size=64 * 1024 * 1024,  # 64 MiB
-    )
+    config = create_channel_config()
 
     if o.scheme == "unix":
         channel = grpclib.client.Channel(path=o.path, config=config, status_details_codec=custom_detail_codec)

modal/_utils/task_command_router_client.py

@@ -1,7 +1,6 @@
 # Copyright Modal Labs 2025
 import asyncio
 import base64
-import io
 import json
 import socket
 import ssl
@@ -9,36 +8,24 @@ import time
 import typing
 import urllib.parse
 import weakref
-from collections.abc import AsyncGenerator, AsyncIterable, Callable
+from collections.abc import AsyncGenerator, Callable
 from contextlib import suppress
-from typing import BinaryIO, TypeVar
+from typing import TypeVar
 
 import grpclib.client
-import grpclib.config
 from grpclib import GRPCError, Status
 from grpclib.exceptions import StreamTerminatedError
 
 from modal.config import logger
-from modal.exception import AuthError, ExecTimeoutError, InternalError, ServiceError, TimeoutError as ModalTimeoutError
+from modal.exception import ExecTimeoutError, TimeoutError as ModalTimeoutError
 from modal_proto import api_pb2, task_command_router_pb2 as sr_pb2
 from modal_proto.task_command_router_grpc import TaskCommandRouterStub
 
 from .._grpc_client import grpc_error_converter
-from .._utils.grpc_utils import ModalChannel
+from .._utils.grpc_utils import ModalChannel, create_channel_config
 from .async_utils import aclosing, retry
 from .grpc_utils import RETRYABLE_GRPC_STATUS_CODES
 
-STREAMING_STDIN_CHUNK_SIZE = 256 * 1024
-
-_STREAMING_STDIN_RESUMABLE_EXCEPTIONS = (
-    AuthError,
-    InternalError,
-    ServiceError,
-    StreamTerminatedError,
-    OSError,
-    asyncio.TimeoutError,
-)
-
 
 @retry(n_attempts=34, base_delay=1, max_delay=10, attempt_timeout=10, total_timeout=310)
 async def _connect_channel(channel: grpclib.client.Channel):
@@ -234,10 +221,7 @@ class TaskCommandRouterClient:
             host,
             port,
             ssl=ssl_context,
-            config=grpclib.config.Configuration(
-                http2_connection_window_size=64 * 1024 * 1024,  # 64 MiB
-                http2_stream_window_size=64 * 1024 * 1024,  # 64 MiB
-            ),
+            config=create_channel_config(),
             closed_error_message="Unable to perform operation on a detached sandbox",
         )
 
@@ -484,118 +468,6 @@ class TaskCommandRouterClient:
                 lambda: self._call_with_auth_retry(self._stub.TaskExecStdinWrite, request)
             )
 
-    async def exec_stdin_write_stream(
-        self,
-        task_id: str,
-        exec_id: str,
-        source: BinaryIO,
-        chunk_size: int = STREAMING_STDIN_CHUNK_SIZE,
-        max_resume_attempts: int = 9,
-    ) -> int:
-        """Stream `source` into the exec's stdin, with bounded resume on transient failures.
-
-        Streams the full contents of `source` in one client-streaming RPC and
-        closes stdin (EOF) on success. On a resumable error, queries
-        `exec_stdin_status` for the server's canonical offset, seeks `source`
-        to that point, and reopens the stream.
-
-        Args:
-            task_id: The task ID of the task running the exec'd command.
-            exec_id: The execution ID of the command to write to.
-            source: A seekable byte source (file, BytesIO, etc.).
-            chunk_size: Bytes per outbound message.
-            max_resume_attempts: Upper bound on resume retries before giving up.
-
-        Returns:
-            The total number of bytes streamed into the exec's stdin.
-        """
-        offset = 0
-        attempt = 0
-        while True:
-
-            async def _chunks():
-                while True:
-                    chunk = source.read(chunk_size)
-                    if not chunk:
-                        return
-                    yield chunk
-
-            try:
-                source.seek(offset)
-                await self._exec_stdin_write_stream(task_id, exec_id, offset, _chunks())
-                return source.tell()
-            except (*_STREAMING_STDIN_RESUMABLE_EXCEPTIONS, AttributeError) as e:
-                # StreamTerminatedError surfaces as AttributeError in grpclib<=0.4.7;
-                # see call_with_retries_on_transient_errors.
-                if isinstance(e, AttributeError) and "_write_appdata" not in str(e):
-                    raise
-                attempt += 1
-                if attempt > max_resume_attempts:
-                    raise
-                status = await self.exec_stdin_status(task_id=task_id, exec_id=exec_id)
-                if status.closed:
-                    pos = source.tell()
-                    if status.num_bytes_written == pos and source.seek(0, io.SEEK_END) == pos:
-                        logger.debug(f"exec_stdin_write_stream completed but response was lost: {e}")
-                        return pos
-                    raise
-                offset = status.num_bytes_written
-                logger.debug(f"exec_stdin_write_stream resuming from offset {offset} after error: {e}")
-
-    async def _exec_stdin_write_stream(
-        self,
-        task_id: str,
-        exec_id: str,
-        start_offset: int,
-        chunks: AsyncIterable[bytes],
-    ) -> sr_pb2.TaskExecStdinWriteStreamResponse | None:
-        """Single client-streaming attempt: Start, Data chunks, then End (EOF).
-
-        Does not retry; `exec_stdin_write_stream` owns resume.
-        """
-        with grpc_error_converter():
-            stream = self._stub.TaskExecStdinWriteStream.open(metadata=self._get_metadata())
-            async with stream as s:
-                start = sr_pb2.TaskExecStdinWriteStreamRequest(
-                    start=sr_pb2.TaskExecStdinWriteStreamStart(
-                        task_id=task_id,
-                        exec_id=exec_id,
-                        offset=start_offset,
-                    ),
-                )
-                await s.send_message(start)
-                async for data in chunks:
-                    if not data:
-                        continue
-                    await s.send_message(sr_pb2.TaskExecStdinWriteStreamRequest(data=data))
-                # The server closes stdin only on this explicit End message. A
-                # stream that breaks before it leaves stdin open for resume.
-                await s.send_message(
-                    sr_pb2.TaskExecStdinWriteStreamRequest(end=sr_pb2.TaskExecStdinWriteStreamEnd()),
-                    end=True,
-                )
-                return await s.recv_message()
-
-    async def exec_stdin_status(self, task_id: str, exec_id: str) -> sr_pb2.TaskExecStdinStatusResponse:
-        """Read the current stdin write status for an exec'd command, to support retries from the right offset.
-
-        Used by streaming clients to find the resume offset after a stream
-        failure. Evicts any in-flight stdin stream for the exec, so it is not
-        safe to call for read-only observability.
-
-        Args:
-            task_id: The task ID of the task running the exec'd command.
-            exec_id: The execution ID of the command to query.
-        Raises:
-            Other errors: If retries are exhausted on transient errors or if there's an error
-              from the RPC itself.
-        """
-        request = sr_pb2.TaskExecStdinStatusRequest(task_id=task_id, exec_id=exec_id)
-        with grpc_error_converter():
-            return await call_with_retries_on_transient_errors(
-                lambda: self._call_with_auth_retry(self._stub.TaskExecStdinStatus, request)
-            )
-
     async def sandbox_stdin_write_v2(
         self, task_id: str, offset: int, data: bytes, eof: bool
     ) -> sr_pb2.SandboxStdinWriteV2Response:
@@ -625,16 +497,7 @@ class TaskCommandRouterClient:
             TimeoutError: If the sandbox does not become ready within `timeout`.
         """
         request = sr_pb2.SandboxWaitUntilReadyTcrRequest(task_id=task_id, timeout=timeout)
-        with grpc_error_converter():
-            try:
-                return await asyncio.wait_for(
-                    call_with_retries_on_transient_errors(
-                        lambda: self._call_with_auth_retry(self._stub.SandboxWaitUntilReady, request, timeout=timeout),
-                    ),
-                    timeout=timeout,
-                )
-            except asyncio.TimeoutError:
-                raise ModalTimeoutError("Timeout expired")
+        return await self._unary_call_with_deadline(self._stub.SandboxWaitUntilReady, request, timeout=timeout)
 
     async def exec_poll(self, task_id: str, exec_id: str, deadline: float | None = None) -> sr_pb2.TaskExecPollResponse:
         """Poll for the exit status of an exec'd command, properly retrying on transient errors.
@@ -944,21 +807,14 @@ class TaskCommandRouterClient:
                 window, the call is cancelled and a `modal.exception.TimeoutError` is raised.
         """
         request = sr_pb2.TaskReloadVolumesRequest(task_id=task_id)
-        with grpc_error_converter():
-            try:
-                return await asyncio.wait_for(
-                    call_with_retries_on_transient_errors(
-                        lambda: self._call_with_auth_retry(self._stub.TaskReloadVolumes, request, timeout=timeout),
-                    ),
-                    timeout=timeout,
-                )
-            except asyncio.TimeoutError:
-                raise ModalTimeoutError("Timeout expired")
+        return await self._unary_call_with_deadline(self._stub.TaskReloadVolumes, request, timeout=timeout)
+
+    async def _unary_call_with_deadline(self, rpc, request, *, timeout: float, **kwargs):
+        """Call a unary RPC with transient-error retries bounded by a total deadline.
 
-    async def _snapshot_with_deadline(self, rpc, request, *, timeout: float, **kwargs):
-        # helper method for snapshot_directory and snapshot_filesystem to handle grpc
-        # deadlines in a consistent way, converting any error to TimeoutError after passing
-        # the total deadline budget
+        Raises `modal.exception.TimeoutError` once the deadline has elapsed, or when
+        the server responds with `DEADLINE_EXCEEDED`.
+        """
         timeout_deadline = time.monotonic() + timeout
 
         def call():
@@ -976,17 +832,21 @@ class TaskCommandRouterClient:
                     exclude_status_codes=[Status.DEADLINE_EXCEEDED, Status.CANCELLED],
                     timeout_deadline=timeout_deadline,
                 )
-        except Exception:
-            if time.monotonic() >= timeout_deadline:
+        except Exception as exc:
+            if time.monotonic() >= timeout_deadline or getattr(exc, "_grpc_status", None) == Status.DEADLINE_EXCEEDED:
                 raise ModalTimeoutError("Timeout expired")
             raise
 
     async def snapshot_directory(
         self, request: sr_pb2.TaskSnapshotDirectoryRequest, *, timeout: float, **kwargs
     ) -> sr_pb2.TaskSnapshotDirectoryResponse:
-        return await self._snapshot_with_deadline(self._stub.TaskSnapshotDirectory, request, timeout=timeout, **kwargs)
+        return await self._unary_call_with_deadline(
+            self._stub.TaskSnapshotDirectory, request, timeout=timeout, **kwargs
+        )
 
     async def snapshot_filesystem(
         self, request: sr_pb2.TaskSnapshotFilesystemRequest, *, timeout: float, **kwargs
     ) -> sr_pb2.TaskSnapshotFilesystemResponse:
-        return await self._snapshot_with_deadline(self._stub.TaskSnapshotFilesystem, request, timeout=timeout, **kwargs)
+        return await self._unary_call_with_deadline(
+            self._stub.TaskSnapshotFilesystem, request, timeout=timeout, **kwargs
+        )

modal/cli/changelog.py

@@ -120,17 +120,17 @@ def changelog(
 
     Examples:
 
-        ```bash
-        modal changelog --since 1.2.0  # Show updates added after a specific version
+    ```bash
+    modal changelog --since 1.2.0  # Show updates added after a specific version
 
-        modal changelog --since 2026-01-01  # Show updates added after a specific date
+    modal changelog --since 2026-01-01  # Show updates added after a specific date
 
-        modal changelog --newer  # Show updates released after the currently installed version
+    modal changelog --newer  # Show updates released after the currently installed version
 
-        modal changelog --last 3  # Show updates included in the 3 most recent releases
+    modal changelog --last 3  # Show updates included in the 3 most recent releases
 
-        modal changelog --for 1.3.1  # Show the changelog for a specific release
-        ```
+    modal changelog --for 1.3.1  # Show the changelog for a specific release
+    ```
 
     Note: when using `--since` or `--last`, only changes up to the currently installed version are shown.
 

modal/cli/environment.py

@@ -242,21 +242,21 @@ async def environment_billing(
 
     Examples:
 
-        ```bash
-        modal environment billing report --start 2025-12-01 --end 2026-01-01
+    ```bash
+    modal environment billing report --start 2025-12-01 --end 2026-01-01
 
-        modal environment billing report --for "last month" --tag-names team,project
+    modal environment billing report --for "last month" --tag-names team,project
 
-        modal environment billing report test_env --for today --resolution h
+    modal environment billing report test_env --for today --resolution h
 
-        modal environment billing report test_env --for "this month" --show-resources
+    modal environment billing report test_env --for "this month" --show-resources
 
-        modal environment billing report prod_env --for yesterday -r h --tz local
+    modal environment billing report prod_env --for yesterday -r h --tz local
 
-        modal environment billing report main_env --for "last month" --csv > report.csv
+    modal environment billing report main_env --for "last month" --csv > report.csv
 
-        modal environment billing report main_env --start 2025-12-01 --json > report.json
-        ```
+    modal environment billing report main_env --start 2025-12-01 --json > report.json
+    ```
 
     """
     if json and csv:

modal/client.pyi

@@ -35,7 +35,7 @@ class _Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.5.2.dev22",
+        version: str = "1.5.3.dev0",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.
@@ -205,7 +205,7 @@ class Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.5.2.dev22",
+        version: str = "1.5.3.dev0",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.

modal/container_process.py

@@ -1,7 +1,7 @@
 # Copyright Modal Labs 2024
 import asyncio
 import platform
-from typing import BinaryIO, Generic, TypeVar
+from typing import Generic, TypeVar
 
 from modal_proto import api_pb2
 
@@ -117,18 +117,6 @@ class _ContainerProcess(Generic[T]):
             )
         return self._returncode
 
-    async def _stdin_write_stream(self, source: BinaryIO) -> int:
-        """mdmd:hidden
-        Stream `source` into the process's stdin and close it (EOF) on success.
-
-        Returns the total bytes streamed.
-        """
-        return await self._command_router_client.exec_stdin_write_stream(
-            task_id=self._task_id,
-            exec_id=self._process_id,
-            source=source,
-        )
-
     async def poll(self) -> int | None:
         """Check if the container process has finished running.
 

modal/container_process.pyi

@@ -53,14 +53,6 @@ class _ContainerProcess(typing.Generic[T]):
 
     @property
     def returncode(self) -> int: ...
-    async def _stdin_write_stream(self, source: typing.BinaryIO) -> int:
-        """mdmd:hidden
-        Stream `source` into the process's stdin and close it (EOF) on success.
-
-        Returns the total bytes streamed.
-        """
-        ...
-
     async def poll(self) -> typing.Optional[int]:
         """Check if the container process has finished running.
 
@@ -113,25 +105,6 @@ class ContainerProcess(typing.Generic[T]):
     @property
     def returncode(self) -> int: ...
 
-    class ___stdin_write_stream_spec(typing_extensions.Protocol):
-        def __call__(self, /, source: typing.BinaryIO) -> int:
-            """mdmd:hidden
-            Stream `source` into the process's stdin and close it (EOF) on success.
-
-            Returns the total bytes streamed.
-            """
-            ...
-
-        async def aio(self, /, source: typing.BinaryIO) -> int:
-            """mdmd:hidden
-            Stream `source` into the process's stdin and close it (EOF) on success.
-
-            Returns the total bytes streamed.
-            """
-            ...
-
-    _stdin_write_stream: ___stdin_write_stream_spec
-
     class __poll_spec(typing_extensions.Protocol):
         def __call__(self, /) -> typing.Optional[int]:
             """Check if the container process has finished running.

modal/functions.pyi

@@ -491,7 +491,7 @@ class Function(
 
     _call_generator: ___call_generator_spec
 
-    class __remote_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER]):
+    class __remote_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER]):
         def __call__(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> ReturnType_INNER:
             """Calls the function remotely, executing it with the given arguments and returning the execution's result.
 
@@ -516,7 +516,7 @@ class Function(
             """
             ...
 
-    remote: __remote_spec[modal._functions.ReturnType, modal._functions.P]
+    remote: __remote_spec[modal._functions.P, modal._functions.ReturnType]
 
     class __remote_gen_spec(typing_extensions.Protocol):
         def __call__(self, /, *args, **kwargs) -> typing.Generator[typing.Any, None, None]:
@@ -566,7 +566,7 @@ class Function(
         """
         ...
 
-    class ___experimental_spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER]):
+    class ___experimental_spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER]):
         def __call__(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]:
             """[Experimental] Calls the function with the given arguments, without waiting for the results.
 
@@ -599,7 +599,7 @@ class Function(
             """
             ...
 
-    _experimental_spawn: ___experimental_spawn_spec[modal._functions.ReturnType, modal._functions.P]
+    _experimental_spawn: ___experimental_spawn_spec[modal._functions.P, modal._functions.ReturnType]
 
     class ___spawn_map_inner_spec(typing_extensions.Protocol[P_INNER]):
         def __call__(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> None: ...
@@ -607,7 +607,7 @@ class Function(
 
     _spawn_map_inner: ___spawn_map_inner_spec[modal._functions.P]
 
-    class __spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER]):
+    class __spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER]):
         def __call__(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]:
             """Calls the function with the given arguments, without waiting for the results.
 
@@ -640,7 +640,7 @@ class Function(
             """
             ...
 
-    spawn: __spawn_spec[modal._functions.ReturnType, modal._functions.P]
+    spawn: __spawn_spec[modal._functions.P, modal._functions.ReturnType]
 
     def get_raw_f(self) -> collections.abc.Callable[..., typing.Any]:
         """Return the inner Python object wrapped by this Modal Function.

modal/parallel_map.py

@@ -654,16 +654,14 @@ async def _map_invocation_inputplane(
     stale_retry_duplicates = 0
     already_complete_duplicates = 0
     retried_outputs = 0
-    input_queue_size = 0
     last_entry_id = ""
 
     # The input-plane server returns this after the first request.
     map_token = None
     map_token_received = asyncio.Event()
 
-    # Single priority queue that holds *both* fresh inputs (timestamp == now)
-    # and future retries (timestamp > now).
-    queue: TimestampPriorityQueue[api_pb2.MapStartOrContinueItem] = TimestampPriorityQueue()
+    input_queue: asyncio.Queue[api_pb2.MapStartOrContinueItem | None] = asyncio.Queue()
+    retry_queue: TimestampPriorityQueue[api_pb2.MapStartOrContinueItem] = TimestampPriorityQueue()
 
     # Maximum number of inputs that may be in-flight (the server sends this in
     # the first response – fall back to the default if we never receive it for
@@ -682,7 +680,7 @@ async def _map_invocation_inputplane(
     map_items_manager = _MapItemsManager(
         retry_policy=retry_policy,
         function_call_invocation_type=api_pb2.FUNCTION_CALL_INVOCATION_TYPE_SYNC,
-        retry_queue=queue,
+        retry_queue=retry_queue,
         sync_client_retries_enabled=True,
         max_inputs_outstanding=MAX_INPUTS_OUTSTANDING_DEFAULT,
         is_input_plane_instance=True,
@@ -729,61 +727,67 @@ async def _map_invocation_inputplane(
             async_map_ordered(input_iter(), create_input, concurrency=BLOB_MAX_PARALLELISM)
         ) as streamer:
             async for q_item in streamer:
-                await queue.put(time.time(), q_item)
+                await input_queue.put(q_item)
 
         # All inputs have been read.
+        await input_queue.put(None)
         update_counters(set_have_all_inputs=True)
         yield
 
-    async def pump_inputs():
+    async def send_map_start_or_continue(request_items: list[api_pb2.MapStartOrContinueItem]):
         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] = [
-                api_pb2.MapStartOrContinueItem(input=qi.input, attempt_token=qi.attempt_token) for qi in batch
-            ]
+        request = api_pb2.MapStartOrContinueRequest(
+            function_id=function.object_id,
+            map_token=map_token or "",
+            parent_input_id=current_input_id() or "",
+            items=request_items,
+        )
 
-            await map_items_manager.add_items_inputplane(request_items)
+        metadata = await client.get_input_plane_metadata(function._input_plane_region)
 
-            # Build request
-            request = api_pb2.MapStartOrContinueRequest(
-                function_id=function.object_id,
-                map_token=map_token,
-                parent_input_id=current_input_id() or "",
-                items=request_items,
-            )
+        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,
+        )
 
-            metadata = await client.get_input_plane_metadata(function._input_plane_region)
+        # match response items to the corresponding request item index
+        response_items_idx_tuple = [
+            (request_items[idx].input.idx, attempt_token) for idx, attempt_token in enumerate(response.attempt_tokens)
+        ]
 
-            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,
-            )
+        map_items_manager.handle_put_continue_response(response_items_idx_tuple)
 
-            # match response items to the corresponding request item index
-            response_items_idx_tuple = [
-                (request_items[idx].input.idx, attempt_token)
-                for idx, attempt_token in enumerate(response.attempt_tokens)
-            ]
+        # Set the function call id and actual retry policy with the data from the first response.
+        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.
+            # Subsequent batches will have the correct user-specified retry policy
+            # set by the updated _MapItemsManager.
+            map_items_manager.update_items_retry_policy(response.retry_policy)
 
-            map_items_manager.handle_put_continue_response(response_items_idx_tuple)
+    async def pump_inputs():
+        async for batch in queue_batch_iterator(input_queue, max_batch_size=MAP_INVOCATION_CHUNK_SIZE):
+            request_items: list[api_pb2.MapStartOrContinueItem] = list(batch)
+            await map_items_manager.add_items_inputplane(request_items)
+            await send_map_start_or_continue(request_items)
+        yield
 
-            # 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 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.
-                # Subsequent batches will have the correct user-specified retry policy
-                # set by the updated _MapItemsManager.
-                map_items_manager.update_items_retry_policy(response.retry_policy)
+    async def retry_inputs():
+        while True:
+            retry_item = await retry_queue.get()
+            if retry_item is None:
+                break
+            request_items = [retry_item]
+            await map_items_manager.add_items_inputplane(request_items)
+            await send_map_start_or_continue(request_items)
         yield
 
     async def check_lost_inputs():
@@ -897,8 +901,8 @@ async def _map_invocation_inputplane(
                 async for item in stream:
                     yield item
         finally:
-            await queue.close()
-            pass
+            await input_queue.put(None)
+            await retry_queue.close()
 
     async def fetch_output(item: api_pb2.FunctionGetOutputsItem) -> tuple[int, Any]:
         try:
@@ -945,7 +949,8 @@ async def _map_invocation_inputplane(
                 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"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}"
+                f"map_items_manager_size={len(map_items_manager)} input_queue_size={input_queue.qsize()} "
+                f"retry_queue_size={retry_queue.qsize()}"
             )
 
         while True:
@@ -960,7 +965,7 @@ async def _map_invocation_inputplane(
     log_task = asyncio.create_task(log_debug_stats())
 
     async with aclosing(
-        async_merge(drain_input_generator(), pump_inputs(), poll_outputs(), check_lost_inputs())
+        async_merge(drain_input_generator(), pump_inputs(), retry_inputs(), poll_outputs(), check_lost_inputs())
     ) as merged:
         async for maybe_output in merged:
             if maybe_output is not None:  # ignore None sentinels