Ray 在 AI 运行时

注释

如果你在Databricks classic计算中使用Ray,配合Databricks运行时机器学习,请参见 Ray在Databricks上的应用。

Ray 是一个用于扩展 Python 工作负载的开放源代码框架。 你可以在 AI 运行时运行 Ray,无需配置或管理底层 GPU 基础设施。

为什么要在 AI 运行时使用 Ray。

Ray 常用于分发 Python 工作负载,如模型训练、批处理推理、数据处理和超参数调优。 在 AI 运行时,Ray 可以访问 Unity 目录卷中的数据,使用 MLflow 跟踪实验,并作为生产工作流程的一部分运行。

你可以在 AI 运行时中使用常见的光线库,如 Ray Core、 Ray Data、 Ray Train 和 Ray Tune 。 现有的 Ray 工作负载可以继续使用库的标准 API。

Ray 如何与 AI Runtime 协同工作

每个 Ray 应用程序都运行在 Ray 集群上,集群可以包含一个或多个节点。 AI 运行时负责节点配置,Ray 则根据可用的 CPU 和 GPU 资源调度任务和角色。 如何启动Ray取决于执行接口:

接口 节点 启动 Ray
Notebook 连接到笔记本的单个节点 从 AI Runtime SDK(databricks.air)调用 ray_init()
Databricks 命令行界面 (CLI) 一个或多个节点的固定集合 在节点 0 上启动头节点,并使用可重用的引导脚本将其余节点作为工作节点加入集群

在笔记本中使用Ray

当你的笔记本连接到 AI 运行时 GPU 计算时,使用 Databricks AI Runtime SDK 中的 ray_init() 在连接的节点上启动 Ray:

from databricks.air import ray_init

ray_init()

ray_init() 调用标准的 ray.init() API,将 Ray 仪表板配置为可通过 Databricks 驱动程序代理访问,并在笔记本输出中打印该仪表板的 URL。 在代码运行时,使用仪表盘检查Ray作业、任务、actor、日志和资源使用情况。

注释

ray_init() 需要环境版本5及以上。 在版本 5 中,请从 serverless_gpu 导入,而不是从 databricks.air 导入。 Databricks AI v6 包含了 Ray。 如果你用的是Standard v6,先安装Ray再调用 ray_init()。 请参阅 设置环境。

笔记本示例

示例 Description
Ray Core Hello World 提交已连接的计算资源上的异步 GPU 任务,在 Ray 仪表板中查看调度情况,并获取结果。
使用 Ray Tune 进行 CIFAR-10 超参数调优 为 PyTorch 图像分类器运行各占用部分 GPU 资源的并发试验,并使用异步逐次减半算法 (ASHA) 提前停止表现不佳的配置。
使用 Ray Data 和 vLLM 进行 Qwen2.5-32B 批量推理 在8个H100 GPU上运行8个持久vLLM副本的多语言批推理,并将结果保存为Parquet,存储在Unity目录卷中。

使用 Databricks CLI 搭配 Ray

Databricks 的 AI 运行时 CLI 命令支持单节点和多节点 Ray 工作负载。 在工作负载配置中,指定固定的 accelerator_type 和总 num_accelerators。 Ray Core 入门教程包含一个可重复使用的引导脚本。 在工作负载中command,设置RAY_ENTRYPOINT为你的 Python 文件路径,然后运行引导脚本。 引导脚本执行以下操作:

  • 在节点 0 上,启动 Ray 头节点并运行已配置的入口点。
  • 在其余每个节点上,它都会启动一个 Ray 工作进程,并在入口点运行期间保持其连接状态。
  • 入口点结束后,它会停止 Ray 进程。

在应用程序中,使用 ray.init(address="auto") 连接到集群。

在标准环境中,将 ray 或所需的额外组件(如 ray[data]、ray[train] 或 ray[tune])添加到工作负载依赖项中。 Databricks AI v6 包含了 Ray。

与Ray一起开始

以下简约示例介绍了主要的Ray库,并共享相同的集群引导模式。

示例 Description
雷核心 在两块A10显卡上启动一个Ray集群,并为每块GPU安排一个Ray核心任务。
雷线列车 在合成数据上训练一个小型多层感知器,每个 GPU 使用一个 Ray Train worker。
光线数据 在Ray集群上运行一个合成数据处理流水线。
雷·图恩 在合成数据上调谐一个小型多层感知器,每个GPU做一次试验。

附加 CLI 示例

示例 Description
ResNet-50 与 Ray Train 微调 通过每个 A10 GPU 在两个节点上配备一个分布式数据并行工作进程,微调 ResNet-50。
Qwen2.5-3B 与 Ray Train 的微调 在一个配备 8 个 H100 GPU 的节点上,使用每个 GPU 配置一个分布式数据并行 worker 来微调 Qwen2.5-3B。
使用 Ray 数据和 vLLM 进行批处理推理 用 Qwen2.5-7B-Instruct 生成回复,并将结果以 Parquet 格式写入 Unity Catalog 卷中。
使用Ray Tune进行LoRA超参数搜索 搜索四块 A10 GPU 上的 Qwen2.5-0.5B LoRA 配置,并用 ASHA 调度器停止性能不佳的试验。
使用Ray Tune进行XGBoost超参数搜索 在四块 A10 GPU 上搜索 XGBoost 超参数,将最佳配置记录到 MLflow,并在 Unity Catalog 中注册模型。

使用数据和 Databricks 功能

Ray 工作负载可以访问 Unity Catalog 卷中的数据,使用 MLflow 记录运行情况,并作为由 Declarative Automation Bundles 定义的 Lakeflow Jobs 任务运行。

访问 Unity 目录中的数据

对于 Unity Catalog 表中的表格数据,可以使用 ray.data.read_databricks_tables() 读取表或通过 Databricks SQL 仓库运行 SQL 查询。 请指定仓库ID、目录和模式,因为AI运行时不提供本地Spark会话。 你可以用SQL在Ray处理结果之前进行连接、聚合和筛选。

对于基于文件和非结构化的数据,可以使用Unity Catalog卷。 Ray 工作进程可以使用其 /Volumes/<catalog>/<schema>/<volume>/... 路径访问文件。 例如,Ray Data 可以使用 ray.data.read_parquet() 读取 Parquet 文件,并使用 Dataset.write_parquet() 写入结果。

跟踪和编排工作负载

使用 MLflow 记录 Ray 工作负载中的参数、指标和工件。 参见 实验跟踪与可观测性。

要调度工作负载,或将其与 CPU 数据准备任务结合使用,可以使用 Lakeflow 作业和声明式自动化捆绑包。 参见 调度 GPU 工作负载和组合任务。

迁移现有的 Ray 工作负载

迁移自管理的Ray工作负载涉及用AI运行时执行模型替换其基础设施配置和服务集成。

评估工作负荷

迁移前,先审查 限制。 确认工作负载能在固定规模的集群上运行,且只有一种加速器类型。

迁移工作负荷

要迁移自行管理的 Ray 工作负载:

  1. 选择笔记本进行交互式单节点开发,或选择Databricks CLI进行单节点或多节点作业。
  2. 将现有集群资源映射到AI运行时加速器类型和加速器数量。 参见硬件。
  3. 在笔记本中,用 ray_init() 替换现有的集群启动方式,或者在 Databricks CLI 工作负载中,使用文档说明的 head 节点和 worker 节点引导配置。
  4. 声明工作负载的 Python 依赖,并更新其数据、检查点、建模和输出路径。
  5. 验证最小适用配置的工作负载。 在测试预期的固定拓扑之前,先确认其数据访问、资源请求和输出。

Limitations

AI 运行时的 Ray 存在以下限制:

  • Ray 集群在工作负载期间有固定节点数。 不支持光线簇自动缩放。
  • 工作负载中的所有节点都使用相同的 accelerator_type。 Ray 集群不能包含单独的 CPU 和 GPU 工作组,也不能独立扩展 CPU 和 GPU 容量。
  • 笔记本会在其所连接的单个节点上启动 Ray。 使用Databricks的CLI来构建多节点Ray集群。
  • 目前仅允许通过 Databricks 驱动代理访问 Ray 仪表盘,适用于笔记本工作负载。

Tip

如果 CPU 准备和 GPU 处理可以作为独立阶段运行,可以使用多任务 Lakeflow 作业,并通过 Unity 目录卷在任务之间传递数据。