本頁面由 Cloud Translation API 翻譯而成。

使用 TPU Trillium 設定 KubeRay

自動駕駛標準

本教學課程說明如何在 Google Kubernetes Engine (GKE) 上，使用 TPU Trillium 設定 KubeRay。瞭解如何設定單一主機和多主機 TPU，包括 TPU Trillium 的必要環境變數和 Pod 規格。

本教學課程適用於平台管理員和營運人員，以及想瞭解如何使用 KubeRay 為單一主機和多主機節點集區設定 TPU Trillium 初始化作業的資料和 AI 專家。本教學課程說明如何使用 Jax 執行指令碼，驗證 TPU 是否已順利初始化。本教學課程不會部署模型。

在 GKE 中設定 KubeRay 之前，請務必先熟悉 GKE 中的 Ray 定義和術語。

總覽

本教學課程說明如何使用 Jax 執行 Python 指令碼，確認使用 KubeRay 初始化 TPU Trillium 成功。Jax 是高效能的數值運算程式庫，支援機器學習工作負載。KubeRay 是 Kubernetes 運算子，可提供統一的方式，在 Kubernetes 上部署、管理及監控 Ray 應用程式。

Trillium TPU (v6e) 需要特定的環境變數和 Pod 規格，與先前的 TPU 世代不同。本教學課程提供必要的設定，協助您在 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 或 Standard 叢集中的 TPU 上設定 KubeRay。建議您使用 Autopilot 叢集，享有全代管 Kubernetes 體驗。如要選擇最適合工作負載的 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 會建立 Autopilot 叢集，並啟用 Ray 運算子外掛程式。外掛程式會在叢集控制層中自動安裝 Ray TPU 網頁掛鉤。

如要與叢集通訊，請設定 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 規格包含 rayClusterSpec，可定義所需的 Ray 叢集設定 (首節點和工作站群組)。
執行特定工作：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 服務，並連線至 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 &

開啟網路瀏覽器並輸入下列網址：
```
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 &

開啟網路瀏覽器並輸入下列網址：
```
http://localhost:8265/#/jobs
```
查看 RayJob 狀態和相關記錄。

Ray 會設定 TPU-{accelerator}-Head 資源，以識別對應 TPU_WORKER_ID=0 值的 Ray 工作站節點。在多主機 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

使用 TPU Trillium 設定 KubeRay 透過集合功能整理內容 你可以依據偏好儲存及分類內容。

總覽

事前準備

啟用 Cloud Shell

建立 GKE 叢集和節點集區

Autopilot

標準

單一主機

多主機

執行 RayJob 自訂資源

單一主機

多主機

在 Ray 資訊主頁中查看 RayJob

單一主機

多主機

清除所用資源

單一主機

多主機

後續步驟

使用 TPU Trillium 設定 KubeRay