litdata 0.2.65__tar.gz → 0.2.67__tar.gz

{litdata-0.2.65/src/litdata.egg-info → litdata-0.2.67}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: litdata
-Version: 0.2.65
+Version: 0.2.67
 Summary: The Deep Learning framework to train, deploy, and ship AI products Lightning fast.
 Home-page: https://github.com/Lightning-AI/litdata
 Download-URL: https://github.com/Lightning-AI/litdata
@@ -44,6 +44,7 @@ Requires-Dist: pillow; extra == "extras"
 Requires-Dist: polars; extra == "extras"
 Requires-Dist: pyarrow<25.0.0; extra == "extras"
 Requires-Dist: tqdm; extra == "extras"
+Requires-Dist: uvloop; sys_platform != "win32" and extra == "extras"
 Requires-Dist: viztracer; extra == "extras"
 Dynamic: author
 Dynamic: author-email
@@ -71,12 +72,13 @@ Dynamic: summary
 &nbsp;
 
 <pre>
-Transform                              Optimize
+Transform                              Optimize / Stream
   
-✅ Parallelize data processing       ✅ Stream large cloud datasets          
-✅ Create vector embeddings          ✅ Accelerate training by 20x           
-✅ Run distributed inference         ✅ Pause and resume data streaming      
-✅ Scrape websites at scale          ✅ Use remote data without local loading
+✅ Parallelize data processing       ✅ Stream raw files with no prep
+✅ Create vector embeddings          ✅ Stream large cloud datasets          
+✅ Run distributed inference         ✅ Accelerate training by 20x           
+✅ Scrape websites at scale          ✅ Pause and resume data streaming      
+                                     ✅ Use remote data without local loading
 </pre>
 
 ---
@@ -92,6 +94,8 @@ Transform                              Optimize
   <a href="#speed-up-model-training">Optimize data</a> •
   <a href="#transform-datasets">Transform data</a> •
   <a href="#key-features">Features</a> •
+  <a href="#stream-raw">Stream raw files</a> •
+  <a href="#resolve-paths">Paths & cloud URLs</a> •
   <a href="#benchmarks">Benchmarks</a> •
   <a href="#start-from-a-template">Templates</a> •
   <a href="#community">Community</a>
@@ -145,6 +149,21 @@ Install all the extras
 pip install 'litdata[extras]'
 ```
 
+On Linux/macOS, `[extras]` includes optional `uvloop` for a faster asyncio event loop used by `StreamingRawDataset` (stdlib asyncio is the fallback when it is not installed).
+
+</details>
+
+<details>
+  <summary>AI agent skill (Cursor, Claude Code, …)</summary>
+
+Install the LitData expert skill so coding agents know the full API, path resolver, optimize/stream recipes, and internals:
+
+```bash
+npx skills add Lightning-AI/litData
+```
+
+Source: [`.claude/skills/litdata/`](.claude/skills/litdata/) in this repository ([skills CLI](https://github.com/vercel-labs/skills)).
+
 </details>
 
 &nbsp;
@@ -154,29 +173,38 @@ pip install 'litdata[extras]'
 # Speed up model training
 Stream datasets directly from cloud storage without local downloads. Choose the approach that fits your workflow:
 
-## Option 1: Start immediately with existing data ⚡⚡
-Stream raw files directly from cloud storage - no pre-optimization needed.
+## Option 1: Stream existing files as-is ⚡⚡ — `StreamingRawDataset`
+
+**No optimize step.** Point LitData at a folder of images, audio, text, or any files (local or cloud) and train with a normal PyTorch `DataLoader`. Downloads are **fully asynchronous** and **batched**; cloud clients include **built-in retries**. You receive **raw `bytes`** — decode, parse, or transform however you want.
+
+Details → [Stream raw files](#stream-raw).
 
 ```python
 from litdata import StreamingRawDataset
 from torch.utils.data import DataLoader
+from PIL import Image
+import io
 
-# Point to your existing cloud data
-dataset = StreamingRawDataset("s3://my-bucket/raw-data/")
-dataloader = DataLoader(dataset, batch_size=32)
+dataset = StreamingRawDataset(
+    "s3://my-bucket/raw-images/",          # or gs://, azure://, /teamspace/s3_connections/..., local path
+    transform=lambda b: Image.open(io.BytesIO(b)).convert("RGB"),  # optional — default is raw bytes
+)
+loader = DataLoader(dataset, batch_size=32, num_workers=8)
 
-for batch in dataloader:
-    # Process raw bytes on-the-fly
-    pass
+for batch in loader:
+    train_step(batch)
 ```
 
 **Key benefits:**
 
-✅ **Instant access:**         Start streaming immediately without preprocessing.    
-✅ **Zero setup time:**        No data conversion or optimization required.    
-✅ **Native format:**          Work with original file formats (images, text, etc.).    
-✅ **Flexible processing:**    Apply transformations on-the-fly during streaming.    
-✅ **Cloud-native:**           Stream directly from S3, GCS, or Azure storage.    
+✅ **Zero preprocess:**     No chunking job — use the files you already have.    
+✅ **Raw bytes, your rules:** Each sample is file `bytes`; decode with PIL, torchaudio, json, or any custom logic (`transform=` optional).    
+✅ **Fully async + batched:** Concurrent downloads via `asyncio` / `__getitems__` (not one-file-at-a-time).    
+✅ **Built-in retries:**     Cloud downloads retry transient failures (adaptive client retries).    
+✅ **Cloud-native:**        S3 / GCS / Azure / Studio connections; same path resolver as optimized streaming.    
+✅ **Grouped samples:**     Override `setup()` to yield image+mask, audio+transcript, etc.    
+✅ **Indexed once:**        `index.json.zstd` cached locally and on the bucket for fast restarts.    
+✅ **Upgrade path:**        When I/O becomes the bottleneck, `optimize` → `StreamingDataset` for max throughput.    
 
 ## Option 2: Optimize for maximum performance ⚡⚡⚡  
 Accelerate model training (20x faster) by optimizing datasets for streaming directly from cloud storage. Work with remote data without local downloads with features like loading data subsets, accessing individual samples, and resumable streaming.
@@ -187,31 +215,33 @@ Transform raw data into optimized chunks for maximum streaming speed.
 This step formats the dataset for fast loading by writing data in an efficient chunked binary format.
 
 ```python
+import io
 import numpy as np
 from PIL import Image
 import litdata as ld
 
 def random_images(index):
-    # Replace with your actual image loading here (e.g., .jpg, .png, etc.)
-    # Recommended: use compressed formats like JPEG for better storage and optimized streaming speed
-    # You can also apply resizing or reduce image quality to further increase streaming speed and save space
-    fake_images = Image.fromarray(np.random.randint(0, 256, (32, 32, 3), dtype=np.uint8))
+    # Replace with your actual image loading (e.g. Image.open("photo.jpg")).
+    # Prefer JPEG: return a JpegImageFile, or re-encode at quality≈95. Plain
+    # Image.fromarray(...) stores uncompressed PIL RAW and can be 10×+ larger.
+    img = Image.fromarray(np.random.randint(0, 256, (32, 32, 3), dtype=np.uint8))
+    buf = io.BytesIO()
+    img.convert("RGB").save(buf, format="JPEG", quality=95)
+    buf.seek(0)
+    jpeg_image = Image.open(buf)  # JpegImageFile → compressed bytes in the chunk
     fake_labels = np.random.randint(10)
 
-    # You can use any key:value pairs. Note that their types must not change between samples, and Python lists must
-    # always contain the same number of elements with the same types
-    data = {"index": index, "image": fake_images, "class": fake_labels}
-
-    return data
+    # Keys/types must stay stable across samples; list lengths/types fixed
+    return {"index": index, "image": jpeg_image, "class": fake_labels}
 
 if __name__ == "__main__":
-    # The optimize function writes data in an optimized format
+    # Exactly one of chunk_bytes or chunk_size
     ld.optimize(
         fn=random_images,                   # the function applied to each input
         inputs=list(range(1000)),           # the inputs to the function (here it's a list of numbers)
         output_dir="fast_data",             # optimized data is stored here
         num_workers=4,                      # the number of workers on the same machine
-        chunk_bytes="64MB"                  # size of each chunk
+        chunk_bytes="64MB"                  # default; see FAQ for larger samples
     )
 ```
 
@@ -229,7 +259,12 @@ Load the data by replacing the PyTorch Dataset and DataLoader with the Streaming
 ```python
 import litdata as ld
 
-dataset = ld.StreamingDataset('s3://my-bucket/fast_data', shuffle=True, drop_last=True)
+dataset = ld.StreamingDataset(
+    's3://my-bucket/fast_data',
+    shuffle=True,
+    drop_last=True,  # important for multi-GPU so every rank sees the same length
+    seed=42,
+)
 
 # Custom collate function to handle the batch (optional)
 def collate_fn(batch):
@@ -239,7 +274,7 @@ def collate_fn(batch):
     }
 
 
-dataloader = ld.StreamingDataLoader(dataset, collate_fn=collate_fn)
+dataloader = ld.StreamingDataLoader(dataset, batch_size=64, collate_fn=collate_fn)
 for sample in dataloader:
     img, cls = sample["image"], sample["class"]
 ```
@@ -304,79 +339,158 @@ ld.map(
 ## Features for optimizing and streaming datasets for model training
 
 <details>
-  <summary> ✅ Stream raw datasets from cloud storage (beta) <a id="stream-raw" href="#stream-raw">🔗</a> </summary>
+  <summary> ✅ Stream raw files as-is (no optimize) — StreamingRawDataset <a id="stream-raw" href="#stream-raw">🔗</a> </summary>
   &nbsp;
 
-Effortlessly stream raw files (images, text, etc.) directly from S3, GCS, and Azure cloud storage without any optimization or conversion. Ideal for workflows requiring instant access to original data in its native format.
+`StreamingRawDataset` streams **your existing files** from local disk or cloud storage with **no conversion step**. It is a map-style `torch.utils.data.Dataset`: use a standard PyTorch `DataLoader` (not `StreamingDataLoader`).
 
-**Prerequisites:**
+**You get raw `bytes`.** LitData does not impose a sample schema — open images with PIL, parse JSONL, decode audio, run your own tokenizer, or pass a `transform=` if you prefer. Grouped items yield `list[bytes]` (e.g. image + mask).
 
-Install the required dependencies to stream raw datasets from cloud storage like **Amazon S3** or **Google Cloud Storage**:
+Downloads are **fully asynchronous** and **batched**: when the DataLoader requests a batch, `__getitems__` fetches those files concurrently with `asyncio.gather`. Cloud clients include **built-in retries** for transient network errors.
 
-```bash
-# for aws s3
-pip install "litdata[extra]" s3fs
+Use it when you want to train or prototype on JPEGs, masks, audio, JSONL, etc. **immediately**. Switch to [`optimize` → `StreamingDataset`](#speed-up-model-training) later if you need maximum cloud training throughput.
 
-# for gcloud storage
-pip install "litdata[extra]" gcsfs
+| | `StreamingRawDataset` | `StreamingDataset` (optimized) |
+|--|----------------------|--------------------------------|
+| Prep | None — point at a folder | One-time `optimize` → `chunk-*.bin` + `index.json` |
+| Item | **Raw file `bytes`** (you decide how to decode) | Deserialized samples (dict/tensor/…) |
+| I/O | Fully async, batched downloads + retries | Chunk prefetch / cache pipeline |
+| Loader | `torch.utils.data.DataLoader` | Prefer `StreamingDataLoader` (shuffle, resume) |
+| Best for | Instant start, full control over bytes | Highest sustained training I/O |
+
+### Install (cloud)
+
+```bash
+pip install "litdata[extra]" s3fs    # Amazon S3
+pip install "litdata[extra]" gcsfs  # Google Cloud Storage
+# Azure / Studio connections: see Paths & cloud URLs
 ```
 
-**Usage Example:**
+### Quick start
+
 ```python
 from torch.utils.data import DataLoader
 from litdata import StreamingRawDataset
+from PIL import Image
+import io
 
-dataset = StreamingRawDataset("s3://bucket/files/")
+def to_image(data: bytes):
+    return Image.open(io.BytesIO(data)).convert("RGB")
+
+dataset = StreamingRawDataset(
+    "s3://my-bucket/images/",   # also: gs://, azure://, /teamspace/s3_connections/..., local path
+    transform=to_image,         # optional; default yields raw bytes
+    storage_options={},         # optional cloud credentials / endpoint
+)
+loader = DataLoader(dataset, batch_size=32, num_workers=8)
 
-# Use with PyTorch DataLoader
-loader = DataLoader(dataset, batch_size=32)
 for batch in loader:
-    # Each item is raw bytes
-    pass
+    train_step(batch)
 ```
 
-> Use `StreamingRawDataset` to stream your data as-is. Use `StreamingDataset` for fastest streaming after optimizing your data.
+### Constructor knobs
 
+| Arg | Default | Purpose |
+|-----|---------|---------|
+| `input_dir` | required | Folder URL/path (same [resolver](#resolve-paths) as optimized streaming) |
+| `cache_dir` | LitData default cache | Where the file index (and optional file cache) live |
+| `cache_files` | `False` | If `True`, keep downloaded files on disk under `cache_dir` (mirror remote layout) |
+| `recompute_index` | `False` | Force re-scan when remote files changed |
+| `transform` | `None` | `fn(bytes) -> Any` or `fn(list[bytes]) -> Any` for grouped items |
+| `storage_options` | `{}` | Cloud client options |
+| `indexer` | `FileIndexer()` | Custom discovery (subclass `BaseIndexer`) |
+| `max_concurrent_downloads` | `None` (adaptive) | Per-worker in-flight downloads. `None` = size-aware budget (bandwidth; Little’s-law only for medians &lt;~8 MiB) split across workers; single-process capped at 128. An explicit `int` is used exactly (no silent clamp) |
+| `max_prefetch` | `16` | Per-worker sequential look-ahead after each batch (default on). When `num_workers > 1`, effective look-ahead is `min(max_prefetch, 64 // num_workers)` so aggregate stays ~64 items. Pass `0` to disable |
+| `prefetch_cache_size` | auto | LRU cap for prefetched items (defaults from `max_prefetch`) |
+| `hedge_delay` | `0` | Seconds before a hedged duplicate GET for a slow download (`0` = off, default; opt-in) |
+| `range_parallel_threshold` | `0` | Objects ≥ this many bytes use parallel ranged GETs (`0` = whole-object only; opt-in) |
+| `item_type` | `"bytes"` | `"bytes"` buffers in RAM; `"path"` returns local cache paths (`cache_files=True` required) |
 
-You can also customize how files are grouped by subclassing `StreamingRawDataset` and overriding the `setup` method. This is useful for pairing related files (e.g., image and mask, audio and transcript) or any custom grouping logic.
+### Group related files (`setup`)
+
+Default: **one file = one sample**. Override `setup` to filter or group (image + mask, audio + transcript, …). Return either a list of `FileMetadata` or a list of groups (`list[list[FileMetadata]]`).
 
 ```python
-from typing import Union
+from collections import defaultdict
 from torch.utils.data import DataLoader
 from litdata import StreamingRawDataset
 from litdata.raw.indexer import FileMetadata
 
 class SegmentationRawDataset(StreamingRawDataset):
-    def setup(self, files: list[FileMetadata]) -> Union[list[FileMetadata], list[list[FileMetadata]]]:
-        # TODO: Implement your custom grouping logic here.
-        # For example, group files by prefix, extension, or any rule you need.
-        # Return a list of groups, where each group is a list of FileMetadata.
-        # Example:
-        #   return [[image, mask], ...]
-        pass
-
-# Initialize the custom dataset
-dataset = SegmentationRawDataset("s3://bucket/files/")
-loader = DataLoader(dataset, batch_size=32)
-for item in loader:
-    # Each item in the batch is a pair: [image_bytes, mask_bytes]
-    pass
+    def setup(self, files: list[FileMetadata]) -> list[list[FileMetadata]]:
+        # Pair img_001.jpg with img_001.png (mask) by stem
+        by_stem: dict[str, dict[str, FileMetadata]] = defaultdict(dict)
+        for f in files:
+            name = f.path.rsplit("/", 1)[-1]
+            stem, _, ext = name.rpartition(".")
+            by_stem[stem][ext.lower()] = f
+        items = []
+        for stem, parts in sorted(by_stem.items()):
+            if "jpg" in parts and "png" in parts:
+                items.append([parts["jpg"], parts["png"]])
+        return items
+
+dataset = SegmentationRawDataset(
+    "s3://bucket/seg/",
+    transform=lambda pair: (pair[0], pair[1]),  # list[bytes]: [image, mask]
+)
+loader = DataLoader(dataset, batch_size=16, num_workers=4)
+for images, masks in loader:
+    ...
 ```
 
-**Smart Index Caching**
+### Index caching (`index.json.zstd`)
 
-`StreamingRawDataset` automatically caches the file index for instant startup. Initial scan, builds and caches the index, then subsequent runs load instantly.
+First open scans the tree and writes a compressed file list:
 
-**Two-Level Cache:**
-- **Local:** Stored in your cache directory for instant access
-- **Remote:** Automatically saved to cloud storage (e.g., `s3://bucket/files/index.json.zstd`) for reuse
+- **Local cache** under your LitData cache dir (fast restart on the same machine)
+- **Remote copy** next to the data when possible (e.g. `s3://bucket/files/index.json.zstd`) so every machine skips the scan
 
-**Force Rebuild:**
 ```python
-# When dataset files have changed
+# After adding/removing files on the bucket:
 dataset = StreamingRawDataset("s3://bucket/files/", recompute_index=True)
 ```
 
+Do **not** confuse this with optimized LitData’s `index.json` (chunk metadata). Raw indexing only lists files.
+
+### How downloads work
+
+1. DataLoader asks for a batch of indices → `__getitems__`.
+2. LitData **asynchronously** downloads those files **in parallel** (`asyncio.gather` + `adownload_fileobj`).
+3. Cloud SDKs apply **retries** on transient failures (e.g. S3 adaptive retries).
+4. Each item is returned as **`bytes`** (or `list[bytes]` if `setup` grouped files), then optional `transform`.
+
+Your training loop stays normal PyTorch — no async/`await` in user code.
+
+```python
+# Default: you own the bytes
+dataset = StreamingRawDataset("s3://bucket/files/")
+raw: bytes = dataset[0]
+# e.g. Image.open(io.BytesIO(raw)), json.loads(raw), np.frombuffer(raw), ...
+```
+
+### Tips
+
+- Prefer `num_workers > 0` so worker processes overlap async batch downloads with training. Scale workers toward host vCPUs for network-bound JPEG-sized objects — avoid saturating every vCPU.
+- On Linux, after any parent-process dataset I/O, use `DataLoader(..., multiprocessing_context="spawn", persistent_workers=True)` — default `fork` can hang S3 clients in workers.
+- Default `max_prefetch=16` enables sequential look-ahead **per DataLoader worker**; shuffled access disables it. Pass `0` to turn off. When `num_workers > 1`, look-ahead and download concurrency both scale down with worker count so aggregate in-flight work stays bounded.
+- Prefer an `s3://` / `gs://` URL or `/teamspace/s3_connections/...` so LitData hits the bucket directly ([resolver](#resolve-paths)) — avoid reading through FUSE.
+- Leave `range_parallel_threshold=0` (default) for typical JPEGs; raise it only for large objects where parallel ranged GETs help.
+- Best for medium/large files. Tiny objects (≲100 KB) are request-overhead bound — pack with [`optimize`](#speed-up-model-training) → `StreamingDataset` when I/O plateaus.
+
+### Throughput
+
+On ImageNet val raw over S3 (50 k JPEGs, batch size 64, spawn workers), throughput gains are clearest at **low worker counts / notebooks** (**+20–80%** at ≤8 workers). At **high workers** (≥16), results are roughly **parity within run-to-run noise**.
+
+| workers | before | after | Δ |
+|--------:|-------:|------:|--:|
+| 0 | 543 | 735 | **+35%** |
+| 2 | 816 | 1475 | **+81%** |
+| 8 | 4841 | 5718 | **+18%** |
+| 16+ | ~6k | ~6k | ~parity |
+
+Useful knobs: `num_workers`, `max_prefetch` (default 16; worker-aware), `download_timeout` (batch-level hang protection). Ranged parallel downloads stay opt-in (`range_parallel_threshold=0`).
+
 </details>
 
 <details>
@@ -412,10 +526,6 @@ storage_options = {
     "aws_secret_access_key": "your_secret_access_key",
 }
 
-dataset = StreamingDataset('s3://my-bucket/my-data', storage_options=storage_options)
-
-
-
 dataset = StreamingDataset('s3://my-bucket/my-data', storage_options=storage_options)
 ```
 
@@ -427,108 +537,198 @@ from litdata import StreamingDataset
 dataset = StreamingDataset('s3://my-bucket/my-data', cache_dir="/path/to/cache")
 ```
 
+Any local path, `s3://` / `gs://` / `r2://` / `azure://` / `hf://`, `local:` network drive, or Lightning `/teamspace/...` connection works — see [Resolve any path or cloud URL](#resolve-paths).
+
 </details>
 
 <details>
-  <summary> ✅ Stream Hugging Face 🤗 datasets <a id="stream-hf" href="#stream-hf">🔗</a> </summary>
-
+  <summary> ✅ Optimize images as JPEG (not raw PIL) <a id="optimize-jpeg" href="#optimize-jpeg">🔗</a> </summary>
 &nbsp;
 
-To use your favorite  Hugging Face dataset with LitData, simply pass its URL to `StreamingDataset`.
+How you return images from `optimize` controls storage size and streaming speed.
 
-<details>
-  <summary>How to get HF dataset URI?</summary>
+| What you return | Serializer | Result |
+|-----------------|------------|--------|
+| `PIL.JpegImageFile` (e.g. `Image.open("x.jpg")`) | JPEG | Compressed bytes — **preferred** |
+| Plain `PIL.Image` / `Image.fromarray(...)` | PIL RAW | Uncompressed pixels — often **10×+ larger** |
 
-https://github.com/user-attachments/assets/3ba9e2ef-bf6b-41fc-a578-e4b4113a0e72
+**Best practice:** store JPEG at **quality ≈ 95** (or keep existing `.jpg` files). Resize when helpful.
+
+```python
+import io
+from PIL import Image
+import litdata as ld
+
+def load_image(path):
+    img = Image.open(path)
+    if not str(path).lower().endswith((".jpg", ".jpeg")):
+        buf = io.BytesIO()
+        img.convert("RGB").save(buf, format="JPEG", quality=95)
+        buf.seek(0)
+        img = Image.open(buf)  # JpegImageFile
+    return {"image": img, "path": path}
+
+if __name__ == "__main__":
+    ld.optimize(fn=load_image, inputs=list_of_paths, output_dir="fast_data", chunk_bytes="64MB", num_workers=8)
+```
+
+Ready-made ImageNet optimize/stream scripts: `benchmarks/litdata/` (`--write_mode jpeg --quality 90`).
 
 </details>
 
-**Prerequisites:**
+<details>
+  <summary> ✅ Custom serializers <a id="serializers" href="#serializers">🔗</a> </summary>
+&nbsp;
 
-Install the required dependencies to stream Hugging Face datasets:
-```sh
-pip install "litdata[extra]" huggingface_hub
+LitData serializes each leaf of your sample with a pluggable registry. Built-ins (tried in order) include: `str`, `bool`, `int`, `float`, `video`, `tifffile`, `pil`, `jpeg`, `jpeg_array`, `bytes`, `numpy` / `tensor` (and no-header variants), and `pickle` (fallback).
 
-# Optional: To speed up downloads on high-bandwidth networks
-pip install hf_transfer
-export HF_HUB_ENABLE_HF_TRANSFER=1
+For images, returning a `JpegImageFile` selects **`jpeg`**; a plain `PIL.Image` selects **`pil`** (raw pixels). See [Optimize images as JPEG](#optimize-jpeg).
+
+Pass custom serializers when **streaming** (and when using the lower-level `Cache` writer):
+
+```python
+from litdata import StreamingDataset
+from litdata.streaming.serializers import Serializer
+
+class MyTypeSerializer(Serializer):
+    def serialize(self, item):
+        return item.to_bytes(), None  # (bytes, optional metadata string)
+
+    def deserialize(self, data: bytes):
+        return MyType.from_bytes(data)
+
+    def can_serialize(self, item) -> bool:
+        return isinstance(item, MyType)
+
+dataset = StreamingDataset(
+    "s3://bucket/data",
+    serializers={"my_type": MyTypeSerializer()},  # merged on top of built-ins
+)
 ```
 
-**Stream Hugging Face dataset:**
+Keys you pass are tried before the defaults (so they win over `pickle`). `optimize()` uses the built-in registry based on the Python types your `fn` returns — prefer JPEG / numpy / tensor leaves for best results.
+
+</details>
+
+<details>
+  <summary> ✅ Stream MosaicML MDS datasets <a id="stream-mds" href="#stream-mds">🔗</a> </summary>
+&nbsp;
+
+If you already have datasets written in [MosaicML Streaming](https://github.com/mosaicml/streaming) MDS (Mosaic Data Shard) format, you can stream them directly with LitData—no re-optimization or conversion required!
+
+LitData's default `PyTreeLoader` natively understands the MDS binary layout, so you can read existing MDS shards using the familiar `StreamingDataset` and `StreamingDataLoader` APIs.
+
+**Assumption:**
+
+Your dataset directory contains MDS shard files (e.g. `shard.00000.mds`, ...) along with an `index.json` describing the shards and their `column_sizes`/`column_names`.
+
+**Stream the MDS dataset:**
 
 ```python
 import litdata as ld
 
-# Define the Hugging Face dataset URI
-hf_dataset_uri = "hf://datasets/leonardPKU/clevr_cogen_a_train/data"
+# point to your MDS dataset stored locally or in the cloud
+
+mds_dataset_uri = "s3://my-bucket/my-mds-data" # or a local path
 
-# Create a streaming dataset
-dataset = ld.StreamingDataset(hf_dataset_uri)
+# LitData automatically detects and deserializes the MDS format
+
+dataset = ld.StreamingDataset(mds_dataset_uri)
 
-# Print the first sample
 print("Sample", dataset[0])
 
-# Stream the dataset using StreamingDataLoader
 dataloader = ld.StreamingDataLoader(dataset, batch_size=4)
 for sample in dataloader:
-    pass 
+  pass
 ```
 
-You don’t need to worry about indexing the dataset or any other setup. **LitData** will **handle all the necessary steps automatically** and `cache` the `index.json` file, so you won't have to index it again.
+**How it works:**
+
+- LitData reads the `format` field from the dataset config. When it's set to `"mds"`, the item loader uses MDS-aware deserialization (`mds_deserialize`) that respects the per-column sizes stored in each shard.
+- Fixed-size columns are read directly, while variable-size columns are prefixed with a `uint32` length header—exactly as in the MosaicML MDS spec.
+- Each sample is reconstructed into its original Python structure via LitData's `data_spec`.
+
+**Key benefits:**
+
+✅ **Zero conversion:**       Reuse existing MDS shards as-is.    
+✅ **Drop-in APIs:**          Use the same `StreamingDataset` / `StreamingDataLoader` you already know.    
+✅ **Cloud-native:**          Stream MDS shards directly from S3, GCS, or Azure.    
+✅ **Easy migration:**        Move from MosaicML Streaming to LitData without re-optimizing.    
 
-This ensures that the next time you stream the dataset, the indexing step is skipped..
+> **Note:** Encrypted data loading is not currently supported for the MDS format.
+
+</details>
+
+<details>
+  <summary> ✅ Stream Hugging Face 🤗 datasets <a id="stream-hf" href="#stream-hf">🔗</a> </summary>
 
 &nbsp;
 
-### Indexing the HF dataset (Optional)
+To use your favorite  Hugging Face dataset with LitData, simply pass its URL to `StreamingDataset`.
 
-If the Hugging Face dataset hasn't been indexed yet, you can index it first using the `index_hf_dataset` method, and then stream it using the code above.
+<details>
+  <summary>How to get HF dataset URI?</summary>
 
-```python
-import litdata as ld
+https://github.com/user-attachments/assets/3ba9e2ef-bf6b-41fc-a578-e4b4113a0e72
 
-hf_dataset_uri = "hf://datasets/leonardPKU/clevr_cogen_a_train/data"
+</details>
+
+**Prerequisites:**
 
-ld.index_hf_dataset(hf_dataset_uri)
+```sh
+pip install 'litdata[extras]' huggingface_hub
+
+# Optional: faster downloads on high-bandwidth networks
+pip install hf_transfer
+export HF_HUB_ENABLE_HF_TRANSFER=1
 ```
 
-- Indexing the Hugging Face dataset ahead of time will make streaming abit faster, as it avoids the need for real-time indexing during streaming.
+**Supported for HF:** datasets stored as **Parquet** only. Gated datasets: set `HF_TOKEN`.
 
-- To use `HF gated dataset`, ensure the `HF_TOKEN` environment variable is set.
+**Stream Hugging Face dataset** (auto-index + auto `ParquetLoader`):
 
-**Note**: For HuggingFace datasets, `indexing` & `streaming` is supported only for datasets in **`Parquet format`**.
+```python
+import litdata as ld
 
-&nbsp;
+hf_dataset_uri = "hf://datasets/leonardPKU/clevr_cogen_a_train/data"
 
-### Full Workflow for Hugging Face Datasets
+dataset = ld.StreamingDataset(hf_dataset_uri)  # indexes on first use; caches index.json locally
+print("Sample", dataset[0])  # dict of columns
 
-For full control over the cache path(`where index.json file will be stored`) and other configurations, follow these steps:
+# With workers on Linux, use spawn (same as other ParquetLoader usage)
+dataloader = ld.StreamingDataLoader(
+    dataset, batch_size=4, num_workers=4, multiprocessing_context="spawn"
+)
+for sample in dataloader:
+    pass
+```
 
-1. Index the Hugging Face dataset first:
+Unlike local/S3 parquet ([stream parquet](#stream-parquet)), `hf://` **automatically** indexes (if needed) and selects `ParquetLoader`.
+
+### Indexing the HF dataset (optional, faster cold start)
 
 ```python
 import litdata as ld
 
-hf_dataset_uri = "hf://datasets/open-thoughts/OpenThoughts-114k/data"
-
-ld.index_parquet_dataset(hf_dataset_uri, "hf-index-dir")
+# Returns the local cache directory that contains index.json
+cache_dir = ld.index_hf_dataset("hf://datasets/leonardPKU/clevr_cogen_a_train/data")
 ```
 
-2. To stream HF datasets now, pass the `HF dataset URI`, the path where the `index.json` file is stored, and `ParquetLoader` as the `item_loader` to the **`StreamingDataset`**:
+Or control the index path explicitly:
 
 ```python
 import litdata as ld
 from litdata.streaming.item_loader import ParquetLoader
 
-hf_dataset_uri = "hf://datasets/open-thoughts/OpenThoughts-114k/data"
-
-dataset = ld.StreamingDataset(hf_dataset_uri, item_loader=ParquetLoader(), index_path="hf-index-dir")
+uri = "hf://datasets/open-thoughts/OpenThoughts-114k/data"
+ld.index_parquet_dataset(uri, "hf-index-dir")  # writes index under hf-index-dir
 
-for batch in ld.StreamingDataLoader(dataset, batch_size=4):
-  pass
+dataset = ld.StreamingDataset(uri, item_loader=ParquetLoader(), index_path="hf-index-dir")
+for batch in ld.StreamingDataLoader(dataset, batch_size=4, multiprocessing_context="spawn"):
+    pass
 ```
 
-&nbsp;
+See also [Stream parquet datasets](#stream-parquet) for `ParquetLoader` knobs, wildcards, and stream-vs-optimize.
 
 ### LitData `Optimize` v/s `Parquet`
 <!-- TODO: Update benchmark -->
@@ -574,6 +774,110 @@ for batch in val_dataloader:
 
 </details>
 
+<details>
+  <summary> ✅ Shuffle, seed, and drop_last <a id="shuffle" href="#shuffle">🔗</a> </summary>
+&nbsp;
+
+Shuffling is **deterministic** and designed for distributed training:
+
+1. Chunks are assigned (and possibly split) across ranks/workers.
+2. Items inside each chunk are permuted.
+
+The permutation depends on `seed`, the epoch, and chunk metadata — the same settings always yield the same order (required for resumable `state_dict`).
+
+```python
+from litdata import StreamingDataset, StreamingDataLoader
+
+train = StreamingDataset(
+    "s3://my-bucket/train",
+    shuffle=True,
+    drop_last=True,  # keep every rank/worker at the same length (default True under DDP)
+    seed=42,         # default is 42; keep stable when resuming
+)
+loader = StreamingDataLoader(train, batch_size=64, num_workers=8)
+
+# shuffle=/drop_last= on the loader override the dataset
+loader = StreamingDataLoader(train, batch_size=64, shuffle=True, drop_last=True)
+```
+
+**Notes**
+
+- Val/test: usually `shuffle=False`, `drop_last=False`.
+- If `drop_last=False` under multi-GPU, LitData warns — collectives can hang when ranks see different lengths.
+- Resume with `loader.state_dict()` / `load_state_dict()`. To deliberately ignore checkpointed shuffle settings, set `force_override_state_dict=True` on the dataset.
+
+</details>
+
+<details>
+  <summary> ✅ FAQ: chunk size &amp; shuffle before optimize <a id="faq-chunk-shuffle" href="#faq-chunk-shuffle">🔗</a> </summary>
+&nbsp;
+
+### What `chunk_bytes` should I use?
+
+Default is **64MB** — a good starting point for typical small/medium samples.
+
+When each datapoint is large (e.g. a few MB), prefer a **larger chunk** (practical range often **256–512MB**) so each chunk holds more samples and **intra-chunk batch randomization** has a bigger pool. Tradeoff: larger chunks take **longer to download** before they can be used.
+
+This is expert guidance (recommended-range mindset), not a published chunk-size sweep.
+
+### Is StreamingDataset shuffle enough if my source data is ordered?
+
+**Not always.** LitData handles **distributed sampling** and **bucket sampling within chunks** automatically (`shuffle=True` randomizes chunk order and item order inside each chunk). That is **not** a substitute for a fully shuffled file-level DataLoader when the source has strong structure (same subject/set contiguous, class blocks, etc.).
+
+If ordered data would make chunked sampling problematic and you cannot embed the grouping as the sample unit:
+
+- Shuffle the list of samples **before** `optimize` so chunks mix well, **or**
+- Use [`StreamingRawDataset`](#stream-raw) (per-file random access via a standard PyTorch `DataLoader` with `shuffle=True`) instead of optimize → `StreamingDataset`.
+
+### FUSE vs LitData (Lightning Studios)
+
+`/teamspace/s3_connections` (and related mounts) are **FUSE** — fine for browsing, not for training I/O. Under load they are very slow and can crash. Pass the same path into LitData (`StreamingRawDataset` / `StreamingDataset` / `optimize`): LitData resolves it and talks **directly** to the bucket ([Resolve any path](#resolve-paths)).
+
+Rough ImageNet order-of-magnitude on a Studio (not hard guarantees; right tuning for raw): FUSE hand-read ~**600** images/s · [`StreamingRawDataset`](#stream-raw) ~**6–7k** · optimized [`StreamingDataset`](#speed-up-model-training) (64MB chunks) ~**11k**.
+
+</details>
+
+<details>
+  <summary> ✅ StreamingDataset & StreamingDataLoader knobs <a id="streaming-kwargs" href="#streaming-kwargs">🔗</a> </summary>
+&nbsp;
+
+**`StreamingDataset`**
+
+| Argument | Default | Description |
+|----------|---------|-------------|
+| `input_dir` | required | Local path, cloud URI, `Dir`, or parquet path (basename wildcards OK) |
+| `cache_dir` | `LITDATA_CACHE_DIR` or `~/.lightning/chunks` | Where chunks are cached |
+| `item_loader` | from index | `TokensLoader`, `ParquetLoader`, … |
+| `shuffle` | `False` | Deterministic shuffle (see [Shuffle](#shuffle)) |
+| `drop_last` | `True` if distributed else `False` | Equal length across ranks |
+| `seed` | `42` | Shuffle / subsample RNG |
+| `serializers` | built-ins | Custom serialize/deserialize map |
+| `max_cache_size` | `"100GB"` | Evict consumed chunks beyond this size |
+| `max_pre_download` | `2` | Chunks each worker may prefetch (raise for throughput; watch disk) |
+| `subsample` | `1.0` | Fraction of data (`0.01`) or upsample (`2.5`) |
+| `encryption` | `None` | `FernetEncryption` / `RSAEncryption` / custom |
+| `storage_options` | `{}` | Cloud client options |
+| `session_options` | `{}` | boto3 session options (S3) |
+| `index_path` | `None` | Parquet/HF `index.json` file or directory |
+| `force_override_state_dict` | `False` | Local ctor args override loaded checkpoint |
+| `transform` | `None` | Callable or list of callables per sample |
+
+Peak disk ≈ `num_workers × max_pre_download × mean_chunk_size`.
+
+**`StreamingDataLoader`**
+
+| Argument | Description |
+|----------|-------------|
+| All usual `torch.utils.data.DataLoader` kwargs | `batch_size`, `num_workers`, `collate_fn`, `pin_memory`, … |
+| `shuffle` / `drop_last` | Forwarded to the streaming dataset |
+| `profile_batches` | `int` / `True` / `False` — viztracer worker trace (see [Profile data loading](#profile-loading)) |
+| `profile_skip_batches` / `profile_dir` | Warm-up skip count; output dir for `result.json` |
+| `multiprocessing_context` | Use **`"spawn"`** (or `"forkserver"`) with `ParquetLoader` + `num_workers>0` on Linux |
+
+Prefer `StreamingDataLoader` over a plain PyTorch `DataLoader` for optimized / combined / parallel datasets (resume + correct batch metadata).
+
+</details>
+
 <details>
   <summary> ✅ Stream from multiple cloud providers <a id="cloud-providers" href="#cloud-providers">🔗</a> </summary>
 
@@ -911,7 +1215,7 @@ if __name__ == "__main__":
 
 Mix and match different sets of data to experiment and create better models.
 
-Combine datasets with `CombinedStreamingDataset`.  As an example, this mixture of [Slimpajama](https://huggingface.co/datasets/cerebras/SlimPajama-627B) & [StarCoder](https://huggingface.co/datasets/bigcode/starcoderdata) was used in the [TinyLLAMA](https://github.com/jzhang38/TinyLlama) project to pretrain a 1.1B Llama model on 3 trillion tokens.
+Combine datasets with `CombinedStreamingDataset`.  As an example, this mixture of [Slimpajama](https://www.cerebras.ai/blog/slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama) & [StarCoder](https://huggingface.co/datasets/bigcode/starcoderdata) was used in the [TinyLLAMA](https://github.com/jzhang38/TinyLlama) project to pretrain a 1.1B Llama model on 3 trillion tokens.
 
 ```python
 from litdata import StreamingDataset, CombinedStreamingDataset, StreamingDataLoader, TokensLoader
@@ -935,7 +1239,12 @@ train_datasets = [
 
 # Mix SlimPajama data and Starcoder data with these proportions:
 weights = (0.693584, 0.306416)
-combined_dataset = CombinedStreamingDataset(datasets=train_datasets, seed=42, weights=weights, iterate_over_all=False)
+combined_dataset = CombinedStreamingDataset(
+    datasets=train_datasets,
+    seed=42,
+    weights=weights,
+    iterate_over_all=False,  # required when passing weights (see below)
+)
 
 train_dataloader = StreamingDataLoader(combined_dataset, batch_size=8, pin_memory=True, num_workers=os.cpu_count())
 
@@ -944,37 +1253,34 @@ for batch in tqdm(train_dataloader):
     pass
 ```
 
-**Batching Methods**
+**`iterate_over_all` vs `weights` (important)**
 
-The `CombinedStreamingDataset` supports two different batching methods through the `batching_method` parameter:
+| Mode | Behavior |
+|------|----------|
+| `iterate_over_all=True` (default) | Iterate until **all** datasets are exhausted. Do **not** pass `weights` — LitData derives them from dataset lengths (raises `ValueError` if you pass both). |
+| `iterate_over_all=False` | Stop when **any** dataset is exhausted. Pass explicit `weights` for your mixture (e.g. TinyLlama). Length may be `None` (variable). |
 
-**Stratified Batching (Default)**:
-With `batching_method="stratified"` (the default), each batch contains samples from multiple datasets according to the specified weights:
+**Batching Methods** (`batching_method`)
+
+**Stratified** (default): each batch mixes samples from multiple datasets according to the weights.
 
 ```python
-# Default stratified batching - batches mix samples from all datasets
 combined_dataset = CombinedStreamingDataset(
-    datasets=[dataset1, dataset2], 
-    batching_method="stratified"  # This is the default
+    datasets=[dataset1, dataset2],
+    batching_method="stratified",  # default
 )
 ```
 
-**Per-Stream Batching**:
-With `batching_method="per_stream"`, each batch contains samples exclusively from a single dataset. This is useful when datasets have different shapes or structures:
+**Per-stream**: each batch comes from only one randomly selected dataset (useful when shapes/dtypes differ).
 
 ```python
-# Per-stream batching - each batch contains samples from only one dataset
 combined_dataset = CombinedStreamingDataset(
-    datasets=[dataset1, dataset2], 
-    batching_method="per_stream"
+    datasets=[dataset1, dataset2],
+    batching_method="per_stream",
 )
-
-# This ensures each batch has consistent structure, helpful for datasets with varying:
-# - Image sizes
-# - Sequence lengths  
-# - Data types
-# - Feature dimensions
 ```
+
+Other knobs: `seed` (default `42`), `force_override_state_dict=True` to let local ctor args override a loaded checkpoint.
 </details>
 
 <details>
@@ -1255,59 +1561,113 @@ The `overwrite` mode will delete the existing data and start from fresh.
   <summary> ✅ Stream parquet datasets <a id="stream-parquet" href="#stream-parquet">🔗</a> </summary>
 &nbsp;
 
-Stream Parquet datasets directly with LitData—no need to convert them into LitData’s optimized binary format! If your dataset is already in Parquet format, you can efficiently index and stream it using `StreamingDataset` and `StreamingDataLoader`.
+Stream existing Parquet files with LitData **without** converting them to LitData chunks — or convert them when you need LitData’s optimized binary format. Hugging Face parquet datasets are covered in [Stream Hugging Face datasets](#stream-hf).
 
-**Assumption:**
+### Stream vs optimize vs map
 
-Your dataset directory contains one or more Parquet files.
+| Goal | Use |
+|------|-----|
+| Train on parquet as-is (no conversion) | `index_parquet_dataset` → `StreamingDataset` + `ParquetLoader` |
+| Faster I/O / tokenize / custom sample shape | `optimize(fn)` that `yield`s rows from parquet ([reduce memory](#reduce-memory)) |
+| Reshard huge parquet files while mapping | `map(..., reader=ParquetReader(cache_folder, num_rows=...))` |
 
-**Prerequisites:**
+Each sample from `ParquetLoader` is a **`dict`** (column name → value).
 
-Install the required dependencies to stream Parquet datasets from cloud storage like **Amazon S3** or **Google Cloud Storage**:
+### Prerequisites
 
 ```bash
-# For Amazon S3
-pip install "litdata[extra]" s3fs
-
-# For Google Cloud Storage
-pip install "litdata[extra]" gcsfs
+pip install 'litdata[extras]'   # includes polars + pyarrow
+# Cloud listing/index extras as needed:
+pip install s3fs    # s3://
+pip install gcsfs   # gs://
 ```
 
-**Index Your Dataset**: 
-
-Index your Parquet dataset to create an index file that LitData can use to stream the dataset.
+### Index a parquet directory
 
 ```python
 import litdata as ld
 
-# Point to your data stored in the cloud
-pq_dataset_uri = "s3://my-bucket/my-parquet-data"  # or "gs://my-bucket/my-parquet-data"
-
-ld.index_parquet_dataset(pq_dataset_uri)
+ld.index_parquet_dataset(
+    "s3://my-bucket/my-parquet-data",  # local path, s3://, gs://, or hf://
+    cache_dir=None,                   # see table below
+    storage_options={},               # cloud credentials / endpoints
+    num_workers=4,                    # parallel metadata reads
+)
 ```
 
-**Stream the Dataset**
+| Scheme | Where `index.json` is written |
+|--------|-------------------------------|
+| Local directory | Next to the files, or under `cache_dir` if set |
+| `s3://` / `gs://` | **Uploaded to the bucket** at `{url}/index.json` (needs write access) |
+| `hf://` | **Local** `cache_dir` (required for HF indexing via this helper) |
+
+**Indexing notes**
+
+- Lists **top-level** `.parquet` files only (not recursive subfolders).
+- All files must share the same schema.
+- Supported for indexing today: local, `s3://`, `gs://`, `hf://` (not `r2://` / `azure://` yet).
+- For HF, prefer `index_hf_dataset(uri)` (returns a local cache dir) or auto-index via `StreamingDataset("hf://...")` — see [HF section](#stream-hf).
 
-Use `StreamingDataset` with `ParquetLoader` to load and stream the dataset efficiently:
+### Stream with `ParquetLoader`
 
+Unlike `hf://`, local/S3/GCS parquet **does not** auto-select the loader — pass `ParquetLoader` explicitly (it must match `index.json`).
 
 ```python
 import litdata as ld
 from litdata.streaming.item_loader import ParquetLoader
 
-# Specify your dataset location in the cloud
-pq_dataset_uri = "s3://my-bucket/my-parquet-data"  # or "gs://my-bucket/my-parquet-data"
+uri = "s3://my-bucket/my-parquet-data"
+dataset = ld.StreamingDataset(
+    uri,
+    item_loader=ParquetLoader(low_memory=True),  # default: row-group streaming
+    # index_path="/path/to/index.json",         # optional if index lives elsewhere
+)
 
-# Set up the streaming dataset
-dataset = ld.StreamingDataset(pq_dataset_uri, item_loader=ParquetLoader())
+# Basename wildcards when the path ends with .parquet:
+# dataset = ld.StreamingDataset("s3://bucket/data/train-*.parquet", item_loader=ParquetLoader())
 
-print("Sample", dataset[0])
+print(dataset[0])  # dict of columns
 
-dataloader = ld.StreamingDataLoader(dataset, batch_size=4)
-for sample in dataloader:
+# Linux + num_workers>0: use spawn (Polars + fork deadlocks)
+dataloader = ld.StreamingDataLoader(
+    dataset,
+    batch_size=4,
+    num_workers=4,
+    multiprocessing_context="spawn",
+)
+for batch in dataloader:
     pass
 ```
 
+### `ParquetLoader` knobs
+
+| Arg | Default | Meaning |
+|-----|---------|---------|
+| `low_memory` | `True` | Stream by row group (lower RAM). `False` loads each whole file into memory (warns). |
+| `pre_load_chunk` | `False` | Prefetch full DataFrame — **only effective when `low_memory=False`**. |
+
+Import: `from litdata.streaming.item_loader import ParquetLoader` (not re-exported at `litdata` top level).
+
+### Reshard parquet for `map` / `optimize`
+
+```python
+from litdata import map
+from litdata.processing.readers import ParquetReader
+
+def process(pq_file, output_dir):
+    # pq_file is a pyarrow.parquet.ParquetFile
+    ...
+
+map(
+    fn=process,
+    inputs=list_of_parquet_paths,
+    output_dir="s3://bucket/out",
+    reader=ParquetReader(cache_folder="/tmp/pq-shards", num_rows=65536),
+)
+```
+
+`ParquetReader` splits inputs that exceed `num_rows` into smaller cached files before your `fn` runs.
+
 </details>
 
 <details>
@@ -1393,17 +1753,60 @@ for batch in dataloader:
   <summary> ✅ Profile data loading speed <a id="profile-loading" href="#profile-loading">🔗</a> </summary>
 &nbsp;
 
-Measure and optimize how fast your data is being loaded, improving efficiency.
+`StreamingDataLoader` can record a **viztracer** Chrome trace of the DataLoader worker loop so you can see where time goes (fetch, deserialize, collate, IPC).
+
+### Prerequisites
+
+```bash
+pip install viztracer
+```
+
+Profiling requires **`num_workers >= 1`** (raises otherwise). On multi-GPU, only **global rank 0** installs the worker profiler.
 
-The `StreamingDataLoader` supports profiling of your data loading process. Simply use the `profile_batches` argument to specify the number of batches you want to profile:
+### Usage
 
 ```python
 from litdata import StreamingDataset, StreamingDataLoader
 
-StreamingDataLoader(..., profile_batches=5)
+dataset = StreamingDataset("s3://my-bucket/my-data", shuffle=True, drop_last=True)
+
+loader = StreamingDataLoader(
+    dataset,
+    batch_size=64,
+    num_workers=4,
+    profile_batches=20,          # record this many batches (int), or True for the whole run
+    profile_skip_batches=5,      # warm up / skip cold-start batches before recording
+    profile_dir="./profiles",    # where to write result.json (default: cwd)
+)
+
+for batch in loader:
+    train_step(batch)
+    # after profile_batches (+ skip) complete, worker 0 saves the trace and prints the path
+```
+
+| Arg | Default | Meaning |
+|-----|---------|---------|
+| `profile_batches` | `False` | `int` → stop after that many **recorded** batches; `True` → profile until the iterator ends; `False` → off |
+| `profile_skip_batches` | `0` | Batches to skip before the tracer starts (useful to skip cache cold-start) |
+| `profile_dir` | current working directory | Directory for `result.json` (overwrites an existing file) |
+
+Only **worker 0** is instrumented. When an `int` is used, the tracer wraps `fetcher.fetch` and stops after `profile_skip_batches + profile_batches` fetch calls. When `True`, tracing runs for the lifetime of that worker loop.
+
+### View the trace
+
+```bash
+# Option A — Chrome
+# open chrome://tracing and load profiles/result.json
+
+# Option B — Perfetto (often better for large traces)
+# open https://ui.perfetto.dev and load the same file
 ```
 
-This generates a Chrome trace called `result.json`. Then, visualize this trace by opening Chrome browser at the `chrome://tracing` URL and load the trace inside.
+### Tips
+
+- Delete or change `profile_dir` between runs — LitData removes an existing `result.json` before starting.
+- Pair with a wiped chunk cache if you care about **cold** epoch behavior (`litdata cache clear`).
+- For deeper LitData internals (download / lock / delete timeline), use `enable_tracer()` + [Litracer](https://github.com/deependujha/litracer) instead — see [Debug & Profile LitData](#debug-profile). That path is complementary: viztracer = DataLoader worker CPU timeline; Litracer = LitData pipeline events.
 
 </details>
 
@@ -1413,7 +1816,9 @@ This generates a Chrome trace called `result.json`. Then, visualize this trace b
 
 Handle large data files efficiently without using too much of your computer's memory.
 
-When processing large files like compressed [parquet files](https://en.wikipedia.org/wiki/Apache_Parquet), use the Python yield keyword to process and store one item at the time, reducing the memory footprint of the entire program.
+**Optimize from parquet** (convert into LitData chunks) when you need tokenization or LitData’s binary format. To **stream parquet without converting**, see [Stream parquet datasets](#stream-parquet).
+
+When processing large parquet files, `yield` one item at a time to keep memory low:
 
 ```python
 from pathlib import Path
@@ -1449,32 +1854,104 @@ outputs = optimize(
   <summary> ✅ Limit local cache space <a id="limit-cache" href="#limit-cache">🔗</a> </summary>
 &nbsp;
 
-Limit the amount of disk space used by temporary files, preventing storage issues.
+Control how much disk the local chunk cache may use. Downloaded chunks are deleted after use once the cache exceeds the limit.
+
+Default `max_cache_size` is **`100GB`**. Peak disk in flight is roughly:
+
+```
+num_workers × max_pre_download × mean_chunk_size
+```
 
-Adapt the local caching limit of the `StreamingDataset`. This is useful to make sure the downloaded data chunks are deleted when used and the disk usage stays low.
+Keep `max_cache_size` comfortably above that peak. For remote datasets, async chunk prefetch often raises `max_pre_download` to **≥4** automatically — see [async prefetch & environment variables](#async-prefetch-env).
 
 ```python
 from litdata import StreamingDataset
 
-dataset = StreamingDataset(..., max_cache_size="10GB")
+dataset = StreamingDataset(
+    "s3://my-bucket/my-data",
+    max_cache_size="10GB",
+    max_pre_download=4,  # chunks each worker may prefetch (default 2; async may floor to 4)
+)
 ```
 
 </details>
 
+<details>
+  <summary> ✅ Async chunk prefetch & environment variables <a id="async-prefetch-env" href="#async-prefetch-env">🔗</a> </summary>
+&nbsp;
+
+### Async chunk prefetch
+
+LitData can overlap **remote chunk downloads** with training using `asyncio` inside each DataLoader worker’s prepare thread. This is **not** an async DataLoader — your loop stays:
+
+```python
+for batch in StreamingDataLoader(dataset, batch_size=64, num_workers=8):
+    train_step(batch)
+```
+
+| Situation | Async prefetch |
+|-----------|----------------|
+| Remote dataset (`s3://`, `gs://`, …) | **On** by default |
+| Local-only dataset | **Off** by default |
+| `LITDATA_ASYNC_CHUNK_PREFETCH=1` | Force on |
+| `LITDATA_ASYNC_CHUNK_PREFETCH=0` | Force off |
+
+When async is on, LitData raises `max_pre_download` to at least **4** so `asyncio.gather` has enough in-flight downloads (override with `LITDATA_ASYNC_MIN_PRE_DOWNLOAD`; set `0` to disable the floor). Peak disk ≈ `num_workers × max_pre_download × chunk_size` — size `max_cache_size` accordingly.
+
+```bash
+# Debugging download/delete races — force synchronous downloads
+export LITDATA_ASYNC_CHUNK_PREFETCH=0
+
+# Keep max_pre_download=2 even with async enabled
+export LITDATA_ASYNC_MIN_PRE_DOWNLOAD=0
+```
+
+### Common environment variables
+
+| Variable | Default | Purpose |
+|----------|---------|---------|
+| `LITDATA_CACHE_DIR` | `~/.lightning/chunks` | Default chunk cache directory |
+| `LITDATA_ASYNC_CHUNK_PREFETCH` | on for remote | `0`/`1` force async chunk download overlap |
+| `LITDATA_ASYNC_MIN_PRE_DOWNLOAD` | `4` | Floor for `max_pre_download` when async is on (`0` = no floor) |
+| `LITDATA_OBSTORE_STREAM_MIN_CHUNK_MIB` | `8` | S3 obstore stream chunk size (MiB) |
+| `MAX_WAIT_TIME` | `120` | Seconds to wait for a chunk before error |
+| `FORCE_DOWNLOAD_TIME` | `30` | Seconds before force re-download of a missing chunk |
+| `LITDATA_DISABLE_VERSION_CHECK` | `0` | `1` skips the upgrade tip |
+| `HF_TOKEN` | — | Gated Hugging Face datasets |
+| `DEBUG_LITDATA` / `PRINT_DEBUG_LOGS` | `0` | Internal debug / stdout logs |
+
+Multi-node `optimize`/`map` on Studios also uses `DATA_OPTIMIZER_*` (set by the platform). Full catalog (debug logs, Studio injects, torchrun): see the LitData skill `reference/env-vars.md` when using agent skills, or the source modules `constants.py` / `async_prefetch.py`.
+
+</details>
+
 <details>
   <summary> ✅ Change cache directory path <a id="cache-directory" href="#cache-directory">🔗</a> </summary>
 &nbsp;
 
-Specify the directory where cached files should be stored, ensuring efficient data retrieval and management. This is particularly useful for organizing your data storage and improving access times.
+Specify where cached chunk files are stored.
 
 ```python
 from litdata import StreamingDataset
 from litdata.streaming.cache import Dir
 
-cache_dir = "/path/to/your/cache"
-data_dir = "s3://my-bucket/my_optimized_dataset"
+# Simple: dedicated cache directory
+dataset = StreamingDataset("s3://my-bucket/my_optimized_dataset", cache_dir="/path/to/your/cache")
+
+# Or when cache path and remote URL should differ:
+dataset = StreamingDataset(input_dir=Dir(path="/path/to/your/cache", url="s3://my-bucket/my_optimized_dataset"))
+```
+
+Global default without passing `cache_dir` every time:
 
-dataset = StreamingDataset(input_dir=Dir(path=cache_dir, url=data_dir))
+```bash
+export LITDATA_CACHE_DIR=/path/to/your/cache
+```
+
+CLI:
+
+```bash
+litdata cache path    # print the active cache directory
+litdata cache clear   # delete cached chunks
 ```
 
 </details>
@@ -1496,114 +1973,110 @@ dataset = StreamingDataset(input_dir="local:/data/shared-drive/some-data")
 </details>
 
 <details>
-  <summary> ✅ Optimize dataset in distributed environment <a id="distributed-optimization" href="#distributed-optimization">🔗</a> </summary>
+  <summary> ✅ Optimize / map across multiple machines (Lightning Studios) <a id="distributed-optimization" href="#distributed-optimization">🔗</a> </summary>
 &nbsp;
 
-Lightning can distribute large workloads across hundreds of machines in parallel. This can reduce the time to complete a data processing task from weeks to minutes by scaling to enough machines.
+On [Lightning Studios](https://lightning.ai/), `num_nodes` and `machine` scale `optimize` / `map` across many machines. This is **not** the same as `num_workers` (processes on one machine).
+
+**How it works**
+
+1. You call `optimize(..., num_nodes=N, machine=...)` (or `map`) inside a Studio.
+2. LitData starts a **data-prep job** that re-runs your script on **N** machines.
+3. Each machine processes a shard of the inputs (`num_nodes × num_workers` total workers). The last node merges chunk indexes into a single `index.json`.
+4. Your local call blocks until the job finishes; open the printed Runs URL to monitor.
 
-To apply the optimize operator across multiple machines, simply provide the num_nodes and machine arguments to it as follows:
+Outside Studio, passing `num_nodes` / `machine` raises an error (create a Studio account to use multi-node).
 
 ```python
-import os
 from litdata import optimize, Machine
 
 def compress(index):
     return (index, index ** 2)
 
-optimize(
-    fn=compress,
-    inputs=list(range(100)),
-    num_workers=2,
-    output_dir="my_output",
-    chunk_bytes="64MB",
-    num_nodes=2,
-    machine=Machine.DATA_PREP, # You can select between dozens of optimized machines
-)
+if __name__ == "__main__":
+    optimize(
+        fn=compress,
+        inputs=list(range(100)),
+        num_workers=8,              # processes per machine
+        output_dir="/teamspace/s3_connections/my-data/optimized-v1",  # durable bucket (recommended)
+        chunk_bytes="64MB",
+        num_nodes=32,               # machines in the job
+        machine=Machine.DATA_PREP,  # or omit to use the current Studio machine type
+    )
 ```
 
-If the `output_dir` is a local path, the optimized dataset will be present in: `/teamspace/jobs/{job_name}/nodes-0/my_output`. Otherwise, it will be stored in the specified `output_dir`.
+**Where outputs land**
 
-Read the optimized dataset:
+| `output_dir` | Result |
+|--------------|--------|
+| `/teamspace/s3_connections/...`, `/teamspace/datasets/...`, `s3://...`, `gs://...` | Written directly to that store (**recommended**) |
+| Local or `/teamspace/studios/this_studio/...` | Remapped to the job’s **artifacts** storage; the Studio UI may also expose it under `/teamspace/jobs/<job>/...` |
 
 ```python
 from litdata import StreamingDataset
 
-output_dir = "/teamspace/jobs/litdata-optimize-2024-07-08/nodes.0/my_output"
-
-dataset = StreamingDataset(output_dir)
-
-print(dataset[:])
+# Prefer the same connection / cloud URL you wrote to:
+dataset = StreamingDataset("/teamspace/s3_connections/my-data/optimized-v1")
 ```
 
+The same `num_nodes` / `machine` pattern works with `map`. See also [Parallelize transforms and data optimization](#parallelize-transforms-and-data-optimization-on-cloud-machines).
+
 </details>
 
 <details>
   <summary> ✅ Encrypt, decrypt data at chunk/sample level <a id="encrypt-decrypt" href="#encrypt-decrypt">🔗</a> </summary>
 &nbsp;
 
-Secure data by applying encryption to individual samples or chunks, ensuring sensitive information is protected during storage.
+Encrypt optimized data at **sample** or **chunk** level. Built-ins: `FernetEncryption` and `RSAEncryption` (`litdata.utilities.encryption`). Requires the `cryptography` package. **Not supported for Mosaic MDS.**
+
+| `level` | Meaning |
+|---------|---------|
+| `"sample"` (default) | Encrypt each sample independently |
+| `"chunk"` | Encrypt whole chunks |
 
-This example shows how to use the `FernetEncryption` class for sample-level encryption with a data optimization function.
+**Fernet (symmetric)**
 
 ```python
-from litdata import optimize
+from litdata import optimize, StreamingDataset
 from litdata.utilities.encryption import FernetEncryption
-import numpy as np
-from PIL import Image
 
-# Initialize FernetEncryption with a password for sample-level encryption
-fernet = FernetEncryption(password="your_secure_password", level="sample")
+fernet = FernetEncryption(password="your_secure_password", level="sample")  # or level="chunk"
 data_dir = "s3://my-bucket/optimized_data"
 
-def random_image(index):
-    """Generate a random image for demonstration purposes."""
-    fake_img = Image.fromarray(np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8))
-    return {"image": fake_img, "class": index}
-
-# Optimize data while applying encryption
-optimize(
-    fn=random_image,
-    inputs=list(range(5)),  # Example inputs: [0, 1, 2, 3, 4]
-    num_workers=1,
-    output_dir=data_dir,
-    chunk_bytes="64MB",
-    encryption=fernet,
-)
-
-# Save the encryption key to a file for later use
-fernet.save("fernet.pem")
-```
-
-Load the encrypted data using the `StreamingDataset` class as follows:
-
-```python
-from litdata import StreamingDataset
-from litdata.utilities.encryption import FernetEncryption
+def fn(index):
+    return {"index": index, "value": index**2}
 
-# Load the encryption key
-fernet = FernetEncryption(password="your_secure_password", level="sample")
-fernet.load("fernet.pem")
+if __name__ == "__main__":
+    optimize(
+        fn=fn,
+        inputs=list(range(5)),
+        num_workers=1,
+        output_dir=data_dir,
+        chunk_bytes="64MB",
+        encryption=fernet,
+    )
+    fernet.save("fernet.pem")  # persist salt/level; keep the password safe
 
-# Create a streaming dataset for reading the encrypted samples
+# Later — load key material with the same password
+fernet = FernetEncryption.load("fernet.pem", password="your_secure_password")
 ds = StreamingDataset(input_dir=data_dir, encryption=fernet)
 ```
 
-Implement your own encryption method: Subclass the `Encryption` class and define the necessary methods:
+**RSA (asymmetric)**
 
 ```python
-from litdata.utilities.encryption import Encryption
+from litdata.utilities.encryption import RSAEncryption
 
-class CustomEncryption(Encryption):
-    def encrypt(self, data):
-        # Implement your custom encryption logic here
-        return data
+rsa = RSAEncryption(password="your_secure_password", level="sample")  # or "chunk"
+optimize(fn=fn, inputs=list(range(5)), output_dir=data_dir, chunk_bytes="64MB", encryption=rsa)
+rsa.save("rsa.pem")
 
-    def decrypt(self, data):
-        # Implement your custom decryption logic here
-        return data
+rsa = RSAEncryption.load("rsa.pem", password="your_secure_password")
+ds = StreamingDataset(input_dir=data_dir, encryption=rsa)
 ```
 
-This allows the data to remain secure while maintaining flexibility in the encryption method.
+**Custom algorithm** — subclass `Encryption` and implement `encrypt` / `decrypt` / `save` / `load` / `state_dict` / `algorithm`.
+
 </details>
 
 <details>
@@ -1675,17 +2148,99 @@ if __name__ == "__main__":
 </details>
 
 <details>
-  <summary> ✅ Lightning AI Data Connections - Direct download and upload <a id="lightning-connections" href="#lightning-connections">🔗</a> </summary>
+  <summary> ✅ Resolve any path or cloud URL (local, S3, GCS, R2, Azure, HF, Studio) <a id="resolve-paths" href="#resolve-paths">🔗</a> </summary>
 
 &nbsp;
 
-[Lightning Studios](https://lightning.ai/) have special directories for data connections that are available to an entire teamspace. LitData functions that reference those directories will experience a significant performance increase as uploads and downloads will happen directly from the bucket that backs the folder.
+LitData **resolves** every dataset path you pass to `StreamingDataset`, `StreamingRawDataset`, `optimize`, `map`, and related APIs. You write one path string; LitData figures out whether to read locally, download from object storage, or (inside [Lightning Studios](https://lightning.ai/)) talk **directly to the bucket** behind a `/teamspace/...` mount instead of going through slow FUSE I/O.
+
+### Supported URI schemes
 
-For example, output artifacts from this code will be directly uploaded to the `my-data-1` s3 bucket.
+| Scheme | Example | Use when |
+|--------|---------|----------|
+| Local path | `./data` or `/data/imagenet` | Files on disk |
+| `s3://` | `s3://my-bucket/optimized` | AWS S3 |
+| `gs://` | `gs://my-bucket/optimized` | Google Cloud Storage |
+| `r2://` | `r2://my-bucket/optimized` | Cloudflare R2 |
+| `azure://` | `azure://container/optimized` | Azure Blob Storage |
+| `hf://` | `hf://datasets/org/name/data` | Hugging Face datasets (parquet) |
+| `local:` | `local:/mnt/nfs/dataset` | Network / shared drive (LitData still caches chunks locally to reduce NAS load) |
 
 ```python
-from litdata import optimize
+from litdata import StreamingDataset, optimize
+
+# Same APIs — only the path changes
+StreamingDataset("s3://my-bucket/fast_data", shuffle=True, drop_last=True)
+StreamingDataset("gs://my-bucket/fast_data")
+StreamingDataset("r2://my-bucket/fast_data", storage_options={...})
+StreamingDataset("azure://my-container/fast_data", storage_options={...})
+StreamingDataset("hf://datasets/org/name/data")
+StreamingDataset("local:/data/shared-drive/some-data")
+StreamingDataset("/var/data/fast_data")  # plain local directory
+```
+
+Pass cloud credentials with `storage_options` (and optional `session_options` for boto3 profiles/regions). See [Stream from multiple cloud providers](#cloud-providers).
+
+### Cache directory vs remote URL
+
+By default LitData caches downloaded chunks under `~/.lightning/chunks` (override with `cache_dir=` or `LITDATA_CACHE_DIR`). When the cache location and the dataset URL must differ, use `Dir`:
+
+```python
+from litdata import StreamingDataset
+from litdata.streaming.resolver import Dir
+
+dataset = StreamingDataset(
+    Dir(path="/fast-ssd/cache/run-1", url="s3://my-bucket/fast_data")
+)
+# Equivalent:
+dataset = StreamingDataset("s3://my-bucket/fast_data", cache_dir="/fast-ssd/cache/run-1")
+```
+
+```bash
+export LITDATA_CACHE_DIR=/fast-ssd/cache
+litdata cache path    # show active cache directory
+litdata cache clear   # wipe cached chunks
+```
+
+### Date/time path templates
+
+Embed a `strftime` pattern in `{...}` and LitData expands it to the current time (useful for versioned `output_dir`s):
+
+```python
+# e.g. on 2025-05-05 → ".../run_2025-05-05"
+optimize(
+    fn=fn,
+    inputs=inputs,
+    output_dir="s3://my-bucket/datasets/run_{%Y-%m-%d}",
+    chunk_bytes="64MB",
+)
+```
+
+### Lightning Studio `/teamspace/...` paths (direct bucket I/O)
+
+In Lightning Studios, data connections appear under `/teamspace/...`. **Prefer these paths in LitData** — optimize/map uploads and StreamingDataset downloads use the **backing object store URL** (and temporary credentials when needed), which is much faster than reading every file through the FUSE mount.
+
+| Path prefix | What LitData does |
+|-------------|-------------------|
+| `/teamspace/studios/this_studio/...` | Local Studio workspace disk (not a cloud URL) |
+| `/teamspace/studios/<other_studio>/...` | Resolves to that Studio’s content bucket (`s3://` or `gs://`) |
+| `/teamspace/s3_connections/<name>/...` | Direct S3 to the connection’s bucket |
+| `/teamspace/gcs_connections/<name>/...` | Direct GCS |
+| `/teamspace/s3_folders/<name>/...` | S3 folder connection |
+| `/teamspace/gcs_folders/<name>/...` | GCS folder connection |
+| `/teamspace/lightning_storage/<name>/...` | Lightning-managed object storage (R2-style) |
+| `/teamspace/datasets/...` | Teamspace datasets mount → project datasets bucket |
+
+```python
+from litdata import StreamingDataset, StreamingRawDataset, optimize
+
+# Stream optimized data from an attached S3 connection (direct bucket download)
+dataset = StreamingDataset("/teamspace/s3_connections/my-data-1/fast_data", shuffle=True, drop_last=True)
+
+# Stream raw files from a connection
+raw = StreamingRawDataset("/teamspace/s3_connections/my-bucket-1/raw")
 
+# Optimize *into* a connection — chunks upload straight to the bucket
 def should_keep(data):
     if data % 2 == 0:
         yield data
@@ -1696,31 +2251,16 @@ if __name__ == "__main__":
         inputs=list(range(1000)),
         output_dir="/teamspace/s3_connections/my-data-1/output",
         chunk_bytes="64MB",
-        num_workers=1
+        num_workers=1,
     )
 ```
 
+**Tips**
 
-Similarly, data will be downloaded directly from the `my-data-1` s3 bucket in this example code.
-
-```python
-from litdata import StreamingRawDataset
-
-if __name__ == "__main__":
-    data_dir = "/teamspace/s3_connections/my-bucket-1/data"
-
-    raw_dataset = StreamingRawDataset(data_dir)
-
-    data = list(raw_dataset)
-    print(data)
-```
+- Version remote outputs (`.../v2`, `.../run_{%Y-%m-%d}`). Optimized datasets are immutable unless you pass `mode="append"` or `mode="overwrite"`.
+- Outside Studio, use `s3://` / `gs://` / … with your own credentials — `/teamspace/...` resolution needs Lightning Studio environment variables.
+- `optimize` / `map` with `num_nodes` launch a Studio **job** (not local multi-process). Prefer a connection / cloud `output_dir`; local / `this_studio` optimize outputs go to job artifacts (UI may show `/teamspace/jobs/...`). Details: [distributed optimization](#distributed-optimization).
 
-References to any of the following directories will work similarly:
-1. `/teamspace/lightning_storage/...`
-2. `/teamspace/s3_connections/...`
-3. `/teamspace/gcs_connections/...`
-4. `/teamspace/s3_folders/...`
-5. `/teamspace/gcs_folders/...`
 </details>
 
 &nbsp;
@@ -1734,31 +2274,113 @@ References to any of the following directories will work similarly:
 
 Apply the same change to different parts of the dataset at once to save time and effort.
 
-The `map` operator can be used to apply a function over a list of inputs.
-
-Here is an example where the `map` operator is used to apply a `resize_image` function over a folder of large images.
+The `map` operator applies a function over a list of inputs. **`fn` must write into `output_dir` and return `None`.** Guard with `if __name__ == "__main__"` when using multiple workers.
 
 ```python
+import os
 from litdata import map
 from PIL import Image
 
-# Note: Inputs could also refer to files on s3 directly.
-input_dir = "my_large_images"
+input_dir = "my_large_images"  # or s3://...
 inputs = [os.path.join(input_dir, f) for f in os.listdir(input_dir)]
 
-# The resize image takes one of the input (image_path) and the output directory.
-# Files written to output_dir are persisted.
 def resize_image(image_path, output_dir):
-  output_image_path = os.path.join(output_dir, os.path.basename(image_path))
-  Image.open(image_path).resize((224, 224)).save(output_image_path)
+    output_image_path = os.path.join(output_dir, os.path.basename(image_path))
+    Image.open(image_path).resize((224, 224)).save(output_image_path)
 
-map(
-    fn=resize_image,
-    inputs=inputs,
-    output_dir="s3://my-bucket/my_resized_images",
-)
+if __name__ == "__main__":
+    map(
+        fn=resize_image,
+        inputs=inputs,
+        output_dir="s3://my-bucket/my_resized_images",
+        num_workers=8,
+    )
 ```
 
+**`map` arguments**
+
+| Argument | Default | Description |
+|----------|---------|-------------|
+| `fn` | required | `fn(input, output_dir) -> None` |
+| `inputs` | required | Sequence or `StreamingDataLoader` |
+| `output_dir` | required | Local or cloud path ([resolver](#resolve-paths)) |
+| `input_dir` | `None` | Root for remote inputs (background download while processing) |
+| `weights` | `None` | Per-input weights to balance workers |
+| `num_workers` | CPU count | Local process workers |
+| `fast_dev_run` | `False` | Process only a few items (`True` → small default, or an int) |
+| `num_nodes` / `machine` | `None` | Scale out on [Lightning Studios](https://lightning.ai/) |
+| `num_downloaders` / `num_uploaders` | auto | I/O concurrency per worker |
+| `reorder_files` | `True` | Pack by file size for balance; `False` preserves order |
+| `error_when_not_empty` | `False` | Error if `output_dir` already has files |
+| `reader` | default | Custom reader for inputs |
+| `batch_size` | `None` | Group inputs into batches for `fn` |
+| `start_method` | spawn† | Multiprocessing start method (†spawn unless IPython) |
+| `optimize_dns` | `None` | Optimized DNS (Studio / cloud) |
+| `storage_options` | `{}` | Cloud credentials / endpoints |
+| `keep_data_ordered` | `True` | `False` = shared work queue (better for uneven/slow workers) |
+
+</details>
+
+<details>
+  <summary> ✅ <code>optimize</code> arguments reference <a id="optimize-kwargs" href="#optimize-kwargs">🔗</a> </summary>
+&nbsp;
+
+Full knob list for `litdata.optimize` (see Quick start for the minimal recipe). **Exactly one of `chunk_bytes` or `chunk_size`.** Use `if __name__ == "__main__"`.
+
+| Argument | Default | Description |
+|----------|---------|-------------|
+| `fn` | required | Maps each input → sample (or `yield` samples / skip bad ones) |
+| `inputs` | `None` | Sequence or `StreamingDataLoader` (ignored if `queue` is set) |
+| `queue` | `None` | `multiprocessing.Queue` of live inputs; send **one** `ALL_DONE` when finished |
+| `output_dir` | `"optimized_data"` | Local or cloud ([resolver](#resolve-paths)); version remote prefixes |
+| `input_dir` | `None` | Remote input root for background download |
+| `weights` | `None` | Per-input weights to balance workers |
+| `chunk_bytes` | `None` | Max bytes per chunk (e.g. `"64MB"`; see [FAQ](#faq-chunk-shuffle) for larger samples) |
+| `chunk_size` | `None` | Max items (or tokens with `TokensLoader`) per chunk |
+| `align_chunking` | `False` | Match single-worker chunk boundaries (needs `chunk_size`; uneven load) |
+| `compression` | `None` | `"zstd"` today |
+| `encryption` | `None` | `FernetEncryption` / `RSAEncryption` / custom ([encrypt](#encrypt-decrypt)) |
+| `num_workers` | CPU count | Local workers |
+| `fast_dev_run` | `False` | Smoke a subset of inputs |
+| `num_nodes` / `machine` | `None` | Multi-node on Lightning Studios |
+| `num_downloaders` / `num_uploaders` | auto | I/O concurrency per worker |
+| `reorder_files` | `True` | Size-based packing; `False` preserves order |
+| `reader` | default | Custom input reader |
+| `batch_size` | `None` | Group inputs for `fn` |
+| `mode` | `None` | `"append"` or `"overwrite"` existing dataset; default treats data as immutable |
+| `use_checkpoint` | `False` | Resume an interrupted optimize from `.checkpoints` |
+| `item_loader` | `None` | e.g. `TokensLoader()` for contiguous tokens |
+| `start_method` | spawn† | Multiprocessing start method |
+| `optimize_dns` | `None` | Optimized DNS |
+| `storage_options` | `{}` | Cloud credentials / endpoints |
+| `keep_data_ordered` | `True` | `False` = shared queue among workers |
+| `verbose` | `True` | Progress logging |
+
+Related features: [shared queue](#shared-queue), [queue input](#queue-input), [append/overwrite](#modify-datasets), [compression](#compression), [TokensLoader / LLM](#llm-training), [filter](#filter-data).
+
+</details>
+
+<details>
+  <summary> ✅ Cloud-optimized <code>walk</code> (list files at scale) <a id="walk" href="#walk">🔗</a> </summary>
+&nbsp;
+
+`litdata.walk` is a threaded, cloud-friendly alternative to `os.walk` for building large `inputs=` lists (especially on Lightning Studios). Yields `(dirpath, dirnames, filenames)` like `os.walk`, but **order is not depth-first**.
+
+```python
+from litdata import walk, optimize
+
+inputs = []
+for root, dirs, files in walk("/teamspace/s3_connections/my-data/raw", max_workers=32):
+    for name in files:
+        if name.endswith(".jpg"):
+            inputs.append(f"{root}/{name}")
+
+if __name__ == "__main__":
+    optimize(fn=load_image, inputs=inputs, output_dir="...", chunk_bytes="64MB")
+```
+
+Prints a warning outside Lightning Studio — it is optimized for that environment; elsewhere prefer `os.walk` or your cloud SDK’s listing API.
+
 </details>
 
 &nbsp;
@@ -1816,8 +2438,7 @@ Speed to stream raw Imagenet 1.2M from different cloud storage providers:
 | AWS S3  | ~6400 +/- 100     | ~3200 +/- 100  |
 | Google Cloud Storage | ~5650 +/- 100     | ~3100 +/- 100  |
 
-> **Note:**
-> Use `StreamingRawDataset` if you want to stream your data as-is. Use `StreamingDataset` if you want the fastest streaming and are okay with optimizing your data first.
+> **Also see:** [`StreamingRawDataset`](#stream-raw) streams existing files with **no optimize step** (great default to start). Use `StreamingDataset` after `optimize` when you need the highest sustained training throughput.
 
 &nbsp;
 
@@ -1842,7 +2463,7 @@ Time to optimize 1.2 million ImageNet images (Faster is better):
 
 ## Parallelize data transforms
 
-Transformations with LitData are linearly parallelizable across machines.
+Transformations with LitData are linearly parallelizable across machines on [Lightning Studios](https://lightning.ai/) (see [distributed optimization](#distributed-optimization) for how the job launch works).
 
 For example, let's say that it takes 56 hours to embed a dataset on a single A10G machine. With LitData,
 this can be speed up by adding more machines in parallel
@@ -1855,20 +2476,20 @@ this can be speed up by adding more machines in parallel
 | ...               | ...            |
 | 64              | 0.875        |
 
-To scale the number of machines, run the processing script on [Lightning Studios](https://lightning.ai/):
-
 ```python
 from litdata import map, Machine
 
 map(
   ...
   num_nodes=32,
-  machine=Machine.DATA_PREP, # Select between dozens of optimized machines
+  machine=Machine.DATA_PREP,  # or omit to inherit the Studio machine
+  # Prefer output_dir on /teamspace/s3_connections/... or s3://...
 )
 ```
 
 ## Parallelize data optimization
-To scale the number of machines for data optimization, use [Lightning Studios](https://lightning.ai/):
+
+Same Studio job launch as `map` — `num_nodes` machines × `num_workers` processes; last node merges the index.
 
 ```python
 from litdata import optimize, Machine
@@ -1876,7 +2497,8 @@ from litdata import optimize, Machine
 optimize(
   ...
   num_nodes=32,
-  machine=Machine.DATA_PREP, # Select between dozens of optimized machines
+  machine=Machine.DATA_PREP,
+  output_dir="/teamspace/s3_connections/my-data/optimized-v1",
 )
 ```
 
@@ -1905,7 +2527,7 @@ Below are templates for real-world applications of LitData at scale.
 | -------------------------------- | ----------------- | ----------------- | -------------- | -------------- |
 | [Benchmark cloud data-loading libraries](https://lightning.ai/lightning-ai/studios/benchmark-cloud-data-loading-libraries) | Image & Label | 10 | 1 | [Imagenet 1M](https://paperswithcode.com/sota/image-classification-on-imagenet?tag_filter=171) |
 | [Optimize GeoSpatial data for model training](https://lightning.ai/lightning-ai/studios/convert-spatial-data-to-lightning-streaming) | Image & Mask | 120 | 32 | [Chesapeake Roads Spatial Context](https://github.com/isaaccorley/chesapeakersc) |
-| [Optimize TinyLlama 1T dataset for training](https://lightning.ai/lightning-ai/studios/prepare-the-tinyllama-1t-token-dataset) | Text | 240 | 32 | [SlimPajama](https://huggingface.co/datasets/cerebras/SlimPajama-627B) & [StarCoder](https://huggingface.co/datasets/bigcode/starcoderdata) |
+| [Optimize TinyLlama 1T dataset for training](https://lightning.ai/lightning-ai/studios/prepare-the-tinyllama-1t-token-dataset) | Text | 240 | 32 | [SlimPajama](https://www.cerebras.ai/blog/slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama) & [StarCoder](https://huggingface.co/datasets/bigcode/starcoderdata) |
 | [Optimize parquet files for model training](https://lightning.ai/lightning-ai/studios/convert-parquets-to-lightning-streaming) | Parquet Files | 12 | 16 | Randomly Generated data |
 
 &nbsp;