使用 GKE 动态切分配置 TPU

本文档介绍了如何在 Google Kubernetes Engine (GKE) 中使用动态切片。借助动态切分,您可以将已配置的 TPU 子块配置为不同的拓扑。此功能可减少重新创建节点池的需求,在发生故障时允许自动恢复,从而增强容错能力,并优化资源利用率。

本文档适用于希望优化 TPU 利用率、缩短预配时间并提高大规模训练和推理工作负载的容错能力的 AI/ML 工程师和平台管理员。

在阅读本文档之前,请确保您熟悉以下内容:

什么是动态切片?

借助动态切分,您可以解耦 TPU 预配,从而灵活管理 Cloud TPU 容量。动态切片涉及以下流程:

  1. 您以称为子块的较小单位来预配资源。子块是 Ironwood (TPU7x) 容量的基本逻辑构建单元。对于 Ironwood (TPU7x),它表示一个由 16 个 TPU 虚拟机组成的组,这些虚拟机具有 4x4x4 拓扑的互连 TPU 芯片。在 TPU 全容量模式和动态切分上下文中,节点池直接映射到子块。
  2. 然后,动态切片会将这些子块拼接成更大的切片。

动态切片的优势

动态切片可帮助您实现以下目标:

  • 缩短配置时间:单独配置子块可加快整体配置速度,因为这样可以最大限度地减少任何单个故障的影响。
  • 缩短恢复时间:如果发生 TPU 芯片故障,最小的故障单元是子块。动态切分可隔离故障子块,以便比重新配置整个大型切片更快地在健康子块上重新调度工作负载。
  • 调整容量:如果您有多种不同的工作负载要求,则无需删除并重新创建节点池以进行拓扑更改(如果没有动态切分,这是必要的)。您可以动态重新配置已预配的节点池,以匹配指定的配置。

动态切片的关键要素

动态切片引入了以下关键概念:

  • 节点池的增量预配:动态切片使用增量预配,这是一种容错的节点池预配模型。此模型会将您的所有 TPU 容量转换为包含 16 个 TPU 虚拟机的节点池。
  • Slice 控制器:在 GKE 控制平面内运行的 Kubernetes 自定义资源控制器,用于管理动态切片。切片控制器用于管理 Slice 自定义资源的生命周期,该资源表示动态切片。切片控制器负责创建、持续监控和删除 Slice。使用调度程序时,调度程序会指导创建和删除 Slice 自定义资源。
  • 切片自定义资源:根据所请求的 TPU 拓扑动态地将这些子块拼接在一起。此过程依赖于 OCS 网络的动态重配置来连接 TPU 节点池,这有助于确保优化性能。您可以通过检查 Slice 自定义资源的状态字段来检查动态切片形成进度或健康状况。

要求

如需在 GKE 中使用动态切分,必须满足以下要求:

  • 使用快速渠道中 1.35.0-gke.274500 版或更高版本的 Standard 集群。
  • 使用 Ironwood (TPU7x) 版本。
  • 为节点使用 Container-Optimized OS 映像。
  • 如需使用增量配置,请使用全容量模式预留。全容量模式是一项由 TPU Cluster Director 启用的功能。

准备工作

在开始之前,请确保您已执行以下任务:

  • 启用 Google Kubernetes Engine API。
  • 启用 Google Kubernetes Engine API
  • 如果您要使用 Google Cloud CLI 执行此任务,请安装初始化 gcloud CLI。 如果您之前安装了 gcloud CLI,请通过运行 gcloud components update 命令来获取最新版本。较早版本的 gcloud CLI 可能不支持运行本文档中的命令。

限制

  • 单个切片必须使用预留中的同一 TPU 块内的子块。如需跨 TPU 块使用子块,请使用 TPU 多切片
  • 动态切片不支持小于 4x4x4 的拓扑。

在 GKE 中使用 Kueue 实现动态切分

本部分介绍了在 GKE 中使用动态切片的流程。

  1. 查看“全容量”模式预留的拓扑和健康状况
  2. 在集群中启用 slice 控制器
  3. 创建 TPU 节点池
  4. 配置 Kueue 以创建 Slice 自定义资源
  5. 使用 Kueue 在动态切片上运行工作负载
  6. 清理.

启用 Slice 控制器

如需使用动态切片,请在集群中启用切片控制器。

  1. 更新集群:

    gcloud container clusters update CLUSTER_NAME \
        --location=LOCATION \
        --enable-slice-controller
    

    替换以下内容:

    • CLUSTER_NAME:您的集群的名称。
    • LOCATION:具有可用 TPU 容量的区域。
  2. 获取凭据,以便您可以使用 kubectl 命令与集群通信:

    gcloud config set container/cluster CLUSTER_NAME
    gcloud container clusters get-credentials CLUSTER_NAME \
        --location=LOCATION
    
  3. 在以下命令的输出中,验证是否存在 slices.accelerator.gke.io 值:

    kubectl get crd slices.accelerator.gke.io
    

    输出类似于以下内容:

    slices.accelerator.gke.io                2026-01-09T23:58:02Z
    

创建具有增量预配的节点池

本部分介绍了如何通过增量配置创建 TPU 节点池。GKE 会将所有 TPU 容量转换为由 16 个 TPU 虚拟机组成的节点池或子块。即使 GKE 无法找到所有 16 个健康虚拟机,它也会通过将节点放置在宿主机上健康的部分,并在修复健康状况不佳的机器时逐步配置这些机器,来配置这些节点池。

您可以将节点池定位为属于以下任一对象:

  • 在全容量模式预留中公开的特定 TPU 块。块定位功能可让 GKE 在指定块内的任何可用子块中创建节点池。
  • TPU 的特定子块或特定 16 节点 TPU 虚拟机组,用于实现更精细的控制。

创建工作负载政策

如需使用 Ironwood (TPU7x) 创建 TPU 切片节点池,您必须先创建一个工作负载政策,并将 accelerator-topology-mode 字段设置为 provision_only。此设置会触发增量配置流程。

创建工作负载政策:

gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
        --project=PROJECT_ID \
        --region=REGION  \
        --type=HIGH_THROUGHPUT \
        --accelerator-topology=4x4x4 \
        --accelerator-topology-mode=provision_only

替换以下内容:

  • WORKLOAD_POLICY_NAME:工作负载政策的名称。
  • PROJECT_ID:您的 Google Cloud 项目 ID。
  • REGION:工作负载政策的区域。

在此命令中,执行以下操作:

  • 始终将 accelerator-topology 字段设置为 4x4x4,以匹配单个子块内的芯片总数。
  • 请务必将 accelerator-topology-mode 字段设置为 provision_only,以确保触发增量配置流程。设置 provision_only 字段后,节点池会预配 TPU 节点,而不形成 ICI 或 OCS 链接。

将节点池的目标设置为属于某个块或子块

您可以在“所有容量”模式预订中定位特定子块或块。

  • 以块为目标:每个节点池都使用指定块中的容量。GKE 会将节点池放置在该块中的可用子块内。您必须创建与要使用的块中的子块数量相同的节点池。
  • 以子块为目标:每个节点池都映射到特定的可用子块。使用子块定位时,如果至少有一个虚拟机的健康状况良好,GKE 就会创建节点池。增量配置可确保所有节点都放置在指定的子块内。

屏蔽

  1. 如需检索预留中的块名称以及块中可用子块的数量,请完成查看“所有容量模式”预留的拓扑和健康状况文档中的以下步骤:

    1. 通过列出所有预留块并复制 name: 字段中的值来确定块的名称。此值是相应块或 BLOCK_NAME 在此文档中的名称。

    2. 通过描述预留块并确定 reservationSubBlockCount 字段中的值,来确定要创建的节点池数量。此值表示可用的子块数量。例如,reservationSubBlockCount: 4 值表示该块有 4 个可用的子块,您需要创建 4 个单独的节点池。

  2. 设置预留路径:

    export RESERVATION_PATH="projects/PROJECT_ID/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME"
    

    替换以下内容:

    • RESERVATION_NAME:TPU 预留的名称。
    • BLOCK_NAME:块的名称。
  3. 为上一步中确定的每个子块创建一个节点池。 例如,如果数量为 4,请运行此命令四次。为每个节点池使用唯一名称。

    gcloud container node-pools create NODE_POOL_NAME \
          --cluster=CLUSTER_NAME \
          --node-locations=ZONE \
          --machine-type=tpu7x-standard-4t \
          --num-nodes=16 \
          --placement-policy=WORKLOAD_POLICY_NAME \
          --reservation-affinity=specific \
          --reservation=${RESERVATION_PATH}
    

    替换以下内容:

    • NODE_POOL_NAME:新节点池的名称。
    • CLUSTER_NAME:GKE 集群的名称。
    • WORKLOAD_POLICY_NAME:您创建的工作负载政策的名称。
    • ZONE:节点池的可用区,例如 us-central1-a

子区块

  1. 如需检索块的名称和可用子块的 ID,请在查看“所有容量模式”预留的拓扑和运行状况文档中完成以下步骤:

    1. 如需确定块的名称,请列出所有预留块,然后复制 name: 字段中的值。此值是相应块或 BLOCK_NAME 在相应文档中的名称。

    2. 如需确定子块的名称,请列出某个块的所有子块,然后复制 reservationSubBlocks 下每个条目的 name: 字段中的值。此值是本文档中子块或 SUBBLOCK_NAME 的名称。

  2. 设置预留路径:

    export RESERVATION_PATH="projects/PROJECT_ID/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME/reservationSubBlocks/SUBBLOCK_NAME"
    

    替换以下内容:

    • RESERVATION_NAME:TPU 预留的名称。
    • BLOCK_NAME:块的名称。
    • SUBBLOCK_NAME:子块的名称。
  3. 创建节点池:

    gcloud container node-pools create NODE_POOL_NAME \
            --project=PROJECT_ID \
            --cluster=CLUSTER_NAME \
            --node-locations=ZONE \
            --machine-type=tpu7x-standard-4t \
            --num-nodes=16 \
            --placement-policy=WORKLOAD_POLICY_NAME \
            --reservation-affinity=specific \
            --reservation=${RESERVATION_PATH}
    

    替换以下内容:

    • NODE_POOL_NAME:新节点池的唯一名称,例如 sub-block-pool-1
    • PROJECT_ID:您的 Google Cloud 项目 ID。
    • CLUSTER_NAME:GKE 集群的名称。
    • ZONE:节点池的可用区,例如 us-central2-b
    • WORKLOAD_POLICY_NAME:您创建的工作负载政策的名称。

在此阶段,节点已创建,但其芯片间互联 (ICI) 链接尚未处于活跃状态。因此,您无法直接在这些节点池上运行工作负载。

如需启用所有必要的 ICI 链接以形成切片并允许调度工作负载,请使用以下方法之一创建动态切片:

  • 创建 Slice 自定义资源。您可以使用 Slice 自定义资源来定义指定的拓扑,然后由 slice 控制器激活该拓扑,而不是使用 Pod。
  • 通过 KueueTAS 安排 GKE 工作负载。 Kueue 会自动处理 Slice 自定义资源的创建和删除。避免手动修改由 Kueue 创建的 Slice 自定义资源。

使用 Kueue 和 TAS 创建动态 slice

在本部分中,您将使用 Kueue 和 TAS 调度 GKE 工作负载。

安装 JobSet 和 Kueue 资源以实现动态切片

  1. 安装 JobSet:

    helm install jobset oci://registry.k8s.io/jobset/charts/jobset \
            --version 0.10.1 \
            --namespace jobset-system \
            --create-namespace \
            --set controller.resources.requests.cpu=4 \
            --set controller.resources.requests.memory=16Gi
    
  2. 安装 Kueue

    helm install kueue oci://registry.k8s.io/kueue/charts/kueue \
            --version 0.16.1 \
            --namespace kueue-system \
            --create-namespace \
            --wait \
            --set controllerManager.replicas=3 \
            --set controllerManager.manager.resources.requests.cpu=16 \
            --set controllerManager.manager.resources.requests.memory=64Gi
    
  3. 安装 Kueue slice 控制器:

    kubectl apply -f https://gist.githubusercontent.com/mwysokin/cd90010d0d375b3bf57c536905692547/raw/506c36dd070f4ac222ba8a5e58ba28bbfcfa8ed3/kueue-slice-controller-v0.8.0-130.yaml
    
  4. 如需为动态切分配置 Kueue,请将以下清单保存为 dynamic-slice-topology.yaml

    apiVersion: kueue.x-k8s.io/v1beta1
    kind: Topology
    metadata:
      name: superslice-topology
    spec:
      levels:
      # Label to identify the physical block a sub-block belongs to.
      # Only sub-blocks from the same block can form a slice.
      - nodeLabel: cloud.google.com/gce-topology-block
      # Label to identify individual TPU sub-blocks (4x4x4 topology).
      - nodeLabel: cloud.google.com/gke-tpu-partition-4x4x4-id
      # Standard Kubernetes label for individual nodes.
      # Required to assign Pods to specific VMs.
      - nodeLabel: kubernetes.io/hostname
    ---
    apiVersion: kueue.x-k8s.io/v1beta1
    kind: ResourceFlavor
    metadata:
      name: superslice-rf
    spec:
      nodeLabels:
        cloud.google.com/gke-tpu-accelerator: tpu7x
      topologyName: superslice-topology
    ---
    apiVersion: kueue.x-k8s.io/v1beta1
    kind: AdmissionCheck
    metadata:
      name: superslice-ac
    spec:
      controllerName: accelerator.gke.io/slice
    ---
    apiVersion: kueue.x-k8s.io/v1beta1
    kind: ClusterQueue
    metadata:
      name: cq
    spec:
      namespaceSelector: {}
      admissionChecks:
      - superslice-ac
      resourceGroups:
      - coveredResources:
        - google.com/tpu
        flavors:
        - name: superslice-rf
          resources:
          - name: google.com/tpu
            nominalQuota: "999999"  # modeling unlimited quota
    ---
    apiVersion: kueue.x-k8s.io/v1beta1
    kind: LocalQueue
    metadata:
      name: lq
      namespace: default
    spec:
      clusterQueue: cq
    
  5. 应用 dynamic-slice-topology.yaml 清单:

    kubectl apply -f dynamic-slice-topology.yaml
    

    在此清单中,您将通过定义以下资源来为动态切分配置 Kueue:

    • Ironwood (TPU7x) 动态切片拓扑 (superslice-topology):拓扑定义了 Kueue 在调度动态切分工作负载时考虑的级别。这些级别如下所示:
      • cloud.google.com/gce-topology-block 标签:此级别是必需的,因为只有来自同一块的子块才能形成切片,因此需要了解哪些子块属于哪些块。
      • cloud.google.com/gke-tpu-partition-4x4x4-id 标签:此级别表示各个 Ironwood (TPU7x) 子块(4x4x4 拓扑)。
      • kubernetes.io/hostname 标签:需要此级别才能将 Pod 分配给特定虚拟机,并观察其标签和污点。
    • Ironwood (TPU7x) SuperSlice ResourceFlavor (superslice-rf):Ironwood (TPU7x) 子块的资源配置,包含 cloud.google.com/gke-tpu-accelerator: tpu7x 标签,用于匹配具有 Ironwood (TPU7x) 机器的节点。
    • SuperSlice AdmissionCheck (superslice-ac):此准入检查会告知 Kueue 在 GKE slice 控制器确认 slice 已变为有效状态之前,不要调度工作负载。首先定义准入检查,然后将其添加到处理动态切片工作负载的 ClusterQueue
    • ClusterQueue (cq) 和 LocalQueue (lq):这些字段用于管理 google.com/tpu 资源。cq ClusterQueue 包含 superslice-ac 准入检查。google.com/tpunominalQuota 可以通过以下两种方式进行配置:
      • 特定配额:将 nominalQuota 设置为与现有容量相匹配,以实现公平共享和配额管理。
      • 无限配额:将 nominalQuota 设置为非常高的值(例如 "999999"),以模拟无限配额。为了专注于 TAS 和动态切分,此配置会绕过 Kueue 的配额管理功能。

定义子区块健康状况选择

除了标准的节点健康状况和就绪状态之外,GKE 还使用 cloud.google.com/gke-tpu-partition-4x4x4-state 标签公开每个子块的特定状态。借助此标签,GKE 可以考虑影响切片形成的因素,例如 TPU 链接的状态。

您可以按如下方式定义 cloud.google.com/gke-tpu-partition-4x4x4-state 标签的值:

  • HEALTHY:基础设施运行正常。
  • DEGRADED:子块的基础设施处于降级状态,例如,由于 OCS 链路降级。子块仍可形成切片,但与正常的子块相比,整体性能可能会降低。如果您可以容忍可能降低的性能,则可以使用节点亲和性将工作负载配置为使用 DEGRADED 子块,如示例 3 所示。
  • UNHEALTHY:子块不健康,无法形成切片。

Kueue Slice Controller 网络钩子会验证工作负载是否包含特定的子块健康要求。如果未指明任何偏好设置,Webhook 会注入默认的节点亲和性。

行为如下:

  • 如果存在以 cloud.google.com/gke-tpu-partition-4x4x4-state 标签为目标的 nodeSelectornodeAffinity,则该 nodeSelectornodeAffinity 保持不变。
  • 如果不存在此类标签配置,Webhook 会注入以下默认节点亲和性,以确保仅使用可用的子块:

    nodeSelector:
      cloud.google.com/gke-tpu-partition-4x4x4-state: "HEALTHY"
    

以下部分包含一些示例,其中配置了 cloud.google.com/gke-tpu-partition-4x4x4-state 标签以指定不同的子块健康状况配置。

使用 Kueue 在动态切片上运行测试工作负载

本部分介绍了如何使用 Kueue 和 TAS 在动态切片上部署工作负载。其中包含三个示例,分别展示了如何创建动态 slice 工作负载和包含多个 slice 的工作负载。工作负载以 JobSet 的形式提交。

示例 1:单个工作负载使用单个动态切片

以下示例介绍了如何使用具有 4x12x16 拓扑(由 12 个子块组成)的切片创建工作负载。Pod 的数量计算如下:(4 * 12 * 16)/ 每个节点 4 个芯片 = 192 个 Pod。

  1. 将以下清单保存为 big-super-slice.yaml

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: big-super-slice
      labels:
        kueue.x-k8s.io/queue-name: lq
      annotations:
    spec:
      replicatedJobs:
        - name: job-jax
          replicas: 1
          template:
            spec:
              parallelism: 192  # pods per slice calculation: 4*12*16 / 4 = 192
              completions: 192
              backoffLimit: 10
              template:
                metadata:
                  annotations:
                    cloud.google.com/gke-tpu-slice-topology: 4x12x16
                spec:
                  tolerations:
                    - key: "google.com/tpu"
                      operator: "Equal"
                      value: "present"
                      effect: "NoSchedule"
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: tpu7x
                    cloud.google.com/gke-tpu-partition-4x4x4-state: "HEALTHY"
                  containers:
                    - name: jax
                      image: python:latest
                      command:
                        - bash
                        - -c
                        - |
                          printenv
                          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(), "Local device count:", jax.local_device_count())'
                      resources:
                        limits:
                          google.com/tpu: 4
                  restartPolicy: Never
    

    在此清单中,以下注释会告知 Kueue 切片特征和拓扑,以配置以下内容:

    • cloud.google.com/gke-tpu-slice-topology:指定 "4x12x16" 作为动态切片拓扑。tpu7x 加速器拓扑的要求包括以下规则:
      • 最小拓扑为 4x4x4
      • 拓扑必须是 AxBxC 格式的三维字符串。例如 4x8x8
      • 每个维度(A、B 和 C)都必须是 4 的倍数。
      • 各维度必须按非降序排序:A <= B <= C。例如,4x8x4 无效,应为 4x4x8
      • 各维度相乘的积 (ABC) 不得超过 9,216。
      • 支持的最大切片拓扑最多可包含 32 个子块。例如,具有 32 个子块的 8x16x16、具有 30 个子块的 8x12x20 或具有 27 个子块的 12x12x12 都在可接受的范围内。
    • cloud.google.com/gke-tpu-accelerator: tpu7x:在运行 Ironwood (TPU7x) 的虚拟机上调度 Pod。
    • kueue.x-k8s.io/queue-name:将 JobSet 分配给 Kueue LocalQueue。
  2. 应用 big-super-slice.yaml 清单:

    kubectl apply -f big-super-slice.yaml
    

    应用清单后,Kueue 会创建一个名为 big-super-sliceJobSet。然后,Kueue 会尝试形成具有 4x12x16 拓扑的单个动态切片。切片处于有效状态后,Kueue 会允许工作负载运行,并将 192 个 Pod 调度到节点上,以形成运行工作负载的动态切片。

示例 2:具有多个副本的工作负载

以下示例演示了如何创建使用两个动态切片的工作负载,每个动态切片由四个子块组成。

  1. 将以下清单保存为 two-super-slices.yaml

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
    name: two-super-slices
    labels:
        kueue.x-k8s.io/queue-name: lq
    annotations:
    spec:
    replicatedJobs:
        - name: job-jax
        replicas: 2
        template:
            spec:
            parallelism: 64  # Pods per slice calculation: (4*8*8) / 4 = 64
            completions: 64
            backoffLimit: 10
            template:
                metadata:
                annotations:
                    cloud.google.com/gke-tpu-slice-topology: 4x8x8
                spec:
                tolerations:
                    - key: "google.com/tpu"
                    operator: "Equal"
                    value: "present"
                    effect: "NoSchedule"
                nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: tpu7x
                    cloud.google.com/gke-tpu-partition-4x4x4-state: "HEALTHY"
                containers:
                    - name: jax
                    image: python:latest
                    command:
                        - bash
                        - -c
                        - |
                        printenv
                        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(), "Local device count:", jax.local_device_count())'
                    resources:
                        limits:
                        google.com/tpu: 4
                restartPolicy: Never
    
  2. 应用 two-super-slices.yaml 清单:

    kubectl apply -f two-super-slices.yaml
    

在此清单中,您在 replicatedJobs 字段中设置了 replicas: 2。 应用清单后,Kueue 会尝试形成两个具有 4x8x8 拓扑的单独切片。Kueue 会为 jobset.spec.replicatedJobs[].replicas 中定义的每个副本创建一个动态切片。如果指定了 n 个副本,Kueue 会为工作负载创建 n 个动态切片,并等待所有切片变为有效状态,然后才允许工作负载运行。

示例 3:具有单个动态切片和 NodeAffinity 的工作负载

从 Kueue 0.15 开始,Kueue 支持 NodeAffinity,用于 TAS 节点选择。此功能可用于允许 HEALTHYDEGRADED 节点都成为动态 slice 的一部分。以下示例展示了如何配置具有单个动态 slice 和 NodeAffinity 的工作负载:

  1. 将以下清单保存为 slice-8x8x8-na.yaml

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: slice-8x8x8-na
      labels:
        kueue.x-k8s.io/queue-name: lq
    spec:
      replicatedJobs:
        - name: rj1
          replicas: 1
          template:
            spec:
              parallelism: 128
              completions: 128
              backoffLimit: 10
              template:
                metadata:
                  annotations:
                    cloud.google.com/gke-tpu-slice-topology: 8x8x8
                spec:
                  tolerations:
                    - key: "google.com/tpu"
                      operator: "Equal"
                      value: "present"
                      effect: "NoSchedule"
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: tpu7x
                  affinity:
                    nodeAffinity:
                      requiredDuringSchedulingIgnoredDuringExecution:
                        nodeSelectorTerms:
                          - matchExpressions:
                              - key: cloud.google.com/gke-tpu-partition-4x4x4-state
                                operator: In
                                values:
                                  - "HEALTHY"
                                  - "DEGRADED"
                  containers:
                    - name: jax
                      image: python:latest
                      command:
                        - bash
                        - -c
                        - |
                          printenv
                          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(), "Local device count:", jax.local_device_count())'
                      resources:
                        limits:
                          google.com/tpu: 4
                  restartPolicy: Never
    
  2. 应用 slice-8x8x8-na.yaml 清单:

    kubectl apply -f slice-8x8x8-na.yaml
    

    应用清单后,Kueue 会创建一个名为 slice-8x8x8-naJobSet。然后,Kueue 尝试形成一个具有 8x8x8 拓扑的动态切片,由于指定了 NodeAffinity,因此允许同时包含 HEALTHYDEGRADED 节点。切片处于有效状态后,Kueue 会接纳工作负载,并将 128 个 Pod 调度到构成动态切片的节点上。

监控 slice 的状态

如需检查动态切片的状态,请运行以下命令:

kubectl describe slice SLICE_NAME

SLICE_NAME 替换为切片的名称。切片名称通常派生自 JobSet 名称和副本索引。对于示例 1,Kueue 创建的切片将具有类似于 default-jobset-big-super-slice-yyyyy-job-jax-0 的名称。

输出类似于以下内容:

Name:         test-slice
Namespace:
Labels:       <none>
Annotations:  <none>
API Version:  accelerator.gke.io/v1beta1
Kind:         Slice
Metadata:
  Creation Timestamp:  2026-02-12T23:44:28Z
  Finalizers:
    accelerator.gke.io/slice-finalizer
  Generation:        1
  Resource Version:  1770939905695871008
  UID:               6dbbfe14-4486-4462-864d-e078d0ca8b5b
Spec:
  Partition Ids:
    5eae6a4f59d59cf30a9bf49de618eb2b
  Topology:  4x4x4
  Type:      tpu7x
Status:
  Conditions:
    Last Transition Time:  2026-02-12T23:45:05Z
    Message:
    Reason:                ACTIVE
    Status:                True
    Type:                  Ready
    Last Transition Time:  2026-02-12T23:45:05Z
    Message:               NodeLabelingCompleted
    Reason:                NodeLabelIsAdded
    Status:                True
    Type:                  NodeLabeled
Events:                    <none>

切片名称遵循以下规则,以确保与底层 Compute Engine 资源命名惯例兼容:

  • 模板{namespace}-jobset-{jobset.metadata.name}-kueueHash[5-character]-{jobset.spec.replicatedJobs[].name}-sliceIndex
  • 长度:名称不超过 54 个字符。控制器会附加一个连字符和一个 8 字符的集群哈希值来创建 Compute Engine 资源名称,这些名称的长度限制为 63 个字符。
  • 格式:名称与正则表达式 ^[a-z]([-a-z0-9]*[a-z0-9])?$ 匹配。该名称具有以下特征:
    • 以小写字母开头。
    • 只能包含小写字母、数字和连字符 (-)。
    • 以小写字母或数字结尾(不能以连字符结尾)。

清理

为避免产生意外费用,请先删除切片,然后再删除节点池。

  1. 删除 JobSet。此操作会触发 Kueue 删除关联的 Slice 自定义资源。

    kubectl delete jobset JOBSET_NAME
    

    JOBSET_NAME 替换为您的 JobSet 的名称,例如 big-super-slice

  2. 删除 TPU 节点池:

    gcloud container node-pools delete NODE_POOL_NAME \
        --cluster=CLUSTER_NAME \
        --location=LOCATION
    

(可选)将动态切片与您自己的调度程序搭配使用

本文档重点介绍如何使用 Kueue 和 TAS。不过,您也可以使用自己的自定义调度程序来管理动态切片。如果您选择使用其他调度程序,请参阅 Slice 自定义资源参考信息。

后续步骤