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文件、需要时的应用代码,以及启动应用的命令。 提交工作量时:
- 工作负载定义会根据支持的模式和配置约束进行验证。
- 如果
code_source配置好,选中的代码会被打包为快照并上传。 - 每个节点都准备好了所需的 GPU 计算资源和环境。
- 工作负载 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 分支或提交。 完整示例请参见 快速入门。
开发和管理工作负载
开发通常遵循以下流程:
- 在本地开发应用程序代码和工作负载 YAML。
- 使用
databricks air run --file <file> --dry-run在本地验证 YAML。 - 通过
databricks air run --file <file>提交工作量。 在运行结束前,将--watch添加到流日志中。 - 使用
databricks air get、databricks air list和databricks air logs来检查工作负载和日志。 - 可选择使用
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 命令组命令。