bugfix: add support for global_ordinal, local_ordinal, world_size in xla (#20872)

* bugfix: add support for global_ordinal, local_ordinal, world_size in xla
* fix: remove set local rank
* Apply suggestions from code review

---------

Co-authored-by: AlexandrByzov <[email protected]>
Co-authored-by: Bhimraj Yadav <[email protected]>
Co-authored-by: Jirka Borovec <[email protected]>

Author: 4 people

src/lightning/fabric/CHANGELOG.md

@@ -16,11 +16,17 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 -
 
+
 ### Changed
 
 - Raise ValueError when seed is `out-of-bounds` or `cannot be cast to int` ([#21029](https://github.com/Lightning-AI/pytorch-lightning/pull/21029))
 
 
+### Fixed
+
+- Fix XLA strategy to add support for `global_ordinal`, `local_ordinal`, `world_size` which came instead of deprecated methods ([#20852](https://github.com/Lightning-AI/pytorch-lightning/issues/20852))
+
+
 - fix: remove extra `name` parameter in accelerator registry decorator ([#20975](https://github.com/Lightning-AI/pytorch-lightning/pull/20975))
 
 

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

@@ -66,6 +66,11 @@ def world_size(self) -> int:
         The output is cached for performance.
 
         """
+        if _XLA_GREATER_EQUAL_2_1:
+            from torch_xla import runtime as xr
+
+            return xr.world_size()
+
         import torch_xla.core.xla_model as xm
 
         return xm.xrt_world_size()
@@ -82,6 +87,11 @@ def global_rank(self) -> int:
         The output is cached for performance.
 
         """
+        if _XLA_GREATER_EQUAL_2_1:
+            from torch_xla import runtime as xr
+
+            return xr.global_ordinal()
+
         import torch_xla.core.xla_model as xm
 
         return xm.get_ordinal()
@@ -98,6 +108,11 @@ def local_rank(self) -> int:
         The output is cached for performance.
 
         """
+        if _XLA_GREATER_EQUAL_2_1:
+            from torch_xla import runtime as xr
+
+            return xr.local_ordinal()
+
         import torch_xla.core.xla_model as xm
 
         return xm.get_local_ordinal()

tests/tests_fabric/plugins/environments/test_xla.py

@@ -97,3 +97,64 @@ def test_detect(monkeypatch):
 
     monkeypatch.setattr(lightning.fabric.accelerators.xla.XLAAccelerator, "is_available", lambda: True)
     assert XLAEnvironment.detect()
+
+
+@mock.patch.dict(os.environ, {}, clear=True)
+@mock.patch("lightning.fabric.accelerators.xla._XLA_GREATER_EQUAL_2_1", True)
+@mock.patch("lightning.fabric.plugins.environments.xla._XLA_GREATER_EQUAL_2_1", True)
+def test_world_size_from_xla_runtime_greater_2_1(xla_available):
+    """Test that world_size uses torch_xla.runtime when XLA >= 2.1."""
+    env = XLAEnvironment()
+
+    with mock.patch("torch_xla.runtime.world_size", return_value=4) as mock_world_size:
+        env.world_size.cache_clear()
+        assert env.world_size() == 4
+        mock_world_size.assert_called_once()
+
+
+@mock.patch.dict(os.environ, {}, clear=True)
+@mock.patch("lightning.fabric.accelerators.xla._XLA_GREATER_EQUAL_2_1", True)
+@mock.patch("lightning.fabric.plugins.environments.xla._XLA_GREATER_EQUAL_2_1", True)
+def test_global_rank_from_xla_runtime_greater_2_1(xla_available):
+    """Test that global_rank uses torch_xla.runtime when XLA >= 2.1."""
+    env = XLAEnvironment()
+
+    with mock.patch("torch_xla.runtime.global_ordinal", return_value=2) as mock_global_ordinal:
+        env.global_rank.cache_clear()
+        assert env.global_rank() == 2
+        mock_global_ordinal.assert_called_once()
+
+
+@mock.patch.dict(os.environ, {}, clear=True)
+@mock.patch("lightning.fabric.accelerators.xla._XLA_GREATER_EQUAL_2_1", True)
+@mock.patch("lightning.fabric.plugins.environments.xla._XLA_GREATER_EQUAL_2_1", True)
+def test_local_rank_from_xla_runtime_greater_2_1(xla_available):
+    """Test that local_rank uses torch_xla.runtime when XLA >= 2.1."""
+    env = XLAEnvironment()
+
+    with mock.patch("torch_xla.runtime.local_ordinal", return_value=1) as mock_local_ordinal:
+        env.local_rank.cache_clear()
+        assert env.local_rank() == 1
+        mock_local_ordinal.assert_called_once()
+
+
+@mock.patch.dict(os.environ, {}, clear=True)
+@mock.patch("lightning.fabric.accelerators.xla._XLA_GREATER_EQUAL_2_1", True)
+@mock.patch("lightning.fabric.plugins.environments.xla._XLA_GREATER_EQUAL_2_1", True)
+def test_setters_readonly_when_xla_runtime_greater_2_1(xla_available):
+    """Test that set_world_size and set_global_rank don't affect values when using XLA runtime >= 2.1."""
+    env = XLAEnvironment()
+
+    with (
+        mock.patch("torch_xla.runtime.world_size", return_value=4),
+        mock.patch("torch_xla.runtime.global_ordinal", return_value=2),
+    ):
+        env.world_size.cache_clear()
+        env.global_rank.cache_clear()
+
+        # Values should come from XLA runtime and not be affected by setters
+        env.set_world_size(100)
+        assert env.world_size() == 4
+
+        env.set_global_rank(100)
+        assert env.global_rank() == 2