kube-batch在AI计算平台的应用

【编者的话】本文介绍了kube-batch在AI计算平台的使用背景、设计原理以及应用过程中遇到的问题和解决方案。

背景

vivo AI研究院的基础计算平台小组从半年前开始，建设了一个AI计算平台，以此解决以下痛点：

统一和高性能的训练环境

痛点：以前算法工程师直接在物理机上进行训练。当使用一台新的机器的时候，算法工程师需要重新安装训练所需要的依赖，这需要花费不少时间，而不同机器的环境一致性也很难保证。不同项目的依赖可能存在冲突，导致物理机不能被共享，降低了资源使用率。其次，物理机只提供CentOS的操作系统，而有些训练场景，在Ubuntu的系统上性能更好。最后，有时候在同一个机器上会有多个训练任务在跑，任务之间会互相争抢资源，降低了训练的效率。

方案：计算平台通过利用这几年流行的容器技术，解决了上述痛点。将训练环境打包成容器镜像后，可以将镜像跑在任何安装了容器的机器上，而同一台机器也可以运行不同的镜像。平台通过对镜像的优化，比如安装编译优化过的TensorFlow，提高训练的性能，并迅速推广到各个项目。容器的隔离和资源限制的特性，可以保证任务互不影响。

大规模的算法分布式训练

痛点：对于某些场景，训练的数据量十分大，在单机上训练的时间十分长。提高算法的训练效率，从而缩短算法的迭代周期，对算法的优化而言至关重要。对于一些场景，必须在一定时间内完成增量训练， 把新模型应用在线上，才能保证线上的效果。

方案：平台提供了一套基于MPI实现的训练框架。单机训练的代码只需要做简单的调整，就可以以分布式的方式跑起来。并且随着训练节点的增加，训练样本数保持了接近线性的增长。

计算资源的高效利用和调度

痛点：之前每个团队都分配了一定数量的物理机，有些团队的资源使用率不高，有些团队的资源十分紧张。分布式训练的任务，需要同时在多台机器上将任务拉起，并且满足不同任务的资源需求，这需要一个成熟的调度系统。

方案：平台使用了这几年炙手可热的Kubernetes（简称k8s），将训练机器加入到Kubernetes集群中，作为一个统一的资源池。各个任务都通过平台向Kubernetes申请资源，由Kubernetes将任务的各个worker调度到各个机器上并启动容器。为了进一步提高资源利用率，我们采用了批调度器kube-batch。 如果你想和更多Kubernetes技术专家交流，可以加我微信liyingjiese，备注『加群』。群里每周都有全球各大公司的最佳实践以及行业最新动态 。

为什么使用kube-batch

Kubernetes原生的调度器，会将需要启动的容器，放到一个优先队列（Priority Queue）里面，每次从队列里面取出一个容器，将其调度到一个节点上。 分布式训练需要所有worker都启动后，训练才能够开始进行。使用原生调度器，可能会出现以下问题：

• 一个任务包含了10个worker，但是集群的资源只满足9个worker。原生调度器会将任务的9个worker调度并启动，而最后一个worker一直无法启动。这样训练一直无法开始，9个已经启动的worker的资源被浪费了。
• 两个任务，各包含10个worker，集群的资源只能启动10个worker。两个任务分别有5个worker被启动了，但两个任务都无法开始训练。10个worker的资源被浪费了。

由此可见，原生调度器对于分布式训练的调度存在问题，影响了资源的利用率。而Kubernetes社区提供了一个批调度器 kube-batch ， 它能够将一个训练任务的多个worker当做一个整体进行调度，只有当任务所有worker的资源都满足，才会将容器在节点上启动。这解决了上述的问题，避免了任务间的资源死锁，提高了资源的利用率。kube-batch还提供了队列的机制，同个队列的任务，会依次运行。不同队列直接可以设置优先级，优先级高的队列中的任务会优先得到调度。队列还可以设置权重，权重高的队列分配到的资源会更多。

kube-batch的原理与实现

如上图所示，kube-batch中有四个模块，分别是Cache、Session、Plugin和Action。

Cache模块

Cache模块封装了对API Server的节点、容器等对象的数据同步逻辑。Kubernetes的数据保存在分布式存储etcd中，所有对数据的查询和操作都通过调用API Server的接口，而非直接操作etcd。在调度时，需要集群中的节点和容器的使用资源和状态等信息。Cache模块通过调用Kubernetes的SDK，通过watch机制监听集群中的节点、容器的状态变化，将信息同步到自己的数据结构中。

Cache模块还封装了对API server的接口的调用。比如Cache.Bind这个接口，会去调用API Server的Bind接口，将容器绑定到指定的节点上。在kube-batch中只有cache模块需要和API Server交互，其他模块只需要调用Cache模块的接口。

Session模块

如图所示，Session模块是将其他三个模块串联起来的一个模块。Kube-batch在每个调度周期开始时，都会新建一个Session对象，这个Session的初始化时，会做以下操作：

• 调用Cache.Snapshot接口，将Cache中节点、任务和队列的信息拷贝一份副本，之后在这个调度周期中使用这份副本进行调度。因为Cache的数据会不断变化，为了保持同个调度周期中的数据一致性，在一开始就拷贝了一份副本。
• 将配置中的各个Plugin初始化，然后调用plugin的OnSessionOpen接口。Plugin在OnSessionOpen中，会初始化自己需要的数据，并将一些回调函数注册到session中。Plugin可以向Session中注册的函数是：
  
  1. jobOrderFns： 决定哪个训练任务优先被处理（调度、回收、抢占）
  2. queueOrderFns：决定哪个训练队列优先被处理
  3. taskOrderFns：决定任务中哪个容器优先被处理
  4. predicateFns： 判断某个节点是否满足容器的基本调度要求。比如容器中指定的节点的标签
  5. nodeOrderFns： 当多个节点满足容器的调度要求时，优先选择哪个节点
  6. preemptableFns： 决定某个容器是否可以被抢占
  7. reclaimableFns ：决定某个容器是否可以被回收
  8. overusedFns： 决定某个队列使用的资源是否超过限额，是的话不再调度对队列中的任务
  9. jobReadyFns：判断某个任务是否已经准备好，可以调用API Server的接口将任务的容器调度到节点
  10. jobPipelinedFns ： 判断某个任务是否处于Pipelined状态
  11. jobValidFns： 判断某个任务是否有效

注意Plugin不需要注册上面所有的函数，而是可以根据自己的需要，注册某几个函数。比如Predict plugin就只注册了predicateFns这个函数到Session中。

初始化成功后，Kube-batch会依次调用不同的Action的Execute方法，并将Session对象作为参数传入。在Execute中，会调用Session的各种方法。这些方法，有些最终会调用到Cache的方法， 有些是调用Plugin注册的方法。

Action模块

Action模块实现了具体的调度的流程。现在有4个不同的Action：

• Allocate：这个Action负责将还未调度的设置了资源限制（Request、Limit）的容器调度到节点上。
• Backfill：这个Action负责将还未调度的的没设置资源限制的容器调度到节点上。
• Reclaim：这个Action负责将任务中满足回收条件的容器删除。
• Preempt：这个Action负责将任务中满足条件的容器抢占。

Action实现了调度的机制（mechanism），Plugin实现了调度的不同策略（policy）。举个例子，在Allocate中，每次会从优先队列中找到一个容器进行调度，这是机制，是由Action决定的。 而在优先队列中容器 排序 的策略，是调用了Session的TaskOrderFn方法，这个方法会调用Plugin注册的方法，因此策略是由Plugin实现。这种机制和策略分离的软件设计，带来了很好的扩展性和灵活性。

Plugin模块

Plugin模块提供了一种可插拔的方式，向调度提供不同的策略的实现。

如图所示，目前最新版本有6个Plugin，它们分别是：

• drf：实现了Dominant Resouce Fairenss算法，这个算法能够有效对多种资源（CPU、Memory、GPU）进行调度。
• gang：实现了gang scheduling的逻辑，即保证任务所需worker同时被启动。
• predict：判断某个节点是否满足容器的基本要求。
• priority： 根据容器和队列设置的PriorityClass决定容器和队列的优先级。
• node order：决定满足调度要求的节点中，哪个节点优先被选择。
• proportion： 根据队列设置的权重决定每个队列分配到的资源。

kube-batch在平台的应用

首先，按照kube-batch项目的文档，将其以deployment的方式部署到Kubernetes集群中。Kubernetes支持多种调度器并存。每个调度器按照约定，只处理指定了自己的容器。容器可以在scheduleName这个字段指定调度器名字。默认是“default-scheduler”，即Kubernetes原生的调度器。平台将训练任务的容器的scheduleName都设置成“kube-batch”，这样这些容器就会被kube-batch所调度。

在使用kube-batch的过程中，我们遇到了一些问题。我们首先在自己使用的版本上将问题修复，保证平台的可用性。然后我们会将问题反馈到社区中，并提供最新版本上的补丁。这样能够反馈社区，和社区共同推进kube-batch项目的发展。

问题一：调度器偶尔会Crash

经过排查日志，我们发现导致调度器Crash的逻辑在proportion这个Plugin中。在给多个队列分配资源时，会有一个变量remaining记录集群剩余的资源容量。这个变量的计算有误，导致程序抛出Panic后Crash了。

给社区的issue： https://github.com/kubernetes- ... s/665

给社区的PR： https://github.com/kubernetes- ... l/666

问题二：当集群资源足够时，任务没有被成功调度

经过日志的排查和代码分析，我们发现这个问题是proportion的Plugin中，给队列分配资源的计算有误。每个队列有两个属性，一个是allocated，表示实际上已经分配了的资源，一个是deserved，表示应该分配的资源。 当allocated大于deserved的时候，就会停止对队列里的任务进行调度。由于对deserved的计算有误，比正确的要少，导致allocated大于deserved，队列中的任务不再被调度。

给社区的issue： https://github.com/kubernetes- ... s/729

给社区的PR： https://github.com/kubernetes- ... l/730

问题三： 大任务阻塞了后续任务的调度

当有一个大任务，即使用了很多资源的任务，因为集群资源不够而处于等待状态时，后面提交的小任务，哪怕集群有足够资源，也无法得到调到。只有把大任务删掉后，才能被成功调度。这是因为在调度周期中，大任务总是优先被处理。在调度大任务时，Session中记录的空闲资源已经分配给这个大任务了，最后发现资源不够无法启动，也不会释放这些资源。所以后续的任务也不会被调度成功。如果修改这个逻辑，就可能导致大的任务永远都无法启动，因为后面一直有小任务提交并被调度。具体的改动方案社区还在讨论中。这个问题的影响可控，改动成本大，因此我们也没自己修复，而是等待社区的方案。

给社区的issue： https://github.com/kubernetes- ... s/561

原文链接： https://mp.weixin.qq.com/s/zXiSC0RWmow8RJ7XLog8JQ