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

有关详细信息，请参阅 GKE 文档，或继续阅读以快速开始。

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

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

创建一个标准 GKE 集群并在 us-central2-b 计算区域中启用 Ray Operator：

gcloud container clusters create kuberay-tpu-cluster \
    --addons=RayOperator \
    --machine-type=n1-standard-8 \
    --cluster-version=1.30 \
    --location=us-central2-b

运行以下命令以向集群添加一个TPU节点池。您也可以从Google Cloud Console创建它：

创建一个具有单主机 TPU 拓扑的节点池，如下所示：

gcloud container node-pools create tpu-pool \
  --zone us-central2-b \
  --cluster kuberay-tpu-cluster \
  --num-nodes 1 \
  --min-nodes 0 \
  --max-nodes 10 \
  --enable-autoscaling \
  --machine-type ct4p-hightpu-4t \
  --tpu-topology 2x2x1

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

gcloud container node-pools create tpu-pool \
  --zone us-central2-b \
  --cluster kuberay-tpu-cluster \
  --num-nodes 2 \
  --min-nodes 0 \
  --max-nodes 10 \
  --enable-autoscaling \
  --machine-type ct4p-hightpu-4t \
  --tpu-topology 2x2x2

--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：配置 kubectl 以连接到集群#

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

gcloud container clusters get-credentials kuberay-tpu-cluster --zone us-central2-b

更多详情，请参阅 GKE 文档。

步骤 3：安装 TPU 初始化 webhook#

GKE 提供了一个 验证和变异 webhook 来处理 TPU Pod 调度和引导用于 JAX 初始化的某些环境变量。Ray TPU webhook 需要至少 v1.1.0 版本的 KubeRay 操作符。GKE 通过 Ray Operator Addon 自动安装 Ray TPU webhook，适用于 GKE 版本 1.30.0-gke.1747000 或更高版本。手动安装 Ray TPU webhook 时，您需要 cert-manager 来处理 TLS 证书注入。您可以使用以下 helm 命令在 GKE 标准和 Autopilot 集群中安装 cert-manager：

[可选] 在没有 Ray 插件的情况下手动安装在 GKE 集群中：#

安装 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：

  1. git clone https://github.com/GoogleCloudPlatform/ai-on-gke

  2. cd ray-on-gke/tpu/kuberay-tpu-webhook

  3. make deploy deploy-cert