Stop threaded mapping workers on cancellation and exceptions

great_ai/utilities/parallel_map/threaded_parallel_map.py

@@ -145,23 +145,25 @@ def threaded_parallel_map(
         for i in range(config.concurrency)
     ]
 
-    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,
-        input_queue=input_queue,
-        output_queue=output_queue,
-        unordered=unordered,
-        ignore_exceptions=ignore_exceptions,
-        check_workers=check_workers,
-    )
-    should_stop.set()
-    for t in threads:
-        t.join(1)
+    try:
+        for t in threads:
+            t.start()
+        yield from manage_communication(
+            input_values=input_values,
+            chunk_size=config.chunk_size,
+            input_queue=input_queue,
+            output_queue=output_queue,
+            unordered=unordered,
+            ignore_exceptions=ignore_exceptions,
+            check_workers=check_workers,
+        )
+    finally:
+        should_stop.set()
+        for t in threads:
+            if t.ident is not None:
+                t.join(1)

tests/utilities/test_threaded_map_regressions.py

@@ -1,7 +1,10 @@
 import subprocess
 import sys
+import threading
+from typing import Generator, cast
 
 import pytest
+from great_ai.utilities import WorkerException, threaded_parallel_map
 
 
 @pytest.mark.parametrize("exception", ["SystemExit", "KeyboardInterrupt"])
@@ -21,3 +24,24 @@ else:
         [sys.executable, "-c", script], capture_output=True, text=True, timeout=10
     )
     assert result.returncode == 0, result.stderr
+
+
+@pytest.mark.parametrize("fail", [False, True])
+def test_thread_workers_stop_after_close_or_error(fail: bool) -> None:
+    def predict(value: int) -> int:
+        if fail:
+            raise ValueError("failed")
+        return value
+
+    before = set(threading.enumerate())
+    iterator = cast(
+        Generator[int, None, None],
+        threaded_parallel_map(predict, range(100), concurrency=2, chunk_size=1),
+    )
+    if fail:
+        with pytest.raises(WorkerException):
+            next(iterator)
+    else:
+        assert next(iterator) == 0
+        iterator.close()
+    assert not set(threading.enumerate()) - before