使用自訂的 Docker 映像檔搭配 AI 執行環境

AI 執行時可以執行儲存在 Artifact Registry 中的自訂 Docker 容器映像檔。 需要時使用自訂圖片：

• 無法使用 environment.dependencies 安裝的系統函式庫或複雜相依性。
• 一個涵蓋開發、研究與生產的可重複環境。
• 由你的平台或資安團隊建立的私密、組織核准映像檔。

先決條件

• 安裝或更新到 Databricks CLI 版本 1.19.0 或更高。 請參閱安裝或更新 Databricks CLI。
• 請工作區管理員啟用 AI 執行時測試版功能 預覽。 相關說明請參見 「管理工作區層級預覽」。
• 在同一工作區設置 Artifact Registry，並取得推送與讀取映像檔的權限。 請參閱「開始使用文物登錄」。
• 在你的本機上安裝並啟動 Docker。

將映像檔推送到 Artifact 登錄檔

在使用 AI Runtime 的自訂映像檔之前，請先將其儲存在您用來提交工作負載的同一個工作區中的 Artifact Registry。

1. 為映像建立或選取 Unity Catalog 的目錄和結構描述。

2. 請按照 Artifact Registry 入門 中的步驟設定 Docker 驗證、授予所需權限，並推送映像檔。

3. 請注意圖片的 Unity 目錄名稱格式如下：

   <catalog>.<schema>.<image>:<tag>
   

   例如： main.ml.training:v1 。 工作負載設定中不要包含登錄主機名稱。

在工作負載中使用 Docker 映像檔

請在你的工作負載 YAML 中指定映像檔的 Unity 目錄名稱：environment.unity_catalog_image

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

此範例會在映像檔中執行來自 train.py 的 /app。 若想單獨上傳應用程式碼而不重建映像檔，請參見 「執行已上傳的應用程式碼」。

使用您自己的 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： rank-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:
  type: snapshot
  snapshot:
    root_path: ./src
command: |-
  cd "$CODE_SOURCE_PATH"
  python3 train.py

root_path 會相對於 workload.yaml 進行解析。 AI 執行時會設定 $CODE_SOURCE_PATH 到已上傳目錄在容器中的路徑。 請參閱 「使用程式碼來源」 以獲得快照選項。

多節點 H100 搭配 RDMA

對於需要在 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 --experimental

Databricks 基礎映像檔

Databricks 在 Docker Hub databricksruntime/air 上發布基礎映像，並預先設定好 CUDA、NCCL 及雲端專用網路（AWS EFA 或 Azure InfiniBand）。

以下 Dockerfile 會將 PyTorch 新增至 Databricks 基礎映像檔。 基礎映像於 /opt/venv 提供 Python，並由 uv 管理。 uv pip install 預設針對該環境。 要使用不同的環境，先在執行 uv pip install前建立並啟用 venv。

要將你的訓練腳本包含在圖片中，請放在 train.py Dockerfile 旁邊。 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。

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。

載入相依性時發生 ssl.SSLError

當函式庫嘗試建立 SSL 上下文時，自訂映像檔在執行時可能會因 OpenSSL 錯誤而失敗，例如：

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

錯誤會在匯入會開啟網路連線的函式庫（例如 huggingface_hub）時發生，並導致這些函式庫無法載入。

這是因為 AI 執行時工作負載運行在啟用聯邦資訊處理標準（FIPS）的主機上。 當映像的密碼函式庫不符合 FIPS 標準時，OpenSSL 無法在 FIPS 模式下初始化，因此建立 SSL 上下文會失敗。

推薦解決方案：

企業、政府、醫療及財務工作負載常依賴 FIPS 140-2 或 140-3 合規以進行 FedRAMP、CMMC 或 HIPAA 審計。 如果你的工作負載必須保持 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 錯誤不再出現。

其他資源

• Databricks CLI 適用於 AI 執行階段的快速入門
• 工作負載 YAML 參考
• air 命令群組