diff --git a/.vscode/settings.json b/.vscode/settings.json index 57681f1..ff78a63 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -48,6 +48,7 @@ "tickvals", "tinydb", "tqdm", + "unchunk", "uvicorn", "Vectorizer", "xmargin", diff --git a/setup.cfg b/setup.cfg index 8236231..932452e 100644 --- a/setup.cfg +++ b/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 = diff --git a/src/great_ai/utilities/__init__.py b/src/great_ai/utilities/__init__.py index 1b119e2..5bb9463 100644 --- a/src/great_ai/utilities/__init__.py +++ b/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 diff --git a/src/great_ai/utilities/parallel_map/__init__.py b/src/great_ai/utilities/parallel_map/__init__.py index e19e313..0e7e247 100644 --- a/src/great_ai/utilities/parallel_map/__init__.py +++ b/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 diff --git a/src/great_ai/utilities/parallel_map/mapper_function.py b/src/great_ai/utilities/parallel_map/mapper_function.py index 4f51b51..02526d2 100644 --- a/src/great_ai/utilities/parallel_map/mapper_function.py +++ b/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, - output_queue: mp.Queue, - should_stop: mp.Event, + input_queue: Union[mp.Queue, queue.Queue], + output_queue: Union[mp.Queue, queue.Queue], + should_stop: Union[mp.Event, threading.Event], serialized_map_function: bytes, ): map_function = dill.loads(serialized_map_function) diff --git a/src/great_ai/utilities/parallel_map/parallel_map.py b/src/great_ai/utilities/parallel_map/parallel_map.py index 8b94359..22042c9 100644 --- a/src/great_ai/utilities/parallel_map/parallel_map.py +++ b/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: diff --git a/src/great_ai/utilities/parallel_map/threaded_parallel_map.py b/src/great_ai/utilities/parallel_map/threaded_parallel_map.py new file mode 100644 index 0000000..5696219 --- /dev/null +++ b/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() diff --git a/src/great_ai/utilities/unchunk.py b/src/great_ai/utilities/unchunk.py new file mode 100644 index 0000000..cf09a44 --- /dev/null +++ b/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 diff --git a/tests/utilities/test_threaded_parallel_map.py b/tests/utilities/test_threaded_parallel_map.py new file mode 100644 index 0000000..84a6eef --- /dev/null +++ b/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 + ) + ) + == [] + )