KAI Scheduler 深入拆解（三）：调度框架内核——Cycle、Snapshot、Session、Action、Statement

Posted on In

如果说上一篇讲的是控制面工作流，那么这一篇要进入 KAI Scheduler 最核心的部分：调度内核是怎么组织的。

我读 KAI 代码时最大的感受是，它并不是把所有逻辑揉进一个“大调度函数”，而是把调度过程拆成了几层非常稳定的抽象：

  • cycle
  • cache
  • snapshot
  • session
  • actions
  • plugins
  • statement / scenario

这些抽象不是为了“显得框架化”，而是为了支撑 KAI 真正需要处理的复杂性：

  • gang scheduling
  • queue fairness
  • reclaim / preempt
  • GPU sharing
  • DRA
  • topology-aware scheduling
  • hierarchical podgroups
  • time-based fairshare

如果没有这套结构，任何一个高级特性都可能让调度逻辑失控。

先看最短的调度主循环

在 pkg/scheduler/scheduler.go 里，KAI 的调度主干其实很短：

1
2
3
4
5
6
7
8
func (s *Scheduler) Run(stopCh <-chan struct{}) {
    s.cache.Run(stopCh)
    s.cache.WaitForCacheSync(stopCh)

    go func() {
        wait.Until(s.runOnce, s.schedulePeriod, stopCh)
    }()
}

真正的重点在 runOnce()：

1
2
3
4
5
6
7
8
9
10
func (s *Scheduler) runOnce() {
    ssn, err := framework.OpenSession(...)
    if err != nil { return }
    defer framework.CloseSession(ssn)

    actions, _ := conf_util.GetActionsFromConfig(s.config)
    for _, action := range actions {
        action.Execute(ssn)
    }
}

这段代码非常能代表 KAI 的设计哲学：

scheduler 自己不直接写死“怎么调度每个 Pod”，它做的是：打开一个 session，然后把控制权交给有顺序的 action 和可插拔的 plugin。

也就是说，KAI 的内核不是一个算法，而是一个 按周期运行的调度执行框架。

第一层：Cycle —— 调度不是事件回调，而是周期性计算

KAI 文档里一直强调“scheduling cycle”，这不是表述习惯，而是真正的架构选择。

每个 cycle 的流程可以简化成：

flowchart LR
    Start([Cycle Start]) --> Cache[Cache Sync]
    Cache --> Snapshot[Take Snapshot]
    Snapshot --> Session[Open Session]
    Session --> Actions[Execute Actions]
    Actions --> Close[Close Session]
    Close --> End([Cycle End])

为什么 KAI 选择 cycle model

对于只做简单 node fitting 的调度器，事件驱动模型也许足够。但 KAI 需要解决的问题更复杂：

  • 一个 PodGroup 能不能整体调度成功？
  • 某个 queue 是不是超出 fair share？
  • 为了让高优 workload 运行，当前要不要 reclaim / preempt？
  • 某个 topology domain 是否更适合这个 workload？
  • 某些资源是否能在一个模拟场景里整体成立？

这些问题都更像“在一个一致性视图上做组合决策”，而不是“处理一个 Pod Add 事件”。

cycle model 的好处正是：

  1. 所有决策基于同一份快照
  2. 复杂动作可以按顺序执行
  3. 调度器可以先模拟，再提交
  4. 高成本策略可以集中在每轮计算里完成

第二层：Cache —— 调度器真正依赖的是聚合态，而不是零散 API 调用

KAI 的 cache 在架构上是非常重要的一层。文档 docs/developer/scheduler-concepts.md 也把它放在很前面。

它的职责并不是简单“缓存 informer 数据”，而是负责：

  • 聚合多个 Kubernetes / KAI 资源对象
  • 维护调度需要的派生状态
  • 生成 snapshot
  • 把调度提交结果写回集群

也就是说，cache 同时承担了：

  • 输入聚合层
  • 状态维护层
  • 提交出口层

在 KAI 里，cache 可以被理解成 scheduler 的“authoritative runtime state”。

为什么这层很关键

因为 scheduler 需要同时看：

  • nodes
  • pods
  • podgroups
  • queues
  • bindrequests
  • storage / DRA / resource claims
  • 历史 usage

如果每次 action 执行都直接查 API server，复杂度和性能都会失控。KAI 的做法是先把这些状态汇总到 cache，再在 snapshot 上运行策略。

第三层：Snapshot —— 每一轮调度都从一致性视图出发

在 KAI 的设计里，snapshot 不是一个调试辅助对象，而是调度 cycle 的正式输入。

openSession() 里最关键的一步就是：

1
2
snapshot, err := cache.Snapshot()
ssn.ClusterInfo = snapshot

也就是说，session 看到的整个世界，都是这一刻从 cache 切出来的静态视图。

snapshot 的意义

snapshot 的核心价值在于三个词：

  • 一致性
  • 可模拟
  • 可复现

一致性

同一轮调度里，所有 action / plugin 面对的是同一个 cluster image。

可模拟

你可以在 session 里虚拟 allocate / evict / reorder，而不立刻影响真实集群。

可复现

如果你想调试某一轮调度，理论上 snapshot 是最接近“输入快照”的东西。这也是 KAI 为什么会有 snapshot-tool 这类工具。

第四层：Session —— 一轮调度真正的执行上下文

如果说 snapshot 是输入，那 session 就是这一轮调度的“运行时容器”。

pkg/scheduler/framework/session.go 里的 Session 结构体非常值得看，因为它几乎把 KAI 的扩展点都明牌写出来了：

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

这说明 KAI 的 session 不是“存点变量”的上下文对象，而是 整个调度框架的插件总线。

Session 负责什么

从设计上看，Session 至少承担四件事：

  1. 保存 snapshot 得到的 ClusterInfo
  2. 承载本轮调度的 callback registry
  3. 提供绑定、驱逐、节点排序、GPU 排序等调度操作
  4. 为 statement/scenario 提供统一语义边界

这层设计的妙处

它把“当前一轮调度会发生什么”从长期状态里剥离出来了。

  • cache 是长期维护的系统状态
  • session 是某一轮的计算上下文

这能让 plugin 在 OnSessionOpen 时做本轮初始化，在 OnSessionClose 时做清理或指标记录，而不用在长期全局状态里互相污染。

第五层：Action —— 流程由 action 决定，不由 plugin 决定

KAI 文档 docs/developer/action-framework.md 给出的 action 顺序是：

  1. Allocate
  2. Consolidate
  3. Reclaim
  4. Preempt
  5. StaleGangEviction

这五个动作本身就已经表达了一种非常清晰的调度哲学：

先尝试无破坏地调度，再尝试优化碎片，再尝试跨 queue 回收资源，再尝试队列内优先级抢占，最后清理 gang 死锁状态。

换句话说，Action 负责的是“调度流程的阶段控制”。

为什么 action 顺序很重要

KAI 不是随便堆几个插件，而是明确规定了“先做什么，后做什么”。

Allocate

优先处理无需干扰即可落地的工作负载。

Consolidate

在已有布局上做整理，减少碎片。

Reclaim

当 queue 之间资源分配不公平时，做跨 queue 回收。

Preempt

当同一 queue 内高优 workload 需要资源时，再做优先级抢占。

StaleGangEviction

处理 gang 语义下的僵死/失效场景。

这个顺序说明 KAI 很重视“尽量把破坏性动作后置”。

第六层：Plugin —— 策略由 plugin 注入，不由 action 写死

如果说 action 决定“什么时候干什么”，那 plugin 决定“具体按什么规则干”。

在 pkg/scheduler/framework/framework.go 里，OpenSession() 会读取配置里的 tiers 和 plugins，然后依次：

  1. 取到 plugin builder
  2. 实例化 plugin
  3. 调用 plugin.OnSessionOpen(ssn)

在 OnSessionOpen 里，plugin 把自己的 callback 注册进 session。

这意味着什么

KAI 的 plugin 不是“调用时才临时执行的一段代码”，而是会在 session 打开时把行为注入调度上下文。例如：

  • queue ordering
  • task ordering
  • node scoring
  • GPU scoring
  • predicates
  • fairness calculation
  • reclaim / preempt validation
  • bind request mutation
  • allocate / deallocate event handlers

一个非常关键的区分

很多调度框架里，action 和 plugin 的边界不够清楚；而 KAI 这里非常明确：

  • Action = 调度流程阶段
  • Plugin = 每个阶段使用的策略函数集合

这让系统具备两个非常好的性质：

  1. 流程稳定

    • cycle 还是那个 cycle，action 还是那些 phase。

  2. 策略可演化

    • node placement、fairness、topology、min runtime、dynamicresources 都可以通过 plugin 扩展。

第七层：Statement / Scenario —— KAI 最有意思的设计之一

KAI 文档里有两个很重要但很容易被忽略的概念：

  • Scenario
  • Statement

我认为这两个概念是 KAI 能处理复杂调度动作的关键。

Scenario：先构造“如果这样调度，会发生什么”

Scenario 本质上是一个假设中的调度状态。

比如某个 reclaim / preempt / consolidate 场景，不是直接去驱逐 Pod，而是先在内存里构造一个可能的结果：

  • 如果把这几个 Pod 挪走会怎样？
  • 如果这个 PodGroup 迁移到另一个 domain 会怎样？
  • 如果先释放某些资源再放入新 workload，会不会整体成立？

这意味着 KAI 不是“想到一个动作就立刻提交”，而是先做 what-if analysis。

Statement：把一组调度操作组织成事务式单元

文档把 Statement 描述成 transaction-like object，我觉得这个词非常准确。

它支持：

  • Checkpoint()
  • Rollback(checkpoint)
  • Allocate(pod, node)
  • Evict(pod, msg)
  • Pipeline(pod, node)
  • Commit()
  • Discard()

也就是说，调度器可以在本轮 session 中：

  1. 先尝试若干虚拟操作
  2. 观察这些操作在 plugin callback、资源状态、fairness 上的影响
  3. 如果不成立就 rollback
  4. 如果成立再 commit

为什么这非常适合 KAI

因为 KAI 做的不只是“找个 node fit 一下”。它经常需要处理下面这种组合问题：

  • 一个 PodGroup 的多个 Pod 必须整体成功
  • 某个 queue reclaim 以后，另一个 queue 是否真的能受益
  • topology domain 切换后，整体 locality 是否更优
  • consolidation 是否会造成新的碎片
  • DRA / resource claim 场景是不是整体可行

没有 statement/scenario，很多逻辑要么只能做局部贪心，要么必须把复杂回滚散落在各处。KAI 用一层统一抽象把这个问题收住了。

Action + Plugin + Statement 这三个东西怎么配合

我觉得可以把它们理解成三层：

flowchart TD
    A[Action] --> B[Statement / Scenario]
    C[Plugin callbacks] --> B
    B --> D[Commit to cache / cluster]

    A1[Allocate / Consolidate / Reclaim / Preempt] -.控制流程.-> A
    C1[Predicate / Order / FairShare / Bind Mutate] -.提供策略.-> C
    B1[Checkpoint / Rollback / Commit] -.提供事务化模拟.-> B
  • action 决定“当前做哪类动作”
  • plugin 决定“按什么规则评估和排序”
  • statement/scenario 决定“怎么安全地模拟和提交”

这三个层次组合在一起，才形成 KAI 的调度内核。

为什么这套框架足以承载高级能力

KAI 后续的很多高级特性，其实都不是“另起炉灶”，而是在这套框架上叠加出来的：

  • 拓扑感知：本质上是 plugin + node/domain ordering
  • fairshare：本质上是 queue ordering / resource accounting / reclaim policy
  • time-based fairshare：本质上是 fairness 数据输入变得更丰富
  • hierarchical podgroups：本质上是 workload model + subgroup ordering + action evaluation 更复杂
  • DRA：本质上是 predicates / bind execution / resource claims 集成

也就是说，KAI 的高阶能力不是靠一堆特判支撑，而是靠一个足够稳的 execution model 支撑。

我对 KAI 调度框架的一个总结

如果只用一句话概括这套内核，我会写成：

KAI 把调度建模成：在一致性 snapshot 上打开一个 session，由 action 驱动流程、由 plugin 注入策略、由 statement/scenario 提供可回滚模拟，然后把成立的结果提交回集群。

这是一个非常适合复杂 workload 调度的模型。

它的好处不是“优雅”，而是很实用：

  • 可以承载越来越复杂的 AI workload 语义
  • 可以在保证流程稳定的同时持续增加新策略
  • 可以把破坏性动作控制在明确的阶段和边界里
  • 可以让调试和离线复现更有抓手

下一篇看什么

框架理解了以后，最自然的下一步就是：

  • 代码里到底是怎么启动的？
  • RunApp() 到 runOnce() 之间发生了什么？
  • action / plugin 注册和配置覆盖怎么配合？
  • BindRequest 是在哪一层创建的？
  • binder 又是怎么接住这个对象的？

下一篇就从源码入口开始，沿着真正的调用链，走一遍 KAI 的一次调度周期。