レガシ Python CLI の Workload YAML リファレンス

重要

このドキュメントは廃止されており、更新されない可能性があります。

databricks-air パッケージでインストールされる Python ベースの air CLI は非推奨となり、アクティブなメンテナンスは終了しました。

新しいワークロードには Databricks CLI を使用します。AI ランタイムで Databricks CLI を使用するを参照してください。

air run --file に渡すワークロード YAML 設定で、トレーニングジョブのエクスペリメント名、コンピュート、コマンド、環境、コードソースを定義します。このページでは、すべてのフィールドについて説明します。

注記

YAML構成のグラウンドトゥルースは、CLI内ヘルプです。air -h config最上位ビューに対してair -h config.<section> を、セクションごとの詳細に対してラン（たとえば、air -h config.environment ）を実行します。

最小限の構成

YAML
experiment_name: my-training
environment:
  dependencies:
    - mlflow
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: echo "Hello World"

次を使って送信：

Bash
air run --file train.yaml -p profile

主要な概念

コアフィールド

ほとんどのトレーニング構成には5つのコンポーネントが含まれます。

experiment_name （必須）：MLflowエクスペリメントを作成または追加します。
environment （オプション） Python の依存関係と基本環境バージョン。
compute (必須): GPUリソース (タイプと数)。
command （必須）：トレーニングを開始するために使用されるbashコマンド。
code_source （オプション）：リモートで利用できるようにしたトレーニングコードへのパス。

サポートされている値とフィールドの制約については、リファレンスを参照してください。

最初のトレーニングジョブ

YAML
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_name simple-trainingという名前のMLflowエクスペリメントを作成します（既に存在する場合は新しいランを追加します）。
environment default環境を使用し、torchとtransformersをインストールします。
compute 1つのH100ノード（8つのH100 GPU）を割り当てます。
code_source フォルダrepoをノードにuploadし、$CODE_SOURCE_PATHで使用できるようにします。
command 8つのH100 GPU全体でtorchrunを介してtrain.pyをランします。ファイルはローカルの/home/username/repo/train.pyに存在します。

一般的なユースケース

環境変数の追加

YAML
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キー、トークン）の使用

YAML
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をインストールします。

YAML
environment:
  version: '4'
  dependencies:
    - torch
    - transformers

環境バージョン

environment.version はオプションであり、ワークロードの管理環境バージョンを選択します。

例は次のとおりです。

"4" または、対応する標準環境バージョンを使用する場合は"5"。
"databricks_ai_v5" は、 ML 固有のパッケージがプレインストールされている Databricks AI 環境バージョン 5 を使用します。(パッケージリスト全体)

次の例では、Databricks AI環境バージョン5を選択しています。

YAML
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)。
インデックスURL : インデックスURL (例: --index-url https://pypi.org/simple)。

YAML
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 をインストールする（ビルドの分離なし）場合に、その独自の依存関係を解決しないときは以下のようになります。

YAML
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 の両方と相互に排他的です — 両方を同じワークロードで使用することはできません。

YAML
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バージョンをスナップショットするには、 branchまたはcommitを含む git:ブロックを追加します。これには、root_pathがGitリポジトリである必要があり、バージョンを意識したスナップショット（キャッシング、git archive）が有効になります。
大規模なリポジトリの場合、 include_pathsを使用してサブセットのスナップショットを作成できます。

最小限の例

YAML
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リポジトリの場合は、git:ブロックを追加して、Branchまたはcommit SHAによってコードバージョンをピン留めします。branchとcommitは相互に排他的です：ブロック内にこれらの一方のみを指定してください。

Branchにピン留め (そのBranchのローカルHEADを使用)：

YAML
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 にピン留め (正確な再現性):

YAML
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リポジトリではないフォルダのスナップショットを作成できます。root_pathがgitリポジトリである必要があるgit:ブロックを省略します。それがない場合、バージョンキャッシングは行われません。実行するたびに新しいtarballがuploadされます。

YAML
code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/my_project
command: $CODE_SOURCE_PATH/train.py

フォルダーフィルタリングと `include_paths`

大規模なモノレポの場合は、特定のフォルダーのみをスナップショットして、uploadとdownloadの時間とスナップショットのサイズを短縮します：

YAML
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 を介して構造化された構成をトレーニングスクリプトに渡します：

YAML
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

スクリプトで読み込みます：

Python
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']

ジョブの信頼性

YAML
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利用ポリシーでの属性の使用を参照してください。

YAML
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	説明	例
`experiment_name`	string	MLflowのエクスペリメント名。	`"my-training-job"`
`mlflow_artifact_location`	string	ランによって記録されたMLflowアーティファクトのルートの場所。任意。	`/Volumes/main/default/mlflow-artifacts/my-training`
`environment.dependencies`	list	pip依存関係仕様のインラインリスト。	`["torch", "transformers"]`
`environment.version`	string	Serverless GPU 環境バージョン。オプション。省略された場合は、defaultの環境が使用されます。環境バージョンを参照してください。	`"4"`、`"5"`、 `"databricks_ai_v5"`
`compute.num_accelerators`	int	GPUの数。選択された `compute.accelerator_type` のノードあたりの GPU の倍数である必要があります。	`1`、`4`、 `8`
`compute.accelerator_type`	string	GPUタイプとノード形状を含むアクセラレータの構成。サポートされている GPU 構成を参照してください。	`"GPU_1xA10"`、`"GPU_1xH100"`、 `"GPU_8xH100"`
`code_source`	dict	コードソース構成。	コードソースの操作を参照してください。
`command`	string	トレーニングを開始するBashコマンド。	`torchrun --nproc_per_node=8 train.py`

フィールド	Type	説明	例
`experiment_name`	string	MLflowのエクスペリメント名。	`"my-training-job"`
`mlflow_artifact_location`	string	ランによって記録されたMLflowアーティファクトのルートの場所。任意。	`/Volumes/main/default/mlflow-artifacts/my-training`
`environment.dependencies`	list	pip依存関係仕様のインラインリスト。	`["torch", "transformers"]`
`environment.version`	string	Serverless GPU 環境バージョン。オプション。省略された場合は、defaultの環境が使用されます。環境バージョンを参照してください。	`"4"`、`"5"`、 `"databricks_ai_v5"`
`compute.num_accelerators`	int	GPUの数。選択された `compute.accelerator_type` のノードあたりの GPU の倍数である必要があります。	`1`、`4`、 `8`
`compute.accelerator_type`	string	GPUタイプとノード形状を含むアクセラレータの構成。サポートされている GPU 構成を参照してください。	`"GPU_1xA10"`、`"GPU_1xH100"`、 `"GPU_8xH100"`
`code_source`	dict	コードソース構成。	コードソースの操作を参照してください。
`command`	string	トレーニングを開始するBashコマンド。	`torchrun --nproc_per_node=8 train.py`

サポートされている GPU 構成

`accelerator_type`	ノードあたりのGPU数	`num_accelerators` 要件	注
`GPU_1xA10`	1	任意の正の整数	シングル A10。開発や小規模なワークロードに適しています。
`GPU_1xH100`	1	`1`	シングルH100。
`GPU_8xH100`	8	8 の正の倍数	分散トレーニングの一般的なフル H100 ノード。

`accelerator_type`	ノードあたりのGPU数	`num_accelerators` 要件	注
`GPU_1xA10`	1	任意の正の整数	シングル A10。開発や小規模なワークロードに適しています。
`GPU_1xH100`	1	`1`	シングルH100。
`GPU_8xH100`	8	8 の正の倍数	分散トレーニングの一般的なフル H100 ノード。

アクセラレータの機能と推奨されるユースケースについては、「ハードウェアオプション」を参照してください。

compute.num_accelerators はワークロードのGPUの合計数です。これは、選択したcompute.accelerator_typeのノードあたりのGPUの倍数である必要があります。

: オプションのフィールド

環境設定

YAML
environment:
  version: '4'
  dependencies:
    - torch
    - transformers
env_variables:
  BATCH_SIZE: '32'
secrets:
  HF_TOKEN: 'my_scope/hf_token'

環境バージョン、依存関係の形式、およびサポートされているインストールフラグについては、環境を参照してください。

カスタムDockerイメージの構成

YAML
environment:
  docker_image:
    url: myorg/myrepo:mytag

environment.dependenciesおよびenvironment.versionと相互に排他的です。使用前に画像をair register imageに登録してください。「従来の Python CLI でカスタム Dockerイメージを使用する」を参照してください。

コードソースの構成

YAML
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.remote git.branch を必要とします (git.commit では効果がありません)。
git: ブロックを省略した場合は、コミットされていない変更を含む作業ツリーがプレーンなtarballとしてパッケージ化されます。

カスタムパラメーター

HYPERPARAMETERS_PATHを介してワークロードに渡されます:

YAML
parameters:
  model:
    name: 'gpt2'
    hidden_size: 768
  training:
    batch_size: 32

MLflow ラン名

YAML
mlflow_run_name: 'experiment-001-baseline'

: MLflowアーティファクトの場所

MLflow エクスペリメントのアーティファクトをカスタムルートロケーションに格納するには、 mlflow_artifact_locationを設定します。このフィールドを省略すると、新しいエクスペリメントは dbfs:/databricks/mlflow-tracking/<experiment-id>/...などの default DBFS ロケーションを使用します。

YAML
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 はそのアーティファクトの場所と一致するか、省略する必要があります。別の場所を使用するには、新しいエクスペリメント名を指定します。

パス解決

ワークロードYAML内のすべてのパスは、絶対パスでない限り、ワークロードYAMLを基準とします。

フォルダ構造:

/home/username/my-project/
├── train.yaml
└── scripts/
    └── train.py

YAML構成：

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

最小限の構成​

主要な概念​

コアフィールド​

最初のトレーニングジョブ​

一般的なユースケース​

環境変数の追加​

シークレット（APIキー、トークン）の使用​

環境​

環境バージョン​

Pythonの依存関係​

カスタム Docker イメージ​

コードソースの操作​

最小限の例​

Gitリポジトリ：Branchまたはcommitでピン留め​

: Git 以外のディレクトリ​

フォルダーフィルタリングと include_paths​

高度な機能​

カスタムハイパーパラメーター​

ジョブの信頼性​

コスト属性​

参照​

コアフィールドのリファレンス​

サポートされている GPU 構成​

: オプションのフィールド​

パス解決​