add single tpu

Co-authored-by: Adrian Wälchli <[email protected]>

Author: justusschock

pytorch_lightning/plugins/training_type/single_tpu.py

@@ -0,0 +1,41 @@
+import io
+import os
+from typing import Optional
+
+import torch
+
+from pytorch_lightning.plugins.training_type.single_device import SingleDevicePlugin
+from pytorch_lightning.utilities import _TPU_AVAILABLE, rank_zero_warn
+
+if _TPU_AVAILABLE:
+    import torch_xla
+    import torch_xla.core.xla_model as xm
+
+
+class SingleTPUPlugin(SingleDevicePlugin):
+    def __init__(self, device: torch.device):
+        super().__init__(device)
+
+        self.tpu_local_core_rank = 0
+        self.tpu_global_core_rank = 0
+
+    def on_tpu(self) -> bool:
+        return True
+
+    def pre_training(self) -> None:
+        if isinstance(self.device, int):
+            self.device = xm.xla_device(self.device)
+
+        self.tpu_local_core_rank = xm.get_local_ordinal()
+        self.tpu_global_core_rank = xm.get_ordinal()
+
+    def post_training(self) -> None:
+        model = self.lightning_module
+
+        if self.on_colab_kaggle:
+            rank_zero_warn("cleaning up... please do not interrupt")
+            self.save_spawn_weights(model)
+
+    @property
+    def on_colab_kaggle(self) -> bool:
+        return bool(os.getenv("COLAB_GPU") or os.getenv("KAGGLE_URL_BASE"))