From 98a52267a2d192e0d5d2d5c62ba874247ddd4896 Mon Sep 17 00:00:00 2001 From: buduoqiu Date: Wed, 29 Jul 2026 06:38:34 +0800 Subject: [PATCH] feat(python): configure streaming transform parallelism (#3699) ## Summary - add a keyword-only `transform_parallelism` option to `StreamingDataset` - preserve CPU auto-detection by default and fall back to one worker when unavailable - apply the configured limit to both the transform executor and concurrency semaphore - document and test explicit, default, fallback, and invalid values ## Testing - `uv run --extra tests --with torch pytest python/tests/test_elastic_dataloader.py -q` (`136 passed`) - `uvx ruff check python/lancedb/streaming.py python/tests/test_elastic_dataloader.py` - `uvx ruff format --check python/lancedb/streaming.py python/tests/test_elastic_dataloader.py` - `git diff --check origin/main...HEAD` Closes #3695 Co-authored-by: buduoqiu --- python/python/lancedb/streaming.py | 29 ++++++++++----- .../python/tests/test_elastic_dataloader.py | 36 +++++++++++++++++++ 2 files changed, 57 insertions(+), 8 deletions(-) diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index 9ffc612d9..525ed3d63 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -59,9 +59,10 @@ class StreamingDataset(IterableDataset): - **Stage 1 (I/O)**: one thread pool with ``num_splits * prefetch_batches`` workers fetches raw ``RecordBatch`` objects from LanceDB in parallel across all splits and places them in a per-split raw-batch queue. - - **Stage 2 (transform)**: a second thread pool with ``os.cpu_count()`` - workers picks up raw batches, applies the transform, and places the - results in a per-split cooked-row queue. + - **Stage 2 (transform)**: a second thread pool with + ``transform_parallelism`` workers picks up raw batches, applies the + transform, and places the results in a per-split cooked-row queue. By + default, the number of workers is determined by ``os.cpu_count()``. The main thread round-robins over the cooked queues, yielding one row per split per cycle. @@ -122,6 +123,10 @@ class StreamingDataset(IterableDataset): are yielded. Receives one batch at a time and must return an iterable whose length equals the number of rows in the batch. When ``None`` (the default) rows are returned as plain Python dicts. + transform_parallelism: + Maximum number of transforms to run concurrently. Must be greater + than zero. When ``None`` (the default), uses ``os.cpu_count()`` or 1 + when the CPU count is unavailable. worker_info_override: If set, used in place of ``torch.utils.data.get_worker_info()`` to determine the DataLoader worker assignment. Intended for unit tests @@ -146,6 +151,7 @@ class StreamingDataset(IterableDataset): shuffle_clump_size: Optional[int] = None, filter: Optional[str] = None, transform: Optional[Callable] = None, + transform_parallelism: Optional[int] = None, connection_factory: Optional[Callable[[str], Any]] = None, worker_info_override=None, ): @@ -159,6 +165,8 @@ class StreamingDataset(IterableDataset): f"num_splits ({num_splits}) must be divisible by " f"world_size ({world_size})" ) + if transform_parallelism is not None and transform_parallelism <= 0: + raise ValueError("transform_parallelism must be greater than 0") self._table = table self._num_splits = num_splits @@ -173,6 +181,7 @@ class StreamingDataset(IterableDataset): self._shuffle_clump_size = shuffle_clump_size self._filter = filter self._transform = transform + self._transform_parallelism = transform_parallelism self._connection_factory = connection_factory self._worker_info_override = worker_info_override @@ -284,7 +293,11 @@ class StreamingDataset(IterableDataset): batch_size = self._read_batch_size max_prefetch = self._prefetch_batches - cpu_workers = os.cpu_count() or 1 + transform_workers = ( + self._transform_parallelism + if self._transform_parallelism is not None + else (os.cpu_count() or 1) + ) final_transform = ( self._transform if self._transform is not None else Transforms.arrow2python ) @@ -296,8 +309,8 @@ class StreamingDataset(IterableDataset): tx_pending = [deque() for _ in range(n)] # Future[list[Any]] cooked = [deque() for _ in range(n)] # rows ready to yield - # Limit simultaneous transforms to cpu_workers across all splits. - tx_semaphore = threading.Semaphore(cpu_workers) + # Limit simultaneous transforms to transform_workers across all splits. + tx_semaphore = threading.Semaphore(transform_workers) # ── Stage 1 helpers ─────────────────────────────────────────────────── @@ -369,7 +382,7 @@ class StreamingDataset(IterableDataset): _advance(i) elif raw_batches[i]: # Acquire a transform slot (may block briefly if all - # cpu_workers are busy with other splits). + # transform_workers are busy with other splits). tx_semaphore.acquire() batch = raw_batches[i].popleft() tx_pending[i].append(tx_pool.submit(_tx_call_guarded, batch)) @@ -383,7 +396,7 @@ class StreamingDataset(IterableDataset): # ── Main loop ───────────────────────────────────────────────────────── with ThreadPoolExecutor(max_workers=n * max_prefetch) as io_pool: - with ThreadPoolExecutor(max_workers=cpu_workers) as tx_pool: + with ThreadPoolExecutor(max_workers=transform_workers) as tx_pool: self._raw_batches_ref = raw_batches self._cooked_ref = cooked self._fetch_head_ref = fetch_head diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index 6a98bed0c..22918082b 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/python/python/tests/test_elastic_dataloader.py @@ -1333,6 +1333,42 @@ def test_transform_none_yields_dicts(lance_table): assert all("id" in item for item in items) +@pytest.mark.parametrize( + ("configured", "detected", "expected"), + [(2, 8, 2), (None, 3, 3), (None, None, 1)], +) +def test_transform_parallelism_configures_executor( + lance_table, monkeypatch, configured, detected, expected +): + """Explicit transform parallelism overrides the detected CPU count.""" + real_executor = streaming.ThreadPoolExecutor + monkeypatch.setattr(streaming.os, "cpu_count", lambda: detected) + + with patch.object(streaming, "ThreadPoolExecutor", wraps=real_executor) as executor: + list( + StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + transform_parallelism=configured, + ) + ) + + assert executor.call_args_list[-1].kwargs["max_workers"] == expected + + +@pytest.mark.parametrize("transform_parallelism", [0, -1]) +def test_transform_parallelism_must_be_positive(lance_table, transform_parallelism): + with pytest.raises( + ValueError, match="transform_parallelism must be greater than 0" + ): + StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + transform_parallelism=transform_parallelism, + ) + + def test_filter_limits_rows(tmp_path): """A filter expression is applied to the permutation so only matching rows are yielded. IDs 0..59 pass ``id < 60``; the other 60 are excluded."""