modal 1.0.6.dev58__py3-none-any.whl → 1.2.3.dev7__py3-none-any.whl

modal/_utils/blob_utils.py

@@ -27,7 +27,6 @@ from modal_proto.modal_api_grpc import ModalClientModal
 
 from ..exception import ExecutionError
 from .async_utils import TaskContext, retry
-from .grpc_utils import retry_transient_errors
 from .hash_utils import UploadHashes, get_upload_hashes
 from .http_utils import ClientSessionRegistry
 from .logger import logger
@@ -85,7 +84,7 @@ async def _upload_to_s3_url(
         ) as resp:
             # S3 signal to slow down request rate.
             if resp.status == 503:
-                logger.warning("Received SlowDown signal from S3, sleeping for 1 second before retrying.")
+                logger.debug("Received SlowDown signal from S3, sleeping for 1 second before retrying.")
                 await asyncio.sleep(1)
 
             if resp.status != 200:
@@ -188,16 +187,10 @@ def get_content_length(data: BinaryIO) -> int:
     return content_length - pos
 
 
-async def _measure_endpoint_latency(item: str) -> int:
-    latency_ms = 0
-    t0 = time.monotonic_ns()
-    async with ClientSessionRegistry.get_session().head(item) as _:
-        latency_ms = (time.monotonic_ns() - t0) // 1_000_000
-    return latency_ms
-
-
-async def _blob_upload_with_fallback(items, blob_ids: list[str], callback) -> tuple[str, bool, int]:
-    r2_latency_ms = 0
+async def _blob_upload_with_fallback(
+    items, blob_ids: list[str], callback, content_length: int
+) -> tuple[str, bool, int]:
+    r2_throughput_bytes_s = 0
     r2_failed = False
     for idx, (item, blob_id) in enumerate(zip(items, blob_ids)):
         # We want to default to R2 95% of the time and S3 5% of the time.
@@ -206,14 +199,13 @@ async def _blob_upload_with_fallback(items, blob_ids: list[str], callback) -> tu
             continue
         try:
             if blob_id.endswith(":r2"):
-                # measure the time it takes to contact the bucket endpoint
-                r2_latency_ms, _ = await asyncio.gather(
-                    _measure_endpoint_latency(item),
-                    callback(item),
-                )
+                t0 = time.monotonic_ns()
+                await callback(item)
+                dt_ns = time.monotonic_ns() - t0
+                r2_throughput_bytes_s = (content_length * 1_000_000_000) // max(dt_ns, 1)
             else:
                 await callback(item)
-            return blob_id, r2_failed, r2_latency_ms
+            return blob_id, r2_failed, r2_throughput_bytes_s
         except Exception as _:
             if blob_id.endswith(":r2"):
                 r2_failed = True
@@ -236,7 +228,7 @@ async def _blob_upload(
         content_sha256_base64=upload_hashes.sha256_base64,
         content_length=content_length,
     )
-    resp = await retry_transient_errors(stub.BlobCreate, req)
+    resp = await stub.BlobCreate(req)
 
     if resp.WhichOneof("upload_types_oneof") == "multiparts":
 
@@ -251,10 +243,11 @@ async def _blob_upload(
                 progress_report_cb=progress_report_cb,
             )
 
-        blob_id, r2_failed, r2_latency_ms = await _blob_upload_with_fallback(
+        blob_id, r2_failed, r2_throughput_bytes_s = await _blob_upload_with_fallback(
             resp.multiparts.items,
             resp.blob_ids,
             upload_multipart_upload,
+            content_length=content_length,
         )
     else:
         from .bytes_io_segment_payload import BytesIOSegmentPayload
@@ -271,16 +264,17 @@ async def _blob_upload(
                 content_md5_b64=upload_hashes.md5_base64,
             )
 
-        blob_id, r2_failed, r2_latency_ms = await _blob_upload_with_fallback(
+        blob_id, r2_failed, r2_throughput_bytes_s = await _blob_upload_with_fallback(
             resp.upload_urls.items,
             resp.blob_ids,
             upload_to_s3_url,
+            content_length=content_length,
         )
 
     if progress_report_cb:
         progress_report_cb(complete=True)
 
-    return blob_id, r2_failed, r2_latency_ms
+    return blob_id, r2_failed, r2_throughput_bytes_s
 
 
 async def blob_upload_with_r2_failure_info(payload: bytes, stub: ModalClientModal) -> tuple[str, bool, int]:
@@ -288,16 +282,16 @@ async def blob_upload_with_r2_failure_info(payload: bytes, stub: ModalClientModa
     logger.debug(f"Uploading large blob of size {size_mib:.2f} MiB")
     t0 = time.time()
     if isinstance(payload, str):
-        logger.warning("Blob uploading string, not bytes - auto-encoding as utf8")
+        logger.debug("Blob uploading string, not bytes - auto-encoding as utf8")
         payload = payload.encode("utf8")
     upload_hashes = get_upload_hashes(payload)
-    blob_id, r2_failed, r2_latency_ms = await _blob_upload(upload_hashes, payload, stub)
+    blob_id, r2_failed, r2_throughput_bytes_s = await _blob_upload(upload_hashes, payload, stub)
     dur_s = max(time.time() - t0, 0.001)  # avoid division by zero
     throughput_mib_s = (size_mib) / dur_s
     logger.debug(
         f"Uploaded large blob of size {size_mib:.2f} MiB ({throughput_mib_s:.2f} MiB/s, total {dur_s:.2f}s). {blob_id}"
     )
-    return blob_id, r2_failed, r2_latency_ms
+    return blob_id, r2_failed, r2_throughput_bytes_s
 
 
 async def blob_upload(payload: bytes, stub: ModalClientModal) -> str:
@@ -305,6 +299,10 @@ async def blob_upload(payload: bytes, stub: ModalClientModal) -> str:
     return blob_id
 
 
+async def format_blob_data(data: bytes, api_stub: ModalClientModal) -> dict[str, Any]:
+    return {"data_blob_id": await blob_upload(data, api_stub)} if len(data) > MAX_OBJECT_SIZE_BYTES else {"data": data}
+
+
 async def blob_upload_file(
     file_obj: BinaryIO,
     stub: ModalClientModal,
@@ -322,7 +320,7 @@ async def _download_from_url(download_url: str) -> bytes:
     async with ClientSessionRegistry.get_session().get(download_url) as s3_resp:
         # S3 signal to slow down request rate.
         if s3_resp.status == 503:
-            logger.warning("Received SlowDown signal from S3, sleeping for 1 second before retrying.")
+            logger.debug("Received SlowDown signal from S3, sleeping for 1 second before retrying.")
             await asyncio.sleep(1)
 
         if s3_resp.status != 200:
@@ -336,7 +334,7 @@ async def blob_download(blob_id: str, stub: ModalClientModal) -> bytes:
     logger.debug(f"Downloading large blob {blob_id}")
     t0 = time.time()
     req = api_pb2.BlobGetRequest(blob_id=blob_id)
-    resp = await retry_transient_errors(stub.BlobGet, req)
+    resp = await stub.BlobGet(req)
     data = await _download_from_url(resp.download_url)
     size_mib = len(data) / 1024 / 1024
     dur_s = max(time.time() - t0, 0.001)  # avoid division by zero
@@ -349,12 +347,12 @@ async def blob_download(blob_id: str, stub: ModalClientModal) -> bytes:
 
 async def blob_iter(blob_id: str, stub: ModalClientModal) -> AsyncIterator[bytes]:
     req = api_pb2.BlobGetRequest(blob_id=blob_id)
-    resp = await retry_transient_errors(stub.BlobGet, req)
+    resp = await stub.BlobGet(req)
     download_url = resp.download_url
     async with ClientSessionRegistry.get_session().get(download_url) as s3_resp:
         # S3 signal to slow down request rate.
         if s3_resp.status == 503:
-            logger.warning("Received SlowDown signal from S3, sleeping for 1 second before retrying.")
+            logger.debug("Received SlowDown signal from S3, sleeping for 1 second before retrying.")
             await asyncio.sleep(1)
 
         if s3_resp.status != 200:
@@ -449,14 +447,24 @@ def get_file_upload_spec_from_fileobj(fp: BinaryIO, mount_filename: PurePosixPat
 _FileUploadSource2 = Callable[[], ContextManager[BinaryIO]]
 
 
+@dataclasses.dataclass
+class FileUploadBlock:
+    # The start (byte offset, inclusive) of the block within the file
+    start: int
+    # The end (byte offset, exclusive) of the block, after having removed any trailing zeroes
+    end: int
+    # Raw (unencoded 32 byte) SHA256 sum of the block, not including trailing zeroes
+    contents_sha256: bytes
+
+
 @dataclasses.dataclass
 class FileUploadSpec2:
     source: _FileUploadSource2
     source_description: Union[str, Path]
 
     path: str
-    # Raw (unencoded 32 byte) SHA256 sum per 8MiB file block
-    blocks_sha256: list[bytes]
+    # 8MiB file blocks
+    blocks: list[FileUploadBlock]
     mode: int  # file permission bits (last 12 bits of st_mode)
     size: int
 
@@ -527,53 +535,102 @@ class FileUploadSpec2:
             source_fp.seek(0, os.SEEK_END)
             size = source_fp.tell()
 
-        blocks_sha256 = await hash_blocks_sha256(source, size, hash_semaphore)
+        blocks = await _gather_blocks(source, size, hash_semaphore)
 
         return FileUploadSpec2(
             source=source,
             source_description=source_description,
             path=mount_filename.as_posix(),
-            blocks_sha256=blocks_sha256,
+            blocks=blocks,
             mode=mode & 0o7777,
             size=size,
         )
 
 
-async def hash_blocks_sha256(
+async def _gather_blocks(
     source: _FileUploadSource2,
     size: int,
     hash_semaphore: asyncio.Semaphore,
-) -> list[bytes]:
+) -> list[FileUploadBlock]:
     def ceildiv(a: int, b: int) -> int:
         return -(a // -b)
 
     num_blocks = ceildiv(size, BLOCK_SIZE)
 
-    def blocking_hash_block_sha256(block_idx: int) -> bytes:
-        sha256_hash = hashlib.sha256()
-        block_start = block_idx * BLOCK_SIZE
+    async def gather_block(block_idx: int) -> FileUploadBlock:
+        async with hash_semaphore:
+            return await asyncio.to_thread(_gather_block, source, block_idx)
+
+    tasks = (gather_block(idx) for idx in range(num_blocks))
+    return await asyncio.gather(*tasks)
 
-        with source() as block_fp:
-            block_fp.seek(block_start)
 
-            num_bytes_read = 0
-            while num_bytes_read < BLOCK_SIZE:
-                chunk = block_fp.read(BLOCK_SIZE - num_bytes_read)
+def _gather_block(source: _FileUploadSource2, block_idx: int) -> FileUploadBlock:
+    start = block_idx * BLOCK_SIZE
+    end = _find_end_of_block(source, start, start + BLOCK_SIZE)
+    contents_sha256 = _hash_range_sha256(source, start, end)
+    return FileUploadBlock(start=start, end=end, contents_sha256=contents_sha256)
 
-                if not chunk:
-                    break
 
-                num_bytes_read += len(chunk)
-                sha256_hash.update(chunk)
+def _hash_range_sha256(source: _FileUploadSource2, start, end):
+    sha256_hash = hashlib.sha256()
+    range_size = end - start
 
-        return sha256_hash.digest()
+    with source() as fp:
+        fp.seek(start)
+
+        num_bytes_read = 0
+        while num_bytes_read < range_size:
+            chunk = fp.read(range_size - num_bytes_read)
+
+            if not chunk:
+                break
+
+            num_bytes_read += len(chunk)
+            sha256_hash.update(chunk)
+
+    return sha256_hash.digest()
+
+
+def _find_end_of_block(source: _FileUploadSource2, start: int, end: int) -> Optional[int]:
+    """Finds the appropriate end of a block, which is the index of the byte just past the last non-zero byte in the
+    block.
+
+    >>> _find_end_of_block(lambda: BytesIO(b"abc123\0\0\0"), 0, 1024)
+    6
+    >>> _find_end_of_block(lambda: BytesIO(b"abc123\0\0\0"), 3, 1024)
+    6
+    >>> _find_end_of_block(lambda: BytesIO(b"abc123\0\0\0"), 0, 3)
+    4
+    >>> _find_end_of_block(lambda: BytesIO(b"abc123\0\0\0a"), 0, 9)
+    6
+    >>> _find_end_of_block(lambda: BytesIO(b"\0\0\0"), 0, 3)
+    0
+    >>> _find_end_of_block(lambda: BytesIO(b"\0\0\0\0\0\0"), 3, 6)
+    3
+    >>> _find_end_of_block(lambda: BytesIO(b""), 0, 1024)
+    0
+    """
+    size = end - start
+    new_end = start
 
-    async def hash_block_sha256(block_idx: int) -> bytes:
-        async with hash_semaphore:
-            return await asyncio.to_thread(blocking_hash_block_sha256, block_idx)
+    with source() as block_fp:
+        block_fp.seek(start)
 
-    tasks = (hash_block_sha256(idx) for idx in range(num_blocks))
-    return await asyncio.gather(*tasks)
+        num_bytes_read = 0
+        while num_bytes_read < size:
+            chunk = block_fp.read(size - num_bytes_read)
+
+            if not chunk:
+                break
+
+            stripped_chunk = chunk.rstrip(b"\0")
+            if stripped_chunk:
+                new_end = start + num_bytes_read + len(stripped_chunk)
+
+            num_bytes_read += len(chunk)
+
+    return new_end
 
 
 def use_md5(url: str) -> bool:

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
@@ -18,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
@@ -39,6 +41,9 @@ from .blob_utils import (
 )
 from .grpc_utils import RETRYABLE_GRPC_STATUS_CODES
 
+if typing.TYPE_CHECKING:
+    import modal._functions
+
 
 class FunctionInfoType(Enum):
     PACKAGE = "package"
@@ -70,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:
@@ -386,9 +395,16 @@ def callable_has_non_self_non_default_params(f: Callable[..., Any]) -> bool:
 
 
 async def _stream_function_call_data(
-    client, stub, 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
 
@@ -406,7 +422,12 @@ async def _stream_function_call_data(
         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:
@@ -475,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:
@@ -532,43 +558,52 @@ def should_upload(
     )
 
 
+# This must be called against the client stub, not the input-plane stub.
 async def _create_input(
     args,
     kwargs,
     stub: ModalClientModal,
     *,
-    max_object_size_bytes: int,
+    function: "modal._functions._Function",
     idx: Optional[int] = None,
-    method_name: Optional[str] = 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")
+
+    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_latency_ms = await blob_upload_with_r2_failure_info(args_serialized, stub)
+        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_latency_ms=r2_latency_ms,
+            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,
@@ -610,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(
@@ -660,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