ワークロード YAML リファレンス

Important

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

air run --fileに渡すワークロード YAML 構成で、トレーニングジョブの実験名、コンピューティング、コマンド、環境、およびコードソースを定義します。このページでは、すべてのフィールドを文書化します。

Note

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_name は、 simple-training という名前の MLflow 実験を作成します (または、既に存在する場合は新しい実行を追加します)。
environmentは、一覧表示されているPython依存関係 (ここでは、torchとtransformers) をインストールします。
compute は、1 つの H100 ノード (8 H100 GPU) を割り当てます。
code_source は、フォルダー repo をノードにアップロードします。これは、 $CODE_SOURCE_PATHで使用できます。
commandは、8 H100 GPU 全体でtrain.py経由でtorchrunを実行します。ファイルはローカル /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 テンプレートを共有する場合、他のユーザーは独自のシークレットを作成するか、参照先のシークレットにアクセスできる必要があります。

依存関係のPython

ワークロードのPython依存関係を、environment.dependenciesのインラインリストとして一覧表示します。

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

environment.version は、サーバーレス GPU 環境のバージョンを選択します。これは省略可能で、既定値は "4" です。

依存関係の形式

依存関係の一覧は、 Databricks 基本環境の仕様に従います。各エントリは pip スタイルのパッケージ仕様です (たとえば、 my-library==6.1)。この一覧には、次のエントリも含まれます。

要件ファイル: requirements.txtなど、-rを使用した既存の-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などのインデックス URL。

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の依存関係に適用されます。フラグを独自の行 (または仕様の前) に配置し、その後に適用される依存関係を指定します。

たとえば、既にインストールされているflash-attn (ビルドの分離なし) に対して、独自の依存関係を解決せずにtorchをインストールするには、次のようにします。

environment:
  version: '4'
  dependencies:
    - torch
    - --no-build-isolation
    - --no-deps
    - flash-attn

Note

--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 パターンなど、詳細については、「カスタム Docker イメージの使用」を参照してください。

コードソースの操作

code_source ブロックは、トレーニングジョブで実行できるようにローカルコードをアップロードします。

root_path は、スナップショットを作成するローカルディレクトリです。既定では、 air は作業ツリー as-is (コミットされていない変更を含む) をプレーンな 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 リポジトリ: ブランチまたはコミットでピン留めする

git リポジトリの場合は、ブランチまたはコミット SHA によってコードバージョンをピン留めする git: ブロックを追加します。 branch と commit は相互に排他的です。ブロック内で 1 つだけ指定します。

ブランチにピン留めする (そのブランチのローカル 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

コミット 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 (省略可能): ブランチ名。ローカル HEAD を使用します。リモートフェッチはありません。 git.commitと相互に排他的です。
git.commit (省略可能): 特定のコミット SHA。 git.branchと相互に排他的です。
git.remote (省略可能): ローカルではなくブランチのリモート HEAD を使用します。 trueに設定してリモートを自動検出するか、リモート名 (たとえば、upstream) に設定して、特定のリモートからフェッチします。 git.branchでのみ有効です。

git: ブロックを省略した場合、airは、コミットされていない変更を含め、作業ツリーをプレーンな tarball としてパッケージ化します。追加のフィールドは必要ありません。

Git 以外のディレクトリ

Git リポジトリではないディレクトリのスナップショットを作成できます。 git: ブロックを省略します。そのためには、root_pathを git リポジトリにする必要があります。これを使用しないと、バージョンキャッシュはありません。すべての実行に対して新しい tarball がアップロードされます。

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

フォルダーのフィルター処理 `include_paths`

大規模な monorepos の場合、特定のフォルダーのみをスナップショットして、アップロードとダウンロードの時間とスナップショットのサイズを減らします。

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

重要なポイント:

フィールドは省略可能です。省略すると、リポジトリ全体が既定で含まれます。
パスは、リポジトリルートに対する相対パスである必要があります (先頭に /はありません)。
.. は許可されていません。親ディレクトリを参照することはできません。

高度な機能

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

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 に解決されます。セットアップについては、「サーバーレス使用ポリシーでの属性の使用」を参照してください。

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

Reference

コアフィールド

フィールド	タイプ	Description	Example
`experiment_name`	文字列	MLflow の実験名。	`"my-training-job"`
`environment.dependencies`	list	pip 依存関係仕様のインラインリスト。	`["torch", "transformers"]`
`environment.version`	文字列	サーバーレス GPU 環境バージョン。省略可能。既定値は `"4"` です。	`"4"`
`compute.num_accelerators`	int	GPU の数。	`1`、`4`、`8`
`compute.accelerator_type`	文字列	GPU の種類。	`"GPU_1xA10"`、`"GPU_8xH100"`
`code_source`	辞書	コードソースの構成。	「コードソースの操作」を参照してください。
`command`	文字列	トレーニングを開始するための Bash コマンド。	`torchrun --nproc_per_node=8 train.py`

サポートされている GPU の種類

`accelerator_type`	ノードあたりの GPU 数	注記
`GPU_1xA10`	1	単一の A10。開発と小規模なワークロードに適しています。
`GPU_1xH100`	1	単一 H100。
`GPU_8xH100`	8	分散トレーニングに一般的な完全な H100 ノード。

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

省略可能なフィールド

環境の構成

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

依存関係の形式、サポートされているインストールフラグ、およびenvironment.versionについては、依存関係Python参照してください。

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

environment:
  docker_image:
    url: myorg/myrepo:mytag

environment.dependenciesとenvironment.versionと相互に排他的です。使用する前に、イメージを air register image に登録します。「カスタム Docker イメージを使用する」を参照してください。

コードソースの構成

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を介してワークロードに渡されます。

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

MLflow 実行名

mlflow_run_name: 'experiment-001-baseline'

パスの解決

ワークロード YAML 内のすべてのパスは、絶対パスでない限り、ワークロード YAML に対して相対的です。

フォルダー構造:

/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

フィードバック

このページはお役に立ちましたか?

Last updated on 2026-06-30

ワークロード YAML リファレンス

最小構成

主要な概念

コア フィールド

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

一般的なユース ケース

環境変数の追加

シークレット (API キー、トークン) を使用する

依存関係のPython

依存関係の形式

サポートされているインストール フラグ

カスタム Docker イメージ

コード ソースの操作

最小限の例

Git リポジトリ: ブランチまたはコミットでピン留めする

Git 以外のディレクトリ

フォルダーのフィルター処理 include_paths

高度な機能

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

ジョブの信頼性

コスト配分

Reference

コア フィールド

サポートされている GPU の種類

省略可能なフィールド

パスの解決

フィードバック

その他のリソース

コアフィールド

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

一般的なユースケース

サポートされているインストールフラグ

コードソースの操作

フォルダーのフィルター処理 `include_paths`

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

コアフィールド