レガシー Python CLI の Workload YAML リファレンス
このドキュメントは廃止されており、更新されない可能性があります。
databricks-air パッケージでインストールされる Python ベースの air CLI は非推奨となり、今後は積極的にメンテナンスされません。
新しいワークロードには Databricks CLI を使用してください。AI ランタイムでの Databricks CLI の使用方法については、Use the Databricks CLI with AI Runtime を参照してください。
air run --file に渡すワークロードの YAML 構成で、トレーニング ジョブのエクスペリメント名、コンピュート、コマンド、環境、コードソースを定義します。このページではすべてのフィールドについてドキュメント化しています。
YAML 設定のグラウンドトゥルースは CLI 内のヘルプです。トップレベルビューについては air -h config を実行し、セクションごとの詳細については air -h config.<section>(例:air -h config.environment)を実行します。
最小限の構成
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
コアコンセプト
コアフィールド
ほとんどのトレーニング構成には、5 つのコンポーネントが含まれます。
experiment_name(必須):MLflowエクスペリメントを作成するか、既存のエクスペリメントに追加します。environment(オプション):Python の依存関係と基本環境のバージョン。compute(必須): GPU リソース(タイプと数)。command(必須):トレーニングの起動に使用されるbashコマンド。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 エクスペリメントを作成します (既に存在する場合は、新しいランを追加します)。environmentuses the default environment and installstorchandtransformers.computeは 1 つの H100 ノード (H100 GPU 8 基) を割り当てます。code_sourceフォルダrepoをノードにuploadします。これは$CODE_SOURCE_PATHで利用できます。command8個のH100 GPU全体で、torchrunを介して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 Templateを共有する場合、他のユーザーは独自のシークレットを作成するか、参照されているシークレットへのアクセス権を持っている必要があります。
環境
environmentブロックを使用してServerless GPU環境を選択し、Pythonの依存関係をインストールします。例えば、次の設定では標準環境バージョン4を選択し、PyTorchとTransformersをインストールします。
environment:
version: '4'
dependencies:
- torch
- transformers
環境バージョン
environment.version はオプションであり、ワークロードの管理対象環境バージョンを選択します。
例は次のとおりです。
"4"または"5"を使用して、対応する Standard 環境バージョンを使用します。"databricks_ai_v5"プレインストールされた ML 固有のパッケージを含む Databricks AI 環境バージョン 5 を使用するため(パッケージ一覧)
次の例では、Databricks AI 環境バージョン 5 を選択しています。
environment:
version: 'databricks_ai_v5'
dependencies: []
environment.version を指定する場合は、インライン リストとして environment.dependencies も指定する必要があります。追加パッケージをインストールする必要がない場合は、空のリストを使用します。
AI ランタイムで使用可能な環境については、環境の設定を参照してください。
Python の依存関係
ワークロードの Python 依存関係を、environment.dependencies の下にインラインリストとして一覧表示します。
依存関係の形式
依存関係リストは 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 URLs :インデックス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.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 パターンなど、詳細については、レガシー Python CLI でのカスタム Docker イメージの使用を参照してください。
コードソースの操作
code_sourceブロックはローカルコードをuploadし、トレーニング ジョブがそれを実行できるようにします。
root_pathスナップショットを作成するローカルディレクトリです。defaultでは、airは作業ツリーをそのまま(コミットされていない変更を含めて)通常の tarball としてパッケージ化します。- 代わりにピン留めされた Git バージョンをスナップショットするには、
git:ブロックにbranchまたはcommitを追加します。これを行うには、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 リポジトリ: Branch または commit でピン留めする
gitリポジトリの場合は、Branch別またはcommit SHA別にコードのバージョンをピン留めする git: ブロックを追加します。branch と commit は相互に排他的です。ブロック内にどちらか一方のみを指定してください。
Pin to a Branch (そのBranchのローカル 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(オプション):Branch 名。ローカルの HEAD を使用します。リモートフェッチは行われません。git.commitと相互に排他的です。git.commit(オプション):特定の commit SHA。git.branchと相互に排他的です。git.remote(オプション):ローカルHEADの代わりに、BranchのリモートHEADを使用します。リモートを自動検出しようとするにはtrueに設定し、特定のリモートからフェッチするにはリモート名(例:upstream)に設定します。git.branchでのみ有効です。
git: ブロックを省略した場合、air は、コミットされていない変更を含め、変更されていない作業ツリーを通常の tarball としてパッケージ化します。追加のフィールドは必要ありません。
Git非対応ディレクトリ
git リポジトリではないディレクトリのスナップショットを作成できます。git: ブロックを省略します。これには root_path が git リポジトリである必要があります。これがないとバージョン キャッシュは行われず、ランごとに新しい tarball が upload されます。
code_source:
type: snapshot
snapshot:
root_path: /home/username/my_project
command: $CODE_SOURCE_PATH/train.py
次を使用したフォルダフィルタリング: include_paths
For large monorepos, スナップショット only specific folders to reduce upload and download time and スナップショット size:
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
主なポイント:
- このフィールドはオプションです。省略した場合は、defaultでリポジトリ全体が含まれます。
- パスはリポジトリのルートからの相対パスである必要があります(先頭に
/は付けないでください)。 ..は許可されていません。親ディレクトリを参照することはできません。
高度な機能
カスタムハイパーパラメーター
構造化された構成を 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
ワークロードが失敗した場合、2回再試行されます。各試行の完了制限時間は90分であるため、総ウォールクロック予算は90×3=270分になります。
コスト配賦
usage_policy_nameを使用して、既存の予算ポリシーにワークロードをアタッチします。ワークロードの起動時に、名前がポリシーのIDに解決されます。セットアップについては、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
リファレンス
コアフィールドのリファレンス
フィールド | Type | 説明 | 例 |
|---|---|---|---|
| string | MLflowのエクスペリメント名。 |
|
| string | ランによって記録されたMLflowアーティファクトのルートロケーション。オプション。 |
|
| list | pip 依存関係仕様のインライン リスト。 |
|
| string | Serverless GPU 環境バージョン。オプション。省略した場合はdefault環境を使用します。環境バージョンを参照してください。 |
|
| int | GPU の数。選択した |
|
| string | GPUタイプやノードシェイプを含むアクセラレータ構成。サポートされているGPU構成を参照してください。 |
|
| dict | コードのソース構成。 | コードソースの操作を参照してください。 |
| string | トレーニングを起動するbashコマンド。 |
|
サポートされているGPU構成
| ノードあたりの GPU 数 |
| 注 |
|---|---|---|---|
| 1 | 正の整数 | 単一の A10、開発および小規模なワークロードに適しています。 |
| 1 |
| シングル H100。 |
| 8 | 8の正の倍数 | フル H100 ノード。分散トレーニングによく使用されます。 |
アクセラレーターの機能と推奨されるユースケースについては、Hardware options を参照してください。
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'
環境バージョン、依存関係のフォーマット、およびサポートされているインストールフラグについては、Environment を参照してください。
カスタムDockerイメージ構成
environment:
docker_image:
url: myorg/myrepo:mytag
environment.dependencies および environment.version と相互に排他的です。使用する前に、画像を air register image に登録する。レガシー Python CLI でカスタム Dockerイメージを使用するについては、See を参照してください。
コードのソース構成
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:ブロック内で正確に 1 つを指定してください。git.remotegit.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 エクスペリメントのアーティファクトをカスタムのルート場所に保存するには、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を指定してください。air CLIは、/VolumesパスをMLflowが使用するdbfs: URIに変換します。
エクスペリメントごとに一意の場所を使用します。MLflow エクスペリメントのアーティファクトの場所は、エクスペリメントの作成時に固定されます。experiment_name が既存のエクスペリメントを特定する場合、mlflow_artifact_location はそのアーティファクトの場所と一致するか、省略される必要があります。別の場所を使用するには、新しいエクスペリメント名義を指定します。
Path resolution
All paths in the workload YAML are relative to the workload YAML unless they are absolute paths.
フォルダ構造:
/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