modAL 1.2.3.dev12__tar.gz → 1.2.3.dev14__tar.gz

{modal-1.2.3.dev12 → modal-1.2.3.dev14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modal
-Version: 1.2.3.dev12
+Version: 1.2.3.dev14
 Summary: Python client library for Modal
 Author-email: Modal Labs <support@modal.com>
 License: Apache-2.0

{modal-1.2.3.dev12 → modal-1.2.3.dev14}/modal/_functions.py

@@ -536,6 +536,7 @@ class _InputPlaneInvocation:
     async def _get_metadata(input_plane_region: str, client: _Client) -> list[tuple[str, str]]:
         if not input_plane_region:
             return []
+        assert client._auth_token_manager, "Client is not open"
         token = await client._auth_token_manager.get_token()
         return [("x-modal-input-plane-region", input_plane_region), ("x-modal-auth-token", token)]
 

{modal-1.2.3.dev12 → modal-1.2.3.dev14}/modal/_load_context.py

@@ -66,8 +66,9 @@ class LoadContext:
 
         Returns a new instance (no in place mutation)"""
 
+        is_valid_client = self._client is not None and not self._client._snapshotted
         return LoadContext(
-            client=await _Client.from_env() if self._client is None else self.client,
+            client=self.client if is_valid_client else await _Client.from_env(),
             environment_name=self._environment_name or config.get("environment") or "",
             app_id=self._app_id,
         )

{modal-1.2.3.dev12 → modal-1.2.3.dev14}/modal/_object.py

@@ -43,7 +43,7 @@ class _Object:
     _prefix_to_type: ClassVar[dict[str, type]] = {}
 
     # For constructors
-    _load: Optional[Callable[[Self, Resolver, LoadContext, Optional[str]], Awaitable[None]]]
+    _load: Optional[Callable[[Self, Resolver, LoadContext, Optional[str]], Awaitable[None]]] = None
     _preload: Optional[Callable[[Self, Resolver, LoadContext, Optional[str]], Awaitable[None]]]
     _rep: str
     _is_another_app: bool
@@ -308,11 +308,17 @@ class _Object:
                 # memory snapshots capture references which must be rehydrated
                 # on restore to handle staleness.
                 logger.debug(f"rehydrating {self} after snapshot")
-                self._is_hydrated = False  # un-hydrate and re-resolve
-                # Set the client on LoadContext before loading
-                root_load_context = LoadContext(client=client)
-                resolver = Resolver()
-                await resolver.load(typing.cast(_Object, self), root_load_context)
+                if self._hydrate_lazily:
+                    logger.debug(f"reloading lazy {self} from server")
+                    self._is_hydrated = False  # un-hydrate and re-resolve
+                    # we don't set an explicit Client here, relying on the default
+                    # env client to be applied by LoadContext.apply_default
+                    root_load_context = LoadContext.empty()
+                    resolver = Resolver()
+                    await resolver.load(typing.cast(_Object, self), root_load_context)
+                else:
+                    logger.debug(f"reloading non-lazy {self} by replacing client")
+                    self._client = client or await _Client.from_env()
                 self._is_rehydrated = True
                 logger.debug(f"rehydrated {self} with client {id(self.client)}")
         elif not self._hydrate_lazily:

{modal-1.2.3.dev12 → modal-1.2.3.dev14}/modal/_utils/grpc_utils.py

@@ -29,6 +29,7 @@ from modal_proto import api_pb2
 from modal_version import __version__
 
 from .._traceback import suppress_tb_frames
+from ..config import config
 from .async_utils import retry
 from .logger import logger
 
@@ -72,6 +73,7 @@ RETRYABLE_GRPC_STATUS_CODES = [
     Status.INTERNAL,
     Status.UNKNOWN,
 ]
+SERVER_RETRY_WARNING_TIME_INTERVAL = 30.0
 
 
 @dataclass
@@ -251,6 +253,50 @@ async def retry_transient_errors(
     return await _retry_transient_errors(fn, req, retry=Retry(max_retries=max_retries))
 
 
+def get_server_retry_policy(exc: Exception) -> Optional[api_pb2.RPCRetryPolicy]:
+    """Get server retry policy."""
+    if not isinstance(exc, GRPCError) or not exc.details:
+        return None
+
+    # Server should not set multiple retry instructions, but if there is more than one, pick the first one
+    for entry in exc.details:
+        if isinstance(entry, api_pb2.RPCRetryPolicy):
+            return entry
+    return None
+
+
+def process_exception_before_retry(
+    exc: Exception,
+    final_attempt: bool,
+    fn_name: str,
+    n_retries: int,
+    delay: float,
+    total_deadline: Optional[float],
+    idempotency_key: str,
+):
+    """Process exception before retry, used by `_retry_transient_errors`."""
+    with suppress_tb_frames(1):
+        if final_attempt:
+            logger.debug(
+                f"Final attempt failed with {repr(exc)} {n_retries=} {delay=} "
+                f"{total_deadline=} for {fn_name} ({idempotency_key[:8]})"
+            )
+            if isinstance(exc, OSError):
+                raise ConnectionError(str(exc))
+            elif isinstance(exc, asyncio.TimeoutError):
+                raise ConnectionError(str(exc))
+            else:
+                raise exc
+
+        if isinstance(exc, AttributeError) and "_write_appdata" not in str(exc):
+            # StreamTerminatedError are not properly raised in grpclib<=0.4.7
+            # fixed in https://github.com/vmagamedov/grpclib/issues/185
+            # TODO: update to newer version (>=0.4.8) once stable
+            raise exc
+
+    logger.debug(f"Retryable failure {repr(exc)} {n_retries=} {delay=} for {fn_name} ({idempotency_key[:8]})")
+
+
 async def _retry_transient_errors(
     fn: typing.Union[
         "modal._grpc_client.UnaryUnaryWrapper[RequestType, ResponseType]",
@@ -273,12 +319,15 @@ async def _retry_transient_errors(
 
     delay = retry.base_delay
     n_retries = 0
+    n_throttled_retries = 0
 
     status_codes = [*RETRYABLE_GRPC_STATUS_CODES, *retry.additional_status_codes]
 
     idempotency_key = str(uuid.uuid4())
 
     t0 = time.time()
+    last_server_retry_warning_time = None
+
     if retry.total_timeout is not None:
         total_deadline = t0 + retry.total_timeout
     else:
@@ -290,14 +339,18 @@ async def _retry_transient_errors(
         attempt_metadata = [
             ("x-idempotency-key", idempotency_key),
             ("x-retry-attempt", str(n_retries)),
+            ("x-throttle-retry-attempt", str(n_throttled_retries)),
             *metadata,
         ]
         if n_retries > 0:
             attempt_metadata.append(("x-retry-delay", str(time.time() - t0)))
+        if n_throttled_retries > 0:
+            attempt_metadata.append(("x-throttle-retry-delay", str(time.time() - t0)))
+
         timeouts = []
         if retry.attempt_timeout is not None:
             timeouts.append(retry.attempt_timeout)
-        if retry.total_timeout is not None and total_deadline is not None:
+        if total_deadline is not None:
             timeouts.append(max(total_deadline - time.time(), retry.attempt_timeout_floor))
         if timeouts:
             timeout = min(timeouts)  # In case the function provided both types of timeouts
@@ -307,6 +360,39 @@ async def _retry_transient_errors(
             with suppress_tb_frames(1):
                 return await fn_callable(req, metadata=attempt_metadata, timeout=timeout)
         except (StreamTerminatedError, GRPCError, OSError, asyncio.TimeoutError, AttributeError) as exc:
+            # Server side instruction for retries
+            if isinstance(exc, GRPCError) and (server_retry_policy := get_server_retry_policy(exc)):
+                server_delay = server_retry_policy.retry_after_secs
+
+                # When `total_deadline` is not defined and server gives instruction for retries, then use
+                # `max_throttle_wait`.
+                if total_deadline is None and (max_throttle_wait := config.get("max_throttle_wait")):
+                    total_deadline = t0 + max_throttle_wait
+
+                final_attempt = (
+                    total_deadline is not None
+                    and time.time() + server_delay + retry.attempt_timeout_floor >= total_deadline
+                )
+                with suppress_tb_frames(1):
+                    process_exception_before_retry(
+                        exc, final_attempt, fn.name, n_retries, server_delay, total_deadline, idempotency_key
+                    )
+
+                now = time.time()
+                if last_server_retry_warning_time is None or (
+                    now - last_server_retry_warning_time >= SERVER_RETRY_WARNING_TIME_INTERVAL
+                ):
+                    last_server_retry_warning_time = now
+                    logger.warning(
+                        f"Warning: Received failure {exc.status}: {exc.message}. "
+                        f"Will retry in {server_delay:0.2f} seconds."
+                    )
+
+                n_throttled_retries += 1
+                await asyncio.sleep(server_delay)
+                continue
+
+            # Client handles retry
             if isinstance(exc, GRPCError) and exc.status not in status_codes:
                 if exc.status == Status.UNAUTHENTICATED:
                     raise AuthError(exc.message)
@@ -320,25 +406,9 @@ async def _retry_transient_errors(
                 final_attempt = False
 
             with suppress_tb_frames(1):
-                if final_attempt:
-                    logger.debug(
-                        f"Final attempt failed with {repr(exc)} {n_retries=} {delay=} "
-                        f"{total_deadline=} for {fn.name} ({idempotency_key[:8]})"
-                    )
-                    if isinstance(exc, OSError):
-                        raise ConnectionError(str(exc))
-                    elif isinstance(exc, asyncio.TimeoutError):
-                        raise ConnectionError(str(exc))
-                    else:
-                        raise exc
-
-                if isinstance(exc, AttributeError) and "_write_appdata" not in str(exc):
-                    # StreamTerminatedError are not properly raised in grpclib<=0.4.7
-                    # fixed in https://github.com/vmagamedov/grpclib/issues/185
-                    # TODO: update to newer version (>=0.4.8) once stable
-                    raise exc
-
-            logger.debug(f"Retryable failure {repr(exc)} {n_retries=} {delay=} for {fn.name} ({idempotency_key[:8]})")
+                process_exception_before_retry(
+                    exc, final_attempt, fn.name, n_retries, delay, total_deadline, idempotency_key
+                )
 
             n_retries += 1
 

{modal-1.2.3.dev12 → modal-1.2.3.dev14}/modal/client.py

@@ -5,7 +5,7 @@ import platform
 import sys
 import urllib.parse
 import warnings
-from collections.abc import AsyncGenerator, AsyncIterator, Collection, Mapping
+from collections.abc import AsyncGenerator, Collection, Mapping
 from typing import Any, ClassVar, Optional, TypeVar, Union
 
 import grpclib.client
@@ -14,7 +14,7 @@ from google.protobuf.message import Message
 from synchronicity.async_wrap import asynccontextmanager
 
 from modal._utils.async_utils import synchronizer
-from modal_proto import api_grpc, api_pb2, modal_api_grpc
+from modal_proto import api_pb2, modal_api_grpc
 from modal_version import __version__
 
 from ._traceback import print_server_warnings
@@ -70,10 +70,10 @@ class _Client:
     _client_from_env: ClassVar[Optional["_Client"]] = None
     _client_from_env_lock: ClassVar[Optional[asyncio.Lock]] = None
     _cancellation_context: TaskContext
-    _cancellation_context_event_loop: asyncio.AbstractEventLoop = None
-    _stub: Optional[api_grpc.ModalClientStub]
-    _auth_token_manager: _AuthTokenManager = None
-    _snapshotted: bool
+    _cancellation_context_event_loop: Optional[asyncio.AbstractEventLoop] = None
+    _stub: Optional[modal_api_grpc.ModalClientModal]
+    _auth_token_manager: Optional[_AuthTokenManager] = None
+    _snapshotted: bool = False
 
     def __init__(
         self,
@@ -90,8 +90,8 @@ class _Client:
         self._credentials = credentials
         self.version = version
         self._closed = False
-        self._stub: Optional[modal_api_grpc.ModalClientModal] = None
-        self._auth_token_manager: Optional[_AuthTokenManager] = None
+        self._stub = None
+        self._auth_token_manager = None
         self._snapshotted = False
         self._owner_pid = None
 
@@ -163,7 +163,7 @@ class _Client:
 
     @classmethod
     @asynccontextmanager
-    async def anonymous(cls, server_url: str) -> AsyncIterator["_Client"]:
+    async def anonymous(cls, server_url: str) -> AsyncGenerator["_Client", None]:
         """mdmd:hidden
         Create a connection with no credentials; to be used for token creation.
         """

{modal-1.2.3.dev12 → modal-1.2.3.dev14}/modal/client.pyi

@@ -5,7 +5,6 @@ import google.protobuf.message
 import grpclib.client
 import modal._utils.async_utils
 import modal._utils.auth_token_manager
-import modal_proto.api_grpc
 import modal_proto.modal_api_grpc
 import synchronicity.combined_types
 import typing
@@ -23,9 +22,9 @@ class _Client:
     _client_from_env: typing.ClassVar[typing.Optional[_Client]]
     _client_from_env_lock: typing.ClassVar[typing.Optional[asyncio.locks.Lock]]
     _cancellation_context: modal._utils.async_utils.TaskContext
-    _cancellation_context_event_loop: asyncio.events.AbstractEventLoop
-    _stub: typing.Optional[modal_proto.api_grpc.ModalClientStub]
-    _auth_token_manager: modal._utils.auth_token_manager._AuthTokenManager
+    _cancellation_context_event_loop: typing.Optional[asyncio.events.AbstractEventLoop]
+    _stub: typing.Optional[modal_proto.modal_api_grpc.ModalClientModal]
+    _auth_token_manager: typing.Optional[modal._utils.auth_token_manager._AuthTokenManager]
     _snapshotted: bool
 
     def __init__(
@@ -33,7 +32,7 @@ class _Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.2.3.dev12",
+        version: str = "1.2.3.dev14",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.
@@ -154,9 +153,9 @@ class Client:
     _client_from_env: typing.ClassVar[typing.Optional[Client]]
     _client_from_env_lock: typing.ClassVar[typing.Optional[asyncio.locks.Lock]]
     _cancellation_context: modal._utils.async_utils.TaskContext
-    _cancellation_context_event_loop: asyncio.events.AbstractEventLoop
-    _stub: typing.Optional[modal_proto.api_grpc.ModalClientStub]
-    _auth_token_manager: modal._utils.auth_token_manager._AuthTokenManager
+    _cancellation_context_event_loop: typing.Optional[asyncio.events.AbstractEventLoop]
+    _stub: typing.Optional[modal_proto.modal_api_grpc.ModalClientModal]
+    _auth_token_manager: typing.Optional[modal._utils.auth_token_manager._AuthTokenManager]
     _snapshotted: bool
 
     def __init__(
@@ -164,7 +163,7 @@ class Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.2.3.dev12",
+        version: str = "1.2.3.dev14",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.

{modal-1.2.3.dev12 → modal-1.2.3.dev14}/modal/config.py

@@ -217,6 +217,13 @@ def _enforce_suffix_rules(x: str) -> str:
     return x
 
 
+def _int_or_none(x: str) -> Optional[int]:
+    if not x:
+        return None
+    x_int = int(x)
+    return x_int if x_int > 0 else None
+
+
 class _Setting(typing.NamedTuple):
     default: typing.Any = None
     transform: typing.Callable[[str], typing.Any] = lambda x: x  # noqa: E731
@@ -258,6 +265,7 @@ _SETTINGS = {
         transform=lambda s: _check_value(["pickle", "cbor"])(s.lower()),
     ),
     "dev_suffix": _Setting("", transform=_enforce_suffix_rules),
+    "max_throttle_wait": _Setting(None, transform=_int_or_none),
 }
 
 

{modal-1.2.3.dev12 → modal-1.2.3.dev14}/modal/functions.pyi

@@ -408,7 +408,7 @@ class Function(
 
     _call_generator: ___call_generator_spec[typing_extensions.Self]
 
-    class __remote_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
+    class __remote_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
         def __call__(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> ReturnType_INNER:
             """Calls the function remotely, executing it with the given arguments and returning the execution's result."""
             ...
@@ -417,7 +417,7 @@ class Function(
             """Calls the function remotely, executing it with the given arguments and returning the execution's result."""
             ...
 
-    remote: __remote_spec[modal._functions.P, modal._functions.ReturnType, typing_extensions.Self]
+    remote: __remote_spec[modal._functions.ReturnType, modal._functions.P, typing_extensions.Self]
 
     class __remote_gen_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, /, *args, **kwargs) -> typing.Generator[typing.Any, None, None]:
@@ -444,7 +444,7 @@ class Function(
         """
         ...
 
-    class ___experimental_spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
+    class ___experimental_spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
         def __call__(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]:
             """[Experimental] Calls the function with the given arguments, without waiting for the results.
 
@@ -468,7 +468,7 @@ class Function(
             ...
 
     _experimental_spawn: ___experimental_spawn_spec[
-        modal._functions.P, modal._functions.ReturnType, typing_extensions.Self
+        modal._functions.ReturnType, modal._functions.P, typing_extensions.Self
     ]
 
     class ___spawn_map_inner_spec(typing_extensions.Protocol[P_INNER, SUPERSELF]):
@@ -477,7 +477,7 @@ class Function(
 
     _spawn_map_inner: ___spawn_map_inner_spec[modal._functions.P, typing_extensions.Self]
 
-    class __spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
+    class __spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
         def __call__(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]:
             """Calls the function with the given arguments, without waiting for the results.
 
@@ -498,7 +498,7 @@ class Function(
             """
             ...
 
-    spawn: __spawn_spec[modal._functions.P, modal._functions.ReturnType, typing_extensions.Self]
+    spawn: __spawn_spec[modal._functions.ReturnType, modal._functions.P, typing_extensions.Self]
 
     def get_raw_f(self) -> collections.abc.Callable[..., typing.Any]:
         """Return the inner Python object wrapped by this Modal Function."""

{modal-1.2.3.dev12 → modal-1.2.3.dev14}/modal.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modal
-Version: 1.2.3.dev12
+Version: 1.2.3.dev14
 Summary: Python client library for Modal
 Author-email: Modal Labs <support@modal.com>
 License: Apache-2.0

{modal-1.2.3.dev12 → modal-1.2.3.dev14}/modal_proto/api.proto

@@ -2642,6 +2642,12 @@ message QueuePutRequest {
   int32 partition_ttl_seconds = 6;
 }
 
+// Retry repolicy used by GRPCError.details for the server to give instructions
+// for the client to retry.
+message RPCRetryPolicy {
+  float retry_after_secs = 1;
+}
+
 // A copy google.rpc.Status for GRPCError.details:
 // https://github.com/googleapis/googleapis/blob/master/google/rpc/status.proto
 // RPCStatus is compatible with google.rpc.Status, so one can encode messages using