TPU Trillium を使用して KubeRay を構成する

このチュートリアルでは、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 では、このドキュメントのコマンドを実行できない場合があります。
  • Ray CLI（バージョン 2.37.0）がインストールされていることを確認します。

Cloud Shell をアクティブにする

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

  1. Google Cloud コンソールに移動します。

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

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

    Cloud Shell セッション

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

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

Autopilot

  1. 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 を自動的にインストールします。

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

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

Standard

  1. 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 の可用性をご覧ください。

    クラスタの作成には数分かかることもあります。

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

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

  3. 単一ホストまたはマルチホストの 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 カスタム リソースを作成するには、次の操作を行います。

単一ホスト

  1. 次の 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

  2. 次のようにマニフェストを適用します。

    kubectl apply -f ray-job.tpu-v6e-singlehost.yaml

  3. 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

  4. 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 -- -------------------------------------

マルチホスト

  1. 次の 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

  2. 次のようにマニフェストを適用します。

    kubectl apply -f ray-job.tpu-v6e-multihost.yaml

  3. 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

  4. 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 インスタンスに接続します。

単一ホスト

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

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

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

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

  3. ヘッドサービスをポート転送して Ray ダッシュボードに接続します。

    kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &

  4. ウェブブラウザを開いて、次の URL を入力します。

    http://localhost:8265/#/jobs

  5. RayJob のステータスと関連するログを表示します。

マルチホスト

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

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

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

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

  3. ヘッドサービスをポート転送して Ray ダッシュボードに接続します。

    kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &

  4. ウェブブラウザを開いて、次の URL を入力します。

    http://localhost:8265/#/jobs

  5. 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

次のステップ