注释
如果你在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 工作负载:
- 选择笔记本进行交互式单节点开发,或选择Databricks CLI进行单节点或多节点作业。
- 将现有集群资源映射到AI运行时加速器类型和加速器数量。 参见硬件。
- 在笔记本中,用
ray_init()替换现有的集群启动方式,或者在 Databricks CLI 工作负载中,使用文档说明的 head 节点和 worker 节点引导配置。 - 声明工作负载的 Python 依赖,并更新其数据、检查点、建模和输出路径。
- 验证最小适用配置的工作负载。 在测试预期的固定拓扑之前,先确认其数据访问、资源请求和输出。
Limitations
AI 运行时的 Ray 存在以下限制:
- Ray 集群在工作负载期间有固定节点数。 不支持光线簇自动缩放。
- 工作负载中的所有节点都使用相同的
accelerator_type。 Ray 集群不能包含单独的 CPU 和 GPU 工作组,也不能独立扩展 CPU 和 GPU 容量。 - 笔记本会在其所连接的单个节点上启动 Ray。 使用Databricks的CLI来构建多节点Ray集群。
- 目前仅允许通过 Databricks 驱动代理访问 Ray 仪表盘,适用于笔记本工作负载。
Tip
如果 CPU 准备和 GPU 处理可以作为独立阶段运行,可以使用多任务 Lakeflow 作业,并通过 Unity 目录卷在任务之间传递数据。