Add comprehensive serialization tests

[ejguan]

🚀 The feature

We have serialization tests in PyTorch Core https://github.com/pytorch/pytorch/blob/b56ba296b1cb5d65a3fe2e33cc1d910481baa644/test/test_datapipe.py#L431

There are a few more things needed:

• Cover all DataPipe
• • Dill test for DataPipes that accept function as input
• Cover DataPipe serialization after iteration starts
• Cover serialization for factory DataPipe (1-to-N) like demux and fork, especially when one of children is not depleted

Motivation, pitch

These tests are going to make sure DataLoader graph can be created. And, for forward compatibility, snapshotting also requires DataPipe serializable.

Alternatives

No response

Additional context

No response

[NivekT]

I just added these PRs related to this issue:

pytorch/pytorch#71456
pytorch/pytorch#71497

Those additions should cover the basic serialization (first checkbox), but we need more work the remaining ones.

Another thing to add may be - if a DataPipe takes a function as argument, it should be serializable with dill?

[ejguan]

Another thing to add may be - if a DataPipe takes a function as argument, it should be serializable with dill?

Please correct me if I misunderstand it. I think the test has been added https://github.com/pytorch/pytorch/blob/f45e217c0158c247772f6491a18662cc972c29cd/test/test_datapipe.py#L445

[NivekT]

Another thing to add may be - if a DataPipe takes a function as argument, it should be serializable with dill?

Please correct me if I misunderstand it. I think the test has been added https://github.com/pytorch/pytorch/blob/f45e217c0158c247772f6491a18662cc972c29cd/test/test_datapipe.py#L445

Yep, the test has been added but we need to make sure all relevant DataPipes are added to the test.

[ejguan]

Another thing to add may be - if a DataPipe takes a function as argument, it should be serializable with dill?

Please correct me if I misunderstand it. I think the test has been added https://github.com/pytorch/pytorch/blob/f45e217c0158c247772f6491a18662cc972c29cd/test/test_datapipe.py#L445

Yep, the test has been added but we need to make sure all relevant DataPipes are added to the test.

Sounds great.

[NivekT]

Note that pytorch/pytorch#73119 addresses these requirements, we still need to do the same for the DataPipes in TorchData.

Also, see #106 for individual DataPipe's test coverage.

[ejguan]

Closing this issue as we rely on #106 to track test coverage.