modal 1.2.1.dev8__py3-none-any.whl → 1.2.2.dev19__py3-none-any.whl

modal/_utils/function_utils.py

@@ -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]]) -> bool:
+    return experimental_options.get("flash", False) if experimental_options else False
+
+
 def is_method_fn(object_qual_name: str):
     # methods have names like Cls.foo.
     if "<locals>" in object_qual_name:

modal/_utils/grpc_utils.py

@@ -8,12 +8,8 @@ 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 typing import Any, Optional, TypeVar
 
 import grpclib.client
 import grpclib.config
@@ -28,6 +24,7 @@ from grpclib.protocol import H2Protocol
 from modal.exception import AuthError, ConnectionError
 from modal_version import __version__
 
+from .._traceback import suppress_tb_frames
 from .async_utils import retry
 from .logger import logger
 
@@ -35,6 +32,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.
@@ -165,7 +163,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,37 +172,66 @@ 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))
+
+
+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
+
+    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 = base_delay
+    delay = retry.base_delay
     n_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
+    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),
@@ -214,16 +241,17 @@ async def retry_transient_errors(
         if n_retries > 0:
             attempt_metadata.append(("x-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 retry.total_timeout is not None and 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_frames(1):
+                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:
@@ -231,45 +259,46 @@ async def retry_transient_errors(
                 else:
                     raise exc
 
-            if max_retries is not None and n_retries >= max_retries:
+            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:
+            with suppress_tb_frames(1):
+                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
 
-            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]})")
 
             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