modal 0.62.16__py3-none-any.whl → 0.72.11__py3-none-any.whl

modal/_utils/blob_utils.py

@@ -5,23 +5,26 @@ import hashlib
 import io
 import os
 import platform
-from contextlib import contextmanager
+import time
+from collections.abc import AsyncIterator
+from contextlib import AbstractContextManager, contextmanager
 from pathlib import Path, PurePosixPath
-from typing import Any, AsyncIterator, BinaryIO, Callable, List, Optional, Union
+from typing import TYPE_CHECKING, Any, BinaryIO, Callable, Optional, Union
 from urllib.parse import urlparse
 
-from aiohttp import BytesIOPayload
-from aiohttp.abc import AbstractStreamWriter
-
 from modal_proto import api_pb2
+from modal_proto.modal_api_grpc import ModalClientModal
 
 from ..exception import ExecutionError
-from .async_utils import retry
+from .async_utils import TaskContext, retry
 from .grpc_utils import retry_transient_errors
-from .hash_utils import UploadHashes, get_sha256_hex, get_upload_hashes
-from .http_utils import http_client_with_tls
+from .hash_utils import UploadHashes, get_upload_hashes
+from .http_utils import ClientSessionRegistry
 from .logger import logger
 
+if TYPE_CHECKING:
+    from .bytes_io_segment_payload import BytesIOSegmentPayload
+
 # Max size for function inputs and outputs.
 MAX_OBJECT_SIZE_BYTES = 2 * 1024 * 1024  # 2 MiB
 
@@ -35,129 +38,59 @@ BLOB_MAX_PARALLELISM = 10
 # read ~16MiB chunks by default
 DEFAULT_SEGMENT_CHUNK_SIZE = 2**24
 
-
-class BytesIOSegmentPayload(BytesIOPayload):
-    """Modified bytes payload for concurrent sends of chunks from the same file.
-
-    Adds:
-    * read limit using remaining_bytes, in order to split files across streams
-    * larger read chunk (to prevent excessive read contention between parts)
-    * calculates an md5 for the segment
-
-    Feels like this should be in some standard lib...
-    """
-
-    def __init__(
-        self,
-        bytes_io: BinaryIO,  # should *not* be shared as IO position modification is not locked
-        segment_start: int,
-        segment_length: int,
-        chunk_size: int = DEFAULT_SEGMENT_CHUNK_SIZE,
-    ):
-        # not thread safe constructor!
-        super().__init__(bytes_io)
-        self.initial_seek_pos = bytes_io.tell()
-        self.segment_start = segment_start
-        self.segment_length = segment_length
-        # seek to start of file segment we are interested in, in order to make .size() evaluate correctly
-        self._value.seek(self.initial_seek_pos + segment_start)
-        assert self.segment_length <= super().size
-        self.chunk_size = chunk_size
-        self.reset_state()
-
-    def reset_state(self):
-        self._md5_checksum = hashlib.md5()
-        self.num_bytes_read = 0
-        self._value.seek(self.initial_seek_pos)
-
-    @contextmanager
-    def reset_on_error(self):
-        try:
-            yield
-        finally:
-            self.reset_state()
-
-    @property
-    def size(self) -> int:
-        return self.segment_length
-
-    def md5_checksum(self):
-        return self._md5_checksum
-
-    async def write(self, writer: AbstractStreamWriter):
-        loop = asyncio.get_event_loop()
-
-        async def safe_read():
-            read_start = self.initial_seek_pos + self.segment_start + self.num_bytes_read
-            self._value.seek(read_start)
-            num_bytes = min(self.chunk_size, self.remaining_bytes())
-            chunk = await loop.run_in_executor(None, self._value.read, num_bytes)
-
-            await loop.run_in_executor(None, self._md5_checksum.update, chunk)
-            self.num_bytes_read += len(chunk)
-            return chunk
-
-        chunk = await safe_read()
-        while chunk and self.remaining_bytes() > 0:
-            await writer.write(chunk)
-            chunk = await safe_read()
-        if chunk:
-            await writer.write(chunk)
-
-    def remaining_bytes(self):
-        return self.segment_length - self.num_bytes_read
+# Files larger than this will be multipart uploaded. The server might request multipart upload for smaller files as
+# well, but the limit will never be raised.
+# TODO(dano): remove this once we stop requiring md5 for blobs
+MULTIPART_UPLOAD_THRESHOLD = 1024**3
 
 
 @retry(n_attempts=5, base_delay=0.5, timeout=None)
 async def _upload_to_s3_url(
     upload_url,
-    payload: BytesIOSegmentPayload,
+    payload: "BytesIOSegmentPayload",
     content_md5_b64: Optional[str] = None,
     content_type: Optional[str] = "application/octet-stream",  # set to None to force omission of ContentType header
 ) -> str:
     """Returns etag of s3 object which is a md5 hex checksum of the uploaded content"""
     with payload.reset_on_error():  # ensure retries read the same data
-        async with http_client_with_tls(timeout=None) as session:
-            headers = {}
-            if content_md5_b64 and use_md5(upload_url):
-                headers["Content-MD5"] = content_md5_b64
-            if content_type:
-                headers["Content-Type"] = content_type
-
-            async with session.put(
-                upload_url,
-                data=payload,
-                headers=headers,
-                skip_auto_headers=["content-type"] if content_type is None else [],
-            ) 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.")
-                    await asyncio.sleep(1)
-
-                if resp.status != 200:
-                    try:
-                        text = await resp.text()
-                    except Exception:
-                        text = "<no body>"
-                    raise ExecutionError(f"Put to url {upload_url} failed with status {resp.status}: {text}")
-
-                # client side ETag checksum verification
-                # the s3 ETag of a single part upload is a quoted md5 hex of the uploaded content
-                etag = resp.headers["ETag"].strip()
-                if etag.startswith(("W/", "w/")):  # see https://www.rfc-editor.org/rfc/rfc7232#section-2.3
-                    etag = etag[2:]
-                if etag[0] == '"' and etag[-1] == '"':
-                    etag = etag[1:-1]
-                remote_md5 = etag
-
-                local_md5_hex = payload.md5_checksum().hexdigest()
-                if local_md5_hex != remote_md5:
-                    raise ExecutionError(
-                        f"Local data and remote data checksum mismatch ({local_md5_hex} vs {remote_md5})"
-                    )
-
-                return remote_md5
+        headers = {}
+        if content_md5_b64 and use_md5(upload_url):
+            headers["Content-MD5"] = content_md5_b64
+        if content_type:
+            headers["Content-Type"] = content_type
+
+        async with ClientSessionRegistry.get_session().put(
+            upload_url,
+            data=payload,
+            headers=headers,
+            skip_auto_headers=["content-type"] if content_type is None else [],
+        ) 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.")
+                await asyncio.sleep(1)
+
+            if resp.status != 200:
+                try:
+                    text = await resp.text()
+                except Exception:
+                    text = "<no body>"
+                raise ExecutionError(f"Put to url {upload_url} failed with status {resp.status}: {text}")
+
+            # client side ETag checksum verification
+            # the s3 ETag of a single part upload is a quoted md5 hex of the uploaded content
+            etag = resp.headers["ETag"].strip()
+            if etag.startswith(("W/", "w/")):  # see https://www.rfc-editor.org/rfc/rfc7232#section-2.3
+                etag = etag[2:]
+            if etag[0] == '"' and etag[-1] == '"':
+                etag = etag[1:-1]
+            remote_md5 = etag
+
+            local_md5_hex = payload.md5_checksum().hexdigest()
+            if local_md5_hex != remote_md5:
+                raise ExecutionError(f"Local data and remote data checksum mismatch ({local_md5_hex} vs {remote_md5})")
+
+            return remote_md5
 
 
 async def perform_multipart_upload(
@@ -165,17 +98,20 @@ async def perform_multipart_upload(
     *,
     content_length: int,
     max_part_size: int,
-    part_urls: List[str],
+    part_urls: list[str],
     completion_url: str,
     upload_chunk_size: int = DEFAULT_SEGMENT_CHUNK_SIZE,
-):
+    progress_report_cb: Optional[Callable] = None,
+) -> None:
+    from .bytes_io_segment_payload import BytesIOSegmentPayload
+
     upload_coros = []
     file_offset = 0
     num_bytes_left = content_length
 
     # Give each part its own IO reader object to avoid needing to
     # lock access to the reader's position pointer.
-    data_file_readers: List[BinaryIO]
+    data_file_readers: list[BinaryIO]
     if isinstance(data_file, io.BytesIO):
         view = data_file.getbuffer()  # does not copy data
         data_file_readers = [io.BytesIO(view) for _ in range(len(part_urls))]
@@ -190,12 +126,13 @@ async def perform_multipart_upload(
             segment_start=file_offset,
             segment_length=part_length_bytes,
             chunk_size=upload_chunk_size,
+            progress_report_cb=progress_report_cb,
         )
         upload_coros.append(_upload_to_s3_url(part_url, payload=part_payload, content_type=None))
         num_bytes_left -= part_length_bytes
         file_offset += part_length_bytes
 
-    part_etags = await asyncio.gather(*upload_coros)
+    part_etags = await TaskContext.gather(*upload_coros)
 
     # The body of the complete_multipart_upload command needs some data in xml format:
     completion_body = "<CompleteMultipartUpload>\n"
@@ -207,25 +144,24 @@ async def perform_multipart_upload(
     bin_hash_parts = [bytes.fromhex(etag) for etag in part_etags]
 
     expected_multipart_etag = hashlib.md5(b"".join(bin_hash_parts)).hexdigest() + f"-{len(part_etags)}"
-    async with http_client_with_tls(timeout=None) as session:
-        resp = await session.post(
-            completion_url, data=completion_body.encode("ascii"), skip_auto_headers=["content-type"]
-        )
-        if resp.status != 200:
-            try:
-                msg = await resp.text()
-            except Exception:
-                msg = "<no body>"
-            raise ExecutionError(f"Error when completing multipart upload: {resp.status}\n{msg}")
-        else:
-            response_body = await resp.text()
-            if expected_multipart_etag not in response_body:
-                raise ExecutionError(
-                    f"Hash mismatch on multipart upload assembly: {expected_multipart_etag} not in {response_body}"
-                )
+    resp = await ClientSessionRegistry.get_session().post(
+        completion_url, data=completion_body.encode("ascii"), skip_auto_headers=["content-type"]
+    )
+    if resp.status != 200:
+        try:
+            msg = await resp.text()
+        except Exception:
+            msg = "<no body>"
+        raise ExecutionError(f"Error when completing multipart upload: {resp.status}\n{msg}")
+    else:
+        response_body = await resp.text()
+        if expected_multipart_etag not in response_body:
+            raise ExecutionError(
+                f"Hash mismatch on multipart upload assembly: {expected_multipart_etag} not in {response_body}"
+            )
 
 
-def get_content_length(data: BinaryIO):
+def get_content_length(data: BinaryIO) -> int:
     # *Remaining* length of file from current seek position
     pos = data.tell()
     data.seek(0, os.SEEK_END)
@@ -234,7 +170,9 @@ def get_content_length(data: BinaryIO):
     return content_length - pos
 
 
-async def _blob_upload(upload_hashes: UploadHashes, data: Union[bytes, BinaryIO], stub) -> str:
+async def _blob_upload(
+    upload_hashes: UploadHashes, data: Union[bytes, BinaryIO], stub, progress_report_cb: Optional[Callable] = None
+) -> str:
     if isinstance(data, bytes):
         data = io.BytesIO(data)
 
@@ -257,9 +195,14 @@ async def _blob_upload(upload_hashes: UploadHashes, data: Union[bytes, BinaryIO]
             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:
-        payload = BytesIOSegmentPayload(data, segment_start=0, segment_length=content_length)
+        from .bytes_io_segment_payload import BytesIOSegmentPayload
+
+        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,
@@ -267,79 +210,103 @@ async def _blob_upload(upload_hashes: UploadHashes, data: Union[bytes, BinaryIO]
             content_md5_b64=upload_hashes.md5_base64,
         )
 
+    if progress_report_cb:
+        progress_report_cb(complete=True)
+
     return blob_id
 
 
-async def blob_upload(payload: bytes, stub) -> str:
+async def blob_upload(payload: bytes, stub: ModalClientModal) -> str:
+    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")
         payload = payload.encode("utf8")
     upload_hashes = get_upload_hashes(payload)
-    return await _blob_upload(upload_hashes, payload, stub)
+    blob_id = 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)." f" {blob_id}")
+    return blob_id
 
 
-async def blob_upload_file(file_obj: BinaryIO, stub) -> str:
-    upload_hashes = get_upload_hashes(file_obj)
-    return await _blob_upload(upload_hashes, file_obj, stub)
+async def blob_upload_file(
+    file_obj: BinaryIO,
+    stub: ModalClientModal,
+    progress_report_cb: Optional[Callable] = None,
+    sha256_hex: Optional[str] = None,
+    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)
 
 
 @retry(n_attempts=5, base_delay=0.1, timeout=None)
-async def _download_from_url(download_url) -> bytes:
-    async with http_client_with_tls(timeout=None) as session:
-        async with session.get(download_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.")
-                await asyncio.sleep(1)
-
-            if resp.status != 200:
-                text = await resp.text()
-                raise ExecutionError(f"Get from url failed with status {resp.status}: {text}")
-            return await resp.read()
-
-
-async def blob_download(blob_id, stub) -> bytes:
-    # convenience function reading all of the downloaded file into memory
+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.")
+            await asyncio.sleep(1)
+
+        if s3_resp.status != 200:
+            text = await s3_resp.text()
+            raise ExecutionError(f"Get from url failed with status {s3_resp.status}: {text}")
+        return await s3_resp.read()
+
+
+async def blob_download(blob_id: str, stub: ModalClientModal) -> bytes:
+    """Convenience function for reading all of the downloaded file into memory."""
+    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)
-
-    return await _download_from_url(resp.download_url)
+    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)")
+    return data
 
 
-async def blob_iter(blob_id, stub) -> AsyncIterator[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)
     download_url = resp.download_url
-    async with http_client_with_tls(timeout=None) as session:
-        async with session.get(download_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.")
-                await asyncio.sleep(1)
+    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.")
+            await asyncio.sleep(1)
 
-            if resp.status != 200:
-                text = await resp.text()
-                raise ExecutionError(f"Get from url failed with status {resp.status}: {text}")
+        if s3_resp.status != 200:
+            text = await s3_resp.text()
+            raise ExecutionError(f"Get from url failed with status {s3_resp.status}: {text}")
 
-            async for chunk in resp.content.iter_any():
-                yield chunk
+        async for chunk in s3_resp.content.iter_any():
+            yield chunk
 
 
 @dataclasses.dataclass
 class FileUploadSpec:
-    source: Callable[[], BinaryIO]
+    source: Callable[[], Union[AbstractContextManager, BinaryIO]]
     source_description: Any
     mount_filename: str
 
     use_blob: bool
     content: Optional[bytes]  # typically None if using blob, required otherwise
     sha256_hex: str
+    md5_hex: str
     mode: int  # file permission bits (last 12 bits of st_mode)
     size: int
 
 
 def _get_file_upload_spec(
-    source: Callable[[], BinaryIO], source_description: Any, mount_filename: PurePosixPath, mode: int
+    source: Callable[[], Union[AbstractContextManager, BinaryIO]],
+    source_description: Any,
+    mount_filename: PurePosixPath,
+    mode: int,
 ) -> FileUploadSpec:
     with source() as fp:
         # Current position is ignored - we always upload from position 0
@@ -348,13 +315,15 @@ def _get_file_upload_spec(
         fp.seek(0)
 
         if size >= LARGE_FILE_LIMIT:
+            # TODO(dano): remove the placeholder md5 once we stop requiring md5 for blobs
+            md5_hex = "baadbaadbaadbaadbaadbaadbaadbaad" if size > MULTIPART_UPLOAD_THRESHOLD else None
             use_blob = True
             content = None
-            sha256_hex = get_sha256_hex(fp)
+            hashes = get_upload_hashes(fp, md5_hex=md5_hex)
         else:
             use_blob = False
             content = fp.read()
-            sha256_hex = get_sha256_hex(content)
+            hashes = get_upload_hashes(content)
 
     return FileUploadSpec(
         source=source,
@@ -362,7 +331,8 @@ def _get_file_upload_spec(
         mount_filename=mount_filename.as_posix(),
         use_blob=use_blob,
         content=content,
-        sha256_hex=sha256_hex,
+        sha256_hex=hashes.sha256_hex(),
+        md5_hex=hashes.md5_hex(),
         mode=mode & 0o7777,
         size=size,
     )
@@ -383,10 +353,11 @@ def get_file_upload_spec_from_path(
 
 
 def get_file_upload_spec_from_fileobj(fp: BinaryIO, mount_filename: PurePosixPath, mode: int) -> FileUploadSpec:
+    @contextmanager
     def source():
         # We ignore position in stream and always upload from position 0
         fp.seek(0)
-        return fp
+        yield fp
 
     return _get_file_upload_spec(
         source,
@@ -403,7 +374,7 @@ def use_md5(url: str) -> bool:
     https://github.com/spulec/moto/issues/816
     """
     host = urlparse(url).netloc.split(":")[0]
-    if host.endswith(".amazonaws.com"):
+    if host.endswith(".amazonaws.com") or host.endswith(".r2.cloudflarestorage.com"):
         return True
     elif host in ["127.0.0.1", "localhost", "172.21.0.1"]:
         return False

modal/_utils/bytes_io_segment_payload.py

@@ -0,0 +1,97 @@
+# Copyright Modal Labs 2024
+
+import asyncio
+import hashlib
+from contextlib import contextmanager
+from typing import BinaryIO, Callable, Optional
+
+# Note: this module needs to import aiohttp in global scope
+# This takes about 50ms and isn't needed in many cases for Modal execution
+# To avoid this, we import it in local scope when needed (blob_utils.py)
+from aiohttp import BytesIOPayload
+from aiohttp.abc import AbstractStreamWriter
+
+# read ~16MiB chunks by default
+DEFAULT_SEGMENT_CHUNK_SIZE = 2**24
+
+
+class BytesIOSegmentPayload(BytesIOPayload):
+    """Modified bytes payload for concurrent sends of chunks from the same file.
+
+    Adds:
+    * read limit using remaining_bytes, in order to split files across streams
+    * larger read chunk (to prevent excessive read contention between parts)
+    * calculates an md5 for the segment
+
+    Feels like this should be in some standard lib...
+    """
+
+    def __init__(
+        self,
+        bytes_io: BinaryIO,  # should *not* be shared as IO position modification is not locked
+        segment_start: int,
+        segment_length: int,
+        chunk_size: int = DEFAULT_SEGMENT_CHUNK_SIZE,
+        progress_report_cb: Optional[Callable] = None,
+    ):
+        # not thread safe constructor!
+        super().__init__(bytes_io)
+        self.initial_seek_pos = bytes_io.tell()
+        self.segment_start = segment_start
+        self.segment_length = segment_length
+        # seek to start of file segment we are interested in, in order to make .size() evaluate correctly
+        self._value.seek(self.initial_seek_pos + segment_start)
+        assert self.segment_length <= super().size
+        self.chunk_size = chunk_size
+        self.progress_report_cb = progress_report_cb or (lambda *_, **__: None)
+        self.reset_state()
+
+    def reset_state(self):
+        self._md5_checksum = hashlib.md5()
+        self.num_bytes_read = 0
+        self._value.seek(self.initial_seek_pos)
+
+    @contextmanager
+    def reset_on_error(self):
+        try:
+            yield
+        except Exception as exc:
+            try:
+                self.progress_report_cb(reset=True)
+            except Exception as cb_exc:
+                raise cb_exc from exc
+            raise exc
+        finally:
+            self.reset_state()
+
+    @property
+    def size(self) -> int:
+        return self.segment_length
+
+    def md5_checksum(self):
+        return self._md5_checksum
+
+    async def write(self, writer: "AbstractStreamWriter"):
+        loop = asyncio.get_event_loop()
+
+        async def safe_read():
+            read_start = self.initial_seek_pos + self.segment_start + self.num_bytes_read
+            self._value.seek(read_start)
+            num_bytes = min(self.chunk_size, self.remaining_bytes())
+            chunk = await loop.run_in_executor(None, self._value.read, num_bytes)
+
+            await loop.run_in_executor(None, self._md5_checksum.update, chunk)
+            self.num_bytes_read += len(chunk)
+            return chunk
+
+        chunk = await safe_read()
+        while chunk and self.remaining_bytes() > 0:
+            await writer.write(chunk)
+            self.progress_report_cb(advance=len(chunk))
+            chunk = await safe_read()
+        if chunk:
+            await writer.write(chunk)
+            self.progress_report_cb(advance=len(chunk))
+
+    def remaining_bytes(self):
+        return self.segment_length - self.num_bytes_read