使用自定义 Docker 映像

Docker 容器服务（DCS）允许将自己的 Docker 容器映像引入 air 工作负载。 如果需要，请使用自定义映像：

• 特定系统库版本。
• 不适合 environment.dependencies 的复杂依赖关系。
• 重现研究成果的确切环境。
• 由组织平台或安全团队构建的标准映像。

先决条件

• 安装 AI 运行时 CLI。
• 对于私有镜像，需要一个有权访问你的镜像的 Docker Hub 账户。

注册镜像

在运行使用自定义镜像的工作负载之前，请先在 air register image 中注册该工作负载。 注册将拉取并缓存 Databricks 平台中的映像。 每个用户必须为每个映像标记注册一次映像。 仅在推送新标签或轮换凭证时才重新注册。 注册需要 2-6 分钟，并且在映像准备就绪之前会阻塞。

公共映像

通过提供 Docker 映像 URL 和 Databricks 配置文件来注册公共映像：

air register image docker.io/nvidia/cuda:12.9.0-devel-ubuntu24.04 -p my-databricks-profile

短格式图像引用也有效。 例如，library/ubuntu:latest。

私有 Docker Hub 镜像

要注册私有 Docker Hub 映像，请首先生成个人访问令牌。 在Docker Hub帐户设置中，单击“个人访问令牌”→“生成新令牌”。 只读访问权限已足够。

选择以下身份验证方法之一：

使用 docker 登录名（建议用于交互式使用）

登录到终端Docker Hub。 系统会提示输入Docker Hub用户名和个人访问令牌：

docker login

这会将你的凭据存储在 ~/.docker/config.json 中。 然后注册该映像——air 会自动读取凭据：

air register image myorg/myrepo:mytag -p my-databricks-profile

使用交互式身份验证

在一个步骤中对 Databricks 机密范围中的凭据进行身份验证和存储：

air register image myorg/myrepo:mytag --interactive-authenticate -p my-databricks-profile

系统会提示输入Docker Hub用户名和个人访问令牌。 凭据存储在你的工作区机密范围中，以备将来注册使用。

使用预先存储的 Databricks 机密（推荐用于 CI/脚本）

将凭据存储在 Databricks 机密中，并直接引用它：

air register image myorg/myrepo:mytag --scope my-secret-scope --key my-docker-key -p my-databricks-profile

在工作负载中使用 Docker 映像

在工作负载 YAML 的 environment.docker_image.url 下指定 Docker 映像：

experiment_name: my-dcs-training
environment:
  docker_image:
    url: myorg/myrepo:mytag
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: python /app/train.py

使用自带的 Docker 映像时，不支持 environment.dependencies 和 environment.version。 在任一字段中指定 environment.docker_image.url 都会触发错误。 如果有其他依赖项，请改为在 Dockerfile 中安装包。

提交工作负荷：

air run --file workload.yaml -p my-databricks-profile

注入到容器的环境变量

AI 运行时在运行时将以下环境变量注入每个容器：

• NUM_NODES — 节点总数。
• LOCAL_WORLD_SIZE — 每个节点的 GPU。
• WORLD_SIZE — 进程总数。
• POD_RANK — 当前节点排名（索引为 0）。 也作为 NODE_RANK 注入。
• LOCAL_ADDR — 本地节点 IP（仅限多节点）。
• MASTER_ADDR — 0 级协调地址（仅用于多节点）。
• MASTER_PORT — 0 级协调端口（仅用于多节点）。

示例

单节点 A10

experiment_name: my-dcs-single-node
environment:
  docker_image:
    url: myorg/myrepo:mytag
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: python3 /app/train.py

支持 RDMA 的多节点 H100

对于需要在 AWS p5 实例上获得完整网络带宽的多节点 H100 作业，请基于一个预配置了 NCCL 和 EFA 的 Databricks 基础映像构建你的映像：

experiment_name: my-dcs-distributed
environment:
  docker_image:
    url: myorg/myrepo:mytag
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 --experimental

Databricks 基础映像

Databricks 在位于 databricksruntime/air 的 Docker Hub 上发布基础镜像，这些镜像已预先配置好 CUDA、NCCL 以及云平台特定的网络配置（AWS EFA 或 Azure InfiniBand）。

除非你的 Dockerfile 需要编译 CUDA 扩展（例如 flash-attn、apex 或自定义内核），否则请使用 runtime 变体。

将 PyTorch 添加到 Databricks 基础镜像的 Dockerfile 示例。 基础映像在 /opt/venv 提供由 uv 管理的 Python。 uv pip install 默认面向该环境;若要使用其他环境，请在运行 uv pip install前创建并激活 venv。

FROM databricksruntime/air:dcs-base-aws-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 /app/train

构建、推送和注册：

docker build -t myorg/myrepo:mytag .
docker push myorg/myrepo:mytag
air register image myorg/myrepo:mytag --interactive-authenticate -p my-databricks-profile

Requirements

• 映像必须托管在Docker Hub上。 不支持 Amazon ECR、Google GCR 和 GitHub GHCR。
• 图像大小必须低于 20 GB。
• WORKDIR 在运行时不生效。 对打包到镜像中的文件使用绝对路径。 例如，使用 python /app/train.py 而不是 python train.py。
• 不能使用 environment.dependencies 或 environment.version 与 environment.docker_image.url. 如果你需要镜像中未包含的额外软件包，则必须将它们添加到 Dockerfile 中。

Troubleshooting

ssl.SSLError: [CRYPTO] 加载依赖项时发生未知错误 (_ssl.c)

当库尝试创建 SSL 上下文时，自定义映像在运行时可能会失败，并出现 OpenSSL 错误，例如：

ssl.SSLError: [CRYPTO] unknown error (_ssl.c:3076)

导入会建立网络连接的库（例如 huggingface_hub）时，会出现此错误，并导致这些库无法加载。

之所以发生这种情况，是因为 air 工作负荷在启用了 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 错误。

另见

• AI 运行时 CLI 快速入门
• 工作负荷 YAML 参考
• AI 运行时 CLI 命令参考