litdata 0.2.64__tar.gz → 0.2.65__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: litdata
-Version: 0.2.64
+Version: 0.2.65
 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

{litdata-0.2.64 → litdata-0.2.65}/src/litdata/__about__.py

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

{litdata-0.2.64 → litdata-0.2.65}/src/litdata/constants.py

@@ -59,6 +59,8 @@ _PRINT_DEBUG_LOGS = bool(int(os.getenv("PRINT_DEBUG_LOGS", "0")))
 _MAX_WAIT_TIME = int(os.getenv("MAX_WAIT_TIME", "120"))
 _FORCE_DOWNLOAD_TIME = int(os.getenv("FORCE_DOWNLOAD_TIME", "30"))
 _LITDATA_DISABLE_VERSION_CHECK = int(os.getenv("LITDATA_DISABLE_VERSION_CHECK", "0"))
+# Experimental async chunk prefetch: set env ``LITDATA_ASYNC_CHUNK_PREFETCH=1``
+# (see ``litdata.streaming.async_prefetch``). Off by default.
 
 # DON'T CHANGE ORDER
 _TORCH_DTYPES_MAPPING = {

litdata-0.2.65/src/litdata/streaming/async_prefetch.py

@@ -0,0 +1,256 @@
+# Copyright The Lightning AI team.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Experimental asyncio helpers for overlapping **remote chunk downloads**.
+
+This is intentionally **not** an async ``StreamingDataLoader``. Training stays on
+the sync ``for batch in loader`` API; decode stays on process workers. Asyncio is
+only useful where we wait on network IO.
+
+On by default for remote datasets; force off with ``LITDATA_ASYNC_CHUNK_PREFETCH=0``,
+or force on locally with ``LITDATA_ASYNC_CHUNK_PREFETCH=1``.
+
+Strategy:
+  * Prefer ``Downloader.adownload_file`` (streaming to disk) when overridden.
+  * Else ``Downloader.adownload_fileobj`` + atomic write.
+  * Otherwise run sync ``download_chunk_from_index`` in ``asyncio.to_thread`` and
+    ``gather`` several chunk indexes — still overlaps latency for blocking cloud SDKs.
+
+Real-S3 note (Studio benches on ~67MB ImageNet chunks): gather concurrency needs
+``max_pre_download >= 4`` to matter; the prepare thread raises the floor when
+async prefetch is enabled (see :func:`async_prefetch_min_pre_download`).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+import os
+import threading
+from typing import TYPE_CHECKING
+
+from litdata.streaming.downloader import Downloader
+
+if TYPE_CHECKING:
+    from litdata.streaming.config import ChunksConfig
+
+logger = logging.getLogger("litdata.streaming.async_prefetch")
+
+# Thread-local event loop so PrepareChunksThread does not pay asyncio.run()
+# startup on every drain batch.
+_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
+
+
+def async_chunk_prefetch_enabled(remote_dir: str | None = None) -> bool:
+    """Return True when async chunk prefetch should run.
+
+    Env ``LITDATA_ASYNC_CHUNK_PREFETCH`` wins when set (``1``/``0``). When unset,
+    default **on** for remote datasets (real-S3: process+async is the fast path)
+    and **off** for local-only caches.
+    """
+    raw = os.getenv("LITDATA_ASYNC_CHUNK_PREFETCH")
+    if raw is not None:
+        return bool(int(raw))
+    return bool(remote_dir)
+
+
+def async_prefetch_min_pre_download() -> int:
+    """Minimum ``max_pre_download`` applied when async chunk prefetch is on.
+
+    Override with ``LITDATA_ASYNC_MIN_PRE_DOWNLOAD`` (default 4). Set to ``0`` to
+    disable the floor and keep the caller's ``max_pre_download`` unchanged.
+    """
+    raw = os.getenv("LITDATA_ASYNC_MIN_PRE_DOWNLOAD")
+    if raw is None:
+        return _DEFAULT_ASYNC_MIN_PRE_DOWNLOAD
+    return max(0, int(raw))
+
+
+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):
+        return max_pre_download
+    floor = async_prefetch_min_pre_download()
+    if floor <= 0 or max_pre_download >= floor:
+        return max_pre_download
+    logger.info(
+        "Async chunk prefetch: raising max_pre_download %s → %s so "
+        "asyncio.gather can overlap remote chunk downloads. "
+        "Set LITDATA_ASYNC_MIN_PRE_DOWNLOAD=0 to keep the original value.",
+        max_pre_download,
+        floor,
+    )
+    return floor
+
+
+def downloader_supports_adownload(downloader: Downloader | None) -> bool:
+    """True when an async download path is available on ``downloader``."""
+    if downloader is None:
+        return False
+    cls = type(downloader)
+    return (
+        cls.adownload_file is not Downloader.adownload_file or cls.adownload_fileobj is not Downloader.adownload_fileobj
+    )
+
+
+def _remote_join(remote_dir: str, filename: str) -> str:
+    """Join cloud URLs without ``os.path.join`` quirks on schemes."""
+    return remote_dir.rstrip("/") + "/" + filename.lstrip("/")
+
+
+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):
+        return
+    # Prefer streaming-to-disk when the backend overrides adownload_file.
+    if type(downloader).adownload_file is not Downloader.adownload_file:
+        await downloader.adownload_file(remote_filepath, local_filepath)
+        return
+    data = await downloader.adownload_fileobj(remote_filepath)
+    if data is None:
+        raise NotImplementedError(
+            f"{type(downloader).__name__}.adownload_fileobj returned None; "
+            "cannot use async chunk prefetch for this backend."
+        )
+    tmp_path = downloader._temp_download_path(local_filepath)
+    try:
+        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)
+    except Exception:
+        with contextlib.suppress(FileNotFoundError, PermissionError):
+            os.remove(tmp_path)
+        raise
+
+
+async def _adownload_chunk_index(config: ChunksConfig, chunk_index: int) -> None:
+    """Async download + decompress for one chunk index (mirrors sync config path)."""
+    assert config._chunks is not None
+    downloader = config._downloader
+    if downloader is None:
+        return
+
+    chunk_filename = config._chunks[chunk_index]["filename"]
+    local_chunkpath = os.path.join(config._cache_dir, chunk_filename)
+    remote_chunkpath = _remote_join(downloader._remote_dir, chunk_filename)
+    lazily_ref_counted = chunk_index not in config._shared_chunk_indexes
+    lock_path = (
+        local_chunkpath.replace(f".{config._compressor_name}", "") if config._compressor_name else local_chunkpath
+    )
+
+    if os.path.exists(local_chunkpath):
+        config.try_decompress(local_chunkpath)
+        if lazily_ref_counted:
+            downloader._increment_local_lock(lock_path, chunk_index)
+        return
+
+    if lazily_ref_counted:
+        downloader._increment_local_lock(lock_path, chunk_index)
+
+    if downloader_supports_adownload(downloader):
+        await _adownload_file_to_path(downloader, remote_chunkpath, local_chunkpath)
+    else:
+        # 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)
+
+
+async def adownload_chunk_indexes(config: ChunksConfig, chunk_indexes: list[int]) -> None:
+    """Download several chunk indexes concurrently (gather)."""
+    if not chunk_indexes:
+        return
+    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])
+
+
+def _thread_event_loop() -> asyncio.AbstractEventLoop:
+    """Return a reusable event loop for the current thread."""
+    loop = getattr(_THREAD_LOOPS, "loop", None)
+    if loop is None or loop.is_closed():
+        loop = asyncio.new_event_loop()
+        _THREAD_LOOPS.loop = loop
+    return loop
+
+
+def close_thread_event_loop() -> None:
+    """Shut down the thread-local loop and its default executor.
+
+    ``asyncio.to_thread`` (used when a downloader has no native async path) parks
+    workers named ``asyncio_N`` on the loop's default executor. Leaving the loop
+    open leaks those threads and trips the session thread-police on Windows.
+
+    Executor shutdown is non-blocking (``wait=False``) so prepare-thread ``finally``
+    never joins forever if a download worker is stuck. Threads are marked daemon and
+    futures cancelled so orphans cannot keep the process alive or poison later
+    DataLoader forks under pytest-xdist.
+    """
+    loop = getattr(_THREAD_LOOPS, "loop", None)
+    _THREAD_LOOPS.loop = None
+    if loop is None or loop.is_closed():
+        return
+    try:
+        if not loop.is_running():
+            with contextlib.suppress(Exception):
+                loop.run_until_complete(loop.shutdown_asyncgens())
+            # Prefer non-blocking teardown over ``shutdown_default_executor()``,
+            # which joins executor threads and can hang prepare-thread exit.
+            executor = getattr(loop, "_default_executor", None)
+            if executor is not None:
+                with contextlib.suppress(Exception):
+                    loop._default_executor = None
+                    for thread in getattr(executor, "_threads", set()):
+                        thread.daemon = True
+                    # cancel_futures is 3.9+; litdata already requires newer Python.
+                    executor.shutdown(wait=False, cancel_futures=True)
+    finally:
+        with contextlib.suppress(Exception):
+            loop.close()
+
+
+def download_chunk_indexes_concurrently(config: ChunksConfig, chunk_indexes: list[int]) -> None:
+    """Sync entry point for ``PrepareChunksThread``: run ``adownload_chunk_indexes``."""
+    if not chunk_indexes:
+        return
+    if len(chunk_indexes) == 1:
+        config.download_chunk_from_index(chunk_indexes[0])
+        return
+    # Reuse a per-thread loop instead of asyncio.run() (which creates+closes a loop
+    # on every prefetch batch — measurable overhead under high chunk churn).
+    try:
+        running = asyncio.get_running_loop()
+    except RuntimeError:
+        running = None
+    if running is not None:
+        # Already inside an event loop (tests / nested callers): block via a bridge.
+        import concurrent.futures
+
+        def _run() -> None:
+            try:
+                _thread_event_loop().run_until_complete(adownload_chunk_indexes(config, chunk_indexes))
+            finally:
+                close_thread_event_loop()
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
+            pool.submit(_run).result()
+        return
+    loop = _thread_event_loop()
+    loop.run_until_complete(adownload_chunk_indexes(config, chunk_indexes))

{litdata-0.2.64 → litdata-0.2.65}/src/litdata/streaming/config.py

@@ -15,10 +15,14 @@ import contextlib
 import logging
 import os
 from collections import defaultdict
+from contextlib import suppress
 from time import sleep, time
 from typing import Any, Optional
 
+from filelock import FileLock, Timeout
+
 from litdata.constants import _INDEX_FILENAME, _MAX_WAIT_TIME
+from litdata.debugger import _get_log_msg
 from litdata.streaming.compression import _COMPRESSORS, Compressor
 from litdata.streaming.downloader import get_downloader
 from litdata.streaming.item_loader import BaseItemLoader, Interval, PyTreeLoader, TokensLoader
@@ -107,7 +111,14 @@ class ChunksConfig:
             self._compressor = _COMPRESSORS[self._compressor_name]
 
         self._skip_chunk_indexes_deletion: list[int] | None = None
+        # Chunk indexes that are shared across workers on this node. Shared chunks are
+        # reference-counted *eagerly* (incremented at iteration start by the reader), so their
+        # lazy download-time increment is skipped to keep the count balanced. See reader.py.
+        self._shared_chunk_indexes: set[int] = set()
         self.zero_based_roi: list[tuple[int, int]] | None = None
+        # Memoizes ``__getitem__`` results per chunk_index (invariant once the config is loaded);
+        # avoids rebuilding the chunk path on every item read.
+        self._chunk_meta_cache: dict[int, tuple[str, int, int]] = {}
         self.filename_to_size_map: dict[str, int] = {}
         for cnk in _original_chunks:
             # since files downloaded while reading will be decompressed, we need to store the name without compression
@@ -119,6 +130,67 @@ class ChunksConfig:
             return True
         return chunk_index not in self._skip_chunk_indexes_deletion
 
+    def _chunk_lock_filepath(self, chunk_index: int) -> str:
+        """The (decompressed) local chunk path whose ``.cnt``/``.lock`` files hold the refcount."""
+        chunk_filepath, _, _ = self[ChunkedIndex(index=-1, chunk_index=chunk_index)]
+        return chunk_filepath
+
+    def remaining_locks(self, chunk_index: int) -> int:
+        """Return the current reference count held on a chunk (0 if none)."""
+        countpath = self._chunk_lock_filepath(chunk_index) + ".cnt"
+        if not os.path.exists(countpath):
+            return 0
+        with suppress(FileNotFoundError), open(countpath) as count_f:
+            try:
+                return int(count_f.read().strip())
+            except Exception:
+                return 1
+        return 0
+
+    def increment_local_lock(self, chunk_index: int) -> None:
+        """Add one reference to a chunk's local lock (a co-reader intends to use it)."""
+        if self._downloader is None:
+            return
+        self._downloader._increment_local_lock(self._chunk_lock_filepath(chunk_index), chunk_index)
+
+    def decrement_local_lock(self, chunk_index: int) -> int:
+        """Remove one reference from a chunk's local lock; return the remaining count.
+
+        Moved here (from ``PrepareChunksThread``) so the reader can release eagerly-acquired locks
+        during teardown without depending on the prefetch thread still being alive.
+        """
+        countpath = self._chunk_lock_filepath(chunk_index) + ".cnt"
+        lock_path = countpath + ".lock"
+        curr_count = 0
+        remove_lock = False
+        with suppress(Timeout, FileNotFoundError), FileLock(lock_path, timeout=3):
+            if os.path.exists(countpath):
+                with open(countpath) as count_f:
+                    try:
+                        curr_count = int(count_f.read().strip())
+                    except Exception:
+                        curr_count = 1
+                curr_count -= 1
+                if curr_count <= 0:
+                    with suppress(FileNotFoundError, PermissionError):
+                        os.remove(countpath)
+                    remove_lock = True
+                else:
+                    with open(countpath, "w+") as count_f:
+                        logger.debug(_get_log_msg({"name": f"decrement_lock_{chunk_index}_to_{curr_count}", "ph": "B"}))
+                        count_f.write(str(curr_count))
+                        logger.debug(_get_log_msg({"name": f"decrement_lock_{chunk_index}_to_{curr_count}", "ph": "E"}))
+            else:
+                remove_lock = True
+        # 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).
+        if remove_lock:
+            with suppress(FileNotFoundError, PermissionError):
+                os.remove(lock_path)
+        return curr_count
+
     @property
     def skip_chunk_indexes_deletion(self) -> list[int] | None:
         return self._skip_chunk_indexes_deletion
@@ -133,22 +205,26 @@ class ChunksConfig:
 
         local_chunkpath = os.path.join(self._cache_dir, chunk_filename)
 
+        # Shared chunks are reference-counted eagerly by the reader (before any reading), so their
+        # download-time increment is skipped here to avoid double-counting. Non-shared chunks keep
+        # the original pay-as-you-download refcounting.
+        lazily_ref_counted = chunk_index not in self._shared_chunk_indexes
+
         if os.path.exists(local_chunkpath):
             self.try_decompress(local_chunkpath)
 
-            if self._downloader is not None and not skip_lock:
+            if self._downloader is not None and not skip_lock and lazily_ref_counted:
                 # We don't want to redownload the base, but we should mark
                 # it as having been requested by something
                 self._downloader._increment_local_lock(
                     local_chunkpath.replace(f".{self._compressor_name}", ""), chunk_index
                 )
-                pass
             return
 
         if self._downloader is None:
             return
 
-        if not skip_lock:
+        if not skip_lock and lazily_ref_counted:
             self._downloader._increment_local_lock(
                 local_chunkpath.replace(f".{self._compressor_name}", ""), chunk_index
             )
@@ -188,37 +264,58 @@ class ChunksConfig:
         if os.path.exists(target_local_chunkpath):
             return
 
-        # Ensure that the compressed file exists and is fully downloaded
+        # 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.
         start_time = time()
-        assert self._chunks is not None
-
-        filename = os.path.basename(local_chunkpath)
-        chunk_index = self._get_chunk_index_from_filename(filename)
-        chunk_bytes = self._chunks[chunk_index]["chunk_size"]
-        exists = os.path.exists(local_chunkpath) and os.stat(local_chunkpath).st_size >= chunk_bytes
-        while not exists:
+        while not os.path.exists(local_chunkpath) and not os.path.exists(target_local_chunkpath):
             sleep(0.1)
-            # Return if the actual file exists
-            if os.path.exists(target_local_chunkpath):
-                return
-            # find the local compressed file
-            exists = os.path.exists(local_chunkpath) and os.stat(local_chunkpath).st_size >= chunk_bytes
-
             if (time() - start_time) > _MAX_WAIT_TIME:
                 raise FileNotFoundError(f"The {local_chunkpath} hasn't been found.")
 
-        with open(local_chunkpath, "rb") as f:
-            data = f.read()
-
-        # delete the files only if they were downloaded
-        if self._downloader is not None:
-            with contextlib.suppress(FileNotFoundError):
-                os.remove(local_chunkpath)
-
-        data = self._compressor.decompress(data)
+        if os.path.exists(target_local_chunkpath):
+            return
 
-        with open(target_local_chunkpath, "wb") as f:
-            f.write(data)
+        decompress_lock = target_local_chunkpath + ".decompress.lock"
+        try:
+            with FileLock(decompress_lock, timeout=_MAX_WAIT_TIME):
+                if os.path.exists(target_local_chunkpath):
+                    return
+
+                with open(local_chunkpath, "rb") as f:
+                    data = f.read()
+
+                # delete the compressed file only if it was downloaded
+                if self._downloader is not None:
+                    with contextlib.suppress(FileNotFoundError):
+                        os.remove(local_chunkpath)
+
+                data = self._compressor.decompress(data)
+
+                assert self._chunks is not None
+                filename = os.path.basename(local_chunkpath)
+                chunk_index = self._get_chunk_index_from_filename(filename)
+                expected_bytes = int(self._chunks[chunk_index]["chunk_bytes"])
+
+                tmp_path = f"{target_local_chunkpath}.tmp.{os.getpid()}"
+                try:
+                    with open(tmp_path, "wb") as f:
+                        f.write(data)
+                    if os.stat(tmp_path).st_size < expected_bytes:
+                        raise OSError(
+                            f"Decompressed chunk {target_local_chunkpath} is smaller than expected "
+                            f"({os.stat(tmp_path).st_size} < {expected_bytes})."
+                        )
+                    os.replace(tmp_path, target_local_chunkpath)
+                except Exception:
+                    with contextlib.suppress(FileNotFoundError, PermissionError):
+                        os.remove(tmp_path)
+                    raise
+        finally:
+            # FileLock leaves its lock file behind; remove after release.
+            with contextlib.suppress(Exception):
+                os.remove(decompress_lock)
 
     @property
     def intervals(self) -> list[Interval]:
@@ -287,7 +384,17 @@ class ChunksConfig:
         )
 
     def __getitem__(self, index: ChunkedIndex) -> tuple[str, int, int]:
-        """Find the associated chunk metadata."""
+        """Find the associated chunk metadata.
+
+        This is called once per item on the read hot path, but its result depends only on
+        ``index.chunk_index`` (the local path, the chunk's begin offset and its byte size are all
+        fixed once the config is loaded). The per-chunk tuple is therefore memoized to avoid
+        rebuilding the path (``os.path.join`` + decompression-suffix stripping) on every item.
+        """
+        cached = self._chunk_meta_cache.get(index.chunk_index)
+        if cached is not None:
+            return cached
+
         assert self._chunks is not None
         chunk = self._chunks[index.chunk_index]
 
@@ -300,7 +407,9 @@ class ChunksConfig:
 
         filesize_bytes = chunk["chunk_bytes"]
 
-        return local_chunkpath, begin, filesize_bytes
+        meta = (local_chunkpath, begin, filesize_bytes)
+        self._chunk_meta_cache[index.chunk_index] = meta
+        return meta
 
     def download_filepath(self, chunk_index: int) -> str:
         """The raw on-disk path that the chunk is downloaded to before any decompression."""

{litdata-0.2.64 → litdata-0.2.65}/src/litdata/streaming/dataloader.py

@@ -41,6 +41,7 @@ from litdata.streaming.combined import CombinedStreamingDataset
 from litdata.streaming.dataset import StreamingDataset
 from litdata.streaming.parallel import ParallelStreamingDataset
 from litdata.streaming.sampler import CacheBatchSampler
+from litdata.streaming.timing import StreamingTimingStats
 from litdata.utilities._pytree import tree_flatten
 from litdata.utilities.base import (
     __NUM_CYCLES_KEY__,
@@ -674,12 +675,13 @@ class StreamingDataLoader(DataLoader):
         self._worker_idx_iter: Any | None = None
         self._latest_worker_idx = 0
         self.restore = False
+        self._prefetch_factor = (2 if num_workers > 0 else None) if prefetch_factor is None else prefetch_factor
         super().__init__(
             dataset,
             *args,
             batch_size=batch_size,
             num_workers=num_workers,
-            prefetch_factor=(2 if num_workers > 0 else None) if prefetch_factor is None else prefetch_factor,
+            prefetch_factor=self._prefetch_factor,
             collate_fn=collate_fn,
             **kwargs,
         )  # type: ignore
@@ -717,9 +719,12 @@ class StreamingDataLoader(DataLoader):
 
         if isinstance(self.dataset, StreamingDataset):
             assert self.batch_size
+            timing = StreamingTimingStats.instance()
             for batch in super().__iter__():
+                t0 = timing.start()
                 self._latest_worker_idx = next(self._worker_idx_iter)  # type: ignore
                 self._num_samples_yielded_streaming += self.batch_size
+                timing.record("dataloader_yield_s", t0)
                 yield batch
         else:
             self.dataset._set_use_streaming_dataloader(True)
@@ -744,6 +749,13 @@ class StreamingDataLoader(DataLoader):
                 else:
                     yield batch
 
+        # NOTE: `restore` is intentionally *not* cleared in a `finally` block here. Breaking out of
+        # this generator early (or letting it get garbage-collected, which throws `GeneratorExit` at
+        # the last `yield`) must leave `restore` untouched: callers that explicitly resumed from a
+        # checkpoint (`load_state_dict`) rely on `restore` staying `True` across such early exits so
+        # that the *next* `__iter__` call keeps skipping `reset_state_dict()` and replays from the
+        # loaded state (see `_StreamingMultiProcessingDataLoaderIter._try_put_index`). `restore` is
+        # only toggled back to `False` here, once the loop above completes a full epoch normally.
         logger.debug(_get_log_msg({"name": "iterating_dataloader", "ph": "E"}))
         self.restore = False