perf: download products and assets in parallel (#1890)

Author: anesson-cs

docs/notebooks/api_user_guide/7_download.ipynb

@@ -87,6 +87,7 @@
 from eodag.utils.stac_reader import fetch_stac_items
 
 if TYPE_CHECKING:
+    from concurrent.futures import ThreadPoolExecutor
     from shapely.geometry.base import BaseGeometry
 
     from eodag.api.product import EOProduct
@@ -1919,6 +1920,7 @@ def download_all(
         search_result: SearchResult,
         downloaded_callback: Optional[DownloadedCallback] = None,
         progress_callback: Optional[ProgressCallback] = None,
+        executor: Optional[ThreadPoolExecutor] = None,
         wait: float = DEFAULT_DOWNLOAD_WAIT,
         timeout: float = DEFAULT_DOWNLOAD_TIMEOUT,
         **kwargs: Unpack[DownloadConf],
@@ -1936,6 +1938,8 @@ def download_all(
                                   size as inputs and handle progress bar
                                   creation and update to give the user a
                                   feedback on the download progress
+        :param executor: (optional) An executor to download EO products of ``search_result`` in parallel
+                                    which will also be reused to download assets of these products in parallel.
         :param wait: (optional) If download fails, wait time in minutes between
                      two download tries of the same product
         :param timeout: (optional) If download fails, maximum time in minutes
@@ -1956,15 +1960,15 @@ def download_all(
         paths = []
         if search_result:
             logger.info("Downloading %s products", len(search_result))
-            # Get download plugin using first product assuming product from several provider
-            # aren't mixed into a search result
+            # Get download plugin using first product assuming all plugins use base.Download.download_all
             download_plugin = self._plugins_manager.get_download_plugin(
                 search_result[0]
             )
             paths = download_plugin.download_all(
                 search_result,
                 downloaded_callback=downloaded_callback,
                 progress_callback=progress_callback,
+                executor=executor,
                 wait=wait,
                 timeout=timeout,
                 **kwargs,
@@ -2026,6 +2030,7 @@ def download(
         self,
         product: EOProduct,
         progress_callback: Optional[ProgressCallback] = None,
+        executor: Optional[ThreadPoolExecutor] = None,
         wait: float = DEFAULT_DOWNLOAD_WAIT,
         timeout: float = DEFAULT_DOWNLOAD_TIMEOUT,
         **kwargs: Unpack[DownloadConf],
@@ -2056,6 +2061,8 @@ def download(
                                   size as inputs and handle progress bar
                                   creation and update to give the user a
                                   feedback on the download progress
+        :param executor: (optional) An executor to download assets of ``product`` in parallel if it has any. If ``None``
+                         , a default executor will be created
         :param wait: (optional) If download fails, wait time in minutes between
                     two download tries
         :param timeout: (optional) If download fails, maximum time in minutes
@@ -2080,7 +2087,11 @@ def download(
             return uri_to_path(product.location)
         self._setup_downloader(product)
         path = product.download(
-            progress_callback=progress_callback, wait=wait, timeout=timeout, **kwargs
+            progress_callback=progress_callback,
+            executor=executor,
+            wait=wait,
+            timeout=timeout,
+            **kwargs,
         )
 
         return path

eodag/api/core.py

@@ -61,6 +61,7 @@
 from eodag.utils.repr import dict_to_html_table
 
 if TYPE_CHECKING:
+    from concurrent.futures import ThreadPoolExecutor
     from shapely.geometry.base import BaseGeometry
 
     from eodag.api.product.drivers.base import DatasetDriver
@@ -122,6 +123,8 @@ class EOProduct:
     search_kwargs: Any
     #: Datetime for download next try
     next_try: datetime
+    #: Stream for requests
+    _stream: requests.Response
 
     def __init__(
         self, provider: str, properties: dict[str, Any], **kwargs: Any
@@ -337,6 +340,7 @@ def register_downloader(
     def download(
         self,
         progress_callback: Optional[ProgressCallback] = None,
+        executor: Optional[ThreadPoolExecutor] = None,
         wait: float = DEFAULT_DOWNLOAD_WAIT,
         timeout: float = DEFAULT_DOWNLOAD_TIMEOUT,
         **kwargs: Unpack[DownloadConf],
@@ -353,6 +357,8 @@ def download(
                                   size as inputs and handle progress bar
                                   creation and update to give the user a
                                   feedback on the download progress
+        :param executor: (optional) An executor to download assets of the product in parallel if it has any. If ``None``
+                         , a default executor will be created
         :param wait: (optional) If download fails, wait time in minutes between
                      two download tries
         :param timeout: (optional) If download fails, maximum time in minutes
@@ -377,17 +383,26 @@ def download(
         )
 
         progress_callback, close_progress_callback = self._init_progress_bar(
-            progress_callback
+            progress_callback, executor
         )
+
         fs_path = self.downloader.download(
             self,
             auth=auth,
             progress_callback=progress_callback,
+            executor=executor,
             wait=wait,
             timeout=timeout,
             **kwargs,
         )
 
+        # shutdown executor if it was not created during parallel product downloads
+        if (
+            executor is not None
+            and executor._thread_name_prefix != "eodag-download-all"
+        ):
+            executor.shutdown(wait=True)
+
         # close progress bar if needed
         if close_progress_callback:
             progress_callback.close()
@@ -408,15 +423,22 @@ def download(
         return fs_path
 
     def _init_progress_bar(
-        self, progress_callback: Optional[ProgressCallback]
+        self,
+        progress_callback: Optional[ProgressCallback],
+        executor: Optional[ThreadPoolExecutor],
     ) -> tuple[ProgressCallback, bool]:
+        # determine position of the progress bar with a counter of executor passings
+        # to avoid bar overwriting in case of parallel downloads
+        count = executor._counter() if executor is not None else 1  # type: ignore
+
         # progress bar init
         if progress_callback is None:
-            progress_callback = ProgressCallback(position=1)
+            progress_callback = ProgressCallback(position=count)
             # one shot progress callback to close after download
             close_progress_callback = True
         else:
             close_progress_callback = False
+            progress_callback.pos = count
             # update units as bar may have been previously used for extraction
             progress_callback.unit = "B"
             progress_callback.unit_scale = True

eodag/api/product/_product.py

@@ -48,6 +48,7 @@
 from urllib.parse import parse_qs
 
 import click
+from concurrent.futures import ThreadPoolExecutor
 
 from eodag.api.collection import CollectionsList
 from eodag.api.core import EODataAccessGateway, SearchResult
@@ -556,6 +557,11 @@ def discover_col(ctx: Context, **kwargs: Any) -> None:
     type=click.Path(dir_okay=True, file_okay=False),
     help="Products or quicklooks download directory (Default: local temporary directory)",
 )
+@click.option(
+    "--max-workers",
+    type=int,
+    help="The maximum number of workers to use for downloading products and assets in parallel",
+)
 @click.pass_context
 def download(ctx: Context, **kwargs: Any) -> None:
     """Download a bunch of products from a serialized search result"""
@@ -601,7 +607,10 @@ def download(ctx: Context, **kwargs: Any) -> None:
 
     else:
         # Download products
-        downloaded_files = satim_api.download_all(search_results, output_dir=output_dir)
+        executor = ThreadPoolExecutor(max_workers=kwargs.pop("max_workers"))
+        downloaded_files = satim_api.download_all(
+            search_results, output_dir=output_dir, executor=executor
+        )
         if downloaded_files and len(downloaded_files) > 0:
             for downloaded_file in downloaded_files:
                 if downloaded_file is None:

eodag/cli.py

@@ -48,14 +48,15 @@
 if TYPE_CHECKING:
     from typing import Any, Optional, Union
 
+    from concurrent.futures import ThreadPoolExecutor
     from mypy_boto3_s3 import S3ServiceResource
     from requests.auth import AuthBase
 
     from eodag.api.product import EOProduct
     from eodag.api.search_result import SearchResult
     from eodag.config import PluginConfig
     from eodag.types.download_args import DownloadConf
-    from eodag.utils import DownloadedCallback, ProgressCallback, Unpack
+    from eodag.utils import ProgressCallback, Unpack
 
 
 logger = logging.getLogger("eodag.apis.ecmwf")
@@ -185,6 +186,7 @@ def download(
         product: EOProduct,
         auth: Optional[Union[AuthBase, S3ServiceResource]] = None,
         progress_callback: Optional[ProgressCallback] = None,
+        executor: Optional[ThreadPoolExecutor] = None,
         wait: float = DEFAULT_DOWNLOAD_WAIT,
         timeout: float = DEFAULT_DOWNLOAD_TIMEOUT,
         **kwargs: Unpack[DownloadConf],
@@ -269,29 +271,6 @@ def download(
         product.location = path_to_uri(product_path)
         return product_path
 
-    def download_all(
-        self,
-        products: SearchResult,
-        auth: Optional[Union[AuthBase, S3ServiceResource]] = None,
-        downloaded_callback: Optional[DownloadedCallback] = None,
-        progress_callback: Optional[ProgressCallback] = None,
-        wait: float = DEFAULT_DOWNLOAD_WAIT,
-        timeout: float = DEFAULT_DOWNLOAD_TIMEOUT,
-        **kwargs: Unpack[DownloadConf],
-    ) -> list[str]:
-        """
-        Download all using parent (base plugin) method
-        """
-        return super(EcmwfApi, self).download_all(
-            products,
-            auth=auth,
-            downloaded_callback=downloaded_callback,
-            progress_callback=progress_callback,
-            wait=wait,
-            timeout=timeout,
-            **kwargs,
-        )
-
     def clear(self) -> None:
         """Clear search context"""
         pass

eodag/plugins/apis/ecmwf.py

@@ -57,12 +57,13 @@
 )
 
 if TYPE_CHECKING:
+    from concurrent.futures import ThreadPoolExecutor
     from mypy_boto3_s3 import S3ServiceResource
     from requests.auth import AuthBase
 
     from eodag.config import PluginConfig
     from eodag.types.download_args import DownloadConf
-    from eodag.utils import DownloadedCallback, Unpack
+    from eodag.utils import Unpack
 
 logger = logging.getLogger("eodag.apis.usgs")
 
@@ -312,6 +313,7 @@ def download(
         product: EOProduct,
         auth: Optional[Union[AuthBase, S3ServiceResource]] = None,
         progress_callback: Optional[ProgressCallback] = None,
+        executor: Optional[ThreadPoolExecutor] = None,
         wait: float = DEFAULT_DOWNLOAD_WAIT,
         timeout: float = DEFAULT_DOWNLOAD_TIMEOUT,
         **kwargs: Unpack[DownloadConf],
@@ -477,26 +479,3 @@ def download_request(
             shutil.move(fs_path, new_fs_path)
             product.location = path_to_uri(new_fs_path)
             return new_fs_path
-
-    def download_all(
-        self,
-        products: SearchResult,
-        auth: Optional[Union[AuthBase, S3ServiceResource]] = None,
-        downloaded_callback: Optional[DownloadedCallback] = None,
-        progress_callback: Optional[ProgressCallback] = None,
-        wait: float = DEFAULT_DOWNLOAD_WAIT,
-        timeout: float = DEFAULT_DOWNLOAD_TIMEOUT,
-        **kwargs: Unpack[DownloadConf],
-    ) -> list[str]:
-        """
-        Download all using parent (base plugin) method
-        """
-        return super(UsgsApi, self).download_all(
-            products,
-            auth=auth,
-            downloaded_callback=downloaded_callback,
-            progress_callback=progress_callback,
-            wait=wait,
-            timeout=timeout,
-            **kwargs,
-        )