Apply sharding based on priority & combine DistInfo and ExtraInfo

[ejguan]

After pytorch/pytorch#88424 is landed, we are able to invoke apply_sharding by sharding levels (distributed or multiprocessing). Then, we are able to give fine-control on sharding by ReadingService.

• For DistributedReadingService, we will only set sharding on the distributed level
• For PrototypeMPReadingService, we will set distributed sharding in the main process and set mp sharding in the worker processes. Previously, we set sharding in each worker process based on both distributed and mp information.
  • worker_init_fn doesn't need DistInfo anymore. As, the DataPipe has been distributed sharded in the main process.
  • Combine DistInfo and ExtraInfo for worker_reset_fn to synchronize the distributed seeds across distributed workers and set worker-local seeds based on both distributed and mp information.

[facebook-github-bot]

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

[NivekT]

I did not know it was possible to land without GH stamp.

[ejguan]

@NivekT It's only possible for TorchData. PyTorch doesn't allow this anymore due to the GH first project