RayRTC：大规模分布式计算学习引擎 Ray 在字节跳动 NLP 场景下的实践-优快云博客

动手点关注干货不迷路 👆

Jialin Liu, Mengyuan Chao, Jian Li, Wei Peng, Sixiang Ma, Wei Xu, Run Yang, Xin Chen

RayRTC 是字节基础架构组与字节 AML 组共同合作，在内部 RTC（Realtime Text Classification）文本训练平台上基于 Ray 进行的下一代 Serverless ML 的探索。RTC 文本分类平台是一个一站式的 NLP 服务平台，包括了数据预处理，标注，模型训练，打分，评估，AutoML 以及模型推理等机器学习全流程。目前字节内各大产品，包括抖音，TikTok，头条，西瓜，番茄等都有使用该平台提供的相关自然语言能力。RayRTC 通过算法与系统的协同设计及 Serverless 等技术为 RTC 提供了性能和资源利用率的极致优化，并由此抽象出一套通用的 Serverless ML 框架，目前已在字节内部机器学习平台上部署上线。

RayRTC 的核心计算引擎是 Ray，最早是 UC Berkeley 的一个针对强化学习所设计的大规模分布式计算框架。Ray 的作者 Robert Nishihara 和 Philipp Moritz 在此基础上成立了 Anyscale 这家公司。开源项目千千万，能成功商业化并在硅谷乃至整个 IT 届产生颠覆性影响的凤毛麟角。Anyscale 的创始人中包括 Ion Stoica，这位罗马尼亚籍教授上一家公司是跟他的学生 Matei Zaharia 以 Spark 技术为基础成立的 Databricks 。Spark 和 Ray 分别诞生于大数据和机器学习时代，前者已经在工业界得到广泛应用，后者也逐渐引起越来越多的公司在不同业务场景进行探索。字节美研计算团队自 2020 年末开始接触 Ray，2021 年开始在不同场景小范围试验。RTC 文本分类平台是第一个大规模上线的 Ray 应用场景，在 RayRTC 的设计过程中，有不少第一手的经验值得分享。本文从 RayRTC 所遇到的实际问题出发，对 Ray 在字节的实践进行介绍。

第一次接触 Ray 的读者可能会问，除了明星创始人团队，深度贴近当前 ML 需求的产学研支持，Ray 这套框架到底有哪些吸引人的地方？

首先是以 Ray 为底座可以非常轻松构建完整机器学习完整生态，如下图所示：

机器学习的研发人员往往不仅需要关注算法本身，在实际的生产环境中，各个环节所涉及的工程量和运维量也不容小觑。不少研究表明，工程师们有 80-90%的时间和精力投入在了算法之外的数据处理，全流程打通等。Ray 社区在近几年的演进中，不断吸收业界领先的理念，积极地与其他开源社区和各大厂商进行合作交流。以 Ray 为计算引擎的上层生态的丰富度是别的开源生态中不常见的。比如大数据处理方面，有 Intel 设计的 RayDP，将 Spark 无缝集成到 Ray 中，通过 Ray 的 Actor 拉起 Spark 的 executor，利用 Ray 的分布式调度实现资源细粒度的调控。这样做的好处在于以 Spark 为大数据引擎的机器学习应用中，通过 Ray 可以将 Spark 产生的 dataframe 以 ML Dataset 的形式直接从内存传给下游的机器学习框架，比如 PyTorch。而 Ray 的生态里的其他组件，比如超参训练（Ray Tune）和推理服务（Ray Serve），则进一步补足了训练阶段后续的一系列工程需求。研发人员可以抛开繁琐的上线部署流程，实现一键分布式以及一键部署。

Ray 的另一个显著优势是其简单通用的 API ，只需在一段函数上加入ray.remote 的装饰器，便可将一个单机程序变成分布式执行单元，如下所示：

#declare a Ray task
@ray.remote
def fun(a):
    return a + 1

#submit and execute a Ray task
fun.remote()

#declare a Ray actor
@ray.remote
class Actor():
    def fun(slef, a):
        return a+1
actor = Actor.remote()
#execute an actor method
actor.fun.remote()

Ray 中最基础的概念包括 Task 和 Actor，分别对应函数和类。函数一般是无状态的，在 Ray 里被封装成 Task，从而被 Ray 的分布式系统进行调度；类一般是有状态的，在 Ray 里被映射成一个 Actor。Actor 的表达性更强，能覆盖大多数的应用程序子模块。基于 Actor 和 Task，Ray 对用户暴露了资源的概念，即每个 actor 或 task 都可以指定运行所需要的资源，这对异构的支持从开发人员的角度变得非常便利。比如：

@ray.remote(num_cpus=1, num_gpus=0.2):
def infer(data):
    return model(data)

当 task 在被提交执行的时候，Ray 的调度器会去找到一个满足指定资源需求的节点。在此同时 Ray 会考虑数据的 locality。比如上述例子中的“data”，实际运行中可能会分布在任意一个远端的节点的内存里，如果 task 不在数据所在的节点上执行，跨节点的数据传输就无法避免。而 Ray 可以让这一类的优化变得透明。框架开发人员也可以利用 Ray 的 API 集成更丰富调度策略，最终提供给用户的是非常简单的 API。Ray 对 Actor 和 Task 还有很多高级的细粒度控制特性，比如支持 gang-scheduling 的 placement group 等，在此不一一赘述。

Ray 另外的优势在于：

高效的数据传递和存储：Ray 通过共享内存实现了一个轻量级的 plasma 分布式 object store。数据通过 Apache Arrow 格式存储。
分布式调度：Ray 的调度是 decentralized，每个节点上的 raylet 都可以进行调度；raylet 通过向 gcs 发送 heart beat 获取全局信息，在本地优先调度不能满足的情况下，快速让位给周边 raylet 进行调度。
多语言的支持：目前已经支持的语言包括：Python, Java, C++。后续 go 的支持以及更通用的多语言架构设计也在进行中。

下图是 RayRTC 的一个早期设计规划图和阶段一核心部分（DP+Training）的 Actor 封装流程图。本文着重讲解阶段一，二的设计和实现。其中在阶段一中所用到的核心组件包括 Ray Actor Pool 和 RaySGD 等。

“DP+Training” Actor 化流程图：

其中主要包括 DataProcessing 和 Training 两个 Stage。每一部分的核心计算逻辑都用 Ray 的 API 封装成为 Actor 或 Task。Actor 提交运行后通过 Ray 的调度被放置到合适的节点上执行。Ray 的集群资源可以通过改造后的 Autoscaler 在字节内的 Yarn/K8S 集群上实现动态扩缩容。

DP 实现过程中，我们利用 Ray 的 ActorPool 解决了一个因为创建 Actor 数目过多而导致的 OOM 问题。Actorpool 本身相当于一个线程池，但 Ray 的 Actorpool 可以被开发者拓展为更高阶的弹性线程池。在 RayRTC 中，给定一组数据，我们需要解决的核心问题之一是使用多少 Ray 的 actor 是比较高效的。这里的高效指：资源使用高效，性能较优且稳定性较好（不能 oom）等。最简单的设计方式是 1 对 1，即对于每一个 HDFS 路径, 都指定一个单独的 DP Actor 来进行处理。但当数据量线性增长时，由于缺少内存管控而很容易出现 OOM。最极端的方式是 n 对 1，即用一个 actor，顺序处理所有数据，这样做显然无法发挥