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 模式)。它当前的资源模型:
- 各团队按集群 TPU 总量分配额度,分高保与低保两档。高保额度之和 ≤ 集群总量;高保 + 低保总额度可略超集群总量。
- 每个团队有一个平台内部队列,用户作业先在团队队列排队,再由平台提交 JobSet 到 K8s 集群。
- 平台只做高保/低保标记,作业提交到集群后不再做任何容量仲裁——集群内不存在抢占或回收机制。
这套模式存在几个结构性问题:
- 高保保证无法兑现:当集群被低保作业占满时,没有任何组件为新提交的高保作业腾出容量,「高保」只是一个标记,兑现依赖人工协调。
- 没有 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 抢占)。
方案
总体职责划分:
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 controller | kubernetes-sigs/jobset,≥ v0.12.0(SubSlicing 指南基线) | 官方 manifest / helm 安装,默认配置即可、无需特殊开关;确认版本与 Kueue 兼容矩阵匹配 |
| Kueue | kubernetes-sigs/kueue,v0.19.2(SubSlicing 指南基线为 ≥ v0.18.2;API v1beta2) | 官方 manifest / helm 安装,见下方 Configuration;controller 多副本 HA(GKE 指引示例为 3 副本,leader election 默认开启);证书用内置轮换,无需 cert-manager |
| 托管 Slice Controller | GKE 控制面托管 | 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 Provisioner | GKE 发布,安装清单随 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 Monitoring | Prometheus 抓取 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 下发),必须明确的开关:
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 | 一个 JobSet | JobSet 的 metadata.labels | 全部 replicatedJob 的全部 pod |
batch/job | 一个 Job | Job 的 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-system 与 kueue-system,过宽):对 pod 类集成它是强豁免——不匹配的命名空间即使打了 queue-name 也永不被接管;对 JobSet、Job 则是弱豁免,带 queue-name 的作业在任何命名空间都会被接管,真正限定它们投递范围的是各 ClusterQueue 的 namespaceSelector(见 4.2)。
纳管范围
两个字段共同决定接管什么:managedJobsNamespaceSelector 圈定命名空间,manageJobsWithoutQueueName: false 表示其中只接管带 queue-name 标签的负载。落到两档上:
| 档位 | 有名为 default 的 LocalQueue? | 未打 queue-name 的负载 |
|---|---|---|
dev(default、falcon-jobs) | 有 | webhook 自动补标,正常排队(见 4.2「免标签投递」) |
production(kubemaker) | 无 | 不被接管、照常运行——该命名空间混部着团队的系统组件,不应纳入配额与排队 |
下发方式:上述 Configuration 不是独立的 CR,而是 kueue-system 命名空间下 kueue-manager-config ConfigMap 中 controller_manager_config.yaml 这一个 data 条目。操作命令:
# 路径 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=50Kueue 以严格模式解码该配置:出现未知字段(例如 v1beta1 才有、v1beta2 已移除的
fairSharing.enable)会导致 controller 启动失败而非被忽略。生产变更前先在 staging 验证 rollout 成功;配置应纳入 GitOps 仓库,与 4.2 的队列对象同源管理。
4.2 Kueue 资源配置(全局一次性 + 每团队增量)
命名空间用 scheduling.antgroup.com/kueue-tier 标签分为两档——生产档(production)与开发档(dev),两档各有独立的 ClusterQueue,互不串号(见下方「命名空间分档与 namespaceSelector」)。Kueue 对象分两类:全局共享、只建一次的,和每新增一个团队/业务线需要增量创建的。
全局共享对象(集群初始化时创建一次)——两档优先级 + 分区树 + 资源风味 + slice 准入检查:
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 为例:
# ──────────── 团队 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 有两条不可绕开的规则:
- 一个资源只能属于一个资源组(
clusterqueue_webhook.go的重复校验); - 一个 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 上,准入直接失败:
couldn't assign flavors to pod set worker: Flavor "cpu-standard" does not support TopologyAwareScheduling因此四种资源必须同组,让 TPU podSet 的全部资源都从同一个 TAS flavor 取得。作为代价,组内每个 flavor 都要声明全部四种资源(cpu-standard 的 google.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-standard | google.com/tpu | 必须为 0 | 纯 CPU 节点池不提供 TPU;填非 0 会凭空造出不存在的 TPU 配额 |
cpu-standard | cpu / 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 逐个点名,而是用同一个标签键、不同的取值给命名空间分档:
# 生产档: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| 档位 | 命名空间 | ClusterQueue | LocalQueue | 可用优先级 |
|---|---|---|---|---|
production | kubemaker | 每团队一个(team-a / team-b / …) | <team>-queue,建在 kubemaker | guaranteed / best-effort |
dev | default、falcon-jobs | 共用一个 dev | default,每个命名空间各建一个 | 默认档(值 10),无需打标签 |
这个设计的三个要点:
- 标签键的存在性与取值分工:
managedJobsNamespaceSelector(4.1)用Exists判定「归不归 Kueue 管」,各 ClusterQueue 的namespaceSelector用取值判定「进哪一档」。新增命名空间只需打标签,队列定义不动。 - 两档同属
tpu-poolcohort,开发档因此能借用生产闲置的容量,不必为自测预留一块常年空转的硬件。安全性由优先级保证: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,写法见本节末:
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-topology | webhook 无从翻译,该 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,写法各不相同——
| podSet | slice 拓扑注解 | gke-tpu-accelerator nodeSelector | TPU 请求 | 落到的 flavor |
|---|---|---|---|---|
pathways-worker | 要(同上例) | 要 | 要 | tpu7x |
pathways-head(纯 CPU) | 不要 | 不要(改用 gke-nodepool 定向到 CPU 池) | 不要,只声明 cpu/memory/ephemeral-storage | cpu-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/job | deployment | |
|---|---|---|
| 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>),不要拼名字,按标签反查:
// 反查: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.priorityClassName与spec.priorityClassSource(v1beta1)在 v1beta2 中已合并为结构化的spec.priorityClassRef(含group/kind/name),解析出的数值仍在spec.priority。按旧字段名读取不会报错,只会静默取到空值——AIStudio 的 watcher、运维脚本与告警规则都要按新字段名编写。判断作业档位建议直接用spec.priority的数值(1000 = 高保 / 100 = 低保),比字符串匹配更稳。
Workload 对象示例一——已准入、运行中(节选,只保留状态判定相关字段):
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 节选):
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 直连)。响应示例(节选):
# 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):
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 增量
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:抢占前,抢占方份额(含新作业)< 受害方当前份额,才允许抢。
对抢占结果的具体影响:
- 跨队受害者从「谁最近准入谁先死」变为「超出公平份额最多的团队先还」。默认排序不看各队借了多少,借 8 芯的队可能替借 100 芯的队挨刀;Fair Sharing 下每一刀都落在当前份额最高的队列上。
- 回收止于份额拉平:不会把某一个团队一撸到底——某队份额被抽到不再最高后,下一刀落到新的份额最高者,直到容量凑够。
- 借用侧同样变公平:多个团队的低保同时排队等闲置时,按份额轮流准入(份额低者优先),不再是先到先得。
- 受害者的优先级门槛不变:受害者仍须严格更低优且其队列在借用中,任何队列 nominal 内的作业不可侵犯——两档模型下低保永不引发牺牲。
- 触发资格放宽:经典语义下借用中的作业从不抢占;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):
| # | 场景 | 预期结果 |
|---|---|---|
| 1 | A 用 128;B 提交 384 芯低保 | 全部准入(B 借走 A 闲置的 128) |
| 2 | 接 1;A 提交 128 芯高保 | B 的低保借用者被驱逐(队内按 LIFO 挑),凑够 128 即止;被驱逐者回队 |
| 3 | A 已满 256 高保,再提 100 芯高保(需借用) | 不动任何队 nominal 内的作业;若他队存在借用中的低保且 FS 份额策略满足,以 InCohortFairSharing 抢占之,否则排队 |
| 4 | A 队内 156 高保 + 100 低保;A 提 100 芯高保 | 队内抢占 A 自己的低保(低保防外不防内) |
| 5 | A 借用中;B 提 100 芯低保(B nominal 有空) | 不抢占,排队(同优先级不抢,低保从不引发牺牲) |
| 6 | A 借 96、B 借 24(均低保);C 提 64 芯高保 | Fair Sharing 生效:受害者全部出自份额最高的 A,B 不受影响 |
4.5 被抢占后的行为:只重试(本期唯一语义)
完整链路:
Kueue 选中受害者
→ JobSet 被置 suspend=true → pod 收到 SIGTERM(宽限期内写 checkpoint)→ pod 删除
→ Slice CR 删除,slice 拆除,分区归还容量池
→ Workload 回到所属队列(Requeued)
→ 配额空出 → 重新准入:TAS 重新装箱选块 → slice 重新成型激活(AdmissionCheck)
→ JobSet 解挂 → 续跑(可能落在另一组分区上)两个必须写进用户文档的事实:
- 抢占不是失败:因 suspend 被删的 pod 不计入 Job 失败,JobSet 的
failurePolicy/maxRestarts对抢占完全不生效、不消耗重启预算。maxRestarts只需按真实故障(进程崩溃、硬件失败)的容忍度设置。 - 续跑可能落在不同硬件上,且需要重新组 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 CR | 7 层分区树(block → 五档几何 → 主机),必需 | 保留但简化为 2 层:cloud.google.com/gke-nodepool(= 一个静态 slice)→ kubernetes.io/hostname |
ResourceFlavor | topologyName: 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 + admissionChecksStrategy | slice 成型门闩,必需 | 不需要(两处一并删除);准入通过即可直接解挂 |
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) | 改为校验「所提拓扑存在对应静态节点池」 |
需要接受的能力损失(即引入动态切片的收益反面):
- 拓扑在建池时固定:每种拓扑需预先规划节点池,任务形态变化要重建节点池(分钟至数十分钟级),不能按作业需求现拼;
- 碎片化不可回收:TAS 仍会在现有节点中做装箱,但无法跨节点池重组硬件——池内剩余不足一个作业所需拓扑时,该容量只能空置;
- 故障恢复更慢:TAS 的 hot swap 仅在同池内有空闲节点时才换得成,换不到就得等节点池重建(而非动态切片的软件级重拼 slice,GKE 侧数据为恢复快约 5 倍),且没有 fast replacement;
- 无 sub-slice 能力:小任务只能独占一个完整拓扑的节点池,或另建小规格静态池;
- 配额按拓扑切分:多种静态拓扑时每种一个 flavor、各自独立配额,不再是一本可自由流动的账(见下方 ② 的说明)。
完整 YAML 示例
① 全局共享对象——与 4.2 相比:Topology 从 7 层简化为 2 层、删除 AdmissionCheck,ResourceFlavor 保留 topologyName 但 nodeLabels 改为匹配静态池的固定拓扑;三个 WorkloadPriorityClass 与 cpu-standard 逐字不变。
由于 ResourceFlavor.nodeLabels 是精确匹配的键值对,一个 flavor 只能代表一种拓扑——集群里有几种静态节点池,就要建几个 flavor。下例覆盖四档拓扑:
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 的节点池上。
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配套的三条规则:
- 作业按 nodeSelector 自动匹配 flavor:声明
gke-tpu-topology: 4x4x8的作业只会与tpu7x-4x4x8flavor 匹配,占用该 flavor 的额度;无需在提交侧指定 flavor 名。 - 配额校验改为按 flavor 分别进行:原「Σ 各队 nominalQuota ≤ 集群总芯数」变为每个 flavor 各校验一次——
Σ 各队在 4x4x8 flavor 的 nominalQuota ≤ 全部 4x4x8 节点池的总芯数,其余拓扑同理。额度下发流水线需相应改造(见测试计划)。另外每档额度应为该拓扑单池芯数的整数倍(作业以整池为单位占用 slice),否则尾数无法被任何作业使用。 - 借用与抢占同样限定在 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 落在不带topologyName的cpu-standard上,无需注解。
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. Topology | 2 层(gke-nodepool → hostname)替换为 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.ioCRD)与 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 双开;按集群逐个推广 | 首次真实回收抢占按预期完成,无误伤工单 |
风险与缓解
- 配额错配(Σnominal > 集群总量):高保保证失效。缓解:额度下发流水线强制校验 + CI 阻断;对账告警兜底。
- 低保作业行为变化冲击(从「从不中断」到「可能被抢占」):存量低保用户没有 checkpoint 习惯,首批被抢占的无 checkpoint 作业会从头重跑,算力浪费、引发不满。缓解:灰度期(策略
Never)提前宣导契约;提交低保作业时平台强提示;抢占历史面板暴露重跑成本;开启抢占后首周专人跟进被抢作业;「被抢即失败」选项列入下期迭代。 - 准入后 slice 成型失败或超时:TAS 在准入时已完成拓扑装箱(配额与放置检查合一,「配额够但放置不下」在准入前即被拦截),但准入后 ICI 布线失败或分区健康劣化仍可能让作业停在「slice 供给中」。缓解:监控
QuotaReserved但未Admitted的时长并告警(Cloud Monitoring 的slice/formation_durations与slice/state指标可直接支撑);托管 Slice Controller 持续维护分区健康标签,TAS 下一轮装箱自动避开不健康分区。 - Kueue webhook/controller 不可用阻塞新提交:缓解:多副本 HA + 就绪探针告警;存量作业运行不受影响,故障窗口只影响新准入。
- 抢占痕迹丢失(事件 1h 过期、条件被覆盖):用户投诉无法解释。缓解:watcher 实时落库为唯一权威历史;watcher 自身监控补齐。
- 升级风险:Kueue 为上游原版、无 fork,随社区版本升级;API 以
v1beta2为基线,升级前在 staging 回放 4.4 场景表。Kueue Slice Provisioner 为 GKE 发布组件,版本与清单以 GKE 侧材料为准,升级前同样在 staging 验证切片链路。节点池升级有专门操作要求:升级前先删除其上的 slice(遗留的 ACTIVE slice 会转FAILED并停留),且必须用 surge 策略max-surge=0(避免超预配)、建议max-unavailable=16(整池并行升级)。 - 依赖 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/ 标签都不用填——集群侧已配好默认队列与默认优先级。
需要的改动
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 节点 |
nodeSelector 的 gke-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 上:
kubectl get workload -n falcon-jobs
kubectl describe workload -n falcon-jobs <workload 名>常见状态:QuotaReserved=False = 仍在排队;Admitted=True = 已启动;Evicted 且 reason 为 Preempted = 被抢占,正在自动重排。