Add error handling to parallel map

This commit is contained in:
Andras Schmelczer 2022-07-01 22:43:08 +02:00
parent cbe843220a
commit a40f456e39
7 changed files with 361 additions and 135 deletions

View file

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

View file

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

View file

@ -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:
try:
map_function = dill.loads(serialized_map_function)
try:
last_chunk = None
while not should_stop.is_set():
if last_chunk is None:
try:
input_chunk = input_queue.get(1)
output_chunk = [(i, map_function(v)) for i, v in input_chunk]
output_queue.put(output_chunk)
input_chunk = input_queue.get_nowait()
last_chunk = []
for i, v in input_chunk:
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:
pass
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

View file

@ -1,11 +1,12 @@
import multiprocessing as mp
import queue
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 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))
return
yield from manage_serial(
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:
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
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()
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 p in processes:
p.join()
p.close()
input_queue.close()
output_queue.close()
progress.close()

View file

@ -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))
return
yield from manage_serial(
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:
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
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
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()
progress.close()

View file

@ -9,10 +9,9 @@ COUNT = int(1e5) + 3
class TestParallelMap(unittest.TestCase):
def test_simple_case_with_progress_bar(self) -> None:
inputs = range(COUNT)
expected = [v**2 for v in range(COUNT)]
assert list(parallel_map(lambda v: v**2, inputs, concurrency=10)) == expected
assert list(parallel_map(lambda v: v**2, range(COUNT), concurrency=4)) == [
v**2 for v in range(COUNT)
]
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)
expected = [v**2 for v in range(COUNT)]
assert (
list(parallel_map(lambda v: v**2, inputs, disable_logging=True))
== expected
assert list(
parallel_map(
lambda v: v**2, range(COUNT), disable_logging=True, concurrency=2
)
) == [v**2 for v in range(COUNT)]
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)) == []

View file

@ -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)
expected = [v**2 for v in range(COUNT)]
assert (
list(threaded_parallel_map(lambda v: v**2, inputs, concurrency=10))
== expected
)
assert list(
threaded_parallel_map(lambda v: v**2, range(COUNT), concurrency=4)
) == [v**2 for v in range(COUNT)]
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)
expected = [v**2 for v in range(COUNT)]
assert (
list(threaded_parallel_map(lambda v: v**2, inputs, disable_logging=True))
== expected
assert list(
threaded_parallel_map(
lambda v: v**2, range(COUNT), disable_logging=True, concurrency=2
)
) == [v**2 for v in range(COUNT)]
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 (