litdata 0.2.66__tar.gz → 0.2.67__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: litdata
-Version: 0.2.66
+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,7 @@ 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> •
@@ -146,6 +149,8 @@ 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>
@@ -168,29 +173,38 @@ Source: [`.claude/skills/litdata/`](.claude/skills/litdata/) in this repository
 # 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.
@@ -227,7 +241,7 @@ if __name__ == "__main__":
         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
     )
 ```
 
@@ -325,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) |
+
+### Group related files (`setup`)
 
-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.
+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>
@@ -715,6 +808,35 @@ loader = StreamingDataLoader(train, batch_size=64, shuffle=True, drop_last=True)
 
 </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;
@@ -1093,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
@@ -2213,7 +2335,7 @@ Full knob list for `litdata.optimize` (see Quick start for the minimal recipe).
 | `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"`) |
+| `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 |
@@ -2316,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;
 
@@ -2406,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;