リソースの可用性に基づいてトレーニング ジョブをスケジュールする

GPU リソースをリクエストするカスタム トレーニング ジョブの場合、Dynamic Workload Scheduler を使用すると、リクエストされた GPU リソースが使用可能になったタイミングに基づいてジョブをスケジュールできます。このページでは、Dynamic Workload Scheduler を使用してカスタム トレーニング ジョブをスケジュールする方法と、Vertex AI でスケジューリング動作をカスタマイズする方法について説明します。

次の状況では、Dynamic Workload Scheduler を使用してカスタム トレーニング ジョブをスケジュールすることをおすすめします。

  • カスタム トレーニング ジョブが A100 または H100 GPU をリクエストし、リクエストされたリソースが利用可能になり次第ジョブを実行する。たとえば、Vertex AI がピーク時間外に GPU リソースを割り当てる場合などです。
  • ワークロードに複数のノードが必要であり、すべての GPU ノードがプロビジョニングされて同時に準備が整うまでワークロードの実行を開始できない。たとえば、分散トレーニング ジョブを作成している場合です。

要件

Dynamic Workload Scheduler を使用するには、カスタム トレーニング ジョブが次の要件を満たしている必要があります。

  • カスタム トレーニング ジョブの timeout は最大 7 日以内にする必要があります。
  • カスタム トレーニング ジョブでは、すべてのワーカープールで同じマシン構成が使用されます。

サポートされているジョブタイプ

CustomJobHyperparameterTuningjobTrainingPipeline など、すべてのカスタム トレーニング ジョブタイプがサポートされています。

カスタム トレーニング ジョブで Dynamic Workload Scheduler を有効にする

カスタム トレーニング ジョブで Dynamic Workload Scheduler を有効にするには、ジョブを作成するときに scheduling.strategy API フィールドを FLEX_START に設定します。

カスタム トレーニング ジョブを作成する方法については、次のリンクをご覧ください。

リソースの可用性を待機する時間を構成する

ジョブがリソースを待機できる時間を scheduling.maxWaitDuration フィールドで構成できます。値が 0 の場合、リクエストされたリソースが使用可能になるまでジョブは無期限に待機します。デフォルト値は 1 日です。

次の例は、customJob で Dynamic Workload Scheduler を有効にする方法を示しています。使用するインターフェースのタブを選択します。

gcloud

Google Cloud CLI を使用してジョブを送信する場合は、config.yaml ファイルに scheduling.strategy フィールドを追加します。

YAML 構成ファイルの例:

workerPoolSpecs:
  machineSpec:
    machineType: n1-highmem-2
  replicaCount: 1
  containerSpec:
    imageUri: gcr.io/ucaip-test/ucaip-training-test
    args:
    - port=8500
    command:
    - start
scheduling:
  strategy: FLEX_START
  maxWaitDuration: 1800s

python

Vertex AI SDK for Python を使用してジョブを送信する場合は、関連する CustomJob 作成方法で scheduling_strategy フィールドを設定します。

from google.cloud.aiplatform_v1.types import custom_job as gca_custom_job_compat

def create_custom_job_with_dws_sample(
    project: str,
    location: str,
    staging_bucket: str,
    display_name: str,
    script_path: str,
    container_uri: str,
    service_account: str,
    experiment: str,
    experiment_run: Optional[str] = None,
) -> None:
    aiplatform.init(project=project, location=location, staging_bucket=staging_bucket, experiment=experiment)

    job = aiplatform.CustomJob.from_local_script(
        display_name=display_name,
        script_path=script_path,
        container_uri=container_uri,
        enable_autolog=True,
    )

    job.run(
        service_account=service_account,
        experiment=experiment,
        experiment_run=experiment_run,
        max_wait_duration=1800,
        scheduling_strategy=gca_custom_job_compat.Scheduling.Strategy.FLEX_START
    )

REST

Vertex AI REST API を使用してジョブを送信する場合は、カスタム トレーニング ジョブを作成するときに scheduling.strategy フィールドと scheduling.maxWaitDuration フィールドを設定します。

リクエストの JSON 本文の例:

{
  "displayName": "MyDwsJob",
  "jobSpec": {
    "workerPoolSpecs": [
      {
        "machineSpec": {
          "machineType": "a2-highgpu-1g",
          "acceleratorType": "NVIDIA_TESLA_A100",
          "acceleratorCount": 1
        },
        "replicaCount": 1,
        "diskSpec": {
          "bootDiskType": "pd-ssd",
          "bootDiskSizeGb": 100
        },
        "containerSpec": {
          "imageUri": "python:3.10",
          "command": [
            "sleep"
          ],
          "args": [
            "100"
          ]
        }
      }
    ],
    "scheduling": {
      "maxWaitDuration": "1800s",
      "strategy": "FLEX_START"
    }
  }
}

割り当て

Dynamic Workload Scheduler を使用してジョブを送信すると、Vertex AI はオンデマンドの Vertex AI 割り当てではなく、プリエンプティブル割り当てを使用します。たとえば、Nvidia H100 GPU の場合は、次のように使用します。

aiplatform.googleapis.com/custom_model_training_nvidia_h100_gpus

Vertex AI は次のリソースを使用します。

aiplatform.googleapis.com/custom_model_training_preemptible_nvidia_h100_gpus

ただし、プリエンプティブ 割り当ては名前でのみ使用されます。リソースはプリエンプティブルではなく、標準リソースのように動作します。

Dynamic Workload Scheduler を使用してジョブを送信する前に、プリエンプティブ クォータが十分な量に増加していることを確認してください。Vertex AI の割り当てと割り当て増加リクエストの手順については、Vertex AI の割り当てと上限をご覧ください。

課金

課金されるのは、ジョブの実行時間のみで、リソースが使用可能になるまでジョブが待機している時間は課金されません。詳細については、料金をご覧ください。

次のステップ