一個 Kueue 測試 cleanup：看懂 TAS 的 rank 與 greedy assignment

在跑 GPU 或 TPU 的分散式訓練時，我們常常希望同一組 Pod 能被放在物理位置比較近的機器上（例如同一個機架或同一個資料中心區塊），因為節點之間距離越近，網路延遲越低，訓練效率就越好。Kueue¹ 提供的 Topology Aware Scheduling（TAS）機制，就是在解決這個問題：它會在 workload 進入佇列時，根據節點的拓撲位置（哪個區塊、哪個機架、哪台主機）提前算好一份「每個 Pod 要放在哪裡」的位置安排表。

但有一個容易被忽略的問題：這份安排表是在「工作開始前」算好的，但等工作真正跑起來之後，現場狀態可能已經變了——機器壞了、Pod 被刷掉了、節點被替換了——原本的安排表就跟實際狀態對不上了。這正是 TAS 設計中最容易踩到坑、但也最關鍵的環節。

最近我在參與 Kueue 專案貢獻時，提交了一個測試 cleanup PR（#9796²）。表面上看只是清掉測試裡不必要的 JobSet 相關程式碼，但在整理的過程中，我發現這個測試案例非常適合用來解釋 TAS 的三個核心概念：

• Rank：Pod 在同一組工作中的固定編號
• TopologyAssignment：工作開始前算好的位置安排表
• Greedy Assignment：當安排表跟現實對不上時的備援機制

本文會從這個實際的測試案例出發，一步一步拆解這些概念，幫你建立對 Kueue TAS 調度機制的完整理解。

Kueue Topology Aware Scheduling (TAS) 概述

Kueue 的 Topology Aware Scheduling（TAS）是一種進階的調度機制，它不只關心「這個 workload 需要多少 CPU / GPU」，還會進一步考慮「這些 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 Tree）。樹中的每個節點稱為一個 topology domain，從上到下代表不同的物理或邏輯階層。以下是一個典型的四層拓撲結構：

每一層對應的 Kubernetes label 如下：

TAS 機制在放置 Pod 時會根據 workload 指定的最低拓撲層級（requiredTopologyLevel），嘗試將同一組 Pod 盡可能放置在同一個 topology domain。例如，當指定最低層級為 rack 時，TAS 會優先將所有 Pod 排入同一個 Rack；如果單一 Rack 容量不足，則向上擴展到同一個 Block 中的多個 Rack。

Rank：Pod 在同一組工作中的固定編號

在 TAS 的語境中，rank 是一個容易被誤解的詞彙。它不是優先權（priority），也不是 NCCL / MPI 中的分散式訓練 rank。在這裡，rank 的含義非常單純：

• 同一組 Pod 中，每個 Pod 的固定編號。

對於有編號的 Indexed Job，Kubernetes 會透過 batch.kubernetes.io/job-completion-index 為每個 Pod 分配一個從 0 到 .spec.completions - 1 的 index，並以 annotation 與 label 的形式附加在 Pod 上（Kubernetes 官方文件⁴），例如：

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

TopologyAssignment：Admission 時的放置計畫

當 Kueue 完成 admission 計算後，會產生一份 TopologyAssignment，將每個 rank 對應到具體的拓撲節點。例如，假設最低層級為 kubernetes.io/hostname ，則一個包含 2 個 Pod 的 workload 會產生：

這份 assignment 本質上是一份「計畫」。它在完成計算當下是正確的，但隨著執行過程狀態的變化（例如 node failure、Pod eviction），現場狀態可能會與這份計畫產生分歧。

問題場景：Replacement Pod 與 Running Pod 的衝突

問題的起源

這個問題最早在 Issue #9210⁵ 中被發現。場景如下：

想像一個已經在跑的工作（workload），Kueue 在 admission 時已經幫它安排好了兩個 Pod 分別要跑在哪台機器上。但過了一段時間，其中一台機器出了問題，需要替換節點（node replacement）。這時候場上的狀態變成：

• Pod A：原本跑在壞掉的節點上，現在正在被終止（terminating）
• Pod A 的替身：一個全新的 replacement Pod 已經建立，正在排隊等著被放行（等待 ungate）
• Pod B：一直好好地跑在另一台機器上，完全沒受影響

這時候問題來了：controller 在決定「Pod A 的替身要放去哪裡」時，如果只看當初 admission 時算好的那份位置安排表，而沒有先確認「現在每台機器上到底有誰在跑」，就可能把替身 Pod 放到一台已經有其他 Pod 在運行的機器上，兩個 Pod 就會撞在同一台節點，違反了 TAS 原本的拓撲放置規則。

以下圖表具體呈現了這個問題場景的三個階段：

圖中的三個階段說明：

1. Step 1：Admission 完成後，Pod A 跟 Pod B 各自在自己的節點上運行，一切正常
2. Step 2：過了一段時間，node-x3 發生故障。Pod A 被終止，Pod B 因某些原因已經跑到了 node-x3（runtime 狀態與原始 assignment 不同），而 node-x1 變成空的。同時，一個新的 replacement Pod 正在等待被放行
3. Step 3：如果 controller 死守原始的 TopologyAssignment（rank 0 → node-x3），就會把 replacement Pod 也放到 node-x3，造成兩個 Pod 撞在同一台節點的衝突

測試案例的具體模擬

對應的測試場景（should fallback to greedy assignment when replacement pod conflicts with already running pod）精確地模擬了這個問題（PR #9796²）。

原始的 TopologyAssignment：

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

但測試刻意建構了一個與原始 assignment 不一致的 runtime 狀態：

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

也就是說，rank 1 的 Pod 實際上已經跑在 node-x3（原本是 rank 0 的預定位置）。當 rank 0 的 replacement Pod 進入調度時，如果 controller 仍然依照原始 assignment 將它放到 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 (新的 rank0 調度到 x1)
┌────────┬────────────┬──────────────┐
│ rank   │ pod        │ node         │
├────────┼────────────┼──────────────┤
│ 1      │ p1-running │ x3           │
│ 0      │ p0-repalce │ x1           │
└────────┴────────────┴──────────────┘

這個場景不是人為捏造的 edge case。它反映的是一個真實的生產環境風險，admission 時算出的 topology assignment 是一份計畫，但現場狀態可能已經不再等於那份計畫。

Greedy Assignment：Fallback 安全機制

當原始的位置安排表跟實際狀態對不上時，TAS 的 topology ungater（負責放行 Pod 的元件）不會硬套舊的安排，而是會自動切換到 greedy assignment，一種根據「現在場上誰在哪裡」來重新分配位置的備援機制。

Greedy assignment 並不是「随便找個地方塞」的粗糙做法。它背後有明確的設計目標：

1. 不能撞在一起：絕對不把替身 Pod 放到已經有其他 Pod 在跑的節點上
2. 先確認現場狀態再決定：舁安排表已經過時，就根據「現在哪台機器是空的」來選位置
3. 保護系統穩定：寧可換一個位置，也不要因為死守過時的安排而把 Pod 放錯地方

TopologyAssignment 與 Greedy Fallback 流程

以下圖表呈現了 TAS 從 admission 到 runtime fallback 的完整流程：

這兩個圖表說明了：

1. Topology Tree：Kueue 將叢集節點組織成 Zone → Block → Rack → Node 的樹狀階層，每一層對應不同的 topology label
2. Admission → Runtime → Greedy Fallback：當 admission 計畫與 runtime 狀態不一致時，controller 會 fallback 到 greedy assignment 以保護 placement correctness

最終的正確 Placement

在上述測試場景中，greedy assignment 產生的最終結果為：

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

重點不在於「維持 rank 0 一定要在 node-x3」，而在於不能因為死守過時的 mapping，導致兩個 Pod 被放置在同一個節點上，違反拓撲放置的正確性。

從這個角度來看，greedy assignment 不是 downgrade，而是一層 safety net——當 rank-based assignment 與現場狀態脫鉤時，用來保護 placement correctness 的安全機制。

測試 Cleanup：為什麼要從 JobSet 改成 Job？

原本的測試寫法有什麼問題？

前面提到的測試案例，原本在程式碼裡用了一些 JobSet 特有的標籤（subgroup labels）。但實際上，這個測試真正在驗證的是「rank 編號對應的節點與實際狀態衝突時，系統會不會正確 fallback」這件事。而這個 rank 編號，用 Kubernetes 原生的 Indexed Job（batch.kubernetes.io/job-completion-index）就能實現，而毋需使用到 JobSet。

此外，整合測試環境通常沒有安裝 JobSet CRD(Issue #9284⁶)。這容易讓讀者誤以為測試在驗證 JobSet 的整合行為，但其實這部分的測試跟可以與 JobSet 解耦。

PR #9796 做了什麼？

PR #9796² 的改動很簡單，但意義明確：

• 把測試裡所有 JobSet 特有的標籤和配置清掉
• 改用 batch.kubernetes.io/job-completion-index 來表達 Pod 的 rank 編號
• 讓測試從頭到尾都是純粹的 Job semantics，不再有任何 JobSet 的痕跡

這個 cleanup 主要是讓測試的描述更直白，「當 Pod 的 rank 編號對應的節點跟實際不同時，TAS 能不能正確切換到 greedy assignment」。現在測試確實只表達了這件事，沒有多餘的雜訊。

結語

這個測試 cleanup 展現了 Kueue TAS 設計中的分層抽象與容錯平衡。TAS 透過 TopologyAssignment 提供理想的位置規劃，同時用 greedy assignment 建立現實世界的容錯機制。真正的挑戰在於：生產環境會因各種非預期因素改變原本的拓撲，設計邏輯中必須有足夠的容錯能力來應對這些變化。

希望這篇文章能幫助你更深入理解 TAS 的設計思路與容錯機制，以及如何透過測試案例的簡化，更清楚地驗證核心機制的正確性。

參考資料