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的原理与实现

1.jpg

如上图所示,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

0 个评论

要回复文章请先登录注册