NVIDIA DSX AI工厂平台整合了多项技术，帮助生态系统设计、构建和运维更高效的AI工厂，其中包括为支持多租户运营和扩展而构建的开源、模块化软件组件。
运营一个AI工厂需要对分布式工作负载有共同的理解：它们的组件、Pod模板、副本数量和状态。在Kubernetes上，每个框架对这些信息的编码方式都不同。JobSet、LeaderWorkerSet等自定义资源往往需要单独的适配器，工具才能解析和管理它们。
NVIDIA Workload Map提供了一套通用的Kubernetes工作负载结构描述方式，用户可以在不同框架之间重复使用。作为NVIDIA DSX AI工厂平台的一部分，Workload Map使集成的调度器、控制器和平台能够解释所支持的工作负载类型，而无需维护各自独立的框架专用解析代码。
本文将展示如何一次性描述一个分布式JobSet工作负载，然后将该定义复用于工作负载检查、团体调度（gang scheduling）以及平台可见性。
为什么每种工作负载类型都需要自己的适配器
在Kubernetes上，真实的工作负载很少是单个Pod。一个分布式训练任务或多节点推理服务器，其本质是一个由资源构成的小型树状结构：若干组Pod，每组都有角色、副本数量、Pod模板，以及一个取决于工作负载类型的状态。
构建调度器、控制器、仪表盘或成本工具的团队，若需要适配这些不同的工作负载，很快就会遇到这个问题。NVIDIA Run:ai团队在构建自家平台时也面临同样的挑战。对于每一个新的自定义资源定义（CRD），开发人员都需要在代码中回答相同的问题：
该工作负载类型的结构是什么？
该资源的状态如何被解读为某个阶段？
如何在不破坏工作负载的情况下修改Pod规格？
开发人员针对每种工作负载类型回答一次这些问题，就会产生一个针对该CRD的适配器。客户引入的每一个新框架都需要再开发一个适配器。这使得适配器开发成为每个框架都要单独排期的路线图项目，而不是一次性的成本投入，这也是为什么对新工作负载类型的支持常常要等到某个版本发布才能实现的原因。
下面的manifest展示了一个用户提交的分布式训练任务示例：一个包含一个leader Job和两个worker Job的JobSet，每个Job中的Pod都请求4个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
该manifest描述了一个leader、两个worker，共计请求12个GPU。要理解这一结构，需要具备JobSet相关知识的代码；而kubectl只会将这些资源分别列出。
图1. kubectl将demo 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引用（Owner References）将各个Job与其JobSet关联起来，但平台仍然需要一份定义才能解读它们各自的角色并汇总资源请求。
如果没有一套共享的工作负载描述，用户必须自己重建该工作负载的结构，并手动统计其GPU请求总量。
一套通用的Kubernetes工作负载结构描述方式
Workload Map（之前名为Karta；该CRD的资源kind仍保留为Karta）以Kubernetes CRD加Go库的形式发布，用于解析工作负载定义。每个定义都指明了该工作负载类型的Pod模板、副本数量和状态字段的位置。
Workload Map为调度和平台使用方提供了一个共享的结构化视图。KAI Scheduler利用这些定义来解析Pod分组和团体调度指令，而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定义来描述其结构和状态，从而添加其他工作负载类型。
关键之处在于，工作负载相关的知识不再固化于编译后的适配器中，而是转化为开发人员可以编写、审查，并像其他manifest一样应用到集群中的数据。本文接下来将以JobSet为例，展示同一份定义如何支持工作负载检查、团体调度以及平台可见性。
为JobSet编写工作负载定义
一份工作负载定义，以Workload Map资源的形式表达，用于描述这种结构。以下片段展示了JobSet定义中的关键部分。完整版本发布在相关代码仓库中。
根组件定义了该工作负载自身的kind，以及如何读取其状态。
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处于ready且active状态，且没有Job失败），并将JobSet原生的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所称的一个实例（instance）。podTemplateSpecPath指向每个Job的Pod模板。scaleDefinition使用jq计算：每个复制Job对应的Pod数量等于replicas乘以parallelism，这也是示例manifest中显式设置parallelism的原因。instanceIdPath从.spec.replicatedJobs[].name中获取每个实例的名称，这正是“leader”和“worker”能够作为不同角色区分开，而非表现为三个匿名Job的原因。
至此，该框架特有的结构已被整合进一份可复用的定义中。下一节将展示不同的使用方如何使用它。
在多种工具中复用同一份描述
同一份JobSet定义支持三种工作流程：检查工作负载、将其一起调度，以及将其作为一个整体工作负载进行管理。
使用Workload Map CLI检查工作负载
Workload Map CLI从集群中读取工作负载，并通过其工作负载定义进行渲染展示。对于demo JobSet，该CLI会加载JobSet描述，解析该实时对象及其Pod，并打印出树状结构。该CLI通过kubeconfig访问集群，并内置了多种工作负载定义。此检查流程无需将Workload Map安装到集群中即可完成。
以下命令在macOS或Linux上通过Homebrew安装该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
在当前Kubernetes上下文中运行demo JobSet的情况下，以下命令可用于检查它。其中的<namespace>占位符需替换为包含该工作负载的命名空间。
kli describe jobset/demo -n <namespace>
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实现团体调度
对于需要调度工作负载的使用方而言，一份定义还可以包含可选的团体调度指令。
optimizationInstructions: gangScheduling: podGroup: name: job subGroups: - componentName: replicatedjob
该定义提供分组信息；KAI Scheduler负责强制执行团体准入（gang admission）并做出调度决策。对于没有原生插件支持的工作负载类型，KAI Scheduler会将Workload Map作为其Pod分组器的回退方案。一份匹配的定义能够提供结构和分组信息，而无需开发新的框架专用插件；同时，原生插件仍保持优先级。
在NVIDIA Run:ai中管理同一工作负载
在NVIDIA Run:ai中，只需一次API调用，即可使用同一份目录定义注册该工作负载类型，无需编写任何额外代码。随后，该平台便可利用这份描述将该JobSet作为一个整体工作负载进行管理。
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工厂的GitHub组织获取。如果你正在Kubernetes上运行AI工作负载，并需要一种一致的方式来解析不同框架的结构，Workload Map有助于弥补这一空缺。
可以先从工作负载目录入手，查找你已使用框架对应的定义，或对其进行调整以支持自定义的工作负载类型。欢迎通过Slack邀请链接加入NVIDIA DSX AI工厂的Slack工作区参与讨论。我们非常期待看到你的创作成果。
