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

modal/client.py

@@ -6,14 +6,7 @@ import sys
 import urllib.parse
 import warnings
 from collections.abc import AsyncGenerator, AsyncIterator, Collection, Mapping
-from typing import (
-    Any,
-    ClassVar,
-    Generic,
-    Optional,
-    TypeVar,
-    Union,
-)
+from typing import Any, ClassVar, Optional, TypeVar, Union
 
 import grpclib.client
 from google.protobuf import empty_pb2
@@ -27,7 +20,8 @@ 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.grpc_utils import ConnectionManager, retry_transient_errors
+from ._utils.auth_token_manager import _AuthTokenManager
+from ._utils.grpc_utils import ConnectionManager
 from .config import _check_config, _is_remote, config, logger
 from .exception import AuthError, ClientClosed
 
@@ -78,6 +72,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 +91,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 +129,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):
@@ -155,7 +151,7 @@ class _Client:
     async def hello(self):
         """Connect to server and retrieve version information; raise appropriate error for various failures."""
         logger.debug(f"Client ({id(self)}): Starting")
-        resp = await retry_transient_errors(self.stub.ClientHello, empty_pb2.Empty())
+        resp = await self.stub.ClientHello(empty_pb2.Empty())
         print_server_warnings(resp.server_warnings)
 
     async def __aenter__(self):
@@ -264,6 +260,14 @@ class _Client:
         # Just used from tests.
         cls._client_from_env = client
 
+    async def get_input_plane_metadata(self, input_plane_region: str) -> list[tuple[str, str]]:
+        assert self._auth_token_manager, "Client must have an instance of auth token manager."
+        token = await self._auth_token_manager.get_token()
+        return [
+            ("x-modal-input-plane-region", input_plane_region),
+            ("x-modal-auth-token", token),
+        ]
+
     async def _call_safely(self, coro, readable_method: str):
         """Runs coroutine wrapped in a task that's part of the client's task context
 
@@ -350,77 +354,3 @@ class _Client:
 
 
 Client = synchronize_api(_Client)
-
-
-class UnaryUnaryWrapper(Generic[RequestType, ResponseType]):
-    # Calls a grpclib.UnaryUnaryMethod using a specific Client instance, respecting
-    # if that client is closed etc. and possibly introducing Modal-specific retry logic
-    wrapped_method: grpclib.client.UnaryUnaryMethod[RequestType, ResponseType]
-    client: _Client
-
-    def __init__(
-        self,
-        wrapped_method: grpclib.client.UnaryUnaryMethod[RequestType, ResponseType],
-        client: _Client,
-        server_url: str,
-    ):
-        self.wrapped_method = wrapped_method
-        self.client = client
-        self.server_url = server_url
-
-    @property
-    def name(self) -> str:
-        return self.wrapped_method.name
-
-    async def __call__(
-        self,
-        req: RequestType,
-        *,
-        timeout: Optional[float] = None,
-        metadata: Optional[_MetadataLike] = None,
-    ) -> ResponseType:
-        if self.client._snapshotted:
-            logger.debug(f"refreshing client after snapshot for {self.name.rsplit('/', 1)[1]}")
-            self.client = await _Client.from_env()
-
-        # Note: We override the grpclib method's channel (see grpclib's code [1]). I think this is fine
-        # since grpclib's code doesn't seem to change very much, but we could also recreate the
-        # grpclib stub if we aren't comfortable with this. The downside is then we need to cache
-        # the grpclib stub so the rest of our code becomes a bit more complicated.
-        #
-        # We need to override the channel because after the process is forked or the client is
-        # snapshotted, the existing channel may be stale / unusable.
-        #
-        # [1]: https://github.com/vmagamedov/grpclib/blob/62f968a4c84e3f64e6966097574ff0a59969ea9b/grpclib/client.py#L844
-        self.wrapped_method.channel = await self.client._get_channel(self.server_url)
-        return await self.client._call_unary(self.wrapped_method, req, timeout=timeout, metadata=metadata)
-
-
-class UnaryStreamWrapper(Generic[RequestType, ResponseType]):
-    wrapped_method: grpclib.client.UnaryStreamMethod[RequestType, ResponseType]
-
-    def __init__(
-        self,
-        wrapped_method: grpclib.client.UnaryStreamMethod[RequestType, ResponseType],
-        client: _Client,
-        server_url: str,
-    ):
-        self.wrapped_method = wrapped_method
-        self.client = client
-        self.server_url = server_url
-
-    @property
-    def name(self) -> str:
-        return self.wrapped_method.name
-
-    async def unary_stream(
-        self,
-        request,
-        metadata: Optional[Any] = None,
-    ):
-        if self.client._snapshotted:
-            logger.debug(f"refreshing client after snapshot for {self.name.rsplit('/', 1)[1]}")
-            self.client = await _Client.from_env()
-        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/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,28 +33,95 @@ class _Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.0.3.dev10",
-    ): ...
+        version: str = "1.2.3.dev7",
+    ):
+        """mdmd:hidden
+        The Modal client object is not intended to be instantiated directly by users.
+        """
+        ...
+
     def is_closed(self) -> bool: ...
     @property
-    def stub(self) -> modal_proto.modal_api_grpc.ModalClientModal: ...
-    async def get_stub(self, server_url: str) -> modal_proto.modal_api_grpc.ModalClientModal: ...
+    def stub(self) -> modal_proto.modal_api_grpc.ModalClientModal:
+        """mdmd:hidden
+        The default stub. Stubs can safely be used across forks / client snapshots.
+
+        This is useful if you want to make requests to the default Modal server in us-east, for example
+        control plane requests.
+
+        This is equivalent to client.get_stub(default_server_url), but it's cached, so it's a bit faster.
+        """
+        ...
+
+    async def get_stub(self, server_url: str) -> modal_proto.modal_api_grpc.ModalClientModal:
+        """mdmd:hidden
+        Get a stub for a specific server URL. Stubs can safely be used across forks / client snapshots.
+
+        This is useful if you want to make requests to a regional Modal server, for example low-latency
+        function calls in us-west.
+
+        This function is O(n) where n is the number of RPCs in ModalClient.
+        """
+        ...
+
     async def _open(self): ...
     async def _close(self, prep_for_restore: bool = False): ...
-    async def hello(self): ...
+    async def hello(self):
+        """Connect to server and retrieve version information; raise appropriate error for various failures."""
+        ...
+
     async def __aenter__(self): ...
     async def __aexit__(self, exc_type, exc, tb): ...
     @classmethod
-    def anonymous(cls, server_url: str) -> typing.AsyncContextManager[_Client]: ...
+    def anonymous(cls, server_url: str) -> typing.AsyncContextManager[_Client]:
+        """mdmd:hidden
+        Create a connection with no credentials; to be used for token creation.
+        """
+        ...
+
     @classmethod
-    async def from_env(cls, _override_config=None) -> _Client: ...
+    async def from_env(cls, _override_config=None) -> _Client:
+        """mdmd:hidden
+        Singleton that is instantiated from the Modal config and reused on subsequent calls.
+        """
+        ...
+
     @classmethod
-    async def from_credentials(cls, token_id: str, token_secret: str) -> _Client: ...
+    async def from_credentials(cls, token_id: str, token_secret: str) -> _Client:
+        """Constructor based on token credentials; useful for managing Modal on behalf of third-party users.
+
+        **Usage:**
+
+        ```python notest
+        client = modal.Client.from_credentials("my_token_id", "my_token_secret")
+
+        modal.Sandbox.create("echo", "hi", client=client, app=app)
+        ```
+        """
+        ...
+
     @classmethod
-    async def verify(cls, server_url: str, credentials: tuple[str, str]) -> None: ...
+    async def verify(cls, server_url: str, credentials: tuple[str, str]) -> None:
+        """mdmd:hidden
+        Check whether can the client can connect to this server with these credentials; raise if not.
+        """
+        ...
+
     @classmethod
-    def set_env_client(cls, client: typing.Optional[_Client]): ...
-    async def _call_safely(self, coro, readable_method: str): ...
+    def set_env_client(cls, client: typing.Optional[_Client]):
+        """mdmd:hidden"""
+        ...
+
+    async def get_input_plane_metadata(self, input_plane_region: str) -> list[tuple[str, str]]: ...
+    async def _call_safely(self, coro, readable_method: str):
+        """Runs coroutine wrapped in a task that's part of the client's task context
+
+        * Raises ClientClosed in case the client is closed while the coroutine is executed
+        * Logs warning if call is made outside of the event loop that the client is running in,
+          and execute without the cancellation context in that case
+        """
+        ...
+
     async def _reset_on_pid_change(self): ...
     async def _get_channel(self, server_url: str) -> grpclib.client.Channel: ...
     async def _call_unary(
@@ -87,6 +156,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__(
@@ -94,15 +164,48 @@ class Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.0.3.dev10",
-    ): ...
+        version: str = "1.2.3.dev7",
+    ):
+        """mdmd:hidden
+        The Modal client object is not intended to be instantiated directly by users.
+        """
+        ...
+
     def is_closed(self) -> bool: ...
     @property
-    def stub(self) -> modal_proto.modal_api_grpc.ModalClientModal: ...
+    def stub(self) -> modal_proto.modal_api_grpc.ModalClientModal:
+        """mdmd:hidden
+        The default stub. Stubs can safely be used across forks / client snapshots.
+
+        This is useful if you want to make requests to the default Modal server in us-east, for example
+        control plane requests.
+
+        This is equivalent to client.get_stub(default_server_url), but it's cached, so it's a bit faster.
+        """
+        ...
 
     class __get_stub_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(self, /, server_url: str) -> modal_proto.modal_api_grpc.ModalClientModal: ...
-        async def aio(self, /, server_url: str) -> modal_proto.modal_api_grpc.ModalClientModal: ...
+        def __call__(self, /, server_url: str) -> modal_proto.modal_api_grpc.ModalClientModal:
+            """mdmd:hidden
+            Get a stub for a specific server URL. Stubs can safely be used across forks / client snapshots.
+
+            This is useful if you want to make requests to a regional Modal server, for example low-latency
+            function calls in us-west.
+
+            This function is O(n) where n is the number of RPCs in ModalClient.
+            """
+            ...
+
+        async def aio(self, /, server_url: str) -> modal_proto.modal_api_grpc.ModalClientModal:
+            """mdmd:hidden
+            Get a stub for a specific server URL. Stubs can safely be used across forks / client snapshots.
+
+            This is useful if you want to make requests to a regional Modal server, for example low-latency
+            function calls in us-west.
+
+            This function is O(n) where n is the number of RPCs in ModalClient.
+            """
+            ...
 
     get_stub: __get_stub_spec[typing_extensions.Self]
 
@@ -119,8 +222,13 @@ class Client:
     _close: ___close_spec[typing_extensions.Self]
 
     class __hello_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(self, /): ...
-        async def aio(self, /): ...
+        def __call__(self, /):
+            """Connect to server and retrieve version information; raise appropriate error for various failures."""
+            ...
+
+        async def aio(self, /):
+            """Connect to server and retrieve version information; raise appropriate error for various failures."""
+            ...
 
     hello: __hello_spec[typing_extensions.Self]
 
@@ -129,19 +237,69 @@ class Client:
     def __exit__(self, exc_type, exc, tb): ...
     async def __aexit__(self, exc_type, exc, tb): ...
     @classmethod
-    def anonymous(cls, server_url: str) -> synchronicity.combined_types.AsyncAndBlockingContextManager[Client]: ...
+    def anonymous(cls, server_url: str) -> synchronicity.combined_types.AsyncAndBlockingContextManager[Client]:
+        """mdmd:hidden
+        Create a connection with no credentials; to be used for token creation.
+        """
+        ...
+
     @classmethod
-    def from_env(cls, _override_config=None) -> Client: ...
+    def from_env(cls, _override_config=None) -> Client:
+        """mdmd:hidden
+        Singleton that is instantiated from the Modal config and reused on subsequent calls.
+        """
+        ...
+
     @classmethod
-    def from_credentials(cls, token_id: str, token_secret: str) -> Client: ...
+    def from_credentials(cls, token_id: str, token_secret: str) -> Client:
+        """Constructor based on token credentials; useful for managing Modal on behalf of third-party users.
+
+        **Usage:**
+
+        ```python notest
+        client = modal.Client.from_credentials("my_token_id", "my_token_secret")
+
+        modal.Sandbox.create("echo", "hi", client=client, app=app)
+        ```
+        """
+        ...
+
     @classmethod
-    def verify(cls, server_url: str, credentials: tuple[str, str]) -> None: ...
+    def verify(cls, server_url: str, credentials: tuple[str, str]) -> None:
+        """mdmd:hidden
+        Check whether can the client can connect to this server with these credentials; raise if not.
+        """
+        ...
+
     @classmethod
-    def set_env_client(cls, client: typing.Optional[Client]): ...
+    def set_env_client(cls, client: typing.Optional[Client]):
+        """mdmd:hidden"""
+        ...
+
+    class __get_input_plane_metadata_spec(typing_extensions.Protocol[SUPERSELF]):
+        def __call__(self, /, input_plane_region: str) -> list[tuple[str, str]]: ...
+        async def aio(self, /, input_plane_region: str) -> list[tuple[str, str]]: ...
+
+    get_input_plane_metadata: __get_input_plane_metadata_spec[typing_extensions.Self]
 
     class ___call_safely_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(self, /, coro, readable_method: str): ...
-        async def aio(self, /, coro, readable_method: str): ...
+        def __call__(self, /, coro, readable_method: str):
+            """Runs coroutine wrapped in a task that's part of the client's task context
+
+            * Raises ClientClosed in case the client is closed while the coroutine is executed
+            * Logs warning if call is made outside of the event loop that the client is running in,
+              and execute without the cancellation context in that case
+            """
+            ...
+
+        async def aio(self, /, coro, readable_method: str):
+            """Runs coroutine wrapped in a task that's part of the client's task context
+
+            * Raises ClientClosed in case the client is closed while the coroutine is executed
+            * Logs warning if call is made outside of the event loop that the client is running in,
+              and execute without the cancellation context in that case
+            """
+            ...
 
     _call_safely: ___call_safely_spec[typing_extensions.Self]
 
@@ -181,43 +339,6 @@ class Client:
         ],
     ) -> collections.abc.AsyncGenerator[typing.Any, None]: ...
 
-class UnaryUnaryWrapper(typing.Generic[RequestType, ResponseType]):
-    wrapped_method: grpclib.client.UnaryUnaryMethod[RequestType, ResponseType]
-    client: _Client
-
-    def __init__(
-        self,
-        wrapped_method: grpclib.client.UnaryUnaryMethod[RequestType, ResponseType],
-        client: _Client,
-        server_url: str,
-    ): ...
-    @property
-    def name(self) -> str: ...
-    async def __call__(
-        self,
-        req: RequestType,
-        *,
-        timeout: typing.Optional[float] = None,
-        metadata: typing.Union[
-            collections.abc.Mapping[str, typing.Union[str, bytes]],
-            collections.abc.Collection[tuple[str, typing.Union[str, bytes]]],
-            None,
-        ] = None,
-    ) -> ResponseType: ...
-
-class UnaryStreamWrapper(typing.Generic[RequestType, ResponseType]):
-    wrapped_method: grpclib.client.UnaryStreamMethod[RequestType, ResponseType]
-
-    def __init__(
-        self,
-        wrapped_method: grpclib.client.UnaryStreamMethod[RequestType, ResponseType],
-        client: _Client,
-        server_url: str,
-    ): ...
-    @property
-    def name(self) -> str: ...
-    def unary_stream(self, request, metadata: typing.Optional[typing.Any] = None): ...
-
 HEARTBEAT_INTERVAL: float
 
 HEARTBEAT_TIMEOUT: float

modal/cloud_bucket_mount.py

@@ -1,6 +1,6 @@
 # Copyright Modal Labs 2022
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, Sequence
 from urllib.parse import urlparse
 
 from modal_proto import api_pb2
@@ -117,9 +117,10 @@ class _CloudBucketMount:
 
     read_only: bool = False
     requester_pays: bool = False
+    force_path_style: bool = False
 
 
-def cloud_bucket_mounts_to_proto(mounts: list[tuple[str, _CloudBucketMount]]) -> list[api_pb2.CloudBucketMount]:
+def cloud_bucket_mounts_to_proto(mounts: Sequence[tuple[str, _CloudBucketMount]]) -> list[api_pb2.CloudBucketMount]:
     """Helper function to convert `CloudBucketMount` to a list of protobufs that can be passed to the server."""
     cloud_bucket_mounts: list[api_pb2.CloudBucketMount] = []
 
@@ -132,7 +133,7 @@ def cloud_bucket_mounts_to_proto(mounts: list[tuple[str, _CloudBucketMount]]) ->
             elif parse_result.hostname.endswith("storage.googleapis.com"):
                 bucket_type = api_pb2.CloudBucketMount.BucketType.GCP
             else:
-                logger.warning(
+                logger.info(
                     "CloudBucketMount received unrecognized bucket endpoint URL. "
                     "Assuming AWS S3 configuration as fallback."
                 )
@@ -159,6 +160,7 @@ def cloud_bucket_mounts_to_proto(mounts: list[tuple[str, _CloudBucketMount]]) ->
             requester_pays=mount.requester_pays,
             key_prefix=key_prefix,
             oidc_auth_role_arn=mount.oidc_auth_role_arn,
+            force_path_style=mount.force_path_style,
         )
         cloud_bucket_mounts.append(cloud_bucket_mount)