class_docstr: Manages a TensorFlow model for training.
class_docstr_tok: ["Manages","a","TensorFlow","model","for","training","."]
docstr: 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.
docstr_tok: ["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","."]
code: def setup_distributed(self, urls, world_rank, world_size):
      """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.
      """
      assert len(urls) == world_size
      tf_config = {
          "cluster": {"worker": urls},
          "task": {"index": world_rank, "type": "worker"},
      }
      os.environ["TF_CONFIG"] = json.dumps(tf_config)
      no_proxy = os.environ.get("no_proxy", "")
      ips = [url.split(":")[0] for url in urls]
      os.environ["no_proxy"] = ",".join(ips) + "," + no_proxy
  
      MultiWorkerMirroredStrategy = _try_import_strategy()
  
      # MultiWorkerMirroredStrategy handles everything for us, from
      # sharding the dataset (or even sharding the data itself if the loader
      # reads files from disk) to merging the metrics and weight updates
      #
      # worker 0 is the "chief" worker and will handle the map-reduce
      # every worker ends up with the exact same metrics and model
      # after model.fit
      #
      # because of this, we only really ever need to query its state
      self.strategy = MultiWorkerMirroredStrategy()
  
      logger.debug("Creating model with MultiWorkerMirroredStrategy")
      with self.strategy.scope():
          self.model = self.model_creator(self.config)
  
      # For use in model.evaluate()
      self.local_model = None
      self.backend = "tf-distributed"
      self.size = world_size
      self.rank = world_rank
  
imports: import logging
  import json
  import os
  import numpy
  import ray
  from contextlib import closing
  import logging
  import socket
  from zoo.orca.data.utils import ray_partitions_get_data_label
label_desc: Use image node_id 5 for calling the TFRunner obj's underlying member method code with example usage: obj.setup_distributed(urls, world_rank, world_size) without return types