Progress bar + minor bugfix

python/custom_routes.py

@@ -19,6 +19,10 @@ import server
 import execution
 import folder_paths
 
+from .upload import upload_file
+from .upload.blob import progress
+
+
 task_status = {}
 
 def random_seed(num_digits=15):
@@ -181,7 +185,7 @@ def load_custom_node(module_path, ignore=set()):
 
             return (module_name, mappings)
         else:
-            print(f"Skip {module_path} module for custom nodes due to the lack of NODE_CLASS_MAPPINGS.")
+            #print(f"Skip {module_path} module for custom nodes due to the lack of NODE_CLASS_MAPPINGS.")
             return (module_name, [])
     except Exception as e:
         print(traceback.format_exc())
@@ -207,9 +211,13 @@ def make_post_request_with_retry(url, data, headers=None, max_retries=3, retry_d
                 raise
 
 def update_dependencies():
+
+    base = folder_paths.base_path
+    custom_nodes_dir = os.path.join(base, "custom_nodes")
+
     # Loop over the files in the specified directory
-    for custom_node in os.listdir("../../custom_nodes/"):
-        filepath = os.path.join("../../custom_nodes", custom_node, "requirements.txt")
+    for custom_node in os.listdir(custom_nodes_dir):
+        filepath = os.path.join(custom_nodes_dir, custom_node, "requirements.txt")
         
         # Check if the file is a requirements.txt file
         if os.path.exists(filepath):
@@ -239,23 +247,19 @@ def update_dependencies():
                 file.writelines(updated_lines)
             
 async def upload_task_execution(task_id, json_response, workflow_id, models_dep, nodes_dep, files):
-    def run():
-        try:
-            exec(base64.b64decode(json_response["data"]), {
-                    "workflow_id": workflow_id,
-                    "models_dep": models_dep,
-                    "nodes_dep": nodes_dep,
-                    "files": files,
-                }
-            )
-            # cleanup temp
-            task_status[task_id] = {"status": "Task completed", "message": "Upload successful"}
-        except Exception as e:
-            task_status[task_id] = {"status": f"Task failed", "message": str(e)}
+    try:
+        await upload_file(
+            workflow_id,
+            models_dep,
+            nodes_dep,
+            files,
+        )
+
+        # cleanup temp
+        task_status[task_id] = {"status": "Task completed", "message": "Upload successful"}
+    except Exception as e:
+        task_status[task_id] = {"status": f"Task failed", "message": str(e)}
 
-    loop = asyncio.get_event_loop()
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        await loop.run_in_executor(executor, run)
 
 @server.PromptServer.instance.routes.post("/comfy-cloud/upload")
 async def upload_dependencies(request):
@@ -282,7 +286,6 @@ async def upload_dependencies(request):
             "token": current_datetime.strftime("%Y-%m-%d %H:%M:%S")
         }
 
-
         url = f"{endpoint}/e"
         response = make_post_request_with_retry(url, data=body, headers=headers)
 
@@ -339,7 +342,8 @@ async def upload_dependencies(request):
 @server.PromptServer.instance.routes.get("/comfy-cloud/upload-status/{task_id}")
 async def get_task_status(request):
     task_id = request.match_info['task_id']
-    status = {"status": "Task failed", "message": str("asd")} #task_status.get(task_id, {"status": "Task failed", "message": "Task not found"})
+    status = task_status.get(task_id, {"status": "Task failed", "message": "Task not found"})
+    status["progress"] = progress
     return web.json_response(status)    
 
 

python/upload/__init__.py

@@ -0,0 +1,184 @@
+import io
+import os
+import asyncio
+
+from .volume import VolumeUploadContextManager
+from .blob import perform_multipart_upload, BytesIOSegmentPayload, _upload_to_s3_url
+from .hash import get_upload_hashes
+
+import folder_paths
+import copy 
+
+import requests
+import json
+import aiohttp
+
+def get_content_length(data):
+    # *Remaining* length of file from current seek position
+    pos = data.tell()
+    data.seek(0, os.SEEK_END)
+    content_length = data.tell()
+    data.seek(pos)
+    return content_length - pos
+
+async def get_file_specs(workflow_id, models_dep, nodes_dep, files):
+    paths = copy.deepcopy(folder_paths.folder_names_and_paths)
+    paths.pop("custom_nodes", None)
+    paths.pop("configs", None)
+
+    base = folder_paths.base_path
+    models_dir = os.path.join(base, "models")
+    custom_nodes_dir = os.path.join(base, "custom_nodes")
+    input_dir = os.path.join(base, "input")
+
+    file_specs = None
+    # Calculate file hashes
+    with VolumeUploadContextManager(None, file_specs, force=True) as batch:
+
+        for name in models_dep:
+            for base_dir in paths.keys():
+                if name in folder_paths.get_filename_list(base_dir):
+                    model_path = os.path.join(models_dir, base_dir, name)
+                    if os.path.exists(model_path):
+                        batch.put_file(model_path, f"/vol/{workflow_id}/comfyui/models/{base_dir}/{name}")
+                    else:
+                        raise Exception(f"Model '{path}' does not exist.")
+
+        # Upload custom nodes
+        for name, value in nodes_dep.items():
+            custom_node_path = os.path.join(custom_nodes_dir, name)
+
+            if os.path.exists(custom_node_path):
+                batch.put_directory(custom_node_path, f"/vol/{workflow_id}/comfyui/custom_nodes/{name}")
+            else:
+                raise Exception(f"Custom node '{path}' does not exist.")
+
+
+        # Upload files
+        for filename in files:
+            input_file_path = os.path.join(input_dir, filename)
+            if os.path.exists(input_file_path):
+                batch.put_file(input_file_path, f"/vol/{workflow_id}/comfyui/input/{filename}")
+            else:
+                raise Exception(f"Input '{path}' does not exist.")
+
+
+        #batch.put_directory(path, "/helloworld")
+        #batch.put_file(path, target_path)
+        x = batch.generate_specs()
+
+
+        file_specs = x
+
+    # If server responds with s3 url, filename, and spec -
+    # then upload
+    s = []
+
+    # Prepare new upload dict
+    async for item in file_specs:
+        obj = item.__dict__
+
+        with item.source() as data:
+            upload_hashes = get_upload_hashes(data)
+
+            if isinstance(data, bytes):
+                data = io.BytesIO(data)
+            
+            content_length = get_content_length(data)
+
+            obj["upload_hashes"] = upload_hashes.__dict__
+            obj["content_length"] = content_length
+            del obj["source"]
+
+        s.append(obj)
+
+    return s
+
+
+
+burl = "https://pytest-12acbe397fbc.herokuapp.com"
+
+async def make_post_request(url, data):
+    async with aiohttp.ClientSession() as session:
+        async with session.post(url, json=data) as response:
+            return await response.json()
+
+# @todo does not support dirs yet
+async def upload_file(workflow_id, models_dep, nodes_dep, files):
+    try:
+
+        # Client generates file specs
+        specs = await get_file_specs(workflow_id, models_dep, nodes_dep, files)
+
+
+        url = f'{burl}/upload-urls'
+        data = {
+            "specs": specs
+        }
+        response_data = await make_post_request(url, data)
+        outputs = response_data["data"]
+
+        # @todo - handle non multipart
+        blob_ids = []
+        for x in outputs:
+
+            is_multipart = x["type"] == "multipart"
+            resp = x["data"]
+
+            if is_multipart:
+
+                # find file
+                filename = resp["filename"]
+                data = open(filename, "rb")
+
+                # Get filename, and search for file
+                # resp obj:
+                # - content_length
+                # - max_part_size
+                # - part_urls 
+                # - completion_url
+                # - filename
+                await perform_multipart_upload(
+                    data,
+                    content_length=resp["content_length"],
+                    max_part_size=resp["max_part_size"],
+                    part_urls=resp["part_urls"],
+                    completion_url=resp["completion_url"],
+                    filename=filename,
+                )
+
+            else:
+                filename = resp["filename"]
+                content_length = resp["content_length"]
+                upload_hashes = resp["upload_hashes"]
+                data = open(filename, "rb")
+
+                payload = BytesIOSegmentPayload(data, segment_start=0, segment_length=content_length, filename = filename)
+                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"],
+                )
+
+            blob_ids.append({
+                "type": x["type"],
+                "blob_id": resp["blob_id"],
+                "use_blob": resp["use_blob"],
+                "sha256_hex": resp["sha256_hex"],
+                "mode": resp["mode"],
+                "content": resp["content"],
+                "mount_filename": resp["mount_filename"],
+            })
+
+        # Once done, send resp.blob_id back to server
+
+        url = f'{burl}/commit'
+        data = {
+            "blob_ids":blob_ids 
+        }
+        response_data = await make_post_request(url, data)
+
+    except Exception as e:
+        print("error", e)
+

python/upload/blob.py

@@ -0,0 +1,359 @@
+import asyncio
+import dataclasses
+import hashlib
+import io
+import os
+import platform
+from contextlib import contextmanager
+from pathlib import Path, PurePosixPath
+from typing import Any, AsyncIterator, BinaryIO, Callable, List, Optional, Union
+from urllib.parse import urlparse
+
+from aiohttp import BytesIOPayload
+from aiohttp.abc import AbstractStreamWriter
+
+from modal.exception import ExecutionError
+from .sync import retry
+
+from modal._utils.hash_utils import get_sha256_hex
+from .net import http_client_with_tls, retry_transient_errors
+
+from tqdm import tqdm
+import logging
+
+progress = {}
+
+logger = logging.getLogger()
+
+
+# Max size for function inputs and outputs.
+MAX_OBJECT_SIZE_BYTES = 2 * 1024 * 1024  # 2 MiB
+
+#  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
+
+# 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,
+        filename: str,
+        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()
+        self.filename = filename
+
+    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 filename(self):
+        return self._filename
+
+    @filename.setter
+    def filename(self, value):
+        self._filename = value
+
+    @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):
+        if self.filename not in progress:
+            progress[self.filename] = {}
+        
+        progress[self.filename]["value"] = self.segment_length - self.num_bytes_read
+        if "max" not in progress[self.filename]:
+            progress[self.filename]["max"] = self.segment_length - self.num_bytes_read
+
+        return self.segment_length - self.num_bytes_read
+
+
+@retry(n_attempts=5, base_delay=0.5, timeout=None)
+async def _upload_to_s3_url(
+    upload_url,
+    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
+
+
+async def perform_multipart_upload(
+    data_file: Union[BinaryIO, io.BytesIO, io.FileIO],
+    *,
+    content_length: int,
+    max_part_size: int,
+    part_urls: List[str],
+    completion_url: str,
+    filename: str,
+    upload_chunk_size: int = DEFAULT_SEGMENT_CHUNK_SIZE,
+):
+    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]
+    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))]
+    else:
+        filename = data_file.name
+        data_file_readers = [open(filename, "rb") for _ in range(len(part_urls))]
+
+    for part_number, (data_file_rdr, part_url) in enumerate(zip(data_file_readers, part_urls), start=1):
+        part_length_bytes = min(num_bytes_left, max_part_size)
+        part_payload = BytesIOSegmentPayload(
+            data_file_rdr,
+            segment_start=file_offset,
+            segment_length=part_length_bytes,
+            chunk_size=upload_chunk_size,
+            filename=filename,
+        )
+        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)
+
+    # The body of the complete_multipart_upload command needs some data in xml format:
+    completion_body = "<CompleteMultipartUpload>\n"
+    for part_number, etag in enumerate(part_etags, 1):
+        completion_body += f"""<Part>\n<PartNumber>{part_number}</PartNumber>\n<ETag>"{etag}"</ETag>\n</Part>\n"""
+    completion_body += "</CompleteMultipartUpload>"
+
+    # etag of combined object should be md5 hex of concatendated md5 *bytes* from parts + `-{num_parts}`
+    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}"
+                )
+
+
+def get_content_length(data: BinaryIO):
+    # *Remaining* length of file from current seek position
+    pos = data.tell()
+    data.seek(0, os.SEEK_END)
+    content_length = data.tell()
+    data.seek(pos)
+    return content_length - pos
+
+
+@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()
+
+
+@dataclasses.dataclass
+class FileUploadSpec:
+    source: Callable[[], BinaryIO]
+    source_description: Any
+    mount_filename: str
+
+    use_blob: bool
+    content: Optional[bytes]  # typically None if using blob, required otherwise
+    sha256_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
+) -> FileUploadSpec:
+    with source() as fp:
+        # Current position is ignored - we always upload from position 0
+        fp.seek(0, os.SEEK_END)
+        size = fp.tell()
+        fp.seek(0)
+
+        if size >= LARGE_FILE_LIMIT:
+            use_blob = True
+            content = None
+            sha256_hex = get_sha256_hex(fp)
+        else:
+            use_blob = False
+            content = fp.read()
+            sha256_hex = get_sha256_hex(content)
+
+    return FileUploadSpec(
+        source=source,
+        source_description=source_description,
+        mount_filename=mount_filename.as_posix(),
+        use_blob=use_blob,
+        content=content,
+        sha256_hex=sha256_hex,
+        mode=mode & 0o7777,
+        size=size,
+    )
+
+
+def get_file_upload_spec_from_path(
+    filename: Path, mount_filename: PurePosixPath, mode: Optional[int] = None
+) -> FileUploadSpec:
+    # 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)
+    return _get_file_upload_spec(
+        lambda: open(filename, "rb"),
+        filename,
+        mount_filename,
+        mode,
+    )
+
+
+def get_file_upload_spec_from_fileobj(fp: BinaryIO, mount_filename: PurePosixPath, mode: int) -> FileUploadSpec:
+    def source():
+        # We ignore position in stream and always upload from position 0
+        fp.seek(0)
+        return fp
+
+    return _get_file_upload_spec(
+        source,
+        str(fp),
+        mount_filename,
+        mode,
+    )
+
+
+def use_md5(url: str) -> bool:
+    """This takes an upload URL in S3 and returns whether we should attach a checksum.
+
+    It's only a workaround for missing functionality in moto.
+    https://github.com/spulec/moto/issues/816
+    """
+    host = urlparse(url).netloc.split(":")[0]
+    if host.endswith(".amazonaws.com"):
+        return True
+    elif host in ["127.0.0.1", "localhost", "172.21.0.1"]:
+        return False
+    else:
+        raise Exception(f"Unknown S3 host: {host}")

python/upload/hash.py

@@ -0,0 +1,57 @@
+import base64
+import dataclasses
+import hashlib
+from typing import IO, Union
+
+HASH_CHUNK_SIZE = 4096
+
+
+def _update(hashers, data: Union[bytes, IO[bytes]]):
+    if isinstance(data, bytes):
+        for hasher in hashers:
+            hasher.update(data)
+    else:
+        pos = data.tell()
+        while 1:
+            chunk = data.read(HASH_CHUNK_SIZE)
+            if not isinstance(chunk, bytes):
+                raise ValueError(f"Only accepts bytes or byte buffer objects, not {type(chunk)} buffers")
+            if not chunk:
+                break
+            for hasher in hashers:
+                hasher.update(chunk)
+        data.seek(pos)
+
+
+def get_sha256_hex(data: Union[bytes, IO[bytes]]) -> str:
+    hasher = hashlib.sha256()
+    _update([hasher], data)
+    return hasher.hexdigest()
+
+
+def get_sha256_base64(data: Union[bytes, IO[bytes]]) -> str:
+    hasher = hashlib.sha256()
+    _update([hasher], data)
+    return base64.b64encode(hasher.digest()).decode("ascii")
+
+
+def get_md5_base64(data: Union[bytes, IO[bytes]]) -> str:
+    hasher = hashlib.md5()
+    _update([hasher], data)
+    return base64.b64encode(hasher.digest()).decode("utf-8")
+
+
+@dataclasses.dataclass
+class UploadHashes:
+    md5_base64: str
+    sha256_base64: str
+
+
+def get_upload_hashes(data: Union[bytes, IO[bytes]]) -> UploadHashes:
+    md5 = hashlib.md5()
+    sha256 = hashlib.sha256()
+    _update([md5, sha256], data)
+    return UploadHashes(
+        md5_base64=base64.b64encode(md5.digest()).decode("ascii"),
+        sha256_base64=base64.b64encode(sha256.digest()).decode("ascii"),
+    )

python/upload/net.py

@@ -0,0 +1,128 @@
+import asyncio
+import contextlib
+import certifi
+import socket
+import ssl
+import time
+import uuid
+
+from aiohttp import ClientSession, ClientTimeout, TCPConnector
+from aiohttp.web import Application
+from aiohttp.web_runner import AppRunner, SockSite
+from grpclib import GRPCError
+from grpclib import GRPCError, Status
+from grpclib.exceptions import StreamTerminatedError
+from typing import (
+    Any,
+    AsyncIterator,
+    Dict,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+)
+
+
+RETRYABLE_GRPC_STATUS_CODES = [
+    Status.DEADLINE_EXCEEDED,
+    Status.UNAVAILABLE,
+    Status.CANCELLED,
+    Status.INTERNAL,
+]
+
+def http_client_with_tls(timeout: Optional[float]) -> ClientSession:
+    """Create a new HTTP client session with standard, bundled TLS certificates.
+
+    This is necessary to prevent client issues on some system where Python does
+    not come pre-installed with specific TLS certificates that are necessary to
+    connect to AWS S3 bucket URLs.
+
+    Specifically: the error "unable to get local issuer certificate" when making
+    an aiohttp request.
+    """
+    ssl_context = ssl.create_default_context(cafile=certifi.where())
+    connector = TCPConnector(ssl=ssl_context)
+    return ClientSession(connector=connector, timeout=ClientTimeout(total=timeout))
+
+
+@contextlib.asynccontextmanager
+async def run_temporary_http_server(app: Application):
+    # Allocates a random port, runs a server in a context manager
+    # This is used in various tests
+    sock = socket.socket()
+    sock.bind(("", 0))
+    port = sock.getsockname()[1]
+    host = f"http://127.0.0.1:{port}"
+
+    runner = AppRunner(app)
+    await runner.setup()
+    site = SockSite(runner, sock=sock)
+    await site.start()
+    try:
+        yield host
+    finally:
+        await runner.cleanup()
+
+
+
+
+
+async def retry_transient_errors(
+    fn,
+    *args,
+    base_delay: float = 0.1,
+    max_delay: float = 1,
+    delay_factor: float = 2,
+    max_retries: Optional[int] = 3,
+    additional_status_codes: list = [],
+    attempt_timeout: Optional[float] = None,  # timeout for each attempt
+    total_timeout: Optional[float] = None,  # timeout for the entire function call
+    attempt_timeout_floor=2.0,  # always have at least this much timeout (only for total_timeout)
+):
+    """Retry on transient gRPC failures with back-off until max_retries is reached.
+    If max_retries is None, retry forever."""
+
+    delay = base_delay
+    n_retries = 0
+
+    status_codes = [*RETRYABLE_GRPC_STATUS_CODES, *additional_status_codes]
+
+    idempotency_key = str(uuid.uuid4())
+
+    t0 = time.time()
+    if total_timeout is not None:
+        total_deadline = t0 + total_timeout
+    else:
+        total_deadline = None
+
+    while True:
+        metadata = [("x-idempotency-key", idempotency_key), ("x-retry-attempt", str(n_retries))]
+        if n_retries > 0:
+            metadata.append(("x-retry-delay", str(time.time() - t0)))
+        timeouts = []
+        if attempt_timeout is not None:
+            timeouts.append(attempt_timeout)
+        if total_timeout is not None:
+            timeouts.append(max(total_deadline - time.time(), attempt_timeout_floor))
+        if timeouts:
+            timeout = min(timeouts)  # In case the function provided both types of timeouts
+        else:
+            timeout = None
+        try:
+            return await fn(*args, metadata=metadata, timeout=timeout)
+        except (StreamTerminatedError, GRPCError, socket.gaierror, asyncio.TimeoutError) as exc:
+            if isinstance(exc, GRPCError) and exc.status not in status_codes:
+                raise exc
+
+            if max_retries is not None and n_retries >= max_retries:
+                raise exc
+
+            if total_deadline and time.time() + delay + attempt_timeout_floor >= total_deadline:
+                # no point sleeping if that's going to push us past the deadline
+                raise exc
+
+            n_retries += 1
+
+            await asyncio.sleep(delay)
+            delay = min(delay * delay_factor, max_delay)
+

python/upload/sync.py

@@ -0,0 +1,81 @@
+import asyncio
+import concurrent.futures
+import functools
+import inspect
+import sys
+import time
+import typing
+from contextlib import asynccontextmanager
+from typing import Any, AsyncGenerator, Callable, Iterator, List, Optional, Set, TypeVar, cast
+
+import synchronicity
+from typing_extensions import ParamSpec
+
+
+synchronizer = synchronicity.Synchronizer()
+
+def synchronize_api(obj, target_module=None):
+    if inspect.isclass(obj):
+        blocking_name = obj.__name__.lstrip("_")
+    elif inspect.isfunction(object):
+        blocking_name = obj.__name__.lstrip("_")
+    elif isinstance(obj, TypeVar):
+        blocking_name = "_BLOCKING_" + obj.__name__
+    else:
+        blocking_name = None
+    if target_module is None:
+        target_module = obj.__module__
+    return synchronizer.create_blocking(obj, blocking_name, target_module=target_module)
+
+def retry(direct_fn=None, *, n_attempts=3, base_delay=0, delay_factor=2, timeout=90):
+    """Decorator that calls an async function multiple times, with a given timeout.
+
+    If a `base_delay` is provided, the function is given an exponentially
+    increasing delay on each run, up until the maximum number of attempts.
+
+    Usage:
+
+    ```
+    @retry
+    async def may_fail_default():
+        # ...
+        pass
+
+    @retry(n_attempts=5, base_delay=1)
+    async def may_fail_delay():
+        # ...
+        pass
+    ```
+    """
+
+    def decorator(fn):
+        @functools.wraps(fn)
+        async def f_wrapped(*args, **kwargs):
+            delay = base_delay
+            for i in range(n_attempts):
+                t0 = time.time()
+                try:
+                    return await asyncio.wait_for(fn(*args, **kwargs), timeout=timeout)
+                except asyncio.CancelledError:
+                    logger.debug(f"Function {fn} was cancelled")
+                    raise
+                except Exception as e:
+                    if i >= n_attempts - 1:
+                        raise
+                    logger.debug(
+                        f"Failed invoking function {fn}: {e}"
+                        f" (took {time.time() - t0}s, sleeping {delay}s"
+                        f" and trying {n_attempts - i - 1} more times)"
+                    )
+                await asyncio.sleep(delay)
+                delay *= delay_factor
+
+        return f_wrapped
+
+    if direct_fn is not None:
+        # It's invoked like @retry
+        return decorator(direct_fn)
+    else:
+        # It's invoked like @retry(n_attempts=...)
+        return decorator
+

python/upload/volume.py

@@ -0,0 +1,144 @@
+import io
+import copy 
+import os
+import base64
+from aiohttp import web
+import asyncio
+
+import concurrent.futures
+import time
+from pathlib import Path, PurePosixPath
+from typing import (
+    IO,
+    AsyncGenerator,
+    AsyncIterator,
+    BinaryIO,
+    Callable,
+    Generator,
+    List,
+    Optional,
+    Sequence,
+    Type,
+    Union,
+    Any, 
+)
+import aiostream
+from .sync import synchronize_api
+
+from .blob import (
+    get_file_upload_spec_from_fileobj,
+    get_file_upload_spec_from_path,
+)
+
+
+class _VolumeUploadContextManager:
+    """Context manager for batch-uploading files to a Volume."""
+
+    _volume_id: str
+    _force: bool
+
+    def __init__(self, volume_id: str, file_specs, force: bool = False):
+        """mdmd:hidden"""
+        self._volume_id = volume_id
+        self._client = None #client
+        self._upload_generators = []
+        self._force = force
+
+        self.files_stream = file_specs
+
+    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._upload_generators:
+                    yield from gen
+
+            async def gen_file_upload_specs(): # -> AsyncGenerator[FileUploadSpec, None]:
+                loop = asyncio.get_event_loop()
+                with concurrent.futures.ThreadPoolExecutor() as exe:
+                    # TODO: avoid eagerly expanding
+                    futs = [loop.run_in_executor(exe, f) for f in gen_upload_providers()]
+                    for fut in asyncio.as_completed(futs):
+                        yield await fut
+
+    async def generate_specs(self):
+        # Flatten all the uploads yielded by the upload generators in the batch
+        def gen_upload_providers():
+            for gen in self._upload_generators:
+                yield from gen
+
+        async def gen_file_upload_specs(): # -> AsyncGenerator[FileUploadSpec, None]:
+            loop = asyncio.get_event_loop()
+            with concurrent.futures.ThreadPoolExecutor() as exe:
+                # TODO: avoid eagerly expanding
+                futs = [loop.run_in_executor(exe, f) for f in gen_upload_providers()]
+                for fut in asyncio.as_completed(futs):
+                    yield await fut
+
+        files_stream = aiostream.stream.iterate(gen_file_upload_specs())
+        self.files_stream = files_stream
+        return files_stream
+
+
+    def put_file(
+        self,
+        local_file: Union[Path, str, BinaryIO],
+        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: get_file_upload_spec_from_path(local_file, PurePosixPath(remote_path), mode)
+            else:
+                yield lambda: get_file_upload_spec_from_fileobj(local_file, PurePosixPath(remote_path), mode or 0o644)
+
+
+        self._upload_generators.append(gen())
+
+    def get_files_stream(self):
+        return self.files_stream
+
+    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_file_spec_provider(subpath):
+            relpath_str = subpath.relative_to(local_path)
+            return lambda: get_file_upload_spec_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_file_spec_provider(subpath)
+
+        self._upload_generators.append(gen())
+
+VolumeUploadContextManager = synchronize_api(_VolumeUploadContextManager)