GPU ワークロードのスケジュールとタスクの構成

備考

プレビュー

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

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

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

仕組み

バンドル には、コードとdatabricks.ymlの設定が含まれています。
ジョブ は、タスクをグループ化し、その実行タイミングを定義します。
ai_runtime_task は、Serverless GPU コンピュート上でコマンドを実行します。

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で線形モデルをトレーニングします。両方のタスクで同じボリュームファイルを使用して、その間でデータを渡します。ジョブは毎日協定世界時 (UTC) 09:00に実行されるようにスケジュールされており、テストするまでスケジュールは停止する。

プロジェクトを作成する

次の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上でスケジュールします。Jobs APIおよびDeclarative Automation Bundlesでのスケジュールを参照してください。
トレーニングのランを追跡します。エクスペリメントの追跡と観察可能性を参照してください。
モデル、オプティマイザー、およびデータパイプラインの状態をチェックポイントに保存します。AI ランタイムでのトレーニングのパフォーマンスと耐障害性の向上を参照してください。

仕組み​

要件​

次を使用したHello worldの例 ai_runtime_task​

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

プロジェクトを作成する​

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

その他のリソース​

仕組み

要件

次を使用したHello worldの例 `ai_runtime_task`

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

プロジェクトを作成する

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

その他のリソース