modal 0.62.16__py3-none-any.whl → 0.72.11__py3-none-any.whl

modal/_utils/grpc_testing.py

@@ -4,7 +4,8 @@ import inspect
 import logging
 import typing
 from collections import Counter, defaultdict
-from typing import Any, Awaitable, Callable, Dict, List, Tuple
+from collections.abc import Awaitable
+from typing import Any, Callable
 
 import grpclib.server
 from grpclib import GRPCError, Status
@@ -26,7 +27,8 @@ def patch_mock_servicer(cls):
         await some_complex_method()
     assert ctx.calls == [("SomeMethod", MyMessage(foo="bar"))]
     ```
-    Also allows to set a predefined queue of responses, temporarily replacing a mock servicer's default responses for a method:
+    Also allows to set a predefined queue of responses, temporarily replacing
+    a mock servicer's default responses for a method:
 
     ```python notest
     with servicer.intercept() as ctx:
@@ -48,10 +50,10 @@ def patch_mock_servicer(cls):
 
     @contextlib.contextmanager
     def intercept(servicer):
-        ctx = InterceptionContext()
+        ctx = InterceptionContext(servicer)
         servicer.interception_context = ctx
         yield ctx
-        ctx.assert_responses_consumed()
+        ctx._assert_responses_consumed()
         servicer.interception_context = None
 
     cls.intercept = intercept
@@ -63,7 +65,7 @@ def patch_mock_servicer(cls):
                 ctx = servicer_self.interception_context
                 if ctx:
                     intercepted_stream = await InterceptedStream(ctx, method_name, stream).initialize()
-                    custom_responder = ctx.next_custom_responder(method_name, intercepted_stream.request_message)
+                    custom_responder = ctx._next_custom_responder(method_name, intercepted_stream.request_message)
                     if custom_responder:
                         return await custom_responder(servicer_self, intercepted_stream)
                     else:
@@ -92,31 +94,36 @@ def patch_mock_servicer(cls):
 
 
 class ResponseNotConsumed(Exception):
-    def __init__(self, unconsumed_requests: List[str]):
+    def __init__(self, unconsumed_requests: list[str]):
         self.unconsumed_requests = unconsumed_requests
         request_count = Counter(unconsumed_requests)
         super().__init__(f"Expected but did not receive the following requests: {request_count}")
 
 
 class InterceptionContext:
-    def __init__(self):
-        self.calls: List[Tuple[str, Any]] = []  # List[Tuple[method_name, message]]
-        self.custom_responses: Dict[str, List[Tuple[Callable[[Any], bool], List[Any]]]] = defaultdict(list)
-        self.custom_defaults: Dict[str, Callable[["MockClientServicer", grpclib.server.Stream], Awaitable[None]]] = {}
-
-    def add_recv(self, method_name: str, msg):
-        self.calls.append((method_name, msg))
+    def __init__(self, servicer):
+        self._servicer = servicer
+        self.calls: list[tuple[str, Any]] = []  # List[Tuple[method_name, message]]
+        self.custom_responses: dict[str, list[tuple[Callable[[Any], bool], list[Any]]]] = defaultdict(list)
+        self.custom_defaults: dict[str, Callable[["MockClientServicer", grpclib.server.Stream], Awaitable[None]]] = {}
 
     def add_response(
         self, method_name: str, first_payload, *, request_filter: Callable[[Any], bool] = lambda req: True
     ):
-        # adds one response to a queue of responses for requests of the specified type
+        """Adds one response payload to an expected queue of responses for a method.
+
+        These responses will be used once each instead of calling the MockServicer's
+        implementation of the method.
+
+        The interception context will throw an exception on exit if not all of the added
+        responses have been consumed.
+        """
         self.custom_responses[method_name].append((request_filter, [first_payload]))
 
     def set_responder(
         self, method_name: str, responder: Callable[["MockClientServicer", grpclib.server.Stream], Awaitable[None]]
     ):
-        """Replace the default responder method. E.g.
+        """Replace the default responder from the MockClientServicer with a custom implementation
 
         ```python notest
         def custom_responder(servicer, stream):
@@ -127,11 +134,31 @@ class InterceptionContext:
             ctx.set_responder("SomeMethod", custom_responder)
         ```
 
-        Responses added via `.add_response()` take precedence.
+        Responses added via `.add_response()` take precedence over the use of this replacement
         """
         self.custom_defaults[method_name] = responder
 
-    def next_custom_responder(self, method_name, request):
+    def pop_request(self, method_name):
+        # fast forward to the next request of type method_name
+        # dropping any preceding requests if there is a match
+        # returns the payload of the request
+        for i, (_method_name, msg) in enumerate(self.calls):
+            if _method_name == method_name:
+                self.calls = self.calls[i + 1 :]
+                return msg
+
+        raise KeyError(f"No message of that type in call list: {self.calls}")
+
+    def get_requests(self, method_name: str) -> list[Any]:
+        if not hasattr(self._servicer, method_name):
+            # we check this to prevent things like `assert ctx.get_requests("ASdfFunctionCreate") == 0` passing
+            raise ValueError(f"{method_name} not in MockServicer - did you spell it right?")
+        return [msg for _method_name, msg in self.calls if _method_name == method_name]
+
+    def _add_recv(self, method_name: str, msg):
+        self.calls.append((method_name, msg))
+
+    def _next_custom_responder(self, method_name, request):
         method_responses = self.custom_responses[method_name]
         for i, (request_filter, response_messages) in enumerate(method_responses):
             try:
@@ -158,7 +185,7 @@ class InterceptionContext:
 
         return responder
 
-    def assert_responses_consumed(self):
+    def _assert_responses_consumed(self):
         unconsumed = []
         for method_name, queued_responses in self.custom_responses.items():
             unconsumed += [method_name] * len(queued_responses)
@@ -166,20 +193,9 @@ class InterceptionContext:
         if unconsumed:
             raise ResponseNotConsumed(unconsumed)
 
-    def pop_request(self, method_name):
-        # fast forward to the next request of type method_name
-        # dropping any preceding requests if there is a match
-        # returns the payload of the request
-        for i, (_method_name, msg) in enumerate(self.calls):
-            if _method_name == method_name:
-                self.calls = self.calls[i + 1 :]
-                return msg
-
-        raise Exception(f"No message of that type in call list: {self.calls}")
-
 
 class InterceptedStream:
-    def __init__(self, interception_context, method_name, stream):
+    def __init__(self, interception_context: InterceptionContext, method_name: str, stream):
         self.interception_context = interception_context
         self.method_name = method_name
         self.stream = stream
@@ -196,7 +212,7 @@ class InterceptedStream:
             return ret
 
         msg = await self.stream.recv_message()
-        self.interception_context.add_recv(self.method_name, msg)
+        self.interception_context._add_recv(self.method_name, msg)
         return msg
 
     async def send_message(self, msg):

modal/_utils/grpc_utils.py

@@ -4,30 +4,37 @@ import contextlib
 import platform
 import socket
 import time
+import typing
 import urllib.parse
 import uuid
+from collections.abc import AsyncIterator
 from typing import (
     Any,
-    AsyncIterator,
-    Dict,
-    List,
     Optional,
-    Type,
     TypeVar,
 )
 
 import grpclib.client
 import grpclib.config
 import grpclib.events
+import grpclib.protocol
+import grpclib.stream
 from google.protobuf.message import Message
 from grpclib import GRPCError, Status
 from grpclib.exceptions import StreamTerminatedError
 from grpclib.protocol import H2Protocol
 
+from modal.exception import AuthError, ConnectionError
 from modal_version import __version__
 
 from .logger import logger
 
+RequestType = TypeVar("RequestType", bound=Message)
+ResponseType = TypeVar("ResponseType", bound=Message)
+
+if typing.TYPE_CHECKING:
+    import modal.client
+
 # Monkey patches grpclib to have a Modal User Agent header.
 grpclib.client.USER_AGENT = "modal-client/{version} ({sys}; {py}/{py_ver})'".format(
     version=__version__,
@@ -54,81 +61,6 @@ class Subchannel:
         return True
 
 
-class ChannelPool(grpclib.client.Channel):
-    """Use multiple channels under the hood. A drop-in replacement for the grpclib Channel.
-
-    The main reason is to get around limitations with TCP connections over the internet,
-    in particular idle timeouts.
-
-    The algorithm is very simple. It reuses the last subchannel as long as it has had less
-    than 64 requests or if it was created less than 30s ago. It closes any subchannel that
-    hits 90s age. This means requests using the ChannelPool can't be longer than 60s.
-    """
-
-    _max_requests: int
-    _max_lifetime: float
-    _max_active: float
-    _subchannels: List[Subchannel]
-
-    def __init__(
-        self,
-        *args,
-        max_requests=64,  # Maximum number of total requests per subchannel
-        max_active=30,  # Don't accept more connections on the subchannel after this many seconds
-        max_lifetime=90,  # Close subchannel after this many seconds
-        **kwargs,
-    ):
-        self._subchannels = []
-        self._max_requests = max_requests
-        self._max_active = max_active
-        self._max_lifetime = max_lifetime
-        super().__init__(*args, **kwargs)
-
-    async def __connect__(self):
-        now = time.time()
-        # Remove any closed subchannels
-        while len(self._subchannels) > 0 and not self._subchannels[-1].connected():
-            self._subchannels.pop()
-
-        # Close and delete any subchannels that are past their lifetime
-        while len(self._subchannels) > 0 and now - self._subchannels[0].created_at > self._max_lifetime:
-            self._subchannels.pop(0).protocol.processor.close()
-
-        # See if we can reuse the last subchannel
-        create_subchannel = None
-        if len(self._subchannels) > 0:
-            if self._subchannels[-1].created_at < now - self._max_active:
-                # Don't reuse subchannel that's too old
-                create_subchannel = True
-            elif self._subchannels[-1].requests > self._max_requests:
-                create_subchannel = True
-            else:
-                create_subchannel = False
-        else:
-            create_subchannel = True
-
-        # Create new if needed
-        # There's a theoretical race condition here.
-        # This is harmless but may lead to superfluous protocols.
-        if create_subchannel:
-            protocol = await self._create_connection()
-            self._subchannels.append(Subchannel(protocol))
-
-        self._subchannels[-1].requests += 1
-        return self._subchannels[-1].protocol
-
-    def close(self) -> None:
-        while len(self._subchannels) > 0:
-            self._subchannels.pop(0).protocol.processor.close()
-
-    def __del__(self) -> None:
-        if len(self._subchannels) > 0:
-            logger.warning("Channel pool not properly closed")
-
-
-_SendType = TypeVar("_SendType")
-_RecvType = TypeVar("_RecvType")
-
 RETRYABLE_GRPC_STATUS_CODES = [
     Status.DEADLINE_EXCEEDED,
     Status.UNAVAILABLE,
@@ -139,9 +71,7 @@ RETRYABLE_GRPC_STATUS_CODES = [
 
 def create_channel(
     server_url: str,
-    metadata: Dict[str, str] = {},
-    *,
-    use_pool: Optional[bool] = None,  # If None, inferred from the scheme
+    metadata: dict[str, str] = {},
 ) -> grpclib.client.Channel:
     """Creates a grpclib.Channel.
 
@@ -150,15 +80,6 @@ def create_channel(
     """
     o = urllib.parse.urlparse(server_url)
 
-    if use_pool is None:
-        use_pool = o.scheme in ("http", "https")
-
-    channel_cls: Type[grpclib.client.Channel]
-    if use_pool:
-        channel_cls = ChannelPool
-    else:
-        channel_cls = grpclib.client.Channel
-
     channel: grpclib.client.Channel
     config = grpclib.config.Configuration(
         http2_connection_window_size=64 * 1024 * 1024,  # 64 MiB
@@ -166,7 +87,7 @@ def create_channel(
     )
 
     if o.scheme == "unix":
-        channel = channel_cls(path=o.path, config=config)  # probably pointless to use a pool ever
+        channel = grpclib.client.Channel(path=o.path, config=config)  # probably pointless to use a pool ever
     elif o.scheme in ("http", "https"):
         target = o.netloc
         parts = target.split(":")
@@ -174,7 +95,7 @@ def create_channel(
         ssl = o.scheme.endswith("s")
         host = parts[0]
         port = int(parts[1]) if len(parts) == 2 else 443 if ssl else 80
-        channel = channel_cls(host, port, ssl=ssl, config=config)
+        channel = grpclib.client.Channel(host, port, ssl=ssl, config=config)
     else:
         raise Exception(f"Unknown scheme: {o.scheme}")
 
@@ -189,23 +110,31 @@ def create_channel(
         logger.debug(f"Sending request to {event.method_name}")
 
     grpclib.events.listen(channel, grpclib.events.SendRequest, send_request)
+
     return channel
 
 
+async def connect_channel(channel: grpclib.client.Channel):
+    """Connects socket (potentially raising errors raising to connectivity."""
+    await channel.__connect__()
+
+
+if typing.TYPE_CHECKING:
+    import modal.client
+
+
 async def unary_stream(
-    method: grpclib.client.UnaryStreamMethod[_SendType, _RecvType],
-    request: _SendType,
+    method: "modal.client.UnaryStreamWrapper[RequestType, ResponseType]",
+    request: RequestType,
     metadata: Optional[Any] = None,
-) -> AsyncIterator[_RecvType]:
-    """Helper for making a unary-streaming gRPC request."""
-    async with method.open(metadata=metadata) as stream:
-        await stream.send_message(request, end=True)
-        async for item in stream:
-            yield item
+) -> AsyncIterator[ResponseType]:
+    # TODO: remove this, since we have a method now
+    async for item in method.unary_stream(request, metadata):
+        yield item
 
 
 async def retry_transient_errors(
-    fn,
+    fn: "modal.client.UnaryUnaryWrapper[RequestType, ResponseType]",
     *args,
     base_delay: float = 0.1,
     max_delay: float = 1,
@@ -215,7 +144,7 @@ async def retry_transient_errors(
     attempt_timeout: Optional[float] = None,  # timeout for each attempt
     total_timeout: Optional[float] = None,  # timeout for the entire function call
     attempt_timeout_floor=2.0,  # always have at least this much timeout (only for total_timeout)
-):
+) -> ResponseType:
     """Retry on transient gRPC failures with back-off until max_retries is reached.
     If max_retries is None, retry forever."""
 
@@ -247,16 +176,35 @@ async def retry_transient_errors(
             timeout = None
         try:
             return await fn(*args, metadata=metadata, timeout=timeout)
-        except (StreamTerminatedError, GRPCError, socket.gaierror, asyncio.TimeoutError) as exc:
+        except (StreamTerminatedError, GRPCError, OSError, asyncio.TimeoutError, AttributeError) as exc:
             if isinstance(exc, GRPCError) and exc.status not in status_codes:
-                raise exc
+                if exc.status == Status.UNAUTHENTICATED:
+                    raise AuthError(exc.message)
+                else:
+                    raise exc
 
             if max_retries is not None and n_retries >= max_retries:
+                final_attempt = True
+            elif total_deadline is not None and time.time() + delay + attempt_timeout_floor >= total_deadline:
+                final_attempt = True
+            else:
+                final_attempt = False
+
+            if final_attempt:
+                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
 
-            if total_deadline and time.time() + delay + attempt_timeout_floor >= total_deadline:
-                # no point sleeping if that's going to push us past the deadline
-                raise exc
+            logger.debug(f"Retryable failure {repr(exc)} {n_retries=} {delay=} for {fn.name}")
 
             n_retries += 1
 
@@ -265,7 +213,12 @@ async def retry_transient_errors(
 
 
 def find_free_port() -> int:
-    """Find a free TCP port, useful for testing."""
+    """
+    Find a free TCP port, useful for testing.
+
+    WARN: if a returned free port is not bound immediately by the caller, that same port
+    may be returned in subsequent calls to this function, potentially creating port collisions.
+    """
     with contextlib.closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
         s.bind(("", 0))
         s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

modal/_utils/hash_utils.py

@@ -2,43 +2,53 @@
 import base64
 import dataclasses
 import hashlib
-from typing import IO, Union
+import time
+from typing import BinaryIO, Callable, Optional, Sequence, Union
 
-HASH_CHUNK_SIZE = 4096
+from modal.config import logger
 
+HASH_CHUNK_SIZE = 65536
 
-def _update(hashers, data: Union[bytes, IO[bytes]]):
+
+def _update(hashers: Sequence[Callable[[bytes], None]], data: Union[bytes, BinaryIO]) -> None:
     if isinstance(data, bytes):
         for hasher in hashers:
-            hasher.update(data)
+            hasher(data)
     else:
+        assert not isinstance(data, (bytearray, memoryview))  # https://github.com/microsoft/pyright/issues/5697
         pos = data.tell()
-        while 1:
+        while True:
             chunk = data.read(HASH_CHUNK_SIZE)
             if not isinstance(chunk, bytes):
                 raise ValueError(f"Only accepts bytes or byte buffer objects, not {type(chunk)} buffers")
             if not chunk:
                 break
             for hasher in hashers:
-                hasher.update(chunk)
+                hasher(chunk)
         data.seek(pos)
 
 
-def get_sha256_hex(data: Union[bytes, IO[bytes]]) -> str:
+def get_sha256_hex(data: Union[bytes, BinaryIO]) -> str:
+    t0 = time.monotonic()
     hasher = hashlib.sha256()
-    _update([hasher], data)
+    _update([hasher.update], data)
+    logger.debug("get_sha256_hex took %.3fs", time.monotonic() - t0)
     return hasher.hexdigest()
 
 
-def get_sha256_base64(data: Union[bytes, IO[bytes]]) -> str:
+def get_sha256_base64(data: Union[bytes, BinaryIO]) -> str:
+    t0 = time.monotonic()
     hasher = hashlib.sha256()
-    _update([hasher], data)
+    _update([hasher.update], data)
+    logger.debug("get_sha256_base64 took %.3fs", time.monotonic() - t0)
     return base64.b64encode(hasher.digest()).decode("ascii")
 
 
-def get_md5_base64(data: Union[bytes, IO[bytes]]) -> str:
+def get_md5_base64(data: Union[bytes, BinaryIO]) -> str:
+    t0 = time.monotonic()
     hasher = hashlib.md5()
-    _update([hasher], data)
+    _update([hasher.update], data)
+    logger.debug("get_md5_base64 took %.3fs", time.monotonic() - t0)
     return base64.b64encode(hasher.digest()).decode("utf-8")
 
 
@@ -47,12 +57,44 @@ class UploadHashes:
     md5_base64: str
     sha256_base64: str
 
+    def md5_hex(self) -> str:
+        return base64.b64decode(self.md5_base64).hex()
+
+    def sha256_hex(self) -> str:
+        return base64.b64decode(self.sha256_base64).hex()
+
+
+def get_upload_hashes(
+    data: Union[bytes, BinaryIO], sha256_hex: Optional[str] = None, md5_hex: Optional[str] = None
+) -> UploadHashes:
+    t0 = time.monotonic()
+    hashers = {}
+
+    if not sha256_hex:
+        sha256 = hashlib.sha256()
+        hashers["sha256"] = sha256
+    if not md5_hex:
+        md5 = hashlib.md5()
+        hashers["md5"] = md5
+
+    if hashers:
+        updaters = [h.update for h in hashers.values()]
+        _update(updaters, data)
 
-def get_upload_hashes(data: Union[bytes, IO[bytes]]) -> UploadHashes:
-    md5 = hashlib.md5()
-    sha256 = hashlib.sha256()
-    _update([md5, sha256], data)
-    return UploadHashes(
-        md5_base64=base64.b64encode(md5.digest()).decode("ascii"),
-        sha256_base64=base64.b64encode(sha256.digest()).decode("ascii"),
+    if sha256_hex:
+        sha256_base64 = base64.b64encode(bytes.fromhex(sha256_hex)).decode("ascii")
+    else:
+        sha256_base64 = base64.b64encode(hashers["sha256"].digest()).decode("ascii")
+
+    if md5_hex:
+        md5_base64 = base64.b64encode(bytes.fromhex(md5_hex)).decode("ascii")
+    else:
+        md5_base64 = base64.b64encode(hashers["md5"].digest()).decode("ascii")
+
+    hashes = UploadHashes(
+        md5_base64=md5_base64,
+        sha256_base64=sha256_base64,
     )
+
+    logger.debug("get_upload_hashes took %.3fs (%s)", time.monotonic() - t0, hashers.keys())
+    return hashes

modal/_utils/http_utils.py

@@ -1,16 +1,18 @@
 # Copyright Modal Labs 2022
 import contextlib
-import socket
-import ssl
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
-import certifi
-from aiohttp import ClientSession, ClientTimeout, TCPConnector
-from aiohttp.web import Application
-from aiohttp.web_runner import AppRunner, SockSite
+# Note: importing aiohttp seems to take about 100ms, and it's not really necessarily,
+# unless we need to work with blobs. So that's why we import it lazily instead.
 
+if TYPE_CHECKING:
+    from aiohttp import ClientSession
+    from aiohttp.web import Application
 
-def http_client_with_tls(timeout: Optional[float]) -> ClientSession:
+from .async_utils import on_shutdown
+
+
+def _http_client_with_tls(timeout: Optional[float]) -> "ClientSession":
     """Create a new HTTP client session with standard, bundled TLS certificates.
 
     This is necessary to prevent client issues on some system where Python does
@@ -20,15 +22,43 @@ def http_client_with_tls(timeout: Optional[float]) -> ClientSession:
     Specifically: the error "unable to get local issuer certificate" when making
     an aiohttp request.
     """
+    import ssl
+
+    import certifi
+    from aiohttp import ClientSession, ClientTimeout, TCPConnector
+
     ssl_context = ssl.create_default_context(cafile=certifi.where())
     connector = TCPConnector(ssl=ssl_context)
     return ClientSession(connector=connector, timeout=ClientTimeout(total=timeout))
 
 
+class ClientSessionRegistry:
+    _client_session: "ClientSession"
+    _client_session_active: bool = False
+
+    @staticmethod
+    def get_session():
+        if not ClientSessionRegistry._client_session_active:
+            ClientSessionRegistry._client_session = _http_client_with_tls(timeout=None)
+            ClientSessionRegistry._client_session_active = True
+            on_shutdown(ClientSessionRegistry.close_session())
+        return ClientSessionRegistry._client_session
+
+    @staticmethod
+    async def close_session():
+        if ClientSessionRegistry._client_session_active:
+            await ClientSessionRegistry._client_session.close()
+            ClientSessionRegistry._client_session_active = False
+
+
 @contextlib.asynccontextmanager
-async def run_temporary_http_server(app: Application):
+async def run_temporary_http_server(app: "Application"):
     # Allocates a random port, runs a server in a context manager
     # This is used in various tests
+    import socket
+
+    from aiohttp.web_runner import AppRunner, SockSite
+
     sock = socket.socket()
     sock.bind(("", 0))
     port = sock.getsockname()[1]

modal/_utils/logger.py

@@ -17,7 +17,8 @@ def configure_logger(logger: logging.Logger, log_level: str, log_format: str):
         json_formatter = jsonlogger.JsonFormatter(
             fmt=(
                 "%(asctime)s %(levelname)s [%(name)s] [%(filename)s:%(lineno)d] "
-                "[dd.service=%(dd.service)s dd.env=%(dd.env)s dd.version=%(dd.version)s dd.trace_id=%(dd.trace_id)s dd.span_id=%(dd.span_id)s] "
+                "[dd.service=%(dd.service)s dd.env=%(dd.env)s dd.version=%(dd.version)s dd.trace_id=%(dd.trace_id)s "
+                "dd.span_id=%(dd.span_id)s] "
                 "- %(message)s"
             ),
             datefmt="%Y-%m-%dT%H:%M:%S%z",