modal 1.1.5.dev66__py3-none-any.whl → 1.3.1.dev8__py3-none-any.whl

modal/_utils/function_utils.py

@@ -8,7 +8,6 @@ from enum import Enum
 from pathlib import Path, PurePosixPath
 from typing import Any, Callable, Literal, Optional
 
-from grpclib import GRPCError
 from grpclib.exceptions import StreamTerminatedError
 
 import modal_proto
@@ -29,9 +28,11 @@ from ..exception import (
     DeserializationError,
     ExecutionError,
     FunctionTimeoutError,
+    InternalError,
     InternalFailure,
     InvalidError,
     RemoteError,
+    ServiceError,
 )
 from ..mount import ROOT_DIR, _is_modal_path, _Mount
 from .blob_utils import (
@@ -39,7 +40,6 @@ from .blob_utils import (
     blob_download,
     blob_upload_with_r2_failure_info,
 )
-from .grpc_utils import RETRYABLE_GRPC_STATUS_CODES
 
 if typing.TYPE_CHECKING:
     import modal._functions
@@ -75,6 +75,10 @@ def is_global_object(object_qual_name: str):
     return "<locals>" not in object_qual_name.split(".")
 
 
+def is_flash_object(experimental_options: Optional[dict[str, Any]], http_config: Optional[api_pb2.HTTPConfig]) -> bool:
+    return bool(experimental_options and experimental_options.get("flash", False)) or http_config is not None
+
+
 def is_method_fn(object_qual_name: str):
     # methods have names like Cls.foo.
     if "<locals>" in object_qual_name:
@@ -125,6 +129,7 @@ class FunctionInfo:
 
     raw_f: Optional[Callable[..., Any]]  # if None - this is a "class service function"
     function_name: str
+    implementation_name: str
     user_cls: Optional[type[Any]]
     module_name: Optional[str]
 
@@ -156,20 +161,16 @@ class FunctionInfo:
         self.raw_f = f
         self.user_cls = user_cls
 
-        if name_override is not None:
-            if not serialized:
-                # We may relax this constraint in the future, but currently we don't track the distinction between
-                # the Function's name inside modal and the name of the object that we need to import in a container.
-                raise InvalidError("Setting a custom `name=` also requires setting `serialized=True`")
-            self.function_name = name_override
-        elif f is None and user_cls:
+        if f is None and user_cls:
             # "service function" for running all methods of a class
-            self.function_name = f"{user_cls.__name__}.*"
+            self.implementation_name = f"{user_cls.__name__}.*"
         elif f and user_cls:
             # Method may be defined on superclass of the wrapped class
-            self.function_name = f"{user_cls.__name__}.{f.__name__}"
+            self.implementation_name = f"{user_cls.__name__}.{f.__name__}"
         else:
-            self.function_name = f.__qualname__
+            self.implementation_name = f.__qualname__
+
+        self.function_name = name_override or self.implementation_name
 
         # If it's a cls, the @method could be defined in a base class in a different file.
         if user_cls is not None:
@@ -436,15 +437,14 @@ async def _stream_function_call_data(
 
                 last_index = chunk.index
                 yield message
-        except (GRPCError, StreamTerminatedError) as exc:
+        except (ServiceError, InternalError, StreamTerminatedError) as exc:
             if retries_remaining > 0:
                 retries_remaining -= 1
-                if isinstance(exc, GRPCError):
-                    if exc.status in RETRYABLE_GRPC_STATUS_CODES:
-                        logger.debug(f"{variant} stream retrying with delay {delay_ms}ms due to {exc}")
-                        await asyncio.sleep(delay_ms / 1000)
-                        delay_ms = min(1000, delay_ms * 10)
-                        continue
+                if isinstance(exc, (ServiceError, InternalError)):
+                    logger.debug(f"{variant} stream retrying with delay {delay_ms}ms due to {exc}")
+                    await asyncio.sleep(delay_ms / 1000)
+                    delay_ms = min(1000, delay_ms * 10)
+                    continue
                 elif isinstance(exc, StreamTerminatedError):
                     continue
             raise
@@ -641,14 +641,13 @@ class FunctionCreationStatus:
         if not self.response:
             self.status_row.finish(f"Unknown error when creating function {self.tag}")
 
-        elif self.response.function.web_url:
+        elif web_url := self.response.handle_metadata.web_url:
             url_info = self.response.function.web_url_info
             requires_proxy_auth = self.response.function.webhook_config.requires_proxy_auth
             proxy_auth_suffix = " 🔑" if requires_proxy_auth else ""
             # Ensure terms used here match terms used in modal.com/docs/guide/webhook-urls doc.
             suffix = _get_suffix_from_web_url_info(url_info)
             # TODO: this is only printed when we're showing progress. Maybe move this somewhere else.
-            web_url = self.response.handle_metadata.web_url
             for warning in self.response.server_warnings:
                 self.status_row.warning(warning)
             self.status_row.finish(

modal/_utils/grpc_testing.py

@@ -45,8 +45,11 @@ def patch_mock_servicer(cls):
     Also patches all unimplemented abstract methods in a mock servicer with default error implementations.
     """
 
-    async def fallback(self, stream) -> None:
-        raise GRPCError(Status.UNIMPLEMENTED, "Not implemented in mock servicer " + repr(cls))
+    def fallback(name: str):
+        async def _fallback(self, stream) -> None:
+            raise GRPCError(Status.UNIMPLEMENTED, f"{name} not implemented in mock servicer " + repr(cls))
+
+        return _fallback
 
     @contextlib.contextmanager
     def intercept(servicer):
@@ -85,7 +88,7 @@ def patch_mock_servicer(cls):
     for name in dir(cls):
         method = getattr(cls, name)
         if getattr(method, "__isabstractmethod__", False):
-            setattr(cls, name, patch_grpc_method(name, fallback))
+            setattr(cls, name, patch_grpc_method(name, fallback(name)))
         elif name[0].isupper() and inspect.isfunction(method):
             setattr(cls, name, patch_grpc_method(name, method))
 

modal/_utils/grpc_utils.py

@@ -1,6 +1,7 @@
 # Copyright Modal Labs 2022
 import asyncio
 import contextlib
+import os
 import platform
 import socket
 import time
@@ -8,26 +9,26 @@ import typing
 import urllib.parse
 import uuid
 from collections.abc import AsyncIterator
-from dataclasses import dataclass
-from typing import (
-    Any,
-    Optional,
-    TypeVar,
-)
+from dataclasses import dataclass, field
+from functools import cache
+from typing import Any, Optional, Sequence, TypeVar
 
 import grpclib.client
 import grpclib.config
 import grpclib.events
-import grpclib.protocol
-import grpclib.stream
 from google.protobuf.message import Message
+from google.protobuf.symbol_database import SymbolDatabase
 from grpclib import GRPCError, Status
+from grpclib.encoding.base import StatusDetailsCodecBase
 from grpclib.exceptions import StreamTerminatedError
 from grpclib.protocol import H2Protocol
 
-from modal.exception import AuthError, ConnectionError
+from modal.exception import ConnectionError
+from modal_proto import api_pb2
 from modal_version import __version__
 
+from .._traceback import suppress_tb_frame
+from ..config import config
 from .async_utils import retry
 from .logger import logger
 
@@ -35,6 +36,7 @@ RequestType = TypeVar("RequestType", bound=Message)
 ResponseType = TypeVar("ResponseType", bound=Message)
 
 if typing.TYPE_CHECKING:
+    import modal._grpc_client
     import modal.client
 
 # Monkey patches grpclib to have a Modal User Agent header.
@@ -70,6 +72,7 @@ RETRYABLE_GRPC_STATUS_CODES = [
     Status.INTERNAL,
     Status.UNKNOWN,
 ]
+SERVER_RETRY_WARNING_TIME_INTERVAL = 30.0
 
 
 @dataclass
@@ -109,6 +112,56 @@ class ConnectionManager:
         self._channels.clear()
 
 
+@cache
+def _sym_db() -> SymbolDatabase:
+    from google.protobuf.symbol_database import Default
+
+    return Default()
+
+
+class CustomProtoStatusDetailsCodec(StatusDetailsCodecBase):
+    """grpclib compatible details codec.
+
+    The server can encode the details using `google.rpc.Status` using grpclib's default codec and this custom codec
+    can decode it into a `api_pb2.RPCStatus`.
+    """
+
+    def encode(
+        self,
+        status: Status,
+        message: Optional[str],
+        details: Optional[Sequence[Message]],
+    ) -> bytes:
+        details_proto = api_pb2.RPCStatus(code=status.value, message=message or "")
+        if details is not None:
+            for detail in details:
+                detail_container = details_proto.details.add()
+                detail_container.Pack(detail)
+        return details_proto.SerializeToString()
+
+    def decode(
+        self,
+        status: Status,
+        message: Optional[str],
+        data: bytes,
+    ) -> Any:
+        sym_db = _sym_db()
+        details_proto = api_pb2.RPCStatus.FromString(data)
+
+        details = []
+        for detail_container in details_proto.details:
+            # If we do not know how to decode an emssage, we'll ignore it.
+            with contextlib.suppress(Exception):
+                msg_type = sym_db.GetSymbol(detail_container.TypeName())
+                detail = msg_type()
+                detail_container.Unpack(detail)
+                details.append(detail)
+        return details
+
+
+custom_detail_codec = CustomProtoStatusDetailsCodec()
+
+
 def create_channel(
     server_url: str,
     metadata: dict[str, str] = {},
@@ -127,7 +180,7 @@ def create_channel(
     )
 
     if o.scheme == "unix":
-        channel = grpclib.client.Channel(path=o.path, config=config)  # probably pointless to use a pool ever
+        channel = grpclib.client.Channel(path=o.path, config=config, status_details_codec=custom_detail_codec)
     elif o.scheme in ("http", "https"):
         target = o.netloc
         parts = target.split(":")
@@ -135,7 +188,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 = grpclib.client.Channel(host, port, ssl=ssl, config=config)
+        channel = grpclib.client.Channel(host, port, ssl=ssl, config=config, status_details_codec=custom_detail_codec)
     else:
         raise Exception(f"Unknown scheme: {o.scheme}")
 
@@ -165,7 +218,7 @@ if typing.TYPE_CHECKING:
 
 
 async def unary_stream(
-    method: "modal.client.UnaryStreamWrapper[RequestType, ResponseType]",
+    method: "modal._grpc_client.UnaryStreamWrapper[RequestType, ResponseType]",
     request: RequestType,
     metadata: Optional[Any] = None,
 ) -> AsyncIterator[ResponseType]:
@@ -174,102 +227,208 @@ async def unary_stream(
         yield item
 
 
+@dataclass(frozen=True)
+class Retry:
+    base_delay: float = 0.1
+    max_delay: float = 1
+    delay_factor: float = 2
+    max_retries: Optional[int] = 3
+    additional_status_codes: list = field(default_factory=list)
+    attempt_timeout: Optional[float] = None  # timeout for each attempt
+    total_timeout: Optional[float] = None  # timeout for the entire function call
+    attempt_timeout_floor: float = 2.0  # always have at least this much timeout (only for total_timeout)
+    warning_message: Optional[RetryWarningMessage] = None
+
+
 async def retry_transient_errors(
-    fn: "modal.client.UnaryUnaryWrapper[RequestType, ResponseType]",
-    *args,
-    base_delay: float = 0.1,
-    max_delay: float = 1,
-    delay_factor: float = 2,
+    fn: "grpclib.client.UnaryUnaryMethod[RequestType, ResponseType]",
+    req: RequestType,
     max_retries: Optional[int] = 3,
-    additional_status_codes: list = [],
-    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)
-    retry_warning_message: Optional[RetryWarningMessage] = None,
-    metadata: list[tuple[str, str]] = [],
+) -> ResponseType:
+    """Minimum API version of _retry_transient_errors that works with grpclib.client.UnaryUnaryMethod.
+
+    Used by modal server.
+    """
+    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,
+    idempotency_key: str,
+):
+    """Process exception before retry, used by `_retry_transient_errors`."""
+    with suppress_tb_frame():
+        if final_attempt:
+            logger.debug(
+                f"Final attempt failed with {repr(exc)} {n_retries=} {delay=} 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
+            # Also be sure to remove the AttributeError from the set of exceptions
+            # we handle in the retry logic once we drop this check!
+            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]",
+        "grpclib.client.UnaryUnaryMethod[RequestType, ResponseType]",
+    ],
+    req: RequestType,
+    retry: Retry,
+    metadata: Optional[list[tuple[str, str]]] = None,
 ) -> ResponseType:
     """Retry on transient gRPC failures with back-off until max_retries is reached.
     If max_retries is None, retry forever."""
+    import modal._grpc_client
 
-    delay = base_delay
+    if isinstance(fn, modal._grpc_client.UnaryUnaryWrapper):
+        fn_callable = fn.direct
+    elif isinstance(fn, grpclib.client.UnaryUnaryMethod):
+        fn_callable = fn  # type: ignore
+    else:
+        raise ValueError("Only modal._grpc_client.UnaryUnaryWrapper and grpclib.client.UnaryUnaryMethod are supported")
+
+    delay = retry.base_delay
     n_retries = 0
+    n_throttled_retries = 0
 
-    status_codes = [*RETRYABLE_GRPC_STATUS_CODES, *additional_status_codes]
+    status_codes = [*RETRYABLE_GRPC_STATUS_CODES, *retry.additional_status_codes]
 
     idempotency_key = str(uuid.uuid4())
 
     t0 = time.time()
-    if total_timeout is not None:
-        total_deadline = t0 + total_timeout
+    last_server_retry_warning_time = None
+
+    if retry.total_timeout is not None:
+        total_deadline = t0 + retry.total_timeout
     else:
         total_deadline = None
 
-    metadata = metadata + [("x-modal-timestamp", str(time.time()))]
+    metadata = (metadata or []) + [("x-modal-timestamp", str(time.time()))]
+
     while True:
         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 attempt_timeout is not None:
-            timeouts.append(attempt_timeout)
-        if total_timeout is not None:
-            timeouts.append(max(total_deadline - time.time(), attempt_timeout_floor))
+        if retry.attempt_timeout is not None:
+            timeouts.append(retry.attempt_timeout)
+        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
         else:
             timeout = None
+
         try:
-            return await fn(*args, metadata=attempt_metadata, timeout=timeout)
+            with suppress_tb_frame():
+                return await fn_callable(req, metadata=attempt_metadata, timeout=timeout)
         except (StreamTerminatedError, GRPCError, OSError, asyncio.TimeoutError, AttributeError) as exc:
-            if isinstance(exc, GRPCError) and exc.status not in status_codes:
-                if exc.status == Status.UNAUTHENTICATED:
-                    raise AuthError(exc.message)
-                else:
-                    raise exc
+            # Note that we only catch AttributeError to handle a specific case that works around a bug
+            # in grpclib<=0.4.7. See above (search for `write_appdata`).
 
-            if max_retries is not None and n_retries >= max_retries:
+            # Server side instruction for retries
+            max_throttle_wait: Optional[int] = config.get("max_throttle_wait")
+            if (
+                max_throttle_wait != 0
+                and isinstance(exc, GRPCError)
+                and (server_retry_policy := get_server_retry_policy(exc))
+            ):
+                server_delay = server_retry_policy.retry_after_secs
+
+                now = time.time()
+
+                # We check if the timeout will be reached **after** the sleep, so we can raise an error early
+                # without needing to actually sleep.
+                total_timeout_will_be_reached = (
+                    retry.total_timeout is not None and (now + server_delay - t0) >= retry.total_timeout
+                )
+                max_throttle_will_be_reached = (
+                    max_throttle_wait is not None and (now + server_delay - t0) >= max_throttle_wait
+                )
+                final_attempt = total_timeout_will_be_reached or max_throttle_will_be_reached
+
+                with suppress_tb_frame():
+                    process_exception_before_retry(
+                        exc, final_attempt, fn.name, n_retries, server_delay, 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 {exc.status}{os.linesep}"
+                        f"{exc.message}{os.linesep}"
+                        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:
+                raise exc
+            if retry.max_retries is not None and n_retries >= retry.max_retries:
                 final_attempt = True
-            elif total_deadline is not None and time.time() + delay + attempt_timeout_floor >= total_deadline:
+            elif total_deadline is not None and time.time() + delay + retry.attempt_timeout_floor >= total_deadline:
                 final_attempt = True
             else:
                 final_attempt = False
 
-            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]})")
+            with suppress_tb_frame():
+                process_exception_before_retry(exc, final_attempt, fn.name, n_retries, delay, idempotency_key)
 
             n_retries += 1
 
             if (
-                retry_warning_message
-                and n_retries % retry_warning_message.warning_interval == 0
+                retry.warning_message
+                and n_retries % retry.warning_message.warning_interval == 0
                 and isinstance(exc, GRPCError)
-                and exc.status in retry_warning_message.errors_to_warn_for
+                and exc.status in retry.warning_message.errors_to_warn_for
             ):
-                logger.warning(retry_warning_message.message)
+                logger.warning(retry.warning_message.message)
 
             await asyncio.sleep(delay)
-            delay = min(delay * delay_factor, max_delay)
+            delay = min(delay * retry.delay_factor, retry.max_delay)
 
 
 def find_free_port() -> int:

modal/_utils/mount_utils.py

@@ -3,7 +3,9 @@ import posixpath
 import typing
 from collections.abc import Mapping, Sequence
 from pathlib import PurePath, PurePosixPath
-from typing import Union
+from typing import Optional, Union
+
+from typing_extensions import TypeGuard
 
 from ..cloud_bucket_mount import _CloudBucketMount
 from ..exception import InvalidError
@@ -76,3 +78,26 @@ def validate_volumes(
             )
 
     return validated_volumes
+
+
+def validate_only_modal_volumes(
+    volumes: Optional[Optional[dict[Union[str, PurePosixPath], _Volume]]],
+    caller_name: str,
+) -> Sequence[tuple[str, _Volume]]:
+    """Validate all volumes are `modal.Volume`."""
+    if volumes is None:
+        return []
+
+    validated_volumes = validate_volumes(volumes)
+
+    # Although the typing forbids `_CloudBucketMount` for type checking, one can still pass a `_CloudBucketMount`
+    # during runtime, so we'll check the type here.
+    def all_modal_volumes(
+        vols: Sequence[tuple[str, Union[_Volume, _CloudBucketMount]]],
+    ) -> TypeGuard[Sequence[tuple[str, _Volume]]]:
+        return all(isinstance(v, _Volume) for _, v in vols)
+
+    if not all_modal_volumes(validated_volumes):
+        raise InvalidError(f"{caller_name} only supports volumes that are modal.Volume")
+
+    return validated_volumes

modal/_utils/name_utils.py

@@ -1,5 +1,6 @@
 # Copyright Modal Labs 2022
 import re
+from collections.abc import Mapping
 
 from ..exception import InvalidError
 
@@ -37,7 +38,7 @@ def is_valid_tag(tag: str, max_length: int = 50) -> bool:
     return bool(re.match(pattern, tag))
 
 
-def check_tag_dict(tags: dict[str, str]) -> dict[str, str]:
+def check_tag_dict(tags: Mapping[str, str]) -> None:
     rules = (
         "\n\nTags may contain only alphanumeric characters, dashes, periods, or underscores, "
         "and must be 63 characters or less."
@@ -49,8 +50,6 @@ def check_tag_dict(tags: dict[str, str]) -> dict[str, str]:
         if not is_valid_tag(value, max_length):
             raise InvalidError(f"Invalid tag value: {value!r}.{rules}")
 
-    return tags
-
 
 def check_object_name(name: str, object_type: str) -> None:
     message = (

modal/_utils/package_utils.py

@@ -1,5 +1,4 @@
 # Copyright Modal Labs 2022
-import importlib
 import importlib.util
 import typing
 from importlib.metadata import PackageNotFoundError, files

modal/_utils/rand_pb_testing.py

@@ -45,7 +45,14 @@ def _fill(msg, desc: Descriptor, rand: Random) -> None:
         if field.containing_oneof is not None and field.name not in oneof_fields:
             continue
         is_message = field.type == FieldDescriptor.TYPE_MESSAGE
-        is_repeated = field.label == FieldDescriptor.LABEL_REPEATED
+
+        # In the C implemenation of protobuf for Python 3.14, it raises a depreation
+        # warning when labels is accessed, but it does not clean up the exception state,
+        # causing an SystemError.
+        if hasattr(field, "is_repeated"):
+            is_repeated = field.is_repeated  # type: ignore
+        else:
+            is_repeated = field.label == FieldDescriptor.LABEL_REPEATED
         if is_message:
             msg_field = getattr(msg, field.name)
             if is_repeated: