Google は AI 技術を使用して、コンテンツをご希望の言語に翻訳しています。AI 翻訳には誤りが含まれる場合があります。

DWS と Kueue を使用して TPU で AI トレーニングを最適化する

Standard

このドキュメントでは、Tensor Processing Unit（TPU）で AI トレーニングの可用性を最大化するように GKE クラスタを構成する方法について説明します。Kueue と呼ばれるオープンソースのジョブキューイングツールと Google Cloudの Dynamic Workload Scheduler（DWS）を使用して、自動フォールバックシステムを構成します。

このドキュメントで定義する構成では、プライマリノードプールとノードのバックアッププールが設定されます。

オンデマンド（プラン A）: このノードプールは、ノードの最初の選択肢です。システムは、まずこれらのオンデマンドマシンでジョブのスケジュール設定を試みます。オンデマンドという用語は、これらのマシンが実行を開始すると、信頼性の高い中断のないアクセスが可能になることを意味します。
DWS Flex Start（プラン B）: バックアップノードプールです。プラン A のマシンが使用できない場合、Kueue スケジューリングプログラムはジョブをこのプラン B プールに自動的に割り当てます。DWS はプラン B のハードウェアを検索しますが、そのハードウェアも使用できない可能性があるため、すぐにアクセスできるとは限りません。ただし、DWS はリクエストを最大 7 日間キューに保持し、マシンが使用可能になるとすぐに自動的に提供します。

このアプローチにより、ジョブがキューで待機する時間を最小限に抑えることができます。つまり、利用可能なリソースを手動で確認したり、異なるマシン用にスクリプトを書き直したりする必要はありません。

構成手順の概要

自動フォールバックシステムを構成するには、いくつかの構成手順を完了する必要があります。この構成は、次の 2 つのカテゴリに分けると便利です。

クラスタ管理者のタスク: GKE クラスタの作成、ノードプールのプロビジョニング、Kueue スケジューリングコントローラのインストールなど、インフラストラクチャの 1 回限りの構成。
AI デベロッパーのタスク: トレーニングジョブの要件の定義やワークロードの送信など、繰り返し行われる日常的なワークフロー。

これらの手順をすべて自分で行う場合でも、この区別を念頭に置くことで、プロセス全体を明確にできます。

システムを構成する前に、実行する構成手順を確認します。

トピック	タスク
インフラストラクチャを構成する（クラスタ管理者）	1. GKE クラスタを作成します 2。ノードプールを作成します。 3. ノードプールで Flex-Start のステータスを確認する
Kueue をインストールして構成する（クラスタ管理者）	1. Kueue をインストールする 2. 構成ルールを定義する
トレーニングジョブを実行する（AI デベロッパー）	1. ConfigMap を作成します。 2. RayJob マニフェストを定義します。 3. ワークロードを送信します 4。RayJob に接続する 5. ログを調べる

主なコンセプト

オンデマンド（プラン A）ノードプール: プライマリの優先度の高いノードプール。ジョブは常に最初にこのプールを使用しようとします。
DWS Flex Start（プラン B）ノードプール: バックアップノードプール。プライマリプールのマシンが使用できない場合、システムはこのプールを自動的に使用して、使用可能なハードウェアを検索します。
Kueue: ジョブキューを管理するスケジューリングプログラム。ジョブリクエストをインターセプトし、使用するノードプール（プラン A またはプラン B）を決定します。
ジョブ: 実行する AI トレーニングワークロード。このドキュメントでは、RayJob マニフェストを使用して定義します。

始める前に

Google Cloud コンソールのプロジェクトセレクタページで、 Google Cloud プロジェクトを選択または作成します。
プロジェクトの選択または作成に必要なロール
- プロジェクトを選択する: プロジェクトの選択に特定の IAM ロールは必要ありません。ロールが付与されているプロジェクトであれば、どのプロジェクトでも選択できます。
- プロジェクトを作成する: プロジェクトを作成するには、resourcemanager.projects.create 権限を含むプロジェクト作成者ロール（roles/resourcemanager.projectCreator）が必要です。詳しくは、ロールを付与する方法をご覧ください。
注: この手順で作成するリソースをそのまま保持する予定でない場合、既存のプロジェクトを選択するのではなく、新しいプロジェクトを作成してください。チュートリアルの終了後にそのプロジェクトを削除すれば、プロジェクトに関連するすべてのリソースを削除できます。

プロジェクトセレクタに移動
Google Cloud プロジェクトで課金が有効になっていることを確認します。
Google Kubernetes Engine API と Cloud TPU API が有効になっていない場合は、有効にします。
API を有効にするために必要なロール
API を有効にするには、serviceusage.services.enable 権限が必要です。プロジェクトを作成した場合は、オーナーロール（roles/owner）を介してこの権限がすでに付与されている可能性があります。それ以外の場合は、Service Usage 管理者ロール（roles/serviceusage.serviceUsageAdmin）を介してこの権限を取得できます。ロールを付与する方法を確認する。
API を有効にする

Google Cloud コンソールで Cloud Shell をアクティブにします。

Cloud Shell をアクティブにする
TPU Flex Start VM を使用するために十分なプリエンプティブル割り当てがあることを確認します。デフォルトの割り当てがニーズを満たしていない場合は、割り当ての引き上げをリクエストします。詳細については、Cloud TPU の割り当てをご覧ください。

環境変数を定義する

このドキュメントで実行するコマンドを簡略化するために、Cloud Shell で環境変数を設定できます。これらの変数には、 Google Cloud プロジェクトの ID、ノードプールの名前、GKE クラスタのロケーションなどの値が格納されます。

これらの変数を定義すると、値を毎回再入力または置き換えるのではなく、変数名（$CLUSTER_NAME など）を参照して、複数のコマンドで再利用できます。このアプローチにより、プロセスがわかりやすくなり、エラーのリスクが軽減されます。

Cloud Shell で次の環境変数を定義するには、次のコマンドを実行します。

export PROJECT_ID=$(gcloud config get project)
export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
export ZONE="us-east5-b"
export REGION="us-east5"
export CLUSTER_NAME="tpu-cluster"
export GKE_VERSION="1.34"
export ONDEMAND_NODEPOOL="on-demand-pool"
export DWS_NODEPOOL="dws-pool"

これらの環境変数の説明は次のとおりです。

PROJECT_ID: 実際の Google Cloud プロジェクト ID。
PROJECT_NUMBER: プロジェクトの一意の識別子番号（例: 123456789012）。
ZONE: クラスタのコンピューティングゾーン（us-east5-b など）。選択したアクセラレータタイプが使用可能なゾーンを選択します。可用性については、Cloud TPU の割り当てまたは GPU の割り当てをご覧ください。
REGION: クラスタリソースを作成するリージョン（us-east5 など）。
CLUSTER_NAME: GKE クラスタに付ける名前。
GKE_VERSION: クラスタの GKE バージョン。バージョン 1.34 以降を使用します。
ONDEMAND_NODEPOOL: 標準オンデマンドノードプールの名前。（これがプラン A のノードプールです）。
DWS_NODEPOOL: DWS Flex Start ノードプールの名前。（これはプラン B のノードプールです）。

インフラストラクチャを構成する（クラスタ管理者）

クラスタ管理者は、フォールバックメカニズムをサポートするように GKE クラスタとノードプールを構成します。

GKE クラスタを作成する

まず、GKE クラスタを作成します。このクラスタは、Kueue コントローラをインストールし、ノードプールを構成して、AI トレーニングジョブを実行する環境です。クラスタを作成して接続する手順は次のとおりです。

クラスタを作成します。
```
gcloud container clusters create ${CLUSTER_NAME} \
  --cluster-version=${GKE_VERSION} \
  --machine-type=n2-standard-16 \
  --location=${ZONE} \
  --enable-image-streaming \
  --addons=RayOperator \
  --project=${PROJECT_ID}
```
このコマンドでは、次のキーフラグを使用します。
- --addons=RayOperator: Ray Operator をクラスタにインストールします。このオペレーターは、このドキュメントで後ほど送信する RayJob ワークロードを管理するために必要です。
- --enable-image-streaming: クラスタがコンテナイメージをより高速に pull できるようにします。この機能により、大規模な AI コンテナイメージの実行開始にかかる時間が大幅に短縮されます。
kubectl CLI がクラスタに接続できるように、クラスタの認証情報を取得します。このコマンドは、デフォルトで ~/.kube/config ディレクトリに保存されている Kubernetes 構成ファイルを更新します。
```
gcloud container clusters get-credentials ${CLUSTER_NAME} \
  --location=${ZONE} \
  --project=${PROJECT_ID}
```

ノードプールを作成する

環境のプライマリノードプールとバックアップノードプール（オンデマンドノードプール（プラン A）と DWS Flex Start ノードプール（プラン B））を作成します。

オンデマンドノードプールを作成する: このプールは、トレーニングジョブのプライマリリソースとして機能します。
```
gcloud container node-pools create ${ONDEMAND_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --reservation-affinity=none \
  --enable-autoscaling \
  --num-nodes=0 \
  --min-nodes=0 \
  --max-nodes=4
```
この例では、優先度の高いマシンは TPU v6e アクセラレータです。このハードウェアは --machine-type=ct6e-standard-4t フラグを使用して指定します。このマシンタイプは、AI モデルに必要なハードウェア（GPU やさまざまな TPU など）に合わせて変更できます。
DWS Flex Start ノードプールを作成する: この例では、プライマリオンデマンドプールに選択したのと同じマシンタイプ（--machine-type=ct6e-standard-4t）を選択します。Plan B ノードプールで別のマシンタイプを使用する必要はありません。この特定のハードウェアをどうしても入手したい場合は、すぐに利用できない場合に別の入手方法に切り替えることをプラン B とします。この代替方法は DWS を使用します。DWS は、最大 7 日間、利用可能なハードウェアを継続的に検索します。
```
gcloud container node-pools create ${DWS_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --enable-autoscaling \
  --enable-queued-provisioning \
  --flex-start \
  --num-nodes=0 \
  --min-nodes=0 \
  --max-nodes=4
```
これらのコマンドでは、次のキーフラグを使用します。
- --num-nodes=0、--min-nodes=0、--max-nodes=4、--enable-autoscaling: この組み合わせにより、ジョブでノードが必要な場合はノードプールをゼロノードからスケールアップし、アイドル状態のときはスケールダウンできるため、費用を節約できます。
- --tpu-topology: TPU チップの物理的な配置を定義します。このレイアウトを指定するのは、チップの物理的な配置が分散トレーニングジョブの実行速度に影響するためです。
- --reservation-affinity=none:（プラン A プールのみ）このフラグは、オンデマンドノードプールが予約済み容量を消費しないようにします。DWS Flex Start ノードプールは予約を使用しません。
- --enable-queued-provisioning と --flex-start:（プラン B プールのみ）これらのフラグを使用すると、DWS は、柔軟な容量が使用可能になったときに、プラン B プールのノードをプロビジョニングできます。

ノードプールで Flex Start のステータスを確認する

DWS Flex-Start ノードプールを調べて、Flex-Start が有効になっていることを確認します。

gcloud container node-pools describe ${DWS_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --format="get(config.flexStart)"

Flex Start が有効になっている場合、出力は True です。

Kueue をインストールして構成する（クラスタ管理者）

このセクションでは、クラスタに Kueue コントローラをインストールします。Kueue は、ジョブキューを管理するスケジューリングプログラムです。ジョブリクエストをインターセプトし、使用するノードプール（オンデマンドまたは DWS Flex Start）を決定して、ジョブを割り当てます。

Kueue をインストールする

次のコマンドを実行して Kueue をインストールします。このコマンドは、公式リポジトリからインストールマニフェストをダウンロードして、クラスタに適用します。

helm install kueue oci://registry.k8s.io/kueue/charts/kueue \
  --namespace kueue-system \
  --create-namespace \
  --set "controllerManager.featureGates[0].name=ElasticJobsViaWorkloadSlices" \
  --set "controllerManager.featureGates[0].enabled=true"

構成ルールを定義する

優先度ルールを定義する YAML マニフェストを作成します。これらのルールは、Kueue にオンデマンドプールを最初に使用し、DWS Flex Start プールを 2 番目に使用するように指示します。

次の内容で dws-tpu-queue.yaml という名前のファイルを作成します。このファイルでは、2 つのリソースのフレーバー（オンデマンドと DWS Flex Start）と、それらを優先するクラスタキューを定義します。この構成ファイルは、Kueue がジョブの処理に使用するロジックを定義します。

ResourceFlavor: このドキュメントの最初で、2 つのノードプールを作成し、環境変数 ${ONDEMAND_NODEPOOL} と ${DWS_NODEPOOL} を使用して名前を割り当てました。これらのノードプールを作成するときに、GKE は、これらの環境変数に選択した名前で、これらのプールのすべてのノードに自動的にラベルを付けました。ResourceFlavor セクションは、これらのラベルを持つノードを探すように Kueue に指示します。
ClusterQueue: マニフェストのこのセクションでは、優先度ルールを定義します。オンデマンドフレーバーが最初にリストされているため、Kueue は最初にオンデマンドマシンをプロビジョニングしようとします。Kueue がこれらのマシンを取得できない場合、代わりに DWS Flex Start マシンのプロビジョニングを試みます。

Quotas: このファイルは、オンデマンドノードプールでジョブが任意の時点で使用できるリソース（CPU、メモリ、TPU チップなど）の合計の上限である割り当てを設定します。ジョブがこの上限に達すると、Kueue は dws-tpu-queue.yaml で構成した DWS Flex Start マシン（プラン B のマシン）を自動的にプロビジョニングしようとします。このマシンには、はるかに高い割り当て上限が設定されています。

apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "default-cpu"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: default-pool
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "on-demand"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: ${ONDEMAND_NODEPOOL}
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "dws"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: ${DWS_NODEPOOL}
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
  name: "cluster-queue"
spec:
  namespaceSelector: {}
  resourceGroups:
    - coveredResources: ["cpu", "memory", "google.com/tpu"]
      flavors:
        - name: "default-cpu" # Used for Ray Head Pod.
          resources:
            - name: "cpu"
              nominalQuota: 10
            - name: "memory"
              nominalQuota: 20Gi
            - name: "google.com/tpu"
              nominalQuota: 0
        - name: "on-demand" # First choice: on-demand node-pool.
          resources:
            - name: "cpu"
              nominalQuota: 40
            - name: "memory"
              nominalQuota: 75Gi
            - name: "google.com/tpu"
              nominalQuota: 16
        - name: "dws" # If on-demand is unavailable, fallback to DWS.
          resources:
            - name: "cpu"
              nominalQuota: 1000000000
            - name: "memory"
              nominalQuota: 1000000000Gi
            - name: "google.com/tpu"
              nominalQuota: 1000000000 # "Infinite" quota
  admissionChecksStrategy:
    admissionChecks:
      - name: "dws-prov"
        onFlavors: [dws]
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  namespace: "default"
  name: "user-queue"
spec:
  clusterQueue: "cluster-queue"
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: AdmissionCheck
metadata:
  name: dws-prov
spec:
  controllerName: kueue.x-k8s.io/provisioning-request
  parameters:
    apiGroup: kueue.x-k8s.io
    kind: ProvisioningRequestConfig
    name: dws-config
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ProvisioningRequestConfig
metadata:
  name: dws-config
spec:
  provisioningClassName: queued-provisioning.gke.io
  managedResources:
    - google.com/tpu

クラスタに構成を適用します。次のコマンドは、envsubst というコマンドラインツールを使用して、ファイル dws-tpu-queue.yaml に表示されるプレースホルダ変数を置き換えます。envsubst は、プレースホルダを前に定義した環境変数の値に置き換えます。
```
envsubst < dws-tpu-queue.yaml | kubectl apply -f -
```

トレーニングジョブを実行する（AI デベロッパー）

AI デベロッパーは、RayJob マニフェストを作成して、トレーニングワークロードを定義して送信します。このマニフェストでリソース要件を指定すると、クラスタ管理者が Kueue と DWS で事前に構成した自動フォールバックシステムが、基盤となるノードプールを処理します。

このセクションでは、次の手順を行います。

Python トレーニングスクリプトを作成します。
このスクリプトを Kubernetes ConfigMap に保存します。
トレーニングスクリプトをノードで実行できるように、ConfigMap をボリュームとしてマウントする RayJob をデプロイします。

これらの手順を実行すると、Ray Train は JAX ワークロードをノード間で自動的に分散し、Kueue は必要なマシンを取得します。

トレーニングスクリプト

次の Python スクリプトをコピーして train.py という名前のファイルに貼り付けます。

import time
import ray
from ray import train
from ray.train import ScalingConfig
from ray.train.v2.jax import JaxTrainer
import jax
import jax.numpy as jnp

def train_func():
    # JaxTrainer handles JAX distributed setup.
    print(f"Local Devices: {jax.local_devices()}")

    # Simple Linear Regression Training Loop.
    key = jax.random.PRNGKey(0)
    x = jax.random.normal(key, (1000, 10))
    w_true = jax.random.normal(key, (10, 1))
    y = jnp.dot(x, w_true)

    # Initialize weights
    w = jnp.zeros((10, 1))
    learning_rate = 0.1

    @jax.jit
    def update(w, x, y):
        y_pred = jnp.dot(x, w)
        loss = jnp.mean((y_pred - y) ** 2)
        grad = jax.grad(lambda w: jnp.mean((jnp.dot(x, w) - y) ** 2))(w)
        return w - learning_rate * grad, loss

    # Training loop
    print("Starting training...")
    for epoch in range(50):
        w, loss = update(w, x, y)
        if epoch % 10 == 0:
            train.report({"loss": loss.item(), "epoch": epoch})
            print(f"Epoch {epoch}: Loss {loss:.4f}")

    print("Training Complete!")
    # Allow metrics to sync before closing
    time.sleep(3)

def main():
    scaling_config = ScalingConfig(
        num_workers=4,
        resources_per_worker={"TPU": 4},
        use_tpu=True,
        topology="4x4",
        accelerator_type="TPU-V6E"
    )

    trainer = JaxTrainer(
        train_loop_per_worker=train_func,
        scaling_config=scaling_config
    )

    result = trainer.fit()
    print(f"Run Result: {result.metrics}")

if __name__ == "__main__":
    main()

トレーニングスクリプトは、高パフォーマンスの数値計算用の Python ライブラリである JAX を使用して、線形回帰モデルをトレーニングします。このスクリプトは、自動フォールバックに DWS と Kueue を使用する方法を示すように設計された簡略化された例であり、データ並列処理やモデル並列処理は実行しません。

トレーニングスクリプトの ScalingConfig セクションで、トレーニングジョブのハードウェア要件が定義されていることに注意してください。このセクションでは、以前に構成したノードプールの物理レイアウトと一致する 4x4 TPU トポロジをリクエストします。

ConfigMap の作成

train.py スクリプトの内容を Kubernetes ConfigMap オブジェクトにアップロードします。これにより、クラスタはスクリプトを保存し、RayJob で使用できるようになります。

kubectl create configmap jax-train-script --from-file=train.py

次のセクションで定義する RayJob は、この ConfigMap をボリュームとしてマウントします。これにより、Ray ソフトウェアがスクリプトファイルを見つけて実行できるように、スクリプトファイルが Ray コンテナ内に表示されます。

RayJob マニフェストを適用する

次の内容で rayjob-tpu-v6e-dws.yaml という名前のファイルを作成します。このマニフェストは、トレーニングジョブを定義し、システムにそのジョブの転送方法を指示します。

apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: rayjob-tpu-v6e-dws-${JOB_ID}
  labels:
    kueue.x-k8s.io/queue-name: user-queue
  annotations:
    kueue.x-k8s.io/elastic-job: "true"
spec:
  shutdownAfterJobFinishes: true
  entrypoint: python /app/train.py
  runtimeEnvYAML: |
    pip:
      - jax[tpu]==0.8.2
      - pandas==2.3.3
  rayClusterSpec:
    enableInTreeAutoscaling: true
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          nodeSelector:
            cloud.google.com/gke-nodepool: default-pool
          containers:
            - name: ray-head
              image: rayproject/ray:2.53.0-py311
              ports:
                - containerPort: 6379
                  name: gcs-server
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
              resources:
                limits:
                  cpu: "2"
                  memory: "4Gi"
                requests:
                  cpu: "2"
                  memory: "4Gi"
              volumeMounts:
                - mountPath: /app
                  name: train-script-volume
          volumes:
            - name: train-script-volume
              configMap:
                name: jax-train-script
    workerGroupSpecs:
      - replicas: 1
        minReplicas: 1
        maxReplicas: 2
        numOfHosts: 4
        groupName: tpu-group
        rayStartParams: {}
        template:
          spec:
            tolerations:
              - key: "google.com/tpu"
                operator: "Exists"
                effect: "NoSchedule"
              - key: "cloud.google.com/gke-queued"
                operator: "Exists"
                effect: "NoSchedule"
            containers:
              - name: ray-worker
                image: rayproject/ray:2.53.0-py311
                resources:
                  limits:
                    cpu: "8"
                    google.com/tpu: "4"
                    memory: "16Gi"
                  requests:
                    cpu: "8"
                    google.com/tpu: "4"
                    memory: "16Gi"
                volumeMounts:
                  - mountPath: /app
                    name: train-script-volume
            volumes:
              - name: train-script-volume
                configMap:
                  name: jax-train-script
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
              cloud.google.com/gke-tpu-topology: 4x4

このマニフェストには、フォールバックシステムを機能させる 3 つの構成が含まれています。

特定のハードウェアをリクエストする: nodeSelector セクションでは、スクリプトに必要なハードウェア（この例では 4x4 トポロジの tpu-v6e-slice）を指定します。
キューを選択します。kueue.x-k8s.io/queue-name ラベルは、ジョブを Kueue に直接転送します。これにより、自動フォールバックロジックが有効になります。
DWS Flex Start ノードを許容する: tolerations セクションでは、ジョブを Plan B ノードプールで実行できます。DWS Flex Start ノードは、通常のワークロードが誤って実行されないように GKE によって特別にマーク（taint）されているため、ジョブで cloud.google.com/gke-queued taint を明示的に許容する必要があります。

ワークロードを送信する

フォールバックシステムが機能することを証明するには、2 つのジョブを送信する必要があります。最初のジョブはプラン A のオンデマンド容量を使用するため、2 番目のジョブはプラン B の DWS Flex Start 容量にフォールバックされます。

次のコマンドを実行して、2 つのジョブを送信します。このコマンドは、for ループと envsubst を使用して、実行ごとに一意のジョブ ID をマニフェストに挿入します。

for i in 1 2; do
  export JOB_ID=$i
  envsubst < rayjob-tpu-v6e-dws.yaml | kubectl apply -f -
  echo "Submitted Job $i"
  sleep 2
done

ジョブを送信すると、システムは次のようにワークロードを処理します。

インターセプト: Kueue はキューラベルを使用してジョブを検出し、一時的に一時停止します。
決定: Kueue は、管理者のルールに照らしてリソースの可用性を評価します。まずプラン A プールがチェックされます。
課題:
- プラン A のリソースは最初のジョブで使用できるため、Kueue はジョブ 1 をプラン A に割り当てます。
- ジョブ 1 がプラン A のリソースを消費するため、Kueue はジョブ 2 をプラン B（DWS Flex Start）プールに自動的に割り当てます。
起動: Kueue は Job の一時停止を解除します。このアクションにより、GKE クラスタオートスケーラーがノードをプロビジョニングし、トレーニングスクリプトを開始します。

RayJob に接続する

最後の検証ステップとして、kubectl port-forward コマンドを使用して Ray ダッシュボードに接続し、ジョブの実行をモニタリングできます。

最初のジョブのステータスを確認するには、次のコマンドを実行します。

kubectl port-forward service/rayjob-tpu-v6e-dws-1-head-svc 8265:8265 &

このコマンドを実行したら、ウェブブラウザを開いて http://localhost:8265 に移動します。Ray ダッシュボードで、ジョブのステータスとレポート指標を表示して、両方のジョブがそれぞれのノードプールで正常に完了したことを確認できます。

次のコマンドを実行して、最初のジョブのログを表示することもできます。

kubectl logs job/rayjob-tpu-v6e-dws-1

トレーニングスクリプトの出力は、次のようになります。出力の末尾付近に Training Complete! と Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded のメッセージが表示されます。

(pid=, ip=10.68.3.4) 5] XLA::TPU program HBM usage: 52.5K / 31.25G
(pid=, ip=10.68.9.4) :2152] XLA::TPU program VMEM usage: 141.0K / 128.00M [repeated 5x across cluster]
(pid=, ip=10.68.9.4) I0320 03:59:34.722540     855 deepsea_compiler_backend.cc:2163] Total hbm usage >= 260.14M: [repeated 5x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777634     888 deepsea_compiler_backend.cc:2167]     reserved           204B [repeated 19x across cluster]
(pid=, ip=10.68.9.4) I0320 03:59:34.722542     855 deepsea_compiler_backend.cc:2163]     program           70.0K  [repeated 5x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777626     888 deepsea_compiler_backend.cc:2163]     arguments            0B  [repeated 12x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777627     888 deepsea_compiler_backend.cc:2163] Output size 0B; shares 0B with arguments. [repeated 14x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777625     888 deepsea_compiler_backend.cc:2163] Total host usage >= 0B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777626     888 deepsea_compiler_backend.cc:2163]     program         unknown size  [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777634     888 deepsea_compiler_backend.cc:2167] Program sflag requirement 224B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167]     scoped              40B [repeated 21x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777636     888 deepsea_compiler_backend.cc:2167] Program vmem requirement 141.0K: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program smem requirement 40B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program host requirement 0B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program hbm requirement 70.0K: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777638     888 deepsea_compiler_backend.cc:2167]     overlays          70.0K [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777638     888 deepsea_compiler_backend.cc:2175] XLA::TPU program SMEM usage: 1.9K / 1.00M (3 parameters) [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777636     888 deepsea_compiler_backend.cc:2167]     HLO temp          76.0K (0.0% utilization: Unpadded (0B) Padded (0B), 100.0% fragmentation (76.0K)) [repeated 14x across cluster]
(RayTrainWorker pid=542, ip=10.68.6.4) Training Complete! [repeated 3x across cluster]
(RayTrainWorker pid=542, ip=10.68.6.4) Epoch 40: Loss 0.0000 [repeated 3x across cluster]
2026-03-20 03:59:51,008 SUCC cli.py:65 -- ------------------------------------------
2026-03-20 03:59:51,008 SUCC cli.py:66 -- Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded
2026-03-20 03:59:51,008 SUCC cli.py:67 -- ------------------------------------------

クリーンアップ

このドキュメントで使用したリソースについて、 Google Cloud アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。

プロジェクトを削除する

注意: プロジェクトを削除すると、次のような影響があります。

プロジェクト内のすべてのものが削除されます。既存のプロジェクトを使用してこのドキュメントのタスクを行った場合、そのプロジェクトを削除すると、プロジェクト内で行った他の作業もすべて削除されます。
カスタムプロジェクト ID が失われます。このプロジェクトを作成したときに、将来使用するカスタムプロジェクト ID を作成した可能性があります。そのプロジェクト ID を使用した URL（例: appspot.com）を保持するには、プロジェクト全体ではなく、プロジェクト内の選択したリソースだけを削除します。

Google Cloud コンソールで [リソースの管理] ページに移動します。

[リソースの管理] に移動

プロジェクトリストで、削除するプロジェクトを選択し、[削除] をクリックします。

ダイアログでプロジェクト ID を入力し、[シャットダウン] をクリックしてプロジェクトを削除します。

個々のリソースの削除

このドキュメントで使用した GGoogle Cloud プロジェクトを保持する場合は、次のコマンドを実行してクラスタを削除します。

gcloud container clusters delete ${CLUSTER_NAME} \
    --location=${ZONE} \
    --project=${PROJECT_ID} \
    --quiet

概要

このドキュメントでは、Ray トレーニング環境を構成してテストしました。この環境では、ハードウェアの可用性を最大化するために、プライマリノードプールとバックアップ DWS プールを使用します。プライマリマシンが使用できない場合に DWS に自動的にフォールバックすることで、トレーニングジョブが待機する時間を最小限に抑えました。

この問題を解決するために、次の手順を実施しました。

GKE クラスタを作成しました: ノードプールとスケジューリングツールをホストする環境を確立しました。
ノードプールを構成しました。オンデマンドノードプール（プラン A）と DWS ノードプール（プラン B）を作成しました。
Kueue をインストールして構成する: Kueue コントローラをデプロイし、最初にプラン A を試行し、プラン B にフォールバックするようにシステムに指示する優先度ルールを適用しました。
ConfigMap を作成した: 簡素化された JAX トレーニングスクリプトをクラスタにデプロイし、テストワークロードとして機能させました。
RayJob マニフェストを定義した: 特定のハードウェアをリクエストし、Kueue コントローラにルーティングし、DWS ノードを許容するようにジョブを構成しました。
ワークロードを送信した: 2 つのジョブを送信して、プラン A のリソースが消費されたときに、Kueue が 2 番目のジョブをプラン B に自動的にルーティングするようにしました。
結果を確認しました: ポート転送を使用して Ray ダッシュボードに接続し、両方のジョブが正常に実行されたことを確認しました。