From 9b4da9328e4a6d2743f36361fa6c4e846bc88890 Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Tue, 28 Jun 2022 18:27:07 +0200 Subject: [PATCH] Improve parallel map --- .vscode/settings.json | 1 + setup.cfg | 8 +- src/great_ai/utilities/parallel_map.py | 52 ------- .../utilities/parallel_map/__init__.py | 1 + .../utilities/parallel_map/get_config.py | 62 ++++++++ .../utilities/parallel_map/mapper_function.py | 23 +++ .../utilities/parallel_map/parallel_map.py | 140 ++++++++++++++++++ .../parallel_map_configuration.py | 25 ++++ tests/utilities/test_parallel_map.py | 30 +++- 9 files changed, 280 insertions(+), 62 deletions(-) delete mode 100644 src/great_ai/utilities/parallel_map.py create mode 100644 src/great_ai/utilities/parallel_map/__init__.py create mode 100644 src/great_ai/utilities/parallel_map/get_config.py create mode 100644 src/great_ai/utilities/parallel_map/mapper_function.py create mode 100644 src/great_ai/utilities/parallel_map/parallel_map.py create mode 100644 src/great_ai/utilities/parallel_map/parallel_map_configuration.py diff --git a/.vscode/settings.json b/.vscode/settings.json index b9d0b65..57681f1 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -23,6 +23,7 @@ "levelname", "levelno", "matplotlib", + "miniters", "Multinomial", "multiprocess", "nbconvert", diff --git a/setup.cfg b/setup.cfg index d1cb178..9b422a9 100644 --- a/setup.cfg +++ b/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 diff --git a/src/great_ai/utilities/parallel_map.py b/src/great_ai/utilities/parallel_map.py deleted file mode 100644 index 14fa621..0000000 --- a/src/great_ai/utilities/parallel_map.py +++ /dev/null @@ -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)) - ) diff --git a/src/great_ai/utilities/parallel_map/__init__.py b/src/great_ai/utilities/parallel_map/__init__.py new file mode 100644 index 0000000..e19e313 --- /dev/null +++ b/src/great_ai/utilities/parallel_map/__init__.py @@ -0,0 +1 @@ +from .parallel_map import parallel_map diff --git a/src/great_ai/utilities/parallel_map/get_config.py b/src/great_ai/utilities/parallel_map/get_config.py new file mode 100644 index 0000000..49e6242 --- /dev/null +++ b/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 diff --git a/src/great_ai/utilities/parallel_map/mapper_function.py b/src/great_ai/utilities/parallel_map/mapper_function.py new file mode 100644 index 0000000..4f51b51 --- /dev/null +++ b/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 diff --git a/src/great_ai/utilities/parallel_map/parallel_map.py b/src/great_ai/utilities/parallel_map/parallel_map.py new file mode 100644 index 0000000..967050f --- /dev/null +++ b/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 diff --git a/src/great_ai/utilities/parallel_map/parallel_map_configuration.py b/src/great_ai/utilities/parallel_map/parallel_map_configuration.py new file mode 100644 index 0000000..8890173 --- /dev/null +++ b/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" + ) diff --git a/tests/utilities/test_parallel_map.py b/tests/utilities/test_parallel_map.py index 0a363df..2077f19 100644 --- a/tests/utilities/test_parallel_map.py +++ b/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 + ), [], )