trainer/proposals/2170-kubeflow-trainer-v2

KEP-2170: Kubeflow Trainer V2 API

Authors

• Andrey Velichkevich - @andreyvelich
• Yuki Iwai - @tenzen-y

Google doc: https://bit.ly/3WzjTlw

Overview

This document discusses the new Kubeflow Training V2 API.

When we built the Kubeflow Training Operator a couple of years ago, Kubernetes lacked better features to support distributed machine learning (ML) training, such as SuccessPolicy and RestartPolicy (PodFailurePolicy in Job). Recently, the Kubernetes community launched the working group Batch, and then the working group actively worked on evolving the batch/v1 Job API and built a new Kubernetes SIGs project: JobSet to manage groups of Jobs.

This document consolidates efforts for the Cloud Native ML Training between Kubeflow and Kubernetes communities.

Motivation

We often implement features similar to batch/v1 Job, such as “suspend”, on the Training Operator side since the Training Operator creates blocks of plain Pod and Service for each rank once Kubeflow Jobs are created. However, if we continue taking the same approach to use lowest level abstractions that introduce redundancy, the maintenance costs will continue to increase.

Replacing repetitive infrastructure layers with JobSet would help to avoid redundancy and reduce developer toil.

Additionally, introducing JobSet as an infrastructure layer would allow us to introduce batch workload features such as the PodFailurePolicy and the PodDisruptionCondition easily.

Please also see the Kubernetes JobSet and Kubeflow Training Operator collaboration document.

User Value

In addition to the above motivation, we will address the following user feedback while implementation:

• Confusion around Workers: https://github.com/kubeflow/training-operator/issues/1790
• Support batch/v1 Job features: https://github.com/kubeflow/training-operator/issues/1718
• ExitCodes for PodFailurePolicy: https://github.com/kubeflow/training-operator/issues/1749
• Migrate to MPI V2 API: https://github.com/kubeflow/training-operator/issues/1906

Personas

We can identify the following personas of Training Operator:

1. DevOps Engineer. They are familiar with Kubernetes concepts and they know how to manage the Kubernetes workloads. Usually, they are not experts in ML frameworks and ML algorithms.
2. MLOps Engineer. They are familiar with ML frameworks and they know how to configure distributed PyTorch settings such as rendezvous backends or MPI configuration. Usually, they are not experts in Kubernetes and ML algorithms.
3. Data Scientists/ML Engineers. They create model architectures and advanced ML algorithms to train models. They prefer to use Python for their work. They are aware of torch.nn APIs, but not with torch.distributed and Kubernetes concepts to scale model training.

Based on the above personas, we should build an API that everyone will benefit from.

Goals

• Introduce the TrainingRuntime and ClusterTrainingRuntime APIs that will store blueprints for model training and LLM fine-tuning using various ML frameworks. These runtimes will be built on top of JobSet APIs with additional functionality for special use-cases. For example, training using MPI orchestration.
• Introduce Kubeflow TrainJob API that allows to reuse these runtimes and quickly start a new training job without understanding complex Kubernetes APIs.
• Enable TrainJob to extend TrainingRuntime defaults through the RuntimePatches API — a typed, multi-owner patch mechanism that lets users, external controllers (e.g. Kueue), and admission webhooks each contribute a named patch entry to configure JobSet/Job metadata and Pod specs (e.g. volumes, scheduling directives, etc.) without conflicting with one another.
• Update Kubeflow Training SDK to allow data scientists quickly create and monitor TrainJobs.
• Create community-supported ClusterTrainingRuntime for distributed training with PyTorch and MPI.
• Create community-supported ClusterTrainingRuntime for LLM fine-tuning for various foundational models (e.g. Mistral, LLama-70b, Gemma-7b).
• Work on the following JobSet improvements:
  • For PyTorch Elastic: https://github.com/kubernetes-sigs/jobset/issues/463
  • For PVC management: https://github.com/kubernetes-sigs/jobset/issues/572
  • For PyTorch Elastic: https://github.com/kubernetes-sigs/jobset/issues/570
• Integrate TrainJob with Kueue and MultiKueue to effectively manage resources for training jobs and orchestrate resources across multiple clusters.

Non-Goals

• Support MPI V1 implementation.
• Distributed training for TensorFlow, XGboost, JAX, and PaddlePaddle will be added after initial implementation.
• Migrate Kubeflow V1 controller to use JobSet.
• Propose the migration mechanisms / ways from Kubeflow Training v1 to v2. We will create dedicated KEP for customers migration.
• Propose the changes to Kubeflow Training Python SDK. After controller implementation, we will propose changes to the kubeflow-training SDK.

Design Details

We propose these APIs:

• TrainJob: A single API which allows data scientists to initiate a training and fine-tuning job from the pre-deployed training runtime. It allows users to tweak configurations for their training jobs such as model parameters, dataset parameters, or trainer configuration. The main goal is to hide unnecessary Kubernetes complexity for data scientists.

• TrainingRuntime and ClusterTrainingRuntime: Set of blueprints for how to start various types of training or fine-tuning jobs. Runtimes are managed by Platform Engineers and allow them to configure infrastructure parameters that are required for the TrainJob. For example, failure policy or gang-scheduling.

User Roles Diagram

The below diagram shows how platform engineers manage TrainingRuntime and how data scientists create TrainJob:

TrainJob can be created using kubectl or Kubeflow Python SDK.

LLM Fine-Tuning Diagram

The below diagram shows which resources will be created for LLM fine-tuning with PyTorch:

Worker and Node Definition

To better understand what "Nodes" and "Worker" mean in the diagram above, the following table explains the naming that each framework or technology uses:

ML Framework or Technology	Definition of a Single Device (GPU)	Definition of a Single VM	Start Command	Reference Docs
Kubernetes	Container Resource Unit	Pod’s Container	Any	Resource units in K8s
PyTorch	Worker (--nproc-per-node)	Node (--nnodes)	torchrun	PyTorch Elastic
MPI (OpenMPI)	Slot (-n)	Node (-host)	mpirun	Reference for OpenMPI
TensorFlow	Worker	Worker Pool Cluster Spec	python	TensorFlow Distributed
Jax	Process jax.local_devices()	Host jax.devices()	python or mpirun	Jax Distributed
PaddlePaddle	Worker	Node	python -m paddle.distributed.launch	Paddle Distributed
XGBoost	Worker	Not Applicable	python	Rabit Tracker for c10d
DeepSpeed	Slot	Node (--num_nodes)	deepspeed	DeepSpeed Distributed

Additionally, check this document for the mpirun command for other MPI implementations: Intel MPI, MPICH, Spectrum MPI.

The TrainJob API

The TrainJob exposes APIs that data scientist can override in TrainingRuntime to create a training job:

type TrainJob struct {
	metav1.TypeMeta `json:",inline"`

	// Standard object's metadata.
	metav1.ObjectMeta `json:"metadata,omitempty"`

	// Specification of the desired TrainJob.
	Spec TrainJobSpec `json:"spec,omitempty"`

	// Current status of TrainJob.
	Status TrainJobStatus `json:"status,omitempty"`
}

const (
	// TrainJobSuspended means the TrainJob is suspended.
	TrainJobSuspended string = "Suspended"

	// TrainJobComplete means that the TrainJob has completed its execution.
	TrainJobComplete string = "Complete"

	// TrainJobFailed means that the actual jobs have failed its execution.
	TrainJobFailed string = "Failed"
)

const (
	// TrainJobSuspendedReason is the "Suspended" condition reason.
	// When the TrainJob is suspended, this is added.
	TrainJobSuspendedReason string = "Suspended"

	// TrainJobResumedReason is the "Suspended" condition reason.
	// When the TrainJob suspension is changed from True to False, this is added.
	TrainJobResumedReason string = "Resumed"
)


// TrainJobSpec represents specification of the desired TrainJob.
type TrainJobSpec struct {
	// Reference to the training runtime.
	// The field is immutable.
	// +kubebuilder:validation:XValidation:rule="self == oldSelf", message="runtimeRef is immutable"
	RuntimeRef RuntimeRef `json:"runtimeRef"`

	// Configuration of the initializer.
	Initializer *Initializer `json:"initializer,omitempty"`

	// Configuration of the trainer.
	Trainer *Trainer `json:"trainer,omitempty"`

	// runtimePatches defines custom patches applied to the TrainJob's Runtime.
	// Patches are keyed by manager to provide clear ownership and avoid conflicts between controllers.
	// +listType=map
	// +listMapKey=manager
	RuntimePatches []RuntimePatch `json:"runtimePatches,omitempty"`

	// Whether the controller should suspend the running TrainJob.
	// Defaults to false.
	// +kubebuilder:default=false
	Suspend *bool `json:"suspend,omitempty"`

	// ManagedBy is used to indicate the controller or entity that manages a TrainJob.
	// The value must be either an empty, `trainer.kubeflow.org/trainjob-controller` or
	// `kueue.x-k8s.io/multikueue`. The built-in TrainJob controller reconciles TrainJob which
	// don't have this field at all or the field value is the reserved string
	// `trainer.kubeflow.org/trainjob-controller`, but delegates reconciling TrainJobs
	// with a 'kueue.x-k8s.io/multikueue' to the Kueue. The field is immutable.
	// Defaults to `trainer.kubeflow.org/trainjob-controller`
	// +kubebuilder:default="trainer.kubeflow.org/trainjob-controller"
	// +kubebuilder:validation:XValidation:rule="self in ['trainer.kubeflow.org/trainjob-controller', 'kueue.x-k8s.io/multikueue']", message="ManagedBy must be trainer.kubeflow.org/trainjob-controller or kueue.x-k8s.io/multikueue if set"
	// +kubebuilder:validation:XValidation:rule="self == oldSelf", message="ManagedBy value is immutable"
	ManagedBy *string `json:"managedBy,omitempty"`
}

type RuntimeRef struct {
	// Name of the runtime being referenced.
	// When namespaced-scoped TrainingRuntime is used, the TrainJob must have
	// the same namespace as the deployed runtime.
	Name string `json:"name"`

	// APIGroup of the runtime being referenced.
	// Defaults to `trainer.kubeflow.org`.
	APIGroup *string `json:"apiGroup,omitempty"`

	// Kind of the runtime being referenced.
	// It must be one of TrainingRuntime or ClusterTrainingRuntime.
	// Defaults to ClusterTrainingRuntime.
	Kind *string `json:"kind,omitempty"`
}

// Initializer represents the desired configuration for the dataset and model initialization.
// It is used to initialize the assets (dataset and pre-trained model) and pre-process data.
type Initializer struct {
	// Configuration of the dataset initialization and pre-processing.
	Dataset *DatasetInitializer `json:"dataset,omitempty"`

	// Configuration of the pre-trained model initialization
	Model *ModelInitializer `json:"model,omitempty"`
}

type TrainJobStatus struct {
	// Conditions for the TrainJob.
	Conditions []metav1.Condition `json:"conditions,omitempty"`

	// JobsStatus tracks the child Jobs in TrainJob.
	JobsStatus []JobStatus `json:"jobsStatus,omitempty"`
}

type JobStatus struct {
	// Name of the child Job.
	Name string `json:"name"`

	// Ready is the number of child Jobs where the number of ready pods and completed pods
	// is greater than or equal to the total expected pod count for the child Job.
	Ready int32 `json:"ready"`

	// Succeeded is the number of successfully completed child Jobs.
	Succeeded int32 `json:"succeeded"`

	// Failed is the number of failed child Jobs.
	Failed int32 `json:"failed"`

	// Active is the number of child Jobs with at least 1 pod in a running or pending state
	// which are not marked for deletion.
	Active int32 `json:"active"`

	// Suspended is the number of child Jobs which are in a suspended state.
	Suspended int32 `json:"suspended"`
}

This table explains the rationale for each TrainJob parameter:

Parameter	What is it ?
RuntimeRef	Reference to the existing TrainingRuntime that is pre-deployed by platform engineers
Initializer	Configuration for the dataset and model initialization.
Trainer	Configuration for the Trainer such as image, number of nodes, accelerators.
RuntimePatches	Structured patches applied to the TrainJob's Runtime, keyed by manager. Used to inject runtime-specific configuration such as user identity, PVC mounts, node selectors, or scheduling constraints. Typically managed by admission webhooks or external controllers (e.g. Kueue) after the user creates the TrainJob via the Python SDK or kubectl.
Suspend and ManagedBy	Scheduling directives for Kueue and MultiKueue

Example of TrainJob

apiVersion: trainer.kubeflow.org/v2alpha1
kind: TrainJob
metadata:
  name: torch-ddp
  namespace: tenant-alpha
spec:
  runtimeRef:
    name: torch-distributed-multi-node
  trainer:
    image: docker.io/custom-training
    command:
      - torchrun train.py
    numNodes: 5
    resourcesPerNode:
      requests:
        nvidia.com/gpu: 2

The container's torchrun command in the above YAML will be converted into:

torchrun --nnodes=5 --nproc-per-node=2 train.py

Additionally, the Kubeflow Training SDK allows the user to create the above TrainJob using the Python API:

def train_func():
    import torch
    class Net(torch.nn.Module):
        """Create the PyTorch Model"""
        ...
    model = Net()

    # Attach model to the distributor
    torch.distributed.init_process_group(backend="nccl")
    model = torch.nn.parallel.DistributedDataParallel(model)

    # Train model
    model.train()

# Use Kubeflow SDK to create TrainJob.
from kubeflow.training import TrainingClient

TrainingClient().train(
    name="torch-ddp",
    func=train_func,
    num_nodes=5,
    resources_per_node={"gpu": 2},
)

Example of LLM Fine-Tuning

This example shows how to create TrainJob to fine-tune LLama 7b:

apiVersion: trainer.kubeflow.org/v2alpha1
kind: TrainJob
metadata:
  name: tune-llama-with-yelp
  namespace: tenant-alpha
spec:
  runtimeRef:
    name: torch-tune-llama-7b
  initializer:
    dataset:
      storageUri: s3://dataset/custom-dataset/yelp-review

The Trainer API

The Trainer represents the APIs that data scientists can use to configure the trainer settings. This trainer is executed on every distributed training Node.

User can override the default parameters for the trainer container of the node Job. The runtime Pod template must contain the following label to identify relationship between PodSpec <-> .trainJob.spec.trainer:

trainer.kubeflow.org/trainjob-ancestor-step: trainer

type Trainer struct {
	// Docker image for the training container.
	Image *string `json:"image,omitempty"`

	// Entrypoint commands for the training container.
	Command []string `json:"command,omitempty"`

	// Arguments to the entrypoint for the training container.
	Args []string `json:"args,omitempty"`

	// List of environment variables to set in the training container.
	// These values will be merged with the TrainingRuntime's trainer environments.
	Env []corev1.EnvVar `json:"env,omitempty"`

	// Number of training nodes.
	NumNodes *int32 `json:"numNodes,omitempty"`

	// Compute resources for each training node.
	ResourcesPerNode *corev1.ResourceRequirements `json:"resourcesPerNode,omitempty"`

	// Number of processes/workers/slots on every training node.
	// For the MPI runtime only int value can be set.
	// For the Torch runtime the value defaults to `auto` and can be overridden with an int.
	NumProcPerNode *int32 `json:"numProcPerNode,omitempty"`
}

The following tables show how TrainingRuntime fields will be overridden with Trainer.

Parameter of Trainer	Parameter of TrainingRuntime
.numNodes	.spec.numNodes

The next table shows parameters used to override the Trainer container. These parameters are derived from the PodSpec of the ReplicatedJob, which includes the corresponding label:

.spec.replicatedJobs[...].template.spec.template.labels[trainer.kubeflow.org/trainjob-ancestor-step: trainer’]

Parameter of Trainer	Parameter of TrainingRuntime.Spec.ReplicatedJob[...].template.spec.template
.image	.spec.containers[name=’trainer’].image
.command	.spec.containers[name=’trainer’].command
.args	.spec.containers[name=’trainer’].args
.env	.spec.containers[name=’trainer’].env
.resourcesPerNode	.spec.containers[name=’trainer’].resources

The Dataset Initializer API

The DatasetInitializer represents the APIs that data scientists can use to configure the dataset location and pre-process data on CPUs.

type DatasetInitializer struct {
	// Storage uri for the dataset provider.
	StorageUri *string `json:"storageUri,omitempty"`

	// List of environment variables to set in the dataset initializer container.
	// These values will be merged with the TrainingRuntime's dataset initializer environments.
	// +listType=map
	// +listMapKey=name
	Env []corev1.EnvVar `json:"env,omitempty"`

	// Reference to the secret with credentials to download dataset.
	// Secret must be created in the TrainJob's namespace.
	SecretRef *corev1.LocalObjectReference `json:"secretRef,omitempty"`
}

Initially we will support the following dataset providers:

• S3: storageUri: s3://bucket-name/path/dataset
• HuggingFace: storageUri: hf://repo-id

User can override the default env variables for the dataset-initializer container of the dataset-initializer Job. The Runtime Pod template must contain the following label to identify relationship between PodSpec <-> TrainJob:

trainer.kubeflow.org/trainjob-ancestor-step: dataset-initializer

For example:

initializer:
  dataset:
    storageUri: s3://datasets/yelp-review
    env:
      - name: ENDPOINT_URL
        value: s3.custom.com

Will be converted to:

replicatedJobs:
  - name: dataset-initializer
    spec:
      template:
        metadata:
          labels:
            trainer.kubeflow.org/trainjob-ancestor-step: dataset-initializer
        spec:
          containers:
            - name: dataset-initializer
              image: docker.io/kubeflow/dataset-initializer
              env:
                - name: STORAGE_URI
                  value: s3://dataset/yelp-review
                - name: ENDPOINT_URL
                  value: s3.custom.com

The Model Initializer API

The ModelInitializer represents the APIs that data scientists can use to configure the pre-trained model input location.

type ModelInitializer struct {
	// Storage uri for the model provider.
	StorageUri *string `json:"storageUri,omitempty"`

	// List of environment variables to set in the model initializer container.
	// These values will be merged with the TrainingRuntime's model initializer environments.
	// +listType=map
	// +listMapKey=name
	Env []corev1.EnvVar `json:"env,omitempty"`

	// Reference to the secret with credentials to download model.
	// Secret must be created in the TrainJob's namespace.
	SecretRef *corev1.LocalObjectReference `json:"secretRef,omitempty"`
}

Initially we will support the following model providers:

• HuggingFace: storageUri: hf://model-path

User can override the default env variables for the model-initializer container of the model-initializer Job. The Runtime Pod template must contain the following label to identify relationship between PodSpec <-> TrainJob:

trainer.kubeflow.org/trainjob-ancestor-step: model-initializer

For example:

initializer:
  model:
    storageUri: hf://bert-based-cased
    env:
      - name: TRANSFORMER_TYPE
        value: AutoModelForCausalLM

Will be converted to:

replicatedJobs:
  - name: model-initializer
    template:
      spec:
        template:
          metadata:
            labels:
              trainer.kubeflow.org/trainjob-ancestor-step: model-initializer
          spec:
            containers:
              - name: model-initializer
                image: docker.io/kubeflow/model-initializer
                env:
                  - name: STORAGE_URI
                    value: hf://bert-based-cased
                  - name: TRANSFORMER_TYPE
                    value: AutoModelForCausalLM

The RuntimePatches API

runtimePatches allows controllers, admission webhooks, and custom clients to attach structured patches to a TrainJob without conflicting with each other. Common use cases include injecting user identity, PVC mounts, node selectors, or scheduling constraints that are determined outside the TrainJob itself – for example by Kueue during resource allocation or by a multi-tenant admission webhook. The patch is applied to the underlying runtime, allowing users to customize Job template before creation.

Each entry is keyed by a unique manager field (the list map key). A manager owns its entry and updates it in place. The time field is set by the Trainer admission webhook on each write and is used for observability only — it is not a list map key.

The trainingRuntimeSpec field is a discriminated union over supported runtime kinds. For ClusterTrainingRuntime and TrainingRuntime-based jobs it exposes a restricted subset of the runtime spec: JobSet-level metadata, per-replicated-job Job template metadata, and a curated set of Pod spec fields.

In the future, we can extend RuntimePatch API with other supported Runtimes (e.g. GroveRuntime). In that case, users will be able to select one of supported specs, for example: TrainingRuntimeSpec or GroveRuntimeSpec

API Design

// RuntimePatch represents a custom patch applied to the TrainJob's training runtime template.
// Patches are keyed by manager to provide clear ownership and avoid conflicts between controllers.
type RuntimePatch struct {
	// manager indicates who owns this patch entry. It can be set by the user, external
	// controllers, or admission webhooks to track ownership and avoid conflicts.
	// For example, Kueue sets this field to "kueue.x-k8s.io/manager".
	// +kubebuilder:validation:MinLength=1
	// +kubebuilder:validation:XValidation:rule="self == oldSelf", message="field is immutable"
	// +required
	Manager string `json:"manager"`

	// time is the timestamp of when this patch was last written.
	// Set by the Trainer admission webhook on each create or update. Used for observability only,
	// not used as a list map key.
	// +optional
	Time *metav1.Time `json:"time,omitempty"`

	// trainingRuntimeSpec defines allowed patches for ClusterTrainingRuntime or TrainingRuntime-based jobs.
	// +optional
	TrainingRuntimeSpec *TrainingRuntimeSpecPatch `json:"trainingRuntimeSpec,omitempty"`
}

// TrainingRuntimeSpecPatch mirrors TrainingRuntimeSpec but only exposes
// the fields managers are permitted to patch.
type TrainingRuntimeSpecPatch struct {
	// template patches the JobSet template.
	// +optional
	Template *JobSetTemplatePatch `json:"template,omitempty"`
}

// JobSetTemplatePatch defines patches for the JobSet template.
// It mirrors JobSetTemplateSpec but only exposes metadata and restricted spec fields.
type JobSetTemplatePatch struct {
	// metadata patches the JobSet object metadata.
	// Only labels and annotations are allowed.
	// +optional
	Metadata *metav1.ObjectMeta `json:"metadata,omitempty"`

	// spec patches the JobSet spec with restricted fields.
	// +optional
	Spec *JobSetSpecPatch `json:"spec,omitempty"`
}

// JobSetSpecPatch defines allowed patches for the JobSet spec.
type JobSetSpecPatch struct {
	// replicatedJobs defines per-job patches, keyed by job name.
	// +listType=map
	// +listMapKey=name
	// +optional
	ReplicatedJobs []ReplicatedJobPatch `json:"replicatedJobs,omitempty"`
}

// ReplicatedJobPatch defines patches for a specific replicated job within the JobSet.
type ReplicatedJobPatch struct {
	// name is the name of the replicated job to patch.
	// +kubebuilder:validation:MinLength=1
	// +required
	Name string `json:"name"`

	// template patches the Job template for this replicated job.
	// +optional
	Template *JobTemplatePatch `json:"template,omitempty"`
}

// JobTemplatePatch defines patches for a Job template within a replicated job.
type JobTemplatePatch struct {
	// metadata patches the Job template metadata.
	// Only labels and annotations are allowed.
	// +optional
	Metadata *metav1.ObjectMeta `json:"metadata,omitempty"`

	// spec patches the Job spec with restricted fields.
	// +optional
	Spec *JobSpecPatch `json:"spec,omitempty"`
}

// JobSpecPatch defines allowed patches for the Job spec.
type JobSpecPatch struct {
	// template patches the Pod template for this Job.
	// +optional
	Template *PodTemplatePatch `json:"template,omitempty"`
}

// PodTemplatePatch defines patches for a Pod template within a Job.
type PodTemplatePatch struct {
	// metadata patches the Pod template metadata.
	// Only labels and annotations are allowed.
	// +optional
	Metadata *metav1.ObjectMeta `json:"metadata,omitempty"`

	// spec patches the Pod spec with the fields managers are permitted to set.
	// +optional
	Spec *PodSpecPatch `json:"spec,omitempty"`
}

// PodSpecPatch contains the Pod spec fields that managers are permitted to patch.
type PodSpecPatch struct {
	// serviceAccountName patches the service account for the Pods in the target job templates.
	// +kubebuilder:validation:XValidation:rule="self == oldSelf", message="field is immutable"
	// +optional
	ServiceAccountName *string `json:"serviceAccountName,omitempty"`

	// volumes patches the Pod's volumes.
	// +listType=map
	// +listMapKey=name
	// +kubebuilder:validation:XValidation:rule="self == oldSelf", message="field is immutable"
	// +optional
	Volumes []corev1.Volume `json:"volumes,omitempty"`

	// initContainers patches the init containers in the target job templates.
	// +listType=map
	// +listMapKey=name
	// +kubebuilder:validation:XValidation:rule="self == oldSelf", message="field is immutable"
	// +optional
	InitContainers []ContainerPatch `json:"initContainers,omitempty"`

	// containers patches specific containers in the target job templates.
	// +listType=map
	// +listMapKey=name
	// +kubebuilder:validation:XValidation:rule="self == oldSelf", message="field is immutable"
	// +optional
	Containers []ContainerPatch `json:"containers,omitempty"`

	// imagePullSecrets patches the image pull secrets for the Pods in the target job templates.
	// +listType=map
	// +listMapKey=name
	// +kubebuilder:validation:XValidation:rule="self == oldSelf", message="field is immutable"
	// +optional
	ImagePullSecrets []corev1.LocalObjectReference `json:"imagePullSecrets,omitempty"`

	// securityContext patches the Pod's security context.
	// More info: https://kubernetes.io/docs/tasks/configure-pod-container/security-context/
	// +kubebuilder:validation:XValidation:rule="self == oldSelf", message="field is immutable"
	// +optional
	SecurityContext *corev1.PodSecurityContext `json:"securityContext,omitempty"`

	// nodeSelector patches the node selector to place Pods on specific nodes.
	// +optional
	NodeSelector map[string]string `json:"nodeSelector,omitempty"`

	// affinity patches the Pod's scheduling affinity.
	// +optional
	Affinity *corev1.Affinity `json:"affinity,omitempty"`

	// tolerations patches the Pod's tolerations.
	// +listType=atomic
	// +optional
	Tolerations []corev1.Toleration `json:"tolerations,omitempty"`

	// schedulingGates patches the scheduling gates for the Pods in the target job templates.
	// More info: https://kubernetes.io/docs/concepts/scheduling-eviction/pod-scheduling-readiness/
	// +listType=map
	// +listMapKey=name
	// +optional
	SchedulingGates []corev1.PodSchedulingGate `json:"schedulingGates,omitempty"`
}

// ContainerPatch represents parameters that can be patched using PodSpecPatch.
type ContainerPatch struct {
	// name for the container. Runtime must have this container.
	// +kubebuilder:validation:MinLength=1
	// +required
	Name string `json:"name,omitempty"`

	// env is the list of environment variables to set in the container.
	// These values will be merged with the Runtime's environments.
	// These values can't be set for container with the name: `node`, `dataset-initializer`, or
	// `model-initializer`. For those containers the envs can only be set via Trainer or Initializer APIs.
	// +listType=map
	// +listMapKey=name
	// +optional
	Env []corev1.EnvVar `json:"env,omitempty"`

	// volumeMounts are the volumes to mount into the container's filesystem.
	// +listType=map
	// +listMapKey=name
	// +optional
	VolumeMounts []corev1.VolumeMount `json:"volumeMounts,omitempty"`

	// securityContext patches the container's security context.
	// More info: https://kubernetes.io/docs/tasks/configure-pod-container/security-context/
	// +optional
	SecurityContext *corev1.SecurityContext `json:"securityContext,omitempty"`
}

The webhook validates that the container names in Containers and InitContainers exist in the Runtime's Job template. The patches are applied during the build phase of the Pipelines Framework in the ComponentBuilder plugin. the RuntimePatches update values in the TrainJob's Runtime template, since they contain the final desired values for the underlying Job.

Example of TrainJob with Patches

This example shows how to patch the user-identity for the sidecar container and add volume to the trainer container.

apiVersion: trainer.kubeflow.org/v2alpha1
kind: TrainJob
metadata:
  name: pytorch-distributed
  namespace: tenant-alpha
spec:
  runtimeRef:
    name: pytorch-distributed-gpu
  trainer:
    image: docker.io/custom-training
  runtimePatches:
    - manager: trainer.kubeflow.org/kubeflow-sdk
      time: "2026-05-01T15:20:00Z"
      trainingRuntimeSpec:
        template:
          spec:
            replicatedJobs:
              - name: node
                template:
                  metadata:
                    labels:
                      custom-label: value
                  spec:
                    template:
                      spec:
                        containers:
                          - name: trainer
                            volumeMounts:
                              - name: user-123-volume
                                mountPath: /workspace
                        volumes:
                          - name: user-123-volume
                            persistentVolumeClaim:
                              claimName: user-123-volume

The same manager can update its patch entry during the lifecycle of a TrainJob. For example, Kueue might update node selectors as resources become available:

apiVersion: trainer.kubeflow.org/v2alpha1
kind: TrainJob
metadata:
  name: pytorch-distributed
  namespace: tenant-alpha
spec:
  runtimeRef:
    name: pytorch-distributed-gpu
  trainer:
    image: docker.io/custom-training
  runtimePatches:
    - manager: kueue.x-k8s.io/manager
      time: "2026-02-05T10:10:00Z"
      trainingRuntimeSpec:
        template:
          spec:
            replicatedJobs:
              - name: node
                template:
                  spec:
                    template:
                      spec:
                        nodeSelector:
                          node-type: gpu-a100
                          zone: us-west-1b

Support Arbitrary Patches (Future Extension)

The long-term goal of the Trainer Extension Framework is to support arbitrary runtimes, including operators that are not part of the upstream Trainer project, such as in-house or third-party runtime implementations. For example, users may operate a custom operator to orchestrate distributed AI workloads that Trainer has no schema knowledge of.

To enable this flexibility, the runtimePatches API supports custom runtimes through an opaque spec patch field. One approach that balances flexibility with API simplicity is to store the spec patch as a raw object:

// +kubebuilder:validation:XValidation:rule="[has(self.trainingRuntimeSpec), has(self.opaqueRuntimeSpec)].filter(x, x).size() <= 1", message="only one spec can be set"
type RuntimePatch struct {
	Manager             string                    `json:"manager"`
	Time                *metav1.Time              `json:"time,omitempty"`
	TrainingRuntimeSpec *TrainingRuntimeSpecPatch `json:"trainingRuntimeSpec,omitempty"`
	// opaqueRuntimeSpec stores the full spec patch for runtimes whose schema is unknown to Trainer.
	// Used when runtimeRef points to a runtime kind not natively supported by Trainer.
	// +optional
	OpaqueRuntimeSpec *runtime.RawExtension `json:"opaqueRuntimeSpec,omitempty"`
}

The following example shows how opaqueRuntimeSpec can represent a patch for a SparkApplication CRD:

runtimePatches:
  - manager: kueue.x-k8s.io/manager
    time: "2026-02-05T10:10:00Z"
    opaqueRuntimeSpec:
      driver: # Spark Driver spec
        template: # Valid PodTemplateSpec
          spec:
            nodeSelector:
              node-type: gpu-a100
      executor: # Spark Executor spec
        template: # Valid PodTemplateSpec
          spec:
            nodeSelector:
              node-type: gpu-a100

Since Trainer has no schema knowledge of the target runtime, opaqueRuntimeSpec content is not validated at admission time beyond confirming it is a valid JSON object. Validation occurs at reconcile time when the patch is applied to the target CR.

State Transition

In this section, we're explaining the TrainJob state transition (.status.conditions). The basic TrainJob state machine is the below. Especially, if we specify the TrainingRuntime or ClusterTrainingRuntime as a runtime, the TrainJob terminal condition (Failed or Complete) is decided based on the JobSet terminal state (status.terminalState) instead of computing from JobSet conditions.

stateDiagram-v2
    #CREATION
    state created_choice <<choice>>
    [*] --> created_choice: TrainJob is created.
    created_choice --> Failed=True: Failed to resolve training runtime reference.

    #INITIALISATION
    state resources_applied_choice <<choice>>
    created_choice --> resources_applied_choice: Apply TrainJob runtime resources.
    resources_applied_choice --> resources_applied_choice: Backoff and retry on error.

    #SUSPENSION
    state suspended_choice <<choice>>
    resources_applied_choice --> suspended_choice: Handle TrainJob suspension.
    suspended_choice --> Suspended=True: TrainJob is suspended.
    Suspended=True --> Suspended=True: Wait for unsuspending.
    Suspended=True --> Suspended=False: TrainJob is resumed.
    suspended_choice --> Suspended=False: TrainJob is not suspended.

    #EXECUTION
    state terminal_choice <<choice>>
    Suspended=False --> terminal_choice: Actual Jobs go to execution phase.
    terminal_choice --> Failed=True: Actual Jobs (e.g., JobSet) failed.
    Failed=True --> [*]

    #COMPLETION
    terminal_choice --> Complete=True: Actual Jobs (e.g., JobSet) completed.
    Complete=True --> [*]

In the above state transition, the Created=False will happen in the following situations and those different situations can be identified by the condition reasons (.status.conditions.[type="Created"].reason).

• JobsBuildFailed: When the TrainJob controller failed to construct objects (resources) using the runtime framework interfaces
• JobsCreationFailed: When the TrainJob controller succeeded to construct objects, but it failed to deploy objects to the cluster.

Additionally, we extend the runtime framework interfaces to allow each plugin to propagate the arbitrary conditions to the TrainJob.

The Training Runtime API

The TrainingRuntime is the pre-created configurations of model training on the cluster, representing as blueprints. For example, Elastic PyTorch training, MPI DeepSpeed configuration, BERT LLM Fine-Tuning.

These blueprints can be deployed within the Training Operator control plane and stored in a Kubeflow public repository that users can apply to their clusters.

Platform or ML engineers can tweak existing blueprints, based on their requirements. For example, using custom configurations.

The Kubeflow Training Operator can maintain more Training Runtimes when the community is ready to support them. For example, runtimes for Jax or MLX. We will support PyTorch and MPI runtimes. After initial implementation, we will support TensorFlow, XGboost, and PaddlePaddle runtimes, but it is out of scope for this KEP.

The TrainingRuntime is immutable, and so to make a change, a new version of the TrainingRuntime must be created and then the user must change the TrainJob to point to the new version. This provides control as to how changes to runtimes propagate to existing training jobs. For example, when training is running for a long time (e.g. 1-2 months).

In the future implementation, we will introduce a revision control mechanism similar to Kubernetes Deployment to control versions of TrainingRuntime and enable rolling updates.

We are going to create two CRDs: TrainingRuntime and ClusterTrainingRuntime. These runtimes have exactly the same APIs, but the first one is the namespace-scoped, the second is the cluster-scoped. User can set the kind and apiGroup parameters in the runtimeRef to use the TrainingRuntime from the TrainJob's namespace, otherwise the ClusterTrainingRuntime will be used.

type ClusterTrainingRuntime struct {
	metav1.TypeMeta `json:",inline"`

	// Standard object's metadata.
	metav1.ObjectMeta `json:"metadata,omitempty"`

	// Specification of the desired ClusterTrainingRuntime.
	Spec TrainingRuntimeSpec `json:"spec,omitempty"`
}

type TrainingRuntime struct {
	metav1.TypeMeta `json:",inline"`

	// Standard object's metadata.
	metav1.ObjectMeta `json:"metadata,omitempty"`

	// Specification of the desired TrainingRuntime.
	Spec TrainingRuntimeSpec `json:"spec,omitempty"`
}

type TrainingRuntimeSpec struct {
	// Configuration for the model training with ML-specific parameters.
	MLPolicy *MLPolicy `json:"mlPolicy,omitempty"`

	// Configuration for the PodGroup to enable gang-scheduling via supported plugins.
	PodGroupPolicy *PodGroupPolicy `json:"podGroupPolicy,omitempty"`

	// JobSet template which will be used by TrainJob.
	Template JobSetTemplateSpec `json:"template"`
}

// JobSetTemplateSpec represents a template of the desired JobSet.
type JobSetTemplateSpec struct {
	// Metadata for custom JobSet's labels and annotations.
	// JobSet name and namespace is equal to the TrainJob's name and namespace.
	metav1.ObjectMeta `json:"metadata,omitempty"`

	// Specification of the desired JobSet which will be created from TrainJob.
	Spec jobsetv1alpha2.JobSetSpec `json:"spec,omitempty"`
}

type MLPolicy struct {
	// Number of training nodes.
	// Defaults to 1.
	NumNodes *int32 `json:"numNodes,omitempty"`

	// Configuration for the runtime-specific parameters, such as Torch or MPI.
	// Only one of its members may be specified.
	MLPolicySource `json:",inline"`
}

// MLPolicySource represents the runtime-specific configuration for various technologies.
// One of the following specs can be set.
type MLPolicySource struct {
	// Configuration for the PyTorch runtime.
	Torch *TorchMLPolicySource `json:"torch,omitempty"`

	// Configuration for the MPI Runtime.
	MPI *MPIMLPolicySource `json:"mpi,omitempty"`
}

The PodGroupPolicy API

The PodGroupPolicy is used to create the appropriate PodGroup for gang-scheduling. It can be used with Volcano or Coscheduling.

type PodGroupPolicy struct {
	// Configuration for gang-scheduling using various plugins.
	PodGroupPolicySource `json:",inline"`
}

// Only one of its members may be specified.
type PodGroupPolicySource struct {
	// Coscheduling plugin from the Kubernetes scheduler-plugins for gang-scheduling.
	Coscheduling *CoschedulingPodGroupPolicySource `json:"coscheduling,omitempty"`
}

// The number of min members in the PodGroupSpec is always equal to the number of nodes.
type CoschedulingPodGroupPolicySource struct {
	// Time threshold to schedule PodGroup for gang-scheduling.
	// If the scheduling timeout is equal to 0, the default value is used.
	// Defaults to 60 seconds.
	ScheduleTimeoutSeconds *int32 `json:"scheduleTimeoutSeconds,omitempty"`
}

The following example shows example of runtime with gang-scheduling using coscheduling plugin. Note: User should add the scheduler name into Pod's .spec.schedulerName if the default scheduler is not the same as PodGroup plugin.

apiVersion: trainer.kubeflow.org/v2alpha1
kind: ClusterTrainingRuntime
metadata:
  name: torch-distributed-multi-node
spec:
  mlPolicy:
    numNodes: 2
    torch: {}
  podGroupPolicy:
    coscheduling:
      scheduleTimeoutSeconds: 100
  template:
    spec:
      replicatedJobs:
        - name: node
          template:
            spec:
              template:
                metadata:
                  labels:
                    trainer.kubeflow.org/trainjob-ancestor-step: trainer
                spec:
                  schedulerName: coscheduling
                  containers:
                    - name: trainer
                      image: docker.io/kubeflow/pytorch-mnist
                      resources:
                        limits:
                          nvidia.com/gpu: 1
                      env:
                        - name: MASTER_ADDR
                          value: "pytorch-node-0-0.pytorch"
                        - name: MASTER_PORT
                          value: 29400
                      command:
                        - torchrun train.py

Training Operator will create the PodGroup using the following spec:

apiVersion: scheduling.x-k8s.io/v1alpha1
kind: PodGroup
metadata:
  name: torch-distributed-multi-node
spec:
  scheduleTimeoutSeconds: 100
  minMember: 5

The TrainJob will be started only when 5 GPUs are available in the cluster.

The TorchMLPolicySource API

The TorchMLPolicySource API represents the configuration for the PyTorch distributed training. This configuration allows platform engineers to explicitly configure torchrun setting.

The distributed parameters are taken from the PyTorch distributed launch run.

The --standalone parameter will be automatically set when numProcPerNode > 0 and numNodes=0.

For the Elastic Training we will always pass the following parameters:

--rdzv-backend=c10d

--rdzv-id will be set automatically.

--rdzv-endpoint will always point to the node-0 Pod.

Since the etcd and etcd-v2 are legacy rendezvous, we won't support them in TorchMLPolicySource. We can introduce them in the future if users will require them.

type TorchMLPolicySource struct {}

The MPIMLPolicySource API

The MPIMLPolicySource API represents the configuration for training using MPI orchestration. E.g. creation of host-files and SSH keys. Using MPI might be more efficient for training on HPC clusters or for some ML frameworks (e.g. MLX distributed with MPI).

We will fully migrate to the MPI Operator V2 functionality as part of this KEP. Check the proposal for the MPI V2 APIs.

type MPIMLPolicySource struct {
	// Number of processes per node.
	// This value is equal to the number of slots for each node in the hostfile.
	NumProcPerNode *int32 `json:"numProcPerNode,omitempty"`

	// Implementation name for the MPI to create the appropriate hostfile.
	// Defaults to OpenMPI.
	MPIImplementation *MPIImplementation `json:"mpiImplementation,omitempty"`

	// Directory where SSH keys are mounted.
	SSHAuthMountPath *string `json:"SSHAuthMountPath,omitempty"`

	// Whether to run training process on the launcher Job.
	// Defaults to false.
	RunLauncherAsNode *bool `json:"runLauncherAsNode,omitempty"`
}

type MPIImplementation string

const (
    MPIImplementationOpenMPI MPIImplementation = "OpenMPI"
    MPIImplementationIntel   MPIImplementation = "Intel"
    MPIImplementationMPICH   MPIImplementation = "MPICH"
)

Supported Runtimes by Community

Kubeflow community are planning to support the following runtimes.

PyTorch Distributed Runtime

Initially, we will maintain only multi-node multi-worker runtime and PyTorch Elastic.

apiVersion: trainer.kubeflow.org/v2alpha1
kind: ClusterTrainingRuntime
metadata:
  name: torch-distributed-multi-node
spec:
  mlPolicy:
    numNodes: 2
    torch: {}
  template:
    spec:
      replicatedJobs:
        - name: node
          template:
            spec:
              template:
                metadata:
                  labels:
                    trainer.kubeflow.org/trainjob-ancestor-step: trainer
                spec:
                  containers:
                    - name: trainer
                      image: docker.io/kubeflow/pytorch-mnist
                      env:
                        - name: MASTER_ADDR
                          value: "pytorch-node-0-0.pytorch"
                        - name: MASTER_PORT
                          value: 29400
                      command:
                        - torchrun train.py

Example of usage:

apiVersion: trainer.kubeflow.org/v2alpha1
kind: TrainJob
metadata:
  name: torch-test
  namespace: tenant-alpha
spec:
  runtimeRef:
    name: torch-distributed-multi-node
  trainer:
    resourcesPerNode:
      requests:
        nvidia.com/gpu: 1
    args:
      - num-epochs=5

Additional PyTorch Runtimes

The following runtimes can be maintained in the future.

Single worker training:

apiVersion: trainer.kubeflow.org/v2alpha1
kind: ClusterTrainingRuntime
metadata:
  name: torch-simple
spec:
  template:
    spec:
      replicatedJobs:
        - name: node
          template:
            spec:
              template:
                metadata:
                  labels:
                    trainer.kubeflow.org/trainjob-ancestor-step: trainer
                spec:
                  containers:
                    - name: trainer
                      image: docker.io/kubeflow/pytorch-mnist
                      command:
                        - torchrun train.py

Single node multi worker training:

apiVersion: trainer.kubeflow.org/v2alpha1
kind: ClusterTrainingRuntime
metadata:
  name: torch-distributed-single-worker
spec:
  mlPolicy: {}
  template:
    spec:
      replicatedJobs:
        - name: node
          template:
            spec:
              template:
                metadata:
                  labels:
                    trainer.kubeflow.org/trainjob-ancestor-step: trainer
                spec:
                  containers:
                    - name: trainer
                      image: docker.io/kubeflow/pytorch-mnist
                      env:
                        - name: MASTER_ADDR
                          value: "pytorch-node-0-0.pytorch"
                        - name: MASTER_PORT
                          value: 29400
                      command:
                        - torchrun train.py

LLM Fine-Tuning Runtimes

In the future, we can consider using the torchtune CLI for Fine-Tuning with PyTorch.

Llama 7b

The following runtime can be used for the Llama 7b model.

apiVersion: trainer.kubeflow.org/v2alpha1
kind: ClusterTrainingRuntime
metadata:
  name: torch-tune-llama-7b
spec:
  mlPolicy:
    numNodes: 1
  template:
    spec:
      replicatedJobs:
        - name: dataset-initializer
          template:
            spec:
              template:
                metadata:
                  labels:
                    trainer.kubeflow.org/trainjob-ancestor-step: dataset-initializer
                spec:
                  containers:
                    - name: dataset-initializer
                      image: docker.io/kubeflow/dataset-initializer
                      env:
                        - name: STORAGE_URI
                          value: hf://tatsu-lab/alpaca
                      volumeMounts:
                        - mountPath: /workspace/dataset
                          name: initializer
                  volumes:
                    - name: initializer
                      persistentVolumeClaim:
                        claimName: initializer
        - name: model-initializer
          template:
            spec:
              template:
                metadata:
                  labels:
                    trainer.kubeflow.org/trainjob-ancestor-step: model-initializer
                spec:
                  containers:
                    - name: model-initializer
                      image: docker.io/kubeflow/model-initializer
                      env:
                        - name: STORAGE_URI
                          value: hf://meta-llama/Llama-2-7b
                        - name: TRANSFORMER_TYPE
                          value: AutoModelForCausalLM
                      volumeMounts:
                        - mountPath: /workspace/model
                          name: initializer
                  volumes:
                    - name: initializer
                      persistentVolumeClaim:
                        claimName: initializer
        - name: node
          dependsOn:
            - name: dataset-initializer
              status: Complete
            - name: model-initializer
              status: Complete
          template:
            spec:
              template:
                metadata:
                  labels:
                    trainer.kubeflow.org/trainjob-ancestor-step: trainer
                spec:
                  containers:
                    - name: trainer
                      image: docker.io/kubeflow/llm-trainer
                      env:
                        - name: MASTER_ADDR
                          value: "pytorch-node-0-0.pytorch"
                        - name: MASTER_PORT
                          value: 29400
                        - name: TRANSFORMER_TYPE
                          value: AutoModelForCausalLM
                        - name: LORA_CONFIG
                          value: |
                            {"peft_type": "LORA", "r": 8, "lora_alpha": 16}
                      command:
                        - torchrun hf_llm_training.py
                      resources:
                        limits:
                          nvidia.com/gpu: 2
                      volumeMounts:
                        - mountPath: /workspace/model
                          name: initializer
                      volumeMounts:
                        - mountPath: /workspace/dataset
                          name: initializer
                  volumes:
                    - name: initializer
                      persistentVolumeClaim:
                        claimName: initializer

Gemma 7b

The following runtime can be used for Gemma fine-tuning.

apiVersion: trainer.kubeflow.org/v2alpha1
kind: ClusterTrainingRuntime
metadata:
  name: torch-tune-gemma-7b
spec:
  mlPolicy:
    numNodes: 1
  template:
    spec:
      replicatedJobs:
        - name: dataset-initializer
          template:
            spec:
              template:
                metadata:
                  labels:
                    trainer.kubeflow.org/trainjob-ancestor-step: dataset-initializer
                spec:
                  containers:
                    - name: dataset-initializer
                      image: docker.io/kubeflow/dataset-initializer
                      env:
                        - name: STORAGE_URI
                          value: hf://tatsu-lab/alpaca
                      volumeMounts:
                        - mountPath: /workspace/dataset
                          name: initializer
                  volumes:
                    - name: initializer
                      persistentVolumeClaim:
                        claimName: initializer
        - name: model-initializer
          template:
            spec:
              template:
                metadata:
                  labels:
                    trainer.kubeflow.org/trainjob-ancestor-step: model-initializer
                spec:
                  containers:
                    - name: model-initializer
                      image: docker.io/kubeflow/model-initializer
                      env:
                        - name: STORAGE_URI
                          value: hf://google/gemma-7b
                        - name: TRANSFORMER_TYPE
                          value: AutoModelForCausalLM
                      volumeMounts:
                        - mountPath: /workspace/model
                          name: initializer
                  volumes:
                    - name: initializer
                      persistentVolumeClaim:
                        claimName: initializer
        - name: node
          dependsOn:
            - name: dataset-initializer
              status: Complete
            - name: model-initializer
              status: Complete
          template:
            spec:
              template:
                spec:
                  containers:
                    - name: trainer
                      image: docker.io/kubeflow/llm-trainer
                      env:
                        - name: MASTER_ADDR
                          value: "pytorch-node-0-0.pytorch"
                        - name: MASTER_PORT
                          value: 29400
                        - name: TRANSFORMER_TYPE
                          value: AutoModelForCausalLM
                        - name: LORA_CONFIG
                          value: |
                            {"peft_type": "LORA", "r": 8, "lora_alpha": 16}
                      command:
                        - torchrun hf_llm_training.py
                      resources:
                        limits:
                          nvidia.com/gpu: 2
                      volumeMounts:
                        - mountPath: /workspace/dataset
                          name: initializer
                        - mountPath: /workspace/model
                          name: initializer
                  volumes:
                    - name: initializer
                      persistentVolumeClaim:
                        claimName: initializer

MPI Runtime

We will re-use the MPI Operator V2 functionality as part of this MPI Runtime. Which means we will use the SSH-based approach to initialize the MPI Job.

The MPI Plugin in Training Operator will be responsible to:

• Build the Secret with the SSH keys.
• Build the ConfigMap with the appropriate hostfile for OpenMPI, IntelMPI, or MPICH. We will support only OpenMPI in the first implementation.

The Secret and ConfigMap will be added to the corresponding JobSet.

The hostfile default location is /etc/mpi/hostfile. For example, for OpenMPI we configure this env variable:

OMPI_MCA_orte_default_hostfile=/etc/mpi/hostfile

The numProcPerNode is equal to the number of slots in the MPI hostfile.

Example of hostfile:

deepspeed-trainer-node-0-0.default.svc slots=5
deepspeed-trainer-node-0-1.default.svc slots=5

Initially, we will introduce support for distributed MLX and DeepSpeed using the MPI Runtime.

Example of simple OpenMPI runtime:

apiVersion: trainer.kubeflow.org/v2alpha1
kind: ClusterTrainingRuntime
metadata:
  name: deepspeed
  namespace: default
spec:
  mlPolicy:
    numNodes: 2
    mpi:
      mpiImplementation: OpenMPI
      numProcPerNode: 5
  template:
    replicatedJobs:
      - name: launcher
        template:
          spec:
            template:
              metadata:
                labels:
                  trainer.kubeflow.org/trainjob-ancestor-step: trainer
              spec:
                containers:
                  - name: mpi-launcher
                    image: docker.io/mpi-launch
                    command:
                      - mpirun launch-job
      - name: node
        template:
          spec:
            template:
              spec:
                containers:
                  - name: trainer
                    image: docker.io/deepspeed-trainer
                    resources:
                      limits:
                        nvidia.com/gpu: 5

TensorFlow Runtime

Will be added after initial implementation for PyTorch.

XGBoost Runtime

Will be added after initial implementation for PyTorch.

Paddle Runtime

Will be added after initial implementation for PyTorch.

Jax Runtime

Will be added after initial implementation for PyTorch.

Pipeline Framework

We introduce the framework as internal mechanism so that we can easily expand mechanism for combination of Runtimes and TrainJob.

The framework is called as Kubeflow Trainer Pipeline Framework, and it has 4 phases as you can see the following overview.

As described in the following, each phase is basically executed step by step although Startup Phase is executed only once during starting trainer-controller-manager:

• Startup Phase: Initialize internal components at once when the kubeflow-trainer-controller-manager starts.
• PreExecution Phase: This phase is executed as a part of admission validating webhooks triggered by TrainJob is created and updated.
• Build Phase: This phase is executed to build child Kubernetes resources and deploy those to the cluster.
• PostExecution Phase: This phase is executed after the Build Phase.

As you can see in the diagram, each phase has 2 types of APIs, Internal API and Extension Point. The Extension Point is exposed and could be added operations within the scope of the Pipeline Framework Plugins Interfaces. These plugins are executed in any order.

On the other hand, the Internal APIs are not exposed and could not add any operations as opposed to the Extension Point.

• Startup Phase:
  • Internal API:
    • Initialize Kubeflow TrainerFrameworkPipeline: Initialize entire Kubeflow TrainerPipelineFramework.
    • TrainJobController: Set up TrainJob controller and register it to Manager.
    • Built-in Webhook Servers: Set up Built-in Admission Webhook Servers and register those to Manager.
    • Start Manager: Start Manager.
  • Extension Point
    • WatchExtension: This registers arbitrary reconciler builders for watching any kind of resources and triggering TrainJob reconciliations.
• PreExecution Phase:
  • Extension Point:
    • CustomValidation: This registers validators for validating any kind of resources to Admission Validating Webhook Servers when TrainJob is created and updated.
• Build Phase:
  • Internal API:
    • ComponentDeployer: This deploys built components (resources) to the cluster which is performed as a part of reconciler.
  • Extension Point:
    • EnforcePodGroupPolicy: This configures PodGroup specific parameters (e.x, specified in TrainingRuntime .spec.podGroupPolicy) to any kind of resources like PodSpec.
    • EnforceMLPolicy: This configure MachineLearning framework specific parameters (e.x, specified in TrainingRuntime .spec.mlPolicy) to any kind of resources like PodSpec.
    • PodNetwork: This identifies Pod-to-Pod communication network endpoints for each Pod and stores them to RuntimeInfo.
    • ComponentBuilder: This builds Kubernetes resources leveraging RuntimeInfo and TrainJob. RuntimeInfo is abstracted objects extracted from runtimes like TrainingRuntime and ClusterTrainingRuntime.
• PostExecution Phase:
  • Internal API:
    • SupendedCondition: Check if TrainJob is suspended state, and then add Suspended condition to TrainJob.
    • CreatedConditon: Check if TrainJob is created state, and then add Created condition to TrainJob.
  • Extension Point:
    • TerminalCondition: Check if TrainJob is terminated state, and then add Complete condition with a propagated terminal reason and message from child Jobs to TrainJob.

Migration from Kubeflow Training V1

These API changes will not be compatible with Training Operator V1 APIs. Thus, existing users have to migrate to the newer APIs. The Kubeflow community will provide instructions on how to migrate existing training jobs to the new APIs.

PyTorchJob Migration

The following example shows how to migrate from PyTorchJob to TrainingRuntime:

apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  name: pytorch-simple
  namespace: kubeflow
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: docker.io/kubeflowkatib/pytorch-mnist:v1beta1-45c5727
              imagePullPolicy: Always
              command:
                - "python3"
                - "/opt/pytorch-mnist/mnist.py"
                - "--epochs=1"
    Worker:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: docker.io/kubeflowkatib/pytorch-mnist:v1beta1-45c5727
              imagePullPolicy: Always
              command:
                - "python3"
                - "/opt/pytorch-mnist/mnist.py"
                - "--epochs=1"

apiVersion: trainer.kubeflow.org/v2alpha1
kind: TrainingRuntime
metadata:
  name: torch-distributed-multi-node
spec:
  mlPolicy:
    numNodes: 2
  template:
    spec:
      replicatedJobs:
        - name: node
          template:
            spec:
              template:
                metadata:
                  labels:
                    trainer.kubeflow.org/trainjob-ancestor-step: trainer
                spec:
                  containers:
                    - name: trainer
                      image: docker.io/kubeflowkatib/pytorch-mnist:v1beta1-45c5727
                      env:
                        - name: MASTER_ADDR
                          value: "pytorch-node-0-0.pytorch"
                        - name: MASTER_PORT
                          value: 29400
                      command:
                        - torchrun train.py

Implementation History

• 2024-07-16 Creation date
• 2025-03-15 Updated the initializer APIs
• 2025-10-09 Added PodTemplateOverrides to TrainJob V2 API
• 2026-02-23 Remove ElasticPolicy from the Torch API – will be implemented in the future KEPs
• 2026-02-23 Removed numProcPerNode from TorchMLPolicySource; TrainJob.Trainer.numProcPerNode now accepts only int values (torch plugin defaults to auto)
• 2025-10-09 Replace PodTemplateOverrides with RuntimePatches API

Alternatives

Alternatives details can be found in this Google document.

Inline JobSet APIs into TrainJob

type TrainJobSpec struct {
    ...

    JobSetSpec *jobsetv1.JobSetSpec `json:",inline"`
}

In that case, TrainJob API will be very complex and users still have to specify every Kubernetes API parameter on job submission.

Use JobSetTemplate as a Training Runtime

Instead of creating the custom CRD for TrainingRuntime, use the JobSetTemplate API to create blueprints for training runtimes.

Platform engineers need to understand all aspect on how to configure parameters for various frameworks (e.g. PyTorch or HuggingFace). Also, it will be hard to implement custom orchestration when it is requires (e.g. MPI or Slurm use-case).

Using CRD for Every Framework (e.g. PyTorchJob)

Instead of TrainJob maintain different CRDs for each framework: PyTorchJob, JaxJob, MPIJob.

Given that ML framework space is growing very fast, it will be very hard to maintain CRD for every framework that users want to run on Kubernetes.

Since frameworks share common functionality for distributed training (data parallelizm or model parallelizm). For some specific use-cases like MPI or Elastic PyTorch, we will leverage MLPolicy parameter.

Allow users to specify arbitrary value in the managedBy field

We can allow users to specify the arbitrary values instead of restricting the .spec.managedBy field in the TrainJob with an empty, 'kubeflow.org/trainjob-controller' or 'kusus.x-k8s.io/multikueue'.

But, the arbitrary values allow users to specify external or in-house customized training-operator, which means that the TrainJobs are reconciled by the controllers without any specification compliance.

Specifically, the arbitrary training-operator could bring bugs for the status transitions. So, we do not support the arbitrary values until we find reasonable use cases that the external controllers need to reconcile the TrainJob.

Note that we should implement the status transitions validations to once we support the arbitrary values in the manageBy field.

Support Multiple API Versions of TrainingRuntime

We can consider to introduce the version field for runtime API version to the .spec.runtimeRef so that we can support multiple API versions of TrainingRuntime.

It could mitigate the pain points when users upgrade the older API Version to newer API Version like alpha to beta. But, we do not aim to support both Alpha and Beta versions or both first Alpha and second Alpha versions in the specific training-operator release. Hence, the version field was not introduced.

type RuntimeRef struct {
	[...]

	// APIVersion is the apiVersion for the runtime.
	// Defaults to the v2alpha1.
	Version *string `json:version,omitempty`

	[...]
}

However, we may want to revisit this discussion when we graduate the API version from Beta to GA because in general, it would be better to support both Beta and GA versions for a while so that we can align with the Kubernetes deprecation policy for mitigating migration obstacles.