AI 运行时可以运行存储在工件注册表中的自定义 Docker 容器镜像。 如果需要，请使用自定义映像：
- 无法使用
environment.dependencies安装的系统库或复杂依赖项。 - 在开发、研究和生产中实现可重复的环境。
- 由你的平台或安全团队创建的私密、组织批准的镜像。
Prerequisites
- 安装或更新到 Databricks CLI 版本 1.19.0 或更高。 请参阅安装或更新 Databricks CLI。
- 请让工作区管理员启用 AI 运行时测试版功能 预览。 有关操作说明，请参见 管理工作区级预览。
- 在同一工作区设置工件注册表，并获得推送和读取镜像的权限。 参见“ 开始使用文物登记册”。
- 在本地机器上安装并启动Docker。
将镜像推送到 Artifact Registry
在使用带有 AI 运行时的自定义镜像之前，先将其存储在你提交工作负载的同一工作区的工件注册表中。
为该映像创建或选择一个 Unity Catalog 目录和架构。
按照工件注册表入门配置 Docker 身份验证、授予所需特权并推送映像。
请注意图片的Unity Catalog名称格式如下：
<catalog>.<schema>.<image>:<tag>例如，
main.ml.training:v1。 工作负载配置中不要包含注册表主机名。
Tip
或者，可以使用 Databricks CLI 中的databricks air images push帮助程序命令。
在工作负载中使用 Docker 镜像
在工作负载 YAML 的 environment.unity_catalog_image 下指定镜像的 Unity Catalog 名称：
experiment_name: my-dcs-training
environment:
unity_catalog_image: main.ml.training:v1
compute:
num_accelerators: 1
accelerator_type: GPU_1xA10
command: python /app/train.py
此示例在映像中从 /app 运行 train.py。 如需单独上传应用代码而不重建镜像，请参见 “运行已上传的应用代码”。
使用自带的 Docker 映像时，不支持 environment.dependencies 和 environment.version。 在任一字段中指定 environment.unity_catalog_image 都会触发错误。 如果有其他依赖项，请改为在 Dockerfile 中安装包。
提交工作负荷：
databricks air run -f workload.yaml -p my-databricks-profile
配置文件必须向存储映像的同一工作区进行身份验证。
注入到容器的环境变量
AI 运行时在运行时将以下环境变量注入每个容器：
-
CODE_SOURCE_PATH：配置完成code_source后，上传应用码的路径。 -
NUM_NODES：节点总数。 -
LOCAL_WORLD_SIZE：每个节点的GPU数。 -
WORLD_SIZE：进程总数。 -
POD_RANK：当前节点排名（0索引）。 也作为NODE_RANK注入。 -
LOCAL_ADDR：本地节点 IP（仅适用于多节点）。 -
MASTER_ADDR：秩-0的协调地址（仅限多节点）。 -
MASTER_PORT：rank-0 协调端口（仅限多节点）。
Examples
以下示例展示了如何运行上传的应用代码和分布式训练，并使用自定义镜像。
运行已上传的应用代码
用 code_source 来单独上传应用码和图片。 你可以编辑并重新提交代码，而无需重建镜像。 在镜像中安装代码的 Python 和系统依赖。
将 train.py 放在 workload.yaml 旁边的本地 src 目录中。 以下配置会上传 src，并在自定义镜像中运行其 train.py：
experiment_name: my-dcs-uploaded-code
environment:
unity_catalog_image: main.ml.training:v1
compute:
num_accelerators: 1
accelerator_type: GPU_1xA10
code_source:
root_path: ./src
command: |-
cd "$CODE_SOURCE_PATH"
python3 train.py
root_path 相对于 workload.yaml 进行解析。 AI 运行时将 $CODE_SOURCE_PATH 设置为容器中上传目录的路径。 请参见code_source，了解代码源选项。
支持 RDMA 的多节点 H100
对于需要在 AWS p5 实例上获得完整网络带宽的多节点 H100 作业，请基于一个预配置了 NCCL 和 EFA 的 Databricks 基础映像构建你的映像：
experiment_name: my-dcs-distributed
environment:
unity_catalog_image: main.ml.training:v1
compute:
num_accelerators: 16 # 2 nodes × 8 H100
accelerator_type: GPU_8xH100
command: |-
torchrun \
--nnodes="${NUM_NODES}" \
--nproc_per_node="${LOCAL_WORLD_SIZE}" \
--node_rank="${POD_RANK}" \
--rdzv_endpoint="${MASTER_ADDR}:${MASTER_PORT}" \
/app/train.py
生成自己的映像
构建自定义映像时，Databricks 建议配合编码代理使用 databricks-ai-runtime 技能，或从 Databricks 基础映像开始。
使用编码代理
安装 databricks-ai-runtime Claude Code 技能以获取分步 Dockerfile 指南，包括从头开始生成、CUDA/NCCL/EFA 兼容性、常见问题和预生成清单。 此技能需要 Databricks CLI 1.0.0 或更高版本。
databricks aitools install --skills databricks-ai-runtime
Databricks 基础映像
Databricks 在位于 databricksruntime/air 的 Docker Hub 上发布基础镜像，这些镜像已预先配置好 CUDA、NCCL 以及云平台特定的网络配置（AWS EFA 或 Azure InfiniBand）。
| 标记 | Variant | CUDA | 何时使用 |
|---|---|---|---|
dcs-base-azure-runtime |
Runtime | 12 | 仅安装预构建的 wheel 包 |
dcs-base-azure-devel |
Devel | 12 | 编译 CUDA 扩展（需要 nvcc） |
下面的 Dockerfile 文件将 PyTorch 添加到 Databricks 基础镜像中。 基础映像在 /opt/venv 提供由 uv 管理的 Python。
uv pip install 默认针对该环境。 要使用不同的环境，先在运行 uv pip install前创建并激活一个venv。
要将你的训练脚本包含在图片中，请在 Dockerfile 旁边放置 train.py 。 Dockerfile 将其复制到 /app/train.py。 如果你上传带有 code_source的应用代码，省略COPY该指令，而是保留train.py在源目录中。
FROM databricksruntime/air:dcs-base-azure-runtime
RUN uv pip install --no-cache \
torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0
RUN uv pip install --no-cache \
transformers==4.45.0 \
accelerate==0.34.0 \
'mlflow>=3.6'
COPY ./train.py /app/train.py
本地构建镜像：
docker build --platform linux/amd64 -t my-training-image:v1 .
然后按照开始使用 Artifact Registry中的说明，为映像添加标记并将其推送到 Artifact Registry。 在工作负载 YAML 中，将生成的 <catalog>.<schema>.<image>:<tag> 名称用作 environment.unity_catalog_image。
Tip
或者，使用 Databricks CLI 中的 databricks air images pushhelp 命令，并按照交互提示操作。
Limitations
- 镜像必须存储在您提交工作负载所在工作区的 Artifact Registry 中。
- 图像大小必须低于 20 GB。
-
WORKDIR在运行时不生效。 对打包到镜像中的文件使用绝对路径。 例如，使用python /app/train.py而不是python train.py。 - 不能使用
environment.dependencies或environment.version与environment.unity_catalog_image. 如果你需要镜像中未包含的额外软件包，则必须将它们添加到 Dockerfile 中。
Troubleshooting
对于与 Artifact Registry 相关的身份验证、权限、镜像推送或镜像发现错误，请参阅 排查 Artifact Registry 问题。
加载依赖项时出现 ssl.SSLError
当库尝试创建 SSL 上下文时，自定义映像在运行时可能会失败，并出现 OpenSSL 错误，例如：
ssl.SSLError: [CRYPTO] unknown error (_ssl.c:3076)
导入会建立网络连接的库（例如 huggingface_hub）时，会出现此错误，并导致这些库无法加载。
这是因为AI运行时工作负载运行在启用联邦信息处理标准（FIPS）的主机上。 当映像的加密库不符合 FIPS 时，OpenSSL 无法在 FIPS 模式下初始化，因此创建 SSL 上下文会失败。
建议的解决方案：
对于 FedRAMP、CMMC 或 HIPAA 审核，企业、政府、医疗保健和金融工作负载通常依赖于 FIPS 140-2 或 140-3 合规性。 如果工作负荷必须符合 FIPS，请使用符合 FIPS 的加密库生成映像。
如果工作负荷不需要 FIPS 符合性，可以通过将 OPENSSL_FORCE_FIPS_MODE 环境变量设置为 0禁用 FIPS 模式。 这样做可能会无提示地中断合规性要求。
若要禁用 FIPS 模式，请在工作负荷 YAML 下 env_variables将其设置为：
env_variables:
OPENSSL_FORCE_FIPS_MODE: '0'
或者，在 Dockerfile 中设置变量，使其适用于使用该映像的每个工作负荷：
ENV OPENSSL_FORCE_FIPS_MODE=0
重新提交工作负荷并确认依赖项加载时不再显示 SSL 错误。