databricks-ai-runtime

Name: databricks-ai-runtime
Author: databricks

github.com/databricksクラウド・DevOps37348 回インストール今日更新

databricks-ai-runtime について

databricks-ai-runtimeは、`databricks air` CLIを使ってDatabricksのサーバーレスコンピュートにGPUトレーニングを送信するためのスキルです。設定、送信、状態確認、ログ監視を扱い、ノード環境変数を使った分散トレーニングにも対応します。

主な機能

01YAMLによるワークロード設定
02GPU実行の送信と確認
03状態監視とログのストリーミング
04シークレット、パラメーター、再試行の設定
05マルチノード分散トレーニング対応

活用シーン

01サーバーレスコンピュートでGPU学習を開始
02実行状態の確認とログによる調査
03学習にハイパーパラメーターとシークレットを渡す
04マルチノード学習コマンドの設定

インストール

コマンドをコピーしてターミナルで実行します。

$npx skills add https://github.com/databricks/databricks-agent-skills --skill databricks-ai-runtime

インストール後、Claude Code で/databricks-ai-runtimeを入力して読み込まれたか確認します。

同梱ファイル

SKILL.md と一緒にインストールされる参考資料やスクリプト。

docker-images.md

agents/

openai.yaml

assets/

databricks.svg

詳細

Databricks AI Runtime（`databricks air`）

Databricks AI Runtime は、GPU トレーニングのワークロードを Databricks のサーバーレスコンピュートに送信するための databricks air CLI コマンド一式です。クラスターやインフラストラクチャを管理する必要なく、環境のセットアップ、分散トレーニングの構成、ワークロードのライフサイクルを管理します。

このスキルでは、日常的な3つの作業フロー、つまり 実行の開始、実行状況の確認、トレーニング中の 実行のモニタリング を扱います。

リファレンスの利用方針

フィールド名、GPU の種類、クラスター上の環境変数、制約については、CLI に組み込まれているヘルプを信頼できる情報源として使用してください。記憶している情報よりも、実行時のヘルプを優先してください。run コマンドで -h にパスを渡すと、設定フィールドを確認できます（パスは独立した引数です）。

databricks air run -h config                              # list the YAML fields
databricks air run -h config.compute                      # one field
databricks air run -h config.compute.accelerator_type     # a nested field

databricks air <subcommand> --help では、各コマンドのフラグを確認できます。このスキルではワークフローと、ヘルプに記載されていないわずかな事項を説明します。

セッションのセットアップ

CLI のバージョンを確認し、コマンドが利用できることを確かめてから、Databricks の設定プロファイルを一覧表示します。

databricks version
databricks air --help
databricks auth profiles --skip-validate    # available Databricks config profiles

例では、Databricks CLI 1.20.0 以降で利用できる、フラット化された code_source フィールドを使用しています。古い CLI では、ワークロード YAML に示す同等のネストされたスナップショット形式を使用してください。CLI 1.20.0 以降では、どちらの形式も使用できます。同じ設定内で、ネストされたフィールドとフラット化されたフィールドを混在させないでください。

すべてのコマンドに -p <profile> を指定してください。スクリプトやプログラムから呼び出す場合は -o json を追加すると、人間向けのテキストではなく、構造化されたエンベロープを取得できます。

成功：{"v": 1, "ts": "...", "data": {...}}
エラー：{"v": 1, "ts": "...", "error": {"code": "...", "kind": "...", "message": "...", "retryable": bool}}

設定の検証エラー（入力エラー）は例外です。-o json を指定していても、標準エラー出力にプレーンテキストで出力され、終了コードは 0 以外になります。そのため、解析した JSON だけでなく、終了コードと標準エラー出力も確認してください。構造化されたエラーが返された場合、kind に対応方法が示されます。

kind	意味	対応方法
`PERMANENT`	認証されていない／権限がない／再試行不可	認証エラーの場合：`databricks auth login --host <workspace-url>`
`NOT_FOUND`	実行 ID／リソースが不正	入力を確認してください
`CONFLICT`	冪等性キーの競合	新しい `--idempotency-key` を使用してください
`TRANSIENT`	再試行可能	再試行してください

サブコマンド

コマンド	機能
`databricks air run`	YAML ファイルからワークロードを送信します。完了するまでログをストリーミングするには `--watch` を追加します。
`databricks air get <JOB_RUN_ID>`	1つの実行のステータス、設定、実行時間を表示します。
`databricks air list`	アクティブな実行を一覧表示します（完了した実行には `--all-status`、全ユーザーの実行には `--all-users`、さらに `--filter`、`--limit` を追加します）。
`databricks air cancel <JOB_RUN_ID> [<JOB_RUN_ID>...]` / `... --all`	1つ以上の実行、または自分のアクティブな実行をすべてキャンセルします。
`databricks air logs <JOB_RUN_ID>`	実行中の実行のログをストリーミングするか、完了した実行のログを取得します。

フラグとデフォルト値は、databricks air <subcommand> --help を実行して確認してください。

ワークロード YAML

experiment_name: my-training-job     # alphanumeric, - and _ only
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10        # run `databricks air run -h config.compute` for valid values
environment:
  version: "5"                       # client image version; prefix "databricks_ai_v" to also load the ML venv
  dependencies:
    - mlflow
code_source:
  root_path: "."                     # relative to workload.yaml; extracted to $CODE_SOURCE_PATH on each node
command: |-
  cd "$CODE_SOURCE_PATH"
  python train.py

databricks air run --file workload.yaml -p <profile> を実行して送信します。フィールドの詳細は databricks air run -h config.<field> で確認できます（これが信頼できる情報源です）。現在のアクセラレーターの種類は GPU_1xA10 と GPU_8xH100 です。databricks air run -h config.compute で確認してください。

古い CLI では、上記の code_source ブロックのみを次の内容に置き換えてください。

code_source:
  type: snapshot
  snapshot:
    root_path: "."

値と動作は同じです。code_source の下に type: snapshot を追加し、すべてのスナップショットフィールド（root_path と、使用する場合は remote_volume、git、include_paths）を code_source.snapshot 内に配置します。ワークロードの残りの部分は変更しないでください。インストールされている CLI でサポートされるフィールドは、databricks air run -h config.code_source で確認してください。

ワークロードの設定

次の省略可能なトップレベルフィールドで、一般的なニーズに対応できます。それぞれの詳細なリファレンスは、databricks air run -h config.<field> を実行して確認してください。

環境変数とシークレット

通常の値は env_variables に設定し、機密性の高い値は Databricks のシークレット secrets から取得します（各エントリは環境変数名を scope/key に対応付けます。値はワーカー上で展開され、ログに記録されることはありません）。

env_variables:
  BATCH_SIZE: "32"
secrets:
  HF_TOKEN: my_scope/hf_token          # -> $HF_TOKEN on every node
  WANDB_API_KEY: my_scope/wandb_key

スコープとキーは事前に作成しておく必要があります。Databricks CLI で一度だけ作成してください。

databricks secrets create-scope my_scope -p <profile>
databricks secrets put-secret my_scope hf_token -p <profile>

シークレットの環境変数名は、env_variables の名前と重複させないでください。

パラメーター（ハイパーパラメーター）

parameters は、ハイパーパラメーターをスクリプトに渡すための適切な方法です。トレーニングコードを編集せずに、実行ごとに値を変更できます。CLI はクラスター上の YAML ファイルに値を書き込み、そのパスを $HYPERPARAMETERS_PATH として公開します。

parameters:
  learning_rate: 0.001
  batch_size: 32
  warmup_steps: 1000
  use_amp: true

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

import os, yaml
with open(os.environ["HYPERPARAMETERS_PATH"]) as f:
    hp = yaml.safe_load(f)
lr = hp["learning_rate"]

値には文字列、数値、ブール値、リスト、ネストされた辞書を指定できます。ハイパーパラメーターをスイープするには、値を変更して（新しい冪等性キーを指定して）再送信します。

再試行

失敗した試行は自動的に再試行されます。max_retries のデフォルト値は 0 以外です（現在の値は databricks air run -h config.max_retries で確認してください）。再試行せずにすぐ失敗させるには、max_retries: 0 を設定します。timeout_minutes を設定した場合、タイムアウトは試行ごとに適用されます。

権限

デフォルトでは、実行を確認できるのは自分だけです。permissions（DABs 互換形式）を使って、送信したジョブへのアクセスを他のユーザーに付与できます。レベルは CAN_VIEW、CAN_MANAGE_RUN、CAN_MANAGE です。

permissions:
  - group_name: data-engineering
    level: CAN_VIEW
  - user_name: [email protected]
    level: CAN_MANAGE

各ノードで提供される環境変数

databricks air は各ノードで以下の環境変数を自動的に設定します。command: とトレーニングコードで使用してください。現在の一覧は databricks air run -h config.command で確認できます。

変数	意味
`CODE_SOURCE_PATH`	ノード上で展開された `code_source` の絶対パス（`cd $CODE_SOURCE_PATH`）。
`MASTER_ADDR` / `MASTER_PORT`	分散処理の調整に使用する rank-0 ノードのアドレス／ポート。
`NUM_NODES`	ノードの合計数。
`NODE_RANK`	現在のノードのランク（0 から `NUM_NODES`-1 まで）。
`WORLD_SIZE`	全ノードの GPU の合計数。
`LOCAL_WORLD_SIZE`	このノード上の GPU 数。
`MLFLOW_RUN_ID`	このジョブの MLflow 実行 ID（メトリクス／アーティファクトを記録します）。
`HYPERPARAMETERS_PATH`	`parameters` YAML ファイルのパス（`parameters:` が定義されている場合のみ設定）。

command: はノードごとに 1 回 実行されるため、プロセス単位のランク（RANK、LOCAL_RANK）はプラットフォームによって設定されません。分散ランチャー（torchrun、accelerate）が GPU ごとに 1 つのプロセスを起動する際にこれらを割り当てます。$NODE_RANK / $NUM_NODES / $MASTER_ADDR / $MASTER_PORT を渡してください。

実行の開始

ユーザーの意図を YAML に反映します。 リクエストからモデル、データセットのパス、GPU の数と種類、トレーニングコマンドを抽出します。重要な項目（データパス、チェックポイントディレクトリ、カスタムイメージ）が曖昧な場合は、設定を作成する前に確認してください。

環境変数を使って command: を記述します。 よくある起動形式は次のとおりです。

フレームワーク	`command:` の形式
torch、単一 GPU	`python train.py`
torch、マルチ GPU（1 ノード）	`torchrun --nproc_per_node=$LOCAL_WORLD_SIZE train.py`
torch、マルチノード	`torchrun --nnodes=$NUM_NODES --nproc_per_node=$LOCAL_WORLD_SIZE --node_rank=$NODE_RANK --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT train.py`
HuggingFace Accelerate	`accelerate launch --config_file accelerate.yaml train.py`
DeepSpeed	`deepspeed --num_gpus=$LOCAL_WORLD_SIZE train.py`
Composer / MosaicML	`composer train.py`

複数行のコマンドは command: |- の下に記述します。アップロードされたコードには $CODE_SOURCE_PATH 経由でアクセスします。GPU の総数 = ノード数 x ノードあたりのアクセラレーター数です。種類ごとのノードあたりの数は databricks air run -h config.compute で確認できるため、num_accelerators はその有効な倍数にしてください。

参照するリソースを事前確認して、ジョブの起動時の失敗を防ぎます。すべての secrets スコープとキーが存在すること（databricks secrets list-scopes / list-secrets <scope>）を確認します。また、environment.unity_catalog_image が設定されている場合は、同じワークスペースの Databricks Artifact Registry にイメージがプッシュされており、送信者の ID に読み取り権限があることを確認してください（docker-images.md を参照）。カスタムイメージと environment.version または environment.dependencies を併用しないでください。
送信前にドライランを実行します。 起動せずに設定を検証できます。
```
databricks air run --file workload.yaml --dry-run -o json -p <profile>
```
設定検証エラーが出た場合は、処理を中止してください。エラーは stderr にプレーンテキストで出力され、終了コードは 0 以外になるため、ユーザーにそのメッセージを提示して YAML を修正します。それ以外のエラーはインフラまたは認証に関するものなので、そのまま報告してください。
送信します。 安定した --idempotency-key を使って、誤って再実行したときに二重送信されないようにします（<user>-<experiment>-<YYYYMMDD>-v<N> の形式で作成し、再送信のたびに vN を増やしてください。コマンドを自動承認可能な状態に保つため、$(uuidgen) のような動的なサブ式は避けてください）。
```
databricks air run --file workload.yaml \
  --idempotency-key <key> -o json -p <profile>
```
実行が終了するまでログをターミナルにストリーミングするには、--watch を追加します。レスポンスから data.run_id を取得して、実行を監視します。

実行の確認

databricks air get <JOB_RUN_ID> -o json -p <profile>   # status, config, timing
databricks air list -o json -p <profile>               # your active runs (add --all-status for finished)

実行の監視

実行中にログをストリーミングします。 最も簡単な方法は、--watch を付けて送信することです。実行が最終状態に達するまで、そのログがターミナルにストリーミングされます。

databricks air run --file workload.yaml --watch -p <profile>

すでに実行中のジョブは、ID を指定してログをストリーミングします。

databricks air logs <JOB_RUN_ID> -p <profile>
# --node N (one node), --lines N (last N lines of a finished run),
# --minutes N (last N minutes), --download-to DIR (save all nodes' logs)

損失、GPU 使用率、メトリクスの履歴は MLflow で確認できます。 CLI がストリーミングするのはログとステータスであり、トレーニングメトリクスではありません。送信時またはステータスのレスポンスに MLflow の実験と実行の URL が含まれるため、それを表示して、Databricks UI で損失曲線、GPU 使用率、システムチャートを確認できるようにします。

実行が失敗した場合

ログを取得して、実際のエラーをユーザーに提示します。

databricks air logs <JOB_RUN_ID> -p <profile>

失敗箇所のトレースバックを表示し、修正方法はユーザーに判断してもらいます。トレーニングの失敗を詳しく調査する場合は、ここで根本原因を推測するのではなく、AI Runtime 独自のデバッグサポートを利用してください。

カスタム Docker イメージ

ワークロードで特定のシステムライブラリ、CUDA 拡張機能（flash-attn、apex、カスタムカーネル）、または environment.dependencies に含められない依存関係が必要な場合は、environment.version の代わりにカスタム Docker イメージを使用してください。

docker-images.md を参照してください。 手順は次のとおりです。

Databricks 提供のベースイメージの使用
Dockerfile のパターン
ビルド前の互換性チェックリスト（CUDA/ドライバー、PyTorch、NCCL、EFA/RDMA）
databricks air images push を使ったイメージのプッシュ、および environment.unity_catalog_image を使った実行

Markdown 対応

ログインしてコメント

基本情報

作者: databricks
リポジトリ: databricks/databricks-agent-skills
カテゴリ: クラウド・DevOps
インストール数: 48
GitHub スター: 373

問題を報告