Deprecate DeviceType in favor of _AcceleratorType (#10503)

Co-authored-by: Carlos Mocholí <[email protected]>

Author: kaushikb11

CHANGELOG.md

@@ -78,6 +78,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Deprecated the `precision_plugin` constructor argument from `Accelerator` ([#10570](https://github.com/PyTorchLightning/pytorch-lightning/pull/10570))
 
 
+- Deprecated `DeviceType` in favor of `_AcceleratorType` ([#10503](https://github.com/PyTorchLightning/pytorch-lightning/pull/10503))
+
+
 - Deprecated the property `Trainer.slurm_job_id` in favor of the new `SLURMEnvironment.job_id()` method ([#10622](https://github.com/PyTorchLightning/pytorch-lightning/pull/10622))
 
 

pytorch_lightning/callbacks/gpu_stats_monitor.py

@@ -29,7 +29,7 @@
 
 import pytorch_lightning as pl
 from pytorch_lightning.callbacks.base import Callback
-from pytorch_lightning.utilities import DeviceType, rank_zero_deprecation, rank_zero_only
+from pytorch_lightning.utilities import _AcceleratorType, rank_zero_deprecation, rank_zero_only
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.parsing import AttributeDict
 from pytorch_lightning.utilities.types import STEP_OUTPUT
@@ -126,7 +126,7 @@ def setup(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", stage: O
         if not trainer.logger:
             raise MisconfigurationException("Cannot use GPUStatsMonitor callback with Trainer that has no logger.")
 
-        if trainer._device_type != DeviceType.GPU:
+        if trainer._device_type != _AcceleratorType.GPU:
             raise MisconfigurationException(
                 "You are using GPUStatsMonitor but are not running on GPU"
                 f" since gpus attribute in Trainer is set to {trainer.gpus}."

pytorch_lightning/callbacks/xla_stats_monitor.py

@@ -21,7 +21,7 @@
 import time
 
 from pytorch_lightning.callbacks.base import Callback
-from pytorch_lightning.utilities import _TPU_AVAILABLE, DeviceType, rank_zero_deprecation, rank_zero_info
+from pytorch_lightning.utilities import _AcceleratorType, _TPU_AVAILABLE, rank_zero_deprecation, rank_zero_info
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
 if _TPU_AVAILABLE:
@@ -70,7 +70,7 @@ def on_train_start(self, trainer, pl_module) -> None:
         if not trainer.logger:
             raise MisconfigurationException("Cannot use XLAStatsMonitor callback with Trainer that has no logger.")
 
-        if trainer._device_type != DeviceType.TPU:
+        if trainer._device_type != _AcceleratorType.TPU:
             raise MisconfigurationException(
                 "You are using XLAStatsMonitor but are not running on TPU"
                 f" since `tpu_cores` attribute in Trainer is set to {trainer.tpu_cores}."

pytorch_lightning/lite/lite.py

@@ -28,7 +28,7 @@
 from pytorch_lightning.lite.wrappers import _LiteDataLoader, _LiteModule, _LiteOptimizer
 from pytorch_lightning.plugins import DDPSpawnPlugin, DeepSpeedPlugin, PLUGIN_INPUT, TPUSpawnPlugin, TrainingTypePlugin
 from pytorch_lightning.trainer.connectors.accelerator_connector import AcceleratorConnector
-from pytorch_lightning.utilities import _StrategyType, DeviceType, move_data_to_device
+from pytorch_lightning.utilities import _AcceleratorType, _StrategyType, move_data_to_device
 from pytorch_lightning.utilities.apply_func import apply_to_collection, convert_to_tensors
 from pytorch_lightning.utilities.data import (
     _auto_add_worker_init_fn,
@@ -448,11 +448,11 @@ def _check_strategy_support(self, strategy: Optional[Union[str, TrainingTypePlug
             )
 
     @staticmethod
-    def _supported_device_types() -> Sequence[DeviceType]:
+    def _supported_device_types() -> Sequence[_AcceleratorType]:
         return (
-            DeviceType.CPU,
-            DeviceType.GPU,
-            DeviceType.TPU,
+            _AcceleratorType.CPU,
+            _AcceleratorType.GPU,
+            _AcceleratorType.TPU,
         )
 
     @staticmethod

pytorch_lightning/loops/optimization/optimizer_loop.py

@@ -30,7 +30,7 @@
 )
 from pytorch_lightning.profiler import BaseProfiler, PassThroughProfiler
 from pytorch_lightning.trainer.progress import OptimizationProgress
-from pytorch_lightning.utilities import AMPType, DeviceType
+from pytorch_lightning.utilities import _AcceleratorType, AMPType
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.finite_checks import detect_nan_parameters
 from pytorch_lightning.utilities.imports import _TPU_AVAILABLE
@@ -378,7 +378,7 @@ def _optimizer_step(
             optimizer,
             opt_idx,
             train_step_and_backward_closure,
-            on_tpu=(self.trainer._device_type == DeviceType.TPU and _TPU_AVAILABLE),
+            on_tpu=(self.trainer._device_type == _AcceleratorType.TPU and _TPU_AVAILABLE),
             using_native_amp=(self.trainer.amp_backend is not None and self.trainer.amp_backend == AMPType.NATIVE),
             using_lbfgs=is_lbfgs,
         )

pytorch_lightning/trainer/connectors/accelerator_connector.py

@@ -61,10 +61,10 @@
     TorchElasticEnvironment,
 )
 from pytorch_lightning.utilities import (
+    _AcceleratorType,
     _StrategyType,
     AMPType,
     device_parser,
-    DeviceType,
     rank_zero_deprecation,
     rank_zero_info,
     rank_zero_warn,
@@ -106,7 +106,7 @@ def __init__(
         plugins,
     ):
         # initialization
-        self._device_type = DeviceType.CPU
+        self._device_type = _AcceleratorType.CPU
         self._distrib_type = None
         self._accelerator_type = None
 
@@ -199,32 +199,32 @@ def _init_deterministic(self, deterministic: bool) -> None:
     def select_accelerator_type(self) -> None:
         if self.distributed_backend == "auto":
             if self.has_tpu:
-                self._accelerator_type = DeviceType.TPU
+                self._accelerator_type = _AcceleratorType.TPU
             elif self.has_ipu:
-                self._accelerator_type = DeviceType.IPU
+                self._accelerator_type = _AcceleratorType.IPU
             elif self.has_gpu:
-                self._accelerator_type = DeviceType.GPU
+                self._accelerator_type = _AcceleratorType.GPU
             else:
                 self._set_devices_to_cpu_num_processes()
-                self._accelerator_type = DeviceType.CPU
-        elif self.distributed_backend == DeviceType.TPU:
+                self._accelerator_type = _AcceleratorType.CPU
+        elif self.distributed_backend == _AcceleratorType.TPU:
             if not self.has_tpu:
                 msg = "TPUs are not available" if not _TPU_AVAILABLE else "you didn't pass `tpu_cores` to `Trainer`"
                 raise MisconfigurationException(f"You passed `accelerator='tpu'`, but {msg}.")
-            self._accelerator_type = DeviceType.TPU
-        elif self.distributed_backend == DeviceType.IPU:
+            self._accelerator_type = _AcceleratorType.TPU
+        elif self.distributed_backend == _AcceleratorType.IPU:
             if not self.has_ipu:
                 msg = "IPUs are not available" if not _IPU_AVAILABLE else "you didn't pass `ipus` to `Trainer`"
                 raise MisconfigurationException(f"You passed `accelerator='ipu'`, but {msg}.")
-            self._accelerator_type = DeviceType.IPU
-        elif self.distributed_backend == DeviceType.GPU:
+            self._accelerator_type = _AcceleratorType.IPU
+        elif self.distributed_backend == _AcceleratorType.GPU:
             if not self.has_gpu:
                 msg = "you didn't pass `gpus` to `Trainer`" if torch.cuda.is_available() else "GPUs are not available"
                 raise MisconfigurationException(f"You passed `accelerator='gpu'`, but {msg}.")
-            self._accelerator_type = DeviceType.GPU
-        elif self.distributed_backend == DeviceType.CPU:
+            self._accelerator_type = _AcceleratorType.GPU
+        elif self.distributed_backend == _AcceleratorType.CPU:
             self._set_devices_to_cpu_num_processes()
-            self._accelerator_type = DeviceType.CPU
+            self._accelerator_type = _AcceleratorType.CPU
 
         if self.distributed_backend in self.accelerator_types:
             self.distributed_backend = None
@@ -250,29 +250,29 @@ def _warn_if_devices_flag_ignored(self) -> None:
         if self.devices is None:
             return
         devices_warning = f"The flag `devices={self.devices}` will be ignored, as you have set"
-        if self.distributed_backend in ("auto", DeviceType.TPU):
+        if self.distributed_backend in ("auto", _AcceleratorType.TPU):
             if self.tpu_cores is not None:
                 rank_zero_warn(f"{devices_warning} `tpu_cores={self.tpu_cores}`")
-        elif self.distributed_backend in ("auto", DeviceType.IPU):
+        elif self.distributed_backend in ("auto", _AcceleratorType.IPU):
             if self.ipus is not None:
                 rank_zero_warn(f"{devices_warning} `ipus={self.ipus}`")
-        elif self.distributed_backend in ("auto", DeviceType.GPU):
+        elif self.distributed_backend in ("auto", _AcceleratorType.GPU):
             if self.gpus is not None:
                 rank_zero_warn(f"{devices_warning} `gpus={self.gpus}`")
-        elif self.distributed_backend in ("auto", DeviceType.CPU):
+        elif self.distributed_backend in ("auto", _AcceleratorType.CPU):
             if self.num_processes != 1:
                 rank_zero_warn(f"{devices_warning} `num_processes={self.num_processes}`")
 
     def _set_devices_if_none(self) -> None:
         if self.devices is not None:
             return
-        if self._accelerator_type == DeviceType.TPU:
+        if self._accelerator_type == _AcceleratorType.TPU:
             self.devices = self.tpu_cores
-        elif self._accelerator_type == DeviceType.IPU:
+        elif self._accelerator_type == _AcceleratorType.IPU:
             self.devices = self.ipus
-        elif self._accelerator_type == DeviceType.GPU:
+        elif self._accelerator_type == _AcceleratorType.GPU:
             self.devices = self.gpus
-        elif self._accelerator_type == DeviceType.CPU:
+        elif self._accelerator_type == _AcceleratorType.CPU:
             self.devices = self.num_processes
 
     def _handle_accelerator_and_strategy(self) -> None:
@@ -386,7 +386,7 @@ def handle_given_plugins(self) -> None:
 
     @property
     def accelerator_types(self) -> List[str]:
-        return ["auto"] + list(DeviceType)
+        return ["auto"] + list(_AcceleratorType)
 
     @property
     def precision_plugin(self) -> PrecisionPlugin:
@@ -424,7 +424,7 @@ def has_cpu(self) -> bool:
 
     @property
     def use_cpu(self) -> bool:
-        return self._accelerator_type == DeviceType.CPU
+        return self._accelerator_type == _AcceleratorType.CPU
 
     @property
     def has_gpu(self) -> bool:
@@ -433,23 +433,23 @@ def has_gpu(self) -> bool:
         gpus = self.parallel_device_ids
         if gpus is not None and len(gpus) > 0:
             return True
-        return self._map_devices_to_accelerator(DeviceType.GPU)
+        return self._map_devices_to_accelerator(_AcceleratorType.GPU)
 
     @property
     def use_gpu(self) -> bool:
-        return self._accelerator_type == DeviceType.GPU and self.has_gpu
+        return self._accelerator_type == _AcceleratorType.GPU and self.has_gpu
 
     @property
     def has_tpu(self) -> bool:
         # Here, we are not checking for TPU availability, but instead if User has passed
         # `tpu_cores` to Trainer for training.
         if self.tpu_cores is not None:
             return True
-        return self._map_devices_to_accelerator(DeviceType.TPU)
+        return self._map_devices_to_accelerator(_AcceleratorType.TPU)
 
     @property
     def use_tpu(self) -> bool:
-        return self._accelerator_type == DeviceType.TPU and self.has_tpu
+        return self._accelerator_type == _AcceleratorType.TPU and self.has_tpu
 
     @property
     def tpu_id(self) -> Optional[int]:
@@ -463,36 +463,36 @@ def has_ipu(self) -> bool:
         # `ipus` to Trainer for training.
         if self.ipus is not None or isinstance(self._training_type_plugin, IPUPlugin):
             return True
-        return self._map_devices_to_accelerator(DeviceType.IPU)
+        return self._map_devices_to_accelerator(_AcceleratorType.IPU)
 
     @property
     def use_ipu(self) -> bool:
-        return self._accelerator_type == DeviceType.IPU and self.has_ipu
+        return self._accelerator_type == _AcceleratorType.IPU and self.has_ipu
 
     def _set_devices_to_cpu_num_processes(self) -> None:
         if self.num_processes == 1:
-            self._map_devices_to_accelerator(DeviceType.CPU)
+            self._map_devices_to_accelerator(_AcceleratorType.CPU)
 
     def _map_devices_to_accelerator(self, accelerator: str) -> bool:
         if self.devices is None:
             return False
-        if accelerator == DeviceType.TPU and _TPU_AVAILABLE:
+        if accelerator == _AcceleratorType.TPU and _TPU_AVAILABLE:
             if self.devices == "auto":
                 self.devices = TPUAccelerator.auto_device_count()
             self.tpu_cores = device_parser.parse_tpu_cores(self.devices)
             return True
-        if accelerator == DeviceType.IPU and _IPU_AVAILABLE:
+        if accelerator == _AcceleratorType.IPU and _IPU_AVAILABLE:
             if self.devices == "auto":
                 self.devices = IPUAccelerator.auto_device_count()
             self.ipus = self.devices
             return True
-        if accelerator == DeviceType.GPU and torch.cuda.is_available():
+        if accelerator == _AcceleratorType.GPU and torch.cuda.is_available():
             if self.devices == "auto":
                 self.devices = GPUAccelerator.auto_device_count()
             self.gpus = self.devices
             self.parallel_device_ids = device_parser.parse_gpu_ids(self.devices)
             return True
-        if accelerator == DeviceType.CPU:
+        if accelerator == _AcceleratorType.CPU:
             if self.devices == "auto":
                 self.devices = CPUAccelerator.auto_device_count()
             if not isinstance(self.devices, int):
@@ -829,7 +829,7 @@ def set_distributed_mode(self, strategy: Optional[str] = None):
         if isinstance(self.distributed_backend, Accelerator):
             return
 
-        is_cpu_accelerator_type = self._accelerator_type and self._accelerator_type == DeviceType.CPU
+        is_cpu_accelerator_type = self._accelerator_type and self._accelerator_type == _AcceleratorType.CPU
         _use_cpu = is_cpu_accelerator_type or self.distributed_backend and "cpu" in self.distributed_backend
 
         if self.distributed_backend is None:
@@ -867,16 +867,16 @@ def set_distributed_mode(self, strategy: Optional[str] = None):
                 self.num_processes = os.cpu_count()
         # special case with TPUs
         elif self.has_tpu and not _use_cpu:
-            self._device_type = DeviceType.TPU
+            self._device_type = _AcceleratorType.TPU
             if isinstance(self.tpu_cores, int):
                 self._distrib_type = _StrategyType.TPU_SPAWN
         elif self.has_ipu and not _use_cpu:
-            self._device_type = DeviceType.IPU
+            self._device_type = _AcceleratorType.IPU
         elif self.distributed_backend and self._distrib_type is None:
             self._distrib_type = _StrategyType(self.distributed_backend)
 
         if self.num_gpus > 0 and not _use_cpu:
-            self._device_type = DeviceType.GPU
+            self._device_type = _AcceleratorType.GPU
 
         _gpu_distrib_types = (_StrategyType.DP, _StrategyType.DDP, _StrategyType.DDP_SPAWN, _StrategyType.DDP2)
         # DP and DDP2 cannot run without GPU
@@ -896,13 +896,13 @@ def set_distributed_mode(self, strategy: Optional[str] = None):
         self.check_interactive_compatibility()
 
         # for DDP overwrite nb processes by requested GPUs
-        if self._device_type == DeviceType.GPU and self._distrib_type in (
+        if self._device_type == _AcceleratorType.GPU and self._distrib_type in (
             _StrategyType.DDP,
             _StrategyType.DDP_SPAWN,
         ):
             self.num_processes = self.num_gpus
 
-        if self._device_type == DeviceType.GPU and self._distrib_type == _StrategyType.DDP2:
+        if self._device_type == _AcceleratorType.GPU and self._distrib_type == _StrategyType.DDP2:
             self.num_processes = self.num_nodes
 
         # Horovod is an extra case...
@@ -965,27 +965,27 @@ def has_horovodrun() -> bool:
     def update_device_type_if_ipu_plugin(self) -> None:
         # This allows the poptorch.Options that are passed into the IPUPlugin to be the source of truth,
         # which gives users the flexibility to not have to pass `ipus` flag directly to Trainer
-        if isinstance(self._training_type_plugin, IPUPlugin) and self._device_type != DeviceType.IPU:
-            self._device_type = DeviceType.IPU
+        if isinstance(self._training_type_plugin, IPUPlugin) and self._device_type != _AcceleratorType.IPU:
+            self._device_type = _AcceleratorType.IPU
 
     def update_device_type_if_training_type_plugin_passed(self) -> None:
         if isinstance(self.strategy, TrainingTypePlugin) or any(
             isinstance(plug, TrainingTypePlugin) for plug in self.plugins
         ):
             if self._accelerator_type is not None:
                 if self.use_ipu:
-                    self._device_type = DeviceType.IPU
+                    self._device_type = _AcceleratorType.IPU
                 elif self.use_tpu:
-                    self._device_type = DeviceType.TPU
+                    self._device_type = _AcceleratorType.TPU
                 elif self.use_gpu:
-                    self._device_type = DeviceType.GPU
+                    self._device_type = _AcceleratorType.GPU
             else:
                 if self.has_ipu:
-                    self._device_type = DeviceType.IPU
+                    self._device_type = _AcceleratorType.IPU
                 elif self.has_tpu:
-                    self._device_type = DeviceType.TPU
+                    self._device_type = _AcceleratorType.TPU
                 elif self.has_gpu:
-                    self._device_type = DeviceType.GPU
+                    self._device_type = _AcceleratorType.GPU
 
     def _set_distrib_type_if_training_type_plugin_passed(self):
         # This is required as when `TrainingTypePlugin` instance is passed to either `strategy`

pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py

@@ -21,7 +21,7 @@
 from pytorch_lightning.plugins.environments.slurm_environment import SLURMEnvironment
 from pytorch_lightning.trainer.connectors.logger_connector.result import _METRICS, _OUT_DICT, _PBAR_DICT
 from pytorch_lightning.trainer.states import RunningStage, TrainerFn
-from pytorch_lightning.utilities import DeviceType, memory
+from pytorch_lightning.utilities import _AcceleratorType, memory
 from pytorch_lightning.utilities.apply_func import apply_to_collection, move_data_to_device
 from pytorch_lightning.utilities.metrics import metrics_to_scalars
 from pytorch_lightning.utilities.warnings import rank_zero_deprecation
@@ -329,7 +329,7 @@ def gpus_metrics(self) -> Dict[str, float]:
         .. deprecated:: v1.5
             Will be removed in v1.7.
         """
-        if self.trainer._device_type == DeviceType.GPU and self.log_gpu_memory:
+        if self.trainer._device_type == _AcceleratorType.GPU and self.log_gpu_memory:
             mem_map = memory.get_memory_profile(self.log_gpu_memory)
             self._gpus_metrics.update(mem_map)
         return self._gpus_metrics