Client now retries .map() failures

modal/_utils/async_utils.py

@@ -12,6 +12,7 @@
 from typing import (
     Any,
     Callable,
+    Generic,
     Optional,
     TypeVar,
     Union,
@@ -26,6 +27,10 @@
 from ..exception import InvalidError
 from .logger import logger
 
+T = TypeVar("T")
+P = ParamSpec("P")
+V = TypeVar("V")
+
 synchronizer = synchronicity.Synchronizer()
 
 
@@ -260,7 +265,59 @@ def run_coro_blocking(coro):
         return fut.result()
 
 
-async def queue_batch_iterator(q: asyncio.Queue, max_batch_size=100, debounce_time=0.015):
+class TimestampPriorityQueue(Generic[T]):
+    """
+    A priority queue that schedules items to be processed at specific timestamps.
+    """
+
+    _MAX_PRIORITY = float("inf")
+
+    def __init__(self, maxsize: int = 0):
+        self.condition = asyncio.Condition()
+        self._queue: asyncio.PriorityQueue[tuple[float, Union[T, None]]] = asyncio.PriorityQueue(maxsize=maxsize)
+
+    async def close(self):
+        await self.put(self._MAX_PRIORITY, None)
+
+    async def put(self, timestamp: float, item: Union[T, None]):
+        """
+        Add an item to the queue to be processed at a specific timestamp.
+        """
+        await self._queue.put((timestamp, item))
+        async with self.condition:
+            self.condition.notify_all()  # notify any waiting coroutines
+
+    async def get(self) -> Union[T, None]:
+        """
+        Get the next item from the queue that is ready to be processed.
+        """
+        while True:
+            async with self.condition:
+                while self.empty():
+                    await self.condition.wait()
+                # peek at the next item
+                timestamp, item = await self._queue.get()
+                now = time.time()
+                if timestamp < now:
+                    return item
+                if timestamp == self._MAX_PRIORITY:
+                    return None
+                # not ready yet, calculate sleep time
+                sleep_time = timestamp - now
+                self._queue.put_nowait((timestamp, item))  # put it back
+                # wait until either the timeout or a new item is added
+                try:
+                    await asyncio.wait_for(self.condition.wait(), timeout=sleep_time)
+                except asyncio.TimeoutError:
+                    continue
+
+    def empty(self) -> bool:
+        return self._queue.empty()
+
+
+async def queue_batch_iterator(
+    q: Union[asyncio.Queue, TimestampPriorityQueue], max_batch_size=100, debounce_time=0.015
+):
     """
     Read from a queue but return lists of items when queue is large
 
@@ -401,11 +458,6 @@ async def wrapper():
     _shutdown_tasks.append(asyncio.create_task(wrapper()))
 
 
-T = TypeVar("T")
-P = ParamSpec("P")
-V = TypeVar("V")
-
-
 def asyncify(f: Callable[P, T]) -> Callable[P, typing.Coroutine[None, None, T]]:
     """Convert a blocking function into one that runs in the current loop's executor."""
 

modal/functions.py

@@ -1,4 +1,5 @@
 # Copyright Modal Labs 2023
+import asyncio
 import dataclasses
 import inspect
 import textwrap
@@ -256,7 +257,10 @@ async def run_function(self) -> Any:
             try:
                 return await self._get_single_output(ctx.input_jwt)
             except (UserCodeException, FunctionTimeoutError) as exc:
-                await user_retry_manager.raise_or_sleep(exc)
+                delay_ms = user_retry_manager.get_delay_ms()
+                if delay_ms is None:
+                    raise exc
+                await asyncio.sleep(delay_ms / 1000)
             except InternalFailure:
                 # For system failures on the server, we retry immediately.
                 pass
@@ -1244,6 +1248,11 @@ async def _map(
         else:
             count_update_callback = None
 
+        if config.get("client_retries"):
+            function_call_invocation_type = api_pb2.FUNCTION_CALL_INVOCATION_TYPE_SYNC
+        else:
+            function_call_invocation_type = api_pb2.FUNCTION_CALL_INVOCATION_TYPE_SYNC_LEGACY
+
         async with aclosing(
             _map_invocation(
                 self,  # type: ignore
@@ -1252,6 +1261,7 @@ async def _map(
                 order_outputs,
                 return_exceptions,
                 count_update_callback,
+                function_call_invocation_type,
             )
         ) as stream:
             async for item in stream: