4. TensorFlow分布式作业如何设置ClusterSpec?

XLearning通过环境变量 TF_CLUSTER_DEF 、 TF_ROLE 、 TF_INDEX 对应的将clusterSpec、job_name、task_index等信息传送给各container(PS或Worker),用户只需在TensorFlow分布式模式程序中,从环境变量中获取对应变量,从而完成ClusterSpec及role、index分配。例如:

  1. import os
  2. import json
  3. cluster_def = json.loads(os.environ["TF_CLUSTER_DEF"])
  4. cluster = tf.train.ClusterSpec(cluster_def)
  5. job_name = os.environ["TF_ROLE"]
  6. task_index = int(os.environ["TF_INDEX"])