diff --git a/fastembed/parallel_processor.py b/fastembed/parallel_processor.py index 4f2973aa..a1d2543b 100644 --- a/fastembed/parallel_processor.py +++ b/fastembed/parallel_processor.py @@ -204,8 +204,6 @@ def semi_ordered_map( 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" if completed: self.join() else: @@ -218,19 +216,26 @@ def semi_ordered_map( # Nothing reads the input pipe anymore, so the feeder thread can be stuck # writing to it, holding every item it hasn't sent. On POSIX, closing our # read end fails that write with EPIPE; Windows also needs to cancel the send. - terminate_broken = getattr(self.input_queue, "_terminate_broken", None) - if terminate_broken is not None: - terminate_broken() - else: - self.input_queue._reader.close() # type: ignore[attr-defined] - self.input_queue.close() - self.output_queue.close() - if self.emergency_shutdown: - self.input_queue.cancel_join_thread() - self.output_queue.cancel_join_thread() - else: - self.input_queue.join_thread() - self.output_queue.join_thread() + if self.input_queue is not None: + terminate_broken = getattr(self.input_queue, "_terminate_broken", None) + if terminate_broken is not None: + terminate_broken() + else: + self.input_queue._reader.close() # type: ignore[attr-defined] + # start() can fail before either queue has been allocated. Clean up only + # initialized queues, without masking the original startup exception. + for queue in (self.input_queue, self.output_queue): + if queue is not None: + # Without a local feeder, close() cannot close the pipe endpoints. + # Workers have already stopped, so neither queue is still in use. + if queue._thread is None: # type: ignore[attr-defined] + queue._reader.close() # type: ignore[attr-defined] + queue._writer.close() # type: ignore[attr-defined] + queue.close() + if self.emergency_shutdown: + queue.cancel_join_thread() + else: + queue.join_thread() def check_worker_health(self) -> None: """ diff --git a/tests/test_parallel_processor.py b/tests/test_parallel_processor.py index c9d4fb79..a55dfa45 100644 --- a/tests/test_parallel_processor.py +++ b/tests/test_parallel_processor.py @@ -1,6 +1,7 @@ import threading from itertools import count from multiprocessing import get_all_start_methods +from unittest.mock import Mock import pytest @@ -16,10 +17,104 @@ def process(self, items): yield from items +@pytest.mark.parametrize("failed_queue", [1, 2]) +def test_queue_creation_failure_preserves_error_and_cleans_up(monkeypatch, failed_queue): + pool = ParallelWorkerPool(2, EchoWorker, start_method="spawn") + queue = Mock() + error = OSError("queue allocation failed") + side_effect = [error] if failed_queue == 1 else [queue, error] + monkeypatch.setattr(pool.ctx, "Queue", Mock(side_effect=side_effect)) + + with pytest.raises(OSError, match="queue allocation failed") as raised: + list(pool.ordered_map([1])) + + assert raised.value is error + assert pool.processes == [] + if failed_queue == 2: + queue.close.assert_called_once_with() + queue.cancel_join_thread.assert_called_once_with() + + +@pytest.mark.parametrize("legacy_cleanup", [False, True]) +def test_second_queue_failure_closes_real_input_queue(monkeypatch, legacy_cleanup): + pool = ParallelWorkerPool(2, EchoWorker, start_method="spawn") + queue = pool.ctx.Queue(pool.queue_size) + if legacy_cleanup: + monkeypatch.setattr(queue, "_terminate_broken", None, raising=False) + monkeypatch.setattr( + pool.ctx, "Queue", Mock(side_effect=[queue, OSError("output queue allocation failed")]) + ) + try: + with pytest.raises(OSError, match="output queue allocation failed"): + list(pool.ordered_map([1])) + assert queue._closed + assert queue._joincancelled + assert queue._reader.closed + assert queue._writer.closed + finally: + queue.close() + queue._reader.close() + queue._writer.close() + + +@pytest.mark.parametrize("failure_stage", ["value", "first_worker", "second_worker"]) +def test_startup_failure_closes_real_queue_endpoints(monkeypatch, failure_stage): + pool = ParallelWorkerPool(2, EchoWorker, start_method="spawn") + error = OSError("pool startup failed") + workers = [] + if failure_stage == "value": + monkeypatch.setattr(pool.ctx, "Value", Mock(side_effect=error)) + else: + create_process = pool.ctx.Process + + def process_factory(*args, **kwargs): + process = create_process(*args, **kwargs) + workers.append(process) + failed_worker = 1 if failure_stage == "first_worker" else 2 + if len(workers) == failed_worker: + monkeypatch.setattr(process, "start", Mock(side_effect=error)) + return process + + monkeypatch.setattr(pool.ctx, "Process", process_factory) + + try: + with pytest.raises(OSError, match="pool startup failed") as raised: + list(pool.ordered_map([1])) + assert raised.value is error + assert pool.processes == [] + assert not any(worker.is_alive() for worker in workers) + for queue in (pool.input_queue, pool.output_queue): + assert queue is not None + assert queue._closed + assert queue._joincancelled + assert queue._reader.closed + assert queue._writer.closed + finally: + for queue in (pool.input_queue, pool.output_queue): + if queue is not None: + queue.close() + queue._reader.close() + queue._writer.close() + + +@pytest.mark.parametrize("items", [[], [1, 2, 3]]) +def test_completed_mapping_stops_workers(items): + pool = ParallelWorkerPool(2, EchoWorker, start_method="spawn") + assert list(pool.ordered_map(items)) == items + assert pool.processes == [] + for queue in (pool.input_queue, pool.output_queue): + assert queue is not None + assert queue._reader.closed + assert queue._writer.closed + + # 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" +@pytest.mark.parametrize( + "start_method", + [method for method in ("spawn", "forkserver") if method in get_all_start_methods()], +) +def test_closing_partially_consumed_iterator_stops_workers(start_method): 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()) @@ -34,6 +129,9 @@ def test_closing_partially_consumed_iterator_stops_workers(): try: assert not closer.is_alive(), "closing the iterator hung" assert not any(worker.is_alive() for worker in workers) + assert pool.output_queue is not None + assert pool.output_queue._reader.closed + assert pool.output_queue._writer.closed finally: for worker in workers: if worker.is_alive():