トレーニングワークロードの本番運用化
プレビュー
この機能は パブリック プレビュー段階です。
DABs を使用して、AI ランタイム のトレーニング ワークロードをコードとして定義します。ソース管理に維持し、環境間でデプロイし、スケジュールし、他のタスクと組み合わせます。このページでは、独自トレーニングのパスについて説明します。このパスでは、ai_runtime_task が Serverless GPU コンピュート上のコードのディレクトリに対して独自のコマンドを実行します。
これは、バンドルを介してServerless GPU上でノートブックを実行するタスクとは異なります。基本的な notebook-on-GPU バンドルの例については、Jobs API と宣言型オートメーションバンドルを参照してください。
要件
- AI Runtimeが有効になっているワークスペース。要件を参照してください。
- バンドルをデプロイするためにインストールおよび構成された Databricks CLI(コマンドラインインターフェイス)。
バンドルでAI Runtimeタスクを定義する
1つのai_runtime_taskがエクスペリメントに名前を付け、code_source_pathでトレーニングコードを指定し、1つのデプロイメント(実行するコマンドとそれを実行するGPU)を宣言します。バンドル内のジョブに追加します。
resources:
jobs:
train:
tasks:
- task_key: train
ai_runtime_task:
experiment: my-experiment
code_source_path: ./src
deployments:
- command_path: ./command.sh
compute:
accelerator_type: GPU_1xA10
accelerator_count: 1
command_path は、タスクが実行するスクリプトです。再試行、タイムアウト、およびアクセス許可は、他の Databricks ジョブと同様にタスクとジョブで設定されるため、既存のバンドルプラクティスを引き継ぐことができます。
ハードウェアアクセラレータを設定する
accelerator_type にはワークロードに必要な GPU を設定し、accelerator_count には GPU の総数を設定します。カウントは、ノードあたりの GPU 数の倍数です。GPU_1xA10 と GPU_1xH100 の場合は 1、GPU_8xH100 の場合は 8 になります。ノードあたりのサイズよりも大きいカウントの場合は、タスクが複数のノードで実行されます。たとえば、GPU_8xH100 と accelerator_count: 16 の場合は 2 つのノードで実行されます。アクセラレータの選択に関するガイダンスについては、ハードウェアオプションを参照してください。
マルチノードランの場合、AI Runtimeはすべてのノードでコマンドを実行し、タスク環境内に標準の分散トレーニング環境変数(NUM_NODES、WORLD_SIZE、LOCAL_WORLD_SIZE、MASTER_ADDR、MASTER_PORT)を設定します。コマンドからそれらを読み込みます(例:torchrunの起動)。バンドルでは設定しません。
環境と依存関係を設定する
ジョブで environments ブロックを宣言し、タスクから environment_key で参照します。AI ランタイムは、コマンドを実行する前にリストされた依存関係をインストールします:
resources:
jobs:
train:
tasks:
- task_key: train
environment_key: default
ai_runtime_task:
# experiment, code_source_path, and deployments as above
environments:
- environment_key: default
spec:
environment_version: '5'
dependencies:
- numpy
利用可能な環境については、環境の設定を参照してください。
トレーニングコードを出荷する
code_source_path トレーニングコードの場所をタスクに伝えます。次の 3 つの形式を受け入れます:
- ローカル ディレクトリ — CLI はそれをパッケージ化し、
databricks bundle deployに upload します。これが最も単純な形式であり、上記の例で使用されているものです。 - 明示的な
tgzアーティファクト — アーティファクトを自分で宣言し、その出力ファイルを指定します。ファイルの一部のみをパッケージ化する必要がある場合や、作業ツリーの代わりにコミットされたGitリビジョンをスナップショットする必要がある場合に使用します。 - ワークスペースまたはボリュームのパス — すでに upload されており、そのまま使用されるコード。
- Local directory
- Packaged artifact
- Workspace or volume path
code_source_pathをディレクトリに向けます。databricks bundle deploy上で、CLIはディレクトリをパッケージ化してuploadし、タスクがそれを展開してコマンドを実行します。
code_source_path: ./src
tgz アーティファクトを宣言し、code_source_path がその出力ファイルを指すようにします。databricks bundle deploy では、CLI は tarball をビルドして upload し、タスクがそれを解凍して、それに対してコマンドを実行します:
artifacts:
code:
type: tgz
path: .
include: [src]
files:
- source: ./dist/code.tgz
resources:
jobs:
train:
tasks:
- task_key: train
ai_runtime_task:
code_source_path: ./dist/code.tgz
作業ツリーからファイルをパッケージ化するには include を使用し、コミットされた Branch または commit のスナップショットを作成するには git を使用します。
すでにuploadされたコードを使用するには、code_source_pathを/Workspace/…または/Volumes/…のパスに設定します。Databricksはパスをそのまま使用し、何もパッケージ化しません。
AI Runtime はコードをディレクトリに抽出し、CODE_SOURCE_PATH 環境変数として公開します。スクリプトを実行する前に、相対パスが解決されるようにコマンドから参照します(例:cd "$CODE_SOURCE_PATH")。
マルチタスク ワークフローを構築する
ai_runtime_taskはDatabricksのジョブタスクであるため、ジョブの他の部分と構成されます。トレーニングの前に準備ステップを実行し、1 つのジョブで GPU タスクと CPU タスクを組み合わせ、タスクごとに異なるアクセラレータを使用できます。
depends_on でタスクの順序を指定する
タスクを順番に実行するには、depends_onを使用します。次のパイプラインでは、準備ノートブックを実行し、その後に準備タスクが成功した場合にのみ起動する GPU トレーニングタスクを実行します:
resources:
jobs:
train_pipeline:
tasks:
- task_key: prep
notebook_task:
notebook_path: ./prep.py
- task_key: train
depends_on:
- task_key: prep
ai_runtime_task:
experiment: my-experiment
code_source_path: ./src
deployments:
- command_path: ./command.sh
compute:
accelerator_type: GPU_1xA10
accelerator_count: 1
GPU と CPU のタスクを組み合わせる
上記のパイプラインでは、トレーニングステップのみに GPU が必要です。データ準備などの非 GPU の作業を別のタスクに維持することで、GPU の時間をトレーニングに集中させることができます。
ai_runtime_task は、Databricks ジョブのタスク値({{tasks.<task_key>.values.<name>}} または dbutils.jobs.taskValues)をサポートしていません。ステップ間でデータを渡すには、Unity Catalog ボリュームやワークスペースファイルなど、両方のタスクが読み取ることができる共有場所にデータを書き込み、各タスクからそのパスを参照します。
ワークロードをスケジュールする
定期的に実行するようにジョブにscheduleを追加します。バンドルのデプロイによって自動的にランが開始されないようにスケジュールを停止した状態で出荷し、準備ができたら停止を解除します。
resources:
jobs:
train_pipeline:
schedule:
quartz_cron_expression: '0 0 9 * * ?'
timezone_id: UTC
pause_status: PAUSED
開発から本番運用へプロモートする
開発から本番運用への昇格は、AI ランタイムタスクがそのまま継承する標準的なバンドル機能です。
バンドルのターゲットとモード
開発ターゲットとともに、mode: productionを使用して本番運用ターゲットを定義します。ターゲットによって、バンドルのデプロイ先とそのリソースの命名方法が制御されます。
targets:
dev:
mode: development
default: true
prod:
mode: production
デプロイしてラン
標準のバンドルコマンドを使用して、ターゲットに対してバンドルをデプロイおよび実行します:
databricks bundle deploy --target dev
databricks bundle run train_pipeline --target dev
次のステップ
- AI ランタイム CLIを使用して、コマンドラインからAI ランタイムワークロードをランおよび管理します。
- トレーニングのランを追跡し、チェックポイントを管理します。エクスペリメントの追跡と観察可能性を参照してください。