WEBKT

深度解析 K8s 调度器扩展框架：编写自定义插件支持复杂 AI 任务

5 0 0 0

在云原生时代，Kubernetes (K8s) 已成为管理容器化应用的事实标准。然而，随着 AI/ML 任务的爆发式增长，默认调度器的“逐个 Pod 调度”逻辑逐渐显露疲态。AI 训练通常涉及分布式计算（如 PyTorch DDP、TensorFlow Parameter Server），对 GPU 资源的亲和性、网络带宽以及任务同步性有极高要求。

为了解决这些问题，K8s 在 v1.15 引入并在后续版本不断完善了 Scheduling Framework。本文将深入探讨这一框架的机制，并介绍如何通过自定义插件实现复杂的 AI 作业调度策略。

为什么默认调度器无法满足 AI 需求？

默认调度器 kube-scheduler 主要针对通用的 Web 服务设计，对于 AI 任务存在以下局限：

  1. Gang Scheduling（成组调度）缺失：分布式训练要求 $N$ 个副本同时启动，否则会因资源死锁（A 等待 B 的资源，B 等待 A 的资源）导致任务停滞。
  2. 拓扑感知能力弱：AI 任务对 GPU 间的 NVLink 通讯非常敏感。默认调度器不感知节点内部的 GPU 拓扑结构。
  3. Binpacking vs Spreading：默认调度倾向于打散 Pod 以保证高可用，但 AI 任务往往希望尽量填满单个节点的 GPU，以减少跨节点通信开销。

K8s 调度框架核心原理

Scheduling Framework 将调度过程拆分为两个循环：调度循环 (Scheduling Cycle) 和 绑定循环 (Binding Cycle)。每个循环中预留了多个扩展点（Extension Points），开发者可以注册自定义插件。

核心扩展点概览

  • PreFilter：预处理 Pod 信息，或检查集群状态。
  • Filter：过滤掉不满足条件的节点（类似旧版的 Predicates）。
  • PostFilter：如果所有节点都无法调度，触发此阶段（常用于抢占逻辑）。
  • PreScore / Score：给节点打分，决定最优路径（类似旧版的 Priorities）。
  • Permit：AI 调度的关键。可以延迟绑定，实现 Gang Scheduling。
  • Bind：执行实际的 API 绑定操作。

实战：编写 AI 调度插件

假设我们要实现一个简单的 Gang Scheduling 插件，其核心逻辑是在 Permit 阶段拦截属于同一个作业的 Pod，直到所有 Pod 到齐后再统一放行。

1. 定义插件结构

在 Go 语言中，插件需要实现 framework.Plugin 接口。

type MyGangSchedulingPlugin struct {
    handle framework.Handle
}

func (p *MyGangSchedulingPlugin) Name() string {
    return "MyGangScheduling"
}

2. 实现 Permit 逻辑

Permit 阶段允许返回 Success、Reject 或 Wait。

func (p *MyGangSchedulingPlugin) Permit(ctx context.Context, state *framework.CycleState, pod *v1.Pod, nodeName string) (*framework.Status, time.Duration) {
    // 1. 获取该 Pod 所属的 Job ID
    jobID := getJobID(pod)
    
    // 2. 检查当前已就绪的 Pod 数量
    readyCount := p.getReadyPodCount(jobID)
    minAvailable := p.getMinAvailable(jobID)

    if readyCount + 1 >= minAvailable {
        // 如果数量够了，通知所有等待中的 Pod
        p.handle.IterateOverWaitingPods(func(wp framework.WaitingPod) {
            if getJobID(wp.GetPod()) == jobID {
                wp.Allow(p.Name())
            }
        })
        return framework.NewStatus(framework.Success), 0
    }

    // 3. 数量不够，进入等待状态
    return framework.NewStatus(framework.Wait), 10 * time.Minute
}

AI 场景下的进阶优化建议

  1. GPU 拓扑感知调度：
    在 Score 阶段，通过读取节点的 GPU 拓扑（如通过 Device Plugin 暴露的 Label 或 CRD），给具有 NVLink 直连且空闲的节点打高分。这能显著提升多卡并行训练的效率。

  2. 动态优先级与抢占：
    对于生产环境，可以使用 PostFilter 插件。当高优先级的训练任务资源不足时，自动驱逐低优先级的实验性任务或离线数据处理任务。

  3. 利用社区力量：
    如果不想从零开发，推荐关注 Kubernetes-sigs/scheduler-plugins。它包含了官方维护的 Coscheduling（Gang Scheduling）、NodeResourceTopology（拓扑感知）等成熟插件。同时，KubeBatch 和 Volcano 也是处理大规模 AI 任务的高性能调度引擎选项。

总结

K8s 调度器框架的插件化设计，赋予了基础设施工程师极大的灵活性。通过在 Permit 阶段控制任务协同，在 Score 阶段优化资源分布，我们可以构建出一个能够理解 AI 负载特征的“智能调度系统”。

在实际落地时，建议先通过 Scheduler Config 启用社区现成插件，当业务出现极端特化需求（如基于 RDMA 拓扑的调度）时，再考虑深度定制开发。这不仅能保持架构的整洁，也能紧跟 K8s 社区的演进步伐。

云原生架构师 KubernetesAI基础设施调度算法

评论点评