GPUワークロードのスケジュール設定とタスクの作成

宣言型オートメーションバンドルを使用して、AIランタイムのGPUワークロードをスケジュールし、アップストリームのデータ準備などの他の作業と組み合わせます。YAMLでワークロードとそのスケジュールを定義し、ジョブとしてデプロイします。トレーニングパイプラインでは、CPUコンピュートで前処理を実行し、データの準備が整った後にGPUトレーニングを起動します。

このガイドは、「Hello world」を出力するワークロードの起動から始まり、スケジュールされた前処理およびトレーニングパイプラインを構築します。すでに AI ランタイムワークロード YAML がある場合は、「 AI ランタイムワークロードをバンドルに変換する」を参照してください。

仕組み

バンドル にはコードとdatabricks.yml構成が含まれています。
ジョブ はタスクをグループ化し、タスクの実行時期を定義します。
ai_runtime_task は、AIランタイムでコマンドを実行します。

depends_on フィールドはタスクを非巡回有向グラフ (DAG) に接続し、schedule フィールドはジョブを一定の頻度で実行するようにスケジュールします。

databricks bundle deploy コードをuploadし、ジョブを作成または更新します。databricks bundle runはジョブのランをすぐに起動します。各ランで、AI ランタイムは GPU コンピュートをプロビジョニングし、コマンドを実行し、名前付き MLflow エクスペリメントにランを記録します。

要件

サポートされているリージョンのワークスペース。要件を参照してください。
最新の Databricks CLI。ワークスペースに対して認証済み。これらの例を使用する前に、既存のインストールを更新してください。
前処理の例として、既存の Unity Catalog ボリューム。ジョブのラン ID には、そのボリュームに対する USE CATALOG、USE SCHEMA、READ VOLUME、および WRITE VOLUME 特権が必要です。Unity Catalog ボリュームの権限を参照してください。

Hello worldの例 `ai_runtime_task`

この例では、1つの A10 GPU でシェルコマンドを実行します。

バンドル用のディレクトリを作成します。そのディレクトリに、以下の内容でcommand.shを作成します：
Bash
```
#!/usr/bin/env bash
set -euo pipefail
echo "Hello world"
```

同じディレクトリに databricks.yml を作成します：

YAML
bundle:
  name: hello-ai-runtime

resources:
  jobs:
    hello:
      name: hello-ai-runtime
      tasks:
        - task_key: hello
          environment_key: default
          ai_runtime_task:
            experiment: hello-ai-runtime
            deployments:
              - command_path: ${workspace.file_path}/command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1
      environments:
        - environment_key: default
          spec:
            environment_version: '6'

targets:
  dev:
    mode: development
    default: true

バンドルディレクトリから、ジョブを検証、デプロイ、ランします。

Bash
databricks bundle validate --target dev
databricks bundle deploy --target dev
databricks bundle run hello --target dev

CLIによって出力されたランのURLを開き、helloタスクの出力を表示します。これにはHello worldが含まれます。ランは hello-ai-runtime MLflow エクスペリメントにも表示されます。

より複雑な例：前処理およびトレーニングパイプラインのスケジューリング

この例では、CPUコンピュートで小さなデータセットを準備し、1つのA10 GPUで線形モデルをトレーニングします。両方のタスクは、それらの間でデータを渡すために同じボリュームファイルを使用します。ジョブは毎日09:00 UTCにスケジュールされていますが、テストするまでスケジュールは停止するされます。

プロジェクトの作成

次のLayoutを持つ別のディレクトリを作成します：

Text
scheduled-training/
├── databricks.yml
├── prep.py
├── command.sh
└── src/
    └── train.py

次の内容で prep.pyを作成します。<catalog>、<schema>、および <volume> をボリュームの名前に置き換えます。このソースノートブックは入力値を正規化し、準備されたデータをボリュームに書き込みます。

Python
# Databricks notebook source
import json
from pathlib import Path

data_path = Path("/Volumes/<catalog>/<schema>/<volume>/scheduled-training/data.json")
inputs = [value / 10 for value in range(10)]
data = {"x": inputs, "y": [2 * value + 1 for value in inputs]}
data_path.parent.mkdir(parents=True, exist_ok=True)
data_path.write_text(json.dumps(data))
print(f"Prepared {len(inputs)} rows at {data_path}")

次の内容で src/train.py を作成します。prep.py と同じボリュームパスを使用します。

Python
import json
from pathlib import Path

import torch

data_path = Path("/Volumes/<catalog>/<schema>/<volume>/scheduled-training/data.json")
data = json.loads(data_path.read_text())
x = torch.tensor(data["x"], device="cuda").reshape(-1, 1)
y = torch.tensor(data["y"], device="cuda").reshape(-1, 1)
model = torch.nn.Linear(1, 1).to("cuda")
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

for _ in range(200):
    optimizer.zero_grad()
    loss = torch.nn.functional.mse_loss(model(x), y)
    loss.backward()
    optimizer.step()

print(f"Trained on {x.device}: loss={loss.item():.4f}")

トレーニングコードをランするcommand.shを作成します。以下のバンドルはsrc/をパッケージ化するため、CODE_SOURCE_PATHは展開されたsrcディレクトリを指します。
Bash
```
#!/usr/bin/env bash
set -euo pipefail
cd "$CODE_SOURCE_PATH"
python train.py
```

次の内容でdatabricks.ymlを作成します。バンドルは、GPU タスク用に src/ をパッケージ化します。prepノートブックはServerless CPUコンピュートで実行され、depends_onはprepが成功した後にのみtrainを起動します。max_concurrent_runs: 1は、このジョブのランがお互いの入力ファイルを上書きするのを防ぎます。

YAML
bundle:
  name: scheduled-training

artifacts:
  code:
    type: tgz
    path: .
    include: [src]
    files:
      - source: ./dist/code.tgz

resources:
  jobs:
    train_pipeline:
      name: scheduled-training
      max_concurrent_runs: 1
      schedule:
        quartz_cron_expression: '0 0 9 * * ?'
        timezone_id: UTC
        pause_status: PAUSED
      tasks:
        - task_key: prep
          notebook_task:
            notebook_path: ./prep.py
        - task_key: train
          depends_on:
            - task_key: prep
          environment_key: training
          ai_runtime_task:
            experiment: scheduled-training
            code_source_path: ./dist/code.tgz
            deployments:
              - command_path: ${workspace.file_path}/command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1
      environments:
        - environment_key: training
          spec:
            environment_version: '6'
            dependencies:
              - torch

targets:
  dev:
    mode: development
    default: true

スケジュールのテストと有効化

scheduled-training/から、バンドルをデプロイして手動ランを起動します。

Bash
databricks bundle validate --target dev
databricks bundle deploy --target dev
databricks bundle run train_pipeline --target dev

CLIによって出力されたランURLを開きます。trainが起動する前にprepが成功したことを確認してください。prep出力は10行の準備済み行を報告し、train出力はTrained on cuda:0と損失を報告します。
databricks.ymlで、スケジュールのpause_statusをPAUSED からUNPAUSEDに変更します。スケジュールブロックは現在次のとおりです。
YAML
```
schedule:
  quartz_cron_expression: '0 0 9 * * ?'
  timezone_id: UTC
  pause_status: UNPAUSED
```
変更をデプロイして毎日のスケジュールをアクティブ化します。
Bash
```
databricks bundle deploy --target dev
```
pause_status: UNPAUSEDを明示的に設定すると、開発ターゲットに対してもこのスケジュールが有効になります。 [ジョブとパイプライン] で、デプロイ済みのジョブを開き、スケジュールが09:00 UTCで有効になっていることを確認します。スケジュールに従ってジョブをランするを参照してください。

スケジュールされたランを一時停止するには、pause_status: PAUSEDを設定して再デプロイします。例のジョブとバンドルによってuploadされたファイルを削除するには、そのバンドルディレクトリから以下をランします。

Bash
databricks bundle destroy --target dev

ボリュームに用意したデータは残ります。scheduled-trainingデータディレクトリが不要になったら、それを削除してください。

その他のリソース

既存の AI ランタイムワークロード YAML をバンドルに変換します。AI ランタイムワークロードをバンドルに変換するを参照してください。
コードのパッケージ化、コンピュート、環境、およびタスクのパラメーターを構成します。「AIランタイムバンドルタスクの構成」を参照してください。
既存のノートブックをGPUでスケジュール設定します。スケジュールされたノートブックジョブの作成と管理を参照してください。
トレーニングのランを追跡します。エクスペリメントの追跡と観察可能性を参照してください。
モデル、オプティマイザー、データパイプラインの状態をチェックポイントします。AI ランタイムでのトレーニングのパフォーマンスと復元力の向上を参照してください。

仕組み​

要件​

Hello worldの例 ai_runtime_task​

より複雑な例：前処理およびトレーニングパイプラインのスケジューリング​

プロジェクトの作成​

スケジュールのテストと有効化​

その他のリソース​

仕組み

要件

Hello worldの例 `ai_runtime_task`

より複雑な例：前処理およびトレーニングパイプラインのスケジューリング

プロジェクトの作成

スケジュールのテストと有効化

その他のリソース