mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-13 23:32:27 +00:00
64513fd34a53980bbb86920212ba2bc172b707bd
8
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 jasonz-lance
|
64513fd34a |
test(python): add 2-phase shuffled read tests
Covers full-coverage/no-duplicates, runt-block and surplus-block discarding, shuffle=False ordering, the max_shuffle_distance displacement bound, filter support, checkpoint resume, and constructor validation for block_size mode. |
||
  lancedb-gatefixer[bot]andXuanwo
|
105fd73bc6 |
fix(python): commit streaming worker checkpoints on consumption (#4023)
## Summary - add `StreamingDataLoader`, which transports worker snapshots with prefetched batches and commits them to the parent dataset only when the trainer receives each batch - preserve exact non-uniform per-split progress and resume lagging splits without replaying already-consumed rows - reject stale parent checkpoints after a standard multi-process `DataLoader` has started, with guidance to use the consumer-aware loader - document the new public loader and merge non-uniform state across ranks ## Root cause PyTorch runs `StreamingDataset.__iter__` in private worker-process copies, while callers invoke `state_dict()` on the parent dataset. Sharing producer counters would still be incorrect because DataLoader prefetch can advance workers beyond batches returned to the trainer. ## Validation - `uv run --extra tests pytest python/tests/test_elastic_dataloader.py -q` (154 passed) - focused non-uniform merge regression (1 passed) - `uv run --project python --extra tests --extra dev ruff format .` - `uv run --project python --extra tests --extra dev ruff check .` - `cd docs && PYTHONPATH=. ../python/.venv/bin/mkdocs build` Fixes #3967 <!-- lance-gatekeeper-fix:v1 agent=572be272619660b97e87fd5c85188341 generation=1 --> --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo <github@xuanwo.io> |
||
  Ayush ChaurasiaandOpenAI Codex
|
242ade8017 |
feat(python): support sequence packing in streaming dataset (#3920)
## How packing works
Consider four tokenized documents:
[1]
[2]
[10, 11, 12, 13, 14, 15, 16, 17]
[20]
With:
```
StreamingDataset(
table,
shuffle=False,
columns=["tokens"],
num_splits=2,
pack_sequences=5,
eos_id=9,
pad_id=0,
blocks_per_epoch=6,
)
```
the documents are assigned to two fixed logical splits. Each split
maintains an independent token buffer, appends eos_id after every
document, and emits blocks of five tokens.
Because blocks_per_epoch=6, each split emits exactly three blocks:
Cycl/e 1:
Split 0: [1, 9, 2, 9, 0] # 9 is eos, 0 is padding
Split 1: [10, 11, 12, 13, 14]
Cycle 2:
Split 0: [0, 0, 0, 0, 0]
Split 1: [15, 16, 17, 9, 20]
Cycle 3:
Split 0: [0, 0, 0, 0, 0]
Split 1: [9, 0, 0, 0, 0]
If a split runs out of tokens early, it emits padded blocks through the
fixed budget. This prevents one rank from finishing before another.
Logical splits are independent of rank and worker ownership. A
checkpoint records each split’s consumed-document count, emitted-block
count, remaining tokens, and document boundaries. Merging
those per-split states allows the same packed stream to resume after the
topology changes.
doc_ids identifies document segments, including continuations across
block boundaries. It is not a padding mask: padding retains the
preceding document ID, so callers must mask padding using a
reserved pad_id.
blocks_per_epoch="auto" is also available. It estimates the budget from
a deterministic bounded sample and warns that the result is approximate.
WIP pre-training tests:
```
┌────────────────────────────────────┬────────────────────────┬─────────────────────────────────────┐
│ │ GPT-2 124M │ GPT-2 medium 354M │
├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
│ Corpus │ 2.4M docs / 12GB table │ 9.67M docs / 45GB table │
│ Tokens (Chinchilla) │ 2.43B │ 7.0B │
├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
│ Data prep (ingest→curate→tokenize) │ ~12 min │ ~51 min │
├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
│ Training wall time │ ~50 min │ 3h 06m │
├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
│ Throughput / MFU │ 1.60M tok/s / 35% │ 684k tok/s / 42.0%, │
├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
│ Final val loss │ 3.230 │ 2.840 │
└────────────────────────────────────┴────────────────────────┴─────────────────────────────────────┘
```
---------
Co-authored-by: OpenAI Codex <codex@openai.com>
|
||
  Weston PaceandClaude Sonnet 4.6
|
1f1d03f306 |
feat(python): add backpressure to StreamingDataset post-transform queue (#3897)
Rename prefetch_batches → io_queue_depth and introduce transform_queue_depth as a symmetric pair: both express "number of batches to buffer per split at this pipeline stage." The old names are still accepted as keyword arguments but log a deprecation warning redirecting callers to the new names. transform_queue_depth caps how many transform-result batches can accumulate per split in the post-transform queue. Without this limit a slow consumer (e.g. a GPU training step) causes cooked rows to pile up unboundedly. The backpressure check in _try_submit_tx counts both already-cooked rows and rows expected from in-flight transforms; it skips proactive transform submission when the combined total reaches the limit. The reactive _ensure_cooked path bypasses the check so the consumer never stalls. --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> |
||
 Jack Ye
|
f39a7a4dd9 |
feat: support remote tables in the data loader (#3981)
`StreamingDataset`, `PermutationBuilder`, and `Permutation` now work
against a `RemoteTable` (LanceDB Cloud and Enterprise), which unblocks
benchmarking the loader against the enterprise cluster cache.
```python
db = lancedb.connect("db://my-db", api_key=..., host_override=...)
ds = StreamingDataset(db.open_table("training"), world_size=8, rank=r)
```
Rows are addressed by `_rowid` exactly as before —
`PermutationReader::load_batch` already built the same `_rowid IN (...)`
filter that `Table::take_row_ids` sends, so the loader's fetch was
always the take path. It just was never allowed to run.
### The guard
`PermutationBuilder.__init__` rejected anything without `_inner`, so a
`RemoteTable` raised `TypeError` before reaching the PyO3 layer — which
already unwraps one via `_table._inner`.
### A bounded schema lookup
`PermutationReader::output_schema` reads the schema off a query plan,
and building a plan on a remote table *executes* the query
(`create_plan` → `execute_query`). With no limit that is `k =
isize::MAX`, so asking a remote table for its output schema pulled the
whole table over HTTP and threw it away — once per assigned split, on
every epoch, since `StreamingDataset.__iter__` constructs a
`Permutation` per split.
One row rather than zero, deliberately: lance gates its limit node on
`self.limit.unwrap_or(0) > 0`, so `Some(0)` means *no limit*.
### Tables with an LSM write spec are refused
A permutation references rows by row id, and rows that have not been
flushed to the base table do not have one yet. The loader could read
around them, but they would then be missing from training with nothing
said about it, so the build refuses such a table up front instead of
half supporting it.
### Fallible identity construction
`PermutationReader::identity` resolved `inner_new` with `unwrap`. That
was near total against a local dataset, but construction counts the base
table — an HTTP round trip for a remote one — so a transient network or
auth failure became a panic across the PyO3 boundary.
### Tests
End-to-end `permutation_builder` and `StreamingDataset` runs against a
mock server, the former torch-free so it runs wherever the suite does,
plus a test that a build succeeds without an LSM write spec and is
refused once one is installed.
|
||
 Sravan AvvaruandClaude Fable 5
|
a615306f39 |
feat(python): add on_transform_error fault tolerance to StreamingDataset (#3763)
Closes #3704 ## Problem Transforms can fail on bad data (e.g. nulls/NaNs from incomplete user surveys). Today any transform exception aborts iteration, and there is no way to skip invalid rows during loading. ## Solution New `on_transform_error` parameter on `StreamingDataset`: - `"raise"` (default, matches current behavior and the convention in tf.data / WebDataset / Ray Data) - `"skip"` — drop the failing rows and continue - `"warn"` — like skip, plus a logged warning per failing batch - a WebDataset-style callable `handler(exc) -> bool`, so users can skip only expected error types Key design points: - **Row-granular skipping**: when a batch fails, the transform is re-run on single-row slices so only the rows that actually fail are dropped (avoids Ray-style whole-block loss). Skips are counted in a new `rows_skipped` property. - **No crash on uneven skips**: the round-robin loop now ends the epoch at the last cycle where every split still has a row, instead of hitting `IndexError` when a split runs dry early. - **Exact resumability under skips**: checkpoints are now position-based. `state_dict` gains `positions_consumed_per_split` (exact for owned splits), and a new `merge_state_dicts` static method combines per-rank states via elementwise max for elastic resume across topology changes. Old checkpoints without the new key still load. Positions equal sample counts when nothing is skipped, so existing behavior is unchanged. - **Guardrail**: transforms returning the wrong number of rows now raise a clear `ValueError` instead of silently corrupting split accounting. ### Answers to the issue's open questions - *Can we do this?* Yes — all transforms funnel through one guarded call in the Stage 2 pipeline. - *What do other libraries do?* tf.data `ignore_errors()`, WebDataset `handler=`, Ray `max_errored_blocks`; MosaicML StreamingDataset offers nothing (skipping conflicts with its determinism model). This design follows the common conventions: raise by default, opt-in skipping, count/log drops. - *Error handling or pre-filtering?* Both: the existing `filter=` remains the recommended tool for predictable bad data (splits are built post-filter, so all guarantees hold — now documented); `on_transform_error` covers failures not expressible as a predicate. - *Impact on splits / elastic determinism?* Per-split sample sequences stay deterministic (skips are data-dependent, not topology-dependent). With unequal bad-row counts across splits the last few global steps of an epoch can differ across topologies (bounded by the skew), which is documented on the parameter. With equal counts per split, full determinism is preserved — covered by a test. ## Testing 15 new tests in `test_elastic_dataloader.py` covering: default raise, invalid values, uniform and uneven skips (including epoch-end truncation), warn logging, selective callable handlers, wrong-row-count guardrail, determinism across runs and across world sizes (1/2/3/4) with skips, exact mid-epoch resume with skips on the same topology, elastic resume via `merge_state_dicts` (ws=2 → ws=1), merge validation, and backward-compat loading of old checkpoints. Note: relying on CI for the test run — my local machine OOMs during the final link of the native extension. The change itself is pure Python. --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> |
||
  buduoqiuandbuduoqiu
|
98a52267a2 |
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 <yaodong-shen@users.noreply.github.com> |
||
|
Weston PaceandClaude Sonnet 4.6
|
c6db80dd0b |
feat: add an elastic dataloader as an iterable dataset (#3509)
# Elastic Streaming Dataloader ## Motivation Training large models on LanceDB tables today requires loading the entire dataset into memory or writing bespoke batching logic. This PR introduces `StreamingDataset`, a PyTorch `IterableDataset` that streams directly from a LanceDB table with two hard guarantees that are difficult to achieve together: **elastic determinism** and **resumability**. ## Goals ### Elastic determinism The dataset partitions the table into a fixed number of *splits* (controlled by `num_splits`, `shuffle_seed`, and `epoch`). Samples are yielded by round-robining over splits one sample per split per cycle. Because the split structure is fixed, the set of samples that makes up each global training step is identical regardless of `world_size` or `num_workers`. You can scale your cluster up or down between runs and the model sees the same data in the same order — no re-sharding, no gradient variance from topology changes. ### Resumability `state_dict()` / `load_state_dict()` capture how many samples each split has consumed. Because all splits are the same size and the round-robin design keeps them in lockstep, the state reduces to a single scalar (`samples_consumed_per_split`) that is topology-independent. A checkpoint saved with 8 GPUs can resume correctly on 4 GPUs or 16 GPUs without any adjustment. ### PyTorch `IterableDataset` / streaming `StreamingDataset` implements the standard PyTorch `IterableDataset` interface, so it drops into any existing `DataLoader` pipeline without modification. Data is fetched lazily from Lance in chunks — only the rows needed for the current batch are ever in memory. Compared to the map dataset this takes more work from pytorch and puts it into the dataset itself (e.g. shuffling, filtering, etc.). We do this because we cannot achieve things like elastic determinism or prefiltering otherwise. ### Multi-worker support DataLoader workers are automatically assigned contiguous sub-blocks of splits (the rank's splits are divided evenly across workers). Each worker is independent: no shared state, no inter-process coordination. The only constraint is that `num_splits` must be divisible by `world_size * num_workers`. That being said, multi-worker is highly discouraged as it relies on multiprocessing which is inefficient. Still, we want to support it. ### Filters as prefilters Filters are applied at *permutation-build time* via `PermutationBuilder.filter()`, not re-evaluated on every fetch. The filtered row IDs are stored in the permutation table so that subsequent reads see only the matching rows. This allows us to avoid loading rows that don't match the filter (which is the default pytorch behavior) ### Prefetching Two parameters control the I/O pipeline: - `read_batch_size` (default 64) — number of rows fetched per `take_offsets` call. Larger values amortise per-request overhead, which is critical on object storage where a single round-trip can cost ~100 ms. - `prefetch_batches` (default 4) — number of batches prefetched in parallel per split via a `ThreadPoolExecutor`. While the model processes the current batch, the next several batches are already in flight, hiding storage latency behind compute. If set correctly then you can get good performance even with num_workers=0 (unless you are bottlenecked on transform). ### Transform parallelism The underlying `Permutation` API supports a `with_transform()` callback for decoding, augmentation, and format conversion. Unfortunately, this is not parallelized. Pytorch typically parallelizes this with num_workers which is multiprocessing which is highly inefficient. For simple transforms we should be able to utilize multithreading and Rust based UDFs. For complex python UDFs we could have a dedicated multiprocessing pipeline for just the transform. Or we could just utilize multithreading. In both cases we would exclude the I/O stage from the multiprocessing because that ends up being very memory hungry and inefficient. --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> |