Fix: Resolve drop_last not passed down from the StreamingDataLoader to the datasets (#147)

Author: tchaton

src/litdata/streaming/combined.py

@@ -117,6 +117,11 @@ def set_shuffle(self, shuffle: bool) -> None:
         for dataset in self._datasets:
             dataset.set_shuffle(shuffle)
 
+    def set_drop_last(self, drop_last: bool) -> None:
+        """Set the current drop_last to the datasets."""
+        for dataset in self._datasets:
+            dataset.set_drop_last(drop_last)
+
     def _check_datasets(self, datasets: List[StreamingDataset]) -> None:
         if any(not isinstance(d, StreamingDataset) for d in datasets):
             raise RuntimeError("The provided datasets should be instances of the StreamingDataset.")

src/litdata/streaming/dataloader.py

@@ -540,6 +540,7 @@ def __init__(
         profile_dir: Optional[str] = None,
         prefetch_factor: Optional[int] = None,
         shuffle: Optional[bool] = None,
+        drop_last: Optional[bool] = False,
         **kwargs: Any,
     ) -> None:  # pyright: ignore
         if not isinstance(dataset, (StreamingDataset, CombinedStreamingDataset)):
@@ -551,6 +552,9 @@ def __init__(
         if shuffle is not None:
             dataset.set_shuffle(shuffle)
 
+        if drop_last is not None:
+            dataset.set_drop_last(drop_last)
+
         shuffle = None
 
         if profile_batches and not _VIZ_TRACKER_AVAILABLE:

src/litdata/streaming/dataset.py

@@ -113,6 +113,9 @@ def __init__(
     def set_shuffle(self, shuffle: bool) -> None:
         self.shuffle = shuffle
 
+    def set_drop_last(self, drop_last: bool) -> None:
+        self.drop_last = drop_last
+
     def set_epoch(self, current_epoch: int) -> None:
         """Set the current epoch to the dataset on epoch starts.
 

status.json

@@ -1,6 +1,6 @@
 import os
 import sys
-from unittest.mock import ANY
+from unittest.mock import ANY, MagicMock
 
 import pytest
 import torch
@@ -53,6 +53,19 @@ def test_combined_dataset_num_samples_yield_iterate_over_all():
     assert len(samples) == 20
 
 
+def test_drop_last_and_shuffle():
+    dataset_mock_1 = MagicMock()
+    dataset_mock_2 = MagicMock()
+
+    dataset = TestCombinedStreamingDataset([dataset_mock_1, dataset_mock_2], 42, iterate_over_all=True)
+    StreamingDataLoader(dataset, shuffle=True, drop_last=True)
+
+    dataset_mock_1.set_shuffle.assert_called()
+    dataset_mock_2.set_shuffle.assert_called()
+    dataset_mock_1.set_drop_last.assert_called()
+    dataset_mock_2.set_drop_last.assert_called()
+
+
 class TestStatefulDataset:
     def __init__(self, size, step):
         self.size = size
@@ -177,6 +190,12 @@ def state_dict(self, **kwargs):
     def set_epoch(self, current_epoch):
         pass
 
+    def set_shuffle(self, _):
+        pass
+
+    def set_drop_last(self, _):
+        pass
+
 
 def test_combined_dataset():
     dataset1 = SimpleDataset(0, 10)

tests/streaming/test_combined.py

@@ -14,6 +14,7 @@ def __init__(self, size, step):
         self.step = step
         self.counter = 0
         self.shuffle = None
+        self.drop_last = None
 
     def set_shuffle(self, shuffle):
         self.shuffle = shuffle
@@ -41,6 +42,9 @@ def load_state_dict(self, state_dict):
     def set_epoch(self, current_epoch):
         pass
 
+    def set_drop_last(self, drop_last):
+        self.drop_last = drop_last
+
 
 class TestCombinedStreamingDataset(CombinedStreamingDataset):
     def _check_datasets(self, datasets) -> None: