机器学习平台技术栈之Ray:架构、原理与关键技术剖析

在人工智能与大模型飞速发展的今天,算力的分布式协调和高效调度成为了决定机器学习平台成败的关键因素。作为新一代的分布式计算框架,Ray 已经从一个简单的分布式 Python 库,演进成了复杂的、支撑 OpenAI 等顶级 AI 实验室运转的 AI 基础设施核心组件。

本文将深入解析 Ray 的核心概念、总体架构设计、关键技术细节(如底层通信、分布式对象存储、调度策略等)及其整个生态系统,帮助读者从本质上理解 Ray 的工作原理。

一、 引言:为什么我们需要 Ray?

过去,针对不同的大数据和机器学习任务,往往需要拼接不同的系统:用 Spark 处理数据,用 Horovod 进行分布式训练,用 Celery 执行异步任务,用 Kubernetes 部署在线推理服务。这就导致了严重的“系统碎片化”问题:

  1. 学习和维护多套系统的成本极高。
  2. 不同系统之间的数据交换需要经过磁盘或低效的网络序列化,导致极大的性能损耗。
  3. Python 原生缺乏细粒度的、高性能的分布式原语。

Ray 的诞生正是为了解决这些痛点。它提供了一种通用的 API,使得开发者可以将单机 Python 代码,只需加上简单的装饰器(@ray.remote),就能无缝扩展到拥有数千个节点的集群上。同时它不仅支持无状态的任务(Task),还支持有状态的服务(Actor),从而能用一套底座统一所有的 AI 计算负载。


二、 Ray 的核心概念与编程模型

要理解 Ray,首先要理解它的三个核心原语:Task(任务)、Actor(参与者)和 Object(对象)

1. Task(无状态任务)

Task 是 Ray 调度的基本单位。任何 Python 函数加上 @ray.remote 装饰器后,都会变成一个远程函数。当你调用这个函数(通过 .remote())时,Ray 会在集群的某个节点上异步执行它,并立即返回一个 Object Ref(对象引用),而不是阻塞等待结果。

2. Actor(有状态参与者)

Task 是无状态的,如果需要在多个操作之间保持状态(例如维护一个神经网络的模型权重,或者持有一个数据库连接),就需要使用 Actor。在 Ray 中,Python 类加上 @ray.remote 后,实例化时就会在集群中创建一个常驻的进程,这个进程就是 Actor。对该 Actor 的方法调用也是异步的,并且保证按顺序执行。

3. Object 与 ObjectRef(分布式对象)

无论是 Task 的返回值,还是通过 ray.put() 显式放入集群的数据,都属于 Ray 的 Object。Object 具有不可变性(Immutable),被存储在 Ray 的分布式内存对象存储(Plasma)中。Ray 会返回一个 ObjectRef,允许跨节点安全地引用该数据,而无需主动拷贝。

概念关系连线图

下面通过 Mermaid 关系图来展示这三个概念如何协同工作:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
graph TD
UserCode[客户端/驱动程序] -->|提交无状态计算| Task1(Ray Task)
UserCode -->|提交无状态计算| Task2(Ray Task)
UserCode -->|创建状态副本| Actor[Ray Actor]

Task1 -->|输出| Obj1[(Object 1)]
Task2 -->|输出| Obj2[(Object 2)]

Actor -->|维护内部状态| State{State}
Actor -->|处理请求并输出| Obj3[(Object 3)]

Obj1 -.->|作为输入引用 ObjectRef| Task2
Obj2 -.->|作为输入引用 ObjectRef| Actor

style UserCode fill:#f9f,stroke:#333,stroke-width:2px
style Task1 fill:#bbf,stroke:#333
style Task2 fill:#bbf,stroke:#333
style Actor fill:#bfb,stroke:#333

三、 Ray 的总体架构设计

Ray 的架构设计融合了多种分布式系统的优点,采用了去中心化与中心化结合的思想。整体由三个主要部分构成:Global Control Store (GCS)、Raylet 和 Client/Worker。

1. 架构示意图

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
graph TB
subgraph Head Node [Head Node (头节点)]
GCS[(Global Control Store)]
Raylet_Head[Raylet]
Driver[Driver Process]
end

subgraph Worker Node 1
Raylet1[Raylet]
Plasma1[(Plasma Object Store)]
Worker1_A[Worker Process]
Worker1_B[Worker Process]
Actor1[Actor Process]

Raylet1 <--> Plasma1
Worker1_A <--> Raylet1
Worker1_B <--> Raylet1
Actor1 <--> Raylet1
end

subgraph Worker Node 2
Raylet2[Raylet]
Plasma2[(Plasma Object Store)]
Worker2_A[Worker Process]

Raylet2 <--> Plasma2
Worker2_A <--> Raylet2
end

GCS <--> Raylet_Head
GCS <--> Raylet1
GCS <--> Raylet2

Plasma1 <..->|对象的 P2P 传输| Plasma2

2. 核心组件解析

2.1 Global Control Store (GCS)

GCS 是 Ray 集群的中央元数据数据库(传统由 Redis 实现,后来演进为自研的高可用服务)。它存储了集群级别的信息,比如:

  • 集群中各节点的健康状态和资源可用量。
  • Actor 的位置注册表。
  • 任务的状态信息(用于 Profiling 和监控)。

GCS 使整个架构更易于实现容错:节点挂掉后,新节点只需查询 GCS 即可重建状态。

2.2 Raylet (节点管理器)

每个计算节点上都运行着一个守护进程,称为 Raylet。Raylet 负责节点内的局部资源管理和任务调度,包含两个子组件:

  • **Node Manager (节点管理器)**:接收来自本地 Worker 或其他 Raylet 的任务,并根据本地资源情况分配给本地的 Worker 进程执行。如果本地资源不足,会将任务请求转发(Spillback)给其他节点。
  • **Object Manager (对象管理器)**:管理当前节点的 Plasma 共享内存,并负责在需要时从其他节点拉取对象数据(P2P传输)。

2.3 Plasma Object Store (共享内存对象存储)

这是 Ray 实现高性能的核心组件。它利用操作系统的共享内存机制。当节点上的不同 Worker 需要读取同一个大型数据集时,可以实现真正的“零拷贝”(Zero-copy)读取。这对机器学习中的海量数据及模型预加载场景至关重要。


四、 关键技术细节剖析

在极其简单的 API 背后,Ray 在底层做了大量极致的优化。以下是几个重点技术的剖析。

1. 分布式的自下而上调度策略 (Bottom-up Scheduling)

不同于 Hadoop/Spark 这种标准的“自上而下”(Master 节点垄断全局调度)的模式,Ray 经历了向“自下而上”调度的演进(在性能要求极高的并发请求下,中心调度器容易成为瓶颈)。

  1. 当 Driver 提交一个 Task,它首先提交给本地节点的 Raylet (Node Manager)。
  2. 本地 Raylet 会首先尝试在本地执行这个 Task(检查是否满足 CPU/GPU 资源和依赖的 Object)。
  3. 如果满足,直接调起本地 Worker 执行(0次跨节点网络通信)。
  4. 如果本地无法执行(例如资源耗尽、或依赖数据极大且在远端),本地 Raylet 才会将求助信息上报,或者将 Task 溢出(Spillback)给有负载能力的远端 Raylet。

优势:极大地降低了任务调度的延迟,单机 Task 的提交延迟可以低至数十微秒级别,吞吐量也突破了中心化调度的天花板。

2. Plasma 对象所有权模型 (Ownership Model) 与 分布式 GC

在早期的 Ray 版本中,元数据维护在 GCS 中,导致极大的通信开销。新的 Ray 版本引入了 **所有权模型 (Ownership Model)**:
谁调用了创建任务的 API,谁就“拥有”这个任务的返回值 ObjectRef。

拥有者(Owner)的 Worker 进程负责:

  • 记录 Object 的位置、状态和生命周期。
  • 如果 Task 失败,由 Owner 决定重试逻辑(而不是由某个全局中心决定)。
  • 管理垃圾回收(Garbage Collection)。当作用域内的 ObjectRef 被 Python 层的引用计数清零后,Owner 会通过内部 RPC 通知所在节点的 Raylet,从 Plasma 内存中安全删除该 Object。

如果对象过大导致内存不足怎么办?Ray Object Manager 支持 **Object Spilling (对象溢出机制)**。当共享内存即将满载,它会自动应用 LRU 策略,将最冷的数据异步序列化并持久化到本地磁盘或者云存储(如 AWS S3),对上层应用完全透明。

3. Actor 的生命周期与容错恢复 (Fault Tolerance)

Ray 的容错机制分为数据的容错和计算的容错:

  • Lineage Reconstruction(血统重建):对于无状态的 Task,如果有节点宕机导致 Object 丢失,只要 Owner 还在,Owner 知道它是如何生成的,直接重新提交 Task 重新计算一次即可。
  • Actor 重启机制:当包含 Actor 的节点发生宕机,GCS 会检测到心跳丢失,在集群内其他存活的节点重新调度启动这个 Actor 进程。但这只重建了 Actor 本身,不保证其运行时的内部状态(如成员变量),这需要开发者结合 Checkpoint 机制来实现真正的业务容错。

五、 Ray 的上层生态系统

Ray 的价值不仅在于底层的 Core,更在于其官方构建的丰富 AI 生态,涵盖整个大模型/机器学习的数据飞轮生命周期。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
mindmap
root((Ray Ecosystem))
Data Processing
[Ray Data]
分布式数据读取
流式Map/Filter处理
Model Training
[Ray Train]
桥接PyTorch DDP
桥接DeepSpeed/Horovod
容错训练
Tuning
[Ray Tune]
超参数搜索
ASHA/PBT等高级统筹算法
Serving
[Ray Serve]
模型多路复用
请求批处理 (Batching)
复杂的推理DAG图 (Ensemble)
RL
[RLlib]
强化学习/PPO算法
支撑AutoGPT等底层逻辑
  1. Ray Data:主要负责最后公里的数据加载。它并不取代 Spark 处理复杂的 SQL,而是用来把已经清洗好的特征数据或者大模型语料,极其高效地以 Pipeline 形式喂入 GPU 进行训练。
  2. Ray Train:把 PyTorch 或 TensorFlow 的分布式训练方案进行了一次薄封装。省去了写一大堆由于节点变更导致的 RANK 更新脚本,并原生集成了 Ray 的容错机制。
  3. Ray Serve:专为复杂模型流设计的在线推理框架。在当今的复杂应用中,推理不仅仅是一次前向传播,可能是调用多个大模型(如文生图、文生文、审核模型),组成一个复杂的 DAG 有向无环图。Ray Serve 可以为图上的不同节点分配不同的资源(比如文本节点用 CPU,图像节点用 GPU),实现资源的最优分配。

六、 生产环境实践:KubeRay 的结合

在实际生产中,Ray 虽然能自己管资源,但企业级数据中心的主流是 Kubernetes (K8s)。因此,如今最主流的部署模式是 KubeRay

KubeRay 是一个 Kubernetes Operator,它用 RayCluster 这个 CRD (Custom Resource Definition) 彻底打通了 K8s。

  • 开发者只需提交一个 YAML 表达 “我要 1 个 Head,10 个带有 GPU 的 Worker”。
  • KubeRay 会拉起对应的 Pods。
  • **自动扩缩容 (Autoscaling)**:Ray 的集群调度器检测到还有大量 Task 排队拿不到 GPU 时,会通知 KubeRay。KubeRay 进而申请拉起新的 Worker Pod,通过 K8s Cluster Autoscaler 进一步向云厂商借虚拟机,实现了业务层直接驱动云底座 IaaS 层的秒级弹性。

七、 总结与大模型时代下的展望

Ray 凭借极其精巧的无中心化调度、零拷贝对象存储以及通用抽象原语,成功地一统了 AI 计算平台的碎片化乱象。在 LLM (大语言模型) 时代,从 OpenAI 训练 GPT 系列所使用的海量分布式强化学习(RLHF),到各种多模态模型集群的并行推理部署,几乎都少不了 Ray 的身影。

理解 Ray 的底层不仅有助于写出高性能的并行 Python 代码,更有助于构建出高可用、高吞吐的现代 AI 基础设施系统。随着 AI 技术演进对算力协同要求的上升,Ray 的技术栈将在很长一段时间内作为 AI Cloud-Native 端的基础标准存在。