# 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>
## Summary
PyTorch's `DataLoader` uses fork-based multiprocessing by default on
Linux, but threads do not survive `fork()`. LanceDB's Python bindings
drive async work through two threaded layers, both of which become inert
in a forked child:
- `BackgroundEventLoop` runs an asyncio loop on a Python
`threading.Thread`.
- `pyo3-async-runtimes::tokio` holds a global multi-threaded tokio
runtime whose worker threads also die on fork — and its runtime lives in
a `OnceLock` that cannot be replaced after first use.
As a result, any `Permutation` (or other async API) used inside a
fork-based `DataLoader` worker hangs indefinitely. This PR makes both
layers fork-safe so `Permutation` works as a `torch.utils.data.Dataset`
with `num_workers > 0`.
## Approach
### Rust — new `python/src/runtime.rs`
Mirrors the pattern used in [Lance's Python
bindings](https://github.com/lance-format/lance/blob/456198cd6f42be07f99617a6d7e39d6209cdf3cc/python/src/lib.rs#L139),
adapted for the async-bridge use case.
- `LanceRuntime` implements `pyo3_async_runtimes::generic::Runtime +
ContextExt`, backed by an `AtomicPtr<tokio::runtime::Runtime>` we own
(sidestepping `pyo3-async-runtimes`'s frozen `OnceLock` global).
- A `pthread_atfork(after_in_child)` handler nulls the pointer; the next
`spawn` rebuilds the runtime in the child. The previous runtime is
intentionally **leaked** — calling `Drop` would try to join now-dead
worker threads and hang.
- `runtime::future_into_py` is a drop-in for
`pyo3_async_runtimes::tokio::future_into_py`. All ~80 call sites in
`arrow.rs` / `connection.rs` / `permutation.rs` / `query.rs` /
`table.rs` are updated to route through it.
- `python/Cargo.toml` adds `libc = "0.2"` and the tokio
`rt-multi-thread` feature.
### Python — `lancedb/background_loop.py`
- Refactors `BackgroundEventLoop.__init__` to a reusable `_start()`
method.
- An `os.register_at_fork(after_in_child=…)` hook calls `LOOP._start()`
to give the singleton a fresh asyncio loop and thread **in place**. This
matters because the rest of the codebase imports `LOOP` via `from
.background_loop import LOOP` — rebinding the module attribute would
leave those references holding the dead loop.
### Python — `lancedb/__init__.py`
Removes the `__warn_on_fork` pre-fork warning (and the now-unused
`import warnings`). Fork is supported.
## Test plan
- [x] New `test_permutation_dataloader_fork_workers` in
`python/tests/test_torch.py`: runs a `Permutation` through
`torch.utils.data.DataLoader(num_workers=2,
multiprocessing_context="fork")` inside a spawn-isolated child with a
30s hang detector. **Pre-fix**: timed out at 36s. **Post-fix**: passes
in ~3.6s.
- [x] New `test_remote_connection_after_fork` in
`python/tests/test_remote_db.py`: forks a child that creates a fresh
`lancedb.connect(...)` against a mock HTTP server and calls
`table_names()`; passes in <1s, validates the runtime reset is
sufficient for fresh remote clients.
- [x] All 62 tests in `test_torch.py` + `test_permutation.py` pass.
- [x] All 35 tests in `test_remote_db.py` pass.
- [x] `test_table.py` (87) + `test_db.py` + `test_query.py` (157, minus
one unrelated `sentence_transformers` import skip) — 244 passing.
- [x] `cargo clippy -p lancedb-python --tests` clean.
- [x] `cargo fmt`, `ruff check`, `ruff format` all clean.
## Known limitation (follow-up)
This PR makes a **freshly-built** `lancedb.connect(...)` work in a
forked child. An **inherited** `Connection` from the parent still
carries an inherited `reqwest::Client` whose hyper connection pool
references socket FDs and TCP/TLS state shared with the parent — using
it from the child after fork is unsafe (especially with HTTP/1.1
keep-alive). The recommended pattern for fork-based `DataLoader` workers
that hit a remote DB is to construct a new connection inside the worker.
Auto-clearing inherited HTTP client pools on fork would require tracking
live `Connection` instances in `lancedb` core and is left for a
follow-up PR.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
## Summary
When pytorch is used with multiprocessing and the mp mode is spawn then
the Permutation needs to be pickled. It could not be pickled because
`Table` and `Connection` are not serializable. This PR adds pickle
support to Permutation without adding general pickle support to `Table`
or `Connection`. To add general support we probably need to start by
adding serialization in the namespace client.
In the meantime this PR enable pickling by adding special cases for:
* In-memory tables (just serialize as Arrow IPC)
* Native tables (serialize the URI)
If a user is not using one of the above cases (e.g. using a remote
connection) then they will need to provide a connection factory that can
be pickled.
## Breaking change
`PermutationBuilder.persist(...)` is removed from the Python bindings;
the permutation table is now always in-memory. The underlying Rust
`PermutationBuilder::persist` API is untouched and can be re-exposed
later if needed. It probably won't make sense to do that until we have a
way to serialize `Table` and `Connection`.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This PR migrates all Rust crates in the workspace to Rust 2024 edition
and addresses the resulting compatibility updates. It also fixes all
clippy warnings surfaced by the workspace checks so the codebase remains
warning-free under the current lint configuration.
Context:
- Scope: workspace edition bump (`2021` -> `2024`) plus follow-up
refactors required by new edition and clippy rules.
- Validation: `cargo fmt --all` and `cargo clippy --quiet --features
remote --tests --examples -- -D warnings` both pass.
There were two issues:
1. The python code needs to get access to the underlying rust table to
setup the permutation reader and the attributes involved in this differ
between the python local table and remote table objects.
~~2. The remote table was sending projection dictionaries as arrays of
tuples and (on LanceDB cloud at least) it does not appear this is how
rest servers are setup to receive them.~~ (this is now fixed as #3023)
~~Leaving as draft as this is built on
https://github.com/lancedb/lancedb/pull/3016~~
I'm working on a lancedb version of pytorch data loading (and hopefully
addressing https://github.com/lancedb/lance/issues/3727).
However, rather than rely on pytorch for everything I'm moving some of
the things that pytorch does into rust. This gives us more control over
data loading (e.g. using shards or a hash-based split) and it allows
permutations to be persistent. In particular I hope to be able to:
* Create a persistent permutation
* This permutation can handle splits, filtering, shuffling, and sharding
* Create a rust data loader that can read a permutation (one or more
splits), or a subset of a permutation (for DDP)
* Create a python data loader that delegates to the rust data loader
Eventually create integrations for other data loading libraries,
including rust & node
Weston Pace
LanceDB Robot
Xuanwo
Jack Ye