Raise an error when a parallel worker exits unexpectedly

great_ai/utilities/parallel_map/manage_communication.py

@@ -1,7 +1,7 @@
 import multiprocessing as mp
 import queue
 import traceback
-from typing import Dict, Iterable, List, TypeVar, Union
+from typing import Callable, Dict, Iterable, List, Optional, Tuple, TypeVar, Union
 
 from ..chunk import chunk
 from ..logger.get_logger import get_logger
@@ -22,19 +22,27 @@ def manage_communication(
     output_queue: Union[mp.Queue, queue.Queue],
     unordered: bool,
     ignore_exceptions: bool,
+    check_workers: Optional[Callable[[], None]] = None,
 ) -> Iterable[V]:
     chunks = iter(chunk(enumerate(input_values), chunk_size=chunk_size))
     indexed_results: Dict[int, V] = {}
     next_output_index = 0
     read_input_length = 0
     is_iteration_over = False
+    pending_chunk: Optional[List[Tuple[int, T]]] = None
 
     while not is_iteration_over or next_output_index < read_input_length:
+        if check_workers is not None:
+            check_workers()
         if not is_iteration_over:
             try:
-                next_chunk = next(chunks)
+                if pending_chunk is None:
-                input_queue.put(next_chunk)
+                    pending_chunk = next(chunks)
-                read_input_length += len(next_chunk)
+                input_queue.put(pending_chunk, timeout=0.1)
+                read_input_length += len(pending_chunk)
+                pending_chunk = None
+            except queue.Full:
+                pass  # Retry the same chunk after checking worker health.
             except StopIteration:
                 is_iteration_over = True
             except Exception as e:
@@ -46,7 +54,7 @@ def manage_communication(
                     raise
 
         try:
-            result_chunk: List[MapResult] = output_queue.get_nowait()
+            result_chunk: List[MapResult] = output_queue.get(timeout=0.01)
             for r in result_chunk:
                 if r.exception is not None:
                     if ignore_exceptions:

great_ai/utilities/parallel_map/parallel_map.py

@@ -16,6 +16,7 @@ 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")
@@ -156,6 +157,14 @@ def parallel_map(
         for i in range(config.concurrency)
     ]
 
+    def check_workers() -> None:
+        for process in processes:
+            if process.exitcode is not None:
+                raise WorkerException(
+                    f"Worker {process.name} exited unexpectedly "
+                    f"with exit code {process.exitcode}"
+                )
+
     completed = False
     try:
         for p in processes:
@@ -168,6 +177,7 @@ def parallel_map(
             output_queue=output_queue,
             unordered=unordered,
             ignore_exceptions=ignore_exceptions,
+            check_workers=check_workers,
         )
         completed = True
     finally:

tests/utilities/test_parallel_map_failures.py

@@ -8,7 +8,9 @@ from pathlib import Path
 import pytest
 
 
-@pytest.mark.parametrize("operation", ["close", "interrupt"])
+@pytest.mark.parametrize(
+    "operation", ["close", "interrupt", "dead_worker", "dead_worker_zero"]
+)
 def test_worker_failure_and_cancellation_finish(tmp_path: Path, operation: str) -> None:
     script = textwrap.dedent("""
         import multiprocessing as mp