此页面由 Cloud Translation API 翻译。

使用 TPU Trillium 配置 KubeRay

Autopilot Standard

本教程介绍了如何在 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 按钮。

一个 Cloud Shell 会话随即会在 Google Cloud 控制台中的新框架内打开，并显示命令行提示符。

创建 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 会创建一个启用了 Ray 运算符插件的 Autopilot 集群。该插件会自动在集群控制平面中安装 Ray TPU Webhook。

如需与集群通信，请配置 kubectl：

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

标准

在 Cloud Shell 中，运行以下命令以创建启用 Ray Operator 插件的 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