k8s之kueue中的workload 和 kubeflow中的TFJob

Kueue中的Workload详解

1. Workload基本概念

Workload是Kueue系统中的核心API资源，代表了一个需要被调度和执行的工作单元。它的主要作用是：

• 封装一个实际的Kubernetes工作负载(如Job、TFJob等)
• 携带该工作负载的资源需求信息
• 在Kueue的队列系统中进行排队和调度

2. Workload的关键特性

2.1 与原生工作负载的关系

Workload不是替代Kubernetes原生工作负载(如Job、Deployment)，而是：

• 包装这些工作负载
• 为它们增加队列管理能力
• 控制它们的执行时机

2.2 主要功能

1. 资源需求声明：指定工作负载需要的CPU、GPU等资源
2. 队列管理：将工作负载分配到特定队列
3. 优先级控制：与PriorityClass结合实现优先级调度
4. 状态跟踪：记录工作负载的排队、调度和执行状态

3. Workload的API结构

典型的Workload资源定义如下：

apiVersion: kueue.x-k8s.io/v1beta1
kind: Workload
metadata:
  name: example-workload
spec:
  queueName: "team-a-queue"
  priorityClassName: "high-priority"
  podSets:
  - name: "main"
    count: 3
    template:
      spec:
        containers:
        - name: job-container
          image: my-image:latest
          resources:
            requests:
              cpu: "1"
              memory: "1Gi"

主要字段说明：

• queueName: 指定工作负载所属的队列
• priorityClassName: 指定优先级类别
• podSets: 定义工作负载的Pod模板和副本数
• admission: (只读字段)记录工作负载何时被准入执行

4. Workload生命周期

Workload会经历以下状态：

1. Pending: 刚创建，等待准入
2. Admitted: 已被队列准入，但资源尚未满足
3. Active: 资源已分配，相关工作负载正在运行
4. Finished: 工作负载完成(成功或失败)

TFJob关键点解析

1. 为什么叫 Job？关键点解析

(1) 一次性执行（Run-to-Completion）

• 核心共性：
  TFJob 和原生 Job 一样，任务执行完成后会自动终止，不会无限运行。例如：

  • 原生 Job：一个数据处理任务完成后，Pod 会退出。
  • TFJob：一个分布式训练任务完成后（如达到指定 epoch 或收敛），所有角色（Worker/PS）的 Pod 也会退出。

• 与长期服务（如 Deployment）的区别：
  像 Deployment 这种工作负载是长期运行的（即使 Pod 崩溃也会重启），而 Job/TFJob 是临时性任务。

(2) 继承原生 Job 的基础特性

• 支持重试和容错：
  如果某个 Pod 失败（如 Worker 崩溃），TFJob 会像原生 Job 一样尝试重启或重新调度（取决于 restartPolicy）。
• 任务状态跟踪：
  提供 succeeded/failed 状态标识，可通过 kubectl get tfjob 查看。

(3) 名称延续 Kubeflow 的命名习惯

Kubeflow 生态中，类似的任务型资源均以 Job 结尾，例如：

• TFJob（TensorFlow）
• PyTorchJob（PyTorch）
• MPIJob（MPI 分布式任务）
  这种命名方式明确了它们的“一次性任务”属性。

---

2. TFJob 的特殊性（与原生 Job 的区别）

虽然 TFJob 属于广义的 Job 范畴，但它针对分布式训练场景做了扩展：

特性	原生 Job	TFJob
任务类型	通用任务（如批处理）	专为 TensorFlow 分布式训练优化
角色定义	所有 Pod 相同（无角色区分）	支持多角色（Chief/Worker/PS）
通信需求	Pod 间通常无需通信	Pod 间需高性能网络通信（如 gRPC）
资源需求	CPU/内存为主	通常需要 GPU/NPU 加速器
使用场景	数据处理、脚本运行	深度学习模型训练

---

Workload 和 TFJob 的关系

(1) 典型交互流程

当用户提交一个 TFJob 并希望由 Kueue 管理时，流程如下：

1. 用户提交 TFJob（或通过某个控制器提交）。
2. Kueue 的准入控制器拦截：
   • 检测到该 TFJob 需要排队（例如，通过 queue-name 注解标记）。
3. Kueue 创建 Workload：
   • Kueue 将 TFJob 的信息（如 Pod 模板、资源需求）封装成一个 Workload 对象，放入指定队列。
4. Workload 等待调度：
   • Kueue 根据 PriorityClass、资源可用性等决定何时允许 Workload 继续。
5. Workload 被批准后，TFJob 真正创建 Pod：
   • 一旦资源足够，Kueue 会解除 TFJob 的阻塞状态，使其正常创建 Pod 执行训练。

(2) 关键关系总结

组件	作用	与对方的关系
TFJob	运行 TensorFlow 分布式训练	被 Workload 封装，受 Kueue 调度
Workload	管理作业的排队和资源分配	持有 TFJob 的元数据，控制其何时执行

---

实际用例

场景：使用 Kueue 管理 TFJob

假设一个集群资源有限，多个团队提交 TFJob，需要通过 Kueue 公平调度：

1. 提交 TFJob 时标记队列：

   apiVersion: kubeflow.org/v1
   kind: TFJob
   metadata:
     name: my-tfjob
     annotations:
       kueue.x-k8s.io/queue-name: "gpu-team-a"  # 指定队列
   spec:
     tfReplicaSpecs:
       Worker:
         replicas: 4
         template:
           spec:
             containers:
             - name: tensorflow
               resources:
                 requests:
                   nvidia.com/gpu: 2
   

2. Kueue 自动创建 Workload：
   • Kueue 会生成一个对应的 Workload，记录该 TFJob 需要的资源（如 8 GPU）。
3. Workload 进入队列：
   • 如果当前集群 GPU 不足，Workload 状态为 Pending，TFJob 不会真正创建 Pod。
4. 资源释放后，Workload 被激活：
   • Kueue 批准 Workload，TFJob 开始运行。

---

技术实现细节

(1) 如何关联 Workload 和 TFJob？

• Kueue 通过 Owner Reference 或 标签/注解 关联两者。
• 例如，Workload 的 spec.podSets 会记录 TFJob 的 Pod 模板。

(2) TFJob 会被修改吗？

• Kueue 不会修改原始 TFJob，而是通过 Kubernetes 的 准入控制（Admission Webhook） 暂停其创建，直到 Workload 被调度。

(3) 其他类似工作负载

• 除了 TFJob，Kueue 同样支持管理：
  • Kubernetes 原生 Job、CronJob
  • Kubeflow 的 PyTorchJob、MPIJob
  • 其他自定义 CRD（需适配）。