Disable auto-detection of Kubeflow environment (#18137)

(cherry picked from commit 41f0425)

Author: awaelchli

src/lightning/fabric/CHANGELOG.md

@@ -7,6 +7,11 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [UnReleased] - 2023-08-DD
 
+### Changed
+
+- Disabled the auto-detection of the Kubeflow environment ([#18137](https://github.com/Lightning-AI/lightning/pull/18137))
+
+
 ### Fixed
 
 - Fixed issue where DDP subprocesses that used Hydra would set hydra's working directory to current directory ([#18145](https://github.com/Lightning-AI/lightning/pull/18145))

src/lightning/fabric/connector.py

@@ -33,7 +33,6 @@
 )
 from lightning.fabric.plugins.environments import (
     ClusterEnvironment,
-    KubeflowEnvironment,
     LightningEnvironment,
     LSFEnvironment,
     MPIEnvironment,
@@ -361,7 +360,6 @@ def _choose_and_init_cluster_environment(self) -> ClusterEnvironment:
         for env_type in (
             SLURMEnvironment,
             TorchElasticEnvironment,
-            KubeflowEnvironment,
             LSFEnvironment,
             MPIEnvironment,
         ):

src/lightning/fabric/plugins/environments/kubeflow.py

@@ -21,7 +21,10 @@
 
 
 class KubeflowEnvironment(ClusterEnvironment):
-    """Environment for distributed training using the `PyTorchJob`_ operator from `Kubeflow`_
+    """Environment for distributed training using the `PyTorchJob`_ operator from `Kubeflow`_.
+
+    This environment, unlike others, does not get auto-detected and needs to be passed to the Fabric/Trainer
+    constructor manually.
 
     .. _PyTorchJob: https://www.kubeflow.org/docs/components/training/pytorch/
     .. _Kubeflow: https://www.kubeflow.org
@@ -42,12 +45,7 @@ def main_port(self) -> int:
 
     @staticmethod
     def detect() -> bool:
-        """Returns ``True`` if the current process was launched using Kubeflow PyTorchJob."""
-        required_env_vars = {"KUBERNETES_PORT", "MASTER_ADDR", "MASTER_PORT", "WORLD_SIZE", "RANK"}
-        # torchelastic sets these. Make sure we're not in torchelastic
-        excluded_env_vars = {"GROUP_RANK", "LOCAL_RANK", "LOCAL_WORLD_SIZE"}
-        env_vars = os.environ.keys()
-        return required_env_vars.issubset(env_vars) and excluded_env_vars.isdisjoint(env_vars)
+        raise NotImplementedError("The Kubeflow environment can't be detected automatically.")
 
     def world_size(self) -> int:
         return int(os.environ["WORLD_SIZE"])

src/lightning/pytorch/CHANGELOG.md

@@ -7,6 +7,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [UnReleased] - 2023-08-DD
 
+### Changed
+
+- Disabled the auto-detection of the Kubeflow environment ([#18137](https://github.com/Lightning-AI/lightning/pull/18137))
+
 ### Fixed
 
 - Fixed a `Missing folder` exception when using a Google Storage URL as a `default_root_dir` ([#18088](https://github.com/Lightning-AI/lightning/pull/18088))

src/lightning/pytorch/trainer/connectors/accelerator_connector.py

@@ -22,7 +22,6 @@
 from lightning.fabric.connector import _convert_precision_to_unified_args, _PRECISION_INPUT, _PRECISION_INPUT_STR
 from lightning.fabric.plugins.environments import (
     ClusterEnvironment,
-    KubeflowEnvironment,
     LightningEnvironment,
     LSFEnvironment,
     MPIEnvironment,
@@ -402,7 +401,6 @@ def _choose_and_init_cluster_environment(self) -> ClusterEnvironment:
         for env_type in (
             SLURMEnvironment,
             TorchElasticEnvironment,
-            KubeflowEnvironment,
             LSFEnvironment,
             MPIEnvironment,
         ):

tests/tests_fabric/plugins/environments/test_kubeflow.py

@@ -74,30 +74,7 @@ def test_attributes_from_environment_variables(caplog):
     assert "setting world size is not allowed" in caplog.text
 
 
-@mock.patch.dict(
-    os.environ,
-    {
-        "KUBERNETES_PORT": "tcp://127.0.0.1:443",
-        "MASTER_ADDR": "1.2.3.4",
-        "MASTER_PORT": "500",
-        "WORLD_SIZE": "20",
-        "RANK": "1",
-    },
-)
 def test_detect_kubeflow():
-    assert KubeflowEnvironment.detect()
-
-
-@mock.patch.dict(
-    os.environ,
-    {
-        "KUBERNETES_PORT": "tcp://127.0.0.1:443",
-        "MASTER_ADDR": "1.2.3.4",
-        "MASTER_PORT": "500",
-        "WORLD_SIZE": "20",
-        "RANK": "1",
-        "GROUP_RANK": "1",
-    },
-)
-def test_detect_torchelastic_over_kubeflow():
-    assert not KubeflowEnvironment.detect()
+    """Test that the KubeflowEnvironment does not support auto-detection."""
+    with pytest.raises(NotImplementedError, match="can't be detected automatically"):
+        KubeflowEnvironment.detect()

tests/tests_fabric/test_connector.py

@@ -588,7 +588,7 @@ def test_strategy_choice_ddp_torchelastic(*_):
 @mock.patch("lightning.fabric.accelerators.cuda.num_cuda_devices", return_value=2)
 @mock.patch("lightning.fabric.accelerators.mps.MPSAccelerator.is_available", return_value=False)
 def test_strategy_choice_ddp_kubeflow(*_):
-    connector = _Connector(accelerator="gpu", devices=2)
+    connector = _Connector(accelerator="gpu", devices=2, plugins=KubeflowEnvironment())
     assert isinstance(connector.accelerator, CUDAAccelerator)
     assert isinstance(connector.strategy, DDPStrategy)
     assert isinstance(connector.strategy.cluster_environment, KubeflowEnvironment)
@@ -607,7 +607,7 @@ def test_strategy_choice_ddp_kubeflow(*_):
     },
 )
 def test_strategy_choice_ddp_cpu_kubeflow():
-    connector = _Connector(accelerator="cpu", devices=2)
+    connector = _Connector(accelerator="cpu", devices=2, plugins=KubeflowEnvironment())
     assert isinstance(connector.accelerator, CPUAccelerator)
     assert isinstance(connector.strategy, DDPStrategy)
     assert isinstance(connector.strategy.cluster_environment, KubeflowEnvironment)

tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py

@@ -244,7 +244,7 @@ def test_accelerator_choice_multi_node_gpu(cuda_count_2, tmpdir, strategy, strat
     assert isinstance(trainer.strategy, strategy_class)
 
 
-def test_accelerator_cpu(cuda_count_0):
+def test_accelerator_cpu(cuda_count_0, mps_count_0):
     trainer = Trainer(accelerator="cpu")
     assert isinstance(trainer.accelerator, CPUAccelerator)
 
@@ -481,7 +481,7 @@ def test_strategy_choice_ddp_torchelastic(_, __, mps_count_0, cuda_count_2):
 @mock.patch("torch.cuda.set_device")
 @mock.patch("lightning.pytorch.strategies.DDPStrategy.setup_distributed", autospec=True)
 def test_strategy_choice_ddp_kubeflow(_, __, mps_count_0, cuda_count_2):
-    trainer = Trainer(fast_dev_run=True, accelerator="gpu", devices=2)
+    trainer = Trainer(fast_dev_run=True, accelerator="gpu", devices=2, plugins=KubeflowEnvironment())
     assert isinstance(trainer.accelerator, CUDAAccelerator)
     assert isinstance(trainer.strategy, DDPStrategy)
     assert isinstance(trainer.strategy.cluster_environment, KubeflowEnvironment)
@@ -501,7 +501,7 @@ def test_strategy_choice_ddp_kubeflow(_, __, mps_count_0, cuda_count_2):
 )
 @mock.patch("lightning.pytorch.strategies.DDPStrategy.setup_distributed", autospec=True)
 def test_strategy_choice_ddp_cpu_kubeflow(cuda_count_0):
-    trainer = Trainer(fast_dev_run=True, accelerator="cpu", devices=2)
+    trainer = Trainer(fast_dev_run=True, accelerator="cpu", devices=2, plugins=KubeflowEnvironment())
     assert isinstance(trainer.accelerator, CPUAccelerator)
     assert isinstance(trainer.strategy, DDPStrategy)
     assert isinstance(trainer.strategy.cluster_environment, KubeflowEnvironment)