Starts the TensorFlowOnSpark cluster and Runs the TensorFlow "main" function on the Spark executors

  Args:
    :sc: SparkContext
    :map_fun: user-supplied TensorFlow "main" function
    :tf_args: ``argparse`` args, or command-line ``ARGV``.  These will be passed to the ``map_fun``.
    :num_executors: number of Spark executors.  This should match your Spark job's ``--num_executors``.
    :num_ps: number of Spark executors which are reserved for TensorFlow PS nodes.  All other executors will be used as TensorFlow worker nodes.
    :tensorboard: boolean indicating if the chief worker should spawn a Tensorboard server.
    :input_mode: TFCluster.InputMode
    :log_dir: directory to save tensorboard event logs.  If None, defaults to a fixed path on local filesystem.
    :driver_ps_nodes: run the PS nodes on the driver locally instead of on the spark executors; this help maximizing computing resources (esp. GPU). You will need to set cluster_size = num_executors + num_ps
    :master_node: name of the "master" or "chief" node in the cluster_template, used for `tf.estimator` applications.
    :reservation_timeout: number of seconds after which cluster reservation times out (600 sec default)
    :queues: *INTERNAL_USE*
    :eval_node: run evaluator node for distributed Tensorflow

  Returns:
    A TFCluster object representing the started cluster.

def run(sc, map_fun, tf_args, num_executors, num_ps, tensorboard=False, input_mode=InputMode.TENSORFLOW,
        log_dir=None, driver_ps_nodes=False, master_node=None, reservation_timeout=600, queues=['input', 'output', 'error'],
        eval_node=False):
  """Starts the TensorFlowOnSpark cluster and Runs the TensorFlow "main" function on the Spark executors

  Args:
    :sc: SparkContext
    :map_fun: user-supplied TensorFlow "main" function
    :tf_args: ``argparse`` args, or command-line ``ARGV``.  These will be passed to the ``map_fun``.
    :num_executors: number of Spark executors.  This should match your Spark job's ``--num_executors``.
    :num_ps: number of Spark executors which are reserved for TensorFlow PS nodes.  All other executors will be used as TensorFlow worker nodes.
    :tensorboard: boolean indicating if the chief worker should spawn a Tensorboard server.
    :input_mode: TFCluster.InputMode
    :log_dir: directory to save tensorboard event logs.  If None, defaults to a fixed path on local filesystem.
    :driver_ps_nodes: run the PS nodes on the driver locally instead of on the spark executors; this help maximizing computing resources (esp. GPU). You will need to set cluster_size = num_executors + num_ps
    :master_node: name of the "master" or "chief" node in the cluster_template, used for `tf.estimator` applications.
    :reservation_timeout: number of seconds after which cluster reservation times out (600 sec default)
    :queues: *INTERNAL_USE*
    :eval_node: run evaluator node for distributed Tensorflow

  Returns:
    A TFCluster object representing the started cluster.
  """
  logging.info("Reserving TFSparkNodes {0}".format("w/ TensorBoard" if tensorboard else ""))

  if driver_ps_nodes and input_mode != InputMode.TENSORFLOW:
    raise Exception('running PS nodes on driver locally is only supported in InputMode.TENSORFLOW')

  if eval_node and input_mode != InputMode.TENSORFLOW:
    raise Exception('running evaluator nodes is only supported in InputMode.TENSORFLOW')

  # compute size of TF cluster and validate against number of Spark executors
  num_master = 1 if master_node else 0
  num_eval = 1 if eval_node else 0
  num_workers = max(num_executors - num_ps - num_eval - num_master, 0)
  total_nodes = num_ps + num_master + num_eval + num_workers

  assert total_nodes == num_executors, "TensorFlow cluster requires {} nodes, but only {} executors available".format(total_nodes, num_executors)
  assert num_master + num_workers > 0, "TensorFlow cluster requires at least one worker or master/chief node"

  # create a cluster template for scheduling TF nodes onto executors
  executors = list(range(num_executors))
  cluster_template = {}

  if num_ps > 0:
    cluster_template['ps'] = executors[:num_ps]
    del executors[:num_ps]
  if master_node:
    cluster_template[master_node] = executors[:1]
    del executors[:1]
  if eval_node:
    cluster_template['evaluator'] = executors[:1]
    del executors[:1]
  if num_workers > 0:
    cluster_template['worker'] = executors[:num_workers]

  logging.info("cluster_template: {}".format(cluster_template))

  # get default filesystem from spark
  defaultFS = sc._jsc.hadoopConfiguration().get("fs.defaultFS")
  # strip trailing "root" slash from "file:///" to be consistent w/ "hdfs://..."
  if defaultFS.startswith("file://") and len(defaultFS) > 7 and defaultFS.endswith("/"):
    defaultFS = defaultFS[:-1]

  # get current working dir of spark launch
  working_dir = os.getcwd()

  # start a server to listen for reservations and broadcast cluster_spec
  server = reservation.Server(num_executors)
  server_addr = server.start()

  # start TF nodes on all executors
  logging.info("Starting TensorFlow on executors")
  cluster_meta = {
    'id': random.getrandbits(64),
    'cluster_template': cluster_template,
    'num_executors': num_executors,
    'default_fs': defaultFS,
    'working_dir': working_dir,
    'server_addr': server_addr
  }
  if driver_ps_nodes:
    nodeRDD = sc.parallelize(range(num_ps, num_executors), num_executors - num_ps)
  else:
    nodeRDD = sc.parallelize(range(num_executors), num_executors)

  if driver_ps_nodes:
    def _start_ps(node_index):
      logging.info("starting ps node locally %d" % node_index)
      TFSparkNode.run(map_fun,
                      tf_args,
                      cluster_meta,
                      tensorboard,
                      log_dir,
                      queues,
                      background=(input_mode == InputMode.SPARK))([node_index])
    for i in cluster_template['ps']:
      ps_thread = threading.Thread(target=lambda: _start_ps(i))
      ps_thread.daemon = True
      ps_thread.start()

  # start TF on a background thread (on Spark driver) to allow for feeding job
  def _start(status):
    try:
      nodeRDD.foreachPartition(TFSparkNode.run(map_fun,
                                                tf_args,
                                                cluster_meta,
                                                tensorboard,
                                                log_dir,
                                                queues,
                                                background=(input_mode == InputMode.SPARK)))
    except Exception as e:
      logging.error("Exception in TF background thread")
      status['error'] = str(e)

  t = threading.Thread(target=_start, args=(tf_status,))
  # run as daemon thread so that in spark mode main thread can exit
  # if feeder spark stage fails and main thread can't do explicit shutdown
  t.daemon = True
  t.start()

  # wait for executors to register and start TFNodes before continuing
  logging.info("Waiting for TFSparkNodes to start")
  cluster_info = server.await_reservations(sc, tf_status, reservation_timeout)
  logging.info("All TFSparkNodes started")

  # print cluster_info and extract TensorBoard URL
  tb_url = None
  for node in cluster_info:
    logging.info(node)
    if node['tb_port'] != 0:
      tb_url = "http://{0}:{1}".format(node['host'], node['tb_port'])

  if tb_url is not None:
    logging.info("========================================================================================")
    logging.info("")
    logging.info("TensorBoard running at:       {0}".format(tb_url))
    logging.info("")
    logging.info("========================================================================================")

  # since our "primary key" for each executor's TFManager is (host, executor_id), sanity check for duplicates
  # Note: this may occur if Spark retries failed Python tasks on the same executor.
  tb_nodes = set()
  for node in cluster_info:
    node_id = (node['host'], node['executor_id'])
    if node_id in tb_nodes:
      msg = '''
Duplicate cluster node id detected (host={0}, executor_id={1})
Please ensure that:
1. Number of executors >= number of TensorFlow nodes
2. Number of tasks per executors is 1
3, TFCluster.shutdown() is successfully invoked when done.
'''.strip()
      raise Exception(msg.format(node_id[0], node_id[1]))
    else:
      tb_nodes.add(node_id)

  # create TFCluster object
  cluster = TFCluster()
  cluster.sc = sc
  cluster.meta = cluster_meta
  cluster.nodeRDD = nodeRDD
  cluster.cluster_info = cluster_info
  cluster.cluster_meta = cluster_meta
  cluster.input_mode = input_mode
  cluster.queues = queues
  cluster.server = server

  return cluster