litdata 0.2.72__tar.gz → 0.2.73__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: litdata
-Version: 0.2.72
+Version: 0.2.73
 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.72 → litdata-0.2.73}/src/litdata/__about__.py

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

{litdata-0.2.72 → litdata-0.2.73}/src/litdata/processing/data_processor.py

@@ -1523,7 +1523,16 @@ class DataChunkRecipe(DataRecipe):
 
         merge_cache = Cache(cache_dir, chunk_bytes=1)
         node_rank = _get_node_rank()
-        merge_cache._merge_no_wait(node_rank if num_nodes > 1 else None, getattr(self, "existing_index", None))
+        # With more than one node, every node would otherwise fold the existing
+        # index into its own {node_rank}-index.json, and the final cross-node
+        # merge below then repeats those chunks once per node. Add the existing
+        # index only once, either here for the single-node case or in the final
+        # merge for the multi-node case.
+        existing_index = getattr(self, "existing_index", None)
+        merge_cache._merge_no_wait(
+            node_rank if num_nodes > 1 else None,
+            existing_index if num_nodes == 1 else None,
+        )
 
         self._merge_and_upload_keys(output_dir, cache_dir, num_nodes, node_rank)
         self._upload_index(output_dir, cache_dir, num_nodes, node_rank)
@@ -1734,7 +1743,9 @@ class DataChunkRecipe(DataRecipe):
                     shutil.copyfile(remote_filepath, node_index_filepath)
 
             merge_cache = Cache(merge_dir, chunk_bytes=1)
-            merge_cache._merge_no_wait()
+            # The per-node index files hold only their own new chunks now, so the
+            # existing index is folded in here, once, as the node files are merged.
+            merge_cache._merge_no_wait(existing_index=getattr(self, "existing_index", None))
             self._upload_index(output_dir, merge_dir, 1, None)
 
 

{litdata-0.2.72 → litdata-0.2.73}/src/litdata/processing/functions.py

@@ -147,6 +147,17 @@ class LambdaMapRecipe(MapRecipe):
         self._contains_device = "device" in params
         self._contains_is_last = "is_last" in params
 
+    def __getstate__(self) -> dict[str, Any]:
+        """Drop the full input sequence when pickling into spawn workers.
+
+        The parent already shards items onto per-worker queues. Workers only need
+        ``prepare_item``; keeping ``_inputs`` in the pickle would duplicate the list
+        once per process.
+        """
+        state = self.__dict__.copy()
+        state["_inputs"] = None
+        return state
+
     def prepare_structure(self, _: str | None) -> Any:
         return self._inputs
 
@@ -210,6 +221,17 @@ class LambdaDataChunkRecipe(DataChunkRecipe):
 
         self.prepare_item = self._prepare_item_generator if self.is_generator else self._prepare_item  # type: ignore
 
+    def __getstate__(self) -> dict[str, Any]:
+        """Drop the full input sequence when pickling into spawn workers.
+
+        The parent already shards items onto per-worker queues. Workers only need
+        ``prepare_item``; keeping ``_inputs`` in the pickle would duplicate the list
+        once per process.
+        """
+        state = self.__dict__.copy()
+        state["_inputs"] = None
+        return state
+
     def check_fn(self) -> None:
         if (
             isinstance(self._fn, (partial, FunctionType))

{litdata-0.2.72 → litdata-0.2.73}/src/litdata/streaming/dataloader.py

@@ -558,25 +558,30 @@ class _StreamingMultiProcessingDataLoaderIter(_MultiProcessingDataLoaderIter):
         )
         self._num_workers = loader.num_workers
         self._cprofile: cProfile.Profile | None = None
-        self._orig_worker_loop: Any = None
 
         distributed_env = _DistributedEnv.detect()
         profile_cprofile = bool(getattr(self._loader, "_profile_cprofile", False))
+        from torch.utils.data._utils import worker
 
-        if distributed_env.global_rank == 0:
-            from torch.utils.data._utils import worker
+        original_worker_loop: Any = None
 
+        if distributed_env.global_rank == 0:
             if self._loader._profile_batches and _VIZ_TRACKER_AVAILABLE:
+                original_worker_loop = worker._worker_loop
                 worker._worker_loop = _ProfileWorkerLoop(
                     self._loader._profile_batches, self._loader._profile_skip_batches, self._loader._profile_dir
                 )
             elif profile_cprofile:
-                self._orig_worker_loop = worker._worker_loop
+                original_worker_loop = worker._worker_loop
                 worker._worker_loop = _CProfileWorkerLoop(_cprofile_output_dir(self._loader._profile_dir))
 
         # Workers fork/spawn here. Enable the parent profiler only after that so
         # the child does not inherit an active cProfile (one profiler per process).
-        super().__init__(loader)
+        try:
+            super().__init__(loader)
+        finally:
+            if original_worker_loop is not None:
+                worker._worker_loop = original_worker_loop
 
         if profile_cprofile and distributed_env.global_rank == 0:
             self._cprofile = cProfile.Profile()
@@ -587,11 +592,6 @@ class _StreamingMultiProcessingDataLoaderIter(_MultiProcessingDataLoaderIter):
             stem = os.path.join(_cprofile_output_dir(self._loader._profile_dir), _CPROFILE_MAIN_STEM)
             _dump_cprofile(self._cprofile, stem)
             self._cprofile = None
-        if self._orig_worker_loop is not None:
-            from torch.utils.data._utils import worker
-
-            worker._worker_loop = self._orig_worker_loop
-            self._orig_worker_loop = None
         super()._shutdown_workers()
 
     def _try_put_index(self) -> None:

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

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