diff --git a/fastembed/parallel_processor.py b/fastembed/parallel_processor.py index bfaaf0c5..ffb3d6ec 100644 --- a/fastembed/parallel_processor.py +++ b/fastembed/parallel_processor.py @@ -1,5 +1,6 @@ import logging import os +import time from collections import defaultdict from copy import deepcopy from enum import Enum @@ -111,7 +112,11 @@ def __init__( self.num_active_workers: BaseValue | None = None def start(self, **kwargs: Any) -> None: + self.emergency_shutdown = False self.input_queue = self.ctx.Queue(self.queue_size) + # An emergency shutdown unblocks the feeder thread with EPIPE (see semi_ordered_map), let it + # exit quietly instead of printing a traceback. ProcessPoolExecutor does the same. + self.input_queue._ignore_epipe = True # type: ignore[attr-defined] self.output_queue = self.ctx.Queue(self.queue_size) ctx_value = self.ctx.Value("i", self.num_workers) @@ -153,6 +158,7 @@ def ordered_map(self, stream: Iterable[Any], *args: Any, **kwargs: Any) -> Itera def semi_ordered_map( self, stream: Iterable[Any], *args: Any, **kwargs: Any ) -> Iterable[tuple[int, Any]]: + completed = False try: self.start(**kwargs) @@ -196,10 +202,23 @@ def semi_ordered_map( raise RuntimeError("Thread unexpectedly terminated") yield out_item read += 1 + completed = True finally: assert self.input_queue is not None, "Input queue is None" assert self.output_queue is not None, "Output queue is None" - self.join() + if completed: + self.join() + else: + # The generator may be closed before the input stream is exhausted (for + # example, when a caller only consumes a prefix of the embeddings). In that + # case workers have not necessarily received their stop signals and a normal + # join would wait forever for processes blocked on the input queue. + self.emergency_shutdown = True + self.join_or_terminate() + # Nothing reads the input pipe anymore, so the feeder thread can be stuck writing to it, + # holding every batch it hasn't sent. Closing our read end fails that write with EPIPE and + # lets the thread exit, same as Queue._terminate_broken() in python 3.12+. + self.input_queue._reader.close() # type: ignore[attr-defined] self.input_queue.close() self.output_queue.close() if self.emergency_shutdown: @@ -227,10 +246,16 @@ def join_or_terminate(self, timeout: int = 1) -> None: @param timeout: @return: """ + # One deadline for the whole pool: workers that can't finish (e.g. after an early close) + # would otherwise each use up the full timeout, one after another. + deadline = time.monotonic() + timeout + for process in self.processes: + process.join(timeout=max(0.0, deadline - time.monotonic())) for process in self.processes: - process.join(timeout=timeout) if process.is_alive(): process.terminate() + for process in self.processes: + process.join(timeout=timeout) self.processes.clear() def join(self) -> None: diff --git a/tests/test_parallel_processor.py b/tests/test_parallel_processor.py new file mode 100644 index 00000000..c9d4fb79 --- /dev/null +++ b/tests/test_parallel_processor.py @@ -0,0 +1,40 @@ +import threading +from itertools import count +from multiprocessing import get_all_start_methods + +import pytest + +from fastembed.parallel_processor import ParallelWorkerPool, Worker + + +class EchoWorker(Worker): + @classmethod + def start(cls, **kwargs): + return cls() + + def process(self, items): + yield from items + + +# semi_ordered_map is closed by garbage collection, so an error in its cleanup is only reported as unraisable +@pytest.mark.filterwarnings("error::pytest.PytestUnraisableExceptionWarning") +def test_closing_partially_consumed_iterator_stops_workers(): + start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn" + pool = ParallelWorkerPool(2, EchoWorker, start_method=start_method) + # the stream never ends, so the workers never get their stop signals + results = pool.ordered_map(count()) + assert next(results) == 0 + workers = list(pool.processes) + + # close() used to wait in join() forever, run it in a thread so a regression fails instead of + # hanging the test session + closer = threading.Thread(target=results.close, daemon=True) + closer.start() + closer.join(timeout=30) + try: + assert not closer.is_alive(), "closing the iterator hung" + assert not any(worker.is_alive() for worker in workers) + finally: + for worker in workers: + if worker.is_alive(): + worker.kill()