菜单

机器学习服务基础设施

相关源文件

目的与范围

本文档概述了 Twitter 的机器学习 (ML) 服务基础设施,该基础设施负责高效地将 ML 模型预测结果交付给推荐系统的各个组件。ML 服务基础设施处理模型的加载、预测请求的处理,并以高性能和高可靠性返回结果。有关特定 ML 模型的详细信息,请参阅排名系统

ML 服务基础设施由两个主要组件组成

  • Navi:一个用 Rust 编写的高性能 ML 服务系统
  • dr_transform:一个数据转换层,用于将数据格式转换为模型可用的格式

来源:README.md22-24 navi/README.md1-3

架构概述

ML 服务基础设施位于模型训练系统和推荐客户端(如 Light Ranker、Heavy Ranker 和 Push Service)之间。它实现了跨不同 ML 框架的高效模型服务,并管理预测请求的数据转换复杂性。

ML 服务架构图

来源:navi/README.md1-10 README.md22-23

Navi 是一个用 Rust 实现的高性能机器学习服务系统,专为 Twitter 的生产环境设计。它提供模型推理能力,并注重性能、稳定性和可用性。

主要功能

  1. 多运行时支持:支持不同 ML 框架的可插拔架构

    • TensorFlow(功能最完善)
    • ONNX Runtime
    • PyTorch(实验性)
  2. gRPC API 兼容性:与 TensorFlow Serving 的 gRPC API 兼容,实现与现有客户端的无缝集成

  3. 性能优化:专注于生产用例高性能的极简设计

Navi 系统组件图

来源:navi/README.md6-9 navi/README.md13-16

数据转换

`dr_transform` 组件是 ML 服务基础设施的关键部分,负责在 Twitter 的内部数据格式与 ML 模型所需的张量格式之间进行数据转换。

关键函数

  • 将 BatchPredictionRequest Thrift 对象转换为 ndarray 格式作为模型输入
  • 与`segdense` 配合使用,以指定如何从 BatchPredictionRequest 中检索特征值
  • 处理不同的输入张量类型(浮点型、整型、字符串等)

数据转换流程图

来源:navi/README.md21-23 navi/dr_transform/Cargo.toml1-23

模型服务工作流

Twitter ML 基础设施中的模型服务工作流包括从模型加载到预测服务的多个步骤

  1. 模型加载:模型从训练管道导出,并由 Navi 的 ModelLoader 加载
  2. 请求处理:客户端服务发送预测请求(通过 gRPC)
  3. 数据转换:BatchPredictionRequests 被转换为张量格式
  4. 模型推理:由相应的运行时(TF/ONNX/PyTorch)执行推理
  5. 响应交付:结果返回给客户端服务

模型以版本化的目录结构组织,每个模型(例如,“web_click”)都有基于时间戳的多个版本。

模型目录结构图

来源:navi/README.md33-38

与客户端服务集成

Twitter 的 ML 服务基础设施被多个客户端服务用于各种推荐任务

客户端服务用途ML 服务组件
轻量级排序器推文候选集的初始评分带有 TensorFlow/ONNX Runtime 的 Navi
重量级排序器已筛选候选集的最终、更详细评分带有 TensorFlow/ONNX Runtime 的 Navi
Push Service通知评分和排名带有 TensorFlow/ONNX Runtime 的 Navi

Light Ranker 执行快速初始评分以筛选出大量候选集,而 Heavy Ranker 则对较小的候选子集执行计算密集型评分。

客户端服务集成图

来源:README.md45-46 pushservice/README.md24-26

运行时部署

Navi 可以通过专用脚本与不同的 ML 框架一起运行

  • `scripts/run_tf2.sh` 用于 TensorFlow
  • `scripts/run_onnx.sh` 用于 ONNX Runtime

生产部署需要创建一个带有版本化模型文件夹的模型目录结构,使用 epoch 时间戳作为版本标识符(例如,1679693908377)。

来源:navi/README.md29-38