将 Databricks CLI 与 AI Runtime 配合使用

Important

Databricks 的 AI 运行时 CLI 命令处于 公开预览阶段。 工作区管理员必须先启用 AI 运行时预览,才能提交工作负载。 参见 要求。

databricks air CLI可在按需无服务器GPU计算平台AI Runtime上运行单节点和多节点GPU工作负载。 用 YAML 定义工作负载的计算、环境和命令,然后使用 databricks air run 提交。

AI 运行时负责 GPU 配置和环境设置,提供运行状态、日志、系统指标和 MLflow 跟踪。 YAML 和 CLI 工作流程可与编码代理自然配合,这些代理可以定义、提交和监控工作负载。

使用CLI来:

  • 在本地环境中开发并迭代Python应用。
  • 在可以进行源代码控制的YAML文件中定义可重复的单节点或多节点工作负载。
  • 从终端或自动化开发流程提交和管理工作负载。

典型工作负载包括LLM微调和后训练、分布式训练、批处理推理、Ray应用和超参数搜索。

要提交工作负载,请按照 快速入门 程序操作。

Databricks 的 AI 运行时 CLI 命令符合与 AI 运行时工作负载相同的合规标准。 有关受支持的标准和异常,请参阅 合规性安全配置文件。

Note

databricks air命令和工作负载配置与air包中的传统基于 Python 的databricks-air CLI 不同。 参见 Python CLI 迁移。

探索可运行的示例

《AI 运行时手册》包含可运行的 Databricks CLI 工作负载,用于模型训练、微调、批处理推理、Ray 和超参数搜索。 每个配方都包含工作负载的 YAML 文件、源代码和运行说明。

工作负载的运行方式

你提供工作负载的YAML文件、需要时的应用代码,以及启动应用的命令。 提交工作量时:

  1. 工作负载定义会根据支持的模式和配置约束进行验证。
  2. 如果 code_source 配置好,选中的代码会被打包为快照并上传。
  3. 每个节点都准备好了所需的 GPU 计算资源和环境。
  4. 工作负载 YAML 中指定的命令在每个节点上运行。

提交时返回一个作业运行ID,用于识别工作负载。 用它来检查状态、检查日志和配置,或者取消工作负载。 你还可以查看作业运行页面和MLflow实验,获取执行细节、系统指标、日志和工件。

该工作流程提供托管GPU配置、可复现环境设置和内置可观察性,无需配置或维护GPU基础设施。

工作负载配置

工作负载 YAML 文件描述了以下组件:

  • 实验:用于跟踪工作负载的 MLflow 实验和运行。
  • 环境:用于托管的环境、依赖或 自定义 Docker 镜像 。
  • 计算:要预配的加速器类型和数量。
  • 代码来源:供该工作负载使用的应用代码的可选快照。
  • Command:启动应用程序的shell命令。
  • 运行时设置:重试、超时、环境变量、秘密和权限等选项。

有关配置字段和示例,请参见 工作负载 YAML 参考。

上传应用代码

使用 code_source 上传工作负载的应用文件。 AI 运行时会让上传的目录在每个节点上可用,并将 $CODE_SOURCE_PATH 设置为其位置。 例如,以下命令运行上传目录中的 train.py:

command: python "$CODE_SOURCE_PATH/train.py"

你可以上传当前的工作树,或者选择 Git 分支或提交。 完整示例请参见 快速入门。

开发和管理工作负载

开发通常遵循以下流程:

  1. 在本地开发应用程序代码和工作负载 YAML。
  2. 使用 databricks air run --file <file> --dry-run 在本地验证 YAML。
  3. 通过 databricks air run --file <file> 提交工作量。 在运行结束前,将 --watch 添加到流日志中。
  4. 使用 databricks air get、 databricks air list和 databricks air logs 来检查工作负载和日志。
  5. 可选择使用databricks air cancel停止正在运行的工作负载。

关于所有命令和标志,请参见 air 命令组。

对于计划运行、多任务工作流或跨环境推广,你可以使用 databricks air convert-to-dabs 作为捆绑包的起点。 参见 调度 GPU 工作负载和组合任务。

用 MLflow 跟踪工作负载

每个提交的工作负载都由Databricks作业运行表示,并通过MLflow实验进行跟踪:

  • Databricks 作业运行:在作业运行页面显示工作负载状态、执行尝试和输出。
  • MLflow 实验:包含自动捕获的系统指标、工作负载配置和日志的 MLflow 运行。 你的应用程序还可以使用 MLflow API 记录自定义参数、指标和构件。

使用作业运行 ID 搭配命令如 databricks air get、 databricks air logs、 databricks air cancel。 参见 实验跟踪与可观测性。

运行单节点和分布式工作负载

AI 运行时会在每个已预配的节点上运行一次已配置的命令。 对于分布式工作负载,它提供诸如 NUM_NODES、NODE_RANK、LOCAL_WORLD_SIZE、MASTER_ADDR 和 MASTER_PORT 等协调变量。

该执行模型与框架无关。 你的命令必须启动并配置分布式框架,例如调用 torchrun、Ray 引导脚本或其他框架启动器。 完整工作负载请参见 AI Runtime Cookbook。

使用AI代理

databricks-ai-runtime代理技能提供指导,帮助 AI 代理(如 Claude Code、Cursor 和 Codex)通过 Databricks CLI 与 AI Runtime 协同工作。 用它来创建工作负载YAML、提交GPU训练作业、检查运行状态、流式日志以及取消运行。 它还包含了自定义 Docker 镜像的指导。

安装Databricks的CLI并配置认证后,执行以下命令:

databricks aitools install --skills-only --skills databricks-ai-runtime

按照提示选择你的 AI 助手和安装范围。 有关安装选项,请参见 aitools 命令组命令。

其他资源