andras/great-ai
andras/great-ai
Watch
1
0
Fork
You've already forked great-ai
 0
Code Issues Pull requests 2 Projects Releases 4 Packages 1 Wiki Activity Actions

Raise an error when a parallel worker exits unexpectedly

Browse source
This commit is contained in:
Andras Schmelczer 2026-10-27 16:52:29 +00:00
commit 3b7584e1b1
1 changed files with 26 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -1,7 +1,7 @@
import multiprocessing as mp
import queue
import traceback
from typing import Dict, Iterable, List, TypeVar, Union
from typing import Callable, Dict, Iterable, List, Optional, Tuple, TypeVar, Union
from ..chunk import chunk
from ..logger.get_logger import get_logger
@ -22,19 +22,27 @@ def manage_communication(
output_queue: Union[mp.Queue, queue.Queue],
unordered: bool,
ignore_exceptions: bool,
check_workers: Optional[Callable[[], None]] = None,
) -> Iterable[V]:
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
pending_chunk: Optional[List[Tuple[int, T]]] = None
while not is_iteration_over or next_output_index < read_input_length:
if check_workers is not None:
check_workers()
if not is_iteration_over:
try:
next_chunk = next(chunks)
input_queue.put(next_chunk)
read_input_length += len(next_chunk)
if pending_chunk is None:
pending_chunk = next(chunks)
input_queue.put(pending_chunk, timeout=0.1)
read_input_length += len(pending_chunk)
pending_chunk = None
except queue.Full:
pass # Retry the same chunk after checking worker health.
except StopIteration:
is_iteration_over = True
except Exception as e:
@ -46,7 +54,7 @@ def manage_communication(
raise
try:
result_chunk: List[MapResult] = output_queue.get_nowait()
result_chunk: List[MapResult] = output_queue.get(timeout=0.01)
for r in result_chunk:
if r.exception is not None:
if ignore_exceptions: