Stop threaded mapping workers on cancellation and exceptions

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