Add threaded parallel_map

.vscode/settings.json

@@ -48,6 +48,7 @@
     "tickvals",
     "tinydb",
     "tqdm",
+    "unchunk",
     "uvicorn",
     "Vectorizer",
     "xmargin",

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = great-ai
-version = 0.0.6
+version = 0.0.7
 author = András Schmelczer
 author_email = andras@scoutinscience.com
 description = 

src/great_ai/utilities/__init__.py

@@ -6,5 +6,6 @@ from .get_sentences import get_sentences
 from .language import english_name_of_language, is_english, predict_language
 from .logger import get_logger
 from .match_names import match_names
-from .parallel_map import parallel_map
+from .parallel_map import parallel_map, threaded_parallel_map
+from .unchunk import unchunk
 from .unique import unique

src/great_ai/utilities/parallel_map/__init__.py

@@ -1 +1,2 @@
 from .parallel_map import parallel_map
+from .threaded_parallel_map import threaded_parallel_map

src/great_ai/utilities/parallel_map/mapper_function.py

@@ -1,13 +1,15 @@
 import multiprocessing as mp
 import queue
+import threading
+from typing import Union
 
 import dill
 
 
 def mapper_function(
-    input_queue: mp.Queue,
+    input_queue: Union[mp.Queue, queue.Queue],
-    output_queue: mp.Queue,
+    output_queue: Union[mp.Queue, queue.Queue],
-    should_stop: mp.Event,
+    should_stop: Union[mp.Event, threading.Event],
     serialized_map_function: bytes,
 ):
     map_function = dill.loads(serialized_map_function)

src/great_ai/utilities/parallel_map/parallel_map.py

@@ -128,7 +128,7 @@ def parallel_map(
                 pass
 
         should_stop.set()
-    except KeyboardInterrupt:
+    except Exception:
         for p in processes:
             p.terminate()
     finally:

src/great_ai/utilities/parallel_map/threaded_parallel_map.py

@@ -0,0 +1,132 @@
+import queue
+import threading
+from typing import Callable, Dict, Iterable, Optional, Sequence, TypeVar, overload
+
+from tqdm.cli 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 threaded_parallel_map(
+    function: Callable[[T], V],
+    input_values: Sequence[T],
+    *,
+    chunk_size: Optional[int],
+    concurrency: Optional[int],
+    disable_logging: bool,
+    unordered: Optional[bool],
+) -> Iterable[V]:
+    ...
+
+
+@overload
+def threaded_parallel_map(
+    function: Callable[[T], V],
+    input_values: Iterable[T],
+    *,
+    chunk_size: int,
+    concurrency: Optional[int],
+    disable_logging: bool,
+    unordered: Optional[bool],
+) -> Iterable[V]:
+    ...
+
+
+def threaded_parallel_map(
+    function,
+    input_values,
+    *,
+    chunk_size=None,
+    concurrency=None,
+    disable_logging=False,
+    unordered=False,
+):
+    config = get_config(
+        function=function,
+        input_values=input_values,
+        chunk_size=chunk_size,
+        concurrency=concurrency,
+        disable_logging=disable_logging,
+    )
+
+    tqdm_options = dict(
+        desc=f"Threaded parallel map {config.function_name}",
+        disable=disable_logging,
+        total=config.input_length,
+        miniters=1,
+        dynamic_ncols=True,
+    )
+
+    if config.concurrency == 1:
+        yield from (function(v) for v in tqdm(input_values, **tqdm_options))
+        return
+
+    input_queue = queue.Queue(0 if config.chunk_count is None else config.chunk_count)
+    output_queue = queue.Queue(0 if config.chunk_count is None else config.chunk_count)
+    should_stop = threading.Event()
+
+    threads = [
+        threading.Thread(
+            name=f"threaded_parallel_map_{config.function_name}_{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 t in threads:
+        t.start()
+
+    progress = tqdm(**tqdm_options)
+
+    chunks = iter(chunk(enumerate(input_values), chunk_size=config.chunk_size))
+    indexed_results: Dict[int, V] = {}
+    next_output_index = 0
+    read_input_length = 0
+    is_iteration_over = False
+    try:
+        while not is_iteration_over or next_output_index < 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()
+                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
+
+    finally:
+        should_stop.set()
+        for t in threads:
+            t.join()
+
+        progress.close()

src/great_ai/utilities/unchunk.py

@@ -0,0 +1,8 @@
+from typing import Iterable, TypeVar
+
+T = TypeVar("T")
+
+
+def unchunk(chunks: Iterable[Iterable[T]]) -> Iterable[T]:
+    for chunk in chunks:
+        yield from chunk

tests/utilities/test_threaded_parallel_map.py

@@ -0,0 +1,74 @@
+import unittest
+
+import pytest
+
+from src.great_ai.utilities import threaded_parallel_map
+
+COUNT = int(1e5) + 3
+
+
+class TestThreadedParallelMap(unittest.TestCase):
+    def test_simple_case_with_progress_bar(self) -> None:
+        inputs = range(COUNT)
+        expected = [v**2 for v in range(COUNT)]
+
+        assert (
+            list(threaded_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 (
+            list(threaded_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)]
+
+        assert (
+            list(threaded_parallel_map(lambda v: v**2, inputs, disable_logging=True))
+            == expected
+        )
+
+    def test_simple_case_invalid_values(self) -> None:
+        inputs = range(COUNT)
+
+        with pytest.raises(AssertionError):
+            list(threaded_parallel_map(lambda v: v**2, inputs, concurrency=0))
+
+        with pytest.raises(AssertionError):
+            list(threaded_parallel_map(lambda v: v**2, inputs, chunk_size=0))
+
+    def test_no_op(self) -> None:
+        assert (
+            list(threaded_parallel_map(lambda v: v**2, [], disable_logging=True))
+            == []
+        )
+
+        assert (
+            list(
+                threaded_parallel_map(
+                    lambda v: v**2, [], disable_logging=True, chunk_size=100
+                )
+            )
+            == []
+        )
+        assert (
+            list(
+                threaded_parallel_map(
+                    lambda v: v**2, [], disable_logging=True, concurrency=100
+                )
+            )
+            == []
+        )