diff --git a/.claude/skills/litdata/SKILL.md b/.claude/skills/litdata/SKILL.md index 96b7a754d..55498f7ac 100644 --- a/.claude/skills/litdata/SKILL.md +++ b/.claude/skills/litdata/SKILL.md @@ -8,8 +8,10 @@ description: >- merge_datasets, index_parquet_dataset, index_hf_dataset), answering how-to questions, choosing raw vs optimize vs parquet/HF/MDS, tuning cache/prefetch/ shuffle/seed, resolving paths (s3/gs/r2/azure/hf/local:/teamspace via - resolver.py), or when navigating/editing src/litdata, tests, CI, or debugging - streaming / optimize / map. + resolver.py), documenting or debugging optimize/map downloaders/uploaders/ + removers, FsProvider vs Downloader, FUSE s3_connections/s3_folders, or + multi-node DATA_OPTIMIZER_* / num_nodes jobs, or when navigating/editing + src/litdata, tests, CI, or debugging streaming / optimize / map. --- # LitData @@ -36,16 +38,20 @@ Useful options: `-g` (user-global), `-a cursor` (Cursor only), `-y` (non-interac Before writing examples or answering how-tos, read the cookbook. Highlights: -| Topic | Remember | -| -------------- | ----------------------------------------------------------------------------------------------------------------------------------------- | -| **Raw files** | `StreamingRawDataset`: raw `bytes`, fully async + batched downloads, retries; torch `DataLoader` — `#stream-raw` / `using-litdata.md` §10 | -| Images | Return **JPEG** (`JpegImageFile` / quality ≈95). Plain `PIL.Image` / `fromarray` → huge PIL RAW | -| Train stream | Optimized: `StreamingDataLoader` + `shuffle=True, drop_last=True, seed=…` | -| Optimize | `if __name__ == "__main__"`; exactly one of `chunk_bytes` \| `chunk_size` | -| Cache | Peak disk ≈ `num_workers × max_pre_download × chunk_size`; default `max_cache_size="100GB"` | -| Async prefetch | Remote downloads overlapped by default; `LITDATA_ASYNC_CHUNK_PREFETCH=0/1`; floor `max_pre` to 4 — `reference/env-vars.md` | -| **Paths** | Studio `/teamspace/s3_connections` & co are **FUSE** — LitData hits S3/GCS/**R2** (`lightning_storage`) directly. `reference/resolver.md` | -| Parquet / HF | Index + `ParquetLoader` (HF auto); `spawn` with workers; `using-litdata.md` §10 | +| Topic | Remember | +| ---------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| **Raw files** | `StreamingRawDataset` + torch `DataLoader` — `#stream-raw` / §10. Prefer cloud URL / connection path over FUSE. Defaults: `max_concurrent_downloads=None` (adaptive Stage 1), `max_prefetch=16` (worker-aware ~64 aggregate), `hedge_delay=0`, `download_timeout=120` (**batch-level**), `range_parallel_threshold=0`. Explicit `int` concurrency = exact permits. | +| Images | Return **JPEG** (`JpegImageFile` / quality ≈95). Plain `PIL.Image` / `fromarray` → huge PIL RAW | +| Train stream | Optimized: `StreamingDataLoader` + `shuffle=True, drop_last=True, seed=…` | +| Optimize | `if __name__ == "__main__"`; exactly one of `chunk_bytes` \| `chunk_size`. Default **64MB**; multi‑MB samples → consider **256–512MB**. **Shuffle the sample list before `optimize()`** when source order matters — README `#faq-chunk-shuffle` | +| Ordered data | Chunk/item shuffle ≠ file-level shuffle. Shuffle before `optimize`, or use `StreamingRawDataset` + `DataLoader(shuffle=True)`. LitData does distributed + within-chunk bucket sampling automatically | +| Cache | Peak disk ≈ `num_workers × max_pre_download × chunk_size`; default `max_cache_size="100GB"` | +| Async prefetch | Remote downloads overlapped by default; `LITDATA_ASYNC_CHUNK_PREFETCH=0/1`; floor `max_pre` to 4 — `reference/env-vars.md` | +| **Paths** | Studio `/teamspace/s3_connections` & co are **FUSE** (convenience only — slow, can crash under load). LitData resolves them and talks **directly** to S3/GCS/**R2**. Never read the mount by hand. `reference/resolver.md` + `reference/data-movement.md` | +| **Optimize I/O** | Processing downloaders/uploaders/removers are **processes** in `data_processor.py` using **FsProvider** — not the streaming `Downloader` ABC. FUSE → `Dir.url` → `/cache/data`. Load `reference/data-movement.md`. | +| **Multi-node** | `num_nodes=` = Lightning Studio job (`_execute`), not torchrun/SLURM. Shard by `DATA_OPTIMIZER_*`; all ranks upload chunks; **last node** merges `{node}-index.json` → `index.json`. Load `reference/multi-node.md`. | +| Throughput | Rough ImageNet Studio order-of-magnitude (not guarantees): FUSE ~**600**/s · Raw (right tuning) ~**6–7k**/s · Optimized 64MB chunks ~**11k**/s — `using-litdata.md` FAQ. Raw benches: medians + provenance SHAs; never cite short-window n=1 against Stage 0 medians. | +| Parquet / HF | Index + `ParquetLoader` (HF auto); `spawn` with workers; `using-litdata.md` §10 | ## Reference map @@ -53,13 +59,16 @@ Before writing examples or answering how-tos, read the cookbook. Highlights: | ----------------------------------------------------------------------------- | --------------------------------------------------- | | **Use the library** (raw, optimize/stream, parquet/HF, serializers, shuffle) | `reference/using-litdata.md` | | **Paths / URLs / Studio mounts / `Dir` / time templates** | `reference/resolver.md` (+ README `#resolve-paths`) | +| **Downloaders / uploaders / removers / FUSE→cloud / optimize I/O** | `reference/data-movement.md` (+ `processing.md`) | +| **Multi-node optimize/map** (`num_nodes`, `DATA_OPTIMIZER_*`, index merge) | `reference/multi-node.md` (+ `processing.md`) | | Read path, shuffle math, item loaders, Combined/Parallel | `reference/streaming.md` | | **Cache / BinaryWriter / BinaryReader / `index.json` / FsProvider / sampler** | `reference/storage-format.md` | | Cache / prefetch / eviction / shared-chunk deletion | `reference/cache-and-chunk-lifecycle.md` | | **Env vars** (async prefetch, cache, debug, `DATA_OPTIMIZER_*`, Studio) | `reference/env-vars.md` | | Fair streaming benchmarks (`benchmarks/` suite) | `reference/benchmarking.md` | +| **Raw adaptive concurrency / look-ahead stages** (clients own rate) | repo `benchmarks/ADAPTIVE_CONCURRENCY.md` | | Lightning Studio env, credentials, free-threading | `reference/lightning-studio.md` | -| Write path / **multi-node** `num_nodes` job launch | `reference/processing.md` | +| Write path orchestration; raw internals; pointers to I/O + multi-node | `reference/processing.md` | | Dev env, PR/CI style | `reference/contributing.md` | | Tests & fixtures | `reference/testing.md` | | Tracing, breakpoints, env knobs | `reference/debugging.md` | @@ -90,7 +99,7 @@ Defined under `streaming/`, `processing/`, `raw/`, `utilities/` — see cookbook - Chunk: `[num_items][offsets][data]`; `index.json` holds chunks + config (`data_format`, `item_loader`, …) — [storage-format.md](reference/storage-format.md). - Item loaders own layout + intervals (`PyTreeLoader`, `TokensLoader`, `ParquetLoader`). -- Write/management I/O = `FsProvider` (s3/gs/r2); training downloads = `Downloader`. Sampler `ChunkedIndex` is read-path; `CacheBatchSampler` is `CacheDataLoader` only. +- Write/management I/O = `FsProvider` (s3/gs/r2); training downloads = `Downloader`. Optimize worker pools use FsProvider via `_download_data_target` / `_upload_fn` — see `data-movement.md`. Sampler `ChunkedIndex` is read-path; `CacheBatchSampler` is `CacheDataLoader` only. - Ranks from env (`_DistributedEnv` / `DATA_OPTIMIZER_*`), not a custom network. - Shuffle deterministic from `seed`+epoch+chunk → resumable (`shuffle.py`, not `sampler.py`). - Design: one less thing to remember; pure PyTorch; backward compatible; test-driven. diff --git a/.claude/skills/litdata/reference/benchmarking.md b/.claude/skills/litdata/reference/benchmarking.md index f383b3a3f..8ced9f7ed 100644 --- a/.claude/skills/litdata/reference/benchmarking.md +++ b/.claude/skills/litdata/reference/benchmarking.md @@ -16,11 +16,26 @@ Studio paths / free-threading → [lightning-studio.md](lightning-studio.md). Pr | `benchmarks/litdata/optimize_imagenet.py` | `optimize` ImageNet (JPEG/PIL write modes) | | `benchmarks/litdata/stream_imagenet.py` | Epoch throughput for an optimized dataset | | `benchmarks/stream_raw_imagenet.py` | `StreamingRawDataset` baseline (no optimize) | +| `benchmarks/bench_raw_before_vs_after.py` | Stage 0 A/B harness for raw cloud download changes | +| `benchmarks/ADAPTIVE_CONCURRENCY.md` | Adaptive concurrency / look-ahead design + Stage 1 | | `benchmarks/ffcv/` | Convert / write / stream with FFCV for format comparison | | `benchmarks/ffcv/README.md` | FFCV install + write/stream steps | Start from `benchmarks/litdata/README.md` for LitData-only runs; use `benchmarks/ffcv/` when comparing formats. All scripts are CLI-based (`--help`). +## Raw streaming Stage 0 protocol + +When measuring or claiming `StreamingRawDataset` cloud download wins (`bench_raw_before_vs_after.py` and friends): + +1. **Window = `max(N batches, T seconds)`** — require **both** floors (default ≥300 batches **and** ≥30s). Not either/or. +2. **Warm** `max(1, num_workers × prefetch_factor)` batches before timing. +3. **Repeats + medians** — prefer interleaved A/B, `n≥5` for grids, report median + spread. Single-run digs are exploratory only. +4. **Append-only artifacts** — write `*.{sha}.{unix_ts}.json` (+ JSONL); never overwrite prior result files. +5. **Provenance** — record `before_sha` / `after_sha` from `git rev-parse` on each PYTHONPATH tree (not only the runner SHA in the filename). Refuse to publish without both. +6. **Trust hierarchy:** provenance-verified confirm cell (known `before_sha`/`after_sha`, protocol floors, n≥3) ≫ full-grid medians with null tree SHAs ≫ short-window or n=1 digs. Never cite a short-window n=1 against Stage 0 medians. + +Design note / Stage 1 formula: `benchmarks/ADAPTIVE_CONCURRENCY.md`. + ### Typical LitData flow ```bash diff --git a/.claude/skills/litdata/reference/data-movement.md b/.claude/skills/litdata/reference/data-movement.md new file mode 100644 index 000000000..1f5543af7 --- /dev/null +++ b/.claude/skills/litdata/reference/data-movement.md @@ -0,0 +1,337 @@ +# Data movement: downloaders, uploaders, removers, path resolution + +Operational reference for how LitData moves bytes between local disk and remote storage during **`optimize` / `map`** and how that relates to the **streaming / raw read** path. Grep landmarks use `src/litdata/` paths. + +**Related:** [processing.md](processing.md) (orchestrator) · [multi-node.md](multi-node.md) · [resolver.md](resolver.md) · [storage-format.md](storage-format.md) (`FsProvider` vs `Downloader`) · [lightning-studio.md](lightning-studio.md) + +______________________________________________________________________ + +## 0. Two different “downloaders” — do not conflate + +| Path | What agents mean by “downloader” | Module / symbols | Transport | +| ---------------------------------------------------------------- | -------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------- | +| **Write / processing** (`optimize`, `map`) | Child **processes** per worker that prefetch input files into a data cache | `processing/data_processor.py`: `_download_data_target`, `_start_downloaders` | **`FsProvider`** (`streaming/fs_provider.py`) for `s3`/`gs`/`r2`; local `shutil.copyfile` otherwise | +| **Read / streaming** (`StreamingDataset`, `StreamingRawDataset`) | **`Downloader` ABC** subclasses selected by URL prefix | `streaming/downloader.py`: `Downloader`, `S3Downloader`, `GCPDownloader`, `R2Downloader`, `AzureDownloader`, `HFDownloader`, `LocalDownloader`, `get_downloader`, `_DOWNLOADERS` | Cloud SDKs / obstore / boto3 per subclass | + +There are **no** classes named `Uploader` or `Remover`. Processing upload/remove are process targets `_upload_fn` and `_remove_target` in `data_processor.py`. + +**Schemes:** + +- Processing I/O via FsProvider: `_SUPPORTED_PROVIDERS = ("s3", "gs", "r2")` in `constants.py`. +- Streaming downloaders also support `azure://`, `hf://`, `local:` (see `_DOWNLOADERS`). + +______________________________________________________________________ + +## 1. End-to-end data movement (`optimize` / `map`) + +``` +User paths (inputs / input_dir / output_dir) + │ + ▼ +_resolve_dir (streaming/resolver.py) + │ Dir(path=…, url=…, data_connection_id=?) + ▼ +DataProcessor (processing/data_processor.py) + │ broadcast_object input/output Dir (only if broadcast_paths / `{%strftime}`) + │ shard items → DataWorkerProcess × num_workers + ▼ +Per worker (BaseWorker._setup): + _collect_paths → rewrite item paths to cache_data_dir when remote/FUSE + _start_downloaders → Process(_download_data_target) × num_downloaders + _start_uploaders → Process(_upload_fn) × num_uploaders + _start_remover → Process(_remove_target) if delete_cached_files + │ + ▼ +ready_to_process_queue → user fn (optimize→Cache/BinaryWriter | map→temp out dir) + │ + ▼ +_try_upload → to_upload_queues → remote/local output_dir + │ + ▼ +remove_queue → delete local intermediates (inputs from data cache; uploaded chunks) + │ + ▼ +DataChunkRecipe._done → merge per-worker indexes → upload index.json + (multi-node: {node_rank}-index.json then last-node merge — see multi-node.md) +``` + +Public knobs (`processing/functions.py` → `DataProcessor`): + +| Knob | Default | Meaning | +| --------------------- | --------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `num_downloaders` | `2` (`DataProcessor`: `num_downloaders or 2`) | Downloader processes **per worker** | +| `num_uploaders` | `1` | Uploader processes **per worker** | +| `delete_cached_files` | `True` on `DataProcessor` | Passed to worker as `remove`; starts remover. **Not** exposed on public `optimize()` / `map()` — stays default True unless you construct `DataProcessor` yourself | +| `input_dir` | Auto via `_get_input_dir(inputs)` or explicit | Resolved `Dir`; drives download + path rewrite | +| `output_dir` | Required | Resolved `Dir`; drives upload | +| `storage_options` | `{}` | Merged with `data_connection_id` via `construct_storage_options` (`processing/utilities.py`) | + +Cache roots (`data_processor.py`): + +| Helper | Env override | Default | +| ----------------------------------------- | ---------------------------------- | ------------------------------------------------ | +| `_get_cache_dir` (chunks) | `DATA_OPTIMIZER_CACHE_FOLDER` | Studio: `/cache/chunks`; else `{tempdir}/chunks` | +| `_get_cache_data_dir` (downloaded inputs) | `DATA_OPTIMIZER_DATA_CACHE_FOLDER` | Studio: `/cache/data`; else `{tempdir}/data` | + +`DataProcessor._cleanup_cache` **rmtrees both** at the start of each `run()` so prior runs cannot poison the job. + +______________________________________________________________________ + +## 2. Path resolution → when remote I/O kicks in + +Canonical tables: [resolver.md](resolver.md). Processing always goes through `_resolve_dir` in `DataProcessor.__init__` and in `optimize`/`map` before constructing the processor. + +### 2.1 `Dir` fields that control movement + +```python +@dataclass +class Dir: + path: str | None # local / FUSE mount path (identity, cache rewrite base) + url: str | None # cloud URL used for FsProvider download/upload + data_connection_id: str | None # temp creds for some Studio connections +``` + +| Situation | `path` | `url` | Processing download behavior | +| ---------------------------------------------------------------------------------------------------------------- | -------------- | ----------------------------------- | -------------------------------------------------------------------------------------------- | +| Plain local dir | abs path | `None` | No cloud download; may `shutil.copyfile` into data cache if path is outside `this_studio` | +| Direct `s3://` / `gs://` / `r2://` | `None` | cloud URL | **Downloader procs skip** (`no_downloaders` when `input_dir.path is None`) — see §3.1 caveat | +| Studio FUSE: `/teamspace/s3_connections/…`, `s3_folders`, `gcs_*`, `lightning_storage`, `datasets`, other studio | FUSE path | backing `s3://` / `gs://` / `r2://` | Downloaders rewrite FUSE→URL and `FsProvider.download_file` | +| `/teamspace/studios/this_studio/…` | workspace path | `None` | Local; LitData does not invent a bucket URL | + +**Agent rule (same as raw streaming):** pass `/teamspace/s3_connections/…` or `s3://…` into LitData. Do **not** train or bulk-copy through FUSE with bare `open()` / `cp`. Resolver + FsProvider talk to the object store directly. + +### 2.2 Studio mount → URL (resolver functions) + +| Mount prefix | Resolver | Typical `url` | +| --------------------------------------- | ---------------------------- | --------------------------------------------------- | +| `/teamspace/s3_connections//…` | `_resolve_s3_connections` | customer S3 (`data_connection.aws.source` + suffix) | +| `/teamspace/s3_folders//…` | `_resolve_s3_folders` | S3 folder connection source + suffix | +| `/teamspace/gcs_connections//…` | `_resolve_gcs_connections` | `gs://…` | +| `/teamspace/gcs_folders//…` | `_resolve_gcs_folders` | `gs://…` | +| `/teamspace/lightning_storage//…` | `_resolve_lightning_storage` | `r2://…` + **always** `data_connection_id` | +| `/teamspace/datasets/…` | `_resolve_datasets` | cluster datasets S3 | +| `/teamspace/studios//…` | `_resolve_studio` | studio content `s3://` or `gs://` | + +Connection name = path segment `[3]`. Credentials: ambient cloud keys, or temp project-role creds when `data_connection_id` is set (`streaming/client.py`). + +### 2.3 How `input_dir` is inferred + +`_get_input_dir(inputs)` (`functions.py`): + +1. Flatten first (or second) input; find filepath-like strings (`_get_indexed_paths` / `_is_remote_file`). +2. Remote scheme → `os.path.dirname(path)` (e.g. `s3://bucket/prefix`). +3. Studio / `/teamspace…` → keep first **four** path segments as root (e.g. `/teamspace/s3_connections/my-conn`). +4. Else → `None` (no shared input root; often `no_downloaders`). + +Workers detect per-item paths with `_is_path(input_dir.path, element)` / `_to_path` — heuristic; items with **zero** paths raise in `_collect_paths`. + +______________________________________________________________________ + +## 3. Processing downloaders (`_download_data_target`) + +### 3.1 When they start + +`BaseWorker.no_downloaders = (input_dir.path is None) or (reader is not None)`. + +- **`no_downloaders` True** → `_start_downloaders` returns immediately; items go straight to `ready_to_process_queue` (or `FakeQueue` when ordered + no downloaders). +- **Important:** pure `s3://…` input resolves to `Dir(path=None, url=…)`, so **`path is None` ⇒ downloaders do not run**. Background download assumes a local/FUSE `path` to rewrite into `cache_data_dir`. For Studio connections, `path` is set (FUSE) **and** `url` is set — that is the path where downloaders matter. +- Custom `reader` (e.g. `StreamingDataLoaderReader`) also disables downloaders; the reader supplies bytes. + +Defaults: `num_downloaders or 2` processes per worker. Each is `multiprocessing.Process(target=_download_data_target, args=(input_dir, cache_data_dir, to_download_queue, ready_to_process_queue, storage_options))`. + +### 3.2 Queue protocol + +1. `_collect_paths` builds `self.paths` and rewrites flattened filepath leaves under `input_dir.path` → `cache_data_dir` (unless path starts with `/teamspace/studios/this_studio`). +2. `_start_downloaders` enqueues `(index, item, paths)` round-robin across `to_download_queues`, then sends `None` sentinel per downloader. +3. Downloader loop: `queue_in.get()` → download/copy → `queue_out.put((index, item, paths))`; on `None`, put `None` and exit. +4. Worker `_loop` consumes `ready_to_process_queue` and runs the recipe. + +### 3.3 Per-path download logic (`_download_data_target`, ~`:128`) + +For each path in the item: + +1. If all paths already exist under the cache rewrite → skip download, forward tuple. +2. If `input_dir.url` is set → `_wait_for_disk_usage_higher_than_threshold("/", 25)` (wait until **>25 GB free** on `/`) so removers can catch up under pressure. +3. Local cache target: `path.replace(input_dir.path, cache_dir)`. +4. If `url` and `path` and the FUSE/local file is **missing**: rewrite path with `path.replace(input_dir.path, input_dir.url)` → cloud URL. +5. If `urlparse(path).scheme in _SUPPORTED_PROVIDERS` (`s3`/`gs`/`r2`): + - `construct_storage_options(storage_options, input_dir)` (injects `data_connection_id`) + - `_get_fs_provider(input_dir.url, …).download_file(remote, local_path)` +6. Elif `os.path.isfile(path)` and not under `this_studio`: `shutil.copyfile` into cache. +7. Else: `ValueError` unsupported URL. + +**Local vs remote summary:** + +| Input | Action | +| ------------------------------------- | ----------------------------------------------------------------------------- | +| FUSE connection + missing local file | Resolve to `url`, FsProvider download into `DATA_OPTIMIZER_DATA_CACHE_FOLDER` | +| Already cached under `cache_data_dir` | No-op, pass through | +| Real local file outside `this_studio` | Copy into data cache | +| `this_studio` local | Leave path as-is (no copy into cache for that prefix) | +| Unsupported scheme | Raise | + +### 3.4 Interaction with user `fn` + +After download, item tree leaves point at **cache_data_dir** paths (rewritten in `_collect_paths`). `prepare_item` / user `fn` should open those local paths — not the original FUSE or `s3://` strings — when downloaders ran. + +______________________________________________________________________ + +## 4. Processing uploaders (`_upload_fn`) + +### 4.1 When they start + +`_start_uploaders` runs unless **both** `output_dir.path` and `output_dir.url` are `None`. + +- Remote `output_dir.url` with scheme in `_SUPPORTED_PROVIDERS` → `FsProvider.upload_file`. +- Local `output_dir.path` → `shutil.copy` into destination (makedirs as needed). +- Else → `ValueError`. + +Default `num_uploaders or 1` per worker. + +### 4.2 Who enqueues uploads + +| Recipe | What gets uploaded | +| ---------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| **`optimize` / `DataChunkRecipe`** | Each closed chunk filepath from `Cache._add_item` / `cache.done()`; optional checkpoint JSON under `.checkpoints` when `use_checkpoint` | +| **`map` / `MapRecipe`** | Every file under a per-item `tempfile.mkdtemp()` after `prepare_item` (user writes into that dir); uploaded as `(tmpdir, filepath)` so relative layout is preserved | + +`_try_upload` no-ops if output_dir has neither path nor url, or data is empty/missing on disk. Round-robins across `to_upload_queues`. + +On worker shutdown (downloaders finished / timeout / `ALL_DONE`): send `None` to each uploader and `join`. + +### 4.3 Upload path construction (`_upload_fn`, ~`:232`) + +1. Ensure `local_filepath` is under `cache_chunks_dir` (join if relative). +2. Remote destination: + - Base = `output_dir.url` + - If path contains `.checkpoints` → nest under `…/.checkpoints` + - Basename-only upload for optimize chunks; map preserves relative path under tmpdir + - `remove_uuid_from_filename` strips UUID from checkpoint names → `checkpoint-.json` +3. After successful upload (or local copy): if `remove_queue` and file exists → `remove_queue.put([local_filepath])` so the **remover deletes the local chunk** after upload. + +### 4.4 Index upload (not the uploader pool) + +Chunk uploaders do **not** own the final index. After all workers finish, `DataChunkRecipe._done` → `_merge_no_wait` → `_upload_index`: + +- Single-node: upload `index.json` from cache to `output_dir`. +- Multi-node: each node uploads `{node_rank}-index.json`; **last node** downloads peers’ indexes, merges, uploads final `index.json` — [multi-node.md](multi-node.md). + +`MapRecipe._done` does not merge LitData chunk indexes (map is side-effect files only). + +______________________________________________________________________ + +## 5. Processing removers (`_remove_target`) + +### 5.1 When they start + +`_start_remover` only if `self.remove` is True. That flag is `DataProcessor.delete_cached_files` (default **True**), passed positionally into `DataWorkerProcess` / `BaseWorker`. + +### 5.2 What gets deleted + +Two producers feed `remove_queue`: + +1. **After each item** (if `remove` and `input_dir.path` and no `reader`): worker puts the item’s **source paths** (original path list) so cached downloads under `cache_data_dir` can be freed. +2. **After each upload**: uploader puts the **local chunk/file** path. + +`_remove_target` (~`:190`): + +- Rewrite paths from `input_dir.path` → `cache_dir` when needed. +- `os.remove` if exists. +- If `input_dir` is falsy: only delete if `keep_path(path)` is True — **refuses** to delete paths containing Studio mount tokens: `s3_connections`, `s3_folders`, `gcs_connections`, `efs_*`, `lightning_storage`, `snowflake_connections` (safety against wiping FUSE mounts). + +On shutdown with `remove`: put `None` sentinel and join remover. + +### 5.3 Post-run check + +`DataChunkRecipe._done`: if `delete_cached_files` and **local** `output_dir.path` is set and any `.bin` still remain in chunk cache → `RuntimeError` (“All the chunks should have been deleted”). Remote outputs rely on uploaders + remover; local outputs expect chunks to have been copied away and removed. + +### 5.4 What removers do **not** do + +- They do **not** delete remote objects. +- They do **not** delete the durable dataset under `output_dir` (except separate overwrite/checkpoint cleanup via FsProvider in `_cleanup_checkpoints` / resolver immutability helpers). +- `DataProcessor._cleanup_cache` wipes entire cache dirs at **start** of the next run, independent of the remover process. + +______________________________________________________________________ + +## 6. Streaming / raw `Downloader` (read path) + +Used when **reading** optimized chunks or raw files — not the optimize worker pool. + +| Piece | Role | +| -------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------- | +| `get_downloader(remote_dir, cache_dir, chunks, storage_options, session_options)` | Prefix match on `_DOWNLOADERS` | +| `Downloader.download_file` / `download_bytes` / `adownload_file` / `adownload_fileobj` | Sync + async APIs | +| Atomic publish | `_temp_download_path` + `_atomic_replace` (tmp includes pid) | +| `register_downloader` / `unregister_downloader` | Extension points | +| `StreamingRawDataset.downloader` | Uses same registry; prefer cloud URL / connection path over FUSE ([using-litdata.md](using-litdata.md) §10) | +| `async_prefetch.py` | Prefers `adownload_file` when overridden | + +**FsProvider vs Downloader** (also [storage-format.md](storage-format.md) §5): + +| | FsProvider | Downloader | +| | \---------- | ---------- | +| Optimize input download / chunk upload / index / merge / empty checks | ✅ | ❌ | +| StreamingDataset chunk prefetch / StreamingRawDataset | ❌ | ✅ | +| Schemes | s3, gs, r2 | + azure, hf, local | + +______________________________________________________________________ + +## 7. Local vs remote — decision table for agents + +| Goal | Prefer | What LitData does | +| -------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------- | +| Optimize files on Studio S3 connection | `input_dir` / paths under `/teamspace/s3_connections/…` | Resolve → downloaders + FsProvider GET into `/cache/data` | +| Optimize from laptop with AWS creds | `s3://bucket/…` in inputs; may need design that doesn’t rely on `path`-based downloaders — verify whether your inputs are local copies or you read via SDK inside `fn` | Pure `s3://` `Dir` has `path=None` → **no** `_download_data_target` pool | +| Write durable chunks | `output_dir=/teamspace/s3_connections/…/vN` or `s3://…` | Uploaders + `_upload_index` via FsProvider | +| Scratch only | local / `this_studio` (small) | Local copy uploaders; multi-node remaps `this_studio` optimize outs to job artifacts ([multi-node.md](multi-node.md)) | +| Raw training I/O | `StreamingRawDataset("s3://…")` or connection path | `Downloader` async; **not** processing downloaders | + +______________________________________________________________________ + +## 8. Error modes & agent checklists + +| Symptom | Likely cause | What to check | +| ------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------- | +| `ValueError: The provided … isn't supported` in downloader/uploader | Scheme outside `_SUPPORTED_PROVIDERS` for processing | Use s3/gs/r2 for optimize I/O; azure/hf are streaming-Downloader-only | +| Auth / 403 on download or upload | Missing keys; RO bucket; connection without write; missing `data_connection_id` for R2 | `storage_options`, Studio connection attach, IAM | +| Hang with remote inputs | Disk wait (`_wait_for_disk_usage_higher_than_threshold` 25 GB); remover stuck; uploader exception only `print`ed | Free space on `/`; `num_workers=1`; watch uploader `print(e)` | +| `The provided item … didn't contain any filepaths` | `_collect_paths` / `_is_path` failed | Pass real paths under `input_dir.path`; set `input_dir` explicitly | +| Chunks left / RuntimeError in `_done` | Uploader failed or `delete_cached_files` + local output mismatch | Inspect cache dirs; uploader errors | +| Index never appears (multi-node) | Last node waiting on peer `{rank}-index.json` | [multi-node.md](multi-node.md) peer wait | +| FUSE “works” in `ls` but training/optimize is slow or crashes | Reading mount directly | Pass path into LitData; confirm `Dir.url` is set | +| Partial / corrupt local file | Crash mid-download (FsProvider path is not always atomic the way `Downloader._atomic_replace` is) | Wipe `DATA_OPTIMIZER_DATA_CACHE_FOLDER` / re-run; prefer connection+resolver path | +| `cloudspaces` in output URL | Rejected in `optimize`/`map` | Use connections / datasets, not studio content URLs | + +**Debug tip:** `num_workers=1`, `fast_dev_run=True`, and inspect `/cache/data` + `/cache/chunks` (or temp equivalents). Worker exceptions land in `error_queue` → main `RuntimeError` + `terminate()` siblings. + +______________________________________________________________________ + +## 9. Grep landmarks + +``` +processing/data_processor.py + _download_data_target _upload_fn _remove_target keep_path + BaseWorker._collect_paths _start_downloaders _start_uploaders _start_remover + _try_upload no_downloaders delete_cached_files + DataChunkRecipe._done _upload_index + +processing/functions.py + optimize map _get_input_dir _resolve_dir(...) + +processing/utilities.py + construct_storage_options remove_uuid_from_filename _get_work_dir + read_index_file_content + +streaming/resolver.py + _resolve_dir _resolve_s3_connections _resolve_s3_folders + _resolve_gcs_* _resolve_lightning_storage _resolve_datasets _execute + +streaming/fs_provider.py + FsProvider _get_fs_provider + +streaming/downloader.py + Downloader get_downloader _DOWNLOADERS register_downloader + +constants.py + _SUPPORTED_PROVIDERS +``` diff --git a/.claude/skills/litdata/reference/lightning-studio.md b/.claude/skills/litdata/reference/lightning-studio.md index f747fe47b..0374ee7bf 100644 --- a/.claude/skills/litdata/reference/lightning-studio.md +++ b/.claude/skills/litdata/reference/lightning-studio.md @@ -144,7 +144,7 @@ Passing `num_nodes=N` (optionally `machine=Machine.DATA_PREP`) from a Studio: **Output tip:** write to `/teamspace/s3_connections/...`, `/teamspace/datasets/...`, or `s3://...` so results land in a durable bucket. Optimize may remap `/teamspace/studios/this_studio/...` outputs to the job artifacts S3 URL; the Studio UI may expose them under `/teamspace/jobs/...`. Ensure **every** node can read inputs and write outputs (attached connections or cloud credentials). -Full launch/env/sharding → [processing.md](processing.md). User recipe → [using-litdata.md](using-litdata.md) §9. +Full launch/env/sharding/index merge → [multi-node.md](multi-node.md). Optimize I/O (downloaders/uploaders) → [data-movement.md](data-movement.md). Orchestration overview → [processing.md](processing.md). User recipe → [using-litdata.md](using-litdata.md) §9. ## Quick Studio smoke diff --git a/.claude/skills/litdata/reference/multi-node.md b/.claude/skills/litdata/reference/multi-node.md new file mode 100644 index 000000000..1bd489d6c --- /dev/null +++ b/.claude/skills/litdata/reference/multi-node.md @@ -0,0 +1,328 @@ +# Multi-node processing (`optimize` / `map`) + +How LitData fans **`optimize`** and **`map`** across machines. This is **Lightning Studio data-prep jobs**, not PyTorch DDP training and not a built-in SLURM launcher. + +**Related:** [processing.md](processing.md) · [data-movement.md](data-movement.md) · [resolver.md](resolver.md) · [lightning-studio.md](lightning-studio.md) · user cookbook [using-litdata.md](using-litdata.md) §9 + +______________________________________________________________________ + +## 0. What “multi-node” means here (and what it is not) + +| Mechanism | Used for | Symbols / env | +| --------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------- | +| **Lightning Studio job** (`num_nodes=N`) | Distributed **optimize/map** | `functions.py` gate → `resolver._execute` → platform sets `DATA_OPTIMIZER_*` | +| **`DATA_OPTIMIZER_*` env** | Rank / world inside each job instance | `_get_num_nodes`, `_get_node_rank`, worker `DATA_OPTIMIZER_GLOBAL_RANK` | +| **`broadcast_object`** | Align dirs when `broadcast_paths` is on (or auto for `{%strftime}` paths) and Lightning app URL present | `utilities/broadcast.py` | +| **Torch distributed / `WORLD_SIZE` / `GLOBAL_RANK` / `NNODES`** | **Training** stream path (`_DistributedEnv.detect`) — **not** how optimize jobs are launched | `utilities/env.py` | +| **SLURM** | **Not** a first-class optimize launcher in this repo | Do not document SLURM as supported for `num_nodes` | + +If `num_nodes` / `machine` are set **outside** Studio (`_IS_IN_STUDIO` false) → `ValueError` (“Only https://lightning.ai/ supports multiple nodes…”). + +Training-time multi-GPU/node streaming (chunk shuffle, sampler) is a **separate** concern — see [streaming.md](streaming.md) / `_DistributedEnv`. Below is **write-path** multi-node only. + +______________________________________________________________________ + +## 1. Launch dual-path (`optimize` / `map`) + +Both APIs share the same gate (`functions.py`): + +``` +if num_nodes is None OR int(DATA_OPTIMIZER_NUM_NODES) > 0: + → construct DataProcessor and run locally on THIS machine +else: + → _execute(...) # resolver.py — create Studio data-prep job, block until done +``` + +### 1.1 Caller Studio (launcher) + +1. User: `optimize(..., num_nodes=N, machine=Machine.DATA_PREP)` (or `map` with same knobs). +2. Gate sees `num_nodes` set and `DATA_OPTIMIZER_NUM_NODES` unset/0 → **`_execute`**. +3. `_execute` (`resolver.py:461`): + - Requires `lightning_sdk` (`_LIGHTNING_SDK_AVAILABLE`). + - `Studio()._studio_api.create_data_prep_machine_job(...)` with: + - `command`: `cd {cwd} &&[LIGHTNING_SKIP_INSTALL=…][LIGHTNING_BRANCH=…] python {' '.join(sys.argv)}` (re-runs the **same script/args**) + - `num_instances=num_nodes` + - `machine=machine or current Studio machine` + - `interruptible=` exists on `_execute` but **optimize/map never pass it** → always default `False` + - Prints job URL (`…/app?app_id=litdata&app_tab=Runs&job_name=…`). + - Polls until `STOPPED` / `COMPLETED`, or raises on `FAILED`. + +### 1.2 Each job instance (worker machine) + +Platform injects env (see §2). Script starts again; gate now sees `DATA_OPTIMIZER_NUM_NODES > 0` → **local** `DataProcessor.run` on that instance’s shard only. No cross-node RPC for work items — pure env + shared object store for indexes/outputs. + +______________________________________________________________________ + +## 2. Environment variables + +### 2.1 Processing ranks (optimize/map workers) + +| Variable | Reader | Role | +| ----------------------------------------- | --------------------------------------------------------------------- | ------------------------------------------------------------------------------------ | +| `DATA_OPTIMIZER_NUM_NODES` | `_get_num_nodes()`; launch gate | World of machines. `>0` means “already inside a distributed job / run DataProcessor” | +| `DATA_OPTIMIZER_NODE_RANK` | `_get_node_rank()` | This machine’s rank in `[0, num_nodes)` | +| `DATA_OPTIMIZER_GLOBAL_RANK` | Set in `BaseWorker._set_environ_variables` | `node_rank * num_workers + worker_index` — used for chunk filenames / writer rank | +| `DATA_OPTIMIZER_NUM_WORKERS` | Set in worker; also `_DistributedEnv._instantiate_in_map_or_optimize` | Local worker count | +| `DATA_OPTIMIZER_CACHE_FOLDER` | `_get_cache_dir` | Chunk cache root (default Studio `/cache/chunks`) | +| `DATA_OPTIMIZER_DATA_CACHE_FOLDER` | `_get_cache_data_dir` | Downloaded input cache (default `/cache/data`) | +| `DATA_OPTIMIZER_TIMEOUT` | Worker `_loop` queue get | Default 300s; shared-queue mode often 200s | +| `DATA_OPTIMIZER_FAST_DEV_RUN` | `_get_fast_dev_run` | Related to fast_dev_run defaults | +| `ENABLE_STATUS_REPORT` / `_ENABLE_STATUS` | Progress | Node 0 may write `status.json` with coarse % | + +### 2.2 Job / artifacts (Studio) + +| Variable | Role | +| ---------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------- | +| `LIGHTNING_SKIP_INSTALL` / `LIGHTNING_BRANCH` | Injected into remote job command string | +| `LIGHTNING_BUCKET_NAME`, `LIGHTNING_CLOUD_PROJECT_ID`, `LIGHTNING_CLOUD_APP_ID`, `LIGHTNING_CLOUD_WORK_ID` | `_get_work_dir()` → artifacts `s3://…/artifacts/{work_id}/content/` | +| `LIGHTNING_CLOUD_URL` | Job URL pretty-print; auth helpers | +| `LIGHTNING_APP_EXTERNAL_URL` | If set, `broadcast_object` uses Lightning broadcast HTTP API | + +### 2.3 Training distributed env (do not confuse) + +`_DistributedEnv.detect()` (`utilities/env.py`): + +- Inside map/optimize workers (`DATA_OPTIMIZER_GLOBAL_RANK` set): builds env from `DATA_OPTIMIZER_*` (`world_size = num_workers * num_nodes`). +- Else: `torch.distributed` if initialized, else `WORLD_SIZE` / `GLOBAL_RANK` / `NNODES`. + +That path drives **streaming** writer/reader rank math when those modules run under optimize; it is **not** a SLURM or torchrun launcher for `num_nodes=`. + +______________________________________________________________________ + +## 3. Work sharding + +``` +world_size = num_nodes * num_workers +``` + +Both `_map_items_to_workers_sequentially` and `_map_items_to_workers_weighted` (`data_processor.py`) pack across **all** ranks, then each node **keeps only** worker ids in: + +``` +[node_rank * num_workers, (node_rank + 1) * num_workers) +``` + +### 3.1 Sequential (`reorder_files=False` or no `input_dir.path` for size packing) + +- Split `user_items` into `world_size` contiguous slices (remainder distributed from the end). +- With `align_chunking=True` + `chunk_size`: assign full chunks of size `chunk_size` per global worker; last worker gets the tail (may be uneven — intentional). + +### 3.2 Weighted / by file size (default when `reorder_files` and `input_dir.path`) + +- `_get_item_filesizes` (threaded `os.path.getsize` — **local/FUSE path based**; TODO in code notes broadcasting sizes from node 0). +- `_pack_greedily` into `world_size` bins; permute within each worker’s list (`np.random.permutation` — seed fixed in `DataProcessor.run`). +- Explicit `weights=` uses same packer with `file_size=False`. + +### 3.3 Queue / shared-queue modes + +- Input `multiprocessing.Queue` or `keep_data_ordered=False`: dynamic consumption; multi-node semantics are weaker / different — checkpointing **unsupported** for Queue inputs. Prefer static list inputs for multi-node jobs. +- `ALL_DONE` sentinel for shared-queue shutdown (`keep_data_ordered=False`). + +### 3.4 Broadcast dirs (`broadcast_paths`) + +`optimize` / `map` / `DataProcessor` take **`broadcast_paths: bool = False`**. + +After resolve, broadcast runs **only when** `broadcast_paths` is effectively on: + +```python +# Auto-on if input_dir or output_dir contains a `{%strftime}` template (detected before resolve). +self.broadcast_paths = broadcast_paths or _has_time_template(input_dir) or _has_time_template(output_dir) + +if self.broadcast_paths: + self.input_dir = broadcast_object("input_dir", self.input_dir, rank=_get_node_rank()) + self.output_dir = broadcast_object("output_dir", self.output_dir, rank=_get_node_rank()) +``` + +| Case | Behavior | +| ------------------------------------ | ------------------------------------------------------------------------------------------------------------ | +| Default (`False`), no `{%…}` in path | **Skip** broadcast — each rank keeps its locally resolved `Dir` (fine for stable `s3://` / connection paths) | +| Path has `{%Y-%m-%d}` (etc.) | **Auto-enable** — ranks must share one expanded timestamp | +| `broadcast_paths=True` | Always broadcast after resolve | + +- If `LIGHTNING_APP_EXTERNAL_URL` is set and broadcast runs: HTTP broadcast until all ranks agree. +- Else `broadcast_object` returns the local `obj` unchanged. +- Multi-node implication when off: ranks must independently resolve to the **same** paths; do not rely on per-rank `datetime.now()` without a shared template + auto-broadcast. + +______________________________________________________________________ + +## 4. Per-node cache and downloads + +- Each node has **its own** `/cache/chunks` and `/cache/data` (or env overrides). **Not** a shared NFS assumption. +- `DataProcessor._cleanup_cache` wipes both at **start** of each node’s run. +- Downloaders/uploaders/removers run **per worker on that node** — see [data-movement.md](data-movement.md). +- Every node needs credentials for **inputs and outputs** (attached connections or keys on all instances). Missing creds on node K → that shard fails; last-node index merge may hang waiting for `{K}-index.json`. + +**FUSE vs direct:** same as single-node — pass `/teamspace/s3_connections/…` so each node resolves to `Dir.url` and downloads via FsProvider, not through FUSE under multi-worker load. + +______________________________________________________________________ + +## 5. Chunk write + upload coordination + +### 5.1 Filenames and ranks + +`BinaryWriter.get_chunk_filename` → `chunk-{rank}-{chunk_index}[.compression].bin` where `rank` is `DATA_OPTIMIZER_GLOBAL_RANK` when set (`writer.py`). + +So **all workers on all nodes** write uniquely named chunks. No “only rank 0 uploads chunks.” + +### 5.2 Who uploads chunks? + +**Every worker’s uploader pool** uploads its own closed chunks to `output_dir` as they are produced (`_try_upload` → `_upload_fn`). There is no barrier that waits for other nodes before uploading bins. + +### 5.3 Who builds vs merges the index? + +After local workers finish, `DataChunkRecipe._done` (`data_processor.py`): + +1. `Cache(...)._merge_no_wait(node_rank if num_nodes > 1 else None, existing_index?)` + - Single-node: merge all per-worker `*-index.json` / rank indexes in cache → `index.json`. + - Multi-node: merge **this node’s** worker indexes → **`{node_rank}-index.json`** (see `BinaryWriter._merge_no_wait`). +2. `_upload_index(output_dir, cache_dir, num_nodes, node_rank)`: + - Upload that node’s index file (`index.json` or `{node_rank}-index.json`) via FsProvider or local copy. + - **If `num_nodes == node_rank + 1` (last node):** + 1. For each peer `0 .. num_nodes-2`: `_wait_for_file_to_exist` on `{peer}-index.json` at `output_dir`, then download into local cache. + 2. `merge_cache._merge_no_wait()` → final `index.json`. + 3. Recurse `_upload_index(..., num_nodes=1, node_rank=None)` to upload the merged index. + +Comment in code: under the Data Optimizer there should be a **single process per node** executing this merge section → no local race on last node. Cross-node coordination is **object-store presence** of peer index files, not a distributed lock. + +### 5.4 Map jobs + +`map` uses the same launch/shard/download/upload worker pools but **no** LitData chunk `index.json` merge. Outputs are whatever files `fn` wrote that uploaders copied. Last-node index logic is optimize-specific (`DataChunkRecipe`). + +______________________________________________________________________ + +## 6. Output directory rules under multi-node + +**Prefer durable remote:** + +- `/teamspace/s3_connections/…`, `/teamspace/datasets/…`, or `s3://…` / `gs://…` / `r2://…` + +**`this_studio` remap (optimize only):** if `output_dir` resolves to a path under `/teamspace/studios/this_studio` **and** `DATA_OPTIMIZER_NUM_NODES > 0`, optimize rewrites to: + +``` +_get_work_dir() + relative_path +→ s3://{LIGHTNING_BUCKET_NAME}/projects/.../artifacts/{work_id}/content/... +``` + +(`functions.py` ~515–524). **`map` does not apply this remap.** Studio UI may also show job mounts under `/teamspace/jobs/…` — LitData does not construct that string itself. + +**Rejected:** resolved URL containing `cloudspaces` → `ValueError` with hint to use connections/datasets. + +**Immutability:** `_assert_dir_has_index_file` / empty checks on remote before write; use `mode="append"|"overwrite"` or versioned prefixes. + +______________________________________________________________________ + +## 7. Resume, checkpoint, append + +### 7.1 `use_checkpoint=True` (optimize) + +- Unsupported for Queue inputs and generator `fn`s (`DataProcessor.run`). +- Saves `.checkpoints/config.json` (`num_workers`, `workers_user_items`) and per-worker checkpoint JSON. +- Writer saves `checkpoint-{rank}-{uuid}.json`; upload strips UUID → `checkpoint-{rank}.json` (`remove_uuid_from_filename`). +- On resume: `_load_checkpoint_config` requires **same** `num_workers` and **identical** `workers_user_items`; trims each worker list to `done_till_index`. +- Remote: download `.checkpoints/` via `FsProvider.download_directory`. +- When **not** using checkpoints, `run` calls `_cleanup_checkpoints` (local rmtree or remote delete of `.checkpoints/`). Successful completion with checkpoints also cleans up at end. + +**Multi-node caveat:** checkpoint compatibility is tied to **local** `workers_user_items` after node slicing. Changing `num_nodes` / `num_workers` / input order breaks resume. Verify behavior in code if mixing checkpoint + multi-node — treat as advanced. + +### 7.2 `mode="append"` + +- Reads existing `index.json` (`read_index_file_content`). +- Advances per-rank `state_dict` from existing `chunk--.bin` filenames so new chunks continue indexes. +- `existing_index` passed into recipe merge so final index concatenates old + new chunks. + +### 7.3 `mode="overwrite"` + +- Resolver / assert helpers delete prior index/chunks (and checkpoints depending on flags) before write. + +______________________________________________________________________ + +## 8. Progress and status + +- Main process aggregates `progress_queue` updates (per-worker counters). tqdm total is **this node’s** item count, not global (status.json tries to scale by `num_nodes` when enabled). +- Node 0 + `_ENABLE_STATUS`: writes `status.json` with `"progress": "{percent}%"`. +- Worker log lines via `msg_queue` → `flush_msg_queue` to avoid breaking tqdm. + +______________________________________________________________________ + +## 9. Failure modes & pitfalls (agent checklist) + +| Pitfall | What happens | Mitigation | +| -------------------------------------- | ---------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------- | +| Peer never writes `{k}-index.json` | Last node **hangs** in `_wait_for_file_to_exist` | Ensure all nodes finish; same `output_dir`; credentials on every node; check failed job instances | +| Duplicate work | Mis-set `DATA_OPTIMIZER_NODE_RANK` / all nodes think they are 0 | Trust platform env; don’t manually override inconsistently | +| NFS shared cache dir | Nodes stomp `/cache` if incorrectly shared | Keep node-local caches; use object store for outputs | +| Reading FUSE for size packing | `_get_item_filesizes` hits mount; slow/wrong under load | Prefer connection paths that exist as files for size, or pass `weights=`; accept TODO that sizes aren’t broadcast from node 0 | +| Only node 0 has AWS keys | Other nodes fail downloads/uploads | Attach connection / inject creds on all instances | +| `keep_data_ordered=False` + multi-node | Shared queue is **per node**, not global | Prefer ordered static sharding for multi-node | +| Index config mismatch across workers | `_merge_no_wait` raises inconsistent `config` | Same `fn` / serializers / compression on all workers | +| Local `output_dir` on multi-node | Machines don’t share disk; merge/upload logic expects reachable `output_dir` | Use remote `output_dir` | +| Assuming torchrun/SLURM | `num_nodes=` outside Studio errors | Use Studio jobs or run your own process manager **and** set `DATA_OPTIMIZER_*` yourself (unsupported DIY — verify carefully) | +| `map` + `this_studio` output | No artifacts remap | Write to connection / `s3://` | +| Uploader swallows exceptions | `_upload_fn` `print(e)` then may still signal remover | Watch logs; missing remote chunks + stuck index | +| Hard kill on worker error | `_exit_on_error` → `terminate()` all local workers | Fix root error with `num_workers=1` / `fast_dev_run` | + +______________________________________________________________________ + +## 10. Mental model diagram + +``` +Studio caller Job instances (NODE 0 .. N-1) +----------------- ---------------------------- +optimize(num_nodes=N) + | + v + _execute (create job, print URL) + | env: DATA_OPTIMIZER_NUM_NODES / NODE_RANK + | for each instance: resolve dirs -> shard items for this node + | downloaders -> fn -> chunk-{global_rank}-*.bin + | uploaders -> output_dir (all ranks upload) + | _done -> upload {node_rank}-index.json + | LAST NODE (node_rank == num_nodes-1): + | wait for peers' {k}-index.json + | merge -> index.json -> upload + v + block until job COMPLETE / FAILED +``` + +______________________________________________________________________ + +## 11. Grep landmarks + +``` +processing/functions.py optimize/map num_nodes gate; this_studio remap +streaming/resolver.py _execute +processing/data_processor.py + _get_num_nodes _get_node_rank + _map_items_to_workers_sequentially _map_items_to_workers_weighted + DataChunkRecipe._done _upload_index + broadcast_object(...) _cleanup_cache _load_checkpoint_config +processing/utilities.py _get_work_dir extract_rank_and_index_from_filename +streaming/writer.py chunk-{rank}-*.bin _merge_no_wait +utilities/broadcast.py broadcast_object +utilities/env.py _DistributedEnv _is_in_map_or_optimize +``` + +## 12. Minimal multi-node recipe (Studio) + +```python +import litdata as ld + +def fn(path): + # read local cached path when downloaders ran; return sample pytree + ... + +if __name__ == "__main__": + ld.optimize( + fn=fn, + inputs=list_of_paths_under_connection, # or walk(...) + input_dir="/teamspace/s3_connections/my-raw", + output_dir="/teamspace/s3_connections/my-opt/v1", # durable, shared + chunk_bytes="64MB", + num_workers=8, + num_nodes=4, # Studio only + # machine=Machine.DATA_PREP, + num_downloaders=2, + num_uploaders=1, + ) +``` + +Ensure the script is restartable with the same args (job re-invokes `sys.argv`). Prefer versioned `output_dir` prefixes; do not rely on FUSE for bulk I/O. diff --git a/.claude/skills/litdata/reference/processing.md b/.claude/skills/litdata/reference/processing.md index e08f2c922..c8f79a8e2 100644 --- a/.claude/skills/litdata/reference/processing.md +++ b/.claude/skills/litdata/reference/processing.md @@ -1,58 +1,30 @@ # The processing (write) pipeline — `optimize` / `map` -All paths under `src/litdata/`. This pipeline fans work across workers (and machines) to transform raw data, and for `optimize` writes it into the litdata chunk format that the streaming pipeline reads. Chunk / `index.json` / `BinaryWriter` / `FsProvider` details → [storage-format.md](storage-format.md). +All paths under `src/litdata/`. This pipeline fans work across workers (and machines) to transform raw data, and for `optimize` writes it into the litdata chunk format that the streaming pipeline reads. -## Public API (`processing/functions.py`) - -User-facing arg tables → [using-litdata.md](using-litdata.md) §9 and README `#optimize-kwargs` / `#map` / `#walk`. - -- **`optimize(...)`** — `functions.py:387`. Runs `fn` per input; flatten via pytree → `chunk-*.bin` + `index.json`. **Exactly one of `chunk_size` / `chunk_bytes`.** Notable: `queue`+`ALL_DONE`, `align_chunking`, `use_checkpoint`, `mode="append"|"overwrite"`, `keep_data_ordered=False` (shared queue), `encryption`, `item_loader=TokensLoader()`, `weights`/`input_dir`, `num_nodes`/`machine`. → `LambdaDataChunkRecipe` / `QueueDataChunkRecipe` → `DataProcessor.run`. -- **`map(...)`** — `functions.py:242`. `fn(input, output_dir) -> None` (side effects only). Same worker/scale knobs + `error_when_not_empty`. → `LambdaMapRecipe`. -- **`merge_datasets(input_dirs, output_dir, max_workers=..., storage_options={})`** — `functions.py:675`. Copy chunks + concat `index.json`; matching `data_format`/compression required. -- **`walk(folder, max_workers=...)`** — `functions.py:621`. Threaded cloud `os.walk` (Studio-optimized); yield order is **not** depth-first. - -## Multi-node launch (`num_nodes` / `machine`) — read this first - -`num_nodes` is **not** local multiprocessing. It only works inside **Lightning Studio** (`_IS_IN_STUDIO`; else `ValueError` in `functions.py`). Dual path for both `map` and `optimize`: - -``` -if num_nodes is None OR DATA_OPTIMIZER_NUM_NODES > 0: - → run DataProcessor on this machine (single-node OR a job worker) -else: - → _execute(...) # resolver.py:461 — create Studio data-prep job, block until done -``` +**Load these when the task touches I/O or scale:** -1. User calls `optimize(..., num_nodes=N, machine=Machine.DATA_PREP)` on a Studio. -2. `_execute` starts a multi-instance job that re-runs `python {' '.join(sys.argv)}` on **N** machines (`resolver.py:483–492`). `machine=None` → current Studio machine. (`interruptible` exists on `_execute` but **optimize/map never pass it** — always `False`; do not document as a public knob.) -3. Platform injects `DATA_OPTIMIZER_NUM_NODES`, `DATA_OPTIMIZER_NODE_RANK`, etc. on each instance. -4. The same script hits the **local** branch (gate sees `DATA_OPTIMIZER_NUM_NODES > 0`) and each node processes only its shard. -5. Caller blocks until the job completes or fails (`FAILED` → `RuntimeError`). Job URL is printed to the Studio Runs UI. +| Topic | Doc | +| ------------------------------------------------------------------------------------ | ----------------------------------------------------- | +| Downloaders / uploaders / removers, FUSE→URL, cache dirs, FsProvider vs `Downloader` | **[data-movement.md](data-movement.md)** (exhaustive) | +| Multi-node Studio jobs, sharding, index merge, checkpoints, pitfalls | **[multi-node.md](multi-node.md)** (exhaustive) | +| Chunk / `index.json` / `BinaryWriter` / `FsProvider` | [storage-format.md](storage-format.md) | +| Path URI tables | [resolver.md](resolver.md) | -**Prefer durable `output_dir`:** `/teamspace/s3_connections/...`, `/teamspace/datasets/...`, or `s3://...`. -If optimize’s `output_dir` is under `/teamspace/studios/this_studio` **and** workers are multi-node, LitData rewrites it to the job artifacts bucket via `_get_work_dir()` (`functions.py:515–524` → `utilities.py:196–205` → `s3://{LIGHTNING_BUCKET_NAME}/projects/.../artifacts/{work_id}/content/...`). **`map` does not apply this remap.** Paths like `/teamspace/jobs/...` are the Studio job mount UI — LitData does not construct that string itself. Rejects outputs whose URL contains `cloudspaces` (use connections/datasets instead). - -### Env vars (multi-node / workers) +## Public API (`processing/functions.py`) -| Var | Role | -| ---------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- | -| `DATA_OPTIMIZER_NUM_NODES` | Launch gate + world size; `>0` means “I’m a job worker / already distributed” | -| `DATA_OPTIMIZER_NODE_RANK` | This node’s rank `[0, num_nodes)` | -| `DATA_OPTIMIZER_GLOBAL_RANK` / `DATA_OPTIMIZER_NUM_WORKERS` | Set inside workers for chunk filenames / writer rank | -| `DATA_OPTIMIZER_CACHE_FOLDER` / `DATA_CACHE_FOLDER` | Cache roots (Studio often `/cache/...`) | -| `DATA_OPTIMIZER_TIMEOUT` | Queue get timeout (default 300s; shared-queue ~200s) | -| `DATA_OPTIMIZER_FAST_DEV_RUN` | Related to `fast_dev_run` defaults | -| `LIGHTNING_SKIP_INSTALL` / `LIGHTNING_BRANCH` | Injected into remote job command | -| `LIGHTNING_BUCKET_NAME`, `LIGHTNING_CLOUD_PROJECT_ID`, `LIGHTNING_CLOUD_APP_ID`, `LIGHTNING_CLOUD_WORK_ID` | `_get_work_dir()` artifacts URL | -| `ENABLE_STATUS_REPORT` | Extra progress reporting | +User-facing arg tables → [using-litdata.md](using-litdata.md) §9 and README `#optimize-kwargs` / `#map` / `#walk`. -### Sharding & index merge (`data_processor.py`) +- **`optimize(...)`** — `functions.py` `optimize`. Runs `fn` per input; flatten via pytree → `chunk-*.bin` + `index.json`. **Exactly one of `chunk_size` / `chunk_bytes`.** Notable: `queue`+`ALL_DONE`, `align_chunking`, `use_checkpoint`, `mode="append"|"overwrite"`, `keep_data_ordered=False` (shared queue), `encryption`, `item_loader=TokensLoader()`, `weights`/`input_dir`, `num_nodes`/`machine`, `num_downloaders`/`num_uploaders`, `broadcast_paths=False` (auto-on for `{%strftime}` paths — see [multi-node.md](multi-node.md) §3.4). → `LambdaDataChunkRecipe` / `QueueDataChunkRecipe` → `DataProcessor.run`. +- **`map(...)`** — `functions.py` `map`. `fn(input, output_dir) -> None` (side effects only). Same worker/scale knobs + `error_when_not_empty` + `broadcast_paths`. → `LambdaMapRecipe`. +- **`merge_datasets(input_dirs, output_dir, max_workers=..., storage_options={})`** — Copy chunks + concat `index.json`; matching `data_format`/compression required. +- **`walk(folder, max_workers=...)`** — Threaded cloud `os.walk` (Studio-optimized); yield order is **not** depth-first. -- `world_size = num_nodes * num_workers`. Items packed across **all** ranks, then each node keeps only its worker slice. **No cross-node RPC** — pure env coordination. -- Each node writes per-rank chunk files + `{rank}-index.json` (node-local). -- **Last node** (`num_nodes == node_rank + 1`) waits for peer index files, merges into final `index.json`, and uploads. Peer wait can **hang** if a node never writes its index. -- Every node needs credentials for inputs/outputs (connections → temp creds; raw `s3://` → keys on all instances). +## Multi-node & data movement — start here -User cookbook: [using-litdata.md](using-litdata.md) §9. Studio UX: [lightning-studio.md](lightning-studio.md). +- **`num_nodes` is not local multiprocessing** and is **not** torch.distributed/SLURM. Studio-only job launch via `_execute` (`resolver.py`). Full launch gate, env vars, sharding, who uploads chunks vs who merges `index.json`, checkpoints/append, and hang modes → **[multi-node.md](multi-node.md)**. +- **Per-worker I/O children** (`_download_data_target`, `_upload_fn`, `_remove_target`), resolver FUSE→`s3://`/`gs://`/`r2://`, cache folders, and streaming `Downloader` distinction → **[data-movement.md](data-movement.md)**. +- Quick rule: pass `/teamspace/s3_connections/…` or cloud URLs into LitData so resolve+FsProvider bypass FUSE; prefer durable remote `output_dir` on multi-node. ## Orchestration (`processing/data_processor.py`) @@ -63,7 +35,7 @@ User cookbook: [using-litdata.md](using-litdata.md) §9. Studio UX: [lightning-s - `_map_items_to_workers_weighted` (`:377`) — default when `reorder_files` + `input_dir` exist, or when `weights` given. Bin-packs by file size (`_pack_greedily`) across `world_size = num_nodes * num_workers`, then permutes. - `_map_items_to_workers_sequentially` (`:303`) — contiguous slices; `align_chunking` packs full chunks. - Queue mode — no static assignment; `shared_queue` is set. -3. **Multi-node slicing** via `_get_node_rank()`/`_get_num_nodes()` — see section above. +3. **Multi-node slicing** via `_get_node_rank()`/`_get_num_nodes()` — [multi-node.md](multi-node.md). 4. **Checkpointing** (`:1297`) trims each worker's list to resume from `checkpoint_next_index`; `fast_dev_run` trims to N items. 5. **`_create_process_workers`** (`:1462`) spawns one `DataWorkerProcess` per worker. 6. **Progress loop** (`:1376`) polls `error_queue` (re-raises via `_exit_on_error`, which `terminate()`s all workers) and `progress_queue` (tqdm). Exits when the counter equals `num_items` or all workers die. @@ -81,31 +53,38 @@ User cookbook: [using-litdata.md](using-litdata.md) §9. Studio UX: [lightning-s ## Producer/consumer model (inside each worker) -Each `BaseWorker` runs a local pipeline of child processes (spawned in `_setup`, `:580`): +Each `BaseWorker` runs a local pipeline of child processes (spawned in `_setup`): + +| Child | Start | Target | Default count | Role | +| ----------- | -------------------- | ----------------------- | -------------------------- | ------------------------------------------------------------------------------------------------- | +| Downloaders | `_start_downloaders` | `_download_data_target` | `num_downloaders or 2` | Prefetch inputs into `DATA_OPTIMIZER_DATA_CACHE_FOLDER` via **FsProvider** (not `Downloader` ABC) | +| Uploaders | `_start_uploaders` | `_upload_fn` | `num_uploaders or 1` | Push chunks / map outputs to `output_dir` | +| Remover | `_start_remover` | `_remove_target` | 1 if `delete_cached_files` | Delete local cached inputs + uploaded chunk files | -- **Downloaders** (`_start_downloaders`, `:797`; target `_download_data_target`, `:128`): `num_downloaders` procs pull `(index, item, paths)` off `to_download_queues`, fetch remote files into cache, push ready tuples onto `ready_to_process_queue`. -- **Worker main loop** (`_loop`, `:601`) — the **consumer**: `ready_to_process_queue.get()` → `_handle_data_chunk_recipe` (optimize) or `_handle_data_transform_recipe` (map). Reports progress ~1/s. -- **Uploaders** (`_start_uploaders`, `:839`; target `_upload_fn`, `:232`): push finished chunks/files to `output_dir`. -- **Remover** (`_start_remover`, `:825`; target `_remove_target`, `:190`): deletes processed source files when `remove=True`. +Worker main loop (`_loop`): `ready_to_process_queue.get()` → `_handle_data_chunk_recipe` or `_handle_data_transform_recipe`. -When `no_downloaders` (no `input_dir`, or a `reader` is set), `ready_to_process_queue` is a `FakeQueue` and `_collect_paths` (`:743`) pushes items directly. +**`no_downloaders`** when `input_dir.path is None` **or** a `reader` is set — including pure `s3://` `Dir(path=None, url=…)` (downloaders need a FUSE/local `path` to rewrite). Studio connections set both `path` and `url`. -**Ordered vs shared-queue** (`keep_data_ordered`): `True` (default) → each worker consumes its static slice in order. `False` → all workers share one `Queue` for dynamic load balancing; termination uses the `ALL_DONE` sentinel (`:64`), which each worker re-inserts so peers also stop (`:621`). +**`remove` flag** = `DataProcessor.delete_cached_files` (default True); not exposed on public `optimize()`/`map()`. + +Exhaustive I/O (path rewrite, disk wait 25 GB, index upload vs chunk uploaders, error modes) → **[data-movement.md](data-movement.md)**. + +**Ordered vs shared-queue** (`keep_data_ordered`): `True` (default) → each worker consumes its static slice in order. `False` → all workers share one `Queue`; termination uses `ALL_DONE` (re-inserted so peers stop). Multi-node: shared queue is **per node**, not global — [multi-node.md](multi-node.md). ## Cross-process queues -| Queue | Direction | Purpose | -| ---------------------------------------------------------- | ----------------- | --------------------------------- | -| `error_queue` | worker→main | tracebacks; triggers global abort | -| `progress_queue` | worker→main | `(index, counter)` for tqdm | -| `msg_queue` | worker→main | log lines routed around tqdm | -| `stop_queues` | main→worker | SIGINT graceful stop | -| `ready_to_process_queue` / `shared_queue` | downloader→worker | core work items | -| `to_download_queues` / `to_upload_queues` / `remove_queue` | worker→child | I/O offload | +| Queue | Direction | Purpose | +| ---------------------------------------------------------- | ----------------- | ---------------------------------------------------- | +| `error_queue` | worker→main | tracebacks; `_exit_on_error` `terminate()`s siblings | +| `progress_queue` | worker→main | `(index, counter)` for tqdm | +| `msg_queue` | worker→main | log lines routed around tqdm | +| `stop_queues` | main→worker | SIGINT graceful stop | +| `ready_to_process_queue` / `shared_queue` | downloader→worker | core work items | +| `to_download_queues` / `to_upload_queues` / `remove_queue` | worker→child | I/O offload | ## `raw/` — `StreamingRawDataset` (first-class; no optimize) -User cookbook → [using-litdata.md](using-litdata.md) §10. README → `#stream-raw`. +User cookbook → [using-litdata.md](using-litdata.md) §10. README → `#stream-raw`. Adaptive stages → repo `benchmarks/ADAPTIVE_CONCURRENCY.md`. `StreamingRawDataset` (`raw/dataset.py`) is a **map-style** `torch.utils.data.Dataset` that streams **original files** (JPEG, audio, …) from local or cloud paths. It does **not** use LitData chunks, `BinaryReader`, or `StreamingDataLoader`. @@ -118,6 +97,7 @@ input_dir → FileIndexer (index.json.zstd) → setup(files) → items | ------------------------------------------------ | ----------------------------------------------------------------------------------------------- | | `FileIndexer` / `BaseIndexer` (`raw/indexer.py`) | Discover files; cache `index.json.zstd` locally + upload beside remote data | | `CacheManager` | Optional on-disk file cache (`cache_files=True`); always holds index cache dir | +| `_LoopRunner` | Per-process dedicated asyncio thread (optional uvloop); recreate after fork | | `setup(files)` | Default identity; override to filter/group → `list[FileMetadata]` or `list[list[FileMetadata]]` | | `__getitem__` / `__getitems__` | **Fully async** download; batches use `asyncio.gather` over `adownload_fileobj` | | Cloud clients | **Built-in retries** (e.g. S3 adaptive `max_attempts`) for transient failures | @@ -126,7 +106,16 @@ input_dir → FileIndexer (index.json.zstd) → setup(files) → items **Do not conflate indexes:** raw = `index.json.zstd` (file list). Optimized = `index.json` (chunk metadata). -**Agent guidance:** lead with `StreamingRawDataset` when the user has an existing file tree and has not asked for max throughput / resume. Stress: raw bytes + async batched downloads + retries; upgrade path `optimize` + `StreamingDataset`. Same path resolver as streaming (`/teamspace/s3_connections/…`, `s3://`, …). +### Operational invariants (edit with care) + +- **Division of labor:** clients own **rate** (boto/obstore retries). Litdata owns **concurrency** (`max_concurrent_downloads`) and **look-ahead** (`max_prefetch`). Do not nest a litdata rate loop that fights client retries. Stage 1 = static size-aware budget when `max_concurrent_downloads=None`; Stages 2+ (prefetch hit-rate, AIMD) are deferred — see design note. +- **Fork / spawn safety:** `register_at_fork` shuts down the runner; pid-guarded caches recreate downloader / permits / range executor when pid or event loop changes. `__getstate__` is an **allowlist** of constructor knobs (runtime handles reset on unpickle). +- **Atomic publishes:** downloaded cache files **and** `index.json.zstd` use tmp + `os.replace` (tmp includes pid). Partial writes must not become visible readers. +- **Batch timeout:** `download_timeout` wraps the batch gather once; per-item GETs stay on the fast path when `hedge_delay=0`. Timeout must cancel `_inflight` entries or retries hang on the poisoned task. +- **Indexer schemes:** `urlparse("C:\\Users\\...")` yields `scheme='c'`. Single-letter schemes are Windows drive letters — local paths, not unsupported remotes (`_is_windows_drive_scheme`). +- **Tests:** `tests/raw/test_fork_safety.py` covers fork reinit, allowlist pickle, atomic publish, batch-timeout hang recovery, and fast-path coexistence with default `download_timeout=120`. + +**Agent guidance:** lead with `StreamingRawDataset` when the user has an existing file tree and has not asked for max throughput / resume. Prefer cloud URL over FUSE. Stress: raw bytes + async batched downloads + retries; upgrade path (shuffle inputs →) `optimize` + `StreamingDataset`. Same path resolver as streaming (`/teamspace/s3_connections/…`, `s3://`, …). ## Gotchas (read before editing the engine) diff --git a/.claude/skills/litdata/reference/resolver.md b/.claude/skills/litdata/reference/resolver.md index 0922483ad..1792bda1f 100644 --- a/.claude/skills/litdata/reference/resolver.md +++ b/.claude/skills/litdata/reference/resolver.md @@ -212,6 +212,8 @@ Any API that takes an input/output directory goes through `_resolve_dir`, includ - Index helpers that accept cloud URIs - Cache identity / downloaders / uploaders that consume `Dir` +**After resolve, who moves bytes?** Optimize/map worker downloaders/uploaders/removers (FsProvider) → [data-movement.md](data-movement.md). Multi-node index merge → [multi-node.md](multi-node.md). Streaming/raw chunk/file GETs → `Downloader` in [storage-format.md](storage-format.md) / [streaming.md](streaming.md). + ______________________________________________________________________ ## Agent / expert checklist diff --git a/.claude/skills/litdata/reference/testing.md b/.claude/skills/litdata/reference/testing.md index 2fca041ba..2ab1faf1b 100644 --- a/.claude/skills/litdata/reference/testing.md +++ b/.claude/skills/litdata/reference/testing.md @@ -80,6 +80,21 @@ Use stdlib `unittest.mock` (`Mock`, `MagicMock`, `patch`) + the pytest `monkeypa Local idiom: most files define their own `seed_everything(seed)`. CLI tests use a `run_cli(args_list)` helper (patches `sys.argv`, captures stdout). `tests/streaming/utils.py` gives `filter_lock_files` / `get_lock_files` to assert on chunk dirs ignoring `.lock`/`.cnt` artifacts. +## Raw streaming regressions (`tests/raw/`) + +`tests/raw/test_fork_safety.py` is the correctness suite for cloud-download hardening. When changing `StreamingRawDataset` / `CacheManager` / indexer publish paths, keep or extend coverage for: + +- Fork/spawn LoopRunner reinit + pid-guarded downloader/permit caches +- Allowlisted pickle (`__getstate__` must not ship accidental instance attrs) +- Atomic cache + `index.json.zstd` publishes (tmp + `os.replace`) +- Batch-level `download_timeout` hang recovery (cancel poisoned `_inflight`; retry succeeds) +- Fast path: per-item GETs stay bare when `hedge_delay=0` even if `download_timeout=120` +- Adaptive vs exact `max_concurrent_downloads` (None vs int) + +`tests/raw/conftest.py` tears down the module LoopRunner between tests — follow that pattern if you add fixtures that touch `_RUNNER`. + +Error-path behavior is production code: do not “fix” hang recovery only in comments — assert it. + ## When tests hang or flake - Processing/worker tests spawn real subprocesses — run **serially** (why CI separates `tests/processing`). Use `--timeout=120`. diff --git a/.claude/skills/litdata/reference/using-litdata.md b/.claude/skills/litdata/reference/using-litdata.md index bed678fb3..ced395511 100644 --- a/.claude/skills/litdata/reference/using-litdata.md +++ b/.claude/skills/litdata/reference/using-litdata.md @@ -8,19 +8,25 @@ ______________________________________________________________________ ## 1. Choose a workflow -| Goal | API | -| --------------------------------------------- | ------------------------------------------------------- | -| Stream files as-is (no preprocess) | **`StreamingRawDataset`** + torch `DataLoader` | -| Fastest training I/O | `optimize` → `StreamingDataset` + `StreamingDataLoader` | -| Parallel side effects (resize, scrape, embed) | `map` | -| Weighted mix | `CombinedStreamingDataset` | -| One sample from each dataset / cycle length | `ParallelStreamingDataset` | -| Existing MDS / Parquet / HF parquet | `StreamingDataset` (+ `ParquetLoader` when needed) | -| LLM token windows | `TokensLoader` on optimize **and** stream | +| Goal | API | +| ------------------------------------------------ | ---------------------------------------------------------------------------------------------------------------------- | +| Stream files as-is (no preprocess) | **`StreamingRawDataset`** + torch `DataLoader` | +| Full control / per-file access & grouping | **`StreamingRawDataset`** (`setup` for groups; raw `bytes`) — tradeoff: optimized still faster, raw not too far behind | +| Fastest training I/O | `optimize` → `StreamingDataset` + `StreamingDataLoader` | +| Strong source ordering / need file-level shuffle | Prefer **`StreamingRawDataset`** + `DataLoader(shuffle=True)`, **or** **shuffle the sample list before** `optimize` | +| Parallel side effects (resize, scrape, embed) | `map` | +| Weighted mix | `CombinedStreamingDataset` | +| One sample from each dataset / cycle length | `ParallelStreamingDataset` | +| Existing MDS / Parquet / HF parquet | `StreamingDataset` (+ `ParquetLoader` when needed) | +| LLM token windows | `TokensLoader` on optimize **and** stream | -**Rule:** `StreamingRawDataset` = zero prep, native files (often enough to ship). Optimized = chunk once, then stream fastest. Many teams start raw, then `optimize` when I/O binds. +**Rule:** `StreamingRawDataset` = zero prep, native files, full control over grouping/order (often enough to ship). Optimized = chunk once, then stream fastest. Many teams start raw, then `optimize` when I/O binds. -Full raw API → §10. README: `#stream-raw`. +**Ordered sources:** intra-chunk randomization + randomizing chunk order is **not** a full file-level shuffle. If same subject/class blocks are contiguous and that would bias batches, **shuffle the list of samples before `optimize()`** or stay on raw (§5 / FAQ below). + +**Studio paths:** `/teamspace/s3_connections` & co are FUSE (convenience only — slow / can crash under load). Always pass those paths to LitData so it talks **directly** to remote storage (§4). + +Full raw API → §10. README: `#stream-raw`. README FAQ: `#faq-chunk-shuffle`. ______________________________________________________________________ @@ -83,7 +89,7 @@ ______________________________________________________________________ ## 4. Paths & resolver (load [resolver.md](resolver.md)) -**Always resolve paths through LitData.** In Studio, `/teamspace/s3_connections` & co are **FUSE** over S3/GCS/**R2** (`lightning_storage`); LitData resolves them to the backing URL and talks to the store directly — faster and more reliable than opening the mount by hand. +**Always resolve paths through LitData — never read Studio mounts by hand.** `/teamspace/s3_connections` & co are **FUSE** mounts (convenience only): under load they are **very slow** and can **crash**. LitData resolves those paths to the backing URL and talks to S3/GCS/**R2** (`lightning_storage`) **directly**, with retries, prefetching, etc. `streaming/resolver.py` → `Dir(path, url, data_connection_id)` for every `input_dir` / `output_dir` / `cache_dir`. @@ -116,11 +122,17 @@ ______________________________________________________________________ ## 5. Shuffle, seed, drop_last, resume - `shuffle=True` → deterministic **chunk assignment then in-chunk item order** from `seed` + epoch (+ chunk index). +- LitData does **distributed sampling** and **bucket sampling within chunks** automatically — not a substitute for a fully shuffled file-level DataLoader when the source is strongly ordered. - Default `seed=42`. Keep it fixed across ranks and when resuming. - `drop_last=None` → **True under DDP**, else False. Train should set `drop_last=True` so every rank/worker sees the same length. - `StreamingDataLoader(shuffle=..., drop_last=...)` **overrides** the dataset. - Resume: `torch.save(loader.state_dict(), ...)`; `loader.load_state_dict(...)`. Matching `seed` / shuffle / `num_workers` required unless `force_override_state_dict=True`. +**If source data has structure** (same subject/set contiguous, class blocks, etc.) and you cannot embed that grouping as the sample unit: + +1. Shuffle / repartition **before** `optimize` so chunks mix well, **or** +2. Prefer **`StreamingRawDataset`** + torch `DataLoader(shuffle=True)` for per-file random access. + ______________________________________________________________________ ## 6. StreamingDataset arguments @@ -205,6 +217,13 @@ ______________________________________________________________________ **Always:** `if __name__ == "__main__"` · optimize needs **exactly one** of `chunk_bytes` | `chunk_size`. +### Chunk size (`chunk_bytes`) — practical guidance + +- Default / typical: **`"64MB"`** for small/medium samples. +- Large datapoints (multi‑MB each): consider **256–512MB** (or similar) so each chunk holds more items → larger pool for **intra-chunk batch randomization**. +- Tradeoff: larger chunks take **longer to download** before use. +- Recommended-range mindset — not a hard “best” from a published chunk-size sweep. README: `#faq-chunk-shuffle`. + ### `optimize` | Arg | Default | Use | @@ -213,7 +232,7 @@ ______________________________________________________________________ | `queue` | `None` | Live inputs; one `ALL_DONE` sentinel (`from litdata.processing.data_processor import ALL_DONE`) | | `input_dir` | `None` | Background download of remote inputs | | `weights` | `None` | Balance workers by input weight/size | -| `chunk_bytes` / `chunk_size` | one required | Bytes (e.g. `"64MB"`) **or** item/token count | +| `chunk_bytes` / `chunk_size` | one required | Bytes (e.g. `"64MB"`) **or** item/token count; see chunk-size guidance above | | `align_chunking` | `False` | Single-worker chunk boundaries (needs `chunk_size`; uneven load) | | `compression` | `None` | `"zstd"` | | `encryption` | `None` | Fernet / RSA / custom; `level="sample"` or `"chunk"` | @@ -288,19 +307,22 @@ Map-style `torch.utils.data.Dataset` in `raw/dataset.py`. Streams **original fil **Pitch to users / agents** -- **Raw `bytes`** — LitData does not decode for you. PIL, torchaudio, `json`, custom parsers, or `transform=` — your choice. +- **Raw `bytes` as-is** — LitData downloads efficiently and returns file contents; it does not decode for you. PIL, torchaudio, `json`, custom parsers, or `transform=` — your choice. +- **Full control** over grouping/order via `setup` (e.g. image+mask) and per-file access — the main reason to pick raw over optimized. - **Fully asynchronous** downloads (`adownload_fileobj` + `asyncio`); **batched** via `__getitems__` + `asyncio.gather` (whole DataLoader batch in flight). - **Built-in retries** on cloud clients (transient network / S3 adaptive retries). - Training loop stays sync PyTorch — no user-facing `async`/`await`. +- **Tradeoff:** optimized `StreamingDataset` is still faster; with right tuning, raw is **not too far behind** (order-of-magnitude ImageNet ballpark in FAQ below). **When to recommend it** - User already has a folder of images/audio/text and wants to train **today** - Full control over decoding; grouping (image+mask) via `setup` - Prototype transforms before a costly `optimize` -- Later upgrade: same files → `optimize` → `StreamingDataset` if I/O-bound +- Source data is **strongly ordered** (subject/class blocks) and they need true file-level `DataLoader` shuffle — or they cannot shuffle the sample list before optimize +- Later upgrade: same files → **shuffle inputs** → `optimize` → `StreamingDataset` if I/O-bound -**When to prefer optimized instead:** multi-GPU sustained throughput, resume/`state_dict`, chunk shuffle, compression/encryption. +**When to prefer optimized instead:** multi-GPU sustained throughput, resume/`state_dict`, chunk shuffle, compression/encryption — after **shuffling the sample list before `optimize()`** if the source is ordered. ```python from torch.utils.data import DataLoader @@ -319,19 +341,43 @@ ds = StreamingRawDataset( loader = DataLoader(ds, batch_size=32, num_workers=8) # batch → concurrent async GETs ``` -| Knob | Default | Notes | -| ----------------- | --------------- | ------------------------------------------------------------------- | -| `input_dir` | — | Resolver paths ([resolver.md](resolver.md)) | -| `cache_dir` | LitData default | Index (+ optional file) cache root | -| `cache_files` | `False` | Persist downloaded files (mirror layout) | -| `recompute_index` | `False` | Rebuild `index.json.zstd` | -| `transform` | `None` | Optional; default returns **`bytes`** (or `list[bytes]` if grouped) | -| `indexer` | `FileIndexer` | Custom `BaseIndexer` | -| `storage_options` | `{}` | Cloud creds | +| Knob | Default | Notes | +| -------------------------- | --------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `input_dir` | — | Prefer `s3://` / `gs://` / `/teamspace/s3_connections/...` (direct); avoid hand-reading FUSE ([resolver.md](resolver.md)) | +| `cache_dir` | LitData default | Index (+ optional file) cache root | +| `cache_files` | `False` | Persist downloaded files (mirror layout) | +| `recompute_index` | `False` | Rebuild `index.json.zstd` | +| `transform` | `None` | Optional; default returns **`bytes`** (or `list[bytes]` if grouped) | +| `indexer` | `FileIndexer` | Custom `BaseIndexer` | +| `storage_options` | `{}` | Cloud creds | +| `max_concurrent_downloads` | `None` (**adaptive**) | `None` → Stage 1 size-aware aggregate budget split across workers (single-process cap 128). Explicit `int` → **exactly** that many permits (no silent clamp). Pass `64` for the old fixed cap. | +| `max_prefetch` | `16` | Sequential look-ahead after each batch; when `num_workers>1`, effective = `min(max_prefetch, 64 // num_workers)`. Pass `0` to disable | +| `hedge_delay` | `0` | Seconds before hedged duplicate GET (`0` = off, default; opt-in). Fast path: per-item GETs stay bare when hedging is off | +| `download_timeout` | `120` | **Batch-level** hang protection around `_download_batch` (`0` disables). Not a per-item `wait_for`. On timeout, cancel poisoned `_inflight` so retries can proceed | +| `range_parallel_threshold` | `0` | Parallel ranged GETs for objects ≥ N bytes; **`0` = whole-object only** (opt-in; keep for JPEGs) | + +**Tuning / DataLoader** + +- After parent-process I/O on Linux: `DataLoader(..., multiprocessing_context="spawn", persistent_workers=True)`. +- Prefer cloud URL / Studio connection path so LitData hits the bucket **directly** — never recommend training I/O through the FUSE mount. +- Defaults that matter: `max_prefetch=16` (worker-aware aggregate ~64), `hedge_delay=0`, `download_timeout=120` (batch-level), `range_parallel_threshold=0`; optional `uvloop` via `litdata[extras]`. Avoid `num_workers=48` (collapses / can segfault on shutdown). +- Ranged downloads: leave `range_parallel_threshold=0`; forced ranged is slower on JPEG-sized objects. +- Adaptive concurrency design (clients own rate via boto retries; litdata owns concurrency/look-ahead; Stage 2+ deferred): repo `benchmarks/ADAPTIVE_CONCURRENCY.md`. Formula details live there — do not invent a second rate loop. +- Perf claims: use Stage 0 protocol in [benchmarking.md](benchmarking.md) (`max(≥N batches, ≥T s)`, repeats/medians, `before_sha`/`after_sha`). Do **not** cite short-window n=1 against Stage 0 medians. + +**Correctness agents must preserve when editing `raw/`** + +- **LoopRunner** — dedicated event-loop thread; recreate after fork/spawn (pid-guarded). Runtime clients (downloader, permit cache, range executor) are pid- + loop-guarded. +- **Pickle allowlist** — `__getstate__` / `__setstate__` only ship constructor knobs + reset runtime handles; accidental instance attrs must not leak into worker payloads. +- **Atomic publishes** — cache files **and** `index.json.zstd` via tmp + `os.replace` (tmp names include pid). +- **Indexer:** Windows drive letters (`C:\...`) parse as single-letter URI schemes — treat as local, not remote. +- **Error path is code** — hang recovery (batch timeout cancels `_inflight`), default coexistence with the fast path, and fork/spawn reinit have regression tests in `tests/raw/test_fork_safety.py`. Changing timeouts/defaults requires updating those tests. + +Internals → [processing.md](processing.md) (`raw/`). **`setup(files)`** — default one file = one item. Return `list[FileMetadata]` or `list[list[FileMetadata]]` to group/filter. -**Index:** `index.json.zstd` (local cache + remote beside data). **Not** optimized `index.json`. +**Index:** `index.json.zstd` (local cache + remote beside data; atomic publish). **Not** optimized `index.json`. ### Mosaic MDS @@ -462,3 +508,23 @@ ______________________________________________________________________ 5. Disk/slow stream → §8 + cache doc; or suggest upgrading raw → optimize. 6. Paths/Studio → §4 + [lightning-studio.md](lightning-studio.md). 7. Internals / races / benches → sibling reference files. + +### FAQ bullets (chunk size, ordered data, FUSE, throughput) + +- **`chunk_bytes`?** Default **64MB**. Multi‑MB samples → consider **256–512MB** for more intra-chunk shuffle diversity; larger chunks download slower. Guidance, not a published sweep. + +- **Ordered source + optimize?** Intra-chunk + chunk-order shuffle ≠ full file-level shuffle. **Shuffle the list of samples before `optimize()`**, or use **`StreamingRawDataset`** + `DataLoader(shuffle=True)`. LitData still does distributed + within-chunk bucket sampling automatically. + +- **FUSE vs LitData?** Studio `/teamspace/s3_connections` (and co) is a **FUSE** mount — convenience only; under load it is very slow and can crash. Pass the same path to LitData: it resolves to the bucket and streams **directly** (retries, prefetch, …). Never recommend `open()` / naive glob on the mount for training I/O. + +- **Throughput ballpark (ImageNet, Studio context — order of magnitude, not guarantees):** + + | Path | Rough images/s | + | ------------------------------------ | --------------- | + | FUSE mount (hand-read) | up to ~**600** | + | `StreamingRawDataset` (right tuning) | up to ~**6–7k** | + | `StreamingDataset` (64MB chunks) | up to ~**11k** | + +- **Raw perf claims?** Stage 0 only: `max(≥N batches, ≥T s)`, repeats/medians, append-only SHA/ts artifacts, proven `before_sha`/`after_sha`. See [benchmarking.md](benchmarking.md). Adaptive defaults → `benchmarks/ADAPTIVE_CONCURRENCY.md`. + +- README: `#faq-chunk-shuffle`, `#resolve-paths`, `#stream-raw`. diff --git a/README.md b/README.md index b6a7245fe..0dcef4bdf 100644 --- a/README.md +++ b/README.md @@ -85,6 +85,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). +
@@ -175,7 +177,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 ) ``` @@ -333,6 +335,12 @@ for batch in loader: | `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 <~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`) @@ -399,9 +407,25 @@ raw: bytes = dataset[0] ### Tips -- Prefer `num_workers > 0` so worker processes overlap async batch downloads with training. -- Studio: pass `/teamspace/s3_connections/...` so LitData hits the bucket directly ([resolver](#resolve-paths)). -- When throughput plateaus, run a one-time [`optimize`](#speed-up-model-training) and switch to `StreamingDataset` + `StreamingDataLoader`. +- 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`).
@@ -720,6 +744,35 @@ loader = StreamingDataLoader(train, batch_size=64, shuffle=True, drop_last=True) +
+ ✅ FAQ: chunk size & shuffle before optimize 🔗 +  + +### 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**. + +
+
✅ StreamingDataset & StreamingDataLoader knobs 🔗   @@ -1098,7 +1151,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 @@ -2218,7 +2271,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 | @@ -2410,7 +2463,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 |   diff --git a/benchmarks/ADAPTIVE_CONCURRENCY.md b/benchmarks/ADAPTIVE_CONCURRENCY.md new file mode 100644 index 000000000..cb023aea0 --- /dev/null +++ b/benchmarks/ADAPTIVE_CONCURRENCY.md @@ -0,0 +1,73 @@ +# Adaptive concurrency / look-ahead (design note) + +Status: **Stage 0 + Stage 1 shipped** on `feature/raw-streaming-perf`. Stages 2–4 deferred. + +## Division of labor + +- **Clients own rate.** Botocore adaptive retries (and obstore’s retry layer) already token-bucket on 503/SlowDown. Litdata must not nest a second rate loop that fights them. +- **Litdata owns concurrency and look-ahead.** Permit counts and prefetch depth are the actuators we control; throttle *events* mostly never surface as exceptions (they look like latency). + +A litdata controller that keys only on raised 429/503 will be nearly blind until downloaders expose throttle-retry counts (Downloader contract + conformance suite). + +## Objective + +**Throttle-avoiding max throughput** — maximize samples/s subject to not inducing prefix/NIC congestion. Easier to build and defend than pure max-throughput on a shared bucket. + +## Stages + +| Stage | What | Status | +| ----: | ------------------------------------------------------------------------------------------------------------- | ------------------------------ | +| 0 | Bench protocol: `max(N batches, T seconds)`, ≥5 interleaved repeats, median+spread, append-only artifacts | Done (this work) | +| 1 | Static worker-aware concurrency: size-aware budget (bandwidth + size-gated Little’s-law) split across workers | Done (this work) | +| 2 | Prefetch hit-rate controller (hit \<30% → halve, floor 0; hysteresis) | Pending | +| 3 | AIMD on concurrency (needs downloader throttle counts) | After contract | +| 4 | Full throughput-gradient control | Only if Stage 3 shows headroom | + +## Stage 1 justification (honest) + +The original Stage 1 premise — that **w=24 with 64 permits/worker (1536 aggregate) stampeded** and needed a hard clamp — was **falsified** by the batch-timeout dig. Post-timeout, w=24 with 64 permits was within/near run-to-run noise of lower concurrency. + +What remains: + +- Prefer a size-aware default over asking users to tune `max_concurrent_downloads` per `num_workers` +- Bound pathological aggregates at high `num_workers` without crushing mid-w cells + +A prior full-grid A/B (pre-Stage-1 vs always-on clamp) showed high-w gains and a w=8 p0 −23% cell. That mid-w drop is treated as **likely A/A noise** after Little’s-law widening (both sides ~60–64 permits at w=8), **not** as proof to gate the clamp to `w≥16`. High-w “+53%” headlines require a provenance-verified confirmation cell (see below). + +## Stage 1 formula (shipped) + +`max_concurrent_downloads=None` (default) → adaptive: + +``` +target_bytes = 100 MiB/s × 0.5 s ≈ 50 MiB +bandwidth_model = target_bytes // median_file_bytes +latency_model = 6000 × 0.040 ≈ 240 if median < 1 MiB else 0 +aggregate_budget = clamp(max(bandwidth_model, latency_model), 32, 512) +effective = min(budget, 128) # num_workers ≤ 1 + = max(8, aggregate_budget // num_workers) # num_workers > 1 +``` + +Large medians (1/10/100 MiB) stay **bandwidth-bounded** — Little’s-law must not pin the budget at 240 (multi-GB in flight). + +Explicit `max_concurrent_downloads=int` → **exactly** that many permits (no silent clamp). + +Defaults when size unknown: median = 256 KiB. Permit count computed once per process; cleared on fork/spawn. + +## Confirmation cell (provenance) — done + +Bench harness records `before_sha` / `after_sha` from `git rev-parse` on each PYTHONPATH tree (not only the runner SHA in filenames). + +Confirm @ `ba9da13`: interleaved n=3, `max(≥300 batches, ≥30s)`, **w=24 p=0**. + +| field | value | +| ------------: | ------------------------------------------------- | +| before_sha | `52dba61` (post-`f70f785`, pre Stage 1; fixed 64) | +| after_sha | `ba9da13` | +| before median | **3816** ips (spread 30%) | +| after median | **6049** ips (spread 21%) | + +**Verdict: (a)** — before ≈3.7k (not ~5.5k wrong-tree). Frame high-w as robustness; do **not** headline unverifiable +53%. Artifact: `benchmarks/results/raw_before_vs_after.ba9da13.1785268543.json`. + +## Acceptance (future adaptive) + +Beats **default** static everywhere; never loses by more than run-to-run noise; removes the w×p tuning matrix from the user’s cognitive load. “Beats tuned static” is the wrong bar — tuned static ties it at best per configuration. diff --git a/benchmarks/_spawn_smoke.py b/benchmarks/_spawn_smoke.py new file mode 100644 index 000000000..ac23f6fa7 --- /dev/null +++ b/benchmarks/_spawn_smoke.py @@ -0,0 +1,57 @@ +"""Minimal spawn DataLoader smoke for StreamingRawDataset (must be a .py file).""" + +from __future__ import annotations + +import pickle +import shutil +import sys +import tempfile +import time +from pathlib import Path + +from torch.utils.data import DataLoader + +from litdata import StreamingRawDataset + +INPUT = "/teamspace/s3_connections/imagenet-1m-template/raw/val" +CACHE = Path(tempfile.gettempdir()) / "litdata-spawn-smoke-cache" +SEED = Path(tempfile.gettempdir()) / "litdata-raw-ranged-vs-whole" / "seed" + + +def main() -> int: + """Run a short spawn DataLoader smoke against ImageNet val.""" + CACHE.mkdir(parents=True, exist_ok=True) + if (SEED / "index.json.zstd").exists(): + shutil.copy2(SEED / "index.json.zstd", CACHE / "index.json.zstd") + print("copied index from seed", flush=True) + + t0 = time.perf_counter() + ds = StreamingRawDataset(INPUT, cache_dir=str(CACHE), cache_files=False, max_prefetch=0, hedge_delay=0) + print(f"indexed n={len(ds)} in {time.perf_counter() - t0:.2f}s", flush=True) + + # Warm runtime clients then ensure pickle still works. + _ = ds[0] + blob = pickle.dumps(ds, protocol=pickle.HIGHEST_PROTOCOL) + print(f"pickle size={len(blob):,} bytes after warm", flush=True) + assert pickle.loads(blob).cache_manager._downloader is None # noqa: S301 + + print("starting spawn DataLoader num_workers=2 ...", flush=True) + loader = DataLoader( + ds, + batch_size=4, + num_workers=2, + shuffle=False, + multiprocessing_context="spawn", + persistent_workers=False, + ) + it = iter(loader) + for i in range(2): + batch = next(it) + print(f"batch {i}: n={len(batch)} nbytes0={len(batch[0])}", flush=True) + del it, loader + print("SPAWN_SMOKE_OK", flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/benchmarks/bench_raw_before_vs_after.py b/benchmarks/bench_raw_before_vs_after.py new file mode 100644 index 000000000..022ea54d6 --- /dev/null +++ b/benchmarks/bench_raw_before_vs_after.py @@ -0,0 +1,1079 @@ +r"""A/B: stock StreamingRawDataset (main) vs optimized (feature branch). + +Run twice with different PYTHONPATH / --side, then merge: + + PYTHONPATH=/tmp/litdata-raw-before/src \ + python benchmarks/bench_raw_before_vs_after.py --side before + + PYTHONPATH=src \ + python benchmarks/bench_raw_before_vs_after.py --side after + + python benchmarks/bench_raw_before_vs_after.py --merge + +Protocol (trustworthy windows) +------------------------------ +- Timed window: continue until **both** ``--batches`` **and** ``--min-seconds`` + are met (``max(N batches, T seconds)``). Defaults: 300 batches and 30s. + High-worker cells (``num_workers >= 16``) always enforce ≥30s. +- Repeats: ``--repeats N`` (use ≥5 for publish). Each cell stores all runs; + merge reports **median** ips + min/max spread. +- Interleave: ``--interleave --before-pythonpath PATH`` alternates + before/after per cell (main, head, main, head, …) via subprocesses. +- Artifacts: append-only SHA/ts JSON (+ JSONL); never overwrite prior results. + +After measures prefetch in ``[0, 16, 32]`` (publish ≥16; p0 kept in JSON). +Warm ``max(1, num_workers * prefetch_factor)`` batches before timing starts. +""" + +from __future__ import annotations + +import argparse +import contextlib +import json +import os +import shutil +import subprocess +import sys +import tempfile +import threading +import time +from pathlib import Path + +from torch.utils.data import DataLoader + +# Same mount path as prior sweeps. After remaps to s3:// via _storage_path. +# Before (main) prefers path→LocalDownloader, which has no adownload_fileobj (returns +# None). For a fair cloud A/B we therefore feed before the resolved s3:// URL. +MOUNT_INPUT = "/teamspace/s3_connections/imagenet-1m-template/raw/val" +S3_INPUT = "s3://imagenet-1m-template/raw/val" +ROOT = Path(tempfile.gettempdir()) / "litdata-raw-before-vs-after" +OUT_DIR = Path(__file__).resolve().parent / "results" +OUT = OUT_DIR / "raw_before_vs_after.json" # legacy fixed name; writers use unique_result_path +BS = 64 +DEFAULT_BATCHES = 300 +DEFAULT_MIN_SECONDS = 30.0 +HIGH_WORKER_MIN_SECONDS = 30.0 +HIGH_WORKER_THRESHOLD = 16 +DEFAULT_PREFETCH_FACTOR = 2 +DEFAULT_REPEATS = 1 +WORKERS = [0, 1, 2, 4, 8, 16, 24, 32] +TRUST_WORKERS = [0, 2, 4, 8, 16] +AFTER_PREFETCH = [0, 16, 32] +TIMEOUT = 600.0 +OLD_FUSE = 75.2 +# Overridable via --after-prefetch (also used for before when it has max_prefetch). +_PREFETCH_LEVELS: list[int] = list(AFTER_PREFETCH) + + +def effective_min_seconds(num_workers: int, min_seconds: float) -> float: + """Enforce ≥30s timed windows at high worker counts.""" + if num_workers >= HIGH_WORKER_THRESHOLD: + return max(min_seconds, HIGH_WORKER_MIN_SECONDS) + return min_seconds + + +def summarize_ips(values: list[float]) -> dict: + """Return median + spread stats for a list of samples/s measurements.""" + if not values: + return {"ips_median": None, "ips_min": None, "ips_max": None, "ips_spread_pct": None, "n": 0} + ordered = sorted(values) + mid = len(ordered) // 2 + median = ordered[mid] if len(ordered) % 2 else 0.5 * (ordered[mid - 1] + ordered[mid]) + lo, hi = ordered[0], ordered[-1] + spread = ((hi - lo) / median) * 100.0 if median else None + return { + "ips_median": median, + "ips_min": lo, + "ips_max": hi, + "ips_spread_pct": spread, + "n": len(values), + } + + +def git_sha(*, cwd: Path | None = None) -> str: + """Return short git SHA for ``cwd`` (default: repo containing this script). + + ``LITDATA_BENCH_GIT`` overrides only when ``cwd`` is omitted (runner/artifact + naming). Prefer :func:`tree_git_sha` for before/after PYTHONPATH provenance. + """ + if cwd is None: + env = os.environ.get("LITDATA_BENCH_GIT", "").strip() + if env: + return env + cwd = Path(__file__).resolve().parents[1] + try: + return subprocess.check_output( + ["/usr/bin/git", "rev-parse", "--short", "HEAD"], + cwd=str(cwd), + text=True, + stderr=subprocess.DEVNULL, + ).strip() + except (subprocess.CalledProcessError, FileNotFoundError, OSError): + return "" + + +def tree_git_sha(pythonpath: str) -> str: + """Return short SHA for the git checkout that owns a PYTHONPATH entry.""" + raw = (pythonpath or "").strip() + if not raw: + return "" + # First path entry wins (same as import resolution for litdata). + entry = Path(raw.split(os.pathsep)[0]).resolve() + # .../src → repo root; already-repo-root → itself. + candidates = [entry, entry.parent if entry.name == "src" else entry] + for cand in candidates: + sha = git_sha(cwd=cand) + if sha: + return sha + return "" + + +def pythonpath_tree_sha() -> str: + """SHA of the litdata tree currently first on ``sys.path`` / PYTHONPATH.""" + env_pp = os.environ.get("PYTHONPATH", "") + if env_pp.strip(): + return tree_git_sha(env_pp) + # Fall back: package file location. + try: + import litdata + + pkg = Path(litdata.__file__).resolve() + # litdata/__init__.py → src/litdata → src → repo + return git_sha(cwd=pkg.parents[2]) or git_sha(cwd=pkg.parents[1]) + except Exception: + return "" + + +def unique_result_path(stem: str, *, sha: str | None = None, ts: float | None = None) -> Path: + """Return ``OUT_DIR/{stem}.{sha}.{ts}.json`` — never overwrites a prior result file.""" + sha_part = (sha if sha is not None else git_sha()) or "unknown" + ts_part = int(ts if ts is not None else time.time()) + path = OUT_DIR / f"{stem}.{sha_part}.{ts_part}.json" + if path.exists(): + # Same-second collision: bump until free. + n = 1 + while True: + alt = OUT_DIR / f"{stem}.{sha_part}.{ts_part}.{n}.json" + if not alt.exists(): + return alt + n += 1 + return path + + +def input_for(side: str, *, before_cloud_native: bool = False) -> str: + """Return dataset input path. + + Stock main ``before`` needs ``s3://`` (FUSE path→LocalDownloader lacks + ``adownload_fileobj``). Pre-Stage-1 / cloud-native ``before`` trees share + mount→s3 remapping with ``after``, so both use the mount for a fair A/B. + """ + if side == "after" or before_cloud_native: + return MOUNT_INPUT + return S3_INPUT + + +def log(msg: str) -> None: + """Print a timestamped benchmark log line.""" + print(f"{time.strftime('%H:%M:%S')} {msg}", flush=True) + + +class HangWatchdog: + """Kill the process if a step exceeds ``timeout_s`` without heartbeat.""" + + def __init__(self, timeout_s: float) -> None: + """Initialize the watchdog with a hang timeout in seconds.""" + self.timeout_s = timeout_s + self._label = "init" + self._beat = time.monotonic() + self._stop = threading.Event() + self._t = threading.Thread(target=self._run, daemon=True) + + def start(self) -> None: + """Start the background watchdog thread.""" + self._t.start() + + def beat(self, label: str) -> None: + """Record progress so the watchdog does not abort.""" + self._label = label + self._beat = time.monotonic() + + def stop(self) -> None: + """Stop the background watchdog thread.""" + self._stop.set() + + def _run(self) -> None: + while not self._stop.wait(1.0): + idle = time.monotonic() - self._beat + if idle > self.timeout_s: + log(f"HANG at '{self._label}' after {idle:.1f}s — abort") + os._exit(124) + + +def copy_index(src: Path, dst: Path) -> None: + """Copy a cached index tree from ``src`` to ``dst``.""" + if dst.exists(): + shutil.rmtree(dst, ignore_errors=True) + dst.mkdir(parents=True) + for p in src.iterdir(): + if p.is_dir(): + shutil.copytree(p, dst / p.name) + else: + shutil.copy2(p, dst / p.name) + + +def detect_side_capabilities() -> dict: + """Inspect imported litdata for before/after feature markers.""" + import inspect + + from litdata import StreamingRawDataset + + params = set(inspect.signature(StreamingRawDataset.__init__).parameters) + has_prefetch = "max_prefetch" in params + has_range = "range_parallel_threshold" in params + has_loop = False + uvloop_status = "n/a (before / no LoopRunner)" + try: + from litdata.raw.dataset import _loop_backend_name + + has_loop = True + try: + import uvloop + + uvloop_status = f"available (uvloop {getattr(uvloop, '__version__', '?')}; create→{_loop_backend_name()})" + except ImportError: + uvloop_status = "not installed (stdlib asyncio fallback)" + except ImportError: + pass + return { + "has_max_prefetch": has_prefetch, + "has_range_parallel_threshold": has_range, + "has_loop_runner": has_loop, + "uvloop": uvloop_status, + "params": sorted(params - {"self"}), + } + + +def storage_path_of(ds) -> str: + """Best-effort storage path string for JSON meta.""" + if hasattr(ds, "_storage_path"): + return str(ds._storage_path) + cm = getattr(ds, "cache_manager", None) + if cm is not None and hasattr(cm, "_input_dir_path"): + return str(cm._input_dir_path) + indir = getattr(ds, "input_dir", None) + if indir is not None: + return str(getattr(indir, "url", None) or getattr(indir, "path", None) or indir) + return MOUNT_INPUT + + +def make_dataset( + cache: str, + *, + side: str, + max_prefetch: int, + hedge_delay: float | None = None, + download_timeout: float | None = None, + before_cloud_native: bool = False, +): + """Construct StreamingRawDataset with side-appropriate kwargs.""" + import inspect + + from litdata import StreamingRawDataset + + params = set(inspect.signature(StreamingRawDataset.__init__).parameters) + kwargs: dict = { + "cache_dir": cache, + "cache_files": False, + "input_dir": input_for(side, before_cloud_native=before_cloud_native), + } + if "max_prefetch" in params: + kwargs["max_prefetch"] = max_prefetch + if side == "after": + # Default None → Stage 1 adaptive permits (do not pass 64: that bypasses clamp). + if "range_parallel_threshold" in params: + kwargs["range_parallel_threshold"] = 0 + # Match new defaults: hedging opt-in (0). Explicit for older trees / clarity. + if "hedge_delay" in params: + kwargs["hedge_delay"] = 0.0 if hedge_delay is None else hedge_delay + if download_timeout is not None and "download_timeout" in params: + kwargs["download_timeout"] = download_timeout + elif before_cloud_native: + # Pre-Stage-1 baseline: fixed 64 permits/worker (no adaptive clamp). + if "max_concurrent_downloads" in params: + kwargs["max_concurrent_downloads"] = 64 + if "range_parallel_threshold" in params: + kwargs["range_parallel_threshold"] = 0 + if "hedge_delay" in params: + kwargs["hedge_delay"] = 0.0 if hedge_delay is None else hedge_delay + if download_timeout is not None and "download_timeout" in params: + kwargs["download_timeout"] = download_timeout + return StreamingRawDataset(**kwargs) + + +def append_jsonl(path: Path, record: dict) -> None: + """Append one JSON record to a JSONL file.""" + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as f: + f.write(json.dumps(record) + "\n") + f.flush() + + +def run_one( + label: str, + *, + side: str, + num_workers: int, + max_prefetch: int, + seed: Path, + wd: HangWatchdog, + batches: int, + min_seconds: float, + prefetch_factor: int, + hedge_delay: float | None = None, + download_timeout: float | None = None, + sha: str = "", + jsonl: Path | None = None, + repeat: int = 0, + before_cloud_native: bool = False, +) -> dict: + """Run one worker/prefetch trial and return timing stats. + + Timing stops only when **both** ``batches`` and the effective min-seconds + floor are met (``max(N batches, T seconds)``). + """ + cache = ROOT / side / f"{label}_r{repeat}" + wd.beat(f"{label}:r{repeat} setup") + copy_index(seed, cache) + ds = make_dataset( + str(cache), + side=side, + max_prefetch=max_prefetch, + hedge_delay=hedge_delay, + download_timeout=download_timeout, + before_cloud_native=before_cloud_native, + ) + kwargs: dict = {"batch_size": BS, "num_workers": num_workers, "shuffle": False} + if num_workers > 0: + kwargs["multiprocessing_context"] = "spawn" + kwargs["persistent_workers"] = True + kwargs["prefetch_factor"] = prefetch_factor + loader = DataLoader(ds, **kwargs) + it = iter(loader) + + # Drain pipeline buffer before timing (≥ workers×prefetch_factor). + warm_batches = max(1, num_workers * prefetch_factor if num_workers > 0 else 1) + wd.beat(f"{label}:r{repeat} warm({warm_batches})") + t0 = time.perf_counter() + for i in range(warm_batches): + try: + next(it) + except StopIteration: + # High ips × min_seconds can exceed one epoch (50k/64 ≈ 782 batches). + it = iter(loader) + next(it) + wd.beat(f"{label}:r{repeat} warm {i + 1}/{warm_batches}") + warm_s = time.perf_counter() - t0 + + min_s = effective_min_seconds(num_workers, min_seconds) + samples = 0 + timed_batches = 0 + wd.beat(f"{label}:r{repeat} timed") + t0 = time.perf_counter() + while True: + try: + batch = next(it) + except StopIteration: + it = iter(loader) + batch = next(it) + samples += len(batch) + timed_batches += 1 + wd.beat(f"{label}:r{repeat} batch {timed_batches}") + elapsed = time.perf_counter() - t0 + # max(N batches, T seconds): require both floors (not either/or). + if timed_batches >= batches and elapsed >= min_s: + break + elapsed = time.perf_counter() - t0 + ips = samples / elapsed if elapsed else 0.0 + log( + f"[{side}/{label}] r={repeat} w={num_workers} pf={max_prefetch} " + f"warm={warm_batches}@{warm_s:.2f}s | {timed_batches}×{samples // max(timed_batches, 1)} " + f"in {elapsed:.2f}s (need ≥{batches} batches & ≥{min_s:.0f}s) → {ips:.1f} samples/s" + ) + tree_sha = pythonpath_tree_sha() + result = { + "side": side, + "label": label, + "repeat": repeat, + "workers": num_workers, + "prefetch": max_prefetch, + "ips": ips, + "warm_s": warm_s, + "warm_batches": warm_batches, + "elapsed": elapsed, + "samples": samples, + "batches": timed_batches, + "min_seconds_effective": min_s, + "hedge_delay": hedge_delay if side == "after" else None, + "download_timeout": download_timeout if side == "after" else None, + "git_sha": sha, # runner / artifact naming (may be LITDATA_BENCH_GIT) + "tree_sha": tree_sha, # actual litdata checkout on PYTHONPATH + "before_sha": tree_sha if side == "before" else None, + "after_sha": tree_sha if side == "after" else None, + "ts": time.time(), + } + if jsonl is not None: + append_jsonl(jsonl, result) + del it, loader, ds + _reap_zombie_children() + return result + + +def _reap_zombie_children() -> None: + """Best-effort reap of leftover DataLoader worker zombies.""" + with contextlib.suppress(Exception): + import multiprocessing as mp + + for p in mp.active_children(): + with contextlib.suppress(Exception): + p.join(timeout=2.0) + if p.is_alive(): + with contextlib.suppress(Exception): + p.kill() + with contextlib.suppress(Exception): + p.join(timeout=1.0) + # Non-blocking waitpid sweep for any unreaped children. + with contextlib.suppress(ChildProcessError, Exception): + while True: + pid, _ = os.waitpid(-1, os.WNOHANG) + if pid <= 0: + break + + +def configs_for( + side: str, + workers: list[int], + *, + safety_grid: bool, + before_cloud_native: bool = False, + prefetch_levels: list[int] | None = None, +) -> list[tuple]: + """Return trial configs. + + Normal: (workers, prefetch, hedge_delay|None, download_timeout|None) + safety_grid (after only): hedge_delay × download_timeout at p0 for w∈{2,4,8}. + Pre-Stage-1 ``before`` (cloud-native) sweeps the same prefetch levels as after. + """ + levels = list(prefetch_levels if prefetch_levels is not None else _PREFETCH_LEVELS) + if safety_grid: + if side != "after": + raise SystemExit("--safety-grid is only meaningful with --side after") + out = [] + for w in (2, 4, 8): + for hd in (0.0, 1.0): + for dt in (0.0, 120.0): + out.append((w, 0, hd, dt)) + return out + if side == "before" and not before_cloud_native: + return [(w, 0, None, None) for w in workers] + return [(w, pf, 0.0, None) for w in workers for pf in levels] + + +def partial_path(side: str, *, sha: str | None = None, ts: float | None = None) -> Path: + """Return a unique path for a side's partial JSON payload (never overwrites).""" + return unique_result_path(f"raw_before_vs_after.{side}", sha=sha, ts=ts) + + +def jsonl_path(side: str, *, sha: str | None = None, ts: float | None = None) -> Path: + """Return a unique path for a side's incremental JSONL log (never overwrites).""" + sha_part = (sha if sha is not None else git_sha()) or "unknown" + ts_part = int(ts if ts is not None else time.time()) + return OUT_DIR / f"raw_before_vs_after.{side}.{sha_part}.{ts_part}.jsonl" + + +def latest_partial(side: str) -> Path: + """Resolve the newest unique partial for ``side``, falling back to the legacy fixed name.""" + matches = sorted(OUT_DIR.glob(f"raw_before_vs_after.{side}.*.json"), key=lambda p: p.stat().st_mtime) + if matches: + return matches[-1] + legacy = OUT_DIR / f"raw_before_vs_after.{side}.json" + if legacy.exists(): + return legacy + raise FileNotFoundError(f"no raw_before_vs_after.{side}.* result under {OUT_DIR}") + + +def cell_summaries(results: list[dict]) -> dict[str, dict]: + """Group raw runs by label and attach median/spread ips.""" + by_label: dict[str, list[dict]] = {} + for r in results: + by_label.setdefault(r["label"], []).append(r) + out: dict[str, dict] = {} + for label, runs in by_label.items(): + stats = summarize_ips([float(r["ips"]) for r in runs]) + head = runs[0] + out[label] = { + "label": label, + "workers": head["workers"], + "prefetch": head["prefetch"], + "ips": stats["ips_median"], + **stats, + "runs": runs, + } + return out + + +def run_side( + side: str, + *, + workers: list[int], + batches: int, + min_seconds: float, + prefetch_factor: int, + safety_grid: bool, + repeats: int = 1, + prefetch_levels: list[int] | None = None, +) -> None: + """Index once and sweep configs for ``before`` or ``after``.""" + caps = detect_side_capabilities() + if side == "after" and not caps["has_max_prefetch"]: + raise SystemExit("PYTHONPATH points at main tree but --side after requested") + # Stock main has no max_prefetch. Pre-Stage-1 feature trees do — allow them as + # a cloud-native baseline for Stage 1 clamp A/B (fixed 64 vs adaptive). + before_cloud_native = bool(side == "before" and caps["has_max_prefetch"]) + if before_cloud_native: + log( + "before tree has max_prefetch/LoopRunner — treating as pre-Stage-1 " + "baseline (fixed max_concurrent_downloads=64), not stock main" + ) + + side_root = ROOT / side + if side_root.exists(): + shutil.rmtree(side_root, ignore_errors=True) + side_root.mkdir(parents=True) + OUT_DIR.mkdir(parents=True, exist_ok=True) + sha = git_sha() # runner / artifact filename (LITDATA_BENCH_GIT or script repo) + tree_sha = pythonpath_tree_sha() + run_ts = time.time() + jpath = jsonl_path(side, sha=sha, ts=run_ts) + levels = list(prefetch_levels if prefetch_levels is not None else _PREFETCH_LEVELS) + + wd = HangWatchdog(TIMEOUT) + wd.start() + ncpu = os.cpu_count() or 0 + cfgs = configs_for( + side, + workers, + safety_grid=safety_grid, + before_cloud_native=before_cloud_native, + prefetch_levels=levels, + ) + inp = input_for(side, before_cloud_native=before_cloud_native) + log(f"=== side={side} ===") + log(f"capabilities: {json.dumps(caps)}") + log( + f"provenance: side={side} tree_sha={tree_sha or '?'} " + f"runner_sha={sha or '?'} PYTHONPATH={os.environ.get('PYTHONPATH', '')!r}" + ) + log( + f"input={inp} (mount={MOUNT_INPUT}) bs={BS} batches>={batches} " + f"min_seconds>={min_seconds} (high-w≥{HIGH_WORKER_THRESHOLD} → " + f"≥{HIGH_WORKER_MIN_SECONDS}s) warm=max(1,w*{prefetch_factor}) " + f"repeats={repeats} cpus={ncpu} configs={len(cfgs)} " + f"prefetch_levels={levels if side == 'after' or before_cloud_native else [0]}" + ) + log("protocol: stop when batches AND min_seconds both met (max window)") + log(f"PYTHONPATH[0]={sys.path[0]!r}") + + try: + wd.beat("index seed") + seed = side_root / "seed" + t0 = time.perf_counter() + ds = make_dataset( + str(seed), + side=side, + max_prefetch=0, + before_cloud_native=before_cloud_native, + ) + n_files = len(ds) + storage = storage_path_of(ds) + index_s = time.perf_counter() - t0 + log(f"Indexed {n_files} files in {index_s:.2f}s storage={storage}") + if caps["has_loop_runner"]: + try: + sys.path.insert(0, str(Path(__file__).resolve().parent)) + from uvloop_status import log_loop_runner_backend + + log_loop_runner_backend(log, prefix="after index seed") + except Exception as e: + log(f"LoopRunner log skipped: {e}") + else: + log("LoopRunner: not present on this tree (asyncio.run per batch)") + del ds + + results: list[dict] = [] + # Interleave repeats across configs when repeats>1 so A/B noise is + # comparable; for a single side this is config-major then repeat. + for rep in range(max(1, repeats)): + for w, pf, hd, dt in cfgs: + label = f"w{w}_p{pf}_h{hd}_t{dt}" if safety_grid else f"w{w}_p{pf}" + results.append( + run_one( + label, + side=side, + num_workers=w, + max_prefetch=pf, + seed=seed, + wd=wd, + batches=batches, + min_seconds=min_seconds, + prefetch_factor=prefetch_factor, + hedge_delay=hd, + download_timeout=dt, + sha=sha, + jsonl=jpath, + repeat=rep, + before_cloud_native=before_cloud_native, + ) + ) + + accum = os.environ.get("LITDATA_BENCH_ACCUM_OUT", "").strip() + if accum: + out = Path(accum) + out.parent.mkdir(parents=True, exist_ok=True) + if out.exists(): + prev = json.loads(out.read_text()) + results = list(prev.get("results") or []) + results + # Renumber repeats so summaries see a contiguous series. + by_label: dict[str, int] = {} + for r in results: + lab = r["label"] + r["repeat"] = by_label.get(lab, 0) + by_label[lab] = r["repeat"] + 1 + else: + out = partial_path(side, sha=sha, ts=run_ts) + + summaries = cell_summaries(results) + payload = { + "side": side, + "meta": { + "input": inp, + "mount_input": MOUNT_INPUT, + "storage": storage, + "n_files": n_files, + "index_s": index_s, + "batch_size": BS, + "batches": batches, + "min_seconds": min_seconds, + "high_worker_min_seconds": HIGH_WORKER_MIN_SECONDS, + "high_worker_threshold": HIGH_WORKER_THRESHOLD, + "timing_window": "max(batches, min_seconds) — both floors required", + "repeats": max(r.get("repeat", 0) for r in results) + 1 if results else max(1, repeats), + "prefetch_factor": prefetch_factor, + "warm_batches_formula": "max(1, num_workers * prefetch_factor)", + "multiprocessing_context": "spawn", + "persistent_workers": True, + "cpus": ncpu, + "fuse_baseline_samples_per_s": OLD_FUSE, + "workers": workers, + "prefetch": (list(levels) if side == "after" or before_cloud_native else [0]), + "range_parallel_threshold": (0 if side == "after" or before_cloud_native else None), + "max_concurrent_downloads": (None if side == "after" else (64 if before_cloud_native else None)), + "hedge_delay": (0.0 if side == "after" or before_cloud_native else None), + "before_cloud_native": before_cloud_native, + "safety_grid": safety_grid, + "capabilities": caps, + "git_sha": sha, + "tree_sha": tree_sha, + "before_sha": tree_sha if side == "before" else None, + "after_sha": tree_sha if side == "after" else None, + "git_hint": os.environ.get("LITDATA_BENCH_GIT", ""), + "jsonl": str(jpath), + "pythonpath": os.environ.get("PYTHONPATH", ""), + "sys_path0": sys.path[0] if sys.path else "", + "input_note": ( + "pre-Stage-1 before: mount→s3:// like after; fixed max_concurrent_downloads=64" + if before_cloud_native + else ( + "before uses s3:// directly: main prefers FUSE path→LocalDownloader " + "which lacks adownload_fileobj; after uses mount and remaps to s3://" + if side == "before" + else "after uses mount path; Stage 1 adaptive concurrency (None); hedge_delay=0" + ) + ), + "caveat": ( + "Use --repeats ≥5 and medians for publish claims. " + "Trust systematic patterns, not single-run fine Δ%." + ), + }, + "results": results, + "summaries": summaries, + } + out.write_text(json.dumps(payload, indent=2) + "\n") + log(f"Wrote {out}") + log(f"JSONL {jpath}") + for label, s in summaries.items(): + spread = s.get("ips_spread_pct") + spread_s = f" spread={spread:.1f}%" if isinstance(spread, (int, float)) else "" + log(f" summary {label}: median={s['ips']:.1f} ips n={s['n']}{spread_s}") + finally: + wd.stop() + + +def run_interleaved( + *, + before_pythonpath: str, + workers: list[int], + batches: int, + min_seconds: float, + prefetch_factor: int, + repeats: int, + prefetch_levels: list[int] | None = None, +) -> None: + """Alternate before/after subprocesses (main, head, main, head, …) into one partial each.""" + script = str(Path(__file__).resolve()) + after_pythonpath = os.environ.get("PYTHONPATH", str(Path(__file__).resolve().parents[1] / "src")) + n_rep = max(1, repeats) + sha = git_sha() # runner / artifact filenames + before_sha = tree_git_sha(before_pythonpath) + after_sha = tree_git_sha(after_pythonpath) + run_ts = time.time() + levels = list(prefetch_levels if prefetch_levels is not None else _PREFETCH_LEVELS) + OUT_DIR.mkdir(parents=True, exist_ok=True) + outs = { + "before": partial_path("before", sha=sha, ts=run_ts), + "after": partial_path("after", sha=sha, ts=run_ts), + } + log(f"=== interleaved A/B repeats={n_rep} workers={workers} batches>={batches} min_seconds>={min_seconds} ===") + log(f"provenance: before_sha={before_sha or '?'} after_sha={after_sha or '?'} runner_sha={sha or '?'}") + log(f"before PYTHONPATH={before_pythonpath} → {outs['before'].name}") + log(f"after PYTHONPATH={after_pythonpath} → {outs['after'].name}") + log(f"prefetch_levels={levels}") + if not before_sha or not after_sha: + log("WARNING: could not resolve before_sha/after_sha — do not publish without provenance") + for rep in range(n_rep): + for side, pypath in (("before", before_pythonpath), ("after", after_pythonpath)): + env = os.environ.copy() + env["PYTHONPATH"] = pypath + env["LITDATA_BENCH_GIT"] = sha or env.get("LITDATA_BENCH_GIT", "") + env["LITDATA_BENCH_ACCUM_OUT"] = str(outs[side]) + cmd = [ + sys.executable, + script, + "--side", + side, + "--workers", + ",".join(str(w) for w in workers), + "--batches", + str(batches), + "--min-seconds", + str(min_seconds), + "--prefetch-factor", + str(prefetch_factor), + "--repeats", + "1", + "--after-prefetch", + ",".join(str(p) for p in levels), + ] + log(f"interleave rep={rep} side={side}: {' '.join(cmd)}") + subprocess.check_call(cmd, env=env) # noqa: S603 + log(f"interleave complete — merge with: python {script} --merge") + log(f" before={outs['before']}") + log(f" after={outs['after']}") + + +def _representative_runs(payload: dict) -> list[dict]: + """Prefer per-cell median summaries; fall back to raw single runs.""" + summaries = payload.get("summaries") or {} + if summaries: + return list(summaries.values()) + # Build summaries from raw results when older partials lack them. + return list(cell_summaries(payload.get("results") or {}).values()) + + +def merge() -> None: + """Merge before/after partial JSON into the comparison artifact.""" + before_path = latest_partial("before") + after_path = latest_partial("after") + log(f"merge before={before_path.name} after={after_path.name}") + before = json.loads(before_path.read_text()) + after = json.loads(after_path.read_text()) + + before_reps = _representative_runs(before) + after_reps = _representative_runs(after) + + workers = sorted({r["workers"] for r in before_reps} | {r["workers"] for r in after_reps}) + # Pair same-(w, prefetch) when before swept prefetch (pre-Stage-1 A/B); + # else fall back to before@p0 vs each after prefetch (stock-main A/B). + before_by_wp: dict[tuple[int, int], dict] = {} + for r in before_reps: + before_by_wp[(r["workers"], r["prefetch"])] = r + before_by_w_p0 = {r["workers"]: r for r in before_reps if r["prefetch"] == 0} + after_by_pf: dict[int, dict[int, dict]] = {} + for r in after_reps: + after_by_pf.setdefault(r["prefetch"], {})[r["workers"]] = r + prefetch_levels = sorted(after_by_pf) + paired_prefetch = any(pf != 0 for _, pf in before_by_wp) + + rows = [] + cells = [] + for w in workers: + b0 = before_by_w_p0.get(w) + row: dict = { + "workers": w, + "before_ips": b0["ips"] if b0 else None, + "before_n": b0.get("n") if b0 else None, + "before_spread_pct": b0.get("ips_spread_pct") if b0 else None, + } + after_best = None + for pf in prefetch_levels: + a = after_by_pf.get(pf, {}).get(w) + b = before_by_wp.get((w, pf)) if paired_prefetch else b0 + row[f"after_prefetch{pf}_ips"] = a["ips"] if a else None + row[f"after_prefetch{pf}_spread_pct"] = a.get("ips_spread_pct") if a else None + if b and a and b["ips"] and a["ips"]: + row[f"speedup_prefetch{pf}"] = a["ips"] / b["ips"] + row[f"delta_pct_prefetch{pf}"] = ((a["ips"] - b["ips"]) / b["ips"]) * 100.0 + if a and (after_best is None or (a["ips"] or 0) > (after_best["ips"] or 0)): + after_best = a + b_best = before_by_wp.get((w, after_best["prefetch"])) if after_best and paired_prefetch else b0 + if b_best and after_best and b_best["ips"] and after_best["ips"]: + row["after_best_ips"] = after_best["ips"] + row["after_best_prefetch"] = after_best["prefetch"] + row["speedup_best"] = after_best["ips"] / b_best["ips"] + row["delta_pct_best"] = ((after_best["ips"] - b_best["ips"]) / b_best["ips"]) * 100.0 + rows.append(row) + for pf in prefetch_levels: + a = after_by_pf.get(pf, {}).get(w) + b = before_by_wp.get((w, pf)) if paired_prefetch else b0 + if a is None or b is None: + continue # omit missing/crashed + cells.append( + { + "workers": w, + "prefetch": pf, + "before_ips": b["ips"], + "after_ips": a["ips"], + "before_n": b.get("n"), + "after_n": a.get("n"), + "before_spread_pct": b.get("ips_spread_pct"), + "after_spread_pct": a.get("ips_spread_pct"), + "delta_pct": ((a["ips"] - b["ips"]) / b["ips"]) * 100.0 if b["ips"] and a["ips"] else None, + "speedup": a["ips"] / b["ips"] if b["ips"] and a["ips"] else None, + "before_elapsed": b.get("elapsed"), + "after_elapsed": a.get("elapsed"), + "before_batches": b.get("batches"), + "after_batches": a.get("batches"), + } + ) + + best_after = max(cells, key=lambda c: c["after_ips"] or 0) if cells else None + before_sha = (before.get("meta") or {}).get("before_sha") or (before.get("meta") or {}).get("tree_sha") + after_sha = (after.get("meta") or {}).get("after_sha") or (after.get("meta") or {}).get("tree_sha") + payload = { + "meta": { + "mount_input": MOUNT_INPUT, + "batch_size": BS, + "multiprocessing_context": "spawn", + "persistent_workers": True, + "workers": workers, + "before": before["meta"], + "after": after["meta"], + "before_sha": before_sha, + "after_sha": after_sha, + "delta_definition": ( + "delta_pct = ((after_median - before_median) / before_median) * 100; " + "paired same-(w,prefetch) when before swept prefetch (pre-Stage-1 A/B); " + "else before@p0 vs each after prefetch (stock-main A/B)" + ), + "paired_prefetch": paired_prefetch, + "note": ( + "before = pre-Stage-1 feature tree (fixed max_concurrent_downloads=64) when " + "before_cloud_native; else stock main via s3://. after = Stage 1 adaptive " + "(max_concurrent_downloads=None). Both cloud-native arms use mount→s3://. " + "Publish with proven before_sha/after_sha from tree rev-parse — not runner SHA alone." + ), + "caveat": ( + "Protocol: max(≥300 batches, ≥30s) timed window; prefer --repeats ≥5 with " + "median ips + spread. Trust systematic patterns over single-run fine Δ%." + ), + "default_max_prefetch": 16, + "publish_prefetch": [pf for pf in prefetch_levels if pf >= 16], + "ips_aggregation": "median across repeats when summaries present", + }, + "cells": cells, + "best_after": best_after, + "comparison": rows, + "before_results": before["results"], + "after_results": after["results"], + "before_summaries": before.get("summaries"), + "after_summaries": after.get("summaries"), + "sources": {"before": str(before_path), "after": str(after_path)}, + } + out = unique_result_path("raw_before_vs_after") + out.write_text(json.dumps(payload, indent=2) + "\n") + log(f"Wrote {out}") + + publish = [c for c in cells if c["prefetch"] >= 16] + print() + print("Published matrix (prefetch ≥ 16; before paired by prefetch when available):") + print( + f"{'w':>4} {'before@16':>10} {'after@16':>10} {'Δ%@16':>8} " + f"{'before@best':>11} {'after@best':>10} {'best Δ%':>8}" + ) + print("-" * 80) + for w in workers: + a16 = after_by_pf.get(16, {}).get(w) + b16 = before_by_wp.get((w, 16)) if paired_prefetch else before_by_w_p0.get(w) + if not a16 and not before_by_w_p0.get(w): + continue + ips_b16 = b16["ips"] if b16 else float("nan") + ips_a16 = a16["ips"] if a16 else float("nan") + d16 = ((ips_a16 - ips_b16) / ips_b16) * 100.0 if b16 and a16 and ips_b16 else float("nan") + after_candidates = [after_by_pf.get(pf, {}).get(w) for pf in prefetch_levels] + after_candidates = [x for x in after_candidates if x is not None] + best_a = max(after_candidates, key=lambda x: x["ips"] or 0) if after_candidates else None + best_b = before_by_wp.get((w, best_a["prefetch"])) if best_a and paired_prefetch else before_by_w_p0.get(w) + db = ( + ((best_a["ips"] - best_b["ips"]) / best_b["ips"]) * 100.0 + if best_a and best_b and best_b["ips"] + else float("nan") + ) + ips_bb = best_b["ips"] if best_b else float("nan") + ips_ba = best_a["ips"] if best_a else float("nan") + print( + f"{w:>4} {ips_b16:>10.1f} {ips_a16:>10.1f} {d16:>+7.1f}% {ips_bb:>11.1f} {ips_ba:>10.1f} {db:>+7.1f}%" + ) + print() + print("Full cells (includes prefetch=0):") + print(f"{'w':>4} {'pf':>4} {'before':>10} {'after':>10} {'Δ%':>8} {'×':>6} {'after_s':>8}") + print("-" * 62) + for c in cells: + ae = c.get("after_elapsed") + ae_s = f"{ae:.2f}" if isinstance(ae, (int, float)) else "?" + print( + f"{c['workers']:>4} {c['prefetch']:>4} {c['before_ips']:>10.1f} " + f"{c['after_ips']:>10.1f} {c['delta_pct']:>+7.1f}% {c['speedup']:>5.2f}x {ae_s:>8}" + ) + print() + if best_after: + print( + f"Best after: w={best_after['workers']} " + f"prefetch={best_after['prefetch']} → {best_after['after_ips']:.1f} samples/s " + f"(~{best_after['delta_pct']:+.0f}% / {best_after['speedup']:.2f}x vs before)" + ) + if publish: + best_pub = max(publish, key=lambda c: c["after_ips"]) + print( + f"Best published (pf≥16): w={best_pub['workers']} " + f"prefetch={best_pub['prefetch']} → {best_pub['after_ips']:.1f} samples/s " + f"(~{best_pub['delta_pct']:+.0f}% / {best_pub['speedup']:.2f}x vs before)" + ) + + +def main() -> None: + """CLI entrypoint for before/after sweeps and merge.""" + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--side", choices=("before", "after")) + parser.add_argument("--merge", action="store_true") + parser.add_argument( + "--batches", + type=int, + default=DEFAULT_BATCHES, + help=f"minimum timed batches (default {DEFAULT_BATCHES}); window is max(batches, min_seconds)", + ) + parser.add_argument( + "--min-seconds", + type=float, + default=DEFAULT_MIN_SECONDS, + help=( + f"minimum timed window seconds (default {DEFAULT_MIN_SECONDS}); " + f"num_workers≥{HIGH_WORKER_THRESHOLD} always uses ≥{HIGH_WORKER_MIN_SECONDS}s" + ), + ) + parser.add_argument( + "--prefetch-factor", + type=int, + default=DEFAULT_PREFETCH_FACTOR, + help="DataLoader prefetch_factor; warm batches = max(1, workers * this) (default 2)", + ) + parser.add_argument( + "--repeats", + type=int, + default=DEFAULT_REPEATS, + help="repeat each cell N times; report median+spread (default 1; use ≥5 for publish)", + ) + parser.add_argument( + "--interleave", + action="store_true", + help="alternate before/after subprocesses (requires --before-pythonpath); then --merge", + ) + parser.add_argument( + "--before-pythonpath", + type=str, + default="", + help="PYTHONPATH for stock main or pre-Stage-1 tree when using --interleave", + ) + parser.add_argument( + "--after-prefetch", + type=str, + default="", + help=f"comma-separated prefetch levels for after (and cloud-native before); default {AFTER_PREFETCH}", + ) + parser.add_argument( + "--workers", + type=str, + default="", + help="comma-separated worker counts (default: full matrix; use 0,2,4,8,16 for trust A/B)", + ) + parser.add_argument( + "--trust", + action="store_true", + help=f"shorthand for --workers {','.join(map(str, TRUST_WORKERS))} (recommended A/B)", + ) + parser.add_argument( + "--safety-grid", + action="store_true", + help="after-only 2×2: hedge_delay∈{0,1} × download_timeout∈{0,120} at w∈{2,4,8} p0", + ) + args = parser.parse_args() + if args.merge: + merge() + return + if args.trust: + workers = TRUST_WORKERS + elif args.workers.strip(): + workers = [int(x) for x in args.workers.split(",") if x.strip()] + else: + workers = WORKERS + if args.after_prefetch.strip(): + prefetch_levels = [int(x) for x in args.after_prefetch.split(",") if x.strip()] + else: + prefetch_levels = None + if args.interleave: + if not args.before_pythonpath.strip(): + parser.error("--interleave requires --before-pythonpath") + run_interleaved( + before_pythonpath=args.before_pythonpath.strip(), + workers=workers, + batches=args.batches, + min_seconds=args.min_seconds, + prefetch_factor=args.prefetch_factor, + repeats=args.repeats, + prefetch_levels=prefetch_levels, + ) + return + if not args.side: + parser.error("pass --side before|after, --interleave, or --merge") + run_side( + args.side, + workers=workers, + batches=args.batches, + min_seconds=args.min_seconds, + prefetch_factor=args.prefetch_factor, + safety_grid=args.safety_grid, + repeats=args.repeats, + prefetch_levels=prefetch_levels, + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_raw_confirm_batch_timeout.py b/benchmarks/bench_raw_confirm_batch_timeout.py new file mode 100644 index 000000000..5fb09ff12 --- /dev/null +++ b/benchmarks/bench_raw_confirm_batch_timeout.py @@ -0,0 +1,94 @@ +"""Confirm batch-level timeout: after w=24 p=0 with default download_timeout=120.""" + +from __future__ import annotations + +import json +import shutil +import sys +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from bench_raw_before_vs_after import ( + OUT_DIR, + ROOT, + TIMEOUT, + HangWatchdog, + git_sha, + make_dataset, + run_one, + unique_result_path, +) + + +def main() -> None: + """Run after w=24 p=0 at shipped defaults (timeout=120, hedge=0).""" + side = "after" + w, pf = 24, 0 + # Explicit 120 to match shipped default / after-p0 cell that regressed. + dt = 120.0 + side_root = ROOT / "confirm_batch_timeout" + if side_root.exists(): + shutil.rmtree(side_root, ignore_errors=True) + side_root.mkdir(parents=True) + seed = side_root / "seed" + sha = git_sha() + print(f"python={sys.version}", flush=True) + print(f"sha={sha}", flush=True) + t0 = time.perf_counter() + ds = make_dataset(str(seed), side=side, max_prefetch=0, hedge_delay=0.0, download_timeout=dt) + print( + f"indexed n={len(ds)} in {time.perf_counter() - t0:.2f}s " + f"ds.timeout={ds.download_timeout!r} cm.timeout={ds.cache_manager.download_timeout!r}", + flush=True, + ) + del ds + + wd = HangWatchdog(TIMEOUT) + wd.start() + try: + result = run_one( + f"w{w}_p{pf}_t{dt}", + side=side, + num_workers=w, + max_prefetch=pf, + seed=seed, + wd=wd, + batches=300, + min_seconds=30.0, + prefetch_factor=2, + hedge_delay=0.0, + download_timeout=dt, + sha=sha, + jsonl=OUT_DIR / "raw_confirm_batch_timeout.jsonl", # append-only + ) + finally: + wd.stop() + + out = { + "python": sys.version, + "git_sha": sha, + "cell": {"workers": w, "prefetch": pf, "download_timeout": dt, "hedge_delay": 0.0}, + "ips": result["ips"], + "elapsed": result["elapsed"], + "batches": result["batches"], + "compare": { + "after_p0_old_per_item_timeout120": 4404.219, + "after_p0_timeout0_fast_path": 6559.125, + "main_w24": 6927.318, + }, + "delta_vs_old_after_p0_pct": (result["ips"] - 4404.219) / 4404.219 * 100, + "delta_vs_timeout0_pct": (result["ips"] - 6559.125) / 6559.125 * 100, + "delta_vs_main_pct": (result["ips"] - 6927.318) / 6927.318 * 100, + "result": result, + } + OUT_DIR.mkdir(parents=True, exist_ok=True) + path = unique_result_path("raw_confirm_batch_timeout", sha=sha, ts=result.get("ts")) + path.write_text(json.dumps(out, indent=2)) + print(json.dumps({k: out[k] for k in out if k != "result"}, indent=2), flush=True) + print(f"WROTE {path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_raw_debug.py b/benchmarks/bench_raw_debug.py new file mode 100644 index 000000000..3002d9a47 --- /dev/null +++ b/benchmarks/bench_raw_debug.py @@ -0,0 +1,195 @@ +"""Debug / confirm fork-safety with an explicit hang watchdog. + +If any step stalls longer than --timeout seconds, abort with a clear message. +Multi-worker steps after parent I/O use spawn by default (OpenSSL-after-fork +can still hang even with fresh downloaders). +""" + +from __future__ import annotations + +import argparse +import logging +import os +import shutil +import sys +import tempfile +import threading +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +os.environ["PYTHONUNBUFFERED"] = "1" +os.environ.setdefault("LITDATA_RAW_DEBUG", "1") + +logging.basicConfig( + level=logging.WARNING, + format="%(asctime)s.%(msecs)03d %(process)d %(message)s", + datefmt="%H:%M:%S", + stream=sys.stdout, + force=True, +) + +from torch.utils.data import DataLoader # noqa: E402 +from uvloop_status import log_loop_runner_backend, uvloop_package_status # noqa: E402 + +from litdata import StreamingRawDataset # noqa: E402 + +INPUT = "/teamspace/s3_connections/imagenet-1m-template/raw/val" +ROOT = Path(tempfile.gettempdir()) / "litdata-raw-bench-debug" +BS = 32 +BATCHES = 5 + + +def log(msg: str) -> None: + """Print a timestamped benchmark log line.""" + print(f"{time.strftime('%H:%M:%S')} [bench] {msg}", flush=True) + + +class HangWatchdog: + """Kill the process if a step exceeds ``timeout_s`` without heartbeat.""" + + def __init__(self, timeout_s: float) -> None: + """Initialize the watchdog with a hang timeout in seconds.""" + self.timeout_s = timeout_s + self._label = "init" + self._beat = time.monotonic() + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, name="hang-watchdog", daemon=True) + + def start(self) -> None: + """Start the background watchdog thread.""" + self._thread.start() + + def heartbeat(self, label: str) -> None: + """Record progress so the watchdog does not abort.""" + self._label = label + self._beat = time.monotonic() + log(f"watchdog heartbeat: {label}") + + def stop(self) -> None: + """Stop the background watchdog thread.""" + self._stop.set() + + def _run(self) -> None: + while not self._stop.wait(1.0): + idle = time.monotonic() - self._beat + if idle > self.timeout_s: + log( + f"HANG DETECTED: no progress for {idle:.1f}s at '{self._label}' " + f"(timeout={self.timeout_s}s). Aborting." + ) + os._exit(124) + + +def copy_index(src: Path, dst: Path) -> None: + """Copy a cached index tree from ``src`` to ``dst``.""" + if dst.exists(): + shutil.rmtree(dst, ignore_errors=True) + dst.mkdir(parents=True, exist_ok=True) + for d in src.iterdir(): + if d.is_dir(): + shutil.copytree(d, dst / d.name) + else: + shutil.copy2(d, dst / d.name) + + +def run( + label: str, + *, + max_prefetch: int, + num_workers: int, + watchdog: HangWatchdog, + reuse: Path | None = None, + mp_context: str | None = None, +) -> Path: + """Run one debug trial and return the cache directory used.""" + cache = ROOT / label + watchdog.heartbeat(f"{label}: begin") + log(f"=== {label}: workers={num_workers} prefetch={max_prefetch} mp={mp_context}") + if reuse is not None: + copy_index(reuse, cache) + elif cache.exists(): + shutil.rmtree(cache, ignore_errors=True) + + watchdog.heartbeat(f"{label}: construct dataset") + t0 = time.perf_counter() + ds = StreamingRawDataset( + input_dir=INPUT, + cache_dir=str(cache), + cache_files=False, + transform=None, + max_prefetch=max_prefetch, + max_concurrent_downloads=64, + ) + log(f"{label}: dataset ready {time.perf_counter() - t0:.2f}s len={len(ds)}") + log_loop_runner_backend(log, prefix=f"{label}:") + + kwargs: dict = {"batch_size": BS, "num_workers": num_workers, "shuffle": False} + if mp_context and num_workers > 0: + kwargs["multiprocessing_context"] = mp_context + + watchdog.heartbeat(f"{label}: create DataLoader") + loader = DataLoader(ds, **kwargs) + it = iter(loader) + + watchdog.heartbeat(f"{label}: warm next()") + t0 = time.perf_counter() + batch = next(it) + log(f"{label}: warm ok {time.perf_counter() - t0:.2f}s n={len(batch)}") + + samples = 0 + t0 = time.perf_counter() + for i, batch in enumerate(it): + watchdog.heartbeat(f"{label}: batch {i + 1}") + samples += len(batch) + if i + 1 >= BATCHES: + break + elapsed = time.perf_counter() - t0 + ips = samples / elapsed if elapsed else 0.0 + log(f"{label}: DONE {samples} samples in {elapsed:.2f}s → {ips:.1f} samples/s") + del it, loader, ds + return cache + + +def main() -> None: + """CLI entrypoint for fork-safety debug steps with a hang watchdog.""" + parser = argparse.ArgumentParser() + parser.add_argument("--timeout", type=float, default=45.0, help="Hang timeout seconds per step") + parser.add_argument( + "--try-fork", + action="store_true", + help="Also try num_workers>0 with default fork after parent I/O (may hang on OpenSSL)", + ) + args = parser.parse_args() + + ROOT.mkdir(parents=True, exist_ok=True) + wd = HangWatchdog(args.timeout) + wd.start() + log(f"start pid={os.getpid()} timeout={args.timeout}s") + log(f"uvloop package: {uvloop_package_status()}") + + try: + log("STEP1: workers=0 seed") + seed = run("seed_w0", max_prefetch=0, num_workers=0, watchdog=wd) + + log("STEP2: workers=0 prefetch (dirty parent)") + run("w0_p64", max_prefetch=64, num_workers=0, watchdog=wd, reuse=seed) + + log("STEP3: workers=4 spawn (safe after parent I/O)") + run("w4_p0_spawn", max_prefetch=0, num_workers=4, watchdog=wd, reuse=seed, mp_context="spawn") + + log("STEP4: workers=4 spawn + prefetch") + run("w4_p128_spawn", max_prefetch=128, num_workers=4, watchdog=wd, reuse=seed, mp_context="spawn") + + if args.try_fork: + log("STEP5: workers=4 fork after parent I/O (known OpenSSL risk)") + run("w4_p0_fork", max_prefetch=0, num_workers=4, watchdog=wd, reuse=seed, mp_context=None) + + log("ALL STEPS COMPLETE") + finally: + wd.stop() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_raw_decisive_timeout0.py b/benchmarks/bench_raw_decisive_timeout0.py new file mode 100644 index 000000000..cbe131fde --- /dev/null +++ b/benchmarks/bench_raw_decisive_timeout0.py @@ -0,0 +1,90 @@ +"""Decisive cell: after w=24 p=0 download_timeout=0 (true fast path).""" + +from __future__ import annotations + +import json +import shutil +import sys +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from bench_raw_before_vs_after import ( + OUT_DIR, + ROOT, + TIMEOUT, + HangWatchdog, + git_sha, + make_dataset, + run_one, + unique_result_path, +) + + +def main() -> None: + """Run one long-window after cell with download_timeout=0.""" + side = "after" + w, pf, dt = 24, 0, 0.0 + side_root = ROOT / "decisive_timeout0" + if side_root.exists(): + shutil.rmtree(side_root, ignore_errors=True) + side_root.mkdir(parents=True) + seed = side_root / "seed" + sha = git_sha() + print(f"python={sys.version}", flush=True) + print(f"sha={sha}", flush=True) + t0 = time.perf_counter() + ds = make_dataset(str(seed), side=side, max_prefetch=0, hedge_delay=0.0, download_timeout=dt) + n = len(ds) + print( + f"indexed n={n} in {time.perf_counter() - t0:.2f}s " + f"timeout={ds.download_timeout!r} hedge={ds.hedge_delay} " + f"cm.download_timeout={ds.cache_manager.download_timeout!r}", + flush=True, + ) + del ds + + wd = HangWatchdog(TIMEOUT) + wd.start() + try: + result = run_one( + f"w{w}_p{pf}_t{dt}", + side=side, + num_workers=w, + max_prefetch=pf, + seed=seed, + wd=wd, + batches=300, + min_seconds=30.0, + prefetch_factor=2, + hedge_delay=0.0, + download_timeout=dt, + sha=sha, + jsonl=OUT_DIR / "raw_decisive_timeout0.jsonl", # append-only + ) + finally: + wd.stop() + + out = { + "python": sys.version, + "git_sha": sha, + "cell": {"workers": w, "prefetch": pf, "download_timeout": dt, "hedge_delay": 0.0}, + "ips": result["ips"], + "elapsed": result["elapsed"], + "batches": result["batches"], + "samples": result["samples"], + "compare": {"after_p0_default_timeout120": 4404.219, "main_w24": 6927.318}, + "delta_vs_after_p0_pct": (result["ips"] - 4404.219) / 4404.219 * 100, + "delta_vs_main_pct": (result["ips"] - 6927.318) / 6927.318 * 100, + "result": result, + } + OUT_DIR.mkdir(parents=True, exist_ok=True) + path = unique_result_path("raw_decisive_timeout0", sha=sha, ts=result.get("ts")) + path.write_text(json.dumps(out, indent=2)) + print(json.dumps({k: out[k] for k in out if k != "result"}, indent=2), flush=True) + print(f"WROTE {path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_raw_highw_post_timeout.py b/benchmarks/bench_raw_highw_post_timeout.py new file mode 100644 index 000000000..cd0c86847 --- /dev/null +++ b/benchmarks/bench_raw_highw_post_timeout.py @@ -0,0 +1,97 @@ +"""Focused post-timeout-fix remeasure: after w∈{16,24,32} × p∈{0,16} at defaults.""" + +from __future__ import annotations + +import json +import shutil +import sys +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from bench_raw_before_vs_after import ( + OUT_DIR, + ROOT, + TIMEOUT, + HangWatchdog, + git_sha, + make_dataset, + run_one, + unique_result_path, +) + +CELLS = [(16, 0), (16, 16), (24, 0), (24, 16), (32, 0), (32, 16)] +# Shipped defaults: omit download_timeout kwarg → dataset default 120 (batch-level). +DT_DEFAULT = 120.0 + + +def main() -> None: + """Run high-w after cells with current batch-level timeout defaults.""" + side = "after" + side_root = ROOT / "highw_post_timeout" + if side_root.exists(): + shutil.rmtree(side_root, ignore_errors=True) + side_root.mkdir(parents=True) + seed = side_root / "seed" + sha = git_sha() + run_ts = time.time() + print(f"python={sys.version}", flush=True) + print(f"sha={sha}", flush=True) + t0 = time.perf_counter() + ds = make_dataset(str(seed), side=side, max_prefetch=0, hedge_delay=0.0) + print( + f"indexed n={len(ds)} in {time.perf_counter() - t0:.2f}s " + f"ds.timeout={ds.download_timeout!r} cm.timeout={ds.cache_manager.download_timeout!r}", + flush=True, + ) + del ds + + wd = HangWatchdog(TIMEOUT) + wd.start() + results = [] + # Unique JSONL per run (never truncate a prior high-w log). + jsonl = OUT_DIR / f"raw_highw_post_timeout.{sha or 'unknown'}.{int(run_ts)}.jsonl" + try: + for w, pf in CELLS: + label = f"w{w}_p{pf}_defaults" + r = run_one( + label, + side=side, + num_workers=w, + max_prefetch=pf, + seed=seed, + wd=wd, + batches=300, + min_seconds=30.0, + prefetch_factor=2, + hedge_delay=0.0, + download_timeout=None, # use shipped default + sha=sha, + jsonl=jsonl, + ) + # Record effective default for clarity in JSON. + r = {**r, "download_timeout": DT_DEFAULT, "download_timeout_note": "shipped default (batch-level)"} + results.append(r) + print(f"DONE {label} ips={r['ips']:.1f}", flush=True) + finally: + wd.stop() + + out = { + "python": sys.version, + "git_sha": sha, + "note": "Focused remeasure after batch-level timeout fix; not a full grid resweep.", + "defaults": {"hedge_delay": 0.0, "download_timeout": DT_DEFAULT, "max_prefetch_cells": [0, 16]}, + "cells": results, + "by_key": {f"w{r['workers']}_p{r['prefetch']}": round(r["ips"], 1) for r in results}, + "jsonl": str(jsonl), + } + OUT_DIR.mkdir(parents=True, exist_ok=True) + path = unique_result_path("raw_highw_post_timeout", sha=sha, ts=run_ts) + path.write_text(json.dumps(out, indent=2)) + print(json.dumps(out["by_key"], indent=2), flush=True) + print(f"WROTE {path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_raw_lru_hitrate.py b/benchmarks/bench_raw_lru_hitrate.py new file mode 100644 index 000000000..62a3eb7be --- /dev/null +++ b/benchmarks/bench_raw_lru_hitrate.py @@ -0,0 +1,254 @@ +"""One-shot LRU hit/miss at w=8 p16 with LITDATA_RAW_DEBUG=1. + +Parses worker ``raw-debug: _download_batch done ... total_hit=... total_miss=...`` +lines (spawn workers log via the root logger to stderr). +""" + +from __future__ import annotations + +import json +import logging +import os +import re +import shutil +import sys +import tempfile +import time +from pathlib import Path + +os.environ["LITDATA_RAW_DEBUG"] = "1" +os.environ["PYTHONUNBUFFERED"] = "1" + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from bench_raw_before_vs_after import git_sha, unique_result_path +from torch.utils.data import DataLoader + +from litdata import StreamingRawDataset +from litdata.raw import dataset as raw_dataset + +raw_dataset._RAW_DEBUG = True + +INPUT = "/teamspace/s3_connections/imagenet-1m-template/raw/val" +OUT_DIR = Path(__file__).resolve().parent / "results" +ROOT = Path(tempfile.gettempdir()) / "litdata-raw-lru-hitrate" +LOG = OUT_DIR / "raw_lru_hitrate.log" +DONE_RE = re.compile( + r"raw-debug: _download_batch done pid=(\d+) n=(\d+) inflight=(\d+) " + r"batch_hit=(\d+) batch_miss=(\d+) total_hit=(\d+) total_miss=(\d+)" +) + + +class _Tee(logging.Handler): + """Append warning+ records to a shared log file.""" + + def __init__(self, path: Path) -> None: + super().__init__(level=logging.WARNING) + self.path = path + self.path.parent.mkdir(parents=True, exist_ok=True) + + def emit(self, record: logging.LogRecord) -> None: + try: + with open(self.path, "a", encoding="utf-8") as f: + f.write(self.format(record) + "\n") + except Exception: + self.handleError(record) + + +def main() -> None: + """Run w=8 p16 and report per-worker prefetch hit rates from debug logs.""" + if ROOT.exists(): + shutil.rmtree(ROOT, ignore_errors=True) + ROOT.mkdir(parents=True) + if LOG.exists(): + LOG.unlink() + + # File handler so spawn workers that reconfigure logging still write somewhere + # when LITDATA_RAW_DEBUG is on; also tee parent stderr. + root = logging.getLogger() + root.setLevel(logging.WARNING) + fmt = logging.Formatter("%(message)s") + sh = logging.StreamHandler(sys.stderr) + sh.setFormatter(fmt) + fh = _Tee(LOG) + fh.setFormatter(fmt) + root.handlers.clear() + root.addHandler(sh) + root.addHandler(fh) + # Ensure litdata.raw.dataset logger propagates. + logging.getLogger("litdata.raw.dataset").setLevel(logging.WARNING) + + print(f"python={sys.version}", flush=True) + cache = ROOT / "cache" + t0 = time.perf_counter() + ds = StreamingRawDataset( + INPUT, + cache_dir=str(cache), + cache_files=False, + max_prefetch=16, + hedge_delay=0, + download_timeout=120, + max_concurrent_downloads=64, + range_parallel_threshold=0, + ) + print(f"indexed n={len(ds)} in {time.perf_counter() - t0:.2f}s", flush=True) + # Force dataset module debug flag in children via env (already set). + loader = DataLoader( + ds, + batch_size=64, + num_workers=8, + shuffle=False, + multiprocessing_context="spawn", + persistent_workers=True, + prefetch_factor=2, + ) + it = iter(loader) + warm = max(1, 8 * 2) + for i in range(warm): + next(it) + print(f"warm {i + 1}/{warm}", flush=True) + t0 = time.perf_counter() + samples = 0 + batches = 0 + for _ in range(80): + batch = next(it) + samples += len(batch) + batches += 1 + elapsed = time.perf_counter() - t0 + ips = samples / elapsed if elapsed else 0.0 + print(f"timed {batches} batches → {ips:.1f} samples/s", flush=True) + del it, loader + + # Allow workers a moment to flush final debug lines. + time.sleep(1.0) + + # Also scrape stderr-captured file; workers may only print to their stderr + # which DataLoader forwards — capture by re-reading if we teed via a wrapper. + text = LOG.read_text(encoding="utf-8") if LOG.exists() else "" + # Fallback: some environments only forward worker stderr to our process stderr, + # which we did not tee. Re-run a tiny in-process sequential prefetch check too. + per_pid: dict[str, dict[str, int]] = {} + for m in DONE_RE.finditer(text): + pid, _n, _inf, _bh, _bm, th, tm = m.groups() + per_pid[pid] = {"total_hit": int(th), "total_miss": int(tm)} + + # In-process sequential: worker-aware schedule uses num_workers=1 → full max_prefetch. + ds_seq = StreamingRawDataset( + INPUT, + cache_dir=str(ROOT / "seq"), + cache_files=False, + max_prefetch=16, + hedge_delay=0, + download_timeout=120, + max_concurrent_downloads=64, + range_parallel_threshold=0, + ) + idx = cache / "index.json.zstd" + if idx.exists(): + (ROOT / "seq").mkdir(parents=True, exist_ok=True) + shutil.copy2(idx, ROOT / "seq" / "index.json.zstd") + ds_seq = StreamingRawDataset( + INPUT, + cache_dir=str(ROOT / "seq"), + cache_files=False, + max_prefetch=16, + hedge_delay=0, + download_timeout=120, + max_concurrent_downloads=64, + range_parallel_threshold=0, + ) + bs = 64 + for start in range(0, bs * 30, bs): + ds_seq.__getitems__(list(range(start, start + bs))) + seq_hits, seq_misses = ds_seq._prefetch_hits, ds_seq._prefetch_misses + seq_total = seq_hits + seq_misses + + # Simulate w=8 stride in-process: each "worker" advances by 8 batches. + ds_w = StreamingRawDataset( + INPUT, + cache_dir=str(ROOT / "w8sim"), + cache_files=False, + max_prefetch=16, + hedge_delay=0, + download_timeout=120, + max_concurrent_downloads=64, + range_parallel_threshold=0, + ) + if idx.exists(): + (ROOT / "w8sim").mkdir(parents=True, exist_ok=True) + shutil.copy2(idx, ROOT / "w8sim" / "index.json.zstd") + ds_w = StreamingRawDataset( + INPUT, + cache_dir=str(ROOT / "w8sim"), + cache_files=False, + max_prefetch=16, + hedge_delay=0, + download_timeout=120, + max_concurrent_downloads=64, + range_parallel_threshold=0, + ) + + # Monkeypatch get_worker_info so _schedule_prefetch thinks num_workers=8. + class _Info: + num_workers = 8 + id = 0 + + real_schedule = ds_w._schedule_prefetch + + def schedule_as_w8(indices: list[int]) -> None: + import torch.utils.data + + real_get = torch.utils.data.get_worker_info + torch.utils.data.get_worker_info = lambda: _Info() # type: ignore[assignment] + try: + real_schedule(indices) + finally: + torch.utils.data.get_worker_info = real_get + + ds_w._schedule_prefetch = schedule_as_w8 # type: ignore[method-assign] + # Worker 0 batches: 0, 8, 16, ... in batch-index space → sample starts 0, 512, 1024, ... + for bi in range(0, 40, 8): + start = bi * bs + ds_w.__getitems__(list(range(start, start + bs))) + w_hits, w_misses = ds_w._prefetch_hits, ds_w._prefetch_misses + w_total = w_hits + w_misses + + worker_hits = sum(v["total_hit"] for v in per_pid.values()) + worker_misses = sum(v["total_miss"] for v in per_pid.values()) + worker_total = worker_hits + worker_misses + + out = { + "python": sys.version, + "cell": {"workers": 8, "max_prefetch": 16, "effective_prefetch": min(16, 64 // 8)}, + "timed": {"batches": batches, "ips": ips, "elapsed": elapsed}, + "spawn_workers_from_log": { + "pids": len(per_pid), + "total_hit": worker_hits, + "total_miss": worker_misses, + "hit_rate": (worker_hits / worker_total) if worker_total else None, + "per_pid": per_pid, + "log_lines": text.count("raw-debug: _download_batch done"), + "note": "None hit_rate means spawn workers did not share the parent log file", + }, + "in_process_sequential_w1": { + "prefetch_hits": seq_hits, + "prefetch_misses": seq_misses, + "hit_rate": (seq_hits / seq_total) if seq_total else 0.0, + }, + "in_process_simulated_w8": { + "note": "stride every 8th batch; effective look-ahead=8", + "prefetch_hits": w_hits, + "prefetch_misses": w_misses, + "hit_rate": (w_hits / w_total) if w_total else 0.0, + }, + } + OUT_DIR.mkdir(parents=True, exist_ok=True) + out["git_sha"] = git_sha() + path = unique_result_path("raw_lru_hitrate", sha=out["git_sha"]) + path.write_text(json.dumps(out, indent=2)) + print(json.dumps(out, indent=2), flush=True) + print(f"WROTE {path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_raw_opt.py b/benchmarks/bench_raw_opt.py new file mode 100644 index 000000000..6d86fb815 --- /dev/null +++ b/benchmarks/bench_raw_opt.py @@ -0,0 +1,150 @@ +"""A/B microbench for StreamingRawDataset optimizations on ImageNet val.""" + +from __future__ import annotations + +import argparse +import inspect +import os +import shutil +import sys +import tempfile +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from torch.utils.data import DataLoader +from tqdm import tqdm +from uvloop_status import log_loop_runner_backend, uvloop_package_status + + +def clear_dir(path: str) -> None: + """Remove ``path`` if it is an existing directory.""" + if os.path.isdir(path): + shutil.rmtree(path, ignore_errors=True) + + +def run_once( + *, + label: str, + input_dir: str, + cache_dir: str, + batch_size: int, + num_workers: int, + num_batches: int, + max_prefetch: int, + clear_cache: bool, +) -> dict: + """Run one StreamingRawDataset throughput trial and return timing stats.""" + from litdata import StreamingRawDataset + + if clear_cache: + clear_dir(cache_dir) + + kwargs: dict = { + "input_dir": input_dir, + "cache_dir": cache_dir, + "cache_files": False, + "recompute_index": False, + "transform": None, + } + sig = inspect.signature(StreamingRawDataset.__init__) + if "max_prefetch" in sig.parameters: + kwargs["max_prefetch"] = max_prefetch + if "max_concurrent_downloads" in sig.parameters: + kwargs["max_concurrent_downloads"] = 64 + + t_index = time.perf_counter() + ds = StreamingRawDataset(**kwargs) + index_s = time.perf_counter() - t_index + n = len(ds) + log_loop_runner_backend(print, prefix=f"[{label}]") + + loader = DataLoader(ds, batch_size=batch_size, num_workers=num_workers, shuffle=False) + it = iter(loader) + + warm_t0 = time.perf_counter() + warm = next(it) + warm_s = time.perf_counter() - warm_t0 + + samples = 0 + t0 = time.perf_counter() + for i, batch in enumerate(tqdm(it, total=num_batches, desc=label, leave=False)): + samples += len(batch) + if i + 1 >= num_batches: + break + elapsed = time.perf_counter() - t0 + ips = samples / elapsed if elapsed > 0 else 0.0 + + storage = getattr(ds, "_storage_path", None) or getattr(ds.cache_manager, "_input_dir_path", "?") + print( + f"[{label}] index={index_s:.2f}s ({n} files) storage={storage!r}\n" + f" warm1={warm_s:.2f}s ({len(warm)} samples) | " + f"{num_batches} batches / {samples} samples in {elapsed:.2f}s " + f"→ {ips:.1f} samples/s (prefetch={max_prefetch}, workers={num_workers}, bs={batch_size})" + ) + return { + "label": label, + "storage_path": storage, + "indexed": n, + "index_s": index_s, + "warm_batch_s": warm_s, + "batches": num_batches, + "samples": samples, + "elapsed_s": elapsed, + "samples_per_s": ips, + "max_prefetch": max_prefetch, + } + + +def main() -> None: + """CLI entrypoint for prefetch A/B microbenchmarks.""" + p = argparse.ArgumentParser() + p.add_argument("--input_dir", default="/teamspace/s3_connections/imagenet-1m-template/raw/val") + p.add_argument("--cache_root", default=str(Path(tempfile.gettempdir()) / "litdata-raw-bench")) + p.add_argument("--batch_size", type=int, default=64) + p.add_argument("--num_workers", type=int, default=4) + p.add_argument("--num_batches", type=int, default=20) + args = p.parse_args() + + os.makedirs(args.cache_root, exist_ok=True) + print(f"uvloop package: {uvloop_package_status()}") + results = [] + + common = { + "input_dir": args.input_dir, + "batch_size": args.batch_size, + "num_workers": args.num_workers, + "num_batches": args.num_batches, + "clear_cache": True, + } + + results.append( + run_once( + label="new-studio-path-prefetch0", + cache_dir=str(Path(args.cache_root) / "new0"), + max_prefetch=0, + **common, + ) + ) + results.append( + run_once( + label="new-studio-path-prefetch128", + cache_dir=str(Path(args.cache_root) / "new128"), + max_prefetch=max(128, 2 * args.batch_size), + **common, + ) + ) + + print("\n=== Summary ===") + for r in results: + print( + f"{r['label']:32s} {r['samples_per_s']:8.1f} samples/s " + f"index={r['index_s']:.2f}s warm={r['warm_batch_s']:.2f}s storage={r['storage_path']}" + ) + if len(results) >= 2 and results[0]["samples_per_s"] > 0: + print(f"\nPrefetch speedup: {results[1]['samples_per_s'] / results[0]['samples_per_s']:.2f}x") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_raw_ranged_vs_whole.py b/benchmarks/bench_raw_ranged_vs_whole.py new file mode 100644 index 000000000..ef4ee9467 --- /dev/null +++ b/benchmarks/bench_raw_ranged_vs_whole.py @@ -0,0 +1,286 @@ +"""Focused ranged vs whole-object compare on fixed StreamingRawDataset tree. + +Protocol: timed window is max(BATCHES, MIN_SECONDS) — both floors required. +Artifacts use SHA/ts-suffixed paths (never overwrite). +""" + +from __future__ import annotations + +import json +import os +import shutil +import sys +import tempfile +import threading +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from bench_raw_before_vs_after import effective_min_seconds, git_sha, unique_result_path +from torch.utils.data import DataLoader +from uvloop_status import log_loop_runner_backend, uvloop_package_status + +from litdata import StreamingRawDataset + +INPUT = "/teamspace/s3_connections/imagenet-1m-template/raw/val" +ROOT = Path(tempfile.gettempdir()) / "litdata-raw-ranged-vs-whole" +OUT_DIR = Path(__file__).resolve().parent / "results" +BS = 64 +BATCHES = 300 +MIN_SECONDS = 30.0 +TIMEOUT = 600.0 +CONFIGS = [(4, 0), (4, 128), (8, 0), (8, 128)] +MODES = [ + ("whole_object", 0), + ("default_32MiB", 33_554_432), + ("force_ranged", 1), +] + + +def log(msg: str) -> None: + """Print a timestamped benchmark log line.""" + print(f"{time.strftime('%H:%M:%S')} {msg}", flush=True) + + +class HangWatchdog: + """Kill the process if a step exceeds ``timeout_s`` without heartbeat.""" + + def __init__(self, timeout_s: float) -> None: + """Initialize the watchdog with a hang timeout in seconds.""" + self.timeout_s = timeout_s + self._label = "init" + self._beat = time.monotonic() + self._stop = threading.Event() + self._t = threading.Thread(target=self._run, daemon=True) + + def start(self) -> None: + """Start the background watchdog thread.""" + self._t.start() + + def beat(self, label: str) -> None: + """Record progress so the watchdog does not abort.""" + self._label = label + self._beat = time.monotonic() + + def stop(self) -> None: + """Stop the background watchdog thread.""" + self._stop.set() + + def _run(self) -> None: + while not self._stop.wait(1.0): + idle = time.monotonic() - self._beat + if idle > self.timeout_s: + log(f"HANG at '{self._label}' after {idle:.1f}s — abort") + os._exit(124) + + +def copy_index(src: Path, dst: Path) -> None: + """Copy a cached index tree from ``src`` to ``dst``.""" + if dst.exists(): + shutil.rmtree(dst, ignore_errors=True) + dst.mkdir(parents=True) + for p in src.iterdir(): + if p.is_dir(): + shutil.copytree(p, dst / p.name) + else: + shutil.copy2(p, dst / p.name) + + +def run(label: str, *, num_workers: int, max_prefetch: int, threshold: int, seed: Path, wd: HangWatchdog) -> dict: + """Run one ranged-vs-whole trial and return timing stats.""" + cache = ROOT / label + wd.beat(f"{label}: setup") + copy_index(seed, cache) + ds = StreamingRawDataset( + INPUT, + cache_dir=str(cache), + cache_files=False, + max_prefetch=max_prefetch, + max_concurrent_downloads=64, + range_parallel_threshold=threshold, + ) + loader = DataLoader( + ds, + batch_size=BS, + num_workers=num_workers, + shuffle=False, + multiprocessing_context="spawn", + persistent_workers=True, + ) + it = iter(loader) + wd.beat(f"{label}: warm") + t0 = time.perf_counter() + _ = next(it) + warm_s = time.perf_counter() - t0 + + min_s = effective_min_seconds(num_workers, MIN_SECONDS) + samples = 0 + timed_batches = 0 + wd.beat(f"{label}: timed") + t0 = time.perf_counter() + while True: + batch = next(it) + samples += len(batch) + timed_batches += 1 + wd.beat(f"{label}: batch {timed_batches}") + elapsed = time.perf_counter() - t0 + if timed_batches >= BATCHES and elapsed >= min_s: + break + elapsed = time.perf_counter() - t0 + ips = samples / elapsed if elapsed else 0.0 + log( + f"[{label}] thr={threshold} w={num_workers} pf={max_prefetch} " + f"warm={warm_s:.2f}s | {timed_batches} batches/{samples} in {elapsed:.2f}s " + f"(need ≥{BATCHES} & ≥{min_s:.0f}s) → {ips:.1f} samples/s" + ) + del it, loader, ds + return { + "label": label, + "mode": label.rsplit("_w", 1)[0], + "range_parallel_threshold": threshold, + "workers": num_workers, + "prefetch": max_prefetch, + "batches": timed_batches, + "min_seconds_effective": min_s, + "ips": ips, + "warm_s": warm_s, + "elapsed": elapsed, + "samples": samples, + } + + +def main() -> None: + """CLI entrypoint for ranged vs whole-object comparisons.""" + if ROOT.exists(): + shutil.rmtree(ROOT, ignore_errors=True) + ROOT.mkdir(parents=True) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + wd = HangWatchdog(TIMEOUT) + wd.start() + try: + log(f"uvloop package: {uvloop_package_status()}") + log( + f"ranged-vs-whole input={INPUT} bs={BS} batches={BATCHES} " + f"mp=spawn persistent_workers configs={len(CONFIGS) * len(MODES)}" + ) + + wd.beat("index seed") + seed = ROOT / "seed" + t0 = time.perf_counter() + ds = StreamingRawDataset(INPUT, cache_dir=str(seed), cache_files=False, max_prefetch=0) + n_files = len(ds) + storage = ds._storage_path + sizes: list[int] = [] + try: + files = getattr(getattr(ds, "_index", None), "files", None) + if files: + for i in range(min(200, len(files))): + meta = files[i] + sz = getattr(meta, "size", None) or (meta.get("size") if isinstance(meta, dict) else None) + if sz is not None: + sizes.append(int(sz)) + except Exception as e: + log(f"size probe skipped: {e}") + log(f"Indexed {n_files} files in {time.perf_counter() - t0:.2f}s storage={storage}") + if sizes: + log( + f"sample sizes (n={len(sizes)}): " + f"min={min(sizes)} avg={sum(sizes) // len(sizes)} max={max(sizes)} " + f"(<< 32MiB → default threshold uses whole-object)" + ) + log_loop_runner_backend(log, prefix="after index seed") + del ds + + results = [] + for mode_name, thr in MODES: + for w, pf in CONFIGS: + label = f"{mode_name}_w{w}_p{pf}" + results.append(run(label, num_workers=w, max_prefetch=pf, threshold=thr, seed=seed, wd=wd)) + + log("\n=== Comparison table (samples/s) ===") + header = ( + f"{'workers':>8} {'prefetch':>8} | " + f"{'whole(thr=0)':>14} {'default(32MiB)':>14} {'force_ranged(1)':>16} | " + f"{'ranged/whole':>12}" + ) + log(header) + log("-" * len(header)) + by = {(r["range_parallel_threshold"], r["workers"], r["prefetch"]): r for r in results} + for w, pf in CONFIGS: + a = by[(0, w, pf)]["ips"] + b = by[(33_554_432, w, pf)]["ips"] + c = by[(1, w, pf)]["ips"] + ratio = c / a if a else float("nan") + log(f"{w:>8} {pf:>8} | {a:>14.1f} {b:>14.1f} {c:>16.1f} | {ratio:>11.2f}x") + + winners = [] + for w, pf in CONFIGS: + rows = [by[(thr, w, pf)] for _, thr in MODES] + best = max(rows, key=lambda r: r["ips"]) + winners.append( + { + "workers": w, + "prefetch": pf, + "best_mode": best["label"].rsplit("_w", 1)[0], + "ips": best["ips"], + } + ) + + mode_means = {} + for mode_name, thr in MODES: + vals = [r["ips"] for r in results if r["range_parallel_threshold"] == thr] + mode_means[mode_name] = sum(vals) / len(vals) + + overall_winner = max(mode_means.items(), key=lambda kv: kv[1])[0] + log("\nMode mean samples/s: " + ", ".join(f"{k}={v:.1f}" for k, v in mode_means.items())) + log(f"Overall winner (mean across configs): {overall_winner}") + + for w, pf in [(8, 0), (8, 128)]: + a = by[(0, w, pf)]["ips"] + c = by[(1, w, pf)]["ips"] + winner = "force_ranged" if c > a else ("whole_object" if a > c else "tie") + log(f"w={w} pf={pf}: whole={a:.1f} vs force_ranged={c:.1f} → {winner} ({max(a, c) / min(a, c):.2f}x)") + + payload = { + "meta": { + "input": INPUT, + "storage": storage, + "n_files": n_files, + "batch_size": BS, + "batches": BATCHES, + "multiprocessing_context": "spawn", + "persistent_workers": True, + "max_concurrent_downloads": 64, + "uvloop": uvloop_package_status(), + "note": ( + "ImageNet val JPEGs are ~50-200KiB; default 32MiB threshold never engages " + "ranged GETs. force_ranged uses threshold=1 to exercise the ranged path." + ), + "sample_sizes": ( + {"n": len(sizes), "min": min(sizes), "avg": sum(sizes) // len(sizes), "max": max(sizes)} + if sizes + else None + ), + "old_sweep_killed": True, + "old_sweep_json": None, + "old_sweep_log": "benchmarks/results/raw_worker_prefetch_sweep.log", + "old_sweep_used_fixed_downloaders": False, + }, + "modes": dict(MODES), + "results": results, + "mode_means": mode_means, + "winners_per_config": winners, + "overall_winner": overall_winner, + } + OUT_DIR.mkdir(parents=True, exist_ok=True) + path = unique_result_path("raw_ranged_vs_whole", sha=git_sha()) + path.write_text(json.dumps(payload, indent=2) + "\n") + log(f"Wrote {path}") + finally: + wd.stop() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_raw_workers.py b/benchmarks/bench_raw_workers.py new file mode 100644 index 000000000..63d665e9e --- /dev/null +++ b/benchmarks/bench_raw_workers.py @@ -0,0 +1,249 @@ +"""Exhaustive worker × prefetch sweep for StreamingRawDataset. + +Uses spawn + persistent_workers after a warm index. Writes a JSON summary for docs. + +Protocol: timed window is max(BATCHES, MIN_SECONDS) — both floors required. +High-worker cells (w≥16) always run ≥30s. Artifacts use SHA/ts-suffixed paths. + +Ranged vs whole-object compare (optional): + LITDATA_RAW_RANGE_PARALLEL_THRESHOLD=0 # whole-object GETs (also the dataset default) + LITDATA_RAW_RANGE_PARALLEL_THRESHOLD=1 # force ranged for any sized object + LITDATA_RAW_RANGE_PARALLEL_THRESHOLD=33554432 # opt in at 32MiB +""" + +from __future__ import annotations + +import json +import os +import shutil +import sys +import tempfile +import threading +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from bench_raw_before_vs_after import effective_min_seconds, git_sha, unique_result_path +from torch.utils.data import DataLoader +from uvloop_status import log_loop_runner_backend, uvloop_package_status + +from litdata import StreamingRawDataset + +INPUT = "/teamspace/s3_connections/imagenet-1m-template/raw/val" +ROOT = Path(tempfile.gettempdir()) / "litdata-raw-worker-sweep" +OUT_DIR = Path(__file__).resolve().parent / "results" +BS = 64 +BATCHES = 300 +MIN_SECONDS = 30.0 +# Up to host vCPUs (4×L4 Studio = 48). +WORKERS = [0, 1, 2, 4, 8, 16, 24, 32, 48] +PREFETCH = [0, 16, 32, 64, 96, 128] +TIMEOUT = 600.0 +OLD_FUSE = 75.2 + +# Optional override for ranged-vs-whole compare; None → dataset default (0 = opt-in off). +_RANGE_ENV = os.getenv("LITDATA_RAW_RANGE_PARALLEL_THRESHOLD") +RANGE_PARALLEL_THRESHOLD: int | None = int(_RANGE_ENV) if _RANGE_ENV is not None else None + + +def log(msg: str) -> None: + """Print a timestamped benchmark log line.""" + print(f"{time.strftime('%H:%M:%S')} {msg}", flush=True) + + +class HangWatchdog: + """Kill the process if a step exceeds ``timeout_s`` without heartbeat.""" + + def __init__(self, timeout_s: float) -> None: + """Initialize the watchdog with a hang timeout in seconds.""" + self.timeout_s = timeout_s + self._label = "init" + self._beat = time.monotonic() + self._stop = threading.Event() + self._t = threading.Thread(target=self._run, daemon=True) + + def start(self) -> None: + """Start the background watchdog thread.""" + self._t.start() + + def beat(self, label: str) -> None: + """Record progress so the watchdog does not abort.""" + self._label = label + self._beat = time.monotonic() + + def stop(self) -> None: + """Stop the background watchdog thread.""" + self._stop.set() + + def _run(self) -> None: + while not self._stop.wait(1.0): + idle = time.monotonic() - self._beat + if idle > self.timeout_s: + log(f"HANG at '{self._label}' after {idle:.1f}s — abort") + os._exit(124) + + +def copy_index(src: Path, dst: Path) -> None: + """Copy a cached index tree from ``src`` to ``dst``.""" + if dst.exists(): + shutil.rmtree(dst, ignore_errors=True) + dst.mkdir(parents=True) + for p in src.iterdir(): + if p.is_dir(): + shutil.copytree(p, dst / p.name) + else: + shutil.copy2(p, dst / p.name) + + +def run(label: str, *, num_workers: int, max_prefetch: int, seed: Path, wd: HangWatchdog) -> dict: + """Run one worker/prefetch trial and return timing stats.""" + cache = ROOT / label + wd.beat(f"{label}: setup") + copy_index(seed, cache) + ds_kwargs: dict = { + "cache_dir": str(cache), + "cache_files": False, + "max_prefetch": max_prefetch, + "max_concurrent_downloads": 64, + } + if RANGE_PARALLEL_THRESHOLD is not None: + ds_kwargs["range_parallel_threshold"] = RANGE_PARALLEL_THRESHOLD + ds = StreamingRawDataset(INPUT, **ds_kwargs) + kwargs: dict = {"batch_size": BS, "num_workers": num_workers, "shuffle": False} + if num_workers > 0: + kwargs["multiprocessing_context"] = "spawn" + kwargs["persistent_workers"] = True + loader = DataLoader(ds, **kwargs) + it = iter(loader) + wd.beat(f"{label}: warm") + t0 = time.perf_counter() + next(it) + warm_s = time.perf_counter() - t0 + + min_s = effective_min_seconds(num_workers, MIN_SECONDS) + samples = 0 + timed_batches = 0 + wd.beat(f"{label}: timed") + t0 = time.perf_counter() + while True: + batch = next(it) + samples += len(batch) + timed_batches += 1 + wd.beat(f"{label}: batch {timed_batches}") + elapsed = time.perf_counter() - t0 + if timed_batches >= BATCHES and elapsed >= min_s: + break + elapsed = time.perf_counter() - t0 + ips = samples / elapsed if elapsed else 0.0 + log( + f"[{label}] w={num_workers} pf={max_prefetch} " + f"warm={warm_s:.2f}s | {timed_batches}×{samples // max(timed_batches, 1)} in {elapsed:.2f}s " + f"(need ≥{BATCHES} batches & ≥{min_s:.0f}s) → {ips:.1f} samples/s ({ips / OLD_FUSE:.1f}x FUSE)" + ) + del it, loader, ds + return { + "label": label, + "workers": num_workers, + "prefetch": max_prefetch, + "batches": timed_batches, + "min_seconds_effective": min_s, + "ips": ips, + "warm_s": warm_s, + "elapsed": elapsed, + "samples": samples, + } + + +def print_matrix(results: list[dict]) -> None: + """Print a workers × prefetch samples/s matrix.""" + by_key = {(r["workers"], r["prefetch"]): r["ips"] for r in results} + header = f"{'w/pf':>6}" + "".join(f"{p:>10}" for p in PREFETCH) + log("\n=== Matrix (samples/s) ===") + log(header) + for w in WORKERS: + row = f"{w:>6}" + for p in PREFETCH: + ips = by_key.get((w, p)) + row += f"{ips:>10.1f}" if ips is not None else f"{'—':>10}" + log(row) + + +def main() -> None: + """CLI entrypoint for the exhaustive worker × prefetch sweep.""" + if ROOT.exists(): + shutil.rmtree(ROOT, ignore_errors=True) + ROOT.mkdir(parents=True) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + wd = HangWatchdog(TIMEOUT) + wd.start() + ncpu = os.cpu_count() or 0 + n_configs = len(WORKERS) * len(PREFETCH) + log(f"uvloop package: {uvloop_package_status()}") + log( + f"Exhaustive sweep input={INPUT} bs={BS} batches={BATCHES} " + f"mp=spawn persistent_workers cpus={ncpu} configs={n_configs}" + ) + log(f"WORKERS={WORKERS}") + log(f"PREFETCH={PREFETCH}") + log( + f"range_parallel_threshold=" + f"{RANGE_PARALLEL_THRESHOLD if RANGE_PARALLEL_THRESHOLD is not None else 'default(32MiB)'}" + ) + + wd.beat("index seed") + seed = ROOT / "seed" + t0 = time.perf_counter() + seed_kwargs: dict = {"cache_dir": str(seed), "cache_files": False, "max_prefetch": 0} + if RANGE_PARALLEL_THRESHOLD is not None: + seed_kwargs["range_parallel_threshold"] = RANGE_PARALLEL_THRESHOLD + ds = StreamingRawDataset(INPUT, **seed_kwargs) + n_files = len(ds) + storage = ds._storage_path + log(f"Indexed {n_files} files in {time.perf_counter() - t0:.2f}s storage={storage}") + log_loop_runner_backend(log, prefix="after index seed") + del ds + + results: list[dict] = [] + try: + for w in WORKERS: + for pf in PREFETCH: + label = f"w{w}_p{pf}" + results.append(run(label, num_workers=w, max_prefetch=pf, seed=seed, wd=wd)) + + print_matrix(results) + best = max(results, key=lambda r: r["ips"]) + log( + f"\nBest: {best['label']} → {best['ips']:.1f} samples/s ({best['ips'] / OLD_FUSE:.1f}x vs FUSE ~{OLD_FUSE})" + ) + + payload = { + "meta": { + "input": INPUT, + "storage": storage, + "n_files": n_files, + "batch_size": BS, + "batches": BATCHES, + "multiprocessing_context": "spawn", + "persistent_workers": True, + "max_concurrent_downloads": 64, + "cpus": ncpu, + "fuse_baseline_samples_per_s": OLD_FUSE, + "workers": WORKERS, + "prefetch": PREFETCH, + "range_parallel_threshold": RANGE_PARALLEL_THRESHOLD, + }, + "results": results, + "best": best, + } + OUT_DIR.mkdir(parents=True, exist_ok=True) + path = unique_result_path("raw_worker_prefetch_sweep", sha=git_sha()) + path.write_text(json.dumps(payload, indent=2) + "\n") + log(f"Wrote {path}") + finally: + wd.stop() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/stage1_ab_resume.sh b/benchmarks/stage1_ab_resume.sh new file mode 100644 index 000000000..dbac81517 --- /dev/null +++ b/benchmarks/stage1_ab_resume.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash +# Append-only Stage 1 A/B resume into LITDATA_BENCH_ACCUM_OUT JSON files. +# Usage example (confirm gate fix): +# BEFORE_ACCUM=... AFTER_ACCUM=... WORKERS=8,16,24 REPEATS_PAIRS=3 \ +# bash benchmarks/stage1_ab_resume.sh +set -euo pipefail +cd "$(dirname "$0")/.." +SCRIPT=benchmarks/bench_raw_before_vs_after.py +BEFORE_PY=${BEFORE_PY:-/tmp/litdata-raw-pre-stage1/src} +AFTER_PY=${AFTER_PY:-src} +BEFORE_ACCUM=${BEFORE_ACCUM:?set BEFORE_ACCUM to existing before JSON} +AFTER_ACCUM=${AFTER_ACCUM:?set AFTER_ACCUM to existing after JSON} +WORKERS=${WORKERS:-2,4,8,16,24} +PREFETCH=${PREFETCH:-0,16} +CATCH_UP_AFTER=${CATCH_UP_AFTER:-0} +REPEATS_PAIRS=${REPEATS_PAIRS:-3} +GIT_SHA=${LITDATA_BENCH_GIT:-$(git rev-parse --short HEAD)} +PY=${PY:-python} + +log() { echo "$(date -u +%H:%M:%S) $*"; } + +run_side() { + local side=$1 pypath=$2 accum=$3 + log "=== resume side=$side accum=$(basename "$accum") ===" + env \ + PYTHONPATH="$pypath" \ + LITDATA_BENCH_GIT="$GIT_SHA" \ + LITDATA_BENCH_ACCUM_OUT="$(pwd)/$accum" \ + "$PY" "$SCRIPT" \ + --side "$side" \ + --workers "$WORKERS" \ + --after-prefetch "$PREFETCH" \ + --batches 300 \ + --min-seconds 30 \ + --repeats 1 +} + +log "Stage 1 A/B resume start pid=$$ sha=$GIT_SHA workers=$WORKERS pairs=$REPEATS_PAIRS" +if [[ "$CATCH_UP_AFTER" == "1" ]]; then + run_side after "$AFTER_PY" "$AFTER_ACCUM" +fi +for ((i=1; i<=REPEATS_PAIRS; i++)); do + log "=== pair $i/$REPEATS_PAIRS ===" + run_side before "$BEFORE_PY" "$BEFORE_ACCUM" + run_side after "$AFTER_PY" "$AFTER_ACCUM" +done +log "=== merge ===" +env PYTHONPATH=src LITDATA_BENCH_GIT="$GIT_SHA" "$PY" "$SCRIPT" --merge +log "Stage 1 A/B resume complete" diff --git a/benchmarks/uvloop_status.py b/benchmarks/uvloop_status.py new file mode 100644 index 000000000..e14f726ee --- /dev/null +++ b/benchmarks/uvloop_status.py @@ -0,0 +1,35 @@ +"""Shared uvloop detection and LoopRunner backend logging for raw benchmarks.""" + +from __future__ import annotations + +from collections.abc import Callable + + +def uvloop_package_status() -> str: + """Return a short string describing uvloop install and preferred loop backend.""" + from litdata.raw.dataset import _loop_backend_name + + try: + import uvloop + except ImportError: + return "not installed (stdlib asyncio fallback)" + version = getattr(uvloop, "__version__", "?") + backend = _loop_backend_name() + return f"available (uvloop {version}; create→{backend})" + + +def log_loop_runner_backend(log_fn: Callable[[str], None], *, prefix: str = "") -> bool: + """Log whether the process-local LoopRunner loop is uvloop-backed.""" + from litdata.raw.dataset import _get_loop_runner, _loop_backend_name + + runner = _get_loop_runner() + loop = runner.loop + loop_type = f"{type(loop).__module__}.{type(loop).__name__}" + active = type(loop).__module__.startswith("uvloop") + tag = f"{prefix} " if prefix else "" + log_fn( + f"{tag}LoopRunner event loop: {loop_type} " + f"(preferred={_loop_backend_name()}, " + f"{'uvloop active' if active else 'stdlib asyncio'}) pid={runner.pid}" + ) + return active diff --git a/docs/source/conf.py b/docs/source/conf.py index 9626a3734..763280a72 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -265,12 +265,12 @@ def _convert_markdown(path_in: str, path_out: str) -> None: # -- Options for intersphinx extension --------------------------------------- -# Example configuration for intersphinx: refer to the Python standard library. -intersphinx_mapping = { - "python": ("https://docs.python.org/3", None), - "torch": ("https://pytorch.org/docs/stable/", None), - "numpy": ("https://numpy.org/doc/stable/", None), -} +# Remote inventory fetches are unused here (no :external: / inventory roles) and +# flake under Sphinx ``-W`` when hosts reset (e.g. docs.python.org ConnectionReset). +# Keep the extension enabled for local experiments; leave the mapping empty so +# ``make linkcheck`` / CI does not depend on third-party ``objects.inv`` availability. +# If re-enabling, prefer current hosts (torch moved to docs.pytorch.org). +intersphinx_mapping = {} # -- Options for todo extension ---------------------------------------------- diff --git a/requirements/extras.txt b/requirements/extras.txt index 2142208a2..f14d565eb 100644 --- a/requirements/extras.txt +++ b/requirements/extras.txt @@ -6,3 +6,4 @@ lightning-sdk==2026.2.6 # Must be pinned to ensure compatibility google-cloud-storage polars fsspec +uvloop; sys_platform != "win32" # optional faster asyncio loop for StreamingRawDataset diff --git a/src/litdata/processing/data_processor.py b/src/litdata/processing/data_processor.py index 3b869ad71..a6c35f4f9 100644 --- a/src/litdata/processing/data_processor.py +++ b/src/litdata/processing/data_processor.py @@ -52,7 +52,7 @@ from litdata.streaming.dataloader import StreamingDataLoader from litdata.streaming.fs_provider import _get_fs_provider, not_supported_provider from litdata.streaming.item_loader import BaseItemLoader -from litdata.streaming.resolver import _resolve_dir +from litdata.streaming.resolver import _has_time_template, _resolve_dir from litdata.utilities._pytree import tree_flatten, tree_unflatten, treespec_loads from litdata.utilities.broadcast import broadcast_object from litdata.utilities.dataset_utilities import load_index_file @@ -1133,6 +1133,7 @@ def __init__( storage_options: dict[str, Any] = {}, keep_data_ordered: bool = True, verbose: bool = True, + broadcast_paths: bool = False, ): """Provides an efficient way to process data across multiple machine into chunks to make training faster. @@ -1162,6 +1163,13 @@ def __init__( storage_options: Storage options for the cloud provider. keep_data_ordered: Whether to use a shared queue for the workers or not. verbose: Whether to print the progress & logs of the workers. Defaults to True. + broadcast_paths: When ``True``, broadcast resolved ``input_dir`` / ``output_dir`` across nodes via + :func:`~litdata.utilities.broadcast.broadcast_object` (Studio multi-node, when + ``LIGHTNING_APP_EXTERNAL_URL`` is set). Defaults to ``False``. Automatically enabled when + ``input_dir`` or ``output_dir`` contains a ``{%strftime}`` time template (so every rank shares + the same expanded path). When ``False`` and no time template is present, each rank keeps its + locally resolved path — fine for stable ``s3://`` / connection paths, but unsafe if ranks + would otherwise expand different timestamps. """ # spawn doesn't work in IPython start_method = start_method or ("fork" if in_notebook() else "spawn") @@ -1176,6 +1184,9 @@ def __init__( multiprocessing.set_start_method(start_method, force=True) + # Detect time templates on the unresolved path strings (before `_resolve_dir` expands them). + self.broadcast_paths = broadcast_paths or _has_time_template(input_dir) or _has_time_template(output_dir) + self.input_dir = _resolve_dir(input_dir) self.output_dir = _resolve_dir(output_dir) @@ -1209,16 +1220,14 @@ def __init__( if self.reader is not None and self.weights is not None: raise ValueError("Either the reader or the weights needs to be defined.") - # Ensure the input dir is the same across all nodes - self.input_dir = broadcast_object("input_dir", self.input_dir, rank=_get_node_rank()) + if self.broadcast_paths: + # Align resolved dirs across nodes (needed when `{%strftime}` expands per-rank). + self.input_dir = broadcast_object("input_dir", self.input_dir, rank=_get_node_rank()) + if self.output_dir: + self.output_dir = broadcast_object("output_dir", self.output_dir, rank=_get_node_rank()) - if self.output_dir: - # Ensure the output dir is the same across all nodes - self.output_dir = broadcast_object("output_dir", self.output_dir, rank=_get_node_rank()) - if verbose: - print( - f"Storing the files under {self.output_dir.path if self.output_dir.path else self.output_dir.url}" - ) + if self.output_dir and verbose: + print(f"Storing the files under {self.output_dir.path if self.output_dir.path else self.output_dir.url}") self.random_seed = random_seed self.verbose = verbose diff --git a/src/litdata/processing/functions.py b/src/litdata/processing/functions.py index 2a01c261f..d0f8e88b8 100644 --- a/src/litdata/processing/functions.py +++ b/src/litdata/processing/functions.py @@ -48,6 +48,7 @@ _assert_dir_has_index_file, _assert_dir_is_empty, _execute, + _has_time_template, _resolve_dir, ) from litdata.utilities._pytree import tree_flatten @@ -259,6 +260,7 @@ def map( optimize_dns: bool | None = None, storage_options: dict[str, Any] = {}, keep_data_ordered: bool = True, + broadcast_paths: bool = False, ) -> None: """Maps a callable over a collection of inputs, possibly in a distributed way. @@ -289,6 +291,8 @@ def map( workload and reduce idle time when some workers finish early. This may lead to unordered processing of items. If True, each worker processes a statically assigned subset of items in order. + broadcast_paths: Broadcast resolved input/output dirs across multi-node ranks. Defaults to ``False``. + Auto-enabled when ``input_dir`` or ``output_dir`` contains a ``{%strftime}`` time template. """ _check_version_and_prompt_upgrade(__version__) @@ -319,6 +323,8 @@ def map( ) if num_nodes is None or int(os.getenv("DATA_OPTIMIZER_NUM_NODES", 0)) > 0: + # Detect before `_resolve_dir` expands `{%strftime}` (Dir objects lose the template). + should_broadcast_paths = broadcast_paths or _has_time_template(output_dir) or _has_time_template(input_dir) _output_dir: Dir = _resolve_dir(output_dir) if _output_dir.url and "cloudspaces" in _output_dir.url: @@ -332,6 +338,7 @@ def map( if not isinstance(inputs, StreamingDataLoader): input_dir = input_dir or _get_input_dir(inputs) + should_broadcast_paths = should_broadcast_paths or _has_time_template(input_dir) resolved_dir = _resolve_dir(input_dir) if isinstance(batch_size, int) and batch_size > 1: @@ -355,6 +362,7 @@ def map( start_method=start_method, storage_options=storage_options, keep_data_ordered=keep_data_ordered, + broadcast_paths=should_broadcast_paths, ) with optimize_dns_context(optimize_dns if optimize_dns is not None else False): @@ -413,6 +421,7 @@ def optimize( storage_options: dict[str, Any] = {}, keep_data_ordered: bool = True, verbose: bool = True, + broadcast_paths: bool = False, ) -> None: """This function converts a dataset into chunks, possibly in a distributed way. @@ -461,6 +470,9 @@ def optimize( processing of items. If True, each worker processes a statically assigned subset of items in order. verbose: Whether to print the progress of the optimization. Defaults to True. + broadcast_paths: Broadcast resolved input/output dirs across multi-node ranks. Defaults to ``False``. + Auto-enabled when ``input_dir`` or ``output_dir`` contains a ``{%strftime}`` time template so ranks + share one expanded path. When off, each rank uses its locally resolved path. """ _check_version_and_prompt_upgrade(__version__) @@ -510,6 +522,8 @@ def optimize( if num_nodes is None or int(os.getenv("DATA_OPTIMIZER_NUM_NODES", 0)) > 0: DATA_OPTIMIZER_NUM_NODES = int(os.getenv("DATA_OPTIMIZER_NUM_NODES", 0)) + # Detect before `_resolve_dir` expands `{%strftime}` (Dir objects lose the template). + should_broadcast_paths = broadcast_paths or _has_time_template(output_dir) or _has_time_template(input_dir) _output_dir: Dir = _resolve_dir(output_dir) if ( @@ -535,7 +549,9 @@ def optimize( if not isinstance(inputs, StreamingDataLoader) and queue is None: assert inputs is not None - resolved_dir = _resolve_dir(input_dir or _get_input_dir(inputs)) + input_dir_resolved = input_dir or _get_input_dir(inputs) + should_broadcast_paths = should_broadcast_paths or _has_time_template(input_dir_resolved) + resolved_dir = _resolve_dir(input_dir_resolved) if isinstance(batch_size, int) and batch_size > 1: inputs = [inputs[pos : pos + batch_size] for pos in range(0, len(inputs), batch_size)] @@ -576,6 +592,7 @@ def optimize( storage_options=storage_options, keep_data_ordered=keep_data_ordered, verbose=verbose, + broadcast_paths=should_broadcast_paths, ) with optimize_dns_context(optimize_dns if optimize_dns is not None else False): diff --git a/src/litdata/raw/dataset.py b/src/litdata/raw/dataset.py index e9e74467b..cf0a3b7de 100644 --- a/src/litdata/raw/dataset.py +++ b/src/litdata/raw/dataset.py @@ -11,13 +11,55 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Streaming raw files with async downloads and optional look-ahead prefetch. + +Concurrency model +----------------- +A per-process ``_LoopRunner`` owns a dedicated event-loop thread. All dataset I/O +is dispatched with ``run_coroutine_threadsafe``, so prefetch continues between +``__getitems__`` calls (including ``num_workers=0`` and notebook nested loops). + +Runtime clients (downloader, semaphore, path-dedupe, prefetch tasks) are keyed by +``(pid, event loop)`` and recreated when either changes. Fork clears the runner +(threads do not survive ``os.fork``); the next call builds a fresh one. + +Cache publishes are atomic via temp file + ``os.replace``. Temp names include +``pid`` and thread id. Cross-process cache writers coordinate with ``O_EXCL`` +lock files (``*.litdata-raw.lock``). ``cache_files=True`` + ``item_type="bytes"`` +uses write-through: bytes return from the network path while a background thread +publishes the cache. + +Known limitations (accepted) +---------------------------- +- ``_close_downloader_best_effort`` cannot deterministically await async ``close()`` + when the bound loop is dead; pooled connections rely on GC in that case. +- A failed ``asyncio.gather`` in ``_download_batch`` may leave sibling resolve + coroutines to finish on the next call or at loop close. +- ``LoopRunner.run(...).result()`` blocks the caller thread (DataLoader worker / + main); cancellation from ``KeyboardInterrupt`` is best-effort. +- ThreadPoolExecutor worker threads (range downloads / write-behind) may still be + alive at ``os.fork``; the child reinitializes locks and drops the runner, but + inherited executor threads are not joined. +""" + +from __future__ import annotations + import asyncio +import atexit +import concurrent.futures +import contextlib import logging import os -from collections.abc import Callable -from functools import lru_cache +import statistics +import threading +import time +from collections import OrderedDict +from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import asynccontextmanager from pathlib import Path -from typing import Any +from typing import Any, Literal, TypeVar +from uuid import uuid4 from torch.utils.data import Dataset @@ -28,6 +70,461 @@ logger = logging.getLogger(__name__) +T = TypeVar("T") +_MISS = object() +_RAW_DEBUG = bool(os.getenv("LITDATA_RAW_DEBUG")) + +# Parallel ranged GETs for large objects (S3/GCS/R2 downloaders with real Range support). +# Opt-in: default 0 disables; pass a positive byte threshold to enable. +_RANGE_PARALLEL_THRESHOLD = 0 +_RANGE_CHUNK_SIZE = 8 * 1024 * 1024 +_TINY_FILE_MEDIAN_BYTES = 100_000 +_LOCK_STALE_SECONDS = 300.0 +_LOCK_SUFFIX = ".litdata-raw.lock" +# Hedge only small / unknown objects; large whole-object GETs must not 2× egress. +_HEDGE_MAX_BYTES = 8 * 1024 * 1024 +# Per-stream floor for hedge delay / batch-timeout sizing (healthy single GET). +# Distinct from ``_ASSUMED_AGGREGATE_BANDWIDTH_BPS`` below (NIC/prefix share used to +# size Stage 1 aggregate concurrency). Do not conflate the two constants. +_HEDGE_ASSUMED_BANDWIDTH_BPS = 25 * 1024 * 1024 # ~25 MB/s per-stream floor +# Cap aggregate sequential look-ahead across DataLoader workers (items total). +# Per-worker effective = min(max_prefetch, budget // num_workers) when num_workers > 1. +_AGGREGATE_PREFETCH_BUDGET = 64 +# Stage 1 static concurrency: aggregate in-flight budget = max(bandwidth model, +# Little's-law / latency model), then split across DataLoader workers. +_ASSUMED_AGGREGATE_BANDWIDTH_BPS = 100 * 1024 * 1024 # ~100 MB/s NIC / prefix share +_CONCURRENCY_PIPELINE_SECONDS = 0.5 # target aggregate bytes ≈ bandwidth × this +_ASSUMED_REQUEST_RATE = 6000.0 # tiny-object target req/s for Little's-law arm +_ASSUMED_REQUEST_LATENCY_S = 0.040 # assumed RTT+TTFB (seconds) +_DEFAULT_MEDIAN_FILE_BYTES = 256 * 1024 +_AGGREGATE_CONCURRENCY_BUDGET_FLOOR = 32 +# Cap keeps high-w stampede (was N×64 → 1536) in check without crushing mid-w +# ImageNet cells to ~128 aggregate (which capped winning w=4/w=8 configs). +_AGGREGATE_CONCURRENCY_BUDGET_CAP = 512 +_MIN_CONCURRENCY_PER_WORKER = 8 +# Unbenchmarked single-process adaptive path: do not open the full ~512 budget. +_SINGLE_PROCESS_CONCURRENCY_CAP = 128 +# Little's-law arm only for sub-MiB objects (request-overhead bound). At ≥1 MiB +# the bandwidth arm alone sizes the budget. Distinct from ``_HEDGE_MAX_BYTES`` +# (8 MiB duplicate-GET hedge policy). +_LATENCY_MODEL_MAX_MEDIAN_BYTES = 1024 * 1024 + +_RUNNER_LOCK = threading.Lock() +_RUNNER: _LoopRunner | None = None + +_WRITE_BEHIND_LOCK = threading.Lock() +_WRITE_BEHIND_FUTURES: set[asyncio.Future[Any]] = set() + + +def _loop_backend_name() -> str: + """Return ``uvloop`` when the package is importable, else ``asyncio``.""" + try: + import uvloop # noqa: F401 + except ImportError: + return "asyncio" + return "uvloop" + + +def _create_event_loop() -> asyncio.AbstractEventLoop: + """Create a new event loop, preferring uvloop when available.""" + try: + import uvloop + except ImportError: + return asyncio.new_event_loop() + return uvloop.new_event_loop() + + +def _consume_task_exception(task: asyncio.Task) -> None: + """Mark task exceptions as retrieved so asyncio does not warn at GC time.""" + if task.cancelled(): + return + with contextlib.suppress(asyncio.InvalidStateError, Exception): + task.exception() + + +def _close_unawaited(coro: Awaitable[Any]) -> None: + """Close a coroutine that will not be awaited (avoids 'was never awaited').""" + close = getattr(coro, "close", None) + if callable(close): + with contextlib.suppress(Exception): + close() + + +def _track_write_behind(fut: asyncio.Future[Any]) -> None: + """Track a ``run_in_executor`` future until the write-behind finishes.""" + with _WRITE_BEHIND_LOCK: + _WRITE_BEHIND_FUTURES.add(fut) + + def _done(f: asyncio.Future[Any]) -> None: + with _WRITE_BEHIND_LOCK: + _WRITE_BEHIND_FUTURES.discard(f) + with contextlib.suppress(Exception): + f.result() + + fut.add_done_callback(_done) + + +def _drain_write_behind_futures() -> None: + """Best-effort wait for in-flight write-behind publishes (shared ~0.5s deadline).""" + with _WRITE_BEHIND_LOCK: + pending = list(_WRITE_BEHIND_FUTURES) + _WRITE_BEHIND_FUTURES.clear() + if not pending: + return + deadline = time.monotonic() + 0.5 + for fut in pending: + while not fut.done(): + remaining = deadline - time.monotonic() + if remaining <= 0: + return + time.sleep(min(0.01, remaining)) + with contextlib.suppress(Exception): + fut.result() + + +atexit.register(_drain_write_behind_futures) + + +class _LoopRunner: + """Process-local asyncio loop running forever on a daemon thread.""" + + def __init__(self) -> None: + self._pid = os.getpid() + self.loop: asyncio.AbstractEventLoop = _create_event_loop() + self._executor = ThreadPoolExecutor(max_workers=32, thread_name_prefix="litdata-raw-pool") + self.loop.set_default_executor(self._executor) + if _RAW_DEBUG: + logger.warning( + "raw-debug: LoopRunner backend=%s pid=%s", + _loop_backend_name(), + self._pid, + ) + self._thread = threading.Thread(target=self._main, name="litdata-raw-aio", daemon=True) + self._started = threading.Event() + self._thread.start() + if not self._started.wait(timeout=10): + raise RuntimeError("Failed to start StreamingRawDataset event-loop thread") + + def _main(self) -> None: + asyncio.set_event_loop(self.loop) + self._started.set() + self.loop.run_forever() + pending = asyncio.all_tasks(self.loop) + for task in pending: + task.cancel() + if pending: + with contextlib.suppress(Exception): + self.loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + with contextlib.suppress(Exception): + self.loop.close() + + @property + def pid(self) -> int: + return self._pid + + def is_alive(self) -> bool: + return self._thread.is_alive() and not self.loop.is_closed() + + def run(self, coro: Coroutine[Any, Any, T]) -> T: + if threading.current_thread() is self._thread: + _close_unawaited(coro) + raise RuntimeError( + "LoopRunner.run() called from the event-loop thread; this would deadlock. " + "Await the coroutine directly instead of calling run()." + ) + if os.getpid() != self._pid: + _close_unawaited(coro) + raise RuntimeError("LoopRunner used after fork; call _get_loop_runner() to recreate") + if not self.is_alive(): + _close_unawaited(coro) + raise RuntimeError("StreamingRawDataset event-loop thread is not running") + return asyncio.run_coroutine_threadsafe(coro, self.loop).result() + + def shutdown_best_effort(self) -> None: + if self.loop.is_closed(): + with contextlib.suppress(Exception): + self._executor.shutdown(wait=False, cancel_futures=True) + return + + def _stop() -> None: + try: + with contextlib.suppress(Exception): + self._executor.shutdown(wait=False, cancel_futures=True) + finally: + # Must always stop the loop — uvloop has no assignable `_default_executor`. + self.loop.stop() + + with contextlib.suppress(Exception): + if self.loop.is_running(): + self.loop.call_soon_threadsafe(_stop) + else: + _stop() + self._thread.join(timeout=2.0) + with contextlib.suppress(Exception): + self._executor.shutdown(wait=False, cancel_futures=True) + + +def _get_loop_runner() -> _LoopRunner: + """Return the process-local loop runner, creating it if needed.""" + global _RUNNER + with _RUNNER_LOCK: + if _RUNNER is None or _RUNNER.pid != os.getpid() or not _RUNNER.is_alive(): + if _RUNNER is not None: + _RUNNER.shutdown_best_effort() + if _RAW_DEBUG: + logger.warning("raw-debug: creating LoopRunner pid=%s", os.getpid()) + _RUNNER = _LoopRunner() + return _RUNNER + + +def _shutdown_runner_before_fork() -> None: + """Stop the loop thread before fork (threads do not survive fork safely).""" + global _RUNNER + with _RUNNER_LOCK: + if _RUNNER is not None: + if _RAW_DEBUG: + logger.warning("raw-debug: shutting down LoopRunner before fork pid=%s", os.getpid()) + _RUNNER.shutdown_best_effort() + _RUNNER = None + + +def _reinit_after_fork() -> None: + """Drop inherited runner/futures and reinit module-level locks in the child.""" + global _RUNNER, _RUNNER_LOCK, _WRITE_BEHIND_LOCK, _WRITE_BEHIND_FUTURES + _RUNNER = None + _RUNNER_LOCK = threading.Lock() + _WRITE_BEHIND_LOCK = threading.Lock() + _WRITE_BEHIND_FUTURES = set() + + +# Backward-compatible alias for tests / callers that still import the old name. +_clear_runner_after_fork = _reinit_after_fork + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(before=_shutdown_runner_before_fork, after_in_child=_reinit_after_fork) + + +def _run_async(coro: Coroutine[Any, Any, T]) -> T: + """Dispatch ``coro`` onto the process-local loop thread and wait for the result.""" + return _get_loop_runner().run(coro) + + +def _pid_alive(pid: int) -> bool: + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +def _unlink_if_dead_pid(path: Path, pid: int) -> None: + if not _pid_alive(pid): + with contextlib.suppress(OSError): + path.unlink() + + +def _sweep_orphan_tmp_files(cache_dir: str) -> None: + """Best-effort cleanup of ``*.tmp.*`` and ``.range-scratch.*`` orphans.""" + root = Path(cache_dir) + if not root.is_dir(): + return + for tmp in root.rglob("*.tmp.*"): + marker = ".tmp." + idx = tmp.name.rfind(marker) + if idx < 0: + continue + suffix = tmp.name[idx + len(marker) :] + pid_str = suffix.split(".", 1)[0] + try: + pid = int(pid_str) + except ValueError: + # Non-pid suffix: drop if mtime is stale. + try: + if time.time() - tmp.stat().st_mtime > _LOCK_STALE_SECONDS: + tmp.unlink() + except OSError: + pass + continue + _unlink_if_dead_pid(tmp, pid) + # Orphan ranged-GET scratch files: ``.range-scratch.....`` + for scratch in root.rglob(".range-scratch.*"): + parts = scratch.name.split(".") + # "", "range-scratch", "", ... + if len(parts) < 3: + continue + try: + pid = int(parts[2]) + except ValueError: + try: + if time.time() - scratch.stat().st_mtime > _LOCK_STALE_SECONDS: + scratch.unlink() + except OSError: + pass + continue + _unlink_if_dead_pid(scratch, pid) + # Stale cross-process lock files from dead pids. + for lock in root.rglob(f"*{_LOCK_SUFFIX}"): + try: + text = lock.read_text().strip() + pid = int(text.split()[0]) + except (OSError, ValueError, IndexError): + try: + if time.time() - lock.stat().st_mtime > _LOCK_STALE_SECONDS: + lock.unlink() + except OSError: + pass + continue + _unlink_if_dead_pid(lock, pid) + + +def _looks_sequential(indices: list[int]) -> bool: + """Return True if indices are a contiguous ascending range (typical DataLoader batch).""" + if len(indices) <= 1: + return True + return all(indices[i] == indices[i - 1] + 1 for i in range(1, len(indices))) + + +def _effective_prefetch(max_prefetch: int, num_workers: int) -> int: + """Per-worker look-ahead capped so aggregate stays near ``_AGGREGATE_PREFETCH_BUDGET``. + + Constructor default ``max_prefetch=16`` stays ergonomic at low workers (w≤4 keeps 16). + At higher worker counts each worker gets a smaller share so total in-flight look-ahead + does not scale as ``num_workers × max_prefetch``. + + TODO(Stage 2): hit-rate controller — if windowed hit rate <30% → halve effective + look-ahead (floor 0); >90% and batch waited on a miss → +1–2. Not implemented yet. + """ + if max_prefetch <= 0: + return 0 + if num_workers <= 1: + return max_prefetch + # TODO(open): if effective < 8, return 0 — pending repeats; priority below downloader conformance. + return min(max_prefetch, max(0, _AGGREGATE_PREFETCH_BUDGET // num_workers)) + + +def _median_file_bytes(files: Sequence[FileMetadata]) -> int | None: + """Return median positive file size from index metadata, or ``None`` if unknown.""" + sizes = [f.size for f in files if f.size > 0] + if not sizes: + return None + return int(statistics.median(sizes)) + + +def _aggregate_concurrency_budget(median_file_bytes: int | None) -> int: + """Aggregate in-flight download slots across all workers (size-aware, clamped). + + Takes the max of two models then clamps to ``[floor, cap]``: + + - **bandwidth**: ``(aggregate_bps × pipeline_s) // median_file_bytes`` — keep + ~50 MiB moving for large objects. + - **latency / Little's law**: ``target_rate × assumed_latency`` (~6000×0.040 ≈ + 240) **only when** ``median < _LATENCY_MODEL_MAX_MEDIAN_BYTES`` (1 MiB) so + tiny-object paths are not request-starved. Medians ≥1 MiB stay + bandwidth-bounded (avoids pinning at 240 slots → multi-GB in flight). + + Per-worker floor of 8 means realized aggregate is ``max(budget, 8 × num_workers)``. + """ + median = median_file_bytes if median_file_bytes and median_file_bytes > 0 else _DEFAULT_MEDIAN_FILE_BYTES + target_bytes = int(_ASSUMED_AGGREGATE_BANDWIDTH_BPS * _CONCURRENCY_PIPELINE_SECONDS) + bandwidth_model = max(1, target_bytes // median) + # Size-gate: Little's-law arm is for request-overhead-bound tiny objects only. + if median < _LATENCY_MODEL_MAX_MEDIAN_BYTES: + latency_model = max(1, int(_ASSUMED_REQUEST_RATE * _ASSUMED_REQUEST_LATENCY_S)) + else: + latency_model = 0 + raw = max(bandwidth_model, latency_model) + return max(_AGGREGATE_CONCURRENCY_BUDGET_FLOOR, min(_AGGREGATE_CONCURRENCY_BUDGET_CAP, raw)) + + +def _effective_concurrency( + max_concurrent_downloads: int | None, + num_workers: int, + median_file_bytes: int | None = None, +) -> int: + """Per-worker download permits for the Stage 1 static clamp. + + - ``max_concurrent_downloads is None`` (default): adaptive — + ``max(floor, budget // num_workers)`` with ``budget`` from + :func:`_aggregate_concurrency_budget`. When ``num_workers <= 1``, returns + ``min(budget, _SINGLE_PROCESS_CONCURRENCY_CAP)`` (unbenchmarked path). + - Explicit ``int``: **exactly** that many permits (no silent clamp). ``<= 0`` + collapses to 1. + """ + if max_concurrent_downloads is not None: + return 1 if max_concurrent_downloads <= 0 else max_concurrent_downloads + budget = _aggregate_concurrency_budget(median_file_bytes) + if num_workers <= 1: + return min(budget, _SINGLE_PROCESS_CONCURRENCY_CAP) + return max(_MIN_CONCURRENCY_PER_WORKER, budget // num_workers) + + +def _num_dataloader_workers() -> int: + """Return DataLoader ``num_workers``, or ``1`` when called outside a worker.""" + try: + from torch.utils.data import get_worker_info + except ImportError: + return 1 + info = get_worker_info() + if info is None: + return 1 + return max(1, int(info.num_workers)) + + +def _consume_prefetch_exception(task: asyncio.Task) -> None: + """Mark prefetch task exceptions as retrieved so asyncio does not warn at GC time.""" + if task.cancelled(): + return + with contextlib.suppress(asyncio.InvalidStateError): + exc = task.exception() + if exc is not None: + logger.debug("prefetch failed; will retry on demand", exc_info=exc) + + +def _effective_hedge_delay(hedge_delay: float, size: int | None) -> float | None: + """Return hedge wait seconds, or ``None`` when hedging should be skipped. + + Unknown / non-positive sizes never hedge (cannot bound duplicate egress). Large + whole-object GETs (``>= _HEDGE_MAX_BYTES``) never hedge either — callers should + hedge per ranged chunk instead. Delay is at least ``3 * size / 25MB/s`` so a + healthy transfer is not spuriously duplicated. + """ + if hedge_delay <= 0: + return None + if size is None or size <= 0: + return None + if size >= _HEDGE_MAX_BYTES: + return None + expected = size / _HEDGE_ASSUMED_BANDWIDTH_BPS + return max(hedge_delay, 3.0 * expected) + + +class _LRUCache: + """Simple ordered LRU cache keyed by dataset index.""" + + def __init__(self, maxsize: int) -> None: + self.maxsize = max(0, maxsize) + self._data: OrderedDict[int, Any] = OrderedDict() + + def get(self, key: int) -> Any: + if self.maxsize <= 0 or key not in self._data: + return _MISS + self._data.move_to_end(key) + return self._data[key] + + def put(self, key: int, value: Any) -> None: + if self.maxsize <= 0: + return + self._data[key] = value + self._data.move_to_end(key) + while len(self._data) > self.maxsize: + self._data.popitem(last=False) + + def __contains__(self, key: object) -> bool: + return isinstance(key, int) and key in self._data + class CacheManager: """Manages file caching for remote datasets, preserving directory structure.""" @@ -38,28 +535,245 @@ def __init__( cache_dir: str | None = None, storage_options: dict | None = None, cache_files: bool = False, + max_concurrent_downloads: int | None = None, + hedge_delay: float = 0.0, + download_timeout: float = 120.0, + range_parallel_threshold: int = _RANGE_PARALLEL_THRESHOLD, + range_chunk_size: int = _RANGE_CHUNK_SIZE, ): self.input_dir = _resolve_dir(input_dir) - self._input_dir_path = str(self.input_dir.path or self.input_dir.url) + self._input_dir_path = _storage_path(self.input_dir) self.cache_files = cache_files + self.max_concurrent_downloads = max_concurrent_downloads + self.hedge_delay = max(0.0, hedge_delay) + self.download_timeout = max(0.0, download_timeout) or None + self.range_parallel_threshold = max(0, range_parallel_threshold) + self.range_chunk_size = max(1, range_chunk_size) + self.lock_wait_timeout = _LOCK_STALE_SECONDS self.cache_dir = self._create_cache_dir(self._input_dir_path, cache_dir) + _sweep_orphan_tmp_files(self.cache_dir) self.storage_options = storage_options or {} + # Index median size (bytes); set by StreamingRawDataset after discovery. + self._median_file_bytes: int | None = None self._downloader: Downloader | None = None + self._downloader_pid: int | None = None + self._downloader_loop: asyncio.AbstractEventLoop | None = None + self._semaphore: asyncio.Semaphore | None = None + self._semaphore_loop: asyncio.AbstractEventLoop | None = None + self._semaphore_permits: int | None = None + # Pid-guarded cache of Stage 1 permit count (avoid hot-path get_worker_info). + self._cached_permits: int | None = None + self._cached_permits_pid: int | None = None + self._path_inflight: dict[str, asyncio.Task] = {} + self._path_inflight_loop: asyncio.AbstractEventLoop | None = None + # Presence hint only: membership does not skip exists checks (stale marks self-heal). + self._present_paths: set[str] = set() + self._range_executor: ThreadPoolExecutor | None = None + self._range_executor_pid: int | None = None + self._hedge_fired = 0 + + def reset_runtime_state(self) -> None: + """Drop process/loop-bound clients (call after fork or when pickling).""" + self._close_downloader_best_effort(self._downloader) + self._downloader = None + self._downloader_pid = None + self._downloader_loop = None + self._semaphore = None + self._semaphore_loop = None + self._semaphore_permits = None + self._cached_permits = None + self._cached_permits_pid = None + self._path_inflight = {} + self._path_inflight_loop = None + self._shutdown_range_executor() + self._hedge_fired = 0 + # Keep _present_paths / _median_file_bytes — index metadata survives fork/spawn. + + def __getstate__(self) -> dict[str, Any]: + """Serialize config only — never downloader/loop/executor/inflight state. + + Allowlisted keys avoid accidental instance attrs (locks, futures) breaking + ``multiprocessing_context='spawn'`` pickling. + """ + return { + "input_dir": self.input_dir, + "_input_dir_path": self._input_dir_path, + "cache_files": self.cache_files, + "max_concurrent_downloads": self.max_concurrent_downloads, + "hedge_delay": self.hedge_delay, + "download_timeout": self.download_timeout, + "range_parallel_threshold": self.range_parallel_threshold, + "range_chunk_size": self.range_chunk_size, + "lock_wait_timeout": self.lock_wait_timeout, + "cache_dir": self.cache_dir, + "storage_options": self.storage_options, + "_median_file_bytes": self._median_file_bytes, + # Runtime — always fresh in the child. + "_downloader": None, + "_downloader_pid": None, + "_downloader_loop": None, + "_semaphore": None, + "_semaphore_loop": None, + "_semaphore_permits": None, + "_cached_permits": None, + "_cached_permits_pid": None, + "_path_inflight": {}, + "_path_inflight_loop": None, + "_present_paths": set(), + "_range_executor": None, + "_range_executor_pid": None, + "_hedge_fired": 0, + } + + def __setstate__(self, state: dict[str, Any]) -> None: + self.__dict__.update(state) + # Belt-and-suspenders: never revive process-bound clients after unpickle. + self._downloader = None + self._downloader_pid = None + self._downloader_loop = None + self._semaphore = None + self._semaphore_loop = None + self._semaphore_permits = None + self._cached_permits = None + self._cached_permits_pid = None + self._path_inflight = {} + self._path_inflight_loop = None + self._present_paths = set(state.get("_present_paths") or ()) + self._range_executor = None + self._range_executor_pid = None + self._hedge_fired = 0 + self._median_file_bytes = state.get("_median_file_bytes") + + def _shutdown_range_executor(self) -> None: + if self._range_executor is not None: + with contextlib.suppress(Exception): + self._range_executor.shutdown(wait=False, cancel_futures=True) + self._range_executor = None + self._range_executor_pid = None + + def _get_range_executor(self) -> ThreadPoolExecutor: + pid = os.getpid() + if self._range_executor is None or self._range_executor_pid != pid: + self._shutdown_range_executor() + # Explicit permit cap when set; otherwise a modest default (adaptive + # Stage 1 budget is applied on the download semaphore, not here). + cap = self.max_concurrent_downloads if self.max_concurrent_downloads is not None else 32 + workers = max(4, min(32, cap)) + self._range_executor = ThreadPoolExecutor( + max_workers=workers, + thread_name_prefix="litdata-raw-range", + ) + self._range_executor_pid = pid + return self._range_executor + + @staticmethod + def _close_downloader_best_effort(downloader: Downloader | None) -> None: + """Best-effort downloader teardown (sync or async ``close`` / ``aclose``). + + Async close coroutines cannot be awaited safely when the bound loop is already + dead (fork reset), so the coroutine object is closed and pooled connections + fall back to GC rather than a deterministic drain. + """ + if downloader is None: + return + for name in ("close", "aclose"): + fn = getattr(downloader, name, None) + if not callable(fn): + continue + with contextlib.suppress(Exception): + result = fn() + if asyncio.iscoroutine(result): + with contextlib.suppress(Exception): + result.close() + return @property def downloader(self) -> Downloader: - """Lazily initialize the downloader.""" - if self._downloader is None: + """Lazily initialize the downloader. + + Recreate when the process id **or** running event loop changes. Clients bound to a + closed parent loop hang both in forked workers and on later main-process access. + """ + pid = os.getpid() + try: + loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop() + except RuntimeError: + loop = None + if self._downloader is None or self._downloader_pid != pid or self._downloader_loop is not loop: + if _RAW_DEBUG and self._downloader is not None: + logger.warning( + "raw-debug: recreating downloader pid=%s->%s loop_changed=%s", + self._downloader_pid, + pid, + self._downloader_loop is not loop, + ) + self._close_downloader_best_effort(self._downloader) self._downloader = get_downloader( remote_dir=self._input_dir_path, cache_dir=self.cache_dir, chunks=[], storage_options=self.storage_options, ) + self._downloader_pid = pid + self._downloader_loop = loop return self._downloader + def _effective_download_permits(self) -> int: + """Worker-aware permit count for the download semaphore (Stage 1 static clamp). + + Computed once per process (pid-guarded cache). Cleared by + ``reset_runtime_state`` / pickle so forked workers recompute. + """ + pid = os.getpid() + if self._cached_permits is not None and self._cached_permits_pid == pid: + return self._cached_permits + permits = _effective_concurrency( + self.max_concurrent_downloads, + _num_dataloader_workers(), + self._median_file_bytes, + ) + self._cached_permits = permits + self._cached_permits_pid = pid + return permits + + def _get_semaphore(self) -> asyncio.Semaphore: + """Return a semaphore bound to the current event loop with effective permits. + + Permit count comes from :meth:`_effective_download_permits` (cached per + process). Loop-keyed like other runtime clients; cleared by + ``reset_runtime_state``. + """ + loop = asyncio.get_running_loop() + permits = self._effective_download_permits() + if self._semaphore is None or self._semaphore_loop is not loop or self._semaphore_permits != permits: + n_workers = _num_dataloader_workers() + budget = ( + _aggregate_concurrency_budget(self._median_file_bytes) + if self.max_concurrent_downloads is None + else None + ) + logger.info( + "adaptive concurrency: median=%s budget=%s workers=%s permits=%s", + self._median_file_bytes, + budget, + n_workers, + permits, + ) + self._semaphore = asyncio.Semaphore(permits) + self._semaphore_loop = loop + self._semaphore_permits = permits + return self._semaphore + + @asynccontextmanager + async def _permit(self, gated: bool = True) -> AsyncIterator[None]: + if gated: + async with self._get_semaphore(): + yield + else: + yield + def _create_cache_dir(self, input_dir: str, cache_dir: str | None = None) -> str: """Create cache directory if it doesn't exist.""" if cache_dir is None: @@ -79,19 +793,445 @@ def get_local_path(self, remote_file_path: str) -> str: local_path.parent.mkdir(parents=True, exist_ok=True) return str(local_path) - async def download_file_async(self, file_path: str) -> bytes: - """Asynchronously download and return file content.""" - if self.cache_files: - local_path = self.get_local_path(file_path) + def _path_is_cached(self, local_path: str) -> bool: + if local_path in self._present_paths: + if os.path.exists(local_path): + return True + # Stale presence mark (file removed under us) — discard and recheck disk. + self._present_paths.discard(local_path) + if os.path.exists(local_path): + self._present_paths.add(local_path) + return True + return False + + def _lock_path(self, local_path: str) -> str: + return f"{local_path}{_LOCK_SUFFIX}" + + @staticmethod + def _pid_alive(pid: int) -> bool: + return _pid_alive(pid) + + def _lock_owner_alive(self, lock_path: str) -> bool: + """Return True if the lock looks held by a live process (or is still being written).""" + try: + text = Path(lock_path).read_text().strip() + except OSError: + return True + if not text: + # Claim race: lock created but pid not written yet — treat as alive unless stale. + try: + return time.time() - os.path.getmtime(lock_path) <= _LOCK_STALE_SECONDS + except OSError: + return True + try: + pid = int(text.split()[0]) + except (ValueError, IndexError): + try: + return time.time() - os.path.getmtime(lock_path) <= _LOCK_STALE_SECONDS + except OSError: + return True + return self._pid_alive(pid) + + def _try_claim_lock(self, local_path: str) -> bool: + """Claim ``local_path.litdata-raw.lock`` with ``O_EXCL``. Return True if we own the download.""" + lock_path = self._lock_path(local_path) + try: + fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + except FileExistsError: + if not self._lock_owner_alive(lock_path): + with contextlib.suppress(OSError): + os.remove(lock_path) + try: + fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + except FileExistsError: + return False + else: + try: + age = time.time() - os.path.getmtime(lock_path) + except OSError: + age = 0.0 + if age > _LOCK_STALE_SECONDS: + with contextlib.suppress(OSError): + os.remove(lock_path) + try: + fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + except FileExistsError: + return False + else: + return False + try: + os.write(fd, f"{os.getpid()}\n".encode()) + finally: + os.close(fd) + return True + + def _release_lock(self, local_path: str) -> None: + with contextlib.suppress(OSError): + os.remove(self._lock_path(local_path)) + + async def _wait_for_cached_file(self, local_path: str, timeout: float | None = None) -> str: + """Poll until another process publishes ``local_path`` or the lock is released/stale.""" + if timeout is None: + timeout = self.lock_wait_timeout + deadline = time.monotonic() + timeout + lock_path = self._lock_path(local_path) + while time.monotonic() < deadline: + if self._path_is_cached(local_path): + return local_path + if not os.path.exists(lock_path): + break + if not self._lock_owner_alive(lock_path): + with contextlib.suppress(OSError): + os.remove(lock_path) + break + await asyncio.sleep(0.05) + if self._path_is_cached(local_path): + return local_path + raise TimeoutError(f"Timed out waiting for cache file {local_path}") + + @staticmethod + def _is_remote_object(file_path: str) -> bool: + return "://" in file_path and not file_path.startswith("file://") + + @staticmethod + def _is_non_retryable_download_error(exc: BaseException) -> bool: + return isinstance( + exc, + ( + asyncio.CancelledError, + NotImplementedError, + ValueError, + TypeError, + PermissionError, + FileNotFoundError, + IsADirectoryError, + TimeoutError, + concurrent.futures.TimeoutError, + ), + ) + + async def _hedged(self, factory: Callable[[], Coroutine[Any, Any, T]], delay: float) -> T: + """Run ``factory``; if slow, start a second request and prefer a non-exception winner. + + Each ``factory`` call is expected to acquire its own semaphore permit. The hedge is + skipped when the semaphore is already exhausted so hedges cannot starve primary work. + + Note: cancelling a losing hedged task that is blocked in ``run_in_executor`` does + not abort the worker thread — the full chunk transfer may still complete and pay + bandwidth even after the asyncio task is cancelled. + """ + first: asyncio.Task[T] = asyncio.create_task(factory()) + if delay <= 0: + return await first + done, _ = await asyncio.wait({first}, timeout=delay) + if done: + return first.result() + + # Only hedge when a concurrency permit is immediately available. + if self._get_semaphore().locked(): + return await first + + second: asyncio.Task[T] = asyncio.create_task(factory()) + self._hedge_fired += 1 + logger.debug("hedge fired count=%s delay=%.3fs", self._hedge_fired, delay) + pending: set[asyncio.Task[T]] = {first, second} + try: + while pending: + finished, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED) + success = [t for t in finished if not t.cancelled() and t.exception() is None] + if success: + for task in pending: + task.cancel() + task.add_done_callback(_consume_task_exception) + for task in finished: + if task not in success: + _consume_task_exception(task) + return success[0].result() + for task in finished: + _consume_task_exception(task) + # Both failed — re-raise from the first attempt. + return first.result() + except BaseException: + for task in (first, second): + if not task.done(): + task.cancel() + task.add_done_callback(_consume_task_exception) + raise + + def _download_budget(self, size: int | None = None, timeout: float | None = None) -> float | None: + """Return a size-aware timeout floor in seconds, or ``None`` when disabled. + + ``download_timeout`` is a floor for sized objects: when ``size`` is known, + the budget is ``max(download_timeout, size / assumed_bandwidth * 3)`` so large + transfers are not cut off by a fixed wall-clock cap. Used by batch-level + hang protection (``max`` over pending indices). Pass an explicit ``timeout`` + to override. + """ + if timeout is not None: + return timeout + base = self.download_timeout + if base is None: + return None + if size is not None and size > 0: + size_floor = size / _HEDGE_ASSUMED_BANDWIDTH_BPS * 3.0 + return max(base, size_floor) + return base + + def _supports_range(self, file_path: str) -> bool: + return file_path.startswith(("s3://", "gs://", "r2://")) + + async def _ranged_download_bytes(self, file_path: str, size: int, *, gated: bool = True) -> bytes: + """Parallel ranged GETs via the sync ``download_bytes`` API (per-chunk hedge + validate).""" + chunk = self.range_chunk_size + ranges = [(start, min(chunk, size - start)) for start in range(0, size, chunk)] + downloader = self.downloader + executor = self._get_range_executor() + base_scratch = os.path.join( + self.cache_dir, + f".range-scratch.{os.getpid()}.{threading.get_ident()}", + ) + chunk_delay = _effective_hedge_delay(self.hedge_delay, chunk) + + async def one(offset: int, length: int) -> tuple[int, bytes]: + async def fetch() -> bytes: + # Unique scratch per attempt so first/hedge never share a path. + scratch = f"{base_scratch}.{offset}.{uuid4().hex}" + try: + async with self._permit(gated): + data = await asyncio.get_running_loop().run_in_executor( + executor, + downloader.download_bytes, + file_path, + offset, + length, + scratch, + ) + if len(data) != length: + raise RuntimeError( + f"Ranged GET short read for {file_path}: offset={offset} expected={length} got={len(data)}" + ) + return data + finally: + with contextlib.suppress(OSError): + os.remove(scratch) + + if chunk_delay is not None and self._is_remote_object(file_path): + data = await self._hedged(fetch, chunk_delay) + else: + data = await fetch() + return offset, data + + parts = await asyncio.gather(*(one(o, n) for o, n in ranges)) + parts.sort(key=lambda x: x[0]) + joined = b"".join(data for _, data in parts) + if len(joined) != size: + raise RuntimeError(f"Ranged download size mismatch for {file_path}: expected={size} got={len(joined)}") + return joined + + async def _fetch_bytes(self, file_path: str, size: int | None = None, *, gated: bool = True) -> bytes: + """Download object bytes (optional range-parallel + size-gated hedging). + + Hang protection is enforced once per batch in + ``StreamingRawDataset._download_batch`` (single ``wait_for`` around the + gather). Defaults (``hedge_delay=0``, ``download_timeout=120``) therefore + take this bare fast path — no per-item ``asyncio.wait_for``. + """ + # Per-chunk hedging happens inside ranged downloads; never hedge the whole object. + if ( + size is not None + and self.range_parallel_threshold > 0 + and size >= self.range_parallel_threshold + and self._supports_range(file_path) + ): + return await self._ranged_download_bytes(file_path, size, gated=gated) + + delay = _effective_hedge_delay(self.hedge_delay, size) if self._is_remote_object(file_path) else None + # Pay-per-use: hedging off/ineligible → bare permit + download (batch enforces timeout). + if delay is None: + async with self._permit(gated): + return await self.downloader.adownload_fileobj(file_path) + + async def once() -> bytes: + async with self._permit(gated): + return await self.downloader.adownload_fileobj(file_path) + + return await self._hedged(once, delay) + + def _schedule_write_behind(self, local_path: str, data: bytes) -> None: + """Atomically publish ``data`` to ``local_path`` on a worker thread.""" + if self._path_is_cached(local_path): + return + + def _write() -> None: if os.path.exists(local_path): + self._present_paths.add(local_path) + return + tmp_path = f"{local_path}.tmp.{os.getpid()}.{threading.get_ident()}" + try: + os.makedirs(os.path.dirname(local_path) or ".", exist_ok=True) + with open(tmp_path, "wb") as f: + f.write(data) + os.replace(tmp_path, local_path) + self._present_paths.add(local_path) + except Exception: + logger.debug("write-behind cache publish failed for %s", local_path, exc_info=True) + finally: + with contextlib.suppress(OSError): + os.remove(tmp_path) + + loop = asyncio.get_running_loop() + fut = loop.run_in_executor(None, _write) + _track_write_behind(fut) + + def _verify_tmp_size(self, tmp_path: str, size: int | None) -> None: + if size is None or size <= 0: + return + actual = os.path.getsize(tmp_path) + if actual < size: + raise RuntimeError(f"Downloaded file truncated: expected>={size} got={actual} path={tmp_path}") + + async def _download_owned(self, file_path: str, local_path: str, size: int | None = None) -> str: + """Download while holding the cross-process lock (caller owns the lock).""" + tmp_path = f"{local_path}.tmp.{os.getpid()}.{threading.get_ident()}" + with contextlib.suppress(OSError): + os.remove(tmp_path) + try: + if self._path_is_cached(local_path): + return local_path + try: + # Hang protection is batch-level; keep the owned path bare. + await self.downloader.adownload_file(file_path, tmp_path) + except Exception as first_exc: + if self._is_non_retryable_download_error(first_exc): + raise + logger.warning( + "adownload_file failed for %s; falling back to bytes path", + file_path, + exc_info=first_exc, + ) + with contextlib.suppress(OSError): + os.remove(tmp_path) + # Caller already holds the download semaphore — avoid nested acquire. + data = await self._fetch_bytes(file_path, size=size, gated=False) + await asyncio.to_thread(Path(tmp_path).write_bytes, data) + self._verify_tmp_size(tmp_path, size) + os.replace(tmp_path, local_path) + self._present_paths.add(local_path) + finally: + with contextlib.suppress(OSError): + os.remove(tmp_path) + return local_path + + async def _download_to_cache(self, file_path: str, local_path: str, size: int | None = None) -> str: + """Download ``file_path`` into ``local_path`` with lock + atomic publish. + + Claim the cross-process lock *after* acquiring a semaphore permit so waiters + do not convoy behind a lock holder that has not started downloading yet. + """ + if self._path_is_cached(local_path): + return local_path + + deadline = time.monotonic() + self.lock_wait_timeout + while time.monotonic() < deadline: + if self._path_is_cached(local_path): + return local_path + + async with self._get_semaphore(): + if self._path_is_cached(local_path): + return local_path + claimed = self._try_claim_lock(local_path) + if claimed: + try: + return await self._download_owned(file_path, local_path, size=size) + finally: + # Sync release — to_thread can be cancelled and leak the lock. + self._release_lock(local_path) + + # Another process holds the lock; wait outside the semaphore. + try: + remaining = max(0.05, deadline - time.monotonic()) + return await self._wait_for_cached_file(local_path, timeout=remaining) + except TimeoutError: + continue + + if self._path_is_cached(local_path): + return local_path + raise TimeoutError(f"Timed out waiting for cache file {local_path}") + + async def _dedupe_path(self, key: str, factory: Callable[[], Coroutine[Any, Any, T]]) -> T: + """Coalesce concurrent work for the same cache key on this event loop.""" + loop = asyncio.get_running_loop() + if self._path_inflight_loop is not loop: + self._path_inflight = {} + self._path_inflight_loop = loop + + task = self._path_inflight.get(key) + if task is None: + task = asyncio.create_task(factory()) + self._path_inflight[key] = task + + def _clear(_t: asyncio.Task, k: str = key) -> None: + self._path_inflight.pop(k, None) + + task.add_done_callback(_clear) + return await task + + async def _ensure_cached_file(self, file_path: str, size: int | None = None) -> str: + """Ensure ``file_path`` is on disk; dedupe concurrent downloads of the same key.""" + local_path = self.get_local_path(file_path) + if self._path_is_cached(local_path): + return local_path + return await self._dedupe_path( + file_path, + lambda: self._download_to_cache(file_path, local_path, size=size), + ) + + async def _download_bytes_write_through(self, file_path: str, size: int | None = None) -> bytes: + local_path = self.get_local_path(file_path) + if self._path_is_cached(local_path): + try: return await asyncio.to_thread(Path(local_path).read_bytes) + except FileNotFoundError: + # Stale presence mark (e.g. file removed under us) — discard and retry once. + self._present_paths.discard(local_path) + if self._path_is_cached(local_path): + return await asyncio.to_thread(Path(local_path).read_bytes) + data = await self._fetch_bytes(file_path, size=size) + self._schedule_write_behind(local_path, data) + return data + async def download_file_async(self, file_path: str, size: int | None = None) -> bytes: + """Asynchronously download and return file content. + + With ``cache_files=True``, uses write-through: network bytes are returned + immediately and the cache file is published atomically in the background. + Concurrent callers for the same path share one in-flight fetch. + """ try: - return await self.downloader.adownload_fileobj(file_path) + if self.cache_files: + return await self._dedupe_path( + f"bytes:{file_path}", + lambda: self._download_bytes_write_through(file_path, size=size), + ) + return await self._fetch_bytes(file_path, size=size) + except Exception as e: + raise RuntimeError(f"Error downloading file {file_path}: {e}") from e + + async def ensure_file_async(self, file_path: str, size: int | None = None) -> str: + """Download to the mirrored cache path and return the local path (no full RAM buffer).""" + if not self.cache_files: + raise ValueError("ensure_file_async requires cache_files=True") + try: + return await self._ensure_cached_file(file_path, size=size) except Exception as e: raise RuntimeError(f"Error downloading file {file_path}: {e}") from e +def _storage_path(input_dir: Dir) -> str: + """Prefer cloud URL over FUSE/local path so downloads hit object storage directly.""" + return str(input_dir.url or input_dir.path) + + class StreamingRawDataset(Dataset): """Base class for streaming raw datasets. @@ -110,12 +1250,20 @@ def __init__( storage_options: dict | None = None, cache_files: bool = False, recompute_index: bool = False, - transform: Callable[[bytes | list[bytes]], Any] | None = None, + transform: Callable[[Any], Any] | None = None, + max_concurrent_downloads: int | None = None, + max_prefetch: int = 16, + prefetch_cache_size: int | None = None, + item_type: Literal["bytes", "path"] = "bytes", + hedge_delay: float = 0.0, + download_timeout: float = 120.0, + range_parallel_threshold: int = _RANGE_PARALLEL_THRESHOLD, + range_chunk_size: int = _RANGE_CHUNK_SIZE, ): """Initialize StreamingRawDataset. Args: - input_dir: Path to dataset root (e.g., 's3://bucket/dataset/'). + input_dir: Path to dataset root (e.g., 's3://bucket/dataset/' or Studio connection path). cache_dir: Directory for caching files (optional). indexer: Custom file indexer (default: FileIndexer). storage_options: Cloud storage options. @@ -124,29 +1272,118 @@ def __init__( If True, forces a re-scan of the input directory and rebuilds the index, ignoring any cached index files. This is useful when the dataset structure or files on the remote storage have changed. - transform: A function to apply to each item. It will receive `bytes` for single-file - items or `List[bytes]` for grouped items. + transform: A function to apply to each item. It receives ``bytes`` / ``list[bytes]`` + when ``item_type="bytes"``, or ``str`` / ``list[str]`` paths when ``item_type="path"``. + Prefer C-level / GIL-releasing transforms, or decode in ``collate_fn``. + max_concurrent_downloads: Per-worker in-flight download permits. + ``None`` (default) applies the Stage 1 adaptive formula (size-aware + aggregate budget from median file size; Little's-law arm only for + medians below 1 MiB, split across workers with a per-worker floor + of 8; single-process capped at 128). An explicit ``int`` sets + **exactly** that many permits with no silent clamp — pass ``64`` + to keep the historical fixed cap. + max_prefetch: Best-effort sequential look-ahead after each batch (default: 16; + roughly ``2×`` a typical batch). Pass ``0`` to disable. Look-ahead is per + DataLoader worker, but when ``num_workers > 1`` the scheduled amount is + capped to ``min(max_prefetch, 64 // num_workers)`` so aggregate look-ahead + stays near 64 items (e.g. w=2→16, w=8→8, w=16→4, w=32→2). Raise + ``max_prefetch`` only helps when it is below that per-worker share. + prefetch_cache_size: LRU entry cap for prefetched items. Defaults to + ``max(max_prefetch * 2, max_prefetch)`` when prefetch is enabled. + item_type: ``"bytes"`` (default) buffers each object in RAM; ``"path"`` downloads to + the cache and returns local path(s). ``item_type="path"`` requires ``cache_files=True``. + hedge_delay: Seconds before starting a hedged duplicate request for a slow GET + (``0`` = off, default). Opt in with a positive delay for object-store p99 + stragglers. Only applied to small objects (~<8MB); large objects use per-chunk + hedging for ranged downloads. + download_timeout: Hang-protection floor in seconds for each batch gather + (``0`` / disabled → no timeout). Defaults to ``120`` and coexists with the + per-item fast path: individual GETs are never wrapped in ``wait_for``; + ``_download_batch`` applies one ``wait_for`` around the gather using + ``max`` of the per-item size-aware floors + (``max(download_timeout, size / ~25MB/s * 3)``). + range_parallel_threshold: Objects at least this large use parallel ranged GETs + when the backend supports Range (``0`` disables; opt in with a positive + byte threshold via the constructor). + range_chunk_size: Part size for ranged parallel downloads. """ + if item_type not in ("bytes", "path"): + raise ValueError(f"item_type must be 'bytes' or 'path', got {item_type!r}") + if item_type == "path" and not cache_files: + raise ValueError("item_type='path' requires cache_files=True") + self.input_dir = _resolve_dir(input_dir) - self.cache_manager = CacheManager(self.input_dir, cache_dir, storage_options, cache_files) + self._storage_path = _storage_path(self.input_dir) + self.cache_files = cache_files + self.item_type = item_type + self.max_concurrent_downloads = max_concurrent_downloads + self.max_prefetch = max(0, max_prefetch) + self.hedge_delay = max(0.0, hedge_delay) + self.download_timeout = max(0.0, download_timeout) + if prefetch_cache_size is None: + prefetch_cache_size = max(self.max_prefetch * 2, self.max_prefetch) if self.max_prefetch > 0 else 0 + self.prefetch_cache_size = max(0, prefetch_cache_size) + + self.cache_manager = CacheManager( + self.input_dir, + cache_dir, + storage_options, + cache_files, + max_concurrent_downloads=max_concurrent_downloads, + hedge_delay=self.hedge_delay, + download_timeout=self.download_timeout, + range_parallel_threshold=range_parallel_threshold, + range_chunk_size=range_chunk_size, + ) self.indexer = indexer or FileIndexer() self.storage_options = storage_options or {} self.transform = transform - # Discover all files in the input directory. + self._prefetch_cache = _LRUCache(self.prefetch_cache_size) + self._inflight: dict[int, asyncio.Task] = {} + self._inflight_loop: asyncio.AbstractEventLoop | None = None + self._prefetch_hits = 0 + self._prefetch_misses = 0 + self._owner_pid = os.getpid() + + # Discover all files — prefer cloud URL over FUSE mount. self.files: list[FileMetadata] = self.indexer.build_or_load_index( - str(self.input_dir.path or self.input_dir.url), + self._storage_path, self.cache_manager.cache_dir, storage_options, recompute_index, ) - logger.info(f"Discovered {len(self.files)} files.") + logger.info("Discovered %s files.", len(self.files)) + median = _median_file_bytes(self.files) + self.cache_manager._median_file_bytes = median + self._maybe_warn_tiny_files(median) # Transform the flat list of files into the desired item structure. self.items: list[FileMetadata] | list[list[FileMetadata]] = self.setup(self.files) if not isinstance(self.items, list): raise TypeError(f"The setup method must return a list, but returned {type(self.items)}") - logger.info(f"Dataset setup with {len(self.items)} items.") + logger.info("Dataset setup with %s items.", len(self.items)) + + def _maybe_warn_tiny_files(self, median: int | None = None) -> None: + """Warn when index median size is tiny (request-overhead bound). + + ``median`` should be the value already computed for Stage 1 concurrency + sizing so we do not rescan sizes just for the warning. + """ + if median is None: + median = self.cache_manager._median_file_bytes + if median is None: + return + n_sized = sum(1 for f in self.files if f.size > 0) + if n_sized < 8: + return + if median < _TINY_FILE_MEDIAN_BYTES: + logger.warning( + "Median file size is %.0f bytes. StreamingRawDataset is often request-overhead " + "bound for tiny objects; consider litdata.optimize() + StreamingDataset for " + "higher sustained throughput.", + median, + ) def setup(self, files: list[FileMetadata]) -> list[FileMetadata] | list[list[FileMetadata]]: """Define the structure of the dataset from the list of discovered files. @@ -163,61 +1400,313 @@ def setup(self, files: list[FileMetadata]) -> list[FileMetadata] | list[list[Fil """ return files - @lru_cache(maxsize=1) def __len__(self) -> int: """Return the number of items in the dataset.""" return len(self.items) + def _ensure_post_fork_state(self) -> None: + """Drop parent-process asyncio/prefetch state after DataLoader fork.""" + pid = os.getpid() + if self._owner_pid == pid: + return + if _RAW_DEBUG: + logger.warning( + "raw-debug: reset dataset state after fork old_pid=%s new_pid=%s", + self._owner_pid, + pid, + ) + self._inflight = {} + self._inflight_loop = None + self._prefetch_cache = _LRUCache(self.prefetch_cache_size) + self._prefetch_hits = 0 + self._prefetch_misses = 0 + self._owner_pid = pid + self.cache_manager.reset_runtime_state() + + def __getstate__(self) -> dict[str, Any]: + """Serialize dataset config + index; strip loop/task/prefetch runtime. + + Allowlisted keys so accidental instance attrs (locks, write-behind refs) + cannot enter the spawn pickle payload. + """ + return { + "input_dir": self.input_dir, + "_storage_path": self._storage_path, + "cache_files": self.cache_files, + "item_type": self.item_type, + "max_concurrent_downloads": self.max_concurrent_downloads, + "max_prefetch": self.max_prefetch, + "hedge_delay": self.hedge_delay, + "download_timeout": self.download_timeout, + "prefetch_cache_size": self.prefetch_cache_size, + "cache_manager": self.cache_manager, + "indexer": self.indexer, + "storage_options": self.storage_options, + "transform": self.transform, + "files": self.files, + "items": self.items, + # Runtime — always fresh in the child (empty cache, no tasks/loops). + "_prefetch_cache": _LRUCache(self.prefetch_cache_size), + "_inflight": {}, + "_inflight_loop": None, + "_prefetch_hits": 0, + "_prefetch_misses": 0, + "_owner_pid": None, + } + + def __setstate__(self, state: dict[str, Any]) -> None: + self.__dict__.update(state) + self._owner_pid = os.getpid() + self._inflight = {} + self._inflight_loop = None + self._prefetch_hits = 0 + self._prefetch_misses = 0 + if not isinstance(self._prefetch_cache, _LRUCache): + self._prefetch_cache = _LRUCache(self.prefetch_cache_size) + # CacheManager may be restored without its __setstate__ on some paths. + cm = getattr(self, "cache_manager", None) + if isinstance(cm, CacheManager): + cm._downloader = None + cm._downloader_pid = None + cm._downloader_loop = None + cm._semaphore = None + cm._semaphore_loop = None + cm._semaphore_permits = None + cm._cached_permits = None + cm._cached_permits_pid = None + cm._path_inflight = {} + cm._path_inflight_loop = None + cm._range_executor = None + cm._range_executor_pid = None + def __getitem__(self, index: int) -> Any: """Get a single item by index.""" if not (0 <= index < len(self)): raise IndexError(f"Index {index} out of range for dataset with length {len(self)}") + self._ensure_post_fork_state() + if _RAW_DEBUG: + logger.warning("raw-debug: __getitem__ pid=%s index=%s", os.getpid(), index) + return _run_async(self._download_batch([index]))[0] + def __getitems__(self, indices: list[int]) -> list[Any]: + """Asynchronously download a batch of items by indices.""" + self._ensure_post_fork_state() + if _RAW_DEBUG: + logger.warning( + "raw-debug: __getitems__ pid=%s n=%s head=%s", + os.getpid(), + len(indices), + indices[:4], + ) + return _run_async(self._download_batch(indices)) + + def _item_size(self, index: int) -> int: item = self.items[index] if isinstance(item, FileMetadata): - return asyncio.run(self._download_and_process_item(item.path)) + return item.size if isinstance(item, list): - file_paths = [fm.path for fm in item] - return asyncio.run(self._download_and_process_group(file_paths)) - raise TypeError(f"Dataset items must be of type FileMetadata or List[FileMetadata], but found {type(item)}") + return sum(fm.size for fm in item) + return 0 - def __getitems__(self, indices: list[int]) -> list[Any]: - """Asynchronously download a batch of items by indices.""" - # asyncio.run() handles loop creation, execution, and teardown cleanly. - return asyncio.run(self._download_batch(indices)) + def _batch_download_budget(self, indices: list[int]) -> float | None: + """Return one hang-protection budget for a pending batch, or ``None`` if disabled. + + Uses the max per-item size-aware floor so large objects are not cut off, while + keeping a single ``wait_for`` around the gather (not per object). + + With a download semaphore smaller than the batch (e.g. ``batch_size=64`` and + fewer Stage 1 permits), downloads span multiple waves — the single-wave + ``max(per-item)`` assumption is no longer the default. ``sum(sizes) / + bandwidth * 3`` (using the per-stream ``_HEDGE_ASSUMED_BANDWIDTH_BPS`` floor) + raises the timeout when aggregate transfer time exceeds that single-item + budget. + """ + base = self.cache_manager.download_timeout + if base is None: + return None + budgets = [self.cache_manager._download_budget(self._item_size(i)) for i in indices] + sized = [b for b in budgets if b is not None] + per_item = max(sized) if sized else base + total_size = sum(self._item_size(i) for i in indices) + if total_size > 0: + return max(per_item, total_size / _HEDGE_ASSUMED_BANDWIDTH_BPS * 3.0) + return per_item + + async def _resolve_index(self, index: int) -> Any: + """Return a cached/inflight/materialized item for ``index``.""" + task = self._inflight.get(index) + if task is not None: + if task.cancelled(): + self._inflight.pop(index, None) + else: + value = await task + self._prefetch_cache.put(index, value) + self._inflight.pop(index, None) + return value + value = await self._materialize_index(index) + self._prefetch_cache.put(index, value) + return value async def _download_batch(self, indices: list[int]) -> list[Any]: - """Asynchronously download and process items.""" - batch_items = [self.items[i] for i in indices] - coros = [] - for item in batch_items: - if isinstance(item, FileMetadata): - coros.append(self._download_and_process_item(item.path)) - elif isinstance(item, list): - file_paths = [fm.path for fm in item] - coros.append(self._download_and_process_group(file_paths)) + """Download/process indices, serving from prefetch cache when possible.""" + if _RAW_DEBUG: + logger.warning("raw-debug: _download_batch start pid=%s n=%s", os.getpid(), len(indices)) + for index in indices: + if not (0 <= index < len(self)): + raise IndexError(f"Index {index} out of range for dataset with length {len(self)}") + + running_loop = asyncio.get_running_loop() + if self._inflight_loop is not running_loop: + self._inflight = {} + self._inflight_loop = running_loop + + results: list[Any] = [None] * len(indices) + pending_positions: dict[int, list[int]] = {} # index -> [pos, ...] + unique_pending: list[int] = [] + batch_hits = 0 + batch_misses = 0 + for pos, index in enumerate(indices): + cached = self._prefetch_cache.get(index) + if cached is not _MISS: + results[pos] = cached + batch_hits += 1 else: - raise TypeError( - f"Dataset items must be of type FileMetadata or List[FileMetadata], but found {type(item)}" - ) - return await asyncio.gather(*coros) + batch_misses += 1 + if index not in pending_positions: + pending_positions[index] = [] + unique_pending.append(index) + pending_positions[index].append(pos) + + if _RAW_DEBUG: + self._prefetch_hits += batch_hits + self._prefetch_misses += batch_misses + + if unique_pending: + # Largest-first so big objects overlap with smaller ones (LPT). + unique_pending.sort(key=self._item_size, reverse=True) + tasks = [asyncio.create_task(self._resolve_index(index)) for index in unique_pending] + budget = self._batch_download_budget(unique_pending) + try: + if budget is None: + fetched = await asyncio.gather(*tasks) + else: + # One wait_for for the whole batch — hang protection without per-item tax. + fetched = await asyncio.wait_for(asyncio.gather(*tasks), timeout=budget) + except (TimeoutError, asyncio.TimeoutError) as exc: + # Plain gather (download_timeout=0): item-level TimeoutError — do not rewrite. + if budget is None: + raise + for task in tasks: + if not task.done(): + task.cancel() + # Cancelling _resolve_index wrappers does not cancel awaited _inflight + # download tasks; hung prefetch entries would otherwise poison retries. + # Prefetch tasks outside unique_pending survive here and self-heal when + # a later batch needs them (one extra budget delay). + for idx in unique_pending: + inflight = self._inflight.pop(idx, None) + if inflight is not None and not inflight.done(): + inflight.cancel() + inflight.add_done_callback(_consume_task_exception) + await asyncio.gather(*tasks, return_exceptions=True) + raise TimeoutError( + f"Batch download timed out after {budget:.1f}s ({len(unique_pending)} pending indices)" + ) from exc + for index, value in zip(unique_pending, fetched): + for pos in pending_positions[index]: + results[pos] = value + + if self.max_prefetch > 0: + self._schedule_prefetch(indices) + if _RAW_DEBUG: + logger.warning( + "raw-debug: _download_batch done pid=%s n=%s inflight=%s " + "batch_hit=%s batch_miss=%s total_hit=%s total_miss=%s", + os.getpid(), + len(indices), + len(self._inflight), + batch_hits, + batch_misses, + self._prefetch_hits, + self._prefetch_misses, + ) + return results + + def _schedule_prefetch(self, indices: list[int]) -> None: + """Best-effort sequential look-ahead into the LRU cache. + + With ``DataLoader(num_workers>1)``, each worker receives every N-th batch, so the + next indices for *this* worker start at ``indices[0] + num_workers * batch_len``. + Scheduled look-ahead uses :func:`_effective_prefetch` so aggregate in-flight cost + stays near ``_AGGREGATE_PREFETCH_BUDGET`` rather than ``num_workers × max_prefetch``. + """ + if self.max_prefetch <= 0 or not indices or not _looks_sequential(indices): + return - async def _download_and_process_group(self, file_paths: list[str]) -> Any: + try: + from torch.utils.data import get_worker_info + + info = get_worker_info() + num_workers = info.num_workers if info is not None else 1 + except Exception: + num_workers = 1 + + effective = _effective_prefetch(self.max_prefetch, num_workers) + if effective <= 0: + return + + batch_len = len(indices) + start = indices[0] + num_workers * batch_len + end = min(start + effective, len(self.items)) + for index in range(start, end): + if index in self._prefetch_cache or index in self._inflight: + continue + task = asyncio.create_task(self._prefetch_index(index)) + task.add_done_callback(_consume_prefetch_exception) + self._inflight[index] = task + + async def _prefetch_index(self, index: int) -> Any: + try: + value = await self._materialize_index(index) + self._prefetch_cache.put(index, value) + return value + finally: + self._inflight.pop(index, None) + + async def _materialize_index(self, index: int) -> Any: + item = self.items[index] + if isinstance(item, FileMetadata): + return await self._download_and_process_item(item.path, size=item.size) + if isinstance(item, list): + file_paths = [fm.path for fm in item] + sizes = [fm.size for fm in item] + return await self._download_and_process_group(file_paths, sizes=sizes) + raise TypeError(f"Dataset items must be of type FileMetadata or List[FileMetadata], but found {type(item)}") + + async def _download_and_process_group( + self, file_paths: list[str], sizes: Sequence[int | None] | None = None + ) -> Any: """Download all files in a group, then apply the transform.""" - download_coros = [self.cache_manager.download_file_async(path) for path in file_paths] - group_data: list[bytes] = await asyncio.gather(*download_coros) + resolved_sizes: list[int | None] = list(sizes) if sizes is not None else [None] * len(file_paths) + if self.item_type == "path": + group_data: list[Any] = await asyncio.gather( + *[self.cache_manager.ensure_file_async(path, size=sz) for path, sz in zip(file_paths, resolved_sizes)] + ) + else: + group_data = await asyncio.gather( + *[self.cache_manager.download_file_async(path, size=sz) for path, sz in zip(file_paths, resolved_sizes)] + ) if self.transform: - # The transform receives a list of bytes, corresponding to the list structure - # of the item defined in setup(). This is true even if the list has only one element. return await asyncio.to_thread(self.transform, group_data) return group_data - async def _download_and_process_item(self, file_path: str) -> Any: + async def _download_and_process_item(self, file_path: str, size: int | None = None) -> Any: """Download a single file and apply the transform.""" - data: bytes = await self.cache_manager.download_file_async(file_path) + if self.item_type == "path": + data: Any = await self.cache_manager.ensure_file_async(file_path, size=size) + else: + data = await self.cache_manager.download_file_async(file_path, size=size) if self.transform: - # The transform receives a single bytes object, corresponding to the - # single FileMetadata object structure of the item. return await asyncio.to_thread(self.transform, data) return data diff --git a/src/litdata/raw/indexer.py b/src/litdata/raw/indexer.py index 06640708d..7911c94a3 100644 --- a/src/litdata/raw/indexer.py +++ b/src/litdata/raw/indexer.py @@ -11,6 +11,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextlib import json import logging import os @@ -26,6 +27,29 @@ logger = logging.getLogger(__name__) _SUPPORTED_PROVIDERS = ("s3", "gs", "azure") _INDEX_FILENAME = "index.json.zstd" +# Warn once per process when remote index upload is denied (read-only creds). +_UPLOAD_DENIED_WARNED_PIDS: set[int] = set() + + +def _is_windows_drive_scheme(scheme: str) -> bool: + r"""True when ``urlparse`` mistook a Windows drive letter for a URI scheme. + + Paths like ``C:\\Users\\...`` parse with ``scheme='c'``. A single-letter scheme is + never a valid URI scheme (RFC 3986 requires >=2 chars), so treat it as local. + """ + return len(scheme) == 1 and scheme.isalpha() + + +def _validate_input_dir_scheme(input_dir: str) -> None: + """Raise if ``input_dir`` uses an unsupported remote scheme. + + Local paths (including Windows drive letters) are allowed. + """ + scheme = urlparse(input_dir).scheme + if scheme and not _is_windows_drive_scheme(scheme) and scheme not in _SUPPORTED_PROVIDERS: + raise ValueError( + f"Unsupported input directory scheme: `{scheme}`. Supported schemes are: {_SUPPORTED_PROVIDERS}" + ) @dataclass @@ -80,12 +104,7 @@ def build_or_load_index( if not _FSSPEC_AVAILABLE: raise ModuleNotFoundError(str(_FSSPEC_AVAILABLE)) - parsed_url = urlparse(input_dir) - if parsed_url.scheme and parsed_url.scheme not in _SUPPORTED_PROVIDERS: - raise ValueError( - f"Unsupported input directory scheme: `{parsed_url.scheme}`. " - f"Supported schemes are: {_SUPPORTED_PROVIDERS}" - ) + _validate_input_dir_scheme(input_dir) if not recompute_index: files = self._load_index_from_cache(input_dir, cache_dir, storage_options) @@ -125,7 +144,13 @@ def _load_index_from_cache( def _build_and_cache_index( self, input_dir: str, cache_dir: str, storage_options: dict[str, Any] | None ) -> list[FileMetadata]: - """Builds a new index and caches it locally and remotely.""" + """Build a new index and cache it locally and remotely. + + By default the index is also uploaded into the dataset's input bucket + (``input_dir/index.json.zstd``). That requires write credentials on the + source bucket; read-only credentials log a warning and continue with the + local cache only. + """ local_index_path = Path(cache_dir) / _INDEX_FILENAME logger.info(f"Building index for {input_dir} at {local_index_path}") files = self.discover_files(input_dir, storage_options) @@ -134,13 +159,22 @@ def _build_and_cache_index( self._save_index_file(str(local_index_path), files, input_dir) - # Upload to remote cache + # Upload to remote cache (source bucket by default — needs write creds). remote_index_path = os.path.join(input_dir, _INDEX_FILENAME) try: self._upload_to_cloud(str(local_index_path), remote_index_path, storage_options) logger.info(f"Uploaded index to remote cache: {remote_index_path}") except Exception as e: - logger.warning(f"Failed to upload index to remote cache: {e}") + pid = os.getpid() + if pid not in _UPLOAD_DENIED_WARNED_PIDS: + _UPLOAD_DENIED_WARNED_PIDS.add(pid) + logger.warning( + "Failed to upload index to remote cache (continuing with local index; " + "further upload failures in this process are silent): %s", + e, + ) + else: + logger.info("Failed to upload index to remote cache: %s", e) logger.info(f"Built index with {len(files)} files from {input_dir} at {local_index_path}") return files @@ -164,7 +198,7 @@ def _load_index_file(self, index_path: str) -> list[FileMetadata] | None: return None def _save_index_file(self, index_path: str, files: list[FileMetadata], source: str) -> None: - """Encodes and saves an index file.""" + """Encode and atomically publish an index file (tmp + ``os.replace``).""" if _PYTHON_GREATER_EQUAL_3_14: from compression import zstd from compression.zstd import ZstdError @@ -172,16 +206,22 @@ def _save_index_file(self, index_path: str, files: list[FileMetadata], source: s import zstd from zstd import Error as ZstdError + tmp_path = f"{index_path}.tmp.{os.getpid()}" try: metadata = { "source": source, "files": [file.to_dict() for file in files], "created_at": time.time(), } - with open(index_path, "wb") as f: + os.makedirs(os.path.dirname(index_path) or ".", exist_ok=True) + with open(tmp_path, "wb") as f: f.write(zstd.compress(json.dumps(metadata).encode("utf-8"))) + os.replace(tmp_path, index_path) except (OSError, ZstdError) as e: logger.warning(f"Error caching index to {index_path}: {e}") + finally: + with contextlib.suppress(OSError): + os.remove(tmp_path) def _download_from_cloud( self, @@ -189,14 +229,21 @@ def _download_from_cloud( local_path: str, storage_options: dict[str, Any] | None, ) -> None: - """Downloads a file from cloud storage.""" + """Download a file from cloud storage via tmp + ``os.replace`` (atomic publish).""" if not _FSSPEC_AVAILABLE: raise ModuleNotFoundError(str(_FSSPEC_AVAILABLE)) import fsspec parsed_url = urlparse(remote_path) fs = fsspec.filesystem(parsed_url.scheme, **(storage_options or {})) - fs.get(remote_path, local_path) + tmp_path = f"{local_path}.tmp.{os.getpid()}" + try: + os.makedirs(os.path.dirname(local_path) or ".", exist_ok=True) + fs.get(remote_path, tmp_path) + os.replace(tmp_path, local_path) + finally: + with contextlib.suppress(OSError): + os.remove(tmp_path) def _upload_to_cloud( self, @@ -227,13 +274,10 @@ def __init__( def discover_files(self, input_dir: str, storage_options: dict[str, Any] | None) -> list[FileMetadata]: """Discover dataset files and return their metadata.""" - parsed_url = urlparse(input_dir) - if parsed_url.scheme and parsed_url.scheme not in _SUPPORTED_PROVIDERS: - raise ValueError( - f"Unsupported input directory scheme: `{parsed_url.scheme}`. " - f"Supported schemes are: {_SUPPORTED_PROVIDERS}" - ) + _validate_input_dir_scheme(input_dir) + parsed_url = urlparse(input_dir) + # Windows drive letters parse as single-letter schemes; treat those as local. if parsed_url.scheme in _SUPPORTED_PROVIDERS: # Cloud storage return self._discover_cloud_files(input_dir, storage_options) diff --git a/src/litdata/streaming/client.py b/src/litdata/streaming/client.py index eebc22ee1..6eca33780 100644 --- a/src/litdata/streaming/client.py +++ b/src/litdata/streaming/client.py @@ -13,6 +13,7 @@ import json import os +import threading from time import time from typing import Any @@ -116,6 +117,17 @@ def __init__( self._client: Any | None = None self._storage_options: dict = storage_options or {} self._session_options: dict = session_options or {} + # Guards lazy create + credential refresh (range GETs hit .client from many threads). + self._client_lock = threading.Lock() + + def __getstate__(self) -> dict[str, Any]: + state = self.__dict__.copy() + state.pop("_client_lock", None) + return state + + def __setstate__(self, state: dict[str, Any]) -> None: + self.__dict__.update(state) + self._client_lock = threading.Lock() def _create_client(self) -> None: # S3 data connections marked available on non-AWS providers can't reach the bucket via the @@ -177,16 +189,18 @@ def _create_client_from_temp_credentials(self, data_connection_id: str) -> None: @property def client(self) -> Any: - if self._client is None: - self._create_client() - self._last_time = time() - - # Re-generate credentials for EC2 - if self._last_time is None or (time() - self._last_time) > self._refetch_interval: - self._create_client() - self._last_time = time() - - return self._client + # boto3 clients are thread-safe for requests; construction/refresh is not. + with self._client_lock: + if self._client is None: + self._create_client() + self._last_time = time() + + # Re-generate credentials for EC2 / temporary Studio creds + if self._last_time is None or (time() - self._last_time) > self._refetch_interval: + self._create_client() + self._last_time = time() + + return self._client class R2Client(S3Client): diff --git a/src/litdata/streaming/downloader.py b/src/litdata/streaming/downloader.py index b612978d8..2ec9e5e58 100644 --- a/src/litdata/streaming/downloader.py +++ b/src/litdata/streaming/downloader.py @@ -16,6 +16,7 @@ import os import shutil import tempfile +import threading from abc import ABC from contextlib import suppress from time import time @@ -56,6 +57,15 @@ def _obstore_stream_min_chunk_size() -> int: class Downloader(ABC): + """Cloud/local chunk downloader. + + Implementers should: + - Publish cache files atomically (temp path + ``os.replace``; see ``_temp_download_path``). + - Be safe for concurrent calls from multiple threads (or document otherwise). + - Prefer real HTTP Range in ``download_bytes`` when the backend supports it. + - Clean up ``.tmp.*`` paths on failure. + """ + def __init__( self, remote_dir: str, @@ -225,9 +235,7 @@ def download_bytes(self, remote_filepath: str, offset: int, length: int, local_c if obj.scheme != "s3": raise ValueError(f"Expected obj.scheme to be `s3`, instead, got {obj.scheme} for remote={remote_filepath}") - if not hasattr(self, "client"): - self._client = S3Client(storage_options=self._storage_options, session_options=self.session_options) - + # self._client is created in __init__; S3Client.client serializes create/refresh. bucket = obj.netloc key = obj.path.lstrip("/") @@ -244,9 +252,6 @@ def download_fileobj(self, remote_filepath: str, fileobj: Any) -> None: if obj.scheme != "s3": raise ValueError(f"Expected obj.scheme to be `s3`, instead, got {obj.scheme} for remote={remote_filepath}") - if not hasattr(self, "_client"): - self._client = S3Client(storage_options=self._storage_options, session_options=self.session_options) - bucket = obj.netloc key = obj.path.lstrip("/") @@ -376,9 +381,7 @@ def download_bytes(self, remote_filepath: str, offset: int, length: int, local_c if obj.scheme != "r2": raise ValueError(f"Expected obj.scheme to be `r2`, instead, got {obj.scheme} for remote={remote_filepath}") - if not hasattr(self, "_client"): - self._client = R2Client(storage_options=self._storage_options, session_options=self.session_options) - + # self._client is created in __init__; R2Client.client serializes create/refresh. bucket = obj.netloc key = obj.path.lstrip("/") @@ -395,9 +398,6 @@ def download_fileobj(self, remote_filepath: str, fileobj: Any) -> None: if obj.scheme != "r2": raise ValueError(f"Expected obj.scheme to be `r2`, instead, got {obj.scheme} for remote={remote_filepath}") - if not hasattr(self, "_client"): - self._client = R2Client(storage_options=self._storage_options, session_options=self.session_options) - bucket = obj.netloc key = obj.path.lstrip("/") @@ -452,10 +452,21 @@ def __init__( raise ModuleNotFoundError(str(_GOOGLE_STORAGE_AVAILABLE)) super().__init__(remote_dir, cache_dir, chunks, storage_options) + self._client: Any | None = None + self._client_lock = threading.Lock() - def download_file(self, remote_filepath: str, local_filepath: str) -> None: - from google.cloud import storage + def _get_client(self) -> Any: + """Return a cached ``google.cloud.storage.Client`` (thread-safe lazy init).""" + if self._client is not None: + return self._client + with self._client_lock: + if self._client is None: + from google.cloud import storage + self._client = storage.Client(**self._storage_options) + return self._client + + def download_file(self, remote_filepath: str, local_filepath: str) -> None: obj = parse.urlparse(remote_filepath) if obj.scheme != "gs": @@ -477,7 +488,7 @@ def download_file(self, remote_filepath: str, local_filepath: str) -> None: if key[0] == "/": key = key[1:] - client = storage.Client(**self._storage_options) + client = self._get_client() bucket = client.bucket(bucket_name) blob = bucket.blob(key) tmp_path = self._temp_download_path(local_filepath) @@ -490,8 +501,6 @@ def download_file(self, remote_filepath: str, local_filepath: str) -> None: raise def download_bytes(self, remote_filepath: str, offset: int, length: int, local_chunkpath: str) -> bytes: - from google.cloud import storage - obj = parse.urlparse(remote_filepath) if obj.scheme != "gs": @@ -500,7 +509,7 @@ def download_bytes(self, remote_filepath: str, offset: int, length: int, local_c bucket_name = obj.netloc key = obj.path.lstrip("/") - client = storage.Client(**self._storage_options) + client = self._get_client() bucket = client.bucket(bucket_name) blob = bucket.blob(key) @@ -511,8 +520,6 @@ def download_bytes(self, remote_filepath: str, offset: int, length: int, local_c def download_fileobj(self, remote_filepath: str, fileobj: Any) -> None: """Download a file from GCS directly to a file-like object.""" - from google.cloud import storage - obj = parse.urlparse(remote_filepath) if obj.scheme != "gs": @@ -521,7 +528,7 @@ def download_fileobj(self, remote_filepath: str, fileobj: Any) -> None: bucket_name = obj.netloc key = obj.path.lstrip("/") - client = storage.Client(**self._storage_options) + client = self._get_client() bucket = client.bucket(bucket_name) blob = bucket.blob(key) @@ -532,11 +539,10 @@ def _get_store(self, bucket: str) -> Any: if not hasattr(self, "_store"): if not _OBSTORE_AVAILABLE: raise ModuleNotFoundError(str(_OBSTORE_AVAILABLE)) - from google.cloud import storage from obstore.auth.google import GoogleCredentialProvider from obstore.store import GCSStore - client = storage.Client(**self._storage_options) + client = self._get_client() credential_provider = GoogleCredentialProvider(credentials=client._credentials) self._store = GCSStore(bucket, credential_provider=credential_provider) return self._store @@ -657,6 +663,18 @@ async def adownload_fileobj(self, remote_filepath: str) -> bytes: class LocalDownloader(Downloader): + async def adownload_fileobj(self, remote_filepath: str) -> bytes: + """Read a local file (sync I/O; avoids leaking default-executor threads in tests).""" + from pathlib import Path + + return Path(remote_filepath).read_bytes() + + async def adownload_file(self, remote_filepath: str, local_filepath: str) -> None: + """Copy a local file into the cache path.""" + if os.path.exists(local_filepath): + return + self.download_file(remote_filepath, local_filepath) + def download_file(self, remote_filepath: str, local_filepath: str) -> None: if not os.path.exists(remote_filepath): raise FileNotFoundError(f"The provided remote_path doesn't exist: {remote_filepath}") @@ -673,10 +691,9 @@ def download_file(self, remote_filepath: str, local_filepath: str) -> None: temp_file_path = local_filepath + ".tmp" shutil.copy(remote_filepath, temp_file_path) os.rename(temp_file_path, local_filepath) - # FileLock doesn't delete its lock file on release — we clean it up manually. - # This must happen after release (Windows can't delete open files) and after the - # work is done (on Linux, deleting an in-use lock file lets other processes lock - # on a new inode, bypassing mutual exclusion). + # FileLock leaves the lock path behind; remove it after release when we held it. + # Delete only after the critical section so other waiters do not race a new inode + # while we still expected exclusive access. if lock_acquired: with contextlib.suppress(Exception): os.remove(lock_path) @@ -741,6 +758,8 @@ def download_file(self, remote_filepath: str, local_filepath: str) -> None: super().download_file(remote_filepath, local_filepath) +# TODO(follow-up): parametrized Downloader conformance suite over _DOWNLOADERS +# (atomic publish, tmp cleanup, download_bytes correctness + concurrent safety). _DOWNLOADERS: dict[str, type[Downloader]] = { "s3://": S3Downloader, "gs://": GCPDownloader, diff --git a/src/litdata/streaming/resolver.py b/src/litdata/streaming/resolver.py index f097f58d7..d2ad8c3c0 100644 --- a/src/litdata/streaming/resolver.py +++ b/src/litdata/streaming/resolver.py @@ -433,6 +433,21 @@ def _get_lightning_cloud_url() -> str: return os.getenv("LIGHTNING_CLOUD_URL", "https://lightning.ai") +_TIME_TEMPLATE_RE = re.compile(r"^.*{%.*}.*$") + + +def _has_time_template(path: str | Path | Dir | None) -> bool: + """Return True if ``path`` contains a LitData ``{%strftime}`` placeholder. + + Used to decide whether multi-node ``broadcast_object`` should align resolved + dirs (each rank would otherwise expand ``datetime.now()`` independently). + Already-resolved :class:`~litdata.streaming.cache.Dir` values return False. + """ + if path is None or isinstance(path, Dir): + return False + return _TIME_TEMPLATE_RE.search(str(path)) is not None + + def _resolve_time_template(path: str) -> str: """Resolves a datetime pattern in the given path string. @@ -449,8 +464,7 @@ def _resolve_time_template(path: str) -> str: Returns: str: The path with the datetime placeholder replaced by the current timestamp. """ - match = re.search("^.*{%.*}.*$", path) - if match is None: + if not _has_time_template(path): return path pattern = path.split("{")[1].split("}")[0] diff --git a/tests/processing/test_data_processor.py b/tests/processing/test_data_processor.py index 505f38b28..719a9e773 100644 --- a/tests/processing/test_data_processor.py +++ b/tests/processing/test_data_processor.py @@ -1742,3 +1742,126 @@ def test_data_processor_end_to_end_with_data_connection_id(tmpdir, monkeypatch): # Note: Due to the complexity of mocking the full pipeline, we mainly verify # that the fs_provider was called, indicating the data_connection_id code paths were executed assert len(calls_made) > 0, "fs_provider should have been called" + + +@pytest.mark.parametrize( + ("output_dir", "broadcast_paths", "expect_broadcast"), + [ + # Default off for ordinary paths + ("/data/out", False, False), + # Explicit True always broadcasts + ("local/out", True, True), + # `{%strftime}` time template auto-enables broadcast + ("local/out_{%Y-%m-%d}", False, True), + ("s3://bucket/run_{%Y-%m-%d_%H-%M-%S}", False, True), + ], +) +def test_data_processor_broadcast_paths(tmpdir, monkeypatch, output_dir, broadcast_paths, expect_broadcast): + """broadcast_paths defaults off; auto-on for `{%strftime}` paths; explicit True forces broadcast.""" + broadcast_mock = mock.MagicMock(side_effect=lambda key, obj, rank: obj) + monkeypatch.setattr(data_processor_module, "broadcast_object", broadcast_mock) + + DataProcessor( + input_dir=str(tmpdir), + output_dir=output_dir, + num_workers=1, + verbose=False, + broadcast_paths=broadcast_paths, + ) + + if expect_broadcast: + assert broadcast_mock.call_count == 2 + assert broadcast_mock.call_args_list[0].args[0] == "input_dir" + assert broadcast_mock.call_args_list[1].args[0] == "output_dir" + else: + broadcast_mock.assert_not_called() + + +def test_data_processor_broadcast_paths_default_false(tmpdir, monkeypatch): + broadcast_mock = mock.MagicMock(side_effect=lambda key, obj, rank: obj) + monkeypatch.setattr(data_processor_module, "broadcast_object", broadcast_mock) + + processor = DataProcessor(input_dir=str(tmpdir), output_dir=str(tmpdir / "out"), num_workers=1, verbose=False) + + assert processor.broadcast_paths is False + broadcast_mock.assert_not_called() + + +def test_optimize_broadcast_paths_auto_on_for_time_template(tmpdir, monkeypatch): + """Optimize detects `{%strftime}` on the unresolved path before `_resolve_dir`.""" + captured: dict[str, Any] = {} + + class CaptureDataProcessor(DataProcessor): + def __init__(self, *args, **kwargs): + captured["broadcast_paths"] = kwargs.get("broadcast_paths") + super().__init__(*args, **kwargs) + + def run(self, data_recipe): + return None + + monkeypatch.setattr(functions, "DataProcessor", CaptureDataProcessor) + monkeypatch.setattr(functions, "_assert_dir_has_index_file", mock.MagicMock()) + + optimize( + fn=lambda x: x, + inputs=[1, 2, 3], + output_dir=str(tmpdir / "out_{%Y-%m-%d}"), + chunk_size=2, + num_workers=1, + verbose=False, + ) + + assert captured["broadcast_paths"] is True + + +def test_optimize_broadcast_paths_default_off(tmpdir, monkeypatch): + captured: dict[str, Any] = {} + + class CaptureDataProcessor(DataProcessor): + def __init__(self, *args, **kwargs): + captured["broadcast_paths"] = kwargs.get("broadcast_paths") + super().__init__(*args, **kwargs) + + def run(self, data_recipe): + return None + + monkeypatch.setattr(functions, "DataProcessor", CaptureDataProcessor) + monkeypatch.setattr(functions, "_assert_dir_has_index_file", mock.MagicMock()) + + optimize( + fn=lambda x: x, + inputs=[1, 2, 3], + output_dir=str(tmpdir / "out"), + chunk_size=2, + num_workers=1, + verbose=False, + ) + + assert captured["broadcast_paths"] is False + + +def test_optimize_broadcast_paths_explicit_true(tmpdir, monkeypatch): + captured: dict[str, Any] = {} + + class CaptureDataProcessor(DataProcessor): + def __init__(self, *args, **kwargs): + captured["broadcast_paths"] = kwargs.get("broadcast_paths") + super().__init__(*args, **kwargs) + + def run(self, data_recipe): + return None + + monkeypatch.setattr(functions, "DataProcessor", CaptureDataProcessor) + monkeypatch.setattr(functions, "_assert_dir_has_index_file", mock.MagicMock()) + + optimize( + fn=lambda x: x, + inputs=[1, 2, 3], + output_dir=str(tmpdir / "out"), + chunk_size=2, + num_workers=1, + verbose=False, + broadcast_paths=True, + ) + + assert captured["broadcast_paths"] is True diff --git a/tests/raw/conftest.py b/tests/raw/conftest.py new file mode 100644 index 000000000..527748830 --- /dev/null +++ b/tests/raw/conftest.py @@ -0,0 +1,14 @@ +"""Fixtures for raw streaming tests.""" + +from __future__ import annotations + +import pytest + + +@pytest.fixture(autouse=True) +def _shutdown_raw_loop_runner(): + """Ensure the process-local LoopRunner does not leak threads across tests.""" + yield + from litdata.raw import dataset as raw_dataset + + raw_dataset._shutdown_runner_before_fork() diff --git a/tests/raw/test_dataset.py b/tests/raw/test_dataset.py index 9eedf41f4..bead4c069 100644 --- a/tests/raw/test_dataset.py +++ b/tests/raw/test_dataset.py @@ -40,16 +40,206 @@ def test_get_local_path(tmp_path): assert local_path.startswith(manager.cache_dir) +@pytest.mark.skipif(condition=sys.platform == "win32", reason="Not supported on windows") +def test_streaming_raw_dataset_default_max_prefetch(tmp_path): + """Default max_prefetch is a positive look-ahead (16).""" + (tmp_path / "file1.jpg").write_bytes(b"x") + dataset = StreamingRawDataset(input_dir=str(tmp_path), cache_files=False) + assert dataset.max_prefetch == 16 + assert dataset.prefetch_cache_size == 32 + + +@pytest.mark.parametrize( + ("num_workers", "max_prefetch", "expected"), + [ + (1, 16, 16), + (0, 16, 16), # treated as single-process (≤1) + (2, 16, 16), # min(16, 64//2) = 16 + (4, 16, 16), # min(16, 64//4) = 16 + (8, 16, 8), # min(16, 64//8) = 8 + (16, 16, 4), # min(16, 64//16) = 4 + (24, 16, 2), # min(16, 64//24) = 2 + (32, 16, 2), # min(16, 64//32) = 2 + (32, 32, 2), # still capped by aggregate budget + (2, 32, 32), # min(32, 64//2) = 32 + (8, 0, 0), + ], +) +def test_effective_prefetch_vs_num_workers(num_workers, max_prefetch, expected): + from litdata.raw.dataset import _effective_prefetch + + assert _effective_prefetch(max_prefetch, num_workers) == expected + + +@pytest.mark.parametrize( + ("num_workers", "max_concurrent", "median_bytes", "expected"), + [ + # Explicit int → exactly that many permits (no silent clamp), any worker count + (0, 64, 100_000, 64), + (1, 64, 100_000, 64), + (24, 64, 100_000, 64), + (32, 64, 100_000, 64), + (2, 4, 100_000, 4), + # Adaptive (None): ~100KB JPEG → bandwidth≈524 → cap 512 + (0, None, 100_000, 128), # single-process cap + (1, None, 100_000, 128), + (2, None, 100_000, 256), # 512//2 + (4, None, 100_000, 128), # 512//4 + (8, None, 100_000, 64), # 512//8 + (16, None, 100_000, 32), # 512//16 + (24, None, 100_000, 21), # 512//24 + (32, None, 100_000, 16), # 512//32 + # Large objects (≥1 MiB): bandwidth-only (no Little's-law pin at 240) + (4, None, 10 * 1024 * 1024, 8), # budget=floor 32, 32//4=8 + (16, None, 10 * 1024 * 1024, 8), # 32//16=2 → floor 8 + # Unknown size uses default median (256KiB) → latency arm (240) + (8, None, None, 30), # 240//8 + ], +) +def test_effective_concurrency_vs_num_workers(num_workers, max_concurrent, median_bytes, expected): + from litdata.raw.dataset import _effective_concurrency + + assert _effective_concurrency(max_concurrent, num_workers, median_bytes) == expected + + +def test_aggregate_concurrency_budget_clamps(): + from litdata.raw.dataset import ( + _AGGREGATE_CONCURRENCY_BUDGET_CAP, + _AGGREGATE_CONCURRENCY_BUDGET_FLOOR, + _ASSUMED_AGGREGATE_BANDWIDTH_BPS, + _ASSUMED_REQUEST_LATENCY_S, + _ASSUMED_REQUEST_RATE, + _CONCURRENCY_PIPELINE_SECONDS, + _aggregate_concurrency_budget, + ) + + latency = int(_ASSUMED_REQUEST_RATE * _ASSUMED_REQUEST_LATENCY_S) # ~240 + target_bytes = int(_ASSUMED_AGGREGATE_BANDWIDTH_BPS * _CONCURRENCY_PIPELINE_SECONDS) + assert _aggregate_concurrency_budget(1) == _AGGREGATE_CONCURRENCY_BUDGET_CAP + # Tiny ImageNet-like: bandwidth wins over latency, then hits cap + assert _aggregate_concurrency_budget(100_000) == _AGGREGATE_CONCURRENCY_BUDGET_CAP + assert ( + _AGGREGATE_CONCURRENCY_BUDGET_FLOOR <= _aggregate_concurrency_budget(None) <= _AGGREGATE_CONCURRENCY_BUDGET_CAP + ) + # Sub-MiB default path still uses Little's-law floor + assert _aggregate_concurrency_budget(256 * 1024) == max( + _AGGREGATE_CONCURRENCY_BUDGET_FLOOR, + min(_AGGREGATE_CONCURRENCY_BUDGET_CAP, max(target_bytes // (256 * 1024), latency)), + ) + + +@pytest.mark.parametrize( + "median_bytes", + [1 * 1024 * 1024, 10 * 1024 * 1024, 100 * 1024 * 1024], +) +def test_aggregate_budget_large_median_bandwidth_bounded(median_bytes): + """Medians ≥1 MiB must not be pinned by the Little's-law arm (~240).""" + from litdata.raw.dataset import ( + _AGGREGATE_CONCURRENCY_BUDGET_FLOOR, + _ASSUMED_AGGREGATE_BANDWIDTH_BPS, + _ASSUMED_REQUEST_LATENCY_S, + _ASSUMED_REQUEST_RATE, + _CONCURRENCY_PIPELINE_SECONDS, + _aggregate_concurrency_budget, + ) + + target_bytes = int(_ASSUMED_AGGREGATE_BANDWIDTH_BPS * _CONCURRENCY_PIPELINE_SECONDS) + bandwidth = max(1, target_bytes // median_bytes) + expected = max(_AGGREGATE_CONCURRENCY_BUDGET_FLOOR, min(512, bandwidth)) + got = _aggregate_concurrency_budget(median_bytes) + assert got == expected + latency = int(_ASSUMED_REQUEST_RATE * _ASSUMED_REQUEST_LATENCY_S) + assert got != latency or bandwidth >= latency # not latency-pinned when bandwidth is smaller + + +def test_effective_download_permits_cached_per_pid(tmp_path): + """Permit math runs once per process, not on every semaphore acquire.""" + (tmp_path / "a.jpg").write_bytes(b"x" * 100_000) + from unittest.mock import patch + + from litdata.raw.dataset import StreamingRawDataset + + ds = StreamingRawDataset(input_dir=str(tmp_path), max_prefetch=0) + cm = ds.cache_manager + with patch("litdata.raw.dataset._num_dataloader_workers", side_effect=[8, 16]) as mock_w: + assert cm._effective_download_permits() == 64 # adaptive: 512//8 + assert cm._effective_download_permits() == 64 # cached — ignores worker change + assert mock_w.call_count == 1 + cm.reset_runtime_state() + with patch("litdata.raw.dataset._num_dataloader_workers", return_value=16): + assert cm._effective_download_permits() == 32 # recomputed: 512//16 + + +def test_effective_download_permits_reset_on_pickle(tmp_path): + """Pickle/spawn clears the pid-guarded permit cache.""" + import pickle + + (tmp_path / "a.jpg").write_bytes(b"x" * 100_000) + from unittest.mock import patch + + from litdata.raw.dataset import StreamingRawDataset + + ds = StreamingRawDataset(input_dir=str(tmp_path), max_prefetch=0) + cm = ds.cache_manager + with patch("litdata.raw.dataset._num_dataloader_workers", return_value=8): + assert cm._effective_download_permits() == 64 + blob = pickle.dumps(cm) + restored = pickle.loads(blob) # noqa: S301 + assert restored._cached_permits is None + assert restored._cached_permits_pid is None + with patch("litdata.raw.dataset._num_dataloader_workers", return_value=16): + assert restored._effective_download_permits() == 32 + + +@pytest.mark.skipif(condition=sys.platform == "win32", reason="Not supported on windows") +def test_schedule_prefetch_uses_effective_budget(tmp_path): + """_schedule_prefetch schedules only the worker-aware effective look-ahead.""" + for i in range(200): + (tmp_path / f"file{i:03d}.jpg").write_bytes(b"x") + + dataset = StreamingRawDataset(input_dir=str(tmp_path), cache_files=False, max_prefetch=16) + + class _Info: + def __init__(self, num_workers: int): + self.num_workers = num_workers + + def _count_scheduled(num_workers: int) -> int: + call_count = {"n": 0} + + def counting_create_task(coro): + call_count["n"] += 1 + coro.close() + + class _Task: + def add_done_callback(self, cb): + return None + + return _Task() + + with ( + patch("torch.utils.data.get_worker_info", return_value=_Info(num_workers)), + patch("asyncio.create_task", side_effect=counting_create_task), + ): + # Sequential batch of 4; start = 0 + num_workers * 4 + dataset._schedule_prefetch([0, 1, 2, 3]) + return call_count["n"] + + # w=16 → effective = min(16, 64//16) = 4 + assert _count_scheduled(16) == 4 + # w=2 → effective = min(16, 64//2) = 16 + assert _count_scheduled(2) == 16 + + @pytest.mark.skipif(condition=sys.platform == "win32", reason="Not supported on windows") def test_streaming_raw_dataset_getitem(tmp_path): """Test single item access.""" test_content = b"test image content" (tmp_path / "file1.jpg").write_bytes(test_content) - dataset = StreamingRawDataset(input_dir=str(tmp_path)) + dataset = StreamingRawDataset(input_dir=str(tmp_path), max_prefetch=0) # Patch async download to return test_content - async def mock_download_file_async(file_path): + async def mock_download_file_async(file_path, size=None): return test_content with patch.object(dataset.cache_manager, "download_file_async", side_effect=mock_download_file_async): @@ -62,7 +252,7 @@ def test_streaming_raw_dataset_getitem_index_error(tmp_path): """Test index error for out of range access.""" (tmp_path / "file1.jpg").write_text("content1") - dataset = StreamingRawDataset(input_dir=str(tmp_path), cache_files=False) + dataset = StreamingRawDataset(input_dir=str(tmp_path), cache_files=False, max_prefetch=0) with pytest.raises(IndexError, match="Index 1 out of range"): dataset[1] @@ -103,7 +293,7 @@ def test_streaming_raw_dataset_getitems(tmp_path): for i, content in enumerate(test_contents): (tmp_path / f"file{i}.jpg").write_bytes(content) - dataset = StreamingRawDataset(input_dir=str(tmp_path), cache_files=False) + dataset = StreamingRawDataset(input_dir=str(tmp_path), cache_files=False, max_prefetch=0) # Mock _download_batch to return test contents async def mock_download_batch(indices): @@ -126,9 +316,9 @@ async def test_download_batch_flat(tmp_path): for file_path, content in test_contents.items(): Path(file_path).write_bytes(content) - dataset = StreamingRawDataset(input_dir=str(tmp_path)) + dataset = StreamingRawDataset(input_dir=str(tmp_path), max_prefetch=0) - async def mock_download_and_process_item(file_path): + async def mock_download_and_process_item(file_path, size=None): return test_contents[file_path] with ( @@ -161,9 +351,9 @@ class GroupedDataset(StreamingRawDataset): def setup(self, files): return [files[i : i + 2] for i in range(0, len(files), 2)] - grouped_dataset = GroupedDataset(input_dir=str(tmp_path)) + grouped_dataset = GroupedDataset(input_dir=str(tmp_path), max_prefetch=0) - async def mock_download_and_process_group(file_paths): + async def mock_download_and_process_group(file_paths, sizes=None): return [test_contents[fp] for fp in file_paths] print(grouped_dataset.items) @@ -185,7 +375,7 @@ def test_thread_safety(tmp_path): for i, content in enumerate(test_contents): (tmp_path / f"file{i}.jpg").write_bytes(content) - dataset = StreamingRawDataset(input_dir=str(tmp_path), cache_files=False) + dataset = StreamingRawDataset(input_dir=str(tmp_path), cache_files=False, max_prefetch=0) # Mock _download_batch to return test contents async def mock_download_batch(indices): @@ -209,7 +399,7 @@ def test_streaming_raw_dataset_getitems_type_error(tmp_path): """Test type error for invalid indices type.""" (tmp_path / "file1.jpg").write_text("content1") - dataset = StreamingRawDataset(input_dir=str(tmp_path), cache_files=False) + dataset = StreamingRawDataset(input_dir=str(tmp_path), cache_files=False, max_prefetch=0) with pytest.raises(TypeError): dataset.__getitems__(0) # Should be a list @@ -220,9 +410,9 @@ def test_streaming_raw_dataset_getitems_index_error(tmp_path): """Test index error for out of range batch access.""" (tmp_path / "file1.jpg").write_text("content1") - dataset = StreamingRawDataset(input_dir=str(tmp_path), cache_files=False) + dataset = StreamingRawDataset(input_dir=str(tmp_path), cache_files=False, max_prefetch=0) - with pytest.raises(IndexError, match="list index out of range"): + with pytest.raises(IndexError, match="out of range"): dataset.__getitems__([0, 1]) @@ -235,10 +425,10 @@ def test_streaming_raw_dataset_transform(tmp_path): def transform(x): return x.decode() + "_transformed" - dataset = StreamingRawDataset(input_dir=str(tmp_path), transform=transform) + dataset = StreamingRawDataset(input_dir=str(tmp_path), transform=transform, max_prefetch=0) # Patch async download to return test_content - async def mock_download_file_async(file_path): + async def mock_download_file_async(file_path, size=None): return test_content with patch.object(dataset.cache_manager, "download_file_async", side_effect=mock_download_file_async): @@ -256,7 +446,7 @@ def test_streaming_raw_dataset_with_dataloader(tmp_path): dataset = StreamingRawDataset(input_dir=str(tmp_path)) # Mock async download to return test content - async def mock_download_async(file_path): + async def mock_download_async(file_path, size=None): index = int(file_path.split("file")[1].split(".")[0]) return test_contents[index] @@ -290,7 +480,7 @@ def test_cache_manager_get_local_path_invalid(): def test_cache_manager_download_file_async_error(): cm = CacheManager(input_dir="s3://bucket/data", cache_dir=None, cache_files=False) - async def fail_download(file_path): + async def fail_download(file_path, *args, **kwargs): raise Exception("fail") cm._downloader = type("Downloader", (), {"adownload_fileobj": fail_download})() @@ -332,7 +522,7 @@ def test_streaming_raw_dataset_transform_none_and_group(tmp_path): ds = StreamingRawDataset(input_dir=str(tmp_path)) # Patch download to return bytes - async def mock_download_file_async(file_path): + async def mock_download_file_async(file_path, size=None): return b"abc" ds.cache_manager.download_file_async = mock_download_file_async diff --git a/tests/raw/test_fork_safety.py b/tests/raw/test_fork_safety.py new file mode 100644 index 000000000..dd93fec77 --- /dev/null +++ b/tests/raw/test_fork_safety.py @@ -0,0 +1,1097 @@ +"""Regression: fork/loop lifecycle, atomic cache publish, prefetch failures.""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +import os +import sys +import threading +import time +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +from torch.utils.data import DataLoader + +from litdata.raw.dataset import ( + _HEDGE_ASSUMED_BANDWIDTH_BPS, + _LOCK_SUFFIX, + CacheManager, + StreamingRawDataset, + _create_event_loop, + _effective_hedge_delay, + _get_loop_runner, + _loop_backend_name, + _run_async, + _sweep_orphan_tmp_files, +) + + +@pytest.mark.skipif(sys.platform == "win32", reason="fork semantics differ on Windows") +def test_spawn_pickle_strips_runtime_state(tmp_path: Path) -> None: + """After warm-up, pickle must succeed and omit loop/task/downloader state.""" + import pickle + + for i in range(8): + (tmp_path / f"f{i}.bin").write_bytes(f"data-{i}".encode()) + + ds = StreamingRawDataset( + str(tmp_path), + cache_dir=str(tmp_path / "cache"), + cache_files=False, + max_prefetch=4, + hedge_delay=0, + ) + _ = ds.__getitems__([0, 1, 2]) + # Accidental instance attrs must not enter the allowlisted pickle payload. + ds._bad_lock = threading.Lock() # type: ignore[attr-defined] + ds.cache_manager._bad_lock = threading.Lock() # type: ignore[attr-defined] + + blob = pickle.dumps(ds, protocol=pickle.HIGHEST_PROTOCOL) + restored = pickle.loads(blob) # noqa: S301 + assert not hasattr(restored, "_bad_lock") + assert not hasattr(restored.cache_manager, "_bad_lock") + assert restored.cache_manager._downloader is None + assert restored.cache_manager._semaphore is None + assert restored.cache_manager._range_executor is None + assert restored._inflight == {} + assert restored._inflight_loop is None + assert restored._owner_pid == os.getpid() + assert restored[0].startswith(b"data-") + + +@pytest.mark.skipif(sys.platform == "win32", reason="fork semantics differ on Windows") +@pytest.mark.parametrize("mp_context", ["spawn", "fork"]) +@pytest.mark.parametrize("cache_files", [False, True]) +def test_parent_worker_parent_lifecycle(tmp_path: Path, mp_context: str, cache_files: bool) -> None: + """Touch dataset in main, iterate with workers, touch again in main.""" + for i in range(8): + (tmp_path / f"f{i}.bin").write_bytes(f"data-{i}".encode()) + + cache = tmp_path / "cache" + ds = StreamingRawDataset( + str(tmp_path), + cache_dir=str(cache), + cache_files=cache_files, + max_prefetch=0, + max_concurrent_downloads=8, + hedge_delay=0, + ) + + first = ds[0] + assert first.startswith(b"data-") + + loader = DataLoader(ds, batch_size=2, num_workers=2, multiprocessing_context=mp_context) + batches = list(loader) + assert len(batches) == 4 + assert len(batches[0]) == 2 + + again = ds[0] + assert again == first + + +@pytest.mark.skipif(sys.platform == "win32", reason="fork semantics differ on Windows") +def test_os_fork_clears_runner_and_lock(tmp_path: Path) -> None: + """Child after os.fork gets a fresh runner; module locks/futures are reinitialized.""" + (tmp_path / "a.bin").write_bytes(b"fork-me") + # cache_files=True so write-behind futures may be in flight around fork. + ds = StreamingRawDataset( + str(tmp_path), + cache_dir=str(tmp_path / "cache"), + cache_files=True, + hedge_delay=0, + max_prefetch=0, + ) + parent_runner = _get_loop_runner() + assert parent_runner.is_alive() + _ = ds[0] + + import litdata.raw.dataset as raw_dataset + + parent_lock = raw_dataset._RUNNER_LOCK + parent_wb_lock = raw_dataset._WRITE_BEHIND_LOCK + rfd, wfd = os.pipe() + pid = os.fork() + if pid == 0: + # Child — register_at_fork already cleared runner + reinit locks. + os.close(rfd) + try: + assert raw_dataset._RUNNER is None + assert raw_dataset._RUNNER_LOCK is not parent_lock + assert raw_dataset._WRITE_BEHIND_LOCK is not parent_wb_lock + assert set() == raw_dataset._WRITE_BEHIND_FUTURES + child_runner = _get_loop_runner() + assert child_runner.pid == os.getpid() + assert child_runner is not parent_runner + val = ds[0] + os.write(wfd, b"ok:" + val) + code = 0 + except Exception as exc: + os.write(wfd, f"err:{exc!r}".encode()) + code = 1 + finally: + os.close(wfd) + os._exit(code) + + os.close(wfd) + with os.fdopen(rfd, "rb") as rf: + msg = rf.read(4096) + _, status = os.waitpid(pid, 0) + assert os.WIFEXITED(status), msg + assert os.WEXITSTATUS(status) == 0, msg + assert msg.startswith(b"ok:") + assert msg[3:] == b"fork-me" + # Parent runner still usable. + assert ds[0] == b"fork-me" + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_loop_runner_deadlock_guard() -> None: + runner = _get_loop_runner() + + async def boom() -> None: + # Calling run() on the loop thread itself must fail fast. + runner.run(asyncio.sleep(0)) + + with pytest.raises(RuntimeError, match="event-loop thread"): + runner.run(boom()) + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_create_event_loop_matches_backend_name() -> None: + loop = _create_event_loop() + try: + name = _loop_backend_name() + if name == "uvloop": + assert type(loop).__module__.startswith("uvloop") + else: + assert not type(loop).__module__.startswith("uvloop") + finally: + loop.close() + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_cache_files_dedupes_concurrent_downloads(tmp_path: Path) -> None: + """Concurrent ensure_file_async for the same key shares one download task.""" + src = tmp_path / "src" + src.mkdir() + (src / "a.bin").write_bytes(b"hello-cache") + cache = tmp_path / "cache" + + cm = CacheManager(str(src), cache_dir=str(cache), cache_files=True, max_concurrent_downloads=4) + remote = str(src / "a.bin") + calls = {"n": 0} + + async def run() -> None: + downloader = cm.downloader + orig = downloader.adownload_file + + async def counting_adownload(remote_filepath: str, local_filepath: str) -> None: + calls["n"] += 1 + await asyncio.sleep(0.05) + await orig(remote_filepath, local_filepath) + + downloader.adownload_file = counting_adownload # type: ignore[method-assign] + paths = await asyncio.gather( + cm.ensure_file_async(remote), + cm.ensure_file_async(remote), + cm.ensure_file_async(remote), + ) + assert len(set(paths)) == 1 + assert Path(paths[0]).read_bytes() == b"hello-cache" + + asyncio.run(run()) + assert calls["n"] == 1 + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_cache_publish_is_atomic_against_direct_writes(tmp_path: Path) -> None: + """Even if the downloader writes non-atomically to its target, final path is atomic.""" + src = tmp_path / "src" + src.mkdir() + payload = b"x" * 64_000 + (src / "big.bin").write_bytes(payload) + cache = tmp_path / "cache" + cm = CacheManager(str(src), cache_dir=str(cache), cache_files=True) + remote = str(src / "big.bin") + local = cm.get_local_path(remote) + seen_partial = {"value": False} + + async def slow_direct_write(remote_filepath: str, local_filepath: str) -> None: + # Intentionally non-atomic write to whatever path CacheManager asks for (the tmp path). + with open(local_filepath, "wb") as f: + f.write(payload[:1000]) + f.flush() + os.fsync(f.fileno()) + # Final publish target must still be absent while the temp is partial. + if not os.path.exists(local): + seen_partial["value"] = True + await asyncio.sleep(0.05) + f.write(payload[1000:]) + + async def run() -> None: + cm.downloader.adownload_file = slow_direct_write # type: ignore[method-assign] + path = await cm.ensure_file_async(remote) + assert path == local + assert Path(local).read_bytes() == payload + assert not list(Path(local).parent.glob("*.tmp.*")) + + asyncio.run(run()) + assert seen_partial["value"] is True + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_prefetch_failure_retries_on_demand(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: + """Failed prefetch must not poison the item; on-demand fetch succeeds without GC warnings.""" + for i in range(6): + (tmp_path / f"f{i}.bin").write_bytes(f"ok-{i}".encode()) + + ds = StreamingRawDataset( + str(tmp_path), + cache_dir=str(tmp_path / "cache"), + cache_files=False, + max_prefetch=4, + prefetch_cache_size=8, + hedge_delay=0, + ) + # Force known order. + ds.items = sorted(ds.items, key=lambda m: m.path) + + fail_once = {"n": 0} + real = ds.cache_manager.download_file_async + + async def flaky(file_path: str, size: int | None = None) -> bytes: + # Fail the first download of index-1's file during prefetch window. + if file_path.endswith("f1.bin") and fail_once["n"] == 0: + fail_once["n"] += 1 + raise RuntimeError("boom-prefetch") + return await real(file_path) + + ds.cache_manager.download_file_async = flaky # type: ignore[method-assign] + + with caplog.at_level(logging.DEBUG, logger="litdata.raw.dataset"): + # Batch [0] schedules prefetch of later indices (worker stride=1). + batch0 = ds.__getitems__([0]) + assert batch0[0] == b"ok-0" + # Allow prefetch tasks to settle (including the failed one + done callback). + time.sleep(0.1) + # On-demand fetch of the previously failed index must succeed. + batch1 = ds.__getitems__([1]) + assert batch1[0] == b"ok-1" + + # No asyncio "never retrieved" noise expected in our logger; callback retrieved it. + joined = "\n".join(r.message for r in caplog.records) + assert "exception was never retrieved" not in joined + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_getitem_from_running_event_loop(tmp_path: Path) -> None: + """``__getitem__`` works when the caller already has a running loop (notebook path).""" + (tmp_path / "a.bin").write_bytes(b"notebook") + ds = StreamingRawDataset( + str(tmp_path), + cache_dir=str(tmp_path / "cache"), + cache_files=False, + hedge_delay=0, + max_prefetch=0, + ) + + async def from_running_loop() -> bytes: + # Nested: running loop + dataset sync API. + return await asyncio.to_thread(lambda: ds[0]) + + assert asyncio.run(from_running_loop()) == b"notebook" + + async def via_run_async_nested() -> bytes: + return _run_async(ds._download_batch([0]))[0] + + assert asyncio.run(via_run_async_nested()) == b"notebook" + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_cache_manager_resets_runtime_state(tmp_path: Path) -> None: + cm = CacheManager(str(tmp_path), cache_dir=str(tmp_path / "c"), cache_files=False) + cm._downloader = object() # type: ignore[assignment] + cm._downloader_pid = 1 + cm._semaphore = object() # type: ignore[assignment] + cm.reset_runtime_state() + assert cm._downloader is None + assert cm._downloader_pid is None + assert cm._semaphore is None + assert cm._path_inflight == {} + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_atomic_publish_never_exposes_short_file(tmp_path: Path) -> None: + """Readers polling the final path must never see a partial publish.""" + src = tmp_path / "src" + src.mkdir() + payload = b"ABCDEFGH" * 8_000 # 64 KiB + (src / "big.bin").write_bytes(payload) + cache = tmp_path / "cache" + cm = CacheManager(str(src), cache_dir=str(cache), cache_files=True) + remote = str(src / "big.bin") + local = cm.get_local_path(remote) + observed_lengths: list[int] = [] + stop = threading.Event() + + def poller() -> None: + while not stop.wait(0.001): + if os.path.exists(local): + with contextlib.suppress(OSError): + observed_lengths.append(os.path.getsize(local)) + + async def chunky_write(remote_filepath: str, local_filepath: str) -> None: + with open(local_filepath, "wb") as f: + for i in range(0, len(payload), 4096): + f.write(payload[i : i + 4096]) + f.flush() + await asyncio.sleep(0.002) + + async def run() -> None: + cm.downloader.adownload_file = chunky_write # type: ignore[method-assign] + path = await cm.ensure_file_async(remote) + assert Path(path).read_bytes() == payload + + t = threading.Thread(target=poller, daemon=True) + t.start() + asyncio.run(run()) + # Allow the poller to observe the atomically published final path. + deadline = time.monotonic() + 2.0 + while not observed_lengths and time.monotonic() < deadline: + time.sleep(0.01) + stop.set() + t.join(timeout=2) + assert observed_lengths, "poller should observe the published file" + assert all(n == len(payload) for n in observed_lengths) + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_two_threads_same_file_cache_publish(tmp_path: Path) -> None: + """Two caller threads sharing the process LoopRunner can cache-publish the same file. + + Note: both threads dispatch onto the single process-local ``_LoopRunner`` loop + (not distinct event loops). Cross-process coordination is covered by the lock tests. + """ + src = tmp_path / "src" + src.mkdir() + payload = b"thread-race-" + os.urandom(32_768) + (src / "shared.bin").write_bytes(payload) + cache = tmp_path / "cache" + cm = CacheManager(str(src), cache_dir=str(cache), cache_files=True) + remote = str(src / "shared.bin") + results: list[bytes] = [] + errors: list[BaseException] = [] + + def worker() -> None: + try: + path = _run_async(cm.ensure_file_async(remote)) + results.append(Path(path).read_bytes()) + except BaseException as exc: + errors.append(exc) + + threads = [threading.Thread(target=worker) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=30) + assert not errors, errors + assert len(results) == 2 + assert all(r == payload for r in results) + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_cancelled_inflight_retries_on_demand(tmp_path: Path) -> None: + """Cancelled prefetch task is dropped; on-demand resolve still returns data.""" + for i in range(4): + (tmp_path / f"f{i}.bin").write_bytes(f"val-{i}".encode()) + + ds = StreamingRawDataset( + str(tmp_path), + cache_dir=str(tmp_path / "cache"), + cache_files=False, + max_prefetch=0, + hedge_delay=0, + ) + ds.items = sorted(ds.items, key=lambda m: m.path) + + async def run() -> bytes: + # Simulate a cancelled prefetch entry for index 1. + async def forever() -> bytes: + await asyncio.sleep(3600) + return b"never" + + task = asyncio.create_task(forever()) + ds._inflight[1] = task + ds._inflight_loop = asyncio.get_running_loop() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + # _resolve_index should treat cancelled task as miss and re-fetch. + return await ds._resolve_index(1) + + assert asyncio.run(run()) == b"val-1" + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_nested_loop_still_serves_items(tmp_path: Path) -> None: + """Caller with a running loop can still fetch via the process LoopRunner thread.""" + for i in range(4): + (tmp_path / f"f{i}.bin").write_bytes(f"n-{i}".encode()) + + ds = StreamingRawDataset( + str(tmp_path), + cache_dir=str(tmp_path / "cache"), + cache_files=False, + max_prefetch=2, + hedge_delay=0, + ) + ds.items = sorted(ds.items, key=lambda m: m.path) + + async def nested() -> list[Any]: + return await asyncio.to_thread(lambda: ds.__getitems__([0])) + + batch = asyncio.run(nested()) + assert batch[0] == b"n-0" + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_sweep_orphan_tmp_files_removes_dead_pid(tmp_path: Path) -> None: + cache = tmp_path / "cache" + cache.mkdir() + dead_pid = os.getpid() + 100000 + dead = cache / f"file.bin.tmp.{dead_pid}.1" + # Likely-dead pid; if somehow alive, skip assertion rather than flaking. + try: + os.kill(dead_pid, 0) + pytest.skip("unexpected live pid collision") + except ProcessLookupError: + pass + except PermissionError: + pytest.skip("pid exists") + dead.write_bytes(b"orphan") + live_pid_tmp = cache / f"other.bin.tmp.{os.getpid()}.99" + live_pid_tmp.write_bytes(b"keep") + stale_lock = cache / f"stale.bin{_LOCK_SUFFIX}" + stale_lock.write_text(f"{dead_pid}\n") + dead_scratch = cache / f".range-scratch.{dead_pid}.1.0.abc123" + dead_scratch.write_bytes(b"scratch-orphan") + live_scratch = cache / f".range-scratch.{os.getpid()}.1.0.def456" + live_scratch.write_bytes(b"keep-scratch") + _sweep_orphan_tmp_files(str(cache)) + assert not dead.exists() + assert live_pid_tmp.exists() + assert not stale_lock.exists() + assert not dead_scratch.exists() + assert live_scratch.exists() + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_write_through_cache_returns_bytes_before_publish(tmp_path: Path) -> None: + """cache_files + bytes path returns network bytes and eventually publishes the file.""" + src = tmp_path / "src" + src.mkdir() + payload = b"write-through-" + os.urandom(4096) + (src / "a.bin").write_bytes(payload) + cm = CacheManager(str(src), cache_dir=str(tmp_path / "cache"), cache_files=True, hedge_delay=0) + remote = str(src / "a.bin") + local = cm.get_local_path(remote) + + data = _run_async(cm.download_file_async(remote)) + assert data == payload + # Write-behind is async; wait briefly for publish. + deadline = time.monotonic() + 5.0 + while not os.path.exists(local) and time.monotonic() < deadline: + time.sleep(0.01) + assert Path(local).read_bytes() == payload + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_write_through_dedupes_concurrent_fetches(tmp_path: Path) -> None: + src = tmp_path / "src" + src.mkdir() + (src / "a.bin").write_bytes(b"dedupe-bytes") + cm = CacheManager(str(src), cache_dir=str(tmp_path / "cache"), cache_files=True, hedge_delay=0) + remote = str(src / "a.bin") + calls = {"n": 0} + + async def run() -> None: + real = cm.downloader.adownload_fileobj + + async def counting(path: str) -> bytes: + calls["n"] += 1 + await asyncio.sleep(0.05) + return await real(path) + + cm.downloader.adownload_fileobj = counting # type: ignore[method-assign] + results = await asyncio.gather( + cm.download_file_async(remote), + cm.download_file_async(remote), + cm.download_file_async(remote), + ) + assert results == [b"dedupe-bytes"] * 3 + + asyncio.run(run()) + assert calls["n"] == 1 + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_loop_runner_survives_across_calls(tmp_path: Path) -> None: + """Prefetch can progress between __getitems__ because the loop thread persists.""" + for i in range(16): + (tmp_path / f"f{i:02d}.bin").write_bytes(f"x-{i}".encode()) + ds = StreamingRawDataset( + str(tmp_path), + cache_dir=str(tmp_path / "cache"), + cache_files=False, + max_prefetch=8, + hedge_delay=0, + ) + ds.items = sorted(ds.items, key=lambda m: m.path) + runner = _get_loop_runner() + assert runner.is_alive() + _ = ds.__getitems__([0, 1]) + time.sleep(0.2) + # Same runner instance should still be alive (not recreated per call). + assert _get_loop_runner() is runner + batch = ds.__getitems__([2, 3]) + assert batch[0].startswith(b"x-") + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_cross_process_lock_claim(tmp_path: Path) -> None: + src = tmp_path / "src" + src.mkdir() + (src / "a.bin").write_bytes(b"locked") + cm = CacheManager(str(src), cache_dir=str(tmp_path / "cache"), cache_files=True, hedge_delay=0) + local = cm.get_local_path(str(src / "a.bin")) + assert cm._try_claim_lock(local) is True + assert Path(cm._lock_path(local)).name.endswith(_LOCK_SUFFIX.lstrip(".")) or cm._lock_path(local).endswith( + _LOCK_SUFFIX + ) + assert cm._try_claim_lock(local) is False + cm._release_lock(local) + assert cm._try_claim_lock(local) is True + cm._release_lock(local) + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_lock_peer_fails_then_claim(tmp_path: Path) -> None: + """If a peer held the lock then failed/released, the waiter can claim and finish.""" + src = tmp_path / "src" + src.mkdir() + payload = b"peer-fail-then-ok" + (src / "a.bin").write_bytes(payload) + cm = CacheManager(str(src), cache_dir=str(tmp_path / "cache"), cache_files=True, hedge_delay=0) + remote = str(src / "a.bin") + local = cm.get_local_path(remote) + # Simulate a peer that claimed the lock then died without publishing. + assert cm._try_claim_lock(local) is True + lock_path = cm._lock_path(local) + dead_pid = os.getpid() + 100000 + try: + os.kill(dead_pid, 0) + pytest.skip("unexpected live pid collision") + except ProcessLookupError: + pass + except PermissionError: + pytest.skip("pid exists") + Path(lock_path).write_text(f"{dead_pid}\n") + assert cm._lock_owner_alive(lock_path) is False + path = _run_async(cm.ensure_file_async(remote)) + assert Path(path).read_bytes() == payload + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_lock_dead_pid_takeover(tmp_path: Path) -> None: + """``_try_claim_lock`` takes over when the recorded owner pid is dead.""" + src = tmp_path / "src" + src.mkdir() + (src / "a.bin").write_bytes(b"takeover") + cm = CacheManager(str(src), cache_dir=str(tmp_path / "cache"), cache_files=True, hedge_delay=0) + local = cm.get_local_path(str(src / "a.bin")) + lock_path = cm._lock_path(local) + dead_pid = os.getpid() + 100000 + try: + os.kill(dead_pid, 0) + pytest.skip("unexpected live pid collision") + except ProcessLookupError: + pass + except PermissionError: + pytest.skip("pid exists") + Path(lock_path).write_text(f"{dead_pid}\n") + assert cm._lock_owner_alive(lock_path) is False + assert cm._try_claim_lock(local) is True + assert Path(lock_path).read_text().strip() == str(os.getpid()) + cm._release_lock(local) + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_hedge_slow_first_wins(tmp_path: Path) -> None: + """Slow first request is beaten by a hedged second request.""" + src = tmp_path / "src" + src.mkdir() + (src / "a.bin").write_bytes(b"hedge-ok") + cm = CacheManager( + str(src), + cache_dir=str(tmp_path / "cache"), + cache_files=False, + hedge_delay=0.05, + max_concurrent_downloads=8, + ) + remote = "s3://bucket/data/a.bin" # remote so hedging is eligible + cm._input_dir_path = "s3://bucket/data" + calls = {"n": 0} + + async def run() -> bytes: + cm._downloader_pid = os.getpid() + cm._downloader_loop = asyncio.get_running_loop() + + async def flaky(path: str) -> bytes: + n = calls["n"] + calls["n"] = n + 1 + if n == 0: + await asyncio.sleep(0.5) + return b"slow" + return b"hedge-ok" + + cm._downloader = SimpleNamespace(adownload_fileobj=flaky) # type: ignore[assignment] + return await cm._fetch_bytes(remote, size=8) + + t0 = time.monotonic() + assert asyncio.run(run()) == b"hedge-ok" + # Hedge delay is 0.05s; allow scheduling jitter on loaded CI. + assert time.monotonic() - t0 < 2.0 + assert calls["n"] >= 2 + assert cm._hedge_fired >= 1 + + +def test_hedge_delay_default_is_zero() -> None: + """Hedging is opt-in (default 0), matching range_parallel_threshold.""" + import inspect + + ds_default = inspect.signature(StreamingRawDataset.__init__).parameters["hedge_delay"].default + cm_default = inspect.signature(CacheManager.__init__).parameters["hedge_delay"].default + assert float(ds_default) == 0.0 + assert float(cm_default) == 0.0 + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +@pytest.mark.parametrize("download_timeout", [0.0, 120.0]) +def test_fetch_bytes_fast_path_when_hedge_off(tmp_path: Path, download_timeout: float) -> None: + """Hedging off takes the bare path even when download_timeout defaults to 120.""" + src = tmp_path / "src" + src.mkdir() + (src / "a.bin").write_bytes(b"fast") + cm = CacheManager( + str(src), + cache_dir=str(tmp_path / "cache"), + cache_files=False, + hedge_delay=0, + download_timeout=download_timeout, + ) + remote = "s3://bucket/data/a.bin" + cm._input_dir_path = "s3://bucket/data" + calls = {"n": 0} + wait_for_calls = {"n": 0} + real_wait_for = asyncio.wait_for + + async def counting_wait_for(awaitable, timeout=None, **kwargs): # type: ignore[no-untyped-def] + wait_for_calls["n"] += 1 + return await real_wait_for(awaitable, timeout=timeout, **kwargs) + + async def run() -> bytes: + cm._downloader_pid = os.getpid() + cm._downloader_loop = asyncio.get_running_loop() + + async def once(path: str) -> bytes: + calls["n"] += 1 + return b"fast" + + cm._downloader = SimpleNamespace(adownload_fileobj=once) # type: ignore[assignment] + return await cm._fetch_bytes(remote, size=4) + + with pytest.MonkeyPatch.context() as mp: + mp.setattr(asyncio, "wait_for", counting_wait_for) + assert asyncio.run(run()) == b"fast" + assert calls["n"] == 1 + assert cm._hedge_fired == 0 + assert wait_for_calls["n"] == 0 + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_download_batch_applies_timeout_once(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Default download_timeout wraps the batch gather once, not each item.""" + for i in range(4): + (tmp_path / f"f{i}.bin").write_bytes(f"data-{i}".encode()) + ds = StreamingRawDataset( + str(tmp_path), + cache_dir=str(tmp_path / "cache"), + cache_files=False, + max_prefetch=0, + hedge_delay=0, + download_timeout=120.0, + ) + assert ds._batch_download_budget([0, 1, 2]) == 120.0 + wait_for_calls = {"n": 0} + real_wait_for = asyncio.wait_for + + async def counting_wait_for(awaitable, timeout=None, **kwargs): # type: ignore[no-untyped-def] + wait_for_calls["n"] += 1 + return await real_wait_for(awaitable, timeout=timeout, **kwargs) + + monkeypatch.setattr(asyncio, "wait_for", counting_wait_for) + items = ds.__getitems__([0, 1, 2]) + assert len(items) == 3 + assert wait_for_calls["n"] == 1 + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_batch_timeout_cancels_hung_inflight_for_recovery(tmp_path: Path) -> None: + """Batch timeout must cancel poisoned ``_inflight`` so a retry can succeed promptly.""" + for i in range(4): + (tmp_path / f"f{i}.bin").write_bytes(f"val-{i}".encode()) + + ds = StreamingRawDataset( + str(tmp_path), + cache_dir=str(tmp_path / "cache"), + cache_files=False, + max_prefetch=0, + hedge_delay=0, + download_timeout=0.25, + ) + ds.items = sorted(ds.items, key=lambda m: m.path) + + hang = {"enabled": True} + real = ds.cache_manager.download_file_async + + async def stub(file_path: str, size: int | None = None) -> bytes: + if hang["enabled"] and file_path.endswith("f1.bin"): + await asyncio.sleep(3600) + return b"never" + return await real(file_path, size=size) + + ds.cache_manager.download_file_async = stub # type: ignore[method-assign] + + async def run() -> None: + # Seed a hung prefetch entry whose download sleeps forever. + task = asyncio.create_task(ds._prefetch_index(1)) + ds._inflight[1] = task + ds._inflight_loop = asyncio.get_running_loop() + await asyncio.sleep(0.05) + + with pytest.raises(TimeoutError, match="Batch download timed out"): + await ds._download_batch([1]) + + # Without cancelling _inflight, retry would await the same hung task and + # pay the full budget again. Healthy stub + prompt success is the recovery. + hang["enabled"] = False + t0 = time.perf_counter() + result = await ds._download_batch([1]) + elapsed = time.perf_counter() - t0 + assert result[0] == b"val-1" + assert elapsed < 1.0 + assert 1 not in ds._inflight + + asyncio.run(run()) + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_download_budget_scales_with_size(tmp_path: Path) -> None: + """Sized objects use download_timeout as a floor, not a hard cap.""" + src = tmp_path / "src" + src.mkdir() + cm = CacheManager( + str(src), + cache_dir=str(tmp_path / "cache"), + cache_files=False, + download_timeout=10.0, + ) + large = 2 * 1024 * 1024 * 1024 # 2 GiB + size_floor = large / _HEDGE_ASSUMED_BANDWIDTH_BPS * 3.0 + budget = cm._download_budget(large) + assert budget is not None + assert budget >= size_floor + assert budget >= 10.0 + assert cm._download_budget(None) == 10.0 + assert cm._download_budget(1024) == 10.0 # tiny → floor is download_timeout + assert cm._download_budget(large, timeout=1.5) == 1.5 # explicit override + + +def test_path_is_cached_discards_stale_present_mark(tmp_path: Path) -> None: + src = tmp_path / "src" + src.mkdir() + cm = CacheManager(str(src), cache_dir=str(tmp_path / "cache"), cache_files=True) + local = str(tmp_path / "cache" / "gone.bin") + Path(local).parent.mkdir(parents=True, exist_ok=True) + Path(local).write_bytes(b"x") + assert cm._path_is_cached(local) + assert local in cm._present_paths + Path(local).unlink() + assert not cm._path_is_cached(local) + assert local not in cm._present_paths + + +def test_hedge_skipped_for_large_file(tmp_path: Path) -> None: + """Files >= 8 MB must not issue a duplicate whole-object GET.""" + assert _effective_hedge_delay(1.0, 8 * 1024 * 1024) is None + assert _effective_hedge_delay(1.0, 16 * 1024 * 1024) is None + assert _effective_hedge_delay(1.0, None) is None + assert _effective_hedge_delay(1.0, 0) is None + src = tmp_path / "src" + src.mkdir() + cm = CacheManager( + str(src), + cache_dir=str(tmp_path / "cache"), + cache_files=False, + hedge_delay=0.01, + range_parallel_threshold=0, # force whole-object path + ) + remote = "s3://bucket/data/big.bin" + calls = {"n": 0} + payload = b"x" * 100 + + async def run(size: int | None) -> bytes: + cm._downloader_pid = os.getpid() + cm._downloader_loop = asyncio.get_running_loop() + + async def once(path: str) -> bytes: + calls["n"] += 1 + await asyncio.sleep(0.05) + return payload + + cm._downloader = SimpleNamespace(adownload_fileobj=once) # type: ignore[assignment] + return await cm._fetch_bytes(remote, size=size) + + assert asyncio.run(run(8 * 1024 * 1024)) == payload + assert calls["n"] == 1 + calls["n"] = 0 + assert asyncio.run(run(None)) == payload + assert calls["n"] == 1 + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_ranged_download_validates_short_part(tmp_path: Path) -> None: + src = tmp_path / "src" + src.mkdir() + cm = CacheManager( + str(src), + cache_dir=str(tmp_path / "cache"), + cache_files=False, + hedge_delay=0, + range_chunk_size=4, + range_parallel_threshold=1, + ) + remote = "s3://bucket/data/obj.bin" + size = 12 + + async def run() -> None: + cm._downloader_pid = os.getpid() + cm._downloader_loop = asyncio.get_running_loop() + + def short_bytes(path: str, offset: int, length: int, scratch: str) -> bytes: + return b"x" * (length - 1) # always short + + cm._downloader = type("D", (), {"download_bytes": staticmethod(short_bytes)})() # type: ignore[assignment] + with pytest.raises(RuntimeError, match="short read"): + await cm._ranged_download_bytes(remote, size) + + asyncio.run(run()) + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_ranged_download_reassembles(tmp_path: Path) -> None: + src = tmp_path / "src" + src.mkdir() + payload = b"abcdefghijklmnopqrstuvwx" # 24 bytes + cm = CacheManager( + str(src), + cache_dir=str(tmp_path / "cache"), + cache_files=False, + hedge_delay=0, + range_chunk_size=8, + range_parallel_threshold=1, + ) + remote = "s3://bucket/data/obj.bin" + scratches: list[str] = [] + call_order: list[int] = [] + + async def run() -> bytes: + cm._downloader_pid = os.getpid() + cm._downloader_loop = asyncio.get_running_loop() + + def ranged(path: str, offset: int, length: int, scratch: str) -> bytes: + scratches.append(scratch) + call_order.append(offset) + return payload[offset : offset + length] + + cm._downloader = type("D", (), {"download_bytes": staticmethod(ranged)})() # type: ignore[assignment] + return await cm._ranged_download_bytes(remote, len(payload)) + + assert asyncio.run(run()) == payload + assert len(scratches) == 3 + assert len(set(scratches)) == 3 # per-chunk scratch paths + # Reassembly must be by ascending offset regardless of completion order. + assert sorted(call_order) == [0, 8, 16] + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_ranged_download_size_mismatch(tmp_path: Path) -> None: + """Joined payload length must match the declared object size.""" + src = tmp_path / "src" + src.mkdir() + cm = CacheManager( + str(src), + cache_dir=str(tmp_path / "cache"), + cache_files=False, + hedge_delay=0, + range_chunk_size=8, + range_parallel_threshold=1, + ) + remote = "s3://bucket/data/obj.bin" + + async def run() -> None: + cm._downloader_pid = os.getpid() + cm._downloader_loop = asyncio.get_running_loop() + + def ok_chunk(path: str, offset: int, length: int, scratch: str) -> bytes: + return b"x" * length + + cm._downloader = type("D", (), {"download_bytes": staticmethod(ok_chunk)})() # type: ignore[assignment] + real_gather = asyncio.gather + + async def gather_truncate(*aws: Any, **kwargs: Any) -> Any: + parts = await real_gather(*aws, **kwargs) + parts = list(parts) + off, data = parts[-1] + parts[-1] = (off, data[:-1]) + return parts + + asyncio.gather = gather_truncate # type: ignore[method-assign, assignment] + try: + with pytest.raises(RuntimeError, match="size mismatch"): + await cm._ranged_download_bytes(remote, 24) + finally: + asyncio.gather = real_gather # type: ignore[method-assign, assignment] + + asyncio.run(run()) + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_ranged_chunk_hedge_uses_distinct_scratch(tmp_path: Path) -> None: + """Concurrent first/hedge attempts for the same chunk must not share scratch paths.""" + src = tmp_path / "src" + src.mkdir() + cm = CacheManager( + str(src), + cache_dir=str(tmp_path / "cache"), + cache_files=False, + hedge_delay=0.05, + range_chunk_size=8, + range_parallel_threshold=1, + max_concurrent_downloads=8, + ) + remote = "s3://bucket/data/obj.bin" + scratches: list[str] = [] + barrier = threading.Barrier(2, timeout=5) + + async def run() -> bytes: + cm._downloader_pid = os.getpid() + cm._downloader_loop = asyncio.get_running_loop() + + def ranged(path: str, offset: int, length: int, scratch: str) -> bytes: + scratches.append(scratch) + # Block both concurrent attempts (first + hedge) until both have started. + with contextlib.suppress(threading.BrokenBarrierError): + barrier.wait() + return b"y" * length + + cm._downloader = type("D", (), {"download_bytes": staticmethod(ranged)})() # type: ignore[assignment] + return await cm._ranged_download_bytes(remote, 8) + + assert asyncio.run(run()) == b"yyyyyyyy" + assert len(scratches) >= 2 + assert len(set(scratches)) == len(scratches) + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_largest_first_preserves_result_order(tmp_path: Path) -> None: + """Downloads may start largest-first, but batch results keep caller index order.""" + sizes = [10, 1000, 50] + for i, n in enumerate(sizes): + (tmp_path / f"f{i}.bin").write_bytes(bytes([i]) * n) + + ds = StreamingRawDataset( + str(tmp_path), + cache_dir=str(tmp_path / "cache"), + cache_files=False, + hedge_delay=0, + max_prefetch=0, + ) + # Stable item order f0, f1, f2 by path. + ds.items = sorted(ds.items, key=lambda m: m.path) + order: list[int] = [] + + real = ds.cache_manager.download_file_async + + async def tracking(file_path: str, size: int | None = None) -> bytes: + idx = int(Path(file_path).stem[1:]) + order.append(idx) + await asyncio.sleep(0.01 * (1 if idx != 1 else 0)) # let scheduling show LPT start + return await real(file_path, size=size) + + ds.cache_manager.download_file_async = tracking # type: ignore[method-assign] + batch = ds.__getitems__([0, 1, 2]) + assert [b[0] for b in batch] == [0, 1, 2] + # Largest (index 1, 1000 bytes) should be started first among the three. + assert order[0] == 1 + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_duplicate_batch_indices_fanout(tmp_path: Path) -> None: + (tmp_path / "a.bin").write_bytes(b"same") + (tmp_path / "b.bin").write_bytes(b"other") + ds = StreamingRawDataset( + str(tmp_path), + cache_dir=str(tmp_path / "cache"), + hedge_delay=0, + max_prefetch=0, + ) + ds.items = sorted(ds.items, key=lambda m: m.path) + calls = {"n": 0} + real = ds.cache_manager.download_file_async + + async def counting(file_path: str, size: int | None = None) -> bytes: + calls["n"] += 1 + return await real(file_path, size=size) + + ds.cache_manager.download_file_async = counting # type: ignore[method-assign] + batch = ds.__getitems__([0, 0, 1, 0]) + assert batch[0] == batch[1] == batch[3] + assert batch[2] == b"other" + # Index 0 materialized once despite three positions. + assert calls["n"] == 2 + + +@pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows") +def test_download_to_cache_logs_fallback(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: + src = tmp_path / "src" + src.mkdir() + (src / "a.bin").write_bytes(b"fallback-ok") + cm = CacheManager(str(src), cache_dir=str(tmp_path / "cache"), cache_files=True, hedge_delay=0) + remote = str(src / "a.bin") + + async def run() -> str: + async def boom(remote_filepath: str, local_filepath: str) -> None: + raise OSError("stream failed") + + cm.downloader.adownload_file = boom # type: ignore[method-assign] + with caplog.at_level(logging.WARNING, logger="litdata.raw.dataset"): + return await cm.ensure_file_async(remote) + + path = asyncio.run(run()) + assert Path(path).read_bytes() == b"fallback-ok" + assert any("falling back to bytes path" in r.message for r in caplog.records) diff --git a/tests/raw/test_indexer.py b/tests/raw/test_indexer.py index 4ff71bfcd..4e67f215b 100644 --- a/tests/raw/test_indexer.py +++ b/tests/raw/test_indexer.py @@ -5,7 +5,13 @@ from litdata import StreamingRawDataset from litdata.constants import _PYTHON_GREATER_EQUAL_3_14 -from litdata.raw.indexer import _INDEX_FILENAME, FileIndexer, FileMetadata +from litdata.raw.indexer import ( + _INDEX_FILENAME, + FileIndexer, + FileMetadata, + _is_windows_drive_scheme, + _validate_input_dir_scheme, +) def test_file_metadata(): @@ -200,6 +206,27 @@ def test_discover_files_unsupported_scheme(): indexer.discover_files("http://unsupported/path", {}) +def test_windows_drive_scheme_treated_as_local(): + r"""urlparse('C:\\\\Users\\\\...') yields scheme='c'; treat as local, not remote.""" + assert _is_windows_drive_scheme("c") + assert _is_windows_drive_scheme("C") + assert not _is_windows_drive_scheme("s3") + assert not _is_windows_drive_scheme("ftp") + assert not _is_windows_drive_scheme("") + + # Same classification urlparse uses for Windows absolute paths (also on Linux). + _validate_input_dir_scheme(r"C:\Users\test\dataset") + _validate_input_dir_scheme("C:/Users/test/dataset") + + with pytest.raises(ValueError, match="Unsupported input directory scheme: `ftp`"): + _validate_input_dir_scheme("ftp://unsupported/path") + + indexer = FileIndexer() + with patch.object(indexer, "_discover_local_files", return_value=[]) as mock_local: + indexer.discover_files(r"C:\Users\test\dataset", {}) + mock_local.assert_called_once_with(r"C:\Users\test\dataset") + + @patch("litdata.raw.indexer.BaseIndexer._upload_to_cloud") @patch("litdata.raw.indexer.BaseIndexer._download_from_cloud", side_effect=FileNotFoundError) @pytest.mark.skipif(condition=sys.platform == "win32", reason="Not supported on windows") @@ -317,6 +344,73 @@ def test_recompute_index_excludes_index_file(tmp_path): assert _INDEX_FILENAME not in f.path +@pytest.mark.skipif(condition=sys.platform == "win32", reason="Not supported on windows") +def test_save_index_file_atomic_publish(tmp_path): + """_save_index_file writes via tmp + os.replace (no direct open of final path).""" + import os + from unittest.mock import patch + + cache_dir = tmp_path / "cache" + cache_dir.mkdir() + index_path = cache_dir / _INDEX_FILENAME + indexer = FileIndexer() + files = [FileMetadata(str(tmp_path / "a.bin"), 10)] + replace_calls: list[tuple[str, str]] = [] + real_replace = os.replace + + def tracking_replace(src, dst): + replace_calls.append((str(src), str(dst))) + return real_replace(src, dst) + + with patch("os.replace", side_effect=tracking_replace): + indexer._save_index_file(str(index_path), files, str(tmp_path)) + + assert index_path.exists() + assert len(replace_calls) == 1 + src, dst = replace_calls[0] + assert src.endswith(f".tmp.{os.getpid()}") + assert dst == str(index_path) + assert indexer._load_index_file(str(index_path)) == files + # No leftover tmp beside the final index. + assert list(cache_dir.glob(f"{_INDEX_FILENAME}.tmp.*")) == [] + + +@pytest.mark.skipif(condition=sys.platform == "win32", reason="Not supported on windows") +def test_download_from_cloud_atomic_publish(tmp_path): + """_download_from_cloud stages to tmp then os.replace into the final path.""" + import os + from unittest.mock import MagicMock, patch + + local_path = tmp_path / "cache" / _INDEX_FILENAME + local_path.parent.mkdir() + replace_calls: list[tuple[str, str]] = [] + real_replace = os.replace + + def tracking_replace(src, dst): + replace_calls.append((str(src), str(dst))) + return real_replace(src, dst) + + def fake_get(remote, dest): + with open(dest, "wb") as f: + f.write(b"index-bytes") + + mock_fs = MagicMock() + mock_fs.get.side_effect = fake_get + indexer = FileIndexer() + with ( + patch("fsspec.filesystem", return_value=mock_fs), + patch("os.replace", side_effect=tracking_replace), + ): + indexer._download_from_cloud("s3://bucket/index.json.zstd", str(local_path), {}) + + assert local_path.read_bytes() == b"index-bytes" + assert len(replace_calls) == 1 + src, dst = replace_calls[0] + assert src.endswith(f".tmp.{os.getpid()}") + assert dst == str(local_path) + assert list(local_path.parent.glob(f"{_INDEX_FILENAME}.tmp.*")) == [] + + def test_load_index_file_handles_corrupted_zstd(tmp_path): """Test that _load_index_file catches ZstdError for corrupted data.""" if _PYTHON_GREATER_EQUAL_3_14: diff --git a/tests/streaming/test_client.py b/tests/streaming/test_client.py index e98f8974e..a2474a518 100644 --- a/tests/streaming/test_client.py +++ b/tests/streaming/test_client.py @@ -1,4 +1,5 @@ import sys +import threading from time import sleep, time from unittest import mock @@ -32,7 +33,8 @@ def test_s3_client_with_storage_options(monkeypatch): config=botocore.config.Config(retries={"max_attempts": 100}), ) - # Create S3Client without storage options + # Create S3Client without storage options (force non-Studio path so IMDS is not used). + monkeypatch.setattr(client, "_IS_IN_STUDIO", False) s3_client = client.S3Client() assert s3_client.client @@ -433,6 +435,53 @@ def test_r2_client_property_refreshes_expired_credentials(monkeypatch): assert second_call_count == first_call_count + 1 +def test_s3_client_refresh_is_serialized_under_threads(monkeypatch): + """Concurrent .client access at a refresh boundary must not race-create clients.""" + in_create = {"n": 0, "max": 0} + counter_lock = threading.Lock() + barrier = threading.Barrier(8) + + boto3_session = mock.MagicMock() + boto3 = mock.MagicMock(Session=boto3_session) + monkeypatch.setattr(client, "boto3", boto3) + monkeypatch.setattr(client, "botocore", mock.MagicMock()) + + s3 = client.S3Client(refetch_interval=0, storage_options={"region_name": "us-east-1"}) + original_create = s3._create_client + + def slow_create() -> None: + with counter_lock: + in_create["n"] += 1 + in_create["max"] = max(in_create["max"], in_create["n"]) + try: + sleep(0.01) + original_create() + finally: + with counter_lock: + in_create["n"] -= 1 + + s3._create_client = slow_create # type: ignore[method-assign] + + errors: list[BaseException] = [] + + def worker() -> None: + try: + barrier.wait(timeout=5) + for _ in range(3): + assert s3.client is not None + except BaseException as exc: + errors.append(exc) + + threads = [threading.Thread(target=worker) for _ in range(8)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=10) + + assert not errors + assert in_create["max"] == 1 + + def test_r2_client_with_session_options(monkeypatch): """Test R2Client with custom session options.""" boto3_session = mock.MagicMock() diff --git a/tests/streaming/test_dataset.py b/tests/streaming/test_dataset.py index 74cca7796..32dcfc687 100644 --- a/tests/streaming/test_dataset.py +++ b/tests/streaming/test_dataset.py @@ -1072,7 +1072,7 @@ def _get_simulated_s3_dataloader(cache_dir, data_dir, shuffle=False): return StreamingDataLoader(dataset, batch_size=2, num_workers=2) -@pytest.mark.skipif(sys.platform == "win32", reason="Not tested on windows and MacOs") +@pytest.mark.skipif(sys.platform in ("win32", "darwin"), reason="Not tested on windows and MacOs") @mock.patch.dict(os.environ, {}, clear=True) @pytest.mark.timeout(90) @pytest.mark.parametrize("shuffle", [True, False]) diff --git a/tests/streaming/test_downloader.py b/tests/streaming/test_downloader.py index af08ae300..ce62fc519 100644 --- a/tests/streaming/test_downloader.py +++ b/tests/streaming/test_downloader.py @@ -148,6 +148,53 @@ def test_r2_downloader_download_bytes_reuses_client(r2_client_mock, tmpdir): ] +@mock.patch("litdata.streaming.downloader.S3Client") +def test_s3_downloader_download_bytes_reuses_client(s3_client_mock, tmpdir): + s3_client_instance = MagicMock() + s3_client_mock.return_value = s3_client_instance + + body = MagicMock() + body.read.return_value = b"hello" + s3_client_instance.client.get_object.return_value = {"Body": body} + + downloader = S3Downloader("s3://random_bucket", str(tmpdir), []) + # __init__ already creates _client; download_bytes must not recreate it. + assert hasattr(downloader, "_client") + client_id = id(downloader._client) + + assert downloader.download_bytes("s3://random_bucket/a.txt", 0, 5, os.path.join(tmpdir, "a.txt")) == b"hello" + assert downloader.download_bytes("s3://random_bucket/a.txt", 5, 5, os.path.join(tmpdir, "a.txt")) == b"hello" + + assert id(downloader._client) == client_id + s3_client_mock.assert_called_once_with(storage_options={}, session_options={}) + assert s3_client_instance.client.get_object.call_args_list == [ + mock.call(Bucket="random_bucket", Key="a.txt", Range="bytes=0-4"), + mock.call(Bucket="random_bucket", Key="a.txt", Range="bytes=5-9"), + ] + + +@mock.patch("litdata.streaming.downloader._GOOGLE_STORAGE_AVAILABLE", True) +def test_gcp_downloader_download_bytes_reuses_client(tmpdir, google_mock): + mock_client = MagicMock() + mock_bucket = MagicMock() + mock_blob = MagicMock() + mock_blob.download_as_bytes.return_value = b"hello" + + google_mock.cloud.storage.Client = MagicMock(return_value=mock_client) + mock_client.bucket = MagicMock(return_value=mock_bucket) + mock_bucket.blob = MagicMock(return_value=mock_blob) + + downloader = GCPDownloader("gs://random_bucket", str(tmpdir), [], {"project": "p"}) + assert downloader.download_bytes("gs://random_bucket/a.txt", 0, 5, os.path.join(tmpdir, "a.txt")) == b"hello" + assert downloader.download_bytes("gs://random_bucket/a.txt", 5, 5, os.path.join(tmpdir, "a.txt")) == b"hello" + + google_mock.cloud.storage.Client.assert_called_once_with(project="p") + assert mock_blob.download_as_bytes.call_args_list == [ + mock.call(start=0, end=4), + mock.call(start=5, end=9), + ] + + @mock.patch("litdata.streaming.downloader._GOOGLE_STORAGE_AVAILABLE", True) def test_gcp_downloader(tmpdir, monkeypatch, google_mock): # Create mock objects diff --git a/tests/streaming/test_resolver.py b/tests/streaming/test_resolver.py index b89eac5ac..00dc5ef44 100644 --- a/tests/streaming/test_resolver.py +++ b/tests/streaming/test_resolver.py @@ -508,6 +508,14 @@ def test_resolve_time_template(): assert resolver._resolve_time_template(path_3) == f"/logs/log_{curr_year}-{curr_month:02d}/important" +def test_has_time_template(): + assert resolver._has_time_template("/logs/log_{%Y-%m-%d}") is True + assert resolver._has_time_template("s3://bucket/run_{%Y-%m-%d_%H-%M-%S}") is True + assert resolver._has_time_template("/logs/my_logfile") is False + assert resolver._has_time_template(None) is False + assert resolver._has_time_template(resolver.Dir(path="/logs/log_2025-05-05")) is False + + @pytest.mark.skipif(sys.platform == "win32", reason="windows isn't supported") def test_src_resolver_gcs_connections(monkeypatch, lightning_cloud_mock): """Test GCS connections resolver."""