Stop parallel workers when iteration is cancelled

great_ai/utilities/parallel_map/parallel_map.py

@@ -16,7 +16,6 @@ import dill
 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")
@@ -135,12 +134,12 @@ def parallel_map(
         else mp.get_context("spawn")
     )
     ctx.freeze_support()
+    serialized_map_function = dill.dumps(func, byref=True, recurse=False)
     manager = ctx.Manager()
     input_queue = manager.Queue(config.concurrency * 2)
     output_queue = manager.Queue(config.concurrency * 2)
 
     should_stop = ctx.Event()
-    serialized_map_function = dill.dumps(func, byref=True, recurse=False)
 
     processes = [
         ctx.Process(
@@ -157,10 +156,11 @@ def parallel_map(
         for i in range(config.concurrency)
     ]
 
-    for p in processes:
+    completed = False
-        p.start()
-
     try:
+        for p in processes:
+            p.start()
+
         yield from manage_communication(
             input_values=input_values,
             chunk_size=config.chunk_size,
@@ -169,18 +169,19 @@ def parallel_map(
             unordered=unordered,
             ignore_exceptions=ignore_exceptions,
         )
-        should_stop.set()
+        completed = True
-    except WorkerException:
-        should_stop.set()
-        raise
-    except Exception:
-        for p in processes:
-            p.terminate()
-            p.kill()
-        raise
     finally:
+        should_stop.set()
+        if not completed:
+            for p in processes:
+                if p.is_alive():
+                    p.terminate()
         for p in processes:
-            p.join()  # terminated processes have to be joined else they remain zombies
+            if p.pid is not None:
+                p.join(timeout=1)
+                if p.is_alive():
+                    p.kill()
+                    p.join()
             p.close()
 
         manager.shutdown()

tests/utilities/test_parallel_map_failures.py

@@ -0,0 +1,63 @@
+import os
+import signal
+import subprocess
+import sys
+import textwrap
+from pathlib import Path
+
+import pytest
+
+
+@pytest.mark.parametrize("operation", ["close", "interrupt"])
+def test_worker_failure_and_cancellation_finish(tmp_path: Path, operation: str) -> None:
+    script = textwrap.dedent("""
+        import multiprocessing as mp
+        import os
+        import time
+        from great_ai.utilities import WorkerException, parallel_map
+
+        def main(operation):
+            if operation.startswith("dead_worker"):
+                code = 0 if operation.endswith("zero") else 3
+                try:
+                    list(parallel_map(lambda _: os._exit(code), range(100), concurrency=1, chunk_size=1))
+                except WorkerException as error:
+                    assert f"exit code {code}" in str(error)
+                else:
+                    raise AssertionError("Unexpected worker exit was ignored")
+            else:
+                iterator = parallel_map(lambda x: 0 if x == 0 else time.sleep(60),
+                                        range(100), concurrency=1, chunk_size=1)
+                assert next(iterator) == 0
+                if operation == "close":
+                    iterator.close()
+                else:
+                    try:
+                        iterator.throw(KeyboardInterrupt)
+                    except KeyboardInterrupt:
+                        pass
+            assert not mp.active_children(), mp.active_children()
+
+        if __name__ == "__main__":
+            main(OPERATION)
+    """).replace("OPERATION", repr(operation))
+    process = subprocess.Popen(
+        [sys.executable, "-c", script],
+        cwd=tmp_path,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        start_new_session=os.name != "nt",
+    )
+    try:
+        stdout, stderr = process.communicate(timeout=15)
+    except subprocess.TimeoutExpired:
+        if os.name == "nt":
+            subprocess.run(
+                ["taskkill", "/F", "/T", "/PID", str(process.pid)], capture_output=True
+            )
+        else:
+            os.killpg(process.pid, signal.SIGKILL)
+        stdout, stderr = process.communicate()
+        pytest.fail(f"Parallel mapping did not finish: {stdout}\n{stderr}")
+    assert process.returncode == 0, f"{stdout}\n{stderr}"