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

Add error handling to parallel map

Browse source
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
Download patch file Download diff file Expand all files Collapse all files

77
src/great_ai/utilities/parallel_map/manage_communication.py Normal file
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()

38
src/great_ai/utilities/parallel_map/manage_serial.py Normal file
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()}"
)

39
src/great_ai/utilities/parallel_map/mapper_function.py
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,
):
map_function = dill.loads(serialized_map_function)
) -> None:
try:
map_function = dill.loads(serialized_map_function)
last_chunk = None
while not should_stop.is_set():
try:
input_chunk = input_queue.get(1)
output_chunk = [(i, map_function(v)) for i, v in input_chunk]
output_queue.put(output_chunk)
except queue.Empty:
pass
if last_chunk is None:
try:
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:
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

75
src/great_ai/utilities/parallel_map/parallel_map.py
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()

68
src/great_ai/utilities/parallel_map/threaded_parallel_map.py
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()