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)