modal 0.74.15__py3-none-any.whl → 0.74.16__py3-none-any.whl

modal/_utils/blob_utils.py

@@ -1,15 +1,25 @@
 # Copyright Modal Labs 2022
 import asyncio
 import dataclasses
+import functools
 import hashlib
-import io
 import os
 import platform
 import time
 from collections.abc import AsyncIterator
 from contextlib import AbstractContextManager, contextmanager
+from io import BytesIO, FileIO
 from pathlib import Path, PurePosixPath
-from typing import TYPE_CHECKING, Any, BinaryIO, Callable, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    BinaryIO,
+    Callable,
+    ContextManager,
+    Optional,
+    Union,
+    cast,
+)
 from urllib.parse import urlparse
 
 from modal_proto import api_pb2
@@ -43,6 +53,9 @@ DEFAULT_SEGMENT_CHUNK_SIZE = 2**24
 # TODO(dano): remove this once we stop requiring md5 for blobs
 MULTIPART_UPLOAD_THRESHOLD = 1024**3
 
+# For block based storage like volumefs2: the size of a block
+BLOCK_SIZE: int = 8 * 1024 * 1024
+
 
 @retry(n_attempts=5, base_delay=0.5, timeout=None)
 async def _upload_to_s3_url(
@@ -94,7 +107,7 @@ async def _upload_to_s3_url(
 
 
 async def perform_multipart_upload(
-    data_file: Union[BinaryIO, io.BytesIO, io.FileIO],
+    data_file: Union[BinaryIO, BytesIO, FileIO],
     *,
     content_length: int,
     max_part_size: int,
@@ -112,9 +125,9 @@ async def perform_multipart_upload(
     # 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]
-    if isinstance(data_file, io.BytesIO):
+    if isinstance(data_file, BytesIO):
         view = data_file.getbuffer()  # does not copy data
-        data_file_readers = [io.BytesIO(view) for _ in range(len(part_urls))]
+        data_file_readers = [BytesIO(view) for _ in range(len(part_urls))]
     else:
         filename = data_file.name
         data_file_readers = [open(filename, "rb") for _ in range(len(part_urls))]
@@ -174,7 +187,7 @@ 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)
+        data = BytesIO(data)
 
     content_length = get_content_length(data)
 
@@ -368,6 +381,125 @@ def get_file_upload_spec_from_fileobj(fp: BinaryIO, mount_filename: PurePosixPat
         mode,
     )
 
+_FileUploadSource2 = Callable[[], ContextManager[BinaryIO]]
+
+@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]
+    mode: int  # file permission bits (last 12 bits of st_mode)
+    size: int
+
+
+    @staticmethod
+    async def from_path(
+        filename: Path,
+        mount_filename: PurePosixPath,
+        mode: Optional[int] = None,
+    ) -> "FileUploadSpec2":
+        # Python appears to give files 0o666 bits on Windows (equal for user, group, and global),
+        # so we mask those out to 0o755 for compatibility with POSIX-based permissions.
+        mode = mode or os.stat(filename).st_mode & (0o7777 if platform.system() != "Windows" else 0o7755)
+
+        def source():
+            return open(filename, "rb")
+
+        return await FileUploadSpec2._create(
+            source,
+            filename,
+            mount_filename,
+            mode,
+        )
+
+
+    @staticmethod
+    async def from_fileobj(
+        source_fp: Union[BinaryIO, BytesIO],
+        mount_filename: PurePosixPath,
+        mode: int
+    ) -> "FileUploadSpec2":
+        try:
+            fileno = source_fp.fileno()
+            def source():
+                new_fd = os.dup(fileno)
+                fp = os.fdopen(new_fd, "rb")
+                fp.seek(0)
+                return fp
+
+        except OSError:
+            # `.fileno()` not available; assume BytesIO-like type
+            source_fp = cast(BytesIO, source_fp)
+            buffer = source_fp.getbuffer()
+            def source():
+                return BytesIO(buffer)
+
+        return await FileUploadSpec2._create(
+            source,
+            str(source),
+            mount_filename,
+            mode,
+        )
+
+
+    @staticmethod
+    async def _create(
+        source: _FileUploadSource2,
+        source_description: Union[str, Path],
+        mount_filename: PurePosixPath,
+        mode: int,
+    ) -> "FileUploadSpec2":
+        # Current position is ignored - we always upload from position 0
+        with source() as source_fp:
+            source_fp.seek(0, os.SEEK_END)
+            size = source_fp.tell()
+
+        blocks_sha256 = await hash_blocks_sha256(source, size)
+
+        return FileUploadSpec2(
+            source=source,
+            source_description=source_description,
+            path=mount_filename.as_posix(),
+            blocks_sha256=blocks_sha256,
+            mode=mode & 0o7777,
+            size=size,
+        )
+
+
+async def hash_blocks_sha256(
+    source: _FileUploadSource2,
+    size: int,
+) -> list[bytes]:
+    def ceildiv(a: int, b: int) -> int:
+        return -(a // -b)
+
+    num_blocks = ceildiv(size, BLOCK_SIZE)
+
+    def hash_block_sha256(block_idx: int) -> bytes:
+        sha256_hash = hashlib.sha256()
+        block_start = block_idx * BLOCK_SIZE
+
+        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)
+
+                if not chunk:
+                    break
+
+                num_bytes_read += len(chunk)
+                sha256_hash.update(chunk)
+
+        return sha256_hash.digest()
+
+    tasks = (asyncio.to_thread(functools.partial(hash_block_sha256, idx)) for idx in range(num_blocks))
+    return await asyncio.gather(*tasks)
+
 
 def use_md5(url: str) -> bool:
     """This takes an upload URL in S3 and returns whether we should attach a checksum.

modal/cli/volume.py

@@ -19,7 +19,7 @@ from modal.cli._download import _volume_download
 from modal.cli.utils import ENV_OPTION, YES_OPTION, display_table, timestamp_to_local
 from modal.client import _Client
 from modal.environments import ensure_env
-from modal.volume import _Volume, _VolumeUploadContextManager
+from modal.volume import _AbstractVolumeUploadContextManager, _Volume
 from modal_proto import api_pb2
 
 volume_cli = Typer(
@@ -198,8 +198,12 @@ async def put(
     if Path(local_path).is_dir():
         with progress_handler.live:
             try:
-                async with _VolumeUploadContextManager(
-                    vol.object_id, vol._client, progress_cb=progress_handler.progress, force=force
+                async with _AbstractVolumeUploadContextManager.resolve(
+                    vol._metadata.version,
+                    vol.object_id,
+                    vol._client,
+                    progress_cb=progress_handler.progress,
+                    force=force
                 ) as batch:
                     batch.put_directory(local_path, remote_path)
             except FileExistsError as exc:
@@ -210,8 +214,12 @@ async def put(
     else:
         with progress_handler.live:
             try:
-                async with _VolumeUploadContextManager(
-                    vol.object_id, vol._client, progress_cb=progress_handler.progress, force=force
+                async with _AbstractVolumeUploadContextManager.resolve(
+                    vol._metadata.version,
+                    vol.object_id,
+                    vol._client,
+                    progress_cb=progress_handler.progress,
+                    force=force
                 ) as batch:
                     batch.put_file(local_path, remote_path)
 

modal/client.pyi

@@ -27,7 +27,7 @@ class _Client:
     _snapshotted: bool
 
     def __init__(
-        self, server_url: str, client_type: int, credentials: typing.Optional[tuple[str, str]], version: str = "0.74.15"
+        self, server_url: str, client_type: int, credentials: typing.Optional[tuple[str, str]], version: str = "0.74.16"
     ): ...
     def is_closed(self) -> bool: ...
     @property
@@ -85,7 +85,7 @@ class Client:
     _snapshotted: bool
 
     def __init__(
-        self, server_url: str, client_type: int, credentials: typing.Optional[tuple[str, str]], version: str = "0.74.15"
+        self, server_url: str, client_type: int, credentials: typing.Optional[tuple[str, str]], version: str = "0.74.16"
     ): ...
     def is_closed(self) -> bool: ...
     @property

modal/functions.pyi

@@ -201,11 +201,11 @@ class Function(
 
     _call_generator_nowait: ___call_generator_nowait_spec[typing_extensions.Self]
 
-    class __remote_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
+    class __remote_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
         def __call__(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> ReturnType_INNER: ...
         async def aio(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> ReturnType_INNER: ...
 
-    remote: __remote_spec[modal._functions.P, modal._functions.ReturnType, typing_extensions.Self]
+    remote: __remote_spec[modal._functions.ReturnType, modal._functions.P, typing_extensions.Self]
 
     class __remote_gen_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, *args, **kwargs) -> typing.Generator[typing.Any, None, None]: ...
@@ -220,19 +220,19 @@ class Function(
         self, *args: modal._functions.P.args, **kwargs: modal._functions.P.kwargs
     ) -> modal._functions.OriginalReturnType: ...
 
-    class ___experimental_spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
+    class ___experimental_spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
         def __call__(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
         async def aio(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
 
     _experimental_spawn: ___experimental_spawn_spec[
-        modal._functions.P, modal._functions.ReturnType, typing_extensions.Self
+        modal._functions.ReturnType, modal._functions.P, typing_extensions.Self
     ]
 
-    class __spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
+    class __spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
         def __call__(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
         async def aio(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
 
-    spawn: __spawn_spec[modal._functions.P, modal._functions.ReturnType, typing_extensions.Self]
+    spawn: __spawn_spec[modal._functions.ReturnType, modal._functions.P, typing_extensions.Self]
 
     def get_raw_f(self) -> collections.abc.Callable[..., typing.Any]: ...
 

modal/volume.py

@@ -10,16 +10,19 @@ import time
 import typing
 from collections.abc import AsyncGenerator, AsyncIterator, Generator, Sequence
 from dataclasses import dataclass
+from io import BytesIO
 from pathlib import Path, PurePosixPath
 from typing import (
     IO,
     Any,
+    Awaitable,
     BinaryIO,
     Callable,
     Optional,
     Union,
 )
 
+from google.protobuf.message import Message
 from grpclib import GRPCError, Status
 from synchronicity.async_wrap import asynccontextmanager
 
@@ -31,7 +34,9 @@ from ._object import EPHEMERAL_OBJECT_HEARTBEAT_SLEEP, _get_environment_name, _O
 from ._resolver import Resolver
 from ._utils.async_utils import TaskContext, aclosing, async_map, asyncnullcontext, synchronize_api
 from ._utils.blob_utils import (
+    BLOCK_SIZE,
     FileUploadSpec,
+    FileUploadSpec2,
     blob_iter,
     blob_upload_file,
     get_file_upload_spec_from_fileobj,
@@ -39,6 +44,7 @@ from ._utils.blob_utils import (
 )
 from ._utils.deprecation import deprecation_error, deprecation_warning, renamed_parameter
 from ._utils.grpc_utils import retry_transient_errors
+from ._utils.http_utils import ClientSessionRegistry
 from ._utils.name_utils import check_object_name
 from .client import _Client
 from .config import logger
@@ -119,6 +125,7 @@ class _Volume(_Object, type_prefix="vo"):
     """
 
     _lock: Optional[asyncio.Lock] = None
+    _metadata: "typing.Optional[api_pb2.VolumeMetadata]"
 
     async def _get_lock(self):
         # To (mostly*) prevent multiple concurrent operations on the same volume, which can cause problems under
@@ -171,10 +178,21 @@ class _Volume(_Object, type_prefix="vo"):
                 version=version,
             )
             response = await resolver.client.stub.VolumeGetOrCreate(req)
-            self._hydrate(response.volume_id, resolver.client, None)
+            self._hydrate(response.volume_id, resolver.client, response.metadata)
 
         return _Volume._from_loader(_load, "Volume()", hydrate_lazily=True)
 
+    def _hydrate_metadata(self, metadata: Optional[Message]):
+        if metadata and isinstance(metadata, api_pb2.VolumeMetadata):
+            self._metadata = metadata
+        else:
+            raise TypeError(
+                "_hydrate_metadata() requires an `api_pb2.VolumeMetadata` to determine volume version"
+            )
+
+    def _get_metadata(self) -> Optional[Message]:
+        return self._metadata
+
     @classmethod
     @asynccontextmanager
     async def ephemeral(
@@ -209,7 +227,7 @@ class _Volume(_Object, type_prefix="vo"):
         async with TaskContext() as tc:
             request = api_pb2.VolumeHeartbeatRequest(volume_id=response.volume_id)
             tc.infinite_loop(lambda: client.stub.VolumeHeartbeat(request), sleep=_heartbeat_sleep)
-            yield cls._new_hydrated(response.volume_id, client, None, is_another_app=True)
+            yield cls._new_hydrated(response.volume_id, client, response.metadata, is_another_app=True)
 
     @staticmethod
     @renamed_parameter((2024, 12, 18), "label", "name")
@@ -481,7 +499,7 @@ class _Volume(_Object, type_prefix="vo"):
         await retry_transient_errors(self._client.stub.VolumeCopyFiles, request, base_delay=1)
 
     @live_method
-    async def batch_upload(self, force: bool = False) -> "_VolumeUploadContextManager":
+    async def batch_upload(self, force: bool = False) -> "_AbstractVolumeUploadContextManager":
         """
         Initiate a batched upload to a volume.
 
@@ -499,7 +517,13 @@ class _Volume(_Object, type_prefix="vo"):
             batch.put_file(io.BytesIO(b"some data"), "/foobar")
         ```
         """
-        return _VolumeUploadContextManager(self.object_id, self._client, force=force)
+        return _AbstractVolumeUploadContextManager.resolve(
+            self._metadata.version,
+            self.object_id,
+            self._client,
+            force=force
+        )
+
 
     @live_method
     async def _instance_delete(self):
@@ -527,7 +551,57 @@ class _Volume(_Object, type_prefix="vo"):
         await retry_transient_errors(obj._client.stub.VolumeRename, req)
 
 
-class _VolumeUploadContextManager:
+Volume = synchronize_api(_Volume)
+
+# TODO(dflemstr): Find a way to add ABC or AbstractAsyncContextManager superclasses while keeping synchronicity happy.
+class _AbstractVolumeUploadContextManager:
+    async def __aenter__(self):
+        ...
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        ...
+
+
+    def put_file(
+        self,
+        local_file: Union[Path, str, BinaryIO, BytesIO],
+        remote_path: Union[PurePosixPath, str],
+        mode: Optional[int] = None,
+    ):
+        ...
+
+    def put_directory(
+        self,
+        local_path: Union[Path, str],
+        remote_path: Union[PurePosixPath, str],
+        recursive: bool = True,
+    ):
+        ...
+
+    @staticmethod
+    def resolve(
+        version: "modal_proto.api_pb2.VolumeFsVersion.ValueType",
+        object_id: str,
+        client,
+        progress_cb: Optional[Callable[..., Any]] = None,
+        force: bool = False
+    ) -> "_AbstractVolumeUploadContextManager":
+
+        if version in [
+            None,
+            api_pb2.VolumeFsVersion.VOLUME_FS_VERSION_UNSPECIFIED,
+            api_pb2.VolumeFsVersion.VOLUME_FS_VERSION_V1
+        ]:
+            return _VolumeUploadContextManager(object_id, client, progress_cb=progress_cb, force=force)
+        elif version == api_pb2.VolumeFsVersion.VOLUME_FS_VERSION_V2:
+            return _VolumeUploadContextManager2(object_id, client, progress_cb=progress_cb, force=force)
+        else:
+            raise RuntimeError(f"unsupported volume version: {version}")
+
+
+AbstractVolumeUploadContextManager = synchronize_api(_AbstractVolumeUploadContextManager)
+
+class _VolumeUploadContextManager(_AbstractVolumeUploadContextManager):
     """Context manager for batch-uploading files to a Volume."""
 
     _volume_id: str
@@ -585,7 +659,7 @@ class _VolumeUploadContextManager:
 
     def put_file(
         self,
-        local_file: Union[Path, str, BinaryIO],
+        local_file: Union[Path, str, BinaryIO, BytesIO],
         remote_path: Union[PurePosixPath, str],
         mode: Optional[int] = None,
     ):
@@ -678,9 +752,204 @@ class _VolumeUploadContextManager:
         )
 
 
-Volume = synchronize_api(_Volume)
 VolumeUploadContextManager = synchronize_api(_VolumeUploadContextManager)
 
+_FileUploader2 = Callable[[], Awaitable[FileUploadSpec2]]
+
+class _VolumeUploadContextManager2(_AbstractVolumeUploadContextManager):
+    """Context manager for batch-uploading files to a Volume version 2."""
+
+    _volume_id: str
+    _client: _Client
+    _force: bool
+    _progress_cb: Callable[..., Any]
+    _uploader_generators: list[Generator[_FileUploader2]]
+
+    def __init__(
+        self, volume_id: str, client: _Client, progress_cb: Optional[Callable[..., Any]] = None, force: bool = False
+    ):
+        """mdmd:hidden"""
+        self._volume_id = volume_id
+        self._client = client
+        self._uploader_generators = []
+        self._progress_cb = progress_cb or (lambda *_, **__: None)
+        self._force = force
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if not exc_val:
+            # Flatten all the uploads yielded by the upload generators in the batch
+            def gen_upload_providers():
+                for gen in self._uploader_generators:
+                    yield from gen
+
+            async def gen_file_upload_specs() -> list[FileUploadSpec2]:
+                uploads = [asyncio.create_task(fut()) for fut in gen_upload_providers()]
+                logger.debug(f"Computing checksums for {len(uploads)} files")
+
+                file_specs = []
+                for file_spec in asyncio.as_completed(uploads):
+                    file_specs.append(await file_spec)
+                return file_specs
+
+            upload_specs = await gen_file_upload_specs()
+            await self._put_file_specs(upload_specs)
+
+
+    def put_file(
+        self,
+        local_file: Union[Path, str, BinaryIO, BytesIO],
+        remote_path: Union[PurePosixPath, str],
+        mode: Optional[int] = None,
+    ):
+        """Upload a file from a local file or file-like object.
+
+        Will create any needed parent directories automatically.
+
+        If `local_file` is a file-like object it must remain readable for the lifetime of the batch.
+        """
+        remote_path = PurePosixPath(remote_path).as_posix()
+        if remote_path.endswith("/"):
+            raise ValueError(f"remote_path ({remote_path}) must refer to a file - cannot end with /")
+
+        def gen():
+            if isinstance(local_file, str) or isinstance(local_file, Path):
+                yield lambda: FileUploadSpec2.from_path(local_file, PurePosixPath(remote_path), mode)
+            else:
+                yield lambda: FileUploadSpec2.from_fileobj(local_file, PurePosixPath(remote_path), mode or 0o644)
+
+        self._uploader_generators.append(gen())
+
+    def put_directory(
+        self,
+        local_path: Union[Path, str],
+        remote_path: Union[PurePosixPath, str],
+        recursive: bool = True,
+    ):
+        """
+        Upload all files in a local directory.
+
+        Will create any needed parent directories automatically.
+        """
+        local_path = Path(local_path)
+        assert local_path.is_dir()
+        remote_path = PurePosixPath(remote_path)
+
+        def create_spec(subpath):
+            relpath_str = subpath.relative_to(local_path)
+            return lambda: FileUploadSpec2.from_path(subpath, remote_path / relpath_str)
+
+        def gen():
+            glob = local_path.rglob("*") if recursive else local_path.glob("*")
+            for subpath in glob:
+                # Skip directories and unsupported file types (e.g. block devices)
+                if subpath.is_file():
+                    yield create_spec(subpath)
+
+        self._uploader_generators.append(gen())
+
+    async def _put_file_specs(self, file_specs: list[FileUploadSpec2]):
+        put_responses = {}
+        # num_blocks_total = sum(len(file_spec.blocks_sha256) for file_spec in file_specs)
+
+        # We should only need two iterations: Once to possibly get some missing_blocks; the second time we should have
+        # all blocks uploaded
+        for _ in range(2):
+            files = []
+
+            for file_spec in file_specs:
+                blocks = [
+                    api_pb2.VolumePutFiles2Request.Block(
+                        contents_sha256=block_sha256,
+                        put_response=put_responses.get(block_sha256)
+                    ) for block_sha256 in file_spec.blocks_sha256
+                ]
+                files.append(api_pb2.VolumePutFiles2Request.File(
+                    path=file_spec.path,
+                    mode=file_spec.mode,
+                    size=file_spec.size,
+                    blocks=blocks
+                ))
+
+            request = api_pb2.VolumePutFiles2Request(
+                volume_id=self._volume_id,
+                files=files,
+                disallow_overwrite_existing_files=not self._force,
+            )
+
+            try:
+                response = await retry_transient_errors(self._client.stub.VolumePutFiles2, request, base_delay=1)
+            except GRPCError as exc:
+                raise FileExistsError(exc.message) if exc.status == Status.ALREADY_EXISTS else exc
+
+            if not response.missing_blocks:
+                break
+
+            await _put_missing_blocks(file_specs, response.missing_blocks, put_responses, self._progress_cb)
+        else:
+            raise RuntimeError("Did not succeed at uploading all files despite supplying all missing blocks")
+
+        self._progress_cb(complete=True)
+
+
+VolumeUploadContextManager2 = synchronize_api(_VolumeUploadContextManager2)
+
+
+async def _put_missing_blocks(
+    file_specs: list[FileUploadSpec2],
+    # TODO(dflemstr): Element type is `api_pb2.VolumePutFiles2Response.MissingBlock` but synchronicity gets confused
+    # by the nested class (?)
+    missing_blocks: list,
+    put_responses: dict[bytes, bytes],
+    progress_cb: Callable[..., Any]
+):
+    async def put_missing_block(
+        # TODO(dflemstr): Type is `api_pb2.VolumePutFiles2Response.MissingBlock` but synchronicity gets confused
+        # by the nested class (?)
+        missing_block
+    ) -> (bytes, bytes):
+        # Lazily import to keep the eager loading time of this module down
+        from ._utils.bytes_io_segment_payload import BytesIOSegmentPayload
+
+        assert isinstance(missing_block, api_pb2.VolumePutFiles2Response.MissingBlock)
+
+        file_spec = file_specs[missing_block.file_index]
+        # TODO(dflemstr): What if the underlying file has changed here in the meantime; should we check the
+        #  hash again just to be sure?
+        block_sha256 = file_spec.blocks_sha256[missing_block.block_index]
+        block_start = missing_block.block_index * BLOCK_SIZE
+        block_length = min(BLOCK_SIZE, file_spec.size - block_start)
+
+        progress_name = f"{file_spec.path} block {missing_block.block_index + 1} / {len(file_spec.blocks_sha256)}"
+        progress_task_id = progress_cb(name=progress_name, size=file_spec.size)
+
+        with file_spec.source() as source_fp:
+            payload = BytesIOSegmentPayload(
+                source_fp,
+                block_start,
+                block_length,
+                progress_report_cb=functools.partial(progress_cb, progress_task_id)
+            )
+
+            async with ClientSessionRegistry.get_session().put(
+                missing_block.put_url,
+                data=payload,
+            ) as response:
+                response.raise_for_status()
+                resp_data = await response.content.read()
+
+        return block_sha256, resp_data
+
+    tasks = [
+        asyncio.create_task(put_missing_block(missing_block))
+        for missing_block in missing_blocks
+    ]
+    for task_result in asyncio.as_completed(tasks):
+        digest, resp = await task_result
+        put_responses[digest] = resp
+
 
 def _open_files_error_annotation(mount_path: str) -> Optional[str]:
     if platform.system() != "Linux":

modal/volume.pyi

@@ -1,6 +1,8 @@
+import _io
 import asyncio.locks
 import collections.abc
 import enum
+import google.protobuf.message
 import modal._object
 import modal._utils.blob_utils
 import modal.client
@@ -36,6 +38,7 @@ class FileEntry:
 
 class _Volume(modal._object._Object):
     _lock: typing.Optional[asyncio.locks.Lock]
+    _metadata: typing.Optional[modal_proto.api_pb2.VolumeMetadata]
 
     async def _get_lock(self): ...
     @staticmethod
@@ -46,6 +49,8 @@ class _Volume(modal._object._Object):
         create_if_missing: bool = False,
         version: typing.Optional[int] = None,
     ) -> _Volume: ...
+    def _hydrate_metadata(self, metadata: typing.Optional[google.protobuf.message.Message]): ...
+    def _get_metadata(self) -> typing.Optional[google.protobuf.message.Message]: ...
     @classmethod
     def ephemeral(
         cls: type[_Volume],
@@ -80,7 +85,7 @@ class _Volume(modal._object._Object):
     async def read_file_into_fileobj(self, path: str, fileobj: typing.IO[bytes]) -> int: ...
     async def remove_file(self, path: str, recursive: bool = False) -> None: ...
     async def copy_files(self, src_paths: collections.abc.Sequence[str], dst_path: str) -> None: ...
-    async def batch_upload(self, force: bool = False) -> _VolumeUploadContextManager: ...
+    async def batch_upload(self, force: bool = False) -> _AbstractVolumeUploadContextManager: ...
     async def _instance_delete(self): ...
     @staticmethod
     async def delete(
@@ -95,44 +100,11 @@ class _Volume(modal._object._Object):
         environment_name: typing.Optional[str] = None,
     ): ...
 
-class _VolumeUploadContextManager:
-    _volume_id: str
-    _client: modal.client._Client
-    _force: bool
-    progress_cb: collections.abc.Callable[..., typing.Any]
-    _upload_generators: list[
-        collections.abc.Generator[collections.abc.Callable[[], modal._utils.blob_utils.FileUploadSpec], None, None]
-    ]
-
-    def __init__(
-        self,
-        volume_id: str,
-        client: modal.client._Client,
-        progress_cb: typing.Optional[collections.abc.Callable[..., typing.Any]] = None,
-        force: bool = False,
-    ): ...
-    async def __aenter__(self): ...
-    async def __aexit__(self, exc_type, exc_val, exc_tb): ...
-    def put_file(
-        self,
-        local_file: typing.Union[pathlib.Path, str, typing.BinaryIO],
-        remote_path: typing.Union[pathlib.PurePosixPath, str],
-        mode: typing.Optional[int] = None,
-    ): ...
-    def put_directory(
-        self,
-        local_path: typing.Union[pathlib.Path, str],
-        remote_path: typing.Union[pathlib.PurePosixPath, str],
-        recursive: bool = True,
-    ): ...
-    async def _upload_file(
-        self, file_spec: modal._utils.blob_utils.FileUploadSpec
-    ) -> modal_proto.api_pb2.MountFile: ...
-
 SUPERSELF = typing.TypeVar("SUPERSELF", covariant=True)
 
 class Volume(modal.object.Object):
     _lock: typing.Optional[asyncio.locks.Lock]
+    _metadata: typing.Optional[modal_proto.api_pb2.VolumeMetadata]
 
     def __init__(self, *args, **kwargs): ...
 
@@ -150,6 +122,8 @@ class Volume(modal.object.Object):
         create_if_missing: bool = False,
         version: typing.Optional[int] = None,
     ) -> Volume: ...
+    def _hydrate_metadata(self, metadata: typing.Optional[google.protobuf.message.Message]): ...
+    def _get_metadata(self) -> typing.Optional[google.protobuf.message.Message]: ...
     @classmethod
     def ephemeral(
         cls: type[Volume],
@@ -256,8 +230,8 @@ class Volume(modal.object.Object):
     copy_files: __copy_files_spec[typing_extensions.Self]
 
     class __batch_upload_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(self, force: bool = False) -> VolumeUploadContextManager: ...
-        async def aio(self, force: bool = False) -> VolumeUploadContextManager: ...
+        def __call__(self, force: bool = False) -> AbstractVolumeUploadContextManager: ...
+        async def aio(self, force: bool = False) -> AbstractVolumeUploadContextManager: ...
 
     batch_upload: __batch_upload_spec[typing_extensions.Self]
 
@@ -303,7 +277,92 @@ class Volume(modal.object.Object):
 
     rename: __rename_spec
 
-class VolumeUploadContextManager:
+class _AbstractVolumeUploadContextManager:
+    async def __aenter__(self): ...
+    async def __aexit__(self, exc_type, exc_val, exc_tb): ...
+    def put_file(
+        self,
+        local_file: typing.Union[pathlib.Path, str, typing.BinaryIO, _io.BytesIO],
+        remote_path: typing.Union[pathlib.PurePosixPath, str],
+        mode: typing.Optional[int] = None,
+    ): ...
+    def put_directory(
+        self,
+        local_path: typing.Union[pathlib.Path, str],
+        remote_path: typing.Union[pathlib.PurePosixPath, str],
+        recursive: bool = True,
+    ): ...
+    @staticmethod
+    def resolve(
+        version: int,
+        object_id: str,
+        client,
+        progress_cb: typing.Optional[collections.abc.Callable[..., typing.Any]] = None,
+        force: bool = False,
+    ) -> _AbstractVolumeUploadContextManager: ...
+
+class AbstractVolumeUploadContextManager:
+    def __init__(self, /, *args, **kwargs): ...
+    def __enter__(self): ...
+    async def __aenter__(self): ...
+    def __exit__(self, exc_type, exc_val, exc_tb): ...
+    async def __aexit__(self, exc_type, exc_val, exc_tb): ...
+    def put_file(
+        self,
+        local_file: typing.Union[pathlib.Path, str, typing.BinaryIO, _io.BytesIO],
+        remote_path: typing.Union[pathlib.PurePosixPath, str],
+        mode: typing.Optional[int] = None,
+    ): ...
+    def put_directory(
+        self,
+        local_path: typing.Union[pathlib.Path, str],
+        remote_path: typing.Union[pathlib.PurePosixPath, str],
+        recursive: bool = True,
+    ): ...
+    @staticmethod
+    def resolve(
+        version: int,
+        object_id: str,
+        client,
+        progress_cb: typing.Optional[collections.abc.Callable[..., typing.Any]] = None,
+        force: bool = False,
+    ) -> AbstractVolumeUploadContextManager: ...
+
+class _VolumeUploadContextManager(_AbstractVolumeUploadContextManager):
+    _volume_id: str
+    _client: modal.client._Client
+    _force: bool
+    progress_cb: collections.abc.Callable[..., typing.Any]
+    _upload_generators: list[
+        collections.abc.Generator[collections.abc.Callable[[], modal._utils.blob_utils.FileUploadSpec], None, None]
+    ]
+
+    def __init__(
+        self,
+        volume_id: str,
+        client: modal.client._Client,
+        progress_cb: typing.Optional[collections.abc.Callable[..., typing.Any]] = None,
+        force: bool = False,
+    ): ...
+    async def __aenter__(self): ...
+    async def __aexit__(self, exc_type, exc_val, exc_tb): ...
+    def put_file(
+        self,
+        local_file: typing.Union[pathlib.Path, str, typing.BinaryIO, _io.BytesIO],
+        remote_path: typing.Union[pathlib.PurePosixPath, str],
+        mode: typing.Optional[int] = None,
+    ): ...
+    def put_directory(
+        self,
+        local_path: typing.Union[pathlib.Path, str],
+        remote_path: typing.Union[pathlib.PurePosixPath, str],
+        recursive: bool = True,
+    ): ...
+    async def _upload_file(
+        self, file_spec: modal._utils.blob_utils.FileUploadSpec
+    ) -> modal_proto.api_pb2.MountFile: ...
+
+class VolumeUploadContextManager(AbstractVolumeUploadContextManager):
     _volume_id: str
     _client: modal.client.Client
     _force: bool
@@ -325,7 +384,7 @@ class VolumeUploadContextManager:
     async def __aexit__(self, exc_type, exc_val, exc_tb): ...
     def put_file(
         self,
-        local_file: typing.Union[pathlib.Path, str, typing.BinaryIO],
+        local_file: typing.Union[pathlib.Path, str, typing.BinaryIO, _io.BytesIO],
         remote_path: typing.Union[pathlib.PurePosixPath, str],
         mode: typing.Optional[int] = None,
     ): ...
@@ -342,4 +401,85 @@ class VolumeUploadContextManager:
 
     _upload_file: ___upload_file_spec[typing_extensions.Self]
 
+class _VolumeUploadContextManager2(_AbstractVolumeUploadContextManager):
+    _volume_id: str
+    _client: modal.client._Client
+    _force: bool
+    _progress_cb: collections.abc.Callable[..., typing.Any]
+    _uploader_generators: list[
+        collections.abc.Generator[
+            collections.abc.Callable[[], typing.Awaitable[modal._utils.blob_utils.FileUploadSpec2]]
+        ]
+    ]
+
+    def __init__(
+        self,
+        volume_id: str,
+        client: modal.client._Client,
+        progress_cb: typing.Optional[collections.abc.Callable[..., typing.Any]] = None,
+        force: bool = False,
+    ): ...
+    async def __aenter__(self): ...
+    async def __aexit__(self, exc_type, exc_val, exc_tb): ...
+    def put_file(
+        self,
+        local_file: typing.Union[pathlib.Path, str, typing.BinaryIO, _io.BytesIO],
+        remote_path: typing.Union[pathlib.PurePosixPath, str],
+        mode: typing.Optional[int] = None,
+    ): ...
+    def put_directory(
+        self,
+        local_path: typing.Union[pathlib.Path, str],
+        remote_path: typing.Union[pathlib.PurePosixPath, str],
+        recursive: bool = True,
+    ): ...
+    async def _put_file_specs(self, file_specs: list[modal._utils.blob_utils.FileUploadSpec2]): ...
+
+class VolumeUploadContextManager2(AbstractVolumeUploadContextManager):
+    _volume_id: str
+    _client: modal.client.Client
+    _force: bool
+    _progress_cb: collections.abc.Callable[..., typing.Any]
+    _uploader_generators: list[
+        collections.abc.Generator[
+            collections.abc.Callable[[], typing.Awaitable[modal._utils.blob_utils.FileUploadSpec2]]
+        ]
+    ]
+
+    def __init__(
+        self,
+        volume_id: str,
+        client: modal.client.Client,
+        progress_cb: typing.Optional[collections.abc.Callable[..., typing.Any]] = None,
+        force: bool = False,
+    ): ...
+    def __enter__(self): ...
+    async def __aenter__(self): ...
+    def __exit__(self, exc_type, exc_val, exc_tb): ...
+    async def __aexit__(self, exc_type, exc_val, exc_tb): ...
+    def put_file(
+        self,
+        local_file: typing.Union[pathlib.Path, str, typing.BinaryIO, _io.BytesIO],
+        remote_path: typing.Union[pathlib.PurePosixPath, str],
+        mode: typing.Optional[int] = None,
+    ): ...
+    def put_directory(
+        self,
+        local_path: typing.Union[pathlib.Path, str],
+        remote_path: typing.Union[pathlib.PurePosixPath, str],
+        recursive: bool = True,
+    ): ...
+
+    class ___put_file_specs_spec(typing_extensions.Protocol[SUPERSELF]):
+        def __call__(self, file_specs: list[modal._utils.blob_utils.FileUploadSpec2]): ...
+        async def aio(self, file_specs: list[modal._utils.blob_utils.FileUploadSpec2]): ...
+
+    _put_file_specs: ___put_file_specs_spec[typing_extensions.Self]
+
+async def _put_missing_blocks(
+    file_specs: list[modal._utils.blob_utils.FileUploadSpec2],
+    missing_blocks: list,
+    put_responses: dict[bytes, bytes],
+    progress_cb: collections.abc.Callable[..., typing.Any],
+): ...
 def _open_files_error_annotation(mount_path: str) -> typing.Optional[str]: ...

{modal-0.74.15.dist-info → modal-0.74.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modal
-Version: 0.74.15
+Version: 0.74.16
 Summary: Python client library for Modal
 Author-email: Modal Labs <support@modal.com>
 License: Apache-2.0

{modal-0.74.15.dist-info → modal-0.74.16.dist-info}/RECORD

@@ -22,7 +22,7 @@ modal/app.py,sha256=1GSPMMYRkG17FQwnA_PleK9YQrPQWXlN_4mpZ5VRA4Q,48468
 modal/app.pyi,sha256=bpC9uN_B4d_UtyXVuhGpXC2RM-IcsytS0ndu-GoKrHQ,27276
 modal/call_graph.py,sha256=1g2DGcMIJvRy-xKicuf63IVE98gJSnQsr8R_NVMptNc,2581
 modal/client.py,sha256=U-YKSw0n7J1ZLREt9cbEJCtmHe5YoPKFxl0xlkan2yc,15565
-modal/client.pyi,sha256=gTGLwo62ctAnv6FICBPjNpwC5GMyQ6CyL5fBhdNM0C4,7593
+modal/client.pyi,sha256=MXDCG1c77_KDP7R66sApXRKz9ZnhLsZMjSETzBqGDzA,7593
 modal/cloud_bucket_mount.py,sha256=YOe9nnvSr4ZbeCn587d7_VhE9IioZYRvF9VYQTQux08,5914
 modal/cloud_bucket_mount.pyi,sha256=30T3K1a89l6wzmEJ_J9iWv9SknoGqaZDx59Xs-ZQcmk,1607
 modal/cls.py,sha256=GvaNl8R5UsH7Vg88WEOyerdjvZEPK7xxi3nqHlyOW_c,33497
@@ -39,7 +39,7 @@ modal/file_io.py,sha256=lcMs_E9Xfm0YX1t9U2wNIBPnqHRxmImqjLW1GHqVmyg,20945
 modal/file_io.pyi,sha256=NTRft1tbPSWf9TlWVeZmTlgB5AZ_Zhu2srWIrWr7brk,9445
 modal/file_pattern_matcher.py,sha256=trosX-Bp7dOubudN1bLLhRAoidWy1TcoaR4Pv8CedWw,6497
 modal/functions.py,sha256=kcNHvqeGBxPI7Cgd57NIBBghkfbeFJzXO44WW0jSmao,325
-modal/functions.pyi,sha256=2J7jPa2tQWhuUy67wg-MWvEx0rwrBNjX6UTEJbAUq_A,14855
+modal/functions.pyi,sha256=I8dNJnvs6FguR8XAUqnMUg6X1214UzeLoAydZBrEa2g,14855
 modal/gpu.py,sha256=Kbhs_u49FaC2Zi0TjCdrpstpRtT5eZgecynmQi5IZVE,6752
 modal/image.py,sha256=I-9_YZL0SSfnuGPywa3-4PlxDmJ-53p7ce3gP74SrOA,92877
 modal/image.pyi,sha256=89zv12C1sFrJs7Es9SnX23_m208ASAdeNGCVTrhjzHI,25632
@@ -78,8 +78,8 @@ modal/snapshot.pyi,sha256=Ypd4NKsjOTnnnqXyTGGLKq5lkocRrUURYjY5Pi67_qA,670
 modal/stream_type.py,sha256=A6320qoAAWhEfwOCZfGtymQTu5AfLfJXXgARqooTPvY,417
 modal/token_flow.py,sha256=APNPQhuqy2avu6IY26U6OiN7erRhpo03Ztt1V60_wts,6776
 modal/token_flow.pyi,sha256=0XV3d-9CGQL3qjPdw3RgwIFVqqxo8Z-u044_mkgAM3o,2064
-modal/volume.py,sha256=JAWeDvoAG95tMBv-fYIERyHsJPS_X_xGpxRRmYtb6j0,30096
-modal/volume.pyi,sha256=kTsXarphjZILXci84LQy7EyC84eXUs5-7D62IM5q3eE,12491
+modal/volume.py,sha256=3c5_aJNJtgpsFRZWBjc0jwn8Zs0jo9V6UDmh6ifrbdA,40145
+modal/volume.pyi,sha256=juOVWGlgz7IeOY4M7jBhbeNRPA9xdGUwvA3AzlZUscQ,17958
 modal/_runtime/__init__.py,sha256=MIEP8jhXUeGq_eCjYFcqN5b1bxBM4fdk0VESpjWR0fc,28
 modal/_runtime/asgi.py,sha256=KNarxvZI9z8fnmZl2vbkWTjnoLXs9kqOahkrbsTLkyc,22429
 modal/_runtime/container_io_manager.py,sha256=0yNO3HTVIM4f338rxJavD8nrRN7KhDpjz1jLux71MRY,43842
@@ -92,7 +92,7 @@ modal/_runtime/user_code_imports.py,sha256=kAv37Pl1TmGKduv0Kozum0xNTD42bDLloSIsT
 modal/_utils/__init__.py,sha256=waLjl5c6IPDhSsdWAm9Bji4e2PVxamYABKAze6CHVXY,28
 modal/_utils/app_utils.py,sha256=88BT4TPLWfYAQwKTHcyzNQRHg8n9B-QE2UyJs96iV-0,108
 modal/_utils/async_utils.py,sha256=b2TJyKY1Hq7df7M-fo3qlFM95mGdo3dCuqRPPcV5hsE,27445
-modal/_utils/blob_utils.py,sha256=jWJovk4g-YNG3CvkvglOds4a6D1M0Tcal_59v7y9VsM,14591
+modal/_utils/blob_utils.py,sha256=WhWaSFLcffVDchMyfKys-j0EFVbm9l7rebzeoa7Z2jM,18214
 modal/_utils/bytes_io_segment_payload.py,sha256=uunxVJS4PE1LojF_UpURMzVK9GuvmYWRqQo_bxEj5TU,3385
 modal/_utils/deprecation.py,sha256=EXP1beU4pmEqEzWMLw6E3kUfNfpmNA_VOp6i0EHi93g,4856
 modal/_utils/docker_utils.py,sha256=h1uETghR40mp_y3fSWuZAfbIASH1HMzuphJHghAL6DU,3722
@@ -132,7 +132,7 @@ modal/cli/run.py,sha256=NX2wWwj8HD6XUhnZRF808Qy9eeouv8KnvyOP57HqIXI,23637
 modal/cli/secret.py,sha256=WB_c-LE9-eDqleLpJxsJ9rZw62Eeza8ZFQFR10vNMEk,4197
 modal/cli/token.py,sha256=mxSgOWakXG6N71hQb1ko61XAR9ZGkTMZD-Txn7gmTac,1924
 modal/cli/utils.py,sha256=hZmjyzcPjDnQSkLvycZD2LhGdcsfdZshs_rOU78EpvI,3717
-modal/cli/volume.py,sha256=c2IuVNO2yJVaXmZkRh3xwQmznlRTgFoJr_BIzzqtVv0,10251
+modal/cli/volume.py,sha256=_PJ5Vn_prkLk_x1Lksx4kZySlKWqIn36T2Edd1-h7Mg,10497
 modal/cli/programs/__init__.py,sha256=svYKtV8HDwDCN86zbdWqyq5T8sMdGDj0PVlzc2tIxDM,28
 modal/cli/programs/run_jupyter.py,sha256=YVvJYu927A4ji72d6i27CKfyZ_uDWteeittARtJnf7E,2775
 modal/cli/programs/vscode.py,sha256=kfvhZQ4bJwtVm3MgC1V7AlygZOlKT1a33alr_uwrewA,3473
@@ -145,7 +145,7 @@ modal/requirements/2024.10.txt,sha256=qD-5cVIVM9wXesJ6JC89Ew-3m2KjEElUz3jaw_MddR
 modal/requirements/PREVIEW.txt,sha256=qD-5cVIVM9wXesJ6JC89Ew-3m2KjEElUz3jaw_MddRo,296
 modal/requirements/README.md,sha256=9tK76KP0Uph7O0M5oUgsSwEZDj5y-dcUPsnpR0Sc-Ik,854
 modal/requirements/base-images.json,sha256=57vMSqzMbLBxw5tFWSaMiIkkVEps4JfX5PAtXGnkS4U,740
-modal-0.74.15.dist-info/licenses/LICENSE,sha256=psuoW8kuDP96RQsdhzwOqi6fyWv0ct8CR6Jr7He_P_k,10173
+modal-0.74.16.dist-info/licenses/LICENSE,sha256=psuoW8kuDP96RQsdhzwOqi6fyWv0ct8CR6Jr7He_P_k,10173
 modal_docs/__init__.py,sha256=svYKtV8HDwDCN86zbdWqyq5T8sMdGDj0PVlzc2tIxDM,28
 modal_docs/gen_cli_docs.py,sha256=c1yfBS_x--gL5bs0N4ihMwqwX8l3IBWSkBAKNNIi6bQ,3801
 modal_docs/gen_reference_docs.py,sha256=cvTgltucqYLLIX84QxAwf51Z5Vc2n6cLxS8VcrxNCAo,6401
@@ -170,9 +170,9 @@ modal_proto/options_pb2_grpc.pyi,sha256=CImmhxHsYnF09iENPoe8S4J-n93jtgUYD2JPAc0y
 modal_proto/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 modal_version/__init__.py,sha256=m94xZNWIjH8oUtJk4l9xfovzDJede2o7X-q0MHVECtM,470
 modal_version/__main__.py,sha256=2FO0yYQQwDTh6udt1h-cBnGd1c4ZyHnHSI4BksxzVac,105
-modal_version/_version_generated.py,sha256=BwTNwunmOMlhLbdcrq2kBshM6ByivCgCW2YhrtCy1ss,149
-modal-0.74.15.dist-info/METADATA,sha256=1EpkkZUHT-Oleu-NfpBAaFlurY2PnvSTOXlkpWh12lg,2474
-modal-0.74.15.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
-modal-0.74.15.dist-info/entry_points.txt,sha256=An-wYgeEUnm6xzrAP9_NTSTSciYvvEWsMZILtYrvpAI,46
-modal-0.74.15.dist-info/top_level.txt,sha256=4BWzoKYREKUZ5iyPzZpjqx4G8uB5TWxXPDwibLcVa7k,43
-modal-0.74.15.dist-info/RECORD,,
+modal_version/_version_generated.py,sha256=iw6S4gIMMk_kxXdUThejt0bL01XM4XObW5lyo2KKxMA,149
+modal-0.74.16.dist-info/METADATA,sha256=SOCGPNZKh82g8EnIIGaDKd-PefSR4ucoYk5mS7JCyoE,2474
+modal-0.74.16.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
+modal-0.74.16.dist-info/entry_points.txt,sha256=An-wYgeEUnm6xzrAP9_NTSTSciYvvEWsMZILtYrvpAI,46
+modal-0.74.16.dist-info/top_level.txt,sha256=4BWzoKYREKUZ5iyPzZpjqx4G8uB5TWxXPDwibLcVa7k,43
+modal-0.74.16.dist-info/RECORD,,

modal_version/_version_generated.py

@@ -1,4 +1,4 @@
 # Copyright Modal Labs 2025
 
 # Note: Reset this value to -1 whenever you make a minor `0.X` release of the client.
-build_number = 15  # git: f5929f1
+build_number = 16  # git: 7430900