将动态切片与自定义调度器搭配使用

本文档介绍了如何通过直接与 Slice 自定义资源互动来使用动态切片。您可以创建切片、监控分区状态并验证切片健康状况。

在按照这些说明操作之前，请确保您了解动态切片的概念。

为何要将动态切片与自定义调度器搭配使用？

如果您有复杂的调度要求，或者想将动态切片与现有调度基础架构集成，请使用自己的调度器来管理 Slice 自定义资源。

如果您希望使用调度程序而不是直接管理 Slice 自定义资源，GKE 可与 Kueue 和拓扑感知调度 (TAS) 集成。如需了解详情，请参阅使用 Kueue 和 TAS 调度动态切片。

工作流程概览

如需将动态切片与自定义调度器搭配使用，请执行本文档中的以下任务：

  1. 启用 slice 控制器。
  2. 创建具有增量预配的节点池。
  3. 根据工作负载要求创建 Slice 自定义资源。将 Slice 自定义资源应用到您的集群。
  4. 监控分区状态和切片健康状况。
  5. 完成后，删除切片。

如需详细了解 Slice 自定义资源的字段和状态，请参阅 Slice 自定义资源参考信息。

要求

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

  • 使用快速渠道中的 Standard 集群，版本如下：
    • 对于动态超级切片配置（拓扑等于或大于 4x4x4），请使用 1.35.2-gke.1842000 或更高版本。
    • 对于动态子切片配置（拓扑小于 4x4x4），请使用版本 1.36.0-gke.3712000 或更高版本。
  • 使用 Ironwood (TPU7x) 版本。
  • 为节点使用 Container-Optimized OS 映像。
  • 如需使用增量配置，请使用全容量模式预留。全容量模式是一项由 TPU Cluster Director 启用的功能。
  • 对于动态子切片，请确保您的节点有待处理的维护事件。 监控实例是否存在待处理的维护事件。 如果您的任何节点的待处理维护事件的结束时间介于 2026 年 9 月 18 日至 2026 年 9 月 30 日之间，您必须先在这些节点上手动触发主机维护事件，然后才能使用子切片。

准备工作

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

  • 启用 Google Kubernetes Engine API。
    • 启用 Google Kubernetes Engine API
  • 如果您要使用 Google Cloud CLI 执行此任务，请安装并初始化 gcloud CLI。如果您之前安装了 gcloud CLI，请通过运行 gcloud components update 命令来获取最新版本。较早版本的 gcloud CLI 可能不支持运行本文档中的命令。
  • 确保您已有 1.35.2-gke.1842000 版或更高版本的 Standard 集群（位于快速渠道中）。如需创建新集群，请参阅创建区域级集群。
  • 确保您在相应区域中拥有足够的 Ironwood (TPU7x) 配额。
  • 如果您计划运行多切片工作负载，请安装 JobSet v0.10.1 或更高版本。
  • 在全容量模式下请求 TPU 容量。

启用 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。
  • 通过 Kueue 和 TAS 安排 GKE 工作负载。 Kueue 会自动处理 Slice 自定义资源的创建和删除。避免手动修改由 Kueue 创建的 Slice 自定义资源。

通过超级切片或子切片形成动态切片

创建节点池后，您可以通过创建 Slice 自定义资源来形成更大的动态超级切片或更小的动态子切片。Slice 自定义资源定义了切片控制器激活的指定拓扑。然后，您的工作负载会在此动态 slice 上进行调度和运行。

动态 slice 分区

分区提供可用于为工作负载形成动态切片的拓扑。分区显示每个节点的所有可用拓扑，包括 2x2x1、2x2x2、2x2x4、2x4x4 和 4x4x4。小于 4x4x4 的拓扑需要 GKE 1.36.0-gke.3712000 版或更高版本。

任何大于 4x4x4 的 slice 都没有分区标签，因为它是通过关联多个 4x4x4 分区创建的。

增量配置节点池中的每个 TPU 节点都必须具有每个指定的分区 ID 和分区状态节点标签。

验证节点和分区状态

  1. 如需从节点池中获取节点名称，请运行以下命令：

    kubectl get nodes -l cloud.google.com/gke-nodepool=${NODE_POOL_NAME}

    结果类似于以下内容：

    NAME                                 STATUS   ROLES    AGE    VERSION
gke-np-status-update-7b4c890c-0jhp   Ready    <none>   2d1h   v1.35.1-gke.1396002
gke-np-status-update-7b4c890c-377r   Ready    <none>   2d1h   v1.35.1-gke.1396002
gke-np-status-update-7b4c890c-gb51   Ready    <none>   2d1h   v1.35.1-gke.1396002

  2. 验证节点的配置模型：

    kubectl describe node NODE_NAME | grep "cloud.google.com/gke-accelerator-topology-mode"

    结果类似于以下内容：

    cloud.google.com/gke-accelerator-topology-mode: PROVISION_ONLY

  3. 检索要定位的拓扑分区对应的节点标签信息：

    kubectl describe node NODE_NAME | grep -E "cloud.google.com/gke-tpu-partition-.*-id"

    将 NODE_NAME 替换为节点池中其中一个节点的名称。

    结果类似于以下内容：

    cloud.google.com/gke-tpu-partition-4x4x4-id=fba785f80d18552357dcdef6d3d16c27
cloud.google.com/gke-tpu-partition-2x4x4-id=e18372d627ac412cb24d5ea8ab8912c9
cloud.google.com/gke-tpu-partition-2x2x4-id=7fbd4e29dc1839217fae41a9dd8211b4
cloud.google.com/gke-tpu-partition-2x2x2-id=a9476d1b02bd4f4e75ffffae3bd23c01
cloud.google.com/gke-tpu-partition-2x2x1-id=0bcfe937d1bb3914a8bdcd94e9f73319

  4. 验证节点是否包含 node.gke.io/created-by-mig 注解：

    kubectl describe node NODE_NAME | grep "node.gke.io/created-by-mig"

    将 NODE_NAME 替换为节点池中其中一个节点的名称。

    结果类似于以下内容：

    node.gke.io/created-by-mig: projects/735972712744/zones/us-central1-ai1a/team/string

    输出包含 node.gke.io/created-by-mig 注解，该注解可让 GKE 控制平面将 Kubernetes 节点与其底层 Compute Engine 资源相关联。

  5. 检索要验证的拓扑分区状态的节点标签信息：

    kubectl describe node NODE_NAME | grep -E "cloud.google.com/gke-tpu-partition-.*-state"

    结果类似于以下内容：

    cloud.google.com/gke-tpu-partition-4x4x4-state=HEALTHY
cloud.google.com/gke-tpu-partition-2x4x4-state=HEALTHY
cloud.google.com/gke-tpu-partition-2x2x4-state=HEALTHY
cloud.google.com/gke-tpu-partition-2x2x2-state=HEALTHY
cloud.google.com/gke-tpu-partition-2x2x1-state=HEALTHY

    cloud.google.com/gke-tpu-partition-[shape]-state 标签（其中 [shape] 对应于分区 ID 的拓扑）用于指示相应分区是否可用于构成动态 slice。此状态标签支持以下值：

    • HEALTHY：分区运行正常且功能齐全。
    • DEGRADED：分区受损，但仍可用于形成动态切片。此状态仅适用于顶级 4x4x4 拓扑。 较小的拓扑没有降级状态。
    • UNHEALTHY：相应分区出现故障，无法用于形成切片。
    • UNSET：由于 GKE slice 控制器初始化失败，状态未定义。
    • INCOMPLETE：分区中的并非所有节点都已预配。

创建 Slice 自定义资源

切片自定义资源会略有不同，具体取决于您创建的是动态超级切片还是动态子切片。

  1. 定义 Slice 自定义资源：

    apiVersion: accelerator.gke.io/v1beta1
kind: Slice
metadata:
  # Name of the slice resource
  name: SLICE_NAME
spec:
  # Specify the type of accelerator for this slice
  type: "tpu7x"
  # Define the desired topology for the accelerator slice
  topology: TOPOLOGY
  partitionIds:
    - PARTITION_ID # Example: a9476d1b02bd4f4e75ffffae3bd23c01
    - PARTITION_ID_2
    # ... add more partition IDs as needed

    替换以下内容：

    • SLICE_NAME：切片的名称。名称必须满足 metadata.name 条件，且不得超过 49 个字符。
    • TOPOLOGY：动态切片的拓扑。 拓扑必须满足以下条件：
      • 对于动态子切分：您可以指定小于 4x4x4 的拓扑，例如 2x2x1、2x2x2、2x2x4 或 2x4x4。 这些较小的拓扑需要 GKE 1.36.0-gke.3712000 版或更高版本。
      • 对于动态超级切片：您可以指定等于或大于 4x4x4 的拓扑。如需配置动态超切片，所请求拓扑的每个维度都必须是 4 的倍数，例如 4A x 4B x 4C。拓扑维度 AxBxC 中的三个值必须按非递减顺序排列 (A ≤ B ≤ C)。例如，4x4x8 有效，但 4x8x4 无效。此顺序有助于确保切片形成一致，并避免出现意外行为。拓扑维度中三个值的乘积（即 A × B × C）不得超过 9,216。
    • PARTITION_ID：一个字符串列表，用于标识构成切片的分区。
      • 对于动态子切片：您必须指定一个分区 ID。
      • 对于动态超级切分：您必须根据芯片总数计算分区数，其中每个分区包含 64 个芯片。spec.partitionIds 列表中的商品数量必须与计算出的分区数量 ((A × B × C) / 64) 完全一致。
      • partitionIds 列表必须满足以下条件：
        • 每个分区都必须映射到预留子块。
        • 所有关联的子块都必须属于同一预留。
        • 所有关联的子块都必须位于同一预留中。
        • 关联的节点池中的所有节点都必须处于 ready 状态。
    • type 字段的值必须为 tpu7x。
    • （可选）如需使切片控制器在切片形成期间自动重试，您可以向切片自定义资源添加 slice.gke.io/retry-on-failure: "true" 注解。如果由于 SliceCreationFailed 状态原因而未创建 slice，控制器将重试，直到成功形成 slice。

    例如，如需创建 4x8x8 切片（动态超级切片），您需要提供四个唯一的分区 ID。

    apiVersion: accelerator.gke.io/v1beta1
kind: Slice
metadata:
    name: test-super-slice-example
    annotations:
      slice.gke.io/retry-on-failure: "true"
spec:
    type: "tpu7x"
    topology: "4x8x8" # (4*8*8)/64 = 4 partitions
    partitionIds:
        - "p0-4x4x4"
        - "p1-4x4x4"
        - "p2-4x4x4"
        - "p3-4x4x4"

    例如，如需创建 2x2x2 切片（动态子切片），您需要提供一个唯一分区 ID。

    apiVersion: accelerator.gke.io/v1beta1
kind: Slice
metadata:
  name: test-sub-slice-example
  annotations:
      slice.gke.io/retry-on-failure: "true"
spec:
  type: "tpu7x"
  topology: "2x2x2"
  partitionIds:
    - "fba785f80d18552357dcdef6d3d16c27" # Only 1 partitionId for sub-slice

  2. 应用 Slice 自定义资源：

    kubectl apply -f test-slice-example.yaml

    此时，GKE 会尝试创建切片。如果出现以下问题之一，切片创建会失败，并且 Slice 自定义资源中的状态原因会更新为 SliceCreationFailed 或 FAILED：

    • 如果自定义资源上所选的节点不存在，则状态原因是 SliceCreationFailed。
    • 如果自定义资源上的任何节点被其他 slice 使用，则状态原因是 SliceCreationFailed。
    • 如果自定义资源上的节点不属于同一预留子块，则状态原因是 FAILED。
    • 如果节点不在同一预留中，则状态原因是 FAILED。
    • 如果拓扑与分区数量不匹配，则状态原因是 SliceCreationFailed。

    如需在状态为 SliceCreationFailed 时自动重试切片创建，请按创建 Slice 自定义资源中所述配置 slice.gke.io/retry-on-failure: "true" 注解。

    如需详细了解 Slice 自定义资源的状态，请参阅 Slice 状态。

监控 Slice 自定义资源的状态

如需检查 Slice 自定义资源的状态，请运行以下命令：

kubectl describe slice SLICE_NAME

将 SLICE_NAME 替换为切片的名称。

输出类似于以下内容：

Name:         test-slice
Namespace:
Labels:       <none>
Annotations:  <none>
API Version:  accelerator.gke.io/v1beta1
Kind:         Slice
Metadata:
  Creation Timestamp:  2026-01-11T23:45:15Z
  Finalizers:
    accelerator.gke.io/slice-finalizer
  Generation:        1
  Resource Version:  1768175347356335006
  UID:               d0b71e5c-be3f-4788-aead-930c7afec4f2
Spec:
  Partition Ids:
    2c79463990ff67c4e3c2648666bfedfa
    ba898ffcac0ad0946e8ff036d771ee53
    [more partition IDs]
  Topology:  8x16x16
  Type:      tpu7x
Status:
  Conditions:
    Last Transition Time:  2026-01-11T23:45:38Z
    Message:               ""
    
    Reason:                FAILED
    
    Status:                False
    Type:                  Ready
Events:

Slice 自定义资源状态中的 reason 字段表示当前生命周期状态。可能的状态因您使用的是动态子切片还是动态超级切片而异。

动态子切片的状况条件

  • SliceNotCreated：控制器执行初始化和资源检查。
    • 如果不满足前提条件，状态会转换为 SliceCreationFailed。
    • 如果验证通过，状态会转换为 ACTIVATING。
  • ACTIVATING：GKE 正在形成切片。
    • 如果成功，状态将转换为 ACTIVE。
    • 如果子块已降级，但切片仍可用，则状态会转换为 ACTIVE_DEGRADED。
    • 如果组建失败，状态会转换为 FAILED。
  • DEACTIVATING：如果 Slice 自定义资源被删除，或者在有效或失败状态下发生严重故障，Slice 会开始拆解。
  • INCOMPLETE：资源完全删除之前的最后一步。

动态超级切片的启用条件

  • SliceNotCreated：切片尚未创建。切片控制器正在初始化并执行切片形成前的预检。
  • SliceCreationFailed：创建失败，因为未满足前提条件（例如，缺少必需的 Compute Engine 资源）或预检失败。如需在此状态下自动重试切片形成，请配置 slice.gke.io/retry-on-failure: "true" 注解。
  • ACTIVATING：切片正在形成（拼接）。
  • ACTIVE：切片已完全形成、健康且已准备好执行工作负载。
  • ACTIVE_DEGRADED：切片已形成，但包含降级的立方体，受 ICI 恢复能力支持。工作负载可以运行，但性能可能会受到影响。动态超级切片可应对单个光路交换机 (OCS) 故障。如果单个 OCS 单元发生故障，则通过该交换机的所有光纤链路都将变得不可用，从而导致超级 pod 中的所有立方体以降级状态运行。
  • DEACTIVATING：切片正在拆解。
  • FAILED：切片不再准备好执行工作负载。如果初始形成失败或有效切片遇到严重的软件或硬件故障，就会出现此状态。
  • INCOMPLETE：没有足够的可用立方体来启动超切片形成。

如需详细了解 Slice 自定义资源的状态，请参阅 Slice 状态。

在动态切片上运行工作负载

当 Slice 自定义资源处于 ACTIVE 状态时，您可以在其上运行工作负载。以下部分包含使用动态切片的工作负载示例。工作负载以作业或 JobSet 的形式提交。

示例 1：单个工作负载使用单个切片

以下示例展示了使用单个 4x4x4 动态超级切片的工作负载。

  1. 将以下示例清单保存为 tpu-job-jax-v7x-64.yaml：

    # Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job-jax-v7x-64
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job-jax-v7x-64
spec:
  backoffLimit: 0
  completions: 16
  parallelism: 16
  completionMode: Indexed
  template:
    metadata:
      annotations:
        cloud.google.com/gke-tpu-slice-topology: 4x4x4
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 4x4x4
        cloud.google.com/gke-tpu-accelerator: tpu7x
        cloud.google.com/gke-tpu-slice: test-slice
      subdomain: headless-svc
      restartPolicy: Never
      containers:
      - name: tpu-job-jax
        env:
        - name: TPU_ACCELERATOR_TYPE
          value: tpu7x-128
        image: python:3.12
        securityContext:
          privileged: false
        command:
        - bash
        - -c
        - |
          set -ex
          pip install -U --pre jax jaxlib libtpu requests -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          pip list
          python -c 'import jax; print("Total TPU devices (cores):", jax.device_count())'
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4

    在此清单中：

    • cloud.google.com/gke-tpu-slice-topology 和 cloud.google.com/gke-tpu-topology 定义了动态切片的拓扑。
    • env.value: tpu7x-128 是 TPU 加速器类型和切片中的总核心数。核心数量的计算方式是将拓扑的维度乘以每个芯片的核心数量。 例如，对于 4x4x4 拓扑，计算公式为 4 × 4 × 4 × 2 = 128，其中 2 是 tpu7x（Ironwood (TPU7x)）的每芯片核心数。 因此，TPU_ACCELERATOR_TYPE 为 tpu7x-128。

  2. 应用 tpu-job-jax-v7x-64.yaml 清单：

    kubectl apply -f tpu-job-jax-v7x-64.yaml

示例 2：使用 JobSet 在多切片节点池上部署工作负载

此示例演示了如何使用 JobSet 在多切片节点池上部署工作负载。

  1. 安装 JobSet：

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/JOBSET_VERSION/manifests.yaml

    将 JOBSET_VERSION 替换为所需的 JobSet 版本。如需使用动态子切片，请使用 JobSet v0.12.0 或更高版本。对于动态超级切片，请使用 JobSet v0.11.1 或更高版本。

  2. 将以下示例清单保存为 tpu-multislice-jax.yaml：

    # Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: tpu-multislice-jax
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-tpu-slice
spec:
  failurePolicy:
    maxRestarts: 3
  replicatedJobs:
  - name: slice-job
    replicas: 2
    template:
      spec:
        parallelism: 16
        completions: 16
        backoffLimit: 0
        completionMode: Indexed
        template:
          metadata:
            annotations:
              # The shape of the slice
              cloud.google.com/gke-tpu-slice-topology: 4x4x4
          spec:
            hostNetwork: true
            dnsPolicy: ClusterFirstWithHostNet
            nodeSelector:
              cloud.google.com/gke-tpu-topology: 4x4x4
              cloud.google.com/gke-tpu-accelerator: tpu7x
              # IMPORTANT: Do NOT put 'cloud.google.com/gke-tpu-slice' here manually.
              # The exclusive-topology annotation handles the slice assignment automatically.
            containers:
            - name: jax-worker
              image: python:3.12
              securityContext:
                privileged: true
              ports:
              - containerPort: 8471
              command:
              - bash
              - -c
              - |
                set -ex
                pip install -U --pre jax jaxlib libtpu requests -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                # Verify JobSet injected the specific slice ID for this worker
                echo "JobSet Index: $JOB_COMPLETION_INDEX"
                python -c 'import jax; print("Total TPU devices:", jax.device_count())'
              resources:
                requests:
                  google.com/tpu: 4
                limits:
                  google.com/tpu: 4

  3. 应用 tpu-multislice-jax.yaml 清单：

    kubectl apply -f tpu-multislice-jax.yaml

    在此清单中：

    • replicatedJobs 下的 replicas: 2 字段表示 JobSet 创建两个单独的作业，每个作业对应一个 4x4x4 TPU 切片。
    • alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-tpu-slice 注解有助于确保每个作业都分配给唯一的 TPU 切片。
    • cloud.google.com/gke-tpu-slice-topology: 4x4x4 注解定义了每个动态切片的拓扑。
    • 在此示例中，未明确设置 TPU_ACCELERATOR_TYPE 环境变量，因为 JobSet 会处理切片分配。JAX 代码会自动检测其分配的切片中的可用 TPU 设备。

删除切片

  1. 删除 slice：

    kubectl patch slice $SLICE_NAME --type json \
  -p='[{"op": "remove", "path": "/metadata/finalizers"}]'

  2. 验证切片是否已删除：

    kubectl get slices

升级节点池

如果您升级配置了增量预配的节点池，则必须使用特定的超额配置参数，以防止出现容量冲突。

如需配置并运行节点池升级，请执行以下操作：

  1. 如果您使用的是动态超级切片，请在执行手动维护之前删除有效切片：

    kubectl delete slice SLICE_NAME

    如果您使用的是动态子切片，则无需先删除有效切片。不过，在升级期间保持活跃状态的子切片会失败并转换为 FAILED 状态。

  2. 更新节点池的升级参数：

    gcloud container node-pools update NODE_POOL_NAME \
    --cluster=CLUSTER_NAME \
    --project=PROJECT_ID \
    --location=LOCATION \
    --max-surge-upgrade=0 \
    --max-unavailable-upgrade=16

    确保 --max-surge-upgrade 字段设置为值 0，以防止 GKE 在升级期间尝试分配额外的 TPU。我们建议您将 --max-unavailable-upgrade=16 字段设置为同时升级完整的 16 节点子块。

  3. 升级节点池：

    gcloud container clusters upgrade CLUSTER_NAME \
    --node-pool=NODE_POOL_NAME \
    --cluster-version=CLUSTER_VERSION \
    --project=PROJECT_ID \
    --location=LOCATION

硬件维护和故障

如果您触发客户发起的维护或发生硬件故障切换，则只会影响目标节点，而不会影响整个节点池。

对于动态超级切片，您必须先删除有效的动态切片，然后才能执行手动维护。如果您使用的是动态子切片，则无需先删除有效切片。如果动态子切片处于有效状态时发生维护或硬件故障，系统会按如下方式处理恢复：

  1. 自动重塑：当关联的主机开始维护时，GKE 会自动重塑有效动态切片。
  2. 切片自定义资源失败：Slice 自定义资源转换为 FAILED 状态。
  3. 调度器观测：调度器观测到 Slice 自定义资源的失败状态。
  4. 重组流程：调度程序会自动尝试在其他可用的健康节点上重新创建 slice 配置。

停用 Slice 控制器

如需停用 slice 控制器，请将其从集群中移除。

  1. 检查 Slice 自定义资源是否为空：

    kubectl get slice -A

  2. 更新集群以停用切片控制器：

    gcloud container clusters update ${CLUSTER_NAME} \
    --location=${REGION} \
    --no-enable-slice-controller

  3. 删除 Slice 自定义资源：

    kubectl delete crd slices.accelerator.gke.io

  4. 验证 Slice 自定义资源是否已删除：

    kubectl get crd | grep slices.accelerator.gke.io

  5. 移除切片控制器添加的标签。需要移除以下标签：

    • cloud.google.com/gke-tpu-slice
    • cloud.google.com/gke-tpu-topology
    1. 如需从特定节点中移除，请更新节点名称
    export NODE_NAME="gke-tpu-bdac9600-3bdg"
kubectl label node $NODE_NAME cloud.google.com/gke-tpu-slice- cloud.google.com/gke-tpu-slice-topology-
    1. 如果您想从集群中的每个节点移除这些标签，请执行以下操作：
    kubectl label nodes --all cloud.google.com/gke-tpu-slice- cloud.google.com/gke-tpu-slice-topology-
    1. 检查节点标签，确认它们为空：
    export NODE_NAME="gke-tpu-bdac9600-3bdg"
kubectl describe node $NODE_NAME | grep "cloud.google.com/gke-tpu-slice"

后续步骤