KAI Scheduler 深入拆解（五）：插件、分片、拓扑感知与时间公平性

Posted on In

前几篇文章里，我一直在强调一件事：KAI Scheduler 的价值不只是“今天已经实现了哪些特性”，更重要的是 它为什么可以持续装下更多特性。

这一篇就专门看这个问题。

如果只看表层功能，KAI 已经足够复杂：

  • hierarchical queues
  • reclaim / preempt
  • topology-aware scheduling
  • hierarchical podgroups
  • DRA
  • GPU sharing
  • time-based fairshare
  • scheduling shards

但真正有意思的地方在于，这些能力并不是一堆互不相干的特判，而是大多能在同一套框架下自然落位：

  • 流程由 action 管
  • 策略由 plugin 管
  • 运行边界由 session / statement 管
  • 平台级定制由 SchedulingShard 管
  • 工作负载语义由 PodGroup 管
  • 历史反馈由 queue controller + Prometheus + usage DB 管

换句话说，KAI 的扩展性不是“多写几个 if”，而是 架构层面为扩展预留了空间。

先看最关键的一点：plugin 不是装饰品，而是调度策略的主入口

pkg/scheduler/plugins/factory.go 把默认插件注册得非常完整：

  • predicates
  • priority
  • nodeplacement
  • nominatednode
  • nodeavailability
  • gpusharingorder
  • gpupack
  • gpuspread
  • resourcetype
  • podaffinity
  • elastic
  • kubeflow
  • ray
  • taskorder
  • subgrouporder
  • dynamicresources
  • topology
  • proportion
  • minruntime
  • snapshot
  • reflectjoborder

这张表最重要的意义不是“插件很多”，而是它说明 KAI 把以下问题都视作 策略问题，而不是写死在主循环里的逻辑：

  • 节点排序
  • GPU 排序
  • queue fairness
  • reclaim / preempt 判断
  • workload integration
  • subgroup 顺序
  • topology 约束
  • dynamic resources
  • bind request mutation

这就是为什么 KAI 的高级特性能持续叠加，而不会轻易把主调度流程搞乱。

Session 暴露的扩展点，基本等于 KAI 的能力边界

在 pkg/scheduler/framework/session.go 里，Session 暴露的函数槽位非常多，最有代表性的包括：

  • JobOrderFns
  • TaskOrderFns
  • SubGroupOrderFns
  • QueueOrderFns
  • NodeOrderFns
  • GpuOrderFns
  • PrePredicateFns
  • PredicateFns
  • CanReclaimResourcesFns
  • ReclaimVictimFilterFns
  • PreemptVictimFilterFns
  • GetQueueAllocatedResourcesFns
  • GetQueueDeservedResourcesFns
  • GetQueueFairShareFns
  • BindRequestMutateFns
  • event handlers

只要看见这些接口，就能立刻理解 KAI 的插件思想：

它不是只开放一个 score 钩子，而是把“调度过程中可能需要定制的所有局部决策”尽量显式化。

这带来了两个直接好处：

  1. 新能力可以落在明确的扩展点上，而不是侵入主流程。
  2. 现有能力可以通过配置和优先级组合，而不是硬编码互斥。

SchedulingShard：KAI 真正的平台配置面

如果说 plugin 是“代码级扩展点”，那 SchedulingShard 就是“平台级策略面”。

我认为这是 KAI 最值得单独夸的一层设计。

在 pkg/apis/kai/v1/schedulingshard_types.go 里，SchedulingShardSpec 暴露的内容包括：

  • Args
  • PlacementStrategy
  • PartitionLabelValue
  • QueueDepthPerAction
  • MinRuntime
  • KValue
  • UsageDBConfig
  • Plugins
  • Actions

这意味着什么？

这意味着 KAI 把很多原本应该藏在源码、flag 或者某个配置文件里的策略选择，提升成了 可声明、可演化、可分片实例化的 Kubernetes API。

为什么这层设计非常强

因为它让“调度器实例”不再是同质的。

你可以想象这样几种 shard：

  • 一个 shard 专门服务 GPU 训练任务，偏向 binpack
  • 一个 shard 服务在线推理，偏向 spread / 更保守 preempt
  • 一个 shard 打开更积极的 time-based fairshare
  • 一个 shard 限制特定 action 的 queue depth
  • 一个 shard 针对某类 workload 打开或关闭特定 plugin

这就把 KAI 从“一个调度器程序”推进成了“一个可按业务场景切片的调度平台”。

分片不是部署技巧，而是语义隔离能力

scheduling-shards.md 里说得很清楚：每个 shard 实际上围绕同一个 partition label key/value 工作。也就是说，scheduler 只会考虑带有同一组分片标签的：

  • Nodes
  • Queues
  • PodGroups

这是一种非常典型的 Kubernetes 风格：

  • 用 label 做资源归属
  • 用独立 scheduler deployment 做计算隔离
  • 用 CRD 做配置入口

这样带来的好处不只是扩容，更重要的是 策略与资源空间同时分片。

也就是说，不同 shard 不只是“不同副本数的 scheduler”，而是：

  • 看到的资源集合不同
  • 使用的策略可以不同
  • action/plugin 组合可以不同
  • fairness / minRuntime / placement 参数可以不同

拓扑感知：KAI 的 topology 不是附属 feature，而是主框架里的自然延伸

很多调度器支持 topology，最终只是把它当作额外 affinity。

KAI 的 topology 明显不是这个级别。

docs/topology/README.md 描述得很清楚，它依赖两个东西：

  1. Topology CRD 定义层级拓扑
  2. PodGroup / workload 声明 required / preferred placement

然后调度时做两层决策：

第一层：domain selection

先选择应该把 workload 放进哪个 topology domain。

这里的策略偏向 bin-packing：优先选择“相对更满但仍能容纳该 workload”的 domain，以减少碎片。

第二层：node ordering within domain

在选中的 domain 内，再按 preferred placement 做节点排序，让 Pod 尽量聚拢到更合适的子域。

这套设计为什么自然？

因为 KAI 的框架本身就已经有：

  • workload-level abstraction（PodGroup）
  • node / GPU ordering 扩展点
  • scenario / statement 模拟能力
  • 多级 subgroup 结构

所以 topology 不需要硬插一条特别路径，而是能顺着既有框架表达出来。

公平性：KAI 的 fairshare 不是单一排序函数，而是 queue 系统的一部分

docs/fairness/README.md 里最重要的一句话是：

  • deserved quota 一定优先满足
  • 剩余资源再按 priority / weight 分 over-quota share

这意味着 KAI 的公平性不是简单“谁资源少谁先跑”，而是建立在 queue hierarchy 上的：

  1. 先分配 deserved quota
  2. 再按 over-quota weight/priority 分多出来的资源
  3. 如果当前实际占用和 fair share 偏离太多，就通过 reclaim 收敛

这件事为什么离不开框架抽象

因为 fairshare 不只是 queue 排序，还会影响：

  • queue ordering
  • reclaim eligibility
  • victim filtering
  • 资源记账
  • action 执行顺序

如果没有 session 里的 queue callback、action 框架和 statement 模型，fairshare 很容易变成一堆局部规则拼接。KAI 的实现方式则更像把它作为一个正式的一等公民。

Time-based fairshare：KAI 开始把“历史行为”放进调度决策

我觉得 time-based fairshare 是 KAI 很有代表性的一个能力，因为它显示出这套架构并不只看“这一刻的资源状态”，而是开始纳入时间维度。

docs/time-based-fairshare/README.md 给出的主线是：

  1. PodGroup Controller 发布 podgroup 资源状态
  2. Queue Controller 聚合 queue usage
  3. Prometheus 存储历史数据
  4. scheduler 通过 usage DB client 拉取 usage
  5. proportion 等 fairness 相关逻辑把历史 usage 纳入 over-quota 资源分配

这个闭环很说明问题：

KAI 不只是个实时调度器，它正在演化成一个“结合当前状态与历史行为”的资源治理系统。

为什么这件事能自然接进去

因为 KAI 之前已经有了：

  • Queue 作为公平性主语
  • Queue Controller 作为聚合器
  • SchedulerCache / usage DB 作为数据接入点
  • KValue、UsageDBConfig 等策略面配置
  • GetQueueFairShareFns 这类 session 级扩展点

所以历史 usage 的接入，不需要推翻已有架构，只需要在已有资源分配框架上把输入变得更丰富。

Config + SchedulingShard 的组合，让 operator 也变成架构的一部分

前面聊得多的是 scheduler 内核，但 KAI 的平台味道其实还来自 operator 层。

在 pkg/apis/kai/v1/config_types.go 里，ConfigSpec 直接声明了整套控制面的组件：

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

也就是说，KAI 不是把 operator 当“安装器”，而是把它做成了这套架构的一部分：

  • Config 负责描述整个平台的全局组件组合
  • SchedulingShard 负责描述单个调度分片的策略配置

这两个 CRD 组合起来，基本就是 KAI 的控制面 API。

为什么我认为 KAI 的扩展路径是健康的

很多系统在功能越来越多以后，会出现一种典型问题：

  • 新特性只能往旧代码里塞特判
  • 不同能力互相穿透
  • 配置开始变得不可推理
  • 主循环越来越重
  • 某个高级功能一加，全局行为就变得不可预测

KAI 目前给我的感受恰好相反。它的高级能力大多还能找到一个比较自然的落点：

1. workload 语义落在 PodGroup

而不是散落在所有 plugin 里。

2. 公平性与资源治理落在 Queue

而不是偷偷写在 node score 里。

3. 执行流程落在 action

而不是被 plugin 任意控制。

4. 局部策略落在 plugin

而不是让 action 变成巨无霸。

5. 中间态执行落在 BindRequest + binder

而不是让 scheduler 同步做完所有副作用。

6. 平台调优落在 SchedulingShard

而不是只能改源码或启动参数。

这就是为什么我觉得 KAI 的架构不仅功能多，而且有继续长下去的空间。

如果想扩展 KAI，最值得看哪几类点

从阅读体验上，我会把 KAI 的扩展入口分成四类。

1. 想加 workload 语义

先看：

  • PodGroup API
  • pod-grouper 的 supported types / grouping plugins
  • subgroup / topology 相关设计文档

2. 想改调度策略

先看：

  • Session 暴露的 callback 槽位
  • plugins/factory.go
  • 现有 plugin 的 OnSessionOpen

3. 想改调度流程

先看：

  • action framework
  • actions/factory.go
  • action 优先级和 SchedulingShard.Actions

4. 想做平台级多租户 / 多场景隔离

先看：

  • SchedulingShard
  • partition label 机制
  • operator 相关 reconciliation

我对 KAI 扩展能力的一个总结

如果只用一句话概括，我会写成：

KAI 把“调度平台”的关键维度分开了：工作负载语义、资源公平性、调度流程、局部策略、执行落地、平台配置，各自都有明确落点，因此高级能力能以组合的方式生长出来。

这跟很多“功能越多越难维护”的 scheduler 最大的不同在于：

  • 它不是靠主循环越来越聪明来进化
  • 而是靠边界越来越清楚来进化

系列收束：怎么继续读这个仓库

如果你是第一次接触 KAI，我会建议按照这条线继续深入：

  1. 先读 docs/developer/scheduler-concepts.md
  2. 再读 docs/developer/action-framework.md
  3. 看 docs/developer/plugin-framework.md
  4. 看 docs/developer/pod-grouper.md 和 docs/operator/scheduling-shards.md
  5. 进入 pkg/scheduler/ 的 session / action / cache 主链路
  6. 最后再去看 topology、fairness、time-based fairshare 等高级能力

因为 KAI 最值得学习的地方，真的不只是某一个调度算法，而是：

它如何把复杂的 AI workload 调度，拆成一套仍然能持续演化的 Kubernetes 控制面架构。

对我来说，这就是这个仓库最有价值的部分。