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.

[ejguan]

I can take a look after I finished landing S3

[NivekT]

I can take a look after I finished landing S3

FYI, this PR is work in progress (basically add graph traversal support for MapDataPipe), but feel free to propose how default serialization should be done differently and I will update the other PR to match what we end up deciding here.

• [DataPipe] Enabling graph traversal for MapDataPipe pytorch#74851

[NicolasHug]

@ejguan @NivekT , going back to @pmeier 's #237 (comment):

• Does torchdata have to support dill?
• If not, would we consider merging Philip's proposed diff for fixing the cyclic reference issue?
• If yes, would you mind providing some details on why dill is needed? Hopefully we can still find a fix for the cyclic reference issue while still supporting dill?

Fixing the cyclic reference issue would allow us to move forward with our preferred design for torchvision new datasets.

[ejguan]

Does torchdata have to support dill?

No. We can revert the custom serialization code to unblock you. cc: @NivekT about if reverting will cause BC breaking?

[NivekT]

The most common use case for dill is the serialization of lambda and similar functions. If we don't support it, any attempt to serialize a DataPipe with a lambda function will fail since pickle can't serialize lambda functions.

For context, prior to that PR, we also used dill to serialize. The difference is that the operations were performed within __getstate__ and __setstate__ of specific DataPipes, and the PR moved that logic into the class IterDataPipe.

Without dill, test_serializable_with_dill will fail in both PyTorch Core and TorchData.

https://github.com/pytorch/pytorch/blob/master/test/test_datapipe.py#L645
https://github.com/pytorch/data/blob/main/test/test_serialization.py#L301

I am not sure how many users use dill and lambda. I would guess the impact on UX is small but it is BC-breaking.

[pmeier]

IIUC, the problem arises only because we are currently trying to support both, correct? Would it be possible to hard depend on dill only always use it? This should also solve the issue and being able to use lambda's and local functions would make our code a lot cleaner.

[NivekT]

I think adding dill as a hard dependency is even more BC-breaking. I'm adding a PR to only use the custom serialization when dill is available. This should mean that your code snippet above will work if dill is not installed.

[NicolasHug]

Thanks @NivekT

This should mean that your code snippet above will work if dill is not installed.

This means that the cyclic reference issue is still happening for users that have dill installed, right? Unfortunately I'm not sure we can expect torchvision's users to not have installed dill.

Do you think there is a way to fix the dill serialization to properly handle the cyclic ref issue?

[NivekT]

Thanks @NivekT

This should mean that your code snippet above will work if dill is not installed.

This means that the cyclic reference issue is still happening for users that have dill installed, right? Unfortunately I'm not sure we can expect torchvision's users to not have installed dill.

Do you think there is a way to fix the dill serialization to properly handle the cyclic ref issue?

Yep, the PR is meant to be a quick fix to unblock your work and I marked that as a TODO. I agree with your point and am looking into it.

[ejguan]

This PR should resolve the recursion problem of traverse pytorch/pytorch#75034

But, for the problem of dill, we still need to fix it.

[NivekT]

I believe this is fixed by pytorch/pytorch#74984