ワークロード YAML リファレンス
プレビュー
この機能は パブリック プレビュー段階です。
databricks air run -f に渡すワークロード YAML 構成で、トレーニングジョブのエクスペリメント名、コンピュート、コマンド、環境、およびコードソースを定義します。このページでは、オンデマンドの A10 および H100 ワークロードの設定について説明します。
CLI は、ワークロード YAML の検証に使用するスキーマと同じスキーマから構成ヘルプを生成します。インストールされているバージョンの完全なフィールド リストを表示するには、databricks air run -h config をランします。セクションごとの詳細については、databricks air run -h config.<section>（例：databricks air run -h config.environment）を使用します。
最小限の構成
experiment_name: my-training
environment:
dependencies:
- mlflow
compute:
num_accelerators: 1
accelerator_type: GPU_1xA10
command: echo "Hello World"
送信方法：
databricks air run -f train.yaml -p profile
基本概念
必須フィールドにより、エクスペリメント、コンピュートリソース、コマンドが特定されます。オプションのフィールドにより、依存関係、コード、およびランの動作が構成されます。
コア項目
ほとんどのトレーニング構成には、5つのコンポーネントが含まれます。
experiment_name(必須): MLflowエクスペリメントを作成するか、既存のエクスペリメントに追加します。1〜100文字のASCII文字、数字、ハイフン、またはアンダースコアを使用してください。environment(オプション): Pythonの依存関係またはベース環境のバージョン。compute(必須)：GPUリソース（タイプおよび数）。command（必須）：トレーニングの起動に使用する、最大1,000行の空ではないシェルコマンドまたはスクリプト。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_namesimple-trainingという名前のMLflowエクスペリメントを作成します（または、既に存在する場合は新しい実行を追加します）。environmentdefault 環境を使用し、torchおよびtransformersをインストールします。computeH100ノードを1つ割り当てます（H100 GPU 8基）。code_sourceフォルダrepoをノードにアップロードし、$CODE_SOURCE_PATHで利用可能です。commandtrain.pyをtorchrun経由で 8基のH100 GPUで実行します。/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 $CODE_SOURCE_PATH/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 $CODE_SOURCE_PATH/train.py
シークレットには scope/key 形式を使用し、Databricks シークレットで設定する必要があります。設定については、シークレットの管理を参照してください。変数名を env_variables と secrets の両方に含めることはできません。
YAMLテンプレートを共有する際、他のユーザーは独自のシークレットを作成するか、参照されているシークレットにアクセスできる必要があります。
環境
environment ブロックを使用して、Serverless GPU 環境を選択し、Python の依存関係をインストールします。例えば、以下の構成では Standard 環境バージョン 4 を選択し、PyTorch と Transformers をインストールします：
environment:
version: '4'
dependencies:
- torch
- transformers
環境バージョン
environment.version はオプションであり、ワークロードの管理対象環境バージョンを選択します。
例は次のとおりです。
"4"または"5"を使用して、対応する Standard 環境バージョンを使用します。"databricks_ai_v5"Databricks AI 環境バージョン 5 を使用します。これには、ML 固有のパッケージがプレインストールされています。(パッケージの全リスト)
次の例では、Databricks AI環境バージョン5を選択します：
environment:
version: 'databricks_ai_v5'
dependencies: []
environment.dependencies environment.version を指定する場合はオプションです。追加のパッケージが必要ない場合は、省略するか空のリストを使用します。依存関係を提供する場合は、requirements ファイルへのスカラパスではなく、文字列のリストを使用してください。
AIランタイムで利用可能な環境の詳細については、「環境の設定」を参照してください。
Python の依存関係
ワークロードのPython依存関係を、environment.dependenciesの下にインラインリストとして記述します。
依存関係の形式
依存関係リストは、Databricksベース環境仕様に従います。各エントリはpip形式のパッケージ指定です（例：my-library==6.1）。リストには、次のエントリも登録できます。
- 要件ファイル ：
-rを使用して既存のrequirements.txtを参照するもので、たとえば-r '/Workspace/Shared/requirements.txt'があります。環境変数（$HOMEなど）は展開されます。 - ホイール： ファイルへの絶対パス
.whl/Workspace/Shared/path/to/simplejson-3.19.3-py3-none-any.whl(例:)。 - インデックスURL ：例えば、
--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
--trusted-host サポートされていません。uv はインデックス URL ごとに信頼を構成するため、代わりに --index-url または --extra-index-url を使用してください。
カスタム Dockerイメージ
管理された環境の代わりとして、environment.unity_catalog_imageを使用してArtifact RegistryからカスタムDockerコンテナーイメージを指定します。この値は、レジストリのホスト名なしで<catalog>.<schema>.<image>:<tag>形式を使用します。environment.unity_catalog_imageは、空の依存関係リストを含めて、environment.dependenciesおよびenvironment.versionの両方と排他的です。
experiment_name: my-dcs-training
environment:
unity_catalog_image: main.ml.training:v1
compute:
num_accelerators: 1
accelerator_type: GPU_1xA10
command: python /app/train.py
カスタムイメージを使用する前に、ワークスペース内でワークロードの送信に使用する Artifact Registry にそのイメージをプッシュしてください。アーティファクト Registry の起動およびAI ランタイムでのカスタム Docker イメージの使用を参照してください。
ソースコードの操作
code_source ブロックがローカルコードをアップロードすることで、トレーニングジョブがそれを実行できます。
root_pathはスナップショットを取得するローカルディレクトリです。git:ブロックがない場合、Databricks CLI は Git の無視ルールを遵守しながら、コミットされていない変更を含めてワークツリーをパッケージ化します。- 代わりにコミットされた Git バージョンのスナップショットを作成するには、
git:またはbranchを指定したcommitブロックを追加します。ディレクトリは Git リポジトリ内にある必要があります。root_pathがサブディレクトリを指している場合、そのサブツリーのみがパッケージ化されます。 - 大規模なリポジトリでは、
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 リポジトリの場合は、Branch またはcommit SHAによってコードのバージョンをピン留めするためにgit:ブロックを追加します。branch と commit は排他的です。ブロック内にどちらか一方のみを指定してください。リビジョンがローカルに存在する必要があります。CLI はリモート リポジトリから取得しません。
ブランチにピン留め (そのブランチのローカルHEADを使用します)：
code_source:
type: snapshot
snapshot:
root_path: /home/username/repo
git:
branch: main # Uses local HEAD of main (no remote fetch)
command: bash $CODE_SOURCE_PATH/train.sh
コミットSHAにピン留め（正確な再現性）：
code_source:
type: snapshot
snapshot:
root_path: /home/username/repo
git:
commit: abc1234567 # Pins specific commit
command: bash $CODE_SOURCE_PATH/train.sh
キーフィールド：
root_path（必須）：リポジトリへのローカルパス、またはスナップショットを取得するサブディレクトリ。git.branch（オプション）：Branch のローカル HEAD を使用します。選択したパス内のコミットされていない変更は、その commit の一部ではないためエラーが発生します。git.commit（オプション）：特定のローカルで利用可能なcommitを使用します。コミットされていない変更は含まれません。git.remoteこのフィールドを省略するか、falseに設定します。trueまたはリモート名を使用したリモートフェッチはサポートされていません。ワークロードを送信する前に、リビジョンをローカルにフェッチしてください。
git: ブロックを省略した場合、Databricks CLI は、コミットされていない変更を含め、無視されるファイルを除外して作業ツリーをパッケージ化します。追加のフィールドは必要ありません。
Gitリポジトリ以外のディレクトリ
Git リポジトリではないディレクトリのスナップショットを作成できます。git: ブロックを省略します。CLI は、Git の除外ルールを尊重しながらディレクトリをパッケージ化します。ワーキングツリーと Git ピン留めされたスナップショットの両方で、スナップショットキャッシュキーが変更されていない場合にuploadされたアーカイブを再利用できます。
code_source:
type: snapshot
snapshot:
root_path: /home/username/my_project
command: python $CODE_SOURCE_PATH/train.py
フォルダ フィルタリング include_paths
大規模なモノレポでは、アップロードとダウンロードの時間、およびスナップショットのサイズを削減するため、特定のフォルダーのみをスナップショットしてください:
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
キーポイント：
- このフィールドはオプションです。省略した場合、CLI は上記で説明したスナップショット モードを使用して
root_path以下のファイルを選択します。空のリストを設定しないでください。 - パスは
root_pathからの相対パスにする必要があり、先頭に/を含めないようにしてください。 ..許可されていません。親ディレクトリを参照することはできません。
スナップショットをボリュームにupload
defaultでは、CLI はスナップショットをワークスペースにuploadします。代わりに Unity Catalog ボリュームを使用するには、code_source.snapshot.remote_volume に /Volumes/ で始まるパスを設定します。
code_source:
type: snapshot
snapshot:
root_path: .
remote_volume: /Volumes/main/ml/training-code
ボリュームへのアクセス権と、そこにファイルを書き込むための権限が必要です。
高度な機能
次のフィールドを使用して、トレーニングパラメーター、リトライ動作、コスト配分を構成します。
カスタム ハイパーパラメーター
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 $CODE_SOURCE_PATH/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 $CODE_SOURCE_PATH/train.py
max_retries: 2
timeout_minutes: 90
max_retries: 2 初回試行後に最大 2 回の再試行が許可されます。The default is 3.再試行を無効にするには、max_retries: 0 を設定します。
timeout_minutes: 90 送信されたジョブランに 90 分のタイムアウトを設定します。再試行ごとに個別の 90 分間のバジェットが割り当てられるわけではありません。値は 1 以上である必要があります。省略した場合は、バックエンドのdefaultが適用されます。
コストアトリビューション
usage_policy_name を使用して、ワークロードを既存の Serverless 使用ポリシーにアタッチします。ワークロードの起動時に、名前がポリシーの ID に解決されます。または、usage_policy_id を既存のポリシーの UUID に設定します。これらのフィールドは相互に排他的です。ポリシー名には1～127文字を使用する必要があります。セットアップについては、Serverless使用ポリシーでの属性の使用を参照してください。
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構成には、以下の表を使用してください。インストールされているCLIで受け付けられる完全なスキーマを確認するには、databricks air run -h configを実行してください。
コアフィールドリファレンス
フィールド | Type | 説明 | 例 |
|---|---|---|---|
| string | 必須の MLflow エクスペリメント名。1 ～ 100 文字の ASCII 文字、数字、ハイフン、またはアンダースコア。 |
|
| string | ランによって記録されたMLflowアーティファクトのルート場所。オプション。 |
|
| list | 依存関係仕様のオプションのリスト。 |
|
| 文字列または整数 | マネージド環境バージョン。オプション。省略した場合は default が使用されます。環境のバージョンを参照してください。 |
|
| int | GPU の数。選択した |
|
| string | GPUタイプやノード形状を含むアクセラレータ構成。サポートされている GPU 構成を参照してください。 |
|
| 辞書 | コードソース構成。 | ソースコードの操作を参照してください。 |
| string | 最大 1,000 行の、必須かつ空ではないシェルコマンドまたはスクリプト。 |
|
サポートされているGPU構成
CLIは、アクセラレータ名を大文字と小文字を区別して一致させます。可用性とクォータはワークスペースによって異なります。
| ノードあたりのGPU |
| 注 |
|---|---|---|---|
| 1 | 任意の正の整数 | シングルA10は、開発および小規模なワークロードに適しています。 |
| 1 |
| 単一のH100。 |
| 8 | 8の正の倍数 | H100フルノード。分散トレーニングの一般的な構成です。 |
| 8 | 8の正の倍数 | フル B300 ノード、GPU ごとに 288 GB の HBM。パブリックプレビュー。AWSのみ。 |
アクセラレータの機能と推奨されるユースケースについては、ハードウェアオプションを参照してください。
compute.num_accelerators は、ワークロードの GPU 合計数です。選択した compute.accelerator_type のノードあたりの GPU 数の倍数である必要があります。
任意フィールド
これらのフィールドは、必要なエクスペリメント、コンピュート、コマンドに加えてランを構成します：
フィールド | Type | 制約と動作 |
|---|---|---|
| 文字列のマップ | プレーンな環境変数。 |
| 文字列のマップ |
|
| map |
|
| 整数 | 負ではない再試行回数。 |
| 整数 | ジョブのラン タイムアウト（分単位）。 |
| string | 送信内容を重複排除するための、64文字以下の空ではないトークン。 |
| string | 1〜100文字のASCII文字、数字、ハイフン、またはアンダースコアで構成されるラン名。default は |
| string | MLflow エクスペリメントのワークスペース ディレクトリ。 |
| string |
|
| オブジェクトのリスト | 各付与には、空でない |
| string | 1～127文字の既存の使用ポリシー名。 |
| string | 既存の使用ポリシーの UUID。 |
環境設定
environment:
version: '4'
dependencies:
- torch
- transformers
env_variables:
BATCH_SIZE: '32'
secrets:
HF_TOKEN: 'my_scope/hf_token'
環境バージョン、依存関係の形式、サポートされているインストールフラグについては、環境を参照してください。
カスタムDockerイメージ
environment:
unity_catalog_image: main.ml.training:v1
このフィールドは、environment.dependencies および environment.version と相互に排他的です。使用する前に、アーティファクトをArtifact Registryにプッシュしてください。AIランタイムでのカスタムDockerイメージの使用をご覧ください。
ラン権限
1 エントリにつき 1 つのプリンシパルを使用して、送信されたジョブへのアクセス権を付与します。例：
permissions:
- group_name: training-team
level: CAN_VIEW
- user_name: trainer@example.com
level: CAN_MANAGE
MLflow エクスペリメント ディレクトリ
共有ワークスペースディレクトリにエクスペリメントを保存します：
mlflow_experiment_directory: /Workspace/Shared/training-experiments
ソースコード構成
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 # Uses the branch's local HEAD
# commit: abc1234567 # Mutually exclusive with 'branch'
remote: false # Optional; remote fetching is not supported
include_paths: # Optional — filter included paths
- src/
- configs/
フィールドの制約：
git.branchとgit.commitは相互に排他的です:git:ブロック内でいずれか1つを指定します。git.remoteを省略するか、falseに設定します。CLIは、リモートからリビジョンを取得しません。git:ブロックを省略した場合、選択したファイルのコミットされていない変更を含め、Git の無視ルールを遵守しながら作業ツリーがパッケージ化されます。
カスタム パラメーター
HYPERPARAMETERS_PATH を介してワークロードに渡されるもの：
parameters:
model:
name: 'gpt2'
hidden_size: 768
training:
batch_size: 32
MLflow実行名
1 ～ 100 文字の ASCII 文字、数字、ハイフン、またはアンダースコアを使用してください。省略した場合、ラン名は experiment_name に default 設定されます。
mlflow_run_name: 'experiment-001-baseline'
MLflowアーティファクトの場所
MLflowエクスペリメントのアーティファクトをカスタムルートの場所に保存するには、mlflow_artifact_location を設定します。このフィールドを省略した場合、新しいエクスペリメントは dbfs:/databricks/mlflow-tracking/<experiment-id>/... などの default DBFS の場所を使用します。
mlflow_artifact_location: /Volumes/main/default/mlflow-artifacts/my-training
DBFS アクセスが制限されている場合、または Unity Catalog を優先する場合は、/Volumes/<catalog>/<schema>/<volume>/... パスまたは同等の dbfs:/Volumes/<catalog>/<schema>/<volume>/... URI を指定します。Databricks CLI は、/Volumes パスを MLflow が使用する dbfs: URI に変換します。
各エクスペリメントに固有の場所を使用してください。MLflowエクスペリメントのアーティファクトの場所は、エクスペリメントの作成時に固定されます。experiment_name が既存のエクスペリメントを識別する場合、mlflow_artifact_location はそのアーティファクトの場所と一致しているか、省略されている必要があります。別の場所を使用するには、新しいエクスペリメント名を指定してください。
パスの解決
相対的な code_source.snapshot.root_path の値は、ワークロード YAML ファイルのディレクトリからの相対パスとして解決されます。include_paths のエントリは root_path から解決されます。command 内のパスは、ローカルマシンではなくリモートランタイム上のファイルを指します。uploadされたコードを参照するには、$CODE_SOURCE_PATH を使用します。
フォルダ構造：
/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