KAI Scheduler 深入拆解（一）：从组件地图理解整体架构

Posted on In

KAI Scheduler 这个项目第一眼很容易被名字误导：看起来像是一个“GPU 版 kube-scheduler”，但真的把仓库翻一遍以后，会发现它更像是一套围绕 AI/ML 工作负载构建的 Kubernetes 原生调度平台。

它当然有 scheduler 这个核心二进制，但仓库真正的架构中心并不是单个调度循环，而是下面这几个层次同时成立：

  1. 工作负载建模：把 Pod、Job、Ray、Kubeflow、JobSet 之类的对象，统一收敛成 PodGroup、Queue 这些调度语义更强的 CRD。
  2. 调度决策：scheduler 基于 snapshot/session/actions/plugins 做周期性调度。
  3. 执行落地：scheduler 不直接 bind Pod，而是创建 BindRequest，交给 binder 异步执行。
  4. 状态与反馈：PodGroup Controller、Queue Controller、Prometheus、usage DB 一起把“当前状态”和“历史使用量”喂回调度器。
  5. 平台配置：operator 通过 Config 和 SchedulingShard 把整套控制面部署出来，并且把每个 shard 的调度策略参数化。

如果只盯着 pkg/scheduler，会看见一个很强的调度内核；但如果把 cmd/、pkg/apis/、pkg/operator/、pkg/binder/、pkg/podgrouper/ 放在一起，才会看见 KAI 真正的系统轮廓。

先从 repo 结构建立“地图感”

这个仓库最值得先看的不是某个具体算法，而是目录本身：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
cmd/
  admission/
  binder/
  operator/
  podgrouper/
  podgroupcontroller/
  queuecontroller/
  scheduler/
  nodescaleadjuster/
  resourcereservation/
  snapshot-tool/
  fairshare-simulator/
  time-based-fairshare-simulator/

pkg/
  apis/
  binder/
  operator/
  podgrouper/
  podgroupcontroller/
  queuecontroller/
  scheduler/
  common/

docs/
  developer/
  operator/
  fairness/
  queues/
  topology/
  time-based-fairshare/

这个结构已经说明了一件事：KAI 的实现方式不是把复杂性都堆进 scheduler 进程，而是把不同职责拆成多个控制器和服务。

cmd/：每个二进制都代表一个架构角色

从 cmd/ 可以直接看出 KAI 的主要运行单元：

  • scheduler：真正做 placement 决策的内核。
  • binder：异步执行 bind、资源准备、回滚。
  • podgrouper：自动把 workload 转换成 PodGroup。
  • podgroupcontroller：回写 PodGroup 的运行态和资源状态。
  • queuecontroller：维护队列层级与资源使用状态。
  • operator：把整套 KAI 控制面部署出来。
  • admission：准入控制和 webhook。
  • nodescaleadjuster：和节点伸缩场景对接。
  • snapshot-tool、fairshare-simulator、time-based-fairshare-simulator：偏调试、验证、离线分析工具。

这跟传统印象中的“scheduler 就是一切”非常不同。KAI 的设计明显在强调：

调度不是一个函数调用，而是一条跨多个控制器的控制面流水线。

pkg/apis/：真正的架构边界写在 CRD 里

很多 Kubernetes 项目表面上是 controller，实际上系统边界都藏在 CRD 里。KAI 也是这样。

对理解全局最关键的几个 API 是：

  • Config
  • SchedulingShard
  • Queue
  • PodGroup
  • BindRequest
  • Topology

这些对象并不是简单的配置载体，它们定义了 KAI 如何把“AI 集群调度”拆成若干可以被控制器协作处理的状态机。

KAI 的整体架构可以拆成五层

我把它整理成下面这个图：

flowchart LR
    User[用户 / 工作负载控制器] --> Admission[Admission]
    Admission --> PodGrouper[Pod Grouper]
    PodGrouper --> PodGroup[PodGroup CRD]
    User --> Queue[Queue CRD]

    Config[Config CRD] --> Operator[Operator]
    Shard[SchedulingShard CRD] --> Operator
    Operator --> Scheduler
    Operator --> Binder
    Operator --> QueueController
    Operator --> PodGroupController
    Operator --> PodGrouper

    Queue --> Scheduler[Scheduler]
    PodGroup --> Scheduler
    Nodes[Node / PVC / DRA / CSI / etc] --> Scheduler

    Scheduler --> BindRequest[BindRequest CRD]
    BindRequest --> Binder[Binder]
    Binder --> Pods[Pods bound to nodes]

    Pods --> PodGroupController[PodGroup Controller]
    Pods --> QueueController[Queue Controller]
    PodGroupController --> PodGroup
    QueueController --> Queue
    QueueController --> Prometheus[Prometheus / Usage DB]
    Prometheus --> Scheduler

这个图最想表达的是三件事：

1. KAI 先把 workload 变成“可调度对象”，再做调度

Kubernetes 原生对象很多：Pod、Job、Deployment、RayCluster、PyTorchJob、MPIJob、JobSet……

但真正做公平性、gang scheduling、拓扑约束、优先级/可抢占性判断时，直接围绕这些原生对象写逻辑会非常散。

KAI 的做法是通过 pod-grouper 把它们统一抽象成 PodGroup，再配合 Queue 建立调度语义。这一步非常关键：

  • PodGroup 代表“应该一起考虑的 workload 单元”
  • Queue 代表“资源应该如何分配给哪个组织/租户/项目”

也就是说，KAI 不是在调度“孤立的 Pod”，而是在调度 workload group 在 queue 体系里的资源位置。

2. scheduler 只负责决策，binder 负责执行

这是我读完整个仓库以后最喜欢的一点。

KAI 没有把“算出节点”与“执行绑定”写成一个同步链路，而是中间插了一个 BindRequest CRD。

这意味着：

  • scheduler 可以更专注在 决策速度和一致性 上；
  • binder 可以独立处理 绑定、副作用、资源声明、GPU sharing、失败回滚；
  • 两个子系统的错误模型、重试模型、并发模型可以分开设计。

这类设计在控制面系统里非常常见，但在 scheduler 里做得这么明确，其实很少见。

3. KAI 不是静态策略，而是可配置的调度平台

很多项目的“可扩展”只是加几个 flag。KAI 明显不是这个级别。

scheduling shard 这一层直接把下面这些内容暴露成平台配置面：

  • partition / node pool
  • placement strategy（binpack / spread）
  • queue depth
  • min runtime
  • plugin/action override
  • historical usage 配置
  • time-based fairshare 参数

也就是说，KAI 的目标不是给你一个固定算法，而是给你一个 围绕 AI 工作负载场景可演化的调度框架。

四个最关键的 CRD：它们决定了系统的“语义骨架”

Config：整套控制面的安装与组合

Config 不是一个简单的 values 文件替代品，它实际上定义了 KAI 控制面里有哪些服务、这些服务怎么部署、Prometheus 是否启用、全局配置如何下发。

从 pkg/apis/kai/v1/config_types.go 可以看到，ConfigSpec 直接包含：

  • PodGrouper
  • Binder
  • Admission
  • Scheduler
  • QueueController
  • PodGroupController
  • NodeScaleAdjuster
  • Prometheus

这很像在 Kubernetes 里做了一层“平台安装 API”。

SchedulingShard：调度器不再是单实例，而是按分片部署

SchedulingShard 是另一个非常有代表性的对象。

它说明 KAI 并不把“集群只有一个 scheduler 实例”当成默认前提，而是支持按 node pool / queue / podgroup label 做逻辑分片。每个 shard 都可以有自己的：

  • partitionLabelValue
  • placement strategy
  • queueDepthPerAction
  • kValue
  • usageDBConfig
  • Plugins
  • Actions

这代表 KAI 的“平台视角”很强：

不是所有 GPU 集群都该共享同一种调度策略。

Queue：资源公平性的组织边界

如果只把 KAI 看成 GPU 调度器，容易低估 Queue 的重要性。实际上 queue 才是公平性、quota、over-quota、reclaim 的基础对象。

KAI 的很多高级能力——尤其是 fairshare、hierarchical queue、reclaim、time-based fairshare——都建立在 queue 体系之上。

PodGroup：KAI 的工作负载主语

我觉得 PodGroup 是这个项目最值得花时间理解的对象。

在 pkg/apis/scheduling/v2alpha2/podgroup_types.go 里，它已经不只是“gang scheduling 的 minMember 容器”，而是逐渐长成了 KAI 的 workload API：

  • MinMember
  • MinSubGroup
  • Queue
  • PriorityClassName
  • Preemptibility
  • TopologyConstraint
  • SubGroups

而在 Status 侧，PodGroup 还维护了 SchedulingConditions 和 ResourcesStatus，让它同时成为 workload 输入模型和运行态观察点。

看到这里就能理解：KAI 实际上在把“AI workload 的调度语义”沉淀进一个独立 API，而不是永远寄生在原生 Pod 字段上。

为什么 scheduler / binder 拆分是全局设计的支点

如果只能记住 KAI 架构中的一件事，我建议记住这一句：

scheduler 决定“应该放哪”，binder 决定“怎么把它真的放上去”。

这个拆分的收益非常大：

  1. 调度循环可以保持短而稳定

    • 不需要在一个 scheduling cycle 里等待所有 bind 副作用完成。

  2. 复杂资源准备可以异步处理

    • 比如 DRA、PVC/CSI、GPU sharing、resource reservation。

  3. 失败与回滚逻辑可以局部化

    • binder 失败不等于 scheduler 内核设计失败；它只是执行面的一次 reconcile 失败。

  4. 状态可观察性更好

    • BindRequest 本身就成了一个中间态对象，便于排障、审计和重试。

这也是为什么我认为 KAI 更像“调度平台”而不是“单个调度器实现”。

我对 KAI 架构的一个总结

读完 repo 以后，我会用下面这句话概括它：

KAI Scheduler 把 AI 集群调度拆成了 工作负载建模、资源分配决策、绑定执行、状态反馈、平台配置 五个层面，并用 CRD + controller + session-based scheduler 的方式把它们拼成一个完整控制面。

这个思路跟很多只在单个 scheduling cycle 上做文章的项目不一样。KAI 真正重视的是：

  • gang scheduling 不是插件细节，而是 workload model；
  • fairshare 不是一个排序函数，而是 queue 系统；
  • topology 不是一个 affinity 小技巧，而是跨 workload 结构的策略面；
  • scheduler 的职责不是把所有事情都做完，而是生成高质量、可执行、可回滚的调度决策。

推荐的阅读顺序

如果你准备继续往下读源码，我建议按这个顺序：

  1. docs/developer/scheduler-concepts.md
  2. docs/developer/action-framework.md
  3. docs/developer/plugin-framework.md
  4. docs/developer/binder.md
  5. docs/developer/pod-grouper.md
  6. docs/operator/scheduling-shards.md
  7. pkg/scheduler/scheduler.go
  8. pkg/scheduler/framework/
  9. pkg/scheduler/cache/cache.go
  10. pkg/binder/controllers/bindrequest_controller.go

下一篇看什么

理解完组件地图以后，下一步就该看这套系统的“主干工作流”：

  • 一个 Pod 进入集群之后，什么时候变成 PodGroup？
  • Queue、PodGroup、node pool label 是怎么汇合到一起的？
  • scheduler 为什么不直接 bind Pod，而是创建 BindRequest？
  • PodGroup Controller 和 Queue Controller 又把什么信息回写给系统？

下一篇就沿着这条链路，从 Pod 到 BindRequest，把 KAI 的控制面工作流走一遍。