Kubeflow Trainer’s TrainJob API splits a distributed training job into a user-owned trainer block and a platform-owned runtime object, so the pod template that actually runs is no longer in the manifest the user submits.

The v1 Training Operator gave each framework its own CRD — PyTorchJob, MPIJob, XGBoostJob — and each job carried a complete pod template per replica role. Every user who submitted a job also chose the base image, the security context, the volume mounts, and the rendezvous wiring. Kubeflow Trainer v2 replaces those with one kind, TrainJob in trainer.kubeflow.org/v1alpha1, plus two runtime kinds: ClusterTrainingRuntime, which is cluster-scoped, and TrainingRuntime, which is namespaced. The project ships runtimes for torch, DeepSpeed, JAX, MLX, XGBoost, and torchtune; the released chart installs them when runtimes.defaultEnabled=true.

A TrainJob names a runtime and supplies only what the user actually owns. The upstream multi-node example is the whole contract:

apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainJob
metadata:
  name: multi-node-example
spec:
  runtimeRef:
    apiGroup: trainer.kubeflow.org
    kind: ClusterTrainingRuntime
    name: torch-distributed
  trainer:
    numNodes: 3
    command:
      - torchrun
      - --no-python
      - python3
      - -c
      - |
        import torch.distributed as dist
        dist.init_process_group(backend="nccl")

Where the pod template lives now

The referenced ClusterTrainingRuntime holds the parts the platform team owns. spec.mlPolicy carries the framework policy — for the shipped torch runtime, numNodes: 1 and a torch: {} block — and spec.template.spec.replicatedJobs carries a JobSet template whose node job supplies the container image. The shipped torch-distributed runtime pins a PyTorch CUDA runtime image there. A TrainJob that sets no image inherits it; a TrainJob that sets one overrides it for that job only.

The trainer block is the user-facing surface. numNodes is the number of training nodes. numProcPerNode is the number of processes, workers, or slots on every training node, typed as an integer in the CRD schema. resourcesPerNode is a standard ResourceRequirements object, so GPU requests go where they always did:

  trainer:
    numNodes: 3
    numProcPerNode: 8
    resourcesPerNode:
      requests:
        cpu: "2"
        memory: "4Gi"
      limits:
        nvidia.com/gpu: 8

Underneath, the controller renders a JobSet. That is the fact that matters for day-2 operations, because the pods are labelled by JobSet, not by TrainJob. Logs and pod listings go through the JobSet name:

kubectl get trainjob multi-node-example
kubectl get pods -l jobset.sigs.k8s.io/jobset-name=multi-node-example
kubectl logs -l jobset.sigs.k8s.io/jobset-name=multi-node-example

One more block sits beside trainer in the TrainJob spec. spec.initializer configures dataset and model initialization as separate containers, each with its own env, so weight-pulling and dataset staging stop being hand-written init containers copied between manifests. That is the same substitution the runtime split performs on the training container: a step that every job needed and every job re-specified becomes a field with a controller behind it. The consequence for debugging is that a job stuck before step 0 may be stuck in an initializer the manifest never names explicitly, which is why the JobSet-labelled pod listing above is the first command rather than the job’s own status.

The rendezvous contract also changed shape. The v1 PyTorchJob injected RANK, WORLD_SIZE, MASTER_ADDR, and MASTER_PORT directly into each container. Trainer injects the PET_* variables that torchrun reads — PET_NNODES, PET_NPROC_PER_NODE, PET_NODE_RANK, PET_MASTER_ADDR, PET_MASTER_PORT — and torchrun is what exposes WORLD_SIZE, RANK, and LOCAL_RANK to the training process. A script that reads RANK from the environment without being launched under torchrun sees nothing.

The queueing seam

The split also gives Kueue a clean insertion point, because admission now toggles one field on a small object instead of gating a framework-specific CRD. The upstream Kueue example sets two things: a label naming the LocalQueue, and suspend: true so the job starts held.

apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainJob
metadata:
  name: kueue-example
  labels:
    kueue.x-k8s.io/queue-name: training-queue
spec:
  runtimeRef:
    apiGroup: trainer.kubeflow.org
    kind: ClusterTrainingRuntime
    name: torch-distributed
  suspend: true
  trainer:
    numNodes: 3

Kueue admits the workload by flipping suspend to false once the LocalQueue’s backing ClusterQueue has capacity. The admitted object is visible as a Kueue Workload:

kubectl get workload -l kueue.x-k8s.io/job-name=kueue-example

spec.managedBy selects which controller owns the TrainJob; it defaults to the built-in trainjob controller and can be set to kueue.x-k8s.io/multikueue to delegate scheduling across clusters. The repository also ships a Volcano integration example, so gang scheduling remains an option at the same seam.

Failure modes

A runtimeRef naming a runtime that is not installed is the first wall. Runtimes are separate objects from the controller, and a chart installed without the default runtimes leaves a cluster where every TrainJob references torch-distributed and nothing exists to render. Check with kubectl get clustertrainingruntimes before debugging the job.

Omitting suspend: true under Kueue is subtler. The TrainJob runs immediately, consuming GPUs that the ClusterQueue believes are available, and the quota system reports capacity it does not have. Nothing errors; the queue is simply bypassed for that job.

Version skew has a documented sharp edge. The v2.3.0 release moved CRDs into the Helm chart template directory and removed runtime finalizers, and the release notes state that an upgrade from v2.0, v2.1, or v2.2 must pass through v2.3 before going any further. Skipping it is not a rollback-and-retry situation on a cluster with live training jobs.

Decision frame

The question when evaluating a move from PyTorchJob to TrainJob is not which API is newer. It is who currently owns the training pod template. If one platform team already reviews every job manifest, the runtime split moves that review from pull requests into a ClusterTrainingRuntime and the user-facing surface collapses to numNodes, numProcPerNode, resourcesPerNode, and a command. If each research team pins its own image, security context, and volumes per job, the runtime objects multiply until there is one per team and the split has bought nothing. The second question is whether the cluster already runs Kueue: the suspend plus queue-label contract is the cheapest part of the migration and the part that changes GPU utilization.