From b7a22ba046ba57072a71b12d16caff000e66f798 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= Date: Tue, 6 Apr 2021 18:37:17 +0200 Subject: [PATCH] CI: fixture for global rank variable reset (#6839) --- tests/conftest.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/conftest.py b/tests/conftest.py index 9bc607e119451..7f6407ecfd82b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -21,6 +21,16 @@ import torch.multiprocessing as mp +@pytest.fixture(scope="function", autouse=True) +def preserve_global_rank_variable(): + """ Ensures that the rank_zero_only.rank global variable gets reset in each test. """ + from pytorch_lightning.utilities.distributed import rank_zero_only + rank = getattr(rank_zero_only, "rank", None) + yield + if rank is not None: + setattr(rank_zero_only, "rank", rank) + + @pytest.fixture(scope="function", autouse=True) def restore_env_variables(): """ Ensures that environment variables set during the test do not leak out. """