Support New AI Frameworks Faster Using NVIDIA Workload Map

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.

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.

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

Owner references connect the Jobs to their JobSet, but a platform still needs a definition to interpret their roles and aggregate resource requests.

Figure 2. NVIDIA Run:ai displays the component Jobs as three separate workloads, each requesting four GPUs, before the JobSet definition is registered

Without a shared workload description, users must reconstruct the workload and total its GPU requests themselves.

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.

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.

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.

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.

Figure 3. Declarative definitions resolve JobSet and LeaderWorkerSet resources into a common workload-tree representation despite their different schemas

The framework-specific structure is now captured in one reusable definition. The next section shows how different consumers use it.

The same JobSet definition supports three workflows: inspecting workloads, scheduling them together, and managing them as one workload.

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.

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.

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

The definition supplies the JobSet-specific structure, so the CLI can inspect the workload without a dedicated parser.

For consumers that schedule workloads, a definition can also include optional gang-scheduling instructions.

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.

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.

Figure 5. NVIDIA Run:ai presents the JobSet as one workload with 12 GPUs and a structure view showing its leader and worker Jobs

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.

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.

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.