Improve parallel map

src/great_ai/utilities/parallel_map/__init__.py

@@ -0,0 +1 @@
+from .parallel_map import parallel_map

src/great_ai/utilities/parallel_map/get_config.py

@@ -0,0 +1,62 @@
+import os
+from math import ceil
+from typing import Callable, Iterable, Optional, Sequence, Union
+
+import dill
+
+from ..logger import get_logger
+from .parallel_map_configuration import ParallelMapConfiguration
+
+logger = get_logger("parallel_map")
+
+
+def get_config(
+    *,
+    function: Callable,
+    input_values: Union[Sequence, Iterable],
+    chunk_length: Optional[int],
+    concurrency: Optional[int],
+) -> ParallelMapConfiguration:
+
+    is_input_sequence = hasattr(input_values, "__len__")
+
+    if concurrency is None:
+        concurrency = len(os.sched_getaffinity(0))
+    assert concurrency >= 1, "At least one mapper process has to be created"
+
+    if chunk_length is None:
+        if is_input_sequence:
+            chunk_length = 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`"
+            )
+    assert chunk_length >= 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)
+        if chunk_count < concurrency:
+            logger.warning(
+                f"Limiting concurrency to {chunk_count} because there are only {chunk_count} chunks"
+            )
+            concurrency = chunk_count
+
+    if concurrency == 1:
+        logger.warning("Running in series, there is no reason for parallelism")
+
+    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,
+        input_length=input_length,
+        serialized_map_function=serialized_map_function,
+    )
+
+    config.pretty_print()
+    return config

src/great_ai/utilities/parallel_map/mapper_function.py

@@ -0,0 +1,23 @@
+import multiprocessing as mp
+import queue
+
+import dill
+
+
+def mapper_function(
+    input_queue: mp.Queue,
+    output_queue: mp.Queue,
+    should_stop: mp.Event,
+    serialized_map_function: bytes,
+):
+    map_function = dill.loads(serialized_map_function)
+    try:
+        while not should_stop.is_set():
+            try:
+                input_chunk = input_queue.get_nowait()
+                output_chunk = [(i, map_function(v)) for i, v in input_chunk]
+                output_queue.put(output_chunk)
+            except queue.Empty:
+                pass
+    except KeyboardInterrupt:
+        return

src/great_ai/utilities/parallel_map/parallel_map.py

@@ -0,0 +1,140 @@
+import multiprocessing as mp
+import queue
+from typing import (
+    Callable,
+    Iterable,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    TypeVar,
+    overload,
+)
+
+from tqdm.auto import tqdm
+
+from ..chunk import chunk
+from .get_config import get_config
+from .mapper_function import mapper_function
+
+T = TypeVar("T")
+V = TypeVar("V")
+
+
+@overload
+def parallel_map(
+    function: Callable[[T], V],
+    input_values: Sequence[T],
+    *,
+    chunk_length: Optional[int],
+    concurrency: Optional[int],
+    disable_progress_bar: bool,
+) -> List[V]:
+    ...
+
+
+@overload
+def parallel_map(
+    function: Callable[[T], V],
+    input_values: Iterable[T],
+    *,
+    chunk_length: int,
+    concurrency: Optional[int],
+    disable_progress_bar: bool,
+) -> List[V]:
+    ...
+
+
+def parallel_map(
+    function,
+    input_values,
+    *,
+    chunk_length=None,
+    concurrency=None,
+    disable_progress_bar=False,
+):
+    config = get_config(
+        function=function,
+        input_values=input_values,
+        chunk_length=chunk_length,
+        concurrency=concurrency,
+    )
+
+    if config.concurrency == 1:
+        return [
+            function(v)
+            for v in tqdm(
+                input_values,
+                desc="Parallel map",
+                disable=disable_progress_bar,
+                total=config.input_length,
+                miniters=1,
+            )
+        ]
+
+    ctx = mp.get_context("spawn")
+    ctx.freeze_support()
+    input_queue = ctx.Queue(0 if config.chunk_count is None else config.chunk_count)
+    output_queue = ctx.Queue(0 if config.chunk_count is None else config.chunk_count)
+    should_stop = ctx.Event()
+
+    processes = [
+        ctx.Process(
+            name=f"parallel_map_{i}",
+            target=mapper_function,
+            kwargs=dict(
+                input_queue=input_queue,
+                output_queue=output_queue,
+                should_stop=should_stop,
+                serialized_map_function=config.serialized_map_function,
+            ),
+        )
+        for i in range(config.concurrency)
+    ]
+
+    for p in processes:
+        p.start()
+
+    progress = tqdm(
+        desc="Parallel map",
+        disable=disable_progress_bar,
+        total=config.input_length,
+        miniters=1,
+    )
+
+    chunks = iter(chunk(enumerate(input_values), chunk_length=config.chunk_length))
+    indexed_results: List[Tuple[int, V]] = []
+    read_input_length = 0
+    is_iteration_over = False
+    try:
+        while not is_iteration_over or len(indexed_results) < read_input_length:
+            if not is_iteration_over:
+                try:
+                    next_chunk = next(chunks)
+                    input_queue.put(next_chunk)
+                    read_input_length += len(next_chunk)
+                except StopIteration:
+                    is_iteration_over = True
+
+            try:
+                result_chunk = output_queue.get_nowait()
+                indexed_results.extend(result_chunk)
+                progress.update(len(result_chunk))
+            except queue.Empty:
+                pass
+        should_stop.set()
+    except KeyboardInterrupt:
+        for p in processes:
+            p.terminate()
+    finally:
+        for p in processes:
+            p.join()
+            p.close()
+        input_queue.close()
+        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

@@ -0,0 +1,25 @@
+from typing import Optional
+
+from pydantic import BaseModel
+
+from ..logger import get_logger
+
+logger = get_logger("parallel_map")
+
+
+class ParallelMapConfiguration(BaseModel):
+    concurrency: int
+    chunk_count: Optional[int]
+    chunk_length: int
+    input_length: Optional[int]
+    serialized_map_function: bytes
+
+    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 count: {self.chunk_count if self.chunk_count else 'unknown'}"
+        )
+        logger.info(
+            f"{prefix} function size: {len(self.serialized_map_function) / 1024:.0f} kB"
+        )