Drop duplicate input deliveries in the container (#49117)

GitOrigin-RevId: 80b19412ff275d7003fbabb1ab357fb71297bc58

Author: parthivapsani

py/modal/_container_entrypoint.py

@@ -145,7 +145,7 @@ async def run_input_async(io_context: IOContext) -> None:
                             item_count += 1
 
                 await container_io_manager._send_outputs.aio(
-                    io_context.output_items_generator_done(started_at, item_count)
+                    io_context, io_context.output_items_generator_done(started_at, item_count)
                 )
             else:
                 value = await io_context.call_function_async()
@@ -182,7 +182,9 @@ def run_input_sync(io_context: IOContext) -> None:
                         container_io_manager._queue_put(generator_queue, value)
                         item_count += 1
 
-                container_io_manager._send_outputs(io_context.output_items_generator_done(started_at, item_count))
+                container_io_manager._send_outputs(
+                    io_context, io_context.output_items_generator_done(started_at, item_count)
+                )
             else:
                 values = io_context.call_function_sync()
                 container_io_manager.push_outputs(io_context, started_at, values)

py/modal/_runtime/container_io_manager.py

@@ -96,6 +96,7 @@ class IOContext:
 
     _cancel_issued: bool = False
     _cancel_callback: Callable[[], None] | None = None
+    _exited: bool = False
 
     def __init__(
         self,
@@ -889,25 +890,35 @@ async def _generate_inputs(
                         if item.kill_switch:
                             logger.debug(f"Task {self.task_id} input kill signal input.")
                             return
-                        inputs.append(
-                            (
-                                item.input_id,
-                                item.retry_count,
-                                item.function_call_id,
-                                item.attempt_token,
-                                item.input,
-                                item.function_call_invocation_type,
+                        live_io_context = self.current_inputs.get(item.input_id)
+                        if live_io_context is not None and (
+                            live_io_context.retry_counts[live_io_context.input_ids.index(item.input_id)]
+                            == item.retry_count
+                        ):
+                            # An expired fetch lease redelivers a still-running attempt under the same
+                            # input id and retry count; a retry of a failed attempt gets a new retry count.
+                            logger.warning(f"Skipping duplicate delivery of input {item.input_id}")
+                        else:
+                            inputs.append(
+                                (
+                                    item.input_id,
+                                    item.retry_count,
+                                    item.function_call_id,
+                                    item.attempt_token,
+                                    item.input,
+                                    item.function_call_invocation_type,
+                                )
                             )
-                        )
                         if item.input.final_input:
                             if request.batch_max_size > 0:
                                 logger.debug(f"Task {self.task_id} Final input not expected in batch input stream")
                             final_input_received = True
                             break
 
-                    # If yielded, allow input slots to be released via exit_context
-                    yield inputs
-                    yielded = True
+                    if inputs:
+                        # If yielded, allow input slots to be released via exit_context
+                        yield inputs
+                        yielded = True
 
                     # TODO(michael): Remove use of max_inputs after worker rollover
                     single_use_container = self.function_def.single_use_containers or self.function_def.max_inputs == 1
@@ -950,7 +961,7 @@ async def run_inputs_outputs(
             # collect all active input slots, meaning all inputs have wrapped up.
             await self._input_slots.close()
 
-    async def _send_outputs(self, outputs: list[api_pb2.FunctionPutOutputsItem]) -> None:
+    async def _send_outputs(self, io_context: IOContext, outputs: list[api_pb2.FunctionPutOutputsItem]) -> None:
         """Send pre-built output items with retry and chunking."""
         # There are multiple outputs for a single IOContext in the case of @modal.batched.
         # Limit the batch size to 20 to stay within message size limits and buffer size limits.
@@ -963,8 +974,7 @@ async def _send_outputs(self, outputs: list[api_pb2.FunctionPutOutputsItem]) ->
                     max_retries=None,  # Retry indefinitely, trying every 1s.
                 ),
             )
-        input_ids = [output.input_id for output in outputs]
-        self.exit_context(input_ids)
+        self.exit_context(io_context)
 
     @asynccontextmanager
     async def handle_input_exception(
@@ -984,7 +994,7 @@ async def handle_input_exception(
             raise
         except (InputCancellation, asyncio.CancelledError):
             outputs = await io_context.output_items_cancellation(started_at)
-            await self._send_outputs(outputs)
+            await self._send_outputs(io_context, outputs)
             logger.warning(f"Successfully canceled input {io_context.input_ids}")
             return
         except BaseException as exc:
@@ -995,13 +1005,24 @@ async def handle_input_exception(
             # print exception so it's logged
             print_exception(*sys.exc_info())
             outputs = await io_context.output_items_exception(started_at, self.task_id, exc)
-            await self._send_outputs(outputs)
+            await self._send_outputs(io_context, outputs)
 
-    def exit_context(self, input_ids: list[str]):
-        for input_id in input_ids:
-            self.current_inputs.pop(input_id)
-
-        self._input_slots.release()
+    def exit_context(self, io_context: IOContext):
+        # A cancellation can land after outputs were already sent, re-entering here for the same
+        # context; exit at most once so the input slot is not released twice.
+        if io_context._exited:
+            return
+        io_context._exited = True
+        try:
+            for input_id in io_context.input_ids:
+                # A retry admitted before its predecessor finished exiting takes over the
+                # tracking entry, so remove it only if this context still owns it.
+                if self.current_inputs.get(input_id) is io_context:
+                    self.current_inputs.pop(input_id)
+                else:
+                    logger.warning(f"Input {input_id} missing from active input tracking")
+        finally:
+            self._input_slots.release()
 
     # skip inspection of user-generated output_data for synchronicity input translation
     @synchronizer.no_io_translation
@@ -1013,7 +1034,7 @@ async def push_outputs(
     ) -> None:
         # The standard output encoding+sending method for successful function outputs
         outputs = await io_context.output_items(started_at, output_data)
-        await self._send_outputs(outputs)
+        await self._send_outputs(io_context, outputs)
 
     @asynccontextmanager
     async def snapshot_context_manager(self) -> AsyncGenerator[None, None]:

py/test/container_test.py

@@ -2157,6 +2157,58 @@ async def _func(x):
     assert not triggered_assertions
 
 
+def test_container_io_manager_drops_duplicate_input_delivery(client, servicer):
+    dummy_container_args = api_pb2.ContainerArguments(
+        function_id="fu-123", function_def=api_pb2.Function(target_concurrent_inputs=2)
+    )
+    from modal._utils.async_utils import synchronizer
+
+    io_manager = ContainerIOManager(dummy_container_args, client)
+    _io_manager = synchronizer._translate_in(io_manager)
+
+    async def _func(x):
+        await asyncio.sleep(x)
+
+    fin_func = FinalizedFunction(
+        _func, is_async=True, is_generator=False, supported_output_formats=[api_pb2.DATA_FORMAT_PICKLE]
+    )
+
+    duplicate_input_id = "in-xyz0"
+    servicer.container_inputs = _get_inputs(((42,), {}), n=4)
+    servicer.container_inputs[1].inputs[0].input_id = duplicate_input_id  # duplicate of a running input
+    servicer.container_inputs[2].inputs[0].input_id = duplicate_input_id  # retry of the running input
+    servicer.container_inputs[2].inputs[0].retry_count = 1
+    servicer.container_inputs[3].inputs[0].input_id = duplicate_input_id  # redelivery after completion
+    yielded: list[IOContext] = []
+
+    for io_context in io_manager.run_inputs_outputs(finalized_functions={"": fin_func}):
+        yielded.append(io_context)
+        if len(yielded) == 2:
+            # The same-id same-retry-count duplicate was dropped and its slot freed, while the
+            # same-id delivery with a new retry count is a genuine retry and runs.
+            assert yielded[0].input_ids == (duplicate_input_id,)
+            assert yielded[0].retry_counts == (0,)
+            assert yielded[1].input_ids == (duplicate_input_id,)
+            assert yielded[1].retry_counts == (1,)
+            io_manager.push_outputs(yielded[0], started_at=0.0, output_data=[None])
+            # The predecessor's exit must not evict the retry's tracking entry (cancel routing).
+            assert _io_manager.current_inputs == {duplicate_input_id: yielded[1]}
+            io_manager.push_outputs(yielded[1], started_at=0.0, output_data=[None])
+            assert duplicate_input_id not in _io_manager.current_inputs
+        elif len(yielded) == 3:
+            # Once the first copy completed, a redelivery of the same input id runs again.
+            assert io_context.input_ids == (duplicate_input_id,)
+            assert io_context.retry_counts == (0,)
+            io_manager.push_outputs(yielded[2], started_at=0.0, output_data=[None])
+            # A repeated exit for an already-exited context must not release its slot twice.
+            active_after = _io_manager._input_slots.active
+            io_manager.exit_context(yielded[2])
+            assert _io_manager._input_slots.active == active_after
+
+    assert len(yielded) == 3
+    assert not _io_manager.current_inputs
+
+
 @pytest.mark.asyncio
 async def test_input_slots():
     slots = InputSlots(10)