[1/n] Add graph function to list DataPipes from graph

[ejguan]

Add a list_dps function to list DataPipes from the graph.

• It's similar to get_all_graph_pipes from pytorch core
• An extra argument of exclude_dps to exclude the DataPipe and its prior graph from the result.

Reason to add this function:

• It's required to set random states differently for DataPipe before/after sharding_filter

graph = traverse_dps(datapipe)
sf_dps = find_dps(graph, ShardingFilter)

# DataPipes prior to `sharding_filter`
p_dps = []
for sf_dp in sf_dps:
    p_dps.extend(list_dps(traverse_dps(sf_dp)))

# DataPipes after `sharding_filter`
a_dps = list_dps(graph, exclude_dps=sf_dps)

Step 1 for #885

[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.

[NivekT]

LGTM!

[NivekT]

Side note, this doesn't guarantee any order and returns a list. This should suffice for the current use case.

I wonder if people will want a Dict (basically just traverse_dps without the IDs). To be clear, this seems fine to me for now.

[ejguan]

If we treat the graph as a tree, there can be in different orders like how do we traverse a tree. Pre-order; post-order, etc. It highly depends on users' preference especially our graph can be more complicated than a tree. I would say let's wait until we have received any use case.

[facebook-github-bot]

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