Update logic for checking TPUs availability (#6767)

* Update logic for checking TPUs availability

* fix flake8

* add fix

Author: kaushikb11

pytorch_lightning/plugins/training_type/tpu_spawn.py

@@ -1,6 +1,7 @@
 import io
 import os
 import re
+import time
 from typing import Any, Dict, Iterable, List, Optional, Union
 
 import torch
@@ -106,6 +107,9 @@ def new_process(self, process_idx: int, trainer, mp_queue) -> None:
         self.__save_end_of_training_weights(self.lightning_module)
         self.transfer_distrib_spawn_state_on_fit_end(results)
 
+        if self.global_rank == 0:
+            time.sleep(2)
+
         self.barrier("end-process")
 
     def __save_end_of_training_weights(self, model: LightningModule) -> None:

pytorch_lightning/utilities/xla_device.py

@@ -17,13 +17,10 @@
 import traceback
 from multiprocessing import Process, Queue
 
-import torch.multiprocessing as mp
-
 from pytorch_lightning.utilities.imports import _XLA_AVAILABLE
 
 if _XLA_AVAILABLE:
     import torch_xla.core.xla_model as xm
-    import torch_xla.distributed.xla_multiprocessing as xmp
 
 #: define waiting time got checking TPU available in sec
 TPU_CHECK_TIMEOUT = 25
@@ -64,23 +61,13 @@ class XLADeviceUtils:
     @pl_multi_process
     def _is_device_tpu() -> bool:
         """
-        Check if device is TPU
+        Check if TPU devices are available
 
         Return:
-            A boolean value indicating if the xla device is a TPU device or not
+            A boolean value indicating if TPU devices are available
         """
 
-        def _fn(_: int, mp_queue):
-            try:
-                device = xm.xla_device()
-                mp_queue.put(device.type == 'xla')
-            except Exception:
-                mp_queue.put(False)
-
-        smp = mp.get_context("spawn")
-        queue = smp.SimpleQueue()
-        xmp.spawn(_fn, args=(queue, ), nprocs=1)
-        return queue.get()
+        return len(xm.get_xla_supported_devices("TPU")) > 0
 
     @staticmethod
     def xla_available() -> bool: