# 定义IP和端口
strps_hosts = 'localhost:1681'
strworker_hosts = 'localhost:1682,localhost:1683'
# 定义角色名称
strjob_name = 'ps'
task_index = 0
# 将字符串转数组
ps_hosts = strps_hosts.split(',')
worker_hosts = strps_hosts.split(',')
cluster_spec = tf.train.ClusterSpec({'ps': ps_hosts, 'worker': worker_hosts})
# 创建server
server = tf.train.Server({'ps':ps_hosts, 'worker':worker_hosts}, job_name=strjob_name, task_index=task_index)

# ps角色使用join进行等待
if strjob_name == 'ps':
  print("wait")
  server.join()

with tf.device(tf.train.replica_device_setter(worker_device='/job:worker/task:%d'%task_index,
                       cluster=cluster_spec)):
  X = tf.placeholder('float')
  Y = tf.placeholder('float')
  # 模型参数
  w = tf.Variable(tf.random_normal([1]), name='weight')
  b = tf.Variable(tf.zeros([1]), name='bias')
  global_step = tf.train.get_or_create_global_step()  # 获取迭代次数
  z = tf.multiply(X, w) + b
  tf.summary('z', z)
  cost = tf.reduce_mean(tf.square(Y - z))
  tf.summary.scalar('loss_function', cost)
  learning_rate = 0.001
  optimizer = tf.train.GradientDescentOptimizer(learning_rate).minimize(cost, global_step=global_step)
  saver = tf.train.Saver(max_to_keep=1)
  merged_summary_op = tf.summary.merge_all() # 合并所有summary
  init = tf.global_variables_initializer()

training_epochs = 2000
display_step = 2
sv = tf.train.Supervisor(is_chief=(task_index == 0),# 0号为chief
             logdir='log/spuer/',
             init_op=init,
             summary_op=None,
             saver=saver,
             global_step=global_step,
             save_model_secs=5)
# 连接目标角色创建session
with sv.managed_session(saver.target) as sess:

print('sess ok')
  print(global_step.eval(session=sess))
  for epoch in range(global_step.eval(session=sess), training_epochs*len(train_x)):
    for (x, y) in zip(train_x, train_y):
      _, epoch = sess.run([optimizer, global_step], feed_dict={X: x, Y: y})
      summary_str = sess.run(merged_summary_op, feed_dict={X: x, Y: y})
      sv.summary_computed(sess, summary_str, global_step=epoch)
      if epoch % display_step == 0:
        loss = sess.run(cost, feed_dict={X:train_x, Y:train_y})
        print("Epoch:", epoch+1, 'loss:', loss, 'W=', sess.run(w), w, 'b=', sess.run(b))
  print(' finished ')
  sv.saver.save(sess, 'log/linear/' + "sv.cpk", global_step=epoch)
sv.stop()

strjob_name = 'worker'
task_index = 1
strjob_name = 'worker'
task_index = 0