modal 1.0.3.dev10__py3-none-any.whl → 1.2.3.dev7__py3-none-any.whl

modal/_utils/async_utils.py

@@ -1,11 +1,14 @@
 # Copyright Modal Labs 2022
 import asyncio
 import concurrent.futures
+import contextlib
 import functools
 import inspect
 import itertools
+import sys
 import time
 import typing
+import warnings
 from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Iterable, Iterator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
@@ -31,6 +34,10 @@ T = TypeVar("T")
 P = ParamSpec("P")
 V = TypeVar("V")
 
+if sys.platform == "win32":
+    # quick workaround for deadlocks on shutdown - need to investigate further
+    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+
 synchronizer = synchronicity.Synchronizer()
 
 
@@ -46,6 +53,10 @@ def synchronize_api(obj, target_module=None):
     return synchronizer.create_blocking(obj, blocking_name, target_module=target_module)
 
 
+# Used for testing to configure the `n_attempts` that `retry` will use.
+RETRY_N_ATTEMPTS_OVERRIDE: Optional[int] = None
+
+
 def retry(direct_fn=None, *, n_attempts=3, base_delay=0, delay_factor=2, timeout=90):
     """Decorator that calls an async function multiple times, with a given timeout.
 
@@ -70,8 +81,13 @@ def retry(direct_fn=None, *, n_attempts=3, base_delay=0, delay_factor=2, timeout
     def decorator(fn):
         @functools.wraps(fn)
         async def f_wrapped(*args, **kwargs):
+            if RETRY_N_ATTEMPTS_OVERRIDE is not None:
+                local_n_attempts = RETRY_N_ATTEMPTS_OVERRIDE
+            else:
+                local_n_attempts = n_attempts
+
             delay = base_delay
-            for i in range(n_attempts):
+            for i in range(local_n_attempts):
                 t0 = time.time()
                 try:
                     return await asyncio.wait_for(fn(*args, **kwargs), timeout=timeout)
@@ -79,12 +95,12 @@ def retry(direct_fn=None, *, n_attempts=3, base_delay=0, delay_factor=2, timeout
                     logger.debug(f"Function {fn} was cancelled")
                     raise
                 except Exception as e:
-                    if i >= n_attempts - 1:
+                    if i >= local_n_attempts - 1:
                         raise
                     logger.debug(
                         f"Failed invoking function {fn}: {e}"
                         f" (took {time.time() - t0}s, sleeping {delay}s"
-                        f" and trying {n_attempts - i - 1} more times)"
+                        f" and trying {local_n_attempts - i - 1} more times)"
                     )
                 await asyncio.sleep(delay)
                 delay *= delay_factor
@@ -120,7 +136,8 @@ class TaskContext:
     _loops: set[asyncio.Task]
 
     def __init__(self, grace: Optional[float] = None):
-        self._grace = grace
+        self._grace = grace  # grace is the time we want for tasks to finish before cancelling them
+        self._cancellation_grace: float = 1.0  # extra graceperiod for the cancellation itself to "bubble up"
         self._loops = set()
 
     async def start(self):
@@ -152,22 +169,29 @@ class TaskContext:
             # still needs to be handled
             # (https://stackoverflow.com/a/63356323/2475114)
             if gather_future:
-                try:
+                with contextlib.suppress(asyncio.CancelledError):
                     await gather_future
-                except asyncio.CancelledError:
-                    pass
 
+            cancelled_tasks: list[asyncio.Task] = []
             for task in self._tasks:
                 if task.done() and not task.cancelled():
                     # Raise any exceptions if they happened.
                     # Only tasks without a done_callback will still be present in self._tasks
                     task.result()
 
-                if task.done() or task in self._loops:  # Note: Legacy code, we can probably cancel loops.
+                if task.done():
                     continue
 
                 # Cancel any remaining unfinished tasks.
                 task.cancel()
+                cancelled_tasks.append(task)
+
+            cancellation_gather = asyncio.gather(*cancelled_tasks, return_exceptions=True)
+            try:
+                await asyncio.wait_for(cancellation_gather, timeout=self._cancellation_grace)
+            except asyncio.TimeoutError:
+                warnings.warn(f"Internal warning: Tasks did not cancel in a timely manner: {cancelled_tasks}")
+
             await asyncio.sleep(0)  # wake up coroutines waiting for cancellations
 
     async def __aexit__(self, exc_type, value, tb):
@@ -274,7 +298,9 @@ class TimestampPriorityQueue(Generic[T]):
 
     def __init__(self, maxsize: int = 0):
         self.condition = asyncio.Condition()
-        self._queue: asyncio.PriorityQueue[tuple[float, Union[T, None]]] = asyncio.PriorityQueue(maxsize=maxsize)
+        self._queue: asyncio.PriorityQueue[tuple[float, int, Union[T, None]]] = asyncio.PriorityQueue(maxsize=maxsize)
+        # Used to tiebreak items with the same timestamp that are not comparable. (eg. protos)
+        self._counter = itertools.count()
 
     async def close(self):
         await self.put(self._MAX_PRIORITY, None)
@@ -283,7 +309,7 @@ class TimestampPriorityQueue(Generic[T]):
         """
         Add an item to the queue to be processed at a specific timestamp.
         """
-        await self._queue.put((timestamp, item))
+        await self._queue.put((timestamp, next(self._counter), item))
         async with self.condition:
             self.condition.notify_all()  # notify any waiting coroutines
 
@@ -296,7 +322,7 @@ class TimestampPriorityQueue(Generic[T]):
                 while self.empty():
                     await self.condition.wait()
                 # peek at the next item
-                timestamp, item = await self._queue.get()
+                timestamp, counter, item = await self._queue.get()
                 now = time.time()
                 if timestamp < now:
                     return item
@@ -304,7 +330,7 @@ class TimestampPriorityQueue(Generic[T]):
                     return None
                 # not ready yet, calculate sleep time
                 sleep_time = timestamp - now
-                self._queue.put_nowait((timestamp, item))  # put it back
+                self._queue.put_nowait((timestamp, counter, 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)
@@ -396,7 +422,7 @@ class _WarnIfGeneratorIsNotConsumed:
         return await self.gen.aclose()
 
 
-synchronize_api(_WarnIfGeneratorIsNotConsumed)
+_BlockingWarnIfGeneratorIsNotConsumed = synchronize_api(_WarnIfGeneratorIsNotConsumed)
 
 
 class _WarnIfNonWrappedGeneratorIsNotConsumed(_WarnIfGeneratorIsNotConsumed):
@@ -647,7 +673,9 @@ class StopSentinelType: ...
 STOP_SENTINEL = StopSentinelType()
 
 
-async def async_merge(*generators: AsyncGenerator[T, None]) -> AsyncGenerator[T, None]:
+async def async_merge(
+    *generators: AsyncGenerator[T, None], cancellation_timeout: float = 10.0
+) -> AsyncGenerator[T, None]:
     """
     Asynchronously merges multiple async generators into a single async generator.
 
@@ -692,8 +720,9 @@ async def async_merge(*generators: AsyncGenerator[T, None]) -> AsyncGenerator[T,
 
     async def producer(generator: AsyncGenerator[T, None]):
         try:
-            async for item in generator:
-                await queue.put(ValueWrapper(item))
+            async with aclosing(generator) as stream:
+                async for item in stream:
+                    await queue.put(ValueWrapper(item))
         except Exception as e:
             await queue.put(ExceptionWrapper(e))
 
@@ -735,15 +764,20 @@ async def async_merge(*generators: AsyncGenerator[T, None]) -> AsyncGenerator[T,
             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
+        unfinished_tasks = [t for t in tasks | {new_output_task} if not t.done()]
+        for t in unfinished_tasks:
+            t.cancel()
+        try:
+            await asyncio.wait_for(
+                asyncio.shield(
+                    # we need to `shield` the `gather` to ensure cooperation with the timeout
+                    # all underlying tasks have been marked as cancelled at this point anyway
+                    asyncio.gather(*unfinished_tasks, return_exceptions=True)
+                ),
+                timeout=cancellation_timeout,
+            )
+        except asyncio.TimeoutError:
+            logger.debug("Timed out while cleaning up async_merge")
 
 
 async def callable_to_agen(awaitable: Callable[[], Awaitable[T]]) -> AsyncGenerator[T, None]:
@@ -761,16 +795,34 @@ async def gather_cancel_on_exc(*coros_or_futures):
         raise
 
 
+async def prevent_cancellation_abortion(coro):
+    # if this is cancelled, it will wait for coro cancellation handling
+    # and then unconditionally re-raises a CancelledError, even if the underlying coro
+    # doesn't re-raise the cancellation itself
+    t = asyncio.create_task(coro)
+    try:
+        return await asyncio.shield(t)
+    except asyncio.CancelledError:
+        if t.cancelled():
+            # coro cancelled itself - reraise
+            raise
+        t.cancel()  # cancel task
+        await t  # this *normally* reraises
+        raise  # if the above somehow resolved, by swallowing cancellation - we still raise
+
+
 async def async_map(
     input_generator: AsyncGenerator[T, None],
     async_mapper_func: Callable[[T], Awaitable[V]],
     concurrency: int,
+    cancellation_timeout: float = 10.0,
 ) -> 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))
+        async with aclosing(input_generator) as stream:
+            async for item in stream:
+                await queue.put(ValueWrapper(item))
 
         for _ in range(concurrency):
             await queue.put(STOP_SENTINEL)
@@ -784,14 +836,17 @@ async def async_map(
         while True:
             item = await queue.get()
             if isinstance(item, ValueWrapper):
-                yield await async_mapper_func(item.value)
+                res = await prevent_cancellation_abortion(async_mapper_func(item.value))
+                yield res
             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 with aclosing(
+        async_merge(*[worker() for i in range(concurrency)], producer(), cancellation_timeout=cancellation_timeout)
+    ) as stream:
         async for item in stream:
             yield item
 

modal/_utils/auth_token_manager.py

@@ -0,0 +1,111 @@
+# Copyright Modal Labs 2025
+import asyncio
+import base64
+import json
+import time
+import typing
+from typing import Any
+
+from modal.exception import ExecutionError
+from modal_proto import api_pb2, modal_api_grpc
+
+from .logger import logger
+
+
+class _AuthTokenManager:
+    """Handles fetching and refreshing of the input plane auth token."""
+
+    # Start refreshing this many seconds before the token expires
+    REFRESH_WINDOW = 5 * 60
+    # If the token doesn't have an expiry field, default to current time plus this value (not expected).
+    DEFAULT_EXPIRY_OFFSET = 20 * 60
+
+    def __init__(self, stub: "modal_api_grpc.ModalClientModal"):
+        self._stub = stub
+        self._token = ""
+        self._expiry = 0.0
+        self._lock: typing.Union[asyncio.Lock, None] = None
+
+    async def get_token(self) -> str:
+        """
+        When called, the AuthTokenManager can be in one of three states:
+        1. Has a valid cached token. It is returned to the caller.
+        2. Has no cached token, or the token is expired. We fetch a new one and cache it. If `get_token` is called
+        concurrently by multiple coroutines, all requests will block until the token has been fetched. But only one
+        coroutine will actually make a request to the control plane to fetch the new token. This ensures we do not hit
+        the control plane with more requests than needed.
+        3. Has a valid cached token, but it is going to expire in the next 5 minutes. In this case we fetch a new token
+        and cache it. If `get_token` is called concurrently, only one request will fetch the new token, and the others
+        will be given the old (but still valid) token - i.e. they will not block.
+        """
+        if not self._token or self._is_expired():
+            # We either have no token or it is expired - block everyone until we get a new token
+            await self._refresh_token()
+        elif self._needs_refresh():
+            # The token hasn't expired yet, but will soon, so it needs a refresh.
+            lock = await self._get_lock()
+            if lock.locked():
+                # The lock is taken, so someone else is refreshing. Continue to use the old token.
+                return self._token
+            else:
+                # The lock is not taken, so we need to fetch a new token.
+                await self._refresh_token()
+
+        return self._token
+
+    async def _refresh_token(self):
+        """
+        Fetch a new token from the control plane. If called concurrently, only one coroutine will make a request for a
+        new token. The others will block on a lock, until the first coroutine has fetched the new token.
+        """
+        lock = await self._get_lock()
+        async with lock:
+            # Double check inside lock - maybe another coroutine refreshed already. This happens the first time we fetch
+            # the token. The first coroutine will fetch the token, while the others block on the lock, waiting for the
+            # new token. Once we have a new token, the other coroutines will unblock and return from here.
+            if self._token and not self._needs_refresh():
+                return
+            resp: api_pb2.AuthTokenGetResponse = await self._stub.AuthTokenGet(api_pb2.AuthTokenGetRequest())
+            if not resp.token:
+                # Not expected
+                raise ExecutionError(
+                    "Internal error: Did not receive auth token from server. Please contact Modal support."
+                )
+
+            self._token = resp.token
+            if exp := self._decode_jwt(resp.token).get("exp"):
+                self._expiry = float(exp)
+            else:
+                # This should never happen.
+                logger.warning("x-modal-auth-token does not contain exp field")
+                # We'll use the token, and set the expiry to 20 min from now.
+                self._expiry = time.time() + self.DEFAULT_EXPIRY_OFFSET
+
+    async def _get_lock(self) -> asyncio.Lock:
+        # Note: this function runs no async code but is marked as async to ensure it's
+        # being run inside the synchronicity event loop and binds the lock to the
+        # correct event loop on Python 3.9 which eagerly assigns event loops on
+        # constructions of locks
+        if self._lock is None:
+            self._lock = asyncio.Lock()
+        return self._lock
+
+    @staticmethod
+    def _decode_jwt(token: str) -> dict[str, Any]:
+        """
+        Decodes a JWT into a dict without verifying signature. We do this manually instead of using a library to avoid
+        adding another dependency to the client.
+        """
+        try:
+            payload = token.split(".")[1]
+            padding = "=" * (-len(payload) % 4)
+            decoded_bytes = base64.urlsafe_b64decode(payload + padding)
+            return json.loads(decoded_bytes)
+        except Exception as e:
+            raise ValueError("Internal error: Cannot parse auth token. Please contact Modal support.") from e
+
+    def _needs_refresh(self):
+        return time.time() >= (self._expiry - self.REFRESH_WINDOW)
+
+    def _is_expired(self):
+        return time.time() >= self._expiry