Add error handling to parallel map

src/great_ai/utilities/parallel_map/manage_communication.py

@@ -0,0 +1,77 @@
+import multiprocessing as mp
+import queue
+import traceback
+from typing import Any, Dict, Iterable, TypeVar, Union
+
+from tqdm.cli import tqdm
+
+from ..chunk import chunk
+from ..logger import get_logger
+
+logger = get_logger("parallel_map")
+
+T = TypeVar("T")
+V = TypeVar("V")
+
+
+def manage_communication(
+    *,
+    tqdm_options: Dict[str, Any],
+    input_values: Iterable[T],
+    chunk_size: int,
+    input_queue: Union[mp.Queue, queue.Queue],
+    output_queue: Union[mp.Queue, queue.Queue],
+    unordered: bool,
+    ignore_exceptions: bool,
+) -> Iterable[V]:
+    progress = tqdm(**tqdm_options)
+    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
+    try:
+        while not is_iteration_over or next_output_index < read_input_length:
+            if not is_iteration_over:
+                try:
+                    next_chunk = next(chunks)
+                    input_queue.put(next_chunk)
+                    read_input_length += len(next_chunk)
+                except StopIteration:
+                    is_iteration_over = True
+                except Exception as e:
+                    if not ignore_exceptions:
+                        raise e
+                    else:
+                        logger.error(
+                            f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}"
+                        )
+
+            try:
+                result_chunk = output_queue.get_nowait()
+                progress.update(len(result_chunk))
+
+                for index, value, exception in result_chunk:
+                    if exception is not None:
+                        e, tb = exception
+                        if not ignore_exceptions:
+                            raise e
+                        else:
+                            logger.error(
+                                f"Exception {e} encountered in worker, traceback:\n{tb}"
+                            )
+                    if unordered:
+                        yield value
+                        next_output_index += 1
+                    else:
+                        indexed_results[index] = value
+
+                if not unordered:
+                    while next_output_index in indexed_results:
+                        yield indexed_results[next_output_index]
+                        del indexed_results[next_output_index]
+                        next_output_index += 1
+            except queue.Empty:
+                pass
+    finally:
+        progress.close()

src/great_ai/utilities/parallel_map/manage_serial.py

@@ -0,0 +1,38 @@
+import traceback
+from typing import Any, Callable, Dict, Iterable, TypeVar
+
+from tqdm.cli import tqdm
+
+from ..logger import get_logger
+
+logger = get_logger("parallel_map")
+
+T = TypeVar("T")
+V = TypeVar("V")
+
+
+def manage_serial(
+    *,
+    function: Callable[[T], V],
+    tqdm_options: Dict[str, Any],
+    input_values: Iterable[T],
+    ignore_exceptions: bool,
+) -> Iterable[V]:
+    try:
+        for v in tqdm(input_values, **tqdm_options):
+            try:
+                yield function(v)
+            except Exception as e:
+                if not ignore_exceptions:
+                    raise e
+                else:
+                    logger.error(
+                        f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}"
+                    )
+    except Exception as e:
+        if not ignore_exceptions:
+            raise e
+        else:
+            logger.error(
+                f"Exception {e} encountered in worker, traceback:\n{traceback.format_exc()}"
+            )

src/great_ai/utilities/parallel_map/mapper_function.py

@@ -1,6 +1,8 @@
 import multiprocessing as mp
 import queue
 import threading
+import traceback
+from time import sleep
 from typing import Union
 
 import dill
@@ -11,15 +13,34 @@ def mapper_function(
     output_queue: Union[mp.Queue, queue.Queue],
     should_stop: Union[mp.Event, threading.Event],
     serialized_map_function: bytes,
-):
+) -> None:
-    map_function = dill.loads(serialized_map_function)
     try:
+        map_function = dill.loads(serialized_map_function)
+
+        last_chunk = None
         while not should_stop.is_set():
-            try:
+            if last_chunk is None:
-                input_chunk = input_queue.get(1)
+                try:
-                output_chunk = [(i, map_function(v)) for i, v in input_chunk]
+                    input_chunk = input_queue.get_nowait()
-                output_queue.put(output_chunk)
+                    last_chunk = []
-            except queue.Empty:
+                    for i, v in input_chunk:
-                pass
+                        result, exception = None, None
+                        try:
+                            result = map_function(v)
+                        except Exception as e:
+                            exception = e, traceback.format_exc()
+                        last_chunk.append((i, result, exception))
+                except queue.Empty:
+                    sleep(
+                        0.1
+                    )  # input_queue.get(0.1) would hang in some cases with multiprocessing
+            else:
+                try:
+                    output_queue.put_nowait(last_chunk)
+                    last_chunk = None
+                except queue.Full:
+                    sleep(
+                        0.1
+                    )  # output_queue.put(0.1) would hang in some cases with multiprocessing
     except (KeyboardInterrupt, BrokenPipeError):
-        return
+        pass

src/great_ai/utilities/parallel_map/parallel_map.py

@@ -1,11 +1,12 @@
 import multiprocessing as mp
-import queue
+from typing import Callable, Iterable, Optional, Sequence, TypeVar, overload
-from typing import Callable, Dict, Iterable, Optional, Sequence, TypeVar, overload
 
-from tqdm.cli import tqdm
+from src.great_ai.utilities.parallel_map.manage_communication import (
+    manage_communication,
+)
 
-from ..chunk import chunk
 from .get_config import get_config
+from .manage_serial import manage_serial
 from .mapper_function import mapper_function
 
 T = TypeVar("T")
@@ -19,8 +20,9 @@ def parallel_map(
     *,
     chunk_size: Optional[int],
     concurrency: Optional[int],
-    disable_logging: bool,
+    disable_logging: Optional[bool],
     unordered: Optional[bool],
+    ignore_exceptions: Optional[bool],
 ) -> Iterable[V]:
     ...
 
@@ -32,8 +34,9 @@ def parallel_map(
     *,
     chunk_size: int,
     concurrency: Optional[int],
-    disable_logging: bool,
+    disable_logging: Optional[bool],
     unordered: Optional[bool],
+    ignore_exceptions: Optional[bool],
 ) -> Iterable[V]:
     ...
 
@@ -46,6 +49,7 @@ def parallel_map(
     concurrency=None,
     disable_logging=False,
     unordered=False,
+    ignore_exceptions=False,
 ):
     config = get_config(
         function=function,
@@ -64,8 +68,12 @@ def parallel_map(
     )
 
     if config.concurrency == 1:
-        yield from (function(v) for v in tqdm(input_values, **tqdm_options))
+        yield from manage_serial(
-        return
+            function=function,
+            tqdm_options=tqdm_options,
+            input_values=input_values,
+            ignore_exceptions=ignore_exceptions,
+        )
 
     ctx = mp.get_context("spawn")
     ctx.freeze_support()
@@ -91,51 +99,20 @@ def parallel_map(
     for p in processes:
         p.start()
 
-    progress = tqdm(**tqdm_options)
-
-    chunks = iter(chunk(enumerate(input_values), chunk_size=config.chunk_size))
-    indexed_results: Dict[int, V] = {}
-    next_output_index = 0
-    read_input_length = 0
-    is_iteration_over = False
     try:
-        while not is_iteration_over or next_output_index < read_input_length:
+        yield from manage_communication(
-            if not is_iteration_over:
+            tqdm_options=tqdm_options,
-                try:
+            input_values=input_values,
-                    next_chunk = next(chunks)
+            chunk_size=config.chunk_size,
-                    input_queue.put(next_chunk)
+            input_queue=input_queue,
-                    read_input_length += len(next_chunk)
+            output_queue=output_queue,
-                except StopIteration:
+            unordered=unordered,
-                    is_iteration_over = True
+            ignore_exceptions=ignore_exceptions,
-
+        )
-            try:
-                result_chunk = output_queue.get_nowait()
-                progress.update(len(result_chunk))
-
-                for index, value in result_chunk:
-                    if unordered:
-                        yield value
-                        next_output_index += 1
-                    else:
-                        indexed_results[index] = value
-
-                if not unordered:
-                    while next_output_index in indexed_results:
-                        yield indexed_results[next_output_index]
-                        del indexed_results[next_output_index]
-                        next_output_index += 1
-            except queue.Empty:
-                pass
-
-        should_stop.set()
-    except Exception:
-        for p in processes:
-            p.terminate()
     finally:
+        should_stop.set()
         for p in processes:
             p.join()
             p.close()
         input_queue.close()
         output_queue.close()
-
-        progress.close()

src/great_ai/utilities/parallel_map/threaded_parallel_map.py

@@ -1,11 +1,10 @@
 import queue
 import threading
-from typing import Callable, Dict, Iterable, Optional, Sequence, TypeVar, overload
+from typing import Callable, Iterable, Optional, Sequence, TypeVar, overload
 
-from tqdm.cli import tqdm
-
-from ..chunk import chunk
 from .get_config import get_config
+from .manage_communication import manage_communication
+from .manage_serial import manage_serial
 from .mapper_function import mapper_function
 
 T = TypeVar("T")
@@ -19,8 +18,9 @@ def threaded_parallel_map(
     *,
     chunk_size: Optional[int],
     concurrency: Optional[int],
-    disable_logging: bool,
+    disable_logging: Optional[bool],
     unordered: Optional[bool],
+    ignore_exceptions: Optional[bool],
 ) -> Iterable[V]:
     ...
 
@@ -32,8 +32,9 @@ def threaded_parallel_map(
     *,
     chunk_size: int,
     concurrency: Optional[int],
-    disable_logging: bool,
+    disable_logging: Optional[bool],
     unordered: Optional[bool],
+    ignore_exceptions: Optional[bool],
 ) -> Iterable[V]:
     ...
 
@@ -46,6 +47,7 @@ def threaded_parallel_map(
     concurrency=None,
     disable_logging=False,
     unordered=False,
+    ignore_exceptions=False,
 ):
     config = get_config(
         function=function,
@@ -64,8 +66,12 @@ def threaded_parallel_map(
     )
 
     if config.concurrency == 1:
-        yield from (function(v) for v in tqdm(input_values, **tqdm_options))
+        yield from manage_serial(
-        return
+            function=function,
+            tqdm_options=tqdm_options,
+            input_values=input_values,
+            ignore_exceptions=ignore_exceptions,
+        )
 
     input_queue = queue.Queue(0 if config.chunk_count is None else config.chunk_count)
     output_queue = queue.Queue(0 if config.chunk_count is None else config.chunk_count)
@@ -88,45 +94,17 @@ def threaded_parallel_map(
     for t in threads:
         t.start()
 
-    progress = tqdm(**tqdm_options)
-
-    chunks = iter(chunk(enumerate(input_values), chunk_size=config.chunk_size))
-    indexed_results: Dict[int, V] = {}
-    next_output_index = 0
-    read_input_length = 0
-    is_iteration_over = False
     try:
-        while not is_iteration_over or next_output_index < read_input_length:
+        yield from manage_communication(
-            if not is_iteration_over:
+            tqdm_options=tqdm_options,
-                try:
+            input_values=input_values,
-                    next_chunk = next(chunks)
+            chunk_size=config.chunk_size,
-                    input_queue.put(next_chunk)
+            input_queue=input_queue,
-                    read_input_length += len(next_chunk)
+            output_queue=output_queue,
-                except StopIteration:
+            unordered=unordered,
-                    is_iteration_over = True
+            ignore_exceptions=ignore_exceptions,
-
+        )
-            try:
-                result_chunk = output_queue.get_nowait()
-                progress.update(len(result_chunk))
-
-                for index, value in result_chunk:
-                    if unordered:
-                        yield value
-                        next_output_index += 1
-                    else:
-                        indexed_results[index] = value
-
-                if not unordered:
-                    while next_output_index in indexed_results:
-                        yield indexed_results[next_output_index]
-                        del indexed_results[next_output_index]
-                        next_output_index += 1
-            except queue.Empty:
-                pass
-
     finally:
         should_stop.set()
         for t in threads:
             t.join()
-
-        progress.close()

tests/utilities/test_parallel_map.py

@@ -9,10 +9,9 @@ COUNT = int(1e5) + 3
 
 class TestParallelMap(unittest.TestCase):
     def test_simple_case_with_progress_bar(self) -> None:
-        inputs = range(COUNT)
+        assert list(parallel_map(lambda v: v**2, range(COUNT), concurrency=4)) == [
-        expected = [v**2 for v in range(COUNT)]
+            v**2 for v in range(COUNT)
-
+        ]
-        assert list(parallel_map(lambda v: v**2, inputs, concurrency=10)) == expected
 
     def test_with_iterable(self) -> None:
         from time import sleep
@@ -30,22 +29,88 @@ class TestParallelMap(unittest.TestCase):
         )
 
     def test_simple_case_without_progress_bar(self) -> None:
-        inputs = range(COUNT)
+        assert list(
-        expected = [v**2 for v in range(COUNT)]
+            parallel_map(
-
+                lambda v: v**2, range(COUNT), disable_logging=True, concurrency=2
-        assert (
+            )
-            list(parallel_map(lambda v: v**2, inputs, disable_logging=True))
+        ) == [v**2 for v in range(COUNT)]
-            == expected
-        )
 
     def test_simple_case_invalid_values(self) -> None:
-        inputs = range(COUNT)
+        with pytest.raises(AssertionError):
+            list(parallel_map(lambda v: v**2, range(COUNT), concurrency=0))
 
         with pytest.raises(AssertionError):
-            list(parallel_map(lambda v: v**2, inputs, concurrency=0))
+            list(parallel_map(lambda v: v**2, range(COUNT), chunk_size=0))
+
+    def test_this_process_exception(self) -> None:
+        def my_generator():
+            yield 1
+            yield 2
+            yield 3
+            assert False
 
         with pytest.raises(AssertionError):
-            list(parallel_map(lambda v: v**2, inputs, chunk_size=0))
+            list(
+                parallel_map(
+                    lambda v: v**2, my_generator(), concurrency=2, chunk_size=2
+                )
+            )
+
+        with pytest.raises(AssertionError):
+            list(
+                parallel_map(
+                    lambda v: v**2, my_generator(), concurrency=1, chunk_size=2
+                )
+            )
+
+    def test_ignore_this_process_exception(self) -> None:
+        def my_generator():
+            yield 1
+            yield 2
+            yield 3
+            yield 1 / 0
+
+        assert list(
+            parallel_map(
+                lambda v: v**2,
+                my_generator(),
+                concurrency=2,
+                chunk_size=2,
+                ignore_exceptions=True,
+            )
+        ) == [1, 4]
+        assert list(
+            parallel_map(
+                lambda v: v**2,
+                my_generator(),
+                concurrency=1,
+                chunk_size=2,
+                ignore_exceptions=True,
+            )
+        ) == [1, 4, 9]
+
+    def test_worker_process_exception(self) -> None:
+        def oh_no(_):
+            raise ValueError("hi")
+
+        with pytest.raises(ValueError):
+            list(parallel_map(oh_no, range(COUNT), concurrency=2))
+
+        with pytest.raises(ValueError):
+            list(parallel_map(oh_no, range(COUNT), concurrency=1))
+
+    def test_ignore_worker_process_exception(self) -> None:
+        def oh_no(_):
+            raise ValueError("hi")
+
+        assert (
+            list(parallel_map(oh_no, range(3), concurrency=2, ignore_exceptions=True))
+            == [None] * 3
+        )
+        assert (
+            list(parallel_map(oh_no, range(3), concurrency=1, ignore_exceptions=True))
+            == [None] * 3
+        )
 
     def test_no_op(self) -> None:
         assert list(parallel_map(lambda v: v**2, [], disable_logging=True)) == []

tests/utilities/test_threaded_parallel_map.py

@@ -7,15 +7,11 @@ from src.great_ai.utilities import threaded_parallel_map
 COUNT = int(1e5) + 3
 
 
-class TestThreadedParallelMap(unittest.TestCase):
+class TestParallelMap(unittest.TestCase):
     def test_simple_case_with_progress_bar(self) -> None:
-        inputs = range(COUNT)
+        assert list(
-        expected = [v**2 for v in range(COUNT)]
+            threaded_parallel_map(lambda v: v**2, range(COUNT), concurrency=4)
-
+        ) == [v**2 for v in range(COUNT)]
-        assert (
-            list(threaded_parallel_map(lambda v: v**2, inputs, concurrency=10))
-            == expected
-        )
 
     def test_with_iterable(self) -> None:
         from time import sleep
@@ -33,22 +29,96 @@ class TestThreadedParallelMap(unittest.TestCase):
         )
 
     def test_simple_case_without_progress_bar(self) -> None:
-        inputs = range(COUNT)
+        assert list(
-        expected = [v**2 for v in range(COUNT)]
+            threaded_parallel_map(
-
+                lambda v: v**2, range(COUNT), disable_logging=True, concurrency=2
-        assert (
+            )
-            list(threaded_parallel_map(lambda v: v**2, inputs, disable_logging=True))
+        ) == [v**2 for v in range(COUNT)]
-            == expected
-        )
 
     def test_simple_case_invalid_values(self) -> None:
-        inputs = range(COUNT)
+        with pytest.raises(AssertionError):
+            list(threaded_parallel_map(lambda v: v**2, range(COUNT), concurrency=0))
 
         with pytest.raises(AssertionError):
-            list(threaded_parallel_map(lambda v: v**2, inputs, concurrency=0))
+            list(threaded_parallel_map(lambda v: v**2, range(COUNT), chunk_size=0))
+
+    def test_this_worker_exception(self) -> None:
+        def my_generator():
+            yield 1
+            yield 2
+            yield 3
+            assert False
 
         with pytest.raises(AssertionError):
-            list(threaded_parallel_map(lambda v: v**2, inputs, chunk_size=0))
+            list(
+                threaded_parallel_map(
+                    lambda v: v**2, my_generator(), concurrency=2, chunk_size=2
+                )
+            )
+
+        with pytest.raises(AssertionError):
+            list(
+                threaded_parallel_map(
+                    lambda v: v**2, my_generator(), concurrency=1, chunk_size=2
+                )
+            )
+
+    def test_ignore_this_worker_exception(self) -> None:
+        def my_generator():
+            yield 1
+            yield 2
+            yield 3
+            yield 1 / 0
+
+        assert list(
+            threaded_parallel_map(
+                lambda v: v**2,
+                my_generator(),
+                concurrency=2,
+                chunk_size=2,
+                ignore_exceptions=True,
+            )
+        ) == [1, 4]
+        assert list(
+            threaded_parallel_map(
+                lambda v: v**2,
+                my_generator(),
+                concurrency=1,
+                chunk_size=2,
+                ignore_exceptions=True,
+            )
+        ) == [1, 4, 9]
+
+    def test_worker_worker_exception(self) -> None:
+        def oh_no(_):
+            raise ValueError("hi")
+
+        with pytest.raises(ValueError):
+            list(threaded_parallel_map(oh_no, range(COUNT), concurrency=2))
+
+        with pytest.raises(ValueError):
+            list(threaded_parallel_map(oh_no, range(COUNT), concurrency=1))
+
+    def test_ignore_worker_worker_exception(self) -> None:
+        def oh_no(_):
+            raise ValueError("hi")
+
+        assert (
+            list(
+                threaded_parallel_map(
+                    oh_no, range(3), concurrency=2, ignore_exceptions=True
+                )
+            )
+            == [None] * 3
+        )
+        assert (
+            list(
+                threaded_parallel_map(
+                    oh_no, range(3), concurrency=1, ignore_exceptions=True
+                )
+            )
+            == [None] * 3
+        )
 
     def test_no_op(self) -> None:
         assert (