litdata

# LitData LitData (`import litdata`) preprocesses and streams datasets for PyTorch training: - **Write** (`optimize` / `map`) → chunked `chunk-*.bin` + `index.json` → `src/litdata/processing/` - **Read** (`StreamingDataset` + `StreamingDataLoader`) → cache → decode → batch → `src/litdata/streaming/` - **Raw** (`StreamingRawDataset`) → stream original files without optimize → `src/litdata/raw/` **To use the library expertly:** always load [reference/using-litdata.md](reference/using-litdata.md) first. Narrative source: repo `README.md`. ## Install this skill From any project (Cursor, Claude Code, and other agents supported by the [skills CLI](https://github.com/vercel-labs/skills)): ```bash npx skills add Lightning-AI/litData ``` Useful options: `-g` (user-global), `-a cursor` (Cursor only), `-y` (non-interactive). In this repo the skill already lives at `.claude/skills/litdata/`. ## Expert usage (load using-litdata.md) Before writing examples or answering how-tos, read the cookbook. Highlights: | Topic | Remember | | ---------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | **Raw files** | `StreamingRawDataset` + torch `DataLoader` — `#stream-raw` / §10. Prefer cloud URL / connection path over FUSE. Defaults: `max_concurrent_downloads=None` (adaptive Stage 1), `max_prefetch=16` (worker-aware ~64 aggregate), `hedge_delay=0`, `download_timeout=120` (**batch-level**), `range_parallel_threshold=0`. Explicit `int` concurrency = exact permits. | | Images / media | Always wrap so a caption is not a path: `Text` / `Image` / `Audio` / `Video` / `File` / … (`types.py`). Class folders: `list_media_folder(root, kind=…)`. Recipes: `examples/modality/`. README `#modality` (`#media-types`). Bare `PIL.Image` / `fromarray` → huge PIL RAW. `litdata_collate` only special-cases graphs — Audio/Video need a custom `collate_fn`. | | Train stream | Optimized: `StreamingDataLoader` + `shuffle=True, drop_last=True, seed=…` | | Optimize | `if __name__ == "__main__"`; exactly one of `chunk_bytes` \| `chunk_size`. Default **64MB**; multi‑MB samples → consider **256–512MB**. **Shuffle the sample list before `optimize()`** when source order matters — README `#faq-chunk-shuffle` | | Ordered data | Chunk/item shuffle ≠ file-level shuffle. Shuffle before `optimize`, or use `StreamingRawDataset` + `DataLoader(shuffle=True)`. LitData does distributed + within-chunk bucket sampling automatically | | Cache | Peak disk ≈ `num_workers × max_pre_download × chunk_size`; default `max_cache_size=None` (75% of free disk, leave ≥50GB). `"100G"` or `0.90` of free. | | **POSIX-fast** | Local/Vast/NFS: mmap in place, no cache copy. `WindowShuffle` only on parallel FS (or `LITDATA_POSIX_FAST=1`). Cap workers / skip WILLNEED from `MemAvailable`. `LITDATA_POSIX_FAST=0` disables. | | **Paths** | Studio `/teamspace/s3_connections` & co are **FUSE** (convenience only — slow, can crash under load). LitData resolves them and talks **directly** to S3/GCS/**R2**. Never read the mount by hand. `reference/resolver.md` + `reference/data-movement.md` | | **Optimize I/O** | Default `keep_data_ordered=False`: per-node shared queue; download/upload/remove are **node-level threads**. Remote I/O uses the streaming **`Downloader`** (`adownload_file` / `aupload_file`). FsProvider is the fallback / index path. Tune `LITDATA_OPTIMIZE_*` / `LITDATA_PREFETCH_BYTES`. Load `reference/data-movement.md`. | | **Multi-node** | `num_nodes=` = Lightning Studio job (`_execute`), not torchrun/SLURM. Shard by `DATA_OPTIMIZER_*`; all ranks upload chunks; **last node** merges `{node}-index.json` → `index.json`. Load `reference/multi-node.md`. | | Throughput | Rough ImageNet Studio order-of-magnitude (not guarantees): FUSE ~**600**/s · Raw (right tuning) ~**6–7k**/s · Optimized 64MB chunks ~**11k**/s — `using-litdata.md` FAQ. Raw benches: medians + provenance SHAs; never cite short-window n=1 against Stage 0 medians. | | **Tracing** | `from litdata.debugger import enable_tracer` then `enable_tracer(level="chunk")` (or `categories=["download","read","delete"]`). Convert with [Lightning-AI/litracer](https://github.com/Lightning-AI/litracer): `litracer --quiet --validate -o trace.json.gz litdata_debug.log`. One line per event; crashes go to stderr + `ph: I`. Full spec: `reference/debugging.md`. | | **S3 workers** | `index.json` is boto3 (no parent tokio). Workers lazy-init obstore; boto3 fallback if parent already started it. Do **not** pass `data_connection_id` / `endpoint_url` into `boto3.Session` — build obstore from `S3Client`/`R2Client`. `FileNotFoundError` after 120s with `num_workers>0` only → [debugging.md](reference/debugging.md) failure modes. | | **Keyed lookup** | `optimize(..., key_fn=...)` writes `keys/`. Read `ds["id"]` / `get_by_key`. Patch locally with `dataset_update`. Needs `polars`. [keyed-lookup.md](reference/keyed-lookup.md). | | Parquet / HF | One-liner: `StreamingDataset("hf://datasets/org/name/data")` (auto-index + `ParquetLoader`). Index persists at `{cache}/hf-index/<hash>/index.json` (no re-scan). Default **prefetches** whole parquet then local `to_pylist()`. Do **not** default `range_read=True` (lost to HF on Hub). `spawn` with workers. Token: `HF_TOKEN`. `using-litdata.md` §10. | ## Reference map | Task | Read | | ----------------------------------------------------------------------------- | --------------------------------------------------- | | **Use the library** (raw, optimize/stream, modality, parquet/HF, serializers) | `reference/using-litdata.md` | | **Paths / URLs / Studio mounts / `Dir` / time templates** | `reference/resolver.md` (+ README `#resolve-paths`) | | **Downloaders / uploaders / removers / FUSE→cloud / optimize I/O** | `reference/data-movement.md` (+ `processing.md`) | | **Multi-node optimize/map** (`num_nodes`, `DATA_OPTIMIZER_*`, index merge) | `reference/multi-node.md` (+ `processing.md`) | | Read path, shuffle math, item loaders, Combined/Parallel | `reference/streaming.md` | | **Cache / BinaryWriter / BinaryReader / `index.json` / FsProvider / sampler** | `reference/storage-format.md` | | Cache / prefetch / eviction / shared-chunk deletion | `reference/cache-and-chunk-lifecycle.md` | | **Env vars** (async prefetch, cache, debug, `DATA_OPTIMIZER_*`, Studio) | `reference/env-vars.md` | | Fair streaming benchmarks (`benchmarks/` suite) | `reference/benchmarking.md` | | Offline streaming what-if (`litsim`) | repo `simulator/README.md` | | **Raw adaptive concurrency / look-ahead stages** (clients own rate) | repo `benchmarks/ADAPTIVE_CONCURRENCY.md` | | Lightning Studio env, credentials, free-threading | `reference/lightning-studio.md` | | Write path orchestration; raw internals; pointers to I/O + multi-node | `reference/processing.md` | | Dev env, PR/CI style | `reference/contributing.md` | | Tests & fixtures | `reference/testing.md` | | Tracing (`enable_tracer`, Litracer, Perfetto), breakpoints, env knobs | `reference/debugging.md` | | **Keyed lookup / `dataset_update` / `key_fn`** | `reference/keyed-lookup.md` | ## Public API (`src/litdata/__init__.py`) | Symbol | Purpose | | ----------------------------------------------------------------------------------------------------------- | -------------------------------------------- | | `StreamingDataset` / `StreamingDataLoader` | Optimized stream + resumable loader | | `CombinedStreamingDataset` / `ParallelStreamingDataset` | Mix or zip streams | | `StreamingRawDataset` | Raw file stream | | `TokensLoader` | Token windows for LLMs | | `optimize` / `map` / `merge_datasets` / `walk` / `complete_dataset` | Write / transform / merge / finish index | | `list_media_folder` / `iter_webdataset_tar` | Class-folder / WebDataset ingest lists |