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

modal/_utils/blob_utils.py

@@ -4,6 +4,7 @@ import dataclasses
 import hashlib
 import os
 import platform
+import random
 import time
 from collections.abc import AsyncIterator
 from contextlib import AbstractContextManager, contextmanager
@@ -26,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
@@ -37,12 +37,15 @@ if TYPE_CHECKING:
 # Max size for function inputs and outputs.
 MAX_OBJECT_SIZE_BYTES = 2 * 1024 * 1024  # 2 MiB
 
+# Max size for async function inputs and outputs.
+MAX_ASYNC_OBJECT_SIZE_BYTES = 8 * 1024  # 8 KiB
+
 #  If a file is LARGE_FILE_LIMIT bytes or larger, it's uploaded to blob store (s3) instead of going through grpc
 #  It will also make sure to chunk the hash calculation to avoid reading the entire file into memory
 LARGE_FILE_LIMIT = 4 * 1024 * 1024  # 4 MiB
 
 # Max parallelism during map calls
-BLOB_MAX_PARALLELISM = 10
+BLOB_MAX_PARALLELISM = 20
 
 # read ~16MiB chunks by default
 DEFAULT_SEGMENT_CHUNK_SIZE = 2**24
@@ -55,6 +58,8 @@ MULTIPART_UPLOAD_THRESHOLD = 1024**3
 # For block based storage like volumefs2: the size of a block
 BLOCK_SIZE: int = 8 * 1024 * 1024
 
+HEALTHY_R2_UPLOAD_PERCENTAGE = 0.95
+
 
 @retry(n_attempts=5, base_delay=0.5, timeout=None)
 async def _upload_to_s3_url(
@@ -79,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:
@@ -182,9 +187,37 @@ def get_content_length(data: BinaryIO) -> int:
     return content_length - pos
 
 
+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.
+        # To ensure the failure path is continuously exercised.
+        if idx == 0 and len(items) > 1 and random.random() > HEALTHY_R2_UPLOAD_PERCENTAGE:
+            continue
+        try:
+            if blob_id.endswith(":r2"):
+                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_throughput_bytes_s
+        except Exception as _:
+            if blob_id.endswith(":r2"):
+                r2_failed = True
+            # Ignore all errors except the last one, since we're out of fallback options.
+            if idx == len(items) - 1:
+                raise
+    raise ExecutionError("Failed to upload blob")
+
+
 async def _blob_upload(
     upload_hashes: UploadHashes, data: Union[bytes, BinaryIO], stub, progress_report_cb: Optional[Callable] = None
-) -> str:
+) -> tuple[str, bool, int]:
     if isinstance(data, bytes):
         data = BytesIO(data)
 
@@ -195,19 +228,26 @@ async def _blob_upload(
         content_sha256_base64=upload_hashes.sha256_base64,
         content_length=content_length,
     )
-    resp = await retry_transient_errors(stub.BlobCreate, req)
-
-    blob_id = resp.blob_id
+    resp = await stub.BlobCreate(req)
+
+    if resp.WhichOneof("upload_types_oneof") == "multiparts":
+
+        async def upload_multipart_upload(part):
+            return await perform_multipart_upload(
+                data,
+                content_length=content_length,
+                max_part_size=part.part_length,
+                part_urls=part.upload_urls,
+                completion_url=part.completion_url,
+                upload_chunk_size=DEFAULT_SEGMENT_CHUNK_SIZE,
+                progress_report_cb=progress_report_cb,
+            )
 
-    if resp.WhichOneof("upload_type_oneof") == "multipart":
-        await perform_multipart_upload(
-            data,
+        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,
-            max_part_size=resp.multipart.part_length,
-            part_urls=resp.multipart.upload_urls,
-            completion_url=resp.multipart.completion_url,
-            upload_chunk_size=DEFAULT_SEGMENT_CHUNK_SIZE,
-            progress_report_cb=progress_report_cb,
         )
     else:
         from .bytes_io_segment_payload import BytesIOSegmentPayload
@@ -215,34 +255,54 @@ async def _blob_upload(
         payload = BytesIOSegmentPayload(
             data, segment_start=0, segment_length=content_length, progress_report_cb=progress_report_cb
         )
-        await _upload_to_s3_url(
-            resp.upload_url,
-            payload,
-            # for single part uploads, we use server side md5 checksums
-            content_md5_b64=upload_hashes.md5_base64,
+
+        async def upload_to_s3_url(url):
+            return await _upload_to_s3_url(
+                url,
+                payload,
+                # for single part uploads, we use server side md5 checksums
+                content_md5_b64=upload_hashes.md5_base64,
+            )
+
+        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
+    return blob_id, r2_failed, r2_throughput_bytes_s
 
 
-async def blob_upload(payload: bytes, stub: ModalClientModal) -> str:
+async def blob_upload_with_r2_failure_info(payload: bytes, stub: ModalClientModal) -> tuple[str, bool, int]:
     size_mib = len(payload) / 1024 / 1024
     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 = 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). {blob_id}")
+    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_throughput_bytes_s
+
+
+async def blob_upload(payload: bytes, stub: ModalClientModal) -> str:
+    blob_id, _, _ = await blob_upload_with_r2_failure_info(payload, stub)
     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,
@@ -251,7 +311,8 @@ async def blob_upload_file(
     md5_hex: Optional[str] = None,
 ) -> str:
     upload_hashes = get_upload_hashes(file_obj, sha256_hex=sha256_hex, md5_hex=md5_hex)
-    return await _blob_upload(upload_hashes, file_obj, stub, progress_report_cb)
+    blob_id, _, _ = await _blob_upload(upload_hashes, file_obj, stub, progress_report_cb)
+    return blob_id
 
 
 @retry(n_attempts=5, base_delay=0.1, timeout=None)
@@ -259,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:
@@ -273,23 +334,25 @@ 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
     throughput_mib_s = size_mib / dur_s
-    logger.debug(f"Downloaded large blob {blob_id} of size {size_mib:.2f} MiB ({throughput_mib_s:.2f} MiB/s)")
+    logger.debug(
+        f"Downloaded large blob {blob_id} of size {size_mib:.2f} MiB ({throughput_mib_s:.2f} MiB/s, total {dur_s:.2f}s)"
+    )
     return data
 
 
 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:
@@ -380,20 +443,31 @@ def get_file_upload_spec_from_fileobj(fp: BinaryIO, mount_filename: PurePosixPat
         mode,
     )
 
+
 _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
 
-
     @staticmethod
     async def from_path(
         filename: Path,
@@ -416,7 +490,6 @@ class FileUploadSpec2:
             hash_semaphore,
         )
 
-
     @staticmethod
     async def from_fileobj(
         source_fp: Union[BinaryIO, BytesIO],
@@ -426,6 +499,7 @@ class FileUploadSpec2:
     ) -> "FileUploadSpec2":
         try:
             fileno = source_fp.fileno()
+
             def source():
                 new_fd = os.dup(fileno)
                 fp = os.fdopen(new_fd, "rb")
@@ -436,6 +510,7 @@ class FileUploadSpec2:
             # `.fileno()` not available; assume BytesIO-like type
             source_fp = cast(BytesIO, source_fp)
             buffer = source_fp.getbuffer()
+
             def source():
                 return BytesIO(buffer)
 
@@ -447,7 +522,6 @@ class FileUploadSpec2:
             hash_semaphore,
         )
 
-
     @staticmethod
     async def _create(
         source: _FileUploadSource2,
@@ -461,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)
 
-        with source() as block_fp:
-            block_fp.seek(block_start)
+    tasks = (gather_block(idx) for idx in range(num_blocks))
+    return await asyncio.gather(*tasks)
 
-            num_bytes_read = 0
-            while num_bytes_read < BLOCK_SIZE:
-                chunk = block_fp.read(BLOCK_SIZE - num_bytes_read)
 
-                if not chunk:
-                    break
+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)
 
-                num_bytes_read += len(chunk)
-                sha256_hash.update(chunk)
 
-        return sha256_hash.digest()
+def _hash_range_sha256(source: _FileUploadSource2, start, end):
+    sha256_hash = hashlib.sha256()
+    range_size = 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 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
 
-    tasks = (hash_block_sha256(idx) for idx in range(num_blocks))
-    return await asyncio.gather(*tasks)
+    with source() as block_fp:
+        block_fp.seek(start)
+
+        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/deprecation.py

@@ -122,3 +122,22 @@ def warn_on_renamed_autoscaler_settings(func: Callable[P, R]) -> Callable[P, R]:
         return func(*args, **kwargs)
 
     return wrapper
+
+
+def warn_if_passing_namespace(
+    namespace: Any,
+    resource_name: str,
+) -> None:
+    """Issue deprecation warning for namespace parameter if non-None value is passed.
+
+    Args:
+        namespace: The namespace parameter value (may be None or actual value)
+        resource_name: Name of the resource type for the warning message
+    """
+    if namespace is not None:
+        deprecation_warning(
+            (2025, 6, 30),
+            f"The `namespace` parameter for `{resource_name}` is deprecated and will be"
+            " removed in a future release. It is no longer needed, so can be removed"
+            " from your code.",
+        )