Don't import torch_xla.debug for torch-xla<1.8 (#10836)

Author: kaushikb11

CHANGELOG.md

@@ -17,7 +17,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed early schedule reset logic in PyTorch profiler that was causing data leak ([#10837](https://github.com/PyTorchLightning/pytorch-lightning/pull/10837))
 - Fixed a bug that caused incorrect batch indices to be passed to the `BasePredictionWriter` hooks when using a dataloader with `num_workers > 0` ([#10870](https://github.com/PyTorchLightning/pytorch-lightning/pull/10870))
 - Fixed an issue with item assignment on the logger on rank > 0 for those who support it ([#10917](https://github.com/PyTorchLightning/pytorch-lightning/pull/10917))
-
+- Fixed importing `torch_xla.debug` for `torch-xla<1.8` ([#10836](https://github.com/PyTorchLightning/pytorch-lightning/pull/10836))
 
 
 ## [1.5.4] - 2021-11-30

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

@@ -18,14 +18,16 @@
 
 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 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):