testing env init

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):
         """