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

有关完整详细信息，请参阅 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 版本。TPU v6e 支持 GKE 版本 1.31.2-gke.1115000 或更高版本。有关 TPU 代系及其最低支持版本，请参阅 GKE 文档。

在本地计算机或 Google Cloud Shell 上运行以下命令。如果从本地计算机运行，请安装 Google Cloud SDK。

创建标准的 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 Console 进行创建。

如下创建具有单主机 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.5.1.
helm install kuberay-operator kuberay/kuberay-operator --version 1.5.1

GKE 提供了一个验证和修改 Webhook，用于处理 TPU Pod 调度并引导在 JAX 初始化中使用的某些环境变量。Ray TPU Webhook 需要 KubeRay Operator 版本至少为 v1.1.0。GKE 在 GKE 版本 1.30.0-gke.1747000 或更高版本中通过 Ray Operator Addon 自动安装 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