diff --git a/great_ai/utilities/parallel_map/threaded_parallel_map.py b/great_ai/utilities/parallel_map/threaded_parallel_map.py index 21d80e5..a6a4d7e 100644 --- a/great_ai/utilities/parallel_map/threaded_parallel_map.py +++ b/great_ai/utilities/parallel_map/threaded_parallel_map.py @@ -15,6 +15,7 @@ from typing import ( 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") @@ -147,6 +148,11 @@ def threaded_parallel_map( for t in threads: t.start() + def check_workers() -> None: + for thread in threads: + if not thread.is_alive(): + raise WorkerException(f"Worker {thread.name} exited unexpectedly") + yield from manage_communication( input_values=input_values, chunk_size=config.chunk_size, @@ -154,6 +160,7 @@ def threaded_parallel_map( output_queue=output_queue, unordered=unordered, ignore_exceptions=ignore_exceptions, + check_workers=check_workers, ) should_stop.set() for t in threads: diff --git a/tests/utilities/test_threaded_map_regressions.py b/tests/utilities/test_threaded_map_regressions.py new file mode 100644 index 0000000..fb9b512 --- /dev/null +++ b/tests/utilities/test_threaded_map_regressions.py @@ -0,0 +1,23 @@ +import subprocess +import sys + +import pytest + + +@pytest.mark.parametrize("exception", ["SystemExit", "KeyboardInterrupt"]) +def test_exited_thread_is_reported(exception: str) -> None: + script = f""" +from great_ai.utilities import threaded_parallel_map, WorkerException +def dead(value): + raise {exception}() +try: + list(threaded_parallel_map(dead, [1], concurrency=1)) +except WorkerException as error: + assert 'exited unexpectedly' in str(error) +else: + raise AssertionError('Worker exit was ignored') +""" + result = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, timeout=10 + ) + assert result.returncode == 0, result.stderr