Important
這項功能目前處於 公開預覽版。
你使用 air run 提交的每個工作負載同時也是 Databricks 作業執行與 MLflow 執行:
- 工作執行(可在工作區的 Jobs 與 Pipelines 頁面顯示)追蹤執行狀態、計算、重試和驅動程式輸出。
- MLflow 執行會追蹤實驗過程:參數、指標、系統指標與工件。
一次提交會產生一個工作執行和一個 MLflow 執行。 重試會產生新的 MLflow 執行。
實驗和執行
兩個工作負載 YAML 欄位會控制執行在 MLflow 中的顯示方式:
experiment_name: my-training # Creates or appends to this MLflow experiment
mlflow_run_name: baseline-lr3e5 # Names the MLflow run for this submission
compute:
num_accelerators: 8
accelerator_type: GPU_8xH100
command: torchrun --nproc_per_node=8 train.py
max_retries: 2
-
experiment_name(必要):如果不存在具有此名稱的 MLflow 實驗,則會建立一個;否則會將新的執行附加到現有的實驗。 一個實驗可以進行多次運算。 -
mlflow_run_name(可選):設定跑名。 若省略,運行名稱將預設為實驗名稱(experiment_name)。 -
max_retries(可選):每次重試都是同一實驗中的一次新的 ML flow 執行,這樣你可以比較兩次嘗試。 原始提交及其重試共用同一個作業執行個體。
在 Jobs、MLflow 和先前的工作負載之間切換
你可以從三個地方進入跑步:
- Jobs:Jobs 執行頁面會列出你的執行紀錄,而每個執行項目都會連結到其對應的 MLflow 執行與實驗。
- MLflow:實驗頁面列出你的 MLflow 實驗。
-
先前的工作負載:
air get run <job-run-id>會輸出可點擊的連結,指向此次執行的作業、實驗和 MLflow 執行。air list runs它會列出你之前的跑道,並讓你篩選找出特定的跑道。
air get run <job-run-id> # Links to the job, experiment, and MLflow run
air list runs # List previous runs; filter to find a specific run
系統指標
GPU、CPU 與記憶體系統的指標會自動在每次執行時被擷取。 不需要設定。 請在 MLflow 執行的 系統指標 標籤中查看。
記錄自訂指標
平台會建立 MLflow 執行作業,並透過 MLFLOW_RUN_ID 環境變數將其 ID 提供給你的訓練流程。 使用 MLflow 追蹤 API 將你自己的參數、指標和產物記錄到該次執行。
在分散式(多節點)工作負載中,每個節點共享相同的機器學習流程執行。 僅記錄秩-0過程的記錄,因此每個指標記錄一次:
import os
import mlflow
# Log from rank 0 only; all nodes share the same MLFLOW_RUN_ID.
if os.environ.get("RANK", "0") == "0":
with mlflow.start_run(run_id=os.environ["MLFLOW_RUN_ID"]):
mlflow.log_param("learning_rate", 3e-4)
for step, loss in enumerate(training_losses):
mlflow.log_metric("train_loss", loss, step=step)
日誌與文物
使用 air logs 串流或下載某次執行的日誌:
air logs <job-run-id> # Stream logs from node 0
air logs <job-run-id> --node 2 # Logs from a specific node
air logs <job-run-id> --download-to ./logs/ # Download instead of streaming
日誌也會以成品的形式在 MLflow 執行中提供。 要持久化模型檢查點,可以將它們寫入 Unity 目錄卷。 關於檢查點模式與管理體積,請參見 實驗追蹤與可觀察性。