実験の追跡と観察可能性
AIランタイムはMLflowと統合されており、エクスペリメントの追跡、可観測性、デバッグツールを提供します。MLflowは、ノートブックで開発する場合、 databricks air CLIでワークロードを送信する場合、またはジョブをスケジュールする場合でも、ランのパラメーター、メトリクス、GPUシステムメトリクス、Log、およびアーティファクトのための単一の場所を提供します。
MLflowが提供する機能
MLflowエクスペリメントは関連するトレーニングのランをグループ化し、各ランはトレーニング設定と結果の永続的な記録を提供します。
- メトリクスとパラメーター ：トレーニング損失、評価メトリクス、学習率、ハイパーパラメーターを記録し、ラン間で比較します。
- システムメトリクス : システムメトリクス tab でマネージドランの GPU、CPU、およびメモリ使用率を表示します。
- Log view : ジョブベースのワークロードのドライバー出力を、ランの Logs tabで表示します。MLflow Logs tabでログを表示するを参照してください。
- アーティファクトとモデル : モデルファイル、設定、およびその他の出力をランに保存します。大規模または耐久性のあるアーティファクトには Unity Catalog ボリュームを使用します。
- 共有とコラボレーション : エクスペリメントへのアクセス権を付与して、他のワークスペース ユーザーがそのランを表示および比較できるようにします。MLflow エクスペリメントを使用してトレーニング ランを整理するを参照してください。
- フレームワーク統合 ：Hugging Face Transformers、PyTorch Lightning、その他のサポート対象ライブラリとの MLflow 統合を使用します。
ジョブベースのワークロード（ databricks air runで送信されるワークロードなど）の場合、 Databricks ジョブ実行は、実行ステータス、コンピュート、試行、およびドライバー出力を記録することによってMLflow を補完します。
MLflowランの作成方法
MLflow実行を起動または参加する方法は、ワークロードを開始する方法によって異なります：
ワークロードのラン方法 | MLflow ランの起動方法 | エクスペリメントを選択する方法 |
|---|---|---|
Databricks CLI ( | 試行ごとに、ランが自動的に作成されます。 | ワークロードYAMLで |
Databricks AI ランタイム SDK ( | 各呼び出しで、ランが自動的に作成されます。 |
|
アタッチされたノートブックで直接ランされたコード | ランは自動的に作成されません。 | ノートブックで |
スケジュール | タスクによってランが自動的に作成されます。 | タスクの構成で |
コードが正しいランを起動または参加した後、MLflowトラッキングAPIを使用してカスタムトレーニングデータを記録します。サポートされているフレームワークの場合は、個別のロギング呼び出しを追加する代わりに、自動ロギングを使用できます。ディープラーニングのパターンについては、 MLflow 3 ディープラーニング ワークフローを参照してください。
MLflow トラッキングを設定するときは、次のプラクティスを使用します。
stepをグローバルバッチやエポックなどの意味のある値に設定します。すべてのバッチではなく、間隔を空けてログに記録します。MLflow 、ランあたりのメトリクス ステップの数が制限されます。 リソース制限を参照してください。- 分散トレーニングの場合、各値が1回だけ記録されるように、ランク0のメトリクスのみをログに記録します。
- ノートブックのワークフローでは、
/Users/<username>/my-experimentや/Workspace/Shared/<team>/my-experimentなどのエクスペリメントの絶対パスを使用します。共有するエクスペリメントは、共有ワークスペースフォルダに配置します。 - 大きくて耐久性のあるチェックポイントを Unity Catalog ボリュームに保存します。チェックポインティング パターンについては、「 AI ランタイムでのトレーニングのパフォーマンスとレジリエンスの向上」を参照してください。
このページの例では、MLflow 3.7 以降を使用しています。
Databricks CLI ワークロードの追跡
databricks air run でワークロードを送信すると、AI ランタイムは実行詳細の Databricks ジョブランとエクスペリメントデータの MLflow ランを作成します。元の送信とその再試行は 1 つのジョブランを共有しますが、各試行は同じエクスペリメントで別の MLflow ランを作成するため、試行を比較できます。
GPU、CPU、およびメモリ システム メトリクスは、マネージド ランごとに自動的にキャプチャされ、MLflow システム メトリクス tabに表示されます。
構成とランへのログ
ワークロードYAMLを使用してエクスペリメントを選択し、オプションでラン名とアーティファクトの場所を構成します。次の例は、必要なMLflowフィールドを示しています：
experiment_name: my-training
mlflow_run_name: baseline-lr3e5
mlflow_artifact_location: /Volumes/main/default/mlflow-artifacts/my-training
フィールドの要件、default、サポートされている値については、Workload YAML リファレンスを参照してください。
AI ランタイムは、MLFLOW_RUN_IDを介してマネージドランIDを公開します。独自のパラメーター、メトリクス、アーティファクトをそのランにLogとして記録するには、 MLflowトラッキングAPIを使用して再開します。
import os
import mlflow
# Log from rank 0 only because 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)
ランとLogsの検索
Logs は、Databricks CLI と MLflow ラン ページの両方から利用できます。ワークロードの検索、ワークロードの調査、またはログのストリームやdownloadを行うには、以下のコマンドを使用してください。
databricks air list --all-status
databricks air get <job-run-id>
databricks air logs <job-run-id>
databricks air logs <job-run-id> --node 2
databricks air logs <job-run-id> --download-to ./logs/
ノートブックワークロードの追跡
ノートブックから、接続されたセッションでコードを直接ランしたり、 Databricks AI ランタイム SDK の @distributed デコレーターを使用して分散トレーニングをランしたりできます。
ノートブックで直接コードをランする
アタッチされているノートブックで直接実行されたコードは、自動的に MLflow ランを作成しません。エクスペリメントを選択し、ランを起動するか、オートロギングを有効にします：
import mlflow
mlflow.set_experiment("/Users/<username>/my-experiment")
with mlflow.start_run(run_name="baseline"):
train()
PyTorch Lightning オートロギングを有効にするには、次のパターンを使用します。
mlflow.pytorch.autolog()
with mlflow.start_run(run_name="lightning-baseline"):
trainer.fit(model, datamodule=datamodule)
Hugging Face Transformersについては、TrainingArgumentsを次のように構成します。
from transformers import TrainingArguments
args = TrainingArguments(
output_dir="/Volumes/<catalog>/<schema>/<volume>/checkpoints",
report_to="mlflow",
run_name="llama7b-sft-lr3e5",
logging_steps=50,
)
Databricks AI ランタイム SDK を使用する
ノートブックからの分散トレーニングには、Databricks AI ランタイム SDK (databricks-sdk-air) の @distributed デコレーターを使用します。SDKは、GPU環境バージョン6以降ではdatabricks.air名前空間で、GPU環境バージョン5以前ではserverless_gpuとして利用可能です。
from databricks.air import distributed
@distributed()
def train():
# Training code
...
train.distributed()
.distributed()を呼び出すたびに、MLflowランが自動的に作成されます。defaultのエクスペリメントは/Users/{WORKSPACE_USER}/{notebook-name}です。別のエクスペリメントを選択するには、.distributed()を呼び出す前に設定します：
mlflow.set_experiment("/Users/<username>/my-experiment")
train.distributed()
アクティブなラン内で.distributed()を呼び出すと、 AI ランタイムはネストされた子ランを作成します。
with mlflow.start_run() as parent_run:
train.distributed()
以前のランを再開するには、.distributed()の前にMLFLOW_RUN_IDを設定します。
import os
os.environ["MLFLOW_RUN_ID"] = "<previous-run-id>"
train.distributed()
セル出力またはノートブックの エクスペリメント パネルのランLinkを使用して、MLflowランを開きます。
Logsを表示
- ノートブック出力 : トレーニングコードからの標準出力とエラーは、ノートブックのセル出力に表示されます。
- MLflow Log :
.distributed(remote=True)呼び出しのようなジョブバックアップ ノートブック ワークロードの場合、MLflow Log tab に、関連するジョブランの出力が表示されます。
ライブGPUリソースの監視
ノートブックの GPU リソース ペインには、シングルノードおよびマルチノード ワークロードのライブ GPU 使用率、メモリ、温度が表示されます。ランに関連付けられた永続的な記録が必要な場合は、MLflow の システム メトリクス tab を使用します。
ウィンドウを開くには、ノートブックを AI ランタイム に接続し、右側のウィンドウで 
GPU リソース をクリックします。
ペインは10秒ごとにメトリクスをポーリングし、最大2時間の履歴を保持します。
更新 をクリックして最新の値を取得します。5分間アクティビティがない場合、ペインは一時停止します。モニタリングを再開するには、再度開きます。
MLflow Logs tabでLogsを表示する
MLflowランページの [Log] tabには、ジョブによってバックアップされたAIランタイムワークロードの出力が表示されます。これを開くには、ランページに移動し、 Logs をクリックします。
新しいランの場合、[ Logs ] tabは、ランの [ アーティファクト ] tabに以前保存されていたログファイルに代わるものであり、非推奨となっています。以前に作成されたランは、[ アーティファクト ] tabにログファイルを保持します。
[ Logs ] tabから、次のことができます。
- 検索 : 検索文字列を含む行のみを表示します。
- ノード、レベル、時間範囲でフィルター ：マルチノードワークロードの単一ノードを選択したり、情報、警告、またはエラーの行のみを表示したり、表示を特定の時間範囲に絞り込んだりできます。
- 実行中のワークロードをフォローする : ランがアクティブな間、新しい出力が到着すると表示されます。上にスクロールしてフォローを停止する。
- コピーとdownload : ロードされた行をコピーするか、現在のフィルターと一致するすべての行をdownloadします。
Log 保持期間
Logs は、 Logs tab で 60 日間利用可能です。
Logs へのアクセスに関するトラブルシューティング
MLflowラン ページの Log tabには、関連付けられたDatabricksジョブの実行からの出力がストリーミングされます。MLflowエクスペリメントへのアクセス権があっても、ジョブへのアクセス権があるとは限りません。
tabに「 これらのLogへのアクセス権がありません 」と表示される場合は、 [管理可能] 権限を持つユーザーまたはワークスペース管理者に依頼して、少なくともジョブに対する [閲覧可能] 権限を付与してください。ジョブへのアクセスの制御を参照してください。
Databricksのグローバル制限
リソース制限を参照してください。