databricks-ai-runtime

Databricks AI Runtime（databricks air）

Databricks AI Runtime 是一组 databricks air CLI 命令，可将 GPU 训练工作负载提交到 Databricks 无服务器计算。它负责管理环境设置、分布式训练配置和工作负载生命周期，无需你管理集群或基础设施。

本技能涵盖三种日常操作：启动运行、查看运行状态以及在训练期间监控运行。

参考规范

CLI 内置帮助是字段名称、GPU 类型、集群上的环境变量和限制条件的权威来源。运行时应优先查阅帮助，而不是依赖记忆。要查询配置字段，可在 run 命令的 -h 参数中传入其路径（路径是一个单独的参数）：

databricks air run -h config                              # list the YAML fields
databricks air run -h config.compute                      # one field
databricks air run -h config.compute.accelerator_type     # a nested field

databricks air <subcommand> --help 会说明各命令的标志。本技能涵盖工作流，以及帮助文档未明确说明的少数事项。

会话设置

检查 CLI 版本，确认命令可用，并列出 Databricks 配置配置文件：

databricks version
databricks air --help
databricks auth profiles --skip-validate    # available Databricks config profiles

示例使用扁平化的 code_source 字段，这些字段适用于 Databricks CLI 1.20.0 及更高版本。对于较旧版本的 CLI，请使用 工作负载 YAML 中展示的等效嵌套快照格式。CLI 1.20.0 及更高版本支持这两种格式。请勿在同一配置中混用嵌套字段和扁平化字段。

每个命令都传入 -p <profile>。脚本或程序化调用时，添加 -o json，以便获取结构化封装，而不是人类可读文本：

• 成功：{"v": 1, "ts": "...", "data": {...}}
• 错误：{"v": 1, "ts": "...", "error": {"code": "...", "kind": "...", "message": "...", "retryable": bool}}

配置验证（输入）错误是例外：即使使用 -o json，它们也会以纯文本形式输出到 stderr，并以非零状态码退出。因此，除了检查解析后的 JSON，还要检查退出码和 stderr。收到结构化错误时，可根据其 kind 判断如何处理：

kind	含义	操作
PERMANENT	未通过身份验证 / 权限被拒绝 / 不可重试	如果是身份验证问题：databricks auth login --host <workspace-url>
NOT_FOUND	运行 ID / 资源无效	检查输入
CONFLICT	幂等键冲突	使用新的 --idempotency-key
TRANSIENT	可重试	重试

子命令

命令	功能
databricks air run	从 YAML 文件提交工作负载。添加 --watch 可持续输出日志，直到运行结束。
databricks air get <JOB_RUN_ID>	显示单次运行的状态、配置和时间信息。
databricks air list	列出你的活动运行（添加 --all-status 可查看已结束的运行；还可添加 --all-users、--filter、--limit）。
databricks air cancel <JOB_RUN_ID> [<JOB_RUN_ID>...] / ... --all	取消一个或多个运行，或取消你所有的活动运行。
databricks air logs <JOB_RUN_ID>	持续输出正在运行的任务的日志，或获取已结束任务的日志。

有关标志和默认值，请运行 databricks air <subcommand> --help。

工作负载 YAML

experiment_name: my-training-job     # alphanumeric, - and _ only
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10        # run `databricks air run -h config.compute` for valid values
environment:
  version: "5"                       # client image version; prefix "databricks_ai_v" to also load the ML venv
  dependencies:
    - mlflow
code_source:
  root_path: "."                     # relative to workload.yaml; extracted to $CODE_SOURCE_PATH on each node
command: |-
  cd "$CODE_SOURCE_PATH"
  python train.py

使用 databricks air run --file workload.yaml -p <profile> 提交。字段详情可通过 databricks air run -h config.<field> 查阅（这是权威来源）。当前支持的加速器类型为 GPU_1xA10 和 GPU_8xH100；请使用 databricks air run -h config.compute 确认。

对于较旧版本的 CLI，仅将上方的 code_source 块替换为：

code_source:
  type: snapshot
  snapshot:
    root_path: "."

字段值和行为保持不变：在 code_source 下添加 type: snapshot，并将所有快照字段（root_path 以及任何 remote_volume、git 或 include_paths）放入 code_source.snapshot。其余工作负载配置保持不变。请运行 databricks air run -h config.code_source，查看已安装的 CLI 支持哪些字段。

配置工作负载

以下可选的顶层字段涵盖常见需求。如需查看其中任意字段的完整参考信息，请运行 databricks air run -h config.<field>。

环境变量和密钥

在 env_variables 下设置普通值；从 Databricks Secrets 的 secrets 下读取敏感值（每个条目将环境变量名称映射到 scope/key，并在工作节点上注入其值，绝不会记录到日志中）：

env_variables:
  BATCH_SIZE: "32"
secrets:
  HF_TOKEN: my_scope/hf_token          # -> $HF_TOKEN on every node
  WANDB_API_KEY: my_scope/wandb_key

作用域和密钥必须已存在；使用 Databricks CLI 创建一次即可：

databricks secrets create-scope my_scope -p <profile>
databricks secrets put-secret my_scope hf_token -p <profile>

密钥环境变量名称不得与 env_variables 中的名称冲突。

参数（超参数）

parameters 是向脚本传递超参数的简便方式，无需修改训练代码即可在不同运行之间更改参数。CLI 会在集群上将它们写入 YAML 文件，并将文件路径作为 $HYPERPARAMETERS_PATH 提供：

parameters:
  learning_rate: 0.001
  batch_size: 32
  warmup_steps: 1000
  use_amp: true

在脚本中读取：

import os, yaml
with open(os.environ["HYPERPARAMETERS_PATH"]) as f:
    hp = yaml.safe_load(f)
lr = hp["learning_rate"]

值可以是字符串、数字、布尔值、列表或嵌套字典。若要扫描超参数，请更改其值后重新提交（并使用新的幂等键）。

重试

失败的尝试会自动重试：max_retries 的默认值非零（使用 databricks air run -h config.max_retries 查看当前值）。设置 max_retries: 0 可在失败时立即终止，不进行重试。设置 timeout_minutes 时，该超时时间适用于每次尝试。

权限

默认情况下，只有你能查看运行。使用 permissions 授予他人访问已提交作业的权限（兼容 DABs 的格式）；权限级别为 CAN_VIEW、 CAN_MANAGE_RUN、CAN_MANAGE：

permissions:
  - group_name: data-engineering
    level: CAN_VIEW
  - user_name: [email protected]
    level: CAN_MANAGE

每个节点上提供的环境变量

databricks air 会在每个节点上自动设置这些变量；可在 command: 和训练代码中使用。使用 databricks air run -h config.command 确认当前列表。

变量	含义
CODE_SOURCE_PATH	节点上已提取的 code_source 的绝对路径（cd $CODE_SOURCE_PATH）。
MASTER_ADDR / MASTER_PORT	用于分布式协调的 rank-0 节点地址/端口。
NUM_NODES	节点总数。
NODE_RANK	当前节点的 rank（0 到 NUM_NODES-1）。
WORLD_SIZE	所有节点上的 GPU 总数。
LOCAL_WORLD_SIZE	此节点上的 GPU 数量。
MLFLOW_RUN_ID	此作业的 MLflow 运行 ID（向其记录指标/制品）。
HYPERPARAMETERS_PATH	parameters YAML 文件的路径（仅在定义了 parameters: 时设置）。

你的 command: 每个节点运行一次，因此平台不会设置进程级 rank（RANK、 LOCAL_RANK）。分布式启动器（torchrun、accelerate）会在为每个 GPU 启动一个进程时分配这些变量；请向其传递 $NODE_RANK / $NUM_NODES / $MASTER_ADDR / $MASTER_PORT。

启动一次运行

1. 将用户的意图转换为 YAML。 从请求中提取模型、数据集路径、GPU 数量和类型，以及训练命令。如果关键信息不明确（数据路径、检查点目录、自定义镜像），请先询问，再生成配置。

2. 编写 command:，使用上面的环境变量。常见的启动方式：

   框架	command: 格式
   纯 torch，单 GPU	python train.py
   torch，多 GPU（单节点）	torchrun --nproc_per_node=$LOCAL_WORLD_SIZE train.py
   torch，多节点	torchrun --nnodes=$NUM_NODES --nproc_per_node=$LOCAL_WORLD_SIZE --node_rank=$NODE_RANK --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT train.py
   HuggingFace Accelerate	accelerate launch --config_file accelerate.yaml train.py
   DeepSpeed	deepspeed --num_gpus=$LOCAL_WORLD_SIZE train.py
   Composer / MosaicML	composer train.py

   多行命令应放在 command: |- 下。通过 $CODE_SOURCE_PATH 访问上传的代码。GPU 总数 = 节点数 x 每节点加速器数量；某种类型的每节点数量可通过 databricks air run -h config.compute 查询，因此 num_accelerators 必须是该数量的有效倍数。

3. 提交前检查所引用的资源，以免作业启动时失败：确认每个 secrets 范围/键都存在（databricks secrets list-scopes / list-secrets <scope>）；如果设置了 environment.unity_catalog_image，还要确认该镜像已推送到同一工作区中的 Databricks Artifact Registry，且提交身份有权读取该镜像（参见 docker-images.md）。不要将自定义镜像与 environment.version 或 environment.dependencies 一起使用。

4. 提交前先试运行。 这会验证配置，但不会启动作业：

   databricks air run --file workload.yaml --dry-run -o json -p <profile>
   

   配置验证错误必须立即停止：错误会以纯文本形式输出到 stderr，且退出码非零，因此应向用户展示该消息并修正 YAML。其他失败属于基础设施/身份验证问题，应向上传递。

5. 提交。 使用稳定的 --idempotency-key，避免意外重新运行导致重复提交（由 <user>-<experiment>-<YYYYMMDD>-v<N> 组成，并在每次重新提交时递增 vN；避免使用 $(uuidgen) 等动态子表达式，以便命令可自动审批）：

   databricks air run --file workload.yaml \
     --idempotency-key <key> -o json -p <profile>
   

   添加 --watch，将日志流式输出到终端，直到运行结束。从响应中获取 data.run_id，然后监控该运行。

查看运行状态

databricks air get <JOB_RUN_ID> -o json -p <profile>   # status, config, timing
databricks air list -o json -p <profile>               # your active runs (add --all-status for finished)

监控运行

运行时流式查看日志。 最简单的方法是在提交时使用 --watch，这样日志会持续流式输出到终端，直到运行进入终态：

databricks air run --file workload.yaml --watch -p <profile>

对于已经启动的运行，可通过 ID 流式查看其日志：

databricks air logs <JOB_RUN_ID> -p <profile>
# --node N (one node), --lines N (last N lines of a finished run),
# --minutes N (last N minutes), --download-to DIR (save all nodes' logs)

损失、GPU 利用率和指标历史记录可在 MLflow 中查看。 CLI 会流式输出日志和状态，不会输出训练指标。提交/状态响应中包含 MLflow 实验和运行 URL；请将其输出，以便用户在 Databricks UI 中查看运行的损失曲线、GPU 利用率和系统图表。

如果运行失败

获取日志，并向用户展示实际错误：

databricks air logs <JOB_RUN_ID> -p <profile>

展示失败的回溯信息，让用户决定如何修复。若要对训练失败进行更深入的引导式诊断，请使用 AI Runtime 自带的调试支持，而不要在此猜测根本原因。

自定义 Docker 镜像

如果工作负载需要特定的系统库、CUDA 扩展（flash-attn、apex、自定义内核）或无法通过 environment.dependencies 满足的依赖项，请使用自定义 Docker 镜像，而不要使用 environment.version。

阅读 docker-images.md，了解以下内容的分步指南：

• 使用 Databricks 提供的基础镜像
• Dockerfile 模式
• 构建前兼容性检查清单（CUDA/驱动程序、PyTorch、NCCL、EFA/RDMA）
• 使用 databricks air images push 推送镜像，并通过 environment.unity_catalog_image 运行镜像