diff --git a/python/custom_routes.py b/python/custom_routes.py index 3636d59..84e157a 100644 --- a/python/custom_routes.py +++ b/python/custom_routes.py @@ -19,62 +19,12 @@ import server import execution import folder_paths -from .upload import upload_file +from .upload import upload_file_specs +from .upload.spec import FileSpecContextManager from .upload.blob import progress - task_status = {} -def random_seed(num_digits=15): - range_start = 10 ** (num_digits - 1) - range_end = (10**num_digits) - 1 - return random.randint(range_start, range_end) - -@server.PromptServer.instance.routes.post("/comfy-cloud/save-log") -async def comfy_cloud_save_log(request): - try: - data = await request.json() - log = data.get("log") - log_id = log["logId"] #str - workflow_id = log.get("workflowId") #str - user_id = log["userId"] #str - - timestamp = datetime.datetime.now() - formatted_timestamp = timestamp.strftime("%Y_%B_%d_%H-%M-%S") - - # Create the full path for the log file - filename = f"log_{formatted_timestamp}.txt" - current_path = os.getcwd() - directory = os.path.join(current_path, "logs") - log_path = os.path.join(directory, filename) - - # Check if the directory exists, if not, create it - if not os.path.exists(directory): - os.makedirs(directory) - - # Create a timestamp for the log entry - timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") - - # Content to be written to the log file - log_content = f"Log entry created at: {timestamp}\n" - log_content += f"log_id: {log_id}\n" - log_content += f"workflow_id: {workflow_id}\n" - log_content += f"user_id: {user_id}\n" - log_content += f"---\n" - - for log_entry in log["logs"]: - log_content += f"{log_entry}\n" - - # Write the content to the log file - with open(log_path, 'w') as log_file: - log_file.write(log_content) - - return web.json_response({ "success": True }, status=200) - - except Exception as e: - print("Error:", e) - return web.json_response({ "error": e }, status=400) - @server.PromptServer.instance.routes.post("/comfy-cloud/validate-input-path") async def validate_input_path(request): @@ -102,6 +52,11 @@ async def comfy_cloud_validate_prompt(request): workflow_api = data.get("workflow_api") + def random_seed(num_digits=15): + range_start = 10 ** (num_digits - 1) + range_end = (10**num_digits) - 1 + return random.randint(range_start, range_end) + for key in workflow_api: if 'inputs' in workflow_api[key] and 'seed' in workflow_api[key]['inputs']: workflow_api[key]['inputs']['seed'] = random_seed() @@ -192,25 +147,10 @@ def load_custom_node(module_path, ignore=set()): print(f"Cannot import {module_path} module for custom nodes:", e) return (module_name, []) -def make_post_request_with_retry(url, data, headers=None, max_retries=3, retry_delay=1): - for attempt in range(1, max_retries + 1): - try: - response = requests.post(url, data=data, headers=headers) - - # Check if the response status code is OK (2xx) - response.raise_for_status() - - return response - except requests.exceptions.RequestException as e: - print(f"Attempt {attempt} failed: {e}") - if attempt < max_retries: - print(f"Retrying in {retry_delay} seconds...") - time.sleep(retry_delay) - else: - print(f"Maximum retries reached. Failed to make POST request.") - raise - def update_dependencies(): + """ + Formats requirements.txt + """ base = folder_paths.base_path custom_nodes_dir = os.path.join(base, "custom_nodes") @@ -246,14 +186,9 @@ def update_dependencies(): with open(filepath, 'w') as file: file.writelines(updated_lines) -async def upload_task_execution(task_id, json_response, workflow_id, models_dep, nodes_dep, files): +async def upload_task_execution(task_id, file_specs, workflow_id): try: - await upload_file( - workflow_id, - models_dep, - nodes_dep, - files, - ) + await upload_file_specs(file_specs, workflow_id) # cleanup temp task_status[task_id] = {"status": "Task completed", "message": "Upload successful"} @@ -265,51 +200,40 @@ async def upload_task_execution(task_id, json_response, workflow_id, models_dep, async def upload_dependencies(request): # Make a request to localhost try: - authorization = request.headers.get('Authorization') json_data = await request.json() - endpoint = json_data["endpoint"] + #endpoint = json_data["endpoint"] workflow_id = json_data["workflow_id"] - headers = None - if authorization: - headers = { - "Authorization": authorization - } - + # dependencies models_dep = json_data["modelsToUpload"] nodes_dep = json_data["nodesToUpload"] files = json_data["filesToUpload"] - current_datetime = datetime.datetime.now() - body = { - "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) + # Paths + base = folder_paths.base_path + models_dir = os.path.join(base, "models") + custom_nodes_dir = os.path.join(base, "custom_nodes") + # Paths - for searching through /model subfolders + paths = copy.deepcopy(folder_paths.folder_names_and_paths) + paths.pop("custom_nodes", None) + paths.pop("configs", None) - # Check if the request was successful (status code 200) - if response.status_code == 200: - # Parse the JSON content - json_response = response.json() - task_id = str(uuid.uuid4()) - task_status[task_id] = {"status": "Task started", "message": None} + # Create upload task + task_id = str(uuid.uuid4()) + task_status[task_id] = {"status": "Task started", "message": None} - # Process custom nodes - update_dependencies() + # Our server uses a custom dependency manager + # that requires a specific format for requirements.txt. + # Loop through all dependent custom nodes and patch. + update_dependencies() - custom_nodes_bytes = {} - - 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") + # Upload code + file_specs = [] + # Add files + with FileSpecContextManager(file_specs) as batch: # Upload models for name in models_dep: found = False @@ -318,23 +242,41 @@ async def upload_dependencies(request): model_path = os.path.join(models_dir, base_dir, name) if os.path.exists(model_path): found = True + batch.put_file(model_path, f"/vol/{workflow_id}/comfyui/models/{base_dir}/{name}") + # Handle model not found + # We have to do this because the code for finding + # models need to go through each sub folder inside + # /models. + # The models folder is comprised of subfolders such + # as checkpoints, controlnets, loras, etc if found is False: - return web.json_response({'success': False, 'message': f"Required model {name} was not found in your ComfyUI/models folder. Make sure you do not have extra_paths.yaml enabled"}, content_type='application/json') + raise Exception(f"Required model {name} was not found in your ComfyUI/models folder. Make sure you do not have extra_paths.yaml enabled") + # Upload custom nodes for name in nodes_dep: # check filepath exists - p = os.path.join(custom_nodes_dir, name) + p = os.path.join(custom_nodes_dir, f"/vol/{workflow_id}/comfyui/custom_nodes/{name}") if not os.path.exists(p): - return web.json_response({'success': False, 'message': f"Required node {name} was not found in your ComfyUI/custom_nodes folder. Make sure you do not have extra_paths.yaml enabled"}, content_type='application/json') - custom_nodes_bytes[name] = "" + raise Exception(f"Required node {name} was not found in your ComfyUI/custom_nodes folder. Make sure you do not have extra_paths.yaml enabled") - asyncio.ensure_future(upload_task_execution(task_id, json_response, workflow_id, models_dep, custom_nodes_bytes, files)) + batch.put_directory(p, name) - return web.json_response({'success': True, 'task_id': task_id}, content_type='application/json') - else: - return web.json_response({'success': False, 'message': "Failed to retrieve exec code"}, status=500, content_type='application/json') + # Upload input 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.") + # Create file specs for upload + file_specs = batch.generate_specs() + + # Queue upload task in background + asyncio.ensure_future(upload_task_execution(task_id, file_specs, workflow_id)) + + return web.json_response({'success': True, 'task_id': task_id}, content_type='application/json') except Exception as e: print("Error", e) return web.json_response({'success': False, 'message': str(e)}, status=500, content_type='application/json') diff --git a/python/upload/__init__.py b/python/upload/__init__.py index 080834b..af91e7b 100644 --- a/python/upload/__init__.py +++ b/python/upload/__init__.py @@ -1,184 +1,62 @@ import io import os import asyncio +import aiostream +from typing import List -from .volume import VolumeUploadContextManager -from .blob import perform_multipart_upload, BytesIOSegmentPayload, _upload_to_s3_url +from .spec import FileSpecContextManager, serialize_spec, FileUploadSpec +from .blob import blob_upload from .hash import get_upload_hashes +from .net import make_post_request -import folder_paths -import copy +base_url = "https://pytest-12acbe397fbc.herokuapp.com" -import requests -import json -import aiohttp +async def upload_file_specs(file_specs: List[FileUploadSpec], workflow_id: str): + """ + Take in a list of specs, and uploads them + """ -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 + # Send dict version of file specs to client + serialize_spec_stream = aiostream.stream.map(file_specs, serialize_spec, task_limit=20) + serialized_specs = await aiostream.stream.list(serialize_spec_stream) + serialized_specs_dict = {} -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.") + # Adds id to list, as well as creates a dict version + for idx, spec in enumerate(serialized_specs): + spec["id"] = idx + serialized_specs_dict[idx] = spec + + # Get a list of blob_ids from the client + url = f'{base_url}/upload-urls' + response_data = await make_post_request(url, { "specs": serialized_specs, "workflow_id": workflow_id }) + upload_data = response_data["data"] - # 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.") + # Create the upload data + # Mix the response data from the server and the file spec + # and pass it to the generator + for spec in upload_data: + if spec["id"] in serialized_specs_dict: + spec["data"] = { + **serialized_specs_dict[spec["id"]], + **spec["data"] + } + def gen_upload_providers(): + for gen in upload_data: + yield gen - #batch.put_directory(path, "/helloworld") - #batch.put_file(path, target_path) - x = batch.generate_specs() + async def _upload_and_commit(spec): + await blob_upload(spec) + # Blob upload has finished. Put the spec + # on the queue to be committed + url = f'{base_url}/commit' + print("Sending commit") + response_data = await make_post_request(url, { "spec": spec }) - 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) + files_stream = aiostream.stream.iterate(gen_upload_providers()) + uploads_stream = aiostream.stream.map(files_stream, _upload_and_commit, task_limit=20) + files = await aiostream.stream.list(uploads_stream) + print("done") diff --git a/python/upload/blob.py b/python/upload/blob.py index 8b286a3..2aecf74 100644 --- a/python/upload/blob.py +++ b/python/upload/blob.py @@ -1,12 +1,10 @@ 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 typing import Any, AsyncIterator, BinaryIO, Callable, Dict, List, Optional, Union from urllib.parse import urlparse from aiohttp import BytesIOPayload @@ -14,12 +12,11 @@ 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 +import uuid progress = {} @@ -119,6 +116,7 @@ class BytesIOSegmentPayload(BytesIOPayload): await writer.write(chunk) def remaining_bytes(self): + #print("Remaining bytes", self.filename, self.segment_length - self.num_bytes_read) if self.filename not in progress: progress[self.filename] = {} @@ -145,6 +143,8 @@ async def _upload_to_s3_url( if content_type: headers["Content-Type"] = content_type + + #print("Uploading to s3", upload_url, payload) async with session.put( upload_url, data=payload, @@ -181,6 +181,8 @@ async def _upload_to_s3_url( return remote_md5 + #print("Done upload to s3") + async def perform_multipart_upload( data_file: Union[BinaryIO, io.BytesIO, io.FileIO], @@ -248,101 +250,61 @@ async def perform_multipart_upload( f"Hash mismatch on multipart upload assembly: {expected_multipart_etag} not in {response_body}" ) +async def blob_upload(spec): + r""" + This spec is not the same as FileUploadSpec. + It is an altered version provided from the server + that includes upload data on top of the existing + FileUploadSpec type + --- + type: str, + id: str, + blob_id: str, + data: { + # byte upload only + upload_url: str, -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 + # multipart upload only + "max_part_size": str + "part_urls": dict + "completion_url": str + # Fields from FileUploadSpec + **FileUploadSpec + } + """ + is_multipart = spec["type"] == "multipart" + resp = spec["data"] -@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) + # Find file + filename = resp["source_description"] + data = open(filename, "rb") - if resp.status != 200: - text = await resp.text() - raise ExecutionError(f"Get from url failed with status {resp.status}: {text}") - return await resp.read() + if is_multipart: + 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: + content_length = resp["content_length"] + upload_hashes = resp["upload_hashes"] -@dataclasses.dataclass -class FileUploadSpec: - source: Callable[[], BinaryIO] - source_description: Any - mount_filename: str + 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"], + ) - 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 + return spec -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. @@ -357,3 +319,5 @@ def use_md5(url: str) -> bool: return False else: raise Exception(f"Unknown S3 host: {host}") + + diff --git a/python/upload/net.py b/python/upload/net.py index c3ac06c..85639a6 100644 --- a/python/upload/net.py +++ b/python/upload/net.py @@ -6,6 +6,7 @@ import ssl import time import uuid +import aiohttp from aiohttp import ClientSession, ClientTimeout, TCPConnector from aiohttp.web import Application from aiohttp.web_runner import AppRunner, SockSite @@ -64,6 +65,10 @@ async def run_temporary_http_server(app: Application): await runner.cleanup() +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() @@ -96,10 +101,12 @@ async def retry_transient_errors( 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: @@ -110,6 +117,7 @@ async def retry_transient_errors( 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 diff --git a/python/upload/spec.py b/python/upload/spec.py new file mode 100644 index 0000000..bdd9c58 --- /dev/null +++ b/python/upload/spec.py @@ -0,0 +1,271 @@ +import io +import copy +import os +import base64 +import platform +from aiohttp import web +import asyncio +import dataclasses + +import concurrent.futures +import time +from pathlib import Path, PurePosixPath +from typing import ( + IO, + AsyncGenerator, + AsyncIterator, + BinaryIO, + Dict, + Callable, + Generator, + Optional, + Sequence, + Type, + Union, + Any, +) +import aiostream +from .sync import synchronize_api +from .blob import LARGE_FILE_LIMIT +from .hash import get_sha256_hex, get_upload_hashes + +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 + + +@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 + +# Example usage: + +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 serialize_spec(spec: FileUploadSpec) -> Dict[Any, Any]: + """ + Turns the spec into a http-transferrable + dictionary + """ + obj = spec.__dict__ + + obj["source_description"] = str(spec.source_description) + + with spec.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"] + + return obj + +class _FileSpecContextManager: + """Context manager for batch-uploading files to a Volume.""" + + def __init__(self, file_specs): + """mdmd:hidden""" + self._upload_generators = [] + + 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: + return 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()] + #print(f"Computing checksums for {len(futs)} files using {exe._max_workers} workers") + 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()] + #print(f"Computing checksums for {len(futs)} files using {exe._max_workers} workers") + for fut in asyncio.as_completed(futs): + yield await fut + + files_stream = aiostream.stream.iterate(gen_file_upload_specs()) + #print(files_stream) + 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 put_model( + self, + local_file: Union[Path, str, BinaryIO], + model_type: str, + 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()) + +FileSpecContextManager = synchronize_api(_FileSpecContextManager) diff --git a/python/upload/sync.py b/python/upload/sync.py index bfa27c4..ca48afc 100644 --- a/python/upload/sync.py +++ b/python/upload/sync.py @@ -57,16 +57,18 @@ def retry(direct_fn=None, *, n_attempts=3, base_delay=0, delay_factor=2, timeout try: return await asyncio.wait_for(fn(*args, **kwargs), timeout=timeout) except asyncio.CancelledError: - logger.debug(f"Function {fn} was cancelled") + #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 diff --git a/python/upload/volume.py b/python/upload/volume.py deleted file mode 100644 index b22bc6b..0000000 --- a/python/upload/volume.py +++ /dev/null @@ -1,144 +0,0 @@ -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)