Refactor TorchElasticEnvironment.detect to use torch.distributed.is_torchelastic_launched (#12376)

* Refactor TorchElasticEnvironment.detect to use native utility from torch.distributed

* fix version and tests

* fix version

* Update tests/accelerators/test_accelerator_connector.py

Co-authored-by: Adrian Wälchli <[email protected]>

Co-authored-by: Adrian Wälchli <[email protected]>

Author: daniellepintz

pytorch_lightning/plugins/environments/torchelastic_environment.py

@@ -15,7 +15,10 @@
 import logging
 import os
 
+import torch.distributed
+
 from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment
+from pytorch_lightning.utilities.imports import _TORCH_GREATER_EQUAL_1_9_1
 from pytorch_lightning.utilities.rank_zero import rank_zero_deprecation, rank_zero_warn
 
 log = logging.getLogger(__name__)
@@ -58,6 +61,8 @@ def main_port(self) -> int:
     @staticmethod
     def detect() -> bool:
         """Returns ``True`` if the current process was launched using the torchelastic command."""
+        if _TORCH_GREATER_EQUAL_1_9_1:
+            return torch.distributed.is_torchelastic_launched()
         required_env_vars = {"RANK", "GROUP_RANK", "LOCAL_RANK", "LOCAL_WORLD_SIZE"}
         return required_env_vars.issubset(os.environ.keys())
 

pytorch_lightning/utilities/imports.py

@@ -92,6 +92,7 @@ def _compare_version(package: str, op: Callable, version: str, use_base_version:
 _TORCH_GREATER_EQUAL_1_8 = _compare_version("torch", operator.ge, "1.8.0")
 _TORCH_GREATER_EQUAL_1_8_1 = _compare_version("torch", operator.ge, "1.8.1")
 _TORCH_GREATER_EQUAL_1_9 = _compare_version("torch", operator.ge, "1.9.0")
+_TORCH_GREATER_EQUAL_1_9_1 = _compare_version("torch", operator.ge, "1.9.1")
 _TORCH_GREATER_EQUAL_1_10 = _compare_version("torch", operator.ge, "1.10.0")
 _TORCH_GREATER_EQUAL_1_11 = _compare_version("torch", operator.ge, "1.11.0")
 

tests/accelerators/test_accelerator_connector.py

@@ -147,6 +147,7 @@ def test_accelerator_choice_ddp2_slurm(*_):
         "RANK": "1",
         "LOCAL_RANK": "1",
         "GROUP_RANK": "0",
+        "TORCHELASTIC_RUN_ID": "1",  # present for torch >= 1.9.1
     },
 )
 @mock.patch("torch.cuda.set_device")
@@ -172,6 +173,7 @@ def test_accelerator_choice_ddp_te(*_):
         "RANK": "1",
         "LOCAL_RANK": "1",
         "GROUP_RANK": "0",
+        "TORCHELASTIC_RUN_ID": "1",
     },
 )
 @mock.patch("torch.cuda.set_device")
@@ -189,7 +191,15 @@ def test_accelerator_choice_ddp2_te(*_):
 
 
 @mock.patch.dict(
-    os.environ, {"WORLD_SIZE": "2", "LOCAL_WORLD_SIZE": "2", "RANK": "1", "LOCAL_RANK": "1", "GROUP_RANK": "0"}
+    os.environ,
+    {
+        "WORLD_SIZE": "2",
+        "LOCAL_WORLD_SIZE": "2",
+        "RANK": "1",
+        "LOCAL_RANK": "1",
+        "GROUP_RANK": "0",
+        "TORCHELASTIC_RUN_ID": "1",
+    },
 )
 @mock.patch("torch.cuda.device_count", return_value=0)
 @mock.patch("pytorch_lightning.strategies.DDPStrategy.setup_distributed", autospec=True)
@@ -786,6 +796,7 @@ def test_strategy_choice_ddp2_slurm(
         "RANK": "1",
         "LOCAL_RANK": "1",
         "GROUP_RANK": "0",
+        "TORCHELASTIC_RUN_ID": "1",
     },
 )
 @mock.patch("torch.cuda.set_device")
@@ -810,6 +821,7 @@ def test_strategy_choice_ddp_te(*_):
         "RANK": "1",
         "LOCAL_RANK": "1",
         "GROUP_RANK": "0",
+        "TORCHELASTIC_RUN_ID": "1",
     },
 )
 @mock.patch("torch.cuda.set_device")
@@ -826,7 +838,15 @@ def test_strategy_choice_ddp2_te(*_):
 
 
 @mock.patch.dict(
-    os.environ, {"WORLD_SIZE": "2", "LOCAL_WORLD_SIZE": "2", "RANK": "1", "LOCAL_RANK": "1", "GROUP_RANK": "0"}
+    os.environ,
+    {
+        "WORLD_SIZE": "2",
+        "LOCAL_WORLD_SIZE": "2",
+        "RANK": "1",
+        "LOCAL_RANK": "1",
+        "GROUP_RANK": "0",
+        "TORCHELASTIC_RUN_ID": "1",
+    },
 )
 @mock.patch("torch.cuda.device_count", return_value=0)
 @mock.patch("pytorch_lightning.strategies.DDPStrategy.setup_distributed", autospec=True)

tests/models/test_gpu.py

@@ -205,6 +205,7 @@ def test_parse_gpu_returns_none_when_no_devices_are_available(mocked_device_coun
         "RANK": "3",
         "WORLD_SIZE": "4",
         "LOCAL_WORLD_SIZE": "2",
+        "TORCHELASTIC_RUN_ID": "1",
     },
 )
 @mock.patch("torch.cuda.device_count", return_value=1)

tests/plugins/environments/test_torchelastic_environment.py

@@ -18,6 +18,7 @@
 import pytest
 
 from pytorch_lightning.plugins.environments import TorchElasticEnvironment
+from tests.helpers.runif import RunIf
 
 
 @mock.patch.dict(os.environ, {})
@@ -70,8 +71,9 @@ def test_attributes_from_environment_variables(caplog):
     assert "setting world size is not allowed" in caplog.text
 
 
-def test_detect():
-    """Test the detection of a torchelastic environment configuration."""
+@RunIf(max_torch="1.9.0")
+def test_detect_before_1_9_1():
+    """Test the detection of a torchelastic environment configuration before 1.9.1."""
     with mock.patch.dict(os.environ, {}):
         assert not TorchElasticEnvironment.detect()
 
@@ -85,3 +87,18 @@ def test_detect():
         },
     ):
         assert TorchElasticEnvironment.detect()
+
+
+@RunIf(min_torch="1.9.1")
+def test_detect_after_1_9_1():
+    """Test the detection of a torchelastic environment configuration after 1.9.1."""
+    with mock.patch.dict(os.environ, {}):
+        assert not TorchElasticEnvironment.detect()
+
+    with mock.patch.dict(
+        os.environ,
+        {
+            "TORCHELASTIC_RUN_ID": "",
+        },
+    ):
+        assert TorchElasticEnvironment.detect()

tests/plugins/test_cluster_integration.py

@@ -49,6 +49,7 @@ def environment_combinations():
         "RANK": "3",
         "WORLD_SIZE": "4",
         "LOCAL_WORLD_SIZE": "2",
+        "TORCHELASTIC_RUN_ID": "1",
     }
     environment = TorchElasticEnvironment()
     yield environment, variables, expected