diff --git a/cirro/cli/cli.py b/cirro/cli/cli.py index 52594f6..b4e9dff 100644 --- a/cirro/cli/cli.py +++ b/cirro/cli/cli.py @@ -6,6 +6,7 @@ from cirro.cli import run_create_pipeline_config, run_validate_folder from cirro.cli import run_ingest, run_download, run_configure, run_list_datasets +from cirro.config import Constants from cirro.cli.controller import handle_error, run_upload_reference, run_list_projects, run_list_files, \ run_resume_upload from cirro.cli.debug import run_debug @@ -73,6 +74,9 @@ def list_datasets(**kwargs): @click.option('--file-limit', help='Maximum number of files to enumerate from the dataset', default=100000, show_default=True) +@click.option('--threads', + help='Number of files to transfer at once (1 disables threading)', + default=Constants.default_transfer_threads, show_default=True, type=int) @click.option('-i', '--interactive', help='Gather arguments interactively', is_flag=True, default=False) @@ -100,6 +104,9 @@ def download(**kwargs): @click.option('-i', '--interactive', help='Gather arguments interactively', is_flag=True, default=False) +@click.option('--threads', + help='Number of files to transfer at once (1 disables threading)', + default=Constants.default_transfer_threads, show_default=True, type=int) @click.option('--include-hidden', help='Include hidden files in the upload (e.g., files starting with .)', is_flag=True, default=False) @@ -122,6 +129,9 @@ def upload(**kwargs): @click.option('-i', '--interactive', help='Gather arguments interactively', is_flag=True, default=False) +@click.option('--threads', + help='Number of files to transfer at once (1 disables threading)', + default=Constants.default_transfer_threads, show_default=True, type=int) @click.option('--include-hidden', help='Include hidden files in the upload (e.g., files starting with .)', is_flag=True, default=False) diff --git a/cirro/cli/controller.py b/cirro/cli/controller.py index 3239f78..10d6f31 100644 --- a/cirro/cli/controller.py +++ b/cirro/cli/controller.py @@ -104,7 +104,8 @@ def run_ingest(input_params: UploadArguments, interactive=False): cirro.datasets.upload_files(project_id=project_id, dataset_id=create_resp.id, directory=directory, - files=files) + files=files, + threads=input_params['threads']) logger.info(f"File content validated by {cirro.configuration.checksum_method_display}") @@ -157,7 +158,8 @@ def run_resume_upload(input_params: ResumeUploadArguments, interactive=False): dataset_id=dataset_id, directory=directory, files=files, - resume=True) + resume=True, + threads=input_params['threads']) logger.info(f"File content validated by {cirro.configuration.checksum_method_display}") @@ -262,7 +264,8 @@ def run_download(input_params: DownloadArguments, interactive=False): dataset_id=dataset_id, download_location=input_params['data_directory'], files=files_to_download, - file_limit=input_params['file_limit']) + file_limit=input_params['file_limit'], + threads=input_params['threads']) def run_list_projects(): diff --git a/cirro/cli/models.py b/cirro/cli/models.py index 80eb614..19b1260 100644 --- a/cirro/cli/models.py +++ b/cirro/cli/models.py @@ -8,6 +8,7 @@ class DownloadArguments(TypedDict): interactive: bool file: Optional[list[str]] file_limit: int + threads: int class UploadArguments(TypedDict): @@ -19,6 +20,7 @@ class UploadArguments(TypedDict): include_hidden: bool interactive: bool file: Optional[list[str]] + threads: int class ResumeUploadArguments(TypedDict): @@ -28,6 +30,7 @@ class ResumeUploadArguments(TypedDict): include_hidden: bool interactive: bool file: Optional[list[str]] + threads: int class ValidateArguments(TypedDict): diff --git a/cirro/clients/s3.py b/cirro/clients/s3.py index 420673a..b7930a9 100644 --- a/cirro/clients/s3.py +++ b/cirro/clients/s3.py @@ -1,16 +1,41 @@ +import os import threading from pathlib import Path -from typing import Callable +from typing import Callable, Optional from boto3 import Session +from boto3.exceptions import S3UploadFailedError +from boto3.s3.transfer import ProgressCallbackInvoker, S3Transfer, TransferConfig, create_transfer_manager from botocore.config import Config +from botocore.exceptions import ClientError from botocore.credentials import RefreshableCredentials from botocore.session import get_session from cirro_api_client.v1.models import AWSCredentials from tqdm import tqdm +from cirro.config import Constants from cirro.models.s3_path import S3Path -from cirro.utils import convert_size + +# boto3 defaults to 10 concurrent requests per transfer; the pool needs headroom +# above that for the transfer manager's submission threads and credential refresh. +# Undersizing it makes urllib3 discard connections and serializes the transfers. +_MAX_POOL_CONNECTIONS = 20 + + +def local_filename(file_path) -> Optional[str]: + """ + Returns the local filesystem path named by `file_path`, or None if it does not + name a local file. + + boto3's managed transfer has to open the file itself, so it only works for real + filesystem paths. Path-like objects backed by something else, such as an s3fs + path, have to be streamed through their own open() instead. + """ + try: + filename = os.fspath(file_path) + except TypeError: + return None + return filename if os.path.isfile(filename) else None def format_creds_for_session(creds: AWSCredentials): @@ -33,42 +58,96 @@ def __call__(self, bytes_amount): class S3Client: - def __init__(self, creds_getter: Callable[[], AWSCredentials] = None, checksum_method: str = None): + def __init__(self, creds_getter: Callable[[], AWSCredentials] = None, checksum_method: str = None, + threads: int = Constants.default_transfer_threads): self._creds_getter = creds_getter + # A single thread means no threading anywhere, so boto3 runs the transfer + # inline rather than handing parts to its own worker pool + self._transfer_config = TransferConfig(use_threads=threads > 1) self._client = self._build_session_client() + self._manager = None + self._transfer_lock = threading.Lock() self._upload_args = dict(ChecksumAlgorithm=checksum_method) self._download_args = dict(ChecksumMode='ENABLED') if checksum_method else dict() def get_aws_client(self): return self._client - def upload_file(self, file_path: Path, bucket: str, key: str): - file_size = file_path.stat().st_size - file_name = file_path.name - - with tqdm(total=file_size, - desc=f'Uploading file {file_name} ({convert_size(file_size)})', - bar_format="{desc} | {percentage:.1f}%|{bar:25} | {rate_fmt}", - unit='B', unit_scale=True, - unit_divisor=1024) as progress: - with file_path.open('rb') as file: - self._client.upload_fileobj(file, bucket, key, - Callback=ProgressPercentage(progress), - ExtraArgs=self._upload_args) - - def download_file(self, local_path: Path, bucket: str, key: str): - file_size = self.get_file_stats(bucket, key)['ContentLength'] - file_name = local_path.name - - with tqdm(total=file_size, - desc=f'Downloading file {file_name} ({convert_size(file_size)})', - bar_format="{desc} | {percentage:.1f}%|{bar:25} | {rate_fmt}", - unit='B', unit_scale=True, - unit_divisor=1024) as progress: - absolute_path = str(local_path.absolute()) - self._client.download_file(bucket, key, absolute_path, - Callback=ProgressPercentage(progress), - ExtraArgs=self._download_args) + def upload_file(self, file_path: Path, bucket: str, key: str, + callback: Callable[[int], None] = None): + """ + Uploads a file to S3, reporting transferred bytes to `callback`. + + Local files are handed to the shared transfer manager by name, which lets + s3transfer read their parts in parallel. Any other Path-like object is + streamed through its own open(). + """ + local_file_path = local_filename(file_path) + + if local_file_path is not None: + self._get_transfer().upload_file( + filename=local_file_path, + bucket=bucket, + key=key, + callback=callback, + extra_args=self._upload_args + ) + return + + with file_path.open('rb') as file_obj: + self._upload_fileobj(file_obj, bucket, key, callback) + + def _upload_fileobj(self, file_obj, bucket: str, key: str, + callback: Callable[[int], None] = None): + """ + S3Transfer only accepts filenames, so a file object goes to the shared manager + directly, mirroring the error translation S3Transfer would have applied. + """ + subscribers = [ProgressCallbackInvoker(callback)] if callback else None + future = self._get_manager().upload(file_obj, bucket, key, self._upload_args, subscribers) + try: + future.result() + except ClientError as e: + raise S3UploadFailedError(f"Failed to upload {key} to {bucket}: {e}") + + def download_file(self, local_path: Path, bucket: str, key: str, + callback: Callable[[int], None] = None): + """ + Downloads a file from S3, reporting transferred bytes to `callback`. + """ + self._get_transfer().download_file( + bucket=bucket, + key=key, + filename=str(local_path.absolute()), + callback=callback, + extra_args=self._download_args + ) + + def _get_manager(self): + """ + A single transfer manager, and the thread pools it owns, is shared by every + transfer on this client rather than rebuilt for each file. It accepts both + filenames and file objects, so both upload paths share these pools. + """ + with self._transfer_lock: + if self._manager is None: + self._manager = create_transfer_manager(self._client, self._transfer_config) + return self._manager + + def _get_transfer(self) -> S3Transfer: + # Wrapping the shared manager costs nothing and adds boto3's error translation + return S3Transfer(manager=self._get_manager()) + + def close(self): + """ + Shuts down the transfer manager's thread pools. Their threads are not + daemons and are only partly reclaimed by garbage collection, so a + long-lived process needs this to avoid accumulating them. + """ + with self._transfer_lock: + if self._manager is not None: + self._manager.shutdown() + self._manager = None def create_object(self, bucket: str, key: str, contents: str, content_type: str): self._client.put_object( @@ -134,7 +213,8 @@ def _build_session_client(self): aws_session_token=creds.session_token ) s3_config = Config( - use_dualstack_endpoint=True + use_dualstack_endpoint=True, + max_pool_connections=_MAX_POOL_CONNECTIONS ) return session.client('s3', region_name=creds.region, config=s3_config) diff --git a/cirro/config.py b/cirro/config.py index 3059772..25583e5 100644 --- a/cirro/config.py +++ b/cirro/config.py @@ -13,6 +13,9 @@ class Constants: config_path = Path(home, 'config.ini').expanduser() default_base_url = 'cirro.bio' default_max_retries = 10 + # Files transferred at once. 1 disables threading entirely, including within + # a single file, which is required in environments without thread support. + default_transfer_threads = 8 class UserConfig(NamedTuple): diff --git a/cirro/file_utils.py b/cirro/file_utils.py index b2b55da..582f8d6 100644 --- a/cirro/file_utils.py +++ b/cirro/file_utils.py @@ -2,13 +2,18 @@ import os import random import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from functools import partial from pathlib import Path, PurePath -from typing import List, Union, Dict +from typing import Callable, List, Union, Dict -from boto3.exceptions import S3UploadFailedError +from boto3.exceptions import RetriesExceededError, S3UploadFailedError from botocore.exceptions import ConnectionError +from tqdm import tqdm from cirro.clients import S3Client +from cirro.clients.s3 import ProgressPercentage +from cirro.config import Constants from cirro.models.file import DirectoryStatistics, File, PathLike from cirro.models.s3_path import S3Path @@ -16,6 +21,16 @@ import win32api import win32con +# Transient failures worth another attempt. Anything else (a missing object, +# denied access) will not succeed on a retry and is surfaced immediately. +_RETRYABLE_ERRORS = (S3UploadFailedError, RetriesExceededError, ConnectionError) + +# Ceiling on the wait between attempts. botocore already retries throttling and +# transient 5xx internally, so this outer loop only needs to cover longer outages. +_MAX_RETRY_DELAY = 60 + +_PROGRESS_BAR_FORMAT = "{desc} | {percentage:.1f}%|{bar:25} | {rate_fmt}" + def filter_files_by_pattern(files: Union[List[File], List[str]], pattern: str) -> Union[List[File], List[str]]: """ @@ -125,6 +140,98 @@ def get_files_stats(files: List[PathLike]) -> DirectoryStatistics: ) +def _transfer_with_retry(transfer: Callable[..., None], + description: str, + callback: ProgressPercentage, + progress: tqdm, + attempts: int): + """ + @private + + Runs one transfer, retrying transient failures. Waiting between attempts only + occupies the current worker, leaving any workers transferring other files + unaffected. + """ + for attempt in range(attempts): + try: + return transfer(callback=callback) + except _RETRYABLE_ERRORS as e: + if attempt == attempts - 1: + raise + delay = min(2 ** attempt, _MAX_RETRY_DELAY) + random.uniform(0, 1) + progress.write(f"Encountered error transferring {description}:\n{str(e)}\n" + f"Retrying in {delay:.0f} seconds " + f"({attempts - (attempt + 1)} attempts remaining)") + time.sleep(delay) + + +def _run_sequentially(transfers: Dict[str, Callable[..., None]], + run_one: Callable[..., None]) -> Dict[str, Exception]: + """ + @private + + Runs the transfers one at a time in the calling thread, spawning nothing. + """ + errors: Dict[str, Exception] = {} + for description, transfer in transfers.items(): + try: + run_one(transfer, description) + except Exception as e: + errors[description] = e + return errors + + +def _run_concurrently(transfers: Dict[str, Callable[..., None]], + run_one: Callable[..., None], + threads: int, + action: str) -> Dict[str, Exception]: + """ + @private + + Runs up to `threads` transfers at a time. + """ + errors: Dict[str, Exception] = {} + with ThreadPoolExecutor(max_workers=threads, thread_name_prefix=f'cirro-{action}') as executor: + futures = {executor.submit(run_one, transfer, description): description + for description, transfer in transfers.items()} + for future in as_completed(futures): + try: + future.result() + except Exception as e: + errors[futures[future]] = e + return errors + + +def _run_transfers(transfers: Dict[str, Callable[..., None]], + progress: tqdm, + threads: int, + max_retries: int, + action: str): + """ + @private + + Runs the given transfers, retrying transient failures and reporting every failure + rather than letting the first one hide the rest. + """ + if threads < 1: + raise ValueError(f"threads must be at least 1, got {threads}") + + # One callback shared by every worker, so its lock serializes bar updates + run_one = partial(_transfer_with_retry, + callback=ProgressPercentage(progress), + progress=progress, + attempts=max(max_retries, 1)) + + if threads == 1: + errors = _run_sequentially(transfers, run_one) + else: + errors = _run_concurrently(transfers, run_one, threads, action) + + if errors: + detail = '\n'.join(f' {description}: {error}' for description, error in errors.items()) + raise RuntimeError(f"Failed to {action} {len(errors)} of {len(transfers)} files:\n{detail}") + + def upload_directory(directory: PathLike, files: List[PathLike], file_path_map: Dict[PathLike, str], @@ -132,7 +239,8 @@ def upload_directory(directory: PathLike, bucket: str, prefix: str, max_retries=10, - resume=False): + resume=False, + threads=Constants.default_transfer_threads): """ @private @@ -147,6 +255,10 @@ def upload_directory(directory: PathLike, prefix (str): S3 prefix max_retries (int): Number of retries resume (bool): Skip files that are already present in S3 under the prefix + threads (int): Number of files to upload at once. 1 disables threading. + + Raises: + RuntimeError: If any file could not be uploaded """ # Ensure all files are of the same type as the directory if not all(isinstance(file, type(directory)) for file in files): @@ -158,6 +270,10 @@ def upload_directory(directory: PathLike, # ListBucket s3:prefix condition on the vended upload credentials. already_uploaded = s3_client.get_file_sizes(bucket, f'{prefix}/') if resume else {} + transfers: Dict[str, Callable[..., None]] = {} + seen_keys = set() + total_size = 0 + for file in files: if isinstance(file, str): file_path = Path(directory, file) @@ -172,55 +288,72 @@ def upload_directory(directory: PathLike, file_relative = file_path.relative_to(directory).as_posix() key = f'{prefix}/{file_relative}' + file_size = file_path.stat().st_size + + # Two sources mapping to one destination cannot both survive, and uploading + # them concurrently would make the winner arbitrary. + if key in seen_keys: + raise ValueError(f"Multiple files map to the same destination path: {file_relative}") + seen_keys.add(key) # When resuming, skip files already uploaded with a matching size. # A size mismatch implies a modified file, so re-upload the file. expected_path = S3Path(f"s3://{bucket}/{key}") - if resume and already_uploaded.get(expected_path) == file_path.stat().st_size: + if resume and already_uploaded.get(expected_path) == file_size: print(f"Uploading {file_relative} skipped as it has already been uploaded.") continue - success = False + transfers[key] = partial(s3_client.upload_file, file_path=file_path, bucket=bucket, key=key) + total_size += file_size - # Retry up to max_retries times - for retry in range(max_retries): + with tqdm(total=total_size, + desc=f'Uploading {len(transfers)} files ({bytes_to_human_readable(total_size)})', + bar_format=_PROGRESS_BAR_FORMAT, + unit='B', unit_scale=True, + unit_divisor=1024) as progress: + _run_transfers(transfers, progress, threads, max_retries, 'upload') - # Try the upload - try: - s3_client.upload_file( - file_path=file_path, - bucket=bucket, - key=key - ) - - success = True - - # Catch the upload error - except (S3UploadFailedError, ConnectionError) as e: - delay = random.uniform(0, 60) + retry * 60 - # Report the error - print(f"Encountered error:\n{str(e)}\n" - f"Retrying in {delay:.0f} seconds ({max_retries - (retry + 1)} attempts remaining)") - time.sleep(delay) - - if success: - break - -def download_directory(directory: str, files: List[str], s3_client: S3Client, bucket: str, prefix: str) -> List[Path]: +def download_directory(directory: str, + files: Union[List[File], List[str]], + s3_client: S3Client, + bucket: str, + prefix: str, + max_retries=10, + threads=Constants.default_transfer_threads) -> List[Path]: """ @private + + Raises: + RuntimeError: If any file could not be downloaded """ - local_paths = [] - for file in files: - key = f'{prefix}/{file}'.lstrip('/') - local_path = Path(directory, file).expanduser() - local_path.parent.mkdir(parents=True, exist_ok=True) - - s3_client.download_file(local_path=local_path, - bucket=bucket, - key=key) - local_paths.append(local_path) + relative_paths = [file if isinstance(file, str) else file.relative_path for file in files] + # File objects carry their size from the dataset listing; plain paths do not + total_size = None if any(isinstance(f, str) for f in files) else sum(f.size for f in files) + + # Resolved up front so the returned order always matches the order requested + local_paths = [Path(directory, relative_path).expanduser() for relative_path in relative_paths] + for parent in {local_path.parent for local_path in local_paths}: + parent.mkdir(parents=True, exist_ok=True) + + transfers: Dict[str, Callable[..., None]] = { + relative_path: partial(s3_client.download_file, + local_path=local_path, + bucket=bucket, + key=f'{prefix}/{relative_path}'.lstrip('/')) + for relative_path, local_path in zip(relative_paths, local_paths) + } + + # Without a total there is no percentage to render + bar_format = _PROGRESS_BAR_FORMAT if total_size is not None else "{desc} | {n_fmt} | {rate_fmt}" + total_described = f' ({bytes_to_human_readable(total_size)})' if total_size is not None else '' + + with tqdm(total=total_size, + desc=f'Downloading {len(files)} files{total_described}', + bar_format=bar_format, + unit='B', unit_scale=True, unit_divisor=1024) as progress: + _run_transfers(transfers, progress, threads, max_retries, 'download') + return local_paths diff --git a/cirro/sdk/dataset.py b/cirro/sdk/dataset.py index 21495f6..a2f0d87 100644 --- a/cirro/sdk/dataset.py +++ b/cirro/sdk/dataset.py @@ -10,6 +10,7 @@ Status, RunAnalysisRequestParams, Tag, ArtifactType, NamedItem, ValidateFileRequirementsRequest from cirro.cirro_client import CirroApi +from cirro.config import Constants from cirro.file_utils import bytes_to_human_readable, filter_files_by_pattern from cirro.models.assets import DatasetAssets from cirro.models.file import PathLike @@ -513,7 +514,8 @@ def list_artifacts(self) -> List[DataPortalFile]: ] ) - def download_files(self, download_location: str = None, glob: str = None) -> None: + def download_files(self, download_location: str = None, glob: str = None, + threads: int = Constants.default_transfer_threads) -> None: """ Download all the files from the dataset to a local directory. @@ -522,12 +524,13 @@ def download_files(self, download_location: str = None, glob: str = None) -> Non glob (str): Optional wildcard expression to filter which files are downloaded (e.g., ``'*.csv'``, ``'data/**/*.tsv.gz'``). If omitted, all files are downloaded. + threads (int): Number of files to download at once. 1 disables threading. """ files = self.list_files() if glob is not None: files = DataPortalFiles(filter_files_by_pattern(list(files), glob)) - files.download(download_location) + files.download(download_location, threads=threads) def run_analysis( self, diff --git a/cirro/sdk/file.py b/cirro/sdk/file.py index 3cdbb4d..f73eae7 100644 --- a/cirro/sdk/file.py +++ b/cirro/sdk/file.py @@ -2,6 +2,7 @@ from typing import List from cirro.cirro_client import CirroApi +from cirro.config import Constants from cirro.models.file import File, PathLike from cirro.sdk.asset import DataPortalAssets, DataPortalAsset from cirro.sdk.exceptions import DataPortalInputError @@ -96,7 +97,7 @@ def download(self, download_location: str = None) -> Path: return self._client.file.download_files( self._file.access_context, download_location, - [self.relative_path] + [self._file] )[0] def validate(self, local_path: PathLike): @@ -133,15 +134,31 @@ class DataPortalFiles(DataPortalAssets[DataPortalFile]): asset_name = "file" - def download(self, download_location: str = None) -> List[Path]: + def download(self, download_location: str = None, + threads: int = Constants.default_transfer_threads) -> List[Path]: """ Download the collection of files to a local directory. + Args: + download_location (str): Path to local directory + threads (int): Number of files to download at once. 1 disables threading. + Returns: List of paths to downloaded files. """ - local_paths = [] - for f in self: - local_paths.append(f.download(download_location)) - return local_paths + if len(self) == 0: + return [] + + if download_location is None: + raise DataPortalInputError("Must provide download location") + + # Downloaded in one call so the S3 client is built once and the files + # transfer concurrently. Every file in a collection shares an access context. + first_file = self[0] + return first_file._client.file.download_files( + first_file._file.access_context, + download_location, + [f._file for f in self], + threads=threads + ) diff --git a/cirro/services/dataset.py b/cirro/services/dataset.py index fb3054b..85c379e 100644 --- a/cirro/services/dataset.py +++ b/cirro/services/dataset.py @@ -10,6 +10,7 @@ from cirro.file_utils import is_hidden_file from cirro.models.assets import DatasetAssets, Artifact from cirro.models.dataset import DatasetValidationResponse +from cirro.config import Constants from cirro.models.file import FileAccessContext, File, PathLike from cirro.services.base import get_all_records from cirro.services.file import FileEnabledService @@ -244,7 +245,8 @@ def upload_files(self, directory: PathLike, files: List[PathLike] = None, file_path_map: Dict[PathLike, str] = None, - resume: bool = False) -> None: + resume: bool = False, + threads: int = Constants.default_transfer_threads) -> None: """ Uploads files to a given dataset from the specified directory. @@ -262,6 +264,7 @@ def upload_files(self, from source path to destination path, used to "re-write" paths within the dataset. resume (bool): If True, skip files already uploaded to the dataset, only uploading the files that are still missing. Used to continue an interrupted upload. + threads (int): Number of files to upload at once. 1 disables threading. ```python from cirro.cirro_client import CirroApi from cirro.file_utils import generate_flattened_file_map @@ -305,7 +308,8 @@ def upload_files(self, directory=directory, files=files, file_path_map=file_path_map, - resume=resume + resume=resume, + threads=threads ) def validate_folder( @@ -372,7 +376,8 @@ def download_files( dataset_id: str, download_location: str, files: Union[List[File], List[str]] = None, - file_limit: int = 100000 + file_limit: int = 100000, + threads: int = Constants.default_transfer_threads ) -> None: """ Downloads files from a dataset @@ -386,6 +391,7 @@ def download_files( download_location (str): Local destination for downloaded files files (typing.List[str]): Optional list of files to download file_limit (int): Maximum number of files to get (default 100,000) + threads (int): Number of files to download at once. 1 disables threading. """ if files is None: files = self.get_assets_listing(project_id, dataset_id, file_limit=file_limit).files @@ -395,7 +401,7 @@ def download_files( first_file = files[0] if isinstance(first_file, File): - files = [file.relative_path for file in files] + # Kept as File objects so their known sizes reach the transfer access_context = first_file.access_context else: dataset = self.get(project_id, dataset_id) @@ -407,7 +413,7 @@ def download_files( access_context = FileAccessContext.download(project_id=project_id, base_url=dataset.s3) - self._file_service.download_files(access_context, download_location, files) + self._file_service.download_files(access_context, download_location, files, threads=threads) def update_samplesheet( self, diff --git a/cirro/services/file.py b/cirro/services/file.py index 9a058f2..02ba900 100644 --- a/cirro/services/file.py +++ b/cirro/services/file.py @@ -1,9 +1,10 @@ import logging import threading +from contextlib import closing from datetime import datetime, timezone from functools import partial from pathlib import Path -from typing import List, Dict +from typing import List, Dict, Union from botocore.client import BaseClient from cirro_api_client import CirroApiClient @@ -11,6 +12,7 @@ from cirro_api_client.v1.models import AWSCredentials, ProjectAccessType from cirro.clients.s3 import S3Client +from cirro.config import Constants from cirro.file_utils import upload_directory, download_directory, get_checksum from cirro.models.file import FileAccessContext, File, PathLike from cirro.services.base import BaseService @@ -160,7 +162,8 @@ def upload_files(self, directory: PathLike, files: List[PathLike], file_path_map: Dict[PathLike, str], - resume: bool = False) -> None: + resume: bool = False, + threads: int = Constants.default_transfer_threads) -> None: """ Uploads a list of files from the specified directory @@ -172,40 +175,45 @@ def upload_files(self, file_path_map (typing.Dict[str|Path, str]): Optional mapping of file paths to upload from source path to destination path, used to "re-write" paths within the dataset. resume (bool): If True, skip files already present in S3 under the destination prefix. + threads (int): Number of files to upload at once. 1 disables threading. """ - s3_client = self._generate_s3_client(access_context) - - upload_directory( - directory=directory, - files=files, - file_path_map=file_path_map, - s3_client=s3_client, - bucket=access_context.bucket, - prefix=access_context.prefix, - max_retries=self.transfer_retries, - resume=resume - ) + with closing(self._generate_s3_client(access_context, threads)) as s3_client: + upload_directory( + directory=directory, + files=files, + file_path_map=file_path_map, + s3_client=s3_client, + bucket=access_context.bucket, + prefix=access_context.prefix, + max_retries=self.transfer_retries, + resume=resume, + threads=threads + ) - def download_files(self, access_context: FileAccessContext, directory: str, files: List[str]) -> List[Path]: + def download_files(self, access_context: FileAccessContext, directory: str, + files: Union[List[File], List[str]], + threads: int = Constants.default_transfer_threads) -> List[Path]: """ Download a list of files to the specified directory Args: access_context (cirro.models.file.FileAccessContext): File access context, use class methods to generate directory (str): download location - files (List[str]): relative path of files to download + files (typing.List[File]|typing.List[str]): File objects, or paths relative to + the access context prefix. File objects carry their size, saving a request per file. + threads (int): Number of files to download at once. 1 disables threading. Returns: List of paths to downloaded files """ - s3_client = self._generate_s3_client(access_context) - - return download_directory( - directory, - files, - s3_client, - access_context.bucket, - access_context.prefix - ) + with closing(self._generate_s3_client(access_context, threads)) as s3_client: + return download_directory( + directory, + files, + s3_client, + access_context.bucket, + access_context.prefix, + threads=threads + ) def is_valid_file(self, file: File, local_file: Path) -> bool: """ @@ -284,13 +292,15 @@ def get_file_stats(self, file: File) -> dict: logger.debug(f"File stats for file {file.relative_path} is {stats}") return stats - def _generate_s3_client(self, access_context: FileAccessContext): + def _generate_s3_client(self, access_context: FileAccessContext, + threads: int = Constants.default_transfer_threads): """ Generates the Cirro-S3 client to perform operations on files """ return S3Client( partial(self.get_access_credentials, access_context), - self.checksum_method + self.checksum_method, + threads ) diff --git a/pyproject.toml b/pyproject.toml index 6d21255..658a114 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "cirro" -version = "1.12.1" +version = "1.13.0" description = "CLI tool and SDK for interacting with the Cirro platform" authors = ["Cirro Bio "] license = "MIT" diff --git a/tests/test_file_utils.py b/tests/test_file_utils.py index 5b60eb9..4754c9f 100644 --- a/tests/test_file_utils.py +++ b/tests/test_file_utils.py @@ -1,8 +1,14 @@ +import tempfile +import threading +import time import unittest from pathlib import Path -from unittest.mock import Mock, call +from unittest.mock import ANY, Mock, call -from cirro.file_utils import upload_directory, get_files_in_directory, get_files_stats +from boto3.exceptions import S3UploadFailedError + +from cirro.file_utils import upload_directory, download_directory, get_files_in_directory, get_files_stats +from cirro.models.file import File, FileAccessContext from cirro.models.s3_path import S3Path @@ -13,6 +19,13 @@ def setUp(self): self.test_bucket = 'project-1a1a' self.test_prefix = 'datasets/1a1a/data' + def _make_files(self, directory: Path, relative_paths): + """Create each relative path under the directory so it can be stat'd.""" + for relative_path in relative_paths: + file_path = Path(directory, relative_path) + file_path.parent.mkdir(parents=True, exist_ok=True) + file_path.write_text(relative_path) + def test_get_file_stats(self): directory = Path(__file__).parent / 'data' files = get_files_in_directory(directory) @@ -23,47 +36,54 @@ def test_get_file_stats(self): self.assertIn('KB', stats.size_friendly) def test_upload_directory_pathlike(self): - test_path = Path('/Users/test/Documents/dataset1') - test_files = [ - Path('/Users/test/Documents/dataset1/test_file.fastq'), - Path('/Users/test/Documents/dataset1/folder1/test_file.fastq'), - ] - upload_directory(directory=test_path, - files=test_files, - file_path_map={}, - s3_client=self.mock_s3_client, - bucket=self.test_bucket, - prefix=self.test_prefix) + with tempfile.TemporaryDirectory() as directory: + test_path = Path(directory) + self._make_files(test_path, ['test_file.fastq', 'folder1/test_file.fastq']) + test_files = [ + test_path / 'test_file.fastq', + test_path / 'folder1' / 'test_file.fastq', + ] + upload_directory(directory=test_path, + files=test_files, + file_path_map={}, + s3_client=self.mock_s3_client, + bucket=self.test_bucket, + prefix=self.test_prefix) # The function should upload files relative to the directory path. self.mock_s3_client.upload_file.assert_has_calls([ - call(file_path=test_files[0], bucket=self.test_bucket, key=f'{self.test_prefix}/test_file.fastq'), - call(file_path=test_files[1], bucket=self.test_bucket, key=f'{self.test_prefix}/folder1/test_file.fastq') + call(file_path=test_files[0], bucket=self.test_bucket, + key=f'{self.test_prefix}/test_file.fastq', callback=ANY), + call(file_path=test_files[1], bucket=self.test_bucket, + key=f'{self.test_prefix}/folder1/test_file.fastq', callback=ANY) ], any_order=True) def test_upload_directory_string(self): - test_path = 'data' - test_files = [ - 'file1.txt', - 'folder1/file2.txt' - ] - upload_directory(directory=test_path, - files=test_files, - file_path_map={}, - s3_client=self.mock_s3_client, - bucket=self.test_bucket, - prefix=self.test_prefix) + with tempfile.TemporaryDirectory() as test_path: + test_files = [ + 'file1.txt', + 'folder1/file2.txt' + ] + self._make_files(Path(test_path), test_files) + upload_directory(directory=test_path, + files=test_files, + file_path_map={}, + s3_client=self.mock_s3_client, + bucket=self.test_bucket, + prefix=self.test_prefix) - # The function should upload files relative to the directory path, - # but also format file_path into the Path object. - self.mock_s3_client.upload_file.assert_has_calls([ - call(file_path=Path(test_path, test_files[0]), - bucket=self.test_bucket, - key=f'{self.test_prefix}/file1.txt'), - call(file_path=Path(test_path, test_files[1]), - bucket=self.test_bucket, - key=f'{self.test_prefix}/folder1/file2.txt') - ], any_order=True) + # The function should upload files relative to the directory path, + # but also format file_path into the Path object. + self.mock_s3_client.upload_file.assert_has_calls([ + call(file_path=Path(test_path, test_files[0]), + bucket=self.test_bucket, + key=f'{self.test_prefix}/file1.txt', + callback=ANY), + call(file_path=Path(test_path, test_files[1]), + bucket=self.test_bucket, + key=f'{self.test_prefix}/folder1/file2.txt', + callback=ANY) + ], any_order=True) def test_upload_directory_different_types(self): test_path = Path('s3://bucket/dataset1') @@ -79,7 +99,6 @@ def test_upload_directory_different_types(self): prefix=self.test_prefix) def test_upload_directory_file_map_included(self): - test_path = 'data' test_files = [ 'file1.txt', 'folder1/file2.txt', @@ -92,25 +111,150 @@ def test_upload_directory_file_map_included(self): # unmapped file3 } - upload_directory(directory=test_path, - files=test_files, - file_path_map=file_path_map, - s3_client=self.mock_s3_client, - bucket=self.test_bucket, - prefix=self.test_prefix) + with tempfile.TemporaryDirectory() as test_path: + self._make_files(Path(test_path), test_files) + upload_directory(directory=test_path, + files=test_files, + file_path_map=file_path_map, + s3_client=self.mock_s3_client, + bucket=self.test_bucket, + prefix=self.test_prefix) - # Check that upload file was called with the mapped key - self.mock_s3_client.upload_file.assert_has_calls([ - call(file_path=Path(test_path, test_files[0]), - bucket=self.test_bucket, - key=f'{self.test_prefix}/mapped_file1.txt'), - call(file_path=Path(test_path, test_files[1]), - bucket=self.test_bucket, - key=f'{self.test_prefix}/mapped_file2.txt'), - call(file_path=Path(test_path, test_files[2]), - bucket=self.test_bucket, - key=f'{self.test_prefix}/folder1/unmapped.txt') - ], any_order=True) + # Check that upload file was called with the mapped key + self.mock_s3_client.upload_file.assert_has_calls([ + call(file_path=Path(test_path, test_files[0]), + bucket=self.test_bucket, + key=f'{self.test_prefix}/mapped_file1.txt', + callback=ANY), + call(file_path=Path(test_path, test_files[1]), + bucket=self.test_bucket, + key=f'{self.test_prefix}/mapped_file2.txt', + callback=ANY), + call(file_path=Path(test_path, test_files[2]), + bucket=self.test_bucket, + key=f'{self.test_prefix}/folder1/unmapped.txt', + callback=ANY) + ], any_order=True) + + def test_upload_directory_duplicate_destination(self): + test_files = ['folder1/file.txt', 'folder2/file.txt'] + + with tempfile.TemporaryDirectory() as test_path: + self._make_files(Path(test_path), test_files) + + # Flattening both files onto one destination path cannot be honored + with self.assertRaises(ValueError): + upload_directory(directory=test_path, + files=test_files, + file_path_map={f: 'file.txt' for f in test_files}, + s3_client=self.mock_s3_client, + bucket=self.test_bucket, + prefix=self.test_prefix) + + def test_upload_directory_runs_files_concurrently(self): + test_files = [f'file{i}.txt' for i in range(4)] + in_flight = [] + peak_in_flight = 0 + lock = threading.Lock() + + def slow_upload(file_path, bucket, key, callback): + nonlocal peak_in_flight + with lock: + in_flight.append(key) + peak_in_flight = max(peak_in_flight, len(in_flight)) + time.sleep(0.05) + with lock: + in_flight.remove(key) + + self.mock_s3_client.upload_file = Mock(side_effect=slow_upload) + + with tempfile.TemporaryDirectory() as test_path: + self._make_files(Path(test_path), test_files) + upload_directory(directory=test_path, + files=test_files, + file_path_map={}, + s3_client=self.mock_s3_client, + bucket=self.test_bucket, + prefix=self.test_prefix, + threads=4) + + self.assertGreater(peak_in_flight, 1) + self.assertEqual(self.mock_s3_client.upload_file.call_count, 4) + + def test_upload_directory_single_thread_spawns_no_threads(self): + test_files = [f'file{i}.txt' for i in range(4)] + thread_names = set() + + def record_thread(file_path, bucket, key, callback): + thread_names.add(threading.current_thread().name) + + self.mock_s3_client.upload_file = Mock(side_effect=record_thread) + + with tempfile.TemporaryDirectory() as test_path: + self._make_files(Path(test_path), test_files) + upload_directory(directory=test_path, + files=test_files, + file_path_map={}, + s3_client=self.mock_s3_client, + bucket=self.test_bucket, + prefix=self.test_prefix, + threads=1) + + # Everything must run in the calling thread, for environments without threads + self.assertEqual(thread_names, {threading.current_thread().name}) + self.assertEqual([t.name for t in threading.enumerate() if t.name.startswith('cirro-')], []) + self.assertEqual(self.mock_s3_client.upload_file.call_count, 4) + + def test_upload_directory_rejects_zero_threads(self): + with tempfile.TemporaryDirectory() as test_path: + self._make_files(Path(test_path), ['file1.txt']) + with self.assertRaises(ValueError): + upload_directory(directory=test_path, + files=['file1.txt'], + file_path_map={}, + s3_client=self.mock_s3_client, + bucket=self.test_bucket, + prefix=self.test_prefix, + threads=0) + + def test_download_directory_single_thread_spawns_no_threads(self): + thread_names = set() + + def record_thread(local_path, bucket, key, callback): + thread_names.add(threading.current_thread().name) + + self.mock_s3_client.download_file = Mock(side_effect=record_thread) + + with tempfile.TemporaryDirectory() as directory: + download_directory(directory=directory, + files=['a.txt', 'b.txt'], + s3_client=self.mock_s3_client, + bucket=self.test_bucket, + prefix=self.test_prefix, + threads=1) + + self.assertEqual(thread_names, {threading.current_thread().name}) + self.assertEqual([t.name for t in threading.enumerate() if t.name.startswith('cirro-')], []) + + def test_upload_directory_reports_failures(self): + test_files = ['file1.txt', 'file2.txt'] + self.mock_s3_client.upload_file = Mock(side_effect=S3UploadFailedError('denied')) + + with tempfile.TemporaryDirectory() as test_path: + self._make_files(Path(test_path), test_files) + + with self.assertRaises(RuntimeError) as raised: + upload_directory(directory=test_path, + files=test_files, + file_path_map={}, + s3_client=self.mock_s3_client, + bucket=self.test_bucket, + prefix=self.test_prefix, + max_retries=1) + + # Every failed file should be named, not just the first + self.assertIn('file1.txt', str(raised.exception)) + self.assertIn('file2.txt', str(raised.exception)) def test_upload_directory_resume_skips_already_uploaded(self): directory = Path(__file__).parent / 'data' / 'example_data_1' @@ -133,7 +277,8 @@ def test_upload_directory_resume_skips_already_uploaded(self): self.mock_s3_client.upload_file.assert_called_once_with( file_path=pending, bucket=self.test_bucket, - key=f'{self.test_prefix}/files.csv') + key=f'{self.test_prefix}/files.csv', + callback=ANY) def test_upload_directory_resume_reuploads_size_mismatch(self): directory = Path(__file__).parent / 'data' / 'example_data_1' @@ -155,7 +300,8 @@ def test_upload_directory_resume_reuploads_size_mismatch(self): self.mock_s3_client.upload_file.assert_called_once_with( file_path=partial, bucket=self.test_bucket, - key=f'{self.test_prefix}/samplesheet.csv') + key=f'{self.test_prefix}/samplesheet.csv', + callback=ANY) def test_upload_directory_resume_disabled_ignores_remote(self): directory = Path(__file__).parent / 'data' / 'example_data_1' @@ -172,7 +318,45 @@ def test_upload_directory_resume_disabled_ignores_remote(self): # Should not get file sizes and should upload all files. self.mock_s3_client.get_file_sizes.assert_not_called() self.mock_s3_client.upload_file.assert_has_calls([ - call(file_path=file1, bucket=self.test_bucket, key=f'{self.test_prefix}/samplesheet.csv'), - call(file_path=file2, bucket=self.test_bucket, key=f'{self.test_prefix}/files.csv'), + call(file_path=file1, bucket=self.test_bucket, + key=f'{self.test_prefix}/samplesheet.csv', callback=ANY), + call(file_path=file2, bucket=self.test_bucket, + key=f'{self.test_prefix}/files.csv', callback=ANY), ], any_order=True) + def _make_dataset_file(self, relative_path: str, size: int) -> File: + access_context = FileAccessContext.download(project_id='project-1a1a', + base_url=f's3://{self.test_bucket}/{self.test_prefix}') + return File(relative_path=relative_path, + size=size, + access_context=access_context, + metadata={}) + + def test_download_directory_uses_known_sizes(self): + files = [self._make_dataset_file('file1.txt', 10), + self._make_dataset_file('folder1/file2.txt', 20)] + + with tempfile.TemporaryDirectory() as directory: + local_paths = download_directory(directory=directory, + files=files, + s3_client=self.mock_s3_client, + bucket=self.test_bucket, + prefix=self.test_prefix, + threads=2) + + # Sizes come from the dataset, so no request is needed to look them up + self.mock_s3_client.get_file_stats.assert_not_called() + + # Returned paths keep the order they were requested in + self.assertEqual(local_paths, [Path(directory, 'file1.txt'), + Path(directory, 'folder1/file2.txt')]) + + # Parent directories are created before the transfer + self.assertTrue(Path(directory, 'folder1').is_dir()) + + self.mock_s3_client.download_file.assert_has_calls([ + call(local_path=local_paths[0], bucket=self.test_bucket, + key=f'{self.test_prefix}/file1.txt', callback=ANY), + call(local_path=local_paths[1], bucket=self.test_bucket, + key=f'{self.test_prefix}/folder1/file2.txt', callback=ANY), + ], any_order=True) diff --git a/tests/test_read_files.py b/tests/test_read_files.py index 2292d0e..7d2e3b3 100644 --- a/tests/test_read_files.py +++ b/tests/test_read_files.py @@ -364,15 +364,18 @@ def setUp(self): self.tsv_file, self.txt_file, ]) - for f in [self.csv_file, self.tsv_file, self.txt_file]: - f.download = Mock(return_value=None) def _downloaded_paths(self): - return [ - f.relative_path - for f in [self.csv_file, self.tsv_file, self.txt_file] - if f.download.called - ] + """ + Relative paths from the single batched download call. Which file's client + receives it depends on which files survive the glob. + """ + for f in [self.csv_file, self.tsv_file, self.txt_file]: + download_files = f._client.file.download_files + if download_files.called: + _, _, files = download_files.call_args.args + return [file.relative_path for file in files] + return [] def test_no_glob_downloads_all(self): self.dataset.download_files(download_location='/tmp') @@ -399,6 +402,14 @@ def test_globstar_filters_by_subdirectory(self): downloaded = self._downloaded_paths() self.assertEqual(downloaded, ['logs/run.log']) + def test_downloads_in_one_batched_call(self): + self.dataset.download_files(download_location='/tmp') + + # The whole collection goes out in one call, rather than one call per file + self.assertEqual(self.csv_file._client.file.download_files.call_count, 1) + self.assertEqual(self.tsv_file._client.file.download_files.call_count, 0) + self.assertEqual(self.txt_file._client.file.download_files.call_count, 0) + class TestPatternToRegex(unittest.TestCase): def _match(self, pattern, path): diff --git a/tests/test_s3_client.py b/tests/test_s3_client.py new file mode 100644 index 0000000..0e365b6 --- /dev/null +++ b/tests/test_s3_client.py @@ -0,0 +1,161 @@ +import io +import tempfile +import unittest +from datetime import datetime, timedelta, timezone +from pathlib import Path +from unittest.mock import Mock, patch + +from cirro.clients.s3 import S3Client, local_filename + + +class RemotePath: + """ + Stands in for a Path-like object backed by something other than the local + filesystem, such as an s3fs path (see #115). + """ + + def __init__(self, uri: str, contents: bytes = b'remote'): + self.uri = uri + self.contents = contents + + def __fspath__(self): + return self.uri + + def open(self, mode='rb'): + return io.BytesIO(self.contents) + + +class StreamOnlyPath: + """A Path-like object that only supports open(), with no filesystem path.""" + + def open(self, mode='rb'): + return io.BytesIO(b'stream') + + +class TestS3Client(unittest.TestCase): + def setUp(self): + # Without a creds getter a standard client is built, which needs no credentials + self.s3_client = S3Client() + + def test_connection_pool_covers_request_concurrency(self): + creds = Mock(access_key_id='key', secret_access_key='secret', session_token='token', + region='us-west-2', + expiration=datetime.now(tz=timezone.utc) + timedelta(hours=1)) + + # A pool smaller than boto3's 10 concurrent requests per transfer would make + # urllib3 discard connections and serialize the transfers + config = S3Client(creds_getter=lambda: creds).get_aws_client().meta.config + self.assertGreater(config.max_pool_connections, 10) + + @patch('cirro.clients.s3.create_transfer_manager') + def test_transfer_manager_is_reused(self, mock_create_manager): + with tempfile.TemporaryDirectory() as directory: + file_path = Path(directory, 'a.txt') + file_path.write_text('hello') + + self.s3_client.upload_file(file_path, 'bucket', 'key-1', callback=Mock()) + self.s3_client.upload_file(file_path, 'bucket', 'key-2', callback=Mock()) + self.s3_client.download_file(Path(directory, 'b.txt'), 'bucket', 'key-1', callback=Mock()) + + # One manager, and the thread pools it owns, for every transfer on this client + self.assertEqual(mock_create_manager.call_count, 1) + + @patch('cirro.clients.s3.create_transfer_manager') + def test_upload_passes_a_filename(self, mock_create_manager): + with tempfile.TemporaryDirectory() as directory: + file_path = Path(directory, 'a.txt') + file_path.write_text('hello') + + self.s3_client.upload_file(file_path, 'bucket', 'key', callback=Mock()) + + # Passing an open file object instead would make s3transfer read the parts + # serially rather than in parallel + upload = mock_create_manager.return_value.upload + self.assertEqual(upload.call_args.args[0], str(file_path)) + + @patch('cirro.clients.s3.create_transfer_manager') + def test_download_with_callback_skips_head_object(self, _mock_create_manager): + self.s3_client.get_file_stats = Mock() + + self.s3_client.download_file(Path('/tmp/a.txt'), 'bucket', 'key', callback=Mock()) + + # The caller already knows the size, so no request is needed to look it up + self.s3_client.get_file_stats.assert_not_called() + + @patch('cirro.clients.s3.create_transfer_manager') + def test_close_shuts_down_the_transfer_manager(self, mock_create_manager): + with tempfile.TemporaryDirectory() as directory: + file_path = Path(directory, 'a.txt') + file_path.write_text('hello') + self.s3_client.upload_file(file_path, 'bucket', 'key', callback=Mock()) + + self.s3_client.close() + + # The pools use non-daemon threads that garbage collection only partly + # reclaims, so they have to be shut down explicitly + mock_create_manager.return_value.shutdown.assert_called_once() + + # A later transfer must not reuse the shut-down manager + self.s3_client.upload_file(file_path, 'bucket', 'key', callback=Mock()) + self.assertEqual(mock_create_manager.call_count, 2) + + @patch('cirro.clients.s3.create_transfer_manager') + def test_close_without_any_transfer_is_a_no_op(self, mock_create_manager): + self.s3_client.close() + self.s3_client.close() + mock_create_manager.assert_not_called() + + def test_single_thread_disables_boto3_threading(self): + # boto3 runs the whole transfer in the calling thread when use_threads is False + self.assertFalse(S3Client(threads=1)._transfer_config.use_threads) + self.assertTrue(S3Client(threads=2)._transfer_config.use_threads) + + def test_local_filename_only_matches_real_local_files(self): + with tempfile.TemporaryDirectory() as directory: + existing = Path(directory, 'a.txt') + existing.write_text('hello') + + self.assertIsNotNone(local_filename(existing)) + self.assertIsNotNone(local_filename(str(existing))) + self.assertIsNone(local_filename(Path(directory, 'missing.txt'))) + self.assertIsNone(local_filename(Path(directory))) + self.assertIsNone(local_filename(RemotePath('s3://bucket/key'))) + self.assertIsNone(local_filename(StreamOnlyPath())) + + @patch('cirro.clients.s3.create_transfer_manager') + def test_local_file_uses_the_transfer_manager(self, mock_create_manager): + with tempfile.TemporaryDirectory() as directory: + file_path = Path(directory, 'a.txt') + file_path.write_text('hello') + self.s3_client.upload_file(file_path, 'bucket', 'key', callback=Mock()) + + # A filename goes through S3Transfer, which reads the parts in parallel + upload = mock_create_manager.return_value.upload + upload.assert_called_once() + self.assertEqual(upload.call_args.args[0], str(file_path)) + + @patch('cirro.clients.s3.create_transfer_manager') + def test_remote_path_is_streamed_through_its_own_open(self, mock_create_manager): + # The transfer manager can only open real files, so a Path-like object backed + # by another filesystem has to be streamed instead (#115) + uploaded = [] + + def capture(stream, *args, **kwargs): + # Read while the stream is still open; upload_file closes it on the way out + uploaded.append(stream.read()) + return Mock() + + mock_create_manager.return_value.upload = Mock(side_effect=capture) + + self.s3_client.upload_file(RemotePath('s3://bucket/source.txt'), 'bucket', 'key', + callback=Mock()) + + # Streamed through the same shared manager, not a per-file one + self.assertEqual(uploaded, [b'remote']) + self.assertEqual(mock_create_manager.call_count, 1) + + @patch('cirro.clients.s3.create_transfer_manager') + def test_stream_only_path_is_supported(self, mock_create_manager): + self.s3_client.upload_file(StreamOnlyPath(), 'bucket', 'key', callback=Mock()) + + mock_create_manager.return_value.upload.assert_called_once()