本教程演示了如何在 Google Kubernetes Engine (GKE) 上编排多个多切片工作负载，以提高资源利用率。您可以部署一个 Jax 工作负载作为示例，在 TPU 多切片上运行该工作负载，并使用 JobSet 和 Kueue 实现 Job 队列。Kueue 会根据可用资源、配额和层次结构来确定 Job 应何时运行，以在团队之间实现公平共享。
本教程适用于希望使用 Kubernetes 的容器编排功能训练 LLM 的机器学习 (ML) 工程师、平台管理员和运维人员。如需详细了解我们在 Google Cloud 内容中提及的常见角色和示例任务，请参阅常见的 GKE 用户角色和任务。
在阅读本页面之前，请确保您熟悉以下内容：
- Cloud TPU 系统架构中的当前 TPU 版本可用性
- GKE 中的 TPU 多切片
目标
- 使用具有三个 v5e TPU 切片的 GKE 集群准备环境。每个 TPU 切片都有一个包含 8 个芯片的
2x4拓扑。因此，总共 24 个 TPU v5e TPU 芯片。 - 创建 Kueue 资源，以确保在工作负载之间公平共享配额。
- 运行多切片工作负载。
准备工作
在开始之前，请确保您已执行以下任务：
- 启用 Google Kubernetes Engine API。 启用 Google Kubernetes Engine API
- 如果您要使用 Google Cloud CLI 执行此任务，请安装并初始化 gcloud CLI。 如果您之前安装了 gcloud CLI，请运行
gcloud components update以获取最新版本。
准备环境
在 Google Cloud 控制台中，启动 Cloud Shell 实例：
打开 Cloud Shell使用
gcloud config set命令设置默认环境变量：gcloud config set project PROJECT_ID请将 PROJECT_ID 替换为您的 Google Cloud 项目 ID。
默认情况下，运行 1.29.2-gke.1521000 版或更高版本的 Autopilot 集群会启用 TPU。Autopilot 集群上的 TPU 在工作负载规范中进行配置。如需了解详情，请参阅使用 JobSet 定义多切片工作负载部分。
创建 GKE 集群
在 Cloud Shell 中创建一个 GKE 集群：
Autopilot
gcloud container clusters create-auto multislice-cluster \
--location=LOCATION \
--cluster-version 1.29.2-gke.1521000 \
--release-channel rapid
在此命令中：
--location标志用于指定集群的 Compute Engine 位置。--cluster-version标志用于指定集群的 Kubernetes 版本。--release-channel标志用于指定集群的发布渠道。在这种情况下，快速渠道支持 GKE 中提供的最新版本。
标准
gcloud container clusters create multislice-cluster \
--location=LOCATION
将 LOCATION 替换为您要在其中创建集群的位置。确保集群有足够的容量用于 ct5lp-hightpu-4t 机器类型。集群创建过程可能需要几分钟时间。
如果您使用 GKE Autopilot 模式，请跳至创建 Kueue 资源部分。默认情况下，运行 1.29.2-gke.1521000 版或更高版本的 Autopilot 集群会启用 TPU。
创建三个 Standard 模式 TPU 切片节点池
在本部分中，您将使用 gcloud beta container node-pools create 命令创建 TPU 节点池。
创建第一个节点池，名为
nodepool1：gcloud beta container node-pools create nodepool1 \ --location=LOCATION \ --cluster=multislice-cluster \ --node-locations=NODE_LOCATION \ --machine-type=ct5lp-hightpu-4t \ --tpu-topology=2x4 \ --project=PROJECT_ID将 NODE_LOCATION 替换为您要在其中创建节点的集群区域中的一个或多个可用区。
创建第二个节点池，名为
nodepool2：gcloud beta container node-pools create nodepool2 \ --location=LOCATION \ --cluster=multislice-cluster \ --node-locations=NODE_LOCATION \ --machine-type=ct5lp-hightpu-4t \ --tpu-topology=2x4 \ --project=PROJECT_ID创建第三个节点池，名为
nodepool3：gcloud beta container node-pools create nodepool3 \ --location=LOCATION \ --cluster=multislice-cluster \ --node-locations=NODE_LOCATION \ --machine-type=ct5lp-hightpu-4t \ --tpu-topology=2x4 \ --project=PROJECT_ID
GKE 会创建三个节点池。每个节点池都是一个单独的 TPU 切片。
在前面的步骤中，您使用了 gcloud beta container node-pools create 命令来创建节点池。这些命令使用以下标志：
--node-locations：GKE 在其中创建节点池的一个或多个可用区的英文逗号分隔列表。--machine-type：用于节点的机器类型。 在此示例中，您使用了ct5lp-hightpu-4t。如需详细了解与 TPU 兼容的机器类型，请使用选择 TPU 版本中的表。--tpu-topology：要用于节点池的 TPU 拓扑。在此示例中，您使用了2x4。如需详细了解 TPU 拓扑，请参阅选择 TPU 拓扑。
创建 Kueue 资源
创建以下
kueue.yaml清单：apiVersion: kueue.x-k8s.io/v1beta1 kind: ResourceFlavor metadata: name: "vlp-24" spec: nodeLabels: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 --- apiVersion: kueue.x-k8s.io/v1beta1 kind: ClusterQueue metadata: name: "cluster-queue" spec: namespaceSelector: {} queueingStrategy: BestEffortFIFO resourceGroups: - coveredResources: ["google.com/tpu"] flavors: - name: "vlp-24" resources: - name: "google.com/tpu" nominalQuota: 24 --- apiVersion: kueue.x-k8s.io/v1beta1 kind: LocalQueue metadata: namespace: default name: multislice-queue spec: clusterQueue: cluster-queue应用
kueue.yaml清单：kubectl apply -f kueue.yaml
GKE 会创建以下 Kueue 资源：
- ResourceFlavor：集群中资源的抽象。在此示例中，GKE 会创建三个具有
2x4拓扑的 TPU 切片。每个 TPU 切片都有一个包含 8 个芯片的2x4拓扑（总共 24 个 TPU 芯片）。 - ClusterQueue：管理工作负载和集群资源的全局队列。
- LocalQueue：对密切相关的工作负载分组，这些工作负载通常由单个租户（用户）运行。每个 LocalQueue 都指向一个 ClusterQueue，系统会从 ClusterQueue 分配资源以运行其工作负载。Kueue 工作负载是一种表示批量工作负载的抽象，在这种情况下，每个工作负载都是一个 JobSet。
使用 JobSet 定义多切片工作负载
在本部分中，您将创建三个 JobSet。Jobset 是一种工作负载 API，可让您将一组 Kubernetes Job 作为一个单元进行管理。JobSet 最常见的应用场景是分布式训练，但您也可以使用它来运行批量工作负载。
以下 JobSet 运行 Jax 工作负载，该工作负载会输出切片中全局数量的 TPU 芯片，然后休眠 60 秒以模拟一些模型训练时间，然后退出。
在集群中安装 JobSet API：
VERSION=v0.8.1 kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/$VERSION/manifests.yaml创建以下
jobsets-multislice.yaml清单：Autopilot
apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: multislice-1slice labels: kueue.x-k8s.io/queue-name: multislice-queue annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: 1 template: spec: parallelism: 2 completions: 2 backoffLimit: 0 template: spec: nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 command: - bash - -c - | pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html python -c 'import jax; print("Global device count:", jax.device_count())' resources: limits: google.com/tpu: 4 --- apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: multislice-2slice labels: kueue.x-k8s.io/queue-name: multislice-queue annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: 2 template: spec: parallelism: 2 completions: 2 backoffLimit: 0 template: spec: nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 command: - bash - -c - | pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html python -c 'import jax; print("Global device count:", jax.device_count())' sleep 60 resources: limits: google.com/tpu: 4 --- apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: multislice-3slice labels: kueue.x-k8s.io/queue-name: multislice-queue annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: 3 template: spec: parallelism: 2 completions: 2 backoffLimit: 0 template: spec: nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 command: - bash - -c - | sleep 60 resources: limits: google.com/tpu: 4标准
apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: multislice-1slice labels: kueue.x-k8s.io/queue-name: multislice-queue annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: 1 template: spec: parallelism: 2 completions: 2 backoffLimit: 0 template: spec: hostNetwork: true dnsPolicy: ClusterFirstWithHostNet nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 securityContext: privileged: true command: - bash - -c - | pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html python -c 'import jax; print("Global device count:", jax.device_count())' resources: limits: google.com/tpu: 4 --- apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: multislice-2slice labels: kueue.x-k8s.io/queue-name: multislice-queue annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: 2 template: spec: parallelism: 2 completions: 2 backoffLimit: 0 template: spec: hostNetwork: true dnsPolicy: ClusterFirstWithHostNet nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 securityContext: privileged: true command: - bash - -c - | pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html python -c 'import jax; print("Global device count:", jax.device_count())' sleep 60 resources: limits: google.com/tpu: 4 --- apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: multislice-3slice labels: kueue.x-k8s.io/queue-name: multislice-queue annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: 3 template: spec: parallelism: 2 completions: 2 backoffLimit: 0 template: spec: hostNetwork: true dnsPolicy: ClusterFirstWithHostNet nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 securityContext: privileged: true command: - bash - -c - | sleep 60 resources: limits: google.com/tpu: 4应用
jobsets-multislice.yaml清单：kubectl apply -f jobsets-multislice.yaml
GKE 会使用以下资源请求创建 Job：
multislice-1sliceJobSet 创建一个 Job，总共需要一个 TPU 切片。multislice-2sliceJobSet 创建两个 Job，总共需要两个 TPU 切片。multislice-3sliceJobSet 创建三个 Job，总共需要三个 TPU 切片。
由于集群只有三个 TPU 切片，因此并非所有 JobSet 可以同时运行。当 Kueue 将所有三个 multislice-3slice JobSet 加入队列时，其 Job 会单独运行以完成。multislice-1slice 和 multislice-2slice 会等待，然后一起运行。
验证 Kueue 是否已允许工作负载
检查 Kueue 中已加入队列的工作负载：
kubectl get workloads输出类似于以下内容：
NAME QUEUE ADMITTED BY AGE jobset-multislice-1slice-2530a multislice-queue 3s jobset-multislice-2slice-ffb02 multislice-queue 4s jobset-multislice-3slice-8c695 multislice-queue cluster-queue 10s
Kueue 会将一个或多个工作负载加入队列，具体取决于它们所需的 TPU 资源。
监控工作负载
Google Cloud 控制台中的 JobSet 和节点池可观测性指标和信息中心已正式发布。
信息中心
如需查看 GKE 上 TPU 多主机节点池的状态，请前往 Cloud Monitoring 提供的 GKE TPU 节点池状态信息中心：
如需了解详情，请参阅监控 TPU 节点和节点池的健康状况指标。
在Google Cloud 控制台的 Kubernetes Engine AI/ML页面中，AI 部署 > 作业标签页显示 JobSet 监控信息中心，其中包含有关 JobSet 及其底层基础设施的健康状况和性能的全面信息，例如 JobSet 状态、副本就绪性、副本状态。该信息中心还包含基础设施指标，包括 CPU、GPU、TPU、内存和存储指标。如需了解详情，请参阅使用指标监控 JobSet 健康状况。
监控哪些 Pod 正在运行
kubectl get pods
输出类似于以下内容：
NAME READY STATUS RESTARTS AGE
multislice-1slice-slice-0-0-pf2ll 1/1 Running 0 1s
multislice-1slice-slice-0-1-55g62 1/1 Running 0 1s
multislice-2slice-slice-0-0-f4hf7 1/1 Running 0 3s
multislice-2slice-slice-0-1-c8kv7 1/1 Running 0 3s
multislice-2slice-slice-1-0-7h46t 1/1 Running 0 3s
multislice-2slice-slice-1-1-lj9hb 1/1 Running 0 3s
multislice-3slice-slice-0-0-wzq9t 0/1 Completed 0 2m31s
multislice-3slice-slice-0-1-zf4dp 0/1 Completed 0 2m30s
multislice-3slice-slice-1-0-hbfn5 0/1 Completed 0 2m31s
multislice-3slice-slice-1-1-45fgl 0/1 Completed 0 2m30s
multislice-3slice-slice-2-0-wjbp4 0/1 Completed 0 2m30s
multislice-3slice-slice-2-1-lwnvs 0/1 Completed 0 2m30s
看到 GKE 先为 multislice-3slice 安排、创建和运行了 pod。然后，GKE 从 multislice-1slice 和 multislice-2slice JobSet 运行了 Pod。
使用指标监控 JobSet 健康状况
如需了解 JobSet 是否按预期运行，或推断其是否被中断，您可以使用 JobSet 指标包中的 Prometheus 指标，例如 kube_jobset_succeeded_replicas。
请注意，只有 GKE 1.32.1-gke.135700 版更高版本支持 Jobset 健康状况指标。在具有受支持版本的新创建的集群中，JobSet 健康状况指标默认处于启用状态。对于升级到受支持版本的现有集群，客户需要手动启用 JobSet 指标包。如需了解详情，请参阅此文档。
在本教程中，请使用以下 PromQL 查询检查 JobSet 完成情况：
kube_jobset_succeeded_replicas{
cluster="multislice-cluster",
jobset_name=~"mulitslice-.*"}
监控 JobSet 正常运行时间、恢复时间 (TTR) 和中断之间的时间 (TBI)
以下指标对于监控 JobSet 的可用性很有用：
kubernetes.io/jobset/uptime：JobSet 处于可用状态的总时间。kubernetes.io/jobset/times_to_recover：JobSet 的恢复期的分布情况。每个样本都表示 JobSet 的停机时间段内的单个恢复事件。kubernetes.io/jobset/times_between_interruptions：JobSet 中上一次中断结束与当前中断开始之间的间隔时间的分布。每个样本都表示上一次中断与当前中断之间的一段时长。
这些指标适用于仅包含一个 GPU 或 TPU 复制作业的 JobSet。指标的计算仅基于该单个复制作业的可用性。所有 GKE 版本均支持这些指标。
如需查看您在本教程中使用的 JobSet 的正常运行时间，请运行以下 PromQL 查询：
avg_over_time(
kubernetes_io:jobset_uptime{
monitored_resource="k8s_entity", entity_type="jobset",
entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}])
如需查看本教程中 JobSet 的 TBI 分布，请运行以下 PromQL 查询：
histogram_quantile(0.50,
sum_over_time(
kubernetes_io:jobset_times_between_interruptions_bucket{
monitored_resource="k8s_entity",entity_type="jobset",
entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))
您可以将查询的间隔时间延长到更长的时间范围（例如 7 天），并计算此时间段内的平均中断时间 (MTBI)：
sum(sum_over_time(
kubernetes_io:jobset_times_between_interruptions_sum{
monitored_resource="k8s_entity",entity_type="jobset",
entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))
/
sum(sum_over_time(
kubernetes_io:jobset_times_between_interruptions_count{
monitored_resource="k8s_entity",entity_type="jobset",
entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))
如需查看 TTR 分布，您可以运行以下 PromQL 查询：
histogram_quantile(0.50,
sum_over_time(
kubernetes_io:jobset_times_to_recover_bucket{
monitored_resource="k8s_entity",entity_type="jobset",
entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))
将查询间隔时间增加到较长的时间范围（例如 7 天）后，您可以计算此时间段内的平均恢复时间 (MTTR)：
sum(sum_over_time(
kubernetes_io:jobset_times_to_recover_sum{
monitored_resource="k8s_entity",entity_type="jobset",
entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))
/
sum(sum_over_time(
kubernetes_io:jobset_times_to_recover_count{
monitored_resource="k8s_entity",entity_type="jobset",
entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))
启用 Kueue 工作负载优先级和抢占
（可选）您可以为 Kueue 工作负载分配优先级，以确定 Kueue 允许已加入队列的工作负载的顺序。
更新
ClusterQueue，使其具有抢占政策：apiVersion: kueue.x-k8s.io/v1beta1 kind: ResourceFlavor metadata: name: "vlp-24" spec: nodeLabels: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 --- apiVersion: kueue.x-k8s.io/v1beta1 kind: ClusterQueue metadata: name: "cluster-queue" spec: namespaceSelector: {} resourceGroups: - coveredResources: ["google.com/tpu"] flavors: - name: "vlp-24" resources: - name: "google.com/tpu" nominalQuota: 24 preemption: reclaimWithinCohort: Any withinClusterQueue: LowerPriority --- apiVersion: kueue.x-k8s.io/v1beta1 kind: LocalQueue metadata: namespace: default name: multislice-queue spec: clusterQueue: cluster-queue为要分配给工作负载的每个不同优先级创建
PriorityClass：apiVersion: scheduling.k8s.io/v1 kind: PriorityClass metadata: name: low-priority value: 100 globalDefault: false description: "This low priority class should be used for some Pods only."为 JobSet 分配
priorityClassName：Autopilot
apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: low-priority labels: kueue.x-k8s.io/queue-name: multislice-queue annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: 1 template: spec: parallelism: 2 completions: 2 backoffLimit: 0 template: spec: nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 priorityClassName: low-priority containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 command: - bash - -c - | sleep 60 resources: limits: google.com/tpu: 4 # Number of TPU chips per worker标准
apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: low-priority labels: kueue.x-k8s.io/queue-name: multislice-queue annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: 1 template: spec: parallelism: 2 completions: 2 backoffLimit: 0 template: spec: hostNetwork: true dnsPolicy: ClusterFirstWithHostNet nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 priorityClassName: low-priority containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 securityContext: privileged: true command: - bash - -c - | sleep 60 resources: limits: google.com/tpu: 4 # Number of TPU chips per worker
GKE 包含抢占政策，用于定义 Kueue 分配可用资源的方式。该政策指定，如果优先级较高的工作负载需要资源，则工作负载可能会被抢占。优先级值较低的工作负载更可能被优先级较高的工作负载抢占。
清理
为避免因本教程中使用的资源导致您的 Google Cloud 账号产生费用，请删除包含这些资源的项目，或者保留项目但删除各个资源。
删除项目
- In the Google Cloud console, go to the Manage resources page.
- In the project list, select the project that you want to delete, and then click Delete.
- In the dialog, type the project ID, and then click Shut down to delete the project.
逐个删除资源
删除 Kueue 资源：
kubectl delete -f jobsets-multislice.yaml kubectl delete -f kueue.yaml删除集群：
gcloud container clusters delete multislice-cluster --region=LOCATION
后续步骤
- 详细了解 Kueue。
- 了解如何在 GKE 上通过命名空间之间的配额共享实现 Job 排队系统。