From 9866b88faf52428318790ccab937a1c4e277e4be Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Fri, 1 Jul 2022 23:50:31 +0200 Subject: [PATCH] Fix exception handling --- .../utilities/parallel_map/parallel_map.py | 6 ++++- .../parallel_map/threaded_parallel_map.py | 26 +++++++++---------- 2 files changed, 17 insertions(+), 15 deletions(-) diff --git a/src/great_ai/utilities/parallel_map/parallel_map.py b/src/great_ai/utilities/parallel_map/parallel_map.py index 971ff1a..22a59fb 100644 --- a/src/great_ai/utilities/parallel_map/parallel_map.py +++ b/src/great_ai/utilities/parallel_map/parallel_map.py @@ -106,10 +106,14 @@ def parallel_map( unordered=unordered, ignore_exceptions=ignore_exceptions, ) - finally: should_stop.set() for p in processes: p.join() p.close() + except: + for p in processes: + p.terminate() + raise + finally: input_queue.close() output_queue.close() diff --git a/src/great_ai/utilities/parallel_map/threaded_parallel_map.py b/src/great_ai/utilities/parallel_map/threaded_parallel_map.py index a6b4b74..bb26060 100644 --- a/src/great_ai/utilities/parallel_map/threaded_parallel_map.py +++ b/src/great_ai/utilities/parallel_map/threaded_parallel_map.py @@ -94,17 +94,15 @@ def threaded_parallel_map( for t in threads: t.start() - try: - yield from manage_communication( - tqdm_options=tqdm_options, - input_values=input_values, - chunk_size=config.chunk_size, - input_queue=input_queue, - output_queue=output_queue, - unordered=unordered, - ignore_exceptions=ignore_exceptions, - ) - finally: - should_stop.set() - for t in threads: - t.join() + yield from manage_communication( + tqdm_options=tqdm_options, + input_values=input_values, + chunk_size=config.chunk_size, + input_queue=input_queue, + output_queue=output_queue, + unordered=unordered, + ignore_exceptions=ignore_exceptions, + ) + should_stop.set() + for t in threads: + t.join()