#使用 Kueue、Checkpointing 和 Inference 共存的可抢占 TrainJobs

当训练作业与在线 InferenceService 工作负载共享一个集群时，你希望同时满足两点：

• Inference 受到保护。 它始终拥有所需的 GPU；队列准入是恒定时间的，并且不会被训练作业阻塞。
• Training 充分利用空闲资源。 只要 inference 低于峰值，训练就可以借用空闲 GPU 并继续推进；一旦 inference 重新收回配额，训练便立即让出资源。

本指南通过一组精简的 asset YAML，将 Kubeflow Trainer v2 + Kueue + HuggingFace Trainer checkpointing 组合起来，实现上述行为。这里的所有内容都已在仓库 e2e harness 中的 c12_kueue_preemption.sh 用例中完成端到端验证。

#前提条件

#cohort：一个 CQ 保留配额，另一个借用配额

核心思路是一个 双 ClusterQueue cohort。Inference 持有 GPU 的名义配额；training 持有 0，但当 inference 空闲时，允许训练借用相同数量的配额。当 inference 工作负载收回其配额时，Kueue 会驱逐借用资源的训练 Workload，而 Trainer v2 会在配额释放后立即从头重建 JobSet。

                     ┌──────────────────────────────┐
                     │  Cohort: c12-shared          │
                     │                              │
   inference label ──┤  c12-inference-cq            │
                     │    nominalQuota: 1 GPU       │
                     │    borrowingLimit: 0         │ ← inference never borrows
                     │    reclaimWithinCohort: Any  │
                     │                              │
   training label  ──┤  c12-training-cq             │
                     │    nominalQuota: 0 GPU       │ ← owns nothing
                     │    borrowingLimit: 1 GPU     │ ← borrows when idle
                     │                              │
                     └──────────────────────────────┘

应用 cohort 和按 namespace 划分的 LocalQueue：

base=https://raw.githubusercontent.com/alauda/aml-docs/master/docs/en/train/guides/assets/kueue/preemption
NS=my-namespace  # edit to the namespace where you submit jobs

# 1. Cluster admin — one ResourceFlavor + two cohort ClusterQueues.
kubectl apply -f $base/cluster-queues.yaml

# 2. Cluster admin — Kueue WorkloadPriorityClasses for inference / training.
kubectl apply -f $base/workload-priorities.yaml

# 3. Namespace admin — LocalQueues pointing at the cohort.
curl -fsSL $base/local-queues.yaml | sed "s/<your-namespace>/$NS/" | kubectl apply -f -

这些 asset 文件分别是：

• cluster-queues.yaml — 包含 cohort、两个 ClusterQueue 以及 ResourceFlavor。修改 nominalQuota 和 borrowingLimit，以匹配你希望借出的 GPU 数量。
• workload-priorities.yaml — 两个 WorkloadPriorityClass 值：c12-inference-prio=1000、c12-training-prio=10。如果没有这些配置，cohort 的回收规则仍会生效，但队列内不会有优先级顺序。
• local-queues.yaml — c12-inference-lq 和 c12-training-lq，每个 ClusterQueue 对应一个。

#让 TrainJob 具备抢占安全性

被抢占的 TrainJob 的 Pod 会被终止（先 SIGTERM，宽限期结束后再 SIGKILL）。要在这种情况下继续运行而不是从头开始，你需要：

1. 在 RWX PVC 上放置 checkpoint 目录。 抢占后的 Pod 可能会落到不同节点上——本地存储不够用。
2. 频繁保存 checkpoint。 save_strategy: steps 加上较小的 save_steps。因抢占而丢失的最大工作量由该时间间隔决定。
3. 在下一次启动时恢复。 HuggingFace Trainer 的 .train(resume_from_checkpoint=<path>) 会自动从 output_dir 中取回 checkpoint-N/。LlamaFactory、training_hub、mini_trainer 以及其他任何基于 Trainer 的 recipe 都能自动继承这一能力——它们都暴露相同的 output_dir / save_strategy / resume_from_checkpoint 配置项。
4. 优雅退出。 将 terminationGracePeriodSeconds 设得足够大，让训练器的信号处理器可以在 SIGKILL 之前刷出最后一个 checkpoint。

training-runtime.yaml 这个 asset 将上述四项打包成一个可运行的 TrainingRuntime。trainer-script 的核心如下：

spec:
  template:
    spec:
      replicatedJobs:
        - name: node
          template:
            spec:
              template:
                spec:
                  terminationGracePeriodSeconds: 60   # let final checkpoint flush
                  volumes:
                    - name: ckpt
                      persistentVolumeClaim: { claimName: c12-ckpt }
                  containers:
                    - name: node
                      env:
                        - { name: CKPT_DIR, value: /mnt/ckpt/run }
                      volumeMounts:
                        - { mountPath: /mnt/ckpt, name: ckpt }
                      command: [bash, -ec]
                      args:
                        - |
                          python - <<'PY'
                          import os, glob
                          from transformers import Trainer, TrainingArguments
                          ckpt_dir = os.environ["CKPT_DIR"]
                          # Auto-detect latest checkpoint so the first run starts clean
                          # and every subsequent (re-admitted) run resumes from it.
                          ckpts = sorted(glob.glob(f"{ckpt_dir}/checkpoint-*"),
                                         key=lambda p: int(p.rsplit("-",1)[1]))
                          resume = ckpts[-1] if ckpts else None
                          args = TrainingArguments(
                              output_dir=ckpt_dir,
                              save_strategy="steps", save_steps=4, save_total_limit=2,
                              # ... rest of your training args
                          )
                          trainer = Trainer(model=..., args=args, train_dataset=...)
                          trainer.train(resume_from_checkpoint=resume)
                          PY

同样的结构也适用于 LlamaFactory（在 lf-sft.yaml 中设置 resume_from_checkpoint: true），以及任何其他基于 Trainer 的 recipe——它们最终都可以归结为“将 output_dir 指向 PVC，设置 save_steps，并把最新 checkpoint 传给 .train()”。

save_steps 应根据你可接受的最坏抢占情况来选择：如果每步耗时五秒，save_steps: 100 可将损失工作量上限定为约 10 分钟。再配合 save_total_limit，避免 PVC 无限制增长。

创建 PVC 并部署 runtime：

curl -fsSL $base/checkpoint-pvc.yaml    | sed "s/<your-namespace>/$NS/" | kubectl apply -f -
curl -fsSL $base/training-runtime.yaml  | sed "s/<your-namespace>/$NS/" | kubectl apply -f -

#提交工作负载

一个训练 TrainJob，带有标签，确保以 training 优先级进入 training 队列：

curl -fsSL $base/trainjob-low-priority.yaml | sed "s/<your-namespace>/$NS/" | kubectl create -f -

一个参与同一 cohort、且处于 inference 优先级的 InferenceService：

curl -fsSL $base/inference-service.yaml | sed "s/<your-namespace>/$NS/" | kubectl create -f -

你应该观察到以下现象：

1. Training 先启动 — training Workload 在 c12-training-cq 上达到 Admitted=True（从 cohort 中的 inference CQ 借用了 GPU 配额）。
2. Inference 到达。 它的 Workload 需要一个当前借给 training 的 GPU。Kueue 的经典抢占会将 training Workload 选为目标并驱逐它：

   status:
     conditions:
       - type: Preempted
         status: "True"
         reason: InCohortReclamation
         message: "Preempted to accommodate a workload ... due to reclamation within the cohort"
       - type: Requeued
         status: "True"
3. Training Pod 终止。 JobSet 发送 SIGTERM；训练器刷出最后一个 checkpoint 并退出。
4. Inference 启动 并且不受阻塞地运行。
5. Inference 完成（或缩容）。Kueue 重新准入 training；Trainer v2 重建 JobSet；训练容器在 PVC 上看到 checkpoint-N/，并从那里继续。

实时查看这一往返过程：

kubectl -n "$NS" get workload -w
kubectl -n "$NS" get trainjob,pods
kubectl -n "$NS" get workload -o jsonpath='{range .items[*]}{.metadata.name}: {range .status.conditions[*]}{.type}={.status} {end}{"\n"}{end}'

#与在线 InferenceService 安全共存

这个双 CQ cohort 是承重部分。再配合几个参数，可以让日常运行更加平稳：

• 按峰值而不是平均值来为 inference CQ 设定容量。 如果按平均值来配额，第一次流量尖峰就会侵蚀训练已经开始消耗的容量——每次抢占都会让训练器暂停。应适当增加 nominalQuota，确保稳定状态下的 inference 准入不会触碰借用配额。
• 在 inference 资源上保持 borrowingLimit: 0。 borrowingLimit 是借用方侧的限制：这样可以防止 inference 工作负载消耗另一个 CQ 的名义配额。它不会阻止 training 借用 inference 的空闲名义配额；如果你需要限制某个 CQ 向 cohort 借出多少资源，请使用 Kueue 的 lendingLimit。
• 在 inference CQ 上使用 reclaimWithinCohort: Any，不要使用 LowerPriority。 使用 LowerPriority 时，只有严格低于 inference 优先级类的工作负载才会被抢占；Any 则允许 inference 无论 training 侧如何配置优先级，都能抢回资源。
• 在 Kueue 配置中为 training 设置 PodsReady 超时。 如果一个被抢占后重新准入的 training Pod 遇到镜像拉取很慢的情况，你不希望它永远占着借来的配额；超时后它会回到队列，并让其他工作负载继续前进。
• 为你发布的每个 InferenceService 都设置 WorkloadPriorityClass，而不只是 cohort 中的那些。 如果缺少该标签，Workload 的优先级会停留在 0，抢占规则就无法提升它。
• 不要在 Kueue 配置中设置 manageJobsWithoutQueueName: true。 一旦开启，所有受限 namespace 中的 Pod/Deployment 都必须带有队列标签，这对集群组件来说非常容易踩坑。
• 让 inference predictor 的资源请求保持为单个工作负载。 如果某个 InferenceService 申请的资源超过了 cohort 的名义 inference 配额，再多的抢占也无法满足它。应改为拆分成多个副本。

#保留并共享：用于 namespace 级保留的对称 cohort

上面的双 CQ 布局是有意设计成 非对称 的——inference 拥有一切，training 负责借用。使用同样原语的另一种形态，可以让每个租户都 保留一个下限，同时在邻居空闲时 借用 cohort 中剩余的资源：

                     ┌──────────────────────────────┐
                     │  Cohort: shared-pool         │
                     │                              │
   ns-a label      ──┤  ns-a-cq                     │
                     │    nominalQuota: 2 GPU       │ ← reserved for ns-a
                     │    borrowingLimit: 4 GPU     │ ← may use 6 if cohort is idle
                     │    reclaimWithinCohort: Any  │
                     │                              │
   ns-b label      ──┤  ns-b-cq                     │
                     │    nominalQuota: 4 GPU       │ ← reserved for ns-b
                     │    borrowingLimit: 2 GPU     │ ← may use 6 if cohort is idle
                     │    reclaimWithinCohort: Any  │
                     │                              │
                     └──────────────────────────────┘
                       total nominal across cohort = 6 GPU

此时每个 ClusterQueue 看起来如下——注意 nominalQuota > 0 且 borrowingLimit > 0，并且设置 reclaimWithinCohort: Any，这样即使邻居借走了资源，所有者也能把自己的保留配额拿回来：

apiVersion: kueue.x-k8s.io/v1beta2
kind: ClusterQueue
metadata: { name: ns-a-cq }
spec:
  cohortName: shared-pool
  namespaceSelector:
    matchLabels: { kueue.x-k8s.io/queue: ns-a }
  resourceGroups:
    - coveredResources: ["nvidia.com/gpualloc"]
      flavors:
        - name: c12-default
          resources:
            - name: nvidia.com/gpualloc
              nominalQuota: 2
              borrowingLimit: 4
  preemption:
    reclaimWithinCohort: Any
    withinClusterQueue: LowerPriority

其行为如下：

• 两个 namespace 都空闲。 Cohort 持有 6 个 GPU 的名义容量，但都未被使用。
• 只有 ns-a 有待处理工作。 ns-a 最多可准入 6 个 GPU（自身 2 个名义配额 + 从 ns-b 的空闲名义配额中借来的 4 个）。
• 随后 ns-b 也有待处理工作。 当前有 4 个 GPU 的保留配额被借给了 ns-a。ns-b 的 Workload 触发 InCohortReclamation；Kueue 驱逐 ns-a 的 Workload，直到 ns-b 能以其保留配额准入。ns-a 的前 2 个 GPU（其自身名义配额）不会被触碰。
• 两个 namespace 都满载。 每个都只会准入到自己的 nominalQuota。由于没有空闲配额可借，因此不会发生借用。

#实际操作中的调节项

• nominalQuota 的总和必须 ≤ 物理容量。 保留配额是保证。如果 cohort 的名义总量超过物理 GPU，那么两个 namespace 可能同时达到各自的保留值，而其中一个会在等待 device plugin，而不是等待 Kueue。
• 根据你希望获得的弹性来设置 borrowingLimit。 borrowingLimit + nominalQuota 是单个 CQ 可准入足迹的上限。如果你想要最大的突发能力，就把它设为整个 cohort 减去你的保留值；如果你想为后到的邻居留出余量，则设小一些。
• 借用的工作负载是可抢占的——要对它做 checkpoint。 任何准入量高于 nominalQuota 的工作负载都运行在借用配额上，所有者一旦收回配额就可能被驱逐。 让 TrainJob 具备抢占安全性 一节中的形态可直接复用：共享 PVC、频繁的 save_steps、terminationGracePeriodSeconds: 60。否则，每次回收都会丢失已经消耗的墙钟时间。
• 使用 borrowWithinCohort 来控制准入时的抢占。 当 borrowWithinCohort.policy: LowerPriority 时，借用准入可以在借出方侧抢占严格更低优先级的工作负载。若不启用，则借用只会发生在真正空闲的配额上——行为更安静，但高优先级作业如果遇到忙碌的邻居 CQ，仍需等待自然空闲容量。
• 不要轻易在同一个 cohort 中混合非对称和对称模式。 带有 borrowingLimit: 0 的 CQ（上面的 inference 模式）仍然可以借出空闲名义配额，但它不会从 cohort 中借回配额。在对称模式中，每个 CQ 都会同时借入和借出自己的名义配额。把这两种形态混在一个 cohort 中是可行的，但心智模型会更复杂；如果你两者都需要，最好使用两个 cohort。

#什么时候选择哪种布局

#验证配置

被抢占的 Workload 上的 condition payload 是你的事实依据：

kubectl -n "$NS" get workload \
  -o jsonpath='{range .items[?(@.status.conditions[?(@.type=="Preempted")].status=="True")]}{.metadata.name}{"\n"}{end}'

一个至少被抢占过一次的 training Workload 会显示 reason: InCohortReclamation。它的替代对象（在 inference 完成后）将是一个新的 Workload，具有相同的 JobSet 血缘关系，但 UID 不同——Trainer v2 会根据 TrainJob 确定性地命名它们，因此 TrainJob 名称在重启之间保持稳定。

如果你想在 HAMI 集群上对整条流程进行可重复的端到端覆盖，可以使用仓库 e2e/ harness 中的 c12_kueue_preemption.sh 用例：它会建立 cohort、提交 TrainJob、触发高优先级抢占者，并对 InCohortReclamation condition 以及 checkpoint 恢复进行断言。

有关完整的 Kueue 配置，请参见 Kueue 文档；有关底层算法，请参见 Preemption 概念页。