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

pl_examples/basic_examples/gpu_template.py

@@ -20,11 +20,20 @@ def main(args):
     # ------------------------
     # 2 INIT TRAINER
     # ------------------------
-    trainer = Trainer.from_argparse_args(args)
+    trainer = Trainer.from_argparse_args(
+        args,
+        distributed_backend='ddp',
+        limit_train_batches=10,
+        limit_val_batches=10,
+        max_epochs=1,
+    )
 
     # ------------------------
     # 3 START TRAINING
     # ------------------------
+    trainer.test(model)
+    trainer.fit(model)
+    trainer.test(model)
     trainer.fit(model)
 
 

pl_examples/basic_examples/gpu_template2.py

@@ -0,0 +1,66 @@
+"""
+Runs a model on a single node across multiple gpus.
+"""
+import os
+from argparse import ArgumentParser
+
+from pytorch_lightning import Trainer, seed_everything, Callback
+from pl_examples.models.lightning_template import LightningTemplateModel
+
+seed_everything(234)
+
+
+class DebugCallback(Callback):
+
+    def on_test_batch_end(self, trainer, pl_module):
+        print('test_batch', trainer.global_rank)
+
+
+def main(args):
+    """ Main training routine specific for this project. """
+    # ------------------------
+    # 1 INIT LIGHTNING MODEL
+    # ------------------------
+    model = LightningTemplateModel(**vars(args))
+
+    # ------------------------
+    # 2 INIT TRAINER
+    # ------------------------
+    trainer = Trainer.from_argparse_args(
+        args,
+        distributed_backend='ddp',
+        limit_train_batches=10,
+        limit_val_batches=10,
+        max_epochs=1,
+        callbacks=[DebugCallback()],
+    )
+
+    # ------------------------
+    # 3 START TRAINING
+    # ------------------------
+    trainer.fit(model)
+    trainer.test(model)
+
+
+def run_cli():
+    # ------------------------
+    # TRAINING ARGUMENTS
+    # ------------------------
+    # these are project-wide arguments
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    parent_parser = ArgumentParser(add_help=False)
+
+    # each LightningModule defines arguments relevant to it
+    parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+    parser = Trainer.add_argparse_args(parser)
+    parser.set_defaults(gpus=2)
+    args = parser.parse_args()
+
+    # ---------------------
+    # RUN TRAINING
+    # ---------------------
+    main(args)
+
+
+if __name__ == '__main__':
+    run_cli()

pytorch_lightning/accelerator_backends/ddp_backend.py

@@ -56,19 +56,19 @@ def train(self, model):
         self.ddp_train(process_idx=self.task_idx, mp_queue=None, model=model)
 
     def spawn_ddp_children(self, model):
+        #
+        assert self.trainer.global_rank == 0
+        self.trainer.set_random_port(force=True)
         port = os.environ['MASTER_PORT']
 
-        master_address = '127.0.0.1' if 'MASTER_ADDR' not in os.environ else os.environ['MASTER_ADDR']
+        master_address = os.environ.get('MASTER_ADDR', '127.0.0.1')
         os.environ['MASTER_PORT'] = f'{port}'
         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 +153,18 @@ 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
+
+        # from torch.distributed import is_initialized
+        # if not is_master or not is_initialized():
+        #     assert not (is_master and self.trainer.global_rank > 0)
+        #     # on rank > 0, we always need to initialize, because these are new processes
         model.init_ddp_connection(
             self.trainer.global_rank,
             self.trainer.world_size,
             self.trainer.is_slurm_managing_tasks
         )
+        # else:
+        #     print('already initialized', os.environ['MASTER_PORT'], os.getpid(), is_master)
 
         # call setup after the ddp process has connected
         self.trainer.call_setup_hook(model)
@@ -225,5 +232,12 @@ def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0
         # clean up memory
         torch.cuda.empty_cache()
 
+        # clean up dist group
+        #if self.use_ddp or self.use_ddp2:
+        # import torch.distributed as torch_distrib
+        # torch_distrib.destroy_process_group()
+
+        # torch.distributed.destroy_process_group()
+
         if self.trainer.global_rank == 0 and self.trainer.distributed_backend not in ['ddp_spawn', 'ddp_cpu']:
             return results

pytorch_lightning/core/decorators.py

@@ -1,8 +1,6 @@
 from functools import wraps
 from typing import Callable
 
-from pytorch_lightning.core.lightning import LightningModule
-
 
 def auto_move_data(fn: Callable) -> Callable:
     """
@@ -40,6 +38,9 @@ def forward(self, x):
     """
     @wraps(fn)
     def auto_transfer_args(self, *args, **kwargs):
+        # local import to prevent circular import issue
+        from pytorch_lightning.core.lightning import LightningModule
+
         if not isinstance(self, LightningModule):
             return fn(self, *args, **kwargs)
 
@@ -48,3 +49,20 @@ def auto_transfer_args(self, *args, **kwargs):
         return fn(self, *args, **kwargs)
 
     return auto_transfer_args
+
+
+def run_once(fn):
+    """
+    Decorate a function or method to make it run only once.
+    Subsequent calls will result in a no-operation.
+    """
+    @wraps(fn)
+    def wrapper(*args, **kwargs):
+        if not wrapper.has_run:
+            wrapper.has_run = True
+            fn(*args, **kwargs)
+
+    wrapper.has_run = False
+    return wrapper
+
+

pytorch_lightning/core/lightning.py

@@ -16,6 +16,7 @@
 from torch.utils.data import DataLoader
 
 from pytorch_lightning import _logger as log
+from pytorch_lightning.core.decorators import run_once
 from pytorch_lightning.core.grads import GradInformation
 from pytorch_lightning.core.hooks import ModelHooks
 from pytorch_lightning.core.memory import ModelSummary
@@ -921,6 +922,7 @@ def _init_slurm_connection(self) -> None:
         root_node = self.trainer.resolve_root_node_address(root_node)
         os.environ['MASTER_ADDR'] = root_node
 
+    #@run_once
     def init_ddp_connection(self, global_rank: int, world_size: int, is_slurm_managing_tasks: bool = True) -> None:
         """
         Override to define your custom way of setting up a distributed environment.
@@ -952,7 +954,7 @@ def init_ddp_connection(self, global_rank: int, world_size: int, is_slurm_managi
                 f"WORLD_SIZE environment variable ({os.environ['WORLD_SIZE']}) "
                 f"is not equal to the computed world size ({world_size}). Ignored."
             )
-
+        print('master port init', os.environ['MASTER_PORT'], os.getpid())
         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}")
         torch_distrib.init_process_group(torch_backend, rank=global_rank, world_size=world_size)

pytorch_lightning/trainer/distrib_data_parallel.py

@@ -171,9 +171,11 @@ 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)
+
+#PID = os.getpid()
+#RNG1 = np.random.RandomState(PID)
+#RANDOM_PORTS = RNG1.randint(10000, 19999, 1000)
+RANDOM_PORTS = list(range(10000, 20000))
 
 
 class TrainerDDPMixin(ABC):
@@ -411,13 +413,15 @@ def set_random_port(self, force=False):
         """
         # 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]
+
+        print('setting port on rank', self.global_rank)
+        default_port = os.environ.get('MASTER_PORT')
 
         # when not forced, use the user port
-        if not force:
-            default_port = os.environ.get('MASTER_PORT', default_port)
+        if force or not default_port:
+            global RANDOM_PORTS
+            default_port = RANDOM_PORTS[-1]
+            RANDOM_PORTS = RANDOM_PORTS[:-1]
 
         os.environ['MASTER_PORT'] = str(default_port)
 

pytorch_lightning/trainer/evaluation_loop.py

@@ -291,6 +291,7 @@ def _evaluate(
 
         # run validation
         for dataloader_idx, dataloader in enumerate(dataloaders):
+            print('here 1')
             dl_outputs = []
 
             # on TPU we have to wrap it under the ParallelLoader
@@ -303,6 +304,7 @@ def _evaluate(
             dl_max_batches = max_batches[dataloader_idx]
 
             for batch_idx, batch in enumerate(dataloader):
+                print('here 2')
                 if batch is None:
                     continue
 
@@ -600,16 +602,19 @@ def __log_evaluation_epoch_metrics(self, eval_results, test_mode):
     def evaluation_forward(self, model, batch, batch_idx, dataloader_idx, test_mode: bool = False):
         # make dataloader_idx arg in validation_step optional
         args = [batch, batch_idx]
-
+        print('here 3')
         if (test_mode and len(self.test_dataloaders) > 1) \
                 or (not test_mode and len(self.val_dataloaders) > 1):
             args.append(dataloader_idx)
 
         # handle DP, DDP forward
         if self.use_ddp or self.use_dp or self.use_ddp2:
+            # SOMETHING GOES WRONG HERE, test loop is stuck
             output = model(*args)
             return output
 
+        print('here 4')
+
         # Horovod
         if self.use_horovod and self.on_gpu:
             batch = self.transfer_batch_to_gpu(batch, hvd.local_rank())
@@ -635,4 +640,6 @@ def evaluation_forward(self, model, batch, batch_idx, dataloader_idx, test_mode:
         else:
             output = model.validation_step(*args)
 
+        print('here 5')
+
         return output

pytorch_lightning/trainer/trainer.py

@@ -1019,7 +1019,7 @@ def fit(
 
         # ddp
         elif self.distributed_backend == 'ddp':
-            self.set_random_port()
+            # self.set_random_port()
             self.accelerator_backend = DDPBackend(self)
             results = self.accelerator_backend.spawn_ddp_children(model)
 
@@ -1296,6 +1296,7 @@ def test(
         self.verbose_test = verbose
 
         if self.global_rank != 0:
+            # do nothing, rank 0 process will launch new processes for testing
             return
 
         # If you supply a datamodule you can't supply train_dataloader or val_dataloaders
@@ -1314,6 +1315,10 @@ def test(
 
         self.teardown('test')
 
+        if torch.distributed.is_initialized():
+            print('destroy in test', self.global_rank, os.getpid())
+            torch.distributed.destroy_process_group()
+
         return results
 
     def __test_using_best_weights(self, ckpt_path, test_dataloaders):
@@ -1347,7 +1352,7 @@ def __test_using_best_weights(self, ckpt_path, test_dataloaders):
 
         # run tests
         self.tested_ckpt_path = ckpt_path
-        self.set_random_port(force=True)
+        #self.set_random_port()
         self.testing = True
         os.environ['PL_TESTING_MODE'] = '1'
         self.model = model
@@ -1370,7 +1375,7 @@ def __test_given_model(self, model, test_dataloaders):
 
         # run test
         # sets up testing so we short circuit to eval
-        self.set_random_port(force=True)
+        #self.set_random_port()
         self.testing = True
         self.model = model
         results = self.fit(model)

pytorch_lightning/trainer/training_loop.py

@@ -1021,7 +1021,8 @@ def run_training_teardown(self):
                 subprocess.Popen.kill(proc)
 
         # clean up dist group
-        if self.use_ddp or self.use_ddp2:
+        if (self.use_ddp or self.use_ddp2):
+            print('destroy on rank ', self.global_rank, os.getpid())
             torch_distrib.destroy_process_group()
 
         # clear mem