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

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