使用 TPUs 启动 Google Cloud GKE 集群以用于 KubeRay#

参阅 GKE 文档获取完整详情，或继续阅读以进行快速入门。

步骤 1：在 GKE 上创建 Kubernetes 集群#

首先，设置以下用于创建 GKE 集群的环境变量

export CLUSTER_NAME=CLUSTER_NAME
export COMPUTE_ZONE=ZONE
export CLUSTER_VERSION=CLUSTER_VERSION

替换以下内容

CLUSTER_NAME: 要创建的 GKE 集群名称。
ZONE: 具有可用 TPU 配额的区域，有关按区域划分的 TPU 可用性列表，请参阅 GKE 文档。
CLUSTER_VERSION: 要使用的 GKE 版本。GKE 1.31.2-gke.1115000 或更高版本支持 TPU v6e。有关 TPU 代及其最低支持版本，请参阅 GKE 文档。

在您的本地机器或 Google Cloud Shell 上运行以下命令。如果在本地机器上运行，请安装 Google Cloud SDK。

创建 Standard GKE 集群并启用 Ray Operator

gcloud container clusters create $CLUSTER_NAME \
    --addons=RayOperator \
    --machine-type=n1-standard-16 \
    --cluster-version=$CLUSTER_VERSION \
    --location=$ZONE

运行以下命令向集群添加 TPU 节点池。您也可以从 Google Cloud 控制台创建它。

按照如下方式创建一个具有单主机 v4 TPU 拓扑的节点池

gcloud container node-pools create v4-4 \
  --zone $ZONE \
  --cluster $CLUSTER_NAME \
  --num-nodes 1 \
  --min-nodes 0 \
  --max-nodes 10 \
  --enable-autoscaling \
  --machine-type ct4p-hightpu-4t \
  --tpu-topology 2x2x1

对于 v4 TPU，ZONE 必须是 us-central2-b。

或者，按照如下方式创建一个多主机节点池

gcloud container node-pools create v4-8 \
  --zone $ZONE \
  --cluster $CLUSTER_NAME \
  --num-nodes 2 \
  --min-nodes 0 \
  --max-nodes 10 \
  --enable-autoscaling \
  --machine-type ct4p-hightpu-4t \
  --tpu-topology 2x2x2

对于 v4 TPU，ZONE 必须是 us-central2-b。

--tpu-topology 标志指定了 TPU Pod 分片的物理拓扑结构。此示例使用了具有 2x2x1 或 2x2x2 拓扑的 v4 TPU 分片。v4 TPU 每个 VM 主机有 4 个芯片，因此 2x2x2 v4 分片总共有 8 个芯片和 2 个 TPU 主机，每个主机调度在自己的节点上。GKE 将多主机 TPU 分片视为原子单元，并使用节点池而不是单个节点进行扩缩。因此，TPU 主机的数量应始终等于 TPU 节点池中的节点数量。有关选择 TPU 拓扑和加速器的更多信息，请参阅 GKE 文档。

GKE 使用 Kubernetes 节点选择器来确保 TPU 工作负载运行在所需的机器类型和拓扑结构上。有关更多详情，请参阅 GKE 文档。

步骤 2：连接到 GKE 集群#

运行以下命令下载 Google Cloud 凭据并配置 Kubernetes CLI 以使用它们。

gcloud container clusters get-credentials $CLUSTER_NAME --zone $ZONE

现在可以通过 kubectl 访问远程 GKE 集群。有关更多详情，请参阅 GKE 文档。

[可选] 在未启用 Ray Operator 插件的 GKE 集群中手动安装 KubeRay 和 TPU webhook：#

在未启用 Ray Operator 插件的集群中，可以使用 helm 手动安装 KubeRay，命令如下

helm repo add kuberay https://ray-project.github.io/kuberay-helm/

# Install both CRDs and KubeRay operator v1.3.0.
helm install kuberay-operator kuberay/kuberay-operator --version 1.3.0

GKE 提供了一个验证和修改 webhook，用于处理 TPU Pod 调度并引导用于 JAX 初始化的某些环境变量。Ray TPU webhook 需要 KubeRay operator 版本至少为 v1.1.0。对于 GKE 版本 1.30.0-gke.1747000 或更高版本，GKE 通过 Ray Operator 插件自动安装 Ray TPU webhook。

手动安装 webhook 时，需要 cert-manager 来处理 TLS 证书注入。您可以使用以下 helm 命令在 GKE Standard 和 Autopilot 集群中安装 cert-manager

安装 cert-manager

helm repo add jetstack https://charts.jetstack.io
helm repo update
helm install --create-namespace --namespace cert-manager --set installCRDs=true --set global.leaderElection.namespace=cert-manager cert-manager jetstack/cert-manager

接下来，部署 Ray TPU 初始化 webhook

git clone https://github.com/GoogleCloudPlatform/ai-on-gke
cd ray-on-gke/tpu/kuberay-tpu-webhook
make deploy deploy-cert