replaced all instances of XLA_AVAILABLE

lezwon · lezwon · commit ecae7d24a4fe · 2020-09-01T00:48:58.000+05:30
diff --git a/pytorch_lightning/accelerators/tpu_backend.py b/pytorch_lightning/accelerators/tpu_backend.py
@@ -22,16 +22,12 @@
 from pytorch_lightning.utilities import rank_zero_info, rank_zero_only, rank_zero_warn
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.accelerators.base_backend import Accelerator
+from pytorch_lightning.utilities.xla_device_utils import TPU_AVAILABLE
 
-try:
-    import torch_xla
+if TPU_AVAILABLE:
     import torch_xla.core.xla_model as xm
     import torch_xla.distributed.xla_multiprocessing as xmp
     import torch_xla.distributed.parallel_loader as xla_pl
-except ImportError:
-    XLA_AVAILABLE = False
-else:
-    XLA_AVAILABLE = True
 
 
 class TPUBackend(Accelerator):
@@ -44,7 +40,7 @@ def __init__(self, trainer):
     def setup(self, model):
         rank_zero_info(f'training on {self.trainer.tpu_cores} TPU cores')
 
-        if not XLA_AVAILABLE:
+        if not TPU_AVAILABLE:
             raise MisconfigurationException('PyTorch XLA not installed.')
 
         # see: https://discuss.pytorch.org/t/segfault-with-multiprocessing-queue/81292/2
@@ -164,7 +160,7 @@ def to_device(self, batch):
         See Also:
             - :func:`~pytorch_lightning.utilities.apply_func.move_data_to_device`
         """
-        if not XLA_AVAILABLE:
+        if not TPU_AVAILABLE:
             raise MisconfigurationException(
                 'Requested to transfer batch to TPU but XLA is not available.'
                 ' Are you sure this machine has TPUs?'
diff --git a/pytorch_lightning/callbacks/early_stopping.py b/pytorch_lightning/callbacks/early_stopping.py
@@ -26,17 +26,16 @@
 from pytorch_lightning import _logger as log
 from pytorch_lightning.callbacks.base import Callback
 from pytorch_lightning.utilities import rank_zero_warn
+from pytorch_lightning.utilities.xla_device_utils import TPU_AVAILABLE
 import os
 
+
 torch_inf = torch.tensor(np.Inf)
 
-try:
-    import torch_xla
+if TPU_AVAILABLE:
     import torch_xla.core.xla_model as xm
-except ImportError:
-    XLA_AVAILABLE = False
-else:
-    XLA_AVAILABLE = True
+    import torch_xla
+
 
 
 class EarlyStopping(Callback):
@@ -201,7 +200,7 @@ def _run_early_stopping_check(self, trainer, pl_module):
         if not isinstance(current, torch.Tensor):
             current = torch.tensor(current, device=pl_module.device)
 
-        if trainer.use_tpu and XLA_AVAILABLE:
+        if trainer.use_tpu and TPU_AVAILABLE:
             current = current.cpu()
 
         if self.monitor_op(current - self.min_delta, self.best_score):
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
@@ -39,13 +39,10 @@
 from pytorch_lightning.utilities.device_dtype_mixin import DeviceDtypeModuleMixin
 from pytorch_lightning.utilities.parsing import AttributeDict, collect_init_args, get_init_args
 from pytorch_lightning.core.step_result import TrainResult, EvalResult
+from pytorch_lightning.utilities.xla_device_utils import TPU_AVAILABLE
 
-try:
+if TPU_AVAILABLE:
     import torch_xla.core.xla_model as xm
-except ImportError:
-    XLA_AVAILABLE = False
-else:
-    XLA_AVAILABLE = True
 
 
 class LightningModule(ABC, DeviceDtypeModuleMixin, GradInformation, ModelIO, ModelHooks, Module):
diff --git a/pytorch_lightning/trainer/data_loading.py b/pytorch_lightning/trainer/data_loading.py
@@ -27,21 +27,16 @@
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.debugging import InternalDebugger
 from pytorch_lightning.utilities.model_utils import is_overridden
-
+from pytorch_lightning.utilities.xla_device_utils import TPU_AVAILABLE
 
 try:
     from apex import amp
 except ImportError:
     amp = None
 
-try:
+if TPU_AVAILABLE:
     import torch_xla
     import torch_xla.core.xla_model as xm
-    import torch_xla.distributed.xla_multiprocessing as xmp
-except ImportError:
-    XLA_AVAILABLE = False
-else:
-    XLA_AVAILABLE = True
 
 try:
     import horovod.torch as hvd
@@ -336,7 +331,7 @@ def request_dataloader(self, dataloader_fx: Callable) -> DataLoader:
             torch_distrib.barrier()
 
         # data download/load on TPU
-        elif self.use_tpu and XLA_AVAILABLE:
+        elif self.use_tpu and TPU_AVAILABLE:
             # all processes wait until data download has happened
             torch_xla.core.xla_model.rendezvous('pl.TrainerDataLoadingMixin.get_dataloaders')
 
diff --git a/pytorch_lightning/trainer/distrib_data_parallel.py b/pytorch_lightning/trainer/distrib_data_parallel.py
@@ -141,6 +141,7 @@ def train_fx(trial_hparams, cluster_manager, _):
 from pytorch_lightning.utilities.cloud_io import atomic_save
 from pytorch_lightning.utilities.distributed import rank_zero_warn, rank_zero_info
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning.utilities.xla_device_utils import TPU_AVAILABLE
 
 try:
     from apex import amp
@@ -155,13 +156,6 @@ def train_fx(trial_hparams, cluster_manager, _):
     HOROVOD_AVAILABLE = True
 
 
-try:
-    import torch_xla
-except ImportError:
-    XLA_AVAILABLE = False
-else:
-    XLA_AVAILABLE = True
-
 
 class TrainerDDPMixin(ABC):
 
@@ -303,7 +297,7 @@ def set_distributed_mode(self, distributed_backend):
 
         rank_zero_info(f'GPU available: {torch.cuda.is_available()}, used: {self.on_gpu}')
         num_cores = self.tpu_cores if self.tpu_cores is not None else 0
-        rank_zero_info(f'TPU available: {XLA_AVAILABLE}, using: {num_cores} TPU cores')
+        rank_zero_info(f'TPU available: {TPU_AVAILABLE}, using: {num_cores} TPU cores')
 
         if torch.cuda.is_available() and not self.on_gpu:
             rank_zero_warn('GPU available but not used. Set the --gpus flag when calling the script.')
diff --git a/pytorch_lightning/trainer/distrib_parts.py b/pytorch_lightning/trainer/distrib_parts.py
@@ -42,12 +42,6 @@
 except ImportError:
     amp = None
 
-try:
-    import torch_xla.core.xla_model as xm
-except ImportError:
-    XLA_AVAILABLE = False
-else:
-    XLA_AVAILABLE = True
 
 try:
     import horovod.torch as hvd
diff --git a/pytorch_lightning/trainer/evaluation_loop.py b/pytorch_lightning/trainer/evaluation_loop.py
@@ -135,14 +135,6 @@
 from pytorch_lightning.core.step_result import EvalResult, Result
 from pytorch_lightning.trainer.evaluate_loop import EvaluationLoop
 
-try:
-    import torch_xla.distributed.parallel_loader as xla_pl
-    import torch_xla.core.xla_model as xm
-except ImportError:
-    XLA_AVAILABLE = False
-else:
-    XLA_AVAILABLE = True
-
 try:
     import horovod.torch as hvd
 except (ModuleNotFoundError, ImportError):
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
@@ -59,6 +59,8 @@
 from pytorch_lightning.utilities.model_utils import is_overridden
 
 # warnings to ignore in trainer
+from pytorch_lightning.utilities.xla_device_utils import TPU_AVAILABLE
+
 warnings.filterwarnings(
     'ignore', message='torch.distributed.reduce_op is deprecated, ' 'please use torch.distributed.ReduceOp instead'
 )
@@ -68,14 +70,8 @@
 except ImportError:
     amp = None
 
-try:
+if TPU_AVAILABLE:
     import torch_xla
-    import torch_xla.core.xla_model as xm
-    import torch_xla.distributed.xla_multiprocessing as xmp
-except ImportError:
-    XLA_AVAILABLE = False
-else:
-    XLA_AVAILABLE = True
 
 try:
     import horovod.torch as hvd
@@ -1131,7 +1127,7 @@ def run_pretrain_routine(self, model: LightningModule):
             torch_distrib.barrier()
 
         # wait for all models to restore weights
-        if self.on_tpu and XLA_AVAILABLE:
+        if self.on_tpu and TPU_AVAILABLE:
             # wait for all processes to catch up
             torch_xla.core.xla_model.rendezvous("pl.Trainer.run_pretrain_routine")
 
@@ -1389,7 +1385,7 @@ def barrier(self, name):
             pass
             # torch_distrib.barrier()
 
-        if self.on_tpu and XLA_AVAILABLE:
+        if self.on_tpu and TPU_AVAILABLE:
             # wait for all processes to catch up
             torch_xla.core.xla_model.rendezvous(f'pl.Trainer.{name}')
 
diff --git a/pytorch_lightning/trainer/training_io.py b/pytorch_lightning/trainer/training_io.py
@@ -118,15 +118,10 @@
 from pytorch_lightning.utilities.cloud_io import load as pl_load
 from pytorch_lightning.utilities.cloud_io import makedirs
 from pytorch_lightning.utilities.upgrade_checkpoint import KEYS_MAPPING as DEPRECATED_CHECKPOINT_KEYS
+from pytorch_lightning.utilities.xla_device_utils import TPU_AVAILABLE
 
-try:
+if TPU_AVAILABLE:
     import torch_xla
-    import torch_xla.core.xla_model as xm
-    import torch_xla.distributed.xla_multiprocessing as xmp
-except ImportError:
-    XLA_AVAILABLE = False
-else:
-    XLA_AVAILABLE = True
 
 try:
     from apex import amp
@@ -209,7 +204,7 @@ def restore_weights(self, model: LightningModule):
             torch_distrib.barrier()
 
         # wait for all models to restore weights
-        if self.on_tpu and XLA_AVAILABLE:
+        if self.on_tpu and TPU_AVAILABLE:
             # wait for all processes to catch up
             torch_xla.core.xla_model.rendezvous("pl.TrainerIOMixin.restore_weights")
 
diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
@@ -182,19 +182,13 @@ def training_step(self, batch, batch_idx):
 from pytorch_lightning.utilities.memory import recursive_detach
 from pytorch_lightning.utilities.parsing import AttributeDict
 from pytorch_lightning.utilities.model_utils import is_overridden
+from pytorch_lightning.utilities.xla_device_utils import TPU_AVAILABLE
 
 try:
     from apex import amp
 except ImportError:
     amp = None
 
-try:
-    import torch_xla.distributed.parallel_loader as xla_pl
-    import torch_xla.core.xla_model as xm
-except ImportError:
-    XLA_AVAILABLE = False
-else:
-    XLA_AVAILABLE = True
 
 try:
     import horovod.torch as hvd
@@ -957,7 +951,7 @@ def call_optimizer_step(self, optimizer, opt_idx, batch_idx, split_batch):
             ).loss
 
             # apply TPU optimizer
-            if self.use_tpu and XLA_AVAILABLE:
+            if self.use_tpu and TPU_AVAILABLE:
                 model.optimizer_step(self.current_epoch, batch_idx,
                                      optimizer, opt_idx, lambda_closure, on_tpu=True)
 
diff --git a/pytorch_lightning/utilities/xla_device_utils.py b/pytorch_lightning/utilities/xla_device_utils.py
@@ -0,0 +1,54 @@
+import functools
+import importlib
+
+TORCHXLA_AVAILABLE = importlib.util.find_spec("torch_xla") is not None
+if TORCHXLA_AVAILABLE:
+    import torch_xla.core.xla_model as xm
+else:
+    xm = None
+
+
+def pl_multi_process(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+
+        from multiprocessing import Process, Queue
+
+        queue = Queue()
+
+        def inner_f():
+            try:
+                queue.put(func(**kwargs))
+            except Exception:
+                import traceback
+
+                traceback.print_exc()
+                queue.put(None)
+
+        proc = Process(target=inner_f, kwargs=kwargs)
+        proc.start()
+        proc.join()
+
+        return queue.get()
+
+    return wrapper
+
+
+def fetch_xla_device_type(device):
+    if xm is not None:
+        return xm.xla_device_hw(device)
+    else:
+        return None
+
+
+@pl_multi_process
+def tpu_device_exists():
+    if xm is not None:
+        device = xm.xla_device()
+        device_type = fetch_xla_device_type(device)
+        return device_type == "TPU"
+    else:
+        return False
+
+
+TPU_AVAILABLE = tpu_device_exists()
diff --git a/tests/models/test_tpu.py b/tests/models/test_tpu.py
@@ -7,26 +7,13 @@
 import tests.base.develop_pipelines as tpipes
 from pytorch_lightning import Trainer
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning.utilities.xla_device_utils import TPU_AVAILABLE
 from tests.base import EvalModelTemplate
 from tests.base.datasets import TrialMNIST
 from tests.base.develop_utils import pl_multi_process_test
 
-try:
+if TPU_AVAILABLE:
     import torch_xla
-    import torch_xla.distributed.xla_multiprocessing as xmp
-
-    def tpu_device_exists(q):
-        device = torch_xla.core.xla_model.xla_device()
-        device_type = torch_xla.core.xla_model.xla_device_hw(device)
-        q.put(device_type == 'TPU')
-
-    queue = Queue()
-    p = Process(target=tpu_device_exists, args=(queue, ))
-    p.start()
-    p.join()
-    TPU_AVAILABLE = queue.get()
-except ImportError:
-    TPU_AVAILABLE = False
 
 
 _LARGER_DATASET = TrialMNIST(download=True, num_samples=2000, digits=(0, 1, 2, 5, 8))
@@ -224,7 +211,6 @@ def test_tpu_misconfiguration():
         Trainer(tpu_cores=[1, 8])
 
 
-# @patch('pytorch_lightning.trainer.trainer.XLA_AVAILABLE', False)
 @pytest.mark.skipif(TPU_AVAILABLE, reason="test requires missing TPU")
 def test_exception_when_no_tpu_found(tmpdir):
     """Test if exception is thrown when xla devices are not available"""
