Graph traversal is broken for custom iter datapipes

[pmeier]

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


class CustomIterDataPipe(IterDataPipe):
    def noop(self, x):
        return x

    def __init__(self):
        self._dp = IterableWrapper([]).map(self.noop)

    def __iter__(self):
        yield from self._dp


traverse(CustomIterDataPipe())

RecursionError: maximum recursion depth exceeded

Without the .map() call it works fine. I don't think this is specific to .map() though. From trying a few datapipes, this always happens if self._dp is composed in some way.

[ejguan]

I believe this is due to the function noop is owned by CustomIterDataPipe. When we call map, MapperIterDataPipe would take reference of CustomIterDataPipe, and CustomIterDataPipe takes reference by self._dp, we are experiencing cyclic reference here.
'As a workaround, we can move noop out of the CustomIterDataPipe. But, the solution would be changing traverse function to prevent cyclic serialization along the same datapipe path (needs a temporary hash table for the Datapipe on the same path)

[pmeier]

One of the reasons we are looking into this architecture is to be able to access instance attributes in these functions. Thus, moving the methods to functions would somewhat make the implementation more verbose. cc @NicolasHug

[ejguan]

Another way would be changing it to staticmethod as this function doesn't really use self.

[NicolasHug]

Thanks for the info @ejguan

Another way would be changing it to staticmethod as this function doesn't really use self.

In general, we'll be using non-static functions that rely on arbitrary self.xyz attributes, so static methods aren't an option for us. We can pass these to functions, for some extra verbosity as @pmeier mentioned.

So that we can plan on how to best move forward on our side, how long do you think a fix would take @ejguan ?

[ejguan]

Mapper ----> CustomIterDataPipe.noop
CustomIterDataPipe -----> Mapper

pickle itself understands the circular references. traverse should behave similarly. We will fix it anyway.

[ejguan]

So that we can plan on how to best move forward on our side, how long do you think a fix would take @ejguan ?

It will be my top priority next week. I need to prepare all release-related stuff this week.

[ejguan]

I was able to provide a work around for traverse function to provide graph to you. But, I found a problem if I let traverse support this behavior.
When multiprocessing step in, pickle is used to serialize/deserialize such object. I originally thought pickle was able to do the serialization even with the circular reference. It turned out pickle didn't work in such case like the example below.

dp = CustomIterDataPipe()
_ = pickle.dumps(dp)  # RecursionError: maximum recursion depth exceeded while calling a Python object

In order to provide solution for it, please correct me if I am wrong about the requirements:

• We want to access these methods as attributes in the output DataPipe (Or you want to have these methods attached to the factory class?)
• These functions can not be static or class method sinice it would depends on other attributes in self
• These functions are used to construct prior DataPipe instances
  In order to achieve so, we have to isolate these methods from the factory function/class:
• Option 1

class _Dataset:
    def noop(self, x):
        return x

class Dataset(IterDataPipe):
    def __init__(self, *args, **kwargs):
        obj = _Dataset(...)
        self._dp = IterableWrapper([]).map(obj.noop)
        self.noop = obj.noop

    def __iter__(self):
        yield from self._dp

• Option 2

class Dataset:
    @classmethod
    def build(cls):
        ds = cls()
        return CustomIterDataPipe(ds)

    def noop(self, x):
        return x

class CustomIterDataPipe(IterDataPipe):
    def __init__(self, ds):
        self._dp = dp.iter.IterableWrapper([]).map(ds.noop)
        self.noop = ds.noop

    def __iter__(self):
        yield from self._dp

cc: @pmeier @NicolasHug

[pmeier]

pickle automatically detects duplicate items and handles them without a recursion error. The problem arises, because we have our custom serialization logic. With this patch

diff --git a/torch/utils/data/datapipes/datapipe.py b/torch/utils/data/datapipes/datapipe.py
index 0a90cec8d6..f1dd0ab88b 100644
--- a/torch/utils/data/datapipes/datapipe.py
+++ b/torch/utils/data/datapipes/datapipe.py
@@ -93,13 +93,7 @@ class IterDataPipe(IterableDataset[T_co], metaclass=_DataPipeMeta):
     def __getstate__(self):
         if IterDataPipe.getstate_hook is not None:
             return IterDataPipe.getstate_hook(self)
-        state_dict = {}
-        for k, v in self.__dict__.items():
-            if callable(v):
-                state_dict[k] = serialize_fn(v)
-            else:
-                state_dict[k] = v
-        return state_dict
+        return self.__dict__
 
     def __setstate__(self, state_dict):
         for k, v in state_dict.items():

serialization works as expected:

import pickle

from torchdata.datapipes.iter import IterDataPipe, IterableWrapper


class CustomIterDataPipe(IterDataPipe):
    def noop(self, x):
        return x + 1

    def __init__(self):
        self._dp = IterableWrapper([1, 2, 4]).map(self.noop)

    def __iter__(self):
        yield from self._dp


assert list(CustomIterDataPipe()) == list(pickle.loads(pickle.dumps(CustomIterDataPipe())))

The logic I patched above was introduced in pytorch/pytorch#72896 cc @NivekT.

[ejguan]

Wondering what would be the visualized graph with such circular reference?

[ejguan]

Yeah. We might need to use a different method to handle dispatch pickling method for callable inside IterDataPipe.
Using copyreg might be an option here (https://docs.python.org/3/library/copyreg.html#module-copyreg). But, we need to make the sure after pickling IterDataPipe, the pickling method for callable function is reset to original ones. cc: @NivekT

[ejguan]

Actually, I am not 100% sure this issue can be resolved by using copyreg. The way we do serialization may need to change. When dill is not available, we use self.__dict__ directly. If not available, we might be able to use the protocol you designed. @NivekT

We could add a test for such circular reference DataPipe and make sure the __reduce__ and other pickle-related functions working as expected.