舊版 Python CLI 的工作負載 YAML 參考

在傳遞至 air run --file 的工作負載 YAML 組態中，定義訓練工作的實驗名稱、運算資源、命令、環境和程式碼來源。 本頁記錄了所有領域。

最小配置

experiment_name: my-training
environment:
  dependencies:
    - mlflow
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: echo "Hello World"

提交時請附上：

air run --file train.yaml -p profile

核心概念

核心欄位

大多數訓練配置包含五個組成部分：

1. experiment_name （必需）：建立或附加於 MLflow 實驗。
2. environment（可選）：Python 相依性與基礎環境版本。
3. compute （必需）：GPU 資源（類型與數量）。
4. command （必備）：用於啟動訓練的 bash 指令或指令。
5. code_source （可選）：可供遠端存取的訓練程式碼路徑。

關於支援的值與欄位限制，請參見 參考文獻。

你的第一份訓練工作

experiment_name: simple-training
environment:
  dependencies:
    - torch
    - transformers
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
command: torchrun --nproc_per_node=8 $CODE_SOURCE_PATH/train.py

在此配置中：

• experiment_name 會建立名為 simple-training 的 MLflow 實驗（如果該實驗已存在，則會新增一個執行）。
• environment 使用預設環境並安裝 torch 和 transformers。
• compute 分配一個 H100 節點（8 顆 H100 GPU）。
• code_source 會將資料夾 repo 上傳到節點，可於 $CODE_SOURCE_PATH 取得。
• command透過 torchrun 在 8 顆 H100 GPU 上執行 train.py。 該檔案在本機位於 /home/username/repo/train.py。

常見使用案例

新增環境變數

experiment_name: training-with-env
environment:
  dependencies:
    - torch
    - transformers
env_variables:
  BATCH_SIZE: '32'
  LEARNING_RATE: '0.001'
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
    git:
      branch: main
command: torchrun --nproc_per_node=8 train.py

使用密鑰（API 金鑰、權杖）

experiment_name: training-with-secrets
environment:
  dependencies:
    - torch
    - transformers
secrets:
  HF_TOKEN: 'my_scope/hf_token'
  WANDB_API_KEY: 'my_scope/wandb'
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
    git:
      branch: main
command: torchrun --nproc_per_node=8 train.py

秘密使用格式 scope/key ，必須在 Databricks 秘密中設定。 請參見 秘密管理 以了解設定。

分享 YAML 範本時，其他使用者必須建立自己的祕密，或具有所引用祕密的存取權限。

環境

用這個environment區塊選擇無伺服器 GPU 環境並安裝 Python 相依關係。 例如，以下配置選擇標準環境版本 4 並安裝 PyTorch 與 Transformers：

environment:
  version: '4'
  dependencies:
    - torch
    - transformers

環境版本

environment.version 為可選選項，並選擇工作負載的管理環境版本。

範例包括：

• "4" 或 "5" 使用對應的標準環境版本。
• "databricks_ai_v5" 可使用 Databricks AI 環境第 5 版，該版本包含預裝的機器學習專用套件。 （完整套裝清單）

以下範例選取了 Databricks AI 環境版本 5：

environment:
  version: 'databricks_ai_v5'
  dependencies: []

如果您指定 environment.version，也必須將 environment.dependencies 作為內嵌清單提供。 如果不需要安裝額外套件，建議使用空清單。

有關 AI 執行環境的資訊，請參閱 「設定您的環境」。

Python 依賴項

請在 environment.dependencies 下方以內嵌清單列出工作負載的 Python 依賴項目。

相依格式

相依清單遵循 Databricks 基礎環境規範。 每個項目都是 pip 風格的封裝規範（例如 my-library==6.1）。 該列表也接受以下條目：

• 需求檔案：使用-r來參照現有的requirements.txt，例如-r '/Workspace/Shared/requirements.txt'。 例如 $HOME 的環境變數會被展開。
• Wheels：.whl 檔案的絕對路徑，例如 /Workspace/Shared/path/to/simplejson-3.19.3-py3-none-any.whl。
• 索引網址：例如 --index-url https://pypi.org/simple，索引網址。

environment:
  version: '4'
  dependencies:
    - --index-url https://pypi.org/simple
    - -r '/Workspace/Shared/requirements.txt'
    - my-library==6.1
    - /Workspace/Shared/path/to/simplejson-3.19.3-py3-none-any.whl

支援的安裝旗標

依賴性是用 UV 安裝的。 以下支援作為清單項目的 pip 風格旗標：

• 套用於整個安裝：--index-url、--extra-index-url 和 --find-links（-f）會設定或延伸套件索引。
• 套用至其後的依賴項：--no-deps、--no-build-isolation、--no-cache-dir和--force-reinstall。 將旗標放在獨立的一行中（或放在規格之前），後面接上其適用的依賴項目。

例如，若要在不進行建置隔離的情況下，針對已安裝的 torch 安裝 flash-attn，且不解析其自身的相依性：

environment:
  version: '4'
  dependencies:
    - torch
    - --no-build-isolation
    - --no-deps
    - flash-attn

自訂 Docker 映像檔

除了 environment.dependencies 之外，你也可以使用 environment.docker_image.url 來指定自訂 Docker 容器映像檔。 environment.docker_image.url 它與 environment.dependencies 和 environment.version 是互斥的——你不能在同一工作負載中使用任何一個。

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

在使用自訂映像檔前，先將其註冊為 air register image。 欲了解完整細節，包括映像需求、Databricks 基礎映像檔及 Dockerfile 模式，請參閱「使用自訂 Docker 映像搭配舊版 Python CLI」。

使用程式碼來源作業

該 code_source 區塊會上傳本地程式碼，讓訓練工作能執行。

• root_path 是快照的本地目錄。 預設情況下，air 會將工作樹按原樣打包為一般 tarball 封存檔（包括任何未提交的變更）。
• 若要改為對已固定的 Git 版本建立快照，請加入包含 branch 或 commit 的 git: 區塊。 這要求 root_path 必須是 Git 儲存庫，並啟用可感知版本的快照功能（快取、git archive）。
• 對大型存放區而言，include_paths 可讓你為其中一個子集建立快照。

最小的例子

experiment_name: simple-training
environment:
  dependencies:
    - torch
    - transformers
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
command: python $CODE_SOURCE_PATH/train.py

在遠端機器上，程式碼會放在 /databricks/code_source/<directory_name>，其中 <directory_name> 是 root_path 的最終路徑元件。 $CODE_SOURCE_PATH 設定為那個絕對路徑，所以請在指令中使用它，而不是硬編碼位置。

Git 存放庫：依分支或提交紀錄釘選

對於 Git 儲存庫，新增 `git:` 區塊即可依分支或提交 SHA 固定程式碼版本。 branch 與 commit 互斥：在區塊中只能指定其中一個。

釘選到分支（使用該分支的本地 HEAD 代碼）：

code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
    git:
      branch: main # Uses local HEAD of main (no remote fetch)
command: train.sh

釘選至 commit SHA（可完全重現）：

code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
    git:
      commit: abc1234567 # Pins specific commit
command: train.sh

主要領域：

• root_path （必填）：你的 git 儲存庫根目錄的本機路徑。
• git.branch （可選）：分支名稱。 使用本地 HEAD；不進行遠端抓取。 與 git.commit互斥。
• git.commit （可選）：特定提交 SHA。 與 git.branch互斥。
• git.remote （可選）：使用分支的遠端 HEAD 而非本地的。 設定為 true 以自動偵測遠端，或設定為遠端名稱（例如 upstream）以從特定遠端抓取。 僅適用於 git.branch。

如果你省略了 git: 區塊，air 會將工作樹打包成一般 tarball，包括任何未提交的變更。 不需要額外的欄位。

非 Git 目錄

你可以快照非 git 倉庫的目錄。 省略 git: 區塊，因為那必須 root_path 是 git 倉庫。 如果沒有它，就不會有版本快取；每次執行都會上傳一個新的 tarball 封存檔。

code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/my_project
command: $CODE_SOURCE_PATH/train.py

使用 include_paths 進行資料夾篩選

對於大型 monorepos，只對特定資料夾進行快照，以減少上傳與下載時間及快照大小：

code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
    include_paths:
      - research/models
      - research/common
      - research/configs
command: python $CODE_SOURCE_PATH/research/models/launch_training.py

關鍵點：

• 這個欄位是可選的。 若省略，則預設包含整個儲存庫。
• 路徑必須相對於儲存庫根節點（無前導 /）。
• .. 不被允許;你不能引用父目錄。

進階功能

自訂超參數

透過 HYPERPARAMETERS_PATH 將結構化設定傳遞給你的訓練指令碼：

experiment_name: parameterized-training
environment:
  dependencies:
    - torch
    - transformers
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
    git:
      branch: main
command: torchrun --nproc_per_node=8 train.py
parameters:
  model:
    name: 'gpt2'
    hidden_size: 768
  training:
    batch_size: 32
    learning_rate: 0.0001

在你的腳本中讀出它們：

import os
import yaml

with open(os.environ['HYPERPARAMETERS_PATH']) as f:
    params = yaml.safe_load(f)

learning_rate = params['training']['learning_rate']
model_name = params['model']['name']

工作可靠性

experiment_name: reliable-training
environment:
  dependencies:
    - torch
    - transformers
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
    git:
      branch: main
command: torchrun --nproc_per_node=8 train.py
max_retries: 2
timeout_minutes: 90

若工作負載失敗，則會重試兩次。 每次嘗試有90分鐘完成，因此牆上計時器總預算為90×3=270分鐘。

成本歸因

透過 usage_policy_name，將工作負載新增至現有預算政策。 當工作負載啟動時，該名稱會解析為政策的 ID。 關於設定，請參見 「屬性使用與無伺服器使用政策」。

experiment_name: my-training
environment:
  dependencies:
    - mlflow
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: echo "Hello World"
usage_policy_name: my team policy

參考文獻

核心場域參考

支援的 GPU 配置

關於加速器的功能與推薦使用情境，請參見 硬體選項。

compute.num_accelerators 是該工作負載的總 GPU 數量。 它必須是所選 compute.accelerator_type 之每個節點 GPU 數量的倍數。

選用欄位

環境配置

environment:
  version: '4'
  dependencies:
    - torch
    - transformers
env_variables:
  BATCH_SIZE: '32'
secrets:
  HF_TOKEN: 'my_scope/hf_token'

關於環境版本、相依格式及支援的安裝標誌，請參見環境。

自訂 Docker 映像設定

environment:
  docker_image:
    url: myorg/myrepo:mytag

與 environment.dependencies 和 environment.version 互斥。 使用前，請先向 air register image 註冊影像。 請參考「使用自訂 Docker 映像搭配舊有的 Python CLI」。

程式碼原始碼配置

code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo # REQUIRED — local path to repo or directory
    git: # Optional (git repos only) — pin to a branch or commit
      branch: main # Branch name; uses local HEAD unless 'remote' is set
      # commit: abc1234567 # Mutually exclusive with 'branch'
      remote: false # Optional — true to auto-detect remote HEAD, or a remote name string
    include_paths: # Optional — filter included paths
      - src/
      - configs/

欄位限制：

• git.branch 與 git.commit 互斥：在區 git: 塊中指定恰好一個。
• git.remote 需要 git.branch （對 git.commit則無影響）。
• 若省略 git: 區塊，工作樹會封裝為一般 tarball，並包含所有尚未提交的變更。

自訂參數

透過 HYPERPARAMETERS_PATH 傳遞至工作負載：

parameters:
  model:
    name: 'gpt2'
    hidden_size: 768
  training:
    batch_size: 32

MLflow 執行名稱

mlflow_run_name: 'experiment-001-baseline'

MLflow 產物位置

設定 mlflow_artifact_location 為將 MLflow 實驗的產物儲存在自訂根位置。 若省略此欄位，新實驗將使用預設的 DBFS 位置，例如 dbfs:/databricks/mlflow-tracking/<experiment-id>/...。

mlflow_artifact_location: /Volumes/main/default/mlflow-artifacts/my-training

如果 DBFS 存取受限，或你偏好 Unity Catalog，請指定 /Volumes/<catalog>/<schema>/<volume>/... 路徑或等效 dbfs:/Volumes/<catalog>/<schema>/<volume>/... 的 URI。 air CLI 會將路徑/Volumes轉換成 dbfs: MLflow 所使用的 URI。

每個實驗都選擇獨特的地點。 MLflow 實驗的產物位置在建立實驗時是固定的。 若 experiment_name 識別為既有實驗，則 mlflow_artifact_location 必須與其成品位置相符，或予以省略。 若要使用不同地點，請指定一個新的實驗名稱。

路徑解析

工作負載 YAML 中的所有路徑都是相對於工作負載 YAML 的，除非它們是絕對路徑。

資料夾結構：

/home/username/my-project/
├── train.yaml
└── scripts/
    └── train.py

YAML 配置：

experiment_name: my-training
environment:
  dependencies:
    - torch
    - transformers
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
code_source:
  type: snapshot
  snapshot:
    root_path: . # Relative to train.yaml
    git:
      branch: main
command: torchrun --nproc_per_node=8 $CODE_SOURCE_PATH/scripts/train.py