litdata 0.2.70__tar.gz → 0.2.72__tar.gz

{litdata-0.2.70/src/litdata.egg-info → litdata-0.2.72}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: litdata
-Version: 0.2.70
+Version: 0.2.72
 Summary: The Deep Learning framework to train, deploy, and ship AI products Lightning fast.
 Home-page: https://github.com/Lightning-AI/litdata
 Download-URL: https://github.com/Lightning-AI/litdata
@@ -34,6 +34,7 @@ Requires-Dist: filelock
 Requires-Dist: numpy
 Requires-Dist: boto3
 Requires-Dist: requests
+Requires-Dist: urllib3>=1.26
 Requires-Dist: tifffile
 Requires-Dist: obstore
 Provides-Extra: extras
@@ -1128,7 +1129,7 @@ Rough ImageNet order-of-magnitude on a Studio (not hard guarantees; right tuning
 | `drop_last` | `True` if distributed else `False` | Equal length across ranks |
 | `seed` | `42` | Shuffle / subsample RNG |
 | `serializers` | built-ins | Custom serialize/deserialize map |
-| `max_cache_size` | `"100GB"` | Evict consumed chunks beyond this size |
+| `max_cache_size` | `None` | Evict consumed chunks beyond this size. Default: 75% of free disk, leaving ≥50GB when possible. Pin with `"100G"` / `"50GB"`, a fraction (`0.90`), or `MAX_CACHE_SIZE`. |
 | `max_pre_download` | `2` | Chunks each worker may prefetch (raise for throughput; watch disk / RAM) |
 | `subsample` | `1.0` | Fraction of data (`0.01`) or upsample (`2.5`) |
 | `encryption` | `None` | `FernetEncryption` / `RSAEncryption` / custom |
@@ -1149,7 +1150,8 @@ On **Vast / NFS / local disk**, POSIX-fast is on by default (`LITDATA_POSIX_FAST
 | All usual `torch.utils.data.DataLoader` kwargs | `batch_size`, `num_workers`, `collate_fn`, `pin_memory`, … |
 | `shuffle` / `drop_last` | Forwarded to the streaming dataset |
 | `profile_batches` | `int` / `True` / `False` — viztracer worker trace (see [Profile data loading](#profile-loading)) |
-| `profile_skip_batches` / `profile_dir` | Warm-up skip count; output dir for `result.json` |
+| `profile_cprofile` | `True` — stdlib cProfile of the main process + worker 0 (see [Profile data loading](#profile-loading)) |
+| `profile_skip_batches` / `profile_dir` | Warm-up skip count; output dir for viztracer / cProfile files |
 | `multiprocessing_context` | Use **`"spawn"`** (or `"forkserver"`) with `ParquetLoader` + `num_workers>0` on Linux |
 
 Prefer `StreamingDataLoader` over a plain PyTorch `DataLoader` for optimized / combined / parallel datasets (resume + correct batch metadata).
@@ -2115,7 +2117,32 @@ Only **worker 0** is instrumented. When an `int` is used, the tracer wraps `fetc
 
 - Delete or change `profile_dir` between runs — LitData removes an existing `result.json` before starting.
 - Pair with a wiped chunk cache if you care about **cold** epoch behavior (`litdata cache clear`).
-- For deeper LitData internals (download / read / delete timeline), use `enable_tracer()` + [Litracer](https://github.com/Lightning-AI/litracer) instead — see [Debug & Profile LitData](#debug-profile). That path is complementary: viztracer = DataLoader worker CPU timeline; Litracer = LitData pipeline events.
+### cProfile (main + worker 0)
+
+Stdlib statistical profile of **the parent process** (queue wait, unpickle, collate handoff) and **worker 0** (fetch, decode, transforms). No extra package.
+
+```python
+loader = StreamingDataLoader(
+    dataset,
+    batch_size=64,
+    num_workers=4,
+    profile_cprofile=True,
+    profile_dir="./profiles",
+)
+
+for batch in loader:
+    train_step(batch)
+# writes profiles/cprofile_main.prof + cprofile_worker0.prof (and .txt summaries)
+```
+
+```bash
+python -m pstats profiles/cprofile_worker0.prof
+# then: sort tottime   stats 30
+```
+
+Do not set `profile_cprofile` and `profile_batches` together — both install a `sys.setprofile` hook. With `num_workers=0` only the main file is written. The parent profiler starts **after** workers spawn so fork does not inherit an active cProfile.
+
+- For deeper LitData internals (download / read / delete timeline), use `enable_tracer()` + [Litracer](https://github.com/Lightning-AI/litracer) instead — see [Debug & Profile LitData](#debug-profile). That path is complementary: viztracer = DataLoader worker CPU timeline; cProfile = function self/cum time; Litracer = LitData pipeline events.
 
 </details>
 
@@ -2165,7 +2192,9 @@ outputs = optimize(
 
 Control how much disk the local chunk cache may use. Downloaded chunks are deleted after use once the cache exceeds the limit.
 
-Default `max_cache_size` is **`100GB`**. Peak disk in flight is roughly:
+Default `max_cache_size` is **`None`**: LitData uses **75% of currently free disk** and still leaves **≥50GB** free when the volume is large enough for checkpoints. Pass `"100G"` / `"50GB"` for a fixed budget, or a float (`0.90`) for that fraction of currently free space. `MAX_CACHE_SIZE` overrides the constructor (size or fraction).
+
+Peak disk in flight is roughly:
 
 ```
 num_workers × max_pre_download × mean_chunk_size
@@ -2205,7 +2234,9 @@ for batch in StreamingDataLoader(dataset, batch_size=64, num_workers=8):
 | `LITDATA_ASYNC_CHUNK_PREFETCH=1` | Force on |
 | `LITDATA_ASYNC_CHUNK_PREFETCH=0` | Force off |
 
-When async is on, LitData raises `max_pre_download` to at least **4** so `asyncio.gather` has enough in-flight downloads (override with `LITDATA_ASYNC_MIN_PRE_DOWNLOAD`; set `0` to disable the floor). Peak disk ≈ `num_workers × max_pre_download × chunk_size` — size `max_cache_size` accordingly.
+When async is on, LitData raises `max_pre_download` to at least **4** so `asyncio.gather` has enough in-flight downloads (override with `LITDATA_ASYNC_MIN_PRE_DOWNLOAD`; set `0` to disable the floor). Concurrent GETs are capped by `LITDATA_ASYNC_DOWNLOAD_CONCURRENCY` (default **8**). The prepare thread drains the prefetch queue up to that gather width whenever a cache slot is free. Peak disk ≈ `num_workers × max_pre_download × chunk_size` — size `max_cache_size` accordingly.
+
+The reader does **not** poll the cache directory for each chunk. After a download finishes, the downloader atomically `os.replace`s the temp file and sets an in-process `Event`. Compressed chunks set that Event only after decompress publishes the readable `.bin`. Other DataLoader workers still fall back to a short filesystem check (Events are per process).
 
 ```bash
 # Debugging download/delete races — force synchronous downloads
@@ -2213,6 +2244,9 @@ export LITDATA_ASYNC_CHUNK_PREFETCH=0
 
 # Keep max_pre_download=2 even with async enabled
 export LITDATA_ASYNC_MIN_PRE_DOWNLOAD=0
+
+# Cap overlapping remote GETs (default 8)
+export LITDATA_ASYNC_DOWNLOAD_CONCURRENCY=4
 ```
 
 ### Common environment variables
@@ -2222,6 +2256,7 @@ export LITDATA_ASYNC_MIN_PRE_DOWNLOAD=0
 | `LITDATA_CACHE_DIR` | `~/.lightning/chunks` | Default chunk cache directory |
 | `LITDATA_ASYNC_CHUNK_PREFETCH` | on for remote | `0`/`1` force async chunk download overlap |
 | `LITDATA_ASYNC_MIN_PRE_DOWNLOAD` | `4` | Floor for `max_pre_download` when async is on (`0` = no floor) |
+| `LITDATA_ASYNC_DOWNLOAD_CONCURRENCY` | `8` | Max in-flight chunk GETs per gather |
 | `LITDATA_OBSTORE_STREAM_MIN_CHUNK_MIB` | `8` | S3 obstore stream chunk size (MiB) |
 | `MAX_WAIT_TIME` | `120` | Seconds to wait for a chunk before error |
 | `FORCE_DOWNLOAD_TIME` | `30` | Seconds before force re-download of a missing chunk |

{litdata-0.2.70 → litdata-0.2.72}/README.md

@@ -1064,7 +1064,7 @@ Rough ImageNet order-of-magnitude on a Studio (not hard guarantees; right tuning
 | `drop_last` | `True` if distributed else `False` | Equal length across ranks |
 | `seed` | `42` | Shuffle / subsample RNG |
 | `serializers` | built-ins | Custom serialize/deserialize map |
-| `max_cache_size` | `"100GB"` | Evict consumed chunks beyond this size |
+| `max_cache_size` | `None` | Evict consumed chunks beyond this size. Default: 75% of free disk, leaving ≥50GB when possible. Pin with `"100G"` / `"50GB"`, a fraction (`0.90`), or `MAX_CACHE_SIZE`. |
 | `max_pre_download` | `2` | Chunks each worker may prefetch (raise for throughput; watch disk / RAM) |
 | `subsample` | `1.0` | Fraction of data (`0.01`) or upsample (`2.5`) |
 | `encryption` | `None` | `FernetEncryption` / `RSAEncryption` / custom |
@@ -1085,7 +1085,8 @@ On **Vast / NFS / local disk**, POSIX-fast is on by default (`LITDATA_POSIX_FAST
 | All usual `torch.utils.data.DataLoader` kwargs | `batch_size`, `num_workers`, `collate_fn`, `pin_memory`, … |
 | `shuffle` / `drop_last` | Forwarded to the streaming dataset |
 | `profile_batches` | `int` / `True` / `False` — viztracer worker trace (see [Profile data loading](#profile-loading)) |
-| `profile_skip_batches` / `profile_dir` | Warm-up skip count; output dir for `result.json` |
+| `profile_cprofile` | `True` — stdlib cProfile of the main process + worker 0 (see [Profile data loading](#profile-loading)) |
+| `profile_skip_batches` / `profile_dir` | Warm-up skip count; output dir for viztracer / cProfile files |
 | `multiprocessing_context` | Use **`"spawn"`** (or `"forkserver"`) with `ParquetLoader` + `num_workers>0` on Linux |
 
 Prefer `StreamingDataLoader` over a plain PyTorch `DataLoader` for optimized / combined / parallel datasets (resume + correct batch metadata).
@@ -2051,7 +2052,32 @@ Only **worker 0** is instrumented. When an `int` is used, the tracer wraps `fetc
 
 - Delete or change `profile_dir` between runs — LitData removes an existing `result.json` before starting.
 - Pair with a wiped chunk cache if you care about **cold** epoch behavior (`litdata cache clear`).
-- For deeper LitData internals (download / read / delete timeline), use `enable_tracer()` + [Litracer](https://github.com/Lightning-AI/litracer) instead — see [Debug & Profile LitData](#debug-profile). That path is complementary: viztracer = DataLoader worker CPU timeline; Litracer = LitData pipeline events.
+### cProfile (main + worker 0)
+
+Stdlib statistical profile of **the parent process** (queue wait, unpickle, collate handoff) and **worker 0** (fetch, decode, transforms). No extra package.
+
+```python
+loader = StreamingDataLoader(
+    dataset,
+    batch_size=64,
+    num_workers=4,
+    profile_cprofile=True,
+    profile_dir="./profiles",
+)
+
+for batch in loader:
+    train_step(batch)
+# writes profiles/cprofile_main.prof + cprofile_worker0.prof (and .txt summaries)
+```
+
+```bash
+python -m pstats profiles/cprofile_worker0.prof
+# then: sort tottime   stats 30
+```
+
+Do not set `profile_cprofile` and `profile_batches` together — both install a `sys.setprofile` hook. With `num_workers=0` only the main file is written. The parent profiler starts **after** workers spawn so fork does not inherit an active cProfile.
+
+- For deeper LitData internals (download / read / delete timeline), use `enable_tracer()` + [Litracer](https://github.com/Lightning-AI/litracer) instead — see [Debug & Profile LitData](#debug-profile). That path is complementary: viztracer = DataLoader worker CPU timeline; cProfile = function self/cum time; Litracer = LitData pipeline events.
 
 </details>
 
@@ -2101,7 +2127,9 @@ outputs = optimize(
 
 Control how much disk the local chunk cache may use. Downloaded chunks are deleted after use once the cache exceeds the limit.
 
-Default `max_cache_size` is **`100GB`**. Peak disk in flight is roughly:
+Default `max_cache_size` is **`None`**: LitData uses **75% of currently free disk** and still leaves **≥50GB** free when the volume is large enough for checkpoints. Pass `"100G"` / `"50GB"` for a fixed budget, or a float (`0.90`) for that fraction of currently free space. `MAX_CACHE_SIZE` overrides the constructor (size or fraction).
+
+Peak disk in flight is roughly:
 
 ```
 num_workers × max_pre_download × mean_chunk_size
@@ -2141,7 +2169,9 @@ for batch in StreamingDataLoader(dataset, batch_size=64, num_workers=8):
 | `LITDATA_ASYNC_CHUNK_PREFETCH=1` | Force on |
 | `LITDATA_ASYNC_CHUNK_PREFETCH=0` | Force off |
 
-When async is on, LitData raises `max_pre_download` to at least **4** so `asyncio.gather` has enough in-flight downloads (override with `LITDATA_ASYNC_MIN_PRE_DOWNLOAD`; set `0` to disable the floor). Peak disk ≈ `num_workers × max_pre_download × chunk_size` — size `max_cache_size` accordingly.
+When async is on, LitData raises `max_pre_download` to at least **4** so `asyncio.gather` has enough in-flight downloads (override with `LITDATA_ASYNC_MIN_PRE_DOWNLOAD`; set `0` to disable the floor). Concurrent GETs are capped by `LITDATA_ASYNC_DOWNLOAD_CONCURRENCY` (default **8**). The prepare thread drains the prefetch queue up to that gather width whenever a cache slot is free. Peak disk ≈ `num_workers × max_pre_download × chunk_size` — size `max_cache_size` accordingly.
+
+The reader does **not** poll the cache directory for each chunk. After a download finishes, the downloader atomically `os.replace`s the temp file and sets an in-process `Event`. Compressed chunks set that Event only after decompress publishes the readable `.bin`. Other DataLoader workers still fall back to a short filesystem check (Events are per process).
 
 ```bash
 # Debugging download/delete races — force synchronous downloads
@@ -2149,6 +2179,9 @@ export LITDATA_ASYNC_CHUNK_PREFETCH=0
 
 # Keep max_pre_download=2 even with async enabled
 export LITDATA_ASYNC_MIN_PRE_DOWNLOAD=0
+
+# Cap overlapping remote GETs (default 8)
+export LITDATA_ASYNC_DOWNLOAD_CONCURRENCY=4
 ```
 
 ### Common environment variables
@@ -2158,6 +2191,7 @@ export LITDATA_ASYNC_MIN_PRE_DOWNLOAD=0
 | `LITDATA_CACHE_DIR` | `~/.lightning/chunks` | Default chunk cache directory |
 | `LITDATA_ASYNC_CHUNK_PREFETCH` | on for remote | `0`/`1` force async chunk download overlap |
 | `LITDATA_ASYNC_MIN_PRE_DOWNLOAD` | `4` | Floor for `max_pre_download` when async is on (`0` = no floor) |
+| `LITDATA_ASYNC_DOWNLOAD_CONCURRENCY` | `8` | Max in-flight chunk GETs per gather |
 | `LITDATA_OBSTORE_STREAM_MIN_CHUNK_MIB` | `8` | S3 obstore stream chunk size (MiB) |
 | `MAX_WAIT_TIME` | `120` | Seconds to wait for a chunk before error |
 | `FORCE_DOWNLOAD_TIME` | `30` | Seconds before force re-download of a missing chunk |

{litdata-0.2.70 → litdata-0.2.72}/requirements.txt

@@ -5,5 +5,6 @@ filelock
 numpy
 boto3
 requests
+urllib3 >=1.26  # Retry(allowed_methods=...) in streaming/client.py
 tifffile
 obstore

{litdata-0.2.70 → litdata-0.2.72}/src/litdata/__about__.py

@@ -14,7 +14,7 @@
 
 import time
 
-__version__ = "0.2.70"
+__version__ = "0.2.72"
 __author__ = "Lightning AI et al."
 __author_email__ = "pytorch@lightning.ai"
 __license__ = "Apache-2.0"

{litdata-0.2.70 → litdata-0.2.72}/src/litdata/streaming/async_prefetch.py

@@ -54,6 +54,8 @@ _THREAD_LOOPS = threading.local()
 # Empirically, async gather on real S3 is bottlenecked when max_pre_download==2
 # (only 1–2 in-flight). Floor to 4 when the feature is enabled unless overridden.
 _DEFAULT_ASYNC_MIN_PRE_DOWNLOAD = 4
+# Cap in-flight GETs so a large drain batch does not open one connection per slot.
+_DEFAULT_ASYNC_DOWNLOAD_CONCURRENCY = 8
 
 
 def async_chunk_prefetch_enabled(remote_dir: str | None = None) -> bool:
@@ -81,6 +83,17 @@ def async_prefetch_min_pre_download() -> int:
     return max(0, int(raw))
 
 
+def async_download_concurrency(n_chunks: int) -> int:
+    """How many remote chunk GETs to run at once inside ``asyncio.gather``.
+
+    Override with ``LITDATA_ASYNC_DOWNLOAD_CONCURRENCY`` (default 8). Always at
+    least 1 and at most ``n_chunks``.
+    """
+    raw = os.getenv("LITDATA_ASYNC_DOWNLOAD_CONCURRENCY")
+    limit = _DEFAULT_ASYNC_DOWNLOAD_CONCURRENCY if raw is None else max(1, int(raw))
+    return max(1, min(limit, max(1, n_chunks)))
+
+
 def apply_async_pre_download_floor(max_pre_download: int, remote_dir: str | None = None) -> int:
     """Raise ``max_pre_download`` when async prefetch needs gather width."""
     if not async_chunk_prefetch_enabled(remote_dir):
@@ -133,6 +146,7 @@ def _remote_join(remote_dir: str, filename: str) -> str:
 async def _adownload_file_to_path(downloader: Downloader, remote_filepath: str, local_filepath: str) -> None:
     """Fetch ``remote_filepath`` asynchronously and publish atomically."""
     if os.path.exists(local_filepath):
+        downloader._notify_published(local_filepath)
         return
     # Prefer streaming-to-disk when the backend overrides adownload_file.
     if type(downloader).adownload_file is not Downloader.adownload_file:
@@ -149,7 +163,7 @@ async def _adownload_file_to_path(downloader: Downloader, remote_filepath: str,
         os.makedirs(os.path.dirname(local_filepath) or ".", exist_ok=True)
         with open(tmp_path, "wb") as f:
             f.write(data)
-        downloader._atomic_replace(tmp_path, local_filepath)
+        downloader._publish_file(tmp_path, local_filepath)
     except Exception:
         with contextlib.suppress(FileNotFoundError, PermissionError):
             os.remove(tmp_path)
@@ -172,7 +186,7 @@ async def _adownload_chunk_index(config: ChunksConfig, chunk_index: int) -> None
     )
 
     if os.path.exists(local_chunkpath):
-        config.try_decompress(local_chunkpath)
+        await asyncio.to_thread(config.try_decompress, local_chunkpath)
         if lazily_ref_counted:
             downloader._increment_local_lock(lock_path, chunk_index)
         return
@@ -186,7 +200,7 @@ async def _adownload_chunk_index(config: ChunksConfig, chunk_index: int) -> None
         # Overlap blocking SDK calls across threads when native async is unavailable.
         await asyncio.to_thread(downloader.download_chunk_from_index, chunk_index)
 
-    config.try_decompress(local_chunkpath)
+    await asyncio.to_thread(config.try_decompress, local_chunkpath)
 
 
 async def adownload_chunk_indexes(config: ChunksConfig, chunk_indexes: list[int]) -> None:
@@ -196,7 +210,17 @@ async def adownload_chunk_indexes(config: ChunksConfig, chunk_indexes: list[int]
     if len(chunk_indexes) == 1:
         await _adownload_chunk_index(config, chunk_indexes[0])
         return
-    await asyncio.gather(*[_adownload_chunk_index(config, idx) for idx in chunk_indexes])
+    limit = async_download_concurrency(len(chunk_indexes))
+    if limit >= len(chunk_indexes):
+        await asyncio.gather(*[_adownload_chunk_index(config, idx) for idx in chunk_indexes])
+        return
+    sem = asyncio.Semaphore(limit)
+
+    async def _one(idx: int) -> None:
+        async with sem:
+            await _adownload_chunk_index(config, idx)
+
+    await asyncio.gather(*[_one(idx) for idx in chunk_indexes])
 
 
 def _thread_event_loop() -> asyncio.AbstractEventLoop:

{litdata-0.2.70 → litdata-0.2.72}/src/litdata/streaming/cache.py

@@ -27,7 +27,7 @@ from litdata.streaming.serializers import Serializer
 from litdata.streaming.writer import BinaryWriter
 from litdata.utilities.encryption import Encryption
 from litdata.utilities.env import _DistributedEnv, _WorkerEnv
-from litdata.utilities.format import _convert_bytes_to_int
+from litdata.utilities.format import _resolve_max_cache_size
 
 logger = logging.Logger(__name__)
 
@@ -43,7 +43,7 @@ class Cache:
         chunk_size: int | None = None,
         chunk_bytes: int | str | None = None,
         item_loader: BaseItemLoader | None = None,
-        max_cache_size: int | str = "100GB",
+        max_cache_size: int | float | str | None = None,
         serializers: dict[str, Serializer] | None = None,
         writer_chunk_index: int | None = None,
         storage_options: dict | None = {},
@@ -64,7 +64,8 @@ class Cache:
             chunk_bytes: The maximum number of bytes within a chunk.
             chunk_size: The maximum number of items within a chunk.
             item_loader: The object responsible to generate the chunk intervals and load an item froma chunk.
-            max_cache_size: The maximum cache size used by the reader when fetching the chunks.
+            max_cache_size: Cache budget. ``None`` uses 75% of free disk (see ``StreamingDataset``).
+                A float such as ``0.90`` is that fraction of currently free space.
             serializers: Provide your own serializers.
             writer_chunk_index: The index of the chunk to start from when writing.
             storage_options: Additional connection options for accessing storage services.
@@ -93,7 +94,7 @@ class Cache:
             self._cache_dir,
             subsampled_files=subsampled_files,
             region_of_interest=region_of_interest,
-            max_cache_size=_convert_bytes_to_int(max_cache_size) if isinstance(max_cache_size, str) else max_cache_size,
+            max_cache_size=_resolve_max_cache_size(max_cache_size, self._cache_dir),
             remote_input_dir=input_dir.url,
             compression=compression,
             encryption=encryption,