Improve async writing orchestration (multithreading).

Author: igorborgest

awswrangler/_utils.py

@@ -7,7 +7,8 @@
 import os
 import random
 import time
-from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union, cast
+from concurrent.futures import FIRST_COMPLETED, Future, wait
+from typing import Any, Callable, Dict, Generator, List, Optional, Sequence, Tuple, Union, cast
 
 import boto3
 import botocore.config
@@ -290,3 +291,21 @@ def get_even_chunks_sizes(total_size: int, chunk_size: int, upper_bound: bool) -
         i_cycled: int = i % len(sizes)
         sizes[i_cycled] += 1
     return tuple(sizes)
+
+
+def get_running_futures(seq: Sequence[Future]) -> Tuple[Future, ...]:  # type: ignore
+    """Filter only running futures."""
+    return tuple(f for f in seq if f.running())
+
+
+def wait_any_future_available(seq: Sequence[Future]) -> None:  # type: ignore
+    """Wait until any future became available."""
+    wait(fs=seq, timeout=None, return_when=FIRST_COMPLETED)
+
+
+def block_waiting_available_thread(seq: Sequence[Future], max_workers: int) -> None:  # type: ignore
+    """Block until any thread became available."""
+    running: Tuple[Future, ...] = get_running_futures(seq=seq)  # type: ignore
+    while len(running) >= max_workers:
+        wait_any_future_available(seq=running)
+        running = get_running_futures(seq=running)

awswrangler/s3/_fs.py

@@ -86,9 +86,9 @@ def __init__(self, use_threads: bool):
         self.closed = False
         self._exec: Optional[concurrent.futures.ThreadPoolExecutor]
         self._results: List[Dict[str, Union[str, int]]] = []
-        cpus: int = _utils.ensure_cpu_count(use_threads=use_threads)
-        if cpus > 1:
-            self._exec = concurrent.futures.ThreadPoolExecutor(max_workers=cpus)
+        self._cpus: int = _utils.ensure_cpu_count(use_threads=use_threads)
+        if self._cpus > 1:
+            self._exec = concurrent.futures.ThreadPoolExecutor(max_workers=self._cpus)
             self._futures: List[Any] = []
         else:
             self._exec = None
@@ -137,6 +137,7 @@ def upload(
     ) -> None:
         """Upload Part."""
         if self._exec is not None:
+            _utils.block_waiting_available_thread(seq=self._futures, max_workers=self._cpus)
             future = self._exec.submit(
                 _UploadProxy._caller,
                 bucket=bucket,

awswrangler/s3/_write_concurrent.py

@@ -16,9 +16,9 @@ class _WriteProxy:
     def __init__(self, use_threads: bool):
         self._exec: Optional[concurrent.futures.ThreadPoolExecutor]
         self._results: List[str] = []
-        cpus: int = _utils.ensure_cpu_count(use_threads=use_threads)
-        if cpus > 1:
-            self._exec = concurrent.futures.ThreadPoolExecutor(max_workers=cpus)
+        self._cpus: int = _utils.ensure_cpu_count(use_threads=use_threads)
+        if self._cpus > 1:
+            self._exec = concurrent.futures.ThreadPoolExecutor(max_workers=self._cpus)
             self._futures: List[Any] = []
         else:
             self._exec = None
@@ -35,6 +35,7 @@ def _caller(
     def write(self, func: Callable[..., List[str]], boto3_session: boto3.Session, **func_kwargs: Any) -> None:
         """Write File."""
         if self._exec is not None:
+            _utils.block_waiting_available_thread(seq=self._futures, max_workers=self._cpus)
             _logger.debug("Submitting: %s", func)
             future = self._exec.submit(
                 _WriteProxy._caller,