このページは Cloud Translation API によって翻訳されました。

TPU Trillium で KubeRay を構成する

Autopilot Standard

このチュートリアルでは、Google Kubernetes Engine（GKE）で TPU Trillium を使用して KubeRay を構成する方法について説明します。TPU Trillium に必要な環境変数と Pod 仕様など、単一ホストとマルチホストの両方の TPU 構成を設定する方法について説明します。

このチュートリアルは、単一ホストノードプールとマルチホストノードプール用に KubeRay を使用して TPU Trillium の初期化を構成する方法を学習するプラットフォーム管理者とオペレーター、データおよび AI スペシャリストを対象としています。このチュートリアルでは、TPU の初期化が成功したことを検証する Jax を使用してスクリプトを実行する方法について説明します。このチュートリアルでは、モデルをデプロイしません。

GKE で KubeRay を構成する前に、GKE での Ray の定義と用語を理解しておいてください。

概要

このチュートリアルでは、KubeRay を使用した TPU Trillium の初期化が成功したことを確認する Jax を使用して Python スクリプトを実行する方法について説明します。Jax は、ML ワークロードをサポートする高パフォーマンスの数値計算ライブラリです。KubeRay は、Kubernetes で Ray アプリケーションをデプロイ、管理、モニタリングするための統一された方法を提供する Kubernetes オペレーターです。

Trillium TPU（v6e）には、以前の TPU 世代とは異なる特定の環境変数と Pod 仕様が必要です。このチュートリアルでは、Trillium TPU で KubeRay を使用してワークロードを正常にデプロイするために必要な構成について説明します。

始める前に

作業を始める前に、次のタスクが完了していることを確認してください。

Google Kubernetes Engine API を有効にする。

Google Kubernetes Engine API の有効化

このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化する。すでに gcloud CLI をインストールしている場合は、gcloud components update を実行して最新のバージョンを取得する。
注: gcloud CLI がすでにインストールされている場合には、必ず compute/region プロパティを設定してください。主にゾーンクラスタを使用する場合は、代わりに compute/zone を設定します。デフォルトのロケーションを設定することで、gcloud CLI のエラー（One of [--zone, --region] must be supplied: Please specify location など）を防止できます。クラスタのロケーションが設定したデフォルトと異なる場合は、特定のコマンドでロケーションの指定が必要になることがあります。

Ray CLI（バージョン 2.37.0）がインストールされていることを確認します。

Cloud Shell をアクティブにする

Cloud Shell には、このチュートリアルで使用する gcloud、helm、kubectl コマンドラインツールがプリインストールされています。

Google Cloud コンソールに移動します。
Google Cloud コンソールウィンドウの上部にある [Cloud Shell をアクティブにする] ボタンをクリックします。

Google Cloud コンソールの新しいフレーム内で Cloud Shell セッションが開き、コマンドラインプロンプトが表示されます。

GKE クラスタとノードプールを作成する

GKE Autopilot クラスタまたは GKE Standard クラスタの TPU で KubeRay を構成できます。フルマネージドの Kubernetes エクスペリエンスを実現するには、Autopilot クラスタを使用することをおすすめします。ワークロードに最適な GKE の運用モードを選択するには、GKE の運用モードについてをご覧ください。

Autopilot

Cloud Shell で、次のコマンドを実行します。
```
gcloud container clusters create-auto CLUSTER_NAME \
    --enable-ray-operator \
    --release-channel=rapid \
    --location=LOCATION
```
次のように置き換えます。
- CLUSTER_NAME: 新しいクラスタの名前。
- LOCATION: TPU Trillium の容量が使用可能なリージョン。詳細については、GKE での TPU の可用性をご覧ください。
GKE は、Ray オペレータアドオンが有効になっている Autopilot クラスタを作成します。アドオンは、クラスタコントロールプレーンに Ray TPU Webhook を自動的にインストールします。

クラスタと通信するには、kubectl を構成します。

gcloud container clusters get-credentials CLUSTER_NAME --location=LOCATION

標準

Cloud Shell で、次のコマンドを実行して、Ray オペレータアドオンを有効にする Standard クラスタを作成します。
```
gcloud container clusters create CLUSTER_NAME \
  --location LOCATION \
  --addons=RayOperator \
  --cluster-version=1.33 \
  --machine-type=n1-standard-16
```
次のように置き換えます。
- CLUSTER_NAME: 新しいクラスタの名前。
- LOCATION: TPU Trillium の容量が使用可能なリージョン。詳細については、GKE での TPU の可用性をご覧ください。
クラスタの作成には数分かかることもあります。

クラスタと通信するには、kubectl を構成します。

gcloud container clusters get-credentials CLUSTER_NAME --location=LOCATION

単一ホストまたはマルチホストの TPU スライスノードプールを作成できます。

単一ホスト

Cloud Shell で、次のコマンドを実行します。

gcloud container node-pools create v6e-4 \
    --location=us-central2-b \
    --cluster=CLUSTER_NAME \
    --machine-type=ct6e-standard-4t \
    --num-nodes=1 \
    --threads-per-core=1 \
    --tpu-topology=2x2

マルチホスト

Cloud Shell で、次のコマンドを実行します。

gcloud container node-pools create v6e-16 \
    --location=us-central2-b \
    --cluster=CLUSTER_NAME \
    --machine-type=ct6e-standard-4t \
    --num-nodes=4 \
    --threads-per-core=1 \
    --tpu-topology=4x4

RayJob カスタムリソースを実行する

RayJob マニフェストを定義すると、KubeRay は次の処理を行います。

RayCluster を作成する: RayJob 仕様には、必要な Ray クラスタ構成（ヘッドグループとワーカーグループ）を定義する rayClusterSpec が含まれています。
特定のジョブを実行する: RayJob 内の entrypoint フィールドは、作成された Ray クラスタ内で実行するコマンドまたはスクリプトを指定します。このチュートリアルでは、entrypoint は TPU Trillium の初期化を検証するように設計された Python スクリプト（tpu_list_devices.py）です。

RayJob カスタムリソースを作成するには、次の操作を行います。

単一ホスト

次の ray-job.tpu-v6e-singlehost.yaml マニフェストを作成します。

apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: v6e-4-job
spec:
  entrypoint: python ai-ml/gke-ray/tpu/tpu_list_devices.py
  runtimeEnvYAML: |
    working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
    pip:
      - jax[tpu]==0.4.33
      - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
  rayClusterSpec:
    rayVersion: '2.43.0'
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          containers:
          -   name: ray-head
              image: rayproject/ray:2.43.0-py310
              ports:
                - containerPort: 6379
                  name: gcs-server
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
              resources:
                limits:
                  cpu: "8"
                  memory: 40G
                requests:
                  cpu: "8"
                  memory: 40G
    workerGroupSpecs:
    -   replicas: 1
        minReplicas: 1
        maxReplicas: 1
        numOfHosts: 1
        groupName: tpu-group
        rayStartParams: {}
        template:
          spec:
            containers:
            -   name: ray-worker
                image: rayproject/ray:2.43.0-py310
                resources:
                  limits:
                    cpu: "24"
                    google.com/tpu: "4"
                    memory: 200G
                  requests:
                    cpu: "24"
                    google.com/tpu: "4"
                    memory: 200G
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
              cloud.google.com/gke-tpu-topology: 2x2

次のようにマニフェストを適用します。
```
kubectl apply -f ray-job.tpu-v6e-singlehost.yaml
```

RayJob が作成され、実行されていることを確認します。

kubectl get rayjobs v6e-4-job

出力は次のようになります。

NAME      JOB STATUS   DEPLOYMENT STATUS   RAY CLUSTER NAME       START TIME  END TIME   AGE
v6e-4-job PENDING      Running             v6e-4-job-raycluster   2024-10-15T23:15:22Z  20s

RayJob の出力を出力します。

kubectl logs -l=job-name=v6e-4-job

出力は次のようになります。

2024-10-15 16:15:40,222 INFO cli.py:300 -- ray job stop v6e-4-job-hzq5q
2024-10-15 16:15:40,246 INFO cli.py:307 -- Tailing logs until the job exits (disable with --no-wait):
2024-10-15 16:15:40,112 INFO job_manager.py:528 -- Runtime env is setting up.
2024-10-15 16:15:50,181 INFO worker.py:1461 -- Using address 10.84.1.25:6379 set in the environment variable RAY_ADDRESS
2024-10-15 16:15:50,181 INFO worker.py:1601 -- Connecting to existing Ray cluster at address: 10.84.1.25:6379...
2024-10-15 16:15:50,186 INFO worker.py:1777 -- Connected to Ray cluster. View the dashboard at 10.84.1.25:8265
['TPU cores:4']
2024-10-15 16:16:12,349 SUCC cli.py:63 -- -------------------------------------
2024-10-15 16:16:12,349 SUCC cli.py:64 -- Job 'v6e-4-job-hzq5q' succeeded
2024-10-15 16:16:12,349 SUCC cli.py:65 -- -------------------------------------

マルチホスト

次の ray-job.tpu-v6e-multihost.yaml マニフェストを作成します。

apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: v6e-16-job
spec:
  entrypoint: python ai-ml/gke-ray/tpu/tpu_list_devices.py
  runtimeEnvYAML: |
    working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
    pip:
      - jax[tpu]==0.4.33
      - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
  rayClusterSpec:
    rayVersion: '2.43.0'
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          containers:
          -   name: ray-head
              image: rayproject/ray:2.43.0-py310
              ports:
                - containerPort: 6379
                  name: gcs-server
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
              resources:
                limits:
                  cpu: "8"
                  memory: 40G
                requests:
                  cpu: "8"
                  memory: 40G
    workerGroupSpecs:
      - replicas: 1
        minReplicas: 1
        maxReplicas: 1
        numOfHosts: 4
        groupName: tpu-group
        rayStartParams: {}
        template:
          spec:
            containers:
              - name: ray-worker
                image: rayproject/ray:2.43.0-py310
                resources:
                  limits:
                    cpu: "24"
                    google.com/tpu: "4"
                    memory: 200G
                  requests:
                    cpu: "24"
                    google.com/tpu: "4"
                    memory: 200G
                env:
                - name: NODE_IP
                  valueFrom:
                    fieldRef:
                      fieldPath: status.hostIP
                - name: VBAR_CONTROL_SERVICE_URL
                  value: $(NODE_IP):8353
                - name: JAX_PLATFORMS
                  value: tpu,cpu
                - name: ENABLE_PJRT_COMPATIBILITY
                  value: "true"
                ports:
                - containerPort: 8081
                  name: mxla
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
              cloud.google.com/gke-tpu-topology: 4x4

次のようにマニフェストを適用します。
```
kubectl apply -f ray-job.tpu-v6e-multihost.yaml
```

v6e-16 RayJob が作成され、実行されていることを確認します。

kubectl get rayjobs v6e-16-job

出力は次のようになります。

NAME         JOB STATUS   DEPLOYMENT STATUS   RAY CLUSTER NAME              START TIME             END TIME   AGE
v6e-16-job                Running             v6e-16-job-raycluster-qr6vk   2024-10-16T19:28:19Z              66s

v6e-16 RayJob の出力を出力します。

kubectl logs -l=job-name=v6e-16-job

出力は次のようになります。

2024-10-16 12:21:33,986 INFO cli.py:300 -- ray job stop v6e-16-job-z44s7
2024-10-16 12:21:34,011 INFO cli.py:307 -- Tailing logs until the job exits (disable with --no-wait):
2024-10-16 12:21:33,826 INFO job_manager.py:528 -- Runtime env is setting up.
2024-10-16 12:21:46,327 INFO worker.py:1461 -- Using address 10.84.1.61:6379 set in the environment variable RAY_ADDRESS
2024-10-16 12:21:46,327 INFO worker.py:1601 -- Connecting to existing Ray cluster at address: 10.84.1.61:6379...
2024-10-16 12:21:46,333 INFO worker.py:1777 -- Connected to Ray cluster. View the dashboard at 10.84.1.61:8265
['TPU cores:16', 'TPU cores:16', 'TPU cores:16', 'TPU cores:16']
2024-10-16 12:22:12,156 SUCC cli.py:63 -- ---------------------------------
2024-10-16 12:22:12,156 SUCC cli.py:64 -- Job 'v6e-16-job-z44s7' succeeded
2024-10-16 12:22:12,156 SUCC cli.py:65 -- ---------------------------------

Ray ダッシュボードで RayJob を表示する

GKE が RayCluster Service を作成したことを確認し、RayCluster インスタンスに接続します。

単一ホスト

RayJob 用に生成された RayCluster の名前を取得します。

export RAYCLUSTER_NAME=$(kubectl get rayjob v6e-4-job -o jsonpath='{.status.rayClusterName}')

RayCluster ヘッドサービスの名前を取得します。

export HEAD_SVC=$(kubectl get svc -l ray.io/cluster=$RAYCLUSTER_NAME,ray.io/node-type=head -o jsonpath='{.items[0].metadata.name}')

ヘッドサービスをポート転送して Ray ダッシュボードに接続します。
```
kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &
```
ウェブブラウザを開いて、次の URL を入力します。
```
http://localhost:8265/#/jobs
```
RayJob のステータスと関連するログを表示します。

マルチホスト

RayJob 用に生成された RayCluster の名前を取得します。

export RAYCLUSTER_NAME=$(kubectl get rayjob v6e-16-job -o jsonpath='{.status.rayClusterName}')

RayCluster ヘッドサービスの名前を取得します。

export HEAD_SVC=$(kubectl get svc -l ray.io/cluster=$RAYCLUSTER_NAME,ray.io/node-type=head -o jsonpath='{.items[0].metadata.name}')

ヘッドサービスをポート転送して Ray ダッシュボードに接続します。
```
kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &
```
ウェブブラウザを開いて、次の URL を入力します。
```
http://localhost:8265/#/jobs
```
RayJob のステータスと関連するログを表示します。

Ray は、TPU_WORKER_ID=0 値に対応する Ray ワーカーノードを識別する TPU-{accelerator}-Head リソースを設定します。マルチホスト TPU グループでは、TPU_WORKER_ID=0 を持つ Ray ノードのリソースに TPU-v6e-16-head: 1.0 が設定されています。この TPU_WORKER_ID 環境変数は、KubeRay のミューテーション GKE Webhook によって設定されます。

クリーンアップ

チュートリアルを完了したら、アカウントで不要な請求が発生しないように、RayJob を削除します。

単一ホスト

kubectl delete rayjobs v6e-4-job

マルチホスト

kubectl delete rayjobs v6e-16-job

次のステップ

Ray on Kubernetes について学習する。
GKE で TPU を使用して vLLM をサービングする方法を確認する。
TPU を使用して GKE で SDXL をサービングする方法を確認する。
詳細については、GKE の TPU についてをご覧ください。

TPU Trillium で KubeRay を構成する コレクションでコンテンツを整理 必要に応じて、コンテンツの保存と分類を行います。

概要

始める前に

Cloud Shell をアクティブにする

GKE クラスタとノードプールを作成する

Autopilot

標準

単一ホスト

マルチホスト

RayJob カスタム リソースを実行する

単一ホスト

マルチホスト

Ray ダッシュボードで RayJob を表示する

単一ホスト

マルチホスト

クリーンアップ

単一ホスト

マルチホスト

次のステップ

TPU Trillium で KubeRay を構成する

RayJob カスタムリソースを実行する