nodes/nathannlu-ComfyUI-Cloud
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Cleanup upload route

Browse Source
This commit is contained in:
Nathan
2024-03-29 15:13:54 -07:00
parent 576f08511b
commit bd677944fa
7 changed files with 442 additions and 521 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/custom_routes.py
+59 -117
View File
@@ -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')
python/upload/__init__.py
+46 -168
View File
@@ -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")
python/upload/blob.py
+55 -91
View File
@@ -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}")
python/upload/net.py
+8
View File
@@ -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
python/upload/spec.py
+271
View File
@@ -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)
python/upload/sync.py
+3 -1
View File
@@ -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
python/upload/volume.py
-144
View File
@@ -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)