fix(sampling): preserve remainder samples across shards

swift/pipelines/sampling/sampling.py

@@ -43,8 +43,10 @@ def _get_dataset(self):
             args.dataset, split_dataset_ratio=0., shuffle=args.dataset_shuffle, **dataset_kwargs)
         logger.info(f'Sampling_dataset: {sampling_dataset}')
         dataset_len = len(sampling_dataset)
-        piece_len = dataset_len // self.total_piece
-        sampling_dataset = sampling_dataset.select(range(piece_len * self.cur_piece, piece_len * (self.cur_piece + 1)))
+        piece_len, remainder = divmod(dataset_len, self.total_piece)
+        start = piece_len * self.cur_piece + min(self.cur_piece, remainder)
+        end = start + piece_len + int(self.cur_piece < remainder)
+        sampling_dataset = sampling_dataset.select(range(start, end))
         return sampling_dataset
 
     def run(self):

tests/sample/test_sampling.py

@@ -0,0 +1,47 @@
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from swift.pipelines.sampling.sampling import SwiftSampling
+
+
+class FakeDataset:
+
+    def __init__(self, values):
+        self.values = values
+
+    def __len__(self):
+        return len(self.values)
+
+    def select(self, indices):
+        return FakeDataset([self.values[i] for i in indices])
+
+
+class TestSampling(unittest.TestCase):
+
+    @patch('swift.pipelines.sampling.sampling.load_dataset')
+    def test_data_range_partitions_cover_dataset(self, mock_load_dataset):
+        args = SimpleNamespace(
+            dataset=['test-dataset'],
+            dataset_shuffle=False,
+            get_dataset_kwargs=lambda: {},
+        )
+
+        for dataset_size, expected_sizes in [(10, [4, 3, 3]), (2, [1, 1, 0])]:
+            with self.subTest(dataset_size=dataset_size):
+                dataset = FakeDataset(list(range(dataset_size)))
+                mock_load_dataset.return_value = dataset, None
+                shards = []
+                for shard_index in range(3):
+                    sampling = SwiftSampling.__new__(SwiftSampling)
+                    sampling.args = args
+                    sampling.cur_piece = shard_index
+                    sampling.total_piece = 3
+                    shards.append(sampling._get_dataset().values)
+
+                self.assertEqual([len(shard) for shard in shards], expected_sizes)
+                self.assertEqual([item for shard in shards for item in shard], list(range(dataset_size)))
+
+
+if __name__ == '__main__':
+    unittest.main()