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

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):