MLflow と [ジョブの実行] ページを使用して実行を追跡する

重要

AI ランタイム CLI は ベータ版です。

air runで送信する各ワークロードは、Databricks ジョブ実行と MLflow 実行の両方です。

  • ジョブの実行 (ワークスペースの [ジョブとパイプライン ] ページに表示されます) は、状態、コンピューティング、再試行、ドライバー出力の実行を追跡します。
  • MLflow 実行では、実験 (パラメーター、メトリック、システム メトリック、成果物) が追跡されます。

1 つの申請で、1 つのジョブ実行と 1 つの MLflow 実行が作成されます。 再試行により、新しい MLflow 実行が作成されます。

実験と実行履歴

2 つのワークロード 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 実験が存在しない場合は作成するか、既存の実験に新しい実行を追加します。 1つの実験には複数の実行が含まれます。
  • mlflow_run_name (省略可能): 実行名を設定します。 省略すると、実行名の既定値は実験名 (experiment_name) になります。
  • max_retries (省略可能): 各再試行は同じ実験で実行される新しい MLflow であるため、試行を比較できます。 元の送信とその再試行は、1 つのジョブ実行を共有します。

メトリックを示す MLflow 実行ページ

次の 3 つの場所から実行できます。

  • ジョブ: [ジョブの実行] ページには実行の一覧が表示され、実行ごとに 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 実行の [システム メトリック] タブ (GPU/CPU/メモリ)

カスタム メトリックをログに記録する

プラットフォームは MLflow 実行を作成し、 MLFLOW_RUN_ID 環境変数を使用してその ID をトレーニング プロセスに公開します。 MLflow 追跡 API を使用して、独自のパラメーター、メトリック、成果物をその実行に記録します。

分散 (マルチノード) ワークロードでは、すべてのノードが同じ MLflow 実行を共有します。 ランク 0 プロセスからのみログを記録するため、各メトリックは 1 回記録されます。

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 カタログ ボリュームに書き込みます。 チェックポイント パターンとボリュームの管理については、「 実験の追跡と可観測性」を参照してください。

その他のリソース

  • Last updated on