andras/great-ai
1
0
Fork 0
Code Issues Pull requests 1 Projects Releases 4 Packages Wiki Activity Actions

Fix hanging bug by adding manager

Browse source
This commit is contained in:
Andras Schmelczer 2022-07-03 16:01:52 +02:00
parent ad1b2a20ef
commit 2f4d203446
10 changed files with 40 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

2
src/great_ai/utilities/__init__.py
View file

@ -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

1
src/great_ai/utilities/parallel_map/__init__.py
View file

@ -1,2 +1,3 @@
from .parallel_map import parallel_map
from .threaded_parallel_map import threaded_parallel_map
from .worker_exception import WorkerException

5
src/great_ai/utilities/parallel_map/manage_communication.py
View file

@ -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}"

17
src/great_ai/utilities/parallel_map/manage_serial.py
View file

@ -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
else:
logger.error(
f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}"
)
except Exception as e:
if not ignore_exceptions:
raise e
raise WorkerException from e
else:
logger.error(
f"Exception {e} encountered in worker, traceback:\n{traceback.format_exc()}"
)
except Exception as e:
if not ignore_exceptions:
raise
else:
logger.error(
f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}"
)

14
src/great_ai/utilities/parallel_map/mapper_function.py
View file

@ -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

23
src/great_ai/utilities/parallel_map/parallel_map.py
View file

@ -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()

2
src/great_ai/utilities/parallel_map/threaded_parallel_map.py
View file

@ -118,4 +118,4 @@ def threaded_parallel_map(
)
should_stop.set()
for t in threads:
t.join()
t.join(1)

2
src/great_ai/utilities/parallel_map/worker_exception.py Normal file
View file

@ -0,0 +1,2 @@
class WorkerException(Exception):
pass

6
tests/utilities/test_parallel_map.py
View file

@ -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:

6
tests/utilities/test_threaded_parallel_map.py
View file

@ -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: