litdata 0.2.70__tar.gz → 0.2.71__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: litdata
-Version: 0.2.70
+Version: 0.2.71
 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
@@ -1128,7 +1128,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 +1149,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 +2116,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 +2191,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 +2233,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 +2243,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 +2255,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.71}/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.71}/src/litdata/__about__.py

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

{litdata-0.2.70 → litdata-0.2.71}/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.71}/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,

{litdata-0.2.70 → litdata-0.2.71}/src/litdata/streaming/compression.py

@@ -13,7 +13,7 @@
 
 import shutil
 from abc import ABC, abstractmethod
-from typing import TypeVar
+from typing import Any, TypeVar
 
 from litdata.constants import _PYTHON_GREATER_EQUAL_3_14, _ZSTD_AVAILABLE
 from litdata.debugger import CAT_DECOMPRESS, trace_span
@@ -54,27 +54,32 @@ class ZSTDCompressor(Compressor):
             raise ModuleNotFoundError(str(_ZSTD_AVAILABLE))
         self.level = level
         self.extension = "zstd"
+        self._zstd: Any | None = None
+
+    def __getstate__(self) -> dict[str, Any]:
+        state = self.__dict__.copy()
+        state["_zstd"] = None
+        return state
+
+    def _zstd_mod(self) -> Any:
+        if self._zstd is None:
+            if _PYTHON_GREATER_EQUAL_3_14:
+                from compression import zstd as mod
+            else:
+                import zstd as mod
+            self._zstd = mod
+        return self._zstd
 
     @property
     def name(self) -> str:
         return f"{self.extension}:{self.level}"
 
     def compress(self, data: bytes) -> bytes:
-        if _PYTHON_GREATER_EQUAL_3_14:
-            from compression import zstd
-        else:
-            import zstd
-
-        return zstd.compress(data, self.level)
+        return self._zstd_mod().compress(data, self.level)
 
     def decompress(self, data: bytes) -> bytes:
-        if _PYTHON_GREATER_EQUAL_3_14:
-            from compression import zstd
-        else:
-            import zstd
-
         with trace_span("decompress", CAT_DECOMPRESS):
-            return zstd.decompress(data)
+            return self._zstd_mod().decompress(data)
 
     def decompress_file(self, src: str, dst: str) -> None:
         if _PYTHON_GREATER_EQUAL_3_14:

{litdata-0.2.70 → litdata-0.2.71}/src/litdata/streaming/config.py

@@ -14,9 +14,10 @@
 import contextlib
 import logging
 import os
+import threading
 from collections import defaultdict
 from contextlib import suppress
-from time import sleep, time
+from time import time
 from typing import Any, Optional
 
 from filelock import FileLock, Timeout
@@ -92,10 +93,15 @@ class ChunksConfig:
         self._length = self._intervals[-1][-1] if len(self._intervals) > 0 else 0
         self._downloader = None
 
+        # In-process "file is visible" Events, set after atomic publish (download or decompress).
+        self._file_published: dict[str, threading.Event] = {}
+        self._file_published_lock = threading.Lock()
+
         if remote_dir:
             self._downloader = get_downloader(
                 remote_dir, cache_dir, self._chunks, self._storage_options, self._session_options
             )
+            self._downloader._on_file_published = self.notify_file_published
 
         self._compressor_name = self._config["compression"]
         self._compressor: Compressor | None = None
@@ -256,6 +262,28 @@ class ChunksConfig:
 
         return self._downloader.download_chunk_bytes_from_index(chunk_index, offset, length)
 
+    def notify_file_published(self, local_filepath: str) -> None:
+        """Mark ``local_filepath`` as atomically published and, if it is a readable chunk, signal ready."""
+        path = os.path.abspath(local_filepath)
+        with self._file_published_lock:
+            event = self._file_published.get(path)
+            if event is None:
+                event = threading.Event()
+                self._file_published[path] = event
+            event.set()
+
+    def wait_file_published(self, local_filepath: str, timeout: float) -> bool:
+        """Wait for an in-process publish Event; fall back to ``exists`` for other workers."""
+        path = os.path.abspath(local_filepath)
+        with self._file_published_lock:
+            event = self._file_published.get(path)
+            if event is None:
+                event = threading.Event()
+                self._file_published[path] = event
+        if event.wait(timeout=timeout):
+            return True
+        return os.path.exists(path)
+
     def _remove_compressed_source(self, local_chunkpath: str, target_local_chunkpath: str) -> None:
         """Drop the compressed download once the decompressed chunk is on disk."""
         if local_chunkpath == target_local_chunkpath:
@@ -273,13 +301,12 @@ class ChunksConfig:
             self._remove_compressed_source(local_chunkpath, target_local_chunkpath)
             return
 
-        # Wait until either the decompressed target appears (another worker finished) or the
-        # compressed source exists. Cloud downloaders publish the compressed path atomically, so
-        # existence of that path means the download is complete — do NOT use chunk_size (item
-        # count) as a byte threshold.
+        # Wait until the downloader publishes the compressed file, or another worker
+        # publishes the decompressed target. Cross-process downloads still fall back to exists.
         start_time = time()
-        while not os.path.exists(local_chunkpath) and not os.path.exists(target_local_chunkpath):
-            sleep(0.1)
+        while not os.path.exists(target_local_chunkpath) and not os.path.exists(local_chunkpath):
+            self.wait_file_published(local_chunkpath, 0.05)
+            self.wait_file_published(target_local_chunkpath, 0.0)
             if (time() - start_time) > _MAX_WAIT_TIME:
                 raise ChunkWaitTimeoutError(local_chunkpath, time() - start_time)
 
@@ -307,6 +334,7 @@ class ChunksConfig:
                             f"({os.stat(tmp_path).st_size} < {expected_bytes})."
                         )
                     os.replace(tmp_path, target_local_chunkpath)
+                    self.notify_file_published(target_local_chunkpath)
                 except Exception:
                     with contextlib.suppress(FileNotFoundError, PermissionError):
                         os.remove(tmp_path)
@@ -464,6 +492,20 @@ class ChunksConfig:
             session_options,
         )
 
+    def __getstate__(self) -> dict[str, Any]:
+        state = self.__dict__.copy()
+        # threading.Lock / Event are not picklable (DataLoader spawn / deepcopy).
+        state["_file_published"] = {}
+        state["_file_published_lock"] = None
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        self._file_published = {}
+        self._file_published_lock = threading.Lock()
+        if self._downloader is not None:
+            self._downloader._on_file_published = self.notify_file_published
+
     def __len__(self) -> int:
         return self._length