diff --git a/src/great_ai/utilities/parallel_map/manage_communication.py b/src/great_ai/utilities/parallel_map/manage_communication.py new file mode 100644 index 0000000..2d36beb --- /dev/null +++ b/src/great_ai/utilities/parallel_map/manage_communication.py @@ -0,0 +1,77 @@ +import multiprocessing as mp +import queue +import traceback +from typing import Any, Dict, Iterable, TypeVar, Union + +from tqdm.cli import tqdm + +from ..chunk import chunk +from ..logger import get_logger + +logger = get_logger("parallel_map") + +T = TypeVar("T") +V = TypeVar("V") + + +def manage_communication( + *, + tqdm_options: Dict[str, Any], + input_values: Iterable[T], + chunk_size: int, + input_queue: Union[mp.Queue, queue.Queue], + output_queue: Union[mp.Queue, queue.Queue], + unordered: bool, + ignore_exceptions: bool, +) -> Iterable[V]: + progress = tqdm(**tqdm_options) + chunks = iter(chunk(enumerate(input_values), chunk_size=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 + except Exception as e: + if not ignore_exceptions: + raise e + else: + logger.error( + f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}" + ) + + try: + result_chunk = output_queue.get_nowait() + progress.update(len(result_chunk)) + + for index, value, exception in result_chunk: + if exception is not None: + e, tb = exception + if not ignore_exceptions: + raise e + else: + logger.error( + f"Exception {e} encountered in worker, traceback:\n{tb}" + ) + 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: + progress.close() diff --git a/src/great_ai/utilities/parallel_map/manage_serial.py b/src/great_ai/utilities/parallel_map/manage_serial.py new file mode 100644 index 0000000..6faf6eb --- /dev/null +++ b/src/great_ai/utilities/parallel_map/manage_serial.py @@ -0,0 +1,38 @@ +import traceback +from typing import Any, Callable, Dict, Iterable, TypeVar + +from tqdm.cli import tqdm + +from ..logger import get_logger + +logger = get_logger("parallel_map") + +T = TypeVar("T") +V = TypeVar("V") + + +def manage_serial( + *, + function: Callable[[T], V], + tqdm_options: Dict[str, Any], + input_values: Iterable[T], + ignore_exceptions: bool, +) -> Iterable[V]: + try: + for v in tqdm(input_values, **tqdm_options): + try: + yield function(v) + except Exception as e: + if not ignore_exceptions: + raise e + else: + logger.error( + f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}" + ) + except Exception as e: + if not ignore_exceptions: + raise e + else: + logger.error( + f"Exception {e} encountered in worker, traceback:\n{traceback.format_exc()}" + ) diff --git a/src/great_ai/utilities/parallel_map/mapper_function.py b/src/great_ai/utilities/parallel_map/mapper_function.py index c4650c1..509487d 100644 --- a/src/great_ai/utilities/parallel_map/mapper_function.py +++ b/src/great_ai/utilities/parallel_map/mapper_function.py @@ -1,6 +1,8 @@ import multiprocessing as mp import queue import threading +import traceback +from time import sleep from typing import Union import dill @@ -11,15 +13,34 @@ def mapper_function( 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) +) -> None: try: + map_function = dill.loads(serialized_map_function) + + last_chunk = None while not should_stop.is_set(): - try: - input_chunk = input_queue.get(1) - output_chunk = [(i, map_function(v)) for i, v in input_chunk] - output_queue.put(output_chunk) - except queue.Empty: - pass + if last_chunk is None: + try: + input_chunk = input_queue.get_nowait() + last_chunk = [] + for i, v in input_chunk: + result, exception = None, None + try: + result = map_function(v) + except Exception as e: + exception = e, traceback.format_exc() + last_chunk.append((i, result, exception)) + except queue.Empty: + sleep( + 0.1 + ) # input_queue.get(0.1) would hang in some cases with multiprocessing + else: + try: + output_queue.put_nowait(last_chunk) + last_chunk = None + except queue.Full: + sleep( + 0.1 + ) # output_queue.put(0.1) would hang in some cases with multiprocessing except (KeyboardInterrupt, BrokenPipeError): - return + pass diff --git a/src/great_ai/utilities/parallel_map/parallel_map.py b/src/great_ai/utilities/parallel_map/parallel_map.py index 22042c9..249c303 100644 --- a/src/great_ai/utilities/parallel_map/parallel_map.py +++ b/src/great_ai/utilities/parallel_map/parallel_map.py @@ -1,11 +1,12 @@ import multiprocessing as mp -import queue -from typing import Callable, Dict, Iterable, Optional, Sequence, TypeVar, overload +from typing import Callable, Iterable, Optional, Sequence, TypeVar, overload -from tqdm.cli import tqdm +from src.great_ai.utilities.parallel_map.manage_communication import ( + manage_communication, +) -from ..chunk import chunk from .get_config import get_config +from .manage_serial import manage_serial from .mapper_function import mapper_function T = TypeVar("T") @@ -19,8 +20,9 @@ def parallel_map( *, chunk_size: Optional[int], concurrency: Optional[int], - disable_logging: bool, + disable_logging: Optional[bool], unordered: Optional[bool], + ignore_exceptions: Optional[bool], ) -> Iterable[V]: ... @@ -32,8 +34,9 @@ def parallel_map( *, chunk_size: int, concurrency: Optional[int], - disable_logging: bool, + disable_logging: Optional[bool], unordered: Optional[bool], + ignore_exceptions: Optional[bool], ) -> Iterable[V]: ... @@ -46,6 +49,7 @@ def parallel_map( concurrency=None, disable_logging=False, unordered=False, + ignore_exceptions=False, ): config = get_config( function=function, @@ -64,8 +68,12 @@ def parallel_map( ) if config.concurrency == 1: - yield from (function(v) for v in tqdm(input_values, **tqdm_options)) - return + yield from manage_serial( + function=function, + tqdm_options=tqdm_options, + input_values=input_values, + ignore_exceptions=ignore_exceptions, + ) ctx = mp.get_context("spawn") ctx.freeze_support() @@ -91,51 +99,20 @@ def parallel_map( for p in processes: p.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 - - should_stop.set() - except Exception: - for p in processes: - p.terminate() + yield from manage_communication( + tqdm_options=tqdm_options, + input_values=input_values, + chunk_size=config.chunk_size, + input_queue=input_queue, + output_queue=output_queue, + unordered=unordered, + ignore_exceptions=ignore_exceptions, + ) finally: + should_stop.set() for p in processes: p.join() p.close() input_queue.close() output_queue.close() - - progress.close() diff --git a/src/great_ai/utilities/parallel_map/threaded_parallel_map.py b/src/great_ai/utilities/parallel_map/threaded_parallel_map.py index 5696219..a6b4b74 100644 --- a/src/great_ai/utilities/parallel_map/threaded_parallel_map.py +++ b/src/great_ai/utilities/parallel_map/threaded_parallel_map.py @@ -1,11 +1,10 @@ import queue import threading -from typing import Callable, Dict, Iterable, Optional, Sequence, TypeVar, overload +from typing import Callable, Iterable, Optional, Sequence, TypeVar, overload -from tqdm.cli import tqdm - -from ..chunk import chunk from .get_config import get_config +from .manage_communication import manage_communication +from .manage_serial import manage_serial from .mapper_function import mapper_function T = TypeVar("T") @@ -19,8 +18,9 @@ def threaded_parallel_map( *, chunk_size: Optional[int], concurrency: Optional[int], - disable_logging: bool, + disable_logging: Optional[bool], unordered: Optional[bool], + ignore_exceptions: Optional[bool], ) -> Iterable[V]: ... @@ -32,8 +32,9 @@ def threaded_parallel_map( *, chunk_size: int, concurrency: Optional[int], - disable_logging: bool, + disable_logging: Optional[bool], unordered: Optional[bool], + ignore_exceptions: Optional[bool], ) -> Iterable[V]: ... @@ -46,6 +47,7 @@ def threaded_parallel_map( concurrency=None, disable_logging=False, unordered=False, + ignore_exceptions=False, ): config = get_config( function=function, @@ -64,8 +66,12 @@ def threaded_parallel_map( ) if config.concurrency == 1: - yield from (function(v) for v in tqdm(input_values, **tqdm_options)) - return + yield from manage_serial( + function=function, + tqdm_options=tqdm_options, + input_values=input_values, + ignore_exceptions=ignore_exceptions, + ) 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) @@ -88,45 +94,17 @@ def threaded_parallel_map( 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 - + yield from manage_communication( + tqdm_options=tqdm_options, + input_values=input_values, + chunk_size=config.chunk_size, + input_queue=input_queue, + output_queue=output_queue, + unordered=unordered, + ignore_exceptions=ignore_exceptions, + ) finally: should_stop.set() for t in threads: t.join() - - progress.close() diff --git a/tests/utilities/test_parallel_map.py b/tests/utilities/test_parallel_map.py index 571ba1f..cab163e 100644 --- a/tests/utilities/test_parallel_map.py +++ b/tests/utilities/test_parallel_map.py @@ -9,10 +9,9 @@ COUNT = int(1e5) + 3 class TestParallelMap(unittest.TestCase): def test_simple_case_with_progress_bar(self) -> None: - inputs = range(COUNT) - expected = [v**2 for v in range(COUNT)] - - assert list(parallel_map(lambda v: v**2, inputs, concurrency=10)) == expected + assert list(parallel_map(lambda v: v**2, range(COUNT), concurrency=4)) == [ + v**2 for v in range(COUNT) + ] def test_with_iterable(self) -> None: from time import sleep @@ -30,22 +29,88 @@ class TestParallelMap(unittest.TestCase): ) def test_simple_case_without_progress_bar(self) -> None: - inputs = range(COUNT) - expected = [v**2 for v in range(COUNT)] - - assert ( - list(parallel_map(lambda v: v**2, inputs, disable_logging=True)) - == expected - ) + assert list( + parallel_map( + lambda v: v**2, range(COUNT), disable_logging=True, concurrency=2 + ) + ) == [v**2 for v in range(COUNT)] def test_simple_case_invalid_values(self) -> None: - inputs = range(COUNT) + with pytest.raises(AssertionError): + list(parallel_map(lambda v: v**2, range(COUNT), concurrency=0)) with pytest.raises(AssertionError): - list(parallel_map(lambda v: v**2, inputs, concurrency=0)) + list(parallel_map(lambda v: v**2, range(COUNT), chunk_size=0)) + + def test_this_process_exception(self) -> None: + def my_generator(): + yield 1 + yield 2 + yield 3 + assert False with pytest.raises(AssertionError): - list(parallel_map(lambda v: v**2, inputs, chunk_size=0)) + list( + parallel_map( + lambda v: v**2, my_generator(), concurrency=2, chunk_size=2 + ) + ) + + with pytest.raises(AssertionError): + list( + parallel_map( + lambda v: v**2, my_generator(), concurrency=1, chunk_size=2 + ) + ) + + def test_ignore_this_process_exception(self) -> None: + def my_generator(): + yield 1 + yield 2 + yield 3 + yield 1 / 0 + + assert list( + parallel_map( + lambda v: v**2, + my_generator(), + concurrency=2, + chunk_size=2, + ignore_exceptions=True, + ) + ) == [1, 4] + assert list( + parallel_map( + lambda v: v**2, + my_generator(), + concurrency=1, + chunk_size=2, + ignore_exceptions=True, + ) + ) == [1, 4, 9] + + def test_worker_process_exception(self) -> None: + def oh_no(_): + raise ValueError("hi") + + with pytest.raises(ValueError): + list(parallel_map(oh_no, range(COUNT), concurrency=2)) + + with pytest.raises(ValueError): + list(parallel_map(oh_no, range(COUNT), concurrency=1)) + + def test_ignore_worker_process_exception(self) -> None: + def oh_no(_): + raise ValueError("hi") + + assert ( + list(parallel_map(oh_no, range(3), concurrency=2, ignore_exceptions=True)) + == [None] * 3 + ) + assert ( + list(parallel_map(oh_no, range(3), concurrency=1, ignore_exceptions=True)) + == [None] * 3 + ) def test_no_op(self) -> None: assert list(parallel_map(lambda v: v**2, [], disable_logging=True)) == [] diff --git a/tests/utilities/test_threaded_parallel_map.py b/tests/utilities/test_threaded_parallel_map.py index 84a6eef..c2e170b 100644 --- a/tests/utilities/test_threaded_parallel_map.py +++ b/tests/utilities/test_threaded_parallel_map.py @@ -7,15 +7,11 @@ from src.great_ai.utilities import threaded_parallel_map COUNT = int(1e5) + 3 -class TestThreadedParallelMap(unittest.TestCase): +class TestParallelMap(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 - ) + assert list( + threaded_parallel_map(lambda v: v**2, range(COUNT), concurrency=4) + ) == [v**2 for v in range(COUNT)] def test_with_iterable(self) -> None: from time import sleep @@ -33,22 +29,96 @@ class TestThreadedParallelMap(unittest.TestCase): ) 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 - ) + assert list( + threaded_parallel_map( + lambda v: v**2, range(COUNT), disable_logging=True, concurrency=2 + ) + ) == [v**2 for v in range(COUNT)] def test_simple_case_invalid_values(self) -> None: - inputs = range(COUNT) + with pytest.raises(AssertionError): + list(threaded_parallel_map(lambda v: v**2, range(COUNT), concurrency=0)) with pytest.raises(AssertionError): - list(threaded_parallel_map(lambda v: v**2, inputs, concurrency=0)) + list(threaded_parallel_map(lambda v: v**2, range(COUNT), chunk_size=0)) + + def test_this_worker_exception(self) -> None: + def my_generator(): + yield 1 + yield 2 + yield 3 + assert False with pytest.raises(AssertionError): - list(threaded_parallel_map(lambda v: v**2, inputs, chunk_size=0)) + list( + threaded_parallel_map( + lambda v: v**2, my_generator(), concurrency=2, chunk_size=2 + ) + ) + + with pytest.raises(AssertionError): + list( + threaded_parallel_map( + lambda v: v**2, my_generator(), concurrency=1, chunk_size=2 + ) + ) + + def test_ignore_this_worker_exception(self) -> None: + def my_generator(): + yield 1 + yield 2 + yield 3 + yield 1 / 0 + + assert list( + threaded_parallel_map( + lambda v: v**2, + my_generator(), + concurrency=2, + chunk_size=2, + ignore_exceptions=True, + ) + ) == [1, 4] + assert list( + threaded_parallel_map( + lambda v: v**2, + my_generator(), + concurrency=1, + chunk_size=2, + ignore_exceptions=True, + ) + ) == [1, 4, 9] + + def test_worker_worker_exception(self) -> None: + def oh_no(_): + raise ValueError("hi") + + with pytest.raises(ValueError): + list(threaded_parallel_map(oh_no, range(COUNT), concurrency=2)) + + with pytest.raises(ValueError): + list(threaded_parallel_map(oh_no, range(COUNT), concurrency=1)) + + def test_ignore_worker_worker_exception(self) -> None: + def oh_no(_): + raise ValueError("hi") + + assert ( + list( + threaded_parallel_map( + oh_no, range(3), concurrency=2, ignore_exceptions=True + ) + ) + == [None] * 3 + ) + assert ( + list( + threaded_parallel_map( + oh_no, range(3), concurrency=1, ignore_exceptions=True + ) + ) + == [None] * 3 + ) def test_no_op(self) -> None: assert (