diff --git a/js/button/actions.js b/js/button/actions.js index dbd400a..e01fe05 100644 --- a/js/button/actions.js +++ b/js/button/actions.js @@ -23,6 +23,60 @@ import { authDialog } from '../auth/index.js'; import { nimbus, local } from '../resource/index.js'; import { endpoint } from '../constants.js'; +import { ComfyCloudDialog } from '../comfy/ui.js'; +import van from '../lib/van.js'; +import { Await } from '../lib/van-ui.js'; + +const {video, h2, img, a, button, div, b, span, source } = van.tags + +const ProgressBar = (progress) => { + const progressPercentage = 100 - (progress.value/progress.max * 100) + return () => div({style: `width: ${progressPercentage}%; height: 24px; background-color: #1D4AFF; transition: all .2s;`}, `${progressPercentage}%`) +} + +const taskId = van.state(null) +const Progress = (dialogInstance) => { + + const activeTab = van.state(0) + const data = van.state(local.pollUploadStatus(taskId.val)) // workflowRun data + + const start = () => dialogInstance.poll = dialogInstance.poll || setInterval(async () => { + const res = await local.pollUploadStatus(taskId.val) + + if(res.status !== "Task started") { + // Stop poll + clearInterval(dialogInstance.poll) + dialogInstance.close() + dialogInstance, dialogInstance.poll = 0; + } + + data.val = Promise.resolve(res) + }, 1000) + + start() + + return () => van.tags.div({style: 'width: 520px'}, + div( + h2("Upload progress"), + ), + Await({ + value: data.val, + container: span, + Loading: () => "Loading", + Error: () => "Request failed.", + }, data => div( + Object.keys(data.progress).length === 0 ? "Loading" : + Object.entries(data.progress).map(([key, val]) => { + return () => div( + b(key), + ProgressBar({value: val.value, max: val.max}) + ) + })) + )) +} + +export const progressDialog = new ComfyCloudDialog(Progress) + export async function onGeneration() { try { @@ -79,8 +133,14 @@ export async function onGeneration() { endpoint, ...dependencies, }) + + const uploadTaskId = res.task_id if(uploadTaskId) { + taskId.val = uploadTaskId + // Open UI window + progressDialog.show(); + await pollSyncDependencies(uploadTaskId) } diff --git a/python/custom_routes.py b/python/custom_routes.py index ac36f69..3636d59 100644 --- a/python/custom_routes.py +++ b/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) diff --git a/python/upload/__init__.py b/python/upload/__init__.py new file mode 100644 index 0000000..080834b --- /dev/null +++ b/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) + diff --git a/python/upload/blob.py b/python/upload/blob.py new file mode 100644 index 0000000..8b286a3 --- /dev/null +++ b/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 = "" + 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 = "\n" + for part_number, etag in enumerate(part_etags, 1): + completion_body += f"""\n{part_number}\n"{etag}"\n\n""" + completion_body += "" + + # 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 = "" + 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}") diff --git a/python/upload/hash.py b/python/upload/hash.py new file mode 100644 index 0000000..888ccff --- /dev/null +++ b/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"), + ) diff --git a/python/upload/net.py b/python/upload/net.py new file mode 100644 index 0000000..c3ac06c --- /dev/null +++ b/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) + diff --git a/python/upload/sync.py b/python/upload/sync.py new file mode 100644 index 0000000..bfa27c4 --- /dev/null +++ b/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 + diff --git a/python/upload/volume.py b/python/upload/volume.py new file mode 100644 index 0000000..b22bc6b --- /dev/null +++ b/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)