fix: clean up partially initialized worker pools

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:
         """

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