import collections
import datetime
import json
import multiprocessing
import os
import subprocess
import sys
import time

_SSHD_BINARY_PATH = "/usr/sbin/sshd"

EnvironmentConfig = collections.namedtuple(
    "EnvironmentConfig",
    ["hosts", "port", "is_chief", "pools", "job_id"])


class DeadlineExceededError(Exception):
  """Indicates an action took too long."""
  pass


def _sub_process_num_gpus(unused):
  del unused
  # This is imported here so that we don't load tensorflow in the parent
  # process. Once the sub-process exits, it releases its allocated GPU memory.
  from tensorflow.python.client import device_lib
  local_device_protos = device_lib.list_local_devices()
  gpus = [x.name for x in local_device_protos if x.device_type == "GPU"]
  return len(gpus)


def _get_available_gpus():
  """Returns the number of GPUs on the machine."""
  pool = multiprocessing.Pool(1)
  result = pool.map(_sub_process_num_gpus, [None])[0]
  pool.close()
  pool.join()
  return result


def parse_environment_config(env_config_str, job_id):
  """Parses environment config and returns a list of hosts as well as the role.

  Returns:
    An EnvironmentConfig.
  """
  if env_config_str:
    ssh_port = -1
    env_config_json = json.loads(env_config_str)
    cluster = env_config_json.get("cluster")
    if not cluster:
      return None, True
    hosts = []
    pools = collections.defaultdict(list)
    for pool_type, tasks_per_type in cluster.items():
      if pool_type == "master":
        pool_type = "chief"
      for host_and_port in tasks_per_type:
        host, port = host_and_port.split(":")
        if host == "127.0.0.1":
          host = "localhost"
        port = int(port)
        if ssh_port == -1:
          ssh_port = port
        elif ssh_port != port:
          raise ValueError("Inconsistent ssh ports across tasks %d != %d." %
                           (ssh_port, port))
        hosts.append(host)
        pools[pool_type].append(host)
    is_chief = False
    has_chief = "chief" in pools
    if (env_config_json["task"]["type"] == "master" or
        env_config_json["task"]["type"] == "chief"):
      is_chief = True
      if int(env_config_json["task"]["index"]) != 0:
        raise ValueError("Only one master node is expected.")
    elif ((not has_chief) and
          (env_config_json["task"]["type"] == "worker") and
          int(env_config_json["task"]["index"]) == 0):
      is_chief = True
      pools["chief"].append(pools["worker"].pop(0))
    elif env_config_json["task"]["type"] != "worker":
      raise ValueError("Unexpected task type for Horovod training: %s." %
                       env_config_json["task"]["type"])
    return EnvironmentConfig(hosts=hosts, port=port, is_chief=is_chief,
                             pools=pools, job_id=job_id)
  else:
    return EnvironmentConfig(hosts=["localhost"], port=2222, is_chief=True,
                             pools={"chief": ["localhost"]}, job_id=job_id)


def start_ssh_server(port, is_chief):
  ssh_server_command = [_SSHD_BINARY_PATH, "-p", str(port)]
  if not is_chief:
    ssh_server_command.append("-D")
  completed = subprocess.call(ssh_server_command)
  if completed != 0:
    raise OSError("SSH server did not start successfully.")


def wait_for_ssh_servers(hosts, port, timeout_seconds):
  deadline_datetime = datetime.datetime.utcnow() + datetime.timedelta(
      seconds=timeout_seconds)
  unavailable_hosts = []
  while datetime.datetime.utcnow() < deadline_datetime:
    unavailable_hosts = []
    for host in hosts:
      ssh_command = ["ssh", "-q", host, "-p", str(port), "true"]
      result = subprocess.call(ssh_command)
      if result != 0:
        unavailable_hosts.append(host)
    if not unavailable_hosts:
      return
    # Retry in 1 second.
    time.sleep(1)

  raise DeadlineExceededError(
      "Timed out while waiting for all hosts to start. "
      "Hosts still not available: %s. TASK_STARTUP_TIMEOUT_SECONDS=%d" %
      (unavailable_hosts, timeout_seconds))


def run_horovod(env_config, jobs_per_host, args):
  env = dict(os.environ)
  del env["TF_CONFIG"]

  num_jobs = len(env_config.hosts) * jobs_per_host
  hosts = ",".join("%s:%d" % (h, jobs_per_host) for h in env_config.hosts)
  horovod_command = [
      "horovodrun", "--ssh-port", str(env_config.port), "-H",
      hosts, "--num-proc", str(num_jobs)
  ]
  horovod_command.extend(args)
  exit_code = subprocess.call(horovod_command, env=env)
  return exit_code


def benchmark_network(env_config):
  if not env_config.pools["worker"]:
    raise ValueError("No workers in the pool to do network benchmarking.")
  iperf_server = ["iperf", "-s", "-p", "6000"]
  server = subprocess.Popen(iperf_server)
  # Wait 10 seconds for the local server to start.
  time.sleep(10)
  iperf_command = ["ssh", "-q", env_config.pools["worker"][0], "-p",
                   str(env_config.port),
                   "iperf", "-p", "6000", "-c", env_config.pools["chief"][0]]
  subprocess.call(iperf_command)
  server.kill()


def copy_files_recursively(src, dest):
  if not dest.startswith("gs://"):
    try:
      os.makedirs(dest)
    except OSError:
      pass
  copy_cmd = ["gsutil", "-m", "rsync", "-r", src, dest]
  exit_code = subprocess.call(copy_cmd)
  if exit_code != 0:
    raise RuntimeError("Error while copying %s to %s" % (src, dest))
  return exit_code


def main():
  env_config_str = os.environ.get("TF_CONFIG")
  job_id = os.environ.get("CLOUD_ML_JOB_ID", "localrun")

  env_config = parse_environment_config(env_config_str, job_id)
  print (env_config, env_config.pools, env_config.hosts, os.environ)
  if os.environ.get("STAGE_GCS_PATH", False):
    copy_files_recursively(
        os.environ.get("STAGE_GCS_PATH"),
        os.environ.get("STAGING_DIR", "/input"))

  start_ssh_server(env_config.port, env_config.is_chief)
  max_num_retries = os.environ.get("NUM_HOROVOD_RETRIES", 1)
  if env_config.is_chief:
    exit_code = 0
    for retry in range(max_num_retries):
      staging_timeout_seconds = int(
          os.environ.get("TASK_STARTUP_TIMEOUT_SECONDS", 600))
      wait_for_ssh_servers(env_config.hosts, env_config.port,
                           staging_timeout_seconds)
      if os.environ.get("BENCHMARK_NETWORK", False):
        benchmark_network(env_config)
      num_gpus = _get_available_gpus()
      # If there are no GPUs, we can just run single process per machine.
      jobs_per_host = max(1, num_gpus)
      args = sys.argv[1:]
      exit_code = run_horovod(env_config=env_config, jobs_per_host=jobs_per_host,
                              args=args)
      if exit_code == 0:
        break
      else:
        print ("Retrying...", retry, "out of", max_num_retries)
    if os.environ.get("GCS_OUTPUT_PATH", False):
      copy_files_recursively(
          os.environ.get("OUTPUT_DIR", "/output"),
          os.path.join(os.environ.get("GCS_OUTPUT_PATH"), job_id))
    sys.exit(exit_code)


if __name__ == "__main__":
  main()

[((12, 20, 14, 53), 'collections.namedtuple', 'collections.namedtuple', ({(13, 4, 13, 23): '"""EnvironmentConfig"""', (14, 4, 14, 52): "['hosts', 'port', 'is_chief', 'pools', 'job_id']"}, {}), "('EnvironmentConfig', ['hosts', 'port', 'is_chief',\n    'pools', 'job_id'])", False, 'import collections\n'), ((27, 24, 27, 55), 'tensorflow.python.client.device_lib.list_local_devices', 'device_lib.list_local_devices', ({}, {}), '()', False, 'from tensorflow.python.client import device_lib\n'), ((34, 9, 34, 32), 'multiprocessing.Pool', 'multiprocessing.Pool', ({(34, 30, 34, 31): '1'}, {}), '(1)', False, 'import multiprocessing\n'), ((96, 14, 96, 49), 'subprocess.call', 'subprocess.call', ({(96, 30, 96, 48): 'ssh_server_command'}, {}), '(ssh_server_command)', False, 'import subprocess\n'), ((134, 14, 134, 55), 'subprocess.call', 'subprocess.call', (), '', False, 'import subprocess\n'), ((142, 11, 142, 41), 'subprocess.Popen', 'subprocess.Popen', ({(142, 28, 142, 40): 'iperf_server'}, {}), '(iperf_server)', False, 'import subprocess\n'), ((144, 2, 144, 16), 'time.sleep', 'time.sleep', ({(144, 13, 144, 15): '(10)'}, {}), '(10)', False, 'import time\n'), ((148, 2, 148, 32), 'subprocess.call', 'subprocess.call', ({(148, 18, 148, 31): 'iperf_command'}, {}), '(iperf_command)', False, 'import subprocess\n'), ((159, 14, 159, 39), 'subprocess.call', 'subprocess.call', ({(159, 30, 159, 38): 'copy_cmd'}, {}), '(copy_cmd)', False, 'import subprocess\n'), ((166, 19, 166, 46), 'os.environ.get', 'os.environ.get', ({(166, 34, 166, 45): '"""TF_CONFIG"""'}, {}), "('TF_CONFIG')", False, 'import os\n'), ((167, 11, 167, 56), 'os.environ.get', 'os.environ.get', ({(167, 26, 167, 43): '"""CLOUD_ML_JOB_ID"""', (167, 45, 167, 55): '"""localrun"""'}, {}), "('CLOUD_ML_JOB_ID', 'localrun')", False, 'import os\n'), ((171, 5, 171, 44), 'os.environ.get', 'os.environ.get', ({(171, 20, 171, 36): '"""STAGE_GCS_PATH"""', (171, 38, 171, 43): '(False)'}, {}), "('STAGE_GCS_PATH', False)", False, 'import os\n'), ((177, 20, 177, 60), 'os.environ.get', 'os.environ.get', ({(177, 35, 177, 56): '"""NUM_HOROVOD_RETRIES"""', (177, 58, 177, 59): '1'}, {}), "('NUM_HOROVOD_RETRIES', 1)", False, 'import os\n'), ((49, 22, 49, 48), 'json.loads', 'json.loads', ({(49, 33, 49, 47): 'env_config_str'}, {}), '(env_config_str)', False, 'import json\n'), ((54, 12, 54, 41), 'collections.defaultdict', 'collections.defaultdict', ({(54, 36, 54, 40): 'list'}, {}), '(list)', False, 'import collections\n'), ((102, 22, 102, 48), 'datetime.datetime.utcnow', 'datetime.datetime.utcnow', ({}, {}), '()', False, 'import datetime\n'), ((102, 51, 103, 30), 'datetime.timedelta', 'datetime.timedelta', (), '', False, 'import datetime\n'), ((105, 8, 105, 34), 'datetime.datetime.utcnow', 'datetime.datetime.utcnow', ({}, {}), '()', False, 'import datetime\n'), ((115, 4, 115, 17), 'time.sleep', 'time.sleep', ({(115, 15, 115, 16): '(1)'}, {}), '(1)', False, 'import time\n'), ((197, 7, 197, 47), 'os.environ.get', 'os.environ.get', ({(197, 22, 197, 39): '"""GCS_OUTPUT_PATH"""', (197, 41, 197, 46): '(False)'}, {}), "('GCS_OUTPUT_PATH', False)", False, 'import os\n'), ((201, 4, 201, 23), 'sys.exit', 'sys.exit', ({(201, 13, 201, 22): 'exit_code'}, {}), '(exit_code)', False, 'import sys\n'), ((109, 15, 109, 43), 'subprocess.call', 'subprocess.call', ({(109, 31, 109, 42): 'ssh_command'}, {}), '(ssh_command)', False, 'import subprocess\n'), ((155, 6, 155, 23), 'os.makedirs', 'os.makedirs', ({(155, 18, 155, 22): 'dest'}, {}), '(dest)', False, 'import os\n'), ((173, 8, 173, 40), 'os.environ.get', 'os.environ.get', ({(173, 23, 173, 39): '"""STAGE_GCS_PATH"""'}, {}), "('STAGE_GCS_PATH')", False, 'import os\n'), ((174, 8, 174, 47), 'os.environ.get', 'os.environ.get', ({(174, 23, 174, 36): '"""STAGING_DIR"""', (174, 38, 174, 46): '"""/input"""'}, {}), "('STAGING_DIR', '/input')", False, 'import os\n'), ((185, 9, 185, 51), 'os.environ.get', 'os.environ.get', ({(185, 24, 185, 43): '"""BENCHMARK_NETWORK"""', (185, 45, 185, 50): '(False)'}, {}), "('BENCHMARK_NETWORK', False)", False, 'import os\n'), ((182, 10, 182, 61), 'os.environ.get', 'os.environ.get', ({(182, 25, 182, 55): '"""TASK_STARTUP_TIMEOUT_SECONDS"""', (182, 57, 182, 60): '600'}, {}), "('TASK_STARTUP_TIMEOUT_SECONDS', 600)", False, 'import os\n'), ((199, 10, 199, 49), 'os.environ.get', 'os.environ.get', ({(199, 25, 199, 37): '"""OUTPUT_DIR"""', (199, 39, 199, 48): '"""/output"""'}, {}), "('OUTPUT_DIR', '/output')", False, 'import os\n'), ((200, 23, 200, 56), 'os.environ.get', 'os.environ.get', ({(200, 38, 200, 55): '"""GCS_OUTPUT_PATH"""'}, {}), "('GCS_OUTPUT_PATH')", False, 'import os\n')]