diff --git a/great_ai/utilities/parallel_map/manage_communication.py b/great_ai/utilities/parallel_map/manage_communication.py index 8476e4c..f9a7cca 100644 --- a/great_ai/utilities/parallel_map/manage_communication.py +++ b/great_ai/utilities/parallel_map/manage_communication.py @@ -1,7 +1,7 @@ import multiprocessing as mp import queue import traceback -from typing import Dict, Iterable, List, TypeVar, Union +from typing import Callable, Dict, Iterable, List, Optional, Tuple, TypeVar, Union from ..chunk import chunk from ..logger.get_logger import get_logger @@ -22,19 +22,27 @@ def manage_communication( output_queue: Union[mp.Queue, queue.Queue], unordered: bool, ignore_exceptions: bool, + check_workers: Optional[Callable[[], None]] = None, ) -> Iterable[V]: 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 + pending_chunk: Optional[List[Tuple[int, T]]] = None while not is_iteration_over or next_output_index < read_input_length: + if check_workers is not None: + check_workers() if not is_iteration_over: try: - next_chunk = next(chunks) - input_queue.put(next_chunk) - read_input_length += len(next_chunk) + if pending_chunk is None: + pending_chunk = next(chunks) + input_queue.put(pending_chunk, timeout=0.1) + read_input_length += len(pending_chunk) + pending_chunk = None + except queue.Full: + pass # Retry the same chunk after checking worker health. except StopIteration: is_iteration_over = True except Exception as e: @@ -46,7 +54,7 @@ def manage_communication( raise try: - result_chunk: List[MapResult] = output_queue.get_nowait() + result_chunk: List[MapResult] = output_queue.get(timeout=0.01) for r in result_chunk: if r.exception is not None: if ignore_exceptions: diff --git a/great_ai/utilities/parallel_map/parallel_map.py b/great_ai/utilities/parallel_map/parallel_map.py index fd141ab..4472642 100644 --- a/great_ai/utilities/parallel_map/parallel_map.py +++ b/great_ai/utilities/parallel_map/parallel_map.py @@ -16,6 +16,7 @@ import dill from .get_config import get_config from .manage_communication import manage_communication from .mapper_function import mapper_function +from .worker_exception import WorkerException T = TypeVar("T") V = TypeVar("V") @@ -156,6 +157,14 @@ def parallel_map( for i in range(config.concurrency) ] + def check_workers() -> None: + for process in processes: + if process.exitcode is not None: + raise WorkerException( + f"Worker {process.name} exited unexpectedly " + f"with exit code {process.exitcode}" + ) + completed = False try: for p in processes: @@ -168,6 +177,7 @@ def parallel_map( output_queue=output_queue, unordered=unordered, ignore_exceptions=ignore_exceptions, + check_workers=check_workers, ) completed = True finally: diff --git a/tests/utilities/test_parallel_map_failures.py b/tests/utilities/test_parallel_map_failures.py index 16a1335..52be7c3 100644 --- a/tests/utilities/test_parallel_map_failures.py +++ b/tests/utilities/test_parallel_map_failures.py @@ -8,7 +8,9 @@ from pathlib import Path import pytest -@pytest.mark.parametrize("operation", ["close", "interrupt"]) +@pytest.mark.parametrize( + "operation", ["close", "interrupt", "dead_worker", "dead_worker_zero"] +) def test_worker_failure_and_cancellation_finish(tmp_path: Path, operation: str) -> None: script = textwrap.dedent(""" import multiprocessing as mp