import tensorflow as tfimport numpy as np
# 创建集群信息，包括ps和worker两种角色。# 集群有两类任务，ps和worker；ps由2个任务组成（一般一个任务是一个机器或者一个分配单元），worker由3个任务组成。ps_hosts = ["xx.xxx.xx.xxxx:oooo", "xx.xxx.xx.xxxx:oooo"]worker_hosts = ["xx.xxx.xx.xxxx:oooo", "xx.xxx.xx.xxxx:oooo", "xx.xxx.xx.xxxx:oooo"]cluster = tf.train.ClusterSpec({"ps": ps_hosts, "worker": worker_hosts})
tf.app.flags.DEFINE_string("job_name", "worker", "One of 'ps', 'worker'")tf.app.flags.DEFINE_integer("task_index", 0, "Index of task within the job")FLAGS = tf.app.flags.FLAGS
def main(_):    server = tf.train.Server(cluster, job_name=FLAGS.job_name, task_index=FLAGS.task_index)    if FLAGS.job_name == "ps":        server.join()    else:        # 会根据job名，将with内的Variable op放到ps tasks，将其他计算op放到worker tasks。默认分配策略是轮询        with tf.device(tf.train.replica_device_setter(                worker_device="/job:worker/task:%d" % FLAGS.task_index,                cluster=cluster)):
            x_data = tf.placeholder(tf.float32, [100])            y_data = tf.placeholder(tf.float32, [100])
            W = tf.Variable(tf.random_uniform([1], -1.0, 1.0))            b = tf.Variable(tf.zeros([1]))            y = W * x_data + b            loss = tf.reduce_mean(tf.square(y - y_data))
            global_step = tf.Variable(0, name="global_step", trainable=False)            optimizer = tf.train.GradientDescentOptimizer(0.1)            train_op = optimizer.minimize(loss, global_step=global_step)
            # The StopAtStepHook handles stopping after running given steps.            hooks = [tf.train.StopAtStepHook(last_step=1000000)]            # 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(master=server.target,                                                   is_chief=(FLAGS.task_index == 0),                                                   # 我们制定task_index为0的任务为主任务，用于负责变量初始化、做checkpoint、保存summary和复原                                                   checkpoint_dir="/tmp/tf_train_logs",                                                   save_checkpoint_secs=None,                                                   hooks=hooks) as mon_sess:                while not mon_sess.should_stop():                    # Run a training step asynchronously.                    # See `tf.train.SyncReplicasOptimizer` for additional details on how to                    # perform *synchronous* training.                    # mon_sess.run handles AbortedError in case of preempted PS.                    train_x = np.random.rand(100).astype(np.float32)                    train_y = train_x * 0.1 + 0.3                    _, step, loss_v, weight, biase = mon_sess.run([train_op, global_step, loss, W, b],                                                                  feed_dict={x_data: train_x, y_data: train_y})                    if step % 100 == 0:                        print("step: %d, weight: %f, biase: %f, loss: %f" % (step, weight, biase, loss_v))                print("Optimization finished.")

if __name__ == "__main__":    tf.app.run()

python dist_tf.py --job_name=ps --task_index=0python dist_tf.py --job_name=ps --task_index=1python dist_tf.py --job_name=worker --task_index=0python dist_tf.py --job_name=worker --task_index=1python dist_tf.py --job_name=worker --task_index=2

# https://stackoverflow.com/questions/55003279/parameter-server-strategy-with-estimatorstensorflowimport tensorflow as tfimport osimport json
NUM_WORKERS = 1IP_ADDRS = ['localhost']PORTS = [12345]
def model_fn(...):    .....
def input_fn(...):    .....

os.environ['TF_CONFIG'] = json.dumps({    'cluster': {        'worker': ['%s:%d' % (IP_ADDRS[w], PORTS[w]) for w in range(NUM_WORKERS)],        'ps': ['%s:%d' % (IP_ADDRS[w], PORTS[w]) for w in range(NUM_WORKERS)]    },    'task': {'type': 'worker', 'index': 0}})
# Method for using ParamterServerStrategystrategy = tf.distribute.experimental.ParameterServerStrategy()
config = tf.estimator.RunConfig(train_distribute=strategy)
classifier = tf.estimator.Estimator(    model_fn=model_fn, model_dir='/tmp/multiworker', config=config)tf.estimator.train_and_evaluate(    classifier,    train_spec=tf.estimator.TrainSpec(input_fn=input_fn),    eval_spec=tf.estimator.EvalSpec(input_fn=input_fn))