docs: add Claude Code skill for the LitData codebase

.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 `<details>` feature blocks are the narrative docs (line refs are cited in the reference files). Repo version: see `src/litdata/__about__.py`.

.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 <path> 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 `<chunk>.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.<random>` then renames. Leftover `<random>`-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 `<chunk>.tmb` tombstones recording which rank deleted a chunk and its `can_delete` verdict.
+4. Inspect the cache dir: leftover `.zstd.bin.<random>` + missing `.bin` for a low-numbered chunk = premature deletion + failed redownload.