Kubeflow实战系列: 利用TFJob运行分布式TensorFlow

最新推荐文章于 2025-07-01 07:30:00 发布

煊琰

最新推荐文章于 2025-07-01 07:30:00 发布

阅读量4k

点赞数

本文介绍如何在Kubernetes上使用TFJob进行TensorFlow分布式训练。通过Kubernetes和TFJob，用户可以轻松地调度计算资源、配置软件及构建分布式训练集群拓扑，解决了数据科学家在分布式训练中面临的难题。

介绍

本系列将介绍如何在阿里云容器服务上运行Kubeflow, 本文介绍如何使用TfJob运行分布式模型训练。

第一篇：阿里云上使用JupyterHub
第二篇：阿里云上小试TFJob
第三篇：利用TFJob运行分布式TensorFlow

TensorFlow分布式训练和Kubernetes

TensorFlow作为现在最为流行的深度学习代码库，在数据科学家中间非常流行，特别是可以明显加速训练效率的分布式训练更是杀手级的特性。但是如何真正部署和运行大规模的分布式模型训练，却成了新的挑战。实际分布式TensorFLow的使用者需要关心3件事情。

寻找足够运行训练的资源，通常一个分布式训练需要若干数量的worker(运算服务器)和ps(参数服务器)，而这些运算成员都需要使用计算资源。
安装和配置支撑程序运算的软件和应用
根据分布式TensorFlow的设计，需要配置ClusterSpec。这个json格式的ClusterSpec是用来描述整个分布式训练集群的架构，比如需要使用两个worker和ps，ClusterSpec应该长成下面的样子，并且分布式训练中每个成员都需要利用这个ClusterSpec初始化tf.train.ClusterSpec对象，建立集群内部通信

cluster = tf.train.ClusterSpec({"worker": ["<VM_1>:2222",
                                           "<VM_2>:2222"],
                                "ps": ["<IP_VM_1>:2223",
                                       "<IP_VM_2>:2223"]})

其中第一件事情是Kubernetes资源调度非常擅长的事情，无论CPU和GPU调度，都是直接可以使用；而第二件事情是Docker擅长的，固化和可重复的操作保存到容器镜像。而自动化的构建ClusterSpec是TFJob解决的问题，让用户通过简单的集中式配置，完成TensorFlow分布式集群拓扑的构建。

应该说烦恼了数据科学家很久的分布式训练问题,通过Kubernetes+TFJob的方案可以得到比较好的解决。

利用Kubernetes和TFJob部署分布式训练

修改TensorFlow分布式训练代码

之前在阿里云上小试TFJob一文中已经介绍了TFJob的定义，这里就不再赘述了。可以知道TFJob里有的角色类型为MASTER, WORKER 和 PS。

举个现实的例子，假设从事分布式训练的TFJob叫做distributed-mnist, 其中节点有1个MASTER, 2个WORKERS和2个PS,ClusterSpec对应的格式如下所示：

{  
    "master":[  
        "distributed-mnist-master-0:2222"
    ],
    "ps":[  
        "distributed-mnist-ps-0:2222",
        "distributed-mnist-ps-1:2222"
    ],
    "worker":[  
        "distributed-mnist-worker-0:2222",
        "distributed-mnist-worker-1:2222"
    ]
}

而tf_operator的工作就是创建对应的5个Pod, 并且将环境变量TF_CONFIG传入到每个Pod中，TF_CONFIG包含三部分的内容，当前集群ClusterSpec，该节点的角色类型，以及id。比如该Pod为worker0，它所收到的环境变量TF_CONFIG为:

{  
   "cluster":{  
      "master":[  
         "distributed-mnist-master-0:2222"
      ],
      "ps":[  
         "distributed-mnist-ps-0:2222"
      ],
      "worker":[  
         "distributed-mnist-worker-0:2222",
         "distributed-mnist-worker-1:2222"
      ]
   },
   "task":{  
      "type":"worker",
      "index":0
   },
   "environment":"cloud"
}

在这里,tf_operator负责将集群拓扑的发现和配置工作完成，免除了使用者的麻烦。对于使用者来说，他只需要在这里代码中使用通过获取环境变量TF_CONFIG中的上下文。

这意味着，用户需要根据和TFJob的规约修改分布式训练代码：

# 从环境变量TF_CONFIG中读取json格式的数据
tf_config_json = os.environ.get("TF_CONFIG", "{}")

# 反序列化成python对象
tf_config = json.loads(tf_config_json)

# 获取Cluster Spec
cluster_spec = tf_config.get("cluster", {})
cluster_spec_object = tf.train.ClusterSpec(cluster_spec)

# 获取角色类型和id， 比如这里的job_name 是 "worker" and task_id 是 0
task = tf_config.get("task", {})
job_name = task["type"]
task_id = task["index"]

# 创建TensorFlow Training Server对象
server_def = tf.train.ServerDef(
    cluster=cluster_spec_object.as_cluster_def(),
    protocol="grpc",
    job_name=job_name,
    task_index=task_id)
server = tf.train.Server(server_def)

# 如果job_name为ps，则调用server.join()
if job_name == 'ps':
    server.join()

# 检查当前进程是否是master， 如果是master，就需要负责创建session和保存summary。
is_chief = (job_name == 'master')


# 通常分布式训练的例子只有ps和worker两个角色，而在TFJob里增加了master这个角色，实际在分布式TensorFlow的编程模型并没有这个设计。而这需要使用TFJob的分布式代码里进行处理，不过这个处理并不复杂，只需要将master也看做worker_device的类型
with tf.device(tf.train.replica_device_setter(
    worker_device="/job:{0}/task:{1}".format(job_name,task_id),
    cluster=cluster_spec)):