diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index 6951fa0bc..c54940ab7 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -24,11 +24,16 @@ import os import random import threading import time +import warnings from collections import deque from concurrent.futures import ThreadPoolExecutor +from copy import deepcopy from multiprocessing import RawArray -from typing import Any, Callable, Iterator, Optional, Union +from typing import Any, Callable, cast, Iterator, Literal, Optional, Union +import pyarrow as pa +import pyarrow.compute as pc +import torch from torch.utils.data import IterableDataset, get_worker_info from .permutation import ( @@ -132,6 +137,39 @@ class StreamingDataset(IterableDataset): 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. + pack_sequences: + Sequence-packing mode: token lists from consecutive documents are + joined with ``eos_id`` and sliced into blocks of this many tokens. + Each item is then a dict of two ``(pack_sequences,)`` LongTensors — + ``input_ids`` and ``doc_ids`` (per-position document index within + the block, for block-diagonal masks or position-id resets). + * Packing happens independently per owned split and preserves per-split + resume state. + * When a split cannot fill a real block for a cycle but an owned sibling + still can, or when only a short tail remains at epoch end, the short buffer + is padded to ``pack_sequences`` with ``pad_id`` so every local cycle emits + one block per owned split. + * ``eos_id``, ``pad_id``, and ``columns`` naming a single integer-list + column are required; incompatible with ``transform``. + eos_id: + Separator token id between packed documents. Required with + ``pack_sequences``, ignored otherwise. + pad_id: + Padding token id used to complete blocks when a split runs out of + real tokens mid-cycle or at epoch end. Required with + ``pack_sequences``, ignored otherwise. It must be reserved for padding: + padding positions retain the preceding document's ``doc_id`` (or zero + in an all-padding block), so callers must mask them separately using + ``input_ids == pad_id``. + blocks_per_epoch: + Total number of packed blocks emitted globally per epoch. Required with + ``pack_sequences``. An integer must be divisible by ``num_splits``. + Every logical split emits exactly ``blocks_per_epoch / num_splits`` + 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 topology. + Pass ``"auto"`` to estimate a corpus-level budget from a bounded sample + of token lists. The estimate may be inaccurate. on_transform_error: What to do when the transform raises an exception: @@ -210,6 +248,10 @@ class StreamingDataset(IterableDataset): filter: Optional[str] = None, transform: Optional[Callable] = None, transform_parallelism: Optional[int] = None, + pack_sequences: Optional[int] = None, + eos_id: Optional[int] = None, + pad_id: Optional[int] = None, + blocks_per_epoch: Optional[Union[int, Literal["auto"]]] = None, on_transform_error: Union[str, Callable[[Exception], bool]] = "raise", transform_queue_depth: Optional[int] = None, connection_factory: Optional[Callable[[str], Any]] = None, @@ -237,6 +279,55 @@ class StreamingDataset(IterableDataset): raise ValueError("io_queue_depth must be greater than 0") if transform_parallelism is not None and transform_parallelism <= 0: raise ValueError("transform_parallelism must be greater than 0") + if pack_sequences is not None: + if pack_sequences <= 0: + raise ValueError("pack_sequences must be greater than 0") + if eos_id is None: + raise ValueError("eos_id is required when pack_sequences is set") + if pad_id is None: + raise ValueError("pad_id is required when pack_sequences is set") + if blocks_per_epoch is None: + raise ValueError( + "blocks_per_epoch is required when pack_sequences is set" + ) + if blocks_per_epoch != "auto": + if not isinstance(blocks_per_epoch, int) or isinstance( + blocks_per_epoch, bool + ): + raise ValueError( + "blocks_per_epoch must be a positive integer or 'auto'" + ) + if blocks_per_epoch <= 0: + raise ValueError("blocks_per_epoch must be greater than 0") + if blocks_per_epoch % num_splits != 0: + raise ValueError( + f"blocks_per_epoch ({blocks_per_epoch}) must be divisible by " + f"num_splits ({num_splits})" + ) + if transform is not None: + raise ValueError("transform cannot be combined with pack_sequences") + if columns is None or len(columns) != 1: + raise ValueError( + "pack_sequences requires columns to name exactly one " + "list-typed column of token ids" + ) + field = table.schema.field(columns[0]) + if not ( + pa.types.is_list(field.type) + or pa.types.is_large_list(field.type) + or pa.types.is_fixed_size_list(field.type) + ): + raise ValueError( + f"pack_sequences requires a list-typed token column; " + f"{columns[0]} has type {field.type}" + ) + if not pa.types.is_integer(field.type.value_type): + raise ValueError( + "pack_sequences requires a token column with integer values; " + f"{columns[0]} has value type {field.type.value_type}" + ) + elif blocks_per_epoch is not None: + raise ValueError("blocks_per_epoch requires pack_sequences") if on_transform_error not in ("raise", "skip", "warn") and not callable( on_transform_error ): @@ -261,11 +352,20 @@ class StreamingDataset(IterableDataset): self._filter = filter self._transform = transform self._transform_parallelism = transform_parallelism + self._pack_sequences = pack_sequences + self._eos_id = eos_id + self._pad_id = pad_id + self._blocks_per_epoch = blocks_per_epoch self._on_transform_error = on_transform_error self._transform_queue_depth = transform_queue_depth self._connection_factory = connection_factory self._worker_info_override = worker_info_override + # Packing resume state: permutation positions and partial-block buffers. + self._pack_consumed: list[int] = [0] * num_splits + self._pack_buffers: dict[int, dict[str, list[int]]] = {} + self._pack_blocks_emitted: list[int] = [0] * num_splits + # Live references to pipeline state, set only while __iter__ is running # in the same process. Used by the observability properties when the # DataLoader runs with num_workers=0. @@ -315,6 +415,9 @@ class StreamingDataset(IterableDataset): else: self._perm_table = builder.split_sequential(fixed=num_splits).execute() + if self._blocks_per_epoch == "auto": + self._blocks_per_epoch = self._estimate_blocks_per_epoch() + # Contiguous block of global split indices assigned to this rank. splits_per_rank = num_splits // world_size rank_start = rank * splits_per_rank @@ -322,6 +425,71 @@ class StreamingDataset(IterableDataset): range(rank_start, rank_start + splits_per_rank) ) + def _estimate_blocks_per_epoch(self) -> int: + """Estimate a fixed packed-block budget from a bounded token sample.""" + # TODO: Replace this fallback with Lance's dedicated exact token-count + # estimation API once it is available. + if self._pack_sequences is None or not self._columns: + raise RuntimeError( + "packing must be configured before estimating its budget" + ) + + 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 + rng = random.Random(self._shuffle_seed) + + warnings.warn( + "blocks_per_epoch='auto' uses an approximate token-count sample; " + "pass an explicit value for exact epoch sizing", + ) + + 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" + ) + + # Sample roughly 1% from each logical split, with at least one row + # per split and a global target cap of 100,000 rows. + sample_rows = min( + split_rows, + max(1, min((split_rows + 99) // 100, sample_cap_per_split)), + ) + sample_offsets = sorted(rng.sample(range(split_rows), sample_rows)) + sample_batch_size = max(1, self._read_batch_size) + for start in range(0, sample_rows, sample_batch_size): + batch = permutation.__getitems__( + sample_offsets[start : start + sample_batch_size] + ) + 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()) + + total_sampled += sample_rows + total_rows += split_rows + + # Pool the samples into one global average. Each document contributes + # one EOS token. + estimated_tokens = ( + (sampled_tokens + total_sampled) * total_rows // total_sampled + ) + blocks = estimated_tokens // pack_len + return max( + self._num_splits, + blocks - blocks % self._num_splits, + ) + def _resolve_my_splits(self) -> list[int]: """Return the split indices this instance should read in __iter__.""" torch_worker_info = get_worker_info() @@ -372,8 +540,14 @@ class StreamingDataset(IterableDataset): ) if self._columns is not None: perm = perm.select_columns(self._columns) - perm = perm.with_transform(lambda batch: batch) - start_pos = self._resume_positions.get(split_idx, self._resume_offset) + perm = perm.with_transform(Transforms.arrow2arrow) + # Both modes resume from absolute permutation positions. Packing + # stores them separately because it also checkpoints partial blocks. + start_pos = ( + self._pack_consumed[split_idx] + if self._pack_sequences is not None + else self._resume_positions.get(split_idx, self._resume_offset) + ) if start_pos > 0: perm = perm.with_skip(start_pos) initial_positions.append(start_pos) @@ -395,9 +569,24 @@ class StreamingDataset(IterableDataset): 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 - ) + final_transform: Callable[[pa.RecordBatch], Any] + if self._pack_sequences is not None: + # Packing consumes raw token lists, one per document. + def arrow_tokens(batch: pa.RecordBatch) -> list[list[int]]: + token_column = batch.column(0) + if token_column.null_count or token_column.flatten().null_count: + raise ValueError( + "pack_sequences does not support null token lists or values" + ) + return cast(list[list[int]], token_column.to_pylist()) + + final_transform = arrow_tokens + else: + final_transform = ( + self._transform + if self._transform is not None + else Transforms.arrow2python + ) # None means no limit; otherwise cap rows per split to # transform_queue_depth batches worth (including in-flight transforms). max_cooked_rows = ( @@ -574,6 +763,82 @@ class StreamingDataset(IterableDataset): else: break # split exhausted + def _update_stats(*, idle: bool = False) -> None: + """Refresh pipeline statistics visible to the parent process.""" + ws = self._worker_stats + ws[0] = sum(split_sizes[j] - fetch_head[j] for j in range(n)) + ws[1] = ( + 0 + if idle + else sum(batch.num_rows for q in raw_batches for _, batch in q) + ) + ws[2] = 0 if idle else sum(len(q) for q in cooked) + ws[3] = sum(local_consumed) + ws[4] = self._bytes_loaded + ws[5] = int(self._fetch_time * 1_000_000) + ws[6] = int(self._transform_time * 1_000_000) + ws[7] = self._rows_skipped + + # Sequence-packing helpers + pack_len = cast(int, self._pack_sequences) + eos_id = cast(int, self._eos_id) + pad_id = cast(int, self._pad_id) + blocks_per_split = ( + cast(int, self._blocks_per_epoch) // self._num_splits + if self._pack_sequences is not None + else 0 + ) + pack_consumed = list(self._pack_consumed) + pack_buffers = deepcopy(self._pack_buffers) + pack_blocks_emitted = list(self._pack_blocks_emitted) + + def _pack_buffer(i: int) -> dict[str, list[int]]: + return pack_buffers.setdefault(my_splits[i], {"tokens": [], "starts": []}) + + def _fill_block(i: int) -> None: + """Fill split i's buffer to one block or exhaust the split.""" + buf = _pack_buffer(i) + while len(buf["tokens"]) < pack_len: + _ensure_cooked(i) + if not cooked[i]: + return + buf["starts"].append(len(buf["tokens"])) + pos, tokens = cooked[i].popleft() + buf["tokens"].extend(tokens) + buf["tokens"].append(eos_id) + pack_consumed[my_splits[i]] = pos + 1 + local_consumed[i] += 1 + _advance(i) + + def _emit_block(i: int) -> dict[str, Any]: + buf = _pack_buffer(i) + tokens, starts = buf["tokens"], buf["starts"] + # doc_ids label document segments within the block; 0 also covers + # the continuation of a document begun in a prior block. + doc_ids = torch.zeros(pack_len, dtype=torch.int64) + doc_starts = [s for s in starts if 0 < s < pack_len] + doc_ids[doc_starts] = 1 + doc_ids.cumsum_(dim=0) # cumulative sum marks document boundaries + block = { + "input_ids": torch.tensor(tokens[:pack_len], dtype=torch.int64), + "doc_ids": doc_ids, + } + del tokens[:pack_len] + # Shift start boundaries for the next call. + buf["starts"] = [s - pack_len for s in starts if s >= pack_len] + return block + + def _commit_pack_state() -> None: + self._pack_consumed = list(pack_consumed) + self._pack_buffers = { + split: { + "tokens": list(buffer["tokens"]), + "starts": list(buffer["starts"]), + } + for split, buffer in pack_buffers.items() + } + self._pack_blocks_emitted = list(pack_blocks_emitted) + # ── Main loop ───────────────────────────────────────────────────────── with ThreadPoolExecutor(max_workers=n * io_queue_depth) as io_pool: @@ -583,10 +848,42 @@ class StreamingDataset(IterableDataset): self._fetch_head_ref = fetch_head self._split_sizes_ref = split_sizes self._local_consumed_ref = local_consumed + try: for i in range(n): _fill_io(i) + if self._pack_sequences is not None: + first_count = pack_blocks_emitted[my_splits[0]] + if any( + pack_blocks_emitted[split] != first_count + for split in my_splits[1:] + ): + raise ValueError( + "Packed checkpoint is not aligned across the splits " + "owned by this iterator; merge every rank " + "state with merge_state_dicts before resuming on a " + "different topology" + ) + + while pack_blocks_emitted[my_splits[0]] < blocks_per_split: + # Each logical split gets one block per cycle. Exhausted + # splits are padded through the fixed global budget. + for i in range(n): + _fill_block(i) + + for i in range(n): + tokens = _pack_buffer(i)["tokens"] + if len(tokens) < pack_len: + tokens.extend([pad_id] * (pack_len - len(tokens))) + block = _emit_block(i) + pack_blocks_emitted[my_splits[i]] += 1 + if i == n - 1: + _commit_pack_state() + _update_stats() + yield block + return + while True: # A cycle only runs if every split can still produce a # row. Without skips all splits exhaust simultaneously @@ -620,21 +917,7 @@ class StreamingDataset(IterableDataset): self._resume_offset = initial_offset + local_consumed[i] for j, split_idx in enumerate(my_splits): self._resume_positions[split_idx] = pos_consumed[j] - ws = self._worker_stats - ws[0] = sum( - split_sizes[j] - fetch_head[j] for j in range(n) - ) - ws[1] = sum( - batch.num_rows - for q in raw_batches - for _, batch in q - ) - ws[2] = sum(len(q) for q in cooked) - ws[3] = sum(local_consumed) - ws[4] = self._bytes_loaded - ws[5] = int(self._fetch_time * 1_000_000) - ws[6] = int(self._transform_time * 1_000_000) - ws[7] = self._rows_skipped + _update_stats() yield row finally: @@ -642,15 +925,7 @@ class StreamingDataset(IterableDataset): # when iteration ends mid-cycle (e.g. a split whose rows # were all skipped before completing a single cycle), so # counters like rows_skipped would otherwise be stale. - ws = self._worker_stats - ws[0] = sum(split_sizes[j] - fetch_head[j] for j in range(n)) - ws[1] = 0 # queue-depth properties document 0 when idle - ws[2] = 0 - ws[3] = sum(local_consumed) - ws[4] = self._bytes_loaded - ws[5] = int(self._fetch_time * 1_000_000) - ws[6] = int(self._transform_time * 1_000_000) - ws[7] = self._rows_skipped + _update_stats(idle=True) self._raw_batches_ref = None self._cooked_ref = None self._fetch_head_ref = None @@ -808,21 +1083,31 @@ class StreamingDataset(IterableDataset): def state_dict(self) -> dict: """Snapshot the dataset's consumption state. - The returned dict is topology-independent: at global step boundaries - every split has been consumed the same number of times (by the - round-robin design), so the per-split count is a single uniform value - that is identical across all ranks and DataLoader workers. - - ``positions_consumed_per_split`` records how far into each split's - permutation iteration has advanced. It only differs from - ``samples_consumed_per_split`` when ``on_transform_error`` skipped - rows, in which case entries are exact for the splits this instance - iterated and a lower bound (the sample count) for splits owned by - other ranks or workers. Combine the state dicts from all ranks with + In row mode, the returned dict is topology-independent at global step + boundaries. ``positions_consumed_per_split`` records how far each + split's permutation has advanced, which can differ from the sample + count when ``on_transform_error`` skips rows. Combine state dicts from + every rank with [merge_state_dicts][lancedb.streaming.StreamingDataset.merge_state_dicts] - to recover the exact value for every split before resuming on a - different topology. + before resuming on a different topology. + + Packed state includes partial token buffers and emitted block counts + for every logical split. When packing is sharded, merge every rank + state with ``merge_state_dicts`` before loading it. """ + if self._pack_sequences is not None: + return { + "shuffle_seed": self._shuffle_seed, + "num_splits": self._num_splits, + "epoch": self._epoch, + "pack_sequences": self._pack_sequences, + "eos_id": self._eos_id, + "pad_id": self._pad_id, + "blocks_per_epoch": self._blocks_per_epoch, + "samples_consumed_per_split": list(self._pack_consumed), + "blocks_emitted_per_split": list(self._pack_blocks_emitted), + "pack_buffers": deepcopy(self._pack_buffers), + } positions = [ self._resume_positions.get(split, self._resume_offset) for split in range(self._num_splits) @@ -840,7 +1125,9 @@ class StreamingDataset(IterableDataset): Raises ``ValueError`` if ``num_splits`` or ``shuffle_seed`` differ from the checkpoint, since a different split structure or shuffle order - makes mid-epoch resumption meaningless. + makes mid-epoch resumption meaningless. Packed checkpoints + pin ``pack_sequences``, ``eos_id``, ``pad_id``, + ``blocks_per_epoch``, and ``epoch``. """ if state["num_splits"] != self._num_splits: raise ValueError( @@ -852,6 +1139,31 @@ class StreamingDataset(IterableDataset): f"shuffle_seed mismatch: checkpoint has {state['shuffle_seed']}, " f"current dataset has {self._shuffle_seed}" ) + + if "pack_buffers" in state or self._pack_sequences is not None: + for key in ( + "pack_sequences", + "eos_id", + "pad_id", + "blocks_per_epoch", + "epoch", + ): + ours = getattr(self, f"_{key}") + if state.get(key) != ours: + raise ValueError( + f"{key} mismatch: checkpoint has {state.get(key)}, " + f"current dataset has {ours}" + ) + self._pack_consumed = [int(c) for c in state["samples_consumed_per_split"]] + self._pack_blocks_emitted = [ + int(c) for c in state["blocks_emitted_per_split"] + ] + self._pack_buffers = { + int(g): {"tokens": list(b["tokens"]), "starts": list(b["starts"])} + for g, b in state["pack_buffers"].items() + } + return + consumed = state["samples_consumed_per_split"] # All entries are equal at step boundaries; use the first. if isinstance(consumed, list): @@ -873,25 +1185,22 @@ class StreamingDataset(IterableDataset): def merge_state_dicts(states: list[dict]) -> dict: """Merge state dicts saved by different ranks into one exact state. - Only needed when ``on_transform_error`` skips rows in multi-rank - training: each rank then knows the exact permutation position only for - its own splits, and records a lower bound for the rest. Because - exactly one rank owns each split, the elementwise maximum across all - ranks' ``positions_consumed_per_split`` recovers the exact position of - every split. Without skipped rows every rank's state is already - identical and merging is a no-op. + For row mode, the elementwise maximum of permutation positions recovers + splits advanced by different ranks after transform failures. For packed + mode, the state that emitted the most blocks for each logical split + supplies that split's permutation position and partial token buffer. Packed + states must cover every rank at the same global step. - Raises ``ValueError`` if the states are empty or were not produced by - the same run (mismatched seed, split count, epoch, or sample counts). + Raises ``ValueError`` if the states are empty, were not produced by + the same run, or do not represent the same global step. The merge is always all-to-all and topology-agnostic: collect the - ``state_dict()`` from every rank of the *previous* run into one list, - merge that whole list, and hand the identical merged result to every - rank of the *next* run — regardless of whether the rank count grew, - shrank, or stayed the same. There is no pairwise or subset merging - step, because each split's exact position is only known to whichever - rank owned that split, and the elementwise maximum needs every rank's - contribution to be correct. + ``state_dict()`` from every rank of the *previous* run into + one list, merge that whole list, and hand the identical merged result + to every rank of the *next* run — regardless of whether the + topology grew, shrank, or stayed the same. There is no pairwise or + subset merging step, because each split's exact state is only known to + whichever iterator owned that split. For example, checkpointing 8 ranks and resuming on 4 (the same pattern applies when growing, e.g. 4 ranks resuming on 8):: @@ -924,13 +1233,73 @@ class StreamingDataset(IterableDataset): if not states: raise ValueError("merge_state_dicts requires at least one state dict") first = states[0] + packed = "pack_buffers" in first + config_keys = ["shuffle_seed", "num_splits", "epoch"] + if packed: + config_keys.extend( + ["pack_sequences", "eos_id", "pad_id", "blocks_per_epoch"] + ) + for state in states[1:]: - for key in ("shuffle_seed", "num_splits", "epoch"): + if ("pack_buffers" in state) != packed: + raise ValueError("cannot merge packed and unpacked state dicts") + for key in config_keys: if state[key] != first[key]: raise ValueError( f"{key} mismatch across state dicts: " f"{state[key]} != {first[key]}" ) + + if packed: + num_splits = first["num_splits"] + for state in states: + for key in ( + "samples_consumed_per_split", + "blocks_emitted_per_split", + ): + if len(state[key]) != num_splits: + raise ValueError( + f"{key} must contain one entry per logical split" + ) + + merged_consumed = [] + merged_emitted = [] + merged_buffers = {} + for split in range(num_splits): + owner = states[0] + owner_progress = ( + owner["blocks_emitted_per_split"][split], + owner["samples_consumed_per_split"][split], + ) + for state in states[1:]: + progress = ( + state["blocks_emitted_per_split"][split], + state["samples_consumed_per_split"][split], + ) + if progress > owner_progress: + owner = state + owner_progress = progress + merged_consumed.append(owner["samples_consumed_per_split"][split]) + merged_emitted.append(owner["blocks_emitted_per_split"][split]) + buffer = owner["pack_buffers"].get( + split, owner["pack_buffers"].get(str(split)) + ) + if buffer is not None: + merged_buffers[split] = deepcopy(buffer) + + if len(set(merged_emitted)) > 1: + raise ValueError( + "packed state dicts were not captured at the same global " + "step or do not cover every rank" + ) + + merged = dict(first) + merged["samples_consumed_per_split"] = merged_consumed + merged["blocks_emitted_per_split"] = merged_emitted + merged["pack_buffers"] = merged_buffers + return merged + + for state in states[1:]: if ( state["samples_consumed_per_split"] != first["samples_consumed_per_split"] diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index 4d860a5de..ff65d100c 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/python/python/tests/test_elastic_dataloader.py @@ -2113,6 +2113,214 @@ def test_shuffle_seed_none_generates_stable_seed(lance_table): assert first == second, "Same resolved seed must produce the same ordering" +# Sequence packing tests + + +def _create_token_table(tmp_path, documents): + db = lancedb.connect(tmp_path) + tokens = pa.array(documents, type=pa.list_(pa.int64())) + return db.create_table("tokens", pa.table({"tokens": tokens})) + + +def _packed_dataset(table, pack_sequences, *, blocks_per_epoch, pad_id=0, **kwargs): + return StreamingDataset( + table, + shuffle=False, + columns=["tokens"], + pack_sequences=pack_sequences, + eos_id=9, + pad_id=pad_id, + blocks_per_epoch=blocks_per_epoch, + **kwargs, + ) + + +def test_pack_sequences_emits_blocks_and_pads_final_tail(tmp_path): + table = _create_token_table(tmp_path, [[1, 2], [3, 4], [5]]) + dataset = _packed_dataset(table, 6, blocks_per_epoch=2) + + blocks = list(dataset) + + assert len(blocks) == 2 + assert blocks[0]["input_ids"].tolist() == [1, 2, 9, 3, 4, 9] + assert blocks[0]["doc_ids"].tolist() == [0, 0, 0, 1, 1, 1] + assert blocks[1]["input_ids"].tolist() == [5, 9, 0, 0, 0, 0] + assert blocks[1]["doc_ids"].tolist() == [0, 0, 0, 0, 0, 0] + assert blocks[0]["input_ids"].dtype == torch.int64 + assert blocks[0]["doc_ids"].dtype == torch.int64 + + +def test_pack_sequences_pads_lagging_splits(tmp_path): + table = _create_token_table( + tmp_path, + [[1], [2], [10, 11, 12, 13, 14, 15, 16, 17], [20]], + ) + dataset = _packed_dataset(table, 5, blocks_per_epoch=6, num_splits=2) + input_ids = [block["input_ids"].tolist() for block in dataset] + # Split 0 has four real tokens including EOS markers, while split 1 has + # eleven. Packing must emit three complete two-split cycles. + assert input_ids == [ + [1, 9, 2, 9, 0], + [10, 11, 12, 13, 14], + [0, 0, 0, 0, 0], + [15, 16, 17, 9, 20], + [0, 0, 0, 0, 0], + [9, 0, 0, 0, 0], + ] + + per_rank = [] + for rank in range(2): + rank_dataset = _packed_dataset( + table, + 5, + blocks_per_epoch=6, + num_splits=2, + world_size=2, + rank=rank, + ) + per_rank.append([block["input_ids"].tolist() for block in rank_dataset]) + + assert [len(blocks) for blocks in per_rank] == [3, 3] + sharded = [block for cycle in zip(*per_rank) for block in cycle] + assert sharded == input_ids + + +def test_pack_sequences_auto_estimates_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], + } + ) + ) + + with pytest.warns(UserWarning, match="approximate token-count sample"): + dataset = _packed_dataset( + table, + 5, + blocks_per_epoch="auto", + num_splits=2, + filter="keep", + ) + + # 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): + table = _create_token_table( + tmp_path, + [[1], [2], [10, 11, 12, 13, 14, 15, 16, 17], [20]], + ) + kwargs = dict(pack_sequences=5, blocks_per_epoch=6, num_splits=2) + reference = list(_packed_dataset(table, **kwargs)) + + datasets = [ + _packed_dataset(table, world_size=2, rank=rank, **kwargs) for rank in range(2) + ] + iterators = [iter(dataset) for dataset in datasets] + first_cycle = [next(iterator) for iterator in iterators] + checkpoint = StreamingDataset.merge_state_dicts( + [dataset.state_dict() for dataset in datasets] + ) + for iterator in iterators: + iterator.close() + + resumed = _packed_dataset(table, **kwargs) + resumed.load_state_dict(checkpoint) + actual_remaining = list(resumed) + + assert [block["input_ids"].tolist() for block in first_cycle] == [ + [1, 9, 2, 9, 0], + [10, 11, 12, 13, 14], + ] + assert checkpoint["blocks_emitted_per_split"] == [1, 1] + assert [block["input_ids"].tolist() for block in actual_remaining] == [ + block["input_ids"].tolist() for block in reference[2:] + ] + assert [block["doc_ids"].tolist() for block in actual_remaining] == [ + block["doc_ids"].tolist() for block in reference[2:] + ] + + +def test_pack_sequences_validates_configuration_and_tokens(tmp_path): + table = _create_token_table(tmp_path, [[1, 2]]) + + with pytest.raises(ValueError, match="pad_id is required"): + StreamingDataset( + table, + shuffle=False, + columns=["tokens"], + pack_sequences=4, + eos_id=9, + ) + + with pytest.raises(ValueError, match="blocks_per_epoch is required"): + StreamingDataset( + table, + shuffle=False, + columns=["tokens"], + pack_sequences=4, + eos_id=9, + pad_id=0, + ) + + with pytest.raises(ValueError, match="must be divisible"): + _packed_dataset(table, 4, blocks_per_epoch=3, num_splits=2) + + with pytest.raises(ValueError, match="positive integer or 'auto'"): + _packed_dataset(table, 4, blocks_per_epoch="estimate") + + checkpoint = _packed_dataset(table, 4, blocks_per_epoch=1).state_dict() + resumed = _packed_dataset(table, 4, blocks_per_epoch=1, pad_id=8) + with pytest.raises(ValueError, match="pad_id mismatch"): + resumed.load_state_dict(checkpoint) + + float_db = lancedb.connect(tmp_path / "float") + float_table = float_db.create_table( + "tokens", + pa.table({"tokens": pa.array([[1.5, 2.5]], type=pa.list_(pa.float64()))}), + ) + with pytest.raises(ValueError, match="token column with integer values"): + _packed_dataset(float_table, 4, blocks_per_epoch=1) + + null_db = lancedb.connect(tmp_path / "null") + null_table = null_db.create_table( + "tokens", + pa.table({"tokens": pa.array([None], type=pa.list_(pa.int64()))}), + ) + with pytest.raises(ValueError, match="does not support null token lists"): + list(_packed_dataset(null_table, 4, blocks_per_epoch=1)) + + null_value_db = lancedb.connect(tmp_path / "null_value") + null_value_table = null_value_db.create_table( + "tokens", + pa.table( + {"tokens": pa.array([[1], [2, None], [3]], type=pa.list_(pa.int64()))} + ), + ) + blocks = list( + _packed_dataset( + null_value_table, + 2, + blocks_per_epoch=2, + on_transform_error="skip", + ) + ) + assert [block["input_ids"].tolist() for block in blocks] == [[1, 9], [3, 9]] + + # --------------------------------------------------------------------------- # Doc examples — each test mirrors the code snippet in index.mdx so that # broken doc examples are caught before they ship.