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

Important

この機能は パブリック プレビュー段階です。 これを使用するには、ワークスペース管理者が プレビュー ページからAIランタイムプレビューを有効にする必要があります。 Manage Azure Databricks プレビューを参照してください。

宣言的自動化バンドルを使ってAIランタイムGPUワークロードをスケジュール化し、上流のデータ準備など他の作業と組み合わせることができます。 YAMLでワークロードとそのスケジュールを定義し、ジョブとしてデプロイします。 トレーニングパイプラインの場合は、CPUの計算で前処理を実行し、データが準備できたらGPUのトレーニングを開始します。

このガイドは「Hello world」を印刷する作業から始まり、スケジュールされた前処理とトレーニングパイプラインを構築します。 すでに AI Runtime workload の YAML をお持ちの場合は、「AI Runtime ワークロードをバンドルに変換する」を参照してください。

どのように機能するのか

  • バンドルにはコードと databricks.yml 設定が含まれています。
  • ジョブはタスクをグループ化し、実行時期を定義します。
  • ai_runtime_task はサーバーレス GPU コンピューティングでコマンドを実行します。

depends_onフィールドはタスクを有向非巡回グラフ(DAG)に結びつけ、scheduleフィールドはジョブを定期的に実行するようにスケジューリングします。

databricks bundle deploy コードをアップロードし、ジョブを作成または更新します。 databricks bundle run はジョブ実行をすぐに開始します。 各実行時に、AI RuntimeはGPUコンピュートをプロビジョニングし、あなたのコマンドを実行し、名前付きのMLflow実験で実行を記録します。

Requirements

Hello Worldの例 ai_runtime_task

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

  1. バンドル用のディレクトリを作成します。 そのディレクトリ内で、以下の内容の command.sh を作成してください。

    #!/usr/bin/env bash
set -euo pipefail
echo "Hello world"

  2. 同じディレクトリ内で databricks.yml を作成してください:

    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

  3. バンドルディレクトリから、検証し、デプロイし、ジョブを実行します:

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

  4. CLI に表示された実行 URL を開き、hello タスクの出力を表示します。 それにはHello worldが含まれています。 このランは hello-ai-runtime MLflow の実験にも表示されます。

より複雑な例: 前処理とトレーニング パイプラインをスケジュールする

この例では、CPU計算で小さなデータセットを作成し、1つのA10 GPU上で線形モデルを訓練します。 両方のタスクは同じボリュームファイルを使ってデータをやり取りします。 作業は毎日09:00 UTCに予定されており、テストするまでスケジュールは一時停止されます。

プロジェクトを作成する

  1. 以下のレイアウトで別のディレクトリを作成します:

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

  2. 以下の内容で prep.py を作成しましょう。 <catalog>、<schema>、<volume>をお使いのボリューム名に置き換えてください。 このソースノートブックは入力値を正規化し、準備済みデータをボリュームに書き込みます。

    # 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}")

  3. 以下の内容で src/train.py を作成しましょう。 prep.pyと同じボリュームパスを使いましょう:

    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}")

  4. command.sh を作成してトレーニング コードを実行します。 以下のバンドルは src/ を同梱しているので、CODE_SOURCE_PATH は抽出された src ディレクトリを指しています:

    #!/usr/bin/env bash
set -euo pipefail
cd "$CODE_SOURCE_PATH"
python train.py

  5. 以下の内容で databricks.yml を作成しましょう。 バンドルには GPU タスク用の src/ が含まれます prepノートブックはサーバーレスCPUコンピュートで動作し、depends_on は train が成功した後でのみ prep を開始します。 max_concurrent_runs: 1 このジョブの実行が互いの入力ファイルを上書きしないようにします:

    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

スケジュールをテストして有効にする

  1. scheduled-training/ からバンドルをデプロイし、手動実行を開始します。

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

  2. CLIで表示されたrunのURLを開きます。 prepが始まる前にtrainが成功しているか確認してください。 prep 出力は準備された行数が 10 であることを報告し、train 出力は Trained on cuda:0 と損失を報告します。

  3. databricks.ymlでは、スケジュールのpause_statusをPAUSEDからUNPAUSEDに変更します。 現在のスケジュールブロックは以下の通りです:

    schedule:
  quartz_cron_expression: '0 0 9 * * ?'
  timezone_id: UTC
  pause_status: UNPAUSED

  4. 日替わりのスケジュールを有効にするために変更をデプロイする:

    databricks bundle deploy --target dev

    pause_status: UNPAUSED を明示的に設定することで、開発ターゲットでもこのスケジュールが有効になります。 Jobs & Pipelinesで、展開済みジョブを開き、09:00 UTCにスケジュールが有効であることを確認してください。 「スケジュールに基づいてジョブを実行する」を参照してください。

予定された実行を一時停止するには、pause_status: PAUSED に設定して再度デプロイしてください。 サンプルジョブとバンドルでアップロードされたファイルを削除するには、バンドルディレクトリから以下を実行します。

databricks bundle destroy --target dev

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

その他のリソース

  • 最終更新: