From fd194ab843dd8c7b91fac527d8fde668e47b7a1b Mon Sep 17 00:00:00 2001 From: William Falcon Date: Mon, 8 Jul 2019 12:27:53 -0400 Subject: [PATCH] added multi-node locked ip search --- pytorch_lightning/models/trainer.py | 71 ++++++++++++++++++++--------- 1 file changed, 49 insertions(+), 22 deletions(-) diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index d7062257..b788abb3 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -12,6 +12,7 @@ import torch.distributed as dist import os import subprocess from time import sleep +import fcntl try: from apex import amp @@ -316,19 +317,24 @@ 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) + # 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) + # determine which process we are and world size - self.proc_rank = proc_rank * len(self.data_parallel_device_ids) + gpu_nb - world_size = self.nb_gpu_nodes * len(self.data_parallel_device_ids) + self.proc_rank = my_node_rank * len(self.data_parallel_device_ids) + gpu_nb # set up server using proc 0's ip address - 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) + dist.init_process_group("nccl", init_method=f'tcp://{root_ip}:12001', rank=self.proc_rank, world_size=world_size) print(f"GPU: {gpu_nb} - Rank: {self.proc_rank}") # copy model to each gpu @@ -339,12 +345,11 @@ class Trainer(TrainerIO): # continue training routine self.__run_pretrain_routine(model) - def __get_root_node_ip(self, proc_rank, nb_gpu_nodes, ip_file_dir): + def __get_root_node_ip(self, nb_gpu_nodes, ip_file_dir, world_size): """ 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: @@ -356,24 +361,46 @@ class Trainer(TrainerIO): # 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] - # save the ip to the file - with open(file=ip_file, mode='w') as f: - f.write(root_ip) + # 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] + to_write = f'{my_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 + # save the ip to the file + # block file so only one process can access at a time + with open(file=ip_file, mode='a') as f: + fcntl.flock(f, fcntl.LOCK_EX) + f.write(to_write) + fcntl.flock(f, fcntl.LOCK_UN) + + # now everyone waits until the file has world_size entries + for i in range(0, 120): + sleep(1.0) + if os.path.exists(ip_file): + lines = list(open(file=ip_file, mode='r')) + if len(lines) == world_size: + 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 def __run_pretrain_routine(self, model): """