From 3d2981872d1371d30b7a16862f545c33e25ba5cb Mon Sep 17 00:00:00 2001 From: Ramzi Alashmali <186906047+RAMZI0TO99@users.noreply.github.com> Date: Sun, 4 Oct 2026 16:11:46 +0300 Subject: [PATCH 1/2] fix: preserve explicit cuda selection in parallel workers --- fastembed/parallel_processor.py | 7 +- tests/test_parallel_cuda_forwarding.py | 121 +++++++++++++++++++++++++ 2 files changed, 127 insertions(+), 1 deletion(-) create mode 100644 tests/test_parallel_cuda_forwarding.py diff --git a/fastembed/parallel_processor.py b/fastembed/parallel_processor.py index 36688d87..1d39c1e4 100644 --- a/fastembed/parallel_processor.py +++ b/fastembed/parallel_processor.py @@ -112,6 +112,11 @@ def __init__( self.num_active_workers: BaseValue | None = None def start(self, **kwargs: Any) -> None: + """Start workers with copies of the supplied initialization options. + + The pool's cuda setting takes precedence over a cuda option in kwargs. + When device IDs are provided, assign them to workers in round-robin order. + """ 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 @@ -125,10 +130,10 @@ def start(self, **kwargs: Any) -> None: for worker_id in range(0, self.num_workers): worker_kwargs = deepcopy(kwargs) + worker_kwargs["cuda"] = self.cuda if self.device_ids: device_id = self.device_ids[worker_id % len(self.device_ids)] worker_kwargs["device_id"] = device_id - worker_kwargs["cuda"] = self.cuda assert hasattr(self.ctx, "Process") process = self.ctx.Process( diff --git a/tests/test_parallel_cuda_forwarding.py b/tests/test_parallel_cuda_forwarding.py new file mode 100644 index 00000000..8fe47cf3 --- /dev/null +++ b/tests/test_parallel_cuda_forwarding.py @@ -0,0 +1,121 @@ +"""Regression tests for forwarding device choices into parallel worker startup.""" + +import os +from multiprocessing import get_all_start_methods +from typing import Any, Iterable +from unittest.mock import Mock + +import pytest + +from fastembed.common.types import Device +from fastembed.parallel_processor import ParallelWorkerPool, Worker + + +START_METHODS = [method for method in ("spawn", "forkserver") if method in get_all_start_methods()] + + +class ReportingWorker(Worker): + """Report initialization options and process IDs from real child processes.""" + + def __init__(self, options: dict[str, Any]) -> None: + """Store the initialization options received by this worker.""" + self.options = options + + @classmethod + def start(cls, **kwargs: Any) -> "ReportingWorker": + """Construct a reporter with the options passed to child startup.""" + return cls(kwargs) + + def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]: + """Yield each indexed item with startup options and the child process ID.""" + for index, item in items: + yield index, (item, self.options, os.getpid()) + + +@pytest.mark.parametrize("start_method", START_METHODS) +@pytest.mark.parametrize("device_ids", [None, [], [2, 5]], ids=["none", "empty", "explicit"]) +@pytest.mark.parametrize("cuda", [False, True, Device.CPU, Device.CUDA, Device.AUTO]) +def test_workers_receive_cuda_selection_and_other_options( + start_method: str, device_ids: list[int] | None, cuda: bool | Device +) -> None: + """Forward CUDA choices and unrelated options with or without explicit device IDs.""" + pool = ParallelWorkerPool( + 1, ReportingWorker, start_method=start_method, device_ids=device_ids, cuda=cuda + ) + other_options = {"model_name": "reporting", "config": {"batch_size": 3}} + + results = list(pool.ordered_map(range(4), **other_options)) + + expected_options = {**other_options, "cuda": cuda} + if device_ids: + expected_options["device_id"] = device_ids[0] + assert [item for item, _, _ in results] == list(range(4)) + for _, options, pid in results: + assert options == expected_options + assert options["cuda"] is cuda + assert pid != os.getpid() + + +def test_device_ids_are_assigned_round_robin_at_worker_creation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Check device assignment at process creation without relying on job scheduling.""" + pool = ParallelWorkerPool( + 5, ReportingWorker, start_method="spawn", device_ids=[2, 5], cuda=Device.CUDA + ) + process_factory = Mock() + process_factory.return_value.is_alive.return_value = False + monkeypatch.setattr(pool.ctx, "Process", process_factory) + other_options = {"model_name": "reporting", "config": {"batch_size": 3}} + + try: + pool.start(**other_options) + + worker_options = [call.kwargs["args"][-1] for call in process_factory.call_args_list] + assert worker_options == [ + {**other_options, "cuda": Device.CUDA, "device_id": device_id} + for device_id in [2, 5, 2, 5, 2] + ] + finally: + pool.join() + for queue in (pool.input_queue, pool.output_queue): + if queue is not None: + queue.close() + queue.join_thread() + + +@pytest.mark.parametrize("start_method", START_METHODS) +@pytest.mark.parametrize("device_ids", [None, [], [2, 5]]) +def test_pool_cuda_selection_takes_precedence_over_worker_options( + start_method: str, device_ids: list[int] | None +) -> None: + """Keep the pool's CUDA choice authoritative over conflicting worker options.""" + pool = ParallelWorkerPool( + 1, ReportingWorker, start_method=start_method, device_ids=device_ids, cuda=False + ) + + results = list(pool.ordered_map(["report"], cuda=True)) + + assert len(results) == 1 + assert results[0][1]["cuda"] is False + + +@pytest.mark.parametrize("start_method", START_METHODS) +def test_two_workers_preserve_order_and_explicit_device_options(start_method: str) -> None: + """Preserve result order and assigned device options with two real workers.""" + pool = ParallelWorkerPool( + 2, ReportingWorker, start_method=start_method, device_ids=[2, 5], cuda=False + ) + + results = list(pool.ordered_map(range(10), model_name="reporting")) + + assert [item for item, _, _ in results] == list(range(10)) + for _, options, pid in results: + assert options == { + "model_name": "reporting", + "cuda": False, + "device_id": options["device_id"], + } + assert options["device_id"] in {2, 5} + assert options["cuda"] is False + assert pid != os.getpid() From c81b6caef8cdb9ecf36f8635ea0d59baf226dd39 Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Mon, 5 Oct 2026 16:10:53 +0700 Subject: [PATCH 2/2] tests: remove too heavy tests, remove a docstring --- fastembed/parallel_processor.py | 5 - tests/test_parallel_cuda_forwarding.py | 121 ------------------------- 2 files changed, 126 deletions(-) delete mode 100644 tests/test_parallel_cuda_forwarding.py diff --git a/fastembed/parallel_processor.py b/fastembed/parallel_processor.py index 1d39c1e4..4f2973aa 100644 --- a/fastembed/parallel_processor.py +++ b/fastembed/parallel_processor.py @@ -112,11 +112,6 @@ def __init__( self.num_active_workers: BaseValue | None = None def start(self, **kwargs: Any) -> None: - """Start workers with copies of the supplied initialization options. - - The pool's cuda setting takes precedence over a cuda option in kwargs. - When device IDs are provided, assign them to workers in round-robin order. - """ 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 diff --git a/tests/test_parallel_cuda_forwarding.py b/tests/test_parallel_cuda_forwarding.py deleted file mode 100644 index 8fe47cf3..00000000 --- a/tests/test_parallel_cuda_forwarding.py +++ /dev/null @@ -1,121 +0,0 @@ -"""Regression tests for forwarding device choices into parallel worker startup.""" - -import os -from multiprocessing import get_all_start_methods -from typing import Any, Iterable -from unittest.mock import Mock - -import pytest - -from fastembed.common.types import Device -from fastembed.parallel_processor import ParallelWorkerPool, Worker - - -START_METHODS = [method for method in ("spawn", "forkserver") if method in get_all_start_methods()] - - -class ReportingWorker(Worker): - """Report initialization options and process IDs from real child processes.""" - - def __init__(self, options: dict[str, Any]) -> None: - """Store the initialization options received by this worker.""" - self.options = options - - @classmethod - def start(cls, **kwargs: Any) -> "ReportingWorker": - """Construct a reporter with the options passed to child startup.""" - return cls(kwargs) - - def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]: - """Yield each indexed item with startup options and the child process ID.""" - for index, item in items: - yield index, (item, self.options, os.getpid()) - - -@pytest.mark.parametrize("start_method", START_METHODS) -@pytest.mark.parametrize("device_ids", [None, [], [2, 5]], ids=["none", "empty", "explicit"]) -@pytest.mark.parametrize("cuda", [False, True, Device.CPU, Device.CUDA, Device.AUTO]) -def test_workers_receive_cuda_selection_and_other_options( - start_method: str, device_ids: list[int] | None, cuda: bool | Device -) -> None: - """Forward CUDA choices and unrelated options with or without explicit device IDs.""" - pool = ParallelWorkerPool( - 1, ReportingWorker, start_method=start_method, device_ids=device_ids, cuda=cuda - ) - other_options = {"model_name": "reporting", "config": {"batch_size": 3}} - - results = list(pool.ordered_map(range(4), **other_options)) - - expected_options = {**other_options, "cuda": cuda} - if device_ids: - expected_options["device_id"] = device_ids[0] - assert [item for item, _, _ in results] == list(range(4)) - for _, options, pid in results: - assert options == expected_options - assert options["cuda"] is cuda - assert pid != os.getpid() - - -def test_device_ids_are_assigned_round_robin_at_worker_creation( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Check device assignment at process creation without relying on job scheduling.""" - pool = ParallelWorkerPool( - 5, ReportingWorker, start_method="spawn", device_ids=[2, 5], cuda=Device.CUDA - ) - process_factory = Mock() - process_factory.return_value.is_alive.return_value = False - monkeypatch.setattr(pool.ctx, "Process", process_factory) - other_options = {"model_name": "reporting", "config": {"batch_size": 3}} - - try: - pool.start(**other_options) - - worker_options = [call.kwargs["args"][-1] for call in process_factory.call_args_list] - assert worker_options == [ - {**other_options, "cuda": Device.CUDA, "device_id": device_id} - for device_id in [2, 5, 2, 5, 2] - ] - finally: - pool.join() - for queue in (pool.input_queue, pool.output_queue): - if queue is not None: - queue.close() - queue.join_thread() - - -@pytest.mark.parametrize("start_method", START_METHODS) -@pytest.mark.parametrize("device_ids", [None, [], [2, 5]]) -def test_pool_cuda_selection_takes_precedence_over_worker_options( - start_method: str, device_ids: list[int] | None -) -> None: - """Keep the pool's CUDA choice authoritative over conflicting worker options.""" - pool = ParallelWorkerPool( - 1, ReportingWorker, start_method=start_method, device_ids=device_ids, cuda=False - ) - - results = list(pool.ordered_map(["report"], cuda=True)) - - assert len(results) == 1 - assert results[0][1]["cuda"] is False - - -@pytest.mark.parametrize("start_method", START_METHODS) -def test_two_workers_preserve_order_and_explicit_device_options(start_method: str) -> None: - """Preserve result order and assigned device options with two real workers.""" - pool = ParallelWorkerPool( - 2, ReportingWorker, start_method=start_method, device_ids=[2, 5], cuda=False - ) - - results = list(pool.ordered_map(range(10), model_name="reporting")) - - assert [item for item, _, _ in results] == list(range(10)) - for _, options, pid in results: - assert options == { - "model_name": "reporting", - "cuda": False, - "device_id": options["device_id"], - } - assert options["device_id"] in {2, 5} - assert options["cuda"] is False - assert pid != os.getpid()