Stop threaded mapping workers on cancellation and exceptions

great_ai/utilities/parallel_map/threaded_parallel_map.py

@@ -145,14 +145,14 @@ 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")
 
+    try:
+        for t in threads:
+            t.start()
         yield from manage_communication(
             input_values=input_values,
             chunk_size=config.chunk_size,
@@ -162,6 +162,8 @@ def threaded_parallel_map(
             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