Skip to content

RFC-0033: AIStudio 训练平台接入 GKE TPU——排队、配额、抢占与 dynamic slicing 集成

概述

本 RFC 描述 AIStudio 训练平台接入 GKE、使用 TPU(All Capacity / dynamic slicing 模式)的整体方案。训练作业以 JobSet 形式提交,单机任务用 Job、常驻服务用 Deployment,三者由 Kueue 统一接管;集群内由上游 Kueue 承担作业级准入——把 AIStudio 现有的「高保/低保」额度模型精确翻译为 nominal/borrowing 配额与优先级策略,完成配额记账、gang 准入和「高保回收低保借用」的自动抢占(驱逐后自动重排续跑);slice 的放置由 Kueue 内置的 TAS(拓扑感知调度)在准入时装箱决策,slice 的动态成型与激活由 GKE 提供的 Kueue Slice Provisioner 经 AdmissionCheck 完成——排队、配额、抢占、放置、供给串在同一条准入链上。大任务(≥ 一个 cube 的 super-slice,已 GA)与小任务(小于一个 cube 的 sub-slice,private preview)在同一套方案内统一支持,共用同一本配额账。AIStudio 保留提交入口、团队内排队与额度管理职责。

背景与动机

AIStudio 计划将训练负载接入 GKE 的 TPU 集群(All Capacity / dynamic slicing 模式)。它当前的资源模型:

  1. 各团队按集群 TPU 总量分配额度,分高保低保两档。高保额度之和 ≤ 集群总量;高保 + 低保总额度可略超集群总量。
  2. 每个团队有一个平台内部队列,用户作业先在团队队列排队,再由平台提交 JobSet 到 K8s 集群。
  3. 平台只做高保/低保标记,作业提交到集群后不再做任何容量仲裁——集群内不存在抢占或回收机制。

这套模式存在几个结构性问题:

  • 高保保证无法兑现:当集群被低保作业占满时,没有任何组件为新提交的高保作业腾出容量,「高保」只是一个标记,兑现依赖人工协调。
  • 没有 gang 语义:多机 TPU JobSet 的 pod 可能只起来一半,剩余 pod 卡在调度,昂贵资源被半个任务空占,甚至多个任务互相死锁。
  • 账本漂移:平台侧按总量记账,与集群真实状态之间存在时间差和失败路径(提交成功但 pod 起不来、任务退出但账未扣)。
  • 与 GKE 动态切片方案脱节:All Capacity 模式下,GKE 官方的动态切片集成以 Kueue 为核心(TAS 负责分区选择,Kueue Slice Provisioner 按准入动态成型 slice);不引入 Kueue,就需要部署并维护另一套独立组件栈,且排队与配额问题依旧无解。

Kueue 的核心能力(排队、cohort 借用、优先级、gang 准入、reclaimWithinCohort 抢占、驱逐自动重排、TAS 拓扑装箱)与「高保/低保」额度模型及动态切片需求精确对应,恰好补上集群内缺失的这层容量仲裁。

目标

  • 团队高保/低保额度以声明式的 ClusterQueue 配置落地,配额记账在集群内权威完成;可通过对账验证与 AIStudio 额度配置一致。
  • 高保作业提交后若容量不足,Kueue 自动驱逐低保借用者;被驱逐作业自动回队、容量空出后自动续跑,全程无需人工介入。
  • 跨队回收按公平份额进行:开启 Fair Sharing,让超出公平份额最多的团队先归还借用,而不是按准入时间随机殃及借得少的团队。
  • JobSet 准入具备 gang 语义:整个 JobSet 的资源要么全部预留、要么整体等待,不再出现半启动状态。
  • TPU slice 按需动态供给且由准入驱动:TAS 在准入时完成拓扑装箱(选定分区与主机),slice 只为已获配额的作业成型,不为排队作业空转硬件。
  • 大小任务统一:super-slice(≥ 4x4x4,GA)与 sub-slice(2x2x1 / 2x2x2 / 2x2x4 / 2x4x4,private preview)共用同一套队列、配额、优先级与抢占方案;小任务优先填充碎片 cube,完整 cube 留给大任务(TAS 最小拓扑装箱),提升整体利用率。
  • AIStudio 与 Kueue 的交互契约(创建/查询/删除/抢占历史)清晰定义并落地,用户在 AIStudio 界面可以看到作业排队位置、被抢占原因与重排状态。
  • 抢占与容量仲裁职责单一归属 Kueue,平台侧不引入任何自研的仲裁逻辑。

验收标准:设计细节 4.4 节的六个抢占场景在 staging 集群逐一复现且行为符合预期;生产灰度期平台账本与 Kueue 账本对账偏差为零。

非目标

  • 「被抢占直接失败退出、不自动重拉」的语义本期不做。本期所有被抢占作业一律自动重排续跑;失败退出选项(基于 Workload.spec.active=false 的外部控制)留待后续迭代,本方案已为其预留数据基础(抢占历史落库)。
  • 不做 MultiKueue 多集群调度;多集群路由仍由 AIStudio 负责。
  • Pathways 运行时的接入设计本期不定稿(训练侧方案见 RFC-0034),待其确定后另行补充。本方案的 ClusterQueue 结构已为混合 CPU/TPU 作业预留(cpu/memory 资源组 + onFlavors,见 4.2);届时的改动面预估见 4.3.6。
  • 不替代 AIStudio 的团队内排队体验、额度审批与用户界面。
  • 不改动 K8s 原生 PriorityClass 体系(Kueue 的 WorkloadPriorityClass 与之无关,不引发 kubelet 层 pod 抢占)。

方案

总体职责划分:

text
AIStudio(平台)             : 团队内排队、额度管理、提交 JobSet(打两个标签)、状态展示、抢占历史落库
Kueue(集群内)              : 作业级准入——配额记账、cohort 借用、gang、优先级、抢占、驱逐重排
Kueue TAS(Kueue 内置)      : 拓扑装箱——按分区容量树选定分区与主机,并把 pod 钉扎到选定节点
Kueue Slice Provisioner(GKE 发布): slice 登记与门闩——翻译拓扑注解、随准入创建/删除 Slice CR(AdmissionCheck)
托管 Slice Controller(GKE 控制面): slice 施工——布线/激活/拆除,维护分区健康标签
kube-scheduler               : pod 到节点的形式绑定(节点已由 TAS 钉扎)

概念映射(平台模型 → Kueue 对象):

AIStudio 概念Kueue 对应物
一个团队一个 ClusterQueue(+ 团队命名空间中的 LocalQueue)
高保额度(Σ ≤ 集群总量)nominalQuota(所有队列挂同一 Cohort)
低保额度(总和可略超集群总量)borrowingLimit(只能借用池内实际闲置的量,不超卖硬件)
作业的高保/低保标记WorkloadPriorityClass(guaranteed=1000 / best-effort=100)
「高保提交后超量则牺牲低保」(现状缺失的能力)preemption.reclaimWithinCohort: LowerPriority:自动驱逐低优借用者——驱逐(suspend)+ 自动重排,不删除作业

关键新增能力:抢占是本方案为集群引入的、现状不存在的行为。其执行方式是 Kueue 把受害 JobSet 重新置为 suspend=true——pod 收到 SIGTERM 后被删,Workload 回到队列,配额空出后自动重新准入、从 checkpoint 续跑。作业对象全程保留,不发生删除。

Kueue 内部流转:一个 JobSet 从提交到运行

理解 Kueue 能力的最快方式是跟着一个 JobSet 走一遍它的内部流程:

几个值得注意的机制点:

  • pod 在准入前根本不存在:配额检查发生在 Workload(资源摘要)上,不放行就没有任何 pod 占用集群——排队不空耗资源的原因就在这里。
  • gang 发生在第 ⑤ 步:QuotaReserved 对整个 JobSet 的所有 podSets 原子生效,不存在「先准入一半」的中间态。
  • 配额与放置一次决策:TAS 把「配额够不够」与「拓扑放不放得下」合并在第 ④ 步的同一次装箱里判定,不存在「账上有量却切不出连续拓扑」的空头准入(残余场景见风险 3)。
  • 队列不是死板 FIFO:默认 BestEffortFIFO——队头大作业暂时放不下时,不阻塞后面放得下的小作业(需要严格顺序可切 StrictFIFO)。
  • Workload 条件时间线:QuotaReserved →(Slice 检查 Ready)→ Admitted →(Evicted → Requeued → Admitted)* → Finished,4.3.2 的 AIStudio 状态机就建立在这条时间线上。
  • Kueue 的职责到「钉扎」为止:TAS 决定 pod 精确落在哪些主机,kube-scheduler 只做形式绑定;slice 的物理布线与激活由托管 Slice Controller 完成,Kueue 不碰硬件。

用户故事

  • 低保用户:提交后若集群有闲置立即开跑;被高保抢占时,AIStudio 界面显示「已被抢占(团队 X 的高保作业回收借用额度),已自动重新排队,当前第 N 位」;容量空出后自动从 checkpoint 续跑,无需重新提交。
  • 高保用户:提交后即便集群被低保占满,系统也会自动腾出容量,作业在分钟级进入调度,且保证额度内必然放得下(Σ高保 ≤ 总量 + 借用者可被回收)。
  • 团队管理员:通过监控面板查看本队 nominal 用量、借用量、被抢次数与等待时长,验证额度划分是否合理。
  • 平台开发:无需自研集群内的配额仲裁与抢占控制器,只维护一个读 Workload 状态的 watcher。

备选方案

方案否决理由
维持现状(集群内无仲裁)高保保证无法兑现、无 gang 语义、账本漂移,且无法使用 GKE 动态切片的官方 Kueue 集成
Volcano / YuniKorn 等批调度器队列与抢占能力相当,但 GKE 动态切片的官方集成(TAS 拓扑装箱 + Slice AdmissionCheck)只面向 Kueue;JobSet 一等公民支持与 K8s SIG 上游维护也以 Kueue 最成熟
GKE 独立组件栈(partition-tree + slice-provisioner + slice-scheduler)GKE 为不使用 Kueue 的用户提供的替代路线:多维护三个组件(含 GCP 凭据),放置决策落在 Kueue 之外的独立调度器上,且只解决 slice 供给——排队/配额/抢占仍需自研
平台自研集群内准入控制器等价于重写一个 Kueue 子集,长期维护成本远高于采用上游

约束与注意事项

  • 一本账 vs 两本账:Kueue 每队只有一本账(nominal + borrowing),不区分「高保用量/低保用量」。额度层面的高低保体现为 nominal vs borrowing,作业层面体现为优先级,实际效果与平台现有语义等价。
  • 借用中作业的抢占资格随 Fair Sharing 变化:经典语义下,需要借用的作业从不触发抢占(不配置 borrowWithinCohort 时的默认行为);开启 Fair Sharing 后此门槛被移除——借用中的作业也可按 DRS 份额策略抢占其他借用队列中严格更低优的作业(任何队列 nominal 内的作业仍绝对安全)。在两档优先级模型下的实际效果:低保借用者仍从不引发牺牲(不存在比低保更低优的受害者);高保作业用满本队 nominal 后并未完全失去特权,份额策略满足时仍可回收他队低保的借用。AIStudio 提交高保作业前仍应校验本队高保余额——nominal 内准入是无条件保证,借用则依赖份额评估。
  • PodDisruptionBudget 挡不住 Kueue 抢占(走 suspend 路径,不经 Eviction API),不要指望用户配 PDB 自保。
  • 低保作业将面临新的行为变化:现状下低保作业提交后从不会被系统中断;引入抢占后可能被驱逐重启,必须满足 checkpoint + SIGTERM 契约(见 4.5)。这是本方案对用户侧最大的语义变化,需提前宣导(见风险 2)。
  • 抢占的受害者排序为:优先级更低者先、同优先级中最近准入者先(LIFO,保护跑得久的作业);且只驱逐「刚好放下抢占者」的最小集合。
  • slice 拓扑取值分两档,AIStudio 提交前须按白名单校验注解取值:super-slice(GA):AxBxC 非降序、每维为 4 的倍数、最小 4x4x4、积 ≤ 9216、最多 32 个子块;sub-slice(private preview):仅 2x2x1 / 2x2x2 / 2x2x4 / 2x4x4。sub-slicing 的启用有三道前置手续(见实施计划阶段〇):项目/预留/子块允许名单、Google 侧软件包 rollout、CTM 维护。TAS 只把具备 Topology 声明的全部层级标签的节点纳入容量树(缺任一层即整个节点被排除),因此 7 层 Topology 必须在五档分区标签齐备之后才能下发——在此之前集群运行 4.6 的静态形态(2 层 Topology),不存在中间过渡态。

设计细节

4.1 组件清单与安装配置

组件来源/版本部署与配置要点
JobSet controllerkubernetes-sigs/jobset,≥ v0.12.0(SubSlicing 指南基线)官方 manifest / helm 安装,默认配置即可、无需特殊开关;确认版本与 Kueue 兼容矩阵匹配
Kueuekubernetes-sigs/kueue,v0.19.2(SubSlicing 指南基线为 ≥ v0.18.2;API v1beta2)官方 manifest / helm 安装,见下方 Configuration;controller 多副本 HA(GKE 指引示例为 3 副本,leader election 默认开启);证书用内置轮换,无需 cert-manager
托管 Slice ControllerGKE 控制面托管gcloud container clusters update <cluster> --enable-slice-controller 开启(安装 slices.accelerator.gke.io CRD,用 kubectl get crd 验证)。物理执行者:消费 Slice CR、ICI 布线/激活/拆除(状态机 ACTIVATING → ACTIVE → DEACTIVATING/FAILED)、把各几何分区健康写为节点 -state 标签。无需也无法自行部署
Kueue Slice ProvisionerGKE 发布,安装清单随 SubSlicing 指南提供(同一份清单覆盖 super + sub)kubectl apply 部署到 slice-controller-system(注意命名陷阱:deployment 名为 slice-controller-controller-manager,与上面的托管组件是两个东西)。Kueue 外部 AdmissionCheck 控制器(controllerName: accelerator.gke.io/slice):webhook 拦截 JobSet 创建、把 gke-tpu-slice-topology 注解翻译成 TAS 约束并注入分区健康亲和;随准入按 TAS 选块结果创建 Slice CR、随驱逐/结束删除;回填检查状态。默认 flags 即可(--activation-timeout=60m、失败重试门控已默认开启)
监控现有监控体系 + Cloud MonitoringPrometheus 抓取 kueue-controller-manager 指标(重点:kueue_cluster_queue_resource_usage 对账、kueue_preempted_workloads_total 抢占频率、kueue_cluster_queue_weighted_share 公平份额);GKE 侧用 Cloud Monitoring 的 kubernetes.io/accelerator/slice/state…/slice/formation_durations…/partition/state(slice 供给时长与分区健康的面板/告警)
AIStudio 适配层平台侧新增提交打标 + Workload watcher(状态映射、抢占历史落库),见 4.3

前提条件(非组件配置,但缺了全链路不转):Standard 集群、Rapid 渠道版本 ≥ 1.36.0-gke.3712000(sub-slicing 基线;若阶段性只跑 super 可放宽至 ≥ 1.35.2-gke.1842000);TPU7x(Ironwood)+ Container-Optimized OS 节点镜像;TPU 容量以 All Capacity 模式预留,节点池走增量供给(workload policy accelerator-topology=4x4x4 + accelerator-topology-mode=provision_only,每池 16 节点 = 一个 cube;增量供给下 cube 内有坏主机也会先建出健康节点,坏机修复后自动补齐);节点自带全部五档几何(4x4x4/2x4x4/2x2x4/2x2x2/2x2x1)的分区 -id/-state 标签——小几何标签在 sub-slicing 手续完成后出现,而 TAS 会把缺失任一层级标签的节点整个排除出容量树,故标签齐备(实施计划阶段〇)是下发 7 层 Topology 的硬前置。

注:Kueue Slice Provisioner 的安装清单与镜像由 GKE 侧发布,内容以 SubSlicing 指南为准;副本数、资源请求按集群规模调整。

Kueue 的 Configuration(通过其 ConfigMap 下发),必须明确的开关:

yaml
apiVersion: config.kueue.x-k8s.io/v1beta2
kind: Configuration
# 只接管带 queue-name 标签的作业(dev 档由默认队列自动补标),
# 未打标签的负载照常运行——生产命名空间里混部着系统组件,见下方「纳管范围」
manageJobsWithoutQueueName: false
integrations:
  frameworks:
  - jobset.x-k8s.io/jobset        # 训练主力形态
  - batch/job                     # 单机/单 Job 任务
  - deployment                    # 常驻服务;会隐式启用 pod 集成(见下方说明)
managedJobsNamespaceSelector:     # pod 类集成(pod/deployment)的作用域,必须收窄:
  matchExpressions:               # 不匹配的命名空间即使打了 queue-name 也永不被接管。
  - key: scheduling.antgroup.com/kueue-tier   # 有此标签即归 Kueue 管;
    operator: Exists                          # 具体进哪档队列由标签的值决定(见 4.2)
metrics:
  enableClusterQueueResources: true # 按队列的资源用量指标,对账与面板依赖它
fairSharing:                        # v1beta2 中出现此配置块即启用公平共享(无 enable 字段;
                                    # Kueue 严格解码配置,写未知字段会导致 controller 启动失败)
  preemptionStrategies: [LessThanOrEqualToFinalShare, LessThanInitialShare]  # 官方默认组合,显式写出
waitForPodsReady:                   # SubSlicing 指南建议启用:配合运行期 slice 故障的快速恢复
  timeout: 10m                      # 准入后到 pod 全 Ready 的时限(slice 已在准入前激活,
                                    # 此窗口主要覆盖镜像拉取与进程启动;按实际镜像大小调校)
  blockAdmission: false             # 不因个别作业未 Ready 而阻塞其他作业准入
  recoveryTimeout: 5m               # 运行期 pod 失联(如 slice 转 FAILED)后的恢复时限,
                                    # 超时驱逐重排,作为 fast replacement 的兜底

TAS 门控(TopologyAwareScheduling)自 Kueue v0.14 起为 Beta 默认开启,通过给 ResourceFlavor 配置 topologyName(见 4.2)即可激活拓扑装箱,无需显式改动。

另需开启 WorkloadPriorityClassDefaulting(v0.19.2 中 alpha、默认关),它是 4.2「免标签投递」的前提。门控不在 Configuration 里,而是 controller 的启动参数——helm 写在 controllerManager.featureGates,manifests 则给容器加 --feature-gates=WorkloadPriorityClassDefaulting=true

三个集成的 Workload 粒度不同,平台侧对接时必须区分:

集成一个 Workload 对应queue-name 标签打在gang 范围
jobset.x-k8s.io/jobset一个 JobSetJobSet 的 metadata.labels全部 replicatedJob 的全部 pod
batch/job一个 JobJob 的 metadata.labels该 Job 的 parallelism 个 pod
deployment一个 Pod(不是一个 Deployment)Deployment 的 metadata.labels(Kueue 自动传递到 pod)单个 pod,无 gang

Deployment 的每 pod 一个 Workload 是上游的既定实现(Kueue 尚不支持把 Deployment 当作单个 Workload),由此带来两个必须知晓的行为:扩容时新 pod 各自排队、逐个准入,配额不足时 Deployment 会以副本数不足的状态运行;缩容时多余 pod 被删、配额立即释放。若某个常驻服务不可接受部分副本,应为其预留足额 nominal 配额,而不是依赖借用。

deployment 集成会隐式启用 pod 集成,Kueue 的 pod webhook 因此会出现在所有 pod 的创建路径上。managedJobsNamespaceSelector 是这里的安全边界,必须显式设置(默认值只排除 kube-systemkueue-system,过宽):对 pod 类集成它是强豁免——不匹配的命名空间即使打了 queue-name 也永不被接管;对 JobSet、Job 则是弱豁免,带 queue-name 的作业在任何命名空间都会被接管,真正限定它们投递范围的是各 ClusterQueue 的 namespaceSelector(见 4.2)。

纳管范围

两个字段共同决定接管什么:managedJobsNamespaceSelector 圈定命名空间,manageJobsWithoutQueueName: false 表示其中只接管带 queue-name 标签的负载。落到两档上:

档位有名为 default 的 LocalQueue?未打 queue-name 的负载
dev(defaultfalcon-jobs)webhook 自动补标,正常排队(见 4.2「免标签投递」)
production(kubemaker)不被接管、照常运行——该命名空间混部着团队的系统组件,不应纳入配额与排队

下发方式:上述 Configuration 不是独立的 CR,而是 kueue-system 命名空间下 kueue-manager-config ConfigMap 中 controller_manager_config.yaml 这一个 data 条目。操作命令:

bash
# 路径 A(manifests 安装):安装前把配置写进清单,一次到位
wget https://github.com/kubernetes-sigs/kueue/releases/download/v0.19.2/manifests.yaml
# 编辑其中 kueue-manager-config ConfigMap 的 data."controller_manager_config.yaml" 条目,
# 填入上述配置内容,然后:
kubectl apply --server-side -f manifests.yaml

# 路径 B(helm 安装):配置写在 values 的 managerConfig.controllerManagerConfigYaml 下
helm upgrade --install kueue oci://registry.k8s.io/kueue/charts/kueue \
  --version 0.19.2 -n kueue-system --create-namespace -f kueue-values.yaml

# 已运行集群上修改配置(GitOps 场景由流水线执行等效操作)
kubectl edit configmap kueue-manager-config -n kueue-system
# ⚠️ ConfigMap 不会被热加载,改完必须重启 controller 才生效
kubectl rollout restart deployment kueue-controller-manager -n kueue-system
kubectl rollout status  deployment kueue-controller-manager -n kueue-system

# 验证:确认生效内容 + controller 正常启动(配置写错时 Pod 会 CrashLoopBackOff)
kubectl get configmap kueue-manager-config -n kueue-system \
  -o jsonpath='{.data.controller_manager_config\.yaml}'
kubectl logs deployment/kueue-controller-manager -n kueue-system --tail=50

Kueue 以严格模式解码该配置:出现未知字段(例如 v1beta1 才有、v1beta2 已移除的 fairSharing.enable)会导致 controller 启动失败而非被忽略。生产变更前先在 staging 验证 rollout 成功;配置应纳入 GitOps 仓库,与 4.2 的队列对象同源管理。

4.2 Kueue 资源配置(全局一次性 + 每团队增量)

命名空间用 scheduling.antgroup.com/kueue-tier 标签分为两档——生产档(production)与开发档(dev),两档各有独立的 ClusterQueue,互不串号(见下方「命名空间分档与 namespaceSelector」)。Kueue 对象分两类:全局共享、只建一次的,和每新增一个团队/业务线需要增量创建的。

全局共享对象(集群初始化时创建一次)——两档优先级 + 分区树 + 资源风味 + slice 准入检查:

yaml
apiVersion: kueue.x-k8s.io/v1beta2
kind: WorkloadPriorityClass
metadata:
  name: guaranteed          # 高保
value: 1000
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: WorkloadPriorityClass
metadata:
  name: best-effort         # 低保
value: 100
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: WorkloadPriorityClass
metadata:
  name: default             # 名字必须是 default——Kueue 以此识别默认优先级,
value: 10                   # 未打优先级标签的作业自动落到这一档(见「免标签投递」)
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: Topology
metadata:
  name: tpu7x-topology      # 分区树声明:TAS 按这七层节点标签构建容量树并装箱
spec:                       # 7 层同时覆盖 super-slice 与 sub-slice(SubSlicing 指南口径);
  levels:                   # 一步到位可避免日后从 3 层迁移(需替换 flavor 并全量 reconcile 存量作业)。
                            # 前置:节点五档 -id 标签已齐备(阶段〇)——TAS 排除缺任一层级标签的节点,
                            # 标签不齐时 7 层拓扑下将没有任何可用节点(连 super-slicing 也无法调度)
  - nodeLabel: cloud.google.com/gce-topology-block          # 物理 block
  - nodeLabel: cloud.google.com/gke-tpu-partition-4x4x4-id  # cube(super-slice 的原子积木)
  - nodeLabel: cloud.google.com/gke-tpu-partition-2x4x4-id  # ↓ 以下四层为 sub-slice 几何
  - nodeLabel: cloud.google.com/gke-tpu-partition-2x2x4-id
  - nodeLabel: cloud.google.com/gke-tpu-partition-2x2x2-id
  - nodeLabel: cloud.google.com/gke-tpu-partition-2x2x1-id
  - nodeLabel: kubernetes.io/hostname                       # 主机
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: ResourceFlavor
metadata:
  name: tpu7x
spec:
  nodeLabels:
    cloud.google.com/gke-tpu-accelerator: tpu7x
  topologyName: tpu7x-topology   # 挂上分区树 → 该 flavor 启用 TAS
  tolerations:                   # TPU 节点带 google.com/tpu=present:NoSchedule 污点。
  - key: google.com/tpu          # 在 flavor 上声明后,Kueue 准入时自动注入到 pod,
    operator: Exists             # 且 TAS 的节点筛选也会算上它——作业模板因此无需自己写。
    effect: NoSchedule           # ⚠️ 挂了 topologyName 的 flavor,tolerations 不可变更,
                                 #    漏配只能删除重建(需先从 ClusterQueue 解引用以释放 finalizer)
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: ResourceFlavor
metadata:
  name: cpu-standard             # 纯 CPU podSet(如 Pathways head)的兜底 flavor:
                                 # 无 topologyName(不走 TAS)、无 nodeLabels。
                                 # 不配 nodeLabels 是安全的——见下方「分流机制」:TPU podSet
                                 # 永远不会落到它上面,CPU 组件的定向由作业模板 nodeSelector 完成
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: AdmissionCheck
metadata:
  name: tpu7x-slice
spec:
  controllerName: accelerator.gke.io/slice   # 由 GKE Kueue Slice Provisioner 实现:
                                             # slice 激活为 ACTIVE 前检查不置 Ready,作业不解挂

每团队一套(新增团队时创建)——1 个 ClusterQueue(配额账本,cluster-scoped)+ 1 个 LocalQueue(投递口,位于共享命名空间,按团队命名)。以团队 A、B 为例:

yaml
# ──────────── 团队 A ────────────
apiVersion: kueue.x-k8s.io/v1beta2
kind: ClusterQueue
metadata:
  name: team-a
spec:
  cohortName: tpu-pool               # 所有团队挂同一资源池,借用只发生在池内
  namespaceSelector:                 # 只接受生产档命名空间(见下方说明);
    matchLabels:                     # 默认值 null = 不匹配任何命名空间,此字段必须显式设置
      scheduling.antgroup.com/kueue-tier: production
  resourceGroups:
  # 单一资源组:TPU 与其伴生的 cpu/memory/ephemeral-storage 必须同组(原因见下方「为什么必须单一资源组」)。
  # 组内 flavor 按顺序尝试,TPU flavor 在前、cpu-standard 兜底;
  # 每个 flavor 必须声明该组的全部资源(API 强制,且顺序与 coveredResources 一致)
  - coveredResources: ["google.com/tpu", "cpu", "memory", "ephemeral-storage"]
    flavors:
    - name: tpu7x                    # TPU podSet:四种资源全部从这一个 TAS flavor 取
      resources:
      - {name: google.com/tpu, nominalQuota: 256, borrowingLimit: 128}   # ← 高保/低保额度
      - {name: cpu, nominalQuota: "100000"}          # TPU 节点的 cpu/mem/disk 随芯片等比例配套,
      - {name: memory, nominalQuota: 100000Gi}       # 不是稀缺资源 → 形式化大配额,只记账不仲裁
      - {name: ephemeral-storage, nominalQuota: 100000Gi}
    - name: cpu-standard             # 纯 CPU podSet 兜底
      resources:
      - {name: google.com/tpu, nominalQuota: 0}      # 必须声明,配 0 表示该 flavor 不提供 TPU
      - {name: cpu, nominalQuota: "2000"}           # 纯 CPU 组件(Pathways head、
      - {name: memory, nominalQuota: 8000Gi}        # Deployment 服务)的真实约束点
      - {name: ephemeral-storage, nominalQuota: 4000Gi}
  admissionChecksStrategy:           # slice 成型门闩:Slice 激活前不解挂
    admissionChecks:                 # (v1beta2 只有 strategy 形式,无平铺 admissionChecks 字段)
    - name: tpu7x-slice
      onFlavors: [tpu7x]             # 判定为 Workload 级:只要有任一 podSet 落在 tpu7x 上,
                                     # 检查即对整个 Workload 生效(gang 完整);
                                     # 全部 podSet 都落 cpu-standard 的纯 CPU 作业不受此门闩影响
  preemption:
    reclaimWithinCohort: LowerPriority   # 跨队回收:只赶走比自己低优的借用者
    withinClusterQueue: LowerPriority    # 队内抢占:高保可赶走本队低保
    # borrowWithinCohort 不配置:需要借用的作业只等闲置、从不抢占
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: LocalQueue
metadata:
  name: team-a-queue                 # 命名约定:<团队名>-queue
  namespace: kubemaker                # 共享命名空间
spec:
  clusterQueue: team-a
---
# ──────────── 团队 B ────────────
apiVersion: kueue.x-k8s.io/v1beta2
kind: ClusterQueue
metadata:
  name: team-b
spec:
  cohortName: tpu-pool               # 与团队 A 同一个 cohort
  namespaceSelector:
    matchLabels:
      scheduling.antgroup.com/kueue-tier: production
  resourceGroups:
  - coveredResources: ["google.com/tpu", "cpu", "memory", "ephemeral-storage"]
    flavors:
    - name: tpu7x
      resources:
      - {name: google.com/tpu, nominalQuota: 256, borrowingLimit: 128}   # 示例中两队同构;
      - {name: cpu, nominalQuota: "100000"}                              # 实际数字由额度系统生成
      - {name: memory, nominalQuota: 100000Gi}
      - {name: ephemeral-storage, nominalQuota: 100000Gi}
    - name: cpu-standard
      resources:
      - {name: google.com/tpu, nominalQuota: 0}
      - {name: cpu, nominalQuota: "2000"}
      - {name: memory, nominalQuota: 8000Gi}
      - {name: ephemeral-storage, nominalQuota: 4000Gi}
  admissionChecksStrategy:
    admissionChecks:
    - name: tpu7x-slice
      onFlavors: [tpu7x]
  preemption:
    reclaimWithinCohort: LowerPriority
    withinClusterQueue: LowerPriority
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: LocalQueue
metadata:
  name: team-b-queue
  namespace: kubemaker                # 与团队 A 在同一命名空间
spec:
  clusterQueue: team-b
---
# ──────────── 开发档(所有非生产命名空间共用一个队列) ────────────
apiVersion: kueue.x-k8s.io/v1beta2
kind: ClusterQueue
metadata:
  name: dev
spec:
  cohortName: tpu-pool               # 与生产同池:可借用生产闲置容量,但随时被回收
  namespaceSelector:
    matchLabels:
      scheduling.antgroup.com/kueue-tier: dev
  resourceGroups:
  - coveredResources: ["google.com/tpu", "cpu", "memory", "ephemeral-storage"]
    flavors:
    - name: tpu7x
      resources:
      - {name: google.com/tpu, nominalQuota: 64, borrowingLimit: 256}  # 保底 1 个 cube + 借用突发
      - {name: cpu, nominalQuota: "100000"}
      - {name: memory, nominalQuota: 100000Gi}
      - {name: ephemeral-storage, nominalQuota: 100000Gi}
    - name: cpu-standard
      resources:
      - {name: google.com/tpu, nominalQuota: 0}
      - {name: cpu, nominalQuota: "500"}
      - {name: memory, nominalQuota: 2000Gi}
      - {name: ephemeral-storage, nominalQuota: 1000Gi}
  admissionChecksStrategy:
    admissionChecks:
    - name: tpu7x-slice
      onFlavors: [tpu7x]
  preemption:
    reclaimWithinCohort: LowerPriority
    withinClusterQueue: LowerPriority
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: LocalQueue
metadata:
  name: default                      # ← 名字必须是 default——Kueue 以此识别默认队列,
  namespace: default                 #    未打 queue-name 标签的作业自动投到这里
spec:                                #    但都指向同一个 dev ClusterQueue、共用一本账
  clusterQueue: dev
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: LocalQueue
metadata:
  name: default
  namespace: falcon-jobs             # 同名不同命名空间
spec:
  clusterQueue: dev

为什么必须单一资源组

Kueue 有两条不可绕开的规则:

  1. 一个资源只能属于一个资源组(clusterqueue_webhook.go 的重复校验);
  2. 一个 podSet 在每个资源组各分配一个 flavor,而声明了 TAS 的 podSet,其用到的每一个 flavor 都必须支持 TAS(tas_flavorassigner.go:172-189)。

TPU worker 除了 google.com/tpu,还会声明 TPU 之外的资源——最少也会有 ephemeral-storage(训练数据与 checkpoint 暂存),混合作业还会有 cpu/memory。只要声明了其中任意一项,若把这些资源单独切一组、只挂非 TAS 的 cpu-standard,TPU podSet 就被迫跨到一个不支持 TAS 的 flavor 上,准入直接失败:

text
couldn't assign flavors to pod set worker: Flavor "cpu-standard" does not support TopologyAwareScheduling

因此四种资源必须同组,让 TPU podSet 的全部资源都从同一个 TAS flavor 取得。作为代价,组内每个 flavor 都要声明全部四种资源(cpu-standardgoogle.com/tpu 配 0)。

组内每个 flavor 的配额值怎么填

这些条目不是可选的装饰:CRD 层的 CEL 校验强制「每个 flavor 声明的资源个数等于 coveredResources 的个数」(clusterqueue_types.go:254),webhook 进一步要求顺序也一致(clusterqueue_webhook.go:339),漏写或错序会被 API Server 直接拒绝。填值原则:

flavor资源该填什么理由
tpu7x-*google.com/tpu真实配额唯一的稀缺资源,团队额度的实际载体
tpu7x-*cpu / memory / ephemeral-storage大到不构成约束(示例用 100000)TPU 节点的 CPU/内存/盘随 TPU 机型捆绑供给,拿到 N 颗 TPU 就必然拿到配套的 CPU。此处若按实际机型填,会凭空引入第二个约束维度,出现「TPU 够但 CPU 配额不足」的假拒绝
cpu-standardgoogle.com/tpu必须为 0纯 CPU 节点池不提供 TPU;填非 0 会凭空造出不存在的 TPU 配额
cpu-standardcpu / memory / ephemeral-storage真实配额这才是 Pathways head 等纯 CPU 组件的实际约束点

分流机制:TAS 注解即路由开关

同组内两类 podSet 靠「是否声明 TAS」天然互斥,无需额外配置:

podSet 类型有 TAS 拓扑注解?遇到 tpu7x(TAS flavor)遇到 cpu-standard(非 TAS)落点
TPU podSet(McJAX worker / Pathways worker)✓ 匹配✗ 跳过(does not support TAS)tpu7x
纯 CPU podSet(Pathways head 等)✗ 跳过(supports only TAS)✓ 匹配cpu-standard

两个推论:

  • cpu-standard 不需要也不应该配 nodeLabels——它绝不会吸走 TPU podSet;CPU 组件到专用节点池的定向由作业模板自带 nodeSelector(如 cloud.google.com/gke-nodepool: <pool>)完成,这样同一个 flavor 可服务多种 CPU 组件。
  • 资源覆盖必须完整:作业请求的每一种资源都要出现在 coveredResources 中,漏掉任意一种(最常见的是 ephemeral-storage)整个 Workload 都不可准入,报 resource <name> unavailable in ClusterQueue。若将来引入 GPU 等其他扩展资源,同样要加进这一组。

命名空间分档与 namespaceSelector

namespaceSelector 默认为 null,语义是「不匹配任何命名空间」——漏写会让所有作业卡在 Workload namespace doesn't match ClusterQueue selector,而不是报配额问题,因此每个 ClusterQueue 都必须显式设置。

本方案不用内置的 kubernetes.io/metadata.name 逐个点名,而是用同一个标签键、不同的取值给命名空间分档:

bash
# 生产档:AIStudio 的正式训练任务
kubectl label namespace kubemaker   scheduling.antgroup.com/kueue-tier=production
# 开发档:自测与外部团队联调
kubectl label namespace default     scheduling.antgroup.com/kueue-tier=dev
kubectl label namespace falcon-jobs scheduling.antgroup.com/kueue-tier=dev
档位命名空间ClusterQueueLocalQueue可用优先级
productionkubemaker每团队一个(team-a / team-b / …)<team>-queue,建在 kubemakerguaranteed / best-effort
devdefaultfalcon-jobs共用一个 devdefault,每个命名空间各建一个默认档(值 10),无需打标签

这个设计的三个要点:

  • 标签键的存在性与取值分工:managedJobsNamespaceSelector(4.1)用 Exists 判定「归不归 Kueue 管」,各 ClusterQueue 的 namespaceSelector 用取值判定「进哪一档」。新增命名空间只需打标签,队列定义不动。
  • 两档同属 tpu-pool cohort,开发档因此能借用生产闲置的容量,不必为自测预留一块常年空转的硬件。安全性由优先级保证:dev(10)低于一切生产作业,生产回收自己 nominal 内的容量走 InCohortReclamation,无条件越过 Fair Sharing 的份额比较,不会被开发档拖住。若某阶段需要硬隔离,把 dev 队列的 cohortName 改成独立值即可。
  • 开发档必须留有 nominal 配额(示例给 64 芯 = 1 个 cube)。配 0 的话开发档会被生产完全饿死,4.4 的六个抢占场景就无法稳定复现。

两条操作约束:

  • 打标签必须先于创建/收紧 ClusterQueue。反过来的话,中间窗口内所有新作业都会被拒(已准入的不受影响)。
  • LocalQueue 必须建在目标命名空间内(Workload 与 LocalQueue 同命名空间)。开发档的两个命名空间因此各需要一个 LocalQueue,即便它们指向同一个 ClusterQueue。

标签键用自有域名:kubernetes.io/k8s.io/ 是 Kubernetes 保留命名空间,kueue.x-k8s.io/managed 已被上游占用(pkg/constants/constants.go:52,用于标记归 Kueue 管的 Pod),两者都不可借用。

免标签投递(开发档)

开发档面向的是已有存量作业的团队,要求他们逐个作业加标签成本高、容易漏。Kueue 提供了两处「名为 default 即为默认」的约定,组合起来可以让存量作业一个字段都不用改:

对象命名约定效果前提
LocalQueue命名空间内名为 default未打 queue-name 的作业自动投到该队列无(默认行为)
WorkloadPriorityClass集群内名为 default未打 priority-class 的作业自动获得该优先级开启 WorkloadPriorityClassDefaulting 门控(见 4.1)

两者都由 webhook 在创建时补打标签,对 JobSet、Job、Deployment 均生效。

默认优先级必须真的建一个 WorkloadPriorityClass,不能依赖「不打标签就是 0」:未打标签时优先级按四层回退取值(reconciler.go:1783-1797),第二层取的是 pod spec 的 priorityClassName。集群里常有高值的 K8s PriorityClass(如 CI 用的 ci-tpu-pool = 1,000,000),pod 一旦带上它,该作业的 Kueue 优先级就比 guaranteed(1000)高三个数量级、反过来抢占生产,且无任何报错。名为 default 的 WorkloadPriorityClass 走第一层,压过该路径。

生产档不使用该机制:kubemaker 下不建名为 default 的 LocalQueue,AIStudio 始终显式指定队列与优先级。

AIStudio 新增团队/业务线时的 K8s 资源清单:在 GitOps 仓库追加 1 个 ClusterQueue + 1 个 LocalQueue(命名约定:<team> / <team>-queue),全局对象无需变动。额度调整 = 修改该团队 ClusterQueue 的 nominalQuota/borrowingLimit 字段。

各队参与 Fair Sharing 的权重(spec.fairSharing.weight)本期不设置、按默认值 1 均分——即所有团队对池内闲置容量的「公平份额」相等;后续如需让大团队占更大闲置份额,可按额度比例设置权重。

注意:生产档内命名空间与团队并非一一对应——多个团队共用 kubemaker。此时团队边界完全由作业上的 kueue.x-k8s.io/queue-name 标签决定(没有命名空间级隔离),AIStudio 提交时必须保证该标签与作业所属团队一致。若希望获得命名空间级隔离,可给每个团队单独建命名空间、打上 kueue-tier: production 标签、并只在其中创建该团队的 LocalQueue。配额下发流水线必须校验 Σ nominalQuota(含 dev 队列)≤ 集群 TPU 总量(见风险 1)。

4.3 AIStudio 交互契约

AIStudio 与本方案的全部交互面 = 业务对象(创建/删除)+ Workload(只读)+ 事件(只读)。业务对象按形态为 JobSet、Job 或 Deployment,三者的对接方式完全一致:创建时打两个标签,状态一律从其 Workload 影子对象读取。核心原则:

AIStudio 只创建和删除业务对象;Workload 是 Kueue 的内部对象,只读、永不修改、永不删除

下文以 JobSet 为例展开(训练主力形态);Job 与 Deployment 的差异集中在 4.3.1 末尾。

4.3.1 创建作业

在现有 JobSet 之上仅需增加两个标签(高保/低保 + 目标队列),其余字段与现状一致。下例为 McJAX 形态(全部 podSet 都用 TPU);Pathways 形态多一个纯 CPU 的 head podSet,写法见本节末:

yaml
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: llm-pretrain-a3f8            # AIStudio 生成,全局唯一,禁止复用近期用过的名字
  namespace: kubemaker                # 所有团队共享的命名空间
  labels:
    kueue.x-k8s.io/queue-name: team-a-queue     # ① 进团队 A 的队列(团队边界由此标签决定)
    kueue.x-k8s.io/priority-class: best-effort  # ② 低保(高保填 guaranteed)
spec:
  failurePolicy:
    maxRestarts: 3        # 只管真实故障(进程崩溃/硬件失败)的重启预算;
                          # Kueue 抢占走 suspend 路径,不消耗此计数(见 4.5)
  replicatedJobs:
  - name: workers
    replicas: 2           # 2 个 slice
    template:
      spec:
        backoffLimit: 0
        completions: 16   # 每 slice 16 台主机(4x4x4 = 64 芯)
        parallelism: 16
        template:
          metadata:
            annotations:
              # 每个 replicatedJob 副本动态成型一个该拓扑的 slice。取值白名单:
              #   super(GA):AxBxC 非降序、每维 4 的倍数、≥4x4x4、积 ≤9216
              #   sub(preview):2x2x1 / 2x2x2 / 2x2x4 / 2x4x4
              cloud.google.com/gke-tpu-slice-topology: 4x4x4
          spec:
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu7x
            # TPU 污点的 toleration 已在 ResourceFlavor 上统一声明并由 Kueue 注入(见 4.2),
            # 此处无需重复;若作业自带也兼容
            containers:
            - name: jax
              image: python:3.12
              command:                  # 冒烟用例:装 JAX 并打印本 slice 认到的芯片数
              - bash                    # 真实训练把这三行换成训练镜像与启动命令即可
              - -c
              - |
                pip install -q "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                python -c 'import jax; print("global:", jax.device_count(), "local:", jax.local_device_count())'
              resources:
                limits:
                  google.com/tpu: '4'          # tpu7x 每主机 4 芯
                requests:
                  ephemeral-storage: 20Gi      # 与 TPU 同属一个资源组,从同一个 flavor
                  google.com/tpu: '4'          # 取额度(见 4.2「为什么必须单一资源组」)
              # 契约:训练进程须处理 SIGTERM(宽限期内写 checkpoint),
              # 启动时从最新 checkpoint 恢复

上例可直接 kubectl apply 验证全链路。预期输出 global: 64 local: 4——每个 slice 是一个独立的 JAX 集群,global 等于该 slice 的芯片数(4x4x4 = 64),local 是本主机的 4 芯;多主机协同所需的环境变量由 GKE 自动注入。两个 replica 各打印一次。

平台只写 gke-tpu-slice-topology 这一个注解。Kueue 实际消费的是 kueue.x-k8s.io/podset-slice-required-topology + podset-slice-size,但这对注解由 Kueue Slice Provisioner 的 mutating webhook 在创建时自动注入,平台不需要自行计算 slice 大小。反过来这也是排障入口:Workload 上看不到这对注解,即说明该 webhook 没被调用(查 slice-controller-system 下的 deployment 日志)。

三条硬约束(平台侧应在提交前校验,否则报错信息很难反推原因):

约束漏掉时的表现
每个 TPU podSet 必须声明 cloud.google.com/gke-tpu-slice-topologywebhook 无从翻译,该 podSet 拿不到 TAS 拓扑请求 → 被所有 TAS flavor 跳过(supports only TopologyAwareScheduling)→ 最终落到 cpu-standard(TPU 额度 0)→ 报「配额不足」,信息里完全不提缺注解
每个 TPU podSet 必须声明 nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu7x该 podSet 对节点类型不作限定,可能被排到非 TPU 节点上;在静态形态下还需同时声明 cloud.google.com/gke-tpu-topology,否则无法确定该落哪一个拓扑 flavor(见 4.6)
TPU 污点必须被容忍TAS 在准入期即按污点筛节点,报 excluded: taint "google.com/tpu=present:NoSchedule": N——pod 还没创建就被拦下。推荐在 ResourceFlavor 上统一声明 tolerations(见 4.2),由 Kueue 准入时注入,作业模板无需重复写;作业自带 toleration 也兼容(两者取并集)

Pathways 形态的差异:JobSet 内含两个 replicatedJob,写法各不相同——

podSetslice 拓扑注解gke-tpu-accelerator nodeSelectorTPU 请求落到的 flavor
pathways-worker(同上例)tpu7x
pathways-head(纯 CPU)不要不要(改用 gke-nodepool 定向到 CPU 池)不要,只声明 cpu/memory/ephemeral-storagecpu-standard

head 到专用 CPU 节点池的定向,由它自己模板里的 nodeSelector: cloud.google.com/gke-nodepool: <pool> 完成。两个 podSet 同属一个 Workload,配额一次性原子预留、被抢占时一起挂起(gang 完整);slice 门闩因 worker 落在 tpu7x 上而对整个 Workload 生效,head 会一起等待 slice 成型。

注意事项:

  • 不要设置 spec.suspend:Kueue 的 webhook 会强制新 JobSet 挂起,准入通过后自动解挂。AIStudio 也不要手工翻转这个字段。
  • 不要设置 pod 的 priorityClassName(K8s 原生优先级)来表达高保/低保——那会引发 kubelet 层抢占,语义完全不同。
  • 命名长度预算:slice 名按 {ns}-jobset-{JobSet名}-{5位hash}-{replicatedJob名}-{序号} 生成,总长 ≤ 59 字符,超长会导致 slice 创建失败。在 kubemaker 命名空间下,JobSet 名 + replicatedJob 名合计约 ≤ 35 字符——AIStudio 生成名字时必须校验。
  • 小任务(sub-slice)写法完全相同:只需把注解改为小拓扑(如 2x2x2)、completions/parallelism 改为对应主机数(2x2x2 = 8 芯 = 2 台主机);队列、优先级、配额记账与大任务无任何差别。
  • 容器声明的 cpu/memory/ephemeral-storage 与 TPU 同属一个资源组,由同一个 flavor 记账:TPU podSet 记入 tpu7x(形式化大配额,不构成准入瓶颈),纯 CPU podSet 记入 cpu-standard(真实配额,见 4.2)。
  • 提交高保作业前,AIStudio 应校验本队高保余额(nominal 剩余量),因为超出 nominal 的高保作业不再有抢占特权(见约束一节)。

Job 与 Deployment 形态的差异。两者与 JobSet 共用同一套队列、配额、优先级与抢占,提交时同样只加 queue-name + priority-class 两个标签,标签打在对象自身的 metadata.labels 上。仅有的差异:

batch/jobdeployment
Workload 粒度一个 Job一个 Pod
gang覆盖 parallelism 个 pod无(逐 pod 准入)
挂起方式spec.suspend(同 JobSet)pod 调度门 kueue.x-k8s.io/admission,Deployment 对象本身不挂起
配额不足时整个 Job 排队,0 pod 落地以副本数不足的状态运行,已准入的 pod 照常服务
TPU 场景与 JobSet 同:需 TAS 拓扑注解 + TPU 请求常驻服务一般为纯 CPU,落 cpu-standard,无需注解

Deployment 的「部分副本运行」是其每 pod 一个 Workload 的直接后果,对常驻服务通常不可接受——因此服务类负载应放在本队 nominal 配额内,不要依赖 cohort 借用(借来的容量会被高保训练作业回收,导致副本被驱逐)。若某服务必须保证副本数,可用 ClusterQueue 的 lendingLimit 限制该队额度被他队借走的比例。

4.3.2 查询作业状态

每个被管理的 JobSet 会有一个 Kueue 自动创建的 Workload 影子对象,作业的排队/准入/抢占状态全部体现在它的 status.conditions 上。AIStudio 通过 client-go 访问(Kueue 随仓库发布 typed client:sigs.k8s.io/kueue/client-go),建议以 informer watch Workload(以 kueue.x-k8s.io/job-uid 标签建索引)代替轮询。Workload 的名字由 Kueue 生成(形如 jobset-<name>-<hash>),不要拼名字,按标签反查:

go
// 反查:JobSet UID → Workload(该标签由 Kueue 自动打上并已建索引)
wls, err := kueueClient.KueueV1beta2().Workloads("kubemaker").List(ctx, metav1.ListOptions{
    LabelSelector: "kueue.x-k8s.io/job-uid=" + string(jobSet.UID),
})

⚠️ v1beta2 的字段更名:优先级相关的 spec.priorityClassNamespec.priorityClassSource(v1beta1)在 v1beta2 中已合并为结构化的 spec.priorityClassRef(含 group/kind/name),解析出的数值仍在 spec.priority。按旧字段名读取不会报错,只会静默取到空值——AIStudio 的 watcher、运维脚本与告警规则都要按新字段名编写。判断作业档位建议直接用 spec.priority 的数值(1000 = 高保 / 100 = 低保),比字符串匹配更稳。

Workload 对象示例一——已准入、运行中(节选,只保留状态判定相关字段):

yaml
apiVersion: kueue.x-k8s.io/v1beta2
kind: Workload
metadata:
  name: jobset-llm-pretrain-a3f8-4f9d2
  namespace: kubemaker
  labels:
    kueue.x-k8s.io/job-uid: 1c8f5b7e-…           # = JobSet 的 UID,反查键
  ownerReferences:
  - apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    name: llm-pretrain-a3f8                      # 随 JobSet 删除被级联 GC
spec:
  queueName: team-a-queue
  priorityClassRef:                              # ⚠️ v1beta2 字段;v1beta1 的
    group: kueue.x-k8s.io                        # priorityClassName / priorityClassSource
    kind: WorkloadPriorityClass                  # 两个字段已被它取代(见下方注意事项)
    name: best-effort
  priority: 100                                  # 由 priorityClassRef 解析得到
  active: true
  podSets:
  - name: workers
    count: 32                                    # 2 slice × 16 主机
status:
  conditions:
  - type: QuotaReserved                          # 配额已整体预留(gang)
    status: "True"
  - type: Admitted                               # 已准入,JobSet 已解挂
    status: "True"
  admissionChecks:
  - name: tpu7x-slice
    state: Ready                                 # slice 已成型激活;未 Ready 时作业停在「slice 供给中」
  admission:
    clusterQueue: team-a                         # 由哪个队的配额准入
    podSetAssignments:
    - name: workers
      flavors:
        google.com/tpu: tpu7x
      resourceUsage:
        google.com/tpu: "128"                    # 记账用量:32 pod × 4 芯
      topologyAssignment:                        # TAS 装箱结果(节选):精确到主机
        levels: [kubernetes.io/hostname]
        domains:
        - values: ["gke-tpu-node-1a2b"]
          count: 1
        # …共 32 台主机;Slice Provisioner 依此创建 Slice CR,托管 Slice Controller 布线成型

示例二——被抢占后、等待重排(status 节选):

yaml
status:
  conditions:
  - type: QuotaReserved
    status: "False"                              # 配额已被收回
  - type: Evicted
    status: "True"
    reason: Preempted                            # ← watcher 落库的触发条件
    message: 'Preempted to accommodate a workload (UID: 9a41…, JobUID: 5be2…) due to
      reclamation within the cohort; preemptor path: /tpu-pool/team-a; preemptee path:
      /tpu-pool/team-b'                          # 消息内含抢占者 JobUID,可反查对方作业
  - type: Preempted
    status: "True"
    reason: InCohortReclamation                  # 队内抢占 InClusterQueue;跨队回收(抢占者在
                                                 # nominal 内)InCohortReclamation;份额抢占
                                                 # (抢占者借用中)InCohortFairSharing
  - type: Requeued                               # 已重新回队,等待下一次准入
    status: "True"
  - type: Admitted
    status: "False"

AIStudio 的状态机映射(watcher 监听 Workload + JobSet 两个对象的条件即可推导全部状态):

AIStudio 展示状态判定依据
排队中Workload 存在,QuotaReserved ≠ True
slice 供给中QuotaReserved=True 且 AdmissionCheck tpu7x-slice 未 Ready(配额已到手,slice 布线成型中)
启动中Workload Admitted=True,JobSet 副本尚未全部 Ready(slice 已激活,pod 拉起中)
运行中JobSet replicatedJobsStatus 全部 Ready
被抢占,重排中Workload Evicted=True(reason=Preempted)随后 Requeued=True;界面附带抢占原因与队列位置
已完成JobSet Completed=True(Workload 同步出现 Finished)
已失败JobSet Failed=True(真实故障耗尽 maxRestarts)

排队位置(「当前第 N 位」)走 Kueue 内置的 Visibility API(on-demand 实时计算,同样有 typed client,也可 REST 直连)。响应示例(节选):

yaml
# GET /apis/visibility.kueue.x-k8s.io/v1beta2/namespaces/kubemaker/localqueues/team-a-queue/pendingworkloads
apiVersion: visibility.kueue.x-k8s.io/v1beta2
kind: PendingWorkloadsSummary
items:
- metadata:
    name: jobset-llm-pretrain-b2d1-8c3a1
    namespace: kubemaker
  priority: 100
  localQueueName: team-a-queue
  positionInClusterQueue: 3          # 在团队 ClusterQueue 全队中的位置
  positionInLocalQueue: 1            # 在本 LocalQueue 中的位置

4.3.3 删除/取消作业

只删 JobSet,client-go 一次调用(默认后台级联 GC):

go
err := jobsetClient.JobsetV1alpha2().JobSets("kubemaker").
    Delete(ctx, "llm-pretrain-a3f8", metav1.DeleteOptions{})

级联链路:JobSet 删除 → Workload 随 ownerReference 被 GC、配额立即释放 → 对应 Slice CR 随之删除(托管 Slice Controller 拆除 slice),分区归还容量池。禁止直接删除 Workload 或 Slice 来「取消」作业。

4.3.4 抢占通知与历史落库

Kueue 在抢占时留下的痕迹及其时效:

信号位置内容时效
Preempted 事件受害者 Workload抢占者 Workload UID / JobUID、原因、双方队列路径与优先级约 1h 过期
Evicted / Preempted / Requeued 条件Workload status同上消息;Preempted 条件的 reason 有三种:队内(InClusterQueue)、跨队回收(InCohortReclamation,抢占者在 nominal 内)、份额抢占(InCohortFairSharing,抢占者借用中)被下次状态变化覆盖
Stopped 事件用户的 JobSet 上驱逐原因(用户 kubectl describe jobset 可见)约 1h 过期

由于事件短命、条件会覆盖,AIStudio watcher 必须在观察到 Evicted(reason=Preempted)时立即落库:记录时间、原因、抢占者 JobUID(可通过 kueue.x-k8s.io/job-uid 标签反查出抢占者的 JobSet 与所属团队——普通用户无跨命名空间权限,这一步由平台以集群视角完成),并推送用户通知。注意判定要按条件类型进行(Evicted 且 reason=Preempted),不要对 Preempted 条件的 reason 做白名单过滤——它有三种取值(见上表),按白名单过滤会漏记份额抢占。该记录同时是后续迭代「被抢即失败」选项的数据基础。

4.3.5 AIStudio 服务账号 RBAC 增量

text
jobsets.jobset.x-k8s.io                 : create / get / list / watch / delete(现有)
jobs.batch                              : create / get / list / watch / delete(现有)
deployments.apps                        : create / get / list / watch / delete(现有)
workloads.kueue.x-k8s.io                : get / list / watch(新增,只读)
events                                  : get / list(新增,只读)
visibility.kueue.x-k8s.io(pendingworkloads): get(新增,队列位置)

4.3.6 前瞻:Pathways 形态作业的接入

Pathways 接入本期不定稿(见非目标)。可行性核实结论:上游曾提供 PathwaysJob CRD 作为提交入口,但其仓库 README 已声明 "PathwaysJob has been deprecated in favor of direct jobsets"。因此未来引入 Pathways 时,提交对象仍是 JobSet,不引入新 CRD 与控制器——4.3.1–4.3.5 的交互契约、Kueue 的准入/gang/TAS/slice 门闩/抢占机制、以及 4.2 已建的全部 CR 都无需改动(head 是纯 CPU podSet,自动落 cpu-standard,与 TPU worker 同 gang 准入)。需要处理的只有四处:

环节改动
提交模板JobSet 内增加 head replicatedJob(纯 CPU 容器)+ worker replicatedJobs(TPU);head 到专用 CPU 池的定向由模板 nodeSelector(如 cloud.google.com/gke-nodepool: <pool>)表达
平台校验 / 参数pod 的 priorityClassName 保持为空(Pathways 文档的 hot-swap 玩法依赖 K8s 原生优先级,违反本方案约束);worker 的 backoffLimit 按 Pathways 弹性指引调高——4.3.1 示例中面向非弹性作业的 backoffLimit: 0 不适用
waitForPodsReady.recoveryTimeout需复核:弹性等待期(故障 slice 等待替换)pod not-ready 属正常状态,超时会把整个 Workload 驱逐、弹性收益归零;须结合 slice 重组时长实测调参
与 dynamic slicing 联动待与 GKE 确认 Pathways 容器 × 动态切片(尤其 sub-slice)的官方验证状态;确认前 Pathways 作业可先跑 4.6 的静态节点池 flavor

4.4 抢占行为规范

规则回顾(上游 Kueue 语义,此处为运维与答疑基准):

  • 触发资格:经典语义下,只有自身不需要借用(准入后总用量 ≤ 本队 nominal)的作业才触发抢占;开启 Fair Sharing 后该门槛被移除,借用中的作业也可触发,能否抢成由份额策略决定(见下节影响 5);
  • 受害者资格:跨队只找借用中(用量 > nominal)的队列,且受害者优先级须严格低于抢占者;队内按优先级;跑在自家 nominal 内的作业,别的队永远抢不到;
  • 受害者选择:开启 Fair Sharing 后(见下),跨队从份额最高的借用队列开始抢;同一队列内部仍按「优先级更低者先 → 同优先级中最近准入者先(LIFO,保护跑得久的作业)」;整体只驱逐「刚好放下抢占者」的最小集合。
  • 放置层行为(大小任务混跑):TAS 采用最小拓扑装箱——sub-slice 任务优先填「刚好只剩所需拓扑」的碎片 cube,完整 cube 留给 super-slice 任务;准入期需要腾位时偏好更低优、更小的作业,并立即为受害者重新找位(SubSlicing 指南行为)。

Fair Sharing 对抢占结果的影响

本方案开启 Fair Sharing(4.1 的 Configuration 中出现 fairSharing 配置块即为启用)。Kueue 为每个 ClusterQueue 实时计算 DRS 份额(dominant resource share:主导资源的借用量占 cohort 可借总量的比例,再除以队列权重),并用两条按序尝试的策略约束每一次抢占:

  • LessThanOrEqualToFinalShare:抢占完成后,抢占方份额(含新作业)≤ 受害方剩余份额,才允许抢;
  • LessThanInitialShare:抢占前,抢占方份额(含新作业)< 受害方当前份额,才允许抢。

对抢占结果的具体影响:

  1. 跨队受害者从「谁最近准入谁先死」变为「超出公平份额最多的团队先还」。默认排序不看各队借了多少,借 8 芯的队可能替借 100 芯的队挨刀;Fair Sharing 下每一刀都落在当前份额最高的队列上。
  2. 回收止于份额拉平:不会把某一个团队一撸到底——某队份额被抽到不再最高后,下一刀落到新的份额最高者,直到容量凑够。
  3. 借用侧同样变公平:多个团队的低保同时排队等闲置时,按份额轮流准入(份额低者优先),不再是先到先得。
  4. 受害者的优先级门槛不变:受害者仍须严格更低优且其队列在借用中,任何队列 nominal 内的作业不可侵犯——两档模型下低保永不引发牺牲。
  5. 触发资格放宽:经典语义下借用中的作业从不抢占;Fair Sharing 下该门槛被移除,借用中的作业也可按份额策略抢占其他借用队列。抢占 reason 相应区分:抢占者在 nominal 内 → InCohortReclamation(FairSharingPreemptWithinNominal 门控,v0.17 起默认开启,此路径无条件越过份额比较);抢占者借用中 → InCohortFairSharing(走 DRS 份额策略)。

对比示例:A 借 96 芯、B 借 24 芯(均为低保),C 在自身 nominal 内提交 64 芯高保——

  • 无 Fair Sharing:按准入时间排序,若 B 的作业更新,可能 B 先被抢光 24 芯、再抢 A 的 40 芯——借得少的 B 反而全军覆没;
  • 有 Fair Sharing:A 份额最高,受害者从 A 中选出;A 归还 64 芯后份额仍不低于 B,B 全程不受影响

各队实时份额可通过 ClusterQueue.status.fairSharing.weightedShare 字段与 kueue_cluster_queue_weighted_share 指标观测,建议纳入团队管理员面板。

场景验收表(staging 必须逐一复现;A/B/C 三队同构,各 nominal 256 / borrowing 128):

#场景预期结果
1A 用 128;B 提交 384 芯低保全部准入(B 借走 A 闲置的 128)
2接 1;A 提交 128 芯高保B 的低保借用者被驱逐(队内按 LIFO 挑),凑够 128 即止;被驱逐者回队
3A 已满 256 高保,再提 100 芯高保(需借用)不动任何队 nominal 内的作业;若他队存在借用中的低保且 FS 份额策略满足,以 InCohortFairSharing 抢占之,否则排队
4A 队内 156 高保 + 100 低保;A 提 100 芯高保队内抢占 A 自己的低保(低保防外不防内)
5A 借用中;B 提 100 芯低保(B nominal 有空)不抢占,排队(同优先级不抢,低保从不引发牺牲)
6A 借 96、B 借 24(均低保);C 提 64 芯高保Fair Sharing 生效:受害者全部出自份额最高的 A,B 不受影响

4.5 被抢占后的行为:只重试(本期唯一语义)

完整链路:

text
Kueue 选中受害者
  → JobSet 被置 suspend=true → pod 收到 SIGTERM(宽限期内写 checkpoint)→ pod 删除
  → Slice CR 删除,slice 拆除,分区归还容量池
  → Workload 回到所属队列(Requeued)
  → 配额空出 → 重新准入:TAS 重新装箱选块 → slice 重新成型激活(AdmissionCheck)
  → JobSet 解挂 → 续跑(可能落在另一组分区上)

两个必须写进用户文档的事实:

  1. 抢占不是失败:因 suspend 被删的 pod 不计入 Job 失败,JobSet 的 failurePolicy/maxRestarts 对抢占完全不生效、不消耗重启预算。maxRestarts 只需按真实故障(进程崩溃、硬件失败)的容忍度设置。
  2. 续跑可能落在不同硬件上,且需要重新组 slice:重排后 TAS 可能选中另一组分区,slice 要重新布线成型,恢复时延高于「回到现成硬件」,属预期行为;应用不得依赖节点本地状态,checkpoint 必须写远端存储(GCS 等)。

抢占之外的运行期故障走同一条重排链路:slice 因硬件故障或主机维护自动拆除时,Slice CR 转 FAILED,Kueue 立即为受影响作业重新装箱选块(fast replacement,无需等待计时器),waitForPodsReady.recoveryTimeout 作为兜底——超时仍未恢复则驱逐重排。对用户而言与被抢占的体验一致:自动回队、自动续跑。

应用侧契约(低保作业提交时由 AIStudio 提示并在文档中约定):处理 SIGTERM、周期性 + 退出前写 checkpoint、启动时自动从最新 checkpoint 恢复、输出采用幂等提交(临时路径 + 原子 rename,或按 run-id 版本化)。

4.6 静态节点池形态(当前集群实际形态)与向 All Capacity 的迁移

当前 GKE 集群不是 All Capacity 模式,TPU 节点池按固定拓扑静态创建(ICI 在建池时布好)。本节即当前阶段的实施基准,4.1–4.5 的动态切片部分是启用 All Capacity 之后的目标形态。

两种形态下,本方案的排队、配额、优先级、抢占、Fair Sharing、gang、AIStudio 交互契约全部一致——这些能力属于 Kueue 本身,与切片方式无关。差异只在「放置与供给」这一层:

配置项动态切片形态(本方案主线)静态节点池形态
集群组件(4.1)Kueue + JobSet + 托管 Slice Controller + Kueue Slice Provisioner仅 Kueue + JobSet(后两者不装,也无需 --enable-slice-controller)
GKE 版本 / 容量前提Rapid ≥ 1.36(sub-slicing);All Capacity + 增量供给节点池无特殊要求,常规 TPU 节点池即可
Topology CR7 层分区树(block → 五档几何 → 主机),必需保留但简化为 2 层:cloud.google.com/gke-nodepool(= 一个静态 slice)→ kubernetes.io/hostname
ResourceFlavortopologyName: tpu7x-topology + 加速器标签保留 topologyName(指向 2 层 Topology);nodeLabels 改为匹配静态池的固定拓扑,例如 cloud.google.com/gke-tpu-accelerator: tpu7x + cloud.google.com/gke-tpu-topology: 4x4x4
TAS 的作用在分区树上选定分区组合(决定 slice 切多大、用哪些 cube)保留,但职责收窄为:把一个 Job 的 pod 装进同一个节点池(gang 放置)、节点故障时 hot swap 换节点、准入时校验放得下
AdmissionCheck + admissionChecksStrategyslice 成型门闩,必需不需要(两处一并删除);准入通过即可直接解挂
cpu-standard flavor 与 cpu/memory 资源组保留保留(与切片方式无关,见 4.2)
作业 YAML(4.3.1)pod 模板注解 gke-tpu-slice-topology: <拓扑>换成 nodeSelector 指定固定拓扑(gke-tpu-accelerator + gke-tpu-topology)+ TAS 注解 kueue.x-k8s.io/podset-required-topology(约束落在同一节点池)
Kueue Configuration(4.1)同主线同主线;waitForPodsReady.recoveryTimeout 可适当收紧(不存在 slice 重组等待)
AIStudio 状态机(4.3.2)含「slice 供给中」态去掉该态,其余映射不变(Admitted 后直接进入「启动中」)
拓扑白名单校验两档白名单(super / sub)改为校验「所提拓扑存在对应静态节点池」

需要接受的能力损失(即引入动态切片的收益反面):

  1. 拓扑在建池时固定:每种拓扑需预先规划节点池,任务形态变化要重建节点池(分钟至数十分钟级),不能按作业需求现拼;
  2. 碎片化不可回收:TAS 仍会在现有节点中做装箱,但无法跨节点池重组硬件——池内剩余不足一个作业所需拓扑时,该容量只能空置;
  3. 故障恢复更慢:TAS 的 hot swap 仅在同池内有空闲节点时才换得成,换不到就得等节点池重建(而非动态切片的软件级重拼 slice,GKE 侧数据为恢复快约 5 倍),且没有 fast replacement;
  4. 无 sub-slice 能力:小任务只能独占一个完整拓扑的节点池,或另建小规格静态池;
  5. 配额按拓扑切分:多种静态拓扑时每种一个 flavor、各自独立配额,不再是一本可自由流动的账(见下方 ② 的说明)。

完整 YAML 示例

① 全局共享对象——与 4.2 相比:Topology 从 7 层简化为 2 层、删除 AdmissionCheck,ResourceFlavor 保留 topologyNamenodeLabels 改为匹配静态池的固定拓扑;三个 WorkloadPriorityClasscpu-standard 逐字不变。

由于 ResourceFlavor.nodeLabels精确匹配的键值对,一个 flavor 只能代表一种拓扑——集群里有几种静态节点池,就要建几个 flavor。下例覆盖四档拓扑:

yaml
apiVersion: kueue.x-k8s.io/v1beta2
kind: WorkloadPriorityClass
metadata:
  name: guaranteed          # 高保(与主线相同)
value: 1000
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: WorkloadPriorityClass
metadata:
  name: best-effort         # 低保(与主线相同)
value: 100
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: WorkloadPriorityClass
metadata:
  name: default             # 默认优先级档(与主线相同)
value: 10
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: Topology
metadata:
  name: nodepool-topology   # 静态形态的分区树只有两层
spec:
  levels:
  - nodeLabel: cloud.google.com/gke-nodepool   # 一个节点池 = 一个静态 slice
  - nodeLabel: kubernetes.io/hostname          # 主机
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: ResourceFlavor
metadata:
  name: tpu7x-4x4x4                            # 64 芯 / 池(16 主机)
spec:
  nodeLabels:
    cloud.google.com/gke-tpu-accelerator: tpu7x
    cloud.google.com/gke-tpu-topology: 4x4x4   # 静态池建池时固定的拓扑
  topologyName: nodepool-topology              # 启用 TAS:池内 gang 放置 + 节点 hot swap
  tolerations:                                 # 与主线相同:TPU 污点在 flavor 上统一声明
  - {key: google.com/tpu, operator: Exists, effect: NoSchedule}
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: ResourceFlavor
metadata:
  name: tpu7x-4x4x8                            # 128 芯 / 池(32 主机)
spec:
  nodeLabels:
    cloud.google.com/gke-tpu-accelerator: tpu7x
    cloud.google.com/gke-tpu-topology: 4x4x8
  topologyName: nodepool-topology
  tolerations: [{key: google.com/tpu, operator: Exists, effect: NoSchedule}]
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: ResourceFlavor
metadata:
  name: tpu7x-4x8x8                            # 256 芯 / 池(64 主机)
spec:
  nodeLabels:
    cloud.google.com/gke-tpu-accelerator: tpu7x
    cloud.google.com/gke-tpu-topology: 4x8x8
  topologyName: nodepool-topology
  tolerations: [{key: google.com/tpu, operator: Exists, effect: NoSchedule}]
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: ResourceFlavor
metadata:
  name: tpu7x-8x8x8                            # 512 芯 / 池(128 主机)
spec:
  nodeLabels:
    cloud.google.com/gke-tpu-accelerator: tpu7x
    cloud.google.com/gke-tpu-topology: 8x8x8
  topologyName: nodepool-topology
  tolerations: [{key: google.com/tpu, operator: Exists, effect: NoSchedule}]
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: ResourceFlavor
metadata:
  name: cpu-standard        # 与主线逐字相同(cpu/memory 记账通道,不配 nodeLabels/topologyName)
# 注意:此形态下不需要 AdmissionCheck 对象

TAS 在此形态下的三项作用:① 池内 gang 放置——把一个 Job 的全部 pod 装进同一个节点池(即同一 slice),避免跨 slice 拆散;② 节点 hot swap——运行期单节点故障时在同池内找替换节点,不重排整个作业(TASFailedNodeReplacement 等门控默认开启);③ 准入即校验放得下——容量树上放不下的作业不会被准入,不会出现「账上有量、实际没位置」。它不再承担「slice 切多大」的决策(那由建池时的拓扑决定)。

② 每团队对象——与 4.2 相比只有两处改动:TPU flavor 换成上面这组静态 flavor、删除整个 admissionChecksStrategy 块;cohort、优先级策略与单一资源组结构全部不变。

关键差异在配额:nominalQuota/borrowingLimit 按 flavor 分别配置,团队的 TPU 额度必须按拓扑拆分,不再是一本可自由流动的账。这不是 Kueue 的限制,而是硬件事实的忠实反映——声明 gke-tpu-topology: 4x4x4 的作业无法落到 8x8x8 的节点池上。

yaml
apiVersion: kueue.x-k8s.io/v1beta2
kind: ClusterQueue
metadata:
  name: team-a
spec:
  cohortName: tpu-pool
  namespaceSelector:
    matchLabels:
      scheduling.antgroup.com/kueue-tier: production
  resourceGroups:
  # 仍是单一资源组;多拓扑体现为组内并列多个 TPU flavor,cpu-standard 兜底。
  # 每个 flavor 都必须完整声明四种资源且与 coveredResources 同序(CEL + webhook 强制)。
  - coveredResources: ["google.com/tpu", "cpu", "memory", "ephemeral-storage"]
    flavors:
    - name: tpu7x-4x4x4            # 团队在小拓扑池上的额度
      resources:
      - {name: google.com/tpu, nominalQuota: 64, borrowingLimit: 64}     # = 1 个 4x4x4 池
      - {name: cpu, nominalQuota: "100000"}          # 形式化大配额,不构成约束
      - {name: memory, nominalQuota: 100000Gi}
      - {name: ephemeral-storage, nominalQuota: 100000Gi}
    - name: tpu7x-4x4x8
      resources:
      - {name: google.com/tpu, nominalQuota: 128, borrowingLimit: 128}   # = 1 个 4x4x8 池
      - {name: cpu, nominalQuota: "100000"}
      - {name: memory, nominalQuota: 100000Gi}
      - {name: ephemeral-storage, nominalQuota: 100000Gi}
    - name: tpu7x-4x8x8
      resources:
      - {name: google.com/tpu, nominalQuota: 256, borrowingLimit: 0}     # = 1 个 4x8x8 池
      - {name: cpu, nominalQuota: "100000"}
      - {name: memory, nominalQuota: 100000Gi}
      - {name: ephemeral-storage, nominalQuota: 100000Gi}
    - name: tpu7x-8x8x8            # 该团队不用大拓扑:配 0,或整段省略
      resources:
      - {name: google.com/tpu, nominalQuota: 0, borrowingLimit: 512}     # 仅允许借用他队闲置的大池
      - {name: cpu, nominalQuota: "100000"}
      - {name: memory, nominalQuota: 100000Gi}
      - {name: ephemeral-storage, nominalQuota: 100000Gi}
    - name: cpu-standard           # 纯 CPU podSet 兜底,TPU 额度为 0
      resources:
      - {name: google.com/tpu, nominalQuota: 0}
      - {name: cpu, nominalQuota: "2000"}            # 这里才是真实约束
      - {name: memory, nominalQuota: 8000Gi}
      - {name: ephemeral-storage, nominalQuota: 4000Gi}
  preemption:
    reclaimWithinCohort: LowerPriority
    withinClusterQueue: LowerPriority
---
apiVersion: kueue.x-k8s.io/v1beta2
kind: LocalQueue
metadata:
  name: team-a-queue
  namespace: kubemaker
spec:
  clusterQueue: team-a

配套的三条规则:

  1. 作业按 nodeSelector 自动匹配 flavor:声明 gke-tpu-topology: 4x4x8 的作业只会与 tpu7x-4x4x8 flavor 匹配,占用该 flavor 的额度;无需在提交侧指定 flavor 名。
  2. 配额校验改为按 flavor 分别进行:原「Σ 各队 nominalQuota ≤ 集群总芯数」变为每个 flavor 各校验一次——Σ 各队在 4x4x8 flavor 的 nominalQuota ≤ 全部 4x4x8 节点池的总芯数,其余拓扑同理。额度下发流水线需相应改造(见测试计划)。另外每档额度应为该拓扑单池芯数的整数倍(作业以整池为单位占用 slice),否则尾数无法被任何作业使用。
  3. 借用与抢占同样限定在 flavor 内:团队 A 可以借用团队 B 闲置的 4x4x8 额度,但 B 闲置的 8x8x8 额度对 A 的 4x4x8 作业毫无帮助。多 flavor 场景下 flavorFungibility(默认 whenCanBorrow: MayStopSearch / whenCanPreempt: TryNextFlavor)决定跨 flavor 的尝试顺序,但只对同时兼容多个 flavor 的作业有意义;本形态下作业按拓扑绑定单一 flavor,该字段基本不起作用,保持默认即可。

拓扑档数越多,这种切分的代价越明显:上例中团队 A 的额度被切成四份,任何一份的闲置都无法支援另一份。这正是动态切片相对静态形态的核心收益——动态切片下只有一个 flavor、一本账,任何拓扑的作业都从同一池子里取容量,不存在「4x4x4 额度用完了但 8x8x8 闲着」的僵局。

③ 作业 YAML——与 4.3.1 相比:gke-tpu-slice-topology 注解换成 TAS 拓扑注解、nodeSelector 中声明固定拓扑,其余全部不变。此形态不部署 Slice Provisioner,没有 webhook 做翻译,而 Kueue 只从 pod 模板的 annotations 读取拓扑请求、无默认值可兜底,因此 TAS 注解必须由平台模板直接写出。注解按作业是否多 slice 二选一:

作业形态注解写法
单 slice(podSet 的全部 pod = 一个 slice)kueue.x-k8s.io/podset-required-topology: cloud.google.com/gke-nodepool
多 slice(一个 podSet 含 N 个 slice,如 JobSet 的 replicas: N)kueue.x-k8s.io/podset-slice-required-topology: cloud.google.com/gke-nodepool + kueue.x-k8s.io/podset-slice-size: "<每 slice 的主机数>"(两者必须成对出现)

⚠️ 多 slice 作业若误用 podset-required-topology,等于要求全部 pod 挤进一个节点池——当 pod 数超过单池主机数时永远不可能满足,作业将一直 pending。JobSet 的多个 replicas 在 Workload 里合并为一个 podSet(count = replicas × parallelism),因此这类作业一律用 slice 级注解。

另:flavor 一旦挂了 topologyName,其上的 podSet 必须声明拓扑注解,否则报 Flavor "<name>" supports only TopologyAwareScheduling。纯 CPU 的 podSet 落在不带 topologyNamecpu-standard 上,无需注解。

yaml
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: llm-pretrain-a3f8
  namespace: kubemaker
  labels:
    kueue.x-k8s.io/queue-name: team-a-queue     # 不变
    kueue.x-k8s.io/priority-class: best-effort  # 不变
spec:
  failurePolicy:
    maxRestarts: 3
  replicatedJobs:
  - name: workers
    replicas: 2           # 2 个 slice → 需要 2 个 4x4x4 静态节点池
    template:
      spec:
        backoffLimit: 0
        completions: 16
        parallelism: 16
        template:
          metadata:
            annotations:
              # ← 替换原 slice 拓扑注解。多 slice 作业(本例 replicas=2)用 slice 级注解:
              #    每 16 个 pod(一个 4x4x4 slice)必须落在同一个节点池内
              kueue.x-k8s.io/podset-slice-required-topology: cloud.google.com/gke-nodepool
              kueue.x-k8s.io/podset-slice-size: "16"
          spec:
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu7x
              cloud.google.com/gke-tpu-topology: 4x4x4   # ← 改为在此声明固定拓扑
            # TPU 污点的 toleration 由 ResourceFlavor 统一声明并注入(见 ①),此处无需重复
            containers:
            - name: jax
              image: python:3.12
              command:                  # 与 4.3.1 相同的冒烟用例
              - bash
              - -c
              - |
                pip install -q "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                python -c 'import jax; print("global:", jax.device_count(), "local:", jax.local_device_count())'
              resources:
                limits:
                  google.com/tpu: '4'
                requests:
                  ephemeral-storage: 20Gi
                  google.com/tpu: '4'

迁移清单:静态形态 → All Capacity + dynamic slicing

届时按下表逐项切换,LocalQueue、WorkloadPriorityClass、cpu/memory/ephemeral-storage 资源组、namespaceSelector、抢占策略与 AIStudio 交互契约均不改动:

步骤操作
1. 容量与集群TPU 容量改以 All Capacity 预留,节点池按增量供给重建(workload policy 4x4x4 + provision_only);集群升到 Rapid ≥ 1.36.0-gke.3712000;完成 sub-slicing preview 三道手续(实施计划阶段〇)
2. 组件gcloud container clusters update --enable-slice-controller 开启托管 Slice Controller;部署 Kueue Slice Provisioner
3. Topology2 层(gke-nodepoolhostname)替换为 7 层(block → 五档几何 → hostname)。必须等节点五档 -id 标签齐备后再下发,否则 TAS 会排除全部节点
4. ResourceFlavor多个静态 flavor(tpu7x-4x4x4 / tpu7x-4x8x8 / …)合并为单个动态 flavor:nodeLabels 只留 gke-tpu-accelerator: tpu7x(去掉 gke-tpu-topology),topologyName 指向 7 层 Topology。资源组结构不变——仍是单组、cpu-standard 仍兜底、分流机制照旧
5. 配额各档拓扑的独立配额合账为单一数字(这是迁移的主要收益:不再有「A 档用完、B 档闲置」);Σ 各队 nominalQuota ≤ 集群总芯数
6. AdmissionCheck新建 controllerName: accelerator.gke.io/slice 的检查,并在各 ClusterQueue 的 admissionChecksStrategy 中以 onFlavors: [<动态 flavor>] 挂载
7. 作业模板拓扑注解从 podset-slice-required-topology + podset-slice-size 换成 cloud.google.com/gke-tpu-slice-topology: <拓扑>(后者由 Slice Provisioner 的 webhook 自动翻译成前两者,平台无需再自行计算 slice 大小);nodeSelector 去掉 gke-tpu-topology 一项。TPU toleration 与资源请求写法不变
8. 状态机AIStudio 增加「slice 供给中」态(QuotaReserved=True 且 AdmissionCheck 未 Ready)

迁移期建议:先在 staging 用同一套 ClusterQueue 名与配额数字演练 3–7 步,确认 Workload 条件序列与 4.4 场景表行为不变,再在生产低负载窗口切换(第 4 步更换 flavor 会触发存量 Workload 重新 reconcile)。

因此静态形态既是当前落地基准,也是动态切片就绪前的完整可用形态:排队/配额/抢占/AIStudio 联调全部可以先行跑通并验收。

测试计划

  • 单元/配置层:额度下发流水线的校验逻辑(Σnominal ≤ 集群总量、borrowingLimit 非负、优先级枚举合法)。
  • 集成(staging 集群):4.4 场景表六个用例逐一自动化复现,断言 Workload 条件序列(Evicted→Requeued→Admitted)、Preempted 条件 reason 按场景取值(InClusterQueue / InCohortReclamation / InCohortFairSharing)、受害者选择(含 Fair Sharing 的份额规则)、驱逐数量最小;抢占-续跑闭环用带 checkpoint 的样例训练任务验证 loss 曲线连续。
  • sub-slicing 用例:小拓扑(如 2x2x2)作业的准入-成型-运行全链路;同一 cube 内多个 sub-slice 共存与配额记账正确;sub-slice 低保作业被高保抢占(复用 4.4 规则,拓扑换成小任务);人为将运行中 slice 置为故障,验证 fast replacement 立即重排与 recoveryTimeout 兜底驱逐。
  • AIStudio 联调:创建/查询/删除/抢占历史四条链路的端到端;状态机映射表逐态核对;队列位置展示与 Visibility API 一致。
  • 对账:灰度期每日核对 Kueue 账本(kueue_cluster_queue_resource_usage)与平台账本,偏差告警。
  • 故障演练:kueue-controller-manager 单副本宕机(预期:存量作业不受影响,仅新准入暂停);watcher 断连重放(预期:落库无缺失,依赖条件而非事件兜底)。
  • 难点:抢占时序断言存在竞态,用条件的 lastTransitionTime 而非事件顺序做判定;slice 成型(ICI 布线)发生在准入之前,「排队中 → slice 供给中 → 启动中」的时延含布线耗时且有波动,测试断言需给足超时。

影响范围

  • AIStudio:提交路径加两个标签与高保余额校验;新增 Workload watcher(状态映射 + 抢占历史落库);状态机与用户通知文案更新(新增「被抢占,重排中」状态)。
  • 集群:新增 Kueue、JobSet controller 及其 CRD/webhook,以及托管 Slice Controller(--enable-slice-controller,含 slices.accelerator.gke.io CRD)与 Kueue Slice Provisioner(官方清单,含 JobSet 创建路径上的 mutating webhook)。注意 Kueue 与 Provisioner 的 webhook 都位于 JobSet 创建路径上,webhook 不可用会阻塞新作业提交(存量不受影响)——以多副本 HA 缓解。未打 queue-name 标签的工作负载因 manageJobsWithoutQueueName: false 不被接管,kubemaker 中混部的系统组件照常运行(见 4.1「纳管范围」)。启用 deployment 集成后 Kueue 的 pod webhook 会出现在 pod 创建路径上,但 managedJobsNamespaceSelector 把接管范围限定在纳管命名空间内,GKE 系统命名空间与其他团队的负载均不受影响(见 4.1)。
  • 各训练团队:作业 YAML 由平台代填标签,用户无感;低保作业新增被抢占的可能(现状从不会被中断),需满足 checkpoint 契约(见 4.5),上线前需宣导。
  • 兼容性/迁移:存量运行中的 JobSet(无 queue-name 标签)不被 Kueue 接管,自然跑完退场;新作业一律走 Kueue。无 CRD 数据迁移。

实施计划

阶段内容出口标准
〇(立即发起,阶段一的硬前置)sub-slicing preview 手续:向 Google 提交 Project ID / Reservation ID 与 zone / 试点子块清单 → 等待软件包 rollout → 按 Reservation API 执行 CTM 维护;集群升级至 Rapid ≥ 1.36.0-gke.3712000允许名单生效,节点出现五档几何分区标签(7 层 Topology 依赖全部层级标签,标签不齐时 TAS 无可用节点)
一(第 1–2 周)staging 安装 JobSet + Kueue + 托管 Slice Controller(--enable-slice-controller)+ Kueue Slice Provisioner(官方清单),下发 4.2 全套配置;手工复现 4.4 六场景与动态切片全链路(super 与 sub 各一条:提交 → slice 成型 → 运行 → 删除 → slice 拆除)六场景与两条切片链路全部符合预期
二(第 3–4 周)AIStudio 适配层开发与联调(打标提交、状态映射、watcher 落库、队列位置)四条交互链路 e2e 通过
三(第 5 周)生产灰度:抢占策略先置双 Never(纯记账观察模式,集群行为与现状一致);同步向低保用户宣导抢占契约连续一周对账零偏差
四(第 6 周)开启抢占:策略改为 LowerPriority 双开;按集群逐个推广首次真实回收抢占按预期完成,无误伤工单

风险与缓解

  1. 配额错配(Σnominal > 集群总量):高保保证失效。缓解:额度下发流水线强制校验 + CI 阻断;对账告警兜底。
  2. 低保作业行为变化冲击(从「从不中断」到「可能被抢占」):存量低保用户没有 checkpoint 习惯,首批被抢占的无 checkpoint 作业会从头重跑,算力浪费、引发不满。缓解:灰度期(策略 Never)提前宣导契约;提交低保作业时平台强提示;抢占历史面板暴露重跑成本;开启抢占后首周专人跟进被抢作业;「被抢即失败」选项列入下期迭代。
  3. 准入后 slice 成型失败或超时:TAS 在准入时已完成拓扑装箱(配额与放置检查合一,「配额够但放置不下」在准入前即被拦截),但准入后 ICI 布线失败或分区健康劣化仍可能让作业停在「slice 供给中」。缓解:监控 QuotaReserved 但未 Admitted 的时长并告警(Cloud Monitoring 的 slice/formation_durationsslice/state 指标可直接支撑);托管 Slice Controller 持续维护分区健康标签,TAS 下一轮装箱自动避开不健康分区。
  4. Kueue webhook/controller 不可用阻塞新提交:缓解:多副本 HA + 就绪探针告警;存量作业运行不受影响,故障窗口只影响新准入。
  5. 抢占痕迹丢失(事件 1h 过期、条件被覆盖):用户投诉无法解释。缓解:watcher 实时落库为唯一权威历史;watcher 自身监控补齐。
  6. 升级风险:Kueue 为上游原版、无 fork,随社区版本升级;API 以 v1beta2 为基线,升级前在 staging 回放 4.4 场景表。Kueue Slice Provisioner 为 GKE 发布组件,版本与清单以 GKE 侧材料为准,升级前同样在 staging 验证切片链路。节点池升级有专门操作要求:升级前先删除其上的 slice(遗留的 ACTIVE slice 会转 FAILED 并停留),且必须用 surge 策略 max-surge=0(避免超预配)、建议 max-unavailable=16(整池并行升级)。
  7. 依赖 private preview 功能(sub-slicing):非 GA、要求 Rapid 渠道 ≥ 1.36,注解取值与行为可能随 GA 调整,支持渠道非标准。缓解:与 GKE 保持沟通、staging 全量回归后再放量;主链路可随时降级——sub-slicing 手续未完成或功能异常时 super-slicing 不受影响,小任务可临时改走静态小节点池 flavor(admissionChecksStrategy.onFlavors 把 slice 检查限定在动态切片 flavor 上),配额仍归同一本账。

安全维度无新增暴露面:AIStudio 仅增加只读 RBAC;Kueue webhook 证书用其内置轮换机制。性能维度:Kueue 准入延迟为秒级,相对训练任务生命周期可忽略,不需专项评估。

原有作业的改造说明

本节面向已经在集群里跑作业、需要迁移到 Kueue 管理的团队(如 falcon),可整节转发。内容只覆盖提交侧的改动,不涉及集群配置。

改动只在 pod 模板上:加一个注解、删两处旧写法。两个 kueue.x-k8s.io/ 标签都不用填——集群侧已配好默认队列与默认优先级。

需要的改动

yaml
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: <作业名>
  namespace: falcon-jobs             # 无需任何 kueue.x-k8s.io/ 标签
spec:
  replicatedJobs:
  - name: workers
    replicas: 2                      # slice 个数
    template:
      spec:
        completions: 16              # 每个 slice 的主机数,与拓扑对应
        parallelism: 16              # 必须与 completions 相同
        template:
          metadata:
            annotations:
              # ★ 新增:声明这个 slice 要多大
              cloud.google.com/gke-tpu-slice-topology: 4x4x4
              # ✗ 删除:JobSet 自带的独占放置,职责已由 Kueue TAS 承担;
              #   且「拓扑域不与别的 Job 共享」与 sub-slicing 的共享 cube 冲突,
              #   留着会让小任务永远排不进去
              # alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
          spec:
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu7x
              # ✗ 删除:动态切片下 slice 按需切出,节点不再绑定固定拓扑,
              #   要多大由上面的注解表达
              # cloud.google.com/gke-tpu-topology: 4x4x4
            containers:
            - name: <容器名>          # 镜像与启动命令保持现有的不变
              resources:
                limits:
                  google.com/tpu: '4'          # 每主机 4 芯,固定值
                requests:
                  ephemeral-storage: 20Gi      # TPU 之外的资源如实声明
                  google.com/tpu: '4'

字段说明

字段怎么填弄错会怎样
cloud.google.com/gke-tpu-slice-topology 注解slice 拓扑,如 4x4x4 / 2x2x2漏了会报「配额不足」,但真实原因是缺注解,提示里完全看不出来
alpha.jobset.sigs.k8s.io/exclusive-topology 注解删除留着与 sub-slicing 的共享 cube 冲突,小任务排不进去
nodeSelector 加速器项cloud.google.com/gke-tpu-accelerator: tpu7x漏了则该 podSet 对节点类型不作限定,可能被排到非 TPU 节点
nodeSelectorgke-tpu-topology删除留着会把作业限制在一个动态切片下不再有意义的标签上
requests 里 TPU 之外的资源(至少 ephemeral-storage)如实声明若声明了某项而 ClusterQueue 未覆盖它,报 resource <名字> unavailable in ClusterQueue,整个作业不可准入
两个 kueue.x-k8s.io/ 标签不用填集群侧已配默认队列与默认优先级;填了也可以,但值必须与集群侧一致

completions / parallelism 与拓扑的换算(每主机 4 芯):4x4x4 = 64 芯 = 16 台主机;2x2x2 = 8 芯 = 2 台主机;2x2x1 = 4 芯 = 1 台主机。两个字段必须相等且等于主机数。

不要做的事

  • 不要设 spec.suspend。Kueue 会自动挂起新作业、准入后自动解挂,手工翻转会打乱状态机。
  • 不要给 pod 设 priorityClassName(K8s 原生优先级)。那是另一套机制,会触发 kubelet 层驱逐,与 Kueue 的优先级无关。开发档作业的优先级由集群侧的默认档统一给定,作业侧不需要表达。
  • 不用写 TPU 的 toleration。集群侧已在 ResourceFlavor 上统一声明并由 Kueue 注入,自带也兼容(取并集)。
  • 不用写 kueue.x-k8s.io/podset-* 开头的注解。那些由 Kueue Slice Provisioner 的 webhook 从 gke-tpu-slice-topology 自动翻译生成。

行为变化

作业提交后不会立刻起 pod——先排队,拿到配额才启动,kubectl get jobset 会看到一段 suspended 状态,属正常现象。

开发档作业优先级最低,生产作业需要容量时会被中断:pod 收到 SIGTERM 后被删,作业自动回队,容量空出后自动重新拉起。因此:

  • 长跑任务需处理 SIGTERM 并写 checkpoint,启动时从最新 checkpoint 恢复;
  • checkpoint 必须写远端存储(GCS 等),重排后可能落在另一组硬件上;
  • 被抢占不计作失败,不消耗 failurePolicy.maxRestarts 的重启预算。

怎么看状态

排队、准入与被抢占的原因都记录在 Workload 对象上,不在 JobSet 上:

bash
kubectl get workload -n falcon-jobs
kubectl describe workload -n falcon-jobs <workload>

常见状态:QuotaReserved=False = 仍在排队;Admitted=True = 已启动;Evicted 且 reason 为 Preempted = 被抢占,正在自动重排。