[大厂实践] Pinterest 基于 Ray 的机器学习数据处理

本文介绍了 Pinterest 如何基于 Ray 满足机器学习数据处理的需求，并为所有 MLOps 组件提供了统一的 AI 运行时，不仅包括数据处理，还包括分布式训练、超参数调整、服务等，并提供了一流的可扩展性支持。原文：Last Mile Data Processing with Ray

Pinterest 的使命是为每个人带去灵感，帮助他们创造热爱的生活。机器学习在这一使命中发挥着至关重要的作用，通过机器学习，可以从平台上数十亿个帖子中精选出高质量灵感，持续提供给 4.6 亿月活用户。在其幕后有数百名机器学习工程师不断改进为 Pinterest 提供支持的各种推荐引擎，使用数百个 GPU 处理 PB 级数据并训练了数千个模型。

最近，我们开始注意到 Pinterest ML 社区有一个有趣的趋势。随着模型架构组件（如 transformers）的标准化，ML 工程师开始表现出越来越强烈的迭代数据集的欲望，包括采样策略、标记、加权以及用于迁移学习和蒸馏的批量推理。

虽然数据集迭代可以产生显著收益，但我们注意到，在过去六个月中，只有少数此类实验得以进行并产品化。这促使我们更深入研究 ML 工程师的开发流程，找出瓶颈，并投资于提高 ML 生命周期中数据集迭代速度的方法。

本文将分享我们对 ML 开发人员效能瓶颈的评估，并深入探讨如何将 Ray（用于扩展 AI/ML 工作负载的开源框架）应用到我们的 ML 平台中，从而将数据集迭代速度从数天提高到数小时，同时将 GPU 利用率提高到 90% 以上。

影响 ML 数据集迭代速度的原因

在 Pinterest，用于推荐模型的 ML 数据集已高度标准化，特征可共享，以 ML 友好的类型表示，并存储在可进行分析查询和大规模训练的 parquet 表中。

然而，即使实现了高度标准化，要快速迭代由数亿用户产生的大规模数据也并非易事。表格中有成千上万个特征，跨越数月的用户参与历史。在某些情况下，PB 级数据流被导入到训练工作中以训练模型。为了尝试新的下采样策略（downsampling strategy），ML 工程师不仅需要找到处理超大规模数据的方法，还需要为生成新数据集付出大量时间。

模式 1：通过工作流模板协调 Apache Spark 作业

Apache Spark 是 ML 工程师处理 PB 级数据最常用的技术之一。ML 工程师使用 Airflow 将一系列 Spark 和 Pytorch 作业串联起来，并将它们打包成 "工作流模板"，以便重复使用，快速生成新的模型训练 DAG。

然而，由于 AI 发展迅速，并非所有数据集的迭代需求都能通过工作流模板得到快速支持，往往需要经过涉及多种语言和框架的漫长过程。ML 工程师必须用 scala / PySpark 编写新的作业并进行测试，必须将作业与工作流系统集成，进行大规模测试、调整，然后发布到生产环境中。而这并不是交互式过程，往往要到最后才能发现错误。

我们发现，在某些情况下，AI 工程师需要花费数周时间，才能利用工作流的新数据集变化训练出模型！这就是我们所说的 "先扩展，后学习" 问题。

模式 2：训练工作中的最后一英里数据处理（Last Mile Data Processing）

由于迭代工作流需要很长时间，一些 ML 工程师开始直接在训练作业中执行数据处理，这就是通常所说的最后一英里数据处理。最后一英里处理可以提高 ML 工程师的工作效率，因为他们可以直接通过 PyTorch 用 Python 编写代码。

不过，这种方法也有挑战。随着 ML 工程师将更多的数据处理工作负载转移到训练任务中，训练吞吐量就会减慢。为了解决这个问题，需要增加更多 CPU 和内存，从而运行更多数据加载器。一旦达到 CPU/内存上限，ML 工程师就会继续通过配置拥有更多 CPU 和内存的昂贵 GPU 机器来纵向扩展机器。这些机器中的 GPU 资源没有得到充分利用，因为训练任务在 CPU 上就遇到了瓶颈。

即使通过分布式训练横向扩展训练工作量，要在训练吞吐量和成本之间找到合适的平衡点也非常具有挑战性。随着数据集越来越大，数据处理逻辑越来越复杂，这些问题就会变得更加突出。为了优化 CPU 和 GPU 资源的使用，需要有能力管理异构类型实例，并以资源感知的方式分配工作量。

解决方案：基于 Ray 进行最后一英里处理

为何选择 Ray

通过对上述两种模式的考察，我们认为横向可扩展的 "最后一英里数据处理"（Last Mile Data Processing）是实现快速高效数据集迭代的方向。理想解决方案应具备三个关键能力：

• 分布式处理：能够在多个节点上高效并行处理大规模数据

• 异构资源管理：能够管理 GPU 和 CPU 等多种资源，确保在最高效的硬件上调度工作负载

• 开发速度快：在一个框架内完成所有工作，用户在创建数据集实验时就不必在多个系统之间切换上下文。

在对各种开源工具进行评估后，我们决定使用 Ray。我们非常高兴看到 Ray 不仅能够满足我们的所有需求，还能为工程师提供独特机会，为所有 MLOps 组件提供统一的 AI 运行时，不仅包括数据处理，还包括分布式训练、超参数调整、服务等，并提供一流的可扩展性支持。

利用 Ray 加速 ML 数据集实验

有了 Ray，ML 工程师在开发过程中首先要启动一个专用的异构 Ray 集群，管理 CPU 和 GPU 资源。这一过程通过统一的训练作业启动器工具实现自动化，该工具还能启动 Ray 驱动程序，管理集群中的数据处理和训练计算。在该驱动程序中，用户还可以调用可编程启动器 API，在多个 GPU 节点上使用 ML 工程师编写的 PyTorch 训练脚本协调分布式训练。

通过在驱动程序中采用 Ray Data，实现了可扩展的最后一英里数据处理。Ray Data 是建立在 Ray 基础上的分布式数据处理库，支持各种数据源和常用数据处理操作符。Ray Data 的关键突破性功能之一是其流式执行能力，使我们能够同时进行数据转换和训练。这意味着：(1) 不需要加载整个数据集就能对其进行处理；(2) 不需要完全完成数据计算就能进行训练。ML 工程师可以在几分钟内收到对其新数据集实验逻辑的反馈。

通过流式执行，可以大大降低 PB 级数据处理所需资源，加快计算速度，并在处理第一个数据块后立即向 ML 工程师提供端到端的即时反馈。此外，为了提高数据处理吞吐量，ML 工程师只需弹性扩展异构 Ray 集群管理的 CPU 资源。

下面的代码段演示了 AI 工程师如何在 jupyter notebook 中交互式尝试 Ray 对训练数据集进行迭代。

基准测试与改进

为了评估将 Ray 用于最后一英里数据处理的优势，我们进行了一组基准测试，在相同的模型架构上训练模型，同时逐步增加最后一英里数据处理的工作负载。

令我们惊讶的是，即使不进行最后一英里数据处理，Ray 数据加载器的训练吞吐量也提高了 20%。与 Torch 数据加载器相比，Ray 数据加载器能更好的处理用户序列特征等超大型特征。

随着数据加载器中加入了更复杂的数据处理和降采样逻辑，改进变得更加明显。在添加了垃圾用户过滤和动态负值降采样后，Ray 数据加载器的速度比基于 Torch 的实现快了 45%。这意味着 ML 工程师现在可以在与以前相同的时间内从训练实验模型中获得 2 倍的学习效果。虽然我们不得不通过增加更多 CPU 节点来横向扩展数据加载器，但缩短的训练时间最终也让我们为这一应用节省了 25% 的成本。

当 ML 工程师通过编写 Spark 作业和工作流执行相同的实验时，他们需要花费 90 个小时来训练新模型。使用 Ray 后，ML 工程师能够将时间缩短到 15 小时，开发速度提高了 6 倍之多！

你好，我是俞凡，在 Motorola 做过研发，现在在 Mavenir 做技术工作，对通信、网络、后端架构、云原生、DevOps、CICD、区块链、AI 等技术始终保持着浓厚的兴趣，平时喜欢阅读、思考，相信持续学习、终身成长，欢迎一起交流学习。为了方便大家以后能第一时间看到文章，请朋友们关注公众号"DeepNoMind"，并设个星标吧，如果能一键三连(转发、点赞、在看)，则能给我带来更多的支持和动力，激励我持续写下去，和大家共同成长进步！