Repository: Lightning-AI/pytorch-lightning
Pull request: [fix] Better support for rank_zero_only setting for SLURM and torchelastic
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
 Please also include relevant motivation and context.
 List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->
Fixes #6797

This is a mitigation for the issue. The environment variable handling is currently split in a few different places:
- in this file
- in the accelerator connector
- in the cluster environment

This needs to be consolidated to better support more custom cluster environments in the future. For now this is a quickfix for SLURM users until we can come up with a better design.

## Before submitting
- [x] 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)
- [x] Did you verify new and existing tests pass locally with your changes?
- [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)

<!-- For CHANGELOG separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review
Anyone in the community is free to review the PR once the tests have passed.
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)
 - [x] Check that all items from **Before submitting** are resolved
 - [x] Make sure the title is self-explanatory and the description concisely explains the PR
 - [x] 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:
LOCAL_RANK not being set in slurm
## 🐛 Bug

A lot of the PTL tooling around multiprocess depends on a specific environment variable: LOCAL_RANK being set correctly, it seems that when running in slurm this isnt set causing it to return the default of 0 for all processes which makes every process do things that should only be done on rank 0, like log stuff. 

Also I'm a little unclear about the name of that variable, if I have multiple nodes, only the *global* rank 0 not the local rank should be logging and saving checkpoints etc.

### To Reproduce

Run in slurm (cant really do it w/ colab), a good way to easily see it is to use the Wandb logger, you'll see that each process makes a new run on the Wandb UI which means that @rank_zero_experiment didnt work properly, and you can confirm this by printing LOCAL_RANK which is defaulted to 0 if unset, it will always give back 0.

### Expected behavior

LOCAL_RANK is set correctly or the rest of the tooling is aware of the global rank of the process

### Environment

Will update if it's really necessary

Hints from issue discussion:
Adding the following to the `SLURMEnvironment` contrustructor does seem to work:

```
os.environ["LOCAL_RANK"] = os.environ["SLURM_PROCID"]
rank_zero_only.rank = int(os.environ["SLURM_PROCID"])
```

but I would like to get some feedback from the developers before proposing this as a solution. 

What's the value you see for `SLURM_LOCALID` ? https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/plugins/environments/slurm_environment.py#L75-L76
That appears to be correct, in `__init__` of `SLURMEnvironment`:

```
def __init__(self):
        super().__init__()
        os.environ["LOCAL_RANK"] = os.environ["SLURM_PROCID"]
        rank_zero_only.rank = int(os.environ["SLURM_PROCID"])
        print(f"SLURM LOCALID: {os.environ['SLURM_LOCALID']}")
```

prints `SLURM LOCALID: 0` on rank 0, `SLURM LOCALID: 3` on rank 3, etc.


I think this is a fairly straightforward fix and it makes sense to delegate this responsibility to the `SLURMConnector` I just want to double check that 
1.  `__init__` is the right place to do it 
2.  that there isnt a better way to get `rank_zero_only.rank` set correctly (this follows a precedent set by the `GPUAccelerator`) 
3. that the semantics of `LOCAL_RANK` are actually supposed to be a global rank (in which case there should be some refactoring around the name of that variable but probably outside the scope of a PR related to this issue) 
@awaelchli what do you think about setting the rank info on the cluster environment / training type plugin during the Trainer init?currently we wait for `setup_environment` in the training plugin to be called, but if someone has launched with SLURM or torchelastic, this information is already available based on the environment variables. Only in the case of spawn do we not know this yet, but we can handle that accordingly
You mean setting the `rank_zero_only.rank` right? If the LOCAL_RANK variable is set, it will [already default to that](https://github.com/PyTorchLightning/pytorch-lightning/blob/bb9ace43334ad50e3758d9cff08ad34216c7d4da/pytorch_lightning/utilities/distributed.py#L48)
I would avoid the following line in the above snippet:
`os.environ["LOCAL_RANK"] = os.environ["SLURM_PROCID"]`
In my opinion Lightning should not write environment variables, only read them, unless there is a very good reason.
If the user needs the local rank, they can access it through the trainer.local_rank and the cluster environment will conveniently translate the slurm local id to that.
We can think of a way to set `rank_zero_only.rank` earlier so that the `rank_zero_only` etc. decorators work as soon as possible. 
> We can think of a way to set `rank_zero_only.rank` earlier so that the `rank_zero_only` etc. decorators work as soon as possible.

It's not that it needs to be set sooner it's that it's not set correctly at all. `LOCAL_RANK` isn't set in a slurm environment, so if you take a look at the line you linked, that uses a default value of 0 if the environment variable isn't set, which means that every process thinks it's the rank 0 process and writes logs/saves checkpoints/etc.

As far as I can tell there's no machinery around getting the slurm assigned rank into `rank_zero.rank` so that's what we would need to fix this.

Also I'm a little concerned about that variable being called LOCAL_RANK which seems to imply that rank 0 on every machine should be doing logging, checkpointing, whatever else. It should only be the global rank zero that does those things. 
yes it defaults to 0 but that's the only assumption you can make if you don't know where else to get the rank information from.
the idea is that later on it gets overwritten, for example here in ddp:

https://github.com/PyTorchLightning/pytorch-lightning/blob/bb9ace43334ad50e3758d9cff08ad34216c7d4da/pytorch_lightning/plugins/training_type/ddp.py#L174

> It's not that it needs to be set sooner it's that it's not set correctly at all.

no I think what you are experiencing is really just because it's not set sooner. I wish I had slurm but I can't test it :( 
In your post you mention wandb creating multiple runs.
I recently fixed this: #6380
Does your version of Lightning include this fix?

Commit message: Update distributed.py

Patch to review:
diff --git a/pytorch_lightning/utilities/distributed.py b/pytorch_lightning/utilities/distributed.py
index bf7a199fc08dc..b36ab3827e674 100644
--- a/pytorch_lightning/utilities/distributed.py
+++ b/pytorch_lightning/utilities/distributed.py
@@ -44,8 +44,18 @@ def wrapped_fn(*args, **kwargs):
     return wrapped_fn
 
 
+# TODO: this should be part of the cluster environment
+def _get_local_rank() -> int:
+    local_rank_keys = ('LOCAL_RANK', 'SLURM_LOCALID')
+    for key in local_rank_keys:
+        rank = os.environ.get(key)
+        if rank is not None:
+            return int(rank)
+    return 0
+
+
 # add the attribute to the function but don't overwrite in case Trainer has already set it
-rank_zero_only.rank = getattr(rank_zero_only, 'rank', int(os.environ.get('LOCAL_RANK', 0)))
+rank_zero_only.rank = getattr(rank_zero_only, 'rank', _get_local_rank())
 
 
 def _warn(*args, **kwargs):


Code review comments:
diff_hunk: @@ -44,8 +44,18 @@ def wrapped_fn(*args, **kwargs):
     return wrapped_fn
 
 
+# TODO: this should be part of the cluster environment
+def _get_local_rank() -> int:
+    local_rank_keys = ('LOCAL_RANK', 'SLURM_LOCALID')
path: pytorch_lightning/utilities/distributed.py
text: @user1:
an nasty way :D
`os.environ.get('LOCAL_RANK', os.environ.get('SLURM_LOCALID', 0))`

@author:
in case we have another cluster to support and we haven't fixed the design issues here, the tuple + loop will be easier to extend. hopefully it doesn't come to that though and we can remove this altogether. 

@user2:
Should be SLURM_PROCID to get the global rank