[源码解析] 深度学习分布式训练框架 horovod (18) --- kubeflow tf-operator

目录

0x00 摘要

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) --- 弹性训练之容错

0x01 背景知识

1.1 Kubernetes

kubernetes,简称K8s,是用8代替8个字符“ubernete”而成的缩写。是一个开源的,用于管理云平台中多个主机上的容器化的应用,Kubernetes的目标是让部署容器化的应用简单并且高效(powerful),Kubernetes提供了应用部署,规划,更新,维护的一种机制。

Kubernetes 是一种越来越受欢迎的深度神经网络训练选项,因为它提供了通过容器使用不同机器学习框架的灵活性,以及按需扩展的敏捷性。

当面临较复杂的模型训练或者数据量大时,单机的计算能力往往无法满足算力要求。通过使用阿里的 AiACC 或者社区的 horovod 等分布式训练框架,仅需修改几行代码,就能将一个单机的训练任务扩展为支持分布式的训练任务。

在 Kubernetes 上常见的是 kubeflow 社区的 tf-operator 支持 Tensorflow PS 模式,或者 mpi-operator 支持 horovod 的 mpi allreduce 模式。

1.2 容器作为调度单元

为什么希望使用容器来作为深度学习系统的调度单元?因为容器拉取/启动快速。隔离资源效果好。抽象来看,可以将容器的image作为job的一部分分发调度执行。当然容器化后会引入gpu,网络等性能的代价。

比如 nvidia gpu 对docker提供了支持,nvidia-docker可以代替docker执行create和run操作。下图就是nvidia-docker架构。

1.3 Kubeflow

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 开始继续执行。

1.4 Tensorflow on Kubeflow

Kubeflow 支持两种不同的 Tensorflow 框架分布式训练方法。

  • 第一种是原生 Tensorflow 架构,它依赖于集中式参数服务器来实现工作线程之间的协调。
  • 第二种是分散式方法,工作线程通过 MPI AllReduce 原语直接相互通信,不使用参数服务器。NVIDIA 的 NCCL 库已经在GPU 上有效地执行了大部分 MPI 原语,而 Uber 的Horovod 让使用 TensorFlow 执行多 GPU 和多节点训练变得轻而易举。与参数服务器相比,第二种方法可以更好地优化带宽更好地扩展

1.5 Operator

Operator 是Kubernetes 之中的概念,主要是用来打包、部署及管理用户的任务。

Operator可以简单理解为 CRD + Controller。

  • CRD(Custom Resource Definition)是 Kubernetes 的扩展类型,用来为用户自定义资源提。
  • Controller 用来让用户操作CRD。

如果用 Java 来比喻,operator 就是 Class,CRD 就是类的成员变量,Controller 就是类成员方法。

1.6 TF-Operator

虽然KubeFlow提供了一大堆组件,涵盖了机器学习的方方面面,但模型训练肯定是KubeFlow最重要的功能。 KubeFlow针对各种各样的机器学习框架提供了训练的能力。方式是定义了各种各样的Operator,其主要是用来管理机器学习或者深度学习里面的任务,比如如何管理维护一个任务的多个节点,如何管理Pod及任务的生命周期,如何进行容错等等。

TF-Operator就是开源社区基于K8S提供的扩展API,提供了TensorFlow的训练能力,从名字也能看出来,这个实现是类似Job的一种方式,其特点如下:

  • 提供TensorFlow原生PS-worker架构 的多机训练
  • 推荐将PS和worker一起启动
  • 通过service做服务发现
  • 在社区中最早期的Operator

因为 TF-Operator 是社区中最早期的Operator,所以我们有必要先看看。

0x02 TensorFlow 分布式

因为 TF-Operator 是为了支持 Tensorflow PS 模式,所以我们首先介绍一下 TensorFlow 分布式。

2.1 Parameter server架构

在Parameter server架构(PS架构)中,集群中的节点被分为两类:参数服务器(parameter server)和工作服务器(worker)。其中参数服务器存放模型的参数,而工作服务器负责计算参数的梯度。在每个迭代过程,工作服务器从参数服务器中获得参数,然后将计算的梯度返回给参数服务器,参数服务器聚合从工作服务器传回的梯度,然后更新参数,并将新的参数广播给工作服务器。

PS-Worker 架构的梯度更新有着 同步更新异步更新 两种方式:

在同步训练中, 所有的Worker设备采用同一个Batch的不同小批(mini-batch)数据来训练,等待所有设备该批次的梯度计算完成后,模型才会根据所有的梯度进行一次参数更新,然后PS将更新后的模型下发到各个设备。

异步训练中,没有设备需要去等待其他设备的梯度计算和参数更新,所有设备独立算并与将梯度结果更新到中心节点(PS)。异步训练总体会训练速度会快很多,但是异步训练的一个很严重的问题是梯度失效问题(stale gradients),刚开始所有设备采用相同的参数来训练,但是异步情况下,某个设备完成一步训练后,可能发现模型参数已经被其它设备更新过了,此时这个设备计算出的梯度就过期了。

2.2 Tensorflow PS-Worker

2.2.1 架构

这里只是大致介绍一下,主要是为了和 TF-Operator 对比。

TF 把Job主要划分为Parameter Server和Worker(因为 TF 版本不同,所以有不同阶段的特别定义,比如 master 或者 chief)。

  • Parameter Job:执行模型相关的作业,包括模型参数存储,分发,汇总,更新;作为分布式训练的服务端,等到各个终端(supervisors)来连接。
  • Worker Job: 在TensorFlow的代码注释中被称为supervisors,执行训练相关的作业,包括推理计算和梯度计算。如果参数的数量太大,一台机器处理不了,这就要需要多个Tasks(动态上理解,主机上的一个进程,从静态的角度理解,Task就是我们写的代码)。
  • Chief supervisors:在众多运算终端中必须选中一个作为主要的运算终端。该终端是在运算终端中最先启动的,它的功能是合并各个终端运算后的学习参数,将其保存再写入。
  • Cluster 是 Jobs 的集合: Cluster(集群) 就是集群系统。

每个具体角色网络标识都是唯一的,即分布在不同IP的机器上(或者同一主机但不同端口号)。

在实际运行中,各个角色的网络构建部分代码必须完全相同,Ps-worker 架构分布式模型的流程大致如下:

  • pull : 各个worker根据数据流图的拓扑结构,从PS拉取最新的模型参数

  • feed: 各worker填充不同的批数据

  • compute: 各worker按照相同的模型参数和不同的批数据计算梯度,得出不同的梯度值

  • push 各worker 将计算得到的梯度值上传给PS

  • update: PS 收集所有worker的梯度值,求平均值,更新模型参数。

2.2.2 代码

具体逻辑如下:

  • Task需要知道集群上都有哪些主机,以及它们都监听什么端口。tf.train.ClusterSpec()就是用来描述这个。
  • 这个Cluster(集群)有两个Job(worker.ps),worker中有三个Task(即,有三个Task执行Tensorflow op操作)
  • ClusterSpec当作参数传入到 tf.train.Server()中,同时指定此TaskJob_nametask_index
  • 由于是相同的代码运行在不同的主机上,所以要传入job_nametask_index加以区分,而ps_hostsworker_hosts对于所有主机来说,都是一样的,用来描述集群的。
  • 一个tf.train.Server包含了本地设备(GPUs,CPUs)的集合,可以连接到到其它task的ip:port(存储在cluster中), 还有一个session target用来执行分布操作。还有最重要的一点就是,它创建了一个服务器,监听port端口,如果有数据传过来,他就会在本地执行(启动session target,调用本地设备执行运算),然后结果返回给调用者。
  • 为了使ps_server能够一直处于监听状态,我们需要使用server.join()。这时,进程就会block在这里.至于为什么ps_server刚创建就join呢,原因是因为下面的代码会将参数指定给ps_server保管,所以ps_server静静的监听就好了。
# 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

0x03 TF-Operator

3.1 TF-Operator 设计思路

了解了 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等)结合使用变得非常容易。

3.2 架构图

具体架构图如下:

3.2.1 什么是Pod

我们从图上来看,先看中间的 pod 概念。

pod 是 k8s调度的最小单元。pod 可以理解为:容器组,同时pod相当于逻辑主机,进入pod后仿佛进入一个linux主机,命令都可用(linux系统下),该“主机”内又有很多容器,进入后又仿佛是又进了一个linux主机。默认情况下,每个容器的文件系统与其他容器完全隔离。每个pod都有自己的ip地址。pod内的容器共享相同的ip和端口空间。

3.2.2 为什么要有 service

首先,每个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集合了。

3.2.3 什么是 controller

因为 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。

3.3 Spec

我们首先给出一个 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

下面我们开始进入代码世界。

3.4 TFJob

首先我们看看 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"`
}

3.5 角色

其次我们看看 TF-Operator 之中,对 TF 角色的对应实现。

3.5.1 定义

首先是角色定义。这里的角色基本对应了 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"
)

3.5.2 创建角色

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
}

3.5.3 如何区分 master

用如下方法区分 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
}

0x04 Contoller

下面就进入正题,看看 Controller 如何实现。

4.1 K8S CRD关键概念

首先我们需要看看 K8S CRD 的一些关键概念。

  • informer:监听apiserver中特定资源变化,然后会存储到一个线程安全的local cache中,最后回调我们自己实现的event handler。

  • local cache:informer实时同步apiserver(也就是etcd)中的数据到内存中存储,可以有效降低apiserver的查询压力,但缺点就是实时性不好,本地会比远程的数据落后一点点但会最终与etcd一致,所以需要根据情况具体分析是走Local cache还是apiserver实时获取数据。

  • Lister:提供了CURD操作访问local cache。

  • controller:一个逻辑概念,就是指调度某种资源的实现而已,需要我们自己开发。Controller做的事情主要包括:

    1. 实现event handler处理资源的CURD操作
    2. 在event handler,可以使用workqueue类库实现相同资源对象的连续event的去重,以及event处理异常后的失败重试,通常是建议使用的。
  • Workqueue:一个单独的类库,是可选使用的,但通常都会使用,原因上面说了。我们需要在实现event handler的时候把发生变化的资源标识放入workqueue,供下面的processor消费。

  • Clientset:默认clientset只能CRUD k8s提供的资源类型,比如deployments,daemonset等;生成的代码为我们自定义的资源(CRD)生成了单独的clientset,从而让我们使用结构化的代码CURD自定义资源。也就是说,想操作内建资源就用k8s自带的clientset,想操作CRD就用生成代码里的clientset。

  • Processor:我们实现的go协程,消费workqueue中的事件,workqueue提供了按资源标识的去重。

4.2 定义

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
}

4.3 入口

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)
}

4.4 syncHandler

syncHandler 的作用是依据 key 来同步 Job,就是从 WorkQueue 之中弄出来一个 job,本地处理。

之前设置有 tc.syncHandler = tc.syncTFJob,所以我们实际来到了 syncTFJob。

  • 如果tfjob的期望值已经实现,那么syncTFJob就会用给定的key来同步tfjob,这意味着它不希望更多的

    pod/service被创建或删除:
  • EnableDynamicWorker 这里会根据不同类型设置。
  • 然后会调用 ReconcileJobs 对具体 job 进行处理。
// 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
}

4.5 ReconcileJobs

reconcileTFJobs检查并更新每个给定TFReplicaSpecreplicas,并且做相应处理,可以认为这里是主控逻辑。

  • 如果 job 结束,则做相应处理,delete所有pod和service。
  • 如果TFJob超过了backofflimit或超过了active deadline,删除所有pod和service,然后将状态设置为failed。
  • 否则 遍历配置文件的TFReplicaSpecs部分,
    • 分别为不同类型的节点启动相应的Pod。
    • 在启动Pod之后,还要为其启动一个Service。
// 如果在创建/删除 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。

4.6 处理 Pod

4.6.1 ReconcilePods

reconcilePods为每个给定的TFReplicaSpec检查和更新pod

具体比如:

  • 初始化 replica 的状态;
  • 如果master pod存在,选择master pod,如果没有master,第一个worker pod被选为master;
  • createNewPod 来创建新的 pod;
  • 或者删除 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
}

4.6.2 createNewPod

createNewPod为给定的indextype创建一个新的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
}

4.6.3 生成配置信息

4.6.3.1 SetClusterSpec

上面函数中的生成配置信息比较重要,所以我们单独摘出来说一下。

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
}
4.6.3.2 genTFConfigJSONStr

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
}
4.6.3.3 genClusterSpec

这里就是从集群信息中获得 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
}

4.6.4 CreatePodsWithControllerRef

得到了集群配置信息之后,就使用集群的配置信息,进行真正启动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)
}

4.6.5 createPods

这里才真正调用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+--+
| | | |
+-----------------+ +-------------+

手机如下:

4.7 处理服务

4.7.1 ReconcileServices

ReconcileServices 为每个给定的TFReplicaSpec检查和更新service,大致如下:

  • 将在创建/删除服务时发生错误时请求tfjob。
  • 获取rt类型的所有service。
    • 或者建立新服务;
    • 或者删除旧服务,目前只允许缩小worker的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
}

4.7.2 CreateNewService

为给定的indextype创建一个新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
}

4.7.3 CreateServicesWithControllerRef

使用集群的配置信息,真正启动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)
}

4.7.4 createServices

此时才真正调用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-Operator 的这种自定义资源对象来描述分布式机器学习的训练任务;
  • 同时实现了 TFJob 的 Controller 来控制容器的生老病死,给用户管理好多个进程之间的关系;

0x05 与普通部署比较

分析到这里,大家可能也有点疑惑,究竟 TF on K8s 和 普通部署有啥区别,优势何处呢?我们下面就具体分析下。

5.1 运行

我们首先看源码中的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

5.2 比较

我们就简单从训练代码看看。

5.2.1 普通 TF

各种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)

5.2.2 TF-Operator

首先,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 则大展身手为你搞定一切。

0x06 总结

综合之前的我们可以得出 TF-Operator 如下优势:

  • 通过 TF-Operator 的这种自定义资源对象来描述分布式机器学习的训练任务;
  • 同时实现了 TFJob 的 Controller 来控制容器的生老病死,给用户管理好多个进程之间的关系;
  • 对于用户,只要创建一个 TFJob 的自定义资源对象,在 Template 配置好相关信息,就相当于描述好一个分布式训练程序的执行过程了。
  • 用户可以把精力集中在模型之上。而devops 则大展身手为你搞定一切;

kubeflow/tf-operator 虽然可以运作,但是依然有很多缺陷。

  • Kubeflow 可以在 Kubernetes 上启动基于 TensorFlow 原生的分布式计算能力的作业。但是 因为后者并不能容错,所以 Kubeflow 并不能无中生有。不能容错,也意味着不 能弹性调度。
  • 使用 kubeflow/tf-operator 执行分布式 TensorFlow 作业,模型迭代必须等待申请的进程全部启动后才能开始。如果集群资源不足以启动所有进程,则当前作业只能等待其他作业释放资源。为了缩短资源等待时间,可以给作业配置专有资源池。
  • 由于资源不共享,集群资源利用率会很低。所以 kubeflow/tf-operator 很难同时兼顾研发效率和集群利用率。

而且,最重要的是:没有和 horovod 联系起来,没有安装 MPI 等软件,所以下文我们看看 MPI-Operator。

0xEE 个人信息

★★★★★★关于生活和技术的思考★★★★★★

微信公众账号:罗西的思考

如果您想及时得到个人撰写文章的消息推送,或者想看看个人推荐的技术资料,敬请关注。

0xFF 参考

tensorflow学习笔记(十九):分布式Tensorflow

在 Kubernetes 上弹性深度学习训练利器-Elastic Training Operator

在阿里云上搭建Kubeflow Pipelines

开发你的机器学习工作流

像Google一样构建机器学习系统3 - 利用MPIJob运行ResNet101

揭秘|一探腾讯基于Kubeflow建立的多租户训练平台背后的技术架构

https://blog.csdn.net/weixin_43970890/article/details/113863716

[KubeFlow] MPI-Operator 深度解读

在 Amazon EKS 上优化分布式深度学习性能的最佳实践

云原生AI平台的加速与实践

云原生的弹性 AI 训练系列之一:基于 AllReduce 的弹性分布式训练实践

MPI on Kubernetes

Kubeflow/tf-operator源码分析

MPI,OpenMPI 与深度学习

通过shell执行kubectl exec并在对应pod容器内执行shell命令

k8s系列 – CRD自定义资源与Controller实现(完结篇)

TensorFlow分布式全套(原理,部署,实例)

Kubernetes Operator最佳实践

ElasticDL:蚂蚁金服开源基于 TensorFlow 的弹性分布式深度学习系统

分布式深度学习系统-容器化资源调度

基于 K8S 构建 AI 平台方案解析

[源码解析] 深度学习分布式训练框架 horovod (18) --- kubeflow tf-operator的更多相关文章

  1. [源码解析] 深度学习分布式训练框架 horovod (19) --- kubeflow MPI-operator

    [源码解析] 深度学习分布式训练框架 horovod (19) --- kubeflow MPI-operator 目录 [源码解析] 深度学习分布式训练框架 horovod (19) --- kub ...

  2. [源码解析] 深度学习分布式训练框架 Horovod (1) --- 基础知识

    [源码解析] 深度学习分布式训练框架 Horovod --- (1) 基础知识 目录 [源码解析] 深度学习分布式训练框架 Horovod --- (1) 基础知识 0x00 摘要 0x01 分布式并 ...

  3. [源码解析] 深度学习分布式训练框架 horovod (2) --- 从使用者角度切入

    [源码解析] 深度学习分布式训练框架 horovod (2) --- 从使用者角度切入 目录 [源码解析] 深度学习分布式训练框架 horovod (2) --- 从使用者角度切入 0x00 摘要 0 ...

  4. [源码解析] 深度学习分布式训练框架 horovod (4) --- 网络基础 &amp; Driver

    [源码解析] 深度学习分布式训练框架 horovod (4) --- 网络基础 & Driver 目录 [源码解析] 深度学习分布式训练框架 horovod (4) --- 网络基础 & ...

  5. [源码解析] 深度学习分布式训练框架 horovod (3) --- Horovodrun背后做了什么

    [源码解析] 深度学习分布式训练框架 horovod (3) --- Horovodrun背后做了什么 目录 [源码解析] 深度学习分布式训练框架 horovod (3) --- Horovodrun ...

  6. [源码解析] 深度学习分布式训练框架 horovod (5) --- 融合框架

    [源码解析] 深度学习分布式训练框架 horovod (5) --- 融合框架 目录 [源码解析] 深度学习分布式训练框架 horovod (5) --- 融合框架 0x00 摘要 0x01 架构图 ...

  7. [源码解析] 深度学习分布式训练框架 horovod (6) --- 后台线程架构

    [源码解析] 深度学习分布式训练框架 horovod (6) --- 后台线程架构 目录 [源码解析] 深度学习分布式训练框架 horovod (6) --- 后台线程架构 0x00 摘要 0x01 ...

  8. [源码解析] 深度学习分布式训练框架 horovod (7) --- DistributedOptimizer

    [源码解析] 深度学习分布式训练框架 horovod (7) --- DistributedOptimizer 目录 [源码解析] 深度学习分布式训练框架 horovod (7) --- Distri ...

  9. [源码解析] 深度学习分布式训练框架 horovod (8) --- on spark

    [源码解析] 深度学习分布式训练框架 horovod (8) --- on spark 目录 [源码解析] 深度学习分布式训练框架 horovod (8) --- on spark 0x00 摘要 0 ...

随机推荐

  1. 004_Intelij 使用,Anonymous Apex

    插件安装好之后,就可以正常添加cloud 项目: 注意:免费使用是30天,为了不去买license,在过期后,去修改下机器的日期,这个日期是在安装后的一个月内,改好日期后,启动Intelij 少一个截 ...

  2. java subList方法小记

    在java中,我们有时候需要对一个list进行拆分处理,这个时候就需要用到list的切割方法subList. 我们知道list里面的序号起始是从0开始的,所以一个含有23条数据的list,它的序号是从 ...

  3. Spider爬虫清洗数据(re方法)

    import re s0 = 'BOY and GIRL' s1 = re.sub(r'BOY|GIRL', 'HUMAN', s0) print s1 # HUMAN and HUMAN 替换方法.

  4. Java实战之02Hibernate-04多表映射

    十.多表映射 0.内容补充:数据完整性 作用:防止用户的误操作. 实体完整性:主键.用于确定表中唯一的一条记录. 域完整性:表中的字段. 数据类型约束: 非空约束: 唯一约束: 参照完整性: 多表设计 ...

  5. web.config详解(配置文件节点说明)

    转载:http://www.zzzj.com/html/20081110/67614.html web.config文件是一个XML文件,它的根结点是<configuration>,在&l ...

  6. RH133读书 笔记(3) - Lab 3 Configuring the kernel

    Lab 3 Configuring the kernel Goal: Develop skills tuning the /proc filesystem. Gain some experience ...

  7. 【Alpha】——Fifth Scrum Meeting

    一.今日站立式会议照片 二.每个人的工作 成员 昨天已完成的工作 今天计划完成的工作 李永豪 测试统计功能 对统计出现的问题进一步完善 郑靖涛 着手编写报表设计 继续报表设计 杨海亮 协助编写统计功能 ...

  8. 在做vue计算属性,v-for处理数组时遇到的一个bug

    bug: You may have an infinite update loop in a component render function 无限循环 需要处理的数组(在 ** ssq **里): ...

  9. Python——使用高德API获取指定城指定类别POI并实现XLSX文件合并

    # 以下内容为原创,转载请注明出处 1 import xlrd # 读xlsx import xlsxwriter # 写xlsx import urllib.request # url请求,Pyth ...

  10. 苹果手机input有圆角阴影的解决方法

    input[type=button], input[type=submit], input[type=file], button { cursor: pointer; -webkit-appearan ...