数据中心 / 云计算
借助 NVIDIA Workload Map 更快支持新型 AI 框架
2026 年 11 月 02 日
NVIDIA DSX AI Factory Platform 汇聚了多项技术,助力生态伙伴设计、构建并运营更高效的 AI 工厂。该平台包含开源的模块化软件组件,专为支持多租户操作与规模扩展而打造。
运营 AI 工厂需要各方对分布式工作负载达成共识:包括其组成部分、Pod 模板、副本数量及运行状态。在 Kubernetes 环境中,不同框架对这些信息的编码方式各不相同。JobSet、LeaderWorkerSet 及其他自定义资源通常需要先经过独立的适配器转换,工具才能正确解析并进行管理。
NVIDIA Workload Map 提供了一套统一的 Kubernetes 工作负载结构描述,用户可在不同框架间重复使用。作为 NVIDIA DSX AI Factory Platform 的一部分,Workload Map 使集成的调度器、控制器和平台能够直接解析所支持的工作负载类型,无需为每种框架单独维护解析代码。
本文将演示如何一次性描述一个分布式 JobSet 工作负载,随后将该定义复用于负载检查、并行调度及平台可视化监控。
为何每种工作负载类型都需要专属适配器
在 Kubernetes 中,真实的工作负载极少仅由单个 Pod 构成。分布式训练作业或多节点推理服务器实际上是一棵小型的资源树:由多个 Pod 组组成,每个 Pod 组都具备特定角色、副本数、Pod 模板,以及依赖于工作负载类型的运行状态。
开发跨工作负载运行的调度器、控制器、仪表盘或成本核算工具的团队会立刻遇到这一难题。NVIDIA Run:ai 团队在构建自身平台时也面临同样的挑战。针对每一个新的自定义资源定义(CRD),开发者都必须在代码中回答相同的问题:
- 该工作负载类型的结构是怎样的?
- 如何将资源的运行状态映射为阶段(phase)?
- 如何在修改 Pod 规格的同时不破坏原有工作负载?
开发者只需针对每种工作负载类型回答一次这些问题,即可生成对应的单 CRD 适配器。客户引入的新框架都需要额外的适配器。这使得适配器开发成为按框架规划的长期任务,而非一次性投入。这也正是新工作负载类型的支持往往要等到新版本发布才能上线的原因。
以下清单文件展示了用户提交分布式训练作业时的原始形态:一个包含主节点 Job 和两个工作节点 Job 的 JobSet,每个 Pod 均申请四张 GPU。
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
name: demo
spec:
replicatedJobs:
- name: leader
replicas: 1
template:
spec:
parallelism: 1
completions: 1
template:
spec:
containers:
- name: leader
image: ghcr.io/example/trainer:latest
resources:
limits:
nvidia.com/gpu: 4
- name: worker
replicas: 2
template:
spec:
parallelism: 1
completions: 1
template:
spec:
containers:
- name: worker
image: ghcr.io/example/trainer:latest
resources:
limits:
nvidia.com/gpu: 4
清单文件描述了一个 Leader、两个 Worker,以及总共 12 张 GPU 的请求。解析该结构需要感知 JobSet 的代码;而 kubectl 会将这些资源分别列出。
图 1. kubectl 将示例 JobSet、其包含的三个 Job 以及对应的 Pod 列为独立资源条目
$ kubectl get jobset,jobs,pods
NAME TERMINALSTATE RESTARTS COMPLETED SUSPENDED AGE
jobset.jobset.x-k8s.io/demo 0 3m46s
NAME STATUS COMPLETIONS DURATION AGE
job.batch/demo-leader-0 Running 0/1 3m46s 3m46s
job.batch/demo-worker-0 Running 0/1 3m46s 3m46s
job.batch/demo-worker-1 Running 0/1 3m46s 3m46s
NAME READY STATUS RESTARTS AGE
pod/demo-leader-0-0-ksb4l 1/1 Running 0 3m47s
pod/demo-worker-0-0-qv955 1/1 Running 0 3m47s
pod/demo-worker-1-0-5smfn 1/1 Running 0 3m47s
Owner 引用将 Job 与其所属的 JobSet 关联起来,但平台仍需一个定义来解析各组件的角色并汇总资源请求。
图 2. 在注册 JobSet 定义之前,NVIDIA Run:ai 将各个组件 Job 显示为三个独立的负载,每个负载均申请四张 GPU
若缺乏统一的负载描述,用户必须自行还原负载结构并手动汇总其 GPU 请求。
Kubernetes 负载结构的通用描述
Workload Map(前身为 Karta;CRD 的资源类型仍为 Karta)以 Kubernetes CRD 配合 Go 库的形式发布,用于解析负载定义。每个定义都指明了如何获取对应负载类型的 Pod 模板、副本数量及状态字段。
Workload Map 为调度器和平台消费者提供共享的结构化视图。KAI Scheduler 利用其定义来解析 Pod 分组与 Gang Scheduling 指令,而 NVIDIA Run:ai 则将其用于负载管理与可视化。所有消费者基于同一份描述进行操作,从而减少了对特定框架专用解析器的依赖。
该 Go 库可将运行时对象转换为统一的结构化视图。借助已有的负载定义,消费者能够执行以下操作:
- 提取 Pod 模板、副本数量、状态及元数据。
- 跨所有实例更新 Pod 规格、标签与注解。
- 理解负载层级关系,例如 JobSet 中复制的 Job 及其下属的 Pod。
Workload Map 内置了一份经过测试的定义目录,涵盖 JobSet、PyTorchJob、RayCluster、RayJob、RayService、KServe、Knative、MPIJob、NVIDIA NIM、LeaderWorkerSet、Milvus、Grove PodCliqueSet 以及 NVIDIA Dynamo。如需支持其他负载类型,只需编写描述其结构与状态的 YAML 定义即可。
核心在于,负载知识已从编译好的适配器中剥离,转为数据开发者可编写、审查并像普通清单一样部署到集群中的配置。本文后续将继续以 JobSet 为例,展示同一份定义如何支持负载检查、Gang Scheduling 及平台可视化。
为 JobSet 编写负载定义
工作负载定义以 Workload Map 资源的形式表达,用于描述该结构。以下摘录展示了 JobSet 定义的关键部分。完整版本随仓库一同发布。
根组件定义了工作负载自身的类型及其状态的读取方式。
apiVersion: run.ai/v1alpha1
kind: Karta
spec:
structureDefinition:
rootComponent:
name: jobset
kind:
group: jobset.x-k8s.io
version: v1alpha2
kind: JobSet
statusDefinition:
conditionsDefinition:
path: .status.conditions
statusMappings:
running:
- byExpression:
expression: "(.status.replicatedJobsStatus // []) | any(.ready > 0 and .active > 0) and all(.failed == 0)"
expectedResult: "true"
completed:
- byConditions:
- type: Completed
status: "True"
failed:
- byConditions:
- type: Failed
status: "True"
statusMappings 是翻译层。由于 JobSet 不直接报告“Running”状态条件,该定义通过 jq 表达式对 replicatedJobsStatus 进行推导(即存在至少一个就绪且活跃的 Job,且无失败 Job),并将原生的 Completed 和 Failed 条件映射到消费者已使用的共享词汇中。当描述 RayJob 或 LeaderWorkerSet 时,它们截然不同的状态也会被映射到相同的类别中。
下一段摘录定义了子组件:即工作负载内部的结构。
childComponents:
- name: replicatedjob
kind: { group: batch, version: v1, kind: Job }
ownerRef: jobset
specDefinition:
podTemplateSpecPath: .spec.replicatedJobs[].template.spec.template
scaleDefinition:
replicasPath: .spec.replicatedJobs[] | .replicas * .template.spec.parallelism
instanceIdPath: .spec.replicatedJobs[].name
JobSet 将角色声明为列表,因此 .spec.replicatedJobs 中每个角色对应一个条目。这些路径中的 [] 表示“遍历该列表中的每个条目”,而每个条目正是 Workload Map 所称的实例。podTemplateSpecPath 指向每个 Job 的 Pod 模板。scaleDefinition 使用 jq 语法:每个副本 Job 的 Pod 数量等于 replicas 乘以 parallelism,这也是示例清单显式设置 parallelism 的原因。instanceIdPath 从 .spec.replicatedJobs[].name 获取每个实例的名称,这使得“leader”和“worker”能够作为独立角色呈现,而不是三个匿名 Job。
图 3. 尽管 JobSet 和 LeaderWorkerSet 资源的架构不同,但声明式定义仍能将它们解析为统一的工作负载树形表示
框架特定的结构现已封装在一个可复用的定义中。下一节将展示不同消费者如何使用它。
跨工具使用同一份描述
同一份 JobSet 定义支持三种工作流:检查工作负载、将它们统一调度,以及将其作为一个整体进行管理。
使用 Workload Map CLI 检查工作负载
Workload Map CLI 从集群中读取工作负载,并通过其工作负载定义进行渲染。以演示用的 JobSet 为例,该 CLI 会加载 JobSet 描述文件,解析实时对象及其 Pod,并输出树状结构。CLI 使用 kubeconfig 访问集群,并内置了工作负载定义。此检查工作流无需将 Workload Map 安装到集群中。
以下命令通过 Homebrew 在 macOS 或 Linux 上安装该 CLI,并确认可执行文件已加入 PATH 环境变量。
brew tap dsx-ai-factory/kli https://github.com/dsx-ai-factory/workload-map
brew install dsx-ai-factory/kli/kli
kli --version
当演示用的 JobSet 在当前 Kubernetes 上下文中运行时,可使用以下命令进行检查。请将占位符替换为包含该工作负载的命名空间。
kli describe jobset/demo -n
JobSet/demo namespace: default definition: jobset-x-k8s-io-jobset-v1alpha2 (catalog) age: 20s
`-- replicatedjob 3/3 ready gpu: 12 runai-cluster-worker-0-0
|-- leader 1/1 ready gpu: 4 runai-cluster-worker-0-0
| `-- demo-leader-0-0-ck6kg Running gpu: 4 runai-cluster-worker-0-0
`-- worker 2/2 ready gpu: 8 runai-cluster-worker-0-0
|-- demo-worker-0-0-j9pdc Running gpu: 4 runai-cluster-worker-0-0
`-- demo-worker-1-0-rrsxz Running gpu: 4 runai-cluster-worker-0-0
Phase: Running
Resources:
COMPONENT REPLICAS GPU CPU MEMORY
leader 1 4 0 0
worker 2 8 0 0
TOTAL 3 12 0 0
输出结果明确了 Leader 和 Worker 角色,并显示示例工作负载总共请求了 12 个 GPU。
图 4。早期的 Karta CLI 视图将 JobSet 解析为 Leader 和 Worker 角色,展示了就绪状态、GPU 总数及节点分布情况。
$ karta workload tree demo
JobSet/demo [Running]
└─ replicatedjob (3/3 replicas) 3/3 ready gpu: 12 runai-cluster-worker-0-2,runai-cluster-worker-0-5,runai-cluster-worker-0-7
├─ leader (1/1 replicas) 1/1 ready gpu: 4 runai-cluster-worker-0-7
│ └─ Pod/demo-leader-0-0-ksb4l Running gpu: 4 runai-cluster-worker-0-7
└─ worker (2/2 replicas) 2/2 ready gpu: 8 runai-cluster-worker-0-2,runai-cluster-worker-0-5
├─ Pod/demo-worker-0-0-qv955 Running gpu: 4 runai-cluster-worker-0-5
└─ Pod/demo-worker-1-0-5smfn Running gpu: 4 runai-cluster-worker-0-2
$ karta workload list
NAMESPACE NAME KIND PHASE COMPONENTS GPU AGE
karta-demo demo JobSet Running leader(1), worker(2) 12 3m
该定义提供了 JobSet 专属的结构信息,因此 CLI 无需专用解析器即可检查工作负载。
启用 KAI Scheduler 的 Gang 调度
对于负责调度工作负载的消费者,定义文件中还可以包含可选的 Gang 调度指令。
optimizationInstructions:
gangScheduling:
podGroup:
name: job
subGroups:
- componentName: replicatedjob
该定义提供分组信息;KAI Scheduler 负责执行任务组准入并做出放置决策。对于缺乏原生插件的工作负载类型,KAI Scheduler 会在其 Pod 分组器中将 Workload Map 作为回退方案。匹配的定义可在不新增框架专属插件的前提下提供结构与分组信息,同时原生插件仍享有优先权。
在 NVIDIA Run:ai 中管理同一工作负载
在 NVIDIA Run:ai 中,仅需一次 API 调用即可借助相同的目录定义注册工作负载类型,无需编写额外代码。平台随后会依据该描述将 JobSet 作为一个整体工作负载进行管理。
图 5. NVIDIA Run:ai 将包含 12 块 GPU 的 JobSet 呈现为单一工作负载,并提供结构视图以展示其 Leader 与 Worker Job。
CLI、KAI Scheduler 与 NVIDIA Run:ai 分别出于不同目的使用该定义:用于检查、调度与平台可见性。对于其他受支持的工作负载类型,对应的目录定义也会通过相同接口提供其结构信息。
Workload Map 在各调度器与平台中的采用情况
NVIDIA Workload Map 采用 Apache 2.0 开源协议授权,并已投入生产环境用于 NVIDIA Run:ai。其声明式定义支持跨受支持的 Kubernetes 工作负载类型的可见性查看、状态解析、提交与变更操作。
该方案的采用范围也已扩展至多集群调度场景。Elotl 的 Nova 调度器利用 Workload Map 来解析分布式工作负载的结构与资源需求,从而进行集群选择。借助预置定义,Nova 减少了对维护框架专属解析逻辑的需求,充分证明了同一工作负载描述模型如何能够服务于不同的调度器与平台消费者。
快速入门
NVIDIA Workload Map 是一个开源组件,可通过 NVIDIA DSX AI Factory GitHub 组织获取。如果您正在 Kubernetes 上运行 AI 工作负载,且需要一种跨框架统一解析其结构的方法,Workload Map 将有效弥合这一差异。
您可以从工作负载目录入手,查找已有框架的定义,或对其进行适配以支持自定义工作负载类型。欢迎通过 Slack 邀请链接加入 NVIDIA DSX AI Factory 的 Slack 频道参与讨论。我们非常期待看到您的创新成果。
