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

modal/_utils/function_utils.py

@@ -1,8 +1,8 @@
 # Copyright Modal Labs 2022
 import asyncio
-import enum
 import inspect
 import os
+import typing
 from collections.abc import AsyncGenerator
 from enum import Enum
 from pathlib import Path, PurePosixPath
@@ -10,7 +10,6 @@ from typing import Any, Callable, Literal, Optional
 
 from grpclib import GRPCError
 from grpclib.exceptions import StreamTerminatedError
-from synchronicity.exceptions import UserCodeException
 
 import modal_proto
 from modal_proto import api_pb2
@@ -19,7 +18,9 @@ from modal_proto.modal_api_grpc import ModalClientModal
 from .._serialization import (
     deserialize,
     deserialize_data_format,
+    get_preferred_payload_format,
     serialize,
+    serialize_data_format as _serialize_data_format,
     signature_to_parameter_specs,
 )
 from .._traceback import append_modal_tb
@@ -33,9 +34,16 @@ from ..exception import (
     RemoteError,
 )
 from ..mount import ROOT_DIR, _is_modal_path, _Mount
-from .blob_utils import MAX_OBJECT_SIZE_BYTES, blob_download, blob_upload
+from .blob_utils import (
+    MAX_ASYNC_OBJECT_SIZE_BYTES,
+    blob_download,
+    blob_upload_with_r2_failure_info,
+)
 from .grpc_utils import RETRYABLE_GRPC_STATUS_CODES
 
+if typing.TYPE_CHECKING:
+    import modal._functions
+
 
 class FunctionInfoType(Enum):
     PACKAGE = "package"
@@ -67,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:
@@ -383,9 +395,19 @@ def callable_has_non_self_non_default_params(f: Callable[..., Any]) -> bool:
 
 
 async def _stream_function_call_data(
-    client, function_call_id: str, variant: Literal["data_in", "data_out"]
+    client,
+    stub,
+    function_call_id: Optional[str],
+    variant: Literal["data_in", "data_out"],
+    attempt_token: Optional[str] = None,
 ) -> AsyncGenerator[Any, None]:
     """Read from the `data_in` or `data_out` stream of a function call."""
+    if not function_call_id and not attempt_token:
+        raise ValueError("function_call_id or attempt_token is required to read from a data stream")
+
+    if stub is None:
+        stub = client.stub
+
     last_index = 0
 
     # TODO(gongy): generalize this logic as util for unary streams
@@ -393,14 +415,19 @@ async def _stream_function_call_data(
     delay_ms = 1
 
     if variant == "data_in":
-        stub_fn = client.stub.FunctionCallGetDataIn
+        stub_fn = stub.FunctionCallGetDataIn
     elif variant == "data_out":
-        stub_fn = client.stub.FunctionCallGetDataOut
+        stub_fn = stub.FunctionCallGetDataOut
     else:
         raise ValueError(f"Invalid variant {variant}")
 
     while True:
-        req = api_pb2.FunctionCallGetDataRequest(function_call_id=function_call_id, last_index=last_index)
+        req = api_pb2.FunctionCallGetDataRequest(
+            function_call_id=function_call_id,
+            last_index=last_index,
+        )
+        if attempt_token:
+            req.attempt_token = attempt_token  # oneof clears function_call_id.
         try:
             async for chunk in stub_fn.unary_stream(req):
                 if chunk.index <= last_index:
@@ -469,7 +496,12 @@ async def _process_result(result: api_pb2.GenericResult, data_format: int, stub,
     elif result.status == api_pb2.GenericResult.GENERIC_STATUS_INTERNAL_FAILURE:
         raise InternalFailure(result.exception)
     elif result.status != api_pb2.GenericResult.GENERIC_STATUS_SUCCESS:
-        if data:
+        if data and data_format in (api_pb2.DATA_FORMAT_PICKLE, api_pb2.DATA_FORMAT_UNSPECIFIED):
+            # *Unspecified data format here but data present usually means that the exception
+            # was created by the server representing an exception that occurred during container
+            # startup (crash looping) that eventually got escalated to input failures.
+            # TaskResult doesn't specify data format, so these results don't have that metadata
+            # the moment.
             try:
                 exc = deserialize(data, client)
             except DeserializationError as deser_exc:
@@ -497,8 +529,9 @@ async def _process_result(result: api_pb2.GenericResult, data_format: int, stub,
                     append_modal_tb(exc, tb_dict, line_cache)
                 except Exception:
                     pass
-            uc_exc = UserCodeException(exc_with_hints(exc))
-            raise uc_exc
+
+            raise exc_with_hints(exc)
+
         raise RemoteError(result.exception)
 
     try:
@@ -511,35 +544,66 @@ async def _process_result(result: api_pb2.GenericResult, data_format: int, stub,
         ) from deser_exc
 
 
+def should_upload(
+    num_bytes: int,
+    max_object_size_bytes: int,
+    function_call_invocation_type: Optional["api_pb2.FunctionCallInvocationType.ValueType"],
+) -> bool:
+    """
+    Determine if the input should be uploaded to blob storage.
+    """
+    return num_bytes > max_object_size_bytes or (
+        function_call_invocation_type == api_pb2.FUNCTION_CALL_INVOCATION_TYPE_ASYNC
+        and num_bytes > MAX_ASYNC_OBJECT_SIZE_BYTES
+    )
+
+
+# This must be called against the client stub, not the input-plane stub.
 async def _create_input(
-    args, kwargs, stub: ModalClientModal, *, idx: Optional[int] = None, method_name: Optional[str] = None
+    args,
+    kwargs,
+    stub: ModalClientModal,
+    *,
+    function: "modal._functions._Function",
+    idx: Optional[int] = None,
+    function_call_invocation_type: Optional["api_pb2.FunctionCallInvocationType.ValueType"] = None,
 ) -> api_pb2.FunctionPutInputsItem:
     """Serialize function arguments and create a FunctionInput protobuf,
     uploading to blob storage if needed.
     """
+    method_name = function._use_method_name
+    max_object_size_bytes = function._max_object_size_bytes
+
     if idx is None:
         idx = 0
-    if method_name is None:
-        method_name = ""  # proto compatible
 
-    args_serialized = serialize((args, kwargs))
+    data_format = get_preferred_payload_format()
+    if not function._metadata:
+        raise ExecutionError("Attempted to call function that has not been hydrated with metadata")
 
-    if len(args_serialized) > MAX_OBJECT_SIZE_BYTES:
-        args_blob_id = await blob_upload(args_serialized, stub)
+    supported_input_formats = function._metadata.supported_input_formats or [api_pb2.DATA_FORMAT_PICKLE]
+    if data_format not in supported_input_formats:
+        data_format = supported_input_formats[0]
 
+    args_serialized = _serialize_data_format((args, kwargs), data_format)
+
+    if should_upload(len(args_serialized), max_object_size_bytes, function_call_invocation_type):
+        args_blob_id, r2_failed, r2_throughput_bytes_s = await blob_upload_with_r2_failure_info(args_serialized, stub)
         return api_pb2.FunctionPutInputsItem(
             input=api_pb2.FunctionInput(
                 args_blob_id=args_blob_id,
-                data_format=api_pb2.DATA_FORMAT_PICKLE,
+                data_format=data_format,
                 method_name=method_name,
             ),
             idx=idx,
+            r2_failed=r2_failed,
+            r2_throughput_bytes_s=r2_throughput_bytes_s,
         )
     else:
         return api_pb2.FunctionPutInputsItem(
             input=api_pb2.FunctionInput(
                 args=args_serialized,
-                data_format=api_pb2.DATA_FORMAT_PICKLE,
+                data_format=data_format,
                 method_name=method_name,
             ),
             idx=idx,
@@ -581,14 +645,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(
@@ -603,6 +666,14 @@ class FunctionCreationStatus:
                     f"Custom domain for {self.tag} => [magenta underline]{custom_domain.url}[/magenta underline]"
                 )
 
+        elif self.response.function.flash_service_urls:
+            for flash_service_url in self.response.function.flash_service_urls:
+                flash_service_url_status_row = self.resolver.add_status_row()
+                flash_service_url_status_row.finish(
+                    f"Created flash service endpoint for {self.tag} => "
+                    f"[magenta underline]{flash_service_url}[/magenta underline]"
+                )
+
         else:
             for warning in self.response.server_warnings:
                 self.status_row.warning(warning)
@@ -623,30 +694,3 @@ class FunctionCreationStatus:
                                 f"Custom domain for {method_definition.function_name} => [magenta underline]"
                                 f"{custom_domain.url}[/magenta underline]"
                             )
-
-
-class IncludeSourceMode(enum.Enum):
-    INCLUDE_NOTHING = False  # can only be set in source, can't be set in config
-    INCLUDE_MAIN_PACKAGE = True  # Default behavior
-
-
-def get_include_source_mode(function_or_app_specific) -> IncludeSourceMode:
-    """Which "automount" behavior should a function use
-
-    function_or_app_specific: explicit value given in the @function or @cls decorator, in an App constructor, or None
-
-    If function_or_app_specific is specified, validate and return the IncludeSourceMode
-    If function_or_app_specific is None, infer it from config
-    """
-    if function_or_app_specific is not None:
-        if not isinstance(function_or_app_specific, bool):
-            raise ValueError(
-                f"Invalid `include_source` value: {function_or_app_specific}. Use one of:\n"
-                f"True - include function's package source\n"
-                f"False - include no Python source (module expected to be present in Image)\n"
-            )
-
-        # explicitly set in app/function
-        return IncludeSourceMode(function_or_app_specific)
-
-    return IncludeSourceMode.INCLUDE_MAIN_PACKAGE

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,12 +24,15 @@ 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
 
 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.
@@ -148,27 +147,14 @@ 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
 
 
+@retry(n_attempts=5, base_delay=0.1)
 async def connect_channel(channel: grpclib.client.Channel):
-    """Connects socket (potentially raising errors raising to connectivity."""
+    """Connect to socket and raise exceptions when there is a connection issue."""
     await channel.__connect__()
 
 
@@ -177,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]:
@@ -186,50 +172,86 @@ 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,
+) -> 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 or []) + [("x-modal-timestamp", str(time.time()))]
+
     while True:
-        metadata = [("x-idempotency-key", idempotency_key), ("x-retry-attempt", str(n_retries))]
+        attempt_metadata = [
+            ("x-idempotency-key", idempotency_key),
+            ("x-retry-attempt", str(n_retries)),
+            *metadata,
+        ]
         if n_retries > 0:
-            metadata.append(("x-retry-delay", str(time.time() - t0)))
+            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=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:
@@ -237,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

modal/_utils/name_utils.py

@@ -1,5 +1,6 @@
 # Copyright Modal Labs 2022
 import re
+from collections.abc import Mapping
 
 from ..exception import InvalidError
 
@@ -31,12 +32,25 @@ def is_valid_environment_name(name: str) -> bool:
     return len(name) <= 64 and re.match(r"^[a-zA-Z0-9][a-zA-Z0-9-_.]+$", name) is not None
 
 
-def is_valid_tag(tag: str) -> bool:
-    """Tags are alphanumeric, dashes, periods, and underscores, and must be 50 characters or less"""
-    pattern = r"^[a-zA-Z0-9._-]{1,50}$"
+def is_valid_tag(tag: str, max_length: int = 50) -> bool:
+    """Tags are alphanumeric, dashes, periods, and underscores, and not longer than the max_length."""
+    pattern = rf"^[a-zA-Z0-9._-]{{1,{max_length}}}$"
     return bool(re.match(pattern, tag))
 
 
+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."
+    )
+    max_length = 63
+    for key, value in tags.items():
+        if not is_valid_tag(key, max_length):
+            raise InvalidError(f"Invalid tag key: {key!r}.{rules}")
+        if not is_valid_tag(value, max_length):
+            raise InvalidError(f"Invalid tag value: {value!r}.{rules}")
+
+
 def check_object_name(name: str, object_type: str) -> None:
     message = (
         f"Invalid {object_type} name: '{name}'."