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

Standard

このドキュメントでは、Tensor Processing Unit（TPU）で AI トレーニングの可用性を最大化するように GKE クラスタを構成する方法について説明します。オープンソースのジョブキューイングツールである Kueue と Google Cloud's 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 ロールは必要ありません。ロールが付与されているプロジェクトを選択できます。
- プロジェクトを作成する: プロジェクトを作成するには、プロジェクト作成者ロール（roles/resourcemanager.projectCreator）が必要です。これには resourcemanager.projects.create 権限が含まれています。詳しくは、ロールを付与する方法をご覧ください。
注: この手順で作成するリソースを保持しない場合は、既存のプロジェクトを選択するのではなく、プロジェクトを作成します。チュートリアルの終了後にそのプロジェクトを削除すれば、プロジェクトに関連するすべてのリソースを削除できます。
プロジェクトセレクタに移動
プロジェクトに対して課金が有効になっていることを確認します Google Cloud 。
Google Kubernetes Engine、Cloud TPU API API を有効にします。
API を有効にするために必要なロール
API を有効にするには、serviceusage.services.enable 権限を含む Service Usage 管理者 IAM ロール（roles/serviceusage.serviceUsageAdmin）が必要です。詳しくは、ロールを付与する方法をご覧ください。
API を有効にする

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

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

環境変数を定義する

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

これらの変数を定義したら、毎回値を再入力または置き換えるのではなく、変数名（$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: Standard オンデマンドノードプールの名前。（これはプラン 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）を選択します。プラン B のノードプールで異なるマシンタイプを使用する必要はありません。この特定のハードウェアが必要なため、プラン B の選択肢にすると、すぐに使用できない場合に別の取得方法に切り替えることになります。この代替方法では、DWS を使用して、最大 7 日間使用可能なハードウェアを継続的に検索します。
```
gcloud container node-pools create ${DWS_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --reservation-affinity=none \
  --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: ノードプールが事前予約済みのハードウェアを使用しないようにします。特定のマシンを予約して、可用性を確保できます。 Google Cloud このフラグを none に設定すると、予約をバイパスして、予約されていないマシンを動的にリクエストするようにシステムに指示します。
- --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 プールを使用します。

次の内容で 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 Flex Start マシン（プラン B マシン）のプロビジョニングを試みます。これは、dws-tpu-queue.yaml で構成したもので、割り当て上限がはるかに高くなっています。

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 セクションでは、ジョブをプラン 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 をそこに割り当てます。
- ジョブ 1 はプラン A のリソースを使用するため、Kueue はジョブ 2 をプラン B（DWS Flex Start）プールに自動的に割り当てます。
起動: Kueue はジョブを中断します。このアクションにより、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 URL）を保持するには、プロジェクト全体ではなく、プロジェクト内の選択したリソースだけを削除します。

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

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

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

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

個々のリソースの削除

このドキュメントで使用した 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 ダッシュボードに接続し、両方のジョブが正常に実行されたことを確認しました。

次のステップ

Ray クラスタとジョブのログと指標を表示する方法を学習する。
複数のスライスにまたがる大規模な分散 TPU ワークロードを管理する方法を学習する。