メインコンテンツまでスキップ
非公開のページ
このページは非公開です。 検索対象外となり、このページのリンクに直接アクセスできるユーザーのみに公開されます。

レガシー 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)を実行します。

最小限の構成​

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 uses the default environment and installs torch and transformers.
  • compute は 1 つの H100 ノード (H100 GPU 8 基) を割り当てます。
  • 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" を使用して、対応する Standard 環境バージョンを使用します。
  • "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)。
  • Index URLs :インデックス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 バージョンをスナップショットするには、git: ブロックに branch または commit を追加します。これを行うには、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リポジトリの場合は、Branch別またはcommit SHA別にコードのバージョンをピン留めする git: ブロックを追加します。branch と commit は相互に排他的です。ブロック内にどちらか一方のみを指定してください。

Pin to a 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 リポジトリではないディレクトリのスナップショットを作成できます。git: ブロックを省略します。これには root_path が git リポジトリである必要があります。これがないとバージョン キャッシュは行われず、ランごとに新しい tarball が upload されます。

YAML
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:

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 ノード。分散トレーニングによく使用されます。

アクセラレーターの機能と推奨されるユースケースについては、Hardware options を参照してください。

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'

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

カスタムDockerイメージ構成

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

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

コードのソース構成

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

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 構成:

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