Don't import torch_xla.debug for torch-xla<1.8

CHANGELOG.md

@@ -199,8 +199,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed the default logging level for batch hooks associated with training from `on_step=False, on_epoch=True` to `on_step=True, on_epoch=False` ([#10756](https://github.com/PyTorchLightning/pytorch-lightning/pull/10756))
 
 
-
--
+- Fixed importing `torch_xla.debug` for `torch-xla<1.8` ([#10836](https://github.com/PyTorchLightning/pytorch-lightning/pull/10836))
 
 
 -

pytorch_lightning/profiler/xla.py

@@ -42,9 +42,10 @@
 from typing import Dict
 
 from pytorch_lightning.profiler.base import BaseProfiler
-from pytorch_lightning.utilities import _TPU_AVAILABLE
+from pytorch_lightning.utilities import _TORCH_GREATER_EQUAL_1_8, _TPU_AVAILABLE
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
-if _TPU_AVAILABLE:
+if _TPU_AVAILABLE and _TORCH_GREATER_EQUAL_1_8:
     import torch_xla.debug.profiler as xp
 
 log = logging.getLogger(__name__)
@@ -65,6 +66,10 @@ class XLAProfiler(BaseProfiler):
     def __init__(self, port: int = 9012) -> None:
         """This Profiler will help you debug and optimize training workload performance for your models using Cloud
         TPU performance tools."""
+        if not _TPU_AVAILABLE:
+            raise MisconfigurationException("`XLAProfiler` is only supported on TPUs")
+        if not _TORCH_GREATER_EQUAL_1_8:
+            raise MisconfigurationException("`XLAProfiler` is only supported with `torch-xla>=1.8`")
         super().__init__(dirpath=None, filename=None)
         self.port = port
         self._recording_map: Dict = {}

tests/profiler/test_xla_profiler.py

@@ -13,19 +13,23 @@
 # limitations under the License.
 import os
 from multiprocessing import Event, Process
+from unittest.mock import patch
 
 import pytest
 
 from pytorch_lightning import Trainer
 from pytorch_lightning.profiler import XLAProfiler
-from pytorch_lightning.utilities import _TPU_AVAILABLE
+from pytorch_lightning.utilities import _TORCH_GREATER_EQUAL_1_8, _TPU_AVAILABLE
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers import BoringModel
 from tests.helpers.runif import RunIf
 
 if _TPU_AVAILABLE:
-    import torch_xla.debug.profiler as xp
     import torch_xla.utils.utils as xu
 
+    if _TORCH_GREATER_EQUAL_1_8:
+        import torch_xla.debug.profiler as xp
+
 
 @RunIf(tpu=True)
 def test_xla_profiler_instance(tmpdir):
@@ -60,3 +64,16 @@ def train_worker():
     p.terminate()
 
     assert os.isfile(os.path.join(logdir, "plugins", "profile", "*", "*.xplane.pb"))
+
+
+@patch("pytorch_lightning.utilities.imports._TPU_AVAILABLE", return_value=False)
+def test_xla_profiler_tpu_not_available_exception(*_):
+    with pytest.raises(MisconfigurationException, match="`XLAProfiler` is only supported on TPUs"):
+        _ = XLAProfiler()
+
+
+@RunIf(max_torch="1.8.0")
+@patch("pytorch_lightning.utilities.imports._TPU_AVAILABLE", return_value=True)
+def test_xla_profiler_torch_lesser_than_1_8_exception(*_):
+    with pytest.raises(MisconfigurationException, match="`XLAProfiler` is only supported with `torch-xla>=1.8`"):
+        _ = XLAProfiler()