Horovod 是一款基于 AllReduce 的分布式训练框架。凭借其对 TensorFlow、PyTorch 等主流深度学习框架的支持,以及通信优化等特点,Horovod 被广泛应用于数据并行的训练中。
前面通过十几篇文章,我们一步一步分析了 Horovod 的方方面面。接下来就是面对 Horovod on K8S 这座大山。
本文以及后几篇文章目的是:借着分析学习 Horovod on K8S 功能,把相关概念梳理一遍,期望可以从中找出设计思路。所以成文方式是:整理学习了很多网上文章,然后自己分析代码。特此对各位作者深表感谢。
本文是 horovod on k8s 的餐前甜点和必备前提,介绍相关概念以及kubeflow 社区的 tf-operator。
本系列其他文章链接如下:
[源码解析] 深度学习分布式训练框架 Horovod (1) --- 基础知识
[源码解析] 深度学习分布式训练框架 horovod (2) --- 从使用者角度切入
[源码解析] 深度学习分布式训练框架 horovod (3) --- Horovodrun背后做了什么
[源码解析] 深度学习分布式训练框架 horovod (4) --- 网络基础 & Driver
[源码解析] 深度学习分布式训练框架 horovod (5) --- 融合框架
[源码解析] 深度学习分布式训练框架 horovod (6) --- 后台线程架构
[源码解析] 深度学习分布式训练框架 horovod (7) --- DistributedOptimizer
[源码解析] 深度学习分布式训练框架 horovod (8) --- on spark
[源码解析] 深度学习分布式训练框架 horovod (9) --- 启动 on spark
[源码解析] 深度学习分布式训练框架 horovod (10) --- run on spark
[源码解析] 深度学习分布式训练框架 horovod (11) --- on spark --- GLOO 方案
[源码解析] 深度学习分布式训练框架 horovod (12) --- 弹性训练总体架构
[源码解析] 深度学习分布式训练框架 horovod (13) --- 弹性训练之 Driver
[源码解析] 深度学习分布式训练框架 horovod (14) --- 弹性训练发现节点 & State
[源码解析] 深度学习分布式训练框架 horovod (15) --- 广播 & 通知
[源码解析] 深度学习分布式训练框架 horovod (16) --- 弹性训练之Worker生命周期
[源码解析] 深度学习分布式训练框架 horovod (17) --- 弹性训练之容错
kubernetes,简称K8s,是用8代替8个字符“ubernete”而成的缩写。是一个开源的,用于管理云平台中多个主机上的容器化的应用,Kubernetes的目标是让部署容器化的应用简单并且高效(powerful),Kubernetes提供了应用部署,规划,更新,维护的一种机制。
Kubernetes 是一种越来越受欢迎的深度神经网络训练选项,因为它提供了通过容器使用不同机器学习框架的灵活性,以及按需扩展的敏捷性。
当面临较复杂的模型训练或者数据量大时,单机的计算能力往往无法满足算力要求。通过使用分布式训练框架,仅需修改几行代码,就能将一个单机的训练任务扩展为支持分布式的训练任务。
在 Kubernetes 上常见的是 kubeflow 社区的 tf-operator 支持 Tensorflow PS 模式,或者 mpi-operator 支持 horovod 的 mpi allreduce 模式。
为什么希望使用容器来作为深度学习系统的调度单元?因为容器拉取/启动快速。隔离资源效果好。抽象来看,可以将容器的image作为job的一部分分发调度执行。当然容器化后会引入gpu,网络等性能的代价。
比如 nvidia gpu 对docker提供了支持,nvidia-docker可以代替docker执行create和run操作。下图就是nvidia-docker架构。
Kubeflow 是一个开源的 Kubernetes 原生平台,用于开发、编排、部署和运行可扩展的便携式机器学习工作负载。Kubeflow 可以在任何Kubernetes 集群上运行。
Kubeflow可以很好的管理多机任务,Kubeflow的名字比较简单,为Kubernetes + TensorFlow,是一个机器学习工具包,是运行在K8s之上的一套技术栈,这套技术栈包含了很多组件,组件之间的关系比较松散,我们可以配合起来用,也可以单独用其中的一部分。
Kubeflow 询问 Kubernetes 计划分配哪几台机器来运行一个分布式作业中的各个进程,随后告 知每个进程,所有其他进程的 IP 地址和 port。从而保证一个作业里各个进程 之间互相知道对方。
为什么需要让所有进程互相知道对方呢?这是 TensorFlow ps-based distribution 方式要求的。TensorFlow 1.x 原生的分布 式训练功能让一个作业中所有进程都执行 TensorFlow 1.x runtime 程序。这些 进程互相通信,互相协调成为一个“分布式 runtime“,来解释执行表示深度学习 计算过程的计算图(graph)。在开始分布式训练之初,graph 被 TensorFlow runtime 拆解成若干子图;每个进程负责执行一个子图 —— 任何一个进程失败 (可能是被更高优先级作业抢占),则整个大图的执行就失败了。所以 TensorFlow 原生的分布式训练能力不是容错的(fault-tolerant)。不过, 它是可以从错误恢复(fault-recoverable)—— TensorFlow API 提供 checkpoint 的能力;如果一个作业失败了,可以重启作业,从最近的 checkpoint 开始继续执行。
Kubeflow 支持两种不同的 Tensorflow 框架分布式训练方法。
Operator 是Kubernetes 之中的概念,主要是用来打包、部署及管理用户的任务。
Operator可以简单理解为 CRD + Controller。
如果用 Java 来比喻,operator 就是 Class,CRD 就是类的成员变量,Controller 就是类成员方法。
虽然KubeFlow
提供了一大堆组件,涵盖了机器学习的方方面面,但模型训练肯定是KubeFlow
最重要的功能。 KubeFlow
针对各种各样的机器学习框架提供了训练的能力。方式是定义了各种各样的Operator
,其主要是用来管理机器学习或者深度学习里面的任务,比如如何管理维护一个任务的多个节点,如何管理Pod及任务的生命周期,如何进行容错等等。
TF-Operator
就是开源社区基于K8S
提供的扩展API
,提供了TensorFlow
的训练能力,从名字也能看出来,这个实现是类似Job
的一种方式,其特点如下:
因为 TF-Operator 是社区中最早期的Operator,所以我们有必要先看看。
因为 TF-Operator 是为了支持 Tensorflow PS 模式,所以我们首先介绍一下 TensorFlow 分布式。
在Parameter server架构(PS架构)中,集群中的节点被分为两类:参数服务器(parameter server)和工作服务器(worker)。其中参数服务器存放模型的参数,而工作服务器负责计算参数的梯度。在每个迭代过程,工作服务器从参数服务器中获得参数,然后将计算的梯度返回给参数服务器,参数服务器聚合从工作服务器传回的梯度,然后更新参数,并将新的参数广播给工作服务器。
PS-Worker 架构的梯度更新有着 同步更新 和 异步更新 两种方式:
在同步训练中, 所有的Worker设备采用同一个Batch的不同小批(mini-batch)数据来训练,等待所有设备该批次的梯度计算完成后,模型才会根据所有的梯度进行一次参数更新,然后PS将更新后的模型下发到各个设备。
异步训练中,没有设备需要去等待其他设备的梯度计算和参数更新,所有设备独立算并与将梯度结果更新到中心节点(PS)。异步训练总体会训练速度会快很多,但是异步训练的一个很严重的问题是梯度失效问题(stale gradients),刚开始所有设备采用相同的参数来训练,但是异步情况下,某个设备完成一步训练后,可能发现模型参数已经被其它设备更新过了,此时这个设备计算出的梯度就过期了。
这里只是大致介绍一下,主要是为了和 TF-Operator 对比。
TF 把Job主要划分为Parameter Server和Worker(因为 TF 版本不同,所以有不同阶段的特别定义,比如 master 或者 chief)。
Task
就是我们写的代码)。每个具体角色网络标识都是唯一的,即分布在不同IP的机器上(或者同一主机但不同端口号)。
在实际运行中,各个角色的网络构建部分代码必须完全相同,Ps-worker 架构分布式模型的流程大致如下:
具体逻辑如下:
Task
需要知道集群上都有哪些主机,以及它们都监听什么端口。tf.train.ClusterSpec()
就是用来描述这个。Cluster
(集群)有两个Job
(worker.ps),worker
中有三个Task
(即,有三个Task
执行Tensorflow op
操作)ClusterSpec
当作参数传入到 tf.train.Server()
中,同时指定此Task
的Job_name
和task_index
。job_name
和task_index
加以区分,而ps_hosts
和worker_hosts
对于所有主机来说,都是一样的,用来描述集群的。# To build a cluster with two ps jobs on hosts ps0 and ps1, and 3 worker
# jobs on hosts worker0, worker1 and worker2.
cluster_spec = {
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]}
# Create a cluster from the parameter server and worker hosts.
cluster = tf.train.ClusterSpec({"ps": ps_hosts, "worker": worker_hosts})
# Create and start a server for the local task.
server = tf.train.Server(cluster,
job_name=FLAGS.job_name,
task_index=FLAGS.task_index)
if FLAGS.job_name == "ps":
server.join()
稍微完整点的代码如下:
def main(_):
ps_hosts = FLAGS.ps_hosts.split(",")
worker_hosts = FLAGS.worker_hosts.split(",")
# Create a cluster from the parameter server and worker hosts.
cluster = tf.train.ClusterSpec({"ps": ps_hosts, "worker": worker_hosts})
# Create and start a server for the local task.
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":
# 找出worker的主节点,即task_index为0的节点
is_chief = (FLAGS.task_index == 0)
# Assigns ops to the local worker by default.
with tf.device(tf.train.replica_device_setter(
worker_device="/job:worker/task:%d" % FLAGS.task_index,
cluster=cluster)):
# Compute
运行如下,可以看出,我们只需要写一个程序,在不同的主机上,传入不同的参数使其运行:
# On ps0.example.com:
$ python trainer.py \
--ps_hosts=ps0.example.com:2222,ps1.example.com:2222 \
--worker_hosts=worker0.example.com:2222,worker1.example.com:2222 \
--job_name=ps --task_index=0
# On ps1.example.com:
$ python trainer.py \
--ps_hosts=ps0.example.com:2222,ps1.example.com:2222 \
--worker_hosts=worker0.example.com:2222,worker1.example.com:2222 \
--job_name=ps --task_index=1
# On worker0.example.com:
$ python trainer.py \
--ps_hosts=ps0.example.com:2222,ps1.example.com:2222 \
--worker_hosts=worker0.example.com:2222,worker1.example.com:2222 \
--job_name=worker --task_index=0
# On worker1.example.com:
$ python trainer.py \
--ps_hosts=ps0.example.com:2222,ps1.example.com:2222 \
--worker_hosts=worker0.example.com:2222,worker1.example.com:2222 \
--job_name=worker --task_index=1
了解了 TF 分布式的大致运作,我们来看看 TF-Operator 设计思路。
以下是从 "Design Doc TFJob K8s CRD" 中翻译的。
目标是使在Kubernetes(K8s)上运行TensorFlow训练(尤其是分布式训练)变得容易。我建议通过创建一个K8s自定义资源描述符(CRD)和关联的控制器来实现这一点。CRD负责管理运行培训作业所需的K8s资源。 Kubernetes通过提供一个流程(而不是以VM为中心)的世界视图,使得管理流程变得更加容易。Kubernetes还为复杂的分布式应用程序提供了基本的构建块。例如,K8s提供对DNS、健康检查、日志收集、度量收集、存储等的内置支持。 在K8s中,控制器负责确保一套Pods是运行状态。Pod是K8s中的基本构建块,它描述了一个或多个应该进行共定位的进程(相同的ip)。K8s配备了许多内置控制器。可以确保N个pod以特定的规范运行。作业控制器可以用来运行二进制文件。 内置控制器不足以运行分布式TensorFlow作业。TensorFlow是一个有状态的应用程序;每个参数服务器和工作者都需要具有唯一的可寻址性,以支持所有不同的分布式培训模式。K8s有一个statefulset。 但是,有状态集用于永久运行的有状态服务(如Redis之类的内存分片缓存服务),而不是用于运行到完成的作业。 因此,今天在K8s上运行分布式TF作业意味着从内置原语中拼凑出一个解决方案。通常,这意味着手动管理多个资源。例如,用户可以为参数服务器创建一个有状态集,为工作者创建一个有状态集,为主服务器创建一个作业。 为了解决内置资源的限制,K8s支持自定义资源(CRD)和控制器。使用CRD,可以很容易地为特定工作负载创建具有所需语义的控制器,同时将用户隐藏在实现中。K8s社区很快就采用了这种模式,贡献了大量的CRD用于各种工作负载。 开发crd和各种控制器的K8s团队的意见是,大多数控制器使用非分布式、多线程设计,可伸缩性不是问题。 TFJob CRD为K8s定义了TFJob资源。 TFJob资源是 TfReplicas 的集合。每个TfReplica对应一个在工作中扮演角色的一组 TensorFlow processes; 我做出了一个明确的决定,不试图隐藏或替换K8s抽象。例如,每个TfReplica都包含一个标准的K8s PodTemplate 以指定要在每个复制副本中运行的进程(包括TF)。我这样做是因为K8s已经提供了一个被广泛采用和理解的API。因此,引入新的概念来代替K8s的概念是令人困惑的。此外,公开PodTemplate 使TFJob用户可以轻松地利用K8s特性。例如,TFJob用户可以使用K8s将卷附加到其TF进程。这使得TF与K8s支持的任何存储系统(如PDs、NFS等)结合使用变得非常容易。
具体架构图如下:
我们从图上来看,先看中间的 pod 概念。
pod 是 k8s调度的最小单元。pod 可以理解为:容器组,同时pod相当于逻辑主机,进入pod后仿佛进入一个linux主机,命令都可用(linux系统下),该“主机”内又有很多容器,进入后又仿佛是又进了一个linux主机。默认情况下,每个容器的文件系统与其他容器完全隔离。每个pod都有自己的ip地址。pod内的容器共享相同的ip和端口空间。
首先,每个Pod都会被分配一个单独的IP地址,而且每个Pod都提供了一个独立的Endpoint(Pod IP + ContainerPort)以被客户端访问,但这种访问仅限于集群内部,外部没法访问集群内部的IP地址,
其次,Pod的生命是有限的,如果Pod重启IP很有可能会发生变化。当 controller 用新 Pod 替代发生故障的 Pod 时,新 Pod 会分配到新的 IP 地址。这样就产生了一个问题:如果一组 Pod 对外提供服务(比如 HTTP),它们的 IP 很有可能发生变化,那么客户端如何找到并访问这个服务呢?
Kubernetes 给出的解决方案是 Service。
Service只是一个抽象概念,Kubernetes Service 从逻辑上代表了一组 Pod,具体是哪些 Pod 则是由 label 来挑选。Service 在逻辑上将一组pod(功能相同)给抽象出来一个统一入口。可以将他简单理解为做了一个服务的负载均衡。
Service 有自己 IP,而且这个 IP 是不变的。客户端只需要访问 Service 的 IP,Kubernetes 则负责建立和维护 Service 与 Pod 的映射关系。无论后端 Pod 如何变化,对客户端不会有任何影响,因为 Service 没有变。所以一般会通过service来访问pod。core-dns会给service分配一个内部的虚拟ip,因此内部服务可以通过这个ip或者是serviceName来访问到pod的服务。
我们给出一个源码中的service 例子。
apiVersion: v1
kind: Service
metadata:
annotations:
prometheus.io/path: /metrics
prometheus.io/scrape: "true"
prometheus.io/port: "8443"
labels:
app: tf-job-operator
name: tf-job-operator
spec:
ports:
- name: monitoring-port
port: 8443
targetPort: 8443
selector:
name: tf-job-operator
type: ClusterIP
现在我们看到已经创建了名为tf-job-operator
的Service,会分配一个Cluster IP,该Service还会持续的监听selector下面的 Pod,会把这些Pod信息更新到一个名为 tf-job-operator 的Endpoints对象上去,这个对象就类似于我们上面说的Pod集合了。
因为 Kubernetes 现有的资源类型无法满足我们的需求,因此需要通过 Custom Resource Definition 的机制进行扩展。
K8S中一切都是resource,比如Deployment,Service等等。
我们可以基于CRD(CustomResourceDefinitions)功能新增resource,比如我想自定义一种Deployment资源,提供不同的部署策略。
我们知道resource可以通过k8s的RESTFUL API进行CURD操作,对于CRD创建的resource也是一样的。
CRD仅仅是定义一种resource,我们还需要实现controller,类似于deployment controller等等,监听对应资源的CURD事件,做出对应的处理,比如部署POD。
其实,TF-Operator 主要就是一个 Controller 的实现,我们下面也主要就是讲解这个 controller。
我们首先给出一个 Job Spec,这样大家可以在后续和代码中对应。样例如下,拥有一个 master,2个 workers,一个 PS。
apiVersion: "kubeflow.org/v1alpha1" # 指定api版本,此值必须在kubectl api-versions中
kind: "TFJob" # 指定创建资源的角色/类型
metadata: # 资源的元数据/属性
name: "example-job"
spec: # 资源规范字段
replicaSpecs: # 声明副本数目
- replicas: 1
tfReplicaType: MASTER
template: # 模版
spec:
containers:
- image: gcr.io/tf-on-k8s-dogfood/tf_sample:dc944ff # 容器使用的镜像地址
name: tensorflow
args:
- --log_dir=gs://my-job/log-dir
restartPolicy: OnFailure
- replicas: 2
tfReplicaType: WORKER
template:
spec:
containers:
- image: gcr.io/tf-on-k8s-dogfood/tf_sample:dc944ff
name: tensorflow
args:
- --log_dir=gs://my-job/log-dir
restartPolicy: OnFailure
- replicas: 1
tfReplicaType: PS
下面我们开始进入代码世界。
首先我们看看 TFJob 的定义,大致可以和上面的 Spec 中找到对应关系,因为本文目的是了解其大略,所以我们就只分析这些即可。
// TFJob represents a TFJob resource.
type TFJob struct {
// Standard Kubernetes type metadata.
metav1.TypeMeta `json:",inline"`
// Standard Kubernetes object's metadata.
// +optional
metav1.ObjectMeta `json:"metadata,omitempty"`
// Specification of the desired state of the TFJob.
// +optional
Spec TFJobSpec `json:"spec,omitempty"`
// Most recently observed status of the TFJob.
// Populated by the system.
// Read-only.
// +optional
Status commonv1.JobStatus `json:"status,omitempty"`
}
// TFJobSpec is a desired state description of the TFJob.
type TFJobSpec struct {
// RunPolicy encapsulates various runtime policies of the distributed training
// job, for example how to clean up resources and how long the job can stay
// active.
RunPolicy commonv1.RunPolicy `json:"runPolicy,inline"`
// SuccessPolicy defines the policy to mark the TFJob as succeeded.
// Default to "", using the default rules.
// +optional
SuccessPolicy *SuccessPolicy `json:"successPolicy,omitempty"`
// A map of TFReplicaType (type) to ReplicaSpec (value). Specifies the TF cluster configuration.
// For example,
// {
// "PS": ReplicaSpec,
// "Worker": ReplicaSpec,
// }
TFReplicaSpecs map[commonv1.ReplicaType]*commonv1.ReplicaSpec `json:"tfReplicaSpecs"`
// // A switch to enable dynamic worker
EnableDynamicWorker bool `json:"enableDynamicWorker,omitempty"`
}
其次我们看看 TF-Operator 之中,对 TF 角色的对应实现。
首先是角色定义。这里的角色基本对应了 Tensorflow 的各个角色,包括很多为了兼容而保留的角色。
// setTypeNamesToCamelCase sets the name of all replica types from any case to correct case.
func setTypeNamesToCamelCase(tfJob *TFJob) {
setTypeNameToCamelCase(tfJob, TFReplicaTypePS)
setTypeNameToCamelCase(tfJob, TFReplicaTypeWorker)
setTypeNameToCamelCase(tfJob, TFReplicaTypeChief)
setTypeNameToCamelCase(tfJob, TFReplicaTypeMaster)
setTypeNameToCamelCase(tfJob, TFReplicaTypeEval)
}
const (
// TFReplicaTypePS is the type for parameter servers of distributed TensorFlow.
TFReplicaTypePS commonv1.ReplicaType = "PS"
// TFReplicaTypeWorker is the type for workers of distributed TensorFlow.
// This is also used for non-distributed TensorFlow.
TFReplicaTypeWorker commonv1.ReplicaType = "Worker"
// TFReplicaTypeChief is the type for chief worker of distributed TensorFlow.
// If there is "chief" replica type, it's the "chief worker".
// Else, worker:0 is the chief worker.
TFReplicaTypeChief commonv1.ReplicaType = "Chief"
// TFReplicaTypeMaster is the type for master worker of distributed TensorFlow.
// This is similar to chief, and kept just for backwards compatibility.
TFReplicaTypeMaster commonv1.ReplicaType = "Master"
// TFReplicaTypeEval is the type for evaluation replica in TensorFlow.
TFReplicaTypeEval commonv1.ReplicaType = "Evaluator"
)
NewTFJobV2 函数就是依据配置的不同,来创建不同的角色。
这里可以看到,生成 job 时候,基本就是按照 spec 的对应字段来处理。
apiVersion: "kubeflow.org/v1alpha1"
kind: "TFJob"
metadata:
name: "example-job"
spec:
replicaSpecs:
下面是函数定义。
func NewTFJobV2(worker, ps, master, cheif, evaluator int) *tfv1.TFJob {
tfJob := &tfv1.TFJob{
TypeMeta: metav1.TypeMeta{
Kind: tfv1.Kind,
},
ObjectMeta: metav1.ObjectMeta{
Name: TestTFJobName,
Namespace: metav1.NamespaceDefault,
},
Spec: tfv1.TFJobSpec{
TFReplicaSpecs: make(map[commonv1.ReplicaType]*commonv1.ReplicaSpec),
},
}
tfv1.SetObjectDefaults_TFJob(tfJob)
if worker > 0 {
worker := int32(worker)
workerReplicaSpec := &commonv1.ReplicaSpec{
Replicas: &worker,
Template: NewTFReplicaSpecTemplate(),
}
tfJob.Spec.TFReplicaSpecs[tfv1.TFReplicaTypeWorker] = workerReplicaSpec
}
if ps > 0 {
ps := int32(ps)
psReplicaSpec := &commonv1.ReplicaSpec{
Replicas: &ps,
Template: NewTFReplicaSpecTemplate(),
}
tfJob.Spec.TFReplicaSpecs[tfv1.TFReplicaTypePS] = psReplicaSpec
}
if master > 0 {
master := int32(master)
masterReplicaSpec := &commonv1.ReplicaSpec{
Replicas: &master,
Template: NewTFReplicaSpecTemplate(),
}
tfJob.Spec.TFReplicaSpecs[tfv1.TFReplicaTypeMaster] = masterReplicaSpec
}
if cheif > 0 {
cheif := int32(cheif)
cheifReplicaSpec := &commonv1.ReplicaSpec{
Replicas: &cheif,
Template: NewTFReplicaSpecTemplate(),
}
tfJob.Spec.TFReplicaSpecs[tfv1.TFReplicaTypeChief] = cheifReplicaSpec
}
if evaluator > 0 {
evaluator := int32(evaluator)
evaluatorReplicaSpec := &commonv1.ReplicaSpec{
Replicas: &evaluator,
Template: NewTFReplicaSpecTemplate(),
}
tfJob.Spec.TFReplicaSpecs[tfv1.TFReplicaTypeChief] = evaluatorReplicaSpec
}
return tfJob
}
用如下方法区分 master。
func (tc *TFController) IsMasterRole(replicas map[commonv1.ReplicaType]*commonv1.ReplicaSpec, rtype commonv1.ReplicaType, index int) bool {
if ContainChieforMasterSpec(replicas) {
return rtype == tfv1.TFReplicaTypeChief || rtype == tfv1.TFReplicaTypeMaster
}
// else check if it is worker with index 0
return rtype == tfv1.TFReplicaTypeWorker && index == 0
}
下面就进入正题,看看 Controller 如何实现。
首先我们需要看看 K8S CRD 的一些关键概念。
TFController 的定义如下,可以看出来几个成员变量各有所用,就分别用到了上述的部分组件。
// TFController is the type for TFJob Controller, which manages
// the lifecycle of TFJobs.
type TFController struct {
common.JobController
// tfJobClientSet is a clientset for CRD TFJob.
tfJobClientSet tfjobclientset.Interface
// To allow injection of sync functions for testing.
syncHandler func(string) (bool, error)
// tfJobInformer is a temporary field for unstructured informer support.
tfJobInformer cache.SharedIndexInformer
// Listers for TFJob, Pod and Service
// tfJobLister can list/get tfjobs from the shared informer's store.
tfJobLister tfjoblisters.TFJobLister
// tfJobInformerSynced returns true if the tfjob store has been synced at least once.
tfJobInformerSynced cache.InformerSynced
}
TF-Operator 逻辑代码的入口是 runWorker,其实就是循环调用 processNextWorkItem。
func (tc *TFController) runWorker() {
for tc.processNextWorkItem() {
}
}
processNextWorkItem
将从WorkQueue
中读取单个工作项,并尝试通过调用syncHandler
来处理它。
// processNextWorkItem will read a single work item off the workqueue and
// attempt to process it, by calling the syncHandler.
func (tc *TFController) processNextWorkItem() bool {
obj, quit := tc.WorkQueue.Get()
if key, ok = obj.(string); !ok {
tc.WorkQueue.Forget(obj)
return true
}
tfJob, err := tc.getTFJobFromKey(key)
// 同步TFJob以将实际状态匹配到所需的状态。
// Sync TFJob to match the actual state to this desired state.
forget, err := tc.syncHandler(key)
}
syncHandler 的作用是依据 key 来同步 Job,就是从 WorkQueue 之中弄出来一个 job,本地处理。
之前设置有 tc.syncHandler = tc.syncTFJob
,所以我们实际来到了 syncTFJob。
tfjob
的期望值已经实现,那么syncTFJob
就会用给定的key
来同步tfjob
,这意味着它不希望更多的
pod/service
被创建或删除:// syncTFJob syncs the tfjob with the given key if it has had its expectations fulfilled, meaning
// it did not expect to see any more of its pods/services created or deleted.
// This function is not meant to be invoked concurrently with the same key.
// 这个函数不能与同一个key同时调用
func (tc *TFController) syncTFJob(key string) (bool, error) {
namespace, name, err := cache.SplitMetaNamespaceKey(key)
sharedTFJob, err := tc.getTFJobFromName(namespace, name)
tfjob := sharedTFJob.DeepCopy()
// Sync tfjob every time if EnableDynamicWorker is true
tfjobNeedsSync := tfjob.Spec.EnableDynamicWorker || tc.satisfiedExpectations(tfjob)
// 为新tfjob设置默认值。
// Set default for the new tfjob.
scheme.Scheme.Default(tfjob)
if tfjobNeedsSync && tfjob.DeletionTimestamp == nil {
// 调用reconcileTFJobs来启动TFJobs
reconcileTFJobsErr = tc.ReconcileJobs(tfjob, tfjob.Spec.TFReplicaSpecs, tfjob.Status, &tfjob.Spec.RunPolicy)
}
return true, err
}
reconcileTFJobs
检查并更新每个给定TFReplicaSpec
的replicas
,并且做相应处理,可以认为这里是主控逻辑。
// 如果在创建/删除 pods/services时发生错误,它将请求tfjob。
// ReconcileJobs checks and updates replicas for each given ReplicaSpec.
// It will requeue the job in case of an error while creating/deleting pods/services.
func (jc *JobController) ReconcileJobs(
job interface{},
replicas map[apiv1.ReplicaType]*apiv1.ReplicaSpec,
jobStatus apiv1.JobStatus,
runPolicy *apiv1.RunPolicy) error {
metaObject, ok := job.(metav1.Object)
jobName := metaObject.GetName()
runtimeObject, ok := job.(runtime.Object)
jobKey, err := KeyFunc(job)
pods, err := jc.Controller.GetPodsForJob(job)
services, err := jc.Controller.GetServicesForJob(job)
oldStatus := jobStatus.DeepCopy()
// 如果TFJob terminated,则delete所有pod和service。
if commonutil.IsSucceeded(jobStatus) || commonutil.IsFailed(jobStatus) {
// If the Job is succeed or failed, delete all pods and services.
jc.DeletePodsAndServices(runPolicy, job, pods)
jc.CleanupJob(runPolicy, jobStatus, job)
return nil
}
// 检索以前的重试次数
// retrieve the previous number of retry
previousRetry := jc.WorkQueue.NumRequeues(jobKey)
activePods := k8sutil.FilterActivePods(pods)
jc.recordAbnormalPods(activePods, runtimeObject)
active := int32(len(activePods))
failed := k8sutil.FilterPodCount(pods, v1.PodFailed)
totalReplicas := k8sutil.GetTotalReplicas(replicas)
prevReplicasFailedNum := k8sutil.GetTotalFailedReplicas(jobStatus.ReplicaStatuses)
if jobExceedsLimit {
// If the Job exceeds backoff limit or is past active deadline
// delete all pods and services, then set the status to failed
jc.DeletePodsAndServices(runPolicy, job, pods); err != nil {
jc.CleanupJob(runPolicy, jobStatus, job); err != nil {
jc.Recorder.Event(runtimeObject, v1.EventTypeNormal, commonutil.JobFailedReason, failureMessage)
commonutil.UpdateJobConditions(&jobStatus, apiv1.JobFailed, commonutil.JobFailedReason, failureMessage)
return jc.Controller.UpdateJobStatusInApiServer(job, &jobStatus)
} else {
// General cases which need to reconcile
if jc.Config.EnableGangScheduling {
minAvailableReplicas := totalReplicas
_, err := jc.SyncPodGroup(metaObject, minAvailableReplicas)
}
// 遍历配置文件的TFReplicaSpecs部分,分别为不同类型的节点启动相应的Pod。
// 在启动Pod之后,还要为其启动一个Service。
// Diff current active pods/services with replicas.
for rtype, spec := range replicas {
err := jc.Controller.ReconcilePods(metaObject, &jobStatus, pods, rtype, spec, replicas)
err = jc.Controller.ReconcileServices(metaObject, services, rtype, spec)
}
}
err = jc.Controller.UpdateJobStatus(job, replicas, &jobStatus)
// No need to update the job status if the status hasn't changed since last time.
if !reflect.DeepEqual(*oldStatus, jobStatus) {
return jc.Controller.UpdateJobStatusInApiServer(job, &jobStatus)
}
return nil
}
目前逻辑如下:
+------------+
| runWorker |
+-----+------+
|
|
v
+--------+------------+
| processNextWorkItem |
+--------+------------+
|
|
v
+----+------+
| syncTFJob |
+----+------+
|
|
v
+-------+--------+
| ReconcileJobs |
+-------+--------+
|
|
v
+--------+---------+
| |
| |
v v
+---------+---------+ +-----+--------+
| | | |
| ReconcileServices | |ReconcilePods |
| | | |
+-------------------+ +--------------+
下面我们分别介绍处理 Pod 和 处理 Service。
reconcilePods
为每个给定的TFReplicaSpec
检查和更新pod
。
具体比如:
// reconcilePods checks and updates pods for each given TFReplicaSpec.
// It will requeue the tfjob in case of an error while creating/deleting pods.
func (tc *TFController) ReconcilePods(
job interface{},
jobStatus *commonv1.JobStatus,
pods []*v1.Pod,
rtype commonv1.ReplicaType,
spec *commonv1.ReplicaSpec,
replicas map[commonv1.ReplicaType]*commonv1.ReplicaSpec,
) error {
tfJob, ok := job.(*tfv1.TFJob)
// Convert ReplicaType to lower string.
rt := strings.ToLower(string(rtype))
// 获取rtype类型的所有pod。
pods, err := tc.FilterPodsForReplicaType(pods, rt)
numReplicas := int(*spec.Replicas)
masterRole := false
initializeReplicaStatuses(jobStatus, rtype)
// GetPodSlices will return enough information here to make decision to add/remove/update resources.
// For example, let's assume we have pods with replica-index 0, 1, 2
// If replica is 4, return a slice with size 4. [[0],[1],[2],[]], a pod with replica-index 3 will be created.
// If replica is 1, return a slice with size 3. [[0],[1],[2]], pod with replica-index 1 and 2 are out of range and will be deleted.
podSlices := tc.GetPodSlices(pods, numReplicas, logger)
for index, podSlice := range podSlices {
if len(podSlice) > 1 {
logger.Warningf("We have too many pods for %s %d", rt, index)
} else if len(podSlice) == 0 {
// 如果master pod存在,选择master pod
// 如果没有master,第一个worker pod被选为master。
// check if this replica is the master role
masterRole = tc.IsMasterRole(replicas, rtype, index)
// TODO: [should change to CreateNewPod]
err = tc.createNewPod(tfJob, rt, strconv.Itoa(index), spec, masterRole, replicas)
} else {
// Check the status of the current pod.
pod := podSlice[0]
// 目前只允许缩减workers
// check if the index is in the valid range, if not, we should kill the pod
if index < 0 || index >= numReplicas {
err = tc.PodControl.DeletePod(pod.Namespace, pod.Name, tfJob)
}
// Check if the pod is retryable.
if spec.RestartPolicy == commonv1.RestartPolicyExitCode {
if pod.Status.Phase == v1.PodFailed && train_util.IsRetryableExitCode(exitCode) {
tc.Recorder.Event(tfJob, corev1.EventTypeWarning, tfJobRestartingReason, msg)
err := commonutil.UpdateJobConditions(jobStatus, commonv1.JobRestarting, tfJobRestartingReason, msg)
tfJobsRestartCount.Inc()
}
}
updateJobReplicaStatuses(jobStatus, rtype, pod)
}
}
return nil
}
createNewPod
为给定的index
和type
创建一个新的pod
:
// createNewPod creates a new pod for the given index and type.
func (tc *TFController) createNewPod(tfjob *tfv1.TFJob, rt, index string, spec *commonv1.ReplicaSpec, masterRole bool,
replicas map[commonv1.ReplicaType]*commonv1.ReplicaSpec) error {
tfjobKey, err := KeyFunc(tfjob)
expectationPodsKey := expectation.GenExpectationPodsKey(tfjobKey, rt)
// Create OwnerReference.
controllerRef := tc.GenOwnerReference(tfjob)
// Set type and index for the worker.
labels := tc.GenLabels(tfjob.Name)
labels[tfReplicaTypeLabel] = rt
labels[tfReplicaIndexLabel] = index
podTemplate := spec.Template.DeepCopy()
// Set name for the template.
podTemplate.Name = common.GenGeneralName(tfjob.Name, rt, index)
if podTemplate.Labels == nil {
podTemplate.Labels = make(map[string]string)
}
for key, value := range labels {
podTemplate.Labels[key] = value
}
// 生成集群的配置信息,这里最关键,看一下实现
if err := tc.SetClusterSpec(tfjob, podTemplate, rt, index); err != nil {
return err
}
// if gang-scheduling is enabled:
// 1. if user has specified other scheduler, we report a warning without overriding any fields.
// 2. if no SchedulerName is set for pods, then we set the SchedulerName to "kube-batch".
if tc.Config.EnableGangScheduling {
if isNonGangSchedulerSet(replicas) {
tc.Recorder.Event(tfjob, v1.EventTypeWarning, podTemplateSchedulerNameReason, errMsg)
} else {
podTemplate.Spec.SchedulerName = gangSchedulerName
}
if podTemplate.Annotations == nil {
podTemplate.Annotations = map[string]string{}
}
podTemplate.Annotations[gangSchedulingPodGroupAnnotation] = tfjob.GetName()
podTemplate.Annotations[volcanoTaskSpecKey] = rt
}
// 使用上面的配置信息,真正启动Pod的创建
err = tc.PodControl.CreatePodsWithControllerRef(tfjob.Namespace, podTemplate, tfjob, controllerRef)
return nil
}
上面函数中的生成配置信息比较重要,所以我们单独摘出来说一下。
setClusterSpec
为给定的podTemplateSpec
生成并设置TF_CONFIG
:
// SetClusterSpec generates and sets TF_CONFIG for the given podTemplateSpec.
func (tc *TFController) SetClusterSpec(job interface{}, podTemplate *v1.PodTemplateSpec, rtype, index string) error {
tfjob, ok := job.(*tfv1.TFJob)
// Generate TF_CONFIG JSON string.
tfConfigStr, err := genTFConfigJSONStr(tfjob, rtype, index)
// Add TF_CONFIG environment variable to tensorflow container in the pod.
for i := range podTemplate.Spec.Containers {
if podTemplate.Spec.Containers[i].Name == tfv1.DefaultContainerName {
if len(podTemplate.Spec.Containers[i].Env) == 0 {
podTemplate.Spec.Containers[i].Env = make([]v1.EnvVar, 0)
}
podTemplate.Spec.Containers[i].Env = append(podTemplate.Spec.Containers[i].Env, v1.EnvVar{
Name: tfConfig,
Value: tfConfigStr,
})
break
}
}
return nil
}
genTFConfigJSONStr 会生成 json 数据。
// genTFConfig will generate the environment variable TF_CONFIG
// {
// "cluster": {
// "ps": ["ps1:2222", "ps2:2222"],
// "worker": ["worker1:2222", "worker2:2222", "worker3:2222"]
// },
// "task": {
// "type": "ps",
// "index": 1
// },
// }
// }
func genTFConfigJSONStr(tfjob *tfv1.TFJob, rtype, index string) (string, error) {
// Configure the TFCONFIG environment variable.
i, err := strconv.ParseInt(index, 0, 32)
if err != nil {
return "", err
}
cluster, err := genClusterSpec(tfjob)
if err != nil {
return "", err
}
var tfConfigJSONByteSlice []byte
if tfjob.Spec.EnableDynamicWorker {
sparseCluster := convertClusterSpecToSparseClusterSpec(cluster, strings.ToLower(rtype), int32(i))
sparseTFConfig := SparseTFConfig{
Cluster: sparseCluster,
Task: TaskSpec{
Type: strings.ToLower(rtype),
Index: int(i),
},
}
tfConfigJSONByteSlice, err = json.Marshal(sparseTFConfig)
} else {
tfConfig := TFConfig{
Cluster: cluster,
Task: TaskSpec{
Type: strings.ToLower(rtype),
Index: int(i),
},
// We need to set environment to cloud otherwise it will default to local which isn't what we want.
// Environment is used by tensorflow.contrib.learn.python.learn in versions <= 1.3
// TODO(jlewi): I don't think it is used in versions TF >- 1.4. So we can eventually get rid of it.
// 我们需要设置环境为cloud,否则它会默认为local,这不是我们想要的。
Environment: "cloud",
}
tfConfigJSONByteSlice, err = json.Marshal(tfConfig)
}
if err != nil {
return "", err
}
return string(tfConfigJSONByteSlice), nil
}
这里就是从集群信息中获得 cluster 信息。
// genClusterSpec will generate ClusterSpec.
func genClusterSpec(tfjob *tfv1.TFJob) (ClusterSpec, error) {
clusterSpec := make(ClusterSpec)
for rtype, spec := range tfjob.Spec.TFReplicaSpecs {
rt := strings.ToLower(string(rtype))
replicaNames := make([]string, 0, *spec.Replicas)
port, err := GetPortFromTFJob(tfjob, rtype)
// 这里循环生成了TF_CONFIG里面的Cluster信息。注意看注释,使用DNS配合Service,解决的还是各个节点IP不固定的问题
for i := int32(0); i < *spec.Replicas; i++ {
// As described here: https://kubernetes.io/docs/concepts/services-networking/dns-pod-service/#a-records.
// Headless service assigned a DNS A record for a name of the form "my-svc.my-namespace.svc.cluster.local".
// And the last part "svc.cluster.local" is called cluster domain
// which maybe different between kubernetes clusters.
// 如下所述:https://kubernetes.io/docs/concepts/services-networking/dns-pos-service/#a-records。
// Headless service为"my-svc.my-namespace.svc.cluster.local"的名称分配一个DNS记录。
// 最后一部分是"svc.cluster.local"被称为cluster domain,在不同的kubernetes集群之间可能存在差异。
hostName := common.GenGeneralName(tfjob.Name, rt, fmt.Sprintf("%d", i))
svcName := hostName + "." + tfjob.Namespace + "." + "svc"
clusterDomain := os.Getenv(EnvCustomClusterDomain)
if len(clusterDomain) > 0 {
svcName += "." + clusterDomain
}
endpoint := fmt.Sprintf("%s:%d", svcName, port)
replicaNames = append(replicaNames, endpoint)
}
clusterSpec[rt] = replicaNames
}
return clusterSpec, nil
}
得到了集群配置信息之后,就使用集群的配置信息,进行真正启动Pod
的创建:
func (r RealPodControl) CreatePods(namespace string, template *v1.PodTemplateSpec, object runtime.Object) error {
return r.createPods("", namespace, template, object, nil)
}
func (r RealPodControl) CreatePodsWithControllerRef(namespace string, template *v1.PodTemplateSpec, controllerObject runtime.Object, controllerRef *metav1.OwnerReference) error {
if err := ValidateControllerRef(controllerRef); err != nil {
return err
}
return r.createPods("", namespace, template, controllerObject, controllerRef)
}
这里才真正调用K8S
接口创建pod
func (r RealPodControl) createPods(nodeName, namespace string, template *v1.PodTemplateSpec, object runtime.Object, controllerRef *metav1.OwnerReference) error {
pod, err := GetPodFromTemplate(template, object, controllerRef)
if len(nodeName) != 0 {
pod.Spec.NodeName = nodeName
}
if labels.Set(pod.Labels).AsSelectorPreValidated().Empty() {
return fmt.Errorf("unable to create pods, no labels")
}
if newPod, err := r.KubeClient.CoreV1().Pods(namespace).Create(pod); err != nil {
return err
} else {
accessor, err := meta.Accessor(object)
}
return nil
}
此时逻辑如下:
+------------------------------+
+------------+ | SetClusterSpec |
| runWorker | | +-------------------------+ |
+-----+------+ | | genTFConfigJSONStr | |
| | | | |
| | | genClusterSpec | |
v | | | |
+--------+------------+ | +-------------------------+ |
| processNextWorkItem | +------------------------------+
+--------+------------+ |
| |
| v
v +------+-------+ +-----------------------------+ +------------+
+----+------+ +----> | createNewPod +----->+ CreatePodsWithControllerRef +------>+ createPods |
| syncTFJob | | +--------------+ +-----------------------------+ +------------+
+----+------+ |
| |
| |
v |
+-------+--------+ |
| ReconcileJobs | |
+-------+--------+ |
| |
| |
v |
+--------+---------+ |
| | |
| | |
v v |
+------+----------+ +----+--------+ |
| | | | |
|ReconcileServices| |ReconcilePods+--+
| | | |
+-----------------+ +-------------+
手机如下:
ReconcileServices 为每个给定的TFReplicaSpec
检查和更新service
,大致如下:
// reconcileServices checks and updates services for each given ReplicaSpec.
// It will requeue the job in case of an error while creating/deleting services.
func (jc *JobController) ReconcileServices(
job metav1.Object,
services []*v1.Service,
rtype apiv1.ReplicaType,
spec *apiv1.ReplicaSpec) error {
// Convert ReplicaType to lower string.
rt := strings.ToLower(string(rtype))
replicas := int(*spec.Replicas)
// Get all services for the type rt.
services, err := jc.FilterServicesForReplicaType(services, rt)
// GetServiceSlices will return enough information here to make decision to add/remove/update resources.
//
// For example, let's assume we have services with replica-index 0, 1, 2
// If replica is 4, return a slice with size 4. [[0],[1],[2],[]], a svc with replica-index 3 will be created.
//
// If replica is 1, return a slice with size 3. [[0],[1],[2]], svc with replica-index 1 and 2 are out of range and will be deleted.
serviceSlices := jc.GetServiceSlices(services, replicas, commonutil.LoggerForReplica(job, rt))
for index, serviceSlice := range serviceSlices {
if len(serviceSlice) > 1 {
} else if len(serviceSlice) == 0 {
err = jc.CreateNewService(job, rtype, spec, strconv.Itoa(index))
} else {
// Check the status of the current svc.
svc := serviceSlice[0]
// check if the index is in the valid range, if not, we should kill the svc
if index < 0 || index >= replicas {
err = jc.ServiceControl.DeleteService(svc.Namespace, svc.Name, job.(runtime.Object))
}
}
}
return nil
}
为给定的index
和type
创建一个新service
:
// createNewService creates a new service for the given index and type.
func (jc *JobController) CreateNewService(job metav1.Object, rtype apiv1.ReplicaType,
spec *apiv1.ReplicaSpec, index string) error {
jobKey, err := KeyFunc(job)
// Convert ReplicaType to lower string.
rt := strings.ToLower(string(rtype))
expectationServicesKey := expectation.GenExpectationServicesKey(jobKey, rt)
err = jc.Expectations.ExpectCreations(expectationServicesKey, 1)
if err != nil {
return err
}
// Append ReplicaTypeLabel and ReplicaIndexLabel labels.
labels := jc.GenLabels(job.GetName())
labels[apiv1.ReplicaTypeLabel] = rt
labels[apiv1.ReplicaIndexLabel] = index
port, err := jc.GetPortFromJob(spec)
if err != nil {
return err
}
service := &v1.Service{
Spec: v1.ServiceSpec{
ClusterIP: "None",
Selector: labels,
Ports: []v1.ServicePort{},
},
}
// Add service port to headless service only if port is set from controller implementation
if port != nil {
svcPort := v1.ServicePort{Name: jc.Controller.GetDefaultContainerPortName(), Port: *port}
service.Spec.Ports = append(service.Spec.Ports, svcPort)
}
service.Name = GenGeneralName(job.GetName(), rt, index)
service.Labels = labels
// Create OwnerReference.
controllerRef := jc.GenOwnerReference(job)
err = jc.ServiceControl.CreateServicesWithControllerRef(job.GetNamespace(), service, job.(runtime.Object), controllerRef)
if err != nil && errors.IsTimeout(err) {
succeededServiceCreationCount.Inc()
return nil
} else if err != nil {
failedServiceCreationCount.Inc()
return err
}
succeededServiceCreationCount.Inc()
return nil
}
使用集群的配置信息,真正启动Service
的创建:
func (r RealServiceControl) CreateServicesWithControllerRef(namespace string, service *v1.Service, controllerObject runtime.Object, controllerRef *metav1.OwnerReference) error {
if err := ValidateControllerRef(controllerRef); err != nil {
return err
}
return r.createServices(namespace, service, controllerObject, controllerRef)
}
此时才真正调用K8S
接口创建service
:
func (r RealServiceControl) createServices(namespace string, service *v1.Service, object runtime.Object, controllerRef *metav1.OwnerReference) error {
if labels.Set(service.Labels).AsSelectorPreValidated().Empty() {
return fmt.Errorf("unable to create Services, no labels")
}
serviceWithOwner, err := GetServiceFromTemplate(service, object, controllerRef)
newService, err := r.KubeClient.CoreV1().Services(namespace).Create(serviceWithOwner)
accessor, err := meta.Accessor(object)
}
此时逻辑拓展如下:
+------------------------------+
+------------+ | SetClusterSpec |
| runWorker | | +-------------------------+ |
+-----+------+ | | genTFConfigJSONStr | |
| | | | |
| | | genClusterSpec | |
v | | | |
+--------+------------+ | +-------------------------+ |
| processNextWorkItem | +------------------------------+
+--------+------------+ |
| |
| v
v +------+-------+ +-----------------------------+ +------------+
+----+------+ +----> | createNewPod +----->+ CreatePodsWithControllerRef +------>+ createPods |
| syncTFJob | | +--------------+ +-----------------------------+ +------------+
+----+------+ |
| |
| |
v | +------------------+ +---------------------------------+ +----------------+
+-------+--------+ | +----> | CreateNewService +---->+ CreateServicesWithControllerRef +--->+ createServices |
| ReconcileJobs | | | +------------------+ +---------------------------------+ +----------------+
+-------+--------+ | |
| | |
| | |
v | |
+--------+---------+ | |
| | | |
| | | |
v v | |
+------+----------+ +----+--------+ | |
| | | | | |
|ReconcileServices| |ReconcilePods+--+ |
| | | | |
+------+----------+ +-------------+ |
| |
+---------------------------------->+
手机如下:
所以我们大致可知,TF-Operator 本质上就是:
分析到这里,大家可能也有点疑惑,究竟 TF on K8s 和 普通部署有啥区别,优势何处呢?我们下面就具体分析下。
我们首先看源码中的Dockerfile内容
FROM tensorflow/tensorflow:1.5.0
ADD . /var/tf_dist_mnist
ENTRYPOINT ["python", "/var/tf_dist_mnist/dist_mnist.py"]
然后看看对应的 spec,分别有2个 PS,4个 Worker。
apiVersion: "kubeflow.org/v1"
kind: "TFJob"
metadata:
name: "dist-mnist-for-e2e-test"
spec:
tfReplicaSpecs:
PS:
replicas: 2
restartPolicy: Never
template:
spec:
containers:
- name: tensorflow
image: kubeflow/tf-dist-mnist-test:1.0
Worker:
replicas: 4
restartPolicy: Never
template:
spec:
containers:
- name: tensorflow
image: kubeflow/tf-dist-mnist-test:1.0
然后再安装example,跑一个分布式的 mnist 训练任务。
cd ./examples/v1/dist-mnist
docker build -f Dockerfile -t kubeflow/tf-dist-mnist-test:1.0 .
kubectl create -f ./tf_job_mnist.yaml
我们就简单从训练代码看看。
各种host 的配置是通过脚本参数来设置的,下面就是读取参数的配置启动。
# 读取参数
ps_spec = FLAGS.ps_hosts.split(',')
worker_spec = FLAGS.worker_hosts.split(',')
# 创建集群
num_worker = len(worker_spec)
cluster = tf.train.ClusterSpec({'ps': ps_spec, 'worker': worker_spec})
server = tf.train.Server(cluster, job_name=FLAGS.job_name, task_index=FLAGS.task_index)
首先,dist_mnist.py中有如下方式获取 cluster 信息。
# If not explicitly specified in the constructor and the TF_CONFIG
# environment variable is present, load cluster_spec from TF_CONFIG.
tf_config = json.loads(os.environ.get('TF_CONFIG') or '{}')
其次,在 TF-Operator 之中有如下,说明 cluster 信息是从这里设置:
tfConfig = "TF_CONFIG"
然后,在 SetClusterSpec 中有如下,就是调用 K8S 接口动态获取配置:
// SetClusterSpec generates and sets TF_CONFIG for the given podTemplateSpec.
func (tc *TFController) SetClusterSpec(job interface{}, podTemplate *v1.PodTemplateSpec, rtype, index string) error {
tfjob, ok := job.(*tfv1.TFJob)
// Do not set TF_CONFIG for local training jobs.
if !isDistributed(tfjob) {
return nil
}
// Generate TF_CONFIG JSON string.
tfConfigStr, err := genTFConfigJSONStr(tfjob, rtype, index)
// Add TF_CONFIG environment variable to tensorflow container in the pod.
for i := range podTemplate.Spec.Containers {
if podTemplate.Spec.Containers[i].Name == tfv1.DefaultContainerName {
if len(podTemplate.Spec.Containers[i].Env) == 0 {
podTemplate.Spec.Containers[i].Env = make([]v1.EnvVar, 0)
}
podTemplate.Spec.Containers[i].Env = append(podTemplate.Spec.Containers[i].Env, v1.EnvVar{
Name: tfConfig,
Value: tfConfigStr,
})
break
}
}
return nil
}
因此可以知道,从用户角度看,就修改了一点代码即可。至于部署服务等,都是由 K8S 接管了。
用户只要在 spec 之中设定需要多少 worker,ps 就成。这样用户就可以把精力集中在模型之上。而devops 则大展身手为你搞定一切。
综合之前的我们可以得出 TF-Operator 如下优势:
kubeflow/tf-operator 虽然可以运作,但是依然有很多缺陷。
而且,最重要的是:没有和 horovod 联系起来,没有安装 MPI 等软件,所以下文我们看看 MPI-Operator。
tensorflow学习笔记(十九):分布式Tensorflow
在 Kubernetes 上弹性深度学习训练利器-Elastic Training Operator
像Google一样构建机器学习系统3 - 利用MPIJob运行ResNet101
揭秘|一探腾讯基于Kubeflow建立的多租户训练平台背后的技术架构
https://blog.csdn.net/weixin_43970890/article/details/113863716
在 Amazon EKS 上优化分布式深度学习性能的最佳实践
云原生的弹性 AI 训练系列之一:基于 AllReduce 的弹性分布式训练实践
通过shell执行kubectl exec并在对应pod容器内执行shell命令