在 KubeCon CN 2023 的「 Open AI 数据 | Open AI Data」专题中,火山引擎软件工程师胡元哲共享了《运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载|Sailing Ray workloads with KubeRay and Kueue in Kubernetes议题。以下是本次演讲的文字稿。

本文将从 Ray 为何得到 AI 研究者们的青睐,在字节怎么运用 KubeRay 来保管 Ray 运用,Kueue 怎么办理和调度 RayJob 三个方面进行介绍。

什么是 Ray

Ray 起源于 UC Berkeley 的 RISElab 实验室,其定位是一个通用的分布式编程结构,能协助用户将自己的程序快速分布式化。Ray Core 供给了 low level 的分布式语法,如 remote func、remote class,上层 Ray AIR 供给了 AI 场景的相关库。

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

Ray 的GitHub repo 现在已有 27K star,其发起者也成立了 Anyscale 公司来办理开源社区以及商业化。在 Anyscale 刚举行的 Ray Summit 2023 上,相关数据显示 Ray 已被 OpenAI/Uber/Amazon/字节跳动/蚂蚁金服等众多企业所运用。根据 Ray,Anyscale 也推出了自己的 LLM 相关商业化产品,并以本钱和易用性等方向作为卖点。

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

上图右侧展现了 Ray cluster 的根本架构:

  • 每个框是一个 Ray 的节点,节点是虚拟的概念,比如在 K8s 集群上,每个节点就对应一个 pod。
  • 一切的节点中,有一个节点的角色不同,便是最左边的 head 节点,它能够了解成整个 Ray cluster 的调度中心,head 节点上有 GCS 存储集群节点的信息、作业信息、actor 的信息等等,head 节点上还有 dashboard 等组件。
  • 除了 head 节点以外的都是 worker 节点,worker 节点主要是承载详细的作业负载。
  • 每个节点上有一个 raylet 看护进程,raylet 也是一个本地调度器,担任 task 的调度以及 worker 的办理,一起 raylet 中还有 object store 组件,担任节点之间 object 的传输,整个 Ray cluster 中的一切 object store 构成一个大的分布式内存。

为了供给简练的分布式编程体会,Ray Core 内部做了非常多作业,比如 actor 调度和 object 的生命周期办理等,上图左边展现了怎么运用 Ray Core 编写一个简略的分布式程序,square 函数和 Counter 类经过 Ray 的语法糖,变成了一些在长途运转的目标,其核算进程会被异步调用并存储在 object store 中,最终经过 ray.get 来获取到本地。

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

除了 Ray Core 供给的底层分布式才能,其上层 Ray AI Runtime(Ray AIR)针对算法场景也完成了一系列东西:

  • ray.data 调集了数据读写、流式处理、shuffle 等功用,给离线推理、数据预处理等场景供给了灵敏 API 和异构的调度功用
  • ray.trainray.tune 能够将 xgboost、pytorch 等练习代码快速改写成根据 Ray 的分布式练习运用
  • ray.serve 是一套在线服务的布置调用结构,支撑复杂模型编排,能够灵敏扩缩实例

能够说,Ray 的生态打破了曩昔 AI 工程中每个模块都是固定范式的传统——

在曩昔,说到数据处理,大家会想到 Spark;说到练习,会想到 Torch DDP、MPI;说到推理,会想到 deployment、service;而 Ray 能够给予你满足的自由度和想象力,能够将 AI 的 pipeline 糅合在一个结构甚至一串代码中完成,其强壮异构调度才能以及友爱的上手调试感触。这也是很多 AI 从业者越来越多地挑选 Ray 的原因。

字节跳动 KubeRay Ray 运用实践

KubeRay 简介

KubeRay 是由字节跳动技能团队牵头,由 AnyScale、蚂蚁金服、微软等公司共同参加建设的开源 Ray 布置集成东西集,目前已成为在 Kubernetes 集群上布置 Ray 运用的事实标准。

假如不运用 KubeRay,直接在物理机来保管 Ray 集群会有什么问题呢?

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

首要,head 和 worker 需求直接经过 ip 和 port 衔接,集群的拉起、节点的增删会比较复杂,可康复才能也较弱。其次,RayJob submit 脚本提交作业的方式在大规划生产环境下很难办理,除此之外,也没有 K8s 生态能够给予你的监控、报警、Ingress、HPA/VPA 等才能。

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

KubeRay 采用了经典的 operator 规划,供给了 RayCluster,RayJob,RayService 这三个 CRD:

  • RayCluster:担任 Ray 集群的搭建
  • RayJob:担任提交作业到一个伴生集群中,并同步状况
  • RaySevice:担任将 RayServe 运用快速布置到云原生环境中

在 operator 完成中,cluster 的 controller 更偏重集群的拉起、康复、与 Ray autoscaler 合作等,Job Service 的 controller 偏重作业提交和状况更新,而且它俩分别对应了离线和在线两个典型场景。

除此之外 KubeRay 还供给了 APIServer 等 client 库来担任 CRD 的增删改差,便利对接上层渠道。

RayCluster

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

假如说 Ray 本身供给了 actor 重启、task 重试等才能来增强代码的高可用性,那么 KubeRay 便是真实让 Ray 在集群维度成为真实高可用的运用。

首要 RayCluster CRD 供给了 pod 的康复才能以及集群粒度的热更新,能够非常便利地办理集群;其次 head 和 worker 经过 service 进行衔接,经过将集群 metadata 挂到长途存储中,合作 service 能够做到无感知的 head 节点康复,一起 Ray autoscaler 能够完成根据集群负载动态弹性集群规划,有用减缩本钱。当然 Cluster CRD 还供给了 metric, 集群状况等的透出。

RayJob

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

RayJob 是生产环境办理 Ray 作业的处理计划,支撑批式调度器,创立伴生 Ray 集群或许挑选已有的 Ray 集群,提交作业,并更新作业状况,最终删去 Ray 集群。在字节跳动,咱们优化了作业状况机转移,增加了超时、等待节点数等功用。

RayService

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

RayService 把 CRD 中的 serve 装备布置到集群上,并经过 service 把 serve agent 的端口透出,完成了 Ray serve 的云原生化。它支撑热更新 Serve 装备,经过 pending cluster 的滚动更新完成 Serve 无感知搬迁。

Ray 在字节跳动的保管

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

在字节跳动,咱们给用户供给了丰富的 Ray 相关生态。首要站内一切的 Ray 集群都由 KubeRay 去办理,咱们根据开源版本做了相关适配和增强来支撑大规划作业提交以及一些额定特性;咱们在渠道层支撑用户创立常驻 Ray 集群用来调试作业,也支撑 single-job 方式让渠道保管创立 RayJob;除此之外还供给了渠道鉴权、historyserver、notebook 等周围的才能。

现在字节跳动内部的相关业务包含了图核算离线推理大模型并行核算等方向,涵盖了离线、在线等场景

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

上图展现了站内某业务在运用常驻集群的场景,其需求是期望尽量利用不同 K8s 集群上的低优 spot 资源供给给用户用于运转、调试作业,一起期望大多数作业感知不到外界资源的颤动。

咱们的计划是在每个 K8s 集群中创立一个大资源量的低优 pod 组成的 Ray 集群,operator 层面会根据每天 quota 的规律性浮动,并合作 Ray autoscaler 主动调整集群规划,尽量减少被 K8s 去主动驱逐 pod 的状况。

一起在上层,用户的脚本会感知每个大集群的剩余资源量决议分发到哪个集群去履行。每个集群内部咱们完成了一个简略的排队功用,收到作业恳求后先将作业放入 dashboard 内部的行列中,经过 placement group 来完成资源 gang 调度,保证作业需求的 GPU、CPU 资源到位后才开始真实运转作业。

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

在用户完毕调试之后,能够经过渠道来保管创立 RayJob CR,KubeRay 会担任集群拉起、作业提交、完毕毁掉。作业运转进程中,Ray 集群的重要信息会以 event 的方式 dump 到长途存储,咱们仿照了 spark history server 的规划思路,用户在作业运转完毕之后能够经过 Ray UI 界面来直接检查前史的作业的日志、metric 等信息。

场景案例

场景一: 图核算

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

在图核算场景,咱们运用 Ray Core 来改造字节跳动内部的图核算引擎,每个图算子经过 Ray Actor 拉起,这些算子会根据初始化的 rank 利用 MPI 进行通信。经过 Ray 的分布式才能和 KubeRay 的编排才能,能够完成端到端的容错,假如 worker 挂掉,能够再次被拉起,从 Pmem 或许 SSD 存储中康复 checkpoint 信息。

场景二:大规划离线推理

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

如图所示,上述作业一起包含数据读取处理和模型推理,一起需求消耗很多核算资源做分布式核算。比较在线推理,离线推理对延迟要求不高,但是对吞吐和资源利用率要求很高。咱们运用 Ray dataset 的流式推理才能来处理这个场景,是因为比较 Spark,Ray 的编程更加灵敏,一起将处理和推理放在异构 actor 并 pipeline,能够做流水线并行、模型并行等操作。咱们还增加了 actor pool 扩缩、端到端容错的一些优化。

这些场景都已在 Anyscale 发表过博客,有爱好能够检查:

Kueue 怎么办理/调度 RayJob

跟着作业规划的增大,怎么有限资源下调度不同优先级的作业,让大家都能安稳有序去运用 GPU 等资源是一个非常重要的问题。除了字节跳动所给出的一些经历,在开源社区侧,另一位共享人殷纳(Kante Yin)也介绍了怎么运用 Kueue 这样一个作业调度器去办理 RayJob 来处理这些问题。

作业办理和调度结构 Kueue

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

Kueue 是去年由 K8s 社区发起的作业办理和调度结构,供给作业层面的行列调度,支撑入队优先级、抢占、资源配额等才能。比较其它拥有行列调度才能的开源组件,Kueue 从规划上期望更多复用 K8s 原生的调度才能,尽量不重复造轮子。Kueue 已经原生支撑了 BatchJob、RayJob、TFJob 等作业类型。

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

从 Kueue 的架构来看,ResourceFalvor 供给了节点的笼统,它经过 nodeLabel 的方式与详细的 node 进行绑定。ClusterQueue 是资源池的笼统,界说这个集群总资源量,ClusterQueue 中存在多个 localQueue,它们之间的资源会共享。一个作业会被提交到一个详细的 localQueue 进行调度。不同 clusterQueue 能够经过 Cohort 的机制共享资源。

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

对于办理员,需求创立 ResouceFlavor、ClusterQueue、LocalQueue 来界说资源和机器之间的区分联系,以及资源池中的 quota 分配。

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

用户提交作业,需求指定自己作业所归于的 localQueue,job 在进入 Kueue 中会进入一个挂起状况,排队进程根据 quota,优先级等信息满足需求后放行,假如总当时资源不行,也有可能触发集群规划的 autoscale 机制。

运用 KubeRay 和 Kueue 在 Kubernetes 中保管 Ray 作业负载

KubeCon 活动现场还展现了相关 Demo:两个优先级不同的 queue 中,跟着优先级和 quota 的改变,来触发多个 RayJob 的抢占和康复流程。

共享人简介

胡元哲,火山引擎批式核算团队软件工程师,主要担任字节站表里 Ray、KubeRay 相关生态建设

殷纳,DaoCloud 高级软件工程师,Kubernetes SIG-Scheduling Maintainer & Kueue Approver