diff --git a/great_ai/utilities/parallel_map/parallel_map.py b/great_ai/utilities/parallel_map/parallel_map.py index f8c8e1f..fd141ab 100644 --- a/great_ai/utilities/parallel_map/parallel_map.py +++ b/great_ai/utilities/parallel_map/parallel_map.py @@ -16,7 +16,6 @@ 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") @@ -135,12 +134,12 @@ def parallel_map( else mp.get_context("spawn") ) ctx.freeze_support() + serialized_map_function = dill.dumps(func, byref=True, recurse=False) manager = ctx.Manager() input_queue = manager.Queue(config.concurrency * 2) output_queue = manager.Queue(config.concurrency * 2) should_stop = ctx.Event() - serialized_map_function = dill.dumps(func, byref=True, recurse=False) processes = [ ctx.Process( @@ -157,10 +156,11 @@ def parallel_map( for i in range(config.concurrency) ] - for p in processes: - p.start() - + completed = False try: + for p in processes: + p.start() + yield from manage_communication( input_values=input_values, chunk_size=config.chunk_size, @@ -169,18 +169,19 @@ def parallel_map( unordered=unordered, ignore_exceptions=ignore_exceptions, ) - should_stop.set() - except WorkerException: - should_stop.set() - raise - except Exception: - for p in processes: - p.terminate() - p.kill() - raise + completed = True finally: + should_stop.set() + if not completed: + for p in processes: + if p.is_alive(): + p.terminate() for p in processes: - p.join() # terminated processes have to be joined else they remain zombies + if p.pid is not None: + p.join(timeout=1) + if p.is_alive(): + p.kill() + p.join() p.close() manager.shutdown() diff --git a/tests/utilities/test_parallel_map_failures.py b/tests/utilities/test_parallel_map_failures.py new file mode 100644 index 0000000..16a1335 --- /dev/null +++ b/tests/utilities/test_parallel_map_failures.py @@ -0,0 +1,63 @@ +import os +import signal +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest + + +@pytest.mark.parametrize("operation", ["close", "interrupt"]) +def test_worker_failure_and_cancellation_finish(tmp_path: Path, operation: str) -> None: + script = textwrap.dedent(""" + import multiprocessing as mp + import os + import time + from great_ai.utilities import WorkerException, parallel_map + + def main(operation): + if operation.startswith("dead_worker"): + code = 0 if operation.endswith("zero") else 3 + try: + list(parallel_map(lambda _: os._exit(code), range(100), concurrency=1, chunk_size=1)) + except WorkerException as error: + assert f"exit code {code}" in str(error) + else: + raise AssertionError("Unexpected worker exit was ignored") + else: + iterator = parallel_map(lambda x: 0 if x == 0 else time.sleep(60), + range(100), concurrency=1, chunk_size=1) + assert next(iterator) == 0 + if operation == "close": + iterator.close() + else: + try: + iterator.throw(KeyboardInterrupt) + except KeyboardInterrupt: + pass + assert not mp.active_children(), mp.active_children() + + if __name__ == "__main__": + main(OPERATION) + """).replace("OPERATION", repr(operation)) + process = subprocess.Popen( + [sys.executable, "-c", script], + cwd=tmp_path, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + start_new_session=os.name != "nt", + ) + try: + stdout, stderr = process.communicate(timeout=15) + except subprocess.TimeoutExpired: + if os.name == "nt": + subprocess.run( + ["taskkill", "/F", "/T", "/PID", str(process.pid)], capture_output=True + ) + else: + os.killpg(process.pid, signal.SIGKILL) + stdout, stderr = process.communicate() + pytest.fail(f"Parallel mapping did not finish: {stdout}\n{stderr}") + assert process.returncode == 0, f"{stdout}\n{stderr}"