modal 0.62.115__py3-none-any.whl → 0.72.13__py3-none-any.whl

modal/parallel_map.py

@@ -1,15 +1,21 @@
 # Copyright Modal Labs 2024
 import asyncio
+import time
 import typing
 from dataclasses import dataclass
-from typing import Any, Callable, Dict, Optional, Set, Tuple
+from typing import Any, Callable, Optional
 
-from aiostream import pipe, stream
-from grpclib import Status
+from grpclib import GRPCError, Status
 
+from modal._runtime.execution_context import current_input_id
 from modal._utils.async_utils import (
-    AsyncOrSyncIteratable,
+    AsyncOrSyncIterable,
+    aclosing,
+    async_map_ordered,
+    async_merge,
+    async_zip,
     queue_batch_iterator,
+    sync_or_async_iter,
     synchronize_api,
     synchronizer,
     warn_if_generator_is_not_consumed,
@@ -23,7 +29,6 @@ from modal._utils.function_utils import (
 )
 from modal._utils.grpc_utils import retry_transient_errors
 from modal.config import logger
-from modal.execution_context import current_input_id
 from modal_proto import api_pb2
 
 if typing.TYPE_CHECKING:
@@ -59,9 +64,12 @@ class _OutputValue:
 
 MAP_INVOCATION_CHUNK_SIZE = 49
 
+if typing.TYPE_CHECKING:
+    import modal.functions
+
 
 async def _map_invocation(
-    function_id: str,
+    function: "modal.functions._Function",
     raw_input_queue: _SynchronizedQueue,
     client: "modal.client._Client",
     order_outputs: bool,
@@ -70,7 +78,7 @@ async def _map_invocation(
 ):
     assert client.stub
     request = api_pb2.FunctionMapRequest(
-        function_id=function_id,
+        function_id=function.object_id,
         parent_input_id=current_input_id() or "",
         function_call_type=api_pb2.FUNCTION_CALL_TYPE_MAP,
         return_exceptions=return_exceptions,
@@ -87,8 +95,8 @@ async def _map_invocation(
         if count_update_callback is not None:
             count_update_callback(num_outputs, num_inputs)
 
-    pending_outputs: Dict[str, int] = {}  # Map input_id -> next expected gen_index value
-    completed_outputs: Set[str] = set()  # Set of input_ids whose outputs are complete (expecting no more values)
+    pending_outputs: dict[str, int] = {}  # Map input_id -> next expected gen_index value
+    completed_outputs: set[str] = set()  # Set of input_ids whose outputs are complete (expecting no more values)
 
     input_queue: asyncio.Queue = asyncio.Queue()
 
@@ -97,23 +105,20 @@ async def _map_invocation(
         idx = num_inputs
         num_inputs += 1
         (args, kwargs) = argskwargs
-        return await _create_input(args, kwargs, client, idx)
+        return await _create_input(args, kwargs, client, idx=idx, method_name=function._use_method_name)
 
     async def input_iter():
         while 1:
             raw_input = await raw_input_queue.get()
             if raw_input is None:  # end of input sentinel
-                return
+                break
             yield raw_input  # args, kwargs
 
     async def drain_input_generator():
         # Parallelize uploading blobs
-        proto_input_stream = stream.iterate(input_iter()) | pipe.map(
-            create_input,  # type: ignore[reportArgumentType]
-            ordered=True,
-            task_limit=BLOB_MAX_PARALLELISM,
-        )
-        async with proto_input_stream.stream() as streamer:
+        async with aclosing(
+            async_map_ordered(input_iter(), create_input, concurrency=BLOB_MAX_PARALLELISM)
+        ) as streamer:
             async for item in streamer:
                 await input_queue.put(item)
 
@@ -126,23 +131,36 @@ async def _map_invocation(
         nonlocal have_all_inputs, num_inputs
         async for items in queue_batch_iterator(input_queue, MAP_INVOCATION_CHUNK_SIZE):
             request = api_pb2.FunctionPutInputsRequest(
-                function_id=function_id, inputs=items, function_call_id=function_call_id
+                function_id=function.object_id, inputs=items, function_call_id=function_call_id
             )
             logger.debug(
                 f"Pushing {len(items)} inputs to server. Num queued inputs awaiting push is {input_queue.qsize()}."
             )
-            resp = await retry_transient_errors(
-                client.stub.FunctionPutInputs,
-                request,
-                max_retries=None,
-                max_delay=10,
-                additional_status_codes=[Status.RESOURCE_EXHAUSTED],
-            )
+            while True:
+                try:
+                    resp = await retry_transient_errors(
+                        client.stub.FunctionPutInputs,
+                        request,
+                        # with 8 retries we log the warning below about every 30 secondswhich isn't too spammy.
+                        max_retries=8,
+                        max_delay=15,
+                        additional_status_codes=[Status.RESOURCE_EXHAUSTED],
+                    )
+                    break
+                except GRPCError as err:
+                    if err.status != Status.RESOURCE_EXHAUSTED:
+                        raise err
+                    logger.warning(
+                        "Warning: map progress is limited. Common bottlenecks "
+                        "include slow iteration over results, or function backlogs."
+                    )
+
             count_update()
             for item in resp.inputs:
                 pending_outputs.setdefault(item.input_id, 0)
             logger.debug(
-                f"Successfully pushed {len(items)} inputs to server. Num queued inputs awaiting push is {input_queue.qsize()}."
+                f"Successfully pushed {len(items)} inputs to server. "
+                f"Num queued inputs awaiting push is {input_queue.qsize()}."
             )
 
         have_all_inputs = True
@@ -158,6 +176,7 @@ async def _map_invocation(
                 timeout=OUTPUTS_TIMEOUT,
                 last_entry_id=last_entry_id,
                 clear_on_success=False,
+                requested_at=time.time(),
             )
             response = await retry_transient_errors(
                 client.stub.FunctionGetOutputs,
@@ -183,8 +202,9 @@ async def _map_invocation(
     async def get_all_outputs_and_clean_up():
         assert client.stub
         try:
-            async for item in get_all_outputs():
-                yield item
+            async with aclosing(get_all_outputs()) as output_items:
+                async for item in output_items:
+                    yield item
         finally:
             # "ack" that we have all outputs we are interested in and let backend clear results
             request = api_pb2.FunctionGetOutputsRequest(
@@ -192,10 +212,11 @@ async def _map_invocation(
                 timeout=0,
                 last_entry_id="0-0",
                 clear_on_success=True,
+                requested_at=time.time(),
             )
             await retry_transient_errors(client.stub.FunctionGetOutputs, request)
 
-    async def fetch_output(item: api_pb2.FunctionGetOutputsItem) -> Tuple[int, Any]:
+    async def fetch_output(item: api_pb2.FunctionGetOutputsItem) -> tuple[int, Any]:
         try:
             output = await _process_result(item.result, item.data_format, client.stub, client)
         except Exception as e:
@@ -206,14 +227,13 @@ async def _map_invocation(
         return (item.idx, output)
 
     async def poll_outputs():
-        outputs = stream.iterate(get_all_outputs_and_clean_up())
-        outputs_fetched = outputs | pipe.map(fetch_output, ordered=True, task_limit=BLOB_MAX_PARALLELISM)  # type: ignore
-
         # map to store out-of-order outputs received
         received_outputs = {}
         output_idx = 0
 
-        async with outputs_fetched.stream() as streamer:
+        async with aclosing(
+            async_map_ordered(get_all_outputs_and_clean_up(), fetch_output, concurrency=BLOB_MAX_PARALLELISM)
+        ) as streamer:
             async for idx, output in streamer:
                 count_update()
                 if not order_outputs:
@@ -228,9 +248,7 @@ async def _map_invocation(
 
         assert len(received_outputs) == 0
 
-    response_gen = stream.merge(drain_input_generator(), pump_inputs(), poll_outputs())
-
-    async with response_gen.stream() as streamer:
+    async with aclosing(async_merge(drain_input_generator(), pump_inputs(), poll_outputs())) as streamer:
         async for response in streamer:
             if response is not None:
                 yield response.value
@@ -243,7 +261,7 @@ def _map_sync(
     kwargs={},  # any extra keyword arguments for the function
     order_outputs: bool = True,  # return outputs in order
     return_exceptions: bool = False,  # propagate exceptions (False) or aggregate them in the results list (True)
-) -> AsyncOrSyncIteratable:
+) -> AsyncOrSyncIterable:
     """Parallel map over a set of inputs.
 
     Takes one iterator argument per argument in the function being mapped over.
@@ -281,11 +299,14 @@ def _map_sync(
     ```
     """
 
-    return AsyncOrSyncIteratable(
+    return AsyncOrSyncIterable(
         _map_async(
             self, *input_iterators, kwargs=kwargs, order_outputs=order_outputs, return_exceptions=return_exceptions
         ),
-        nested_async_message="You can't iter(Function.map()) or Function.for_each() from an async function. Use async for ... Function.map.aio() or Function.for_each.aio() instead.",
+        nested_async_message=(
+            "You can't iter(Function.map()) or Function.for_each() from an async function. "
+            "Use async for ... Function.map.aio() or Function.for_each.aio() instead."
+        ),
     )
 
 
@@ -314,7 +335,7 @@ async def _map_async(
 
     async def feed_queue():
         # This runs in a main thread event loop, so it doesn't block the synchronizer loop
-        async with stream.zip(*[stream.iterate(it) for it in input_iterators]).stream() as streamer:
+        async with aclosing(async_zip(*[sync_or_async_iter(it) for it in input_iterators])) as streamer:
             async for args in streamer:
                 await raw_input_queue.put.aio((args, kwargs))
         await raw_input_queue.put.aio(None)  # end-of-input sentinel
@@ -322,11 +343,13 @@ async def _map_async(
     feed_input_task = asyncio.create_task(feed_queue())
 
     try:
-        # note that `map()` and `map.aio()` are not synchronicity-wrapped, since they accept executable code in the form of
+        # note that `map()` and `map.aio()` are not synchronicity-wrapped, since
+        # they accept executable code in the form of
         # iterators that we don't want to run inside the synchronicity thread.
         # Instead, we delegate to `._map()` with a safer Queue as input
-        async for output in self._map.aio(raw_input_queue, order_outputs, return_exceptions):  # type: ignore[reportFunctionMemberAccess]
-            yield output
+        async with aclosing(self._map.aio(raw_input_queue, order_outputs, return_exceptions)) as map_output_stream:
+            async for output in map_output_stream:
+                yield output
     finally:
         feed_input_task.cancel()  # should only be needed in case of exceptions
 
@@ -363,7 +386,7 @@ async def _starmap_async(
 
     async def feed_queue():
         # This runs in a main thread event loop, so it doesn't block the synchronizer loop
-        async with stream.iterate(input_iterator).stream() as streamer:
+        async with aclosing(sync_or_async_iter(input_iterator)) as streamer:
             async for args in streamer:
                 await raw_input_queue.put.aio((args, kwargs))
         await raw_input_queue.put.aio(None)  # end-of-input sentinel
@@ -383,7 +406,7 @@ def _starmap_sync(
     kwargs={},
     order_outputs: bool = True,
     return_exceptions: bool = False,
-) -> AsyncOrSyncIteratable:
+) -> AsyncOrSyncIterable:
     """Like `map`, but spreads arguments over multiple function arguments.
 
     Assumes every input is a sequence (e.g. a tuple).
@@ -400,9 +423,12 @@ def _starmap_sync(
         assert list(my_func.starmap([(1, 2), (3, 4)])) == [3, 7]
     ```
     """
-    return AsyncOrSyncIteratable(
+    return AsyncOrSyncIterable(
         _starmap_async(
             self, input_iterator, kwargs=kwargs, order_outputs=order_outputs, return_exceptions=return_exceptions
         ),
-        nested_async_message="You can't run Function.map() or Function.for_each() from an async function. Use Function.map.aio()/Function.for_each.aio() instead.",
+        nested_async_message=(
+            "You can't run Function.map() or Function.for_each() from an async function. "
+            "Use Function.map.aio()/Function.for_each.aio() instead."
+        ),
     )

modal/parallel_map.pyi

@@ -1,87 +1,75 @@
 import modal._utils.async_utils
 import modal.client
+import modal.functions
 import typing
 import typing_extensions
 
 class _SynchronizedQueue:
-    async def init(self):
-        ...
-
-    async def put(self, item):
-        ...
-
-    async def get(self):
-        ...
-
+    async def init(self): ...
+    async def put(self, item): ...
+    async def get(self): ...
 
 class SynchronizedQueue:
-    def __init__(self, /, *args, **kwargs):
-        ...
+    def __init__(self, /, *args, **kwargs): ...
 
     class __init_spec(typing_extensions.Protocol):
-        def __call__(self):
-            ...
-
-        async def aio(self, *args, **kwargs):
-            ...
+        def __call__(self): ...
+        async def aio(self): ...
 
     init: __init_spec
 
     class __put_spec(typing_extensions.Protocol):
-        def __call__(self, item):
-            ...
-
-        async def aio(self, *args, **kwargs):
-            ...
+        def __call__(self, item): ...
+        async def aio(self, item): ...
 
     put: __put_spec
 
     class __get_spec(typing_extensions.Protocol):
-        def __call__(self):
-            ...
-
-        async def aio(self, *args, **kwargs):
-            ...
+        def __call__(self): ...
+        async def aio(self): ...
 
     get: __get_spec
 
-
 class _OutputValue:
     value: typing.Any
 
-    def __init__(self, value: typing.Any) -> None:
-        ...
-
-    def __repr__(self):
-        ...
-
-    def __eq__(self, other):
-        ...
-
-
-def _map_invocation(function_id: str, raw_input_queue: _SynchronizedQueue, client: modal.client._Client, order_outputs: bool, return_exceptions: bool, count_update_callback: typing.Union[typing.Callable[[int, int], None], None]):
-    ...
-
-
-def _map_sync(self, *input_iterators, kwargs={}, order_outputs: bool = True, return_exceptions: bool = False) -> modal._utils.async_utils.AsyncOrSyncIteratable:
-    ...
-
-
-def _map_async(self, *input_iterators: typing.Union[typing.Iterable[typing.Any], typing.AsyncIterable[typing.Any]], kwargs={}, order_outputs: bool = True, return_exceptions: bool = False) -> typing.AsyncGenerator[typing.Any, None]:
-    ...
-
-
-def _for_each_sync(self, *input_iterators, kwargs={}, ignore_exceptions: bool = False):
-    ...
-
-
-async def _for_each_async(self, *input_iterators, kwargs={}, ignore_exceptions: bool = False):
-    ...
-
-
-def _starmap_async(self, input_iterator: typing.Union[typing.Iterable[typing.Sequence[typing.Any]], typing.AsyncIterable[typing.Sequence[typing.Any]]], kwargs={}, order_outputs: bool = True, return_exceptions: bool = False) -> typing.AsyncIterable[typing.Any]:
-    ...
-
-
-def _starmap_sync(self, input_iterator: typing.Iterable[typing.Sequence[typing.Any]], kwargs={}, order_outputs: bool = True, return_exceptions: bool = False) -> modal._utils.async_utils.AsyncOrSyncIteratable:
-    ...
+    def __init__(self, value: typing.Any) -> None: ...
+    def __repr__(self): ...
+    def __eq__(self, other): ...
+
+def _map_invocation(
+    function: modal.functions._Function,
+    raw_input_queue: _SynchronizedQueue,
+    client: modal.client._Client,
+    order_outputs: bool,
+    return_exceptions: bool,
+    count_update_callback: typing.Optional[typing.Callable[[int, int], None]],
+): ...
+def _map_sync(
+    self, *input_iterators, kwargs={}, order_outputs: bool = True, return_exceptions: bool = False
+) -> modal._utils.async_utils.AsyncOrSyncIterable: ...
+def _map_async(
+    self,
+    *input_iterators: typing.Union[typing.Iterable[typing.Any], typing.AsyncIterable[typing.Any]],
+    kwargs={},
+    order_outputs: bool = True,
+    return_exceptions: bool = False,
+) -> typing.AsyncGenerator[typing.Any, None]: ...
+def _for_each_sync(self, *input_iterators, kwargs={}, ignore_exceptions: bool = False): ...
+async def _for_each_async(self, *input_iterators, kwargs={}, ignore_exceptions: bool = False): ...
+def _starmap_async(
+    self,
+    input_iterator: typing.Union[
+        typing.Iterable[typing.Sequence[typing.Any]], typing.AsyncIterable[typing.Sequence[typing.Any]]
+    ],
+    kwargs={},
+    order_outputs: bool = True,
+    return_exceptions: bool = False,
+) -> typing.AsyncIterable[typing.Any]: ...
+def _starmap_sync(
+    self,
+    input_iterator: typing.Iterable[typing.Sequence[typing.Any]],
+    kwargs={},
+    order_outputs: bool = True,
+    return_exceptions: bool = False,
+) -> modal._utils.async_utils.AsyncOrSyncIterable: ...