Improve parallel map API

src/great_ai/great_ai/context.py

@@ -103,7 +103,7 @@ def configure(
 
     logger.info("Settings: configured ✅")
     for k, v in get_context().to_flat_dict().items():
-        logger.info(f"🔩 {k}: {v}")
+        logger.info(f"  ⚙️ {k}: {v}")
 
     if not is_production and not disable_se4ml_banner:
         logger.warning(

src/great_ai/utilities/chunk.py

@@ -3,13 +3,13 @@ from typing import Iterable, List, TypeVar
 T = TypeVar("T")
 
 
-def chunk(values: Iterable[T], chunk_length: int) -> Iterable[T]:
+def chunk(values: Iterable[T], chunk_size: int) -> Iterable[T]:
-    assert chunk_length >= 1
+    assert chunk_size >= 1
 
     result: List[T] = []
     for v in values:
         result.append(v)
-        if len(result) == chunk_length:
+        if len(result) == chunk_size:
             yield result
             result = []
 

src/great_ai/utilities/parallel_map/get_config.py

@@ -14,8 +14,9 @@ def get_config(
     *,
     function: Callable,
     input_values: Union[Sequence, Iterable],
-    chunk_length: Optional[int],
+    chunk_size: Optional[int],
     concurrency: Optional[int],
+    disable_logging: bool,
 ) -> ParallelMapConfiguration:
 
     is_input_sequence = hasattr(input_values, "__len__")
@@ -24,18 +25,18 @@ def get_config(
         concurrency = len(os.sched_getaffinity(0))
     assert concurrency >= 1, "At least one mapper process has to be created"
 
-    if chunk_length is None:
+    if chunk_size is None:
         if is_input_sequence:
-            chunk_length = max(1, ceil(len(input_values) / concurrency / 10))
+            chunk_size = max(1, ceil(len(input_values) / concurrency / 10))
         else:
             raise ValueError(
-                "The argument for `values` does not implement `__len__`, therefore, you must provide a `chunk_length`"
+                "The argument for `values` does not implement `__len__`, therefore, you must provide a `chunk_size`"
             )
-    assert chunk_length >= 1, "Chunks have to contain at least one element"
+    assert chunk_size >= 1, "Chunks have to contain at least one element"
 
     chunk_count: Optional[int] = None
     if is_input_sequence:
-        chunk_count = ceil(len(input_values) / chunk_length)
+        chunk_count = ceil(len(input_values) / chunk_size)
         if chunk_count < concurrency:
             logger.warning(
                 f"Limiting concurrency to {chunk_count} because there are only {chunk_count} chunks"
@@ -47,16 +48,17 @@ def get_config(
 
     input_length = len(input_values) if is_input_sequence else None
 
-    serialized_map_function = dill.dumps(function, byref=True, recurse=True)
-
-    logger.info("Parallel map: configured ✅")
     config = ParallelMapConfiguration(
         concurrency=concurrency,
         chunk_count=chunk_count,
-        chunk_length=chunk_length,
+        chunk_size=chunk_size,
         input_length=input_length,
-        serialized_map_function=serialized_map_function,
+        serialized_map_function=dill.dumps(function, byref=True, recurse=True),
+        function_name=function.__name__,
     )
 
+    if not disable_logging:
+        logger.info("Parallel map: configured ✅")
         config.pretty_print()
+
     return config

src/great_ai/utilities/parallel_map/parallel_map.py

@@ -1,17 +1,8 @@
 import multiprocessing as mp
 import queue
-from typing import (
+from typing import Callable, Dict, Iterable, Optional, Sequence, TypeVar, overload
-    Callable,
-    Iterable,
-    List,
-    Optional,
-    Sequence,
-    Tuple,
-    TypeVar,
-    overload,
-)
 
-from tqdm.auto import tqdm
+from tqdm.cli import tqdm
 
 from ..chunk import chunk
 from .get_config import get_config
@@ -26,10 +17,11 @@ def parallel_map(
     function: Callable[[T], V],
     input_values: Sequence[T],
     *,
-    chunk_length: Optional[int],
+    chunk_size: Optional[int],
     concurrency: Optional[int],
-    disable_progress_bar: bool,
+    disable_logging: bool,
-) -> List[V]:
+    unordered: Optional[bool],
+) -> Iterable[V]:
     ...
 
 
@@ -38,10 +30,11 @@ def parallel_map(
     function: Callable[[T], V],
     input_values: Iterable[T],
     *,
-    chunk_length: int,
+    chunk_size: int,
     concurrency: Optional[int],
-    disable_progress_bar: bool,
+    disable_logging: bool,
-) -> List[V]:
+    unordered: Optional[bool],
+) -> Iterable[V]:
     ...
 
 
@@ -49,30 +42,32 @@ def parallel_map(
     function,
     input_values,
     *,
-    chunk_length=None,
+    chunk_size=None,
     concurrency=None,
-    disable_progress_bar=False,
+    disable_logging=False,
+    unordered=False,
 ):
     config = get_config(
         function=function,
         input_values=input_values,
-        chunk_length=chunk_length,
+        chunk_size=chunk_size,
         concurrency=concurrency,
+        disable_logging=disable_logging,
+    )
+
+    tqdm_options = dict(
+        desc=f"Parallel map {config.function_name}",
+        disable=disable_logging,
+        total=config.input_length,
+        miniters=1,
+        dynamic_ncols=True,
     )
 
     if config.concurrency == 1:
-        return [
+        yield from (function(v) for v in tqdm(input_values, **tqdm_options))
-            function(v)
+        return
-            for v in tqdm(
-                input_values,
-                desc="Parallel map",
-                disable=disable_progress_bar,
-                total=config.input_length,
-            )
-        ]
 
-    start_methods = mp.get_all_start_methods()
+    ctx = mp.get_context("spawn")
-    ctx = mp.get_context("fork") if "fork" in start_methods else mp.get_context("spawn")
     ctx.freeze_support()
 
     input_queue = ctx.Queue(0 if config.chunk_count is None else config.chunk_count)
@@ -81,7 +76,7 @@ def parallel_map(
 
     processes = [
         ctx.Process(
-            name=f"parallel_map_{i}",
+            name=f"parallel_map_{config.function_name}_{i}",
             target=mapper_function,
             kwargs=dict(
                 input_queue=input_queue,
@@ -96,18 +91,15 @@ def parallel_map(
     for p in processes:
         p.start()
 
-    progress = tqdm(
+    progress = tqdm(**tqdm_options)
-        desc="Parallel map",
-        disable=disable_progress_bar,
-        total=config.input_length,
-    )
 
-    chunks = iter(chunk(enumerate(input_values), chunk_length=config.chunk_length))
+    chunks = iter(chunk(enumerate(input_values), chunk_size=config.chunk_size))
-    indexed_results: List[Tuple[int, V]] = []
+    indexed_results: Dict[int, V] = {}
+    next_output_index = 0
     read_input_length = 0
     is_iteration_over = False
     try:
-        while not is_iteration_over or len(indexed_results) < read_input_length:
+        while not is_iteration_over or next_output_index < read_input_length:
             if not is_iteration_over:
                 try:
                     next_chunk = next(chunks)
@@ -118,10 +110,23 @@ def parallel_map(
 
             try:
                 result_chunk = output_queue.get_nowait()
-                indexed_results.extend(result_chunk)
                 progress.update(len(result_chunk))
+
+                for index, value in result_chunk:
+                    if unordered:
+                        yield value
+                        next_output_index += 1
+                    else:
+                        indexed_results[index] = value
+
+                if not unordered:
+                    while next_output_index in indexed_results:
+                        yield indexed_results[next_output_index]
+                        del indexed_results[next_output_index]
+                        next_output_index += 1
             except queue.Empty:
                 pass
+
         should_stop.set()
     except KeyboardInterrupt:
         for p in processes:
@@ -134,7 +139,3 @@ def parallel_map(
         output_queue.close()
 
         progress.close()
-
-    results = [v for _, v in sorted(indexed_results)]
-
-    return results

src/great_ai/utilities/parallel_map/parallel_map_configuration.py

@@ -10,13 +10,14 @@ logger = get_logger("parallel_map")
 class ParallelMapConfiguration(BaseModel):
     concurrency: int
     chunk_count: Optional[int]
-    chunk_length: int
+    chunk_size: int
     input_length: Optional[int]
     serialized_map_function: bytes
+    function_name: str
 
     def pretty_print(self, prefix="  ⚙️ "):
         logger.info(f"{prefix} concurrency: {self.concurrency}")
-        logger.info(f"{prefix} chunk length: {self.chunk_length}")
+        logger.info(f"{prefix} chunk size: {self.chunk_size}")
         logger.info(
             f"{prefix} chunk count: {self.chunk_count if self.chunk_count else 'unknown'}"
         )

tests/utilities/test_parallel_map.py

@@ -1,5 +1,7 @@
 import unittest
 
+import pytest
+
 from src.great_ai.utilities import parallel_map
 
 COUNT = int(1e5) + 3
@@ -10,7 +12,7 @@ class TestParallelMap(unittest.TestCase):
         inputs = range(COUNT)
         expected = [v**2 for v in range(COUNT)]
 
-        assert parallel_map(lambda v: v**2, inputs, concurrency=10) == expected
+        assert list(parallel_map(lambda v: v**2, inputs, concurrency=10)) == expected
 
     def test_with_iterable(self) -> None:
         from time import sleep
@@ -23,38 +25,42 @@ class TestParallelMap(unittest.TestCase):
         expected = [v**3 for v in range(10)]
 
         assert (
-            parallel_map(lambda x: x**3, my_generator(), chunk_length=1) == expected
+            list(parallel_map(lambda x: x**3, my_generator(), chunk_size=1))
+            == expected
         )
 
     def test_simple_case_without_progress_bar(self) -> None:
         inputs = range(COUNT)
         expected = [v**2 for v in range(COUNT)]
 
-        self.assertEqual(
+        assert (
-            parallel_map(lambda v: v**2, inputs, disable_progress_bar=True), expected
+            list(parallel_map(lambda v: v**2, inputs, disable_logging=True))
+            == expected
         )
 
     def test_simple_case_invalid_values(self) -> None:
         inputs = range(COUNT)
 
-        self.assertRaises(
+        with pytest.raises(AssertionError):
-            AssertionError, parallel_map, lambda v: v**2, inputs, concurrency=0
+            list(parallel_map(lambda v: v**2, inputs, concurrency=0))
-        )
+
-        self.assertRaises(
+        with pytest.raises(AssertionError):
-            AssertionError, parallel_map, lambda v: v**2, inputs, chunk_length=0
+            list(parallel_map(lambda v: v**2, inputs, chunk_size=0))
-        )
 
     def test_no_op(self) -> None:
-        assert parallel_map(lambda v: v**2, [], disable_progress_bar=True) == []
+        assert list(parallel_map(lambda v: v**2, [], disable_logging=True)) == []
-        self.assertEqual(
+
-            parallel_map(
+        assert (
-                lambda v: v**2, [], disable_progress_bar=True, chunk_length=100
+            list(
-            ),
+                parallel_map(lambda v: v**2, [], disable_logging=True, chunk_size=100)
-            [],
             )
-        self.assertEqual(
+            == []
-            parallel_map(
+        )
-                lambda v: v**2, [], disable_progress_bar=True, concurrency=100
+        assert (
-            ),
+            list(
-            [],
+                parallel_map(
+                    lambda v: v**2, [], disable_logging=True, concurrency=100
+                )
+            )
+            == []
         )