--- title: "Loading Data for Model Training" sidebarTitle: "Data Loading & Shuffles" description: "Stream, shuffle, transform, and resume model training data with LanceDB." icon: boxes-stacked availability: oss: available enterprise: available summary: >- Streaming, shuffling and resuming training data work with local, object-storage and deployment tables in the Python SDK. --- LanceDB provides an excellent data backend for training machine learning models. For basic use cases, a `Table` can be fed directly into a data loader. For a more complete solution, a `StreamingDataset` adds prefetching, elastic determinism, resumability, multithreaded transformations, and more. ## Basic Data Loading Most model training frameworks feed data into a model by sequencing through successive batches, a process often referred to as **data loading**. The simplest way to perform data loading is to iterate through a LanceDB table in a loop. ```py Python icon=Python import lancedb db = lancedb.connect("file://some/db/path") table = db.open_table("some_table") for batch in table: print(batch.to_pydict()) ``` ## Advanced Data Loading Advanced data loading scenarios can be handled by the `StreamingDataset`. This object allows you to apply preprocessing before training, pipelined reads, and more. Use this when the training data does not fit in memory, or when you need deterministic global batches across cluster sizes. ```py Python icon=Python import lancedb import torch from lancedb.streaming import StreamingDataset db = lancedb.connect("file://some/db/path") table = db.open_table("some_table") dataset = StreamingDataset(table, shuffle_seed=42) dataloader = torch.utils.data.DataLoader( dataset, batch_size=128, num_workers=0, ) for batch in dataloader: train_step(batch) ``` The `StreamingDataset` wraps a LanceDB `Table` and yields plain Python dictionaries by default. Whether iterating directly or through a PyTorch `DataLoader`, only one iterator can be simultaneously active on a `StreamingDataset` instance. `StreamingDataset` is built on the `permutation` API and works with both LanceDB OSS and LanceDB Enterprise. The underlying table data can live on local disk or object storage. To understand the wide range of scenarios handled by `StreamingDataset`, consider a model trained with stochastic gradient descent (SGD) and distributed data parallelism (DDP). In this example, we need to load batches onto multiple GPUs across multiple servers. After each batch is processed, the GPUs exchange weights and the next batch is loaded. We fix the following terminology: - **World size** - The number of GPUs being loaded simultaneously. For example, if we have $2$ servers with $4$ GPUs each, the world size is $8$. - **Rank** - The identifier of an individual GPU. It is an integer in the range `[0, world_size)`. Each rank is allocated a disjoint portion of the source data. - **Batch size** - The number of rows processed by one GPU in a single training step. For example, $128$. - **Global batch size** - The total number of rows processed across all GPUs in a training step (sometimes called a "minibatch" in machine learning literature). For example, if we have $8$ GPUs with batch size $128$, the global batch size is $1024$. Other concepts, such as read batch size and `num_workers`, are introduced in the relevant sections below. ### Filtering The `StreamingDataset` object allows you to train models on only the subset of rows that satisfy a user-defined filter, specified using the `filter` parameter. This can be useful for removing large rows, badly-sourced data, and a wide range of other applications. You can also use the `columns` parameter to load only a subset of columns. ```py Python icon=Python ds = StreamingDataset( table, shuffle_seed=42, columns=["image", "label"], # skip all other columns filter="category = 'train'", # only training rows ) ``` ### Shuffling Rows In model training, model quality can be hurt by unintentionally learning patterns in the ordering of the source data. Shuffling the dataset before training alleviates this problem. However, fetching rows with scattered sparse reads across object storage can be dramatically slower than fetching contiguous batches. `StreamingDataset` offers two different implementations to read data in shuffled order, toggled by the presence of the `block_size` parameter. `StreamingDataset` shuffle seeds combine `shuffle_seed` with `epoch`; each epoch produces a different permutation while runs with the same inputs remain reproducible. Set `shuffle_seed = None` to generate a random seed upon pipeline initialization. #### One-Phase Shuffle Setting `block_size = None` corresponds to the naive $1$-phase shuffled read: the rows of the source dataset are shuffled according to a pseudorandom permutation, then fed in batches to the accelerator. This gives the best possible randomness, but can exhibit low throughput. ```py Python icon=Python # 1-phase shuffled read: rows shuffled individually. for epoch in range(num_epochs): ds = StreamingDataset( table, shuffle_seed=42, epoch=epoch, # changes the permutation each epoch ) for sample in ds: train_step(sample) #Evaluation: deterministic sequential order, no shuffle. eval_ds = StreamingDataset(table, shuffle=False) for sample in eval_ds: eval_step(sample) ``` #### Two-Phase Shuffle Setting `block_size` to an integer corresponds to selecting $2$-phase shuffled read. In this implementation, the source dataset is first partitioned into fixed-size `blocks`, which are then assigned randomly to splits. Within each split, rows are then further shuffled, with each row moving at most `max_shuffle_distance` distance away from its original position. Reading blocks allows storage I/O to constitute mostly large, contiguous reads rather than single row fetches. This yields much higher throughput while slightly compromising randomness. ```py Python icon=Python # 2-phase shuffled read: rows are shuffled block-by-block instead of # individually, trading a bit of randomness for much cheaper I/O. for epoch in range(num_epochs): ds = StreamingDataset( table, shuffle_seed=42, epoch=epoch, block_size=1024, # enables 2-phase window_blocks=4, # blocks resident per split at once max_shuffle_distance=256, # how far a row may move from its # original position ) for sample in ds: train_step(sample) # Evaluation: deterministic sequential order, no shuffle. eval_ds = StreamingDataset(table, shuffle=False, block_size=1024) for sample in eval_ds: eval_step(sample) ``` ### Parallelism & Elasticity When loading data, `StreamingDataset` partitions the the source dataset among disjoint parallel threads called **splits**. Each rank is allocated a contiguous group of splits, and each DataLoader worker associated to a rank gets a contiguous subgroup of its rank's splits. If using `shuffle = false`, data is allocated to splits naively in ordered contiguous blocks of equal size. If `shuffle = true`, in one-phase shuffling, the rows of the source dataset are distributed randomly among splits such that each split receives the same number of rows. In two-phase shuffling, the dataset is first partitioned into blocks, which are then distributed randomly among splits such that each split receives the same number of blocks. The number of splits defaults to `world_size` if `num_splits` is omitted during initialization. This is not recommended, as a later run using a different world size will then construct a different layout and structure of global batches. To get **elastic determinism**, choose a fixed `num_splits` that is compatible with every topology you plan to use. For example, a run with `num_splits=8` and $4$ GPUs assigns $2$ splits to each rank. The ranks pull from their splits in round-robin order, producing the same global batches as a run with $8$ GPUs and the same `num_splits`, `shuffle_seed`, and `epoch`. The following constraints apply: - `num_splits` must be divisible by `world_size`. - With DataLoader worker processes, `num_splits` must be divisible by `world_size * num_workers`. - For the same samples to form each global training step across topologies, `global_batch_size` must be a multiple of `num_splits`. - The filtered row count must be at least `num_splits`. If it is not evenly divisible by `num_splits`, up to `num_splits - 1` surplus rows are dropped so that every split has the same length. For example, to switch between $8$ and $6$ GPUs, use a common multiple such as $24$ for `num_splits`. Highly composite values can support many layouts; for example, `num_splits = 48` is compatible with $10$ distinct possible GPU counts. Remember to include `num_workers` when checking divisibility. ```py Python icon=Python import torch.distributed as dist dist.init_process_group("nccl") rank = dist.get_rank() world_size = dist.get_world_size() # num_splits=48 is divisible by many integers, # so this dataset works unchanged as you add or remove GPUs. ds = StreamingDataset( table, num_splits=48, shuffle_seed=42, epoch=current_epoch, rank=rank, world_size=world_size, ) for sample in ds: train_step(sample) ``` ### Fetching Traditional datasets have often been designed to yield a single sample at a time upon iteration. However, in modern use cases, accessing data on object storage one row at a time introduces too much per-call overhead. Thus, `StreamingDataset` fetches data in read batches, whose size is controlled by the `read_batch_size` parameter (defaults = `64`). In addition to batching requests, the prefetching mechanism reads ahead in the background. While one batch is being processed, `StreamingDataset` reads several subsequent batches. The `prefetch_batches` parameter specifies the number of batches simultaneously kept in flight per split (default = `4`). A larger value provides more buffering against jittery workloads, but also increases memory use and I/O concurrency. ```py Python icon=Python ds = StreamingDataset( table, shuffle_seed=42, read_batch_size=256, # rows fetched per LanceDB call per split prefetch_batches=8, # batches to keep in flight per split ) ``` ### Transforming Many model training workloads require performing a transformation to the data after it is loaded from object storage and before it is passed to the GPU. For example, we may need to decode images, tokenize text, or normalize vectors. A transformation function can be provided using the `transform` parameter. A transformation function must receive a PyArrow `RecordBatch` and return an iterable. For example, the default transform function returns a list of Python dictionaries. Furthermore, the input expected by PyTorch must match the collation function used in the PyTorch data loader. PyTorch's default collation function accepts a Python dictionary and several other types, and is thus compatible with the default `StreamingDataset` transform function. If your custom transform function returns something different (for example, Arrow record batches), you may need to provide a compatible PyTorch collation function. ```py Python icon=Python import pyarrow as pa def normalize(batch: pa.RecordBatch) -> list[dict]: # This pure-Python loop holds the GIL and is shown for illustration only. # In practice, prefer a library like torchvision or numpy that releases the # GIL so the ThreadPoolExecutor can run transforms in parallel. rows = batch.to_pylist() for row in rows: row["image"] = [v / 255.0 for v in row["image"]] return rows ds = StreamingDataset(table, shuffle_seed=42, transform=normalize) ``` `StreamingDataset` manages expensive transforms through a `ThreadPoolExecutor` whose worker count equals the number of available CPUs. Furthermore, transformations are applied to batches (instead of individual samples) to amortize per-batch overhead. ##### DataLoader Workers The thread-based transformation model that `StreamingDataset` uses by default is only effective when the transform function releases the GIL. This is true for many Python scientific libraries, including NumPy, PyArrow, and TorchVision, but not for pure-Python transforms. PyTorch can launch multiple DataLoader worker processes per rank, and `StreamingDataset` uses `get_worker_info()` to give each worker a non-overlapping group of splits. Multiprocessing adds pickling and transfer overhead and increases memory use. Start with `num_workers=0`, which keeps loading in the rank's main process, and add workers only after confirming an unavoidable GIL bottleneck. If you use workers, `num_splits` must be divisible by `world_size * num_workers`. Use the `forkserver` or `spawn` multiprocessing context because LanceDB uses internal threads. ```py Python icon=Python dataloader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, num_workers=2, multiprocessing_context="forkserver", persistent_workers=True, ) ``` By default, `StreamingDataset` serializes enough table state to reopen the table in each worker. For custom connection setup, pass a picklable `connection_factory` callable that accepts a table name and returns an open table. This avoids serializing connection credentials into worker state. ### Observability Optimizing data loader performance can be tricky, as runtime bottlenecks are often difficult to locate across various stages of the I/O and CPU pipelines. To help distinguish bottlenecks, `StreamingDataset` exposes several pipeline counters: `raw_queue_depth` shows the number of loaded rows waiting to be transformed, and `prefetch_queue_depth` is the number of transformed rows ready to be consumed. `unscanned_rows` and `consumed_rows` show how far the current iterator has progressed. If the `prefetch_queue_depth` is consistently zero but the `raw_queue_depth` is not, then you likely have a CPU transformation bottleneck. You should investigate GIL bottlenecks or look for ways to optimize your transformation. This can often be done by batching the compute work. If both `prefetch_queue_depth` and `raw_queue_depth` are consistently zero while the consumer is waiting, I/O is the likely bottleneck. A larger read batch size or clumped shuffling could help. ```py Python icon=Python import threading, time ds = StreamingDataset(table, shuffle_seed=42) def log_pipeline_health(): while True: print( f"unscanned={ds.unscanned_rows} " f"raw={ds.raw_queue_depth} " f"cooked={ds.prefetch_queue_depth} " f"consumed={ds.consumed_rows}" ) time.sleep(1.0) monitor = threading.Thread(target=log_pipeline_health, daemon=True) monitor.start() for sample in ds: train_step(sample) print( f"bytes loaded: {ds.bytes_loaded} " f"fetch time: {ds.fetch_time:.2f}s " f"transform time: {ds.transform_time:.2f}s" ) ``` `bytes_loaded` measures raw Arrow buffer bytes before transformation. `bytes_loaded`, `fetch_time`, and `transform_time` are cumulative across iterations of the same dataset instance; because work runs concurrently, the summed stage times can exceed wall-clock time. The queue depths report rows currently waiting in the pipeline, and the progress counters reflect the latest iterator snapshot. ### Checkpointing `StreamingDataset` provides checkpointing in case your model training fails partway through a run. Use `StreamingDataset.state_dict()` to capture the current number of samples consumed from each split in a Python dictionary, and `load_state_dict()` to restore the last saved position. A stored `StreamingDataset` state reflects only round-robin cycles that have been completed across all splits; samples completed partway through a cycle at the time of saving can be replayed. To ensure efficient saving, make sure `global_batch_size` is a multiple of `num_splits`, as in the example below. ```py Python icon=Python import torch global_batch_size = 48 # one sample per split in each global step batch_size = global_batch_size // world_size dataset = StreamingDataset( table, num_splits=48, shuffle_seed=42, epoch=current_epoch, rank=rank, world_size=world_size, ) dataloader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, num_workers=0, ) for step, batch in enumerate(dataloader): train_step(batch) if (step + 1) % checkpoint_interval == 0: torch.save( {"model": model.state_dict(), "dataset": dataset.state_dict()}, f"checkpoint_{step + 1}.pt", ) # --- resuming after a crash --- checkpoint = torch.load("checkpoint_100.pt") model.load_state_dict(checkpoint["model"]) dataset_state = checkpoint["dataset"] dataset = StreamingDataset( table, num_splits=dataset_state["num_splits"], shuffle_seed=dataset_state["shuffle_seed"], epoch=dataset_state["epoch"], rank=rank, world_size=world_size, # may differ from the run that saved the checkpoint ) dataset.load_state_dict(dataset_state) dataloader = torch.utils.data.DataLoader( dataset, batch_size=global_batch_size // world_size, num_workers=0, ) for batch in dataloader: train_step(batch) ``` `load_state_dict()` rejects a checkpoint whose `num_splits` or `shuffle_seed` differs from the new dataset. For an exact resume, also use the same table snapshot, `epoch`, `shuffle`, `filter`, and other data-selection settings. The `world_size` and number of DataLoader workers may change as long as the split and batch-size divisibility constraints still hold. ##### Checkpointing with Multiple DataLoader Workers Python's built-in `torch.utils.data.DataLoader` produces a safe checkpoint when `num_workers = 0`. With `num_workers > 0`, calling `state_dict()` often raises `RuntimeError`. In this case, use `StreamingDataLoader` for multi-worker prefetch and exact resumability. It carries a state snapshot alongside every internal batch and commits it to the parent `StreamingDataset` only when that batch is returned to the trainer. The trainer receives the same collated batch it would receive from a standard `DataLoader`. ```py Python icon=Python from lancedb.streaming import StreamingDataset, StreamingDataLoader dataset = StreamingDataset( table, num_splits=48, shuffle_seed=42, epoch=current_epoch, rank=rank, world_size=world_size, ) dataloader = StreamingDataLoader( dataset, batch_size=batch_size, num_workers=4, ) for step, batch in enumerate(dataloader): train_step(batch) if (step + 1) % checkpoint_interval == 0: torch.save( {"model": model.state_dict(), "dataset": dataset.state_dict()}, f"checkpoint_{step + 1}.pt", ) ``` `StreamingDataLoader` accepts the same arguments as `torch.utils.data.DataLoader`, with a few restrictions: - `dataset` must be a `StreamingDataset` (subclasses that override `__iter__` are not supported). - `in_order = True` is required so that consumer-committed checkpoints stay deterministic. - `persistent_workers = True` is not supported, because prefetched worker copies cannot be restored from parent-committed state. - `drop_last = True` is not supported, because incomplete tails discarded by worker replicas cannot be checkpointed topology-independently. With more than one worker, call `state_dict()` at a complete logical step boundary where every split assigned to the rank has the same consumed-sample count. Calling it mid-step raises `RuntimeError` asking you to consume more batches first. ##### Resuming Across Different Topologies When training across ranks, each rank owns its own subset of splits and only its own splits have exact progress. To resume on a different `world_size`, collect the `state_dict()` from every rank of the previous run and merge them with `StreamingDataset.merge_state_dicts` before calling `load_state_dict()` on the new run. The merge is topology-agnostic: pass the full list of per-rank states in, and hand the identical merged dict to every rank of the resumed job, regardless of whether the topology grew, shrank, or stayed the same. ```py Python icon=Python # On each rank of the previous run: state = dataset.state_dict() # ...gather `state` from every rank into a single list `states` (for example # via torch.distributed all-gather or by saving one file per rank). merged = StreamingDataset.merge_state_dicts(states) # On each rank of the resumed run (world_size may differ): dataset = StreamingDataset( table, num_splits=merged["num_splits"], shuffle_seed=merged["shuffle_seed"], epoch=merged["epoch"], rank=rank, world_size=new_world_size, ) dataset.load_state_dict(merged) ``` ## Permutations In specific use cases, you may need the flexibility to shuffle, split, and select data without using the full `StreamingDataset` object. In these cases, use `Permutation`, the lower-level functionality on which `StreamingDataset` is built. A `Permutation` defines a customizable row ordering, and supports map-style access through `__getitem__()` and batched access through `__getitems__()`. ##### Version Pinning A `Permutation` is pinned to the version of its base table at the time it was constructed, and every read (including future reads from a `StreamingDataset` that wraps it, and every DataLoader worker after a `fork`) resolves against that pinned version. This makes iteration deterministic across compactions and worker forks: the permutation only addresses rows that existed when it was built, and rows appended to the base table afterwards are not visible through the existing permutation. To include newly appended rows, build a new `Permutation` (and any `StreamingDataset` that depends on it) against the updated table. Permutations written before this behavior shipped continue to read as they did before.