Kube-batch 在 vivo AI 计算平台的应用

阅读数:5494 2019 年 7 月 4 日 08:25

Kube-batch在vivo AI计算平台的应用

半年前,vivo 研究院开始建设 AI 计算平台,以期解决在统一和高性能训练环境、大规模分布式训练、计算资源的高效利用和调度方面的痛点。在这个过程中,整个团队发现 K8s 的原生调度器存在诸多问题,最后决定采用社区提供的批调度器 kube-batch,本文讲解了 vivo 基础计算小组的 kube-batch 实践历程和解决的问题。

Kube-batch在vivo AI计算平台的应用

vivo AI 研究院的基础计算平台小组从半年前开始建设 AI 计算平台,以解决如下痛点:

统一和高性能的训练环境

痛点:

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

方案:

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

大规模算法分布式训练

痛点:

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

方案:

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

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

痛点:

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

方案:

平台使用了这几年炙手可热的 Kubernetes(简称 k8s),将训练机器加入 k8s 集群。作为统一资源池,各任务都通过平台向 k8s 申请资源,由 k8s 将任务的各个 worker 调度到各个机器上并启动容器。为了进一步提高资源利用率,我们采用了批调度器 kube-batch。

为什么使用 kube-batch

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

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

  • 两个任务,各包含 10 个 worker, 集群资源只能启动 10 个 worker。两个任务分别有 5 个 worker 被启动,但两个任务都无法开始训练。10 个 worker 的资源被浪费了。

由此可见,原生调度器对于分布式训练的调度存在问题,影响了资源的利用率。而 k8s 社区提供了一个批调度器 kube-batch , 它能够将一个训练任务的多个 worker 当做整体进行调度,只有当任务所有 worker 的资源都满足,才会将容器在节点上启动。这解决了上述问题,避免了任务间的资源死锁,提高了资源利用率。

kube-batch 还提供了队列机制,同个队列的任务会依次运行。不同队列直接可以设置优先级,优先级高的队列中的任务会优先得到调度。队列还可以设置权重,权重高的队列分配到的资源会更多。

kube-batch 的原理与实现

Kube-batch在vivo AI计算平台的应用

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

Cache 模块

Cache 模块封装了对 API Server 节点、容器等对象的数据同步逻辑。K8s 的数据保存在分布式存储 etcd 中,所有对数据的查询和操作都通过调用 API Server 接口,而非直接操作 etcd。在调度时,需要集群中节点和容器的使用资源和状态等信息。Cache 模块通过调用 K8s 的 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 的方式部署到 k8s 集群中。k8s 支持多种调度器并存。每个调度器按照约定,只处理指定了自己的容器。容器可以在 scheduleName 这个字段指定调度器名字。默认是”default-scheduler“, 即 k8s 原生的调度器。平台将训练任务的容器的 scheduleName 都设置成”kube-batch“, 这样这些容器就会被 kube-batch 所调度。

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

问题一:调度器偶尔会 crash

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

给社区的 issue:  https://github.com/kubernetes-sigs/kube-batch/issues/665
给社区的 PR: https://github.com/kubernetes-sigs/kube-batch/pull/666

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

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

给社区的 issue: https://github.com/kubernetes-sigs/kube-batch/issues/729
给社区的 PR: https://github.com/kubernetes-sigs/kube-batch/pull/730

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

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

给社区的 issue: https://github.com/kubernetes-sigs/kube-batch/issues/561

作者介绍:

吴梓洋,曾就职于 Oracle, Rancher 等公司,目前是 vivo AI 研究院计算平台组的资深工程师。也是 kube-batch, tf-operator 等项目的 contributor。

评论

发布