ref: unify slurm and TE under backendPlugin 3/n (#4581)

Author: williamFalcon

pytorch_lightning/accelerators/__init__.py

@@ -21,6 +21,5 @@
 from pytorch_lightning.accelerators.tpu_accelerator import TPUAccelerator
 from pytorch_lightning.accelerators.horovod_accelerator import HorovodAccelerator
 from pytorch_lightning.accelerators.ddp_hpc_accelerator import DDPHPCAccelerator
-from pytorch_lightning.accelerators.ddp_cpu_torchelastic_accelerator import DDPCPUTorchElasticAccelerator
-from pytorch_lightning.accelerators.ddp_cpu_slurm_accelerator import DDPCPUSLURMAccelerator
+from pytorch_lightning.accelerators.ddp_cpu_hpc_accelerator import DDPCPUHPCAccelerator
 from pytorch_lightning.accelerators.accelerator import Accelerator

pytorch_lightning/accelerators/accelerator_connector.py

@@ -220,7 +220,7 @@ def select_accelerator(self):
             )
 
         elif use_ddp_cpu_slurm:
-            accelerator_backend = accelerators.DDPCPUSLURMAccelerator(
+            accelerator_backend = accelerators.DDPCPUHPCAccelerator(
                 self.trainer,
                 cluster_env,
                 self.trainer.plugin_connector.ddp_plugin
@@ -234,7 +234,7 @@ def select_accelerator(self):
             )
 
         elif use_ddp_cpu_torch_elastic:
-            accelerator_backend = accelerators.DDPCPUTorchElasticAccelerator(
+            accelerator_backend = accelerators.DDPCPUHPCAccelerator(
                 self.trainer,
                 cluster_env,
                 self.trainer.plugin_connector.ddp_plugin

pytorch_lightning/accelerators/ddp_cpu_slurm_accelerator.py renamed to pytorch_lightning/accelerators/ddp_cpu_hpc_accelerator.py

@@ -37,12 +37,7 @@
     HYDRA_AVAILABLE = True
 
 
-# -------------------------------------------
-# !!!!!!!!!!!!!! NOTE !!!!!!!!!!!!!!!!!!!!!!
-# TEMP CLASS WHILE WE DECOUPLE TE FROM DDP
-# !!!!!!!!!!!!!! NOTE !!!!!!!!!!!!!!!!!!!!!!
-# -------------------------------------------
-class DDPCPUSLURMAccelerator(Accelerator):
+class DDPCPUHPCAccelerator(Accelerator):
 
     def __init__(self, trainer, cluster_environment=None, ddp_plugin=None):
         super().__init__(trainer, cluster_environment, ddp_plugin)

pytorch_lightning/accelerators/ddp_cpu_torchelastic_accelerator.py

@@ -229,7 +229,7 @@ def test_accelerator_choice_ddp_cpu_te(tmpdir):
     class CB(Callback):
         def on_fit_start(self, trainer, pl_module):
             assert trainer.use_ddp
-            assert isinstance(trainer.accelerator_backend, accelerators.DDPCPUTorchElasticAccelerator)
+            assert isinstance(trainer.accelerator_backend, accelerators.DDPCPUHPCAccelerator)
             assert isinstance(trainer.accelerator_backend.cluster_environment, TorchElasticEnvironment)
             assert trainer.accelerator_backend.task_idx == 10
             assert trainer.accelerator_backend.cluster_environment.local_rank() == trainer.accelerator_backend.task_idx
@@ -260,7 +260,7 @@ def test_accelerator_choice_ddp_cpu_slurm(tmpdir):
     class CB(Callback):
         def on_fit_start(self, trainer, pl_module):
             assert trainer.use_ddp
-            assert isinstance(trainer.accelerator_backend, accelerators.DDPCPUSLURMAccelerator)
+            assert isinstance(trainer.accelerator_backend, accelerators.DDPCPUHPCAccelerator)
             assert isinstance(trainer.accelerator_backend.cluster_environment, SLURMEnvironment)
             raise SystemExit()
 
@@ -295,7 +295,7 @@ def master_address(self):
     class CB(Callback):
         def on_fit_start(self, trainer, pl_module):
             assert trainer.use_ddp
-            assert isinstance(trainer.accelerator_backend, accelerators.DDPCPUSLURMAccelerator)
+            assert isinstance(trainer.accelerator_backend, accelerators.DDPCPUHPCAccelerator)
             assert isinstance(trainer.accelerator_backend.cluster_environment, CustomCluster)
             raise SystemExit()
 
@@ -353,7 +353,7 @@ def on_fit_start(self, trainer, pl_module):
 def test_dist_backend_accelerator_mapping(tmpdir):
     class CB(Callback):
         def on_fit_start(self, trainer, pl_module):
-            assert isinstance(trainer.accelerator_backend, accelerators.DDPCPUSLURMAccelerator)
+            assert isinstance(trainer.accelerator_backend, accelerators.DDPCPUHPCAccelerator)
             raise SystemExit()
 
     model = BoringModel()