云原生 AI 的资源调度和 AI 工作流引擎设计分享

2022-12-16
北京
本文字数：8775 字
阅读完需：约 29 分钟

本文整理自同名线上分享，是 12 月份「百度百舸 - 云原生 AI」技术公开课的第一期。

这次分享讲解了单机单卡、单机多卡、多机多卡等场景下云原生 AI 的资源调度和管理方法，介绍了可以帮助 AI 工程师屏蔽底层资源复杂性、无缝对接 AI 任务与 AI 资源的 AI 工作流引擎 PaddleFlow 的技术架构和产品细节，这些内容能够帮读者理解如何在 AI 工程中提升资源效能和工程效能等问题。

今天给大家分享的主题是《云原生 AI 的资源调度和 AI 工作流引擎设计分享》，主要有 3 部分内容：

第一部分会先简单介绍下什么是云原生 AI；

第二部分会重点介绍下云原生 AI 下的资源管理与调度；

第三部分会给大家介绍下百度自研的 AI 工作流引擎 PaddleFlow。

1. 云原生 AI 简介

在过去的一年中，我们先后走访了泛互、智驾、生科、高校等几个行业的 10 多家企业，收集了大家所关心的痛点问题的 Top10。

从这些问题中，我们看到大多数客户的问题，都集中在 2 个方面，首先是资源效能，比如资源利用率、异构芯片调度与虚拟化、国产化芯片，容器网络等；其次是工程效能，例如大模型落地、训练/推理任务的效率、AI 镜像启动速度等等。

除了以上客户所关心的痛点问题，我们也对传统的 AI 工程做了分析，这里列举了比较突出的 4 个问题：

资源分配的不均衡，或者规划的不合理，往往会导致高优任务无法高效的运行；
资源碎片较多，导致在集群有空余资源的情况下，某些任务依旧无法运行，这个问题在 AI 训练的场景中尤为明显；
集群 GPU 资源利用率低，当然这个问题与前两个密切相关；
AI 工程效率低下，这其实也是一个综合性问题，比如分布式训练任务编排复杂度较高、训练数据加载缓慢、训练算子实现不合理导致 GPU 利用率低，以及 AI 作业的镜像往往比较大启动会很缓慢，这些都会导致 AI 工程效率低下。

以上问题，在云原生技术的帮助下，能够得到极大的改善。

我们首先来介绍下什么是云原生与云原生 AI。

云原生是目前云计算领域中最火热的一个话题，我这里直接给出了 CNCF 官方的定义：云原生是构建应用程序一类技术的统称，通过云原生技术可以构建出可弹性扩展的应用程序，这些应用程序可以被运行在不同环境当中，比如公有云、私有云、混合云等新型动态环境中。

从 CNCF 的定义中，我们可以看到两个关键词，可弹性扩展和动态环境，也就是说云原生帮应用解决了资源的高效利用和无缝迁移的问题。

通过云原生的几大技术，应用的开发者不用再考虑底层的运行环境，可以轻易实现快速部署、按需弹性扩伸缩的应用程序。

云原生 AI 则是在 AI 场景中，利用云原生技术，形成以容器服务为核心，以云原生技术作为基础架构的 AI 工程解决方案，无缝的整合了云的计算、存储、负载均衡等服务，同时贯穿了 AI 任务的全生命周期。

云原生 AI 中涉及到了几大技术板块，包括 AI 作业编排、AI 任务加速，AI 异构资源调度与虚拟化、AI 数据加速等。

下面我们就给大家介绍下百度智能云在云原生 AI 领域的一些具体实践。

百度智能云的云原生 AI 是构建在百度百舸· AI 异构计算平台中。百度百舸是一套专注于 AI 工程化建设，提供软硬一体的异构计算平台，最新的 2.0 版本包含了 AI 计算、AI 存储、AI 加速、AI 容器等四大套件。

百度百舸可以为多个业务场景提供专业的解决方案。这其中 AI 容器与 AI 加速，以及 AI 计算的部分能力，都是由云原生 AI 来提供的。

那么下面我们就来看下云原生 AI 的整体架构。

百度智能云的云原生 AI 从底至上对 AI 工程提供了多层次的端到端解决方案。

首先是资源管理层，云原生 AI 提供了异构芯片管理、高性能 RDMA 网络接入、高性能存储接入的能力，其中异构芯片管理中，又包含了双引擎的 GPU 容器虚拟化（技术详情参见文末“传送门”）、remoteGPU、昆仑芯虚拟化等技术。
然后是 AI 调度层，我们通过对多种业务场景的分析，以及客户交流，逐渐将多种调度算法沉淀到云原生 AI 产品中，为 AI 任务提供了高效的、高性能的运行环境。
再往上就是 AI 任务管理层，在这里我们提供了支持多种分布式训练任务的 Operator 部署接口，支持复杂工程作业编排的工作流引擎，以及训练加速、推理加速、镜像加速、数据加速等四大加速能力。其中 AIAK 加速套件是我们面向与推理、训练场景提供高阶加速能力，这次「百度百舸 - 云原生 AI」技术公开课的第二期和第三期会分别进行详细的介绍。

从上面的架构图中可以看到，我们在为 AI 工程的全面提速，提供了端到端的解决方案。

首先在 AI 任务启动前，我们为集群管理员提供了丰富的资源配额管理入口，使得用户可以对集群资源进行很好的规划，同时 AI 工程师可以通过工作流引擎或者分布式训练任务 Operator 对作业进行快捷的编排部署。
然后在 AI 任务启动中，通过我们的 AI 调度器与镜像加速模块，为 AI 任务提供了预加速的能力，在高效的利用集群资源的同时，AI 镜像的按需加载能力使得 AI 任务启动时间提升十倍以上。
最后在 AI 任务启动后，也就是 AI 任务的运行时，通过 AIAK 的加速套件以及我们自研的数据加速组件，使得 AI 任务有数倍的加速效果。

2. 云原生 AI 下的资源调度

我们先来看下云原生 AI 的资源管理和调度的整体逻辑图。

首先在集群中我们会有一个整体的资源视图。在这个资源视图中，可以管理多种异构计算资源，以及自定义资源。集群管理员可以站在全局视角，对各种资源进行合理的容量规划。

当 AI 任务调度时，调度器会根据全局资源视图与租户配额，选择最优节点来进行调度。

在资源配额管理方面，我们实现了一个基于队列的资源配额管理模块，我们称之为资源队列。

资源队列中可以支持多种资源配额，如下图所示，除了常规的 CPU 和内存外，还支持 GPU 卡、GPU 显存以及自定义资源的配额管理。

资源队列分为两种类型。第一种是超发队列，当队列中的资源配额被消耗完毕时，仍然可以将任务超发至队列中，不过这里有个前提是，超发的任务需要用户使用标签或者我们的 Console 界面来显示指定，当其他的资源队列资源不足时，调度器会优先驱逐超发任务，来保证资源供给。另一种是非超发队列，当资源队列剩余资源不足时，则无法调度新的任务到资源队列中。

资源队列支持多租户。从用户视角来看，合理的规划资源队列，可以解决多租户间资源争抢，以及资源分配不合理的问题。

下面我们就来看下云原生 AI 是如何基于资源队列来进行资源调度的。

首先介绍一下 PodGroup 的概念。PodGroup 是一组强关联 Pod 的集合，主要用于批处理工作负载场景，比如 TensorFlow 中的一组 PS 和 Worker，这些 Pod 会用到相同的资源执行相同的任务，并且通常都是同时起且同时停。所以在 AI 场景中，调度器将以 PodGroup 作为调度单元。需要补充的是 PodGroup 存在最小资源量的概念，调度器将根据 PodGroup 最小资源量来判断是否满足调度条件。

调度器在启动之后，便会周期性的开启一个调度会话，同时将当前集群的整体资源视图保存在会话的快照中，然后依次执行入队->资源分配->资源回收->资源抢占->回填等几个 action，每个 action 在执行时又会根据插件集合中不同的调度算法，对准备调度的 PodGroup 进行资源分配。

我们先来介绍下这几个 action 的含义。

首先 PodGroup 在创建时，都会被指定一个资源队列。当调度器在执行第一个 action 入队操作时，首先会遍历集群中的所有队列，然后将处于 pending 状态的 PodGroup 尝试放到队列中，这时如果队列剩余资源不满足 PodGroup 的最小资源量时则无法入队。
调度器执行的第二个 action 是资源分配，同样的，调度器会先从快照中取出所有资源队列进行遍历，然后再将资源队列中 PodGroup 取出进行遍历，如果 PodGroup 中有 pending 的 Pod，便会进行类似 K8s 默认调度器的预选与优选操作，最终选出一个最优节点分配给 Pod。
调度器执行的第三个 action 是资源回收，也就是将面前我们提到的超发队列中的超发任务资源回收回来，给待调度的任务分配资源，
第四个 action 是资源抢占，资源抢占会根据同一个队列中任务的优先级进行抢占操作。所以资源回收和资源抢占的一个核心区别是，资源回收发生在队列间，资源抢占发生在队列内。
最后一个 action 是回填，只要集群中有资源就会分配给待调度的 Pod，回填的主要场景，是防止集群中资源被大任务占用，小任务难以分配到资源的场景，回填可以使小任务被快速调度，从而提升集群整体资源利用率。以上是 AI 调度器的大体逻辑。

下面我们将介绍下云原生 AI 调度器中的一些核心调度算法，以及他们所解决的问题。

首先是 Gang 调度，这个其实在 AI 场景中比较常见的一种调度算法。在 AI 训练场景中，如果某些训练的 Worker Pod 没有被调度成功，已调度的 Worker 会继续空等，造成资源浪费、甚至资源死锁。所以 Gang 调度会根据业务所设置的 PodGroup 最小资源量，当资源满足的情况下才会真正调度。

在上面介绍调度器逻辑时我们提到了资源抢占的逻辑，那同样的道理，如果只是抢占了一个任务中的部分 Pod，也会带来 2 个结果，要么被抢占的任务失败，要么处于空等状态，同样也会造成资源浪费。所以我们基于 Gang 调度的逻辑，也实现了 Gang 抢占的逻辑，即发生抢占时，会将被抢占的 PodGroup 缩容至最小资源量。

然后是 Binpack 调度插件，Binpack 主要解决的是集群资源碎片的问题。从下图中可以看到 Pod 所分配的 GPU 资源，都按照集中调度的逻辑，将节点的资源占满后，才会分配第二个节点。

Binpack 的逻辑同样适用在我们自研的 GPU 虚拟化的场景中，多个共享 GPU 的任务在调度时，也是会优先占满一张卡，再占第二张卡。这样通过 Binpack 插件可以有效的提升集群资源利用率。

然后是 GPU 在离线混部调度。在百度集团内部的 GPU 集群中，往往会同时存在两种类型的业务：

一种是在线业务，这种业务对延时敏感度较高，SLA 也比较严苛，而且通常会长期占用着 GPU，但是运行时也会有着明显的波峰和波谷；
一种是近线业务，大家可以将其理解为常驻的离线业务，它的特点是对延时不敏感，对吞吐要求相对较高，但也有一定 SLA 要求。

从提升资源利用率的角度来讲，我们希望当在线业务处于波谷时，能够将算力让出给到近线业务，处于波峰时，能压制近线业务，将算力重新还回在线业务。

因此我们开发了 GPU 在离线混部功能，这个功能分两部分调度逻辑：

一是在 K8s 层， AI 调度器将按照亲和性，尽量让在线与离线业务混部，同时在线与在线业务反亲和调度；
二是底层，cGPU 将根据在离线任务实际负载，调整算力。离线任务使用在线任务空闲算力，同时通过水位线设置保证在线任务的 SLA。

我们通过下图左边示例来具体了解下这块儿的逻辑。

首先集群中有 3 个单卡的 GPU 节点，node 1 上的 GPU 有一个在线业务和一个离线业务在运行，红色虚线为 SLA 水位线。在 T1 时刻，在线业务的 GPU 算力使用量超过了 SLA 水位线，所以离线任务被 pending 住了。虽然离线业务没有退出进程，但是此时已经分配不到算力。而在 T2 时刻，在线业务的 GPU 使用量下降到了 SLA 水位线之下，此时离线任务被分配到了算力，并且可以用满在线业务的剩余算力。

然后我们在看下 node2 和 node3，如果此时再有一个离线任务被调度时，调度器则会按照在离线的亲和性逻辑，将离线业务调度到 node2 上。

GPU 拓扑架构感知调度，是我们针对单机多卡、多机多卡训练时的一个专项调度优化。

我们知道 NVIDIA GPU 卡从 Volta 架构开始，主流的训练卡就已经支持了 NVLink。相较于 PCIe，NVLink 有着巨大的带宽优势。

从最佳实践来讲，单机多卡训练时应该优先将同一 NUMA node 下 NVLink 数最多的 GPU 调度给同一个 Pod。而多机多卡的训练时，除了需要考虑 NVLink 数量的条件外，还需要考虑网卡与 GPU 卡的物理位置关系，避免出现同一张卡网分配给了不同的 Pod 的情况。因为在这种情况下，可能会出现多 Pod 间相互带宽影响的情况。

基于上述考量，我们实现了 GPU 拓扑架构感知调度算法，首先是感知 GPU 间的拓扑关系，其次是感知网卡与 GPU 的拓扑关系。

最后我们介绍下 Tor 架构感知调度算法。这个主要是针对大规模预训练模型场景的一个调度优化。

对于一些超大模型，可能会需要用到上千卡来进行训练，这会导致 GPU 节点的规模可能会到几百，这些节点毫无疑问会处于不同的 Tor 交换机下，下图左边是一个比较常见的 Spine-Leaf 网络架构图。

通常这种场景下，模型会采用混合并行策略，即数据并行叠加模型并行的网络结构进行训练，会将训练任务拆分成 n 组单元，n 组单元内的 Pod 会先进行数据同步，再到单元间进行数据同步，那么如果这些训练 Pod 分散到不同交换机的节点上，那有可能同组单元的 Pod 要经过最顶层的 Spine 交换机进行通信，这势必会严重影响训练性能，极端情况下会引起严重的网络拥塞。

基于这种情况，我们实现了 Tor 架构感知调度，调度器会依据交换机收敛比，尽可能将同一单元内的 m 个 Pod 调度到同一个 Tor 下的 Node 上，提升大规模分布式训练的网络通讯性能。

以上是我们云原生 AI 资源调度管理方面的一些介绍，接下来我们来介绍下云原生 AI 另一个核心组件，AI 工作流引擎- PaddleFlow。

3. AI 工作流引擎 – PaddleFlow

云原生技术为我们极大的提升资源效能的同时，也引入了比较高的学习成本，对于 AI 工程师来说，如何快速接入云原生环境中是个问题。因此便有了 AI 工作流引擎的出现。

AI 工作流引擎可以作为一个很好的桥梁，让 AI 工程师使用更加简单、熟悉的语义来编排 AI 工程，并且形成标准、可服用的工程模版，来提升作业效能；同时工作流引擎也为 AI 工程师屏蔽了很多底层细节，使得 AI 任务与 AI 资源可以无缝的对接。

基于上述背景，百度开发了专门面向 AI 场景的工作流引擎 —— PaddleFlow。

Paddleflow 是百度智能云 CCE（Cloud Container Engine）云原生 AI 产品中的核心组件，向上可以承接各种 AI 中台，向下承载了云原生 AI 中的各种核心能力，包括资源调度、GPU 虚拟化、RDMA 组件、以及分布式缓存系统等等。

首先 PaddleFlow 是一个构建在 K8s 之上的工作流引擎，天生具有云原生的特性，借助 K8s 的技术生态，PaddleFlow 可以轻易的接入各种异构环境。

其次 PaddleFlow 提供了计算、存储的统一接入方式，内置了多种数据处理和算法库，同时支持业界各种主流的分布式训练框架。

目前 PaddleFlow 在百度厂内已经有广泛的实践，基于我们上面介绍的云原生 AI 资源调度管理引擎，已经可以承载万级算力卡的调度，每日运行的训练作业数也已达万级，能够支持千亿级参数模型的训练任务。

最后，PaddleFlow 是一个轻量易用的工作流引擎，我们在 CCE 上提供了一键部署的产品入口，结合云端负载均衡器、RDS 服务，为用户提供了一个稳定高可用的运行环境。

PaddleFlow 目前已经在 GitHub 上开源，感兴趣的同学也可以在 GitHub 上下载 PaddlFlow 来部署到自己的 K8s 环境中试用。

下面我们来看下 PaddleFlow 的整体架构。

前面我们提到 AI 工作流引擎是 AI 工程与云原生之间的一个桥梁，主要解决了作业效能与资源效能的两个核心问题，所以我们将 PaddleFlow 整体分为了两部分。

第一部分是面向于 AI 开发工程师或者 AI 平台的批量作业调度系统，也就是下图中的绿色部分。在这部分中，我们通过内置的深度学习引擎与传统的机器学习引擎，可以支持各种主流 AI 框架，并且通过工作流调度内核 Pipeline Core，使得 AI 任务可以高效的运作起来，再往上用户接口层，我们也提供了命令行/SDK/WebUI 的多种接入方式，使得 AI 任务可以灵活接入。
第二部分则是面向底层资源的管理调度系统，也就是下图中的红色部分，这里我们对计算资源与存储资源进行了进一步的抽象，形成了资源调度与数据访问两大内核模块。可以看到资源调度部分，内部使用的都是我们在上文第二部分介绍的调度算法，而数据访问这部分则是结合 PaddleFlowFS 分布式缓存文件系统，向上屏蔽了底层的存储差异，为用户提供了高效的数据访问能力，而对用户则是暴露 POSIX 接口，任务无需代码改造就可接入。

同时在这一层我们也实现了位置感知的能力，使得计算实例与存储实例可以协同调度，减短数据访问路径，从而加速数据访问性能。

下面我们重点介绍一下几个核心模块的设计，首先是工作流调度的核心模块之一：Pipeline Core。

在这里我们首要解决的问题是作业编排的易用性。当用户需要构建一个工作流时，只需要预先准备好自己的代码、数据、镜像，同时定义好工作流，便可以一键提交运行。考虑到 AI 开发者大多数都是使用 Python 语言进行开发，我们为用户提供了 Python DSL 和 Yaml 两种语义的编排方式，AI 工程师可以利用内置的多种管理模块与条件表达式，来快速构建出一个复杂的工作流。

同时我们支持开发者在开发机上直接挂载远端数据到本地，方便查看结果和调试数据。

针对训练场景我们做了丰富的定制化功能。例如，一个复杂的训练工作流运行时，可能会频繁出现的某个任务失败的情况，这时如果我们想要重新运行失败的任务时，不需要从工作流的开头任务重新训练，基于 PaddleFlow 断点续跑的机制，继续在失败的任务节点重新运行工作流。而对于已经成功完成的任务，PaddleFlow 会自动缓存中间数据，重新运行工作流时会自动跳过已经成功的任务节点。

针对开发实验场景，AI 工程师可能需要频繁的进行调参测试，这个过程可能需要频繁的对数据进行归档、任务间数据传递、历史版本数据对比等等。针对这种情况，我们提供了零入侵的输入输出管理机制，为归档、对比、复现等工作提供基础数据。

在作业编排方面，我们相比业界比较流行的工作流引擎 Argo，能够提供更丰富的 DAG 调度能力，例如我们支持子 DAG 运行。

这里我们给出了一个简单的工作流编排示例：

用户在工作流定义文件中，先定义了一个名为 preprocess 的数据预处理任务作为工作流入口，然后定义了一个名为 train 的训练任务进行训练，任务定义中通过 deps 字段描述了它的前置依赖任务 preprocess，最后又定义了一个名 validate 的任务对训练后的模型进行验证，这便是一个最简单的训练工作流。

从 Yaml 中可以看到我们相对于 K8s 的复杂命令进行了抽象，一次书写可多次调用。AI 工程师可以通过简单的编排，就可以运行一个完整的工作流，同时我们在工作流中预置了很多常用的算子，可以直接在 command 中调用。

当工作流被定义好并提交到 PaddleFlow Server 后，PaddleFlow Server 会先对工作流进行解析，然后按照 DAG 的执行顺序进行调度，将任务以 Pod 的形式提交到 K8s 环境中。

进入到 K8s 之后便是上文第二部分内容的逻辑了，但是在 AI 工作流引擎这个场景中，我们进一步扩展出了层级调度的能力。

层级调度是 PaddleFlow 基于云原生 AI 调度实现的特色能力。我们对云原生 AI 的资源队列进行了进一步的扩展，实现了层次化的弹性 Quota 设计。

当集群中有多个用户时，为了保障用户有足够的资源使用，管理员会将集群的资源固定分配给不同的用户。

传统的方法是通过 K8s 原生的 ResourceQuota 方式进行固定资源的分配。因为不同的用户使用资源的周期和方式不同，可能会出现在某一时刻，一些用户的资源不够用，而另一些用户的资源闲置。如果出现多个类似的情况，则整个集群会有很多资源浪费，导致集群的整体资源利用率下降。

针对上述问题，我们实现了层级队列以及层级调度，这里有几个关键设计点：

层次化的队列设计。这种层次化的队列设计保证了子队列可以使用父队列设置的全部资源。这样通过层次化的管理，更容易合理分配和限制资源的使用。
容量保证。队列上都会设置一个资源的占比，这样可以保证每个队列都不会占用整个集群的资源。
弹性分配。空闲的资源可以被分配给任何队列。当多个队列出现争用的时候，则会按照比例或其他策略进行平衡。
多租户租用。通过队列的容量限制，多个用户就可以共享同一个集群，同事保证每个队列分配到自己的容量，提高利用率。

这些设计特点带来了诸多优势，比如做到了灵活的资源管理，更适合企业商用多个层级进行资源隔离和分配的场景；资源容量的保证，通过设置资源上下限，保障任务最低可使用资源并限制资源上限；按需抢占回收，保障不同的优先级和资源空闲下任务之间可以复用资源。

PaddleFlow 的核心能力之二是如何统一存储抽象，能够支持 AI 业务快速接入更多的存储系统，减少代码改造成本。

AI 存储面向的场景非常复杂。举个例子 AI 场景中会有图片、文本结构化的存储需求，大文件读写、小文件随机读写、日志 writeback 等等，没有一种存储系统能够同时满足所有需求。对此我们进行了专门的对比问题：传统存储系统存在 AI 训练数据不够用，meta 访问一般存在问题，远程访问性能较差经常超时等问题。

考虑到私有化用户一般都已有文件系统，并且无更多机器预算数据迁移，因此我们采用增加一个统一访问协议的存储中间件，开放二次开发能力，并针对 AI 场景重点优化，这样做的好处是轻量、运维简单。

下面我们来看下具体的设计，

下图左侧是数据访问内核 PaddleFlowFS 的逻辑架构图。

首先是接口层，我们实现了 Fuse Client，能够支持训练场景最常见的 API，在训练第二轮 epoch 时直接读缓存，提升训练效率。除此之外，我们也提供了 SDK 与 CSI Driver 两种接口，以便用户更灵活的接入。
接下来是 VFS 抽象层，在这一层，我们实现了对文件系统的抽象定义，包括文件读写接口、文件属性接口、目录树操作接口，使得可以对接多种远端存储，包括 HDFS 文件系统、类 S3 对象存储系统等等。
然后在 VFS 抽象层与远端存储中间，我们还实现了一个缓存层，可以将高频访问的数据，存储到本地内存和磁盘上，提供两层缓存能力，利用 Cache Locality 技术，实现带宽和网络延迟亚毫秒级别保障，同时我们为 HDFS 与 S3 存储建立了目录树缓存，能够支持百万 QPS。

下图右侧是一张实测的性能对比图，可以看到 PaddleFlowFS 相比 HDFS Fuse 与 S3FS，首次读性能提升 5-10 倍，而第二次经缓存读也有 30%+ 的性能提升。