[Datasets] Iterating through DatasetPipeline fails with ZeroDivisionError

[amogkam]

What happened + What you expected to happen

I expect the below code snippet to not fail.

Instead it fails with the following error:

self = <ray.data.dataset_pipeline.DatasetPipeline.repeat.<locals>.RepeatIterator object at 0x16a801a30>

    def __next__(self) -> Dataset[T]:
        # Still going through the original pipeline.
        if self._original_iter:
            try:
                make_ds = next(self._original_iter)
                self._results.append(make_ds)

                def gen():
                    res = make_ds()
                    res._set_epoch(0)
                    return res

                return gen
            except StopIteration:
                self._original_iter = None
                # Calculate the cursor limit.
                if times:
                    self._max_i = len(self._results) * (times - 1)
                else:
                    self._max_i = float("inf")
        # Going through a repeat of the pipeline.
        if self._i < self._max_i:
>           make_ds = self._results[self._i % len(self._results)]
E           ZeroDivisionError: integer division or modulo by zero

Versions / Dependencies

master

Reproduction script

pipe = ray.data.range(6, parallelism=6).window(blocks_per_window=2).repeat()
assert pipe.schema() == int
pipe = pipe.map_batches(lambda x: x)
next(pipe.iter_epochs())

Issue Severity

High: It blocks me from completing my task.

[jianoaix]

A mitigation is switching the order of pipe.schema() and pipe.map_batches():

pipe = ray.data.range(6, parallelism=6).window(blocks_per_window=2).repeat()
pipe = pipe.map_batches(lambda x: x)
assert pipe.schema() == int
next(pipe.iter_epochs())

The issue seems to be the peek() (used in pipe.schema()) advanced the base_iterable of pipeline, which was then used to create another pipeline (with pipe.map_batches()).

[amogkam]

In this case, we're using the schema to determine what batch format to use for the map_batches call so I won't be able to switch the order.