modal 0.62.16__py3-none-any.whl → 0.72.11__py3-none-any.whl

modal/_utils/async_utils.py

@@ -3,25 +3,34 @@ import asyncio
 import concurrent.futures
 import functools
 import inspect
-import sys
+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__
@@ -86,14 +95,24 @@ def retry(direct_fn=None, *, n_attempts=3, base_delay=0, delay_factor=2, timeout
 
 
 class TaskContext:
-    """Simple thing to make sure we don't have stray tasks.
+    """A structured group that helps manage stray tasks.
+
+    This differs from the standard library `asyncio.TaskGroup` in that it cancels all tasks still
+    running after exiting the context manager, rather than waiting for them to finish.
+
+    A `TaskContext` can have an optional `grace` period in seconds, which will wait for a certain
+    amount of time before cancelling all remaining tasks. This is useful for allowing tasks to
+    gracefully exit when they determine that the context is shutting down.
 
     Usage:
+
+    ```python notest
     async with TaskContext() as task_context:
-        task = task_context.create(coro())
+        task = task_context.create_task(coro())
+    ```
     """
 
-    _loops: Set[asyncio.Task]
+    _loops: set[asyncio.Task]
 
     def __init__(self, grace: Optional[float] = None):
         self._grace = grace
@@ -130,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
 
@@ -140,15 +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
 
-                if sys.version_info >= (3, 11):
-                    already_cancelling = task.cancelling() > 0
-                    if not already_cancelling:
-                        logger.warning(f"Canceling remaining unfinished task: {task}")
-
+                # 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()
@@ -165,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")
@@ -194,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):
@@ -234,8 +263,10 @@ def run_coro_blocking(coro):
 async def queue_batch_iterator(q: asyncio.Queue, max_batch_size=100, debounce_time=0.015):
     """
     Read from a queue but return lists of items when queue is large
+
+    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:
@@ -257,44 +288,102 @@ async def queue_batch_iterator(q: asyncio.Queue, max_batch_size=100, debounce_ti
 
 
 class _WarnIfGeneratorIsNotConsumed:
-    def __init__(self, gen, gen_f):
+    def __init__(self, gen, function_name: str):
         self.gen = gen
-        self.gen_f = gen_f
+        self.function_name = function_name
         self.iterated = False
         self.warned = False
 
     def __aiter__(self):
         self.iterated = True
-        return self.gen
+        return self.gen.__aiter__()
 
     async def __anext__(self):
         self.iterated = True
         return await self.gen.__anext__()
 
+    async def asend(self, value):
+        self.iterated = True
+        return await self.gen.asend(value)
+
     def __repr__(self):
         return repr(self.gen)
 
     def __del__(self):
         if not self.iterated and not self.warned:
             self.warned = True
-            name = self.gen_f.__name__
             logger.warning(
-                f"Warning: the results of a call to {name} was not consumed, so the call will never be executed."
-                f" Consider a for-loop like `for x in {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)
 
 
-def warn_if_generator_is_not_consumed(gen_f):
+class _WarnIfNonWrappedGeneratorIsNotConsumed(_WarnIfGeneratorIsNotConsumed):
+    # used for non-synchronicity-wrapped generators and iterators
+    def __iter__(self):
+        self.iterated = True
+        return iter(self.gen)
+
+    def __next__(self):
+        self.iterated = True
+        return self.gen.__next__()
+
+    def send(self, value):
+        self.iterated = True
+        return self.gen.send(value)
+
+
+def warn_if_generator_is_not_consumed(function_name: Optional[str] = None):
     # https://gist.github.com/erikbern/01ae78d15f89edfa7f77e5c0a827a94d
-    @functools.wraps(gen_f)
-    def f_wrapped(*args, **kwargs):
-        gen = gen_f(*args, **kwargs)
-        return _WarnIfGeneratorIsNotConsumed(gen, gen_f)
+    def decorator(gen_f):
+        presented_func_name = function_name if function_name is not None else gen_f.__name__
 
-    return f_wrapped
+        @functools.wraps(gen_f)
+        def f_wrapped(*args, **kwargs):
+            gen = gen_f(*args, **kwargs)
+            if inspect.isasyncgen(gen):
+                return _WarnIfGeneratorIsNotConsumed(gen, presented_func_name)
+            else:
+                return _WarnIfNonWrappedGeneratorIsNotConsumed(gen, presented_func_name)
+
+        return f_wrapped
+
+    return decorator
+
+
+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`
+
+    nested_async_message is raised as an InvalidError if the async variant is called
+    from an already async context, since that would otherwise deadlock the event loop
+    """
+
+    def __init__(self, async_iterable: typing.AsyncGenerator[Any, None], nested_async_message):
+        self._async_iterable = async_iterable
+        self.nested_async_message = nested_async_message
+
+    def __aiter__(self):
+        return self._async_iterable
+
+    def __iter__(self):
+        try:
+            with Runner() as runner:
+                yield from run_async_gen(runner, self._async_iterable)
+        except NestedEventLoops:
+            raise InvalidError(self.nested_async_message)
 
 
 _shutdown_tasks = []
@@ -314,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]]:
@@ -339,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.
@@ -384,3 +440,326 @@ async def asyncnullcontext(*args, **kwargs):
         pass
     """
     yield
+
+
+YIELD_TYPE = typing.TypeVar("YIELD_TYPE")
+SEND_TYPE = typing.TypeVar("SEND_TYPE")
+
+
+def run_async_gen(
+    runner: Runner,
+    gen: typing.AsyncGenerator[YIELD_TYPE, SEND_TYPE],
+) -> typing.Generator[YIELD_TYPE, SEND_TYPE, None]:
+    """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
+    exc: Optional[BaseException] = None
+    while True:
+        try:
+            if exc:
+                next_yield = runner.run(gen.athrow(exc))
+            else:
+                next_yield = runner.run(gen.asend(next_send))  # type: ignore[arg-type]
+        except KeyboardInterrupt as e:
+            raise e from None
+        except StopAsyncIteration:
+            break  # typically a graceful exit of the async generator
+        try:
+            next_send = yield next_yield
+            exc = None
+        except BaseException as err:
+            exc = err
+
+
+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