## 机器学习平台 JDLP 长成记
文/徐新坤，吕江昭，郭卫龙
>京东容器平台经过几年的发展，高效支撑京东全部业务系统。积累了丰富的数据中心基础设施建设、应用调度、业务系统高可用、弹性伸缩等方面的宝贵经验。更重要的是京东容器平台可以集中提供65万核 CPU-Cores 的计算能力。自然会全力 Support 目前最具影响力的机器学习领域需求。以此京东商城基础平台部集 群技术团队与机器学习团队联合推出基于 Kubernetes 研发的机器学习平台 JDLP。皆在为研发团队提供具有充足 CPU+GPU 计算能力的统一云端机器学习平台，服务众多业务方。让机器学习计算平台资源按需随手可得，并统一提供训练任务强隔离、高可用、弹性伸缩等能力和服务，让业务更关注在算法和业务需求上。
训练脚本：用于进行训练的脚本，一般使用 python 写成。
训练数据集：用于进行训练的数据集合。
训练模型：训练的最终结果数据。可以根据训练模型提供如 Google 翻译等类似的服务。
用户：使用机器学习平台进行训练的用户。
外部用户：使用训练模型得到如谷歌翻译之类的服务的用户。
rs：Kubernetes 中的 replica set。rs 可以指定副本数量。Kubernetes 会自动监控符合条件的 Pod 数量，当 Pod 由于退出或者其他原因导致数量不足时，会自动进行重启或者重建，补齐 Pod 数量以提供服务。
svc：Kubernetes 中的 service。
pod：Kubernetes 中调度的基本单位。由一个或多个容器组成。在本平台的实践中，每个 pod 只有一个容器。因此下文中所指容器与 pod 为同一概念。
### 基础架构
Tensorkube 是京东基于 Tensorflow+Kubernetes 研发的云端机器学习平台。负责整个平台训练任务的编排、serving 服务的管理等。用户通过 Tensorkube 实现训练任务的提交、查询、删除以及 serving 服务的管理（如图1）。
图1 JDLP 基础架构
机器学习平台 Tensorkube 立足于容器，基于
Kubernetes，充分利用了 Kubernetes 的 replica set 的故障恢复和 service 的域名服务功能，结合 Tensorkube 对于任务的编排，提供了任务从训练到提供 servering 服务的整套功能。
它主要提供了任务训练和 servering 两大服务。任务训练的主要目标是根据用户的训练脚本，利用 GPU 或 CPU 资源进行训练，最终生成训练模型。servering 的主要服务是根据已有的训练模型，对外部用户提供实际的服务。
### 训练任务
#### 训练脚本
典型的分布式训练包括 ps 和 worker 两种类型的进程。ps 主要保存训练中的相关参数，而 worker 则是实际执行训练任务。为适应分布式的训练，对于使用的训练脚本需要进行一定的规范，一个典型的训练脚本如 mnist.py 所示。其中比较重要的几个启动参数包括 job_name，task_index，ps_hosts，worker_hosts。
1. job_name：任务的类型，主要用以区别该进程是作为 ps 还是 worker 提供服务的。
2. task_index：任务编号。
3. ps_hosts：用于声明 ps 服务的多个地址。
4. worker_hosts：用于声明 worker 服务的多个地址。
用户提供的训练脚本需要使用该规范。Tensorkube 平台将利用以上参数对训练任务进行编排，以提供伸缩和故障恢复的服务保障。
### 训练镜像
使用 Dockerfile 制作训练镜像，充分利用 docker 容器的分层特点，将训练脚本制作形成镜像，以便分发。一个典型的 Dockerfile 如下。
```
FROM tensorflow/tensorflow:nightly
COPY train.py /
```
Tensorkube 将用户的训练脚本和生成的 Dockerfile 传输给 jenkins，并触发 jenkins 的自动构建。jenkins 最终将容器镜像（假设镜像名为 train01:v1）推送给镜像中心 Harbor。
### 训练任务的编排
对于一个 job，我们将 Pod 分为了三种类型，tensorboard、ps 和 worker。tensorboard 主要用于在训练过程中实时查看训练的进度。一个 job 中只有一个 tensorboard 容器。ps 和 worker 的数量可以由用户指定（最小为1）。
为提升训练效率，worker 容器使用 GPU 资源进行训练。tensorboard、ps 使用 CPU 资源。
worker 容器默认使用 GPU 资源，在集群 GPU 资源不足时，也可由用户指定使用 CPU 资源进行训练。
tensorboard、ps 和 worker 的容器均使用上一节中为该任务生成的训练容器镜像（train01:v1）。三者容器的不同之处在于启动命令不同。根据 ps 和 worker 的数量，分别为其生成启动的参数。
例如，我们设定 ps 2个，worker 3个，则 ps-0 的启动命令设定为：
```
/usr/bin/python /train.py --task_index=0 --job_name=ps --worker_hosts=train01-worker-0:5000,train01-worker-1:5000,train01-worker-2:5000 --ps_hosts=train01-ps-0:5000,train01-ps-1:5000
```
我们为 tensorboard、ps 和 worker 的每个容器分别建立了 rs 和 service。rs 将维持容器的可用性。当容器失效时，将会自动重启或重建。
以下是我们在实际集群的一个样例：
[root@A01 tensorflow]# kubectl get rs（表1）
[root@A01 ~]# kubectl get pod（表2）
[root@A01 tensorflow]# kubectl get service（表3）
每个 rs 只控制一种类型的一个容器，且副本数为1。例如 train01-ps-0 的 rs 副本数为1，其控制的 pod 为 train01-ps-0-jl1on。当 train01-ps-0-jl1on
的 pod 失效时，train01-ps-0 的 rs 将会对 pod 重启或重建。
同时，由于每个 ps 和 worker 均建立了 service，因此可以直接利用域名解析服务。因此在集群内部发向 train01-worker-0:5000 的请求将被最终直接定向到 worker-0 的容器中。
另外，worker 训练过程中的训练数据结果(中间训练模型)将保存在 JSS 中。
当 ps 和 worker 的容器就绪后，即自动开始训练任务。
### servering 服务
训练任务的最终输出结果是训练模型。Tensorkube 将训练模型数据制作成 servering 镜像，用以提供 servering 服务。
servering 服务一般是无状态服务，因此同一个任务的 servering 服务使用同一个 servering 镜像，且使用相同的命令启动。Tensorkube 使用 rs 保证 servering 服务容器的可用性，这一过程类似于无状态的 tomcat 服务，这里不再详述。
### 任务调度
Kubernetes 的调器度（scheduler）作为一种插件（plugins）形式存在，是 Kubernetes 的重要组成部分。所有的 Pod 创建都要经过 scheduler，tensorflow 中的 ps 和 worker 也不例外。
首先，GPU 资源如何衡量和计算是一个颇具争议的话题。现在已经有相关技术支持多任务共享一个 GPU 设备的工作模式，但是仍然不够理想，对应用程序的开发有比较多的要求和限制。当前较为稳妥的做法是不允许多任务共享同一 GPU 设备，一个 GPU 一次只能分配给一个 Pod 而不能“切分”成若干部分进行分配，无法做到像 CPU 资源那样精确到小数。
从大量实际应用来看，一个 Pod 最多获得一个 GPU 设备，一个 GPU 设备最多只能同时分配给一个 Pod，这种方案能够满足大部分应用情形，同时大大降低复杂性、应用稳定性和安全性。纵观目前的硬件发展水平，高性能的
GPU 的并行计算能力超过 CPU 一个数量级，大多数情况下给一个 Pod 分配一个 GPU 设备已经绰绰有余。而如果分配多个 GPU 给一个 Pod，为了充分利用这些 GPU 资源会导致应用程序的复杂性提升。如果一个 GPU 设备同时分配给多个 Pod，可以认为多个应用程序同时使用同一块 GPU，而由于 GPU 本身硬件限制，一旦应用使用的总显存达到设备最大显存，就会出现 OOM（Out Of Memory）错误导致程序崩溃。为了简化管理和优化资源利用，我们最终选择了一对一的折衷方案。
GPU 资源的调度算法在社区中并不算成熟，我们根据实际使用过程中的经验积累，设计并开发了一系列定制调度算法（如图3）。
图3 定制调度算法
### GPU/CPU 资源筛选
GPU 资源筛选主要根据 node 上 GPU 资源是否满足进行初步筛选过滤。每个 node 上的 GPU 设备数量是一定的，亦即可用于分配调度的数量是一定的。kubelet 收集节点上的 GPU 总量及使用信息，并更新到 kube-apiserver 中，进而在 etcd 中持久化存储。Scheduler 根据 kube-apiserver 中各个节点的 GPU 信息判断该 node 上可用的 GPU 资源是否满足 Pod 要求，如果满足则把 node 信息保存。经过筛选后，从而得到符合 GPU 资源的节点（node）列表。
CPU 的资源筛选与此类似。当 worker 指定使用 CPU 资源进行训练，或者对 tensorboard 和 ps 进行调度时，同样先根据 kube-apiserver 中各个节点的 CPU 信息以判断该节点上可用的 CPU 资源满足要求。经过筛选后，从而得到符合 CPU 资源的节点列表。
### 根据 node 上其他条件是否满足做进一步筛选
除了满足基本的资源要求，还有一些筛选算法可以用于进一步筛选，如是否有端口冲突、指定宿主机范围、node 负载大小等。
### 亲和性调度
利用 Kubernetes 亲和性/反亲和性机制，可以实现灵活的 ps 和 worker Pod 调度。首先，利用 node affinity 可以选择具有特定标签的 node，并且语法更加灵活。如规则设置为“soft”或者“preference”时，即使所有 node 都不满足亲和性条件也能成功调度。其次，利用 inter-pod affinity 可以根据 node 上正在运行的 Pod 的标签（而不是 node 的标签），选择是否将
Pod 调度到该 node 上。亲和性/反亲和性在 Pod 的 annotation 字段进行设置。
实际应用中，可以根据需要对 ps/worker Pod 的 node affinity 进行设置。如根据服务器综合性能高低（CPU 型号/内存速度/机器新旧程度/是否 SSD 盘等指标），预先为所有 node 添加标签，将计算密集的 worker Pod 调度到机器性能较高的 node 上，因为 ps Pod 仅存储训练任务相关参数，不承担实际运算工作。同时，为提升 ps 和 worker 之间的数据共享效率，优先将 ps Pod 和 worker Pod 调度到相同的 node 上。平台通过设置 inter-pod affinity 来达到预期效果。
对满足要求的 node 计算优先级，择优选取。
平台在经过以上多个筛选后，通过多种优先级算法（prioritizers），将满足要求的 node 进行排序。优先级算法根据不同规则和标准计算 node 得分，包括该 Pod 使用的镜像是否已经在 node 上、亲和性/反亲和性、均衡资源使用等。平台最终选择一个最优的节点，将调度信息发送至 apiserver。
### 弹性伸缩
tensorflow 训练任务的弹性伸缩一般都是对 worker Pod 数量进行弹性伸缩。如何实现训练任务的弹性伸缩是一个具有挑战性的问题。这里指的弹性伸缩和 Kubernetes 自身提供的弹性伸缩（Horizontal Pod Autoscaling，HPA）有所不同。Kubernetes 的 HPA 是针对无状态服务的，而此处讨论的弹性伸缩针对的 worker Pod 启动时需要传入一些必要参数，如 ps_hosts，worker_hosts，job_name，task_index 等。我们可以认为 worker 是有状态服务，不能单纯使用 HPA 进行管理。
我们通过对 tensorflow 深入定制开发，加上和 Kubernetes 底层机制的开发配合，较好实现了训练任务的弹性伸缩。需要扩容时，首先通过一条 grpc 调用通知 ps 服务器，增加一台或多台 worker（如图4），同时将新增 worker 的域名和端口等参数传递给 ps。ps 服务器获知即将有 worker 加入训练集群，更新相关参数等待
worker 启动和加入，同时确保原训练任务不被中断。然后向 Tensorkube 发送请求增加相应的 worker。
Tensorkube 获取 worker 的启动参数后调用 Kubernetes 接口创建 worker，同时更新该任务所有相关 rs 中的容器启动参数，以便当某个容器故障重启时使用更新后的参数进行重启或重建。
worker 启动后会自动接收 ps 分发的训练任务。值得注意的是，新加入的 worker 会根据当前的训练进度自动执行接下来的训练，而不是从头开始。
当需要缩容时，Tensorkube 将直接销毁对应 worker 的 rs，ps 检测到 worker 数量减少便会自动将任务分配到正在运行的 worker 上，不再向失效的 worker 下发任务，从而完成缩容。
### 故障与恢复
Tensorkube 支持任一训练任务或 servering 服务中任意容器的故障自动恢复。
#### 训练任务故障
**ps 故障**
当 ps 容器故障后，worker 容器也会相继退出，训练将会暂时中止。而后 Kubernetes 将监视到 ps 和 worker 容器退出，并尝试重启或者重建（如果 ps/worker 所在的节点不可用，将会触发重建，并调度到可用节点进行创建）ps 和 worker 容器。ps 和 worker 容器在启动时，从 JSS 上拉取之前保存的中间训练结果，并恢复训练。
**worker 或 tensorboard 故障**
当 worker 或 tensorboard 故障后，不会影响当前正在进行的训练。而后 Kubernetes 将会监视到 tensorboard 或 worker 容器退出，会尝试重启或者重建。当 worker 容器恢复后，将会自动加入到任务中，由 ps 分配训练任务以继续训练。
**servering 服务故障**
servering 是无状态服务，因此当任意一个 servering 的容器失效时，Kubernetes 将会自动维护 servering 服务的容器数量，进行重启或者重建。