Kubeflow Trainer v2.4 Adds Workload-Aware Scheduling for TrainJob

Why distributed training needs group scheduling

Standard Kubernetes scheduling uses a greedy approach, attempting to schedule pods one at a time. According to Kubeflow, that model creates problems for distributed training, where workers may need to start together and be placed close to one another.

In one example from the project, a training job requiring 64 GPUs may find capacity for 63 GPUs and remain idle while it waits for the final GPU. Another example involves a data initializer and 64 GPU workers. If the workers consume the available quota, the initializer cannot run, while the workers cannot start without the data, creating a dependency deadlock.

Kubeflow Trainer v2.4 uses Kubernetes Workload API building blocks, including Workload and PodGroup, to support gang scheduling and topology-aware scheduling. The TrainJob controller creates these resources so that the scheduler can admit pods as a unified group. The project says this can help prevent expensive GPUs from sitting unused and reduce communication latency.

Three levels map workloads to infrastructure

Trainer v2.4 introduces a three-tier scheduling hierarchy. At Level 1, the TrainJob level, a PodGroup represents the entire training job. The physical target is a cloud zone, using topology.kubernetes.io/zone. This is designed for uniform workloads such as PyTorch DDP, where all workers should be scheduled, placed and preempted as one unit.

At Level 2, the ReplicatedJob level, each configured entry can receive its own PodGroup. This supports different policies within one pipeline. For example, a dataset initializer can use a basic policy and run individually, while an MPI launcher and GPU workers can use a gang policy and be coordinated within the same zone.

At Level 3, the Job or per-replica level targets rack or Top-of-Rack switch placement through topology.kubernetes.io/rack. Kubeflow describes this level as supporting maximum network throughput and failure isolation. In v2.4, the targeted ReplicatedJob can have at most one replica because per-replica PodGroups require Kubernetes v1.38.

Platform administrators configure the policy

The scheduling configuration is set in the runtime by platform administrators. Data scientists can then specify a value such as numNodes in their TrainJob without manually creating the underlying workload and PodGroup resources.

For a distributed PyTorch example, the source shows a TrainJob with numNodes: 100 and a request for four NVIDIA GPUs per node. The TrainJob controller converts the requested node count into a gang scheduling minimum, creating a Workload and an automatically generated PodGroup associated with the resulting pod specification.

Kubeflow says this design hides low-level scheduling resources from practitioners while giving platform engineers control over placement and scheduling behavior. The release therefore connects TrainJob configuration with the physical topology of AI data centers, with the aim of coordinating distributed workloads and avoiding the resource waste and dependency deadlocks described by the project.