modal 1.0.6.dev32__tar.gz → 1.0.6.dev34__tar.gz

{modal-1.0.6.dev32 → modal-1.0.6.dev34}/PKG-INFO

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

{modal-1.0.6.dev32 → modal-1.0.6.dev34}/modal/_functions.py

@@ -401,9 +401,7 @@ class _InputPlaneInvocation:
             parent_input_id=current_input_id() or "",
             input=input_item,
         )
-        metadata: list[tuple[str, str]] = []
-        if input_plane_region and input_plane_region != "":
-            metadata.append(("x-modal-input-plane-region", input_plane_region))
+        metadata = await _InputPlaneInvocation._get_metadata(input_plane_region, client)
         response = await retry_transient_errors(stub.AttemptStart, request, metadata=metadata)
         attempt_token = response.attempt_token
 
@@ -419,9 +417,7 @@ class _InputPlaneInvocation:
                 timeout_secs=OUTPUTS_TIMEOUT,
                 requested_at=time.time(),
             )
-            metadata: list[tuple[str, str]] = []
-            if self.input_plane_region and self.input_plane_region != "":
-                metadata.append(("x-modal-input-plane-region", self.input_plane_region))
+            metadata = await self._get_metadata(self.input_plane_region, self.client)
             await_response: api_pb2.AttemptAwaitResponse = await retry_transient_errors(
                 self.stub.AttemptAwait,
                 await_request,
@@ -457,6 +453,12 @@ class _InputPlaneInvocation:
                     await_response.output.result, await_response.output.data_format, control_plane_stub, self.client
                 )
 
+    @staticmethod
+    async def _get_metadata(input_plane_region: str, client: _Client) -> list[tuple[str, str]]:
+        if not input_plane_region:
+            return []
+        token = await client._auth_token_manager.get_token()
+        return [("x-modal-input-plane-region", input_plane_region), ("x-modal-auth-token", token)]
 
 # Wrapper type for api_pb2.FunctionStats
 @dataclass(frozen=True)

modal-1.0.6.dev34/modal/_utils/auth_token_manager.py

@@ -0,0 +1,114 @@
+# 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 .grpc_utils import retry_transient_errors
+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):
+        """
+        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 retry_transient_errors(
+                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

{modal-1.0.6.dev32 → modal-1.0.6.dev34}/modal/_utils/grpc_utils.py

@@ -149,21 +149,7 @@ def create_channel(
 
         logger.debug(f"Sending request to {event.method_name}")
 
-    async def recv_initial_metadata(initial_metadata: grpclib.events.RecvInitialMetadata) -> None:
-        # If we receive an auth token from the server, include it in all future requests.
-        # TODO(nathan): This isn't perfect because the metadata isn't propagated when the
-        # process is forked and a new channel is created. This is OK for now since this
-        # token is only used by the experimental input plane
-        if token := initial_metadata.metadata.get("x-modal-auth-token"):
-            metadata["x-modal-auth-token"] = str(token)
-
-    async def recv_trailing_metadata(trailing_metadata: grpclib.events.RecvTrailingMetadata) -> None:
-        if token := trailing_metadata.metadata.get("x-modal-auth-token"):
-            metadata["x-modal-auth-token"] = str(token)
-
     grpclib.events.listen(channel, grpclib.events.SendRequest, send_request)
-    grpclib.events.listen(channel, grpclib.events.RecvInitialMetadata, recv_initial_metadata)
-    grpclib.events.listen(channel, grpclib.events.RecvTrailingMetadata, recv_trailing_metadata)
 
     return channel
 

{modal-1.0.6.dev32 → modal-1.0.6.dev34}/modal/client.py

@@ -27,6 +27,7 @@ from modal_version import __version__
 from ._traceback import print_server_warnings
 from ._utils import async_utils
 from ._utils.async_utils import TaskContext, synchronize_api
+from ._utils.auth_token_manager import _AuthTokenManager
 from ._utils.grpc_utils import ConnectionManager, retry_transient_errors
 from .config import _check_config, _is_remote, config, logger
 from .exception import AuthError, ClientClosed
@@ -78,6 +79,7 @@ class _Client:
     _cancellation_context: TaskContext
     _cancellation_context_event_loop: asyncio.AbstractEventLoop = None
     _stub: Optional[api_grpc.ModalClientStub]
+    _auth_token_manager: _AuthTokenManager = None
     _snapshotted: bool
 
     def __init__(
@@ -96,6 +98,7 @@ class _Client:
         self.version = version
         self._closed = False
         self._stub: Optional[modal_api_grpc.ModalClientModal] = None
+        self._auth_token_manager: Optional[_AuthTokenManager] = None
         self._snapshotted = False
         self._owner_pid = None
 
@@ -133,9 +136,9 @@ class _Client:
         self._cancellation_context = TaskContext(grace=0.5)  # allow running rpcs to finish in 0.5s when closing client
         self._cancellation_context_event_loop = asyncio.get_running_loop()
         await self._cancellation_context.__aenter__()
-
         self._connection_manager = ConnectionManager(client=self, metadata=metadata)
         self._stub = await self.get_stub(self.server_url)
+        self._auth_token_manager = _AuthTokenManager(self.stub)
         self._owner_pid = os.getpid()
 
     async def _close(self, prep_for_restore: bool = False):
@@ -424,3 +427,4 @@ class UnaryStreamWrapper(Generic[RequestType, ResponseType]):
         self.wrapped_method.channel = await self.client._get_channel(self.server_url)
         async for response in self.client._call_stream(self.wrapped_method, request, metadata=metadata):
             yield response
+

{modal-1.0.6.dev32 → modal-1.0.6.dev34}/modal/client.pyi

@@ -4,6 +4,7 @@ import collections.abc
 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
@@ -24,6 +25,7 @@ class _Client:
     _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
     _snapshotted: bool
 
     def __init__(
@@ -31,7 +33,7 @@ class _Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.0.6.dev32",
+        version: str = "1.0.6.dev34",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.
@@ -153,6 +155,7 @@ class Client:
     _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
     _snapshotted: bool
 
     def __init__(
@@ -160,7 +163,7 @@ class Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.0.6.dev32",
+        version: str = "1.0.6.dev34",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.

{modal-1.0.6.dev32 → modal-1.0.6.dev34}/modal/sandbox.py

@@ -578,7 +578,9 @@ class _Sandbox(_Object, type_prefix="sb"):
 
     async def _get_task_id(self) -> str:
         while not self._task_id:
-            resp = await self._client.stub.SandboxGetTaskId(api_pb2.SandboxGetTaskIdRequest(sandbox_id=self.object_id))
+            resp = await retry_transient_errors(
+                self._client.stub.SandboxGetTaskId, api_pb2.SandboxGetTaskIdRequest(sandbox_id=self.object_id)
+            )
             self._task_id = resp.task_id
             if not self._task_id:
                 await asyncio.sleep(0.5)

{modal-1.0.6.dev32 → modal-1.0.6.dev34}/modal.egg-info/PKG-INFO

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

{modal-1.0.6.dev32 → modal-1.0.6.dev34}/modal.egg-info/SOURCES.txt

@@ -101,6 +101,7 @@ modal/_runtime/user_code_imports.py
 modal/_utils/__init__.py
 modal/_utils/app_utils.py
 modal/_utils/async_utils.py
+modal/_utils/auth_token_manager.py
 modal/_utils/blob_utils.py
 modal/_utils/bytes_io_segment_payload.py
 modal/_utils/deprecation.py

{modal-1.0.6.dev32 → modal-1.0.6.dev34}/modal_version/__init__.py

@@ -1,4 +1,4 @@
 # Copyright Modal Labs 2025
 """Supplies the current version of the modal client library."""
 
-__version__ = "1.0.6.dev32"
+__version__ = "1.0.6.dev34"