avoid calling xm.xla_device() on multi-tpu

Author: DuYicong515

pytorch_lightning/trainer/trainer.py

@@ -2066,7 +2066,11 @@ def num_nodes(self) -> int:
     @property
     def device_ids(self) -> List[int]:
         """List of device indexes per node."""
-        devices = getattr(self.strategy, "parallel_devices", [self.strategy.root_device])
+        devices = (
+            self.strategy.parallel_devices
+            if isinstance(self.strategy, ParallelStrategy)
+            else [self.strategy.root_device]
+        )
         device_ids = []
         for idx, device in enumerate(devices):
             if isinstance(device, torch.device):

tests/accelerators/test_tpu.py

@@ -94,7 +94,6 @@ def test_accelerator_cpu_with_tpu_cores_flag():
 
 
 @RunIf(tpu=True)
-@pl_multi_process_test
 @pytest.mark.parametrize(["accelerator", "devices"], [("auto", 8), ("auto", "auto"), ("tpu", None)])
 def test_accelerator_tpu(accelerator, devices):
     assert TPUAccelerator.is_available()
@@ -104,13 +103,12 @@ def test_accelerator_tpu(accelerator, devices):
     assert isinstance(trainer.strategy, TPUSpawnStrategy)
     assert trainer.num_devices == 8
     with pytest.deprecated_call(
-        match= "`Trainer.tpu_cores` is deprecated in v1.6 and will be removed in v1.8. "
+        match="`Trainer.tpu_cores` is deprecated in v1.6 and will be removed in v1.8. "
         "Please use `Trainer.devices` instead."
     ):
         trainer.tpu_cores == 8
 
 
-
 @RunIf(tpu=True)
 def test_accelerator_tpu_with_tpu_cores_priority():
     """Test for checking `tpu_cores` flag takes priority over `devices`."""
@@ -124,7 +122,6 @@ def test_accelerator_tpu_with_tpu_cores_priority():
 
 
 @RunIf(tpu=True)
-@pl_multi_process_test
 def test_set_devices_if_none_tpu():
     trainer = Trainer(accelerator="tpu", tpu_cores=8)
     assert trainer.num_devices == 8