From 154cc52fa6ca46b50dc4ce2388a802957cbd7640 Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Thu, 30 Jun 2022 16:54:01 +0200 Subject: [PATCH] Improve parallel map API --- src/great_ai/great_ai/context.py | 2 +- src/great_ai/utilities/chunk.py | 6 +- .../utilities/parallel_map/get_config.py | 26 +++--- .../utilities/parallel_map/parallel_map.py | 91 ++++++++++--------- .../parallel_map_configuration.py | 5 +- tests/utilities/test_parallel_map.py | 48 +++++----- 6 files changed, 94 insertions(+), 84 deletions(-) diff --git a/src/great_ai/great_ai/context.py b/src/great_ai/great_ai/context.py index 7e8f2e2..c7a4c2f 100644 --- a/src/great_ai/great_ai/context.py +++ b/src/great_ai/great_ai/context.py @@ -103,7 +103,7 @@ def configure( logger.info("Settings: configured ✅") for k, v in get_context().to_flat_dict().items(): - logger.info(f"🔩 {k}: {v}") + logger.info(f" ⚙️ {k}: {v}") if not is_production and not disable_se4ml_banner: logger.warning( diff --git a/src/great_ai/utilities/chunk.py b/src/great_ai/utilities/chunk.py index 05ac663..694407f 100644 --- a/src/great_ai/utilities/chunk.py +++ b/src/great_ai/utilities/chunk.py @@ -3,13 +3,13 @@ from typing import Iterable, List, TypeVar T = TypeVar("T") -def chunk(values: Iterable[T], chunk_length: int) -> Iterable[T]: - assert chunk_length >= 1 +def chunk(values: Iterable[T], chunk_size: int) -> Iterable[T]: + assert chunk_size >= 1 result: List[T] = [] for v in values: result.append(v) - if len(result) == chunk_length: + if len(result) == chunk_size: yield result result = [] diff --git a/src/great_ai/utilities/parallel_map/get_config.py b/src/great_ai/utilities/parallel_map/get_config.py index 49e6242..d7d8f8a 100644 --- a/src/great_ai/utilities/parallel_map/get_config.py +++ b/src/great_ai/utilities/parallel_map/get_config.py @@ -14,8 +14,9 @@ def get_config( *, function: Callable, input_values: Union[Sequence, Iterable], - chunk_length: Optional[int], + chunk_size: Optional[int], concurrency: Optional[int], + disable_logging: bool, ) -> ParallelMapConfiguration: is_input_sequence = hasattr(input_values, "__len__") @@ -24,18 +25,18 @@ def get_config( concurrency = len(os.sched_getaffinity(0)) assert concurrency >= 1, "At least one mapper process has to be created" - if chunk_length is None: + if chunk_size is None: if is_input_sequence: - chunk_length = max(1, ceil(len(input_values) / concurrency / 10)) + chunk_size = 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`" + "The argument for `values` does not implement `__len__`, therefore, you must provide a `chunk_size`" ) - assert chunk_length >= 1, "Chunks have to contain at least one element" + assert chunk_size >= 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) + chunk_count = ceil(len(input_values) / chunk_size) if chunk_count < concurrency: logger.warning( f"Limiting concurrency to {chunk_count} because there are only {chunk_count} chunks" @@ -47,16 +48,17 @@ def get_config( 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, + chunk_size=chunk_size, input_length=input_length, - serialized_map_function=serialized_map_function, + serialized_map_function=dill.dumps(function, byref=True, recurse=True), + function_name=function.__name__, ) - config.pretty_print() + if not disable_logging: + logger.info("Parallel map: configured ✅") + config.pretty_print() + return config diff --git a/src/great_ai/utilities/parallel_map/parallel_map.py b/src/great_ai/utilities/parallel_map/parallel_map.py index a00a176..8b94359 100644 --- a/src/great_ai/utilities/parallel_map/parallel_map.py +++ b/src/great_ai/utilities/parallel_map/parallel_map.py @@ -1,17 +1,8 @@ import multiprocessing as mp import queue -from typing import ( - Callable, - Iterable, - List, - Optional, - Sequence, - Tuple, - TypeVar, - overload, -) +from typing import Callable, Dict, Iterable, Optional, Sequence, TypeVar, overload -from tqdm.auto import tqdm +from tqdm.cli import tqdm from ..chunk import chunk from .get_config import get_config @@ -26,10 +17,11 @@ def parallel_map( function: Callable[[T], V], input_values: Sequence[T], *, - chunk_length: Optional[int], + chunk_size: Optional[int], concurrency: Optional[int], - disable_progress_bar: bool, -) -> List[V]: + disable_logging: bool, + unordered: Optional[bool], +) -> Iterable[V]: ... @@ -38,10 +30,11 @@ def parallel_map( function: Callable[[T], V], input_values: Iterable[T], *, - chunk_length: int, + chunk_size: int, concurrency: Optional[int], - disable_progress_bar: bool, -) -> List[V]: + disable_logging: bool, + unordered: Optional[bool], +) -> Iterable[V]: ... @@ -49,30 +42,32 @@ def parallel_map( function, input_values, *, - chunk_length=None, + chunk_size=None, concurrency=None, - disable_progress_bar=False, + disable_logging=False, + unordered=False, ): config = get_config( function=function, input_values=input_values, - chunk_length=chunk_length, + chunk_size=chunk_size, concurrency=concurrency, + disable_logging=disable_logging, + ) + + tqdm_options = dict( + desc=f"Parallel map {config.function_name}", + disable=disable_logging, + total=config.input_length, + miniters=1, + dynamic_ncols=True, ) if config.concurrency == 1: - return [ - function(v) - for v in tqdm( - input_values, - desc="Parallel map", - disable=disable_progress_bar, - total=config.input_length, - ) - ] + yield from (function(v) for v in tqdm(input_values, **tqdm_options)) + return - start_methods = mp.get_all_start_methods() - ctx = mp.get_context("fork") if "fork" in start_methods else mp.get_context("spawn") + ctx = mp.get_context("spawn") ctx.freeze_support() input_queue = ctx.Queue(0 if config.chunk_count is None else config.chunk_count) @@ -81,7 +76,7 @@ def parallel_map( processes = [ ctx.Process( - name=f"parallel_map_{i}", + name=f"parallel_map_{config.function_name}_{i}", target=mapper_function, kwargs=dict( input_queue=input_queue, @@ -96,18 +91,15 @@ def parallel_map( for p in processes: p.start() - progress = tqdm( - desc="Parallel map", - disable=disable_progress_bar, - total=config.input_length, - ) + progress = tqdm(**tqdm_options) - chunks = iter(chunk(enumerate(input_values), chunk_length=config.chunk_length)) - indexed_results: List[Tuple[int, V]] = [] + 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 len(indexed_results) < read_input_length: + while not is_iteration_over or next_output_index < read_input_length: if not is_iteration_over: try: next_chunk = next(chunks) @@ -118,10 +110,23 @@ def parallel_map( try: result_chunk = output_queue.get_nowait() - indexed_results.extend(result_chunk) 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 + should_stop.set() except KeyboardInterrupt: for p in processes: @@ -134,7 +139,3 @@ def parallel_map( 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 index e8fd14c..384c642 100644 --- a/src/great_ai/utilities/parallel_map/parallel_map_configuration.py +++ b/src/great_ai/utilities/parallel_map/parallel_map_configuration.py @@ -10,13 +10,14 @@ logger = get_logger("parallel_map") class ParallelMapConfiguration(BaseModel): concurrency: int chunk_count: Optional[int] - chunk_length: int + chunk_size: int input_length: Optional[int] serialized_map_function: bytes + function_name: str 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 size: {self.chunk_size}") logger.info( f"{prefix} chunk count: {self.chunk_count if self.chunk_count else 'unknown'}" ) diff --git a/tests/utilities/test_parallel_map.py b/tests/utilities/test_parallel_map.py index 2077f19..571ba1f 100644 --- a/tests/utilities/test_parallel_map.py +++ b/tests/utilities/test_parallel_map.py @@ -1,5 +1,7 @@ import unittest +import pytest + from src.great_ai.utilities import parallel_map COUNT = int(1e5) + 3 @@ -10,7 +12,7 @@ class TestParallelMap(unittest.TestCase): inputs = range(COUNT) expected = [v**2 for v in range(COUNT)] - assert parallel_map(lambda v: v**2, inputs, concurrency=10) == expected + assert list(parallel_map(lambda v: v**2, inputs, concurrency=10)) == expected def test_with_iterable(self) -> None: from time import sleep @@ -23,38 +25,42 @@ class TestParallelMap(unittest.TestCase): expected = [v**3 for v in range(10)] assert ( - parallel_map(lambda x: x**3, my_generator(), chunk_length=1) == expected + list(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)] - self.assertEqual( - parallel_map(lambda v: v**2, inputs, disable_progress_bar=True), expected + assert ( + list(parallel_map(lambda v: v**2, inputs, disable_logging=True)) + == expected ) def test_simple_case_invalid_values(self) -> None: inputs = range(COUNT) - self.assertRaises( - AssertionError, parallel_map, lambda v: v**2, inputs, concurrency=0 - ) - self.assertRaises( - AssertionError, parallel_map, lambda v: v**2, inputs, chunk_length=0 - ) + with pytest.raises(AssertionError): + list(parallel_map(lambda v: v**2, inputs, concurrency=0)) + + with pytest.raises(AssertionError): + list(parallel_map(lambda v: v**2, inputs, chunk_size=0)) def test_no_op(self) -> None: - assert parallel_map(lambda v: v**2, [], disable_progress_bar=True) == [] - self.assertEqual( - parallel_map( - lambda v: v**2, [], disable_progress_bar=True, chunk_length=100 - ), - [], + assert list(parallel_map(lambda v: v**2, [], disable_logging=True)) == [] + + assert ( + list( + parallel_map(lambda v: v**2, [], disable_logging=True, chunk_size=100) + ) + == [] ) - self.assertEqual( - parallel_map( - lambda v: v**2, [], disable_progress_bar=True, concurrency=100 - ), - [], + assert ( + list( + parallel_map( + lambda v: v**2, [], disable_logging=True, concurrency=100 + ) + ) + == [] )