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

Fix typing and minor issues

Browse source
This commit is contained in:
Andras Schmelczer 2022-07-07 14:12:44 +02:00
parent 2db2253578
commit 72ab627a34
54 changed files with 635 additions and 589 deletions
Download patch file Download diff file Expand all files Collapse all files

32
great_ai/utilities/parallel_map/manage_communication.py
View file

@ -1,10 +1,11 @@
import multiprocessing as mp
import queue
import traceback
from typing import Dict, Iterable, TypeVar, Union
from typing import Dict, Iterable, List, TypeVar, Union
from ..chunk import chunk
from ..logger import get_logger
from .map_result import MapResult
from .worker_exception import WorkerException
logger = get_logger("parallel_map")
@ -27,6 +28,7 @@ def manage_communication(
next_output_index = 0
read_input_length = 0
is_iteration_over = False
while not is_iteration_over or next_output_index < read_input_length:
if not is_iteration_over:
try:
@ -36,30 +38,30 @@ def manage_communication(
except StopIteration:
is_iteration_over = True
except Exception as e:
if not ignore_exceptions:
raise
else:
if ignore_exceptions:
logger.error(
f"Exception {e} encountered in input, traceback:\n{traceback.format_exc()}"
)
else:
raise
try:
result_chunk = output_queue.get_nowait()
for index, value, exception in result_chunk:
if exception is not None:
e, tb = exception
if not ignore_exceptions:
raise WorkerException from e
else:
result_chunk: List[MapResult] = output_queue.get_nowait()
for r in result_chunk:
if r.exception is not None:
if ignore_exceptions:
logger.error(
f"Exception {e} encountered in worker, traceback:\n{tb}"
f"Exception {r.exception} encountered in worker, traceback:\n{r.worker_traceback}"
)
else:
raise WorkerException from r.exception
if unordered:
yield value
yield r.value
next_output_index += 1
else:
indexed_results[index] = value
indexed_results[r.order] = r.value
if not unordered:
while next_output_index in indexed_results: