Main process hanging on TPU Pod

[Liyang90]

Bug description

On a TPU pod with more than 1 hosts, the main process on non-master hosts hangs at https://github.com/Lightning-AI/lightning/blob/beced489040f76e7eee2f4a82d29823834b77327/src/lightning/pytorch/strategies/launchers/xla.py#L89
The root cause is the process_id is used as local_rank, but it is actually a global rank id: https://github.com/Lightning-AI/lightning/blob/beced489040f76e7eee2f4a82d29823834b77327/src/lightning/pytorch/strategies/launchers/xla.py#L106

Because of that, on non-master hosts, the spawned processes don't put anything in the return_queue: https://github.com/Lightning-AI/lightning/blob/beced489040f76e7eee2f4a82d29823834b77327/src/lightning/pytorch/strategies/launchers/xla.py#L112-L113
so the main process hangs when trying to get item from the queue.

A potential fix would be to use cluster_environment to provide local_rank instead.

@JackCaoG @steventk-g @will-cromar

How to reproduce the bug

No response

Error messages and logs

# Error messages and logs here please

Environment

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):

More info

No response

cc @JackCaoG @steventk-g @Liyang90 @carmocca @justusschock @awaelchli

[awaelchli]

That's a nice find, thanks! Didn't know it is actually the global rank, because in torch.multiprocessing it would be the local rank.