适用于 AI Runtime 的 Databricks CLI 快速入门

Important

此功能目前以公共预览版提供。

用于 AI Runtime 的 Databricks CLI 允许你从终端提交、监控和管理无服务器 GPU 工作负载。 本快速入门展示了如何配置和运行 AI 运行时工作负载、流式日志以及检查运行过程。

Requirements

  • 一个启用 AI 运行时的 Azure Databricks 工作区。 参见 要求。
  • 您需要安装最新的 Databricks CLI，并已通过身份验证并连接到您的工作区。 请参阅安装或更新 Databricks CLI 和 Databricks CLI 的身份验证。
  • 已安装 Git，并可在你的PATH上运行 Python 示例。 CLI使用Git来打包本地代码。

确认这些 air 命令在Databricks CLI中可用：

databricks air --help

Tip

要获取编写工作负载配置和管理运行的帮助，可以为你的 AI 智能体安装 databricks-ai-runtime 技能。 参见 使用人工智能代理。

运行最小工作负载

创建一个 workload.yaml 文件，请求一个 A10 GPU，并运行 nvidia-smi 以打印 GPU 名称：

experiment_name: air-quickstart
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: nvidia-smi --query-gpu=name --format=csv,noheader

提交工作负载并使用 databricks air run 流式传输其日志：

databricks air run --file workload.yaml --watch

CLI 返回作业运行 ID 以及指向您工作区中该运行的链接。 --watch 保持终端连接，并流式传输日志，直到工作负载完成。

Submitting experiment: air-quickstart
Submitted workload with Job Run ID: 783940125630184
View job run at: https://example.cloud.databricks.com/jobs/runs/783940125630184?o=1234567890123456

Monitoring run and streaming logs...
...
Environment setup completed at Thu Oct 1 14:23:45 UTC 2026
NVIDIA A10

自己运行Python代码

要上传和运行本地 Python 代码，请创建一个包含以下文件的air-quickstart目录：

air-quickstart/
├── workload.yaml
└── main.py

Tip

要从可运行项目开始，而不是手动创建这些文件，可以克隆 AI Runtime cookbook repository，并切换到 CLI getting-started 目录：

git clone https://github.com/databricks-solutions/ai-runtime-cookbook.git
cd ai-runtime-cookbook/getting-started/air-cli

使用以下最简的 Python 代码验证 AI 运行时是否在 GPU 上运行上传的代码：

import torch

print("Hello from Python on AI Runtime")
print(f"GPU: {torch.cuda.get_device_name(0)}")

为了扩展示例，你可以用自己的训练、批量推理或其他人工智能应用替换 main.py。

对workload.yaml使用以下配置：

experiment_name: python-quickstart
environment:
  version: 6
  dependencies:
    - torch==2.11.0
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xH100
code_source:
  root_path: .
command: python "$CODE_SOURCE_PATH/main.py"

该配置使用以下设置，运行 main.py 在一块 H100 显卡上：

  • environment.version 选择 标准环境版本6。
  • environment.dependencies 在所选环境中安装 PyTorch 2.11.0。
  • root_path: . 上传包含 workload.yaml的目录。 运行时，$CODE_SOURCE_PATH 指向已上传的目录。

从包含 workload.yaml 的目录提交工作负载：

databricks air run --file workload.yaml --watch

流式日志包括：

...
Hello from Python on AI Runtime
GPU: NVIDIA H100 80GB HBM3

要运行完整的训练、微调、批处理推理和 Ray 工作负载，可以克隆 AI Runtime Cookbook 并运行一个示例。

要配置环境、代码源、计算及其他工作负载设置，请参见 工作负载 YAML 参考。

检查一次运行

CLI在提交工作负载时打印作业运行ID。 用它检查运行：

databricks air get <job-run-id>

databricks air get 显示运行状态、配置和时序。 它还链接到作业运行，以及（如果有的话）连接MLflow运行。

要在提交后查看流式日志或恢复已中断的日志流，请运行：

databricks air logs <job-run-id>

CLI还提供命令，用于列出和检查运行、流式传输日志以及取消当前运行。 请参见air命令参考。

Reference

其他资源

  • 上次更新时间为