From e61b280d5ac59c8ab757ea5ae610416e47945285 Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Sun, 3 Jul 2022 16:01:52 +0200 Subject: [PATCH] Fix hanging bug by adding manager --- src/great_ai/utilities/__init__.py | 2 +- .../utilities/parallel_map/__init__.py | 1 + .../parallel_map/manage_communication.py | 5 ++-- .../utilities/parallel_map/manage_serial.py | 9 ++++---- .../utilities/parallel_map/mapper_function.py | 14 ++++------- .../utilities/parallel_map/parallel_map.py | 23 +++++++++++++------ .../parallel_map/threaded_parallel_map.py | 2 +- .../parallel_map/worker_exception.py | 2 ++ tests/utilities/test_parallel_map.py | 6 ++--- tests/utilities/test_threaded_parallel_map.py | 6 ++--- 10 files changed, 40 insertions(+), 30 deletions(-) create mode 100644 src/great_ai/utilities/parallel_map/worker_exception.py diff --git a/src/great_ai/utilities/__init__.py b/src/great_ai/utilities/__init__.py index 5bb9463..b4348f2 100644 --- a/src/great_ai/utilities/__init__.py +++ b/src/great_ai/utilities/__init__.py @@ -6,6 +6,6 @@ from .get_sentences import get_sentences from .language import english_name_of_language, is_english, predict_language from .logger import get_logger from .match_names import match_names -from .parallel_map import parallel_map, threaded_parallel_map +from .parallel_map import WorkerException, parallel_map, threaded_parallel_map from .unchunk import unchunk from .unique import unique diff --git a/src/great_ai/utilities/parallel_map/__init__.py b/src/great_ai/utilities/parallel_map/__init__.py index 0e7e247..7b6ceec 100644 --- a/src/great_ai/utilities/parallel_map/__init__.py +++ b/src/great_ai/utilities/parallel_map/__init__.py @@ -1,2 +1,3 @@ from .parallel_map import parallel_map from .threaded_parallel_map import threaded_parallel_map +from .worker_exception import WorkerException diff --git a/src/great_ai/utilities/parallel_map/manage_communication.py b/src/great_ai/utilities/parallel_map/manage_communication.py index 1558bd2..48f4db0 100644 --- a/src/great_ai/utilities/parallel_map/manage_communication.py +++ b/src/great_ai/utilities/parallel_map/manage_communication.py @@ -5,6 +5,7 @@ from typing import Dict, Iterable, TypeVar, Union from ..chunk import chunk from ..logger import get_logger +from .worker_exception import WorkerException logger = get_logger("parallel_map") @@ -36,7 +37,7 @@ def manage_communication( is_iteration_over = True except Exception as e: if not ignore_exceptions: - raise e + raise else: logger.error( f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}" @@ -49,7 +50,7 @@ def manage_communication( if exception is not None: e, tb = exception if not ignore_exceptions: - raise e + raise WorkerException from e else: logger.error( f"Exception {e} encountered in worker, traceback:\n{tb}" diff --git a/src/great_ai/utilities/parallel_map/manage_serial.py b/src/great_ai/utilities/parallel_map/manage_serial.py index ff0677d..ae3f8c5 100644 --- a/src/great_ai/utilities/parallel_map/manage_serial.py +++ b/src/great_ai/utilities/parallel_map/manage_serial.py @@ -2,6 +2,7 @@ import traceback from typing import Callable, Iterable, TypeVar from ..logger import get_logger +from .worker_exception import WorkerException logger = get_logger("parallel_map") @@ -21,15 +22,15 @@ def manage_serial( yield function(v) except Exception as e: if not ignore_exceptions: - raise e + raise WorkerException from e else: logger.error( - f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}" + f"Exception {e} encountered in worker, traceback:\n{traceback.format_exc()}" ) except Exception as e: if not ignore_exceptions: - raise e + raise else: logger.error( - f"Exception {e} encountered in worker, traceback:\n{traceback.format_exc()}" + f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}" ) diff --git a/src/great_ai/utilities/parallel_map/mapper_function.py b/src/great_ai/utilities/parallel_map/mapper_function.py index 7ecdf27..c294a0a 100644 --- a/src/great_ai/utilities/parallel_map/mapper_function.py +++ b/src/great_ai/utilities/parallel_map/mapper_function.py @@ -2,7 +2,6 @@ import multiprocessing as mp import queue import threading import traceback -from time import sleep from typing import Callable, Union import dill @@ -19,7 +18,7 @@ def mapper_function( map_function = dill.loads(map_function) last_chunk = None - while not should_stop.is_set(): + while not should_stop.wait(0.1): if last_chunk is None: try: input_chunk = input_queue.get_nowait() @@ -32,16 +31,13 @@ def mapper_function( 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: + pass + + if last_chunk is not None: 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 + pass except (KeyboardInterrupt, BrokenPipeError): pass diff --git a/src/great_ai/utilities/parallel_map/parallel_map.py b/src/great_ai/utilities/parallel_map/parallel_map.py index 9934046..0c088c0 100644 --- a/src/great_ai/utilities/parallel_map/parallel_map.py +++ b/src/great_ai/utilities/parallel_map/parallel_map.py @@ -7,6 +7,7 @@ from .get_config import get_config from .manage_communication import manage_communication from .manage_serial import manage_serial from .mapper_function import mapper_function +from .worker_exception import WorkerException T = TypeVar("T") V = TypeVar("V") @@ -87,13 +88,17 @@ def parallel_map( ignore_exceptions=ignore_exceptions, ) - ctx = mp.get_context("spawn") + ctx = ( + mp.get_context("fork") + if "fork" in mp.get_all_start_methods() + else mp.get_context("spawn") + ) ctx.freeze_support() + manager = ctx.Manager() + input_queue = manager.Queue(config.concurrency * 2) + output_queue = manager.Queue(config.concurrency * 2) - input_queue = ctx.Queue(config.concurrency * 2) - output_queue = ctx.Queue(config.concurrency * 2) should_stop = ctx.Event() - serialized_map_function = dill.dumps(function, byref=True, recurse=True) processes = [ @@ -124,13 +129,17 @@ def parallel_map( ignore_exceptions=ignore_exceptions, ) should_stop.set() - except: + except WorkerException: + should_stop.set() + raise + except Exception: for p in processes: p.terminate() + p.kill() 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() + + manager.shutdown() 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 53fcab2..122b692 100644 --- a/src/great_ai/utilities/parallel_map/threaded_parallel_map.py +++ b/src/great_ai/utilities/parallel_map/threaded_parallel_map.py @@ -118,4 +118,4 @@ def threaded_parallel_map( ) should_stop.set() for t in threads: - t.join() + t.join(1) diff --git a/src/great_ai/utilities/parallel_map/worker_exception.py b/src/great_ai/utilities/parallel_map/worker_exception.py new file mode 100644 index 0000000..92babc3 --- /dev/null +++ b/src/great_ai/utilities/parallel_map/worker_exception.py @@ -0,0 +1,2 @@ +class WorkerException(Exception): + pass diff --git a/tests/utilities/test_parallel_map.py b/tests/utilities/test_parallel_map.py index 07a14c8..7900d60 100644 --- a/tests/utilities/test_parallel_map.py +++ b/tests/utilities/test_parallel_map.py @@ -2,7 +2,7 @@ import unittest import pytest -from src.great_ai.utilities import parallel_map +from src.great_ai.utilities import WorkerException, parallel_map COUNT = int(1e5) + 3 @@ -91,10 +91,10 @@ class TestParallelMap(unittest.TestCase): def oh_no(_): raise ValueError("hi") - with pytest.raises(ValueError): + with pytest.raises(WorkerException): list(parallel_map(oh_no, range(COUNT), concurrency=2)) - with pytest.raises(ValueError): + with pytest.raises(WorkerException): list(parallel_map(oh_no, range(COUNT), concurrency=1)) def test_ignore_worker_process_exception(self) -> None: diff --git a/tests/utilities/test_threaded_parallel_map.py b/tests/utilities/test_threaded_parallel_map.py index 0bf84a0..d421e90 100644 --- a/tests/utilities/test_threaded_parallel_map.py +++ b/tests/utilities/test_threaded_parallel_map.py @@ -2,7 +2,7 @@ import unittest import pytest -from src.great_ai.utilities import threaded_parallel_map +from src.great_ai.utilities import WorkerException, threaded_parallel_map COUNT = int(1e5) + 3 @@ -91,10 +91,10 @@ class TestParallelMap(unittest.TestCase): def oh_no(_): raise ValueError("hi") - with pytest.raises(ValueError): + with pytest.raises(WorkerException): list(threaded_parallel_map(oh_no, range(COUNT), concurrency=2)) - with pytest.raises(ValueError): + with pytest.raises(WorkerException): list(threaded_parallel_map(oh_no, range(COUNT), concurrency=1)) def test_ignore_worker_worker_exception(self) -> None: