diff --git a/src/great_ai/utilities/parallel_map/parallel_map.py b/src/great_ai/utilities/parallel_map/parallel_map.py index 5d058a8..7ce61ff 100644 --- a/src/great_ai/utilities/parallel_map/parallel_map.py +++ b/src/great_ai/utilities/parallel_map/parallel_map.py @@ -123,13 +123,13 @@ def parallel_map( ignore_exceptions=ignore_exceptions, ) should_stop.set() - for p in processes: - p.join() - p.close() except: for p in processes: p.terminate() raise finally: + for p in processes: + p.join() # terminated processes have to be joined else they remain zombies + p.close() input_queue.close() output_queue.close() diff --git a/tests/utilities/test_parallel_map.py b/tests/utilities/test_parallel_map.py index fa42161..07a14c8 100644 --- a/tests/utilities/test_parallel_map.py +++ b/tests/utilities/test_parallel_map.py @@ -61,31 +61,31 @@ class TestParallelMap(unittest.TestCase): ) ) - # def test_ignore_this_process_exception(self) -> None: - # def my_generator(): - # yield 1 - # yield 2 - # yield 3 - # yield 1 / 0 + 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] + 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(_):