Repository: Lightning-AI/pytorch-lightning
Pull request: Fix environment variable order for global rank determination
This resolves the issue #11380 and the underlying root bug described here https://github.com/PyTorchLightning/pytorch-lightning/issues/11380#issuecomment-1009376167. 

Fixes #11380 

### Does your PR introduce any breaking changes? If yes, please list them.

<!-- FILL IN or None -->

## Before submitting

- [ ] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- [x] Did you make sure to **update the documentation** with your changes? (if necessary)
- [x] Did you write any **new necessary tests**? (not for typos and docs)
- [ ] Did you verify new and **existing tests pass** locally with your changes?
- [x] Did you list all the **breaking changes** introduced by this pull request?
- [x] Did you **update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)**? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [x] Is this pull request ready for review? (if not, please submit in draft mode)
- [ ] Check that all items from **Before submitting** are resolved
- [ ] Make sure the title is self-explanatory and the description concisely explains the PR
- [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?

Make sure you had fun coding 🙃


Problem statement:
WandbLogger in DDP setting creates multiple runs in W&B
When using the `WandbLogger` in DDP mode with `n>1` GPUs, `n` runs in Weights&Biases will be created. Only the first run will have the actual logs, all other runs will be empty, essentially cluttering up the W&B project.

### Expected behavior

Only one run will be created, even when using DDP on multiple GPUs.

### Environment
Example config, but this issue also happens on other configurations with different GPU models:
* CUDA:
	- GPU:
		- NVIDIA A100-SXM-80GB
		- NVIDIA A100-SXM-80GB
	- available:         True
	- version:           11.0
* Packages:
	- numpy:             1.22.0
	- pyTorch_debug:     False
	- pyTorch_version:   1.10.0
	- pytorch-lightning: 1.5.8
	- tqdm:              4.62.3
* System:
	- OS:                Linux
	- architecture:
		- 64bit
		- ELF
	- processor:         x86_64
	- python:            3.9.9
	- version:           # 91-Ubuntu SMP Thu Jul 15 19:09:17 UTC 2021


### Additional Info
Since this bug only happens in DDP with multiple GPUs, I can't reproduce it in Colab.

cc @awaelchli @morganmcg1 @AyushExel @borisdayma @scottire

Hints from issue discussion:
I *think* I can "fix" this by manually calling `wand.init` only on the `rank0` process once before creating the `WandbLogger`, but it would be much more in the spirit of pytorch-lightning if the user didn't have to think about these things.

EDIT: Some more information:
My WandbLogger looks like this and I am calling `wandb.watch` afterwards:
```python
wandb_logger = WandbLogger(
        project=<project>, entity=<entity>, name=<name>, tags=<tags>, log_model=True)
wandb_logger.watch(model, log="all", log_graph=False)
```
I **do not** make any manual calls to `wandb` in my project, I only use the integration exposed through `pytorch-lightning`.
Hey @konstantinjdobler, Morgan from W&B here, thanks for the issue. Yep like you found, calling wand.init on the rank0 process is one way to work with DDP today.

We also have an experimental feature called `wandb service` that improves wandb's handling of multiprocessing. You can read more about [here](https://github.com/wandb/client/blob/master/docs/dev/wandb-service-user.md). It might address your issue if you'd like to try?

Note that as part of it you'll have to install this branch: 

`pip install --force git+https://github.com/wandb/pytorch-lightning.git@wandb-service-attach`

Commit message: Fix rank environment variable priority

Patch to review:
diff --git a/pytorch_lightning/utilities/distributed.py b/pytorch_lightning/utilities/distributed.py
index 7e817b14c4325..884bcd157268c 100644
--- a/pytorch_lightning/utilities/distributed.py
+++ b/pytorch_lightning/utilities/distributed.py
@@ -55,7 +55,7 @@ def wrapped_fn(*args: Any, **kwargs: Any) -> Optional[Any]:
 
 # TODO: this should be part of the cluster environment
 def _get_rank() -> int:
-    rank_keys = ("RANK", "SLURM_PROCID", "LOCAL_RANK")
+    rank_keys = ("RANK", "LOCAL_RANK", "SLURM_PROCID")
     for key in rank_keys:
         rank = os.environ.get(key)
         if rank is not None:


Code review comments:
diff_hunk: @@ -55,7 +55,7 @@ def wrapped_fn(*args: Any, **kwargs: Any) -> Optional[Any]:
 
 # TODO: this should be part of the cluster environment
 def _get_rank() -> int:
-    rank_keys = ("RANK", "SLURM_PROCID", "LOCAL_RANK")
+    rank_keys = ("RANK", "LOCAL_RANK", "SLURM_PROCID")
path: pytorch_lightning/utilities/distributed.py
text: @user3:
This was discussed in the original PR: https://github.com/PyTorchLightning/pytorch-lightning/pull/6802#discussion_r607593905

cc @user1 @user2 for visibility

@author:
From https://github.com/PyTorchLightning/pytorch-lightning/pull/6802#discussion_r608400925:
> RANK = torchelastic
SLURM_PROCID = slurm
LOCAL_RANK = parity with existing setup though I think it's not right
> I set local rank last as RANK and SLURM_PROCID correspond to global rank already. 

At least in my case (running a *single* slurm job from which I then start a DDP training with `pytorch-lightning`), this is not the case, as `SLURM_PROCID` is always set to 0, on all processes started through DDP. 

I'm no Slurm expert but at least on clusters shared between many people, I don't think my case is an uncommon pattern.

@user1:
@author by that do you mean you don't launch all process with slurm first? And pytorch Lightning is the entity creating the processes?

@author:
Exactly! I have to use Slurm to *provision* resources on my cluster (e.g. a `srun` command specifying 4 GPUs in a partition which spawns a sub-shell in which I then execute the Python training script containing pytorch-lightning code like so: `python train.py <args...> --gpus 4`. By specifying `--gpus 4` I use the regular pytorch-lightning DDP implementation. 

This pattern actually works really well in terms of code clarity (I don't need any extra code that only works on slurm clusters, e.g. I can run the exact same code on a cluster without slurm or without DDP) and also provides frictionless access to resources on a shared cluster. Especially the point of code clarity has been one of the major selling points of pytorch-lightning for me.

@user1:
@author to confirm:
1. are you also setting SLURM_NTASKS? https://github.com/PyTorchLightning/pytorch-lightning/blob/948cfd24de4f64a2980395581f15544e5e37eab0/pytorch_lightning/plugins/environments/slurm_environment.py#L82-L85 
2. Are you setting LOCAL_RANK? https://github.com/PyTorchLightning/pytorch-lightning/blob/948cfd24de4f64a2980395581f15544e5e37eab0/pytorch_lightning/plugins/environments/lightning_environment.py#L43-L50

If you answered no to both of those questions, a slightly better fix would be to iterate through the different environments lightning supports, and call detect: 
https://github.com/PyTorchLightning/pytorch-lightning/blob/948cfd24de4f64a2980395581f15544e5e37eab0/pytorch_lightning/trainer/connectors/accelerator_connector.py#L810-L812

for the first one that's available, instantiate it and call `env.global_rank()` 

then we can finally resolve the TODO that's listed here and I think that should fix the bug

@user2:
Did you try forcing `Trainer(plugins=LightningEnvironment())` (import pytorch_lightning.plugins LightningEnvironment) to bypass our automatic detection of the slurm environment? 

@user1:
@user2 i think at least one  issue here is that the rank_zero utilities are operating independently of the cluster environment 

@author:
@user1 Thanks for the help! Keep in mind, I am not manually setting any of these, this is done automatically either by Slurm or pytorch-lightning
> 1. are you also setting SLURM_NTASKS? 

SLURM_NTASKS is indeed set to 1 already when I start the training script

> 2. Are you setting LOCAL_RANK? 

LOCAL_RANK is appropriately set **by pytorch-lightning** in the processes launched by DDP (e.g. LOCAL_RANK is 1 in the first additional process launched by DDP and so on). I'm not manually setting it and it is not set when I first start the training script.

> If you answered no to both of those questions, a slightly better fix would be to iterate through the different environments lightning supports, and call detect:

Both are set in my case, so not sure how to go forward. At the moment, I just fix this issue for me by monkey patching pytorch-lightning with the change proposed in my pull request. I think the crux of the issue is that I am using Slurm, but **not orchestrating the multiprocessing with Slurm**, but because Slurm environment variables are available, pytorch-lightning thinks that Slurm is taking care of the multiprocessing.

Interestingly, when calling `trainer.fit` it all works out correctly again and pytorch-lightning starts the appropriate additional processes, it's just that `rank_zero` gets "initialized" incorrectly because the SLURM_PROCID variable is present in the environment.



@user1:
> This pattern actually works really well in terms of code clarity (I don't need any extra code that only works on slurm clusters, e.g. I can run the exact same code on a cluster without slurm or without DDP) and also provides frictionless access to resources on a shared cluster. Especially the point of code clarity has been one of the major selling points of pytorch-lightning for me.

Lightning has first-class support for SLURM, such you could create the processes externally (through SLURM) without any code changes. Is this a viable option for you?

The issue now is Lightning assumes the presence of certain environment variables indicate the intention for how someone is using them. That is, if someone has set SLURM_NTASKS, or SLURM_PROCID, then we assume they're set and set correctly

@author:
I mean I could probably create the processes externally (through SLURM) but for me that defeats the purpose of using Lightning: I don't want to deal with creating any processes, I want to let Lightning handle it.

LOCAL_RANK is set by Lightning, right? The question is if there exists a scenario where LOCAL_RANK would be set when the multiprocessing is actually handled externally. If not, to me it makes sense to check LOCAL_RANK first, to avoid issues in my scenario.

@user2:
> which spawns a sub-shell in which I then execute the Python training script containing pytorch-lightning code like so: python train.py <args...> --gpus 4. By specifying --gpus 4 I use the regular pytorch-lightning DDP implementation.

So that means if you would unset the environment variable, it would work as expected too:
```
unset SLURM_NTASKS
unset SLURM_PROCID
python train.py <args...> --gpus 4
```
Then you don't need to monkey-patch Lightning as you said.

I guess we can still go ahead with the PR but I'm also concerned with the same questions as @user1 

@author:
I just tested unsetting the environment variables and it indeed does work (I still prefer the monkey patch, since I don't have to mess with the environment). The problem is solved for me, since I know about it now. 

In some sense, since I opened the issue other people with the same problem can now find the solution through a google search, but these fixes go very deep into the under-the-hood behavior of lightning and are quite unintuitive for the original problem of multiple runs being created in Weights&Biases (one for each GPU). Of course, it is up to you if believe any changes are worth it to accommodate users that use Slurm only to allocate resources and let lightning handle multiprocessing.

@user2:
@user1 I don't see any danger merging this. Slurm doesn't use the LOCAL_RANK variable so there is no collision there. Also, Lightning doesn't set LOCAL_RANK if the slurm cluster environment is used. Meaning, existing slurm users who let their jobs get scheduled won't be affected by this change.

I recommend to go ahead if it's ok with you.