メインコンテンツまでスキップ
に最終更新

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

備考

ベータ版

AI Runtime 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つのコンポーネントが含まれます。

  1. experiment_name必須。MLflowエクスペリメントを作成または追加します。
  2. environment：オプション。Python の依存関係とベース環境
  3. compute必須。GPUリソース（タイプと数）
  4. command必須。トレーニングの起動に使用されるbashコマンド。
  5. 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 リストされたPythonの依存関係（ここでは、torchとtransformers）をインストールします。
  • compute H100ノードを1つ割り当てます（H100 GPU 8基）。
  • code_source フォルダrepoをノードにアップロードし、$CODE_SOURCE_PATHで利用可能です。
  • command train.py を torchrun 経由で 8基のH100 GPUで実行します。/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 Secrets で設定する必要があります。設定については、「シークレット管理」を参照してください。

YAMLテンプレートを共有する際、他のユーザーは独自のシークレットを作成するか、参照されているシークレットにアクセスできる必要があります。

Python の依存関係​

environment.dependenciesの下に、ワークロードのPython依存関係をインラインリストとして記述してください：

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

environment.version サーバレス GPU 環境のバージョンを選択しますオプションであり、デフォルトは"4"です。

依存関係の形式​

依存関係リストは、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。
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 を使用してください。

ソースコードの操作​

code_source ブロックがローカルコードをアップロードすることで、トレーニングジョブがそれを実行できます。

  • root_path スナップショットのローカルディレクトリです。デフォルトでは、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リポジトリ：ブランチまたはコミットによるピン留め​

Gitリポジトリの場合、ブランチまたはコミットSHAでコードバージョンをピン留めするために、git:ブロックを追加します。branch と commit は相互に排他的です。ブロック内では、いずれか一方のみを指定してください。

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

コミット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 （オプション）特定のコミットSHAです。git.branchと相互に排他的です。
  • git.remote （オプション）ローカルのHEADではなく、ブランチのリモートHEADを使用します。true に設定するとリモートが自動検出されます。または、特定の「リモート」から取得する場合は、リモート名（例: upstream）に設定してください。git.branchでのみ有効です。

git:ブロックを省略すると、airは、コミットされていない変更点を含む作業ツリーをプレーンなtarボールとしてパッケージ化します。追加のフィールドは不要です。

Gitリポジトリ以外のディレクトリ​

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

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

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

大規模なモノレポでは、アップロードとダウンロードの時間、およびスナップショットのサイズを削減するため、特定のフォルダーのみをスナップショットしてください:

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

キーポイント：

  • フィールドはオプションです。省略した場合、リポジトリ全体がデフォルトで含まれます。
  • パスはリポジトリのルートに相対的である必要があります（先頭に/なし）。
  • .. 許可されていません；親ディレクトリは参照できません。

高度な機能​

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

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_id経由でワークロードをアタッチしてください。セットアップについては、サーバレス利用ポリシーでの属性の使用を参照してください。

YAML
experiment_name: my-training
environment:
  dependencies:
    - mlflow
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: echo "Hello World"
usage_policy_id: abcd123-25b8-3e87-9a2c-f86eb19d101c

リファレンス​

コア項目​

フィールド

Type

説明

例

experiment_name

string

MLflowのエクスペリメント名。

"my-training-job"

environment.dependencies

list

pip 依存関係の仕様のインラインリスト。

["torch", "transformers"]

environment.version

string

サーバレス GPU 環境バージョン。オプション。デフォルトは"4"です。

"4"

compute.num_accelerators

int

GPU数。

1、4、 8

compute.accelerator_type

string

GPUタイプ。

"GPU_1xA10", "GPU_8xH100"

code_source

辞書

コードソース構成。

ソースコードの操作を参照してください。

command

string

トレーニングを起動するBash コマンド

torchrun --nproc_per_node=8 train.py

サポートされている GPU タイプ​

accelerator_type

ノードあたりのGPU

注

GPU_1xA10

1

シングル A10 は、開発や小規模なワークロードに最適です。

GPU_1xH100

1

単一のH100。

GPU_8xH100

8

フル H100 ノードは、分散トレーニングに一般的に使用されます。

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

任意フィールド​

環境設定

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

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

ソースコード構成

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'

パスの解決​

ワークロード 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
このページの見出し