360 lines
12 KiB
Python
360 lines
12 KiB
Python
import asyncio
|
|
import dataclasses
|
|
import hashlib
|
|
import io
|
|
import os
|
|
import platform
|
|
from contextlib import contextmanager
|
|
from pathlib import Path, PurePosixPath
|
|
from typing import Any, AsyncIterator, BinaryIO, Callable, List, Optional, Union
|
|
from urllib.parse import urlparse
|
|
|
|
from aiohttp import BytesIOPayload
|
|
from aiohttp.abc import AbstractStreamWriter
|
|
|
|
from modal.exception import ExecutionError
|
|
from .sync import retry
|
|
|
|
from modal._utils.hash_utils import get_sha256_hex
|
|
from .net import http_client_with_tls, retry_transient_errors
|
|
|
|
from tqdm import tqdm
|
|
import logging
|
|
|
|
progress = {}
|
|
|
|
logger = logging.getLogger()
|
|
|
|
|
|
# Max size for function inputs and outputs.
|
|
MAX_OBJECT_SIZE_BYTES = 2 * 1024 * 1024 # 2 MiB
|
|
|
|
# If a file is LARGE_FILE_LIMIT bytes or larger, it's uploaded to blob store (s3) instead of going through grpc
|
|
# It will also make sure to chunk the hash calculation to avoid reading the entire file into memory
|
|
LARGE_FILE_LIMIT = 4 * 1024 * 1024 # 4 MiB
|
|
|
|
# Max parallelism during map calls
|
|
BLOB_MAX_PARALLELISM = 10
|
|
|
|
# read ~16MiB chunks by default
|
|
DEFAULT_SEGMENT_CHUNK_SIZE = 2**24
|
|
|
|
|
|
class BytesIOSegmentPayload(BytesIOPayload):
|
|
"""Modified bytes payload for concurrent sends of chunks from the same file.
|
|
|
|
Adds:
|
|
* read limit using remaining_bytes, in order to split files across streams
|
|
* larger read chunk (to prevent excessive read contention between parts)
|
|
* calculates an md5 for the segment
|
|
|
|
Feels like this should be in some standard lib...
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
bytes_io: BinaryIO, # should *not* be shared as IO position modification is not locked
|
|
segment_start: int,
|
|
segment_length: int,
|
|
filename: str,
|
|
chunk_size: int = DEFAULT_SEGMENT_CHUNK_SIZE,
|
|
):
|
|
# not thread safe constructor!
|
|
super().__init__(bytes_io)
|
|
self.initial_seek_pos = bytes_io.tell()
|
|
self.segment_start = segment_start
|
|
self.segment_length = segment_length
|
|
# seek to start of file segment we are interested in, in order to make .size() evaluate correctly
|
|
self._value.seek(self.initial_seek_pos + segment_start)
|
|
assert self.segment_length <= super().size
|
|
self.chunk_size = chunk_size
|
|
self.reset_state()
|
|
self.filename = filename
|
|
|
|
def reset_state(self):
|
|
self._md5_checksum = hashlib.md5()
|
|
self.num_bytes_read = 0
|
|
self._value.seek(self.initial_seek_pos)
|
|
|
|
@contextmanager
|
|
def reset_on_error(self):
|
|
try:
|
|
yield
|
|
finally:
|
|
self.reset_state()
|
|
|
|
@property
|
|
def filename(self):
|
|
return self._filename
|
|
|
|
@filename.setter
|
|
def filename(self, value):
|
|
self._filename = value
|
|
|
|
@property
|
|
def size(self) -> int:
|
|
return self.segment_length
|
|
|
|
def md5_checksum(self):
|
|
return self._md5_checksum
|
|
|
|
async def write(self, writer: AbstractStreamWriter):
|
|
loop = asyncio.get_event_loop()
|
|
|
|
async def safe_read():
|
|
read_start = self.initial_seek_pos + self.segment_start + self.num_bytes_read
|
|
self._value.seek(read_start)
|
|
num_bytes = min(self.chunk_size, self.remaining_bytes())
|
|
chunk = await loop.run_in_executor(None, self._value.read, num_bytes)
|
|
|
|
await loop.run_in_executor(None, self._md5_checksum.update, chunk)
|
|
self.num_bytes_read += len(chunk)
|
|
return chunk
|
|
|
|
chunk = await safe_read()
|
|
while chunk and self.remaining_bytes() > 0:
|
|
await writer.write(chunk)
|
|
chunk = await safe_read()
|
|
if chunk:
|
|
await writer.write(chunk)
|
|
|
|
def remaining_bytes(self):
|
|
if self.filename not in progress:
|
|
progress[self.filename] = {}
|
|
|
|
progress[self.filename]["value"] = self.segment_length - self.num_bytes_read
|
|
if "max" not in progress[self.filename]:
|
|
progress[self.filename]["max"] = self.segment_length - self.num_bytes_read
|
|
|
|
return self.segment_length - self.num_bytes_read
|
|
|
|
|
|
@retry(n_attempts=5, base_delay=0.5, timeout=None)
|
|
async def _upload_to_s3_url(
|
|
upload_url,
|
|
payload: BytesIOSegmentPayload,
|
|
content_md5_b64: Optional[str] = None,
|
|
content_type: Optional[str] = "application/octet-stream", # set to None to force omission of ContentType header
|
|
) -> str:
|
|
"""Returns etag of s3 object which is a md5 hex checksum of the uploaded content"""
|
|
with payload.reset_on_error(): # ensure retries read the same data
|
|
async with http_client_with_tls(timeout=None) as session:
|
|
headers = {}
|
|
if content_md5_b64 and use_md5(upload_url):
|
|
headers["Content-MD5"] = content_md5_b64
|
|
if content_type:
|
|
headers["Content-Type"] = content_type
|
|
|
|
async with session.put(
|
|
upload_url,
|
|
data=payload,
|
|
headers=headers,
|
|
skip_auto_headers=["content-type"] if content_type is None else [],
|
|
) as resp:
|
|
|
|
# S3 signal to slow down request rate.
|
|
if resp.status == 503:
|
|
logger.warning("Received SlowDown signal from S3, sleeping for 1 second before retrying.")
|
|
await asyncio.sleep(1)
|
|
|
|
if resp.status != 200:
|
|
try:
|
|
text = await resp.text()
|
|
except Exception:
|
|
text = "<no body>"
|
|
raise ExecutionError(f"Put to url {upload_url} failed with status {resp.status}: {text}")
|
|
|
|
# client side ETag checksum verification
|
|
# the s3 ETag of a single part upload is a quoted md5 hex of the uploaded content
|
|
etag = resp.headers["ETag"].strip()
|
|
if etag.startswith(("W/", "w/")): # see https://www.rfc-editor.org/rfc/rfc7232#section-2.3
|
|
etag = etag[2:]
|
|
if etag[0] == '"' and etag[-1] == '"':
|
|
etag = etag[1:-1]
|
|
remote_md5 = etag
|
|
|
|
local_md5_hex = payload.md5_checksum().hexdigest()
|
|
if local_md5_hex != remote_md5:
|
|
raise ExecutionError(
|
|
f"Local data and remote data checksum mismatch ({local_md5_hex} vs {remote_md5})"
|
|
)
|
|
|
|
return remote_md5
|
|
|
|
|
|
async def perform_multipart_upload(
|
|
data_file: Union[BinaryIO, io.BytesIO, io.FileIO],
|
|
*,
|
|
content_length: int,
|
|
max_part_size: int,
|
|
part_urls: List[str],
|
|
completion_url: str,
|
|
filename: str,
|
|
upload_chunk_size: int = DEFAULT_SEGMENT_CHUNK_SIZE,
|
|
):
|
|
upload_coros = []
|
|
file_offset = 0
|
|
num_bytes_left = content_length
|
|
|
|
# Give each part its own IO reader object to avoid needing to
|
|
# lock access to the reader's position pointer.
|
|
data_file_readers: List[BinaryIO]
|
|
if isinstance(data_file, io.BytesIO):
|
|
view = data_file.getbuffer() # does not copy data
|
|
data_file_readers = [io.BytesIO(view) for _ in range(len(part_urls))]
|
|
else:
|
|
filename = data_file.name
|
|
data_file_readers = [open(filename, "rb") for _ in range(len(part_urls))]
|
|
|
|
for part_number, (data_file_rdr, part_url) in enumerate(zip(data_file_readers, part_urls), start=1):
|
|
part_length_bytes = min(num_bytes_left, max_part_size)
|
|
part_payload = BytesIOSegmentPayload(
|
|
data_file_rdr,
|
|
segment_start=file_offset,
|
|
segment_length=part_length_bytes,
|
|
chunk_size=upload_chunk_size,
|
|
filename=filename,
|
|
)
|
|
upload_coros.append(_upload_to_s3_url(part_url, payload=part_payload, content_type=None))
|
|
num_bytes_left -= part_length_bytes
|
|
file_offset += part_length_bytes
|
|
|
|
part_etags = await asyncio.gather(*upload_coros)
|
|
|
|
# The body of the complete_multipart_upload command needs some data in xml format:
|
|
completion_body = "<CompleteMultipartUpload>\n"
|
|
for part_number, etag in enumerate(part_etags, 1):
|
|
completion_body += f"""<Part>\n<PartNumber>{part_number}</PartNumber>\n<ETag>"{etag}"</ETag>\n</Part>\n"""
|
|
completion_body += "</CompleteMultipartUpload>"
|
|
|
|
# etag of combined object should be md5 hex of concatendated md5 *bytes* from parts + `-{num_parts}`
|
|
bin_hash_parts = [bytes.fromhex(etag) for etag in part_etags]
|
|
|
|
expected_multipart_etag = hashlib.md5(b"".join(bin_hash_parts)).hexdigest() + f"-{len(part_etags)}"
|
|
async with http_client_with_tls(timeout=None) as session:
|
|
resp = await session.post(
|
|
completion_url, data=completion_body.encode("ascii"), skip_auto_headers=["content-type"]
|
|
)
|
|
if resp.status != 200:
|
|
try:
|
|
msg = await resp.text()
|
|
except Exception:
|
|
msg = "<no body>"
|
|
raise ExecutionError(f"Error when completing multipart upload: {resp.status}\n{msg}")
|
|
else:
|
|
response_body = await resp.text()
|
|
if expected_multipart_etag not in response_body:
|
|
raise ExecutionError(
|
|
f"Hash mismatch on multipart upload assembly: {expected_multipart_etag} not in {response_body}"
|
|
)
|
|
|
|
|
|
def get_content_length(data: BinaryIO):
|
|
# *Remaining* length of file from current seek position
|
|
pos = data.tell()
|
|
data.seek(0, os.SEEK_END)
|
|
content_length = data.tell()
|
|
data.seek(pos)
|
|
return content_length - pos
|
|
|
|
|
|
@retry(n_attempts=5, base_delay=0.1, timeout=None)
|
|
async def _download_from_url(download_url) -> bytes:
|
|
async with http_client_with_tls(timeout=None) as session:
|
|
async with session.get(download_url) as resp:
|
|
# S3 signal to slow down request rate.
|
|
if resp.status == 503:
|
|
logger.warning("Received SlowDown signal from S3, sleeping for 1 second before retrying.")
|
|
await asyncio.sleep(1)
|
|
|
|
if resp.status != 200:
|
|
text = await resp.text()
|
|
raise ExecutionError(f"Get from url failed with status {resp.status}: {text}")
|
|
return await resp.read()
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class FileUploadSpec:
|
|
source: Callable[[], BinaryIO]
|
|
source_description: Any
|
|
mount_filename: str
|
|
|
|
use_blob: bool
|
|
content: Optional[bytes] # typically None if using blob, required otherwise
|
|
sha256_hex: str
|
|
mode: int # file permission bits (last 12 bits of st_mode)
|
|
size: int
|
|
|
|
|
|
def _get_file_upload_spec(
|
|
source: Callable[[], BinaryIO], source_description: Any, mount_filename: PurePosixPath, mode: int
|
|
) -> FileUploadSpec:
|
|
with source() as fp:
|
|
# Current position is ignored - we always upload from position 0
|
|
fp.seek(0, os.SEEK_END)
|
|
size = fp.tell()
|
|
fp.seek(0)
|
|
|
|
if size >= LARGE_FILE_LIMIT:
|
|
use_blob = True
|
|
content = None
|
|
sha256_hex = get_sha256_hex(fp)
|
|
else:
|
|
use_blob = False
|
|
content = fp.read()
|
|
sha256_hex = get_sha256_hex(content)
|
|
|
|
return FileUploadSpec(
|
|
source=source,
|
|
source_description=source_description,
|
|
mount_filename=mount_filename.as_posix(),
|
|
use_blob=use_blob,
|
|
content=content,
|
|
sha256_hex=sha256_hex,
|
|
mode=mode & 0o7777,
|
|
size=size,
|
|
)
|
|
|
|
|
|
def get_file_upload_spec_from_path(
|
|
filename: Path, mount_filename: PurePosixPath, mode: Optional[int] = None
|
|
) -> FileUploadSpec:
|
|
# Python appears to give files 0o666 bits on Windows (equal for user, group, and global),
|
|
# so we mask those out to 0o755 for compatibility with POSIX-based permissions.
|
|
mode = mode or os.stat(filename).st_mode & (0o7777 if platform.system() != "Windows" else 0o7755)
|
|
return _get_file_upload_spec(
|
|
lambda: open(filename, "rb"),
|
|
filename,
|
|
mount_filename,
|
|
mode,
|
|
)
|
|
|
|
|
|
def get_file_upload_spec_from_fileobj(fp: BinaryIO, mount_filename: PurePosixPath, mode: int) -> FileUploadSpec:
|
|
def source():
|
|
# We ignore position in stream and always upload from position 0
|
|
fp.seek(0)
|
|
return fp
|
|
|
|
return _get_file_upload_spec(
|
|
source,
|
|
str(fp),
|
|
mount_filename,
|
|
mode,
|
|
)
|
|
|
|
|
|
def use_md5(url: str) -> bool:
|
|
"""This takes an upload URL in S3 and returns whether we should attach a checksum.
|
|
|
|
It's only a workaround for missing functionality in moto.
|
|
https://github.com/spulec/moto/issues/816
|
|
"""
|
|
host = urlparse(url).netloc.split(":")[0]
|
|
if host.endswith(".amazonaws.com"):
|
|
return True
|
|
elif host in ["127.0.0.1", "localhost", "172.21.0.1"]:
|
|
return False
|
|
else:
|
|
raise Exception(f"Unknown S3 host: {host}")
|