Clean up environment access in plugins

pytorch_lightning/plugins/environments/cluster_environment.py

@@ -18,6 +18,9 @@
 class ClusterEnvironment(ABC):
     """ Specification of a cluster environment. """
 
+    def __init__(self):
+        super().__init__()
+
     @abstractmethod
     def creates_children(self) -> bool:
         """ Whether the environment creates the subprocesses or not. """
@@ -34,6 +37,18 @@ def master_port(self) -> int:
     def world_size(self) -> Optional[int]:
         """ The number of processes across all devices and nodes. """
 
+    @abstractmethod
+    def set_world_size(self, size: int) -> None:
+        pass
+
+    @abstractmethod
+    def global_rank(self) -> Optional[int]:
+        """ The rank (index) of the currently running process across all nodes and devices. """
+
+    @abstractmethod
+    def set_global_rank(self, rank: int) -> None:
+        pass
+
     @abstractmethod
     def local_rank(self) -> int:
         """ The rank (index) of the currently running process inside of the current node. """

pytorch_lightning/plugins/environments/lightning_environment.py

@@ -34,6 +34,8 @@ class LightningEnvironment(ClusterEnvironment):
     def __init__(self):
         super().__init__()
         self._master_port = None
+        self._global_rank = 0
+        self._world_size = None
 
     def creates_children(self) -> bool:
         return False
@@ -47,7 +49,16 @@ def master_port(self) -> int:
         return int(self._master_port)
 
     def world_size(self) -> Optional[int]:
-        return None
+        return self._world_size
+
+    def set_world_size(self, size: int) -> None:
+        self._world_size = size
+
+    def global_rank(self) -> Optional[int]:
+        return self._global_rank
+
+    def set_global_rank(self, rank: int) -> None:
+        self._global_rank = rank
 
     def local_rank(self) -> int:
         return int(os.environ.get("LOCAL_RANK", 0))

pytorch_lightning/plugins/environments/slurm_environment.py

@@ -15,6 +15,7 @@
 import logging
 import os
 import re
+from typing import Optional
 
 from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment
 
@@ -70,7 +71,16 @@ def master_port(self) -> int:
         return int(default_port)
 
     def world_size(self):
-        return None
+        return int(os.environ["SLURM_NTASKS"])
+
+    def set_world_size(self, size: int) -> None:
+        pass
+
+    def global_rank(self) -> Optional[int]:
+        return int(os.environ["SLURM_PROCID"])
+
+    def set_global_rank(self, rank: int) -> None:
+        pass
 
     def local_rank(self) -> int:
         return int(os.environ['SLURM_LOCALID'])

pytorch_lightning/plugins/environments/torchelastic_environment.py

@@ -51,6 +51,15 @@ def world_size(self) -> Optional[int]:
         world_size = os.environ.get('WORLD_SIZE')
         return int(world_size) if world_size is not None else world_size
 
+    def set_world_size(self, size: int) -> None:
+        pass
+
+    def global_rank(self) -> Optional[int]:
+        return int(os.environ["RANK"])
+
+    def set_global_rank(self, rank: int) -> None:
+        pass
+
     def local_rank(self) -> int:
         return int(os.environ['LOCAL_RANK'])
 

pytorch_lightning/plugins/training_type/ddp.py

@@ -78,11 +78,11 @@ def __init__(
         self._ddp_kwargs = kwargs
         self._has_spawned_children = False
         self.task_idx = None
-        self.node_rank = 0
         self.num_processes = len(parallel_devices) if parallel_devices is not None else parallel_devices
         self._ddp_comm_state = ddp_comm_state
         self._ddp_comm_hook = ddp_comm_hook
         self._ddp_comm_wrapper = ddp_comm_wrapper
+        rank_zero_only.rank = self.global_rank
 
     @property
     def root_device(self):
@@ -193,7 +193,7 @@ def setup_distributed(self):
         # set up server using proc 0's ip address
         # try to init for 20 times at max in case ports are taken
         # where to store ip_table
-        self.init_ddp_connection(self.global_rank, self.world_size)
+        self.init_ddp_connection()
 
         # on world_size=0 let everyone know training is starting
         if self.is_global_zero and not torch.distributed.is_initialized():
@@ -214,10 +214,9 @@ def _check_can_spawn_children(self):
             )
 
     def set_world_ranks(self):
-        self.local_rank = self.task_idx
-        self.node_rank = self.cluster_environment.node_rank()
-        self.global_rank = self.node_rank * self.num_processes + self.local_rank
-        self.world_size = self.num_nodes * self.num_processes
+        # self.local_rank = self.task_idx  # TODO ?? hmm ??
+        self.cluster_environment.set_global_rank(self.node_rank * self.num_processes + self.local_rank)
+        self.cluster_environment.set_world_size(self.num_nodes * self.num_processes)
 
     def pre_configure_ddp(self):
         # if unset, default `find_unused_parameters` `True`
@@ -260,11 +259,11 @@ def determine_ddp_device_ids(self):
             return None
         return [self.root_device.index]
 
-    def init_ddp_connection(self, global_rank: int, world_size: int) -> None:
-        os.environ["MASTER_ADDR"] = str(self.cluster_environment.master_address())
+    def init_ddp_connection(self, global_rank: Optional[int] = None, world_size: Optional[int] = None) -> None:
+        global_rank = global_rank if global_rank is not None else self.cluster_environment.global_rank()
+        world_size = world_size if world_size is not None else self.cluster_environment.world_size()
+        os.environ["MASTER_ADDR"] = self.cluster_environment.master_address()
         os.environ["MASTER_PORT"] = str(self.cluster_environment.master_port())
-        os.environ["WORLD_SIZE"] = str(self.cluster_environment.world_size())
-
         if not torch.distributed.is_initialized():
             log.info(f"initializing ddp: GLOBAL_RANK: {global_rank}, MEMBER: {global_rank + 1}/{world_size}")
             torch_distrib.init_process_group(self.torch_distributed_backend, rank=global_rank, world_size=world_size)

pytorch_lightning/plugins/training_type/ddp_spawn.py

@@ -61,11 +61,15 @@ def __init__(
         self._ddp_kwargs = kwargs
         self.dist = LightningDistributed()
         self.num_processes = len(parallel_devices) if parallel_devices is not None else 0
-        self.node_rank = 0
         self.mp_queue = None
         self._ddp_comm_state = ddp_comm_state
         self._ddp_comm_hook = ddp_comm_hook
         self._ddp_comm_wrapper = ddp_comm_wrapper
+        self._local_rank = 0
+
+    @property
+    def local_rank(self):
+        return self._local_rank
 
     def __getstate__(self):
         """ Makes this plugin pickleable without destroying the queue in the current process. """
@@ -96,11 +100,9 @@ def setup(self, model):
         self.mp_queue = smp.SimpleQueue()
 
     def set_world_ranks(self, process_idx):
-        self.local_rank = process_idx
-        self.node_rank = self.cluster_environment.node_rank()
-        self.task_idx = self.cluster_environment.local_rank()
-        self.global_rank = self.node_rank * self.num_processes + self.local_rank
-        self.world_size = self.num_nodes * self.num_processes
+        self._local_rank = process_idx
+        self.cluster_environment.set_global_rank(self.node_rank * self.num_processes + self.local_rank)
+        self.cluster_environment.set_world_size(self.num_nodes * self.num_processes)
 
     @property
     def mp_spawn_kwargs(self):
@@ -213,11 +215,12 @@ def configure_ddp(self):
         )
         self._register_ddp_hooks()
 
-    def init_ddp_connection(self, global_rank: int, world_size: int) -> None:
+    def init_ddp_connection(self, global_rank: Optional[int], world_size: Optional[int]) -> None:
         # TODO: this code is duplicated in DDP and DDPSpawn, make this a function
-        os.environ["MASTER_ADDR"] = str(self.cluster_environment.master_address())
+        global_rank = global_rank if global_rank is not None else self.cluster_environment.global_rank()
+        world_size = world_size if world_size is not None else self.cluster_environment.world_size()
+        os.environ["MASTER_ADDR"] = self.cluster_environment.master_address()
         os.environ["MASTER_PORT"] = str(self.cluster_environment.master_port())
-        os.environ["WORLD_SIZE"] = str(self.cluster_environment.world_size())
 
         if not torch.distributed.is_initialized():
             log.info(f"initializing ddp: GLOBAL_RANK: {global_rank}, MEMBER: {global_rank + 1}/{world_size}")

pytorch_lightning/plugins/training_type/dp.py

@@ -27,6 +27,22 @@ class DataParallelPlugin(ParallelPlugin):
     def __init__(self, parallel_devices: Optional[List[torch.device]]):
         super().__init__(parallel_devices=parallel_devices, cluster_environment=None)
 
+    @property
+    def global_rank(self):
+        return 0
+
+    @property
+    def local_rank(self):
+        return 0
+
+    @property
+    def node_rank(self):
+        return 0
+
+    @property
+    def world_size(self):
+        return 1
+
     def setup(self, model):
         # model needs to be moved to the device before it is wrapped
         model.to(self.root_device)

pytorch_lightning/plugins/training_type/horovod.py

@@ -31,24 +31,31 @@ class HorovodPlugin(ParallelPlugin):
 
     def __init__(self, parallel_devices: Optional[List[torch.device]] = None):
         super().__init__(parallel_devices=parallel_devices, cluster_environment=None)
+        rank_zero_only.rank = self.global_rank
+
+    @property
+    def global_rank(self):
+        return hvd.rank()
+
+    @property
+    def local_rank(self):
+        return hvd.local_rank()
+
+    @property
+    def world_size(self):
+        return hvd.size()
 
     @property
     def root_device(self):
         return self.parallel_devices[self.local_rank]
 
     @property
     def distributed_sampler_kwargs(self):
-        distributed_sampler_kwargs = dict(num_replicas=hvd.size(), rank=hvd.rank())
+        distributed_sampler_kwargs = dict(num_replicas=self.world_size, rank=self.global_rank)
         return distributed_sampler_kwargs
 
     def setup(self, model):
         self._model = model
-
-        self.global_rank = hvd.rank()
-        self.local_rank = hvd.local_rank()
-        self.world_size = hvd.size()
-        rank_zero_only.rank = self.global_rank
-
         self.model_to_device()
 
     def pre_dispatch(self):
@@ -63,14 +70,14 @@ def _unpack_lightning_optimizer(opt):
         # increased total batch size
         for optimizer in optimizers:
             for param_group in optimizer.param_groups:
-                param_group["lr"] *= hvd.size()
+                param_group["lr"] *= self.world_size
 
         # Horovod: adjust base LR used by schedulers to match scaled optimizer initial LR
         lr_schedulers = self.lightning_module.trainer.lr_schedulers
         for scheduler in lr_schedulers:
             scheduler = scheduler["scheduler"]
             if isinstance(scheduler, _LRScheduler):
-                scheduler.base_lrs = [lr * hvd.size() for lr in scheduler.base_lrs]
+                scheduler.base_lrs = [lr * self.world_size for lr in scheduler.base_lrs]
 
         # Horovod: broadcast parameters & optimizer state to ensure consistent initialization
         hvd.broadcast_parameters(self.lightning_module.state_dict(), root_rank=0)

pytorch_lightning/plugins/training_type/parallel.py

@@ -36,9 +36,6 @@ def __init__(
         super().__init__()
         self.parallel_devices = parallel_devices
         self.cluster_environment = cluster_environment
-        self.global_rank = 0
-        self.world_size = 1
-        self.local_rank = 0
 
     @property
     @abstractmethod
@@ -53,6 +50,22 @@ def on_gpu(self):
     def lightning_module(self):
         return unwrap_lightning_module(self._model)
 
+    @property
+    def global_rank(self):
+        return self.cluster_environment.global_rank() if self.cluster_environment is not None else 0
+
+    @property
+    def local_rank(self):
+        return self.cluster_environment.local_rank() if self.cluster_environment is not None else 0
+
+    @property
+    def node_rank(self):
+        return self.cluster_environment.node_rank() if self.cluster_environment is not None else 0
+
+    @property
+    def world_size(self):
+        return self.cluster_environment.world_size() if self.cluster_environment is not None else 1
+
     @property
     def is_global_zero(self) -> bool:
         return self.global_rank == 0

pytorch_lightning/plugins/training_type/rpc_sequential.py

@@ -100,8 +100,8 @@ def __init__(
 
     def init_ddp_connection(
         self,
-        global_rank: int,
-        world_size: int,
+        global_rank: Optional[int] = None,
+        world_size: Optional[int] = None,
     ) -> None:
         if self.lightning_module.trainer.amp_backend is not None:
             raise MisconfigurationException(
@@ -110,10 +110,10 @@ def init_ddp_connection(
 
         if self._skip_init_connections():
             return
-        super().init_ddp_connection(
-            global_rank=global_rank,
-            world_size=world_size,
-        )
+
+        global_rank = global_rank if global_rank is not None else self.cluster_environment.global_rank()
+        world_size = world_size if world_size is not None else self.cluster_environment.world_size()
+        super().init_ddp_connection(global_rank, world_size)
         super().init_rpc_connection(global_rank=global_rank, world_size=world_size)
         model = self.lightning_module
         self.gpus_per_model = self._infer_check_num_gpus()

pytorch_lightning/trainer/connectors/accelerator_connector.py

@@ -302,8 +302,8 @@ def root_gpu(self) -> Optional[int]:
 
     @property
     def is_using_torchelastic(self) -> bool:
-        te_flags_passed = "WORLD_SIZE" in os.environ and ("GROUP_RANK" in os.environ or "NODE_RANK" in os.environ)
-        return te_flags_passed
+        required_env_vars = ("RANK", "GROUP_RANK", "LOCAL_RANK")
+        return all(v in os.environ for v in required_env_vars)
 
     def select_precision_plugin(self) -> PrecisionPlugin:
         # set precision type