跳到主要内容
返回时间线
NVIDIA Technical Blog来源发表:

新框架接入 Kubernetes 不再等适配器:NVIDIA 把作业结构写成一份可复用 YAML

核心概要

NVIDIA 把 Kubernetes 分布式作业的结构描述从编译进工具的适配器换成一份可读定义:Workload Map 是一个自定义资源加 Go 库,为 JobSet、PyTorchJob、RayCluster 等框架给出统一结构视图。此前每接入一种自定义资源都要写一个适配器,解释角色、把状态折算成阶段、改 pod 模板而不弄坏作业,新类型支持常要等一个发布。定义以 YAML 写明模板位置、副本数与状态映射,调度器与平台共用;示例 JobSet 解析出一个 leader、两个 worker,共 12 块 GPU。

Source-provided article image: Support New AI Frameworks Faster Using NVIDIA Workload Map
Figure 2 · NVIDIA 开发者技术博客

深度剖析

Kubernetes 上的分布式训练作业现在可以用一份共享描述来读,而不只是靠每种框架各自的解析代码。 此前的做法是每接入一种自定义资源就写一个适配器,回答作业结构是什么、状态怎样折算成阶段、改动 pod 模板时如何不破坏作业,适配器开发因此成了按框架排期的路线图条目,新工作负载类型的支持常常要等一个发布。 同一份 JobSet 定义被多个消费方复用:命令行工具把实况对象渲染成 leader 与 worker 的角色树并汇总出示例作业 12 块 GPU 的请求,KAI Scheduler 在没有原生插件的类型上用它兜底做 pod 分组与成组调度,NVIDIA Run:ai 用一次 API 调用注册该工作负载类型而无需额外代码。

每种工作负载类型由一份 YAML 定义描述:pod 模板在哪个字段、副本数怎么算、状态该映射到共享词汇里的哪个阶段。 承载这套描述的是一个 Kubernetes 自定义资源加一个 Go 库,库把集群里的实况对象按定义解析成统一结构视图;JobSet 并不直接汇报正在运行的状态,定义就用一段对副本状态求值的表达式把它推出来,再把原生的已完成与失败条件映射到消费方已有的阶段名。 目录中已带有 JobSet、PyTorchJob、RayCluster、KServe、MPIJob、LeaderWorkerSet、Milvus、NVIDIA Dynamo 等预置并经测试的定义,新增类型只需写一份描述其结构与状态的 YAML。

启示与展望

已演示的范围是 Kubernetes 上受支持的工作负载类型:正文走通的是一个 leader 加两个 worker 的 JobSet,定义以 Apache 2.0 发布,并在 NVIDIA Run:ai 的生产环境中使用,Elotl 的多集群调度器 Nova 也用它解读作业结构与资源需求来选择集群;不在目录里的类型通过写一份 YAML 定义接入。

本篇以一个 JobSet 走通全流程,角色树、12 块 GPU 的汇总与成组分组都出自同一份定义,其他框架的预置定义没有给出同等细度的输出。KAI Scheduler 把这类定义当作没有原生插件时的兜底,原生插件仍然优先,因此分组与成组调度的行为取决于目标类型是否已有插件。