Detect exited threaded mapping workers

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:

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