Switch S3 client to boto3 (#572)

ponyisi · web-flow · commit 6517857ec415 · 2025-03-26T11:32:07.000-05:00
* Limit simultaneous downloads; retry failed downloads

* Better handling of fail_if_incomplete, access to set concurrency

* Switch S3 client to boto3
diff --git a/pyproject.toml b/pyproject.toml
@@ -38,7 +38,7 @@ dependencies = [
     "pydantic>=2.6.0",
     "aiohttp-retry>=2.8.3",
     "httpx>=0.24",
-    "miniopy-async>=1.20.1",
+    "aioboto3>=14.1.0",
     "tinydb>=4.7",
     "google-auth>=2.17",
     "typer>=0.12.1",
@@ -47,7 +47,6 @@ dependencies = [
     "importlib_metadata; python_version <= '3.9'",
     "typing_extensions; python_version <= '3.10'", # compatible versions controlled through pydantic
     "rich>=13.0.0",  # databinder
-    "aiofile",  # compatible versions controlled through miniopy-async
     "make-it-sync",  # compatible versions controlled through func_adl
     "ruamel.yaml>=0.18.7",
     "filelock>=3.12.0",
@@ -80,6 +79,7 @@ test = [
     "pandas>=2.0.2",
     "pyarrow>=12.0.0",
     "pre-commit>=4.0.1",
+    "pytest-aioboto3>=0.6.0",
 ]
 docs = [
     "sphinx>=7.0.1, <8.2.0",
diff --git a/servicex/minio_adapter.py b/servicex/minio_adapter.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2022, IRIS-HEP
+# Copyright (c) 2022-2025, IRIS-HEP
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
@@ -30,13 +30,21 @@
 from pathlib import Path
 from typing import List
 
-import aiohttp
+import asyncio
 from tenacity import retry, stop_after_attempt, wait_random_exponential
 
-from miniopy_async import Minio
+import aioboto3
 
 from servicex.models import ResultFile, TransformStatus
 
+_semaphore = asyncio.Semaphore(8)
+
+
+def init_download_semaphore(concurrency: int = 8):
+    "Update the number of concurrent connections"
+    global _semaphore
+    _semaphore = asyncio.Semaphore(concurrency)
+
 
 def _sanitize_filename(fname: str):
     "No matter the string given, make it an acceptable filename on all platforms"
@@ -56,10 +64,11 @@ def __init__(
         secret_key: str,
         bucket: str,
     ):
-        self.minio = Minio(
-            endpoint_host, access_key=access_key, secret_key=secret_key, secure=secure
+        self.minio = aioboto3.Session(
+            aws_access_key_id=access_key, aws_secret_access_key=secret_key
         )
 
+        self.endpoint_host = ("https://" if secure else "http://") + endpoint_host
         self.bucket = bucket
 
     @classmethod
@@ -76,16 +85,18 @@ def for_transform(cls, transform: TransformStatus):
         stop=stop_after_attempt(3), wait=wait_random_exponential(max=60), reraise=True
     )
     async def list_bucket(self) -> List[ResultFile]:
-        objects = await self.minio.list_objects(self.bucket)
-        return [
-            ResultFile(
-                filename=obj.object_name,
-                size=obj.size,
-                extension=obj.object_name.split(".")[-1],
-            )
-            for obj in objects
-            if not obj.is_dir
-        ]
+        async with self.minio.resource("s3", endpoint_url=self.endpoint_host) as s3:
+            bucket = await s3.Bucket(self.bucket)
+            objects = bucket.objects.all()
+            return [
+                ResultFile(
+                    filename=obj.key,
+                    size=await obj.size,
+                    extension=obj.key.split(".")[-1],
+                )
+                async for obj in objects
+                if not obj.key.endswith("/")
+            ]
 
     @retry(
         stop=stop_after_attempt(3), wait=wait_random_exponential(max=60), reraise=True
@@ -103,24 +114,31 @@ async def download_file(
             )
         )
 
-        async with aiohttp.ClientSession(
-            timeout=aiohttp.ClientTimeout(total=10 * 60)
-        ) as session:
-            _ = await self.minio.fget_object(
-                bucket_name=self.bucket,
-                object_name=object_name,
-                file_path=path.as_posix(),
-                session=session,
-            )
+        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())
+                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:
-        return await self.minio.get_presigned_url(
-            bucket_name=self.bucket, object_name=object_name, method="GET"
-        )
+        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}
+                )
 
     @classmethod
     def hash_path(cls, file_name):
diff --git a/servicex/query_core.py b/servicex/query_core.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2024, IRIS-HEP
+# Copyright (c) 2024-2025, IRIS-HEP
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
@@ -204,14 +204,14 @@ def transform_complete(task: Task):
                     f'ServiceX Exception for request ID {self.request_id} ({self.title})"',
                     exc_info=task.exception(),
                 )
-                self.cache.delete_record_by_request_id(self.request_id)
-                if download_files_task:
-                    download_files_task.cancel("Transform failed")
+                if self.fail_if_incomplete:
+                    self.cache.delete_record_by_request_id(self.request_id)
+                    if download_files_task:
+                        download_files_task.cancel("Transform failed")
                 raise task.exception()
 
             if self.current_status.status in DONE_STATUS:
                 if self.current_status.files_failed:
-                    self.cache.delete_record_by_request_id(self.request_id)
                     titlestr = (
                         f'"{self.current_status.title}" '
                         if self.current_status.title is not None
@@ -220,7 +220,8 @@ def transform_complete(task: Task):
                     errorstr = (
                         f"Transform {titlestr}completed with failures: "
                         f"{self.current_status.files_failed}/"
-                        f"{self.current_status.files} files failed. Will not cache."
+                        f"{self.current_status.files} files failed."
+                        f"{'Will not cache.' if self.fail_if_incomplete else ''}"
                     )
                     failedfiles = (
                         self.servicex.url
@@ -231,6 +232,7 @@ def transform_complete(task: Task):
                         "A list of failed files is at [bold red on white]"
                         f"[link={failedfiles}]this link[/link][/bold red on white]"
                     )
+                    logger.error(errorstr)
                     logger.error(errorstr2)
                     logger.error(
                         f"Transform Request id: {self.current_status.request_id}"
@@ -246,7 +248,10 @@ def transform_complete(task: Task):
                             f"More information of '{self.title}' [bold red on white][link={kibana_link}]HERE[/link][/bold red on white]"  # NOQA: E501
                         )
                     if self.fail_if_incomplete:
+                        self.cache.delete_record_by_request_id(self.request_id)
                         raise ServiceXException(errorstr)
+                    else:
+                        logger.error("Will continue to download what is available")
                 else:
                     logger.info("Transforms completed successfully")
             else:  # pragma: no cover
diff --git a/servicex/servicex_client.py b/servicex/servicex_client.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2022, IRIS-HEP
+# Copyright (c) 2022-2025, IRIS-HEP
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
@@ -228,6 +228,7 @@ def deliver(
     fail_if_incomplete: bool = True,
     ignore_local_cache: bool = False,
     progress_bar: ProgressBarFormat = ProgressBarFormat.default,
+    concurrency: int = 8,
 ):
     r"""
     Execute a ServiceX query.
@@ -249,9 +250,13 @@ def deliver(
             will have its own progress bars; :py:const:`ProgressBarFormat.compact` gives one
             summary progress bar for all transformations; :py:const:`ProgressBarFormat.none`
             switches off progress bars completely.
+    :param concurrency: specify how many downloads to run in parallel (default is 8).
     :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
+
+    init_download_semaphore(concurrency)
     config = _load_ServiceXSpec(spec)
 
     if ignore_local_cache or config.General.IgnoreLocalCache:
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -222,3 +222,9 @@ def codegen_list():
         "uproot": "http://servicex-code-gen-uproot:8000",
         "uproot-raw": "http://servicex-code-gen-uproot-raw:8000",
     }
+
+
+# This exists to force an async event loop to be created
+@fixture
+async def with_event_loop():
+    pass
diff --git a/tests/test_databinder.py b/tests/test_databinder.py
@@ -286,7 +286,7 @@ def test_invalid_dataset_identifier():
         )
 
 
-def test_submit_mapping(transformed_result, codegen_list):
+def test_submit_mapping(transformed_result, codegen_list, with_event_loop):
     from servicex import deliver
 
     spec = {
@@ -316,7 +316,9 @@ def test_submit_mapping(transformed_result, codegen_list):
         assert list(results["sampleA"]) == ["1.parquet"]
 
 
-def test_submit_mapping_signed_urls(transformed_result_signed_url, codegen_list):
+def test_submit_mapping_signed_urls(
+    transformed_result_signed_url, codegen_list, with_event_loop
+):
     from servicex import deliver
 
     spec = {
@@ -347,7 +349,7 @@ def test_submit_mapping_signed_urls(transformed_result_signed_url, codegen_list)
         ]
 
 
-def test_submit_mapping_failure(transformed_result, codegen_list):
+def test_submit_mapping_failure(transformed_result, codegen_list, with_event_loop):
     from servicex import deliver
 
     spec = {
@@ -378,7 +380,7 @@ def test_submit_mapping_failure(transformed_result, codegen_list):
                 pass
 
 
-def test_submit_mapping_failure_signed_urls(codegen_list):
+def test_submit_mapping_failure_signed_urls(codegen_list, with_event_loop):
     from servicex import deliver
 
     spec = {
@@ -634,7 +636,7 @@ def run_query(input_filenames=None):
         _load_ServiceXSpec(path2)
 
 
-def test_funcadl_query(transformed_result, codegen_list):
+def test_funcadl_query(transformed_result, codegen_list, with_event_loop):
     from servicex import deliver
     from servicex.query import FuncADL_Uproot  # type: ignore
 
@@ -664,7 +666,7 @@ def test_funcadl_query(transformed_result, codegen_list):
         deliver(spec, config_path="tests/example_config.yaml")
 
 
-def test_query_with_codegen_override(transformed_result, codegen_list):
+def test_query_with_codegen_override(transformed_result, codegen_list, with_event_loop):
     from servicex import deliver
     from servicex.query import FuncADL_Uproot  # type: ignore
 
@@ -748,7 +750,7 @@ def test_databinder_load_dict():
     )
 
 
-def test_python_query(transformed_result, codegen_list):
+def test_python_query(transformed_result, codegen_list, with_event_loop):
     from servicex import deliver
     from servicex.query import PythonFunction  # type: ignore
 
@@ -782,7 +784,7 @@ def run_query(input_filenames=None):
         deliver(spec, config_path="tests/example_config.yaml")
 
 
-def test_uproot_raw_query(transformed_result, codegen_list):
+def test_uproot_raw_query(transformed_result, codegen_list, with_event_loop):
     from servicex import deliver
     from servicex.query import UprootRaw  # type: ignore
 
@@ -810,7 +812,7 @@ def test_uproot_raw_query(transformed_result, codegen_list):
         deliver(spec, config_path="tests/example_config.yaml")
 
 
-def test_uproot_raw_query_parquet(transformed_result, codegen_list):
+def test_uproot_raw_query_parquet(transformed_result, codegen_list, with_event_loop):
     from servicex import deliver
     from servicex.query import UprootRaw  # type: ignore
 
@@ -897,7 +899,7 @@ def test_generic_query(codegen_list):
             )
 
 
-def test_deliver_progress_options(transformed_result, codegen_list):
+def test_deliver_progress_options(transformed_result, codegen_list, with_event_loop):
     from servicex import deliver, ProgressBarFormat
     from servicex.query import UprootRaw  # type: ignore
 
diff --git a/tests/test_minio_adapter.py b/tests/test_minio_adapter.py
