class RingAllReduceService(arrpc.RingAllReduceServiceServicer):
    def __init__(self, vars_init_fn, scatter_fn, gather_fn):
        # 参数初始化回调函数，由外部trainer传入
        self.vars_init_fn = vars_init_fn
        # scatter回调函数，由外部的trainer传入
        self.scatter_fn = scatter_fn
        # gather回调函数，由外部的trainer传入
        self.gather_fn = gather_fn


    def VariableWeightsInit(self, varibale_weights_req, context):
        '''
        变量节点初始化。接收上一个worker发送来的初始值并更新自身的变量节点值
        '''
        variable_weights_cache = DistCommon._deserialize_proto_variable_weights(
            varibale_weights_req)
        self.vars_init_fn(variable_weights_cache)
        return common_pb2.VariableWeightsReqResp()

    def Recieve(self, send_req, context):
        stage = send_req.stage
        # 从gRPC请求中解析出发送来的节点和梯度
        node_gradients_dict = DistCommon._deserialize_proto_node_gradients(send_req.node_gradients)
        # 接收到左邻居的请求，根据当前阶段的不同，执行不同的回调函数
        if stage == arpb.RingAllReduceReq.SCATTER:
            acc_no = send_req.node_gradients.acc_no
            self.scatter_fn(node_gradients_dict, acc_no)
        elif stage == arpb.RingAllReduceReq.GATHER:
            self.gather_fn(node_gradients_dict)
        else:
            print('[ALLREDUCE] Invalid ring all-reduce stage: {}, it should be either SCATTER or GATHER'.format(stage))
        return arpb.RingAllReduceResp()

class DistTrainerRingAllReduce(Trainer):
    '''
    Ring All-Reduce模式的分布式训练
    '''
    def __init__(self, *args, **kargs):
        Trainer.__init__(self, *args, **kargs)

        # 读取集群配置信息和自身信息
        self.cluster_conf = kargs['cluster_conf']
        self.worker_index = kargs['worker_index']

        self.workers = self.cluster_conf['workers']
        self.worker_num = len(self.workers)
        self.host = self.workers[self.worker_index]

        self.step = self.worker_num - 1

        # 根据集群的环状拓扑结构确定右邻居
        self.target_host = self.workers[(
            self.worker_index + 1) % self.worker_num]

        # 本节点是否已被初始化
        self.is_init = False
        self.init_cond = threading.Condition()

        self.cur_partion_index = self.worker_index
        self.partition = []

        # 获取所有可训练节点
        self.variables = get_trainable_variables_from_graph()

        # 根据worker的总数量，对即将更新的变量节点列表进行等长切分
        self._partition_variables()

        # 用于控制梯度的发送和接收
        self.is_recieved = False
        self.recieved_gradients = None
        self.recieved_acc_no = None
        self.cond = threading.Condition()

        # 创建本节点的梯度接收服务
        allreduce.RingAllReduceServer(
            self.host, self.worker_index,
            self._variable_weights_init_callback,
            self._scatter_callback,
            self._gather_callback).serve()

        # 创建连接目标节点的梯度发送client
        self.client = allreduce.RingAllReduceClient(self.target_host)


    def _variable_weights_init(self):

        var_weights_dict = dict()
        for node in default_graph.nodes:
            if isinstance(node, Variable) and node.trainable:
                var_weights_dict[node.name] = node.value
        print('[INIT] Send variable init weights to worker ', self.target_host)

        # 第一个节点不需要等待，使用默认值更新给下一个节点
        if self.worker_index == 0:
            self.client.variable_weights_init(var_weights_dict)
        else:
            self.init_cond.acquire()
            while not self.is_init:
                self.init_cond.wait()
            self.init_cond.release()
            self.client.variable_weights_init(var_weights_dict)


    def _variable_weights_init_callback(self, var_weights_dict):

        # 第一个节点不需要接收上一个节点的初始值
        if self.worker_index != 0:
            print('[INIT] Variables initializing weights from last worker node...')
            for var_name, weights in var_weights_dict.items():
                update_node_value_in_graph(var_name, weights)
        # 已初始化完成，通知发送流程
        self.init_cond.acquire()
        self.is_init = True
        self.init_cond.notify_all()
        self.init_cond.release()


    def _optimizer_update(self):
        # 共执行 N-1 次scatter操作，把本worker的梯度切片发送给下一个worker
        # 同时接收左邻居发送过来的梯度，累加到自己的对应切片上
        for scatter_index in range(self.step):
            gradients_part = self._get_gradients_partition()
            cur_acc_no = self.optimizer.acc_no if scatter_index == 0 else self.recieved_acc_no
            # 把自身的一个数据分块发送给右邻居
            self.client.send(gradients_part, cur_acc_no, 'scatter')
            # 等待接收并处理完左邻居节点的数据
            self._wait_for_recieve('scatter')

        # 然后执行 N-1 次all-gather操作，把本worker的梯度切片发送给下一个worker
        # 同时接收上一个worker发送过来的梯度并替换自己的对应切片
        for gather_index in range(self.step):
            gradients_part = self._get_gradients_partition()
            self.client.send(gradients_part, 0, 'gather')
            self._wait_for_recieve('gather')

        self.optimizer.update()


    def _partition_variables(self):
        '''
        根据worker的总数量，对即将更新的权值变量列表进行等长切分
        '''
        var_num = len(self.variables)
        part_length = math.ceil(var_num / self.worker_num)
        assert part_length > 0

        start = 0
        end = start + part_length
        for i in range(self.worker_num - 1):
            self.partition.append((start, end))
            start = end
            end = start + part_length

        self.partition.append((start, var_num))


    def _get_gradients_partition(self):
        '''
        获取下一个梯度切片
        '''
        start, end = self.partition[self.cur_partion_index]
        part_variables = self.variables[start:end]
        self.cur_partion_index = (
            self.cur_partion_index + self.step) % self.worker_num
        part_gradients = dict()
        for var in part_variables:
            part_gradients[var] = self.optimizer.acc_gradient[var]
        return part_gradients


    def _scatter_callback(self, node_gradients_dict, acc_no):
        '''
        Scatter 阶段的回调函数，接收上一个worker发送过来的梯度和样本数
        '''
        if self.cond.acquire():
            while self.is_recieved:
                self.cond.wait()

            # 把接收到的梯度缓存下来
            self.recieved_gradients = node_gradients_dict
            self.recieved_acc_no = acc_no
            self.is_recieved = True

            # 通知主流程，把接收到的梯度更新到优化器
            self.cond.notify_all()
            self.cond.release()
        else:
            self.cond.wait()

    def _gather_callback(self, node_gradients_dict):
        '''
        All-gather 阶段的回调函数，接收上一个worker发送来的梯度
        '''
        if self.cond.acquire():
            while self.is_recieved:
                self.cond.wait()

            self.recieved_gradients = node_gradients_dict
            self.is_recieved = True

            # 通知主流程，把接收到的梯度更新到优化器
            self.cond.notify_all()
            self.cond.release()
        else:
            self.cond.wait()


    def _wait_for_recieve(self, stage):
        '''
        等待梯度，并把接收到的梯度更新到优化器中
        '''
        if self.cond.acquire():
            while not self.is_recieved:
                self.cond.wait()

            # 如果是scatter阶段则累加梯度，同时累加样本数
            if stage == 'scatter':
                self.optimizer.apply_gradients(
                    self.recieved_gradients,  summarize=True, acc_no=self.recieved_acc_no)

            # 如果是all-gather阶段则覆盖梯度，样本数保持不变
            else:
                self.optimizer.apply_gradients(
                    self.recieved_gradients, summarize=False, acc_no=self.optimizer.acc_no)

            self.is_recieved = False

            # 梯度已被更新，通知接收流程继续接收新的梯度
            self.cond.notify_all()
            self.cond.release()
        else:
            self.cond.wait()

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