From a33b025a054f30c37728acd2f846b417d1c52bd6 Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Fri, 2 Apr 2021 15:44:33 -0700
Subject: [PATCH 01/19] Update distributed.py

---
 pytorch_lightning/utilities/distributed.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/pytorch_lightning/utilities/distributed.py b/pytorch_lightning/utilities/distributed.py
index bf7a199fc08dc..8da00069cfcc8 100644
--- a/pytorch_lightning/utilities/distributed.py
+++ b/pytorch_lightning/utilities/distributed.py
@@ -43,9 +43,17 @@ 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):

From a711a1a6a9630a3403fcc1cc4afe8ca926215b79 Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Fri, 2 Apr 2021 15:52:31 -0700
Subject: [PATCH 02/19] Update distributed.py

---
 pytorch_lightning/utilities/distributed.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/pytorch_lightning/utilities/distributed.py b/pytorch_lightning/utilities/distributed.py
index 8da00069cfcc8..b36ab3827e674 100644
--- a/pytorch_lightning/utilities/distributed.py
+++ b/pytorch_lightning/utilities/distributed.py
@@ -43,6 +43,7 @@ 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')
@@ -52,6 +53,7 @@ def _get_local_rank() -> int:
             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', _get_local_rank())
 

From 739791b86e23c627c8b50beff672d4c2adf745f5 Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Sun, 4 Apr 2021 23:00:25 -0700
Subject: [PATCH 03/19] Update distributed.py

---
 pytorch_lightning/utilities/distributed.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/pytorch_lightning/utilities/distributed.py b/pytorch_lightning/utilities/distributed.py
index b36ab3827e674..27859a022e641 100644
--- a/pytorch_lightning/utilities/distributed.py
+++ b/pytorch_lightning/utilities/distributed.py
@@ -45,9 +45,9 @@ def wrapped_fn(*args, **kwargs):
 
 
 # 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:
+def _get_rank() -> int:
+    rank_keys = ('LOCAL_RANK', 'SLURM_PROCID')
+    for key in rank_keys:
         rank = os.environ.get(key)
         if rank is not None:
             return int(rank)
@@ -55,7 +55,7 @@ def _get_local_rank() -> int:
 
 
 # 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', _get_local_rank())
+rank_zero_only.rank = getattr(rank_zero_only, 'rank', _get_rank())
 
 
 def _warn(*args, **kwargs):

From 74a08046eee82e42a1de99f954f2e00b1b1e3a1b Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Sun, 4 Apr 2021 23:01:32 -0700
Subject: [PATCH 04/19] Update distributed.py

---
 pytorch_lightning/utilities/distributed.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pytorch_lightning/utilities/distributed.py b/pytorch_lightning/utilities/distributed.py
index 27859a022e641..b4793889f1299 100644
--- a/pytorch_lightning/utilities/distributed.py
+++ b/pytorch_lightning/utilities/distributed.py
@@ -46,7 +46,7 @@ def wrapped_fn(*args, **kwargs):
 
 # TODO: this should be part of the cluster environment
 def _get_rank() -> int:
-    rank_keys = ('LOCAL_RANK', 'SLURM_PROCID')
+    rank_keys = ('RANK', 'SLURM_PROCID', 'LOCAL_RANK')
     for key in rank_keys:
         rank = os.environ.get(key)
         if rank is not None:

From f9c3c793448099bb2d5bfb2ffe1ed76ba5c5ef92 Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Sun, 4 Apr 2021 23:06:57 -0700
Subject: [PATCH 05/19] Update CHANGELOG.md

---
 CHANGELOG.md | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 81846809fbf85..b80140724712f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -170,6 +170,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Set better defaults for `rank_zero_only.rank` when training is launched with SLURM and torchelastic ([#6802](https://github.com/PyTorchLightning/pytorch-lightning/pull/6802/))
+
+
 - Made the `Plugin.reduce` method more consistent across all Plugins to reflect a mean-reduction by default ([#6011](https://github.com/PyTorchLightning/pytorch-lightning/pull/6011))
 
 
@@ -196,7 +199,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed torch distributed not available in setup hook for DDP ([#6506](https://github.com/PyTorchLightning/pytorch-lightning/pull/6506))
 
-
 ## [1.2.6] - 2021-03-30
 
 ### Changed

From 8610461f9a9dc5a961d0b1627289630332cb85bb Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Sun, 4 Apr 2021 23:56:06 -0700
Subject: [PATCH 06/19] Update distributed.py

---
 tests/utilities/distributed.py | 46 ++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/tests/utilities/distributed.py b/tests/utilities/distributed.py
index 80c0246ce6c57..bc980afc58250 100644
--- a/tests/utilities/distributed.py
+++ b/tests/utilities/distributed.py
@@ -16,6 +16,7 @@
 import sys
 from pathlib import Path
 from subprocess import TimeoutExpired
+from unittest import mock
 
 import pytorch_lightning
 
@@ -42,3 +43,48 @@ def call_training_script(module_file, cli_args, method, tmpdir, timeout=60):
         p.kill()
         std, err = p.communicate()
     return std, err
+
+
+@mock.patch.dict(os.environ, {"SLURM_PROCID": "0"})
+def test_rank_zero_slurm():
+    """ Test that SLURM environment variables are properly checked for rank_zero_only. """
+    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
+    rank_zero_only.rank = _get_rank()
+
+    @rank_zero_only
+    def foo():
+        # The return type is optional because on non-zero ranks it will not be called
+        return 1
+
+    x = foo()
+    assert x == 1
+
+
+@mock.patch.dict(os.environ, {"RANK": "0"})
+def test_rank_zero_torchelastic():
+    """ Test that torchelastic environment variables are properly checked for rank_zero_only. """
+    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
+    rank_zero_only.rank = _get_rank()
+
+    @rank_zero_only
+    def foo():
+        # The return type is optional because on non-zero ranks it will not be called
+        return 1
+
+    x = foo()
+    assert x == 1
+
+
+@mock.patch.dict(os.environ, {"RANK": "1", "SLURM_PROCID": "2", "LOCAL_RANK": "3"})
+def test_rank_zero_none_set():
+    """ Test that function is not called when rank environment variables are not global zero. """
+
+    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
+    rank_zero_only.rank = _get_rank()
+
+    @rank_zero_only
+    def foo():
+        return 1
+
+    x = foo()
+    assert x is None

From 51c94132112eba6157165635c70f3b01efcf143b Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Mon, 5 Apr 2021 03:07:15 -0700
Subject: [PATCH 07/19] Update tests/utilities/distributed.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com>
---
 tests/utilities/distributed.py | 23 ++++++++++++++---------
 1 file changed, 14 insertions(+), 9 deletions(-)

diff --git a/tests/utilities/distributed.py b/tests/utilities/distributed.py
index bc980afc58250..248bbe11a9dc0 100644
--- a/tests/utilities/distributed.py
+++ b/tests/utilities/distributed.py
@@ -75,16 +75,21 @@ def foo():
     assert x == 1
 
 
-@mock.patch.dict(os.environ, {"RANK": "1", "SLURM_PROCID": "2", "LOCAL_RANK": "3"})
-def test_rank_zero_none_set():
+@pytest.mark.parametrize("rank_key,rank", [
+    ("RANK", "1"),
+    ("SLURM_PROCID", "2"),
+    ("LOCAL_RANK", "3"),
+])
+def test_rank_zero_none_set(rank_key, rank):
     """ Test that function is not called when rank environment variables are not global zero. """
 
-    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
-    rank_zero_only.rank = _get_rank()
+    with mock.patch.dict(os.environ, {rank_key: rank}):
+        from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
+        rank_zero_only.rank = _get_rank()
 
-    @rank_zero_only
-    def foo():
-        return 1
+        @rank_zero_only
+        def foo():
+            return 1
 
-    x = foo()
-    assert x is None
+        x = foo()
+        assert x is None

From 3a6a4b7045a802b5fd06f38aa52de670e809d574 Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Mon, 5 Apr 2021 03:13:35 -0700
Subject: [PATCH 08/19] make new test_distributed.py

---
 tests/utilities/distributed.py      | 46 --------------------
 tests/utilities/test_distributed.py | 67 +++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 46 deletions(-)
 create mode 100644 tests/utilities/test_distributed.py

diff --git a/tests/utilities/distributed.py b/tests/utilities/distributed.py
index bc980afc58250..80c0246ce6c57 100644
--- a/tests/utilities/distributed.py
+++ b/tests/utilities/distributed.py
@@ -16,7 +16,6 @@
 import sys
 from pathlib import Path
 from subprocess import TimeoutExpired
-from unittest import mock
 
 import pytorch_lightning
 
@@ -43,48 +42,3 @@ def call_training_script(module_file, cli_args, method, tmpdir, timeout=60):
         p.kill()
         std, err = p.communicate()
     return std, err
-
-
-@mock.patch.dict(os.environ, {"SLURM_PROCID": "0"})
-def test_rank_zero_slurm():
-    """ Test that SLURM environment variables are properly checked for rank_zero_only. """
-    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
-    rank_zero_only.rank = _get_rank()
-
-    @rank_zero_only
-    def foo():
-        # The return type is optional because on non-zero ranks it will not be called
-        return 1
-
-    x = foo()
-    assert x == 1
-
-
-@mock.patch.dict(os.environ, {"RANK": "0"})
-def test_rank_zero_torchelastic():
-    """ Test that torchelastic environment variables are properly checked for rank_zero_only. """
-    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
-    rank_zero_only.rank = _get_rank()
-
-    @rank_zero_only
-    def foo():
-        # The return type is optional because on non-zero ranks it will not be called
-        return 1
-
-    x = foo()
-    assert x == 1
-
-
-@mock.patch.dict(os.environ, {"RANK": "1", "SLURM_PROCID": "2", "LOCAL_RANK": "3"})
-def test_rank_zero_none_set():
-    """ Test that function is not called when rank environment variables are not global zero. """
-
-    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
-    rank_zero_only.rank = _get_rank()
-
-    @rank_zero_only
-    def foo():
-        return 1
-
-    x = foo()
-    assert x is None
diff --git a/tests/utilities/test_distributed.py b/tests/utilities/test_distributed.py
new file mode 100644
index 0000000000000..75f1454be6199
--- /dev/null
+++ b/tests/utilities/test_distributed.py
@@ -0,0 +1,67 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+from unittest import mock
+
+import pytest
+
+
+@mock.patch.dict(os.environ, {"SLURM_PROCID": "0"})
+def test_rank_zero_slurm():
+    """ Test that SLURM environment variables are properly checked for rank_zero_only. """
+    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
+    rank_zero_only.rank = _get_rank()
+
+    @rank_zero_only
+    def foo():
+        # The return type is optional because on non-zero ranks it will not be called
+        return 1
+
+    x = foo()
+    assert x == 1
+
+
+@mock.patch.dict(os.environ, {"RANK": "0"})
+def test_rank_zero_torchelastic():
+    """ Test that torchelastic environment variables are properly checked for rank_zero_only. """
+    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
+    rank_zero_only.rank = _get_rank()
+
+    @rank_zero_only
+    def foo():
+        # The return type is optional because on non-zero ranks it will not be called
+        return 1
+
+    x = foo()
+    assert x == 1
+
+
+@pytest.mark.parametrize("rank_key,rank", [
+    ("RANK", "1"),
+    ("SLURM_PROCID", "2"),
+    ("LOCAL_RANK", "3"),
+])
+def test_rank_zero_none_set(rank_key, rank):
+    """ Test that function is not called when rank environment variables are not global zero. """
+
+    with mock.patch.dict(os.environ, {rank_key: rank}):
+        from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
+        rank_zero_only.rank = _get_rank()
+
+        @rank_zero_only
+        def foo():
+            return 1
+
+        x = foo()
+        assert x is None

From 10a40966f70048e7b8c8fa3b60973208c22f21ce Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Fri, 2 Apr 2021 15:44:33 -0700
Subject: [PATCH 09/19] Update distributed.py

---
 pytorch_lightning/utilities/distributed.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/pytorch_lightning/utilities/distributed.py b/pytorch_lightning/utilities/distributed.py
index bf7a199fc08dc..8da00069cfcc8 100644
--- a/pytorch_lightning/utilities/distributed.py
+++ b/pytorch_lightning/utilities/distributed.py
@@ -43,9 +43,17 @@ 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):

From d54135ff7cf42e296653b3f28a0e9521893ce80d Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Fri, 2 Apr 2021 15:52:31 -0700
Subject: [PATCH 10/19] Update distributed.py

---
 pytorch_lightning/utilities/distributed.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/pytorch_lightning/utilities/distributed.py b/pytorch_lightning/utilities/distributed.py
index 8da00069cfcc8..b36ab3827e674 100644
--- a/pytorch_lightning/utilities/distributed.py
+++ b/pytorch_lightning/utilities/distributed.py
@@ -43,6 +43,7 @@ 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')
@@ -52,6 +53,7 @@ def _get_local_rank() -> int:
             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', _get_local_rank())
 

From a425b019981f183f1f009ca3bb90d63b17687166 Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Sun, 4 Apr 2021 23:00:25 -0700
Subject: [PATCH 11/19] Update distributed.py

---
 pytorch_lightning/utilities/distributed.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/pytorch_lightning/utilities/distributed.py b/pytorch_lightning/utilities/distributed.py
index b36ab3827e674..27859a022e641 100644
--- a/pytorch_lightning/utilities/distributed.py
+++ b/pytorch_lightning/utilities/distributed.py
@@ -45,9 +45,9 @@ def wrapped_fn(*args, **kwargs):
 
 
 # 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:
+def _get_rank() -> int:
+    rank_keys = ('LOCAL_RANK', 'SLURM_PROCID')
+    for key in rank_keys:
         rank = os.environ.get(key)
         if rank is not None:
             return int(rank)
@@ -55,7 +55,7 @@ def _get_local_rank() -> int:
 
 
 # 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', _get_local_rank())
+rank_zero_only.rank = getattr(rank_zero_only, 'rank', _get_rank())
 
 
 def _warn(*args, **kwargs):

From b372c777594e2ecbdfb7ef70ad7afc323ea2afc3 Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Sun, 4 Apr 2021 23:01:32 -0700
Subject: [PATCH 12/19] Update distributed.py

---
 pytorch_lightning/utilities/distributed.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pytorch_lightning/utilities/distributed.py b/pytorch_lightning/utilities/distributed.py
index 27859a022e641..b4793889f1299 100644
--- a/pytorch_lightning/utilities/distributed.py
+++ b/pytorch_lightning/utilities/distributed.py
@@ -46,7 +46,7 @@ def wrapped_fn(*args, **kwargs):
 
 # TODO: this should be part of the cluster environment
 def _get_rank() -> int:
-    rank_keys = ('LOCAL_RANK', 'SLURM_PROCID')
+    rank_keys = ('RANK', 'SLURM_PROCID', 'LOCAL_RANK')
     for key in rank_keys:
         rank = os.environ.get(key)
         if rank is not None:

From 6e452ba752acc871352c82125b4d6635230bed1a Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Sun, 4 Apr 2021 23:06:57 -0700
Subject: [PATCH 13/19] Update CHANGELOG.md

---
 CHANGELOG.md | 19 +++----------------
 1 file changed, 3 insertions(+), 16 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1d2629d3928f0..f47078ade259b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -176,6 +176,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Set better defaults for `rank_zero_only.rank` when training is launched with SLURM and torchelastic ([#6802](https://github.com/PyTorchLightning/pytorch-lightning/pull/6802/))
+
+
 - Sanitize `None` params during pruning ([#6836](https://github.com/PyTorchLightning/pytorch-lightning/pull/6836))
 
 
@@ -205,22 +208,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed torch distributed not available in setup hook for DDP ([#6506](https://github.com/PyTorchLightning/pytorch-lightning/pull/6506))
 
-
-- Fixed TPU Colab hang issue, post training ([#6816](https://github.com/PyTorchLightning/pytorch-lightning/pull/6816))
-
-
-- Enforce an epoch scheduler interval when using SWA ([#6588](https://github.com/PyTorchLightning/pytorch-lightning/pull/6588))
-
-
-- Fixed an issue with `IterableDataset` when `__len__` is not defined ([#6828](https://github.com/PyTorchLightning/pytorch-lightning/pull/6828))
-
-
-- Fixed `EarlyStopping` logic when `min_epochs` or `min_steps` requirement is not met ([#6705](https://github.com/PyTorchLightning/pytorch-lightning/pull/6705))
-
-
-- Fixed a bug where `TensorBoardLogger` would give a warning and not log correctly to a symbolic link `save_dir` ([#6730](https://github.com/PyTorchLightning/pytorch-lightning/pull/6730))
-
-
 ## [1.2.6] - 2021-03-30
 
 ### Changed

From 592907462b12aec8a6a19f04d1ac3b0b8970a2aa Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Sun, 4 Apr 2021 23:56:06 -0700
Subject: [PATCH 14/19] Update distributed.py

---
 tests/utilities/distributed.py | 46 ++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/tests/utilities/distributed.py b/tests/utilities/distributed.py
index 80c0246ce6c57..bc980afc58250 100644
--- a/tests/utilities/distributed.py
+++ b/tests/utilities/distributed.py
@@ -16,6 +16,7 @@
 import sys
 from pathlib import Path
 from subprocess import TimeoutExpired
+from unittest import mock
 
 import pytorch_lightning
 
@@ -42,3 +43,48 @@ def call_training_script(module_file, cli_args, method, tmpdir, timeout=60):
         p.kill()
         std, err = p.communicate()
     return std, err
+
+
+@mock.patch.dict(os.environ, {"SLURM_PROCID": "0"})
+def test_rank_zero_slurm():
+    """ Test that SLURM environment variables are properly checked for rank_zero_only. """
+    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
+    rank_zero_only.rank = _get_rank()
+
+    @rank_zero_only
+    def foo():
+        # The return type is optional because on non-zero ranks it will not be called
+        return 1
+
+    x = foo()
+    assert x == 1
+
+
+@mock.patch.dict(os.environ, {"RANK": "0"})
+def test_rank_zero_torchelastic():
+    """ Test that torchelastic environment variables are properly checked for rank_zero_only. """
+    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
+    rank_zero_only.rank = _get_rank()
+
+    @rank_zero_only
+    def foo():
+        # The return type is optional because on non-zero ranks it will not be called
+        return 1
+
+    x = foo()
+    assert x == 1
+
+
+@mock.patch.dict(os.environ, {"RANK": "1", "SLURM_PROCID": "2", "LOCAL_RANK": "3"})
+def test_rank_zero_none_set():
+    """ Test that function is not called when rank environment variables are not global zero. """
+
+    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
+    rank_zero_only.rank = _get_rank()
+
+    @rank_zero_only
+    def foo():
+        return 1
+
+    x = foo()
+    assert x is None

From 21a656bebe4059475144125df5ed901069eebc67 Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Mon, 5 Apr 2021 03:07:15 -0700
Subject: [PATCH 15/19] Update tests/utilities/distributed.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com>
---
 tests/utilities/distributed.py | 23 ++++++++++++++---------
 1 file changed, 14 insertions(+), 9 deletions(-)

diff --git a/tests/utilities/distributed.py b/tests/utilities/distributed.py
index bc980afc58250..248bbe11a9dc0 100644
--- a/tests/utilities/distributed.py
+++ b/tests/utilities/distributed.py
@@ -75,16 +75,21 @@ def foo():
     assert x == 1
 
 
-@mock.patch.dict(os.environ, {"RANK": "1", "SLURM_PROCID": "2", "LOCAL_RANK": "3"})
-def test_rank_zero_none_set():
+@pytest.mark.parametrize("rank_key,rank", [
+    ("RANK", "1"),
+    ("SLURM_PROCID", "2"),
+    ("LOCAL_RANK", "3"),
+])
+def test_rank_zero_none_set(rank_key, rank):
     """ Test that function is not called when rank environment variables are not global zero. """
 
-    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
-    rank_zero_only.rank = _get_rank()
+    with mock.patch.dict(os.environ, {rank_key: rank}):
+        from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
+        rank_zero_only.rank = _get_rank()
 
-    @rank_zero_only
-    def foo():
-        return 1
+        @rank_zero_only
+        def foo():
+            return 1
 
-    x = foo()
-    assert x is None
+        x = foo()
+        assert x is None

From 9713d40b9865cfd3debb65a0737fd8c9f7d31b1b Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Mon, 5 Apr 2021 03:13:35 -0700
Subject: [PATCH 16/19] make new test_distributed.py

---
 tests/utilities/distributed.py      | 51 ----------------------
 tests/utilities/test_distributed.py | 67 +++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 51 deletions(-)
 create mode 100644 tests/utilities/test_distributed.py

diff --git a/tests/utilities/distributed.py b/tests/utilities/distributed.py
index 248bbe11a9dc0..80c0246ce6c57 100644
--- a/tests/utilities/distributed.py
+++ b/tests/utilities/distributed.py
@@ -16,7 +16,6 @@
 import sys
 from pathlib import Path
 from subprocess import TimeoutExpired
-from unittest import mock
 
 import pytorch_lightning
 
@@ -43,53 +42,3 @@ def call_training_script(module_file, cli_args, method, tmpdir, timeout=60):
         p.kill()
         std, err = p.communicate()
     return std, err
-
-
-@mock.patch.dict(os.environ, {"SLURM_PROCID": "0"})
-def test_rank_zero_slurm():
-    """ Test that SLURM environment variables are properly checked for rank_zero_only. """
-    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
-    rank_zero_only.rank = _get_rank()
-
-    @rank_zero_only
-    def foo():
-        # The return type is optional because on non-zero ranks it will not be called
-        return 1
-
-    x = foo()
-    assert x == 1
-
-
-@mock.patch.dict(os.environ, {"RANK": "0"})
-def test_rank_zero_torchelastic():
-    """ Test that torchelastic environment variables are properly checked for rank_zero_only. """
-    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
-    rank_zero_only.rank = _get_rank()
-
-    @rank_zero_only
-    def foo():
-        # The return type is optional because on non-zero ranks it will not be called
-        return 1
-
-    x = foo()
-    assert x == 1
-
-
-@pytest.mark.parametrize("rank_key,rank", [
-    ("RANK", "1"),
-    ("SLURM_PROCID", "2"),
-    ("LOCAL_RANK", "3"),
-])
-def test_rank_zero_none_set(rank_key, rank):
-    """ Test that function is not called when rank environment variables are not global zero. """
-
-    with mock.patch.dict(os.environ, {rank_key: rank}):
-        from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
-        rank_zero_only.rank = _get_rank()
-
-        @rank_zero_only
-        def foo():
-            return 1
-
-        x = foo()
-        assert x is None
diff --git a/tests/utilities/test_distributed.py b/tests/utilities/test_distributed.py
new file mode 100644
index 0000000000000..75f1454be6199
--- /dev/null
+++ b/tests/utilities/test_distributed.py
@@ -0,0 +1,67 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+from unittest import mock
+
+import pytest
+
+
+@mock.patch.dict(os.environ, {"SLURM_PROCID": "0"})
+def test_rank_zero_slurm():
+    """ Test that SLURM environment variables are properly checked for rank_zero_only. """
+    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
+    rank_zero_only.rank = _get_rank()
+
+    @rank_zero_only
+    def foo():
+        # The return type is optional because on non-zero ranks it will not be called
+        return 1
+
+    x = foo()
+    assert x == 1
+
+
+@mock.patch.dict(os.environ, {"RANK": "0"})
+def test_rank_zero_torchelastic():
+    """ Test that torchelastic environment variables are properly checked for rank_zero_only. """
+    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
+    rank_zero_only.rank = _get_rank()
+
+    @rank_zero_only
+    def foo():
+        # The return type is optional because on non-zero ranks it will not be called
+        return 1
+
+    x = foo()
+    assert x == 1
+
+
+@pytest.mark.parametrize("rank_key,rank", [
+    ("RANK", "1"),
+    ("SLURM_PROCID", "2"),
+    ("LOCAL_RANK", "3"),
+])
+def test_rank_zero_none_set(rank_key, rank):
+    """ Test that function is not called when rank environment variables are not global zero. """
+
+    with mock.patch.dict(os.environ, {rank_key: rank}):
+        from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
+        rank_zero_only.rank = _get_rank()
+
+        @rank_zero_only
+        def foo():
+            return 1
+
+        x = foo()
+        assert x is None

From 4d07b3ad16c0d39920b0c6d1c156a18fa008d5f8 Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Wed, 7 Apr 2021 01:53:30 -0700
Subject: [PATCH 17/19] Update test_distributed.py

---
 tests/utilities/test_distributed.py | 35 ++++++++++-------------------
 1 file changed, 12 insertions(+), 23 deletions(-)

diff --git a/tests/utilities/test_distributed.py b/tests/utilities/test_distributed.py
index 75f1454be6199..d38150abe8338 100644
--- a/tests/utilities/test_distributed.py
+++ b/tests/utilities/test_distributed.py
@@ -12,42 +12,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+from typing import Mapping
 from unittest import mock
 
 import pytest
 
 
-@mock.patch.dict(os.environ, {"SLURM_PROCID": "0"})
-def test_rank_zero_slurm():
+@pytest.mark.parameterize("env_vars", [{"RANK": "0"}, {"SLURM_PROCID": "0"}])
+def test_rank_zero_known_cluster_envs(env_vars: Mapping[str, str]):
     """ Test that SLURM environment variables are properly checked for rank_zero_only. """
     from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
     rank_zero_only.rank = _get_rank()
 
-    @rank_zero_only
-    def foo():
-        # The return type is optional because on non-zero ranks it will not be called
-        return 1
-
-    x = foo()
-    assert x == 1
-
-
-@mock.patch.dict(os.environ, {"RANK": "0"})
-def test_rank_zero_torchelastic():
-    """ Test that torchelastic environment variables are properly checked for rank_zero_only. """
-    from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
-    rank_zero_only.rank = _get_rank()
+    with mock.patch.dict(os.environ, env_vars):
+        from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
+        rank_zero_only.rank = _get_rank()
 
-    @rank_zero_only
-    def foo():
-        # The return type is optional because on non-zero ranks it will not be called
-        return 1
+        @rank_zero_only
+        def foo():  # The return type is optional because on non-zero ranks it will not be called
+            return 1
 
-    x = foo()
-    assert x == 1
+        x = foo()
+        assert x == 1
 
 
-@pytest.mark.parametrize("rank_key,rank", [
+@pytest.mark.parameterize("rank_key,rank", [
     ("RANK", "1"),
     ("SLURM_PROCID", "2"),
     ("LOCAL_RANK", "3"),

From e4043d98526a78d7836cc15ebf9949d82b8d54dd Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Wed, 7 Apr 2021 01:53:51 -0700
Subject: [PATCH 18/19] Update test_distributed.py

---
 tests/utilities/test_distributed.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/utilities/test_distributed.py b/tests/utilities/test_distributed.py
index d38150abe8338..879a1cb9c4cd5 100644
--- a/tests/utilities/test_distributed.py
+++ b/tests/utilities/test_distributed.py
@@ -18,7 +18,7 @@
 import pytest
 
 
-@pytest.mark.parameterize("env_vars", [{"RANK": "0"}, {"SLURM_PROCID": "0"}])
+@pytest.mark.parametrize("env_vars", [{"RANK": "0"}, {"SLURM_PROCID": "0"}])
 def test_rank_zero_known_cluster_envs(env_vars: Mapping[str, str]):
     """ Test that SLURM environment variables are properly checked for rank_zero_only. """
     from pytorch_lightning.utilities.distributed import _get_rank, rank_zero_only
@@ -36,7 +36,7 @@ def foo():  # The return type is optional because on non-zero ranks it will not
         assert x == 1
 
 
-@pytest.mark.parameterize("rank_key,rank", [
+@pytest.mark.parametrize("rank_key,rank", [
     ("RANK", "1"),
     ("SLURM_PROCID", "2"),
     ("LOCAL_RANK", "3"),

From c64ffffd9ff2779d31c44aed8548df078660bffb Mon Sep 17 00:00:00 2001
From: ananthsub <ananth.subramaniam@gmail.com>
Date: Wed, 7 Apr 2021 01:54:36 -0700
Subject: [PATCH 19/19] Update CHANGELOG.md

---
 CHANGELOG.md | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f47078ade259b..ed9b2d1586ea4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -208,6 +208,22 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed torch distributed not available in setup hook for DDP ([#6506](https://github.com/PyTorchLightning/pytorch-lightning/pull/6506))
 
+
+- Fixed TPU Colab hang issue, post training ([#6816](https://github.com/PyTorchLightning/pytorch-lightning/pull/6816))
+
+
+- Enforce an epoch scheduler interval when using SWA ([#6588](https://github.com/PyTorchLightning/pytorch-lightning/pull/6588))
+
+
+- Fixed an issue with `IterableDataset` when `__len__` is not defined ([#6828](https://github.com/PyTorchLightning/pytorch-lightning/pull/6828))
+
+
+- Fixed `EarlyStopping` logic when `min_epochs` or `min_steps` requirement is not met ([#6705](https://github.com/PyTorchLightning/pytorch-lightning/pull/6705))
+
+
+- Fixed a bug where `TensorBoardLogger` would give a warning and not log correctly to a symbolic link `save_dir` ([#6730](https://github.com/PyTorchLightning/pytorch-lightning/pull/6730))
+
+
 ## [1.2.6] - 2021-03-30
 
 ### Changed