def setup_distributed(self, mode, cluster): """Sets up TensorFLow distributed environment and initializes the model. Args: urls (str): the URLs that each node uses to connect. world_rank (int): the index of the runner. world_size (int): the total number of runners. """ self.cluster = cluster tc = BarrierTaskContext().get() self.rank = self._get_rank(cluster, tc) print("cluster is: ", cluster) import os os.environ["TF_CONFIG"] = json.dumps({ 'cluster': { 'worker': cluster }, 'task': { 'type': 'worker', 'index': self.rank } }) ips = set([node.split(":")[0] for node in cluster]) os.environ["no_proxy"] = ",".join(ips) self.strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy( ) # For use in model.evaluate() self.local_model = None self.backend = "tf-distributed"
def _get_rank(cluster_info): # As task placement may not be identical between two different jobs, # we cannot simply index cluster_info using partitionId to get current # ip and port. # The approach here is to first get all tasks' ip in this job and compute # a local rank by counting how many tasks has the same ip but with lower id. # We then use the local rank to find the right slot in cluster_info to find # the right global_rank. tc = BarrierTaskContext().get() infos = tc.getTaskInfos() idx = tc.partitionId() local_ip = infos[idx].address.split(":")[0] local_rank = 0 for i in range(0, idx): if infos[i].address.startswith(local_ip): local_rank += 1 global_rank = -1 local_count = 0 for node in cluster_info: if node.startswith(local_ip): local_count += 1 global_rank += 1 if local_count == local_rank + 1: break return global_rank
def find_ip_and_port(pre_iter): tc = BarrierTaskContext().get() address = tc.getTaskInfos()[tc.partitionId()].address.split(":")[0] with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s: s.bind(("", 0)) s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) tc.barrier() free_ip_port = f"{address}:{s.getsockname()[1]}" return [free_ip_port]
def find_ip_and_port(pre_iter): tc = BarrierTaskContext().get() free_port = find_free_port(tc) return [free_port]