トレーニング ワークロードを本番運用化

Important

この機能は パブリック プレビュー段階です

DABを使ってAIランタイムのトレーニングワークロードをコードとして定義します。 ソース管理に留め、環境間でデプロイし、スケジュールを組み、他のタスクと組み合わせて管理しましょう。 このページでは、ai_runtime_task がサーバーレス GPU コンピュート上のディレクトリに対して自分のコマンドを実行する「自分でトレーニングをする」パスについて説明しています。

ヒント

  • 宣言的自動化バンドルを使って、トレーニングワークロードをコードとして定義し、環境間で展開し、スケジュールを組みましょう。
  • ai_runtime_taskはコードディレクトリに対して自分のコマンドを実行します(「自分のトレーニングを持ち込む」と)。
  • GPUとCPUのタスクをマルチタスクジョブで組み合わせます。

これはサーバーレスGPUでノートブックをバンドル経由で動かすのとは別の作業です。 基本的なノートブック・オン-GPUバンドルの例については、「 Schedule with the Jobs API」および「Declarative Automation Bundles」を参照してください。

Requirements

  • AIランタイムが有効になったワークスペース。 要件を参照してください。
  • バンドルをデプロイするための Databricks CLI(コマンドライン インターフェース)がインストールおよび構成済みであること。

バンドル内でAIランタイムタスクを定義します

ai_runtime_taskは実験の名前を付け、code_source_pathでトレーニングコードを指し示し、1つのデプロイメントを宣言します:実行コマンドと実行するGPUです。 バンドル内のジョブにそれを追加してください:

resources:
  jobs:
    train:
      tasks:
        - task_key: train
          ai_runtime_task:
            experiment: my-experiment
            code_source_path: ./dist/code.tgz
            deployments:
              - command_path: ./command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1

code_source_path パッケージされたトレーニングコードを指し示し、 command_path タスクが実行するスクリプトです。 再試行、タイムアウト、権限は、Azure Databricksのジョブと同じようにタスクやジョブで設定されるため、既存のバンドル慣行が引き継がれます。 コードのパッケージ化や参照方法については、「 トレーニングコードを出荷」をご覧ください。

ai_runtime_task フィールド

フィールド タイプ Description
experiment String Required. そのランの MLflow 実験名 「 実験の追跡と観測可能性」を参照してください。
code_source_path String 実行する訓練コード:パッケージ化された tgz アーティファクトの出力ファイル、またはすでにアップロードされたコードへの /Workspace または /Volumes パスです。 「 トレーニングコードを発送する」をご覧ください。
deployments 順番 Required. コマンドとそれを実行する計算を記述した単一のデプロイメントです。 各エントリーには command_pathcompute、そして任意の nameが含まれています。
deployments[].command_path String Required. タスクのスクリプトは各ノードで実行されます。
deployments[].compute.accelerator_type String Required. GPUの種類、例えば GPU_1xA10GPU_1xH100GPU_8xH100などです。
deployments[].compute.accelerator_count Integer Required. すべてのノードにまたがるGPUの総数、つまり accelerator_typeにエンコードされたノードあたりの数の倍数です。
deployments[].name String デプロイメントのオプション名で、ログやUIで使われます。
docker_image_url String 管理環境ではなく、コマンドを実行するためのオプションのカスタムDockerイメージ。 「カスタム Docker イメージを使用する」を参照してください。
mlflow_run String MLflow実行用のオプション表示名です。
mlflow_experiment_directory String 実験を作成するオプションのワークスペースディレクトリ。 /Workspaceで始まる必要があります。 デフォルトユーザーディレクトリを持たないサービスプリンシパルとして実行するときに設定してください。
mlflow_artifact_location String MLflowアーティファクトのための任意のルート位置、例えば /Volumes/<catalog>/<schema>/<volume>/… パスなど。 既存の実験の遺物の位置と一致しなければ、除外されます。

やり直し、タイムアウト、権限、環境(environment_key)をタスクや作業に設定し、 ai_runtime_taskの中ではなく。 完全なタスク参照については、 AI Runtimeタスクを参照してください。

ハードウェアアクセラレータの設定

accelerator_typeは必要なGPUに設定し、accelerator_countはGPUの総数にします。 ノードあたりのGPU数の倍数で、GPU_1xA10GPU_1xH1001GPU_8xH1008です。 ノードあたりのサイズより大きいカウントは複数のノードにまたがるタスクを実行します。例えば、accelerator_count: 16GPU_8xH100は2つのノードで実行されます。 アクセラレーターの選択に関する指針については、「 ハードウェアオプション」をご覧ください。

Note

マルチノード実行の場合、AI Runtimeはすべてのノードでコマンドを実行し、タスク環境内の標準的な分散訓練環境変数(NUM_NODESWORLD_SIZELOCAL_WORLD_SIZEMASTER_ADDRMASTER_PORT)を入力します。 コマンドから読み取る(例えば torchrun ローンチ)。バンドル内で設定するわけではありません。

環境と依存関係を設定します

ジョブで environments ブロックを宣言し、environment_key でタスクから参照します。 AI Runtimeは、コマンドを実行する前にリストに記載された依存関係をインストールします:

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: '6'
            dependencies:
              - numpy

利用可能な環境については、「 環境の設定」をご覧ください。

トレーニングコードを発送してください

code_source_path は、学習コードがどこにあるかをタスクに指定します。 それは2つの形態のいずれかを取る。

  • パッケージ化された tgz アーティファクト — アーティファクトを宣言し、その出力ファイルを code_source_path 指す。 Azure Databricksターボールを作り、databricks bundle deployにアップロードします。 これはローカルプロジェクトディレクトリやコミットされたGitリビジョンからコードを送る方法です。
  • ワークスペースやボリュームパス — すでにアップロードされ、as-is使われているコード。

パッケージされたアーティファクト

tgzアーティファクトを宣言し、code_source_pathその出力ファイルを指し示します。 databricks bundle deployでは、CLIがtarballを作成しアップロードし、タスクがそれを抽出してあなたのコマンドを実行します。

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を使って作業木からファイルをパッケージ化したり、gitでコミットされたブランチやコミットをスナップショットで取得したりできます。

ワークスペースまたはボリュームパス

すでにアップロードされているコードを使う場合は、 code_source_path/Workspace/…/Volumes/… パスに設定してください。 Azure Databricks はそのパスをそのまま使用し、一切パッケージ化しません。

tgzアーティファクトフィールド:

フィールド Description
type tgz buildコマンドを実行する代わりに、ソースファイルからgzip化されたtarballを構築します。
path パッケージする基本ディレクトリ。 include パスやアーカイブのエントリ名はそれに相対的なものです。
include パッケージ化する path のサブパスの一覧。 path のすべてをパッケージ化しない。 .gitignore は尊重されます; バンドル全体の sync.includesync.exclude は適用されません。 buildコマンドの代替手段です。
git 作業ツリーの代わりにコミット済みのGit参照をスナップショットしてください。 git.branch または git.commit を設定します(両方が設定されている場合は commit が優先されます)。 buildコマンドの代替手段です。
files[].source 築かれたタールボールの道。 これにcode_source_pathを向けてください。

Note

AI Runtimeはコードをディレクトリに抽出し、 CODE_SOURCE_PATH 環境変数として公開します。 相対パスが正しく解決されるように、スクリプトを実行する前に、コマンドからそれを参照してください。たとえば cd "$CODE_SOURCE_PATH" のようにします。

完全なコード例

この例はローカルプロジェクトの単一のA10 GPUでトレーニングしており、事前のCLI設定はAzure Databricks CLIのインストールと設定以外にありません。 このプロジェクトには3つのファイルがあります:

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

command.shcommand_pathによって指定された入口です。 抽出されたコードディレクトリに切り替え、トレーニングスクリプトを実行します:

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

databricks.yml バンドル名を付け、 src/tgz アーティファクトとしてパッケージ化し、 ai_runtime_taskとして実行し、タスク環境に numpy をインストールし、開発および本番の目標を定義します。

bundle:
  name: my-training

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

resources:
  jobs:
    train:
      name: my-training
      tasks:
        - task_key: train
          environment_key: default
          ai_runtime_task:
            experiment: /Users/me@example.com/my-training
            code_source_path: ./dist/code.tgz
            deployments:
              - command_path: ./command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1
      environments:
        - environment_key: default
          spec:
            environment_version: '6'
            dependencies:
              - numpy

targets:
  dev:
    mode: development
    default: true
  prod:
    mode: production

バンドルをデプロイしてジョブを実行してください。 databricks bundle deploy tgzアーティファクトを作成・アップロードし、ジョブを作成する。databricks bundle run開始:

databricks bundle deploy --target dev
databricks bundle run train --target dev

マルチタスクワークフローを構築する

ai_runtime_task は Azure Databricks のジョブ タスクであるため、ジョブの他の部分と組み合わせて使用できます。 トレーニング前に準備ステップを実行し、GPUとCPUのタスクを1つのジョブで統合し、タスクごとに異なるアクセラレータを使うことができます。

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: ./dist/code.tgz
            deployments:
              - command_path: ./command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1

GPUとCPUのタスクを組み合わせる

上記のパイプラインでは、GPUが必要なのはトレーニングステップだけです。 データ準備などのGPU以外の作業を別々のタスクに分けておくことで、GPUの時間はトレーニングに集中できます。

Note

ai_runtime_taskはAzure Databricksジョブタスク値({{tasks.<task_key>.values.<name>}}またはdbutils.jobs.taskValues)をサポートしません。 ステップ間でデータを渡すには、Unityカタログのボリュームやワークスペースファイルなど、両方のタスクが読み取れる共有場所に書き込み、そのパスを参照します。

課題のスケジュールを組む

ジョブに 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 Runtime CLIを使ってコマンドラインからAIランタイムワークロードを実行し管理できます。
  • トレーニングランの記録やチェックポイントの管理。 「 実験の追跡と観測可能性」を参照してください。