fabric tests

Author: justusschock

src/lightning/fabric/strategies/launchers/xla.py

@@ -42,14 +42,14 @@ class _XLALauncher(_Launcher):
     def __init__(self, strategy: Union["XLAStrategy", "XLAFSDPStrategy"]) -> None:
         super().__init__()
         _raise_enterprise_not_available()
-        from pytorch_lightning_enterprise.strategies.launchers.xla import _XLALauncher as EnterpriseXLALauncher
+        from pytorch_lightning_enterprise.strategies.xla.launcher import XLALauncherFabric as EnterpriseXLALauncher
 
         self.xla_impl = EnterpriseXLALauncher(strategy=strategy)
 
     @property
     @override
     def is_interactive_compatible(self) -> bool:
-        return self.xla_impl.is_interactive_compatible()
+        return self.xla_impl.is_interactive_compatible
 
     @override
     def launch(self, function: Callable, *args: Any, **kwargs: Any) -> Any:
@@ -65,3 +65,11 @@ def launch(self, function: Callable, *args: Any, **kwargs: Any) -> Any:
 
         """
         return self.xla_impl.launch(function=function, *args, **kwargs)
+
+    @property
+    def _start_method(self) -> str:
+        return self.xla_impl._start_method
+
+    @_start_method.setter
+    def _start_method(self, start_method: str) -> None:
+        self.xla_impl._start_method = start_method

tests/tests_fabric/conftest.py

@@ -153,6 +153,8 @@ def mock_xla_available(monkeypatch: pytest.MonkeyPatch, value: bool = True) -> N
     monkeypatch.setitem(sys.modules, "torch_xla.core.xla_model", Mock())
     monkeypatch.setitem(sys.modules, "torch_xla.experimental", Mock())
     monkeypatch.setitem(sys.modules, "torch_xla.distributed.fsdp.wrap", Mock())
+    monkeypatch.setitem(sys.modules, "torch_xla._internal", Mock())
+    monkeypatch.setitem(sys.modules, "torch_xla._internal.tpu", Mock())
 
     # Then patch the _XLA_AVAILABLE flags in various modules
     monkeypatch.setattr(pytorch_lightning_enterprise.utils.imports, "_XLA_AVAILABLE", value)
@@ -161,6 +163,12 @@ def mock_xla_available(monkeypatch: pytest.MonkeyPatch, value: bool = True) -> N
     monkeypatch.setattr(lightning.fabric.accelerators.xla, "_XLA_AVAILABLE", value)
     monkeypatch.setattr(lightning.fabric.accelerators.xla, "_XLA_GREATER_EQUAL_2_1", value)
     monkeypatch.setattr(lightning.fabric.accelerators.xla, "_XLA_GREATER_EQUAL_2_5", value)
+    # Patch in the modules where they're used after import
+    monkeypatch.setattr("pytorch_lightning_enterprise.accelerators.xla._XLA_AVAILABLE", value)
+    monkeypatch.setattr("pytorch_lightning_enterprise.accelerators.xla._XLA_GREATER_EQUAL_2_1", value)
+    monkeypatch.setattr("pytorch_lightning_enterprise.accelerators.xla._XLA_GREATER_EQUAL_2_5", value)
+    monkeypatch.setattr("pytorch_lightning_enterprise.plugins.environments.xla._XLA_AVAILABLE", value)
+    monkeypatch.setattr("pytorch_lightning_enterprise.plugins.environments.xla._XLA_GREATER_EQUAL_2_1", value)
 
 
 @pytest.fixture
@@ -172,6 +180,11 @@ def mock_tpu_available(monkeypatch: pytest.MonkeyPatch, value: bool = True) -> N
     mock_xla_available(monkeypatch, value)
     monkeypatch.setattr(lightning.fabric.accelerators.xla.XLAAccelerator, "is_available", lambda: value)
     monkeypatch.setattr(lightning.fabric.accelerators.xla.XLAAccelerator, "auto_device_count", lambda *_: 8)
+    # Also mock the enterprise XLAAccelerator methods
+    import pytorch_lightning_enterprise.accelerators.xla
+
+    monkeypatch.setattr(pytorch_lightning_enterprise.accelerators.xla.XLAAccelerator, "is_available", lambda: value)
+    monkeypatch.setattr(pytorch_lightning_enterprise.accelerators.xla.XLAAccelerator, "auto_device_count", lambda *_: 8)
 
 
 @pytest.fixture

tests/tests_fabric/test_connector.py

@@ -242,8 +242,10 @@ class TestStrategy(DDPStrategy):
         ),
     ],
 )
-@mock.patch("lightning.fabric.plugins.environments.lsf.LSFEnvironment._read_hosts", return_value=["node0", "node1"])
-@mock.patch("lightning.fabric.plugins.environments.lsf.LSFEnvironment._get_node_rank", return_value=0)
+@mock.patch(
+    "pytorch_lightning_enterprise.plugins.environments.lsf.LSFEnvironment._read_hosts", return_value=["node0", "node1"]
+)
+@mock.patch("pytorch_lightning_enterprise.plugins.environments.lsf.LSFEnvironment._get_node_rank", return_value=0)
 def test_fallback_from_ddp_spawn_to_ddp_on_cluster(_, __, env_vars, expected_environment):
     with mock.patch.dict(os.environ, env_vars, clear=True):
         connector = _Connector(strategy="ddp_spawn", accelerator="cpu", devices=2)
@@ -341,8 +343,8 @@ def test_cuda_accelerator_can_not_run_on_system(_):
 
 
 @pytest.mark.skipif(XLAAccelerator.is_available(), reason="test requires missing TPU")
-@mock.patch("lightning.fabric.accelerators.xla._XLA_AVAILABLE", True)
-@mock.patch("lightning.fabric.accelerators.xla._using_pjrt", return_value=True)
+@mock.patch("pytorch_lightning_enterprise.accelerators.xla._XLA_AVAILABLE", True)
+@mock.patch("pytorch_lightning_enterprise.accelerators.xla._using_pjrt", return_value=True)
 def test_tpu_accelerator_can_not_run_on_system(_):
     with pytest.raises(RuntimeError, match="XLAAccelerator` can not run on your system"):
         _Connector(accelerator="tpu", devices=8)

tests/tests_fabric/utilities/test_throughput.py

@@ -45,7 +45,13 @@ def test_get_available_flops(xla_available):
     ):
         assert get_available_flops(torch.device("cuda"), torch.bfloat16) is None
 
-    from torch_xla.experimental import tpu
+    # Import from the right module based on _XLA_GREATER_EQUAL_2_1
+    from lightning.fabric.accelerators.xla import _XLA_GREATER_EQUAL_2_1
+
+    if _XLA_GREATER_EQUAL_2_1:
+        from torch_xla._internal import tpu
+    else:
+        from torch_xla.experimental import tpu
 
     assert isinstance(tpu, Mock)