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

modal/_utils/async_utils.py

@@ -3,25 +3,34 @@ import asyncio
 import concurrent.futures
 import functools
 import inspect
+import itertools
 import time
 import typing
+from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Iterable, Iterator
 from contextlib import asynccontextmanager
-from typing import Any, AsyncGenerator, Callable, Iterator, List, Optional, Set, TypeVar, cast
+from dataclasses import dataclass
+from typing import (
+    Any,
+    Callable,
+    Optional,
+    TypeVar,
+    Union,
+    cast,
+)
 
 import synchronicity
-from typing_extensions import ParamSpec
+from synchronicity.async_utils import Runner
+from synchronicity.exceptions import NestedEventLoops
+from typing_extensions import ParamSpec, assert_type
 
 from ..exception import InvalidError
 from .logger import logger
 
 synchronizer = synchronicity.Synchronizer()
-# atexit.register(synchronizer.close)
 
 
 def synchronize_api(obj, target_module=None):
-    if inspect.isclass(obj):
-        blocking_name = obj.__name__.lstrip("_")
-    elif inspect.isfunction(object):
+    if inspect.isclass(obj) or inspect.isfunction(obj):
         blocking_name = obj.__name__.lstrip("_")
     elif isinstance(obj, TypeVar):
         blocking_name = "_BLOCKING_" + obj.__name__
@@ -103,7 +112,7 @@ class TaskContext:
     ```
     """
 
-    _loops: Set[asyncio.Task]
+    _loops: set[asyncio.Task]
 
     def __init__(self, grace: Optional[float] = None):
         self._grace = grace
@@ -140,7 +149,6 @@ class TaskContext:
             if gather_future:
                 try:
                     await gather_future
-                # pre Python3.8, CancelledErrors were a subclass of exception
                 except asyncio.CancelledError:
                     pass
 
@@ -150,11 +158,12 @@ class TaskContext:
                     # Only tasks without a done_callback will still be present in self._tasks
                     task.result()
 
-                if task.done() or task in self._loops:
+                if task.done() or task in self._loops:  # Note: Legacy code, we can probably cancel loops.
                     continue
 
                 # Cancel any remaining unfinished tasks.
                 task.cancel()
+            await asyncio.sleep(0)  # wake up coroutines waiting for cancellations
 
     async def __aexit__(self, exc_type, value, tb):
         await self.stop()
@@ -171,28 +180,32 @@ class TaskContext:
         task.add_done_callback(self._tasks.discard)
         return task
 
-    def infinite_loop(self, async_f, timeout: Optional[float] = 90, sleep: float = 10) -> asyncio.Task:
-        function_name = async_f.__qualname__
+    def infinite_loop(
+        self, async_f, timeout: Optional[float] = 90, sleep: float = 10, log_exception: bool = True
+    ) -> asyncio.Task:
+        if isinstance(async_f, functools.partial):
+            function_name = async_f.func.__qualname__
+        else:
+            function_name = async_f.__qualname__
 
         async def loop_coro() -> None:
             logger.debug(f"Starting infinite loop {function_name}")
-            while True:
-                t0 = time.time()
+            while not self.exited:
                 try:
                     await asyncio.wait_for(async_f(), timeout=timeout)
-                    # pre Python3.8, CancelledErrors were a subclass of exception
-                except asyncio.CancelledError:
-                    raise
-                except Exception:
-                    time_elapsed = time.time() - t0
-                    logger.exception(f"Loop attempt failed for {function_name} (time_elapsed={time_elapsed})")
+                except Exception as exc:
+                    if log_exception and isinstance(exc, asyncio.TimeoutError):
+                        # Asyncio sends an empty message in this case, so let's use logger.error
+                        logger.error(f"Loop attempt for {function_name} timed out")
+                    elif log_exception:
+                        # Propagate the exception to the logger
+                        logger.exception(f"Loop attempt for {function_name} failed")
                 try:
                     await asyncio.wait_for(self._exited.wait(), timeout=sleep)
                 except asyncio.TimeoutError:
                     continue
-                # Only reached if self._exited got set.
-                logger.debug(f"Exiting infinite loop for {function_name}")
-                break
+
+            logger.debug(f"Exiting infinite loop for {function_name}")
 
         t = self.create_task(loop_coro())
         t.set_name(f"{function_name} loop")
@@ -200,29 +213,39 @@ class TaskContext:
         t.add_done_callback(self._loops.discard)
         return t
 
-    async def wait(self, *tasks):
-        # Waits until all of tasks have finished
-        # This is slightly different than asyncio.wait since the `tasks` argument
-        # may be a subset of all the tasks.
-        # If any of the task context's task raises, throw that exception
-        # This is probably O(n^2) sadly but I guess it's fine
-        unfinished_tasks = set(tasks)
-        while True:
-            unfinished_tasks &= self._tasks
-            if not unfinished_tasks:
-                break
-            try:
-                done, pending = await asyncio.wait_for(
-                    asyncio.wait(self._tasks, return_when=asyncio.FIRST_COMPLETED), timeout=30.0
-                )
-            except asyncio.TimeoutError:
-                continue
-            for task in done:
-                task.result()  # Raise exception if needed
-                if task in unfinished_tasks:
-                    unfinished_tasks.remove(task)
-                if task in self._tasks:
-                    self._tasks.remove(task)
+    @staticmethod
+    async def gather(*coros: Awaitable) -> Any:
+        """Wait for a sequence of coroutines to finish, concurrently.
+
+        This is similar to `asyncio.gather()`, but it uses TaskContext to cancel all remaining tasks
+        if one fails with an exception other than `asyncio.CancelledError`. The native `asyncio`
+        function does not cancel remaining tasks in this case, which can lead to surprises.
+
+        For example, if you use `asyncio.gather(t1, t2, t3)` and t2 raises an exception, then t1 and
+        t3 would continue running. With `TaskContext.gather(t1, t2, t3)`, they are cancelled.
+
+        (It's still acceptable to use `asyncio.gather()` if you don't need cancellation — for
+        example, if you're just gathering quick coroutines with no side-effects. Or if you're
+        gathering the tasks with `return_exceptions=True`.)
+
+        Usage:
+
+        ```python notest
+        # Example 1: Await three coroutines
+        created_object, other_work, new_plumbing = await TaskContext.gather(
+            create_my_object(),
+            do_some_other_work(),
+            fix_plumbing(),
+        )
+
+        # Example 2: Gather a list of coroutines
+        coros = [a.load() for a in objects]
+        results = await TaskContext.gather(*coros)
+        ```
+        """
+        async with TaskContext() as tc:
+            results = await asyncio.gather(*(tc.create_task(coro) for coro in coros))
+        return results
 
 
 def run_coro_blocking(coro):
@@ -243,7 +266,7 @@ async def queue_batch_iterator(q: asyncio.Queue, max_batch_size=100, debounce_ti
 
     Treats a None value as end of queue items
     """
-    item_list: List[Any] = []
+    item_list: list[Any] = []
 
     while True:
         if q.empty() and len(item_list) > 0:
@@ -290,10 +313,18 @@ class _WarnIfGeneratorIsNotConsumed:
         if not self.iterated and not self.warned:
             self.warned = True
             logger.warning(
-                f"Warning: the results of a call to {self.function_name} was not consumed, so the call will never be executed."
-                f" Consider a for-loop like `for x in {self.function_name}(...)` or unpacking the generator using `list(...)`"
+                f"Warning: the results of a call to {self.function_name} was not consumed, "
+                "so the call will never be executed."
+                f" Consider a for-loop like `for x in {self.function_name}(...)` or "
+                "unpacking the generator using `list(...)`"
             )
 
+    async def athrow(self, exc):
+        return await self.gen.athrow(exc)
+
+    async def aclose(self):
+        return await self.gen.aclose()
+
 
 synchronize_api(_WarnIfGeneratorIsNotConsumed)
 
@@ -331,7 +362,7 @@ def warn_if_generator_is_not_consumed(function_name: Optional[str] = None):
     return decorator
 
 
-class AsyncOrSyncIteratable:
+class AsyncOrSyncIterable:
     """Compatibility class for non-synchronicity wrapped async iterables to get
     both async and sync interfaces in the same way that synchronicity does (but on the main thread)
     so they can be "lazily" iterated using either `for _ in x` or `async for _ in x`
@@ -340,7 +371,7 @@ class AsyncOrSyncIteratable:
     from an already async context, since that would otherwise deadlock the event loop
     """
 
-    def __init__(self, async_iterable: typing.AsyncIterable[Any], nested_async_message):
+    def __init__(self, async_iterable: typing.AsyncGenerator[Any, None], nested_async_message):
         self._async_iterable = async_iterable
         self.nested_async_message = nested_async_message
 
@@ -349,9 +380,9 @@ class AsyncOrSyncIteratable:
 
     def __iter__(self):
         try:
-            for output in run_generator_sync(self._async_iterable):  # type: ignore
-                yield output
-        except NestedAsyncCalls:
+            with Runner() as runner:
+                yield from run_async_gen(runner, self._async_iterable)
+        except NestedEventLoops:
             raise InvalidError(self.nested_async_message)
 
 
@@ -372,6 +403,7 @@ def on_shutdown(coro):
 
 T = TypeVar("T")
 P = ParamSpec("P")
+V = TypeVar("V")
 
 
 def asyncify(f: Callable[P, T]) -> Callable[P, typing.Coroutine[None, None, T]]:
@@ -397,40 +429,6 @@ async def iterate_blocking(iterator: Iterator[T]) -> AsyncGenerator[T, None]:
         yield cast(T, obj)
 
 
-class ConcurrencyPool:
-    def __init__(self, concurrency_limit: int):
-        self.semaphore = asyncio.Semaphore(concurrency_limit)
-
-    async def run_coros(self, coros: typing.Iterable[typing.Coroutine], return_exceptions=False):
-        async def blocking_wrapper(coro):
-            # Not using async with on the semaphore is intentional here - if return_exceptions=False
-            # manual release prevents starting extraneous tasks after exceptions.
-            try:
-                await self.semaphore.acquire()
-            except asyncio.CancelledError:
-                coro.close()  # avoid "coroutine was never awaited" warnings
-
-            try:
-                res = await coro
-                self.semaphore.release()
-                return res
-            except BaseException as e:
-                if return_exceptions:
-                    self.semaphore.release()
-                raise e
-
-        # asyncio.gather() is weird - it doesn't cancel outstanding awaitables on exceptions when
-        # return_exceptions=False --> wrap the coros in tasks are cancel them explicitly on exception.
-        tasks = [asyncio.create_task(blocking_wrapper(coro)) for coro in coros]
-        g = asyncio.gather(*tasks, return_exceptions=return_exceptions)
-        try:
-            return await g
-        except BaseException as e:
-            for t in tasks:
-                t.cancel()
-            raise e
-
-
 @asynccontextmanager
 async def asyncnullcontext(*args, **kwargs):
     """Async noop context manager.
@@ -448,21 +446,11 @@ YIELD_TYPE = typing.TypeVar("YIELD_TYPE")
 SEND_TYPE = typing.TypeVar("SEND_TYPE")
 
 
-class NestedAsyncCalls(Exception):
-    pass
-
-
-def run_generator_sync(
+def run_async_gen(
+    runner: Runner,
     gen: typing.AsyncGenerator[YIELD_TYPE, SEND_TYPE],
 ) -> typing.Generator[YIELD_TYPE, SEND_TYPE, None]:
-    try:
-        asyncio.get_running_loop()
-    except RuntimeError:
-        pass  # no event loop - this is what we expect!
-    else:
-        raise NestedAsyncCalls()
-    loop = asyncio.new_event_loop()  # set up new event loop for the map so we can use async logic
-
+    """Convert an async generator into a sync one"""
     # more or less copied from synchronicity's implementation:
     next_send: typing.Union[SEND_TYPE, None] = None
     next_yield: YIELD_TYPE
@@ -470,14 +458,308 @@ def run_generator_sync(
     while True:
         try:
             if exc:
-                next_yield = loop.run_until_complete(gen.athrow(exc))
+                next_yield = runner.run(gen.athrow(exc))
             else:
-                next_yield = loop.run_until_complete(gen.asend(next_send))  # type: ignore[arg-type]
+                next_yield = runner.run(gen.asend(next_send))  # type: ignore[arg-type]
+        except KeyboardInterrupt as e:
+            raise e from None
         except StopAsyncIteration:
-            break
+            break  # typically a graceful exit of the async generator
         try:
             next_send = yield next_yield
             exc = None
         except BaseException as err:
             exc = err
-    loop.close()
+
+
+class aclosing(typing.Generic[T]):  # noqa
+    # backport of Python contextlib.aclosing from Python 3.10
+    def __init__(self, agen: AsyncGenerator[T, None]):
+        self.agen = agen
+
+    async def __aenter__(self) -> AsyncGenerator[T, None]:
+        return self.agen
+
+    async def __aexit__(self, exc, exc_type, tb):
+        await self.agen.aclose()
+
+
+async def sync_or_async_iter(iter: Union[Iterable[T], AsyncIterable[T]]) -> AsyncGenerator[T, None]:
+    if hasattr(iter, "__aiter__"):
+        agen = typing.cast(AsyncGenerator[T, None], iter)
+        try:
+            async for item in agen:
+                yield item
+        finally:
+            if hasattr(agen, "aclose"):
+                # All AsyncGenerator's have an aclose method
+                # but some AsyncIterable's don't necessarily
+                await agen.aclose()
+    else:
+        assert hasattr(iter, "__iter__"), "sync_or_async_iter requires an Iterable or AsyncGenerator"
+        # This intentionally could block the event loop for the duration of calling __iter__ and __next__,
+        # so in non-trivial cases (like passing lists and ranges) this could be quite a foot gun for users #
+        # w/ async code (but they can work around it by always using async iterators)
+        for item in typing.cast(Iterable[T], iter):
+            yield item
+
+
+@typing.overload
+def async_zip(g1: AsyncGenerator[T, None], g2: AsyncGenerator[V, None], /) -> AsyncGenerator[tuple[T, V], None]:
+    ...
+
+
+@typing.overload
+def async_zip(*generators: AsyncGenerator[T, None]) -> AsyncGenerator[tuple[T, ...], None]:
+    ...
+
+
+async def async_zip(*generators):
+    tasks = []
+    try:
+        while True:
+            try:
+
+                async def next_item(gen):
+                    return await gen.__anext__()
+
+                tasks = [asyncio.create_task(next_item(gen)) for gen in generators]
+                items = await asyncio.gather(*tasks)
+                yield tuple(items)
+            except StopAsyncIteration:
+                break
+    finally:
+        cancelled_tasks = []
+        for task in tasks:
+            if not task.done():
+                task.cancel()
+                cancelled_tasks.append(task)
+        try:
+            await asyncio.gather(*cancelled_tasks)
+        except asyncio.CancelledError:
+            pass
+
+        first_exception = None
+        for gen in generators:
+            try:
+                await gen.aclose()
+            except BaseException as e:
+                if first_exception is None:
+                    first_exception = e
+                logger.exception(f"Error closing async generator: {e}")
+        if first_exception is not None:
+            raise first_exception
+
+
+@dataclass
+class ValueWrapper(typing.Generic[T]):
+    value: T
+
+
+@dataclass
+class ExceptionWrapper:
+    value: Exception
+
+
+class StopSentinelType:
+    ...
+
+
+STOP_SENTINEL = StopSentinelType()
+
+
+async def async_merge(*generators: AsyncGenerator[T, None]) -> AsyncGenerator[T, None]:
+    """
+    Asynchronously merges multiple async generators into a single async generator.
+
+    This function takes multiple async generators and yields their values in the order
+    they are produced. If any generator raises an exception, the exception is propagated.
+
+    Args:
+        *generators: One or more async generators to be merged.
+
+    Yields:
+        The values produced by the input async generators.
+
+    Raises:
+        Exception: If any of the input generators raises an exception, it is propagated.
+
+    Usage:
+    ```python
+    import asyncio
+    from modal._utils.async_utils import async_merge
+
+    async def gen1():
+        yield 1
+        yield 2
+
+    async def gen2():
+        yield "a"
+        yield "b"
+
+    async def example():
+        values = set()
+        async for value in async_merge(gen1(), gen2()):
+            values.add(value)
+
+        return values
+
+    # Output could be: {1, "a", 2, "b"} (order may vary)
+    values = asyncio.run(example())
+    assert values == {1, "a", 2, "b"}
+    ```
+    """
+    queue: asyncio.Queue[Union[ValueWrapper[T], ExceptionWrapper]] = asyncio.Queue(maxsize=len(generators) * 10)
+
+    async def producer(generator: AsyncGenerator[T, None]):
+        try:
+            async for item in generator:
+                await queue.put(ValueWrapper(item))
+        except Exception as e:
+            await queue.put(ExceptionWrapper(e))
+
+    tasks = {asyncio.create_task(producer(gen)) for gen in generators}
+    new_output_task = asyncio.create_task(queue.get())
+
+    try:
+        while tasks:
+            done, _ = await asyncio.wait(
+                [*tasks, new_output_task],
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+
+            if new_output_task in done:
+                item = new_output_task.result()
+                if isinstance(item, ValueWrapper):
+                    yield item.value
+                else:
+                    assert_type(item, ExceptionWrapper)
+                    raise item.value
+
+                new_output_task = asyncio.create_task(queue.get())
+
+            finished_producers = done & tasks
+            tasks -= finished_producers
+            for finished_producer in finished_producers:
+                # this is done in order to catch potential raised errors/cancellations
+                # from within worker tasks as soon as they happen.
+                await finished_producer
+
+        while not queue.empty():
+            item = await new_output_task
+            if isinstance(item, ValueWrapper):
+                yield item.value
+            else:
+                assert_type(item, ExceptionWrapper)
+                raise item.value
+
+            new_output_task = asyncio.create_task(queue.get())
+
+    finally:
+        if not new_output_task.done():
+            new_output_task.cancel()
+        for task in tasks:
+            if not task.done():
+                try:
+                    task.cancel()
+                    await task
+                except asyncio.CancelledError:
+                    pass
+
+
+async def callable_to_agen(awaitable: Callable[[], Awaitable[T]]) -> AsyncGenerator[T, None]:
+    yield await awaitable()
+
+
+async def gather_cancel_on_exc(*coros_or_futures):
+    input_tasks = [asyncio.ensure_future(t) for t in coros_or_futures]
+    try:
+        return await asyncio.gather(*input_tasks)
+    except BaseException:
+        for t in input_tasks:
+            t.cancel()
+        await asyncio.gather(*input_tasks, return_exceptions=False)  # handle cancellations
+        raise
+
+
+async def async_map(
+    input_generator: AsyncGenerator[T, None],
+    async_mapper_func: Callable[[T], Awaitable[V]],
+    concurrency: int,
+) -> AsyncGenerator[V, None]:
+    queue: asyncio.Queue[Union[ValueWrapper[T], StopSentinelType]] = asyncio.Queue(maxsize=concurrency * 2)
+
+    async def producer() -> AsyncGenerator[V, None]:
+        async for item in input_generator:
+            await queue.put(ValueWrapper(item))
+
+        for _ in range(concurrency):
+            await queue.put(STOP_SENTINEL)
+
+        if False:
+            # Need it to be an async generator for async_merge
+            # but we don't want to yield anything
+            yield
+
+    async def worker() -> AsyncGenerator[V, None]:
+        while True:
+            item = await queue.get()
+            if isinstance(item, ValueWrapper):
+                yield await async_mapper_func(item.value)
+            elif isinstance(item, ExceptionWrapper):
+                raise item.value
+            else:
+                assert_type(item, StopSentinelType)
+                break
+
+    async with aclosing(async_merge(*[worker() for _ in range(concurrency)], producer())) as stream:
+        async for item in stream:
+            yield item
+
+
+async def async_map_ordered(
+    input_generator: AsyncGenerator[T, None],
+    async_mapper_func: Callable[[T], Awaitable[V]],
+    concurrency: int,
+    buffer_size: Optional[int] = None,
+) -> AsyncGenerator[V, None]:
+    semaphore = asyncio.Semaphore(buffer_size or concurrency)
+
+    async def mapper_func_wrapper(tup: tuple[int, T]) -> tuple[int, V]:
+        return (tup[0], await async_mapper_func(tup[1]))
+
+    async def counter() -> AsyncGenerator[int, None]:
+        for i in itertools.count():
+            await semaphore.acquire()
+            yield i
+
+    next_idx = 0
+    buffer = {}
+
+    async with aclosing(async_map(async_zip(counter(), input_generator), mapper_func_wrapper, concurrency)) as stream:
+        async for output_idx, output_item in stream:
+            buffer[output_idx] = output_item
+
+            while next_idx in buffer:
+                yield buffer[next_idx]
+                semaphore.release()
+                del buffer[next_idx]
+                next_idx += 1
+
+
+async def async_chain(*generators: AsyncGenerator[T, None]) -> AsyncGenerator[T, None]:
+    try:
+        for gen in generators:
+            async for item in gen:
+                yield item
+    finally:
+        first_exception = None
+        for gen in generators:
+            try:
+                await gen.aclose()
+            except BaseException as e:
+                if first_exception is None:
+                    first_exception = e
+                logger.exception(f"Error closing async generator: {e}")
+        if first_exception is not None:
+            raise first_exception