diff --git a/examples/new_project_templates/trainer_gpu_cluster_template.py b/examples/new_project_templates/trainer_gpu_cluster_template.py index df3c7d16..98cf87af 100644 --- a/examples/new_project_templates/trainer_gpu_cluster_template.py +++ b/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) diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index e7142880..0e741dd6 100644 --- a/pytorch_lightning/models/trainer.py +++ b/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) \ No newline at end of file