[DataPipe] Add RoundRobinDemux

[ejguan]

The only reason I added this DataPipe without directly using enumereate().demux().drop_index() is this RoundRobinDemux should provide a valid length. So, this PR needs pytorch/pytorch#89216 landed

And, this DataPipe will be used for Proto2RS.

Side change: Move Unzip to combining.py as well.

[ejguan]

Failed tests should be fixed the the PR in PyTorch Core landed.

[facebook-github-bot]

@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[ejguan]

Failed tests should be fixed by #905

[wenleix]

LGTM % some noob/minor comments

[wenleix]

curious: why demux is PyTorch core (https://github.com/pytorch/pytorch/blob/7a2930b357a4e62bb0bab53bb0d23c607b6ede38/torch/utils/data/datapipes/iter/combining.py#L305-L306) but round_robin_demux is in TorchData?

[ejguan]

Historical reason. We put lots of DataPipe in core but we plan to move them to torchdata eventually. To reduce the amount of work, I think it's better to put round_robin_demux to torchdata directly.

[wenleix]

As a Python noob, TIL I learnt X(...) in Python may not return a instance of X

https://www.geeksforgeeks.org/__new__-in-python/

... because other class constructors can be called by new method or it can simply return other objects as an instance of this class.

Will _ChildDataPipe.__init__ be called for each instance in the list? ~

[ejguan]

Yeah.

[wenleix]

curious: can we just do

_DemultiplexerIterDataPipe(
    datapipe, 
    num_instances, 
    partial(num_instances, _round_robin_fn), 
    False, 
    buffer_size
)

where _round_robin_fn is:

def _round_robin_fn(num_instances: int, idx_data) -> int:
        idx, _ = idx_data
        return idx % num_instances

Is that because you want to override _DemultiplexerIterDataPipe.get_length_by_instance ?

[ejguan]

Yeah. demux doesn't provide a valid length but this DataPipe should have a deterministic result. Let me know if you have better way to do so.

[wenleix]

noob question : what's "deterministic result" here mean?

Especially I was somehow under the impression that in general try to avoid query len(datapipe) for large datapipe , since it probably needs to read all data and it's expensive?

[ejguan]

If the prior DataPipe provides a length, we should be able to get the length ahead of time.
For demux which uses a classify_fn, there is no way for us to know the length and the length can vary across epoches.

in general try to avoid query len(datapipe) for large datapipe , since it probably needs to read all data and it's expensive?

It's true. But, we do provide DataPipe that can manually assign length to the DataPipe

data/torchdata/datapipes/iter/util/header.py

Line 62 in 67afe0c

@functional_datapipe("set_length")

[wenleix]

Unrelated question: Why demux and unzip data pipes are in combining.py ? (shouldn't they be in splitting.py)?

[ejguan]

Why demux and unzip data pipes are in combining.py ? (shouldn't they be in splitting.py)?

Aha. We categorized DataPipe into a few sectors. See: https://pytorch.org/data/beta/index.html We put them together as they are the counter party of mux, zip, etc., and we don't have such a mount of DataPipes.

When #293 is done, we probably need to make them into combining.py and splitting.py.

[facebook-github-bot]

@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@ejguan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.