server = tf.train.Server(cluster, job_name="worker", task_index=0)

server = tf.train.Server(cluster, job_name="worker", task_index=1)

server = tf.train.Server(cluster, job_name="ps", task_index=0)

import tensorflow as tf

with tf.device("/job:ps/task:0"):
  w = tf.get_variable([[1., 2., 3.], [1., 3., 5.]])

input_data = ...
inputs = tf.split(input_data, num_workers)
outputs = []

for i in range(num_workers):
  with tf.device("/job:ps/task:%s" % str(i)):   
    outputs.append(tf.matmul(inputs[i], w))

output = tf.concat(outputs, axis=0)
with tf.Session() as sess:
  sess.run(tf.global_variables_initializer())
  print sess.run(output)

with tf.device("/job:ps/task:0"):
  w = tf.get_variable(name='w', shape=[784, 10])
  b = tf.get_variable(name='b', shape=[10])

x = tf.placeholder(tf.float32, shape=[None, 784])
y = tf.placeholder(tf.int32, shape=[None])
logits = tf.matmul(x, w) + b
loss = ...
train_op = ...

with tf.Session() as sess:
  for _ in range(10000):
    sess.run(train_op, feed_dict=...)

import tensorflow as tf

ps_hosts = FLAGS.ps_hosts.split(",")
worker_hosts = FLAGS.worker_hosts.split(",")
cluster = tf.train.ClusterSpec({"ps": ps_hosts, "worker": worker_hosts})
# 启动 ps 或worker 
server = tf.train.Server(cluster, job_name=FLAGS.job_name, task_index=FLAGS.task_index)

if FLAGS.job_name == 'ps':
  server.join()
elif FLAGS.job_name == "worker":  
  # 开始扮演 Client
  with tf.device(tf.train.replica_device_setter(worker_device="/job:worker/task:%d" % FLAGS.task_index,cluster=cluster)):
    # Build model...
    loss = ...
    train_op = ...

  with tf.train.MonitoredTrainingSession(master="/job:worker/task:0",is_chief=(FLAGS.task_index == 0),checkpoint_dir="/tmp/train_logs") as mon_sess:
    while not mon_sess.should_stop():
      mon_sess.run(train_op)

mirrored_strategy = tf.distribute.MirroredStrategy()
config = tf.estimator.RunConfig(train_distribute=mirrored_strategy, eval_distribute=mirrored_strategy)
regressor = tf.estimator.LinearRegressor(
    feature_columns=[tf.feature_column.numeric_column('feats')],
    optimizer='SGD',
    config=config)

my_strategy = xxxStrategy(cluster_resolver,variable_partitioner,...)
with my_strategy.scope():
  @tf.function
  def distribute_train_epoch(dataset):
    def replica_fn(input):
      # process input and return result
      return result
    total_result = 0
    for x in dataset:
      per_replica_result = my_strategy.run(replica_fn, args=(x,))
      total_result += my_strategy.reduce(tf.distribute.ReduceOp.SUM,per_replica_result, axis=None)
    return total_result

  dist_dataset = my_strategy.experimental_distribute_dataset(dataset)
  for _ in range(EPOCHS):
    train_result = distribute_train_epoch(dist_dataset)

class Estimator(object):
  def _train_model_distributed(self, input_fn, hooks, saving_listeners):
    ...
    # Configures the strategy class with `cluster_spec` 核心逻辑是 _initialize_multi_worker
    self._config._train_distribute.configure(self._config.session_config)
    return self._actual_train_model_distributed(self._config._train_distribute, input_fn, hooks, saving_listeners)
  def _actual_train_model_distributed(self, strategy, input_fn, hooks,saving_listeners):                               
    worker_hooks = []
    with tf.Graph().as_default() as g:
      with strategy.scope(), scale_ctx:     # 都在Strategy.scope 范围下
        # 调用input_fn来得到训练特征和标签
        iterator, input_hooks = self._get_iterator_from_input_fn(input_fn, ModeKeys.TRAIN, strategy)
        worker_hooks.extend(input_hooks)
        global_step_tensor = self._create_and_assert_global_step(g)
        tf.compat.v1.add_to_collection(training_util.GLOBAL_STEP_READ_KEY,strategy.extended.read_var(global_step_tensor))
        features, labels = estimator_util.parse_iterator_result(iterator.get_next())
        # 调用model_fn来得到训练图
        grouped_estimator_spec = strategy.extended.call_for_each_replica(self._call_model_fn,args=(features,labels,  ModeKeys.TRAIN,self.config))
        loss = strategy.reduce(_get_loss_reduce_op_for_reporting(),grouped_estimator_spec.loss,axis=None)
        distributed_train_op = grouped_estimator_spec.train_op
        # 各种hook
        training_hooks = get_hooks_from_the_first_device(grouped_estimator_spec.training_hooks)
        training_chief_hooks = get_hooks_from_the_first_device(grouped_estimator_spec.training_chief_hooks)
        # 进入training loop
        estimator_spec = model_fn_lib.EstimatorSpec(mode=grouped_estimator_spec.mode,loss=loss,train_op=strategy.group(distributed_train_op),training_hooks=training_hooks,training_chief_hooks=training_chief_hooks,scaffold=scaffold)
        return self._train_with_estimator_spec(estimator_spec, worker_hooks,hooks, global_step_tensor,saving_listeners)
  def _train_with_estimator_spec(self, estimator_spec, worker_hooks, hooks,global_step_tensor, saving_listeners):
    # 配置Hook
    worker_hooks.extend(hooks)
    worker_hooks.append(tf.compat.v1.train.NanTensorHook(estimator_spec.loss))
    worker_hooks.extend(estimator_spec.training_hooks)

    if (self._config.cluster_spec and type(self._train_distribution).__name__ in ('CollectiveAllReduceStrategy','CollectiveAllReduceStrategyV1','MultiWorkerMirroredStrategy')):
      return self._train_with_estimator_spec_distributed(estimator_spec, worker_hooks, saving_listeners)
    ...
    # 使用MonitoredTrainingSession进行Training loop
    with training.MonitoredTrainingSession(master=self._config.master,is_chief=self._config.is_chief,checkpoint_dir=self._model_dir,...)
      while not mon_sess.should_stop():
        _, loss = mon_sess.run([estimator_spec.train_op, estimator_spec.loss])
    return loss                          

# 指定如何对变量进行分区
variable_partitioner = ( tf.distribute.experimental.partitioners.MinSizePartitioner(min_shard_bytes=(256 << 10),max_shards=NUM_PS))
strategy = tf.distribute.experimental.ParameterServerStrategy(cluster_resolver,variable_partitioner=variable_partitioner)

  def __init__(self, cluster_resolver, variable_partitioner=None):
    self._cluster_resolver = cluster_resolver
    self._verify_args_and_config(cluster_resolver)
    self._cluster_coordinator = None
    self._connect_to_cluster(coordinator_name="chief")  # 连接到集群
    self._extended = ParameterServerStrategyV2Extended(self, cluster_resolver,variable_partitioner)                                           
    super(ParameterServerStrategyV2, self).__init__(self._extended)
    distribute_lib.distribution_strategy_gauge.get_cell("V2").set("ParameterServerStrategy")
    self._should_use_with_coordinator = True
    # Used while constructing distributed iterators.
    self._canonicalize_devices = False

class DistributedDatasetsFromFunctionV1(DistributedDatasetsFromFunction):
  def _make_initializable_iterator(self, shared_name=None):
    return self._get_iterator()
  def _make_one_shot_iterator(self):
    return self._get_iterator()
  def _get_iterator(self):
    iterators = _create_iterators_per_worker(self._datasets,self._input_workers, True,self._options)                                                                                  
    iterator = DistributedIteratorV1(self._input_workers, iterators,self._strategy,self._enable_get_next_as_optional)                                
    iterator._element_spec = self._element_spec 
    return iterator

class DistributedDatasetV1(DistributedDataset):
  def make_one_shot_iterator(self):
    return self._make_one_shot_iterator()
  def _make_one_shot_iterator(self):
    return self._get_iterator()
  def _get_iterator(self):
    worker_iterators = _create_iterators_per_worker(self._cloned_datasets,self._input_workers)                                          
    iterator = DistributedIteratorV1(self._input_workers, worker_iterators,self._strategy)                         
    iterator.element_spec = self.element_spec  
    return iterator

class StrategyBase(object):
  def run(self, fn, args=(), kwargs=None, options=None):      # Invokes `fn` on each replica, with the given arguments.
    ...
  def reduce(self, reduce_op, value, axis):   # Reduce `value` across replicas and return result on current device.
    ...
  def extended(self):     # Strategy 很多操作直接转给了 StrategyExtendV1/V2
    return self._extended
  def scope(self):
    return self._extended._scope(self) 
  def make_dataset_iterator(self, dataset):
  def make_input_fn_iterator(self,input_fn,replication_mode=InputReplicationMode.PER_WORKER):
  def experimental_distribute_dataset(self, dataset, options=None):
  def distribute_datasets_from_function(self, dataset_fn, options=None):
  def experimental_run(self, fn, input_iterator=None):
    with self.scope():
      args = (input_iterator.get_next(),) if input_iterator is not None else ()
    return self.run(fn, args=args)

# tensorflow/tensorflow/python/distribute/distribute_lib.py
class StrategyExtendedV2(object):
  def _scope(self, strategy):
    def creator_with_resource_vars(*args, **kwargs):
      _require_strategy_scope_extended(self)
      kwargs["use_resource"] = True
      kwargs["distribute_strategy"] = strategy
      return self._create_variable(*args, **kwargs)
    def distributed_getter(getter, *args, **kwargs):
      return getter(*args, **kwargs)
  return _CurrentDistributionContext(strategy,variable_scope.variable_creator_scope(creator_with_resource_vars),
        variable_scope.variable_scope(variable_scope.get_variable_scope(),custom_getter=distributed_getter), self._default_device)
# tensorflow/tensorflow/python/ops/variable_scope.py
def variable_creator_scope(variable_creator):
  with ops.get_default_graph()._variable_creator_scope(variable_creator): 
    yield
# tensorflow/tensorflow/python/framework/ops.py
class Graph(object):
  def _variable_creator_scope(self, creator, priority=100):
    old = self._variable_creator_stack
    new = list(old)
    new.append((priority, creator))
    new.sort(key=lambda item: item[0])
    self._thread_local._variable_creator_stack = new
    ...

# tensorflow/tensorflow/python/ops/variables.py
class Variable(six.with_metaclass(VariableMetaclass,trackable.Trackable)):
  # Variable 是 VariableMetaclass 子类
  ...
class VariableMetaclass(type):
  def _variable_v1_call(cls,initial_value=None,trainable=None,...):
    previous_getter = lambda **kwargs: default_variable_creator(None, **kwargs)
    # previous_getter 来自 ops.get_default_graph()
    for _, getter in ops.get_default_graph()._variable_creator_stack: 
      previous_getter = _make_getter(getter, previous_getter)
    return previous_getter(initial_value=initial_value,trainable=trainable,...)
  def __call__(cls, *args, **kwargs):
    if cls is VariableV1:
      return cls._variable_v1_call(*args, **kwargs)
    elif cls is Variable:
      return cls._variable_v2_call(*args, **kwargs)
    ...                                     

class _CurrentDistributionContext(object):
  def __enter__(self):
    ...
    _push_per_thread_mode(self._context)  # ==> distribution_strategy_context._push_per_thread_mode
  def __exit__(self, exception_type, exception_value, traceback):
    ...
    _pop_per_thread_mode()
# tensorflow/python/distribute/distribution_strategy_context.py
def _push_per_thread_mode(context):
  ops.get_default_graph()._distribution_strategy_stack.append(context) 
def _pop_per_thread_mode():
  ops.get_default_graph()._distribution_strategy_stack.pop(-1)
def get_strategy():
  return _get_per_thread_mode().strategy
def _get_per_thread_mode():
  try:
    return ops.get_default_graph()._distribution_strategy_stack[-1]  # pylint: disable=protected-access
  except (AttributeError, IndexError):
    return _get_default_replica_mode()

server = tf.train.Server(...)
is_cheif = FLAGS.task_index == 0
with tf.train.MonitoredTrainingSession(master=server.target,is_chief=is_chief) as sess:
  while not sess.should_stop():
    sess.run(train_op)