OpenShift AI’s TrainJob API turns distributed training into a recoverable Kubernetes workload

Red Hat has validated Kubeflow Trainer v2 on OpenShift AI 3.4 from the TrainJob API down to multi-node NVIDIA GPU execution, checkpoint recovery and JobSet failure policies. The result is a clearer boundary between the simple path most training teams need and the lower-level controls platform engineers still have to own.

Trainer v2 replaces framework-specific resources such as PyTorchJob, TFJob and MPIJob with one TrainJob API. A TrainJob references a ClusterTrainingRuntime and states the code, node count and resources per node. The runtime generates the torchrun configuration, environment variables and pod coordination; JobSet supplies Kubernetes-native orchestration underneath.

What Red Hat tested

The first phase compared a two-node TrainJob with a raw JobSet. Both allocated NVIDIA GPUs, created working inter-pod networking through headless services and produced the expected NCCL all-reduce result. The difference was configuration: TrainJob generated the distributed setup, while the raw JobSet required explicit master address, port, rank, world size and torchrun arguments.

The second phase trained a PyTorch fraud-detection model over 1.29 million simulated transactions on one, two and four NVIDIA A10G GPUs. Red Hat reports 1.94-times speedup on two GPUs and 3.8-times on four, or 95% scaling efficiency at four GPUs. Training time fell from more than 15 minutes to four while model quality remained stable.

The test also exercised failure paths. Periodic and shutdown checkpoints were stored in MinIO. Suspending a two-GPU TrainJob removed both pods, freed the GPUs and resumed from the saved step when the job restarted. A forced pod deletion caused JobSet to recreate the training group and recover from the latest checkpoint.

Where JobSet still matters

TrainJob is the concise interface for a homogeneous distributed run. Raw JobSet becomes useful when a workload has multiple roles or policies. Red Hat combined a CPU-only data-preparation job with two GPU training pods, using dependency ordering so accelerators were not allocated before data was ready.

JobSet also allowed different failure actions by stage: fail the entire workload when data preparation breaks, but restart training after a transient pod or node failure. Coordinator labels removed hard-coded rendezvous addresses, and volume-claim policies handled scratch storage lifecycle.

Storage remains a deployment decision. Red Hat’s cluster only offered ReadWriteOnce volumes, so multi-node sharing used MinIO; a ReadWriteMany class could support a shared automatically provisioned volume instead.

What to try first

OpenShift AI teams should begin with TrainJob and the built-in torch-distributed runtime, then move down to JobSet only when they need heterogeneous stages, custom startup order or stage-specific failure policies. Before production use, verify checkpoint durability, gang scheduling through Kueue, available volume modes and GPU telemetry under the same node-failure conditions the cluster must survive.

The important change is not merely fewer lines of YAML. Distributed training now has a common Kubernetes contract with explicit recovery and resource-management behavior.