[fix] Better support for rank_zero_only setting for SLURM and torchelastic

[ananthsub]

What does this PR do?

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

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

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. In short, see the following bullet-list:

• 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 🙃

[pep8speaks]

Hello @ananthsub! Thanks for updating this PR.

There are currently no PEP 8 issues detected in this Pull Request. Cheers! 🍻

Comment last updated at 2021-04-07 08:54:44 UTC

[codecov]

Codecov Report

Merging #6802 (e49bd32) into master (b7a22ba) will decrease coverage by 48%.
The diff coverage is n/a.

❗ Current head e49bd32 differs from pull request most recent head f9d0cd8. Consider uploading reports for the commit f9d0cd8 to get more accurate results

@@           Coverage Diff            @@
##           master   #6802     +/-   ##
========================================
- Coverage      92%     44%    -48%     
========================================
  Files         193     192      -1     
  Lines       12271   12191     -80     
========================================
- Hits        11234    5346   -5888     
- Misses       1037    6845   +5808     

[awaelchli]

note, the tests won't run because the file distributed.py is not prefixed with test_

[tchaton]

Can you move this directly to SLURM cluster environment ?

What if RANK, SLURM_PROCID or LOCAL_RANK are different ? Should we take the latest or did you order rank_keys based on priority ?

Best,
T.C

[ananthsub]

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. The linked issue has more discussion, but I think we should make global rank and world size properties of the cluster environment. So the cluster environment becomes the source of truth propagating from Cluster environment => training type plugin => accelerator => trainer.

The main issue now is the global rank isn't set on trainer initialization. If the cluster environment is marked as creating children, then we can leave the initialization of these fields for later, but both torchelastic and slurm have this data already available in the environment variables, and we should expose that as soon as possible (on Trainer init) for users to read this state.

Currently, this waits for trainer.fit() to be called, going through the accelerator setup flow for these properties to be initialized on the training type plugin.