Is it possible to access the global rank and world size outside of `LightningModule`?
Hi,
I'm using an IterableDataset and need to manually split up the data depending on the GPU (I'm using DDP). The lightning docs say that replace_sampler_ddp doesn't automatically handle IterableDatasets, so I need to do this myself.
Is it possible to access the global rank or world size in a dataset (aka are there any utility functions available for this purpose)? I saw _get_rank here, but ideally I wouldn't use a private helper function.
I tried using:
file_paths_filtered = file_paths[
    torch.distributed.get_rank() :: torch.distributed.get_world_size()
]

inside the IterableDataset, but it fails because the default process group has not been setup yet (which makes sense):
RuntimeError: Default process group has not been initialized, please make sure to call init_process_group.

Is there any way to access the global rank/world size? If not, how would you recommend I check which GPU the dataset is on without causing an issue with init_process_group not being called yet (should I maybe call this manually)?
Issue ended up being torch.distributed.get_rank() and torch.distributed.get_world_size() aren't available when on a single GPU, which I was testing on locally before sending to multi-GPU server.
The following works:
if torch.distributed.is_available() and torch.distributed.is_initialized():
    file_paths = file_paths[
        torch.distributed.get_rank() :: torch.distributed.get_world_size()
    ]