Add cpu device parser to validate cpu devices (#12160)

Author: rohitgr7

pytorch_lightning/accelerators/cpu.py

@@ -16,8 +16,8 @@
 import torch
 
 from pytorch_lightning.accelerators.accelerator import Accelerator
+from pytorch_lightning.utilities import device_parser
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from pytorch_lightning.utilities.rank_zero import rank_zero_warn
 from pytorch_lightning.utilities.types import _DEVICE
 
 
@@ -39,19 +39,16 @@ def get_device_stats(self, device: _DEVICE) -> Dict[str, Any]:
         return {}
 
     @staticmethod
-    def parse_devices(devices: Union[int, str, List[int]]) -> Union[int, str, List[int]]:
+    def parse_devices(devices: Union[int, str, List[int]]) -> int:
         """Accelerator device parsing logic."""
+        devices = device_parser.parse_cpu_cores(devices)
         return devices
 
     @staticmethod
     def get_parallel_devices(devices: Union[int, str, List[int]]) -> List[torch.device]:
         """Gets parallel devices for the Accelerator."""
-        if isinstance(devices, int):
-            return [torch.device("cpu")] * devices
-        rank_zero_warn(
-            f"The flag `devices` must be an int with `accelerator='cpu'`, got `devices={devices!r}` instead."
-        )
-        return []
+        devices = device_parser.parse_cpu_cores(devices)
+        return [torch.device("cpu")] * devices
 
     @staticmethod
     def auto_device_count() -> int:

pytorch_lightning/trainer/connectors/accelerator_connector.py

@@ -512,7 +512,7 @@ def _set_parallel_devices_and_init_accelerator(self) -> None:
             self._parallel_devices = self.accelerator.get_parallel_devices(self._devices_flag)
 
     def _set_devices_flag_if_auto_passed(self) -> None:
-        if self._devices_flag == "auto" or not self._devices_flag:
+        if self._devices_flag == "auto" or self._devices_flag is None:
             self._devices_flag = self.accelerator.auto_device_count()
 
     def _choose_and_init_cluster_environment(self) -> ClusterEnvironment:

pytorch_lightning/utilities/device_parser.py

@@ -79,8 +79,9 @@ def parse_gpu_ids(gpus: Optional[Union[int, str, List[int]]]) -> Optional[List[i
     Returns:
         a list of gpus to be used or ``None`` if no GPUs were requested
 
-    If no GPUs are available but the value of gpus variable indicates request for GPUs
-    then a MisconfigurationException is raised.
+    Raises:
+        MisconfigurationException:
+            If no GPUs are available but the value of gpus variable indicates request for GPUs
     """
     # Check that gpus param is None, Int, String or List
     _check_data_type(gpus)
@@ -137,6 +138,29 @@ def parse_tpu_cores(tpu_cores: Optional[Union[int, str, List[int]]]) -> Optional
     return tpu_cores
 
 
+def parse_cpu_cores(cpu_cores: Union[int, str, List[int]]) -> int:
+    """Parses the cpu_cores given in the format as accepted by the ``devices`` argument in the
+    :class:`~pytorch_lightning.trainer.Trainer`.
+
+    Args:
+        cpu_cores: An int > 0.
+
+    Returns:
+        an int representing the number of processes
+
+    Raises:
+        MisconfigurationException:
+            If cpu_cores is not an int > 0
+    """
+    if isinstance(cpu_cores, str) and cpu_cores.strip().isdigit():
+        cpu_cores = int(cpu_cores)
+
+    if not isinstance(cpu_cores, int) or cpu_cores <= 0:
+        raise MisconfigurationException("`devices` selected with `CPUAccelerator` should be an int > 0.")
+
+    return cpu_cores
+
+
 def _normalize_parse_gpu_string_input(s: Union[int, str, List[int]]) -> Union[int, List[int]]:
     if not isinstance(s, str):
         return s

pytorch_lightning/utilities/parsing.py

@@ -243,7 +243,7 @@ def save_hyperparameters(
         if isinstance(init_args[k], nn.Module):
             rank_zero_warn(
                 f"Attribute {k!r} is an instance of `nn.Module` and is already saved during checkpointing."
-                " It is recommended to ignore them using `self.save_hyperparameters(ignore=[k!r])`."
+                f" It is recommended to ignore them using `self.save_hyperparameters(ignore=[{k}!r])`."
             )
 
     if not args:

tests/accelerators/test_accelerator_connector.py

@@ -570,14 +570,6 @@ def test_set_devices_if_none_cpu():
     assert trainer.num_devices == 3
 
 
-def test_devices_with_cpu_only_supports_integer():
-
-    with pytest.warns(UserWarning, match="The flag `devices` must be an int"):
-        trainer = Trainer(accelerator="cpu", devices="1,3")
-    assert isinstance(trainer.accelerator, CPUAccelerator)
-    assert trainer.num_devices == 1
-
-
 @pytest.mark.parametrize("training_type", ["ddp2", "dp"])
 def test_unsupported_strategy_types_on_cpu(training_type):
     with pytest.warns(UserWarning, match="is not supported on CPUs, hence setting `strategy='ddp"):

tests/accelerators/test_cpu.py

@@ -11,6 +11,7 @@
 from pytorch_lightning.plugins.io.torch_plugin import TorchCheckpointIO
 from pytorch_lightning.plugins.precision.precision_plugin import PrecisionPlugin
 from pytorch_lightning.strategies import SingleDeviceStrategy
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers.boring_model import BoringModel
 
 
@@ -66,3 +67,10 @@ def load_checkpoint(self, checkpoint_path: Union[str, Path]) -> Dict[str, Any]:
     for func in (trainer.test, trainer.validate, trainer.predict):
         plugin.setup_called = False
         func(model, ckpt_path=checkpoint_path)
+
+
+@pytest.mark.parametrize("devices", ([3], -1, 0))
+def test_invalid_devices_with_cpu_accelerator(devices):
+    """Test invalid device flag raises MisconfigurationException with CPUAccelerator."""
+    with pytest.raises(MisconfigurationException, match="should be an int > 0"):
+        Trainer(accelerator="cpu", devices=devices)