生產化訓練工作負載

Important

這項功能目前處於 公開預覽版。

使用 DAB 來定義 AI 執行時訓練工作負載的程式碼。 將其納入原始碼版本控制、部署到不同環境中、設定排程,並與其他任務組合使用。 本頁說明自備訓練流程:由 ai_runtime_task 在無伺服器 GPU 運算環境中,針對某個程式碼目錄執行你自己的命令。

Tip

  • 使用宣告式自動化套件(Declarative Automation Bundles)將訓練工作負載定義為程式碼,部署於各環境,並排程。
  • ai_runtime_task 會對程式碼目錄執行您自己的命令(自備訓練)。
  • 將 GPU 與 CPU 任務合併成多工工作。

這和用無伺服器 GPU 搭配套件運行筆記型電腦是不同的任務。 如需基本 notebook-on-GPU 組合範例,請參閱使用 Jobs API 和宣告式自動化組合進行排程。

Requirements

  • 一個啟用 AI 執行環境的工作空間。 請參閱 需求說明。
  • 已安裝並設定完成 Databricks CLI(命令列介面),可用來部署 bundles。

在套件中定義 AI 執行階段任務

ai_runtime_task 會為一個實驗命名、使用 code_source_path 指向你的訓練程式碼,並宣告一項部署:要執行的命令,以及執行該命令所使用的 GPU。 將它新增至你組合包中的作業:

resources:
  jobs:
    train:
      tasks:
        - task_key: train
          ai_runtime_task:
            experiment: my-experiment
            code_source_path: ./dist/code.tgz
            deployments:
              - command_path: ./command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1

code_source_path 指向你打包的訓練程式碼,是 command_path 任務執行的腳本。 重試、逾時和權限在任務與作業上的設定方式,與任何 Azure Databricks 作業相同,因此你現有的套件組合做法都可直接沿用。 關於如何打包和引用你的程式碼,請參見 「交付你的訓練程式碼」。

ai_runtime_task 欄位

Field 類型 Description
experiment String Required. 此次執行的 MLflow 實驗名稱。 參見 實驗追蹤與可觀測性。
code_source_path String 要執行的訓練程式碼:封裝後的 tgz 成品輸出檔案,或指向已上傳程式碼的 /Workspace 或 /Volumes 路徑。 請參閱 「寄送你的訓練代碼」。
deployments 排序 Required. 描述該命令及執行該命令所需運算資源的單一部署。 每個項目包含 command_path、compute,以及一個可選的 name。
deployments[].command_path String Required. 任務在每個節點上執行的指令碼。
deployments[].compute.accelerator_type String Required. 例如 GPU_1xA10GPU 類型,例如 、 GPU_1xH100、 或 GPU_8xH100。
deployments[].compute.accelerator_count Integer Required. 所有節點的 GPU 總數——為編碼於 accelerator_type的每個節點數量的倍數。
deployments[].name String 部署的可選名稱,用於日誌和使用者介面。
docker_image_url String 可選用的自訂 Docker 映像檔,用來執行指令,而不是在受管理的環境中執行。 請參見 使用自訂 Docker 映像。
mlflow_run String MLflow 執行作業的選用顯示名稱。
mlflow_experiment_directory String 一個可選的工作區目錄,實驗就是在它下建立的。 必須以 /Workspace開頭。 當以沒有預設使用者目錄的服務主體執行時,請設定此值。
mlflow_artifact_location String MLflow 產物的可選根位置,例如 /Volumes/<catalog>/<schema>/<volume>/… 路徑。 必須與現有實驗的遺物位置相符,否則將被省略。

在任務和作業上設定重試、逾時、權限和環境(environment_key),不要在 ai_runtime_task 內部設定。 完整任務參考請參見 AI 執行時任務。

設定硬體加速器

設定 accelerator_type 你的工作負載所需的 GPU,以及 accelerator_count 總 GPU 數量。 該數量為每個節點的 GPU 數量倍數:對於 GPU_1xA10 和 GPU_1xH100,為 1;對於 GPU_8xH100,為 8。 若計數大於每個節點大小,則任務會跨多個節點執行——例如 GPU_8xH100accelerator_count: 16 ,在兩個節點上執行。 關於選擇加速器的指引,請參見 硬體選項。

備註

對於多節點執行,AI 執行時會在每個節點執行你的指令,並在任務環境中填充標準分散式訓練環境變數——、NUM_NODES、MASTER_ADDRWORLD_SIZELOCAL_WORLD_SIZEMASTER_PORT、 、 。 請從你執行的指令中讀取它們(例如 torchrun 啟動);你不會在套件中設定它們。

設定環境與相依性

在作業上宣告一個 environments 區塊,並在任務中使用 environment_key 參照該區塊。 AI 執行時會在你的指令執行前安裝上述的相依關係:

resources:
  jobs:
    train:
      tasks:
        - task_key: train
          environment_key: default
          ai_runtime_task:
            # experiment, code_source_path, and deployments as above
      environments:
        - environment_key: default
          spec:
            environment_version: '6'
            dependencies:
              - numpy

關於可用的環境,請參見 「設定你的環境」。

寄送你的訓練代碼

code_source_path 告訴任務你的訓練程式碼在哪裡。 它有兩種形式之一:

  • 已封裝的tgz成品 — 宣告成品,並將code_source_path其指向輸出檔。 Azure Databricks 會建立 tar 封存檔並將其上傳至 databricks bundle deploy。 以下說明如何從本機專案目錄或已提交的 Git 修訂版本部署程式碼。
  • 工作區或磁碟區路徑——已上傳並照原樣使用的程式碼。

封裝工件

宣告一個 tgz 工件並指向 code_source_path 其輸出檔案。 在 databricks bundle deploy 上,CLI 會建立 tar 封存檔並將其上傳,而該任務會將其解壓縮,然後對其執行你的指令:

artifacts:
  code:
    type: tgz
    path: .
    include: [src]
    files:
      - source: ./dist/code.tgz
resources:
  jobs:
    train:
      tasks:
        - task_key: train
          ai_runtime_task:
            code_source_path: ./dist/code.tgz

使用 include 將工作樹中的檔案打包,或使用 git 為已提交的分支或提交建立快照。

工作區或磁碟區路徑

要使用已上傳的程式碼,請設 code_source_path 為 a /Workspace/… 或 /Volumes/… path。 Azure Databricks 使用路徑 as-is,且不打包任何東西。

神器場:tgz

Field Description
type tgz 會從原始檔建立 gzip 壓縮的 tarball,而不是執行 build 指令。
path 要打包的基底目錄。 include 路徑以及封存檔中的項目名稱都是相對於它的。
include 要打包的 path 子路徑清單。 略過將 path 的全部內容打包。 遵循.gitignore;整個套件的 sync.include 和 sync.exclude 不適用。 build 指令的替代方案。
git 為已提交的 Git 參照建立快照,而不是工作樹。 設定 git.branch 或 git.commit(若兩者都已設定,則以 commit 為準)。 build 指令的替代方案。
files[].source 建造瀝青球的路徑。 將 code_source_path 對準這個。

備註

AI 執行時會將你的程式碼解壓到一個目錄,並以 CODE_SOURCE_PATH 環境變數的形式暴露。 在指令中引用它,讓相對路徑得以正確解析,例如在執行指令碼前先加入 cd "$CODE_SOURCE_PATH"。

完整範例

此範例是在本地專案的單一 A10 GPU 上訓練,事前沒有任何 CLI 設定,僅安裝並設定 Azure Databricks CLI。 該計畫包含三個檔案:

my-training/
├── databricks.yml
├── command.sh
└── src/
    └── train.py

command.sh 是由 command_path 命名的入口點。 它會切換到解壓出來的程式碼目錄,並執行訓練腳本:

#!/usr/bin/env bash
set -euo pipefail
cd "$CODE_SOURCE_PATH"
python train.py

databricks.yml為 bundle 命名、將 src/ 封裝為 tgz 成品、將其作為 ai_runtime_task 執行、在任務環境中安裝 numpy,並定義開發與生產目標:

bundle:
  name: my-training

artifacts:
  code:
    type: tgz
    path: .
    include: [src]
    files:
      - source: ./dist/code.tgz

resources:
  jobs:
    train:
      name: my-training
      tasks:
        - task_key: train
          environment_key: default
          ai_runtime_task:
            experiment: /Users/me@example.com/my-training
            code_source_path: ./dist/code.tgz
            deployments:
              - command_path: ./command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1
      environments:
        - environment_key: default
          spec:
            environment_version: '6'
            dependencies:
              - numpy

targets:
  dev:
    mode: development
    default: true
  prod:
    mode: production

部署套件並執行工作。 databricks bundle deploy 建立並上傳 tgz 成品,並建立該作業;databricks bundle run 啟動該作業:

databricks bundle deploy --target dev
databricks bundle run train --target dev

建立多任務工作流程

ai_runtime_task 是 Azure Databricks 的作業工作,因此可與作業的其餘部分組合。 你可以在訓練前執行準備步驟,將 GPU 和 CPU 任務合併在同一個工作中,並使用不同加速器。

使用 depends_on 排列任務順序

用 depends_on 來依序執行任務。 以下流程會執行一個準備筆記本,接著執行一個只有在準備任務成功後才開始的 GPU 訓練任務:

resources:
  jobs:
    train_pipeline:
      tasks:
        - task_key: prep
          notebook_task:
            notebook_path: ./prep.py
        - task_key: train
          depends_on:
            - task_key: prep
          ai_runtime_task:
            experiment: my-experiment
            code_source_path: ./dist/code.tgz
            deployments:
              - command_path: ./command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1

結合 GPU 與 CPU 任務

在上述流程中,只有訓練步驟需要 GPU。 將非 GPU 工作如資料準備分開處理,能讓 GPU 時間專注於訓練。

備註

ai_runtime_task 不支援 Azure Databricks 的作業工作值({{tasks.<task_key>.values.<name>}} 或 dbutils.jobs.taskValues)。 要在步驟間傳遞資料,將資料寫入兩個任務都能讀取的共用位置,例如 Unity 目錄卷或工作區檔案,並從每個任務中引用該路徑。

安排工作量

在作業中加入 schedule,讓它依排程定期執行。 發佈時先將排程設為暫停,以免部署套件後自行開始執行;等您準備就緒後,再解除暫停:

resources:
  jobs:
    train_pipeline:
      schedule:
        quartz_cron_expression: '0 0 9 * * ?'
        timezone_id: UTC
        pause_status: PAUSED

從開發到生產的推廣

從開發環境推送到正式環境是標準套件功能,AI Runtime 任務會原封不動地繼承此功能。

套件目標與模式

使用 mode: production 定義生產目標,並搭配您的開發目標。 目標會決定套件部署到何處,以及其資源如何命名:

targets:
  dev:
    mode: development
    default: true
  prod:
    mode: production

部署和運行

使用標準的套件指令,將套件部署到目標並在其上執行:

databricks bundle deploy --target dev
databricks bundle run train_pipeline --target dev

下一步