diff --git a/.claude/skills/litdata/SKILL.md b/.claude/skills/litdata/SKILL.md new file mode 100644 index 000000000..96c2bb7db --- /dev/null +++ b/.claude/skills/litdata/SKILL.md @@ -0,0 +1,75 @@ +--- +name: litdata +description: Work on the LitData codebase — understand its architecture, contribute code, write/run tests, and debug/profile runtime issues. Use when navigating or editing src/litdata, answering "how does X work / where does X live", getting a change merge-ready, adding or running tests, or diagnosing a slow/hanging/nondeterministic StreamingDataset or optimize/map run. Covers streaming (read) and processing (write) pipelines, cloud backends, chunk format, shuffling, CLI, tests, CI, and tracing. +--- + +# LitData + +LitData (`import litdata`) preprocesses and streams datasets for PyTorch training. Two pipelines that mirror each other: + +- **Write** (`optimize` / `map`): distributed workers turn raw data into a chunked binary format (`chunk-*.bin` + `index.json`) and upload it. → `src/litdata/processing/` +- **Read** (`StreamingDataset` + `StreamingDataLoader`): stream those chunks back cloud → local cache → decoded items → batches, with deterministic shuffling and resumable state. → `src/litdata/streaming/` + +A lighter path, `StreamingRawDataset` (`src/litdata/raw/`), streams *un-optimized* original files, skipping the write step. + +## Read this first, then open the matching reference file + +Keep this SKILL.md as the map; load a `reference/` file only for the task at hand. + +| Your task | Read | +| ----------------------------------------------------------------------- | ---------------------------------------- | +| Understand the read path, chunk format, shuffling, resume, item loaders | `reference/streaming.md` | +| Cache↔Writer/Reader, PrepareChunksThread, shared-chunk deletion races | `reference/cache-and-chunk-lifecycle.md` | +| Understand the write path (`optimize`/`map`), worker model, raw indexer | `reference/processing.md` | +| Set up dev env, coding style, branch/PR flow, lint/type/CI gates | `reference/contributing.md` | +| Write or run tests, fixtures, mocking cloud, gating | `reference/testing.md` | +| Trace/profile, set worker breakpoints, env knobs, diagnose failures | `reference/debugging.md` | + +## Public API (`src/litdata/__init__.py`) + +| Symbol | Purpose | Defined in | +| -------------------------------------------- | --------------------------------------------- | ------------------------------------------------------- | +| `StreamingDataset` | Read optimized chunks (`IterableDataset`) | `streaming/dataset.py:51` | +| `StreamingDataLoader` | `DataLoader` subclass with resumable state | `streaming/dataloader.py:559` | +| `CombinedStreamingDataset` | Sample one of N datasets per step (by weight) | `streaming/combined.py:40` | +| `ParallelStreamingDataset` | Pull one sample from every dataset per step | `streaming/parallel.py:44` | +| `StreamingRawDataset` | Stream raw files, no optimize step | `raw/dataset.py:95` | +| `TokensLoader` | Item loader for tokenized/NLP data | `streaming/item_loader.py:402` | +| `optimize` | Turn a dataset into litdata chunks | `processing/functions.py:387` | +| `map` | Run a fn over inputs for side effects | `processing/functions.py:242` | +| `merge_datasets` | Merge several optimized datasets | `processing/functions.py:675` | +| `walk` | Cloud-optimized `os.walk` | `processing/functions.py:621` | +| `train_test_split` | Split a `StreamingDataset` by chunk ROIs | `utilities/train_test_split.py:14` | +| `index_parquet_dataset` / `index_hf_dataset` | Index parquet / HF for streaming | `streaming/writer.py:578`, `utilities/hf_dataset.py:13` | +| `breakpoint` | Multiprocessing-safe pdb (works in workers) | `utilities/breakpoint.py:33` | + +## Package map + +- `streaming/` — read pipeline · `processing/` — write pipeline (`optimize`/`map`) · `raw/` — raw streaming + file indexer +- `cli/` — the `litdata` command; dispatch is `__main__.py:app` → `parser.parse_args` iterates `COMMAND_REGISTRY` (`cli/commands.py:68`). Commands: `cache clear`, `cache path`, `optimize` (stub). Add one by appending a registrar. +- `utilities/` — `env.py` (rank detection), `encryption.py`, `subsample.py`, `train_test_split.py`, `parquet.py`, `hf_dataset.py`, `dataset_utilities.py` (`get_default_cache_dir`), `_pytree.py` (vendored; excluded from lint/type) +- `constants.py` — `_*_AVAILABLE` optional-dep flags, env-var knobs, dtype maps, default chunk size (`1<<26` = 64 MB) · `debugger.py` — structured tracing + +## Concepts shared across both pipelines + +- **Chunk format** (`writer.py:218`): `[num_items:uint32][offset_array:uint32[N+1]][item_data:bytes]`. Per-worker `{rank}.index.json` files are merged into one `index.json` holding a `chunks` list + `config` (compression, `data_format`, `data_spec` treespec, `item_loader` class). +- **Item loaders own byte layout AND interval math** (`item_loader.py`): `PyTreeLoader` (default), `TokensLoader` (mmap token windows), `ParquetLoader`. The class in `index.json` must match the reader's (`config.py:361`). +- **Distribution is env-var driven, not networked** — rank read from `_DistributedEnv`/`_WorkerEnv` (read) or `DATA_OPTIMIZER_*` (write). +- **Determinism**: read-path shuffling is seeded by `seed`+`epoch`+`num_chunks`+`chunk_index`; same inputs ⇒ same order ⇒ resumable. +- **Pluggable registries**: `Downloader` (read, by URL prefix, `downloader.py`), `FsProvider` (write/management, `fs_provider.py`), serializers (`serializers.py`), compressors (`compression.py`). `resolver.py` turns a path/URL/teamspace path into a `Dir`. + +## Design principles (CONTRIBUTING.md) — honor when editing + +"One less thing to remember." No abstractions on top of pure PyTorch. Simple, readable internal code (many users aren't engineers). Backward-compatible APIs with deprecation warnings. Test-driven: reproduce a bug as a failing test, then fix. + +## Quick commands + +```bash +make setup # dev env (uv install + pre-commit) +pre-commit run --all-files # lint + format + hooks +mypy # type check (files=["src"]) +pytest tests/path/test_x.py::test_name -v --capture=no # one test +litdata cache path # print cache dir · litdata cache clear # wipe it +``` + +Runnable end-to-end examples live in `examples/`; the README's `
` feature blocks are the narrative docs (line refs are cited in the reference files). Repo version: see `src/litdata/__about__.py`. diff --git a/.claude/skills/litdata/reference/cache-and-chunk-lifecycle.md b/.claude/skills/litdata/reference/cache-and-chunk-lifecycle.md new file mode 100644 index 000000000..705bdacd4 --- /dev/null +++ b/.claude/skills/litdata/reference/cache-and-chunk-lifecycle.md @@ -0,0 +1,77 @@ +# Cache, Writer/Reader, PrepareChunksThread & the chunk lifecycle + +This is the deepest and most race-prone part of LitData. Read it before touching anything under `streaming/{cache,reader,writer,config,item_loader,downloader}.py`. + +## Cache: the façade over Writer and Reader + +`Cache` (`cache.py:35`) is the single object a `StreamingDataset` talks to. It is **bidirectional**: + +- **Write side** — `Cache.__setitem__`/`_add_item` → `BinaryWriter` (`writer.py:50`). Used by the *caching* path (`CacheDataLoader`) and by `optimize`'s `Cache`. `Cache.done()` flushes, `Cache.merge()` concatenates per-rank `{rank}.index.json` into one `index.json`. +- **Read side** — `Cache.__getitem__` (`cache.py:151`) → `BinaryReader.read` (`reader.py:438`). Used by `StreamingDataset`. + +`Cache.filled` (`cache.py:120`) decides write-vs-read: it checks for `index.json` + expected worker count. A missing `index.json` when reading is the "did you optimize?" error (`dataset.py:322`). + +The chunk binary format (`[num_items:uint32][offset:uint32[N+1]][data]`, `writer.py:218`) and item-loader coupling are covered in [streaming.md](streaming.md). This file focuses on the **runtime read machinery** and the **cross-worker chunk lifecycle**, which streaming.md only summarizes. + +## BinaryReader.read — the per-item hot path + +`BinaryReader.read(index)` (`reader.py:438`): + +1. Lazily loads `ChunksConfig` (`_try_load_config`) — downloads `index.json` on first call. +2. `config[index]` (`config.py:289`) → `(local_chunkpath, begin, filesize_bytes)`. For compressed data `local_chunkpath` is the **decompressed** `.bin` path. +3. If remote/compressed: `setup_thread_and_download_chunk` starts (once) the per-worker `PrepareChunksThread` and enqueues this worker's chunks for download. +4. `item_loader.load_item_from_chunk(...)` (`item_loader.py:172`) opens the `.bin`, seeks the offset table, reads `[begin,end)`, deserializes. **If the `.bin` isn't present yet it busy-waits** (`item_loader.py:206-223`): sleeps 0.1s, at `FORCE_DOWNLOAD_TIME` (30s) calls `force_download`, and at `MAX_WAIT_TIME` (120s) raises `FileNotFoundError: The hasn't been found.` +5. On chunk transition (`index.chunk_index != self._last_chunk_index`): `_decrement_local_lock(last)` + `delete([last])` (enqueues deletion of the chunk just finished). +6. On `index.is_last_index`: close handle, decrement, delete, stop the thread, reset per-epoch state. + +`rank` (`reader.py:435`) = `global_rank * worker_world_size + worker_rank` — the flat worker identity used everywhere. + +## PrepareChunksThread — one daemon thread per worker + +`PrepareChunksThread` (`reader.py:50`) is what makes streaming overlap with training. Each `BinaryReader` owns one. It services three queues in its `run()` loop (`reader.py:276`): + +- `_force_download_queue` (checked first each iteration) → `_force_download` (`reader.py:240`): delete-then-redownload a chunk a reader is blocked on. Attached to the item loader at `reader.py:411` so `item_loader.force_download` feeds this thread. +- `_to_download_queue` → `download_chunk_from_index` for each assigned chunk, bounded by `max_pre_download` (default 2) via `_pre_download_counter`. `_END_TOKEN` stops it. +- `_to_delete_queue` → `_maybe_delete_chunks` → `_apply_delete`, gated by `max_cache_size` and `_can_delete_chunk` (`reader.py:228`). + +**Eviction is only active when `max_cache_size` is set** (or `MAX_CACHE_SIZE` env). `_delete_chunks_when_processed` (`reader.py:84`) is true when a node's slice doesn't fit in the cache; then chunks are deleted aggressively as they're consumed. Otherwise deletion waits until the folder exceeds `max_cache_size` (`_get_folder_size`, `reader.py:581`, which is deliberately robust to deletion races and ignores `.cnt`/`.lock`/`.zstd.bin`). + +## Distributed sampling — how workers get disjoint (mostly) chunk slices + +Covered in depth in [streaming.md](streaming.md) ("Shuffling & sharding"). The essentials that matter for the lifecycle: + +- `Shuffle.get_chunks_and_intervals_per_workers` returns `workers_chunks` / `workers_intervals` — **flat lists indexed by `rank * num_workers + worker`**, length `world_size * num_workers`. +- `_associate_chunks_and_intervals_to_workers` (`utilities/shuffle.py:65`) gives each worker an item budget and **splits a chunk's interval across worker boundaries** when it straddles them. Consequence: **a single chunk can be assigned to several workers**, each reading a different sub-interval. This is the source of all the deletion complexity below. +- `dataset.__iter__` slices out this worker's `worker_chunks`/`worker_intervals` (`dataset.py:378`) and computes the skip-deletion mapping for the node. + +## The shared-chunk deletion problem (and the three mechanisms guarding it) + +Because a chunk can be shared by multiple workers on a node (all pointing at the same cache file), **one worker deleting a chunk after it finishes can pull the file out from under another worker that still needs it** → `FileNotFoundError: chunk-N-M.bin hasn't been found` raised at `item_loader.py:223` (PyTree) / `:515` (Tokens). Three mechanisms exist to prevent this: + +### 1. Static skip-deletion list (`skip_chunk_indexes_deletion` / `can_delete`) + +Computed in `dataset.__iter__` via `_find_chunks_per_workers_on_which_to_skip_deletion` (`utilities/shuffle.py:147`): for each shared chunk it determines the single worker that reads it **last in consumption order**, and tells every *other* sharing worker to skip deleting it. Stored on `ChunksConfig.skip_chunk_indexes_deletion`; queried via `config.can_delete(chunk_index)` (`config.py:117`). This is the primary, race-free guard (it is derived from the deterministic shuffle assignment, not from wall-clock state). **It must be consulted in `reader._apply_delete`** — historically it was computed but ignored (see the bug note below). + +### 2. Cross-worker reference counting (`.cnt` + `.cnt.lock` files) + +`downloader._increment_local_lock` (`downloader.py:55`) bumps `.bin.cnt` under a `FileLock` when a worker will read a chunk; `reader._decrement_local_lock` (`reader.py:111`) decrements when a worker finishes it and removes the file at zero. `_apply_delete` refuses to delete while `_remaining_locks > 0` (`reader.py:180`). This is the dynamic guard. **Its correctness depends on every sharing worker incrementing before any sharer decrements to zero** — see the increment-lag hazard below. + +### 3. Priority force-redownload (`_force_download_queue`) + +When a reader blocks on a missing `.bin` for `FORCE_DOWNLOAD_TIME`, `item_loader.force_download` enqueues the chunk onto its own thread's `_force_download_queue`; `_force_download` (`reader.py:240`) deletes any stale copy and re-downloads under `FileLock(..., timeout=0)`. This is a last-resort recovery, not prevention. + +### Failure modes / invariants to preserve + +- **`can_delete` must gate `_apply_delete`.** If its result is computed but unused, mechanism 1 is dead and only the racy refcount protects shared chunks. +- **Increment-lag:** increments happen lazily in the prefetch thread (bounded by `max_pre_download`), while decrement+delete happen when a worker finishes a chunk. If a fast worker finishes a shared chunk before a slow co-worker has prefetched (incremented) it, the count is 0 and the chunk is deleted prematurely. Mechanism 1 (static list) is what actually closes this hole; the refcount alone does not. +- **The skip list must be set on resume too.** It is computed in `dataset.__iter__`; if it is only set on the fresh-epoch branch, resumed runs have no mechanism-1 protection. +- **Node slicing:** `workers_chunks` is indexed by `rank * num_workers + worker`, so a node's workers start at `first_rank_this_node * num_workers`. Using `* (node_size * num_workers)` overshoots for multi-node. +- **`skip_lock=True`** (force-redownload path) intentionally bypasses both the skip list and the refcount — a worker re-fetching its *own* needed chunk must be allowed through. +- **s3transfer temp files:** boto3 downloads to `chunk-*.zstd.bin.` then renames. Leftover ``-suffixed files in the cache are interrupted/duplicated downloads — a symptom of delete/redownload churn, not the root cause. `_get_folder_size` counts `.bin`-containing temp files but ignores `.zstd.bin`. + +### Debugging this class of failure + +1. Reproduce with `num_workers` > 1 and `max_cache_size` set small enough to force eviction (that's when deletion runs). +2. `enable_tracer()` — the `increment_lock_*` / `decrement_lock_*` / `delete_chunk_*` events show the refcount timeline per chunk; a `delete_chunk_X` with a later `read_chunk_X` on another worker is the race. +3. `DEBUG_LITDATA=1` writes `.tmb` tombstones recording which rank deleted a chunk and its `can_delete` verdict. +4. Inspect the cache dir: leftover `.zstd.bin.` + missing `.bin` for a low-numbered chunk = premature deletion + failed redownload. diff --git a/.claude/skills/litdata/reference/contributing.md b/.claude/skills/litdata/reference/contributing.md new file mode 100644 index 000000000..433b0142a --- /dev/null +++ b/.claude/skills/litdata/reference/contributing.md @@ -0,0 +1,75 @@ +# Contributing to LitData + +Full reference: `CONTRIBUTING.md`. This is the fast path. + +## Environment setup + +Requires Python ≥ 3.10 and `make`. LitData uses `uv` for installs. + +```bash +make setup # install-dependencies + install-pre-commit (recommended one-shot) +``` + +Or manually: + +```bash +uv pip install -e ".[extras]" -r requirements/test.txt +pre-commit install +``` + +The package is `src/`-layout (`package_dir={"": "src"}`), installed editable. CLI entry point: `litdata.__main__:app`. + +## Design principles (honor when writing code) + +LitData is used by researchers, not just engineers: + +- **"One less thing to remember."** Simplify the API; minimize what the user must track. +- **No abstractions on top of pure PyTorch.** +- **Simple, readable internal code** over clever tricks. +- **Backward-compatible APIs** with clear deprecation warnings. +- **Thorough tests** — valued even more than features. Reproduce a bug as a failing test, then fix. + +## Branch & PR conventions + +- Branch off `main` (never work on `main` directly). Name: `/_`, type ∈ `bugfix|feature|docs|tests`. +- Features: open a GitHub issue first. Ask "is this NECESSARY?" — LitData rejects PRs that only add engineering complexity. +- Link the issue in the PR; update/add tests and docs. Use `[wip]` / `[blocked by #N]` title tags when relevant. +- Output uses f-strings, except logging which stays lazy `%`-style: `logging.info("Hello %s!", name)`. +- Changelog: `src/litdata/CHANGELOG.md` (exempt from the large-file pre-commit hook). + +## Before you push — run these locally + +```bash +pre-commit run --all-files # lint + format + all hooks (auto-fixes most issues) +ruff check . # add --fix to auto-fix +ruff format . +mypy # config: files=["src"], disallow_untyped_defs=true +pytest tests/path/test_x.py::test_name -v --capture=no # see testing.md +``` + +## Style rules enforced (`pyproject.toml [tool.ruff]`) + +- `line-length = 120`, `target-version = "py310"`. +- Rule sets: `E, F, W`, `S` (bandit), `UP` (pyupgrade), `I` (isort), `C4`, `D` (pydocstyle, **Google convention**), `PT`, `RET`, `SIM`, `NPY201`, `RUF100`. Max mccabe complexity 10. +- `assert` allowed (`S101` ignored). Docstring rules `D1xx` relaxed in `src/**`, `tests/**`, `examples/**`. +- `src/litdata/debugger.py` and `utilities/_pytree.py` fully excluded from ruff; several files excluded from mypy (`[tool.mypy] exclude`). +- New public code needs type annotations (`disallow_untyped_defs`). + +## CI a PR must pass (`.github/workflows/`) + +**`ci-testing.yml`** — job `pytester` (matrix, `fail-fast: false`): + +- OS: `ubuntu-22.04`, `macos-14`, `windows-2022`. Python: `3.10–3.14` (windows excludes 3.13/3.14). +- Two phases: fast tests parallel (`pytest tests --ignore=tests/processing --ignore=tests/raw -n 2 --dist=loadgroup ...`), then processing/raw sequential. `--dist=loadgroup` is **required** (see testing.md). +- `--timeout=120` per test; FFmpeg installed for video tests; `UV_TORCH_BACKEND=cpu`. +- `testing-guardian` job is the required gate — fails the PR if any leg fails. + +**`ci-checks.yml`** — reusable Lightning workflows: `check-typing` (mypy), `check-schema`, `check-package` (import test), `check-docs`, `check-md-links`. + +**pre-commit.ci** runs hooks on the PR and auto-fixes. Hooks: ruff + ruff-format, codespell, mdformat, prettier (json/yml/toml), pyproject-fmt, validate-pyproject, large-file guard (max 350 kb). + +## Adding examples / docs + +- Runnable examples → `examples/` (subfolders `getting_started/`, `multi_modal/`, ...). Keep minimal; use public or synthetic data. +- README is the primary narrative doc; feature sections are `
` blocks. `README.md` is exempt from `mdformat`/`trailing-whitespace` hooks — keep its formatting consistent manually. +- Sphinx docs build from `docs/source` (`make docs`). diff --git a/.claude/skills/litdata/reference/debugging.md b/.claude/skills/litdata/reference/debugging.md new file mode 100644 index 000000000..c3456e709 --- /dev/null +++ b/.claude/skills/litdata/reference/debugging.md @@ -0,0 +1,88 @@ +# Debugging & profiling LitData + +## Breakpoints inside worker processes + +Normal `breakpoint()` doesn't work in DataLoader or `optimize`/`map` worker subprocesses (no stdin). Use LitData's multiprocessing-safe pdb, which reopens `sys.stdin` under a lock (`utilities/breakpoint.py:33`, exported as `litdata.breakpoint`): + +```python +from litdata.utilities.breakpoint import breakpoint # or: import litdata; litdata.breakpoint() +breakpoint() # works inside a worker +``` + +Run with `num_workers=0` first when you can — it keeps everything in the main process so a plain debugger works. + +## Structured tracing + Litracer/Perfetto (`debugger.py`) + +`enable_tracer()` (`debugger.py:124`) logs the streaming pipeline to a file, which you convert into a Chrome/Perfetto trace. + +```python +import litdata as ld +from litdata.debugger import enable_tracer + +enable_tracer() # WARNING: delete an existing litdata_debug.log first, or traces mix + +if __name__ == "__main__": + dataset = ld.StreamingDataset("s3://my-bucket/my-data", shuffle=True) + for batch in ld.StreamingDataLoader(dataset, batch_size=64): + ... +``` + +Then: + +```bash +python main.py # writes litdata_debug.log +go install github.com/deependujha/litracer@latest # or download a release binary +litracer litdata_debug.log -o litdata_trace.json -w 100 +# open litdata_trace.json in ui.perfetto.dev (preferred) or chrome://tracing +``` + +`enable_tracer(flush_interval=5, item_loader=True, iterating_dataset=True, getitem_dataset_for_chunk_index=True)` sets the `LITDATA_LOG_*` env vars and returns a singleton `LitDataLogger`. `TimedFlushFileHandler` flushes every N seconds from a daemon thread. `env_info()` auto-injects distributed + worker rank/world-size into every event. `ChromeTraceColors` holds trace phase colors. + +## Environment-variable knobs (`constants.py` + `debugger.py`) + +| Env var | Effect | Default | +| ------------------------------- | ------------------------------------------------ | --------------------- | +| `LITDATA_CACHE_DIR` | Override the chunk cache directory | `~/.lightning/chunks` | +| `DEBUG_LITDATA` | Enable internal debug behavior (`_DEBUG`) | `0` | +| `PRINT_DEBUG_LOGS` | Print debug logs to stdout | `0` | +| `MAX_WAIT_TIME` | Max seconds to wait for a chunk download | `120` | +| `FORCE_DOWNLOAD_TIME` | Force re-download threshold (s) | `30` | +| `LITDATA_LOG_FILE` | Trace log file path | `litdata_debug.log` | +| `LITDATA_LOG_LEVEL` | Trace log level | `DEBUG` | +| `LITDATA_LOG_ITERATING_DATASET` | Include `iterating_dataset` events | `True` | +| `LITDATA_LOG_GETITEM` | Include `getitem_dataset_for_chunk_index` events | `True` | +| `LITDATA_LOG_ITEM_LOADER` | Include `item_loader` events | `True` | +| `LITDATA_DISABLE_VERSION_CHECK` | Skip the upgrade prompt | `0` | +| `ENABLE_STATUS_REPORT` | Emit status reports | `0` | + +`optimize`/`map` cross-node coordination: `DATA_OPTIMIZER_NODE_RANK`, `DATA_OPTIMIZER_NUM_NODES`, `DATA_OPTIMIZER_GLOBAL_RANK`, `DATA_OPTIMIZER_NUM_WORKERS`. + +Cache CLI: `litdata cache path` (print dir) · `litdata cache clear` (rmtree it). + +## Common failure modes + +**Streaming (read) — see streaming.md** + +- **"did you optimize?" / `FileNotFoundError`** → no `index.json` at the path (`dataset.py:322`); the data wasn't optimized, or you pointed at raw files (use `StreamingRawDataset`). +- **Hang/timeout on chunk download** → raise `MAX_WAIT_TIME`; check credentials/`storage_options`; confirm the URL prefix matches a registered `Downloader`. +- **Cache grows without bound** → set `max_cache_size` (default `"100GB"`); eviction is gated by `.cnt` refcount + `FileLock`. Stale `.lock`/`.cnt` files from a crash can block deletion — `litdata cache clear`. +- **Deadlock: parquet + workers** → `ParquetLoader` under `fork` with `num_workers>0` raises by design (`dataloader.py:629`); use `spawn`/`forkserver`. +- **Nondeterministic / wrong order after resume** → shuffling is seeded by `seed`+`epoch`+`num_chunks`+`chunk_index`; `_validate_state_dict` (`dataset.py:607`) hard-fails if `shuffle`/`num_workers`/`seed`/`input_dir`/`item_loader`/`drop_last` changed between save and resume (override with `force_override_state_dict`). +- **`state_dict()` raises** → must be called in the main process, not a worker (`dataset.py:574`). +- **Uneven batches across ranks** → `drop_last` defaults to `True` in distributed for a reason; don't force `False` there. + +**Processing (write) — see processing.md** + +- **`optimize` hangs, no error** → a worker exception is swallowed into `error_queue` then re-raised by the main loop; if it can't build a traceback string the loop waits until all workers die. Run with `num_workers=1` (or `fast_dev_run=True`) to surface the real exception. +- **Peers hang with `keep_data_ordered=False`** → termination relies on the `ALL_DONE` sentinel being re-inserted by each worker; a custom early-exit that drops it hangs the pool. +- **`ValueError` on `prepare_structure`** → must return `list | StreamingDataLoader | multiprocessing.Queue` (`:1247`). +- **`map` fn silently does nothing** → `map` recipes' `prepare_item` must return `None` (`:913`) and must write outputs to `output_dir`. +- **Pickling error under spawn** → worker args must be picklable; avoid closures/weakrefs captured on `self`. + +## Fast triage checklist + +1. Reproduce with `num_workers=0` + a small local dataset → isolates worker/multiprocessing issues from logic bugs. +2. `fast_dev_run=True` (optimize/map) for a ~10-item smoke run. +3. `enable_tracer()` + Litracer/Perfetto to see where time goes (download vs decode vs collate). +4. `litdata cache clear` to rule out stale cache/lock state. +5. `PRINT_DEBUG_LOGS=1` for stdout logs without the trace pipeline. diff --git a/.claude/skills/litdata/reference/processing.md b/.claude/skills/litdata/reference/processing.md new file mode 100644 index 000000000..531c355cb --- /dev/null +++ b/.claude/skills/litdata/reference/processing.md @@ -0,0 +1,106 @@ +# The processing (write) pipeline — `optimize` / `map` + +All paths under `src/litdata/`. This pipeline fans work across workers (and machines) to transform raw data, and for `optimize` writes it into the litdata chunk format that the streaming pipeline reads. + +## Public API (`processing/functions.py`) + +- **`optimize(fn, inputs=None, output_dir="optimized_data", chunk_size=None, chunk_bytes=None, compression=None, encryption=None, num_workers=None, num_nodes=None, machine=None, num_downloaders=None, num_uploaders=None, reader=None, mode=None, use_checkpoint=False, item_loader=None, start_method=None, storage_options={}, keep_data_ordered=True, ...)`** — `functions.py:387`. Runs `fn` per input; the return value is flattened via pytree and serialized into `chunk-*.bin` + `index.json`. **Requires exactly one of `chunk_size` / `chunk_bytes`.** `mode="append"`/`"overwrite"` reuse/replace an existing index. Builds `LambdaDataChunkRecipe` (or `QueueDataChunkRecipe` if a `queue` is given) → `DataProcessor.run`. +- **`map(fn, inputs, output_dir, ...)`** — `functions.py:242`. Applies `fn(input, output_dir)` for side effects (no chunking). **`fn` must write to `output_dir` and return `None`.** Builds `LambdaMapRecipe`. +- **`merge_datasets(input_dirs, output_dir, max_workers=os.cpu_count(), storage_options={})`** — `functions.py:675`. Merges already-optimized datasets by copying chunk files and concatenating their `index.json` chunk lists. Validates matching `data_format`/`compression`. +- **`walk(folder, max_workers=os.cpu_count())`** — `functions.py:621`. Cloud-optimized `os.walk` using a `ThreadPoolExecutor`; order is not depth-first. + +## Orchestration (`processing/data_processor.py`) + +`DataProcessor.run(recipe)` (`data_processor.py:1226`) is the single entrypoint both `optimize` and `map` call. Lifecycle: + +1. **`recipe.prepare_structure(input_dir)`** (`:1244`) returns the work list — must be a `list`, a `StreamingDataLoader`, or a `multiprocessing.Queue` (else raises, `:1247`). +2. **Item→worker assignment** (`:1258-1287`) → `workers_user_items: list[list]` (one sublist per local worker), via: + - `_map_items_to_workers_weighted` (`:377`) — default when `reorder_files` + `input_dir` exist, or when `weights` given. Bin-packs by file size (`_pack_greedily`) across `world_size = num_nodes * num_workers`, then permutes. + - `_map_items_to_workers_sequentially` (`:303`) — contiguous slices; `align_chunking` packs full chunks. + - Queue mode — no static assignment; `shared_queue` is set. +3. **Multi-node slicing** is baked into both mappers via `_get_node_rank()`/`_get_num_nodes()` (env `DATA_OPTIMIZER_NODE_RANK`/`_NUM_NODES`); each returns only its node's worker slice. **No cross-node RPC** — coordination is purely env-var driven. +4. **Checkpointing** (`:1297`) trims each worker's list to resume from `checkpoint_next_index`; `fast_dev_run` trims to N items. +5. **`_create_process_workers`** (`:1462`) spawns one `DataWorkerProcess` per worker. +6. **Progress loop** (`:1376`) polls `error_queue` (re-raises via `_exit_on_error`, which `terminate()`s all workers) and `progress_queue` (tqdm). Exits when the counter equals `num_items` or all workers die. +7. **`recipe._done(...)`** (`:1432`) merges caches / writes the index; index merge + platform registration run only on the last node. + +## Key classes + +- **`DataProcessor`** (`:1114`) — the main-process orchestrator. Owns `error_queue`, `msg_queue`, `progress_queue`, `stop_queues`, `shared_queue`. +- **`BaseWorker`** (`:481`) — the processing unit. `run()` (`:570`) = `_setup()` → `_loop()` → `_terminate()`, wrapping all exceptions into `error_queue`. +- **`DataWorkerProcess`** (`:927`) — `BaseWorker` + `multiprocessing.Process`; what actually runs in a child process. +- **`DataRecipe`** (`:948`) abstract: `prepare_structure` + `prepare_item` + `_done`. + - **`DataChunkRecipe`** (`:981`) — for `optimize`; chunk_size/chunk_bytes (defaults to 64 MB, `:995`), compression, encryption; `_done` merges + uploads index. + - **`MapRecipe`** (`:1100`) — for `map`; `prepare_item` writes to output and must return `None`. +- **`FakeQueue`** (`:461`) — drop-in for `Queue` when no downloading is needed (in-process, avoids serialization). + +## Producer/consumer model (inside each worker) + +Each `BaseWorker` runs a local pipeline of child processes (spawned in `_setup`, `:580`): + +- **Downloaders** (`_start_downloaders`, `:797`; target `_download_data_target`, `:128`): `num_downloaders` procs pull `(index, item, paths)` off `to_download_queues`, fetch remote files into cache, push ready tuples onto `ready_to_process_queue`. +- **Worker main loop** (`_loop`, `:601`) — the **consumer**: `ready_to_process_queue.get()` → `_handle_data_chunk_recipe` (optimize) or `_handle_data_transform_recipe` (map). Reports progress ~1/s. +- **Uploaders** (`_start_uploaders`, `:839`; target `_upload_fn`, `:232`): push finished chunks/files to `output_dir`. +- **Remover** (`_start_remover`, `:825`; target `_remove_target`, `:190`): deletes processed source files when `remove=True`. + +When `no_downloaders` (no `input_dir`, or a `reader` is set), `ready_to_process_queue` is a `FakeQueue` and `_collect_paths` (`:743`) pushes items directly. + +**Ordered vs shared-queue** (`keep_data_ordered`): `True` (default) → each worker consumes its static slice in order. `False` → all workers share one `Queue` for dynamic load balancing; termination uses the `ALL_DONE` sentinel (`:64`), which each worker re-inserts so peers also stop (`:621`). + +## Cross-process queues + +| Queue | Direction | Purpose | +| ---------------------------------------------------------- | ----------------- | --------------------------------- | +| `error_queue` | worker→main | tracebacks; triggers global abort | +| `progress_queue` | worker→main | `(index, counter)` for tqdm | +| `msg_queue` | worker→main | log lines routed around tqdm | +| `stop_queues` | main→worker | SIGINT graceful stop | +| `ready_to_process_queue` / `shared_queue` | downloader→worker | core work items | +| `to_download_queues` / `to_upload_queues` / `remove_queue` | worker→child | I/O offload | + +## `raw/` — `StreamingRawDataset` and the indexer + +`StreamingRawDataset` (`raw/dataset.py:95`) is a plain `torch.utils.data.Dataset` streaming **original files** (no optimize step). Item structure is user-defined via `setup(files)` (`:151`) returning `list[FileMetadata]` or `list[list[FileMetadata]]` (grouped items). Uses async batched downloads (`__getitems__`/`_download_batch`, `asyncio.gather`). `CacheManager` (`:32`) mirrors remote structure into a local cache (opt-in). Pass `recompute_index=True` to force a rebuild (`indexer.py:58`). + +The **indexer** (`raw/indexer.py`) replaces the optimize pass: `BaseIndexer.build_or_load_index` (`:53`) tries a local cached index, then a remote one (`input_dir/index.json.zstd`), rebuilding via `discover_files` only if neither exists or `recompute_index=True`. `FileIndexer` (`:217`) recursively lists files (fsspec `fs.find` / `Path.rglob`). Note the index filename `index.json.zstd` differs from the optimized `index.json` — don't conflate them. + +## Gotchas (read before editing the engine) + +01. `prepare_structure` must return `list | StreamingDataLoader | multiprocessing.Queue` (`:1247`). +02. `map` recipes' `prepare_item` must return `None` (`:913`); `optimize` recipes' return value is serialized. +03. Start method forced globally: `set_start_method(..., force=True)` (`:1177`), `fork` in notebooks else `spawn`. Under `spawn`, worker args must be picklable. +04. Distribution is env-var driven (`DATA_OPTIMIZER_*`); mappers slice by `world_size = num_nodes * num_workers`. +05. `keep_data_ordered=False` switches to a shared queue + `ALL_DONE` sentinel; early-exit paths must keep re-inserting `ALL_DONE` or peers hang. Timeout is lower (200s vs 300s). +06. `FakeQueue` is not a real queue — no inter-process/blocking semantics. +07. Path detection is heuristic (`_is_path`/`_to_path`, `functions.py:439`); an item with zero detected file paths raises (`:775`). +08. `align_chunking` requires `chunk_size` (not bytes) (`:497`, `:1275`). +09. Index merge + platform registration only on the last node (`num_nodes == node_rank + 1`). +10. Errors are swallowed into `error_queue` then re-raised in the main loop; a worker exception without a traceback string can hang the loop. `_exit_on_error` hard-`terminate()`s siblings — no graceful flush. **To surface a hidden error, run with `num_workers=1` or `fast_dev_run=True`.** +11. Checkpointing unsupported for Queue inputs (`:1298`) or generator `fn`s (`:1303`). + +## Runnable examples (from README) + +```python +# Optimize a dataset into chunks (README:144) +import litdata as ld +def fn(index): + return {"index": index, "image": ..., "class": ...} +if __name__ == "__main__": + ld.optimize(fn=fn, inputs=list(range(1000)), output_dir="fast_data", + num_workers=4, chunk_bytes="64MB") + +# Map: transform files in parallel, write to output_dir (README:205) +import os, litdata as ld +from PIL import Image +def resize_image(image_path, output_dir): + out = os.path.join(output_dir, os.path.basename(image_path)) + Image.open(image_path).resize((224, 224)).save(out) +ld.map(fn=resize_image, inputs=inputs, output_dir="output_dir") + +# Encrypt at sample level (README:1491) +from litdata.utilities.encryption import FernetEncryption +enc = FernetEncryption(password="secret", level="sample") +ld.optimize(fn=fn, inputs=..., output_dir="enc_data", chunk_bytes="64MB", encryption=enc) +``` + +README feature sections: LLM pre-training / tokenization (721), filter illegal data (784), shared queue for optimize (607), queue as input (664), merge (995), distributed optimize (1436), encryption (1478), Lightning data connections (1615). diff --git a/.claude/skills/litdata/reference/streaming.md b/.claude/skills/litdata/reference/streaming.md new file mode 100644 index 000000000..56d5f605c --- /dev/null +++ b/.claude/skills/litdata/reference/streaming.md @@ -0,0 +1,126 @@ +# The streaming (read) pipeline + +All paths under `src/litdata/streaming/` unless noted. This is the inverse of the write pipeline: **cloud storage → download → local chunk file → mmap/seek + deserialize → collate → batch.** + +## Key classes + +| Layer | Class | File | +| ------------------------------------------ | -------------------------------------- | -------------------------------- | +| User dataset (`IterableDataset`) | `StreamingDataset` | `dataset.py:51` | +| DataLoader (subclasses torch `DataLoader`) | `StreamingDataLoader` | `dataloader.py:559` | +| Read/write facade | `Cache` | `cache.py:35` | +| Reader + background download thread | `BinaryReader` / `PrepareChunksThread` | `reader.py:312` / `reader.py:50` | +| Index/chunk metadata | `ChunksConfig` | `config.py:33` | +| Chunk→item decoding | `BaseItemLoader` + subclasses | `item_loader.py` | +| Backend download | `Downloader` subclasses | `downloader.py` | +| Chunk→worker assignment / shuffle | `Shuffle` (`NoShuffle`/`FullShuffle`) | `shuffle.py` | +| Item index | `ChunkedIndex` (dataclass) | `sampler.py:24` | + +## Runtime flow (single item) + +1. `StreamingDataLoader.__init__` (`dataloader.py:614`) pushes `batch_size`/`num_workers`/`shuffle`/`drop_last` into the dataset via `set_*` methods, then calls torch `DataLoader.__init__`. Because `StreamingDataset` is an `IterableDataset`, torch uses its `__iter__`/`__next__` (no external sampler). +2. `StreamingDataLoader.__iter__` (`dataloader.py:687`) sets the epoch, resets per-epoch counters, delegates to `super().__iter__()`, and wraps the iterator to track `_num_samples_yielded_streaming` for checkpointing. +3. `StreamingDataset.__iter__` (`dataset.py:357`) detects `_WorkerEnv`/`_DistributedEnv`, builds the `Cache` (`_create_cache`, `dataset.py:271`), builds the `Shuffle` (`_create_shuffler`, `dataset.py:330`), and computes **this worker's** chunk list + intervals via `shuffler.get_chunks_and_intervals_per_workers(...)`. Slices by `worker_rank = global_rank * worker_world_size + worker_rank` (`dataset.py:377`). +4. `StreamingDataset.__next__` (`dataset.py:504`) pops the next global index from `upcoming_indexes` (refilled per-chunk, shuffled in-chunk by `shuffler.__call__`), wraps it in a `ChunkedIndex`, calls `__getitem__`. +5. `__getitem__` (`dataset.py:481`) → `Cache.__getitem__` (`cache.py:151`) → `BinaryReader.read` (`reader.py:438`). +6. `BinaryReader.read` lazily loads `ChunksConfig` (`_try_load_config`, `reader.py:384`), resolves `(chunk_filepath, begin, filesize_bytes)` via `ChunksConfig.__getitem__` (`config.py:289`), queues the `PrepareChunksThread` (`reader.py:398`), then asks the item loader to decode. It also drives chunk lifetime: decrement lock + queue deletion of the previous chunk, and on `is_last_index` stops the thread. +7. Batching: torch collates via `StreamingDataLoaderCollateFn` (`dataloader.py:502`), which unwraps the `__SAMPLES__`/`__NUM_SAMPLES_YIELDED__` envelope. + +**Prefetching/eviction**: `PrepareChunksThread` (`reader.py:50`) is a per-worker daemon thread. It pre-downloads up to `max_pre_download` chunks, pre-loads them, and deletes consumed chunks once `max_cache_size` is exceeded (`_maybe_delete_chunks`, `reader.py:206`). Cross-worker refcounting uses `.cnt` files + `FileLock`. + +**On-demand mode**: if `on_demand_bytes` is true and the data is remote + PyTreeLoader + no compression/encryption, `read` fetches only the item's byte range (HTTP range GET) via `read_item_bytes` (`reader.py:535`) → `ChunksConfig.download_chunk_bytes_from_index` (`config.py:160`). Iterating flips `on_demand_bytes=False` (`dataset.py:365`) so full chunks are cached during training; it flips back on `StopIteration`. + +## On-disk chunk format + +Binary layout (documented at `writer.py:218`): + +``` ++-----------+----------------+-----------+ +| num_items | offset_array | item_data | +| uint32 | uint32[N+1] | bytes | ++-----------+----------------+-----------+ +``` + +Writer `BinaryWriter` (`writer.py:50`): `add_item`/`serialize` flatten the sample via pytree, serialize each leaf, record `data_format` (serializer names) and `data_spec` (treespec) once, delegate final byte layout to the item loader's `encode_data`. Chunks roll over when `chunk_bytes` or `chunk_size` is exceeded (`writer.py:381`). Filenames: `chunk-{rank}-{idx}.bin`. Each worker writes `{rank}.index.json`; `merge` (`writer.py:455`) concatenates all per-rank indexes into one `index.json` and asserts configs match. + +Reader `ChunksConfig` (`config.py:33`): loads `index.json`, deserializes `data_spec`, calls `item_loader.setup` + `generate_intervals`, maps global index → `(chunk_index, offset)` (`config.py:266`). + +## Item loaders (`item_loader.py`) — own byte layout AND interval math + +- **`PyTreeLoader`** (`item_loader.py:129`) — default. `load_item_from_chunk` seeks the offset table, reads `[begin,end)`, `deserialize` splits by the `data_format` header and reconstructs via `tree_unflatten`. Supports sample/chunk encryption and MDS (Mosaic) format. +- **`TokensLoader`** (`item_loader.py:402`) — NLP. Requires `block_size`. Uses `np.memmap` over the chunk, returns fixed-size token windows with `torch.frombuffer`. `dim` in chunk metadata = token count. +- **`ParquetLoader`** (`item_loader.py:608`) — reads parquet chunks with polars/pyarrow. Used for `hf://` and parquet-indexed datasets. + +`ChunksConfig._validate_item_loader` (`config.py:361`) raises if the passed loader class ≠ `config["item_loader"]`. + +## Backends + +- **`downloader.py`** — read path. `Downloader` ABC (`:41`); subclasses `S3Downloader`, `R2Downloader`, `GCPDownloader`, `AzureDownloader`, `HFDownloader`, `LocalDownloader`. Selected by URL prefix via `_DOWNLOADERS` registry + `get_downloader` (`:658`); register new ones with `register_downloader` (`:632`). +- **`fs_provider.py`** — write/management path (uploads, existence, deletes). `FsProvider` ABC; `S3FsProvider`, `R2FsProvider`, `GCPFsProvider`. Selected via `_get_fs_provider` (`:329`). Covers fewer backends than downloader. +- **`resolver.py`** — path/URL resolution (NOT I/O backend selection). `_resolve_dir` (`:50`) → `Dir(path, url, data_connection_id)`. Resolves Lightning teamspace paths (`/teamspace/...`) to bucket URLs via the Lightning SDK. +- **`client.py`** — `S3Client`/`R2Client` wrap boto3 with credential refresh + temporary project-role credentials from the Lightning control plane. +- **`compression.py`** — `Compressor` ABC + `ZSTDCompressor`, registered into `_COMPRESSORS`. Decompression in `ChunksConfig.try_decompress` (`config.py:182`). +- **`serializers.py`** — `Serializer` ABC; ordered `_SERIALIZERS` dict (`:553`, order matters because `can_serialize` is tried top-to-bottom). Includes `str/bool/int/float/video/tifffile/pil/jpeg/numpy/tensor/pickle` (catch-all). Users pass custom serializers via `StreamingDataset(serializers=...)`. + +## Shuffling & sharding — two independent stages + +**Stage A — chunk→worker assignment** (`Shuffle.get_chunks_and_intervals_per_workers`): + +- `NoShuffle` (`shuffle.py:60`) keeps order; `FullShuffle` (`shuffle.py:83`) permutes deterministically with `np.random.RandomState([seed, seed_shift])` where `seed_shift = 1` for multi-node else `current_epoch`. +- Both call `_associate_chunks_and_intervals_to_workers` (`utilities/shuffle.py:65`), which greedily fills each worker's item budget and **splits chunk intervals** across worker boundaries — so one chunk can be shared by multiple workers (and re-downloaded by each). +- Multi-node epoch>1 adds an intra-node reshuffle (`_intra_node_chunk_shuffle`). + +**Stage B — in-chunk item order** (`Shuffle.__call__`): `FullShuffle` permutes with `np.random.RandomState([seed, num_chunks, current_epoch, chunk_index])` (`shuffle.py:140`), deterministic per (chunk, epoch). + +`sampler.py` (`ChunkedIndex`, `CacheBatchSampler`) is used by the **write/caching path** (`CacheDataLoader`), NOT the read path — don't confuse it with `shuffle.py`. + +## Resume / checkpointing + +`_replay_sampling` (`dataset.py:755`) + `_replay_chunks_sampling` (`dataset.py:778`) reconstruct where each worker stopped from `num_samples_yielded`/`num_workers`/`batch_size`. `_resume` (`dataset.py:425`) re-shuffles the current chunk and skips consumed indexes. `_validate_state_dict` (`dataset.py:607`) hard-fails if shuffle/num_workers/seed/input_dir/item_loader/drop_last changed (unless `force_override_state_dict`). `state_dict` raises if called inside a worker (`dataset.py:574`). + +## CombinedStreamingDataset & ParallelStreamingDataset + +Both subclass `_BaseStreamingDatasetWrapper` (`utilities/base.py:27`) which fans `set_*`/`load_state_dict` to all wrapped datasets and carries the `__SAMPLES__`/`__NUM_SAMPLES_YIELDED__`/`__NUM_CYCLES__` envelope. + +- **`CombinedStreamingDataset`** (`combined.py:40`) — samples **one** dataset per step by weight (seeded `random.Random`). `iterate_over_all=True` iterates until all exhausted, re-normalizing weights; `batching_method` is `"stratified"` (mix within a batch) or `"per_stream"` (whole batch from one dataset). Weights default to inverse dataset length. +- **`ParallelStreamingDataset`** (`parallel.py:44`) — pulls **one sample from every dataset per step**, yields a tuple or `transform(samples[, rngs])`. Supports cycling (`length` = None/int/`inf`), resumable per-worker RNGs seeded via SHA-256 from `(seed, worker_rank, samples_yielded, cycles)`, and `reset_rngs`. + +## Gotchas (read before editing) + +- **Item loader coupling**: `encode_data` (write) and `load_item_from_chunk`/`generate_intervals` (read) must stay consistent. Header size `len(data_format)*4` and the `(1+(index-begin))*4` offset formula are load-bearing. +- **Chunk lifetime is fragile**: deletion gated by `.cnt` refcount + `FileLock`; `skip_chunk_indexes_deletion` (`config.py:117`) prevents deleting a chunk a same-node worker still needs. Deletion is deferred until after the next item loads (avoids mmap segfaults). Windows requires `close()` before `os.remove`. +- **ParquetLoader + fork deadlock**: `StreamingDataLoader` raises if `ParquetLoader` is used with `num_workers>0` under `fork` — use `spawn`/`forkserver` (`dataloader.py:629`). +- **`num_workers=0` counts as 1 worker** internally (`dataset.py:200`). +- **Datasets are immutable**: resolver refuses to write into a non-empty output dir unless `append`/`overwrite` (`resolver.py:311`). +- **Two dataloaders exist**: `CacheDataLoader` (`dataloader.py:260`, write/cache path) vs `StreamingDataLoader` (read path). Don't cross-wire them. +- Missing `index.json` in `_create_cache` raises a "did you optimize?" error (`dataset.py:322`). + +## Runnable examples (from README) + +```python +# Stream optimized data (README:328) +from litdata import StreamingDataset, StreamingDataLoader +dataset = StreamingDataset('s3://my-bucket/my-data', shuffle=True, drop_last=True) +dataloader = StreamingDataLoader(dataset, batch_size=64) +for batch in dataloader: + ... + +# Custom S3-compatible endpoint (README:342) +dataset = StreamingDataset('s3://my-bucket/my-data', storage_options={ + "endpoint_url": "...", "aws_access_key_id": "...", "aws_secret_access_key": "..."}) + +# Custom cache dir + cache cap (README:360; dataset.py:63 default "100GB") +dataset = StreamingDataset('s3://my-bucket/my-data', cache_dir="/path/to/cache", + max_cache_size="50GB") + +# Combine datasets with weights (README:846) +from litdata import CombinedStreamingDataset +combined = CombinedStreamingDataset(datasets=[ds1, ds2], weights=(0.5, 0.5), + iterate_over_all=False) + +# Pause & resume (README:573) — StreamingDataLoader exposes state_dict()/load_state_dict() +state = dataloader.state_dict() # in the main process +dataloader.load_state_dict(state) # resume exactly where you stopped +``` + +README feature sections (from `grep -n '' README.md`): multi-GPU/multi-node (483), multiple providers (515), pause/resume (573), combine (846), parallel streaming (918), cycle (968), subsets (1115), parquet (1192), compression (1251), on-demand access (1284), transforms (1302), cache limits (1386). diff --git a/.claude/skills/litdata/reference/testing.md b/.claude/skills/litdata/reference/testing.md new file mode 100644 index 000000000..c2aa83c86 --- /dev/null +++ b/.claude/skills/litdata/reference/testing.md @@ -0,0 +1,88 @@ +# Testing LitData + +Framework: **pytest** (`requirements/test.txt`). Config in `pyproject.toml [tool.pytest.ini_options]`. + +## Running tests + +```bash +# Single test (the fast dev loop) +pytest tests/streaming/test_dataset.py::test_name -v --capture=no + +# A whole file / dir +pytest tests/streaming/ -v + +# How CI runs it (authoritative). Two phases because processing tests must be serial: +pytest tests --ignore=tests/processing --ignore=tests/raw \ + -n 2 --dist=loadgroup --cov=litdata --durations=0 --timeout=120 --capture=no --verbose +pytest tests/processing tests/raw --cov=litdata --cov-append --durations=0 --timeout=120 --capture=no --verbose +``` + +**`--dist=loadgroup` is required, not optional, when running with `-n>1`.** `tests/conftest.py:165` (`pytest_collection_modifyitems`) pins every test using the `clean_pq_index_cache` fixture onto a single xdist worker via `xdist_group`, because that fixture `shutil.rmtree`s `~/.lightning/chunks` and would otherwise race across workers. + +The `Makefile test` target is stale (`--flake8` flag, points at `src`) — prefer the CI commands above. + +### Pytest config gotchas (`pyproject.toml`) + +- `addopts` includes `--doctest-modules` → **docstrings in `src/` are collected as tests**; a malformed docstring example fails CI. +- `--strict-markers` → only registered markers allowed. Sole custom marker: `cloud` (real integration tests, not in the standard matrix). +- `filterwarnings = ["error::FutureWarning"]` → **FutureWarnings are hard errors.** +- `xfail_strict = true` → an `xfail` that unexpectedly passes fails the suite. + +## Test layout + +Tests mirror `src/litdata/`: `src/litdata/streaming/dataset.py` → `tests/streaming/test_dataset.py`. Dirs: `tests/streaming/`, `tests/processing/`, `tests/utilities/`, `tests/raw/`, plus top-level `tests/test_cli.py`, `test_debugger.py`, `test_helper.py`, `test_imports.py`, `test_requirements.py`. Every dir has `__init__.py`; cross-test helpers import as a package (`from tests.streaming.utils import filter_lock_files`). + +Template (from `CONTRIBUTING.md`): + +```python +def test_explain_what_is_being_tested(tmpdir): + """One-line description of what/why.""" + cache_dir = os.path.join(tmpdir, "cache_dir") + assert ... +``` + +## Gating tests — there is NO `RunIf` helper + +`CONTRIBUTING.md` shows a `@RunIf(min_cuda_gpus=1)` template copied from PyTorch Lightning, but **no `RunIf` class exists in this repo.** Don't use it. The real pattern is `pytest.mark.skipif` + the `_*_AVAILABLE` `RequirementCache` constants in `src/litdata/constants.py`: + +```python +from litdata.constants import _ZSTD_AVAILABLE, _PIL_AVAILABLE + +# optional dependency, at param level +pytest.param("zstd", marks=pytest.mark.skipif(not _ZSTD_AVAILABLE, reason="Requires: ['zstd']")) + +# at function level +@pytest.mark.skipif(not _PIL_AVAILABLE or sys.platform == "win32", reason="Requires: ['pil']") + +# OS gating (very common) +@pytest.mark.skipif(sys.platform == "win32", reason="Not tested on windows") + +# python-version gating +@pytest.mark.skipif(sys.version_info >= (3, 12), reason="Multiprocessing issues on 3.12+") +``` + +Available flags: `_ZSTD_AVAILABLE`, `_BOTO3_AVAILABLE`, `_FSSPEC_AVAILABLE`, `_CRYPTOGRAPHY_AVAILABLE`, `_GOOGLE_STORAGE_AVAILABLE`, `_AZURE_STORAGE_AVAILABLE`, `_POLARS_AVAILABLE`, `_PYARROW_AVAILABLE`, `_PIL_AVAILABLE`, `_TORCH_VISION_AVAILABLE`, `_AV_AVAILABLE`, `_OBSTORE_AVAILABLE`, `_HF_HUB_AVAILABLE`, `_LIGHTNING_SDK_AVAILABLE`, `_VIZ_TRACKER_AVAILABLE`. + +## Mocking cloud — no moto / mock_aws + +Two techniques: + +1. **Fake modules injected into `sys.modules`** via `monkeypatch.setitem`, provided as `conftest.py` fixtures: `google_mock`, `fsspec_mock`, `obstore_mock`, `azure_mock`, `lightning_cloud_mock` (sets `rest_client.LightningClient = Mock()`), `lightning_sdk_mock`, `huggingface_hub_mock`, `huggingface_hub_fs_mock` (full fake `HfFileSystem` backed by local parquet), `fsspec_pq_mock`. These let import-guarded cloud code run without the real SDK. +2. **`monkeypatch.setattr` on the module's `boto3`/`botocore`** for S3 — e.g. `tests/streaming/test_client.py` patches `client.boto3`/`client.botocore` with `MagicMock()` and asserts on `.client.assert_called_with(...)`. `boto3` is a hard dep, so it's stubbed, not installed-away. + +Use stdlib `unittest.mock` (`Mock`, `MagicMock`, `patch`) + the pytest `monkeypatch` fixture. + +## Useful fixtures (`tests/conftest.py`) + +- **Autouse/session**: `teardown_process_group` (destroys torch.distributed group), `_thread_police` (fails a test leaking a non-daemon "zombie" thread — special-cases `PrepareChunksThread`, `QueueFeederThread`, `pytest_timeout`). +- **Autouse/per-test**: `disable_signals` (no-ops `signal.signal` so worker signal handlers don't break tests). +- **Data**: `mosaic_mds_index_data`, `prepare_combined_dataset` (two on-disk 50-item `Cache` datasets), `combined_dataset` (a `CombinedStreamingDataset`), `pq_data`, `write_pq_data` (5 polars parquet files), `clean_pq_index_cache` (wipes the default cache dir — the one that triggers xdist grouping). + +Local idiom: most files define their own `seed_everything(seed)`. CLI tests use a `run_cli(args_list)` helper (patches `sys.argv`, captures stdout). `tests/streaming/utils.py` gives `filter_lock_files` / `get_lock_files` to assert on chunk dirs ignoring `.lock`/`.cnt` artifacts. + +## When tests hang or flake + +- Processing/worker tests spawn real subprocesses — run **serially** (why CI separates `tests/processing`). Use `--timeout=120`. +- A leaked thread → `_thread_police` `AssertionError: Test left zombie thread`. Ensure background threads (`PrepareChunksThread`) are stopped. +- `spawn` requires picklable args; a test failing only under spawn usually means an unpicklable closure. +- To debug inside a worker: `from litdata.utilities.breakpoint import breakpoint; breakpoint()` (see debugging.md).