address comments for testing

Author: DuYicong515

tests/accelerators/test_tpu.py

@@ -310,3 +310,20 @@ def test_mp_device_dataloader_attribute(_):
 def test_warning_if_tpus_not_used():
     with pytest.warns(UserWarning, match="TPU available but not used. Set `accelerator` and `devices`"):
         Trainer()
+
+
+@RunIf(tpu=True)
+@pytest.mark.parametrize(
+    ["trainer_kwargs", "expected_device_ids"],
+    [
+        ({"accelerator": "tpu", "devices": 1}, [0]),
+        ({"accelerator": "tpu", "devices": 8}, list(range(8))),
+        ({"accelerator": "tpu", "devices": "8"}, list(range(8))),
+        ({"accelerator": "tpu", "devices": [2]}, [2]),
+        ({"accelerator": "tpu", "devices": "2,"}, [2]),
+    ],
+)
+def test_trainer_config_device_ids(monkeypatch, trainer_kwargs, expected_device_ids):
+    trainer = Trainer(**trainer_kwargs)
+    assert trainer.device_ids == expected_device_ids
+    assert trainer.num_devices == len(expected_device_ids)

tests/deprecated_api/test_remove_1-8.py

@@ -20,6 +20,7 @@
 import torch
 from torch import optim
 
+import pytorch_lightning
 from pytorch_lightning import Callback, Trainer
 from pytorch_lightning.loggers import CSVLogger, LightningLoggerBase
 from pytorch_lightning.plugins.precision.precision_plugin import PrecisionPlugin
@@ -695,19 +696,28 @@ def on_load_checkpoint(self, checkpoint):
     [
         ({}, 1),
         ({"devices": 1}, 1),
+        ({"devices": 1}, 1),
+        ({"devices": "1"}, 1),
+        ({"devices": 2}, 2),
         ({"accelerator": "gpu", "devices": 1}, 1),
-        ({"strategy": "ddp", "devices": 1}, 1),
-        ({"strategy": "ddp", "accelerator": "gpu", "devices": 1}, 1),
-        ({"strategy": "ddp", "devices": 2}, 2),
-        ({"strategy": "ddp", "accelerator": "gpu", "devices": 2}, 2),
-        ({"strategy": "ddp", "accelerator": "gpu", "devices": [2]}, 1),
-        ({"strategy": "ddp", "accelerator": "gpu", "devices": [0, 2]}, 2),
+        ({"accelerator": "gpu", "devices": 2}, 2),
+        ({"accelerator": "gpu", "devices": "2"}, 2),
+        ({"accelerator": "gpu", "devices": [2]}, 1),
+        ({"accelerator": "gpu", "devices": "2,"}, 1),
+        ({"accelerator": "gpu", "devices": [0, 2]}, 2),
+        ({"accelerator": "gpu", "devices": "0, 2"}, 2),
+        ({"accelerator": "ipu", "devices": 1}, 1),
+        ({"accelerator": "ipu", "devices": 2}, 2),
     ],
 )
-def test_v1_8_0_trainer_devices(monkeypatch, trainer_kwargs, expected_devices):
+def test_trainer_config_device_ids(monkeypatch, trainer_kwargs, expected_devices):
     if trainer_kwargs.get("accelerator") == "gpu":
         monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
         monkeypatch.setattr(torch.cuda, "device_count", lambda: 4)
+    elif trainer_kwargs.get("accelerator") == "ipu":
+        monkeypatch.setattr(pytorch_lightning.accelerators.ipu.IPUAccelerator, "is_available", lambda _: True)
+        monkeypatch.setattr(pytorch_lightning.strategies.ipu, "_IPU_AVAILABLE", lambda: True)
+
     trainer = Trainer(**trainer_kwargs)
     with pytest.deprecated_call(
         match="`Trainer.devices` was deprecated in v1.6 and will be removed in v1.8."

tests/trainer/test_trainer.py

@@ -30,6 +30,7 @@
 from torch.optim import SGD
 from torch.utils.data import DataLoader, IterableDataset
 
+import pytorch_lightning
 import tests.helpers.utils as tutils
 from pytorch_lightning import Callback, LightningDataModule, LightningModule, Trainer
 from pytorch_lightning.accelerators import CPUAccelerator, GPUAccelerator
@@ -2153,19 +2154,28 @@ def test_dataloaders_are_not_loaded_if_disabled_through_limit_batches(running_st
     [
         ({}, [0]),
         ({"devices": 1}, [0]),
+        ({"devices": 1}, [0]),
+        ({"devices": "1"}, [0]),
+        ({"devices": 2}, [0, 1]),
         ({"accelerator": "gpu", "devices": 1}, [0]),
-        ({"strategy": "ddp", "devices": 1}, [0]),
-        ({"strategy": "ddp", "accelerator": "gpu", "devices": 1}, [0]),
-        ({"strategy": "ddp", "devices": 2}, [0, 1]),
-        ({"strategy": "ddp", "accelerator": "gpu", "devices": 2}, [0, 1]),
-        ({"strategy": "ddp", "accelerator": "gpu", "devices": [2]}, [2]),
-        ({"strategy": "ddp", "accelerator": "gpu", "devices": [0, 2]}, [0, 2]),
+        ({"accelerator": "gpu", "devices": 2}, [0, 1]),
+        ({"accelerator": "gpu", "devices": "2"}, [0, 1]),
+        ({"accelerator": "gpu", "devices": [2]}, [2]),
+        ({"accelerator": "gpu", "devices": "2,"}, [2]),
+        ({"accelerator": "gpu", "devices": [0, 2]}, [0, 2]),
+        ({"accelerator": "gpu", "devices": "0, 2"}, [0, 2]),
+        ({"accelerator": "ipu", "devices": 1}, [0]),
+        ({"accelerator": "ipu", "devices": 2}, [0, 1]),
     ],
 )
 def test_trainer_config_device_ids(monkeypatch, trainer_kwargs, expected_device_ids):
     if trainer_kwargs.get("accelerator") == "gpu":
         monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
         monkeypatch.setattr(torch.cuda, "device_count", lambda: 4)
+    elif trainer_kwargs.get("accelerator") == "ipu":
+        monkeypatch.setattr(pytorch_lightning.accelerators.ipu.IPUAccelerator, "is_available", lambda _: True)
+        monkeypatch.setattr(pytorch_lightning.strategies.ipu, "_IPU_AVAILABLE", lambda: True)
+
     trainer = Trainer(**trainer_kwargs)
     assert trainer.device_ids == expected_device_ids
-    assert len(trainer.device_ids) == trainer.num_devices
+    assert trainer.num_devices == len(expected_device_ids)