modAL 1.1.2__tar.gz → 1.1.2.dev1__tar.gz

{modal-1.1.2 → modal-1.1.2.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modal
-Version: 1.1.2
+Version: 1.1.2.dev1
 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.2
+Requires-Dist: synchronicity~=0.10.1
 Requires-Dist: toml
 Requires-Dist: typer>=0.9
 Requires-Dist: types-certifi

{modal-1.1.2 → modal-1.1.2.dev1}/modal/__main__.py

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

{modal-1.1.2 → modal-1.1.2.dev1}/modal/_container_entrypoint.py

@@ -185,9 +185,8 @@ def call_function(
 ):
     async def run_input_async(io_context: IOContext) -> None:
         started_at = time.time()
-        reset_context = execution_context._set_current_context_ids(
-            io_context.input_ids, io_context.function_call_ids, io_context.attempt_tokens
-        )
+        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)
         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
@@ -196,14 +195,9 @@ 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(
-                    current_function_call_id,
-                    current_attempt_token,
+                    function_call_ids[0],
                     io_context.finalized_function.data_format,
                     generator_queue,
                 ):
@@ -236,9 +230,8 @@ def call_function(
 
     def run_input_sync(io_context: IOContext) -> None:
         started_at = time.time()
-        reset_context = execution_context._set_current_context_ids(
-            io_context.input_ids, io_context.function_call_ids, io_context.attempt_tokens
-        )
+        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)
         with container_io_manager.handle_input_exception(io_context, started_at):
             res = io_context.call_finalized_function()
 
@@ -248,14 +241,10 @@ 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(
-                    current_function_call_id,
-                    current_attempt_token,
+                    function_call_ids[0],
                     io_context.finalized_function.data_format,
                     generator_queue,
                 ):

{modal-1.1.2 → modal-1.1.2.dev1}/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, AsyncIterator, Callable, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union
 
 import typing_extensions
 from google.protobuf.message import Message
@@ -71,8 +71,6 @@ 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,
@@ -80,7 +78,6 @@ from .parallel_map import (
     _map_invocation_inputplane,
     _map_sync,
     _spawn_map_async,
-    _spawn_map_invocation,
     _spawn_map_sync,
     _starmap_async,
     _starmap_sync,
@@ -217,11 +214,7 @@ class _Invocation:
         return _Invocation(stub, function_call_id, client, retry_context)
 
     async def pop_function_call_outputs(
-        self,
-        index: int = 0,
-        timeout: Optional[float] = None,
-        clear_on_success: bool = False,
-        input_jwts: Optional[list[str]] = None,
+        self, timeout: Optional[float], clear_on_success: bool, input_jwts: Optional[list[str]] = None
     ) -> api_pb2.FunctionGetOutputsResponse:
         t0 = time.time()
         if timeout is None:
@@ -239,8 +232,6 @@ 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,
@@ -274,7 +265,6 @@ 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,
@@ -318,16 +308,14 @@ class _Invocation:
 
             await self._retry_input()
 
-    async def poll_function(self, timeout: Optional[float] = None, *, index: int = 0):
+    async def poll_function(self, timeout: Optional[float] = None):
         """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(
-            index=index,
-            timeout=timeout,
-            clear_on_success=False,
+            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
@@ -360,47 +348,11 @@ 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."""
+    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."""
 
     stub: ModalClientModal
 
@@ -510,7 +462,7 @@ class _InputPlaneInvocation:
                 _stream_function_call_data(
                     self.client,
                     self.stub,
-                    function_call_id=None,
+                    "",
                     variant="data_out",
                     attempt_token=self.attempt_token,
                 ),
@@ -651,7 +603,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: int = 300,
+        timeout: Optional[int] = None,
         min_containers: Optional[int] = None,
         max_containers: Optional[int] = None,
         buffer_containers: Optional[int] = None,
@@ -1178,8 +1130,6 @@ 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
@@ -1324,8 +1274,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
 
             self._hydrate(response.function_id, resolver.client, response.handle_metadata)
 
-        environment_rep = f", environment_name={environment_name!r}" if environment_name else ""
-        rep = f"modal.Function.from_name('{app_name}', '{name}'{environment_rep})"
+        rep = f"Function.from_name('{app_name}', '{name}')"
         return cls._from_loader(_load_remote, rep, is_another_app=True, hydrate_lazily=True)
 
     @classmethod
@@ -1594,22 +1543,6 @@ 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:
@@ -1856,7 +1789,6 @@ 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"):
@@ -1871,28 +1803,12 @@ 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)
 
-    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()`.
+    async def get(self, timeout: Optional[float] = None) -> ReturnType:
+        """Get the result of the function call.
 
         This function waits indefinitely by default. It takes an optional
         `timeout` argument that specifies the maximum number of seconds to wait,
@@ -1900,37 +1816,7 @@ class _FunctionCall(typing.Generic[ReturnType], _Object, type_prefix="fc"):
 
         The returned coroutine is not cancellation-safe.
         """
-        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
+        return await self._invocation().poll_function(timeout=timeout)
 
     async def get_call_graph(self) -> list[InputInfo]:
         """Returns a structure representing the call graph from a given root
@@ -1984,15 +1870,7 @@ class _FunctionCall(typing.Generic[ReturnType], _Object, type_prefix="fc"):
         if client is None:
             client = await _Client.from_env()
 
-        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
+        fc: _FunctionCall[Any] = _FunctionCall._new_hydrated(function_call_id, client, None)
         return fc
 
     @staticmethod

{modal-1.1.2 → modal-1.1.2.dev1}/modal/_object.py

@@ -191,20 +191,9 @@ 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,
-        rep: Optional[str] = None,
+        cls, object_id: str, client: _Client, handle_metadata: Optional[Message], is_another_app: bool = False
     ) -> Self:
         obj_cls: type[Self]
         if cls._type_prefix is not None:
@@ -221,7 +210,7 @@ class _Object:
 
         # Instantiate provider
         obj = _Object.__new__(obj_cls)
-        rep = rep or f"modal.{obj_cls.__name__.strip('_')}.from_id({object_id!r})"
+        rep = f"Object({object_id})"  # TODO(erikbern): dumb
         obj._init(rep, is_another_app=is_another_app)
         obj._hydrate(object_id, client, handle_metadata)
 

{modal-1.1.2 → modal-1.1.2.dev1}/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,  # mdmd:line-hidden
+    _warn_parentheses_missing=None,
     *,
     # 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,  # mdmd:line-hidden
+    _warn_parentheses_missing=None,
     *,
     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,  # mdmd:line-hidden
+    _warn_parentheses_missing=None,
     *,
     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,  # mdmd:line-hidden
+    _warn_parentheses_missing=None,
     *,
     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,  # mdmd:line-hidden
+    _warn_parentheses_missing=None,
     *,
     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,  # mdmd:line-hidden
+    _warn_parentheses_missing=None,
     *,
     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,  # mdmd:line-hidden
+    _warn_parentheses_missing=None,
     *,
     max_batch_size: int,
     wait_ms: int,
@@ -758,7 +758,7 @@ def _batched(
 
 
 def _concurrent(
-    _warn_parentheses_missing=None,  # mdmd:line-hidden
+    _warn_parentheses_missing=None,
     *,
     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.2 → modal-1.1.2.dev1}/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_attempt_token, current_function_call_id
+from .execution_context import current_function_call_id
 
 FIRST_MESSAGE_TIMEOUT_SECONDS = 5.0
 
@@ -106,7 +106,6 @@ 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)
@@ -143,7 +142,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, attempt_token)
+            message_gen = container_io_manager.get_data_in.aio(function_call_id)
             first_message_task = asyncio.create_task(message_gen.__anext__())
 
             try:

{modal-1.1.2 → modal-1.1.2.dev1}/modal/_runtime/container_io_manager.py

@@ -39,6 +39,7 @@ 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
 
 
@@ -65,7 +66,6 @@ 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,7 +77,6 @@ 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,
@@ -86,7 +85,6 @@ 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
@@ -97,11 +95,11 @@ class IOContext:
         cls,
         client: _Client,
         finalized_functions: dict[str, "modal._runtime.user_code_imports.FinalizedFunction"],
-        inputs: list[tuple[str, int, str, str, api_pb2.FunctionInput]],
+        inputs: list[tuple[str, int, 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, attempt_tokens, function_inputs = zip(*inputs)
+        input_ids, retry_counts, function_call_ids, 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.
@@ -123,7 +121,6 @@ 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,
@@ -303,7 +300,11 @@ class _ContainerIOManager:
         self.function_def = container_args.function_def
         self.checkpoint_id = container_args.checkpoint_id or None
 
-        self.input_plane_server_url = container_args.input_plane_server_url
+        # 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.calls_completed = 0
         self.total_user_time = 0.0
@@ -483,21 +484,18 @@ class _ContainerIOManager:
             else {"data": data}
         )
 
-    async def get_data_in(self, function_call_id: str, attempt_token: Optional[str]) -> AsyncIterator[Any]:
+    async def get_data_in(self, function_call_id: 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, variant="data_in", attempt_token=attempt_token
-        ):
+        async for data in _stream_function_call_data(self._client, stub, function_call_id, "data_in"):
             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],
@@ -518,8 +516,6 @@ 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)
@@ -529,7 +525,7 @@ class _ContainerIOManager:
 
     @asynccontextmanager
     async def generator_output_sender(
-        self, function_call_id: str, attempt_token: str, data_format: int, message_rx: asyncio.Queue
+        self, function_call_id: 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()
@@ -558,7 +554,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, attempt_token, index, data_format, serialized_messages)
+                await self.put_data_out(function_call_id, index, data_format, serialized_messages)
                 index += len(serialized_messages)
 
         task = asyncio.create_task(generator_output_task())
@@ -594,7 +590,7 @@ class _ContainerIOManager:
         self,
         batch_max_size: int,
         batch_wait_ms: int,
-    ) -> AsyncIterator[list[tuple[str, int, str, str, api_pb2.FunctionInput]]]:
+    ) -> AsyncIterator[list[tuple[str, int, str, api_pb2.FunctionInput]]]:
         request = api_pb2.FunctionGetInputsRequest(function_id=self.function_id)
         iteration = 0
         while self._fetching_inputs:
@@ -629,9 +625,7 @@ 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.attempt_token, item.input)
-                        )
+                        inputs.append((item.input_id, item.retry_count, item.function_call_id, 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")