Dataloader does not stop if one of the zipped DataPipes has a perpetual cycle

[magehrig]

🐛 Describe the bug

The Zipper is supposed to stop as soon as the shortest input DataPipe is exhausted. This is not the case if one of the DataPipes is constructed with a perpetual cycle.

Minimal Example:

from torch.utils.data import DataLoader
from torchdata.datapipes.iter import IterableWrapper, IterDataPipe


class TestDataset(IterDataPipe):
    def __init__(self):
        self.data = IterableWrapper(list(range(10)))

    def __iter__(self):
        inf_cycle = IterableWrapper([42]).cycle(count=None)
        datapipe = self.data.zip(inf_cycle)
        return iter(datapipe)


if __name__ == '__main__':
    dataset = TestDataset()
    loader = DataLoader(dataset=dataset, batch_size=None, num_workers=0)

    for batch in loader:
        print(batch)

Prints the following but does not exit the program:

[0, 42]
[1, 42]
[2, 42]
[3, 42]
[4, 42]
[5, 42]
[6, 42]
[7, 42]
[8, 42]
[9, 42]

Versions

PyTorch version: 1.14.0.dev20221026+cpu
Is debug build: False
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~18.04) 9.4.0
Clang version: 8.0.0-3~ubuntu18.04.2 (tags/RELEASE_800/final)
CMake version: version 3.20.0
Libc version: glibc-2.27

Python version: 3.9.13 | packaged by conda-forge | (main, May 27 2022, 16:56:21)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-4.15.0-194-generic-x86_64-with-glibc2.27
Is CUDA available: False
CUDA runtime version: 11.1.105
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration:
GPU 0: Quadro RTX 8000
GPU 1: Quadro RTX 8000
GPU 2: Quadro RTX 8000
GPU 3: Quadro RTX 8000
GPU 4: Quadro RTX 8000
GPU 5: Quadro RTX 8000
GPU 6: Quadro RTX 8000
GPU 7: Quadro RTX 8000

Nvidia driver version: 515.65.01
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.22.4
[pip3] torch==1.14.0.dev20221026+cpu
[pip3] torchdata==0.6.0.dev20221026
[pip3] torchvision==0.13.1
[conda] blas                      2.116                       mkl    conda-forge
[conda] blas-devel                3.9.0            16_linux64_mkl    conda-forge
[conda] cudatoolkit               10.2.89             h713d32c_10    conda-forge
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libblas                   3.9.0            16_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            16_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            16_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            16_linux64_mkl    conda-forge
[conda] magma                     2.5.4                h5da55e3_2    conda-forge
[conda] mkl                       2022.1.0           h84fe81f_915    conda-forge
[conda] mkl-devel                 2022.1.0           ha770c72_916    conda-forge
[conda] mkl-include               2022.1.0           h84fe81f_915    conda-forge
[conda] numpy                     1.22.4           py39hc58783e_0    conda-forge
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch                     1.14.0.dev20221026+cpu          pypi_0    pypi
[conda] torchdata                 0.6.0.dev20221026          pypi_0    pypi
[conda] torchvision               0.13.1               py39_cu102    pytorch

[ejguan]

Thank you for opening this issue. Confirmed this is bug for zip operation, which will deplete all DataPipe for the sake of buffer clean up and infinite DataPipe will prevent zip exiting.
https://github.com/pytorch/pytorch/blob/aaba0bd30641c56db1dc0550b81fbc458db46276/torch/utils/data/datapipes/iter/combining.py#L545-L548

cc: @VitalyFedyunin I am not sure how to let zip know if prior DataPipe is an infinite one. I would imagine there is going to be the similar problem for streaming or online training case.

[magehrig]

Version 0.4.1 does not seem to have that bug

[ejguan]

Yeah. It's introduced after 0.4 to clean up remaining opened file streams in the DataPipes within zip.

As discussed with Vitaly, we think it's better to add a method for DataPipe to clean up itself rather than relying on zip to do the job. And, zip should invoke the corresponding clean-up method for each DataPipe in the finally clause.

[NivekT]

I think we will need to change the clean up for zip in particular as well. Calling list(dp) on infinite DataPipe will not work. Maybe some sort of time out or limit?

[ejguan]

Minimum repro script:

dp = IterableWrapper(list(range(10)))
inf_dp = IterableWrapper(list(range(10))).cycle(count=None)
dp = dp.zip(inf_dp)
for d in dp:
    print(d)

[peterdudfield]

+1 to this please

could the clean have an options to be turned off e.g. (this is in here)

    def __init__(self, *datapipes: IterDataPipe, clean_up: Optional[bool] = True):
        if not all(isinstance(dp, IterDataPipe) for dp in datapipes):
            raise TypeError("All inputs are required to be `IterDataPipe` "
                            "for `ZipIterDataPipe`.")
        super().__init__()
        self.datapipes = datapipes  # type: ignore[assignment]
        self.length = None
        self.clean_up = clean_up

        
    def __iter__(self) -> Iterator[Tuple[T_co]]:
        iterators = [iter(datapipe) for datapipe in self.datapipes]
        try:
            for data in zip(*iterators):
                yield data
        finally:
            if clean_up:
                 unused = []
                for iterator in iterators:
                    try:
                        unused += list(iterator)
                    except RuntimeError:  # Some iterators may have been invalidated by single iterator constraints
                        pass

                # TODO(VitalyFedyunin): This should be Exception or warning when torchdata.debug is enabled
                for item in unused:
                    StreamWrapper.close_streams(item)

I can make a PR for this if you like @ejguan

[ejguan]

@peterdudfield Thanks for trying to fix it.
To prevent such problem happening, I think a method like cleanup is required for DataPipe. I will send a PR to fix it.