Make accessing WorkerInfo from within a DataPipe more convenient

🚀 The feature

import torchdata.datapipes as dp
from torch.utils.data.datapipes.iter.sharding import SHARDING_PRIORITIES
from torchdata.dataloader2 import MultiProcessingReadingService, DataLoader2

my_worker_info = None

def abc(x):
    return x * my_worker_info.worker_id

def worker_init(dp, worker_info):
    global my_worker_info 
    my_worker_info = worker_info
    return dp
    
    
pipe = dp.iter.IterableWrapper(range(10))
pipe = pipe.map(abc)
pipe = pipe.sharding_filter(SHARDING_PRIORITIES.MULTIPROCESSING)

rs = MultiProcessingReadingService(num_workers=2, worker_init_fn=worker_init)
dl = DataLoader2(pipe, reading_service=rs)

for x in dl:
    print(x)

Output:

0
1
0
3
0
5
0
7
0
9

This seems to be the only way to my knowledge to access the WorkerInfo from within a DataPipe when using Dataloader2. Global state is obviously awkward and becomes a problem for larger coebases that aren't toy examples. It would be good if there was a more convenient way (and also uniform way wrt Dataloader) a kin to get_worker_info.

Traversing the graph and calling set_worker_info if available would be a good option for this IMO.

Motivation, pitch

I want to easily access the current WorkerInfo from my datapipe.

Alternatives

No response

Additional context

No response

Start by tracing how DataLoader2 and MultiProcessingReadingService expose worker context to DataPipes, and compare that with DataLoader's get_worker_info API. Evaluate the proposed graph traversal and set_worker_info approach; done means DataPipes can access the current WorkerInfo without global state and the access is uniform with DataLoader.