modal 1.0.3.dev10__py3-none-any.whl → 1.2.3.dev7__py3-none-any.whl

modal/_runtime/container_io_manager.pyi

@@ -9,13 +9,25 @@ import synchronicity.combined_types
 import typing
 import typing_extensions
 
-class UserException(Exception): ...
-class Sentinel: ...
+class UserException(Exception):
+    """Used to shut down the task gracefully."""
+
+    ...
+
+class Sentinel:
+    """Used to get type-stubs to work with this object."""
+
+    ...
 
 class IOContext:
+    """Context object for managing input, function calls, and function executions
+    in a batched or single input context.
+    """
+
     input_ids: list[str]
     retry_counts: list[int]
     function_call_ids: list[str]
+    attempt_tokens: list[str]
     function_inputs: list[modal_proto.api_pb2.FunctionInput]
     finalized_function: modal._runtime.user_code_imports.FinalizedFunction
     _cancel_issued: bool
@@ -26,32 +38,56 @@ 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[modal_proto.api_pb2.FunctionInput],
         is_batched: bool,
         client: modal.client._Client,
-    ): ...
+    ):
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
+
     @classmethod
     async def create(
         cls,
         client: modal.client._Client,
         finalized_functions: dict[str, modal._runtime.user_code_imports.FinalizedFunction],
-        inputs: list[tuple[str, int, str, modal_proto.api_pb2.FunctionInput]],
+        inputs: list[tuple[str, int, str, str, modal_proto.api_pb2.FunctionInput]],
         is_batched: bool,
     ) -> IOContext: ...
     def set_cancel_callback(self, cb: collections.abc.Callable[[], None]): ...
     def cancel(self): ...
     def _args_and_kwargs(self) -> tuple[tuple[typing.Any, ...], dict[str, list[typing.Any]]]: ...
-    def call_finalized_function(self) -> typing.Any: ...
-    def validate_output_data(self, data: typing.Any) -> list[typing.Any]: ...
+    def _generator_output_format(self) -> int: ...
+    def _prepare_batch_output(self, data: typing.Any) -> list[typing.Any]: ...
+    def call_function_sync(self) -> list[typing.Any]: ...
+    async def call_function_async(self) -> list[typing.Any]: ...
+    def call_generator_sync(self) -> typing.Generator[typing.Any, None, None]: ...
+    def call_generator_async(self) -> collections.abc.AsyncGenerator[typing.Any, None]: ...
+    async def output_items_cancellation(self, started_at: float): ...
+    def _determine_output_format(self, input_format: int) -> int: ...
+    async def output_items_exception(
+        self, started_at: float, task_id: str, exc: BaseException
+    ) -> list[modal_proto.api_pb2.FunctionPutOutputsItem]: ...
+    def output_items_generator_done(
+        self, started_at: float, items_total: int
+    ) -> list[modal_proto.api_pb2.FunctionPutOutputsItem]: ...
+    async def output_items(
+        self, started_at: float, data: list[typing.Any]
+    ) -> list[modal_proto.api_pb2.FunctionPutOutputsItem]: ...
 
 class InputSlots:
+    """A semaphore that allows dynamically adjusting the concurrency."""
+
     active: int
     value: int
     waiter: typing.Optional[asyncio.Future]
     closed: bool
 
-    def __init__(self, value: int) -> None: ...
+    def __init__(self, value: int) -> None:
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
+
     async def acquire(self) -> None: ...
     def _wake_waiter(self) -> None: ...
     def release(self) -> None: ...
@@ -59,11 +95,18 @@ class InputSlots:
     async def close(self) -> None: ...
 
 class _ContainerIOManager:
+    """Synchronizes all RPC calls and network operations for a running container.
+
+    TODO: maybe we shouldn't synchronize the whole class.
+    Then we could potentially move a bunch of the global functions onto it.
+    """
+
     task_id: str
     function_id: str
     app_id: str
     function_def: modal_proto.api_pb2.Function
     checkpoint_id: typing.Optional[str]
+    input_plane_server_url: typing.Optional[str]
     calls_completed: int
     total_user_time: float
     current_input_id: typing.Optional[str]
@@ -81,7 +124,6 @@ class _ContainerIOManager:
     _is_interactivity_enabled: bool
     _fetching_inputs: bool
     _client: modal.client._Client
-    _GENERATOR_STOP_SENTINEL: typing.ClassVar[Sentinel]
     _singleton: typing.ClassVar[typing.Optional[_ContainerIOManager]]
 
     def _init(self, container_args: modal_proto.api_pb2.ContainerArguments, client: modal.client._Client): ...
@@ -90,9 +132,15 @@ class _ContainerIOManager:
     @staticmethod
     def __new__(
         cls, container_args: modal_proto.api_pb2.ContainerArguments, client: modal.client._Client
-    ) -> _ContainerIOManager: ...
+    ) -> _ContainerIOManager:
+        """Create and return a new object.  See help(type) for accurate signature."""
+        ...
+
     @classmethod
-    def _reset_singleton(cls): ...
+    def _reset_singleton(cls):
+        """Only used for tests."""
+        ...
+
     async def hello(self): ...
     async def _run_heartbeat_loop(self): ...
     async def _heartbeat_handle_cancellations(self) -> bool: ...
@@ -100,46 +148,82 @@ class _ContainerIOManager:
     def stop_heartbeat(self): ...
     def dynamic_concurrency_manager(self) -> typing.AsyncContextManager[None]: ...
     async def _dynamic_concurrency_loop(self): ...
-    def serialize_data_format(self, obj: typing.Any, data_format: int) -> bytes: ...
-    async def format_blob_data(self, data: bytes) -> dict[str, typing.Any]: ...
-    def get_data_in(self, function_call_id: str) -> collections.abc.AsyncIterator[typing.Any]: ...
+    def get_data_in(
+        self, function_call_id: str, attempt_token: typing.Optional[str]
+    ) -> collections.abc.AsyncIterator[typing.Any]:
+        """Read from the `data_in` stream of a function call."""
+        ...
+
     async def put_data_out(
-        self, function_call_id: str, start_index: int, data_format: int, serialized_messages: list[typing.Any]
-    ) -> None: ...
-    async def generator_output_task(
-        self, function_call_id: str, data_format: int, message_rx: asyncio.queues.Queue
-    ) -> None: ...
-    async def _queue_create(self, size: int) -> asyncio.queues.Queue: ...
-    async def _queue_put(self, queue: asyncio.queues.Queue, value: typing.Any) -> None: ...
+        self,
+        function_call_id: str,
+        attempt_token: str,
+        start_index: int,
+        data_format: int,
+        serialized_messages: list[typing.Any],
+    ) -> None:
+        """Put data onto the `data_out` stream of a function call.
+
+        This is used for generator outputs, which includes web endpoint responses. Note that this
+        was introduced as a performance optimization in client version 0.57, so older clients will
+        still use the previous Postgres-backed system based on `FunctionPutOutputs()`.
+        """
+        ...
+
+    def generator_output_sender(
+        self, function_call_id: str, attempt_token: str, data_format: int, message_rx: asyncio.queues.Queue
+    ) -> typing.AsyncContextManager[None]:
+        """Runs background task that feeds generator outputs into a function call's `data_out` stream."""
+        ...
+
+    async def _queue_create(self, size: int) -> asyncio.queues.Queue:
+        """Create a queue, on the synchronicity event loop (needed on Python 3.8 and 3.9)."""
+        ...
+
+    async def _queue_put(self, queue: asyncio.queues.Queue, value: typing.Any) -> None:
+        """Put a value onto a queue, using the synchronicity event loop."""
+        ...
+
     def get_average_call_time(self) -> float: ...
     def get_max_inputs_to_fetch(self): ...
     def _generate_inputs(
         self, batch_max_size: int, batch_wait_ms: int
-    ) -> collections.abc.AsyncIterator[list[tuple[str, int, str, modal_proto.api_pb2.FunctionInput]]]: ...
+    ) -> collections.abc.AsyncIterator[list[tuple[str, int, str, str, modal_proto.api_pb2.FunctionInput]]]: ...
     def run_inputs_outputs(
         self,
         finalized_functions: dict[str, modal._runtime.user_code_imports.FinalizedFunction],
         batch_max_size: int = 0,
         batch_wait_ms: int = 0,
     ) -> collections.abc.AsyncIterator[IOContext]: ...
-    async def _push_outputs(
-        self,
-        io_context: IOContext,
-        started_at: float,
-        data_format: int,
-        results: list[modal_proto.api_pb2.GenericResult],
-    ) -> None: ...
-    def serialize_exception(self, exc: BaseException) -> bytes: ...
-    def serialize_traceback(self, exc: BaseException) -> tuple[typing.Optional[bytes], typing.Optional[bytes]]: ...
-    def handle_user_exception(self) -> typing.AsyncContextManager[None]: ...
-    def handle_input_exception(self, io_context: IOContext, started_at: float) -> typing.AsyncContextManager[None]: ...
+    async def _send_outputs(self, started_at: float, outputs: list[modal_proto.api_pb2.FunctionPutOutputsItem]) -> None:
+        """Send pre-built output items with retry and chunking."""
+        ...
+
+    def handle_user_exception(self) -> typing.AsyncContextManager[None]:
+        """Sets the task as failed in a way where it's not retried.
+
+        Used for handling exceptions from container lifecycle methods at the moment, which should
+        trigger a task failure state.
+        """
+        ...
+
+    def handle_input_exception(self, io_context: IOContext, started_at: float) -> typing.AsyncContextManager[None]:
+        """Handle an exception while processing a function input."""
+        ...
+
     def exit_context(self, started_at, input_ids: list[str]): ...
-    async def push_outputs(
-        self, io_context: IOContext, started_at: float, data: typing.Any, data_format: int
-    ) -> None: ...
+    async def push_outputs(self, io_context: IOContext, started_at: float, output_data: list[typing.Any]) -> None: ...
     async def memory_restore(self) -> None: ...
-    async def memory_snapshot(self) -> None: ...
-    async def volume_commit(self, volume_ids: list[str]) -> None: ...
+    async def memory_snapshot(self) -> None:
+        """Message server indicating that function is ready to be checkpointed."""
+        ...
+
+    async def volume_commit(self, volume_ids: list[str]) -> None:
+        """Perform volume commit for given `volume_ids`.
+        Only used on container exit to persist uncommitted changes on behalf of user.
+        """
+        ...
+
     async def interact(self, from_breakpoint: bool = False): ...
     @property
     def target_concurrency(self) -> int: ...
@@ -148,20 +232,41 @@ class _ContainerIOManager:
     @property
     def input_concurrency_enabled(self) -> int: ...
     @classmethod
-    def get_input_concurrency(cls) -> int: ...
+    def get_input_concurrency(cls) -> int:
+        """Returns the number of usable input slots.
+
+        If concurrency is reduced, active slots can exceed allotted slots. Returns the larger value
+        in this case.
+        """
+        ...
+
     @classmethod
-    def set_input_concurrency(cls, concurrency: int): ...
+    def set_input_concurrency(cls, concurrency: int):
+        """Edit the number of input slots.
+
+        This disables the background loop which automatically adjusts concurrency
+        within [target_concurrency, max_concurrency].
+        """
+        ...
+
     @classmethod
     def stop_fetching_inputs(cls): ...
 
 SUPERSELF = typing.TypeVar("SUPERSELF", covariant=True)
 
 class ContainerIOManager:
+    """Synchronizes all RPC calls and network operations for a running container.
+
+    TODO: maybe we shouldn't synchronize the whole class.
+    Then we could potentially move a bunch of the global functions onto it.
+    """
+
     task_id: str
     function_id: str
     app_id: str
     function_def: modal_proto.api_pb2.Function
     checkpoint_id: typing.Optional[str]
+    input_plane_server_url: typing.Optional[str]
     calls_completed: int
     total_user_time: float
     current_input_id: typing.Optional[str]
@@ -179,15 +284,19 @@ class ContainerIOManager:
     _is_interactivity_enabled: bool
     _fetching_inputs: bool
     _client: modal.client.Client
-    _GENERATOR_STOP_SENTINEL: typing.ClassVar[Sentinel]
     _singleton: typing.ClassVar[typing.Optional[ContainerIOManager]]
 
-    def __init__(self, /, *args, **kwargs): ...
+    def __init__(self, /, *args, **kwargs):
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
+
     def _init(self, container_args: modal_proto.api_pb2.ContainerArguments, client: modal.client.Client): ...
     @property
     def heartbeat_condition(self) -> asyncio.locks.Condition: ...
     @classmethod
-    def _reset_singleton(cls): ...
+    def _reset_singleton(cls):
+        """Only used for tests."""
+        ...
 
     class __hello_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, /): ...
@@ -229,45 +338,92 @@ class ContainerIOManager:
 
     _dynamic_concurrency_loop: ___dynamic_concurrency_loop_spec[typing_extensions.Self]
 
-    def serialize_data_format(self, obj: typing.Any, data_format: int) -> bytes: ...
-
-    class __format_blob_data_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(self, /, data: bytes) -> dict[str, typing.Any]: ...
-        async def aio(self, /, data: bytes) -> dict[str, typing.Any]: ...
-
-    format_blob_data: __format_blob_data_spec[typing_extensions.Self]
-
     class __get_data_in_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(self, /, function_call_id: str) -> typing.Iterator[typing.Any]: ...
-        def aio(self, /, function_call_id: str) -> collections.abc.AsyncIterator[typing.Any]: ...
+        def __call__(
+            self, /, function_call_id: str, attempt_token: typing.Optional[str]
+        ) -> typing.Iterator[typing.Any]:
+            """Read from the `data_in` stream of a function call."""
+            ...
+
+        def aio(
+            self, /, function_call_id: str, attempt_token: typing.Optional[str]
+        ) -> collections.abc.AsyncIterator[typing.Any]:
+            """Read from the `data_in` stream of a function call."""
+            ...
 
     get_data_in: __get_data_in_spec[typing_extensions.Self]
 
     class __put_data_out_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
-            self, /, function_call_id: str, start_index: int, data_format: int, serialized_messages: list[typing.Any]
-        ) -> None: ...
+            self,
+            /,
+            function_call_id: str,
+            attempt_token: str,
+            start_index: int,
+            data_format: int,
+            serialized_messages: list[typing.Any],
+        ) -> None:
+            """Put data onto the `data_out` stream of a function call.
+
+            This is used for generator outputs, which includes web endpoint responses. Note that this
+            was introduced as a performance optimization in client version 0.57, so older clients will
+            still use the previous Postgres-backed system based on `FunctionPutOutputs()`.
+            """
+            ...
+
         async def aio(
-            self, /, function_call_id: str, start_index: int, data_format: int, serialized_messages: list[typing.Any]
-        ) -> None: ...
+            self,
+            /,
+            function_call_id: str,
+            attempt_token: str,
+            start_index: int,
+            data_format: int,
+            serialized_messages: list[typing.Any],
+        ) -> None:
+            """Put data onto the `data_out` stream of a function call.
+
+            This is used for generator outputs, which includes web endpoint responses. Note that this
+            was introduced as a performance optimization in client version 0.57, so older clients will
+            still use the previous Postgres-backed system based on `FunctionPutOutputs()`.
+            """
+            ...
 
     put_data_out: __put_data_out_spec[typing_extensions.Self]
 
-    class __generator_output_task_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(self, /, function_call_id: str, data_format: int, message_rx: asyncio.queues.Queue) -> None: ...
-        async def aio(self, /, function_call_id: str, data_format: int, message_rx: asyncio.queues.Queue) -> None: ...
+    class __generator_output_sender_spec(typing_extensions.Protocol[SUPERSELF]):
+        def __call__(
+            self, /, function_call_id: str, attempt_token: str, data_format: int, message_rx: asyncio.queues.Queue
+        ) -> synchronicity.combined_types.AsyncAndBlockingContextManager[None]:
+            """Runs background task that feeds generator outputs into a function call's `data_out` stream."""
+            ...
 
-    generator_output_task: __generator_output_task_spec[typing_extensions.Self]
+        def aio(
+            self, /, function_call_id: str, attempt_token: str, data_format: int, message_rx: asyncio.queues.Queue
+        ) -> typing.AsyncContextManager[None]:
+            """Runs background task that feeds generator outputs into a function call's `data_out` stream."""
+            ...
+
+    generator_output_sender: __generator_output_sender_spec[typing_extensions.Self]
 
     class ___queue_create_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(self, /, size: int) -> asyncio.queues.Queue: ...
-        async def aio(self, /, size: int) -> asyncio.queues.Queue: ...
+        def __call__(self, /, size: int) -> asyncio.queues.Queue:
+            """Create a queue, on the synchronicity event loop (needed on Python 3.8 and 3.9)."""
+            ...
+
+        async def aio(self, /, size: int) -> asyncio.queues.Queue:
+            """Create a queue, on the synchronicity event loop (needed on Python 3.8 and 3.9)."""
+            ...
 
     _queue_create: ___queue_create_spec[typing_extensions.Self]
 
     class ___queue_put_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(self, /, queue: asyncio.queues.Queue, value: typing.Any) -> None: ...
-        async def aio(self, /, queue: asyncio.queues.Queue, value: typing.Any) -> None: ...
+        def __call__(self, /, queue: asyncio.queues.Queue, value: typing.Any) -> None:
+            """Put a value onto a queue, using the synchronicity event loop."""
+            ...
+
+        async def aio(self, /, queue: asyncio.queues.Queue, value: typing.Any) -> None:
+            """Put a value onto a queue, using the synchronicity event loop."""
+            ...
 
     _queue_put: ___queue_put_spec[typing_extensions.Self]
 
@@ -277,10 +433,10 @@ class ContainerIOManager:
     class ___generate_inputs_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self, /, batch_max_size: int, batch_wait_ms: int
-        ) -> typing.Iterator[list[tuple[str, int, str, modal_proto.api_pb2.FunctionInput]]]: ...
+        ) -> typing.Iterator[list[tuple[str, int, str, str, modal_proto.api_pb2.FunctionInput]]]: ...
         def aio(
             self, /, batch_max_size: int, batch_wait_ms: int
-        ) -> collections.abc.AsyncIterator[list[tuple[str, int, str, modal_proto.api_pb2.FunctionInput]]]: ...
+        ) -> collections.abc.AsyncIterator[list[tuple[str, int, str, str, modal_proto.api_pb2.FunctionInput]]]: ...
 
     _generate_inputs: ___generate_inputs_spec[typing_extensions.Self]
 
@@ -302,50 +458,54 @@ class ContainerIOManager:
 
     run_inputs_outputs: __run_inputs_outputs_spec[typing_extensions.Self]
 
-    class ___push_outputs_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(
-            self,
-            /,
-            io_context: IOContext,
-            started_at: float,
-            data_format: int,
-            results: list[modal_proto.api_pb2.GenericResult],
-        ) -> None: ...
-        async def aio(
-            self,
-            /,
-            io_context: IOContext,
-            started_at: float,
-            data_format: int,
-            results: list[modal_proto.api_pb2.GenericResult],
-        ) -> None: ...
+    class ___send_outputs_spec(typing_extensions.Protocol[SUPERSELF]):
+        def __call__(self, /, started_at: float, outputs: list[modal_proto.api_pb2.FunctionPutOutputsItem]) -> None:
+            """Send pre-built output items with retry and chunking."""
+            ...
 
-    _push_outputs: ___push_outputs_spec[typing_extensions.Self]
+        async def aio(self, /, started_at: float, outputs: list[modal_proto.api_pb2.FunctionPutOutputsItem]) -> None:
+            """Send pre-built output items with retry and chunking."""
+            ...
 
-    def serialize_exception(self, exc: BaseException) -> bytes: ...
-    def serialize_traceback(self, exc: BaseException) -> tuple[typing.Optional[bytes], typing.Optional[bytes]]: ...
+    _send_outputs: ___send_outputs_spec[typing_extensions.Self]
 
     class __handle_user_exception_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(self, /) -> synchronicity.combined_types.AsyncAndBlockingContextManager[None]: ...
-        def aio(self, /) -> typing.AsyncContextManager[None]: ...
+        def __call__(self, /) -> synchronicity.combined_types.AsyncAndBlockingContextManager[None]:
+            """Sets the task as failed in a way where it's not retried.
+
+            Used for handling exceptions from container lifecycle methods at the moment, which should
+            trigger a task failure state.
+            """
+            ...
+
+        def aio(self, /) -> typing.AsyncContextManager[None]:
+            """Sets the task as failed in a way where it's not retried.
+
+            Used for handling exceptions from container lifecycle methods at the moment, which should
+            trigger a task failure state.
+            """
+            ...
 
     handle_user_exception: __handle_user_exception_spec[typing_extensions.Self]
 
     class __handle_input_exception_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self, /, io_context: IOContext, started_at: float
-        ) -> synchronicity.combined_types.AsyncAndBlockingContextManager[None]: ...
-        def aio(self, /, io_context: IOContext, started_at: float) -> typing.AsyncContextManager[None]: ...
+        ) -> synchronicity.combined_types.AsyncAndBlockingContextManager[None]:
+            """Handle an exception while processing a function input."""
+            ...
+
+        def aio(self, /, io_context: IOContext, started_at: float) -> typing.AsyncContextManager[None]:
+            """Handle an exception while processing a function input."""
+            ...
 
     handle_input_exception: __handle_input_exception_spec[typing_extensions.Self]
 
     def exit_context(self, started_at, input_ids: list[str]): ...
 
     class __push_outputs_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(self, /, io_context: IOContext, started_at: float, data: typing.Any, data_format: int) -> None: ...
-        async def aio(
-            self, /, io_context: IOContext, started_at: float, data: typing.Any, data_format: int
-        ) -> None: ...
+        def __call__(self, /, io_context: IOContext, started_at: float, output_data: list[typing.Any]) -> None: ...
+        async def aio(self, /, io_context: IOContext, started_at: float, output_data: list[typing.Any]) -> None: ...
 
     push_outputs: __push_outputs_spec[typing_extensions.Self]
 
@@ -356,14 +516,28 @@ class ContainerIOManager:
     memory_restore: __memory_restore_spec[typing_extensions.Self]
 
     class __memory_snapshot_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(self, /) -> None: ...
-        async def aio(self, /) -> None: ...
+        def __call__(self, /) -> None:
+            """Message server indicating that function is ready to be checkpointed."""
+            ...
+
+        async def aio(self, /) -> None:
+            """Message server indicating that function is ready to be checkpointed."""
+            ...
 
     memory_snapshot: __memory_snapshot_spec[typing_extensions.Self]
 
     class __volume_commit_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(self, /, volume_ids: list[str]) -> None: ...
-        async def aio(self, /, volume_ids: list[str]) -> None: ...
+        def __call__(self, /, volume_ids: list[str]) -> None:
+            """Perform volume commit for given `volume_ids`.
+            Only used on container exit to persist uncommitted changes on behalf of user.
+            """
+            ...
+
+        async def aio(self, /, volume_ids: list[str]) -> None:
+            """Perform volume commit for given `volume_ids`.
+            Only used on container exit to persist uncommitted changes on behalf of user.
+            """
+            ...
 
     volume_commit: __volume_commit_spec[typing_extensions.Self]
 
@@ -380,13 +554,34 @@ class ContainerIOManager:
     @property
     def input_concurrency_enabled(self) -> int: ...
     @classmethod
-    def get_input_concurrency(cls) -> int: ...
+    def get_input_concurrency(cls) -> int:
+        """Returns the number of usable input slots.
+
+        If concurrency is reduced, active slots can exceed allotted slots. Returns the larger value
+        in this case.
+        """
+        ...
+
     @classmethod
-    def set_input_concurrency(cls, concurrency: int): ...
+    def set_input_concurrency(cls, concurrency: int):
+        """Edit the number of input slots.
+
+        This disables the background loop which automatically adjusts concurrency
+        within [target_concurrency, max_concurrency].
+        """
+        ...
+
     @classmethod
     def stop_fetching_inputs(cls): ...
 
-def check_fastapi_pydantic_compatibility(exc: ImportError) -> None: ...
+def check_fastapi_pydantic_compatibility(exc: ImportError) -> None:
+    """Add a helpful note to an exception that is likely caused by a pydantic<>fastapi version incompatibility.
+
+    We need this becasue the legacy set of container requirements (image_builder_version=2023.12) contains a
+    version of fastapi that is not forwards-compatible with pydantic 2.0+, and users commonly run into issues
+    building an image that specifies a more recent version only for pydantic.
+    """
+    ...
 
 MAX_OUTPUT_BATCH_SIZE: int
 

modal/_runtime/execution_context.py

@@ -72,22 +72,38 @@ def current_function_call_id() -> Optional[str]:
         return None
 
 
-def _set_current_context_ids(input_ids: list[str], function_call_ids: list[str]) -> Callable[[], None]:
-    assert len(input_ids) == len(function_call_ids) and len(input_ids) > 0
+def current_attempt_token() -> Optional[str]:
+    # This ContextVar isn't useful to expose to users.
+    try:
+        return _current_attempt_token.get()
+    except LookupError:
+        return None
+
+
+def _set_current_context_ids(
+    input_ids: list[str], function_call_ids: list[str], attempt_tokens: list[str]
+) -> Callable[[], None]:
+    assert len(input_ids) == len(function_call_ids) == len(attempt_tokens) and input_ids
+
     input_id = input_ids[0]
     function_call_id = function_call_ids[0]
+    attempt_token = attempt_tokens[0]
+
     input_token = _current_input_id.set(input_id)
     function_call_token = _current_function_call_id.set(function_call_id)
+    attempt_token_token = _current_attempt_token.set(attempt_token)
 
     def _reset_current_context_ids():
         _current_input_id.reset(input_token)
         _current_function_call_id.reset(function_call_token)
+        _current_attempt_token.reset(attempt_token_token)
 
     return _reset_current_context_ids
 
 
 _current_input_id: ContextVar = ContextVar("_current_input_id")
 _current_function_call_id: ContextVar = ContextVar("_current_function_call_id")
+_current_attempt_token: ContextVar = ContextVar("_current_attempt_token")
 
 _is_currently_importing = False  # we set this to True while a container is importing user code