modal 1.0.4.dev12__py3-none-any.whl → 1.0.5__py3-none-any.whl

modal/parallel_map.py

@@ -1,6 +1,7 @@
 # Copyright Modal Labs 2024
 import asyncio
 import enum
+import inspect
 import time
 import typing
 from asyncio import FIRST_COMPLETED
@@ -9,6 +10,7 @@ from typing import Any, Callable, Optional
 
 from grpclib import Status
 
+import modal.exception
 from modal._runtime.execution_context import current_input_id
 from modal._utils.async_utils import (
     AsyncOrSyncIterable,
@@ -26,6 +28,7 @@ from modal._utils.async_utils import (
     warn_if_generator_is_not_consumed,
 )
 from modal._utils.blob_utils import BLOB_MAX_PARALLELISM
+from modal._utils.deprecation import deprecation_warning
 from modal._utils.function_utils import (
     ATTEMPT_TIMEOUT_GRACE_PERIOD,
     OUTPUTS_TIMEOUT,
@@ -89,6 +92,7 @@ async def _map_invocation(
     client: "modal.client._Client",
     order_outputs: bool,
     return_exceptions: bool,
+    wrap_returned_exceptions: bool,
     count_update_callback: Optional[Callable[[int, int], None]],
     function_call_invocation_type: "api_pb2.FunctionCallInvocationType.ValueType",
 ):
@@ -278,17 +282,19 @@ async def _map_invocation(
                 )
             )
             map_done_task = asyncio.create_task(map_done_event.wait())
-            done, pending = await asyncio.wait([get_response_task, map_done_task], return_when=FIRST_COMPLETED)
-            if get_response_task in done:
-                map_done_task.cancel()
-                response = get_response_task.result()
-            else:
-                assert map_done_event.is_set()
-                # map is done, cancel the pending call
+            try:
+                done, pending = await asyncio.wait([get_response_task, map_done_task], return_when=FIRST_COMPLETED)
+                if get_response_task in done:
+                    map_done_task.cancel()
+                    response = get_response_task.result()
+                else:
+                    assert map_done_event.is_set()
+                    # map is done - no more outputs, so return early
+                    return
+            finally:
+                # clean up tasks, in case of cancellations etc.
                 get_response_task.cancel()
-                # not strictly necessary - don't leave dangling task
-                await asyncio.gather(get_response_task, return_exceptions=True)
-                return
+                map_done_task.cancel()
 
             last_entry_id = response.last_entry_id
             now_seconds = int(time.time())
@@ -339,7 +345,13 @@ async def _map_invocation(
             output = await _process_result(item.result, item.data_format, client.stub, client)
         except Exception as e:
             if return_exceptions:
-                output = e
+                if wrap_returned_exceptions:
+                    # Prior to client 1.0.4 there was a bug where return_exceptions would wrap
+                    # any returned exceptions in a synchronicity.UserCodeException. This adds
+                    # deprecated non-breaking compatibility bandaid for migrating away from that:
+                    output = modal.exception.UserCodeException(e)
+                else:
+                    output = e
             else:
                 raise e
         return (item.idx, output)
@@ -399,7 +411,7 @@ async def _map_invocation(
         async_merge(drain_input_generator(), pump_inputs(), poll_outputs(), retry_inputs())
     ) as streamer:
         async for response in streamer:
-            if response is not None:
+            if response is not None:  # type: ignore[unreachable]
                 yield response.value
     log_debug_stats_task.cancel()
     await log_debug_stats_task
@@ -411,6 +423,7 @@ async def _map_helper(
     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)
+    wrap_returned_exceptions: bool = True,
 ) -> typing.AsyncGenerator[Any, None]:
     """Core implementation that supports `_map_async()`, `_starmap_async()` and `_for_each_async()`.
 
@@ -423,7 +436,6 @@ async def _map_helper(
     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)
     """
-
     raw_input_queue: Any = SynchronizedQueue()  # type: ignore
     await raw_input_queue.init.aio()
 
@@ -441,12 +453,41 @@ async def _map_helper(
     # 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 with aclosing(
-        async_merge(self._map.aio(raw_input_queue, order_outputs, return_exceptions), feed_queue())
+        async_merge(
+            self._map.aio(raw_input_queue, order_outputs, return_exceptions, wrap_returned_exceptions), feed_queue()
+        )
     ) as map_output_stream:
         async for output in map_output_stream:
             yield output
 
 
+def _maybe_warn_about_exceptions(func_name: str, return_exceptions: bool, wrap_returned_exceptions: bool):
+    if return_exceptions and wrap_returned_exceptions:
+        deprecation_warning(
+            (2025, 6, 27),
+            (
+                f"Function.{func_name} currently leaks an internal exception wrapping type "
+                "(modal.exceptions.UserCodeException) when `return_exceptions=True` is set. "
+                "In the future, this will change, and the underlying exception will be returned directly.\n"
+                "To opt into the future behavior and silence this warning, add `wrap_returned_exceptions=False`:\n\n"
+                f"    f.{func_name}(..., return_exceptions=True, wrap_returned_exceptions=False)"
+            ),
+        )
+
+
+def _invoked_from_sync_wrapper() -> bool:
+    """Check whether the calling function was called from a sync wrapper."""
+    # This is temporary: we only need it to avoind double-firing the wrap_returned_exceptions warning.
+    # (We don't want to push the warning lower in the stack beacuse then we can't attribute to the user's code.)
+    try:
+        frame = inspect.currentframe()
+        caller_function_name = frame.f_back.f_back.f_code.co_name
+        # Embeds some assumptions about how the current calling stack works, but this is just temporary.
+        return caller_function_name == "asend"
+    except Exception:
+        return False
+
+
 @warn_if_generator_is_not_consumed(function_name="Function.map.aio")
 async def _map_async(
     self: "modal.functions.Function",
@@ -456,10 +497,18 @@ async def _map_async(
     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)
+    wrap_returned_exceptions: bool = True,  # wrap returned exceptions in modal.exception.UserCodeException
 ) -> typing.AsyncGenerator[Any, None]:
+    if not _invoked_from_sync_wrapper():
+        _maybe_warn_about_exceptions("map.aio", return_exceptions, wrap_returned_exceptions)
     async_input_gen = async_zip(*[sync_or_async_iter(it) for it in input_iterators])
     async for output in _map_helper(
-        self, async_input_gen, kwargs=kwargs, order_outputs=order_outputs, return_exceptions=return_exceptions
+        self,
+        async_input_gen,
+        kwargs=kwargs,
+        order_outputs=order_outputs,
+        return_exceptions=return_exceptions,
+        wrap_returned_exceptions=wrap_returned_exceptions,
     ):
         yield output
 
@@ -472,13 +521,17 @@ async def _starmap_async(
     kwargs={},
     order_outputs: bool = True,
     return_exceptions: bool = False,
+    wrap_returned_exceptions: bool = True,
 ) -> typing.AsyncIterable[Any]:
+    if not _invoked_from_sync_wrapper():
+        _maybe_warn_about_exceptions("starmap.aio", return_exceptions, wrap_returned_exceptions)
     async for output in _map_helper(
         self,
         sync_or_async_iter(input_iterator),
         kwargs=kwargs,
         order_outputs=order_outputs,
         return_exceptions=return_exceptions,
+        wrap_returned_exceptions=wrap_returned_exceptions,
     ):
         yield output
 
@@ -488,7 +541,12 @@ async def _for_each_async(self, *input_iterators, kwargs={}, ignore_exceptions:
     # rather than iterating over the result
     async_input_gen = async_zip(*[sync_or_async_iter(it) for it in input_iterators])
     async for _ in _map_helper(
-        self, async_input_gen, kwargs=kwargs, order_outputs=False, return_exceptions=ignore_exceptions
+        self,
+        async_input_gen,
+        kwargs=kwargs,
+        order_outputs=False,
+        return_exceptions=ignore_exceptions,
+        wrap_returned_exceptions=False,
     ):
         pass
 
@@ -500,6 +558,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)
+    wrap_returned_exceptions: bool = True,
 ) -> AsyncOrSyncIterable:
     """Parallel map over a set of inputs.
 
@@ -537,10 +596,16 @@ def _map_sync(
         print(list(my_func.map(range(3), return_exceptions=True)))
     ```
     """
+    _maybe_warn_about_exceptions("map", return_exceptions, wrap_returned_exceptions)
 
     return AsyncOrSyncIterable(
         _map_async(
-            self, *input_iterators, kwargs=kwargs, order_outputs=order_outputs, return_exceptions=return_exceptions
+            self,
+            *input_iterators,
+            kwargs=kwargs,
+            order_outputs=order_outputs,
+            return_exceptions=return_exceptions,
+            wrap_returned_exceptions=wrap_returned_exceptions,
         ),
         nested_async_message=(
             "You can't iter(Function.map()) from an async function. Use async for ... in Function.map.aio() instead."
@@ -620,6 +685,7 @@ def _starmap_sync(
     kwargs={},
     order_outputs: bool = True,
     return_exceptions: bool = False,
+    wrap_returned_exceptions: bool = True,
 ) -> AsyncOrSyncIterable:
     """Like `map`, but spreads arguments over multiple function arguments.
 
@@ -637,9 +703,15 @@ def _starmap_sync(
         assert list(my_func.starmap([(1, 2), (3, 4)])) == [3, 7]
     ```
     """
+    _maybe_warn_about_exceptions("starmap", return_exceptions, wrap_returned_exceptions)
     return AsyncOrSyncIterable(
         _starmap_async(
-            self, input_iterator, kwargs=kwargs, order_outputs=order_outputs, return_exceptions=return_exceptions
+            self,
+            input_iterator,
+            kwargs=kwargs,
+            order_outputs=order_outputs,
+            return_exceptions=return_exceptions,
+            wrap_returned_exceptions=wrap_returned_exceptions,
         ),
         nested_async_message=(
             "You can't `iter(Function.starmap())` from an async function. "
@@ -716,6 +788,7 @@ class _MapItemContext:
         Return True if input state was changed to COMPLETE, otherwise False.
         """
         # If the item is already complete, this is a duplicate output and can be ignored.
+
         if self.state == _MapItemState.COMPLETE:
             logger.debug(
                 f"Received output for input marked as complete. Must be duplicate, so ignoring. "
@@ -761,11 +834,12 @@ class _MapItemContext:
             delay_ms = 0
 
         # None means the maximum number of retries has been reached, so output the error
-        if delay_ms is None:
+        if delay_ms is None or item.result.status == api_pb2.GenericResult.GENERIC_STATUS_TERMINATED:
             self.state = _MapItemState.COMPLETE
             return _OutputType.FAILED_COMPLETION
 
         self.state = _MapItemState.WAITING_TO_RETRY
+
         await retry_queue.put(now_seconds + (delay_ms / 1000), item.idx)
 
         return _OutputType.RETRYING

modal/parallel_map.pyi

@@ -12,6 +12,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 +20,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,11 +44,21 @@ 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."""
+        ...
 
 def _map_invocation(
     function: modal._functions._Function,
@@ -52,6 +66,7 @@ def _map_invocation(
     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,
 ): ...
@@ -61,13 +76,33 @@ def _map_helper(
     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 +113,91 @@ 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 _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 +205,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 +260,10 @@ class _MapItemContext:
         input: modal_proto.api_pb2.FunctionInput,
         retry_manager: modal.retries.RetryManager,
         sync_client_retries_enabled: bool,
-    ): ...
+    ):
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
+
     def handle_put_inputs_response(self, item: modal_proto.api_pb2.FunctionPutInputsResponseItem): ...
     async def handle_get_outputs_response(
         self,
@@ -138,7 +271,13 @@ 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 handle_retry_response(self, input_jwt: str): ...
 
@@ -150,12 +289,18 @@ class _MapItemsManager:
         retry_queue: modal._utils.async_utils.TimestampPriorityQueue,
         sync_client_retries_enabled: bool,
         max_inputs_outstanding: int,
-    ): ...
+    ):
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
+
     async def add_items(self, items: list[modal_proto.api_pb2.FunctionPutInputsItem]): ...
     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 get_input_jwts_waiting_for_output(self) -> list[str]:
+        """Returns a list of input_jwts 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_inputs_response(self, items: list[modal_proto.api_pb2.FunctionPutInputsResponseItem]): ...