From ec53c82ef8de9d6811605bcafa76ad33f7a04c28 Mon Sep 17 00:00:00 2001 From: Weston Pace Date: Mon, 19 Oct 2026 11:09:27 -0700 Subject: [PATCH] fix(python): bound StreamingDataset's raw_batches queue against io_queue_depth (#4530) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - `StreamingDataset`'s Stage 1 (I/O) queue (`raw_batches`) had no cap: `_fill_io` only limited *in-flight* fetch futures (`io_pending`) to `io_queue_depth`, but every completed fetch was unconditionally drained into `raw_batches` and the freed slot immediately refilled, so I/O could race arbitrarily far ahead of a slow transform stage with no bound on buffered memory. - `transform_queue_depth`'s admission check (`_try_submit_tx`) couldn't compensate either, since it only counts `cooked` rows and in-flight transform futures — it never looks at `raw_batches`. - Fix: `_fill_io` now counts both `io_pending` and `raw_batches` against the `io_queue_depth` budget, matching the parameter's documented "batches to keep in flight per split" contract. ## Test plan - [x] Added `test_raw_queue_depth_bounded_when_transform_is_the_bottleneck` (deterministic, via synchronous I/O) asserting `raw_queue_depth <= io_queue_depth * read_batch_size` even when the transform stage is the bottleneck. - [x] Full `python/python/tests/test_elastic_dataloader.py` suite passes (190/190), confirming no change to sample ordering, counts, or checkpoint/resumability behavior. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Sonnet 5 --- python/python/lancedb/streaming.py | 24 ++++-- .../python/tests/test_elastic_dataloader.py | 86 +++++++++++++++++++ 2 files changed, 105 insertions(+), 5 deletions(-) diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index 76b3702a1..91ca6ce28 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -261,10 +261,15 @@ class StreamingDataset(IterableDataset): storage) at the cost of higher memory usage per split buffer. Defaults to ``DEFAULT_READ_BATCH_SIZE`` (64). io_queue_depth: - Number of I/O batches to keep in flight per split. Higher values - overlap storage latency with transform and training compute at the cost - of more memory and threads. Must be greater than zero. Defaults to - ``DEFAULT_PREFETCH_BATCHES`` (4). + Number of fetched batches to keep ahead of the transform stage per + split: this bounds the combined count of in-flight fetches and + already-fetched batches still waiting for a transform slot. Higher + values overlap storage latency with transform and training compute + at the cost of more memory and threads, but I/O stops running ahead + once this many batches are buffered, applying backpressure when the + transform stage (or ``transform_queue_depth``) is the bottleneck. + Must be greater than zero. Defaults to ``DEFAULT_PREFETCH_BATCHES`` + (4). columns: Optional list of column names to read. When set, only those columns are fetched from storage; all others are omitted. ``None`` (the @@ -836,8 +841,17 @@ class StreamingDataset(IterableDataset): io_pending[i].append((abs_start, io_pool.submit(_io_call, perm_i, indices))) def _fill_io(i: int) -> None: + # Budget covers both in-flight fetches (io_pending) and batches + # that have already landed but are still waiting for a transform + # slot (raw_batches). Without raw_batches in the count, a slow + # transform stage never applies backpressure to Stage 1: each + # completed fetch is drained into raw_batches and immediately + # replaced by a new one, so I/O keeps racing arbitrarily far + # ahead of the transform stage instead of merely overlapping + # storage latency with it. while ( - len(io_pending[i]) < io_queue_depth and fetch_head[i] < split_sizes[i] + len(io_pending[i]) + len(raw_batches[i]) < io_queue_depth + and fetch_head[i] < split_sizes[i] ): _submit_io(i) diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index da27e5bfc..cc6874e6c 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/python/python/tests/test_elastic_dataloader.py @@ -2036,6 +2036,92 @@ def test_transform_queue_depth_bounds_cooked_rows(lance_table): ) +def test_raw_queue_depth_bounded_when_transform_is_the_bottleneck(tmp_path): + """io_queue_depth bounds raw_batches even when transform is the bottleneck. + + Regression test for a bug where Stage 1 (I/O) applied no backpressure at + all: _fill_io only capped io_pending (in-flight fetches) at + io_queue_depth, but every completed fetch was drained into raw_batches + unconditionally (_drain_io) and immediately replaced by a new fetch + (_fill_io again), so I/O kept racing arbitrarily far ahead of a slow + transform stage -- raw_batches had no cap of its own. + transform_queue_depth did not help either: its admission check in + _try_submit_tx only counts cooked rows and in-flight transform futures + (`len(cooked[i]) + len(tx_pending[i]) * batch_size`); it never looks at + len(raw_batches[i]), so it cannot throttle Stage 1. + + The fix makes _fill_io count both io_pending and raw_batches against the + io_queue_depth budget, so raw_queue_depth should never exceed + io_queue_depth * read_batch_size per split, regardless of how slow the + transform stage is. + + I/O is forced to run synchronously (same technique as + test_transform_queue_depth_does_not_admit_at_capacity_minus_one) so the + assertion is deterministic regardless of machine load, rather than + relying on a real-time race between I/O and transform threads. Only the + transform runs on a real background thread (with a short sleep) so the + consumer genuinely blocks waiting for it, letting Stage 1 attempt to run + ahead between blocking waits. + """ + import concurrent.futures as cf + from concurrent.futures import ThreadPoolExecutor + from unittest.mock import patch + import pyarrow as pa + import time + + n_rows = 200 + read_batch_size = 10 + io_queue_depth = 2 + transform_queue_depth = 2 + max_allowed = io_queue_depth * read_batch_size + + db = lancedb.connect(tmp_path) + table = db.create_table("t", pa.table({"id": list(range(n_rows))})) + + def slow_transform(batch: pa.RecordBatch) -> list: + time.sleep(0.005) + return batch.column("id").to_pylist() + + original_submit = ThreadPoolExecutor.submit + + def synchronous_io_submit(self, fn, *args, **kwargs): + if getattr(fn, "__name__", "") == "_io_call": + fut = cf.Future() + try: + fut.set_result(fn(*args, **kwargs)) + except Exception as exc: + fut.set_exception(exc) + return fut + return original_submit(self, fn, *args, **kwargs) + + with patch.object(ThreadPoolExecutor, "submit", synchronous_io_submit): + ds = StreamingDataset( + table, + num_splits=1, + shuffle_seed=SHUFFLE_SEED, + read_batch_size=read_batch_size, + io_queue_depth=io_queue_depth, + transform_queue_depth=transform_queue_depth, + transform_parallelism=1, + transform=slow_transform, + ) + + peak_raw = 0 + ids = [] + for item in ds: + ids.append(item) # slow_transform yields bare ids, not dicts + if ds.raw_queue_depth > peak_raw: + peak_raw = ds.raw_queue_depth + + assert peak_raw <= max_allowed, ( + f"raw_queue_depth peaked at {peak_raw}, expected <= {max_allowed} " + "(io_queue_depth * read_batch_size)" + ) + # The backpressure fix must not come at the cost of dropping or + # duplicating rows. + assert sorted(ids) == list(range(n_rows)) + + def test_transform_queue_depth_does_not_admit_at_capacity_minus_one(tmp_path): """Admission requires a full read_batch_size of free space, not just one slot.