diff --git a/fastembed/parallel_processor.py b/fastembed/parallel_processor.py index 4f2973aa..416e481e 100644 --- a/fastembed/parallel_processor.py +++ b/fastembed/parallel_processor.py @@ -176,7 +176,7 @@ def semi_ordered_map( out_item = None else: try: - out_item = self.output_queue.get(timeout=processing_timeout) + out_item = self._get_output() except Empty as e: self.join_or_terminate() raise e @@ -196,7 +196,7 @@ def semi_ordered_map( while read < pushed: self.check_worker_health() - out_item = self.output_queue.get(timeout=processing_timeout) + out_item = self._get_output() if out_item == QueueSignals.error: self.join_or_terminate() raise RuntimeError("Thread unexpectedly terminated") @@ -232,6 +232,19 @@ def semi_ordered_map( self.input_queue.join_thread() self.output_queue.join_thread() + def _get_output(self) -> Any: + """Wait for output, checking worker health without resetting the overall timeout.""" + assert self.output_queue is not None, "Output queue was not initialized" + deadline = time.monotonic() + processing_timeout + while True: + remaining = max(0.0, deadline - time.monotonic()) + try: + return self.output_queue.get(timeout=min(1.0, remaining)) + except Empty: + self.check_worker_health() + if time.monotonic() >= deadline: + raise + def check_worker_health(self) -> None: """ Checks if any worker process has terminated unexpectedly diff --git a/tests/test_worker_crash_wait.py b/tests/test_worker_crash_wait.py new file mode 100644 index 00000000..8ea3941c --- /dev/null +++ b/tests/test_worker_crash_wait.py @@ -0,0 +1,96 @@ +import os +import time +from multiprocessing import get_all_start_methods +from pathlib import Path +from queue import Empty + +import pytest + +import fastembed.parallel_processor as processor +from fastembed.parallel_processor import ParallelWorkerPool, Worker + + +class ControlledWorker(Worker): + @classmethod + def start(cls, marker, crash=False, delay=0, **kwargs): + worker = cls() + worker.marker = marker + worker.crash = crash + worker.delay = delay + return worker + + def process(self, items): + for item in items: + deadline = time.monotonic() + 30 + while not Path(self.marker).exists(): + if time.monotonic() >= deadline: + raise RuntimeError("Parent did not reach its result wait") + time.sleep(0.01) + if self.crash: + os._exit(23) + time.sleep(self.delay) + yield item + + +@pytest.fixture +def make_pool(monkeypatch, tmp_path): + workers = [] + monkeypatch.setattr(processor, "max_internal_batch_size", 1) + + def create(start_method="spawn"): + marker = tmp_path / "waiting" + pool = ParallelWorkerPool(1, ControlledWorker, start_method=start_method) + original_start = pool.start + + def start(**kwargs): + original_start(marker=str(marker), **kwargs) + workers.extend(pool.processes) + queue = pool.output_queue + assert queue is not None + original_get = queue.get + + def get(block=True, timeout=None): + if block and timeout is not None: + marker.touch() + return original_get(block=block, timeout=timeout) + + monkeypatch.setattr(queue, "get", get) + + monkeypatch.setattr(pool, "start", start) + return pool + + yield create + for worker in workers: + if worker.is_alive(): + worker.kill() + worker.join(timeout=5) + + +@pytest.mark.parametrize( + "start_method", + [method for method in ("spawn", "forkserver") if method in get_all_start_methods()], +) +@pytest.mark.parametrize("backpressure", [False, True]) +def test_worker_crash_during_result_wait(make_pool, monkeypatch, start_method, backpressure): + monkeypatch.setattr(processor, "processing_timeout", 10) + pool = make_pool(start_method) + # A second item fills the input window, exercising the backpressure read rather than the drain. + items = range(2 if backpressure else 1) + with pytest.raises(RuntimeError, match="terminated unexpectedly with code 23"): + list(pool.ordered_map(items, crash=True)) + assert pool.processes == [] + + +def test_healthy_worker_can_finish_after_a_poll(make_pool, monkeypatch): + monkeypatch.setattr(processor, "processing_timeout", 10) + pool = make_pool() + assert list(pool.ordered_map([7], delay=1.2)) == [7] + assert pool.processes == [] + + +def test_healthy_worker_keeps_the_overall_timeout(make_pool, monkeypatch): + monkeypatch.setattr(processor, "processing_timeout", 1.2) + pool = make_pool() + with pytest.raises(Empty): + list(pool.ordered_map([7], delay=5)) + assert pool.processes == []