NVIDIA DSX AI Factory Platform brings together technologies to help the ecosystem design, build, and operate more efficient AI factories, including open source, modular software components built to support multi-tenant operations and scaling.
Operating an AI factory requires a shared understanding of distributed workloads: their components, pod templates, replica counts, and status. On Kubernetes, each framework encodes that information differently. JobSet, LeaderWorkerSet, and other custom resources often require separate adapters before a tool can interpret and manage them.
NVIDIA Workload Map provides a common description of Kubernetes workload structure that users can reuse across frameworks. As part of the NVIDIA DSX AI Factory Platform, Workload Map enables integrated schedulers, controllers, and platforms to interpret supported workload types without maintaining separate framework-specific parsing code.
This post shows how to describe a distributed JobSet workload once, then reuse that definition for workload inspection, gang scheduling, and platform visibility.
Why every workload type needs its own adapter
On Kubernetes, a real workload is rarely a single pod. A distributed training job or a multi-node inference server is a small tree of resources: groups of pods, each with a role, a replica count, a pod template, and a status that depends on the workload type.
Teams building schedulers, controllers, dashboards, or cost tools that must work across these workloads encounter this immediately. The NVIDIA Run:ai team faced the same challenge while creating its platform. For each new custom resource definition (CRD), developers answer the same questions in code:
- What’s the structure of the workload type?
- How is the resource’s status interpreted as a phase?
- How is the pod spec changed without breaking the workload?
A developer answers those questions once per workload type and produces a per-CRD adapter. Each new framework a customer brings requires another adapter. That makes adapter development a roadmap item per framework, not a one-time cost, which is why support for a new workload type so often waits for a release.
The following manifest shows a distributed training job as a user submits it: a JobSet with a leader Job and two worker Jobs, each with a pod requesting four GPUs.
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
The manifest describes one leader, two workers, and a total request of 12 GPUs. Interpreting that structure requires JobSet-aware code; kubectl lists the resources separately.
Figure 1. kubectl lists the demo JobSet, its three Jobs, and their pods as separate resource entries
$ 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 references connect the Jobs to their JobSet, but a platform still needs a definition to interpret their roles and aggregate resource requests.
Without a shared workload description, users must reconstruct the workload and total its GPU requests themselves.
A common description of Kubernetes workload structure
Workload Map (formerly Karta; the CRD’s resource kind remains Karta) ships as a Kubernetes CRD plus a Go library that interprets workload definitions. Each definition specifies where to find pod templates, replica counts, and status fields for its workload type.
Workload Map supplies a shared structural view to scheduling and platform consumers. KAI Scheduler uses its definitions to interpret pod grouping and gang-scheduling instructions, while NVIDIA Run:ai uses them for workload management and visibility. Each consumer acts on the same description, reducing the need for separate framework-specific parsers.
The Go library resolves a live object into a uniform structural view. With a workload definition in hand, a consumer can:
- Extract pod templates, replica counts, status, and metadata.
- Update pod specs, labels, and annotations across every instance.
- Understand the workload hierarchy, for example, a JobSet’s replicated Jobs and their pods.
Workload Map includes a catalog of prebuilt, tested definitions for JobSet, PyTorchJob, RayCluster, RayJob, RayService, KServe, Knative, MPIJob, NVIDIA NIM, LeaderWorkerSet, Milvus, Grove PodCliqueSet, and NVIDIA Dynamo. Additional workload types can be added by writing a YAML definition that describes their structure and status.
The key is that workload knowledge moves out of compiled adapters and into data developers can write, review, and apply to a cluster like any other manifest. The rest of this post describes the JobSet example, then shows how the same definition supports workload inspection, gang scheduling, and platform visibility.
Writing a workload definition for JobSet
A workload definition, expressed as a Workload Map resource, describes that shape. The following excerpt shows key parts of the JobSet definition. The full version ships in the repository.
The root component defines the workload’s own kind and how its status is read.
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 is the translation layer. JobSet does not report a plain “Running” condition, so the definition derives it with a jq expression over replicatedJobsStatus (any Job ready and active, none failed) and maps the native Completed and Failed conditions onto the shared vocabulary a consumer already uses. When a RayJob or a LeaderWorkerSet is described, its very different status maps to the same categories.
The next excerpt defines the child component: the structure inside the workload.
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
A JobSet declares its roles as a list, so .spec.replicatedJobs holds one entry per role. The [] in these paths means “for each entry in that list,” and each entry is what Workload Map calls an instance. podTemplateSpecPath points to each Job’s pod template. scaleDefinition uses jq: the number of pods per replicated Job equals replicas times parallelism, which is why the example manifest sets parallelism explicitly. instanceIdPath takes each instance’s name from .spec.replicatedJobs[].name, which is how “leader” and “worker” show up as distinct roles instead of three anonymous Jobs.
The framework-specific structure is now captured in one reusable definition. The next section shows how different consumers use it.
Using one description across tools
The same JobSet definition supports three workflows: inspecting workloads, scheduling them together, and managing them as one workload.
Inspecting workloads with the Workload Map CLI
The Workload Map CLI reads workloads from a cluster and renders them through their workload definitions. For the demo JobSet, the CLI loads the JobSet description, resolves the live object and its pods, and prints the tree. The CLI uses a kubeconfig to access the cluster and includes built-in workload definitions. This inspection workflow does not require installing Workload Map into the cluster.
The following commands install the CLI with Homebrew on macOS or Linux and confirm that the executable is on the PATH.
brew tap dsx-ai-factory/kli https://github.com/dsx-ai-factory/workload-map
brew install dsx-ai-factory/kli/kli
kli --version
With the demo JobSet running in the current Kubernetes context, the following command inspects it. The <namespace> placeholder must be replaced with the namespace containing the workload.
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
The output identifies the leader and worker roles and shows a total request of 12 GPUs for the example workload.
Figure 4. An earlier Karta CLI view resolves the JobSet into leader and worker roles, showing readiness, GPU totals, and node placement
$ 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
The definition supplies the JobSet-specific structure, so the CLI can inspect the workload without a dedicated parser.
Enabling gang scheduling with KAI Scheduler
For consumers that schedule workloads, a definition can also include optional gang-scheduling instructions.
optimizationInstructions:
gangScheduling:
podGroup:
name: job
subGroups:
- componentName: replicatedjob
The definition supplies grouping information; KAI Scheduler enforces gang admission and makes placement decisions. KAI Scheduler uses Workload Map as a fallback in its pod grouper for workload types without a native plugin. A matching definition provides the structure and grouping information without requiring a new framework-specific plugin, while native plugins retain precedence.
Managing the same workload in NVIDIA Run:ai
In NVIDIA Run:ai, a single API call registers the workload type using the same catalog definition with no additional code. The platform then uses that description to manage the JobSet as one workload.
The CLI, KAI Scheduler, and NVIDIA Run:ai use the same definition for different purposes: inspection, scheduling, and platform visibility. For another supported workload type, the corresponding catalog definition supplies its structure through the same interface.
Workload Map adoption across schedulers and platforms
NVIDIA Workload Map is available under the Apache 2.0 license and is used in production within NVIDIA Run:ai. Its declarative definitions support workload visibility, status interpretation, submission, and mutation across supported Kubernetes workload types.
Adoption also extends to multi-cluster scheduling. Elotl’s Nova uses Workload Map to interpret workload structure and resource requirements when selecting a cluster for distributed workloads. By using prebuilt definitions, Nova reduces the need to maintain framework-specific parsing logic, demonstrating how the same workload-description model can support different scheduling and platform consumers.
Getting started
NVIDIA Workload Map is an open source component available through the NVIDIA DSX AI Factory GitHub organization. If you are running AI workloads on Kubernetes and need a consistent way to interpret their structure across frameworks, Workload Map helps close the gap.
Start with the workload catalog to find a definition for a framework you already use, or adapt one to support a custom workload type. Join the conversation on the NVIDIA DSX AI Factory Slack workspace using the Slack invitation link. We can’t wait to see what you build.
