diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index 086084e6..289efe3b 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -347,17 +347,7 @@ class Trainer(TrainerIO): ip_file = os.path.join(ip_file_dir, ip_table_name) os.makedirs(ip_file_dir, exist_ok=True) - root_node = os.environ['SLURM_NODELIST'].split(' ')[0] - print('-'*100) - print('SLURM ROOT NODE: ', root_node) - print('-'*100) - - os.environ['MASTER_ADDR'] = root_node - os.environ['MASTER_PORT'] = '12006' - - dist.init_process_group("nccl", rank=self.proc_rank, - world_size=self.world_size) - # self.__init_tcp_connection(ip_file_dir) + self.__init_tcp_connection(ip_file_dir) # CHOOSE OPTIMIZER # filter out the weights that were done on gpu so we can load on good old cpus @@ -382,16 +372,20 @@ class Trainer(TrainerIO): # continue training routine self.__run_pretrain_routine(model) - def __init_tcp_connection(self, path, port=12000, tries=0): + def __init_tcp_connection(self, port=12000, tries=0): if tries > 20: raise RuntimeError('Failed to connect using 20 different ip addresses') try: - dist.init_process_group("nccl", init_method=f'file://{path}:{port}', rank=self.proc_rank, world_size=self.world_size) + root_node = os.environ['SLURM_NODELIST'].split(' ')[0] + os.environ['MASTER_ADDR'] = root_node + os.environ['MASTER_PORT'] = '12006' + dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size) + except RuntimeError as e: # port taken warnings.warn(f'port {port} taken, trying port {port}...') - self.__init_tcp_connection(path, port + 1, tries + 1) + self.__init_tcp_connection(port + 1, tries + 1) def __get_root_node_ip(self, world_gpu_nb, nb_gpu_nodes): """