modal 1.5.4.dev9__tar.gz → 1.5.4.dev10__tar.gz

{modal-1.5.4.dev9 → modal-1.5.4.dev10}/PKG-INFO

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

modal-1.5.4.dev10/modal/_utils/auth_token_manager.py

@@ -0,0 +1,181 @@
+# Copyright Modal Labs 2025
+import asyncio
+import base64
+import json
+import random
+import time
+from typing import Any
+
+from modal._utils.grpc_utils import DEFAULT_MAX_RETRIES, Retry
+from modal.exception import AuthError, ExecutionError, PermissionDeniedError
+from modal_proto import api_pb2, modal_api_grpc
+
+from .logger import logger
+
+AUTH_DENIED_EXCEPTIONS = (AuthError, PermissionDeniedError)
+
+
+class _AuthTokenManager:
+    """Handles fetching and refreshing of the input plane auth token."""
+
+    # Fraction of the token's lifetime to use before refreshing.
+    REFRESH_FRACTION = 0.5
+    # Width of the random band added on top of REFRESH_FRACTION, so that clients issued
+    # same-lifetime tokens at the same moment (a fan-out of containers starting together)
+    # spread their refreshes out instead of refreshing in lockstep every cycle. The band
+    # only ever delays a refresh, so the margin before expiry stays at no less than
+    # 1 - REFRESH_FRACTION - REFRESH_JITTER of the lifetime.
+    REFRESH_JITTER = 0.1
+    # Bound each AuthTokenGet attempt so a refresh can't hang.
+    AUTH_TOKEN_TIMEOUT = 5.0
+    AUTH_TOKEN_RETRY_TOTAL_TIMEOUT = 3 * AUTH_TOKEN_TIMEOUT
+    # After a failed refresh, wait before hitting the server again, growing exponentially with
+    # consecutive failures between these bounds (in seconds).
+    FAILURE_BACKOFF_BASE = 0.5
+    FAILURE_BACKOFF_MAX = 60.0
+    # 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._refresh_at = 0.0
+        self._retry_after = 0.0
+        self._backoff = self.FAILURE_BACKOFF_BASE
+        self._refresh_task: asyncio.Task | 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 that has used up its refresh fraction of its lifetime. 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._try_refresh_token()
+        elif self._should_refresh():
+            # The token is still valid, but has used up its refresh fraction of its lifetime.
+            if self._refresh_task is not None:
+                # A refresh is in flight, so someone else is refreshing. Continue to use the old token.
+                return self._token
+            else:
+                # Nobody is refreshing, so we need to fetch a new token.
+                await self._try_refresh_token()
+
+        return self._token
+
+    async def _try_refresh_token(self):
+        try:
+            await self._refresh_token()
+        except AUTH_DENIED_EXCEPTIONS:
+            raise
+        except Exception as e:
+            if not self._token:
+                raise
+            logger.warning("Auth token refresh failed; falling back to cached token: %s", e)
+
+    async def _refresh_token(self):
+        """
+        Fetch a new token from the control plane. Concurrent callers share a single request: the first caller starts it
+        and the others await the same task, so a burst of callers results in one RPC whether it succeeds or fails. The
+        task is owned by the event loop rather than by the caller that started it, so it runs to completion even if that
+        caller goes away.
+        """
+        if self._refresh_task is None:
+            # Maybe another coroutine refreshed already, or we're cooling down after a failure.
+            if self._token and not self._should_refresh():
+                return
+            # Note: the task is created in async code so it binds to the synchronicity event loop.
+            self._refresh_task = asyncio.create_task(self._fetch_token())
+            self._refresh_task.add_done_callback(self._refresh_task_done)
+        # Shielded so a cancelled caller doesn't cancel the refresh for everyone else.
+        await asyncio.shield(self._refresh_task)
+
+    def _refresh_task_done(self, task: "asyncio.Task"):
+        if self._refresh_task is task:
+            self._refresh_task = None
+        if not task.cancelled():
+            # Retrieve the exception, if any, so a failure isn't logged as unhandled when no caller is waiting.
+            task.exception()
+
+    async def _fetch_token(self):
+        """Make the AuthTokenGet request and cache its token, or arm the failure backoff."""
+        try:
+            resp: api_pb2.AuthTokenGetResponse = await self._stub.AuthTokenGet(
+                api_pb2.AuthTokenGetRequest(),
+                # No cached token to fall back on, so a failure is user-visible: retry transient errors.
+                # Otherwise one attempt, and _retry_after handles the cooldown.
+                retry=Retry(
+                    attempt_timeout=self.AUTH_TOKEN_TIMEOUT,
+                    total_timeout=self.AUTH_TOKEN_RETRY_TOTAL_TIMEOUT,
+                    max_retries=0 if self._token else DEFAULT_MAX_RETRIES,
+                ),
+            )
+            if not resp.token:
+                # Not expected
+                raise ExecutionError(
+                    "Internal error: Did not receive auth token from server. Please contact Modal support."
+                )
+            if exp := self._decode_jwt(resp.token).get("exp"):
+                expiry = float(exp)
+            else:
+                # This should never happen.
+                logger.warning("x-modal-auth-token does not contain exp field")
+                expiry = time.time() + self.DEFAULT_EXPIRY_OFFSET
+            self._token = resp.token
+            self._set_expiry(expiry)
+            self._retry_after = 0.0
+            self._backoff = self.FAILURE_BACKOFF_BASE
+        except AUTH_DENIED_EXCEPTIONS:
+            raise
+        except Exception as e:
+            # Back off (exponentially on consecutive failures) so we don't hammer a struggling server.
+            self._retry_after = time.time() + self._backoff
+            self._backoff = min(self._backoff * 2, self.FAILURE_BACKOFF_MAX)
+            raise
+
+    @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 _in_backoff(self):
+        return time.time() < self._retry_after
+
+    def _should_refresh(self):
+        # Fetch only when the token is stale/expiring and we're not in a post-failure backoff.
+        return self._needs_refresh() and not self._in_backoff()
+
+    def _set_expiry(self, expiry: float) -> None:
+        """Record a new expiry and schedule the proactive refresh a jittered fraction of the way through.
+
+        The delay is measured from the remaining lifetime as this client sees it, so a client whose clock is
+        skewed relative to the issuer still keeps a margin proportional to the lifetime it observes.
+        """
+        now = time.time()
+        self._expiry = expiry
+        remaining = max(expiry - now, 0.0)
+        fraction = self.REFRESH_FRACTION + random.random() * self.REFRESH_JITTER
+        self._refresh_at = now + remaining * fraction
+
+    def _needs_refresh(self):
+        return time.time() >= self._refresh_at
+
+    def _is_expired(self):
+        return time.time() >= self._expiry

{modal-1.5.4.dev9 → modal-1.5.4.dev10}/modal/_utils/task_command_router_client.py

@@ -292,9 +292,10 @@ class TaskCommandRouterClient:
     ) -> "TaskCommandRouterClient":
         """Initialize a TaskCommandRouterClient for a V2 sandbox.
 
-        `access` is the router access returned by SandboxCreateV2, which lets a
-        freshly created sandbox connect without a round-trip. When it is absent —
-        a re-attached sandbox, or the server could not mint a token — we ask for it.
+        `access` is the router access returned by SandboxCreateV2/SandboxRestoreV2,
+        which lets a freshly scheduled sandbox connect without a round-trip. When it
+        is absent — a re-attached sandbox, or the server could not mint a token — we
+        ask for it.
         """
         if access is not None:
             url, jwt = access.url, access.jwt

{modal-1.5.4.dev9 → modal-1.5.4.dev10}/modal/client.pyi

@@ -37,7 +37,7 @@ class _Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.5.4.dev9",
+        version: str = "1.5.4.dev10",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.
@@ -207,7 +207,7 @@ class Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.5.4.dev9",
+        version: str = "1.5.4.dev10",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.

{modal-1.5.4.dev9 → modal-1.5.4.dev10}/modal/functions.pyi

@@ -546,7 +546,7 @@ class Function(
 
     _call_generator: ___call_generator_spec
 
-    class __remote_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER]):
+    class __remote_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER]):
         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.
 
@@ -571,7 +571,7 @@ class Function(
             """
             ...
 
-    remote: __remote_spec[modal._functions.ReturnType, modal._functions.P]
+    remote: __remote_spec[modal._functions.P, modal._functions.ReturnType]
 
     class __remote_gen_spec(typing_extensions.Protocol):
         def __call__(self, /, *args, **kwargs) -> typing.Generator[typing.Any, None, None]:
@@ -621,7 +621,7 @@ class Function(
         """
         ...
 
-    class ___experimental_spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER]):
+    class ___experimental_spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER]):
         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.
 
@@ -654,7 +654,7 @@ class Function(
             """
             ...
 
-    _experimental_spawn: ___experimental_spawn_spec[modal._functions.ReturnType, modal._functions.P]
+    _experimental_spawn: ___experimental_spawn_spec[modal._functions.P, modal._functions.ReturnType]
 
     class ___spawn_map_inner_spec(typing_extensions.Protocol[P_INNER]):
         def __call__(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> None: ...
@@ -662,7 +662,7 @@ class Function(
 
     _spawn_map_inner: ___spawn_map_inner_spec[modal._functions.P]
 
-    class __spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER]):
+    class __spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER]):
         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.
 
@@ -695,7 +695,7 @@ class Function(
             """
             ...
 
-    spawn: __spawn_spec[modal._functions.ReturnType, modal._functions.P]
+    spawn: __spawn_spec[modal._functions.P, modal._functions.ReturnType]
 
     def get_raw_f(self) -> collections.abc.Callable[..., typing.Any]:
         """Return the inner Python object wrapped by this Modal Function.

{modal-1.5.4.dev9 → modal-1.5.4.dev10}/modal/sandbox.py

@@ -2288,6 +2288,9 @@ class _Sandbox(_Object, type_prefix="sb"):
 
         sandbox = await _Sandbox.from_id(restore_resp.sandbox_id, client)
         sandbox._task_id = restore_resp.task_id
+        # Seeded after from_id, whose hydration resets the field.
+        if restore_resp.HasField("command_router_access"):
+            sandbox._init_command_router_access = restore_resp.command_router_access
         return sandbox
 
     @property

{modal-1.5.4.dev9 → modal-1.5.4.dev10}/modal.egg-info/PKG-INFO

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

modal-1.5.4.dev10/modal_proto/__init__.py

@@ -0,0 +1,11 @@
+# Copyright Modal Labs 2024
+"""Modal gRPC types for Python.
+
+Warning: direct usage of Modal's gRPC API is discouraged,
+and no support or compatibility guarantees are provided.
+Message fields or entire RPCs may be changed or removed
+without notice. We recommend using official SDKs instead.
+"""
+from pkgutil import extend_path
+
+__path__ = extend_path(__path__, __name__)  # type: ignore[has-type]

{modal-1.5.4.dev9 → modal-1.5.4.dev10}/modal_proto/api_grpc.py

@@ -123,6 +123,10 @@ class ModalClientBase(abc.ABC):
     async def AuthTokenGet(self, stream: 'grpclib.server.Stream[modal_proto.api_pb2.AuthTokenGetRequest, modal_proto.api_pb2.AuthTokenGetResponse]') -> None:
         pass
 
+    @abc.abstractmethod
+    async def AutoscalerOverrideHistory(self, stream: 'grpclib.server.Stream[modal_proto.api_pb2.AutoscalerOverrideHistoryRequest, modal_proto.api_pb2.AutoscalerOverrideHistoryResponse]') -> None:
+        pass
+
     @abc.abstractmethod
     async def BlobCreate(self, stream: 'grpclib.server.Stream[modal_proto.api_pb2.BlobCreateRequest, modal_proto.api_pb2.BlobCreateResponse]') -> None:
         pass
@@ -1117,6 +1121,12 @@ class ModalClientBase(abc.ABC):
                 modal_proto.api_pb2.AuthTokenGetRequest,
                 modal_proto.api_pb2.AuthTokenGetResponse,
             ),
+            '/modal.client.ModalClient/AutoscalerOverrideHistory': grpclib.const.Handler(
+                self.AutoscalerOverrideHistory,
+                grpclib.const.Cardinality.UNARY_UNARY,
+                modal_proto.api_pb2.AutoscalerOverrideHistoryRequest,
+                modal_proto.api_pb2.AutoscalerOverrideHistoryResponse,
+            ),
             '/modal.client.ModalClient/BlobCreate': grpclib.const.Handler(
                 self.BlobCreate,
                 grpclib.const.Cardinality.UNARY_UNARY,
@@ -2533,6 +2543,12 @@ class ModalClientStub:
             modal_proto.api_pb2.AuthTokenGetRequest,
             modal_proto.api_pb2.AuthTokenGetResponse,
         )
+        self.AutoscalerOverrideHistory = grpclib.client.UnaryUnaryMethod(
+            channel,
+            '/modal.client.ModalClient/AutoscalerOverrideHistory',
+            modal_proto.api_pb2.AutoscalerOverrideHistoryRequest,
+            modal_proto.api_pb2.AutoscalerOverrideHistoryResponse,
+        )
         self.BlobCreate = grpclib.client.UnaryUnaryMethod(
             channel,
             '/modal.client.ModalClient/BlobCreate',