modal 1.1.1.dev44__tar.gz → 1.1.2__tar.gz

{modal-1.1.1.dev44 → modal-1.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modal
-Version: 1.1.1.dev44
+Version: 1.1.2
 Summary: Python client library for Modal
 Author-email: Modal Labs <support@modal.com>
 License: Apache-2.0
@@ -22,7 +22,7 @@ Requires-Dist: click~=8.1
 Requires-Dist: grpclib<0.4.9,>=0.4.7
 Requires-Dist: protobuf!=4.24.0,<7.0,>=3.19
 Requires-Dist: rich>=12.0.0
-Requires-Dist: synchronicity~=0.10.1
+Requires-Dist: synchronicity~=0.10.2
 Requires-Dist: toml
 Requires-Dist: typer>=0.9
 Requires-Dist: types-certifi

{modal-1.1.1.dev44 → modal-1.1.2}/modal/__main__.py

@@ -37,7 +37,6 @@ def main():
 
         from grpclib import GRPCError, Status
         from rich.panel import Panel
-        from rich.text import Text
 
         if isinstance(exc, GRPCError):
             status_map = {
@@ -69,7 +68,7 @@ def main():
                 content = f"{content}\n\nNote: {' '.join(notes)}"
 
         console = make_console(stderr=True)
-        panel = Panel(Text(content), title=title, title_align="left", border_style="red")
+        panel = Panel(content, title=title, title_align="left", border_style="red")
         console.print(panel, highlight=False)
         sys.exit(1)
 

{modal-1.1.1.dev44 → modal-1.1.2}/modal/_container_entrypoint.py

@@ -185,8 +185,9 @@ def call_function(
 ):
     async def run_input_async(io_context: IOContext) -> None:
         started_at = time.time()
-        input_ids, function_call_ids = io_context.input_ids, io_context.function_call_ids
-        reset_context = execution_context._set_current_context_ids(input_ids, function_call_ids)
+        reset_context = execution_context._set_current_context_ids(
+            io_context.input_ids, io_context.function_call_ids, io_context.attempt_tokens
+        )
         async with container_io_manager.handle_input_exception.aio(io_context, started_at):
             res = io_context.call_finalized_function()
             # TODO(erikbern): any exception below shouldn't be considered a user exception
@@ -195,9 +196,14 @@ def call_function(
                     raise InvalidError(f"Async generator function returned value of type {type(res)}")
 
                 # Send up to this many outputs at a time.
+                current_function_call_id = execution_context.current_function_call_id()
+                assert current_function_call_id is not None  # Set above.
+                current_attempt_token = execution_context.current_attempt_token()
+                assert current_attempt_token is not None  # Set above, but can be empty string.
                 generator_queue: asyncio.Queue[Any] = await container_io_manager._queue_create.aio(1024)
                 async with container_io_manager.generator_output_sender(
-                    function_call_ids[0],
+                    current_function_call_id,
+                    current_attempt_token,
                     io_context.finalized_function.data_format,
                     generator_queue,
                 ):
@@ -230,8 +236,9 @@ def call_function(
 
     def run_input_sync(io_context: IOContext) -> None:
         started_at = time.time()
-        input_ids, function_call_ids = io_context.input_ids, io_context.function_call_ids
-        reset_context = execution_context._set_current_context_ids(input_ids, function_call_ids)
+        reset_context = execution_context._set_current_context_ids(
+            io_context.input_ids, io_context.function_call_ids, io_context.attempt_tokens
+        )
         with container_io_manager.handle_input_exception(io_context, started_at):
             res = io_context.call_finalized_function()
 
@@ -241,10 +248,14 @@ def call_function(
                     raise InvalidError(f"Generator function returned value of type {type(res)}")
 
                 # Send up to this many outputs at a time.
+                current_function_call_id = execution_context.current_function_call_id()
+                assert current_function_call_id is not None  # Set above.
+                current_attempt_token = execution_context.current_attempt_token()
+                assert current_attempt_token is not None  # Set above, but can be empty string.
                 generator_queue: asyncio.Queue[Any] = container_io_manager._queue_create(1024)
-
                 with container_io_manager.generator_output_sender(
-                    function_call_ids[0],
+                    current_function_call_id,
+                    current_attempt_token,
                     io_context.finalized_function.data_format,
                     generator_queue,
                 ):

{modal-1.1.1.dev44 → modal-1.1.2}/modal/_functions.py

@@ -9,7 +9,7 @@ import warnings
 from collections.abc import AsyncGenerator, Sequence, Sized
 from dataclasses import dataclass
 from pathlib import PurePosixPath
-from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+from typing import TYPE_CHECKING, Any, AsyncIterator, Callable, Optional, Union
 
 import typing_extensions
 from google.protobuf.message import Message
@@ -71,6 +71,8 @@ from .mount import _get_client_mount, _Mount
 from .network_file_system import _NetworkFileSystem, network_file_system_mount_protos
 from .output import _get_output_manager
 from .parallel_map import (
+    _experimental_spawn_map_async,
+    _experimental_spawn_map_sync,
     _for_each_async,
     _for_each_sync,
     _map_async,
@@ -78,6 +80,7 @@ from .parallel_map import (
     _map_invocation_inputplane,
     _map_sync,
     _spawn_map_async,
+    _spawn_map_invocation,
     _spawn_map_sync,
     _starmap_async,
     _starmap_sync,
@@ -214,7 +217,11 @@ class _Invocation:
         return _Invocation(stub, function_call_id, client, retry_context)
 
     async def pop_function_call_outputs(
-        self, timeout: Optional[float], clear_on_success: bool, input_jwts: Optional[list[str]] = None
+        self,
+        index: int = 0,
+        timeout: Optional[float] = None,
+        clear_on_success: bool = False,
+        input_jwts: Optional[list[str]] = None,
     ) -> api_pb2.FunctionGetOutputsResponse:
         t0 = time.time()
         if timeout is None:
@@ -232,6 +239,8 @@ class _Invocation:
                 clear_on_success=clear_on_success,
                 requested_at=time.time(),
                 input_jwts=input_jwts,
+                start_idx=index,
+                end_idx=index,
             )
             response: api_pb2.FunctionGetOutputsResponse = await retry_transient_errors(
                 self.stub.FunctionGetOutputs,
@@ -265,6 +274,7 @@ class _Invocation:
         # waits indefinitely for a single result for the function, and clear the outputs buffer after
         item: api_pb2.FunctionGetOutputsItem = (
             await self.pop_function_call_outputs(
+                index=0,
                 timeout=None,
                 clear_on_success=True,
                 input_jwts=[expected_jwt] if expected_jwt else None,
@@ -308,14 +318,16 @@ class _Invocation:
 
             await self._retry_input()
 
-    async def poll_function(self, timeout: Optional[float] = None):
+    async def poll_function(self, timeout: Optional[float] = None, *, index: int = 0):
         """Waits up to timeout for a result from a function.
 
         If timeout is `None`, waits indefinitely. This function is not
         cancellation-safe.
         """
         response: api_pb2.FunctionGetOutputsResponse = await self.pop_function_call_outputs(
-            timeout=timeout, clear_on_success=False
+            index=index,
+            timeout=timeout,
+            clear_on_success=False,
         )
         if len(response.outputs) == 0 and response.num_unfinished_inputs == 0:
             # if no unfinished inputs and no outputs, then function expired
@@ -348,11 +360,47 @@ class _Invocation:
                 if items_total is not None and items_received >= items_total:
                     break
 
+    async def enumerate(self, start_index: int, end_index: int):
+        """Iterate over the results of the function call in the range [start_index, end_index)."""
+        limit = 49
+        current_index = start_index
+        while current_index < end_index:
+            # batch_end_indx is inclusive, so we subtract 1 to get the last index in the batch.
+            batch_end_index = min(current_index + limit, end_index) - 1
+            request = api_pb2.FunctionGetOutputsRequest(
+                function_call_id=self.function_call_id,
+                timeout=0,
+                last_entry_id="0-0",
+                clear_on_success=False,
+                requested_at=time.time(),
+                start_idx=current_index,
+                end_idx=batch_end_index,
+            )
+            response: api_pb2.FunctionGetOutputsResponse = await retry_transient_errors(
+                self.stub.FunctionGetOutputs,
+                request,
+                attempt_timeout=ATTEMPT_TIMEOUT_GRACE_PERIOD,
+            )
+
+            outputs = list(response.outputs)
+            outputs.sort(key=lambda x: x.idx)
+            for output in outputs:
+                if output.idx != current_index:
+                    break
+                result = await _process_result(output.result, output.data_format, self.stub, self.client)
+                yield output.idx, result
+                current_index += 1
+
+            # We're missing current_index, so we need to poll the function for the next result
+            if len(outputs) < (batch_end_index - current_index + 1):
+                result = await self.poll_function(index=current_index)
+                yield current_index, result
+                current_index += 1
+
 
 class _InputPlaneInvocation:
     """Internal client representation of a single-input call to a Modal Function using the input
-    plane server API. As of 4/22/2025, this class is experimental and not used in production.
-    It is OK to make breaking changes to this class."""
+    plane server API."""
 
     stub: ModalClientModal
 
@@ -462,7 +510,7 @@ class _InputPlaneInvocation:
                 _stream_function_call_data(
                     self.client,
                     self.stub,
-                    "",
+                    function_call_id=None,
                     variant="data_out",
                     attempt_token=self.attempt_token,
                 ),
@@ -603,7 +651,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         memory: Optional[Union[int, tuple[int, int]]] = None,
         proxy: Optional[_Proxy] = None,
         retries: Optional[Union[int, Retries]] = None,
-        timeout: Optional[int] = None,
+        timeout: int = 300,
         min_containers: Optional[int] = None,
         max_containers: Optional[int] = None,
         buffer_containers: Optional[int] = None,
@@ -1130,6 +1178,8 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                     target_concurrent_inputs=options.target_concurrent_inputs,
                     batch_max_size=options.batch_max_size,
                     batch_linger_ms=options.batch_wait_ms,
+                    scheduler_placement=options.scheduler_placement,
+                    cloud_provider_str=options.cloud,
                 )
             else:
                 options_pb = None
@@ -1274,7 +1324,8 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
 
             self._hydrate(response.function_id, resolver.client, response.handle_metadata)
 
-        rep = f"Function.from_name('{app_name}', '{name}')"
+        environment_rep = f", environment_name={environment_name!r}" if environment_name else ""
+        rep = f"modal.Function.from_name('{app_name}', '{name}'{environment_rep})"
         return cls._from_loader(_load_remote, rep, is_another_app=True, hydrate_lazily=True)
 
     @classmethod
@@ -1543,6 +1594,22 @@ Use the `Function.get_web_url()` method instead.
                 async for item in stream:
                     yield item
 
+    @live_method
+    async def _spawn_map(self, input_queue: _SynchronizedQueue) -> "_FunctionCall[ReturnType]":
+        self._check_no_web_url("spawn_map")
+        if self._is_generator:
+            raise InvalidError("A generator function cannot be called with `.spawn_map(...)`.")
+
+        assert self._function_name
+        function_call_id, num_inputs = await _spawn_map_invocation(
+            self,
+            input_queue,
+            self.client,
+        )
+        metadata = api_pb2.FunctionCallFromIdResponse(function_call_id=function_call_id, num_inputs=num_inputs)
+        fc: _FunctionCall[ReturnType] = _FunctionCall._new_hydrated(function_call_id, self.client, metadata)
+        return fc
+
     async def _call_function(self, args, kwargs) -> ReturnType:
         invocation: Union[_Invocation, _InputPlaneInvocation]
         if self._input_plane_url:
@@ -1789,6 +1856,7 @@ Use the `Function.get_web_url()` method instead.
     starmap = MethodWithAio(_starmap_sync, _starmap_async, synchronizer)
     for_each = MethodWithAio(_for_each_sync, _for_each_async, synchronizer)
     spawn_map = MethodWithAio(_spawn_map_sync, _spawn_map_async, synchronizer)
+    experimental_spawn_map = MethodWithAio(_experimental_spawn_map_sync, _experimental_spawn_map_async, synchronizer)
 
 
 class _FunctionCall(typing.Generic[ReturnType], _Object, type_prefix="fc"):
@@ -1803,12 +1871,28 @@ class _FunctionCall(typing.Generic[ReturnType], _Object, type_prefix="fc"):
     """
 
     _is_generator: bool = False
+    _num_inputs: Optional[int] = None
 
     def _invocation(self):
         return _Invocation(self.client.stub, self.object_id, self.client)
 
-    async def get(self, timeout: Optional[float] = None) -> ReturnType:
-        """Get the result of the function call.
+    def _hydrate_metadata(self, metadata: Optional[Message]):
+        if not metadata:
+            return
+        assert isinstance(metadata, api_pb2.FunctionCallFromIdResponse)
+        self._num_inputs = metadata.num_inputs
+
+    @live_method
+    async def num_inputs(self) -> int:
+        """Get the number of inputs in the function call."""
+        # Should have been hydrated.
+        assert self._num_inputs is not None
+        return self._num_inputs
+
+    async def get(self, timeout: Optional[float] = None, *, index: int = 0) -> ReturnType:
+        """Get the result of the index-th input of the function call.
+        `.spawn()` calls have a single output, so only specifying `index=0` is valid.
+        A non-zero index is useful when your function has multiple outputs, like via `.spawn_map()`.
 
         This function waits indefinitely by default. It takes an optional
         `timeout` argument that specifies the maximum number of seconds to wait,
@@ -1816,7 +1900,37 @@ class _FunctionCall(typing.Generic[ReturnType], _Object, type_prefix="fc"):
 
         The returned coroutine is not cancellation-safe.
         """
-        return await self._invocation().poll_function(timeout=timeout)
+        return await self._invocation().poll_function(timeout=timeout, index=index)
+
+    @live_method_gen
+    async def iter(self, *, start: int = 0, end: Optional[int] = None) -> AsyncIterator[ReturnType]:
+        """Iterate in-order over the results of the function call.
+
+        Optionally, specify a range [start, end) to iterate over.
+
+        Example:
+        ```python
+        @app.function()
+        def my_func(a):
+            return a ** 2
+
+
+        @app.local_entrypoint()
+        def main():
+            fc = my_func.spawn_map([1, 2, 3, 4])
+            assert list(fc.iter()) == [1, 4, 9, 16]
+            assert list(fc.iter(start=1, end=3)) == [4, 9]
+        ```
+
+        If `end` is not provided, it will iterate over all results.
+        """
+        num_inputs = await self.num_inputs()
+        if end is None:
+            end = num_inputs
+        if start < 0 or end > num_inputs:
+            raise ValueError(f"Invalid index range: {start} to {end} for {num_inputs} inputs")
+        async for _, item in self._invocation().enumerate(start_index=start, end_index=end):
+            yield item
 
     async def get_call_graph(self) -> list[InputInfo]:
         """Returns a structure representing the call graph from a given root
@@ -1870,7 +1984,15 @@ class _FunctionCall(typing.Generic[ReturnType], _Object, type_prefix="fc"):
         if client is None:
             client = await _Client.from_env()
 
-        fc: _FunctionCall[Any] = _FunctionCall._new_hydrated(function_call_id, client, None)
+        async def _load(self: _FunctionCall, resolver: Resolver, existing_object_id: Optional[str]):
+            request = api_pb2.FunctionCallFromIdRequest(function_call_id=function_call_id)
+            resp = await retry_transient_errors(resolver.client.stub.FunctionCallFromId, request)
+            self._hydrate(function_call_id, resolver.client, resp)
+
+        rep = f"FunctionCall.from_id({function_call_id!r})"
+        fc: _FunctionCall[Any] = _FunctionCall._from_loader(_load, rep, hydrate_lazily=True)
+        # We already know the object ID, so we can set it directly
+        fc._object_id = function_call_id
         return fc
 
     @staticmethod

{modal-1.1.1.dev44 → modal-1.1.2}/modal/_object.py

@@ -191,9 +191,20 @@ class _Object:
     def _is_id_type(cls, object_id) -> bool:
         return cls._get_type_from_id(object_id) == cls
 
+    @classmethod
+    def _repr(cls, name: str, environment_name: Optional[str] = None) -> str:
+        public_cls = cls.__name__.strip("_")
+        environment_repr = f", environment_name={environment_name!r}" if environment_name else ""
+        return f"modal.{public_cls}.from_name({name!r}{environment_repr})"
+
     @classmethod
     def _new_hydrated(
-        cls, object_id: str, client: _Client, handle_metadata: Optional[Message], is_another_app: bool = False
+        cls,
+        object_id: str,
+        client: _Client,
+        handle_metadata: Optional[Message],
+        is_another_app: bool = False,
+        rep: Optional[str] = None,
     ) -> Self:
         obj_cls: type[Self]
         if cls._type_prefix is not None:
@@ -210,7 +221,7 @@ class _Object:
 
         # Instantiate provider
         obj = _Object.__new__(obj_cls)
-        rep = f"Object({object_id})"  # TODO(erikbern): dumb
+        rep = rep or f"modal.{obj_cls.__name__.strip('_')}.from_id({object_id!r})"
         obj._init(rep, is_another_app=is_another_app)
         obj._hydrate(object_id, client, handle_metadata)
 

{modal-1.1.1.dev44 → modal-1.1.2}/modal/_partial_function.py

@@ -282,7 +282,7 @@ class _MethodDecoratorType:
 
 # TODO(elias): fix support for coroutine type unwrapping for methods (static typing)
 def _method(
-    _warn_parentheses_missing=None,
+    _warn_parentheses_missing=None,  # mdmd:line-hidden
     *,
     # Set this to True if it's a non-generator function returning
     # a [sync/async] generator object
@@ -337,7 +337,7 @@ def _parse_custom_domains(custom_domains: Optional[Iterable[str]] = None) -> lis
 
 
 def _fastapi_endpoint(
-    _warn_parentheses_missing=None,
+    _warn_parentheses_missing=None,  # mdmd:line-hidden
     *,
     method: str = "GET",  # REST method for the created endpoint.
     label: Optional[str] = None,  # Label for created endpoint. Final subdomain will be <workspace>--<label>.modal.run.
@@ -400,7 +400,7 @@ def _fastapi_endpoint(
 
 
 def _web_endpoint(
-    _warn_parentheses_missing=None,
+    _warn_parentheses_missing=None,  # mdmd:line-hidden
     *,
     method: str = "GET",  # REST method for the created endpoint.
     label: Optional[str] = None,  # Label for created endpoint. Final subdomain will be <workspace>--<label>.modal.run.
@@ -468,7 +468,7 @@ def _web_endpoint(
 
 
 def _asgi_app(
-    _warn_parentheses_missing=None,
+    _warn_parentheses_missing=None,  # mdmd:line-hidden
     *,
     label: Optional[str] = None,  # Label for created endpoint. Final subdomain will be <workspace>--<label>.modal.run.
     custom_domains: Optional[Iterable[str]] = None,  # Deploy this endpoint on a custom domain.
@@ -525,7 +525,7 @@ def _asgi_app(
 
 
 def _wsgi_app(
-    _warn_parentheses_missing=None,
+    _warn_parentheses_missing=None,  # mdmd:line-hidden
     *,
     label: Optional[str] = None,  # Label for created endpoint. Final subdomain will be <workspace>--<label>.modal.run.
     custom_domains: Optional[Iterable[str]] = None,  # Deploy this endpoint on a custom domain.
@@ -645,7 +645,7 @@ def _web_server(
 
 
 def _enter(
-    _warn_parentheses_missing=None,
+    _warn_parentheses_missing=None,  # mdmd:line-hidden
     *,
     snap: bool = False,
 ) -> Callable[[Union[_PartialFunction, NullaryMethod]], _PartialFunction]:
@@ -696,7 +696,7 @@ def _exit(_warn_parentheses_missing=None) -> Callable[[NullaryMethod], _PartialF
 
 
 def _batched(
-    _warn_parentheses_missing=None,
+    _warn_parentheses_missing=None,  # mdmd:line-hidden
     *,
     max_batch_size: int,
     wait_ms: int,
@@ -758,7 +758,7 @@ def _batched(
 
 
 def _concurrent(
-    _warn_parentheses_missing=None,
+    _warn_parentheses_missing=None,  # mdmd:line-hidden
     *,
     max_inputs: int,  # Hard limit on each container's input concurrency
     target_inputs: Optional[int] = None,  # Input concurrency that Modal's autoscaler should target

{modal-1.1.1.dev44 → modal-1.1.2}/modal/_runtime/asgi.py

@@ -16,7 +16,7 @@ from modal.config import logger
 from modal.exception import ExecutionError, InvalidError
 from modal.experimental import stop_fetching_inputs
 
-from .execution_context import current_function_call_id
+from .execution_context import current_attempt_token, current_function_call_id
 
 FIRST_MESSAGE_TIMEOUT_SECONDS = 5.0
 
@@ -106,6 +106,7 @@ def asgi_app_wrapper(asgi_app, container_io_manager) -> tuple[Callable[..., Asyn
             raise ExecutionError("Unpexected state in ASGI scope")
         scope["state"] = state
         function_call_id = current_function_call_id()
+        attempt_token = current_attempt_token()
         assert function_call_id, "internal error: function_call_id not set in asgi_app() scope"
 
         messages_from_app: asyncio.Queue[dict[str, Any]] = asyncio.Queue(1)
@@ -142,7 +143,7 @@ def asgi_app_wrapper(asgi_app, container_io_manager) -> tuple[Callable[..., Asyn
             # This initial message, "http.request" or "websocket.connect", should be sent
             # immediately after starting the ASGI app's function call. If it is not received, that
             # indicates a request cancellation or other abnormal circumstance.
-            message_gen = container_io_manager.get_data_in.aio(function_call_id)
+            message_gen = container_io_manager.get_data_in.aio(function_call_id, attempt_token)
             first_message_task = asyncio.create_task(message_gen.__anext__())
 
             try:

{modal-1.1.1.dev44 → modal-1.1.2}/modal/_runtime/container_io_manager.py

@@ -39,7 +39,6 @@ from modal.exception import ClientClosed, InputCancellation, InvalidError, Seria
 from modal_proto import api_pb2
 
 if TYPE_CHECKING:
-    import modal._runtime.asgi
     import modal._runtime.user_code_imports
 
 
@@ -66,6 +65,7 @@ class IOContext:
     input_ids: list[str]
     retry_counts: list[int]
     function_call_ids: list[str]
+    attempt_tokens: list[str]
     function_inputs: list[api_pb2.FunctionInput]
     finalized_function: "modal._runtime.user_code_imports.FinalizedFunction"
 
@@ -77,6 +77,7 @@ class IOContext:
         input_ids: list[str],
         retry_counts: list[int],
         function_call_ids: list[str],
+        attempt_tokens: list[str],
         finalized_function: "modal._runtime.user_code_imports.FinalizedFunction",
         function_inputs: list[api_pb2.FunctionInput],
         is_batched: bool,
@@ -85,6 +86,7 @@ class IOContext:
         self.input_ids = input_ids
         self.retry_counts = retry_counts
         self.function_call_ids = function_call_ids
+        self.attempt_tokens = attempt_tokens
         self.finalized_function = finalized_function
         self.function_inputs = function_inputs
         self._is_batched = is_batched
@@ -95,11 +97,11 @@ class IOContext:
         cls,
         client: _Client,
         finalized_functions: dict[str, "modal._runtime.user_code_imports.FinalizedFunction"],
-        inputs: list[tuple[str, int, str, api_pb2.FunctionInput]],
+        inputs: list[tuple[str, int, str, str, api_pb2.FunctionInput]],
         is_batched: bool,
     ) -> "IOContext":
         assert len(inputs) >= 1 if is_batched else len(inputs) == 1
-        input_ids, retry_counts, function_call_ids, function_inputs = zip(*inputs)
+        input_ids, retry_counts, function_call_ids, attempt_tokens, function_inputs = zip(*inputs)
 
         async def _populate_input_blobs(client: _Client, input: api_pb2.FunctionInput) -> api_pb2.FunctionInput:
             # If we got a pointer to a blob, download it from S3.
@@ -121,6 +123,7 @@ class IOContext:
             cast(list[str], input_ids),
             cast(list[int], retry_counts),
             cast(list[str], function_call_ids),
+            cast(list[str], attempt_tokens),
             finalized_function,
             cast(list[api_pb2.FunctionInput], function_inputs),
             is_batched,
@@ -300,11 +303,7 @@ class _ContainerIOManager:
         self.function_def = container_args.function_def
         self.checkpoint_id = container_args.checkpoint_id or None
 
-        # We could also have the worker pass this in explicitly.
-        self.input_plane_server_url = None
-        for obj in container_args.app_layout.objects:
-            if obj.object_id == self.function_id:
-                self.input_plane_server_url = obj.function_handle_metadata.input_plane_url
+        self.input_plane_server_url = container_args.input_plane_server_url
 
         self.calls_completed = 0
         self.total_user_time = 0.0
@@ -484,18 +483,21 @@ class _ContainerIOManager:
             else {"data": data}
         )
 
-    async def get_data_in(self, function_call_id: str) -> AsyncIterator[Any]:
+    async def get_data_in(self, function_call_id: str, attempt_token: Optional[str]) -> AsyncIterator[Any]:
         """Read from the `data_in` stream of a function call."""
         stub = self._client.stub
         if self.input_plane_server_url:
             stub = await self._client.get_stub(self.input_plane_server_url)
 
-        async for data in _stream_function_call_data(self._client, stub, function_call_id, "data_in"):
+        async for data in _stream_function_call_data(
+            self._client, stub, function_call_id, variant="data_in", attempt_token=attempt_token
+        ):
             yield data
 
     async def put_data_out(
         self,
         function_call_id: str,
+        attempt_token: str,
         start_index: int,
         data_format: int,
         serialized_messages: list[Any],
@@ -516,6 +518,8 @@ class _ContainerIOManager:
             data_chunks.append(chunk)
 
         req = api_pb2.FunctionCallPutDataRequest(function_call_id=function_call_id, data_chunks=data_chunks)
+        if attempt_token:
+            req.attempt_token = attempt_token  # oneof clears function_call_id.
 
         if self.input_plane_server_url:
             stub = await self._client.get_stub(self.input_plane_server_url)
@@ -525,7 +529,7 @@ class _ContainerIOManager:
 
     @asynccontextmanager
     async def generator_output_sender(
-        self, function_call_id: str, data_format: int, message_rx: asyncio.Queue
+        self, function_call_id: str, attempt_token: str, data_format: int, message_rx: asyncio.Queue
     ) -> AsyncGenerator[None, None]:
         """Runs background task that feeds generator outputs into a function call's `data_out` stream."""
         GENERATOR_STOP_SENTINEL = Sentinel()
@@ -554,7 +558,7 @@ class _ContainerIOManager:
                     else:
                         serialized_messages.append(serialize_data_format(message, data_format))
                         total_size += len(serialized_messages[-1]) + 512  # 512 bytes for estimated framing overhead
-                await self.put_data_out(function_call_id, index, data_format, serialized_messages)
+                await self.put_data_out(function_call_id, attempt_token, index, data_format, serialized_messages)
                 index += len(serialized_messages)
 
         task = asyncio.create_task(generator_output_task())
@@ -590,7 +594,7 @@ class _ContainerIOManager:
         self,
         batch_max_size: int,
         batch_wait_ms: int,
-    ) -> AsyncIterator[list[tuple[str, int, str, api_pb2.FunctionInput]]]:
+    ) -> AsyncIterator[list[tuple[str, int, str, str, api_pb2.FunctionInput]]]:
         request = api_pb2.FunctionGetInputsRequest(function_id=self.function_id)
         iteration = 0
         while self._fetching_inputs:
@@ -625,7 +629,9 @@ class _ContainerIOManager:
                         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.input))
+                        inputs.append(
+                            (item.input_id, item.retry_count, item.function_call_id, item.attempt_token, item.input)
+                        )
                         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")