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

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

要件

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

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

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

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

カスタムトレーニングジョブで 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 の割り当てと上限をご覧ください。

課金

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

次のステップ

カスタムトレーニングジョブのコンピューティングリソースの構成の詳細を確認する。
カスタムトレーニングジョブで分散トレーニングを使用する方法の詳細を確認する。
その他のスケジュール設定オプションを確認する。

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

おすすめのユースケース

要件