CI-260 - Transfer dataset files concurrently

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)

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():

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):

cirro/clients/s3.py

@@ -1,16 +1,22 @@
 import threading
 from pathlib import Path
-from typing import Callable
+from typing import Callable, Optional
 
 from boto3 import Session
+from boto3.s3.transfer import S3Transfer, TransferConfig
 from botocore.config import Config
 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 format_creds_for_session(creds: AWSCredentials):
@@ -33,42 +39,69 @@ 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._transfer: Optional[S3Transfer] = 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`.
+        """
+        # Pass the path rather than an open file object: s3transfer only reads parts
+        # in parallel when it can open the file itself.
+        self._get_transfer().upload_file(
+            filename=str(file_path),
+            bucket=bucket,
+            key=key,
+            callback=callback,
+            extra_args=self._upload_args
+        )
+
+    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_transfer(self) -> S3Transfer:
+        """
+        A single transfer manager, and the thread pools it owns, is shared by every
+        transfer on this client rather than rebuilt for each file.
+        """
+        with self._transfer_lock:
+            if self._transfer is None:
+                self._transfer = S3Transfer(self._client, self._transfer_config)
+            return self._transfer
+
+    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._transfer is not None:
+                self._transfer.__exit__(None, None, None)
+                self._transfer = None
 
     def create_object(self, bucket: str, key: str, contents: str, content_type: str):
         self._client.put_object(
@@ -134,7 +167,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)
 

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):