fix: terminate interrupted parallel workers

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:

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()