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

modal/parallel_map.pyi

@@ -1,9 +1,11 @@
 import asyncio
 import asyncio.events
+import asyncio.queues
 import collections.abc
 import enum
 import modal._functions
 import modal._utils.async_utils
+import modal._utils.grpc_utils
 import modal.client
 import modal.functions
 import modal.retries
@@ -12,6 +14,7 @@ import typing
 import typing_extensions
 
 class _SynchronizedQueue:
+    """mdmd:hidden"""
     async def init(self): ...
     async def put(self, item): ...
     async def get(self): ...
@@ -19,7 +22,10 @@ class _SynchronizedQueue:
 SUPERSELF = typing.TypeVar("SUPERSELF", covariant=True)
 
 class SynchronizedQueue:
-    def __init__(self, /, *args, **kwargs): ...
+    """mdmd:hidden"""
+    def __init__(self, /, *args, **kwargs):
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
 
     class __init_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, /): ...
@@ -40,34 +46,156 @@ class SynchronizedQueue:
     get: __get_spec[typing_extensions.Self]
 
 class _OutputValue:
+    """_OutputValue(value: Any)"""
+
     value: typing.Any
 
-    def __init__(self, value: typing.Any) -> None: ...
-    def __repr__(self): ...
-    def __eq__(self, other): ...
+    def __init__(self, value: typing.Any) -> None:
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
+
+    def __repr__(self):
+        """Return repr(self)."""
+        ...
+
+    def __eq__(self, other):
+        """Return self==value."""
+        ...
+
+class InputPreprocessor:
+    """Constructs FunctionPutInputsItem objects from the raw-input queue, and puts them in the processed-input queue."""
+    def __init__(
+        self,
+        client: modal.client._Client,
+        *,
+        raw_input_queue: _SynchronizedQueue,
+        processed_input_queue: asyncio.queues.Queue,
+        function: modal._functions._Function,
+        created_callback: collections.abc.Callable[[int], None],
+        done_callback: collections.abc.Callable[[], None],
+    ):
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
+
+    def input_iter(self): ...
+    def create_input_factory(self): ...
+    def drain_input_generator(self): ...
+
+class InputPumper:
+    """Reads inputs from a queue of FunctionPutInputsItems, and sends them to the server."""
+    def __init__(
+        self,
+        client: modal.client._Client,
+        *,
+        input_queue: asyncio.queues.Queue,
+        function: modal._functions._Function,
+        function_call_id: str,
+        max_batch_size: int,
+        map_items_manager: typing.Optional[_MapItemsManager] = None,
+    ):
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
+
+    def pump_inputs(self): ...
+    @property
+    def _function_inputs_retry(self) -> modal._utils.grpc_utils.Retry: ...
+
+class SyncInputPumper(InputPumper):
+    """Reads inputs from a queue of FunctionPutInputsItems, and sends them to the server."""
+    def __init__(
+        self,
+        client: modal.client._Client,
+        *,
+        input_queue: asyncio.queues.Queue,
+        retry_queue: modal._utils.async_utils.TimestampPriorityQueue,
+        function: modal._functions._Function,
+        function_call_jwt: str,
+        function_call_id: str,
+        map_items_manager: _MapItemsManager,
+    ):
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
+
+    def retry_inputs(self): ...
+
+class AsyncInputPumper(InputPumper):
+    """Reads inputs from a queue of FunctionPutInputsItems, and sends them to the server."""
+    def __init__(
+        self,
+        client: modal.client._Client,
+        *,
+        input_queue: asyncio.queues.Queue,
+        function: modal._functions._Function,
+        function_call_id: str,
+    ):
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
 
+    def pump_inputs(self): ...
+
+async def _spawn_map_invocation(
+    function: modal._functions._Function, raw_input_queue: _SynchronizedQueue, client: modal.client._Client
+) -> tuple[str, int]: ...
 def _map_invocation(
     function: modal._functions._Function,
     raw_input_queue: _SynchronizedQueue,
     client: modal.client._Client,
     order_outputs: bool,
     return_exceptions: bool,
+    wrap_returned_exceptions: bool,
     count_update_callback: typing.Optional[collections.abc.Callable[[int, int], None]],
     function_call_invocation_type: int,
 ): ...
+def _map_invocation_inputplane(
+    function: modal._functions._Function,
+    raw_input_queue: _SynchronizedQueue,
+    client: modal.client._Client,
+    order_outputs: bool,
+    return_exceptions: bool,
+    wrap_returned_exceptions: bool,
+    count_update_callback: typing.Optional[collections.abc.Callable[[int, int], None]],
+) -> typing.AsyncGenerator[typing.Any, None]:
+    """Input-plane implementation of a function map invocation.
+
+    This is analogous to `_map_invocation`, but instead of the control-plane
+    `FunctionMap` / `FunctionPutInputs` / `FunctionGetOutputs` RPCs it speaks
+    the input-plane protocol consisting of `MapStartOrContinue`, `MapAwait`, and `MapCheckInputs`.
+    """
+    ...
+
 def _map_helper(
     self: modal.functions.Function,
     async_input_gen: typing.AsyncGenerator[typing.Any, None],
     kwargs={},
     order_outputs: bool = True,
     return_exceptions: bool = False,
-) -> typing.AsyncGenerator[typing.Any, None]: ...
+    wrap_returned_exceptions: bool = True,
+) -> typing.AsyncGenerator[typing.Any, None]:
+    """Core implementation that supports `_map_async()`, `_starmap_async()` and `_for_each_async()`.
+
+    Runs in an event loop on the main thread. Concurrently feeds new input to the input queue and yields available
+    outputs to the caller.
+
+    Note that since the iterator(s) can block, it's a bit opaque how often the event
+    loop decides to get a new input vs how often it will emit a new output.
+
+    We could make this explicit as an improvement or even let users decide what they
+    prefer: throughput (prioritize queueing inputs) or latency (prioritize yielding results)
+    """
+    ...
+
+def _maybe_warn_about_exceptions(func_name: str, return_exceptions: bool, wrap_returned_exceptions: bool): ...
+def _invoked_from_sync_wrapper() -> bool:
+    """Check whether the calling function was called from a sync wrapper."""
+    ...
+
 def _map_async(
     self: modal.functions.Function,
     *input_iterators: typing.Union[typing.Iterable[typing.Any], typing.AsyncIterable[typing.Any]],
     kwargs={},
     order_outputs: bool = True,
     return_exceptions: bool = False,
+    wrap_returned_exceptions: bool = True,
 ) -> typing.AsyncGenerator[typing.Any, None]: ...
 def _starmap_async(
     self,
@@ -78,14 +206,118 @@ def _starmap_async(
     kwargs={},
     order_outputs: bool = True,
     return_exceptions: bool = False,
+    wrap_returned_exceptions: bool = True,
 ) -> typing.AsyncIterable[typing.Any]: ...
 async def _for_each_async(self, *input_iterators, kwargs={}, ignore_exceptions: bool = False) -> None: ...
 def _map_sync(
-    self, *input_iterators, kwargs={}, order_outputs: bool = True, return_exceptions: bool = False
-) -> modal._utils.async_utils.AsyncOrSyncIterable: ...
-async def _spawn_map_async(self, *input_iterators, kwargs={}) -> None: ...
-def _spawn_map_sync(self, *input_iterators, kwargs={}) -> None: ...
-def _for_each_sync(self, *input_iterators, kwargs={}, ignore_exceptions: bool = False): ...
+    self,
+    *input_iterators,
+    kwargs={},
+    order_outputs: bool = True,
+    return_exceptions: bool = False,
+    wrap_returned_exceptions: bool = True,
+) -> modal._utils.async_utils.AsyncOrSyncIterable:
+    """Parallel map over a set of inputs.
+
+    Takes one iterator argument per argument in the function being mapped over.
+
+    Example:
+    ```python
+    @app.function()
+    def my_func(a):
+        return a ** 2
+
+
+    @app.local_entrypoint()
+    def main():
+        assert list(my_func.map([1, 2, 3, 4])) == [1, 4, 9, 16]
+    ```
+
+    If applied to a `app.function`, `map()` returns one result per input and the output order
+    is guaranteed to be the same as the input order. Set `order_outputs=False` to return results
+    in the order that they are completed instead.
+
+    `return_exceptions` can be used to treat exceptions as successful results:
+
+    ```python
+    @app.function()
+    def my_func(a):
+        if a == 2:
+            raise Exception("ohno")
+        return a ** 2
+
+
+    @app.local_entrypoint()
+    def main():
+        # [0, 1, UserCodeException(Exception('ohno'))]
+        print(list(my_func.map(range(3), return_exceptions=True)))
+    ```
+    """
+    ...
+
+async def _experimental_spawn_map_async(self, *input_iterators, kwargs={}) -> modal._functions._FunctionCall: ...
+async def _spawn_map_helper(
+    self: modal.functions.Function, async_input_gen, kwargs={}
+) -> modal._functions._FunctionCall: ...
+def _experimental_spawn_map_sync(self, *input_iterators, kwargs={}) -> modal._functions._FunctionCall:
+    """mdmd:hidden
+    Spawn parallel execution over a set of inputs, returning as soon as the inputs are created.
+
+    Unlike `modal.Function.map`, this method does not block on completion of the remote execution but
+    returns a `modal.FunctionCall` object that can be used to poll status and retrieve results later.
+
+    Takes one iterator argument per argument in the function being mapped over.
+
+    Example:
+    ```python
+    @app.function()
+    def my_func(a, b):
+        return a ** b
+
+
+    @app.local_entrypoint()
+    def main():
+        fc = my_func.spawn_map([1, 2], [3, 4])
+    ```
+    """
+    ...
+
+async def _spawn_map_async(self, *input_iterators, kwargs={}) -> None:
+    """This runs in an event loop on the main thread. It consumes inputs from the input iterators and creates async
+    function calls for each.
+    """
+    ...
+
+def _spawn_map_sync(self, *input_iterators, kwargs={}) -> None:
+    """Spawn parallel execution over a set of inputs, exiting as soon as the inputs are created (without waiting
+    for the map to complete).
+
+    Takes one iterator argument per argument in the function being mapped over.
+
+    Example:
+    ```python
+    @app.function()
+    def my_func(a):
+        return a ** 2
+
+
+    @app.local_entrypoint()
+    def main():
+        my_func.spawn_map([1, 2, 3, 4])
+    ```
+
+    Programmatic retrieval of results will be supported in a future update.
+    """
+    ...
+
+def _for_each_sync(self, *input_iterators, kwargs={}, ignore_exceptions: bool = False):
+    """Execute function for all inputs, ignoring outputs. Waits for completion of the inputs.
+
+    Convenient alias for `.map()` in cases where the function just needs to be called.
+    as the caller doesn't have to consume the generator to process the inputs.
+    """
+    ...
+
 def _starmap_sync(
     self,
     input_iterator: typing.Iterable[typing.Sequence[typing.Any]],
@@ -93,7 +325,25 @@ def _starmap_sync(
     kwargs={},
     order_outputs: bool = True,
     return_exceptions: bool = False,
-) -> modal._utils.async_utils.AsyncOrSyncIterable: ...
+    wrap_returned_exceptions: bool = True,
+) -> modal._utils.async_utils.AsyncOrSyncIterable:
+    """Like `map`, but spreads arguments over multiple function arguments.
+
+    Assumes every input is a sequence (e.g. a tuple).
+
+    Example:
+    ```python
+    @app.function()
+    def my_func(a, b):
+        return a + b
+
+
+    @app.local_entrypoint()
+    def main():
+        assert list(my_func.starmap([(1, 2), (3, 4)])) == [3, 7]
+    ```
+    """
+    ...
 
 class _MapItemState(enum.Enum):
     # The input is being sent the server with a PutInputs request, but the response has not been received yet.
@@ -130,7 +380,12 @@ class _MapItemContext:
         input: modal_proto.api_pb2.FunctionInput,
         retry_manager: modal.retries.RetryManager,
         sync_client_retries_enabled: bool,
-    ): ...
+        is_input_plane_instance: bool = False,
+    ):
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
+
+    def handle_map_start_or_continue_response(self, attempt_token: str): ...
     def handle_put_inputs_response(self, item: modal_proto.api_pb2.FunctionPutInputsResponseItem): ...
     async def handle_get_outputs_response(
         self,
@@ -138,9 +393,17 @@ class _MapItemContext:
         now_seconds: int,
         function_call_invocation_type: int,
         retry_queue: modal._utils.async_utils.TimestampPriorityQueue,
-    ) -> _OutputType: ...
+    ) -> _OutputType:
+        """Processes the output, and determines if it is complete or needs to be retried.
+
+        Return True if input state was changed to COMPLETE, otherwise False.
+        """
+        ...
+
     async def prepare_item_for_retry(self) -> modal_proto.api_pb2.FunctionRetryInputsItem: ...
+    def set_retry_policy(self, retry_policy: modal_proto.api_pb2.FunctionRetryPolicy): ...
     def handle_retry_response(self, input_jwt: str): ...
+    async def create_map_start_or_continue_item(self, idx: int) -> modal_proto.api_pb2.MapStartOrContinueItem: ...
 
 class _MapItemsManager:
     def __init__(
@@ -150,16 +413,32 @@ class _MapItemsManager:
         retry_queue: modal._utils.async_utils.TimestampPriorityQueue,
         sync_client_retries_enabled: bool,
         max_inputs_outstanding: int,
-    ): ...
+        is_input_plane_instance: bool = False,
+    ):
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
+
+    def set_retry_policy(self, retry_policy: modal_proto.api_pb2.FunctionRetryPolicy): ...
     async def add_items(self, items: list[modal_proto.api_pb2.FunctionPutInputsItem]): ...
+    async def add_items_inputplane(self, items: list[modal_proto.api_pb2.MapStartOrContinueItem]): ...
     async def prepare_items_for_retry(
         self, retriable_idxs: list[int]
     ) -> list[modal_proto.api_pb2.FunctionRetryInputsItem]: ...
-    def get_input_jwts_waiting_for_output(self) -> list[str]: ...
+    def update_items_retry_policy(self, retry_policy: modal_proto.api_pb2.FunctionRetryPolicy): ...
+    def get_input_jwts_waiting_for_output(self) -> list[str]:
+        """Returns a list of input_jwts for inputs that are waiting for output."""
+        ...
+
+    def get_input_idxs_waiting_for_output(self) -> list[tuple[int, str]]:
+        """Returns a list of input_idxs for inputs that are waiting for output."""
+        ...
+
     def _remove_item(self, item_idx: int): ...
     def get_item_context(self, item_idx: int) -> _MapItemContext: ...
+    def handle_put_continue_response(self, items: list[tuple[int, str]]): ...
     def handle_put_inputs_response(self, items: list[modal_proto.api_pb2.FunctionPutInputsResponseItem]): ...
     def handle_retry_response(self, input_jwts: list[str]): ...
+    async def handle_check_inputs_response(self, response: list[tuple[int, bool]]): ...
     async def handle_get_outputs_response(
         self, item: modal_proto.api_pb2.FunctionGetOutputsItem, now_seconds: int
     ) -> _OutputType: ...

modal/partial_function.py

@@ -3,7 +3,6 @@
 from ._partial_function import (
     _asgi_app,
     _batched,
-    _build,
     _concurrent,
     _enter,
     _exit,
@@ -25,7 +24,6 @@ fastapi_endpoint = synchronize_api(_fastapi_endpoint, target_module=__name__)
 asgi_app = synchronize_api(_asgi_app, target_module=__name__)
 wsgi_app = synchronize_api(_wsgi_app, target_module=__name__)
 web_server = synchronize_api(_web_server, target_module=__name__)
-build = synchronize_api(_build, target_module=__name__)
 enter = synchronize_api(_enter, target_module=__name__)
 exit = synchronize_api(_exit, target_module=__name__)
 batched = synchronize_api(_batched, target_module=__name__)