import tensorflow as tf
import horovod.tensorflow as hvd

# Initialize Horovod
hvd.init()
# Pin GPU to be used to process local rank (one GPU per process)
config = tf.ConfigProto()
config.gpu_options.visible_device_list = str(hvd.local_rank())
# Build model...
loss = ...
opt = tf.train.AdagradOptimizer(0.01 * hvd.size())
# Add Horovod Distributed Optimizer
opt = hvd.DistributedOptimizer(opt)
# Add hook to broadcast variables from rank 0 to all other processes during initialization.
hooks = [hvd.BroadcastGlobalVariablesHook(0)]
# Make training operation
train_op = opt.minimize(loss)
# Save checkpoints only on worker 0 to prevent other workers from corrupting them.
checkpoint_dir = '/tmp/train_logs' if hvd.rank() == 0 else None
# The MonitoredTrainingSession takes care of session initialization, restoring from a checkpoint, saving to a checkpoint, and closing when done or an error occurs.
with tf.train.MonitoredTrainingSession(checkpoint_dir=checkpoint_dir,config=config,hooks=hooks) as mon_sess:
  while not mon_sess.should_stop():
    # Perform synchronous training.
    mon_sess.run(train_op)

# horovod/horovod/runner/launch.py
def _run(args):
  if _is_elastic(args):
    return _run_elastic(args)
  else:
    return _run_static(args)
def _run_static(args):
  ...
  all_host_names, _ = hosts.parse_hosts_and_slots(args.hosts)
  ...
  command = args.command
  _launch_job(args, settings, nics, command)
def _launch_job(args, settings, nics, command):
    env = os.environ.copy()
    config_parser.set_env_from_args(env, args)
    def gloo_run_fn():
      driver_ip = network.get_driver_ip(nics)
      gloo_run(settings, nics, env, driver_ip, command)
    def mpi_run_fn():
      mpi_run(settings, nics, env, command)
    run_controller(args.use_gloo, gloo_run_fn,args.use_mpi, mpi_run_fn,args.use_jsrun, js_run_fn,args.verbose)
if __name__ == '__main__':
    run_commandline()   # ==> _run 

# horovod/horovod/runner/gloo_run.py
def gloo_run(settings, nics, env, server_ip, command):
    # Each thread will use ssh command to launch the job on each remote host. If an error occurs in one thread, entire process will be terminated. Otherwise, threads will keep running and ssh session.
    exec_command = _exec_command_fn(settings)
    launch_gloo(command, exec_command, settings, nics, env, server_ip)
def launch_gloo(command, exec_command, settings, nics, env, server_ip):
    # Make the output directory if it does not exist
    if settings.output_filename:
        _mkdir_p(settings.output_filename)
    # start global rendezvous server and get port that it is listening on
    # 建立 RendezvousServer，这个会被底层 Gloo C++ 环境使用到
    rendezvous = RendezvousServer(settings.verbose)
    # allocate processes into slots
    # 来根据host进行分配slot，就是horovod的哪个rank应该在哪个host上的哪个slot之上运行
    hosts = parse_hosts(settings.hosts)
    host_alloc_plan = get_host_assignments(hosts, settings.num_proc)
    # start global rendezvous server and get port that it is listening on
    global_rendezv_port = rendezvous.start()
    rendezvous.init(host_alloc_plan)
    # 获取到可执行命令
    run_command = get_run_command(command, server_ip, nics, global_rendezv_port)
    # 得到在slot之上可执行的 slot command
    slot_info_to_command = _slot_info_to_command_fn(run_command, env)
    event = register_shutdown_event()
    # 依据 slot_info_to_command_fn 构建 args_list，这个 list 之中，每一个arg就是一个 slot command
    args_list = [[slot_info_to_command(slot_info), slot_info, [event]]
                 for slot_info in host_alloc_plan]
    # If an error occurs in one thread, entire process will be terminated.
    # Otherwise, threads will keep running.
    # 多线程执行，在每一个 exec_command 之上执行每一个 arg（slot command）
    res = threads.execute_function_multithreaded(exec_command,args_list,block_until_all_done=True)
    for name, value in sorted(res.items(), key=lambda item: item[1][1]):
        exit_code, timestamp = value

# horovod/horovod/tensorflow/__init__.py
def DistributedOptimizer(optimizer, name=None, use_locking=False, device_dense='',...):
  ...
  return hvd_k.DistributedOptimizer(optimizer=optimizer,name=name,device_dense=device_dense,device_sparse=device_sparse,...)
# horovod/horovod/tensorflow/keras/__init__.py
def DistributedOptimizer(optimizer, name=None,device_dense='', device_sparse='',...):
  ...
  return _impl.create_distributed_optimizer(keras=keras,optimizer=optimizer,name=name,device_dense=device_dense,device_sparse=device_sparse,...)
# horovod/horovod/_keras/__init__.py
def create_distributed_optimizer(keras, optimizer, name, device_dense, device_sparse,...):
  class _DistributedOptimizer(keras.optimizers.Optimizer):
    def __init__(self, **kwargs):
      super(self.__class__, self).__init__(**kwargs)
      ...
      self._allreduce_grads = hvd._make_allreduce_grads_fn(self._name,device_dense,device_sparse,...)
    def _compute_gradients(self, loss, var_list, grad_loss=None, tape=None):
      tape = tf.GradientTape() if tape is None else tape
      # 计算梯度
      grads_and_vars = super(self.__class__, self)._compute_gradients(loss,var_list,grad_loss,tape=tape)
      grads, weights = list(zip(*grads_and_vars))
      # 利用 AllReduce 来得到各个进程平均后的梯度
      allreduced_grads = self._allreduce(grads, weights)
      return list(zip(allreduced_grads, weights))
    def _allreduce(self, grads, vars):
      ...
      return self._allreduce_grads(grads, vars)
    def apply_gradients(self, *args, **kwargs):
      if self._agg_helper:
        ...
      else:
        results = super(self.__class__, self).apply_gradients(*args, **kwargs)
      return results
# horovod/horovod/tensorflow/__init__.py
def _make_allreduce_grads_fn(name, device_dense, device_sparse,compression, sparse_as_dense,...):
    groups = vars_to_refs(groups) if isinstance(groups, list) else groups
    # 弯弯绕绕最后执行_allreduce
    return _make_cached_allreduce_grads_fn(name, device_dense, device_sparse,compression, sparse_as_dense,...)
# horovod/horovod/tensorflow/mpi_ops.py
def _load_library(name):
    filename = resource_loader.get_path_to_datafile(name)
    library = load_library.load_op_library(filename)
    return library
MPI_LIB = _load_library('mpi_lib' + get_ext_suffix())
def _allreduce(tensor, name=None, op=Sum, prescale_factor=1.0, postscale_factor=1.0,...):  
    # 调用的就是 HorovodAllreduceOp      
    return MPI_LIB.horovod_allreduce(tensor, name=name, reduce_op=op,...)

import tensorflow as tf
import horovod.tensorflow as hvd

hvd.init()
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
  tf.config.experimental.set_memory_growth(gpu, True)
if gpus:
  tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')

dataset = ...
model = ...

optimizer = tf.optimizers.Adam(lr * hvd.size())

@tf.function
def train_one_batch(data, target, allreduce=True):
    with tf.GradientTape() as tape:
        probs = model(data, training=True)
        loss = tf.losses.categorical_crossentropy(target, probs)
    if allreduce:
        tape = hvd.DistributedGradientTape(tape)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

# Initialize model and optimizer state so we can synchronize across workers
data, target = get_random_batch()
train_one_batch(data, target, allreduce=False)

# 使用 @hvd.elastic.run 对 train 做了一个封装
@hvd.elastic.run
def train(state):
    for state.epoch in range(state.epoch, epochs):
        for state.batch in range(state.batch, batches_per_epoch):
            data, target = get_random_batch()
            train_one_batch(data, target)
            if state.batch % batches_per_commit == 0:
                state.commit()
        state.batch = 0

def on_state_reset():
    optimizer.lr.assign(lr * hvd.size())

# 这里是新修改处，传入了一个 TensorFlowKerasState   
state = hvd.elastic.TensorFlowKerasState(model, optimizer, batch=0, epoch=0)
state.register_reset_callbacks([on_state_reset])
train(state)

# horovod/tensorflow/elastic.py 
def run(func):
  from tensorflow.python.framework.errors_impl import UnknownError
  def wrapper(state, *args, **kwargs):
    try:
      return func(state, *args, **kwargs)
    except UnknownError as e:
      if 'HorovodAllreduce' in e.message or 'HorovodAllgather' in e.message or 'HorovodBroadcast' in e.message:
        raise HorovodInternalError(e)
  return run_fn(wrapper, _reset)
# horovod/common/elastic.py
def run_fn(func, reset):
  @functools.wraps(func)
  def wrapper(state, *args, **kwargs):
    notification_manager.init()
    notification_manager.register_listener(state)
    skip_sync = False
    try:
      while True:
        if not skip_sync:
          state.sync()                      # 当重置时候，用户也会进行必要的同步，具体是广播变量 和 存模型 两步
        try:
          return func(state, *args, **kwargs)
        except HorovodInternalError:        # 训练出错
          state.restore()                   # 进行恢复，就是重新加载模型，具体加载就是利用 TensorFlowKerasState 的 model, optimizer 这两个成员变量。
          skip_sync = False
        except HostsUpdatedInterrupt as e:  # driver发现一个节点被标记为新增或者移除时，将发送一个通知到 所有worker，worker 抛出 HostsUpdatedInterrupt
          skip_sync = e.skip_sync
        reset()
        state.on_reset()                    # 执行用户设置的 reset callback
    finally:
        notification_manager.remove_listener(state)
  return wrapper
def _reset():
    shutdown()
    init()          # 重新建立 MPI 相关 context