レガシーPython CLI用のワークロードYAML参照

Important

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

databricks-airパッケージにインストールされたPythonベースのair CLIは現在非推奨となり、活発なメンテナンスは行われていません。

新しいワークロードにはDatabricksのCLIを使いましょう。 「 Data bricks CLIとAIランタイムを使いましょう」を参照してください。

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

  1. experiment_name (必須): MLflow の実験を作成するか、既存の実験に追加します。
  2. environment(オプション):Python依存関係と基本環境バージョン。
  3. compute (必須):GPUリソース(タイプとカウント)。
  4. command (必須):訓練を開始するために使われるバッシュコマンド。
  5. 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 デフォルトの環境を使用し、 torch と transformersをインストールします。
  • compute は、1 つの H100 ノード (8 H100 GPU) を割り当てます。
  • code_source は、フォルダー repo をノードにアップロードします。これは、 $CODE_SOURCE_PATHで使用できます。
  • commandは、8基の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 テンプレートを共有する場合、他のユーザーは独自のシークレットを作成するか、参照先のシークレットにアクセスできる必要があります。

環境

environmentブロックを使ってサーバーレスGPU環境を選択し、Python依存関係をインストールしてください。 例えば、以下の構成は標準環境バージョン4を選択し、PyTorchとTransformersをインストールします。

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

環境のバージョン

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

その例は次のとおりです。

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

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

environment:
  version: 'databricks_ai_v5'
  dependencies: []

environment.versionを指定する場合は、インラインリストとしてenvironment.dependenciesも提供しなければなりません。 追加のパッケージをインストールする必要がない場合は、空のリストを使いましょう。

AIランタイムで利用可能な環境に関する情報は「 環境の設定」をご覧ください。

Python の依存関係

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

依存関係形式

依存関係の一覧は、 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のベースイメージ、Dockerファイルのパターンなどの詳細については、「レガシーPython CLIでカスタム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"
mlflow_artifact_location 文字列 実行によって記録されたMLflowアーティファクトのルート位置。 Optional. /Volumes/main/default/mlflow-artifacts/my-training
environment.dependencies list pip 依存関係仕様のインライン リスト。 ["torch", "transformers"]
environment.version 文字列 サーバーレス GPU 環境バージョン。 Optional. 省略する場合はデフォルトの環境を使用します。 環境 バージョンを参照してください。 "4"、"5"、"databricks_ai_v5"
compute.num_accelerators int GPU の数。 選択した compute.accelerator_typeのノードごとにGPUの倍数でなければなりません。 1、4、8
compute.accelerator_type 文字列 GPUの種類やノード形状を含むアクセラレータの設定。 対応GPU構成をご覧ください。 "GPU_1xA10"、"GPU_1xH100"、"GPU_8xH100"
code_source 辞書 コード ソースの構成。 「コード ソースの操作」を参照してください。
command 文字列 トレーニングを開始するための 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 ノード。

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

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イメージを使う」をご覧ください。

コード ソースの構成

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'

MLflowのアーティファクトの位置

mlflow_artifact_location をカスタムルート位置に MLflow 実験用のアーティファクトを保存するように設定します。 このフィールドを省略すると、新しい実験はデフォルトのDBFS位置(例えば dbfs:/databricks/mlflow-tracking/<experiment-id>/...)を使用します。

mlflow_artifact_location: /Volumes/main/default/mlflow-artifacts/my-training

DBFSへのアクセスが制限されている場合やUnity Catalogを好む場合は、 /Volumes/<catalog>/<schema>/<volume>/... パスまたはURI dbfs:/Volumes/<catalog>/<schema>/<volume>/... 同等のものを指定してください。 airCLIは、MLflowが使用するdbfs:URIに/Volumesパスを変換します。

各実験ごとにユニークな場所を使いましょう。 MLflow実験のアーティファクトの位置は、実験が作成される際に固定されます。 experiment_name既存の実験を特定した場合、mlflow_artifact_locationその遺物の位置と一致しなければならず、そうでなければ除外されます。 別の場所を使う場合は、新しい実験名を指定します。

パス解決

ワークロード 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