using slurm flag to fine node nb

examples/new_project_templates/trainer_gpu_cluster_template.py

@@ -42,9 +42,6 @@ def main(hparams, cluster, results_dict):
     :param hparams:
     :return:
     """
-    path = 'emv_' + os.environ['SLURM_LOCALID'] + '_id_' + os.environ['SLURM_NODEID']
-    os.makedirs(os.path.join(hparams.test_tube_save_path, path), exist_ok=True)
-
     # ------------------------
     # 1 INIT LIGHTNING MODEL
     # ------------------------
@@ -55,12 +52,17 @@ def main(hparams, cluster, results_dict):
     # ------------------------
     # 2 INIT TEST TUBE EXP
     # ------------------------
+    # when using grid search, it's possible for all models to start at once
+    # and use the same test tube experiment version
+    relative_node_id = int(os.environ['SLURM_NODEID'])
+    sleep(relative_node_id + 1)
+
+    # init experiment
     exp = Experiment(
         name=hyperparams.experiment_name,
         save_dir=hyperparams.test_tube_save_path,
         autosave=False,
-        description='test demo',
-        version=cluster.hpc_exp_number
+        description='test demo'
     )
 
     exp.argparse(hparams)

pytorch_lightning/models/trainer.py

@@ -12,7 +12,6 @@ import torch.distributed as dist
 import os
 import subprocess
 from time import sleep
-import fcntl
 
 try:
     from apex import amp
@@ -303,23 +302,13 @@ class Trainer(TrainerIO):
 
         # when GPU is called, spawn off a single worker for each gpu
         if self.on_gpu:
-            # every process writes their process id + ip to a shared file
-            my_ip = subprocess.run(['hostname', '-I'], stdout=subprocess.PIPE).stdout.decode('utf-8')
-            my_ip = my_ip.split(' ')[0]
-
-            r = np.random.uniform(0, 1, 1)[0]
-
-            test_name = f'{os.getpid()}_{my_ip}_{r}'
-            ip_dir = os.path.join(self.exp_save_path, '.ips', test_name)
-            os.makedirs(ip_dir, exist_ok=True)
-
             rank = 0
             self.experiment = self.experiment.get_meta_copy()
-            mp.spawn(self.dp_train, nprocs=len(self.data_parallel_device_ids), args=(rank, model))
+            mp.spawn(self.dp_train, nprocs=len(self.data_parallel_device_ids), args=(model))
         else:
             self.__run_pretrain_routine(model)
 
-    def dp_train(self, gpu_nb, proc_rank, model):
+    def dp_train(self, gpu_nb, model):
         """
         Entry point into a DP thread
         :param gpu_nb:
@@ -327,24 +316,22 @@ class Trainer(TrainerIO):
         :param cluster_obj:
         :return:
         """
-        # world size = all nodes x nb_gpus on each node
-        # 2 nodes, 8 gpus each = 16 world size
-        world_size = self.nb_gpu_nodes * len(self.data_parallel_device_ids)
+        # node rank using relative slurm id
+        node_rank = int(os.environ['SLURM_NODEID'])
 
         # recover original exp before went into process
         self.experiment = self.experiment.get_non_ddp_exp()
 
         # show progbar only on prog_rank 0
-        self.prog_bar = self.prog_bar and proc_rank == 0
-
-        # determine my node rank and the root ip
-        my_node_rank, root_ip = self.__get_root_node_ip(self.nb_gpu_nodes, self.exp_save_path, world_size)
+        self.prog_bar = self.prog_bar and node_rank == 0 and gpu_nb == 0
 
         # determine which process we are and world size
-        self.proc_rank = my_node_rank * len(self.data_parallel_device_ids) + gpu_nb
+        self.proc_rank = node_rank * len(self.data_parallel_device_ids) + gpu_nb
+        world_size = self.nb_gpu_nodes * len(self.data_parallel_device_ids)
 
         # set up server using proc 0's ip address
-        dist.init_process_group("nccl", init_method=f'tcp://{root_ip}:12001', rank=self.proc_rank, world_size=world_size)
+        ip = self.__get_root_node_ip(self.proc_rank, self.nb_gpu_nodes, self.exp_save_path)
+        dist.init_process_group("nccl", init_method=f'tcp://{ip}:12001', rank=self.proc_rank, world_size=world_size)
         print(f"GPU: {gpu_nb} - Rank: {self.proc_rank}")
 
         # copy model to each gpu
@@ -355,57 +342,41 @@ class Trainer(TrainerIO):
         # continue training routine
         self.__run_pretrain_routine(model)
 
-    def __get_root_node_ip(self, nb_gpu_nodes, ip_file_dir, world_size):
+    def __get_root_node_ip(self, proc_rank, nb_gpu_nodes, ip_file_dir):
         """
         Resolves the ip address of proc 0.
         Proc 0 writes address to a file. Every other process waits until the ip is available before it starts
 
+        :param proc_rank:
         :param nb_gpu_nodes:
         :param ip_file_dir:
         :return:
         """
         # on one node we use localhost
         if nb_gpu_nodes == 1:
-            return 0, '127.0.0.1'
+            return '127.0.0.1'
 
         # on multi-node, every node rank > 0 waits until rank 0
         # saves the ip to disk
         ip_file = os.path.join(ip_file_dir, '.ip_meta')
+        if proc_rank == 0:
+            # get the proc 0 IP
+            root_ip = subprocess.run(['hostname', '-I'], stdout=subprocess.PIPE).stdout.decode('utf-8')
+            root_ip = root_ip.split(' ')[0]
 
-        # every process writes their process id + ip to a shared file
-        my_ip = subprocess.run(['hostname', '-I'], stdout=subprocess.PIPE).stdout.decode('utf-8')
-        my_ip = my_ip.split(' ')[0]
+            # save the ip to the file
+            with open(file=ip_file, mode='w') as f:
+                f.write(root_ip)
 
-        # save the ip to the file
-        # block file so only one process can access at a time
-        ip_dir = os.path.join(ip_file_dir, '.ips', my_ip)
-        os.makedirs(ip_dir, exist_ok=True)
-
-        # now everyone waits until the file has world_size entries
-        for i in range(0, 120*10):
-            sleep(1.0)
-            nb_folders = [x for x in os.listdir(os.path.join(ip_file_dir, '.ips')) if '.' in x]
-            if nb_folders == nb_gpu_nodes:
-                break
-
-        # the ip_table is written at this point
-        # now every process reads it and decides what rank they are based on their node
-        ip_table = list(open(file=ip_file, mode='r'))
-        my_node_rank, root_ip = self.__determine_my_node_rank(ip_table, my_ip)
-
-        return my_node_rank, root_ip
-
-    def __determine_my_node_rank(self, ip_pid_lines, my_ip):
-        # the node rank is the index of my_ip in the world-size ip table
-        # when de-duped and sorted
-        unique_ips = list(set(ip_pid_lines))
-        unique_ips.sort()
-
-        my_node_rank = unique_ips.index(my_ip)
-
-        # use the first ip as the root ip everyone will connect to
-        root_ip = unique_ips[0]
-        return my_node_rank, root_ip
+            return root_ip
+        else:
+            # wait up to 120 seconds until proc 0 writes
+            # once written, read proc 0's address and use it to configure server
+            for i in range(0, 120):
+                sleep(1.0)
+                if os.path.exists(ip_file):
+                    ip = list(open(file=ip_file, mode='r'))[0]
+                    return ip
 
     def __run_pretrain_routine(self, model):
         """
@@ -676,4 +647,4 @@ class Trainer(TrainerIO):
         # model checkpointing
         print('save callback...')
         if self.proc_rank == 0:
-            self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic)
+            self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic)