通过一个 Kueue 测试清理，看懂 TAS 的 Rank 与 Greedy Assignment

在运行 GPU 或 TPU 分布式训练时，我们通常希望同一组 Pod 被放到物理位置更接近的机器上，例如同一个机架或同一个数据中心区块。节点之间越接近，网络延迟通常越低，训练效率也越高。Kueue¹ 提供的 Topology Aware Scheduling（TAS）正是为了解决这个问题：它会在工作负载进入准入阶段时，基于拓扑信息提前计算出一份“每个 Pod 应该放在哪里”的放置计划。

但这里有一个很容易被忽略的问题：这份放置计划是在工作真正运行之前算好的。等到 Pod 实际创建或替换时，集群的现场状态可能已经变了。节点可能故障、Pod 可能被驱逐、机器也可能被替换。此时，原始计划就可能与真实运行状态脱节。这正是 TAS 中最微妙、也最关键的一环。

我最近在参与 Kueue 贡献时，提交了一个很小的测试清理 PR：#9796²。表面上看，它只是把测试里不必要的 JobSet 相关代码移除了；但在整理过程中，我发现这个测试非常适合用来解释 TAS 的三个核心概念：

• Rank：同一工作负载内每个 Pod 的稳定编号
• TopologyAssignment：工作开始前计算出的放置计划
• Greedy Assignment：当原始计划与运行时现实不一致时启用的回退机制

本文就从这个实际测试案例出发，一步一步说明 Kueue TAS 的设计思路。

Kueue Topology Aware Scheduling（TAS）概览

Kueue 的 Topology Aware Scheduling（TAS）是一种进阶调度能力。它不只关心“这个 workload 需要多少 CPU / GPU / TPU”，也关心“这些 Pod 应该放在拓扑结构中的哪个位置”。

根据 Kueue 官方文档³，TAS 通过节点上的 topology labels 来表示数据中心中的层级结构，例如：

• cloud.provider.com/topology-block：数据中心区块
• cloud.provider.com/topology-rack：机架
• kubernetes.io/hostname：主机名，最低层级

在 workload 进入 admission 阶段时，Kueue 会根据这些拓扑标签计算出一个 TopologyAssignment。你可以把它理解为一份“Pod 到拓扑域的放置计划表”。

Topology Tree：Kueue 的拓扑树

Kueue 会结合 ResourceFlavor、Topology 配置与 Nodes，把集群中的节点组织成一棵拓扑树。树中的每个节点都代表一个 topology domain，每一层对应一个物理或逻辑边界。一个典型的四层拓扑结构如下：

各层对应的 Kubernetes Label 如下：

在放置 Pod 时，TAS 会根据 workload 指定的 requiredTopologyLevel，尽可能把同一组 Pod 放在更紧凑的 topology domain 中。比如当最低层级要求是 rack 时，TAS 会优先尝试把所有 Pod 放进同一个机架；如果单个机架容量不足，再向上扩展到同一个 block 下的多个 rack。

Rank：工作负载中 Pod 的稳定编号

在 TAS 的上下文里，rank 很容易被误解。它不是优先级，也不是 NCCL 或 MPI 里的分布式训练 rank。这里的 rank 含义很简单：

• 同一工作负载内，每个 Pod 的稳定编号。

对于带编号的 Indexed Job，Kubernetes 会通过 batch.kubernetes.io/job-completion-index 给每个 Pod 分配从 0 到 .spec.completions - 1 的编号，并以 annotation 与 label 的形式暴露出来⁴。例如：

workload
└─ Indexed Job
   ├─ pod A (job-0) -> rank 0
   └─ pod B (job-1) -> rank 1

TopologyAssignment：准入阶段计算出的计划

当 Kueue 完成 admission 计算后，会生成一份 TopologyAssignment，把每个 rank 映射到具体的拓扑域。假设最低拓扑层级是 kubernetes.io/hostname，一个包含两个 Pod 的 workload 可能得到如下结果：

这份 assignment 本质上是一份“计划”。它在计算完成的当下是正确的，但运行过程中一旦发生节点故障、Pod 驱逐等事件，现场状态就可能偏离这份计划。

问题场景：Replacement Pod 与 Running Pod 发生冲突

问题是怎么出现的？

这个问题最早是在 Issue #9210⁵ 中被发现的。场景如下：

假设某个 workload 已经在运行，Kueue 在 admission 时已经为它规划好了两个 Pod 分别要跑在哪台机器上。过一段时间后，其中一台机器故障，需要进行节点替换。此时现场状态会变成：

• Pod A：原本跑在故障节点上，现在正在终止（terminating）
• Pod A 的替身：一个新的 replacement Pod 已经被创建出来，正在等待 ungate
• Pod B：继续在别的地方运行，没有直接受影响

问题出在 controller 决定“Pod A 的替身该放到哪里”时。如果它只看 admission 阶段生成的原始 TopologyAssignment，却没有先确认“当前每台机器上到底有哪些 Pod 正在运行”，它就可能把 replacement Pod 放到一个已经有别的运行中 Pod 的节点上。这样就会让两个 Pod 撞到同一个节点，违反 TAS 原本想保证的放置正确性。

下面这张图展示了这个问题的三个阶段：

三个阶段分别是：

1. Step 1：admission 完成后，Pod A 和 Pod B 分别运行在各自计划好的节点上
2. Step 2：一段时间后，node-x3 出现故障。Pod A 正在终止，Pod B 已经实际运行在 node-x3 上，而 node-x1 变成空闲。同时，一个 replacement Pod 正在等待被 ungate
3. Step 3：如果 controller 死守旧的 assignment（rank 0 -> node-x3），它也会把 replacement Pod 放到 node-x3，从而与已经运行中的 Pod 发生冲突

测试如何精确模拟这个场景

测试 should fallback to greedy assignment when replacement pod conflicts with already running pod 精确复现了这一情况²。

原始的 TopologyAssignment 是：

rank 0 -> node-x3
rank 1 -> node-x1

但测试刻意构造了一个与原始 assignment 不一致的运行时状态：

node-x3: rank 1 (running)
node-x1: empty

也就是说，rank 1 的 Pod 实际已经跑在 node-x3 上，而这原本是 rank 0 的计划位置。如果 controller 仍坚持把 rank 0 的 replacement Pod 调度到 node-x3，就会与正在运行的 rank 1 Pod 冲突。

Kueue 原始计划中的预期放置
┌────────┬────────────┬──────────────┐
│ rank   │ pod        │ node         │
├────────┼────────────┼──────────────┤
│ 0      │ p0         │ x3           │
│ 1      │ p1         │ x1           │
└────────┴────────────┴──────────────┘

实际运行时状态（p1-running 已经跑到 x3）
┌────────┬──────────────┬──────────────┐
│ rank   │ pod          │ node         │
├────────┼──────────────┼──────────────┤
│ 1      │ p1-running   │ x3           │
└────────┴──────────────┴──────────────┘

Greedy assignment 的结果（新的 rank 0 Pod 改放到 x1）
┌────────┬────────────────┬──────────────┐
│ rank   │ pod            │ node         │
├────────┼────────────────┼──────────────┤
│ 1      │ p1-running     │ x3           │
│ 0      │ p0-replacement │ x1           │
└────────┴────────────────┴──────────────┘

这并不是人为捏造出来的边界场景，而是一个真实的生产风险：admission 时算出的 topology assignment 只是计划，运行时的现场状态却可能已经不是原来的样子。

Greedy Assignment：运行时的安全回退机制

当原始放置计划与真实运行状态不一致时，TAS 的 topology ungater 不会生硬地套用旧 mapping，而是会自动回退到 greedy assignment，基于当前运行时状态重新计算放置结果。

Greedy assignment 并不是“随便找个地方塞进去”的粗放策略。它背后有明确的设计目标：

1. 不能发生碰撞：绝不把 replacement Pod 放到已经有同一 workload 其他运行中 Pod 的节点上
2. 先看实时状态再决定：如果旧 assignment 已经过时，就根据当前真正空闲的位置来选节点
3. 优先保证放置正确性：宁可选择不同但合法的位置，也不要强行遵守过时映射导致错误放置

TopologyAssignment 与 Greedy Fallback 流程

下面这张图展示了 TAS 从 admission 规划到运行时回退的完整流程：

它说明了两个关键点：

1. Topology Tree：Kueue 会基于 topology labels，把集群节点组织成 Zone -> Block -> Rack -> Node 的层级结构
2. Admission -> Runtime -> Greedy Fallback：一旦 admission 计划与 runtime 现场状态出现偏差，controller 就会回退到 greedy assignment，以保证放置正确性

最终正确的 Placement

在这个测试场景中，greedy assignment 给出的最终结果是：

p1-running      (rank 1) -> node-x3
p0-replacement  (rank 0) -> node-x1

重点不在于“rank 0 一定要不惜代价留在 node-x3 上”，而在于 controller 不能死守过时的 rank-to-node 映射，最终把两个 Pod 放到同一个节点上。

从这个角度看，greedy assignment 不是退化（downgrade），而是一层安全网：当基于 rank 的 assignment 已经和运行时现实脱节时，它用来保护 placement correctness。

为什么这个测试要从 JobSet 清理成 Job？

原始测试的问题在哪里？

这个测试原本用了几处 JobSet 特有的 subgroup labels。但它真正要验证的事情其实很简单：当 Pod 的 rank 对应节点和实际运行状态不一致时，TAS 是否会正确回退到 greedy assignment？

而这个 rank 身份，本来就可以直接用 Kubernetes 原生 Indexed Job 的 batch.kubernetes.io/job-completion-index 表达，不需要依赖 JobSet。

另外还有一个很现实的原因：集成测试环境通常并不会安装 JobSet CRD，这一点在 Issue #9284⁶ 中也有提到。把 JobSet 相关细节留在测试里，容易让读者误以为测试在验证 JobSet 集成行为，但实际上核心逻辑与 JobSet 无关。

PR #9796 做了什么？

PR #9796² 做的是一组小而明确的清理：

• 删除测试里 JobSet 特有的 labels 和配置
• 使用 batch.kubernetes.io/job-completion-index 表达 Pod 的 rank
• 让整个测试只依赖原生 Job 语义，不再与 JobSet 耦合

清理之后，测试表达的重点就非常清楚了：当基于 rank 的放置与真实运行状态冲突时，TAS 能否正确切换到 greedy assignment？现在这个测试只验证这件事，没有额外噪音。

结语

这个测试清理很好地体现了 Kueue TAS 在抽象设计与容错机制之间的平衡。TopologyAssignment 负责给出理想的放置计划，greedy assignment 则负责在运行时处理现实世界中的状态漂移。

分布式调度真正困难的地方，不是只在某个时刻算出一份“完美计划”，而是在生产环境不可避免变化时，依然保证放置正确性。TAS 通过把“计划中的放置”和“运行时的回退机制”分开处理，解决了这个问题。

希望这篇文章能帮助你更清晰地理解 TAS 的设计思路，特别是 rank、TopologyAssignment 与 greedy assignment 三者是如何配合工作的，也帮助你看懂为什么通过简化测试，可以更准确地验证真正的核心机制。

参考资料