[WIP] Fix Trainer.test in ddp before running Trainer.fit

pytorch_lightning/accelerators/ddp_backend.py

@@ -11,20 +11,22 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License
-
+import atexit
 import os
+import socket
+
+import torch
+import torch.distributed
 import subprocess
 import sys
-from os.path import abspath
 from time import sleep
-from typing import Optional
-
 import numpy as np
-import torch
+from os.path import abspath
 
+from pytorch_lightning.utilities import NATIVE_AMP_AVALAIBLE
+from pytorch_lightning.utilities.distributed import rank_zero_only, rank_zero_debug
 from pytorch_lightning import _logger as log
-from pytorch_lightning.utilities import AMPType
-from pytorch_lightning.utilities.distributed import rank_zero_only
+from typing import Optional
 
 try:
     from hydra.utils import to_absolute_path, get_original_cwd
@@ -37,14 +39,17 @@
 try:
     from apex import amp
 except ImportError:
-    amp = None
+    APEX_AVAILABLE = False
+else:
+    APEX_AVAILABLE = True
 
 
 class DDPBackend(object):
 
     def __init__(self, trainer):
         self.trainer = trainer
         self.task_idx = None
+        self.distributed_connection = DistributedConnection(trainer)
 
     def slurm_setup(self):
         self.task_idx = int(os.environ['SLURM_LOCALID'])
@@ -56,19 +61,15 @@ def train(self, model):
         self.ddp_train(process_idx=self.task_idx, mp_queue=None, model=model)
 
     def spawn_ddp_children(self, model):
-        port = os.environ['MASTER_PORT']
+        assert self.trainer.global_rank == 0
 
-        master_address = '127.0.0.1' if 'MASTER_ADDR' not in os.environ else os.environ['MASTER_ADDR']
-        os.environ['MASTER_PORT'] = f'{port}'
+        master_address = os.environ.get('MASTER_ADDR', '127.0.0.1')
         os.environ['MASTER_ADDR'] = f'{master_address}'
 
         # allow the user to pass the node rank
         node_rank = '0'
-        if 'NODE_RANK' in os.environ:
-            node_rank = os.environ['NODE_RANK']
-        if 'GROUP_RANK' in os.environ:
-            node_rank = os.environ['GROUP_RANK']
-
+        node_rank = os.environ.get('NODE_RANK', node_rank)
+        node_rank = os.environ.get('GROUP_RANK', node_rank)
         os.environ['NODE_RANK'] = node_rank
         os.environ['LOCAL_RANK'] = '0'
 
@@ -153,11 +154,8 @@ def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0
         # try to init for 20 times at max in case ports are taken
         # where to store ip_table
         model.trainer = self.trainer
-        model.init_ddp_connection(
-            self.trainer.global_rank,
-            self.trainer.world_size,
-            self.trainer.is_slurm_managing_tasks
-        )
+
+        self.distributed_connection.reset_connection(self.trainer, model)
 
         # call setup after the ddp process has connected
         self.trainer.call_setup_hook(model)
@@ -176,6 +174,8 @@ def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0
         self.trainer.lr_schedulers = lr_schedulers
         self.trainer.optimizer_frequencies = optimizer_frequencies
 
+        print('here 1')
+
         # call sync_bn before .cuda(), configure_apex and configure_ddp
         if self.trainer.sync_batchnorm:
             model = model.configure_sync_batchnorm(model)
@@ -193,15 +193,20 @@ def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0
                 available_gpus = os.environ['CUDA_VISIBLE_DEVICES'].split(',')
                 gpu_idx = int(available_gpus[self.trainer.local_rank])
 
+            print('here 2')
             self.trainer.root_gpu = gpu_idx
             torch.cuda.set_device(self.trainer.root_gpu)
             model.cuda(self.trainer.root_gpu)
 
+        print('here 3')
+
         # set model properties before going into wrapper
         self.trainer.copy_trainer_model_properties(model)
 
-        # AMP - run through amp wrapper before going to distributed DP
-        if self.trainer.amp_type == AMPType.APEX:
+        # AMP
+        # run through amp wrapper before going to distributed DP
+        # TODO: remove with dropping NVIDIA AMP support
+        if self.trainer.use_amp and not NATIVE_AMP_AVALAIBLE:
             model, optimizers = model.configure_apex(amp, model, self.trainer.optimizers, self.trainer.amp_level)
             self.trainer.optimizers = optimizers
             self.trainer.reinit_scheduler_properties(self.trainer.optimizers, self.trainer.lr_schedulers)
@@ -212,12 +217,18 @@ def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0
         else:  # includes ddp_cpu
             device_ids = None
 
+        print('here 4')
+
         # allow user to configure ddp
         model = model.configure_ddp(model, device_ids)
 
+        print('here 5')
+
         # continue training routine
         results = self.trainer.run_pretrain_routine(model)
 
+        print('here 6')
+
         # get original model
         model = self.trainer.get_model()
 
@@ -229,3 +240,98 @@ def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0
 
         if self.trainer.global_rank == 0 and self.trainer.distributed_backend not in ['ddp_spawn', 'ddp_cpu']:
             return results
+
+
+class DistributedConnection:
+
+    def __init__(self, trainer):
+        super().__init__()
+        self.trainer = trainer
+        if trainer.num_nodes == 1:
+            # select or forcibly set an initial port before ddp connection is initialized
+            self._set_master_port(port=self._get_master_port())
+
+    def reset_connection(self, trainer, model):
+        if not torch.distributed.is_initialized():
+            print('init ddp', 'rank', trainer.global_rank, 'port', self._get_master_port())
+            model.init_ddp_connection(trainer.global_rank, trainer.world_size, trainer.is_slurm_managing_tasks)
+
+    def reset_connection_old(self, trainer, model):
+
+        if not torch.distributed.is_initialized():
+            print('init ddp', 'rank', trainer.global_rank, 'port', self._get_master_port())
+            model.init_ddp_connection(trainer.global_rank, trainer.world_size, trainer.is_slurm_managing_tasks)
+            print('init ddp', 'rank', trainer.global_rank, 'port', self._get_master_port(), 'done')
+
+        new_port = torch.tensor([int(self._get_master_port())], dtype=torch.int, device='cuda')
+        if torch.distributed.is_initialized() and trainer.global_rank == 0:
+            print(trainer.global_rank, "DDP connection already initialized. Reinitializing on new port...")
+
+            #model.init_ddp_connection(trainer.global_rank, trainer.world_size, trainer.is_slurm_managing_tasks)
+
+            # torch.distributed.barrier()
+
+
+            #if trainer.global_rank == 0:
+            port = find_open_network_port()
+            new_port[0] = port
+
+        torch.distributed.broadcast(new_port, src=0)
+        new_port = int(new_port.item())
+        print('recv new port', 'rank', trainer.global_rank, 'port', new_port)
+
+        if int(self._get_master_port()) != new_port:
+            print('need to update port')
+            torch.distributed.destroy_process_group()  # destroy connections on old port
+            print('destroy group', 'rank', trainer.global_rank, 'port', self._get_master_port())
+            print('set port', 'rank', trainer.global_rank, 'port', self._get_master_port())
+            self._set_master_port(port=new_port)
+
+            model.init_ddp_connection(trainer.global_rank, trainer.world_size, trainer.is_slurm_managing_tasks)
+
+        print('exit')
+
+        # s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        # #print('shutdown', self._get_master_address(), int(self._get_master_port()))
+        # s.connect((self._get_master_address(), int(self._get_master_port())))
+        # s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        # #s.shutdown(socket.SHUT_RDWR)
+        # s.close()
+        # #sleep(10)
+
+        def exit_handler():
+            if torch.distributed.is_initialized() and trainer.global_rank > 0:
+                print('destroying on ', trainer.global_rank)
+                torch.distributed.destroy_process_group()
+
+        atexit.register(exit_handler)
+
+    def _get_master_port(self):
+        return os.environ.get('MASTER_PORT')
+
+    def _get_master_address(self):
+        return os.environ.get('MASTER_ADDR')
+
+    def _set_master_port(self, port: int = None):
+        """
+        Sets the `MASTER_PORT` environment variable in single-node DDP training.
+
+        Args:
+            port: If provided, sets the environment variable MASTER_PORT, and otherwhise
+                an attempt is made to find an unused open port.
+
+        Return:
+            The port that was set.
+        """
+        assert self.trainer.num_nodes == 1, 'random port can only be called from single node training'
+        os.environ['MASTER_PORT'] = str(port or find_open_network_port())
+        return port
+
+
+def find_open_network_port():
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    s.bind(("", 0))
+    s.listen(1)
+    port = s.getsockname()[1]
+    s.close()
+    return port

pytorch_lightning/accelerators/ddp_spawn_backend.py

@@ -15,25 +15,26 @@
 import torch
 import torch.multiprocessing as mp
 
-from pytorch_lightning import _logger as log
-from pytorch_lightning.utilities import AMPType
+from pytorch_lightning.accelerators.ddp_backend import DistributedConnection
 from pytorch_lightning.utilities.distributed import rank_zero_only
+from pytorch_lightning import _logger as log
 
 try:
     from apex import amp
 except ImportError:
-    amp = None
+    APEX_AVAILABLE = False
+else:
+    APEX_AVAILABLE = True
 
 
 class DDPSpawnBackend(object):
 
     def __init__(self, trainer):
         self.trainer = trainer
         self.mp_queue = None
+        self.distributed_connection = DistributedConnection(trainer)
 
     def setup(self):
-        self.trainer.set_random_port()
-
         # pass in a state q
         smp = mp.get_context('spawn')
         self.mp_queue = smp.SimpleQueue()
@@ -94,11 +95,8 @@ def ddp_train(self, process_idx, mp_queue, model):
         # try to init for 20 times at max in case ports are taken
         # where to store ip_table
         model.trainer = self.trainer
-        model.init_ddp_connection(
-            self.trainer.global_rank,
-            self.trainer.world_size,
-            self.trainer.is_slurm_managing_tasks
-        )
+
+        self.distributed_connection.reset_connection(self.trainer, model)
 
         # call setup after the ddp process has connected
         self.trainer.call_setup_hook(model)
@@ -132,9 +130,11 @@ def ddp_train(self, process_idx, mp_queue, model):
         # set model properties before going into wrapper
         self.trainer.copy_trainer_model_properties(model)
 
-        # AMP -
+        # AMP
         # run through amp wrapper before going to distributed DP
-        if self.trainer.amp_type == AMPType.APEX:
+        # TODO: remove with dropping NVIDIA AMP support
+        native_amp_available = hasattr(torch.cuda, "amp") and hasattr(torch.cuda.amp, "autocast")
+        if self.trainer.use_amp and not native_amp_available:
             model, optimizers = model.configure_apex(amp, model, self.trainer.optimizers, self.trainer.amp_level)
             self.trainer.optimizers = optimizers
             self.trainer.reinit_scheduler_properties(self.trainer.optimizers, self.trainer.lr_schedulers)

pytorch_lightning/core/lightning.py

@@ -954,7 +954,7 @@ def init_ddp_connection(self, global_rank: int, world_size: int, is_slurm_managi
             )
 
         torch_backend = "nccl" if self.trainer.on_gpu else "gloo"
-        log.info(f"initializing ddp: GLOBAL_RANK: {global_rank}, MEMBER: {global_rank+1}/{world_size}")
+        log.info(f"initializing ddp: GLOBAL_RANK: {global_rank}, MEMBER: {global_rank+1}/{world_size}, ADDR: {os.environ['MASTER_ADDR']}")
         torch_distrib.init_process_group(torch_backend, rank=global_rank, world_size=world_size)
 
     def configure_sync_batchnorm(self, model: 'LightningModule') -> 'LightningModule':

pytorch_lightning/trainer/distrib_data_parallel.py

@@ -131,13 +131,7 @@ def train_fx(trial_hparams, cluster_manager, _):
 import re
 from abc import ABC, abstractmethod
 from distutils.version import LooseVersion
-from typing import Union, List, Optional, Callable, Tuple
-import subprocess
-import sys
-from time import sleep
-import numpy as np
-from os.path import abspath
-from pkg_resources import parse_version
+from typing import Union, List, Optional, Tuple
 
 import torch
 from pytorch_lightning import _logger as log
@@ -168,10 +162,6 @@ def train_fx(trial_hparams, cluster_manager, _):
 else:
     XLA_AVAILABLE = True
 
-PID = os.getpid()
-RNG1 = np.random.RandomState(PID)
-RANDOM_PORTS = RNG1.randint(10000, 19999, 1000)
-
 
 class TrainerDDPMixin(ABC):
 
@@ -397,22 +387,6 @@ def set_nvidia_flags(self, is_slurm_managing_tasks, data_parallel_device_ids):
         # don't make this debug... this is good UX
         rank_zero_info(f'CUDA_VISIBLE_DEVICES: [{os.environ["CUDA_VISIBLE_DEVICES"]}]')
 
-    def set_random_port(self, force=False):
-        """
-        When running DDP NOT managed by SLURM, the ports might collide
-        """
-        # pick a random port first
-        assert self.num_nodes == 1, 'random port can only be called from single node training'
-        global RANDOM_PORTS
-        default_port = RANDOM_PORTS[-1]
-        RANDOM_PORTS = RANDOM_PORTS[:-1]
-
-        # when not forced, use the user port
-        if not force:
-            default_port = os.environ.get('MASTER_PORT', default_port)
-
-        os.environ['MASTER_PORT'] = str(default_port)
-
     def transfer_distrib_spawn_state_on_fit_end(self, model, mp_queue, results):
         if self.distributed_backend.lower() not in ['ddp_spawn', 'ddp_cpu', 'tpu']:
             return