Add support for grain.IterDataset in sampling (apple#980)

Author: RsEnts

axlearn/common/input_grain.py

@@ -131,7 +131,7 @@ def sample_from_datasets(
     sources: Sequence[Dataset],
     weights: Sequence[float],
 ) -> Dataset:
-    """Mixes one or more data sources.
+    """Mixes one or more repeated data sources.
 
     Different from `input_tf_data.sample_from_datasets`, the mixing is deterministic:
     https://github.com/google/grain/blob/ddf825c68b6d2c811f9e599d7fb7ae7572affd8c/grain/_src/python/dataset/transformations/mix.py#L222
@@ -148,22 +148,24 @@ def sample_from_datasets(
         A Dataset for the mixed data source.
     """
 
-    # Without repeat, mixing stops as soon as the first dataset is exhausted.
-    def maybe_repeat(ds: Dataset):
-        if not isinstance(ds, grain.MapDataset):
-            raise ValueError(
-                f"{sample_from_datasets.__name__} requires {grain.MapDataset.__name__}"
-            )
-        # Only repeat if not already infinite.
-        if len(ds) != sys.maxsize:
-            ds = ds.repeat()
-        return ds
-
-    # TODO(markblee): Support mixing grain.IterDataset.
-    return grain.MapDataset.mix(
-        datasets=[maybe_repeat(source) for source in sources],
-        weights=weights,
-    )
+    def _ensure_repeated(sources: Sequence[Dataset]):
+        # There is no easy way to check if a grain.IterDataset is repeated.
+        for source in sources:
+            if isinstance(source, grain.MapDataset) and len(source) != sys.maxsize:
+                raise ValueError(
+                    f"sample_from_datasets requires each dataset to be repeated, {source} is not."
+                )
+            if isinstance(source, grain.IterDataset):
+                logging.info(
+                    "Sampling from grain.IterDataset, please make sure your dataset is repeated."
+                )
+
+    _ensure_repeated(sources)
+    # If any of the datasets are grain.IterDataset, we should use grain.IterDataset.mix().
+    if any(isinstance(ds, grain.IterDataset) for ds in sources):
+        return grain.IterDataset.mix(datasets=sources, weights=weights)
+
+    return grain.MapDataset.mix(datasets=sources, weights=weights)
 
 
 def default_pad_example_fn(example: utils.Nested[Any]) -> utils.Nested[Any]:

axlearn/common/input_grain_test.py

@@ -85,37 +85,65 @@ def check(starting_step, ds):
         dict(
             sources=[slice(0, 10, 2), slice(1, 5, 2)],
             weights=[1, 1],
+            is_iter_dataset=[False, False],
             take=10,
             expected=[0, 1, 2, 3, 4, 1, 6, 3, 8, 1],
         ),
         dict(
             sources=[slice(0, 10, 2), slice(1, 5, 2)],
             weights=[2, 1],
+            is_iter_dataset=[False, False],
             take=10,
             expected=[0, 2, 1, 4, 6, 3, 8, 0, 2, 1],
         ),
         dict(
             sources=[slice(0, 10, 2), slice(1, 5, 2)],
             weights=[1, 1e-9],
+            is_iter_dataset=[False, False],
             take=10,
             expected=[0, 2, 4, 6, 8, 0, 2, 4, 6, 8],
         ),
+        # IterDataset
+        dict(
+            sources=[slice(0, 10, 2), slice(1, 5, 2)],
+            weights=[1, 1],
+            is_iter_dataset=[True, True],
+            take=10,
+            expected=[0, 1, 2, 3, 4, 1, 6, 3, 8, 1],
+        ),
+        # Mixture of IterDataset and MapDataset.
+        dict(
+            sources=[slice(0, 10, 2), slice(1, 5, 2)],
+            weights=[1, 1],
+            is_iter_dataset=[True, False],
+            take=10,
+            expected=[0, 1, 2, 3, 4, 1, 6, 3, 8, 1],
+        ),
     )
     def test_sample_from_datasets(
         self,
         sources: list[slice],
         weights: list[int],
+        is_iter_dataset: list[bool],
         take: Optional[int],
         expected: list[int],
     ):
+        sources = [
+            range_dataset(start=src.start, stop=src.stop, step=src.step).repeat() for src in sources
+        ]
+        sources = [
+            source.to_iter_dataset() if should_convert else source
+            for source, should_convert in zip(sources, is_iter_dataset)
+        ]
         ds = sample_from_datasets(
-            sources=[
-                range_dataset(start=src.start, stop=src.stop, step=src.step) for src in sources
-            ],
+            sources=sources,
             weights=weights,
         )
-        ds = ds.slice(slice(0, take))
-        self.assertCountEqual(expected, list(ds))
+        ds_iter = iter(ds)
+        result = []
+        for _ in range(take):
+            result.append(next(ds_iter))
+        self.assertCountEqual(expected, list(result))
 
     def test_sample_from_datasets_errors(self):
         ds = range_dataset(start=0, stop=2)
@@ -124,17 +152,12 @@ def test_sample_from_datasets_errors(self):
         repeated_ds = sample_from_datasets(sources=[ds], weights=[1]).slice(slice(0, 4))
         self.assertEqual([0, 1, 0, 1], list(repeated_ds))
 
-        # Make sure that non-map dataset raises.
-        with self.assertRaisesRegex(ValueError, "MapDataset"):
-            ds = ds.to_iter_dataset()
-            sample_from_datasets(sources=[ds], weights=[1])
-
     def test_shuffle_dataset(self):
         # Test without repeat.
         ds = sample_from_datasets(
             sources=[
-                range_dataset(start=0, stop=10, step=2),
-                range_dataset(start=1, stop=5, step=2),
+                range_dataset(start=0, stop=10, step=2).repeat(),
+                range_dataset(start=1, stop=5, step=2).repeat(),
             ],
             weights=[2, 1],
         )
@@ -174,7 +197,7 @@ def test_slice_dataset(self, s: slice, expected: list[int]):
 
     def test_batch(self):
         # [0, 1, 2, 3, 4].
-        ds = range_dataset(start=0, stop=5, seed=123)
+        ds = range_dataset(start=0, stop=5, seed=123).repeat()
         # [1, 2, 3, 4, 5].
         other_ds = ds.map(_PlusOne())
         # [0, 1, 2, 1, 3, 4, 2, 5, 1, 3, ...].