滴滴基于 Ray 的 XGBoost 大规模分布式训练实践

背景介绍

作为机器学习模型的核心代表，XGBoost 在滴滴众多策略算法业务场景中发挥着至关重要的作用。因此，保障并持续提升 XGBoost 模型的离线训练及在线推理稳定性一直是机器学习平台的重点工作。同时，面对多样化的业务场景定制需求和数据规模从万到亿级的跨度，XGBoost 的训练效率和灵活性也成为我们需要重点关注的问题。

由于平台历史架构原因，平台 XGBoost 模型训练仍是开源 XGBoost On Spark 技术方案。但是在易用性和稳定性上一直存在问题：

1. 缺乏任务容错和弹性训练。XGBoost On Spark 不支持 Task 粒度容错，导致在集群负载较高的情况下，任务失败概率非常高。

2. 在亿级数据规模下，Spark 相关性能优化参数组合较多，对于算法同学门槛较高，问题排查也较为困难。

3. 基于 JVM 语言生态实现，与主流的 AI 生态较难整合。XGBoost On Spark 仍然以 Scala 生态为主，大多数算法同学仍主要以 Python 语言为主，对于一些定制化需求（自定义 Loss、自定义评估指标）等场景上手难度较大。

4. XGBoost On Spark 版本升级困难。当前依赖的 XGBoost 仍为0.82，社区 XGBoost 最新已经到2.0，由于各种包版本依赖问题，很多社区新特性或者性能优化不能直接使用。

5. 超参数搜索能力不足。对于有调参诉求的用户，只能够支持围绕 Spark 生态构建调参策略，与业内优秀的调参算法能力集成难度较大。

因此，鉴于以上 XGBoost On Spark 存在的问题，以及平台整体架构方向的综合考虑，我们先后调研了业内多种 MLOps 技术方案。最终，基于公司数据引擎在 Ray 方向的建设情况，我们尝试构建基于 Ray 云原生方案的 XGBoost 训练方式，结合 Ray 强大的 AI 生态，来最大化降低平台策略算法开发及业务侧算法同学维护和使用门槛。Ray 涵盖了常见特征处理、模型训练、模型评估、模型离线预测等功能，可以带来以下直接收益：

1. Ray 原生支持 XGBoost、LightGBM、Tensorflow、Pytorch 等算法实践中常用框架，同时支持一定的容错机制。

2. 拥抱 Ray AI 生态技术栈，包含数据处理、模型训练、AutoML、超参数搜索。完全 Python 化，接入门槛低。

3. 借用 Ray Tune 模块，能快速验证及接入业内更加丰富的超参搜索框架和优秀的算法策略。

4. 借助 KubeRay，支持 Job 级别容器粒度的隔离机制以及训练容器定制化诉求。

最终，经过努力，各项收益中可衡量的关键指标都有不错的提升：

1. 整体训练效率层面，相对原版 XGBoost On Spark 有2-6倍左右的提升。

2. 在稳定性层面，经过压测因框架或者集群问题带来的失败率在（1.6%），远低于原 XGBoost On Spark 的因集群或者框架导致的失败率（4.6%）。

技术方案

整体架构

整体架构图

整体的技术架构如上图所示，整体的流程描述如下：

• 服务端解析训练任务参数构建 RayJob 实例，并通过 watch RayJob 的状态来监听任务的实时状态变化。

• 每个训练任务创建独立的 ray 集群，通过定制化容器，做到 Pod 粒度隔离。

• 采用 Volcano 做资源调度，实现 Gang 调度策略，并结合 Ray autoscaler 的能力，实现任务的动态资源申请。同时借助 gang 调度机制，在保证任务获得足够资源的同时，避免训练多个训练任务同时抢占集群资源，从而导致集群所有任务 hang 住的问题。

• 通过将 Ray job 的日志写入到挂载的宿主机盘上，通过引擎侧提供的 Log Server Daemon 来实时获取训练日志。

XGBoost on Ray 训练性能

作为 XGBoost On Spark 的替代方案，训练性能是我们首先需要考虑的事情。因此，我们对基于 Spark 和基于 Ray 不同滴滴业务场景中不同数量级的不同数据地进行了评测，结果如下图所示。

 显示了基于Ray（红色）和基于Spark（蓝色）的XGBoost 训练效率对比

在测试中我们均使用相近的物理资源量（CPU/内存），得益于 XGBoost 版本的升级以及我们多次优化，从结果上看，Ray 在小规模数据集上的训练效率提升更为明显，在从百万至亿条样本（1M -242M）的不同业务数据中获得了约2-6 倍性能提升。

值得注意的是，图中虽然同一个实验下采用了相同的参数，但不同业务数据集具有不同特性（如特征数、稀疏度），并且在不同业务场景下采用不同最优参数（如树个数、树深度），因此训练时间并不严格随着数据集样本数的增加而线性增长。

XGBoost on Ray 容错能力

支持容错能力是