Tools

AI-generated text

NVIDIA Workload Map: a shared Kubernetes workload description for AI platforms

NVIDIA introduced Workload Map (formerly Karta), an open-source Kubernetes CRD and Go library that provides a reusable, declarative description of distributed workload structure for AI platforms.

NVIDIA Workload Map: a shared Kubernetes workload description for AI platforms

NVIDIA Workload Map (formerly Karta) is an open-source Kubernetes CRD plus a Go library that provides declarative, reusable descriptions of distributed AI workload structure. Its purpose is to give schedulers, controllers and platforms a common structural view so they no longer need separate per-framework adapters to interpret pod templates, replica counts and status fields.

Why this is needed

On Kubernetes, real AI workloads are usually trees of resources rather than single pods: distributed training jobs or multi-node inference services contain roles, replica counts, pod templates and status information. Different frameworks (for example JobSet, LeaderWorkerSet, PyTorchJob, RayCluster) expose these structures with different CRD schemas, so platforms historically implemented an adapter for each new CRD. That creates repeated development and maintenance work for every framework a customer wants to support.

The example manifest in the article defines a JobSet with one leader Job and two worker Jobs; each pod requests 4 GPUs, producing a total request of 12 GPUs. kubectl lists the JobSet, the Jobs and the pods as separate resources, so an external platform must reassemble the logical workload and compute totals unless it understands the JobSet schema.

How Workload Map works

Workload Map ships as a Karta CRD and a Go library. Every workload definition specifies where to find pod templates, how to compute replica counts and which fields represent status. The library resolves a live Kubernetes object into a uniform tree-like structural view so consumers can:

  • extract pod templates, replica counts, status and metadata,
  • update pod specs, labels and annotations across every instance,
  • understand workload hierarchy (for example a JobSet’s replicated Jobs and their pods).

The project provides a catalog of prebuilt definitions for many workload types, including JobSet, PyTorchJob, RayCluster, RayJob, RayService, KServe, Knative, MPIJob, NVIDIA NIM, LeaderWorkerSet, Milvus, Grove PodCliqueSet and NVIDIA Dynamo. New types can be added by writing a YAML definition that describes their structure and status.

The central idea is to move workload knowledge out of compiled, framework-specific adapters and into data (YAML manifests) that developers can write, review and apply to clusters.

JobSet example: writing a workload definition

A Karta resource’s root component names the workload kind and explains how to read its status. The definition’s statusMappings translate native CRD fields into shared phases a consumer understands. For JobSet, which does not expose a simple "Running" condition, the Workload Map definition derives Running with a jq expression over .status.replicatedJobsStatus (for example: any job ready and active and none failed). Completed and Failed native conditions are mapped to the same shared vocabulary so different frameworks’ statuses become comparable.

The childComponents section describes the workload’s internal structure. JobSet declares roles as a list under .spec.replicatedJobs; Workload Map uses paths such as .spec.replicatedJobs[].template.spec.template for the podTemplateSpecPath and computes replicas with a jq expression (.replicas * .template.spec.parallelism). instanceIdPath reads each instance’s name (e.g. "leader" and "worker") so roles appear with identifiers rather than as anonymous Jobs.

Captured this way, framework-specific layout becomes a reusable definition that any consumer can interpret.

Using one description across tools

  • Workload Map CLI: the CLI loads a workload definition from the catalog, resolves the live object and its pods and prints a workload tree. The CLI can be installed with Homebrew and does not require Workload Map to be installed in the cluster. For the demo JobSet the CLI shows leader and worker roles and an aggregate request of 12 GPUs.

  • KAI Scheduler: definitions can include optional gang-scheduling instructions (optimizationInstructions → gangScheduling) that provide grouping information. KAI Scheduler enforces gang admission and placement using these grouping rules and falls back to Workload Map when a native plugin for a workload type is not available.

  • NVIDIA Run:ai: a single API call registers a workload type using the same catalog definition with no additional code. Run:ai then treats the JobSet as a single workload, using the description for management and visibility (for example presenting the JobSet as one workload with 12 GPUs and a structure view showing leader and worker Jobs).

The CLI, KAI Scheduler and NVIDIA Run:ai use the same definition for inspection, scheduling and platform visibility respectively.

Adoption and licensing

NVIDIA Workload Map is available under the Apache 2.0 license and is used in production inside NVIDIA Run:ai. The declarative definitions support visibility, status interpretation, submission and mutation across the supported Kubernetes workload types. Multi-cluster schedulers also adopt the approach: Elotl’s Nova relies on Workload Map to interpret workload structure and resource needs when selecting clusters for distributed workloads, avoiding framework-specific parsing logic.

Getting started

Workload Map is hosted in the NVIDIA DSX AI Factory GitHub organization. Users running AI workloads on Kubernetes can start with the workload catalog to find or adapt a definition for a framework they use. The project also maintains a Slack workspace for community discussion.

Workload Map aims to reduce duplicated adapter work by making workload descriptions declarative, reviewable and cluster-applicable, so different schedulers and platforms can rely on the same structural knowledge.