Add support for exact iteration (#139)

Author: tchaton

src/litdata/streaming/combined.py

@@ -80,15 +80,27 @@ def __init__(
         self._use_streaming_dataloader = False
         self._num_samples_yielded: Optional[List[int]] = None
         self._current_epoch = 0
+        self.num_workers = 1
+        self.batch_size = 1
 
-    def __len__(self) -> Optional[int]:
+    def get_len(self, num_workers: int, batch_size: int) -> Optional[int]:
+        self.num_workers = num_workers
+        self.batch_size = batch_size
         if self._iterate_over_all:
             return self._get_total_length()
         return None
 
+    def __len__(self) -> Optional[int]:
+        return self.get_len(1, 1)
+
     # total length of the datasets
     def _get_total_length(self) -> int:
-        return sum(len(d) for d in self._datasets)
+        return sum(self._get_len(d) for d in self._datasets)
+
+    def _get_len(self, d: Any) -> int:
+        if isinstance(d, StreamingDataset):
+            return d.get_len(self.num_workers, self.batch_size)
+        return len(d)
 
     def set_epoch(self, current_epoch: int) -> None:
         """Set the current epoch to the datasets on epoch starts.

src/litdata/streaming/dataloader.py

@@ -615,6 +615,16 @@ def __iter__(self) -> Any:
 
         self.restore = False
 
+    def __len__(self) -> int:
+        if self._dataset_kind == _DatasetKind.Iterable:
+            length = self._IterableDataset_len_called = self.dataset.get_len(self.num_workers, self.batch_size)
+            if self.batch_size is not None:  # IterableDataset doesn't allow custom sampler or batch_sampler
+                from math import ceil
+
+                return length // self.batch_size if self.drop_last else ceil(length / self.batch_size)
+            return length
+        return len(self._index_sampler)
+
     def state_dict(self) -> Dict[str, Any]:
         if isinstance(self.dataset, StreamingDataset):
             assert self.batch_size

src/litdata/streaming/dataset.py

@@ -107,6 +107,8 @@ def __init__(
         self.shuffler: Optional[Shuffle] = None
         self.serializers = serializers
         self._state_dict: Optional[Dict[str, Any]] = None
+        self.num_workers: Optional[int] = None
+        self.batch_size: Optional[int] = None
 
     def set_shuffle(self, shuffle: bool) -> None:
         self.shuffle = shuffle
@@ -157,10 +159,16 @@ def _create_shuffler(self, cache: Cache) -> Shuffle:
         return FullShuffle(cache, seed, drop_last) if self.shuffle else NoShuffle(cache, seed, drop_last)
 
     def __len__(self) -> int:
+        return self.get_len(1, 1)
+
+    def get_len(self, num_workers: int, batch_size: int) -> int:
+        self.num_workers = num_workers
+        self.batch_size = batch_size
+        worker_env = _WorkerEnv.detect()
         if self.shuffler is None:
-            cache = self._create_cache(worker_env=_WorkerEnv.detect())
+            cache = self._create_cache(worker_env=worker_env)
             self.shuffler = self._create_shuffler(cache)
-        return self.shuffler.get_len(self.distributed_env, self.current_epoch)
+        return self.shuffler.get_len(self.distributed_env, num_workers, batch_size, self.current_epoch)
 
     def __iter__(self) -> "StreamingDataset":
         # When the StreamingDataset is used within map or optimize, let's refetch the distributed env.
@@ -178,7 +186,7 @@ def __iter__(self) -> "StreamingDataset":
             self.current_epoch = state["current_epoch"]
 
         chunks_per_replica, intervals_per_replica = self.shuffler.get_chunks_and_intervals_per_ranks(
-            self.distributed_env, self.current_epoch
+            self.distributed_env, self.worker_env.world_size, self.batch_size or 1, self.current_epoch
         )
         chunks_replica = chunks_per_replica[self.distributed_env.global_rank % self.distributed_env.world_size]
         intervals_replica = intervals_per_replica[self.distributed_env.global_rank % self.distributed_env.world_size]
@@ -187,10 +195,6 @@ def __iter__(self) -> "StreamingDataset":
         if self._state_dict:
             self._resume(chunks_replica, intervals_replica)
         else:
-            chunks_per_replica, intervals_per_replica = self.shuffler.get_chunks_and_intervals_per_ranks(
-                self.distributed_env, self.current_epoch
-            )
-
             # Find the chunks shared across multiple ranks.
             # For each shared chunk, find the rank to use the chunk last and prevent deletion
             # for the other ranks.

src/litdata/streaming/shuffle.py

@@ -31,8 +31,10 @@ def __init__(self, cache: Cache, seed: int, drop_last: bool):
         self.drop_last = drop_last
 
     @lru_cache(maxsize=10)
-    def get_len(self, distributed_env: _DistributedEnv, current_epoch: int) -> int:
-        _, intervals_per_ranks = self.get_chunks_and_intervals_per_ranks(distributed_env, current_epoch)
+    def get_len(self, distributed_env: _DistributedEnv, num_workers: int, batch_size: int, current_epoch: int) -> int:
+        _, intervals_per_ranks = self.get_chunks_and_intervals_per_ranks(
+            distributed_env, num_workers, batch_size, current_epoch
+        )
 
         if self.drop_last:
             items_per_process = [
@@ -46,7 +48,9 @@ def get_len(self, distributed_env: _DistributedEnv, current_epoch: int) -> int:
         return sum((interval[-1] - interval[0]) for interval in intervals_per_ranks[distributed_env.global_rank])
 
     @abstractmethod
-    def get_chunks_and_intervals_per_ranks(self, distributed_env: _DistributedEnv, current_epoch: int) -> Any:
+    def get_chunks_and_intervals_per_ranks(
+        self, distributed_env: _DistributedEnv, num_workers: int, batch_size: int, current_epoch: int
+    ) -> Any:
         pass
 
     @abstractmethod
@@ -59,14 +63,16 @@ class NoShuffle(Shuffle):
     is True."""
 
     @lru_cache(maxsize=10)
-    def get_chunks_and_intervals_per_ranks(self, distributed_env: _DistributedEnv, current_epoch: int) -> Any:
+    def get_chunks_and_intervals_per_ranks(
+        self, distributed_env: _DistributedEnv, num_workers: int, batch_size: int, current_epoch: int
+    ) -> Any:
         # 1. Get the intervals
         chunk_intervals = self.cache.get_chunk_intervals()
         indexes = range(len(chunk_intervals))
 
         # 2. Compute the items budget of each rank
         chunks_per_ranks, intervals_per_ranks = _associate_chunks_and_internals_to_ranks(
-            distributed_env, indexes, chunk_intervals, self.drop_last
+            distributed_env, indexes, chunk_intervals, self.drop_last, num_workers, batch_size
         )
 
         return chunks_per_ranks, intervals_per_ranks
@@ -94,7 +100,9 @@ class FullShuffle(Shuffle):
     """
 
     @lru_cache(maxsize=10)
-    def get_chunks_and_intervals_per_ranks(self, distributed_env: _DistributedEnv, current_epoch: int) -> Any:
+    def get_chunks_and_intervals_per_ranks(
+        self, distributed_env: _DistributedEnv, num_workers: int, batch_size: int, current_epoch: int
+    ) -> Any:
         # 1. Get the intervals
         chunk_intervals = self.cache.get_chunk_intervals()
 
@@ -113,7 +121,7 @@ def get_chunks_and_intervals_per_ranks(self, distributed_env: _DistributedEnv, c
 
         # 3. Compute the items budget of each rank
         chunks_per_ranks, intervals_per_ranks = _associate_chunks_and_internals_to_ranks(
-            distributed_env, shuffled_indexes, shuffled_chunk_intervals, self.drop_last
+            distributed_env, shuffled_indexes, shuffled_chunk_intervals, self.drop_last, num_workers, batch_size
         )
 
         # For the first epoch, no need of further shuffling
@@ -126,7 +134,7 @@ def get_chunks_and_intervals_per_ranks(self, distributed_env: _DistributedEnv, c
         shuffled_chunk_intervals = np.asarray(chunk_intervals)[shuffled_indexes].tolist()
 
         chunks_per_ranks, intervals_per_ranks = _associate_chunks_and_internals_to_ranks(
-            distributed_env, shuffled_indexes, shuffled_chunk_intervals, self.drop_last
+            distributed_env, shuffled_indexes, shuffled_chunk_intervals, self.drop_last, num_workers, batch_size
         )
 
         return chunks_per_ranks, intervals_per_ranks

src/litdata/utilities/shuffle.py

@@ -46,6 +46,8 @@ def _associate_chunks_and_internals_to_ranks(
     indexes: Any,
     chunk_intervals: Any,
     drop_last: bool,
+    num_workers: int = 1,
+    batch_size: int = 1,
 ) -> Tuple[List[List[int]], List[Any]]:
     num_items = sum([(interval[-1] - interval[0]) for interval in chunk_intervals])
     num_items_per_ranks: List[int] = [
@@ -54,6 +56,10 @@ def _associate_chunks_and_internals_to_ranks(
         else num_items // distributed_env.world_size
         for rank in range(distributed_env.world_size)
     ]
+    if drop_last:
+        ratio = num_workers * batch_size
+        num_items_per_ranks = [ratio * int(item // ratio) for item in num_items_per_ranks]
+
     chunks_per_ranks: List[List[int]] = [[] for _ in range(distributed_env.world_size)]
     intervals_per_ranks: List[List[List[int]]] = [[] for _ in range(distributed_env.world_size)]
 

tests/processing/test_data_processor.py

@@ -1208,7 +1208,7 @@ def fetch_from_dataset(batch, output_dir):
             f.write("Hello World!")
 
 
-@pytest.mark.skipif(sys.platform == "win32", reason="skip windows")
+@pytest.mark.skipif(sys.platform == "win32" or sys.platform == "darwin", reason="skip windows")
 def test_streaming_dataset_in_map(tmpdir):
     seed_everything(42)
 

tests/streaming/test_combined.py

@@ -262,7 +262,7 @@ def test_combined_dataset_with_dataloader_and_one_worker(batch_size):
     }
 
 
-@pytest.mark.skipif(sys.platform == "win32" or sys.platform == "darwin", reason="too slow in CI")
+@pytest.mark.skipif(sys.platform == "win32", reason="too slow in CI")
 def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
     data_dir_1 = os.path.join(tmpdir, "data_1")
     data_dir_2 = os.path.join(tmpdir, "data_2")

tests/streaming/test_dataset.py

@@ -171,7 +171,7 @@ def test_streaming_dataset_distributed_no_shuffle(drop_last, tmpdir, compression
     assert len(process_2_2) == 50 + int(not drop_last)
 
     _, intervals_per_ranks = dataset.shuffler.get_chunks_and_intervals_per_ranks(
-        dataset.distributed_env, dataset.current_epoch
+        dataset.distributed_env, 1, 1, dataset.current_epoch
     )
 
     assert process_1_1 == process_1_2