Distributed Training with TrainJob

TrainJob is the Kubernetes-native API for submitting distributed training and fine-tuning workloads on OpenShift. You define what to train, how many nodes to use, and which framework to run — TrainJob handles Pod creation, worker coordination, failure recovery, and cleanup.

TrainJob is provided by the KubeFlow Training Operator v2. The operator runs in the background; as a user, you interact only with TrainJob and the training runtimes that configure it.

This module walks through creating TrainJobs for PyTorch and MPI workloads, using training runtimes for reusable configurations, and integrating with Kueue for admission control.

What TrainJob Replaces

Without TrainJob, running distributed training on Kubernetes requires you to:

  • Manually create Pods, Services, and ConfigMaps for each training worker

  • Handle coordination between workers (who is the master? what are the peer addresses?)

  • Detect and handle partial failures (one worker crashes — do you restart or abort?)

  • Manage cleanup after training completes

This is the same problem that traditional HPC schedulers solve with job submission scripts and mpirun. TrainJob provides an equivalent abstraction for Kubernetes: you submit a job definition, and the system handles the rest.

How TrainJob Works

Submitting a TrainJob

A TrainJob is a short custom resource that describes your training workload:

  • Runtime reference — Which training runtime to use (PyTorch distributed, MPI, etc.).

  • Trainer — Your container image, the number of nodes, resources per node, and any environment variables or hyperparameters.

That is the minimum. You do not need to define individual Pods, configure peer discovery, or set up SSH keys. The system builds all of that from your TrainJob spec and the referenced runtime.

apiVersion: kubeflow.org/v2alpha1
kind: TrainJob
metadata:
  name: my-training
spec:
  runtimeRef:
    name: torch-distributed    (1)
  trainer:
    image: my-registry/my-training:latest
    numNodes: 4                (2)
    resourcesPerNode:          (3)
      requests:
        cpu: "2"
        memory: "4Gi"
1 References a ClusterTrainingRuntime that defines how PyTorch distributed training is set up.
2 The operator creates 4 worker Pods and configures torchrun with the correct --nnodes and rendezvous endpoint.
3 Each worker Pod gets these resource requests.

Training Runtimes

A ClusterTrainingRuntime (cluster-scoped) or TrainingRuntime (namespace-scoped) defines a reusable template for how training jobs run. It specifies:

  • The framework (PyTorch distributed, MPI, etc.)

  • Pod templates for workers (and master/launcher if applicable)

  • Default resource requests and limits

  • Gang scheduling integration

Runtimes are typically created by administrators. As a user, you reference a runtime by name in your TrainJob and override only what you need — image, node count, resources, and environment variables.

This separation keeps TrainJob specs simple and portable. The same TrainJob definition can run on different clusters that each have their own runtime configurations.

Supported Frameworks

TrainJob supports multiple distributed training frameworks through different runtimes:

  • PyTorch Distributed — Uses torchrun / torch.distributed for data-parallel and model-parallel training. The system creates workers and configures rendezvous (peer discovery) so that torch.distributed.init_process_group() works automatically.

  • MPI — Uses a launcher/worker pattern. The system creates an MPI launcher Pod that runs mpirun/mpiexec across worker Pods. This is the same pattern used in HPC: one launcher, many workers, MPI for communication.

  • TensorFlow — Parameter server or multi-worker strategies.

  • Other frameworks — Extensible to additional runtimes via CRDs.

MPI support in TrainJob v2 is being actively developed. Red Hat Engineering is working on MVP MPI support for customer use cases of MPI on OpenShift.

Exercise 1: Creating a Distributed PyTorch TrainJob

In this exercise you will create a distributed PyTorch training job using TrainJob.

Prerequisites

Ensure the Training Operator is installed on your cluster.

  1. Verify the operator is running:

    oc get pods -n kubeflow

    You should see the training-operator Pod in Running state.

  2. Create a namespace for training:

    oc new-project training-demo --skip-config-update || oc project training-demo

Create a PyTorch TrainJob

  1. Create a simple distributed PyTorch training job:

    cat <<EOF | oc apply -f -
apiVersion: kubeflow.org/v2alpha1
kind: TrainJob
metadata:
  name: pytorch-distributed-demo
  namespace: training-demo
spec:
  runtimeRef:
    name: torch-distributed
  trainer:
    image: docker.io/kubeflowkatib/pytorch-mnist:v0.17.0
    numNodes: 4
    resourcesPerNode:
      requests:
        cpu: "1"
        memory: "2Gi"
EOF

    This TrainJob tells the system to run pytorch-mnist across 4 nodes using the torch-distributed runtime. Behind the scenes, the operator will:

    • Create 4 worker Pods (node-0 through node-3).

    • Configure torchrun with the correct --nnodes, --nproc_per_node, and rendezvous endpoint.

    • Coordinate startup so all workers connect before training begins.

    • Monitor for completion or failure.

  2. Watch the training Pods:

    oc get pods -n training-demo -l training.kubeflow.org/trainjob-name=pytorch-distributed-demo -w

    Press Ctrl+C once all Pods are Running.

  3. Check the TrainJob status:

    oc get trainjob pytorch-distributed-demo -n training-demo -o yaml

    The status section shows the training progress and conditions.

  4. View the logs from the workers:

    oc logs -n training-demo -l training.kubeflow.org/trainjob-name=pytorch-distributed-demo --tail=20

Exercise 2: Customizing TrainJobs with Runtime Overrides

TrainJobs reference a runtime for defaults, then override specific fields. This lets you change the image, node count, resources, and environment variables without redefining the entire training configuration.

  1. Examine the available ClusterTrainingRuntimes:

    oc get clustertrainingruntimes

  2. Look at what a runtime provides:

    oc get clustertrainingruntime torch-distributed -o yaml

    This shows the Pod templates, framework configuration, and defaults that your TrainJob inherits.

  3. Create a TrainJob that overrides the defaults:

    cat <<EOF | oc apply -f -
apiVersion: kubeflow.org/v2alpha1
kind: TrainJob
metadata:
  name: custom-training
  namespace: training-demo
spec:
  runtimeRef:
    name: torch-distributed
  trainer:
    image: docker.io/kubeflowkatib/pytorch-mnist:v0.17.0
    numNodes: 2
    resourcesPerNode:
      requests:
        cpu: "2"
        memory: "4Gi"
    env:
      - name: EPOCHS
        value: "5"
      - name: BATCH_SIZE
        value: "64"
EOF

    This TrainJob uses the same torch-distributed runtime but changes the node count to 2, increases resources per node, and passes hyperparameters as environment variables.

  4. Monitor the job:

    oc get trainjob -n training-demo
oc get pods -n training-demo -l training.kubeflow.org/trainjob-name=custom-training

Exercise 3: MPI TrainJobs

MPI (Message Passing Interface) is the standard communication library for HPC applications. TrainJob can run MPI workloads using the launcher/worker pattern — identical to how mpirun works on a traditional HPC cluster.

MPI support in TrainJob v2 may require the MPI runtime to be available. Check your cluster’s available ClusterTrainingRuntimes.

  1. Create an MPI-based TrainJob:

    cat <<EOF | oc apply -f -
apiVersion: kubeflow.org/v2alpha1
kind: TrainJob
metadata:
  name: mpi-job-demo
  namespace: training-demo
spec:
  runtimeRef:
    name: mpi-distributed
  trainer:
    image: mpioperator/mpi-pi:latest
    numNodes: 4
    resourcesPerNode:
      requests:
        cpu: "1"
        memory: "1Gi"
EOF

    From your perspective, this looks just like the PyTorch TrainJob — a runtime reference, an image, and a node count. The runtime determines the execution model. For MPI, the system creates:

    • One launcher Pod that runs mpirun (or mpiexec).

    • Four worker Pods that the launcher can SSH into or connect to via MPI.

    • The necessary SSH keys and hostfile configuration.

  2. Watch the MPI job progress:

    oc get pods -n training-demo -l training.kubeflow.org/trainjob-name=mpi-job-demo

  3. View the launcher logs to see MPI output:

    oc logs -n training-demo -l training.kubeflow.org/trainjob-name=mpi-job-demo,training.kubeflow.org/replica-type=launcher

Exercise 4: Running TrainJobs with Kueue

TrainJob integrates with Kueue for admission control. When Kueue is installed, TrainJobs are submitted to a LocalQueue just like regular Jobs — add the queue label and Kueue handles the rest.

This is especially important for training workloads because distributed training requires all workers to start together. Kueue ensures that a TrainJob is only admitted when the ClusterQueue has enough quota for all workers simultaneously, providing gang-scheduling-like behavior through admission control.

This exercise assumes you have already set up Kueue with a ClusterQueue and LocalQueue as described in the Kueue module. If you have not, complete that module first or use the existing queues on your cluster.

  1. Verify that a LocalQueue exists in your namespace:

    oc get localqueues -n training-demo

    If no LocalQueue exists, create one that points to your ClusterQueue:

    cat <<EOF | oc apply -f -
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  name: training-queue
  namespace: training-demo
spec:
  clusterQueue: cluster-queue
EOF

  2. Submit a TrainJob to the queue by adding the kueue.x-k8s.io/queue-name label:

    cat <<EOF | oc apply -f -
apiVersion: kubeflow.org/v2alpha1
kind: TrainJob
metadata:
  name: queued-training
  namespace: training-demo
  labels:
    kueue.x-k8s.io/queue-name: training-queue
spec:
  runtimeRef:
    name: torch-distributed
  trainer:
    image: docker.io/kubeflowkatib/pytorch-mnist:v0.17.0
    numNodes: 4
    resourcesPerNode:
      requests:
        cpu: "2"
        memory: "4Gi"
EOF

    The only difference from the earlier TrainJobs is the kueue.x-k8s.io/queue-name label. Everything else about TrainJob works the same — Kueue controls when the job is admitted, not how it runs.

  3. Check the Kueue workload status:

    oc get workloads -n training-demo

    You should see a workload corresponding to your TrainJob. If the ClusterQueue has sufficient quota, the workload will be admitted and the training Pods will start. If not, it will remain pending until quota is available.

  4. Watch the TrainJob and its Pods:

    oc get trainjob -n training-demo
oc get pods -n training-demo -l training.kubeflow.org/trainjob-name=queued-training

  5. To see how Kueue manages quota across multiple TrainJobs, submit a second job while the first is still running:

    cat <<EOF | oc apply -f -
apiVersion: kubeflow.org/v2alpha1
kind: TrainJob
metadata:
  name: queued-training-2
  namespace: training-demo
  labels:
    kueue.x-k8s.io/queue-name: training-queue
spec:
  runtimeRef:
    name: torch-distributed
  trainer:
    image: docker.io/kubeflowkatib/pytorch-mnist:v0.17.0
    numNodes: 4
    resourcesPerNode:
      requests:
        cpu: "2"
        memory: "4Gi"
EOF

  6. Check both workloads:

    oc get workloads -n training-demo
oc get trainjob -n training-demo

    If the ClusterQueue does not have enough quota for both jobs at once, the second TrainJob will wait in the queue until the first completes and releases its resources.

Clean Up

  1. Remove all training resources:

    oc delete trainjobs --all -n training-demo

Summary

TrainJob provides a simple, consistent API for submitting distributed training workloads on OpenShift.

In this module you:

  • Learned what TrainJob replaces — manual Pod creation, worker coordination, failure handling, and cleanup for distributed training.

  • Created a distributed PyTorch TrainJob that ran across multiple worker nodes.

  • Customized TrainJobs with runtime overrides — changing node counts, resources, and hyperparameters while inheriting runtime defaults.

  • Submitted an MPI TrainJob using the launcher/worker pattern familiar to HPC users.

  • Submitted TrainJobs through Kueue for admission control, and observed how Kueue queues training workloads when quota is constrained.

In the next section, you will learn about Apache Spark on OpenShift for distributed data processing workloads.