ref: adding compute environments (2/n) (#3842)

* ref: adding compute environments (2/n)

* ref: adding compute environments (2/n)

* ref: adding compute environments (2/n)

* ref: adding compute environments (2/n)

Author: williamFalcon

pytorch_lightning/accelerators/accelerator_connector.py

@@ -7,6 +7,8 @@
 from pytorch_lightning.utilities.distributed import rank_zero_warn, rank_zero_info
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning import _logger as log
+from pytorch_lightning.cluster_environments.slurm_environment import SLURMEnvironment
+from pytorch_lightning.cluster_environments.torchelastic_environment import TorchElasticEnvironment
 
 try:
     import torch_xla
@@ -40,9 +42,12 @@ def on_trainer_init(
             sync_batchnorm,
             benchmark,
             replace_sampler_ddp,
-            deterministic
+            deterministic,
+            cluster_environment
     ):
         self.trainer.deterministic = deterministic
+        self.cluster_environment = cluster_environment
+
         torch.backends.cudnn.deterministic = self.trainer.deterministic
         if self.trainer.deterministic:
             # fixing non-deterministic part of horovod
@@ -123,6 +128,22 @@ def on_trainer_init(
 
         self.trainer.replace_sampler_ddp = replace_sampler_ddp
 
+    def _select_environment(self):
+        env = None
+
+        # in priority: user environment, torchelastic (which is a generic environment), slurm
+        if self.cluster_environment is not None:
+            env = self.cluster_environment
+        elif self._is_using_torchelastic():
+            env = TorchElasticEnvironment()
+        elif self.trainer.is_slurm_managing_tasks:
+            env = SLURMEnvironment()
+        return env
+
+    def _is_using_torchelastic(self):
+        te_flags_passed = 'WORLD_SIZE' in os.environ and ('GROUP_RANK' in os.environ or 'NODE_RANK' in os.environ)
+        return te_flags_passed
+
     def select_accelerator(self):
         if self.trainer.accelerator_backend is not None:
             return self.trainer.accelerator_backend

pytorch_lightning/accelerators/base_backend.py

@@ -23,8 +23,9 @@
 
 class Accelerator(object):
 
-    def __init__(self, trainer):
+    def __init__(self, trainer, cluster_environment=None):
         self.trainer = trainer
+        self.cluster_environment = cluster_environment
         self.dist = AttributeDict(rank=0, device=None)
 
     def setup(self, model):

pytorch_lightning/accelerators/cpu_backend.py

@@ -20,8 +20,8 @@
 
 class CPUBackend(Accelerator):
 
-    def __init__(self, trainer):
-        super().__init__(trainer)
+    def __init__(self, trainer, cluster_environment=None):
+        super().__init__(trainer, cluster_environment)
 
     def setup(self, model):
         # run through amp wrapper

pytorch_lightning/accelerators/ddp2_backend.py

@@ -35,8 +35,8 @@
 
 class DDP2Backend(Accelerator):
 
-    def __init__(self, trainer):
-        super().__init__(trainer)
+    def __init__(self, trainer, cluster_environment=None):
+        super().__init__(trainer, cluster_environment)
         self.task_idx = None
         self.dist = LightningDistributed()
 

pytorch_lightning/accelerators/ddp_backend.py

@@ -42,8 +42,8 @@
 
 class DDPBackend(Accelerator):
 
-    def __init__(self, trainer):
-        super().__init__(trainer)
+    def __init__(self, trainer, cluster_environment=None):
+        super().__init__(trainer, cluster_environment)
         self.task_idx = None
         self._has_spawned_children = False
         self.interactive_ddp_procs = []

pytorch_lightning/accelerators/ddp_cpu_spawn_backend.py

@@ -38,8 +38,8 @@
 
 class DDPCPUSpawnBackend(Accelerator):
 
-    def __init__(self, trainer, nprocs):
-        super().__init__(trainer)
+    def __init__(self, trainer, nprocs, cluster_environment=None):
+        super().__init__(trainer, cluster_environment)
         self.mp_queue = None
         self.nprocs = nprocs
         self.dist = LightningDistributed()

pytorch_lightning/accelerators/ddp_slurm_backend.py

@@ -40,8 +40,8 @@
 # -------------------------------------------
 class DDPSLURMBackend(Accelerator):
 
-    def __init__(self, trainer):
-        super().__init__(trainer)
+    def __init__(self, trainer, cluster_environment=None):
+        super().__init__(trainer, cluster_environment)
         self.task_idx = None
         self._has_spawned_children = False
         self.dist = LightningDistributed()

pytorch_lightning/accelerators/ddp_spawn_backend.py

@@ -40,8 +40,8 @@
 
 class DDPSpawnBackend(Accelerator):
 
-    def __init__(self, trainer, nprocs):
-        super().__init__(trainer)
+    def __init__(self, trainer, nprocs, cluster_environment=None):
+        super().__init__(trainer, cluster_environment)
         self.mp_queue = None
         self.nprocs = nprocs
         self.dist = LightningDistributed()

pytorch_lightning/accelerators/ddp_torchelastic_backend.py

@@ -40,8 +40,8 @@
 # -------------------------------------------
 class DDPTorchElasticBackend(Accelerator):
 
-    def __init__(self, trainer):
-        super().__init__(trainer)
+    def __init__(self, trainer, cluster_environment=None):
+        super().__init__(trainer, cluster_environment)
         self.task_idx = None
         self._has_spawned_children = False
         self.dist = LightningDistributed()

pytorch_lightning/accelerators/dp_backend.py

@@ -25,8 +25,8 @@
 
 class DataParallelBackend(Accelerator):
 
-    def __init__(self, trainer):
-        super().__init__(trainer)
+    def __init__(self, trainer, cluster_environment=None):
+        super().__init__(trainer, cluster_environment)
         self.model_autocast_original_forward = None
         self.dist = LightningDistributed()