diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index a94b3c89f..8fc863156 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -165,12 +165,10 @@ class StreamingDataset(IterableDataset): blocks: exhausted splits emit padding, while tokens beyond the budget are left out of the epoch. This fixed per-split budget keeps packed iteration and checkpoints independent of rank and worker topology. - Pass ``"auto"`` to estimate the budget by reading only the token column - for approximately 1% of the rows, with a target cap of 100,000 rows and - at least one row from every logical split. Sampling is batched and evenly - distributed across logical splits. This is a convenience for smaller - workloads. For large-scale training, materialize a ``token_count`` column, - calculate an explicit block budget once, and pass that integer instead. + Pass ``"auto"`` to calculate a corpus-level budget by scanning only the + token column in bounded batches. For large-scale training, materialize a + ``token_count`` column, calculate an explicit block budget once, and pass + that integer instead. on_transform_error: What to do when the transform raises an exception: @@ -398,7 +396,7 @@ class StreamingDataset(IterableDataset): ) def _estimate_blocks_per_epoch(self) -> int: - """Estimate a fixed packed-block budget from a bounded token sample.""" + """Calculate a fixed packed-block budget from the complete token column.""" if self._pack_sequences is None or not self._columns: raise RuntimeError( "packing must be configured before estimating its budget" @@ -406,64 +404,45 @@ class StreamingDataset(IterableDataset): pack_len = self._pack_sequences token_column = self._columns[0] - sample_cap_per_split = max(1, 100_000 // self._num_splits) - sampled_tokens = 0 - total_sampled = 0 - total_rows = 0 warnings.warn( - "blocks_per_epoch='auto' is a convenience estimate that samples full " - "token lists and repeats for every process constructing the dataset. It " - "is recommended to materialize a token_count column, calculate an " - "explicit blocks_per_epoch once, and pass that integer instead.", + "blocks_per_epoch='auto' scans the complete token column on every " + "distributed rank. For large-scale training, materialize a token_count " + "column, calculate blocks_per_epoch once, and pass that integer instead.", # TODO: Link to docs example that shows how to do this ) - for split in range(self._num_splits): - permutation = Permutation.from_tables( - self._table, self._perm_table, split=split - ) - permutation = permutation.select_columns([token_column]) - permutation = permutation.with_transform(Transforms.arrow2arrow) - split_rows = permutation.num_rows - if split_rows == 0: - raise ValueError( - "blocks_per_epoch='auto' cannot estimate an empty dataset" - ) + fragments = list(self._table.to_lance().get_fragments()) - # Roughly 1% per logical split, with at least one row from each - # split and a global cap of approximately 100,000 rows. - sample_rows = min( - split_rows, - max(1, min((split_rows + 99) // 100, sample_cap_per_split)), + def fragment_total(fragment) -> tuple[int, int]: + token_count = 0 + row_count = 0 + scanner = fragment.scanner( + columns=[token_column], + filter=self._filter, + batch_size=2048, ) - sample_batch_size = max(1, self._read_batch_size) - for start in range(0, sample_rows, sample_batch_size): - stop = min(start + sample_batch_size, sample_rows) - # approx mid of each bucket - batch_offsets = [ - ((2 * index + 1) * split_rows) // (2 * sample_rows) - for index in range(start, stop) - ] - batch = permutation.__getitems__(batch_offsets) + for batch in scanner.to_batches(): lengths = pc.list_value_length(batch.column(0)) if lengths.null_count: raise ValueError("pack_sequences does not support null token lists") - sampled_tokens += int(pc.sum(lengths).as_py()) + token_count += int(pc.sum(lengths).as_py()) + row_count += batch.num_rows + return token_count, row_count - total_sampled += sample_rows - total_rows += split_rows + token_count = 0 + row_count = 0 + with ThreadPoolExecutor(max_workers=min(16, len(fragments))) as pool: + for fragment_tokens, fragment_rows in pool.map(fragment_total, fragments): + token_count += fragment_tokens + row_count += fragment_rows - # Pool all samples into one global average. Each document contributes - # one EOS token. Round down so estimation - # favors leaving a short tail unused over emitting all-padding blocks. - estimated_tokens = ( - (sampled_tokens + total_sampled) * total_rows // total_sampled - ) - estimated_blocks = estimated_tokens // pack_len + # Each document contributes one EOS token. Round down to a complete + # split cycle, favoring a short unused tail over all-padding blocks. + blocks = (token_count + row_count) // pack_len blocks_per_epoch = max( self._num_splits, - estimated_blocks - estimated_blocks % self._num_splits, + blocks - blocks % self._num_splits, ) return blocks_per_epoch diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index a1e336891..3bf8d5979 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/python/python/tests/test_elastic_dataloader.py @@ -1998,28 +1998,37 @@ def test_pack_sequences_pads_lagging_splits(tmp_path): assert sharded == input_ids -def test_pack_sequences_auto_estimates_in_bounded_batches(tmp_path): - table = _create_token_table(tmp_path, [[1, 2, 3, 4]] * 1_000) - batch_sizes = [] - getitems = streaming.Permutation.__getitems__ +def test_pack_sequences_auto_scans_filtered_token_column(tmp_path): + db = lancedb.connect(tmp_path) + table = db.create_table( + "tokens", + pa.table( + { + "tokens": pa.array([[1] * 4, [2] * 9], type=pa.list_(pa.int64())), + "keep": [True, False], + } + ), + ) + table.add( + pa.table( + { + "tokens": pa.array([[3] * 4, [4] * 9], type=pa.list_(pa.int64())), + "keep": [True, False], + } + ) + ) - def record_batch_size(permutation, indices): - batch_sizes.append(len(indices)) - return getitems(permutation, indices) + with pytest.warns(UserWarning, match="complete token column"): + dataset = _packed_dataset( + table, + 5, + blocks_per_epoch="auto", + num_splits=2, + filter="keep", + ) - with patch.object(streaming.Permutation, "__getitems__", record_batch_size): - with pytest.warns(UserWarning, match="materialize a token_count column"): - dataset = _packed_dataset( - table, - 5, - blocks_per_epoch="auto", - num_splits=2, - read_batch_size=3, - ) - - assert dataset.state_dict()["blocks_per_epoch"] == 1_000 - assert sum(batch_sizes) == 10 - assert max(batch_sizes) <= 3 + # Two kept documents contain 8 tokens plus 2 EOS tokens: two blocks. + assert dataset.state_dict()["blocks_per_epoch"] == 2 def test_pack_sequences_checkpoint_resumes_on_new_topology(tmp_path):