微信存在大量AI计算的应用场景,主要分为三种:流量分发、产品运营和内容创作。流量分发场景中的 AI 计算主要用于搜索、广告、推荐场景的核心特征生产,产品运营相关的 AI 计算主要用于产品功能相关和内容运营相关(低质、优质、生态建设),由于大模型的兴起,AIGC 相关的文生图、图生图、AI 特效等内容创作场景的 AI 计算也有了较多的落地。目前AI 计算几乎覆盖了微信的所有业务场景。
首先:编写分布式计算既简单又直观。开发者不必了解所有通信和调度细节,也不必对此进行推理。借助 Ray 的简单原语,可以将任何 Python 函数或类转换为分布式执行:只需添加一个装饰器,就大功告成了。Ray 的分布式API 很简单,所有复杂性都由 Ray 的执行框架处理。函数将被安排为无状态任务执行,而类将是一个有状态的远程服务。
如果使用 ray 来做 ocr 推理,只需要添加装饰器@remote,指定模型使用的 cpu 和 gpu 资源数,通过一个python 脚本即可完成ocr应用的部署,效率提升至少一个数量级。
▲ 图 4:Ray AIR 如何以简单的方式统一 ML 库
其次:大多数流行的 ML 库都与 Ray 有很强的集成性,而且 Ray 的原生库也整合了这些库。例如,开发者可以轻松地将 XGBBoost 与 Ray Train 结合使用,可以轻松地将 HuggingFace 与 Ray Serve 结合使用。或者,可以轻松地将 PyTorch 和 TensorFlow 与 Ray Train 结合使用。简而言之,它拥有丰富的集成生态系统,不仅与 ML 库集成,还与其他工具和框架集成。
第三:开发人员可以使用笔记本电脑进行开发。当你想将其扩展到 Ray 集群时,只需更改一行代码或不更改任何代码即可轻松完成。
总的来说,Ray提供了高性能的分布式框架和简单的分布式原语,提供了统一的分布式底盘。Ray融合不同计算范式,与众多开源组件便捷地结合从而实现对现有流程的提效。同时,Ray有完善的生态,数据处理、训练、推理和服务等AI基础设施需要的主流框架都可以很方便地在Ray上进行集成,大量知名企业选用 Ray开发 AI 计算。综上,我们选择了Ray 作为微信 AI 计算平台的分布式底座。
我们基于 Ray 计算底座,解决了上述三个核心问题,构建出适合 AI 计算平台:AstraRay,在微信内进行了大规模 AI 应用部署的实践。
AstraRay 相比社区版本Ray(KubeRay) 有以下改进:
4、AstraRay平台架构概览
▲ 图 7:kuberay 架构
▲ 图 8:KubeRay 提交任务流程
业界使用社区成熟的 KubeRay 方案,通过 Ray 和 K8s 结合,提供了易用、高可用、高伸缩的云原生 Ray 集群服务,可以满足中小规模 AI 应用的需求。但它有集群规模小(最大仅支持数千个节点),异构资源扩展困难(单个 ray 集群只能部署在一个 k8s 集群,不支持联邦k8s 集群)和伸缩慢(受限于 K8s 的扩缩容速度)的问题,不适合微信内超大规模 AI 应用的需求。
▲ 图 9:AstraRay 整体架构
我们在落地 Ray 的过程中遇到了三个核心技术挑战:
1)百万级 pod 的集群管理:在视频号业务场景中,有超过百万核的超级应用,已经远超 K8s 集群上限,我们希望单个 Ray 应用能支持百万级别的 pod 的扩展;
2)不稳定资源下构建稳定服务:由于 AI 计算的资源消耗大,为了降低成本,我们大量使用了低成本、闲置,但稳定性差的计算资源。我们希望可以在不稳定资源上提供可靠稳定的服务;
3)降低应用部署的复杂度:微信内 AI 应用遇到模型、硬件、模块三种维度的异构问题,部署复杂度高。我们希望使用统一的应用维度来简化应用部署,即将 O(n^3) 复杂度降低为 O(1)。
我们支持了Conda作为Python运行环境的隔离和打包,与Ray本身的Conda不同在于:Ray的Conda要先拉起Ray,而 Ray 的worker节点要求和Ray的头节点使用相同的版本,导致应用无法切换Python版本。而我们通过在启动Ray之前初始化运行环境,使每个应用自定义不同的Python版本。
为了提升 Ray 应用的扩展能力,我们通过starlink实现了Ray联邦集群架构,每个Ray应用可以拥有多个Ray集群,单个Ray集群都拥有完整的功能。用户可以调整单个Ray集群的大小,在单个Ray集群内进行Actor的资源分配,提升应用处理能力,提升资源利用率,实现垂直扩展能力;可以通过扩容Ray 集群数量,实现水平扩展。
我们还在 Ray 联邦集群架构基础上,增强了 Ray集群的容灾能力,具体策略为:当head node下线,则水平重新扩容一个Ray集群。当worker node下线,则在这个Ray集群重新拉起一个worker。通过上述策略,我们使用不稳定的低优资源的情况下,Ray自身架构引起的失败影响可以降低到最低。
AI 时代的来临对微信后台的基础设施带来了许多挑战。我们引入业界先进的Ray作为基座,适配了微信的基础环境,提供了方便快捷的AI应用开发范式。同时,在Ray的基础上,简化了Ray本身集群管理的难度,并使用低成本的闲置资源节省了大量的机器成本。AstraRay作为一个刚诞生一年的项目,为微信的AI应用的工程化提供了坚实基础,并且在持续不断的优化,为将来更多AI应用在微信落地做好了准备。