在 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 场景的相关库。
Ray 的GitHub repo 如今已有 27K star,其发起者也成立了 Anyscale 公司来管理开源社区以及商业化。在 Anyscale 刚举办的 Ray Summit 2023 上,相关数据显示 Ray 已被 OpenAI/Uber/Amazon/字节跳动/蚂蚁金服等众多企业所使用。基于 Ray,Anyscale 也推出了自己的 LLM 相关商业化产品,并以成本和易用性等方向作为卖点。
上图右侧展示了 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 来获取到本地。