Refactor XLADeviceUtils

Author: carmocca

docs/source-lit/conf.py

@@ -406,8 +406,6 @@ def find_source():
 from pytorch_lightning.cli import LightningCLI
 from pytorch_lightning.utilities import (
     _APEX_AVAILABLE,
-    _XLA_AVAILABLE,
-    _TPU_AVAILABLE,
     _TORCHVISION_AVAILABLE,
     _TORCH_GREATER_EQUAL_1_10,
 )

docs/source-pytorch/conf.py

@@ -394,8 +394,6 @@ def package_list_from_file(file):
 from pytorch_lightning.cli import _JSONARGPARSE_SIGNATURES_AVAILABLE as _JSONARGPARSE_AVAILABLE
 from pytorch_lightning.utilities import (
     _APEX_AVAILABLE,
-    _XLA_AVAILABLE,
-    _TPU_AVAILABLE,
     _TORCHVISION_AVAILABLE,
     _TORCH_GREATER_EQUAL_1_10,
 )

src/lightning_lite/accelerators/tpu.py

@@ -11,18 +11,27 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, List, Optional, Union
+import functools
+import queue as q
+import traceback
+from multiprocessing import Process, Queue
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import torch
+from lightning_utilities.core.imports import RequirementCache
 
 from lightning_lite.accelerators.accelerator import Accelerator
 from lightning_lite.utilities.device_parser import _check_data_type
-from lightning_lite.utilities.imports import _TPU_AVAILABLE
 
 
 class TPUAccelerator(Accelerator):
     """Accelerator for TPU devices."""
 
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        if not _XLA_AVAILABLE:
+            raise ModuleNotFoundError(str(_XLA_AVAILABLE))
+        super().__init__(*args, **kwargs)
+
     def setup_device(self, device: torch.device) -> None:
         pass
 
@@ -47,8 +56,10 @@ def auto_device_count() -> int:
         return 8
 
     @staticmethod
+    @functools.lru_cache(maxsize=1)
     def is_available() -> bool:
-        return _TPU_AVAILABLE
+        # check `_XLA_AVAILABLE` again to avoid launching processes
+        return _XLA_AVAILABLE and _is_device_tpu()
 
     @classmethod
     def register_accelerators(cls, accelerator_registry: Dict) -> None:
@@ -59,6 +70,64 @@ def register_accelerators(cls, accelerator_registry: Dict) -> None:
         )
 
 
+# define TPU availability timeout in seconds
+TPU_CHECK_TIMEOUT = 60
+
+
+def _inner_f(queue: Queue, func: Callable, *args: Any, **kwargs: Any) -> None:  # pragma: no cover
+    try:
+        queue.put(func(*args, **kwargs))
+    except Exception:
+        traceback.print_exc()
+        queue.put(None)
+
+
+def _multi_process(func: Callable) -> Callable:
+    @functools.wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> Union[bool, Any]:
+        queue: Queue = Queue()
+        proc = Process(target=_inner_f, args=(queue, func, *args), kwargs=kwargs)
+        proc.start()
+        proc.join(TPU_CHECK_TIMEOUT)
+        try:
+            return queue.get_nowait()
+        except q.Empty:
+            traceback.print_exc()
+            return False
+
+    return wrapper
+
+
+@_multi_process
+def _is_device_tpu() -> bool:
+    """Check if TPU devices are available. Runs XLA device check within a separate process.
+
+    Return:
+        A boolean value indicating if TPU devices are available
+    """
+    if not _XLA_AVAILABLE:
+        return False
+    import torch_xla.core.xla_model as xm
+
+    # For the TPU Pod training process, for example, if we have
+    # TPU v3-32 with 4 VMs, the world size would be 4 and as
+    # we would have to use `torch_xla.distributed.xla_dist` for
+    # multiple VMs and TPU_CONFIG won't be available, running
+    # `xm.get_xla_supported_devices("TPU")` won't be possible.
+    return (xm.xrt_world_size() > 1) or bool(xm.get_xla_supported_devices("TPU"))
+
+
+_XLA_AVAILABLE = RequirementCache("torch_xla")
+
+
+def tpu_distributed() -> bool:
+    if not TPUAccelerator.is_available():
+        return False
+    import torch_xla.core.xla_model as xm
+
+    return xm.xrt_world_size() > 1
+
+
 def parse_tpu_cores(tpu_cores: Optional[Union[int, str, List[int]]]) -> Optional[Union[int, List[int]]]:
     """
     Parses the tpu_cores given in the format as accepted by the

src/lightning_lite/plugins/environments/xla_environment.py

@@ -13,13 +13,10 @@
 # limitations under the License.
 import logging
 import os
+from typing import Any
 
+from lightning_lite.accelerators.tpu import _XLA_AVAILABLE, TPUAccelerator
 from lightning_lite.plugins.environments.cluster_environment import ClusterEnvironment
-from lightning_lite.utilities.imports import _TPU_AVAILABLE
-
-if _TPU_AVAILABLE:
-    import torch_xla.core.xla_env_vars as xenv
-    import torch_xla.core.xla_model as xm
 
 log = logging.getLogger(__name__)
 
@@ -31,36 +28,53 @@ class XLAEnvironment(ClusterEnvironment):
     `here <https://github.com/pytorch/xla/blob/master/torch_xla/core/xla_env_vars.py>`_.
     """
 
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        if not _XLA_AVAILABLE:
+            raise ModuleNotFoundError(str(_XLA_AVAILABLE))
+        super().__init__(*args, **kwargs)
+
     @property
     def creates_processes_externally(self) -> bool:
         return False
 
     @property
     def main_address(self) -> str:
+        import torch_xla.core.xla_env_vars as xenv
+
         return os.environ[xenv.TPU_MESH_CTLER_ADDR]
 
     @property
     def main_port(self) -> int:
+        import torch_xla.core.xla_env_vars as xenv
+
         return int(os.environ[xenv.TPU_MESH_CTLER_PORT])
 
     @staticmethod
     def detect() -> bool:
-        return _TPU_AVAILABLE
+        return TPUAccelerator.is_available()
 
     def world_size(self) -> int:
+        import torch_xla.core.xla_model as xm
+
         return xm.xrt_world_size()
 
     def set_world_size(self, size: int) -> None:
         log.debug("XLAEnvironment.set_world_size was called, but setting world size is not allowed. Ignored.")
 
     def global_rank(self) -> int:
+        import torch_xla.core.xla_model as xm
+
         return xm.get_ordinal()
 
     def set_global_rank(self, rank: int) -> None:
         log.debug("XLAEnvironment.set_global_rank was called, but setting global rank is not allowed. Ignored.")
 
     def local_rank(self) -> int:
+        import torch_xla.core.xla_model as xm
+
         return xm.get_local_ordinal()
 
     def node_rank(self) -> int:
+        import torch_xla.core.xla_env_vars as xenv
+
         return int(os.environ.get(xenv.HOST_ORDINAL, 0))

src/lightning_lite/plugins/io/xla_plugin.py

@@ -16,21 +16,24 @@
 
 from lightning_utilities.core.apply_func import apply_to_collection
 
+from lightning_lite.accelerators.tpu import _XLA_AVAILABLE
 from lightning_lite.plugins.io.torch_plugin import TorchCheckpointIO
 from lightning_lite.utilities.cloud_io import get_filesystem
-from lightning_lite.utilities.imports import _OMEGACONF_AVAILABLE, _TPU_AVAILABLE
+from lightning_lite.utilities.imports import _OMEGACONF_AVAILABLE
 from lightning_lite.utilities.types import _PATH
 
-if _TPU_AVAILABLE:
-    import torch_xla.core.xla_model as xm
-
 if _OMEGACONF_AVAILABLE:
     from omegaconf import DictConfig, ListConfig, OmegaConf
 
 
 class XLACheckpointIO(TorchCheckpointIO):
     """CheckpointIO that utilizes :func:`xm.save` to save checkpoints for TPU training strategies."""
 
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        if not _XLA_AVAILABLE:
+            raise ModuleNotFoundError(str(_XLA_AVAILABLE))
+        super().__init__(*args, **kwargs)
+
     def save_checkpoint(self, checkpoint: Dict[str, Any], path: _PATH, storage_options: Optional[Any] = None) -> None:
         """Save model/training states as a checkpoint file through state-dump and file-write.
 
@@ -55,4 +58,6 @@ def save_checkpoint(self, checkpoint: Dict[str, Any], path: _PATH, storage_optio
         # Ref: https://github.com/pytorch/xla/issues/2773
         if _OMEGACONF_AVAILABLE:
             checkpoint = apply_to_collection(checkpoint, (DictConfig, ListConfig), OmegaConf.to_container)
+        import torch_xla.core.xla_model as xm
+
         xm.save({k: v for k, v in checkpoint.items() if k != "callbacks"}, path)

src/lightning_lite/strategies/launchers/xla.py

@@ -19,16 +19,10 @@
 import torch.multiprocessing as mp
 from torch.multiprocessing import ProcessContext
 
+from lightning_lite.accelerators.tpu import _XLA_AVAILABLE
 from lightning_lite.strategies.launchers.multiprocessing import _GlobalStateSnapshot, _MultiProcessingLauncher
-from lightning_lite.utilities import _TPU_AVAILABLE
 from lightning_lite.utilities.apply_func import move_data_to_device
 
-if _TPU_AVAILABLE:
-    import torch_xla.core.xla_model as xm
-    import torch_xla.distributed.xla_multiprocessing as xmp
-else:
-    xm, xmp = None, None
-
 if TYPE_CHECKING:
     from lightning_lite.strategies import Strategy
 
@@ -50,6 +44,8 @@ class _XLALauncher(_MultiProcessingLauncher):
     """
 
     def __init__(self, strategy: "Strategy") -> None:
+        if not _XLA_AVAILABLE:
+            raise ModuleNotFoundError(str(_XLA_AVAILABLE))
         super().__init__(strategy=strategy, start_method="fork")
 
     @property
@@ -103,6 +99,8 @@ def _save_spawn(
 ) -> Optional[ProcessContext]:
     """Wraps the :func:`torch_xla.distributed.xla_multiprocessing.spawn` with added teardown logic for the worker
     processes."""
+    import torch_xla.core.xla_model as xm
+    import torch_xla.distributed.xla_multiprocessing as xmp
 
     @wraps(fn)
     def wrapped(rank: int, *_args: Any) -> None:

src/lightning_lite/strategies/xla.py

@@ -13,35 +13,30 @@
 # limitations under the License.
 import io
 import os
-from typing import Any, Dict, List, Mapping, Optional, Sequence, Union
+from typing import Any, Dict, List, Mapping, Optional, Sequence, TYPE_CHECKING, Union
 
 import torch
 from torch import Tensor
 from torch.nn import Module
 from torch.utils.data import DataLoader
 
 from lightning_lite.accelerators import Accelerator
+from lightning_lite.accelerators.tpu import _XLA_AVAILABLE
 from lightning_lite.plugins.environments import XLAEnvironment
 from lightning_lite.plugins.io.checkpoint_plugin import CheckpointIO
 from lightning_lite.plugins.io.xla_plugin import XLACheckpointIO
 from lightning_lite.plugins.precision import Precision
 from lightning_lite.strategies.ddp_spawn import DDPSpawnStrategy
 from lightning_lite.strategies.launchers.xla import _XLALauncher
 from lightning_lite.strategies.strategy import TBroadcast
-from lightning_lite.utilities import _TPU_AVAILABLE
 from lightning_lite.utilities.apply_func import apply_to_collection
 from lightning_lite.utilities.data import has_len
 from lightning_lite.utilities.distributed import ReduceOp
 from lightning_lite.utilities.rank_zero import rank_zero_only
 from lightning_lite.utilities.types import _PATH
 
-if _TPU_AVAILABLE:
-    import torch_xla.core.xla_env_vars as xenv
-    import torch_xla.core.xla_model as xm
-    from torch_xla.core.xla_model import rendezvous
+if TYPE_CHECKING and _XLA_AVAILABLE:
     from torch_xla.distributed.parallel_loader import MpDeviceLoader
-else:
-    xm, xmp, MpDeviceLoader, rendezvous = [None] * 4
 
 
 class XLAStrategy(DDPSpawnStrategy):
@@ -71,6 +66,8 @@ def __init__(
     def root_device(self) -> torch.device:
         if not self._launched:
             raise RuntimeError("Accessing the XLA device before processes have spawned is not allowed.")
+        import torch_xla.core.xla_model as xm
+
         return xm.xla_device()
 
     @property
@@ -89,6 +86,8 @@ def distributed_sampler_kwargs(self) -> Dict[str, int]:
 
     @property
     def is_distributed(self) -> bool:
+        import torch_xla.core.xla_env_vars as xenv
+
         # HOST_WORLD_SIZE is not set outside the xmp.spawn process
         return (xenv.HOST_WORLD_SIZE in os.environ) and self.world_size != 1
 
@@ -106,8 +105,10 @@ def setup_module(self, module: Module) -> Module:
     def module_to_device(self, module: Module) -> None:
         module.to(self.root_device)
 
-    def process_dataloader(self, dataloader: DataLoader) -> MpDeviceLoader:
+    def process_dataloader(self, dataloader: DataLoader) -> "MpDeviceLoader":
         XLAStrategy._validate_dataloader(dataloader)
+        from torch_xla.distributed.parallel_loader import MpDeviceLoader
+
         dataloader = MpDeviceLoader(dataloader, self.root_device)
         # Mimic interface to torch.utils.data.DataLoader
         dataloader.dataset = dataloader._loader.dataset
@@ -126,6 +127,7 @@ def reduce(
                 "Currently, the XLAStrategy only supports `sum`, `mean`, `avg` for the reduce operation, got:"
                 f" {reduce_op}"
             )
+        import torch_xla.core.xla_model as xm
 
         output = xm.mesh_reduce("reduce", output, sum)
 
@@ -136,7 +138,9 @@ def reduce(
 
     def barrier(self, name: Optional[str] = None, *args: Any, **kwargs: Any) -> None:
         if self.is_distributed:
-            rendezvous(name)
+            import torch_xla.core.xla_model as xm
+
+            xm.rendezvous(name)
 
     def broadcast(self, obj: TBroadcast, src: int = 0) -> TBroadcast:
         if not self.is_distributed:
@@ -145,6 +149,8 @@ def broadcast(self, obj: TBroadcast, src: int = 0) -> TBroadcast:
         torch.save(obj, buffer)
         data = bytearray(buffer.getbuffer())
         data_tensor = torch.tensor(data, device=self.root_device, dtype=torch.float)
+        import torch_xla.core.xla_model as xm
+
         data = xm.all_gather(data_tensor)
         buffer = io.BytesIO(data.cpu().byte().numpy())
         obj = torch.load(buffer)
@@ -162,6 +168,8 @@ def all_gather(self, tensor: Tensor, group: Optional[Any] = None, sync_grads: bo
         """
         if isinstance(tensor, Tensor) and tensor.dim() == 0:
             tensor = tensor.unsqueeze(0)
+        import torch_xla.core.xla_model as xm
+
         return xm.all_gather(tensor)
 
     def save_checkpoint(

src/lightning_lite/utilities/__init__.py

@@ -29,8 +29,6 @@
     _TORCH_GREATER_EQUAL_1_10,
     _TORCH_GREATER_EQUAL_1_11,
     _TORCH_GREATER_EQUAL_1_12,
-    _TPU_AVAILABLE,
-    _XLA_AVAILABLE,
 )
 from lightning_lite.utilities.rank_zero import (  # noqa: F401
     rank_zero_deprecation,