Improve parallel map

.vscode/settings.json

@@ -23,6 +23,7 @@
     "levelname",
     "levelno",
     "matplotlib",
+    "miniters",
     "Multinomial",
     "multiprocess",
     "nbconvert",

setup.cfg

@@ -20,12 +20,11 @@ packages = find:
 include_package_data = True
 python_requires = >=3.8
 install_requires =
-    unidecode >= 1.3.0
-    multiprocess >= 0.70.0.0
-    tqdm >= 4.0.0
     scikit-learn
     matplotlib
     numpy
+    tqdm >= 4.0.0
+    unidecode >= 1.3.0
     syntok >= 1.4.0
     langcodes[data] >= 3.3.0
     langdetect >= 1.0.9
@@ -40,7 +39,8 @@ install_requires =
     uvicorn[standard] >= 0.18.0
     watchdog >= 2.1.0
     pymongo >= 3.0.0
-    aiohttp >= 3.8.0
+    dill >= 0.3.5.0
+    aiohttp[speedups] >= 3.8.0
 
 [options.package_data]
 * = *.conf, *.css

src/great_ai/utilities/parallel_map.py

@@ -1,52 +0,0 @@
-from math import ceil
-from typing import Any, Callable, Iterable, List, Optional
-
-import multiprocess as mp
-from tqdm.cli import tqdm
-
-from .logger import get_logger
-
-logger = get_logger("parallel_map")
-
-
-def parallel_map(
-    function: Callable[[Any], Any],
-    values: Iterable[Any],
-    chunk_size: Optional[int] = None,
-    concurrency: Optional[int] = None,
-    disable_progress: bool = False,
-) -> List[Any]:
-    if concurrency is None:
-        concurrency = mp.cpu_count()
-
-    assert concurrency > 0
-    assert chunk_size is None or chunk_size > 0
-
-    values = list(values)
-
-    if not chunk_size:
-        chunk_size = max(1, ceil(len(values) / concurrency / 10))
-
-    chunk_count = ceil(len(values) / chunk_size)
-    if chunk_count < concurrency:
-        logger.warning(
-            f"Limiting concurrency to {chunk_count} because there are only {chunk_count} chunks"
-        )
-        concurrency = chunk_count
-
-    logger.info(
-        f"Starting parallel map (concurrency: {concurrency}, chunk size: {chunk_size})"
-    )
-
-    if concurrency == 1 or len(values) <= chunk_size:
-        logger.warning("Running in series, there is no reason for parallelism")
-        iterable = values if disable_progress else tqdm(values)
-        return [function(v) for v in iterable]
-
-    with mp.Pool(processes=concurrency) as pool:
-        if disable_progress:
-            return pool.map(function, values, chunksize=chunk_size)
-
-        return list(
-            tqdm(pool.imap(function, values, chunksize=chunk_size), total=len(values))
-        )

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"
+        )

tests/utilities/test_parallel_map.py

@@ -10,14 +10,28 @@ class TestParallelMap(unittest.TestCase):
         inputs = range(COUNT)
         expected = [v**2 for v in range(COUNT)]
 
-        assert parallel_map(lambda v: v**2, inputs) == expected
+        assert parallel_map(lambda v: v**2, inputs, concurrency=10) == expected
+
+    def test_with_iterable(self) -> None:
+        from time import sleep
+
+        def my_generator():
+            for i in range(10):
+                yield i
+                sleep(0.1)
+
+        expected = [v**3 for v in range(10)]
+
+        assert (
+            parallel_map(lambda x: x**3, my_generator(), chunk_length=1) == expected
+        )
 
     def test_simple_case_without_progress_bar(self) -> None:
         inputs = range(COUNT)
         expected = [v**2 for v in range(COUNT)]
 
         self.assertEqual(
-            parallel_map(lambda v: v**2, inputs, disable_progress=True), expected
+            parallel_map(lambda v: v**2, inputs, disable_progress_bar=True), expected
         )
 
     def test_simple_case_invalid_values(self) -> None:
@@ -27,16 +41,20 @@ class TestParallelMap(unittest.TestCase):
             AssertionError, parallel_map, lambda v: v**2, inputs, concurrency=0
         )
         self.assertRaises(
-            AssertionError, parallel_map, lambda v: v**2, inputs, chunk_size=0
+            AssertionError, parallel_map, lambda v: v**2, inputs, chunk_length=0
         )
 
     def test_no_op(self) -> None:
-        assert parallel_map(lambda v: v**2, [], disable_progress=True) == []
+        assert parallel_map(lambda v: v**2, [], disable_progress_bar=True) == []
         self.assertEqual(
-            parallel_map(lambda v: v**2, [], disable_progress=True, chunk_size=100),
+            parallel_map(
+                lambda v: v**2, [], disable_progress_bar=True, chunk_length=100
+            ),
             [],
         )
         self.assertEqual(
-            parallel_map(lambda v: v**2, [], disable_progress=True, concurrency=100),
+            parallel_map(
+                lambda v: v**2, [], disable_progress_bar=True, concurrency=100
+            ),
             [],
         )