Related Resources
AI 平台/部署

NVIDIA 开源 Run:ai 调度程序以推动社区协作

2025年 4月 1日
By and

今天，NVIDIA 宣布推出 KAI Scheduler 的开源版本，这是一种 Kubernetes-native GPU 调度解决方案，现已在 Apache 2.0 许可证下提供。KAI Scheduler 最初在 Run:ai 平台中开发，现在可供社区使用，同时继续作为 NVIDIA Run:ai 平台 的一部分打包和交付。该计划强调了 NVIDIA 致力于推动开源和企业 AI 基础设施的发展，打造积极协作的社区，鼓励贡献、反馈和创新。

在本文中，我们概述了 KAI Scheduler 的技术细节，强调了其对 IT 和 ML 团队的价值，并解释了调度周期和操作。

KAI 调度程序的优势

管理 GPU 和 CPU 上的 AI 工作负载带来了传统资源调度器通常无法解决的一系列挑战。调度程序专为解决以下问题而开发：

  • 管理不断变化的 GPU 需求
  • 减少计算访问的等待时间
  • 资源保障或 GPU 分配
  • 无缝连接 AI 工具和框架

管理不断变化的 GPU 需求

AI 工作负载可能会迅速变化。例如，您可能只需要一个 GPU 来进行交互式工作 (例如，用于数据探索) ，然后突然需要多个 GPU 来进行分布式训练或多项实验。传统的调度程序难以应对这种可变性。

KAI 调度程序不断重新计算公平份额值，并实时调整配额和限制，自动匹配当前的工作负载需求。这种动态方法有助于确保高效的 GPU 分配，而无需管理员不断进行手动干预。

减少计算访问的等待时间

对于 ML 工程师来说，时间至关重要。调度程序通过将 gang scheduling、GPU sharing 和 hierarchical queuing system 相结合来减少等待时间，使您能够提交批量 jobs，然后离开，确信 tasks 将在资源可用时立即启动，并符合 priorities 和 fairness。

为了进一步优化资源使用，即使在需求波动的情况下，调度程序也对 GPU 和 CPU 工作负载采用了两种有效的策略：

  • Bin-packing 和整合：通过消除资源碎片化 (将较小的任务打包成部分使用的 GPU 和 CPU) ，并通过在节点之间重新分配任务来解决节点碎片化问题，从而更大限度地提高计算利用率。
  • 扩散 ：在节点或 GPU 和 CPU 之间均匀分布工作负载，以最大限度地减少每节点负载，并最大限度地提高每个工作负载的资源可用性。

资源保障或 GPU 分配

在共享集群中，一些研究人员在早期获得的 GPU 数量超过了必要数量，以确保始终可用。这种做法可能会导致资源未得到充分利用，即使其他团队仍有未使用的配额。

KAI 调度程序通过执行资源保证来解决这一问题。它可确保 AI 从业者团队获得分配的 GPU，同时动态地将空闲资源重新分配给其他工作负载。这种方法可以防止资源占用，并提高集群的整体效率。

无缝连接 AI 工具和框架

将 AI 工作负载与各种 AI 框架连接起来可能非常困难。传统上，团队需要通过一系列手动配置来将工作负载与 Kubeflow、Ray、Argo 和 Training Operator 等工具绑定在一起。这种复杂性会延迟原型设计。

KAI Scheduler 通过内置的 podgrouper 来解决这一问题，该程序可自动检测并连接这些工具和框架，从而降低配置复杂性并加速开发。

The scheduling process diagram shows a workload assigned to a queue, which generates pods grouped into a podgroup. The podgroup is sent to the scheduler, which considers workloads from multiple queues. The scheduler operates in an infinite loop, performing a series of steps: taking a cluster snapshot (GPUs and CPUs), dividing resources, executing scheduling actions (allocation, consolidation, reclamation, and preemption), and updating the cluster status. The system ensures efficient resource allocation and management within defined execution orders.
图 1。KAI Scheduler 工作流

核心调度实体

KAI Scheduler 有两个主要实体：podgroups 和 queues。

Pod 组

Pod 组是用于调度的原子单元，表示必须作为单个单元执行的一个或多个相互依赖的 Pod，也称为 gang scheduling 。这一概念对于分布式 AI 训练框架如 TensorFlow 或 PyTorch 至关重要。

以下是 Podgroup 的关键属性：

  • 最低成员数 ：指定必须一起调度的 pod 的最低数量。如果所有成员均不可用所需资源，则 pod 组将保持待处理状态。
  • 队列关联 ：每个 podgroup 都与特定的调度队列关联，并将其与更广泛的资源分配策略关联。
  • 优先级类： 确定相对于其他 podgroups 的调度顺序，从而影响作业优先级。

队列

队列是实现资源公平性的基础实体。每个队列都有指导其资源分配的特定属性：

  • 配额：保证向队列分配的基准资源分配（Quota）。
  • 超额配额权重： 影响所有队列中超出基准配额的额外剩余资源的分配方式。
  • 限制： 定义队列可以消耗的最大资源，例如 GPU、DPU 等。
  • 队列优先级 ：确定相对于其他队列的调度顺序，从而影响队列优先级。

架构和调度周期

调度程序的核心是持续捕获 Kubernetes 集群的状态、计算最优资源分配并应用有针对性的调度操作。

The diagram shows the internal workflow of the KAI scheduler responsible for managing GPU and CPU resources. The scheduler operates in an infinite loop, continuously executing four key processes: taking a cluster snapshot, dividing available resources, performing scheduling actions, and updating the cluster status. Scheduling actions are executed in a defined order: allocation, consolidation, reclamation, and preemption.
图 2。调度周期

该过程分为以下阶段：

  • 集群快照
  • 资源分配和公平分配计算
  • 调度操作

集群快照

调度周期始于捕获 Kubernetes 集群的完整快照。在此阶段，系统会记录节点、pod 组和队列的当前状态。

Kubernetes 对象可转换为内部数据结构，从而防止中期不一致，并确保所有调度决策均基于集群的稳定视图。

资源分配和公平分配计算

有了准确的快照，就可以通过资源分割算法为每个调度队列计算公平份额。该算法的结果是一个数值，该值表示队列应获得的最佳资源量，以最大限度地提高集群中所有队列之间的公平性。

除法算法具有以下阶段：

  • 当之无愧配额： 每个调度队列最初都会分配与其基准配额（baseline quota）相等的资源。这可保证每个部门、项目或团队都能获得应有的最少资源。
  • 超额配额 ：任何剩余资源将按照超出配额权重的比例分配给未满足要求的队列。这种迭代过程会微调每个队列的公平比例，动态适应工作负载需求。

调度操作

计算公平份额后，调度程序会应用一系列有针对性的操作，使当前分配与计算出的最佳状态保持一致：

  • 分配 ：待处理作业的评估基于分配与公平分配比率。适合可用资源的作业会立即绑定，而需要当前释放资源的作业则会排队等待流水线分配。
  • 整合：对于训练工作负载，调度程序会构建剩余待处理训练作业的有序队列。 它会对这些作业进行迭代，并尝试通过将当前分配的 pods 移动到其他节点来分配它们。此过程可最大限度地减少资源碎片化，为待处理作业释放连续块。
  • 回收： 为确保公平性，调度程序可识别消耗超过合理比例的队列。然后，它根据定义明确的策略驱逐选定的作业，确保服务不足的队列获得必要的资源。
  • 抢占先机 ：在同一队列中，低优先级作业可以优先用于高优先级的待处理作业，确保关键工作负载不会因资源争夺而资源乏。

示例场景

设想一个集群有三个节点，每个节点配备八个 GPU，总共 24 个 GPU。目前正在进行两个项目：project-a 和 project-b。它们分别与 queue-a (中优先级队列) 和 queue-b (高优先级队列) 相关联。假设所有队列均在公平分配范围内。图 3 显示了当前状态。

A diagram shows a cluster with three nodes, each containing eight GPUs. Jobs are assigned from two priority queues: a high-priority queue and a medium-priority queue . Node 1 runs Training Jobs 1 and 2, with the high-priority job using four GPUs and the medium-priority job using two GPUs. Node 2 is occupied by Training Job 3, a medium-priority job using six out of eight GPUs. Node 3 hosts an interactive job from the medium-priority queue, using five out of eight GPUs.
图 3。集群状态示例
  • 节点 1： 训练作业 1 使用高优先级队列中的四个 GPU，训练作业 2 使用两个 GPU。
  • 节点 2： 训练任务 3 使用六个 GPU。
  • 节点 3：交互式作业 (包括数据探索或脚本调试) 使用 5 个 GPU。

现在有两个待处理的作业提交到队列中 (图 4) ：

  • 训练作业 A 需要在单个节点上使用四个连续的 GPU。
  • 训练作业 B 需要向高优先级队列提交三个连续的 GPU。
A diagram shows that two training jobs, Training Job A and Training Job B, are submitted with GPU requests of four and three, respectively. Training Job A belongs to the medium-priority queue (queue-a), while Training Job B is associated with the high-priority queue (queue-b).
图 4。包含两项训练作业的示例场景

分配阶段

调度程序按优先级对待处理作业进行排序。在本例中，Training Job B 会被优先处理，因为其队列在紧急程度方面排名较高。当然，作业订购流程比简单的优先级排序更复杂。调度程序使用高级策略来确保公平性和效率。

After capturing the cluster snapshot, the scheduler begins processing the submitted jobs. It first attempts the allocation action, prioritizing the highest-priority queue. As a result, Training Job B is scheduled on Node 3.
图 5。训练任务 B 的分配阶段
  • 对于训练作业 B，节点 3 符合资格，因为它有三个连续的免费 GPU。调度程序将训练作业 B 绑定到节点 3，节点 3 将完全占用。
  • 接下来，调度程序尝试分配训练任务 A。但是，没有一个节点提供四个连续的免费 GPU。

巩固阶段

由于无法通过分配操作来调度 Training Job A，因此调度程序进入整合阶段。它会检查节点，确定是否可以重新分配正在运行的 pods，以形成用于 Training Job A 的连续块。

To accommodate Training Job A, the scheduler attempts to allocate four GPUs. However, no node currently has four available GPUs. As a result, the scheduler enters the consolidation phase, relocating Training Job 2 from Node 1 to Node 2. This frees up space on Node 1, enabling Training Job A to be scheduled there.
图 6。Training Job A 的分配和整合阶段
  • 在节点 1 中，除了 Training Job 1 (由于优先级较高，必须保持不变) 之外，Training Job 2 占用两个 GPUs。调度程序会将 Training Job 2 从节点 1 迁移到节点 2。
  • 此次整合后，节点 1 的可用 GPU 数量从两个增加到四个连续的 GPU。
  • 借助这种新布局，Training Job A 现在可以分配到 Node 1，从而满足四个连续的 GPU 的要求。

在这种情况下，整合行动已经足够了，既不需要 reclaim，也不需要 preempt。但是，如果队列中有一项额外的待处理作业低于其 fair share，且没有可用的整合路径，则调度程序的 reclaim 或 preempt 操作将开始执行，即从过度分配的队列或同一队列中的低优先级作业中回收资源，以恢复平衡。

状态更新

队列的状态会更新，并且整个周期会从头开始，以确保日程安排能够跟上新工作负载的发展。

社区协作

KAI 调度程序 不仅仅是一个原型。它是 NVIDIA Run:ai 平台 核心的强大引擎，受到许多企业的信任，并为关键的 AI 运营提供支持。凭借其可靠的跟踪记录，KAI 调度程序为 AI 工作负载编排设定了黄金标准。

我们邀请企业、初创公司、研究实验室和开源社区在您自己的环境中试验 scheduler，并贡献 learnings。

加入我们的 /NVIDIA/KAI-scheduler GitHub 资源库，为项目添加一颗星，进行安装，并分享您的真实体验。在我们不断突破 AI 基础架构的极限时，您的反馈对我们来说非常宝贵。

KubeCon 2025

NVIDIA 很高兴参加 4 月 1 日至 4 日在英国伦敦举行的 KubeCon 。欢迎在 S750 展位与我们互动，进行实时对话和演示。有关我们工作的更多信息，包括我们的 16 场演讲，请参阅 NVIDIA 在 KubeCon 上的会议。

 

相关资源

标签

AI 平台/部署 | 数据中心/云端 | AI Platforms / Deployment | Deep dive | General | Intermediate Technical | kubernetes | Run:ai | 数据中心/云 | 数据科学

关于作者

Ronen Dar 在收购 Run:ai 后担任 NVIDIA 软件系统副总裁，他在 2018 - 24 年期间共同创立并担任首席技术官领导了 Run:ai。在此之前，Ronen 在新泽西州贝尔实验室完成了博士后研究，并在特拉维夫大学获得了博士学位。
Ekin Karabulut 是一位数据科学家兼开发者倡导者，他曾在 Run：ai 任职，现就职于 NVIDIA，探索如何在不同的生产场景中高效使用大型模型。此前，她研究了联邦学习对隐私的影响，专注于分布式训练技术，并对研究和行业环境中 GPU 使用效率低下感到着迷。她成立了 AI 基础设施俱乐部，现居德国慕尼黑。