Upgrade local download concurrency control to use aioboto3's built-in handling (#579)

* Use aioboto3's built-in concurrency handling

* Add test coverage

Author: ponyisi

servicex/minio_adapter.py

@@ -30,20 +30,20 @@
 from pathlib import Path
 from typing import List
 
-import asyncio
 from tenacity import retry, stop_after_attempt, wait_random_exponential
 
 import aioboto3
+from boto3.s3.transfer import TransferConfig
 
 from servicex.models import ResultFile, TransformStatus
 
-_semaphore = asyncio.Semaphore(8)
+_transferconfig = TransferConfig(max_concurrency=10)
 
 
-def init_download_semaphore(concurrency: int = 8):
+def init_s3_config(concurrency: int = 10):
     "Update the number of concurrent connections"
-    global _semaphore
-    _semaphore = asyncio.Semaphore(concurrency)
+    global _transferconfig
+    _transferconfig = TransferConfig(max_concurrency=concurrency)
 
 
 def _sanitize_filename(fname: str):
@@ -114,31 +114,29 @@ async def download_file(
             )
         )
 
-        async with _semaphore:
-            async with self.minio.resource("s3", endpoint_url=self.endpoint_host) as s3:
-                obj = await s3.Object(self.bucket, object_name)
-                remotesize = await obj.content_length
-                if path.exists():
-                    # if file size is the same, let's not download anything
-                    # maybe move to a better verification mechanism with e-tags in the future
-                    localsize = path.stat().st_size
-                    if localsize == remotesize:
-                        return path.resolve()
-                await obj.download_file(path.as_posix())
+        async with self.minio.resource("s3", endpoint_url=self.endpoint_host) as s3:
+            obj = await s3.Object(self.bucket, object_name)
+            remotesize = await obj.content_length
+            if path.exists():
+                # if file size is the same, let's not download anything
+                # maybe move to a better verification mechanism with e-tags in the future
                 localsize = path.stat().st_size
-                if localsize != remotesize:
-                    raise RuntimeError(f"Download of {object_name} failed")
+                if localsize == remotesize:
+                    return path.resolve()
+            await obj.download_file(path.as_posix(), Config=_transferconfig)
+            localsize = path.stat().st_size
+            if localsize != remotesize:
+                raise RuntimeError(f"Download of {object_name} failed")
         return path.resolve()
 
     @retry(
         stop=stop_after_attempt(3), wait=wait_random_exponential(max=60), reraise=True
     )
     async def get_signed_url(self, object_name: str) -> str:
-        async with _semaphore:
-            async with self.minio.client("s3", endpoint_url=self.endpoint_host) as s3:
-                return await s3.generate_presigned_url(
-                    "get_object", Params={"Bucket": self.bucket, "Key": object_name}
-                )
+        async with self.minio.client("s3", endpoint_url=self.endpoint_host) as s3:
+            return await s3.generate_presigned_url(
+                "get_object", Params={"Bucket": self.bucket, "Key": object_name}
+            )
 
     @classmethod
     def hash_path(cls, file_name):

servicex/servicex_client.py

@@ -228,7 +228,7 @@ def deliver(
     fail_if_incomplete: bool = True,
     ignore_local_cache: bool = False,
     progress_bar: ProgressBarFormat = ProgressBarFormat.default,
-    concurrency: int = 8,
+    concurrency: int = 10,
 ):
     r"""
     Execute a ServiceX query.
@@ -254,9 +254,9 @@ def deliver(
     :return: A dictionary mapping the name of each :py:class:`Sample` to a :py:class:`.GuardList`
             with the file names or URLs for the outputs.
     """
-    from .minio_adapter import init_download_semaphore
+    from .minio_adapter import init_s3_config
 
-    init_download_semaphore(concurrency)
+    init_s3_config(concurrency)
     config = _load_ServiceXSpec(spec)
 
     if ignore_local_cache or config.General.IgnoreLocalCache:

tests/test_minio_adapter.py

@@ -160,6 +160,25 @@ async def test_download_short_filename_change(minio_adapter, populate_bucket):
     result.unlink()  # it should exist, from above ...
 
 
+@pytest.mark.parametrize("populate_bucket", ["test.txt"], indirect=True)
+@pytest.mark.asyncio
+async def test_download_repeat(minio_adapter, populate_bucket):
+    import asyncio
+
+    print(await minio_adapter.list_bucket())
+    result = await minio_adapter.download_file("test.txt", local_dir="/tmp/foo")
+    assert str(result).endswith("test.txt")
+    assert result.exists()
+    t0 = result.stat().st_mtime_ns
+    await asyncio.sleep(4)  # hopefully long enough for Windows/FAT32 ... ?
+
+    result2 = await minio_adapter.download_file("test.txt", local_dir="/tmp/foo")
+    assert result2.exists()
+    assert result2 == result
+    assert t0 == result2.stat().st_mtime_ns
+    result.unlink()  # it should exist, from above ...
+
+
 @pytest.mark.parametrize("populate_bucket", ["test.txt"], indirect=True)
 @pytest.mark.asyncio
 async def test_download_file_retry(download_patch, minio_adapter, populate_bucket):