Repository: kubeflow/trainer
Pull request #2548: Construct Trainer based on trainer.kubeflow.org/trainjob-ancestor-step label
Author: login: tenzen-y
Labels: name: size/XL
name: lgtm
name: approved
Created: 2025-03-19T11:54:15Z | Merged: 2025-03-20T14:01:58Z

<!--  Thanks for sending a pull request! Here are some tips for you:
1. If this is your first time, check our contributor guidelines: https://www.kubeflow.org/docs/about/contributing
2. To know more about Kubeflow Trainer, check the developer guide:
    https://github.com/kubeflow/trainer/blob/master/CONTRIBUTING.md
3. If you want *faster* PR reviews, check how: https://git.k8s.io/community/contributors/guide/pull-requests.md#best-practices-for-faster-reviews
-->

**What this PR does / why we need it**:

As @andreyvelich described in https://github.com/kubeflow/trainer/pull/2523, we construct actual Trainers based on `trainer.kubeflow.org/trainjob-ancestor-step` label. Additionally, I enabled runLauncherAsNode by default for MPI runtime.
By this change, MLEs can override commands, args, images.... via TrainJob.

```yaml
apiVersion: trainer.kubeflow.org/v1alpha1
kind: TrainJob
metadata:
  name: mpirun
  namespace: default
spec:
  runtimeRef:
    name: mpi-distributed
  trainer:
    numNodes: 10
    command:
    - "mpirun"
    args:
    - "--tag-output"
    - "/home/mpiuser/pi"
```

```shell
$ kubectl get pod
NAME                        READY   STATUS      RESTARTS   AGE
mpirun-launcher-0-0-t5k6p   0/2     Completed   0          2s
mpirun-node-0-0-5nzbd       0/1     Running     0          2s
mpirun-node-0-1-gwxnq       0/1     Running     0          2s
mpirun-node-0-2-4rnbv       0/1     Running     0          2s
mpirun-node-0-3-2xvj2       0/1     Running     0          2s
mpirun-node-0-4-5k7zb       0/1     Running     0          2s
mpirun-node-0-5-llrz9       0/1     Running     0          2s
mpirun-node-0-6-hkndl       0/1     Running     0          2s
mpirun-node-0-7-dht76       0/1     Running     0          2s
mpirun-node-0-8-j88q5       0/1     Running     0          2s. # There are only 9 nodes since we use RunLauncherAsNode
$
$ k get pod mpirun-launcher-0-0-t5k6p -ojsonpath='{.spec.containers[0].args}'
# The `--tag-output` was propagated from TrainJob.
["--tag-output","/home/mpiuser/pi"]
$ 
$ kubectl logs mpirun-launcher-0-0-t5k6p
[1,0]<stdout>:Workers: 10.                                  # World size is 10 (launcher + 9 nodes)
[1,0]<stdout>:Rank 0 on host mpirun-launcher-0-0
[1,5]<stdout>:Rank 5 on host mpirun-node-0-4
[1,8]<stdout>:Rank 8 on host mpirun-node-0-7
[1,3]<stdout>:Rank 3 on host mpirun-node-0-2
[1,9]<stdout>:Rank 9 on host mpirun-node-0-8
[1,7]<stdout>:Rank 7 on host mpirun-node-0-6
[1,1]<stdout>:Rank 1 on host mpirun-node-0-0
[1,2]<stdout>:Rank 2 on host mpirun-node-0-1
[1,4]<stdout>:Rank 4 on host mpirun-node-0-3
[1,6]<stdout>:Rank 6 on host mpirun-node-0-5
[1,0]<stdout>:pi is approximately 3.1414545600000001
```

**Which issue(s) this PR fixes** _(optional, in `Fixes #<issue number>, #<issue number>, ...` format, will close the issue(s) when PR gets merged)_:
Fixes #

**Checklist:**

- [ ] [Docs](https://www.kubeflow.org/docs/components/trainer/) included if any changes are user facing



Diff:
diff --git a/hack/e2e-setup-cluster.sh b/hack/e2e-setup-cluster.sh
index 2ec2cadd13..61bde9f3b5 100755
--- a/hack/e2e-setup-cluster.sh
+++ b/hack/e2e-setup-cluster.sh
@@ -30,7 +30,7 @@ TIMEOUT="5m"
 
 # Kubeflow Trainer images.
 # TODO (andreyvelich): Support initializers images.
-CONTROLLER_MANAGER_CI_IMAGE_NAME="kubeflow/trainer-controller-manager"
+CONTROLLER_MANAGER_CI_IMAGE_NAME="ghcr.io/kubeflow/trainer/trainer-controller-manager"
 CONTROLLER_MANAGER_CI_IMAGE_TAG="test"
 CONTROLLER_MANAGER_CI_IMAGE="${CONTROLLER_MANAGER_CI_IMAGE_NAME}:${CONTROLLER_MANAGER_CI_IMAGE_TAG}"
 echo "Build Kubeflow Trainer images"
diff --git a/manifests/base/runtimes/pretraining/mpi_distributed.yaml b/manifests/base/runtimes/pretraining/mpi_distributed.yaml
index de21fe6f29..a77a0f3c0c 100644
--- a/manifests/base/runtimes/pretraining/mpi_distributed.yaml
+++ b/manifests/base/runtimes/pretraining/mpi_distributed.yaml
@@ -12,10 +12,11 @@ spec:
       numProcPerNode: 1
       mpiImplementation: OpenMPI
       sshAuthMountPath: /home/mpiuser/.ssh
+      runLauncherAsNode: true
   template:
     spec:
       network:
-        publishNotReadyAddresses: false
+        publishNotReadyAddresses: true
       successPolicy:
         operator: All
         targetReplicatedJobs:
@@ -23,25 +24,28 @@ spec:
       replicatedJobs:
         - name: launcher
           template:
+            metadata:
+              labels:
+                trainer.kubeflow.org/trainjob-ancestor-step: trainer
             spec:
               template:
                 spec:
                   containers:
-                    - name: launcher
-                      image: mpioperator/mpi-pi:openmpi
-                      securityContext:
-                        runAsUser: 1000
-                      command:
-                        - mpirun
-                      args:
-                        - /home/mpiuser/pi
-        - name: trainer-node
+                  - name: node
+                    image: mpioperator/mpi-pi:openmpi
+                    securityContext:
+                      runAsUser: 1000
+                    command:
+                      - mpirun
+                    args:
+                      - /home/mpiuser/pi
+        - name: node
           template:
             spec:
               template:
                 spec:
                   containers:
-                    - name: trainer-node
+                    - name: node
                       image: mpioperator/mpi-pi:openmpi
                       securityContext:
                         runAsUser: 1000
diff --git a/manifests/base/runtimes/pretraining/torch_distributed.yaml b/manifests/base/runtimes/pretraining/torch_distributed.yaml
index 9b678ea33b..8a81d8d9bc 100644
--- a/manifests/base/runtimes/pretraining/torch_distributed.yaml
+++ b/manifests/base/runtimes/pretraining/torch_distributed.yaml
@@ -12,13 +12,16 @@ spec:
   template:
     spec:
       replicatedJobs:
-        - name: trainer-node
+        - name: node
           template:
+            metadata:
+              labels:
+                trainer.kubeflow.org/trainjob-ancestor-step: trainer
             spec:
               template:
                 spec:
                   containers:
-                    - name: trainer
+                    - name: node
                       image: pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime
                       command:
                         - /bin/bash
diff --git a/pkg/constants/constants.go b/pkg/constants/constants.go
index dc23079364..304ac4434f 100644
--- a/pkg/constants/constants.go
+++ b/pkg/constants/constants.go
@@ -34,11 +34,12 @@ const (
 	// ModelMountPath is the volumeMount path for model.
 	ModelMountPath string = "/workspace/model"
 
-	// JobTrainerNode is the Job name for the trainer node.
-	JobTrainerNode string = "trainer-node"
+	// AncestorTrainer is the ancestor name for Trainer, which is mostly used for the value of
+	// 'trainer.kubeflow.org/trainjob-ancestor-step'.
+	AncestorTrainer string = "trainer"
 
-	// ContainerTrainer is the container name for the trainer.
-	ContainerTrainer string = "trainer"
+	// Node is the name of the Job and container for the trainer node
+	Node string = "node"
 
 	// ContainerTrainerPort is the default port for the trainer nodes communication.
 	ContainerTrainerPort int32 = 29500
diff --git a/pkg/runtime/core/clustertrainingruntime_test.go b/pkg/runtime/core/clustertrainingruntime_test.go
index 943259d435..06ae1b52c4 100644
--- a/pkg/runtime/core/clustertrainingruntime_test.go
+++ b/pkg/runtime/core/clustertrainingruntime_test.go
@@ -51,7 +51,7 @@ func TestClusterTrainingRuntimeNewObjects(t *testing.T) {
 				testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeClusterTrainingRuntimeWrapper("test-runtime").Spec).
 					JobSetSpec(
 						testingutil.MakeJobSetWrapper("", "").
-							DependsOn(constants.JobTrainerNode,
+							DependsOn(constants.Node,
 								[]jobsetv1alpha2.DependsOn{
 									{
 										Name:   constants.DatasetInitializer,
@@ -73,7 +73,7 @@ func TestClusterTrainingRuntimeNewObjects(t *testing.T) {
 							WithNumNodes(100).
 							Obj(),
 					).
-					Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 					PodGroupPolicyCoschedulingSchedulingTimeout(120).
 					Obj(),
 			).Obj(),
@@ -92,14 +92,14 @@ func TestClusterTrainingRuntimeNewObjects(t *testing.T) {
 					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
 					Suspend(true).
 					PodLabel(schedulerpluginsv1alpha1.PodGroupLabel, "test-job").
-					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.JobTrainerNode, constants.JobLauncher).
+					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node, constants.JobLauncher).
 					Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
 					Completions(1, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
 					NumNodes(100).
 					Container(constants.DatasetInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 					Container(constants.ModelInitializer, constants.ModelInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
-					Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
-					DependsOn(constants.JobTrainerNode,
+					Container(constants.Node, constants.Node, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
+					DependsOn(constants.Node,
 						[]jobsetv1alpha2.DependsOn{
 							{
 								Name:   constants.DatasetInitializer,
diff --git a/pkg/runtime/core/trainingruntime.go b/pkg/runtime/core/trainingruntime.go
index fe83cd8c44..297f48642e 100644
--- a/pkg/runtime/core/trainingruntime.go
+++ b/pkg/runtime/core/trainingruntime.go
@@ -150,11 +150,18 @@ func (r *TrainingRuntime) newRuntimeInfo(
 		// TODO: Support multiple replicas ('.template.spec.replicatedJobs[*].replicas') for replicated Jobs.
 		// REF: https://github.com/kubeflow/trainer/issues/2318
 		count := ptr.Deref(rJob.Template.Spec.Parallelism, 1)
-		if *rJob.Name == constants.JobTrainerNode && mlPolicy != nil {
-			count = ptr.Deref(mlPolicy.NumNodes, 1)
+		var ancestor *string
+		if metadata := rJob.Template.ObjectMetaApplyConfiguration; metadata != nil && metadata.Labels != nil {
+			if labelAncestor, ok := metadata.Labels[constants.LabelTrainJobAncestor]; ok {
+				if labelAncestor == constants.AncestorTrainer && mlPolicy != nil {
+					count = ptr.Deref(mlPolicy.NumNodes, 1)
+				}
+				ancestor = &labelAncestor
+			}
 		}
 		opts = append(opts, runtime.WithPodSet(
 			*rJob.Name,
+			ancestor,
 			count,
 			*jobSetTemplateSpec.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.DeepCopy(),
 			rJob.Template.Spec.Template.Spec),
diff --git a/pkg/runtime/core/trainingruntime_test.go b/pkg/runtime/core/trainingruntime_test.go
index 9d6e4c70bf..44f6d8df02 100644
--- a/pkg/runtime/core/trainingruntime_test.go
+++ b/pkg/runtime/core/trainingruntime_test.go
@@ -67,7 +67,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 						PodGroupPolicyCoschedulingSchedulingTimeout(120).
 						Container(constants.DatasetInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 						Container(constants.ModelInitializer, constants.ModelInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
-						Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+						Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 						Obj(),
 				).Obj(),
 			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
@@ -89,13 +89,13 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 					Label("conflictLabel", "override").
 					Annotation("conflictAnnotation", "override").
 					PodLabel(schedulerpluginsv1alpha1.PodGroupLabel, "test-job").
-					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.JobTrainerNode, constants.JobLauncher).
+					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node, constants.JobLauncher).
 					Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
 					Completions(1, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
 					NumNodes(30).
 					Container(constants.DatasetInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 					Container(constants.ModelInitializer, constants.ModelInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
-					Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 					Obj(),
 				testingutil.MakeSchedulerPluginsPodGroup(metav1.NamespaceDefault, "test-job").
 					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
@@ -116,8 +116,8 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 							WithNumNodes(100).
 							Obj(),
 					).
-					Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
-					Env(constants.JobTrainerNode, constants.ContainerTrainer,
+					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+					Env(constants.Node, constants.Node,
 						[]corev1.EnvVar{
 							{
 								Name:  "TRAIN_JOB",
@@ -155,12 +155,12 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 			wantObjs: []runtime.Object{
 				testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
 					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
-					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.JobTrainerNode).
+					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node).
 					Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
 					Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
 					NumNodes(100).
-					Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
-					Env(constants.JobTrainerNode, constants.ContainerTrainer,
+					Container(constants.Node, constants.Node, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
+					Env(constants.Node, constants.Node,
 						[]corev1.EnvVar{
 							{
 								Name:  "TRAIN_JOB",
@@ -189,7 +189,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 					).
 					Container(constants.DatasetInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 					Container(constants.ModelInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
-					Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 					Obj(),
 			).Obj(),
 			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
@@ -235,13 +235,13 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 			wantObjs: []runtime.Object{
 				testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
 					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
-					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.JobTrainerNode).
+					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node).
 					Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
 					Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
 					NumNodes(100).
 					Container(constants.DatasetInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 					Container(constants.ModelInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
-					Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 					Env(constants.DatasetInitializer, constants.DatasetInitializer,
 						[]corev1.EnvVar{
 							{
@@ -306,7 +306,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 					).
 					JobSetSpec(
 						testingutil.MakeJobSetWrapper("", "").
-							DependsOn(constants.JobTrainerNode,
+							DependsOn(constants.Node,
 								[]jobsetv1alpha2.DependsOn{
 									{
 										Name:   constants.DatasetInitializer,
@@ -321,7 +321,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 							Obj().
 							Spec,
 					).
-					Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 					Obj(),
 			).Obj(),
 			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
@@ -337,13 +337,13 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 			wantObjs: []runtime.Object{
 				testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
 					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
-					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.JobTrainerNode, constants.JobLauncher).
+					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node, constants.JobLauncher).
 					Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
 					Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
 					NumNodes(30).
-					Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 					ContainerTrainerPorts([]corev1.ContainerPort{{ContainerPort: constants.ContainerTrainerPort}}).
-					Env(constants.JobTrainerNode, constants.ContainerTrainer,
+					Env(constants.Node, constants.Node,
 						[]corev1.EnvVar{
 							{
 								Name:  constants.TorchEnvNumNodes,
@@ -363,7 +363,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 							},
 							{
 								Name:  constants.TorchEnvMasterAddr,
-								Value: fmt.Sprintf("test-job-%s-0-0.test-job", constants.JobTrainerNode),
+								Value: fmt.Sprintf("test-job-%s-0-0.test-job", constants.Node),
 							},
 							{
 								Name:  constants.TorchEnvMasterPort,
@@ -371,7 +371,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 							},
 						}...,
 					).
-					DependsOn(constants.JobTrainerNode,
+					DependsOn(constants.Node,
 						[]jobsetv1alpha2.DependsOn{
 							{
 								Name:   constants.DatasetInitializer,
@@ -398,8 +398,8 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 							).
 							Obj(),
 					).
-					Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
-					Env(constants.JobTrainerNode, constants.ContainerTrainer,
+					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+					Env(constants.Node, constants.Node,
 						[]corev1.EnvVar{
 							{
 								Name:  "TRAIN_JOB",
@@ -437,13 +437,13 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 			wantObjs: []runtime.Object{
 				testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
 					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
-					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.JobTrainerNode, constants.JobLauncher).
+					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node, constants.JobLauncher).
 					Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
 					Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
 					NumNodes(100).
-					Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
+					Container(constants.Node, constants.Node, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
 					ContainerTrainerPorts([]corev1.ContainerPort{{ContainerPort: constants.ContainerTrainerPort}}).
-					Env(constants.JobTrainerNode, constants.ContainerTrainer,
+					Env(constants.Node, constants.Node,
 						[]corev1.EnvVar{
 							{
 								Name:  "TRAIN_JOB",
@@ -471,7 +471,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 							},
 							{
 								Name:  constants.TorchEnvMasterAddr,
-								Value: fmt.Sprintf("test-job-%s-0-0.test-job", constants.JobTrainerNode),
+								Value: fmt.Sprintf("test-job-%s-0-0.test-job", constants.Node),
 							},
 							{
 								Name:  constants.TorchEnvMasterPort,
@@ -490,21 +490,22 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 			ObjCmpOpts: cmp.Options{
 				cmp.Comparer(testingutil.MPISecretDataComparer),
 			},
-			trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").RuntimeSpec(
-				testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Spec).
-					WithMLPolicy(
-						testingutil.MakeMLPolicyWrapper().
-							WithNumNodes(1).
-							WithMLPolicySource(*testingutil.MakeMLPolicySourceWrapper().
-								MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), ptr.To(false)).
+			trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").
+				RuntimeSpec(
+					testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Spec).
+						WithMLPolicy(
+							testingutil.MakeMLPolicyWrapper().
+								WithNumNodes(1).
+								WithMLPolicySource(*testingutil.MakeMLPolicySourceWrapper().
+									MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), ptr.To(false)).
+									Obj(),
+								).
 								Obj(),
-							).
-							Obj(),
-					).
-					LauncherReplica().
-					Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
-					Obj(),
-			).
+						).
+						LauncherReplica().
+						Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+						Obj(),
+				).
 				Obj(),
 			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
 				UID("uid").
@@ -520,8 +521,8 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 				testingutil.MakeConfigMapWrapper(fmt.Sprintf("test-job%s", constants.MPIHostfileConfigMapSuffix), metav1.NamespaceDefault).
 					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
 					WithData(map[string]string{
-						constants.MPIHostfileName: `test-job-trainer-node-0-0.test-job slots=8
-test-job-trainer-node-0-1.test-job slots=8
+						constants.MPIHostfileName: `test-job-node-0-0.test-job slots=8
+test-job-node-0-1.test-job slots=8
 `,
 					}).
 					Obj(),
@@ -537,7 +538,7 @@ test-job-trainer-node-0-1.test-job slots=8
 				testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
 					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
 					LauncherReplica().
-					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.JobTrainerNode, constants.JobLauncher).
+					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node, constants.JobLauncher).
 					Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
 					Completions(1, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
 					NumNodes(2).
@@ -580,7 +581,7 @@ test-job-trainer-node-0-1.test-job slots=8
 							},
 						},
 					).
-					Volumes(constants.JobTrainerNode,
+					Volumes(constants.Node,
 						corev1.Volume{
 							Name: constants.MPISSHAuthVolumeName,
 							VolumeSource: corev1.VolumeSource{
@@ -608,6 +609,9 @@ test-job-trainer-node-0-1.test-job slots=8
 						corev1.VolumeMount{Name: constants.MPISSHAuthVolumeName, MountPath: "/root/.ssh"},
 						corev1.VolumeMount{Name: constants.MPIHostfileVolumeName, MountPath: constants.MPIHostfileDir},
 					).
+					VolumeMounts(constants.Node, constants.Node,
+						corev1.VolumeMount{Name: constants.MPISSHAuthVolumeName, MountPath: "/root/.ssh"},
+					).
 					Env(constants.JobLauncher, constants.ContainerLauncher,
 						corev1.EnvVar{
 							Name:  constants.OpenMPIEnvHostFileLocation,
@@ -626,7 +630,7 @@ test-job-trainer-node-0-1.test-job slots=8
 							Value: constants.OpenMPIEnvDefaultValueRSHArgs,
 						},
 					).
-					Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 					Obj(),
 			},
 		},
diff --git a/pkg/runtime/framework/core/framework_test.go b/pkg/runtime/framework/core/framework_test.go
index d7dcc0e599..6b047e4b6e 100644
--- a/pkg/runtime/framework/core/framework_test.go
+++ b/pkg/runtime/framework/core/framework_test.go
@@ -174,18 +174,18 @@ func TestRunEnforceMLPolicyPlugins(t *testing.T) {
 			registry: fwkplugins.NewRegistry(),
 			runtimeInfo: runtime.NewInfo(
 				runtime.WithMLPolicySource(testingutil.MakeMLPolicyWrapper().Obj()),
-				runtime.WithPodSet(constants.DatasetInitializer, 1, corev1.PodSpec{}, corev1ac.PodSpec().
+				runtime.WithPodSet(constants.DatasetInitializer, ptr.To(constants.DatasetInitializer), 1, corev1.PodSpec{}, corev1ac.PodSpec().
 					WithContainers(
 						corev1ac.Container().WithName(constants.DatasetInitializer),
 					),
 				),
-				runtime.WithPodSet(constants.ModelInitializer, 1, corev1.PodSpec{}, corev1ac.PodSpec().
+				runtime.WithPodSet(constants.ModelInitializer, ptr.To(constants.ModelInitializer), 1, corev1.PodSpec{}, corev1ac.PodSpec().
 					WithContainers(
 						corev1ac.Container().WithName(constants.ModelInitializer),
 					),
 				),
-				runtime.WithPodSet(constants.JobTrainerNode, 10, corev1.PodSpec{}, corev1ac.PodSpec().
-					WithContainers(corev1ac.Container().WithName(constants.ContainerTrainer)),
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 10, corev1.PodSpec{}, corev1ac.PodSpec().
+					WithContainers(corev1ac.Container().WithName(constants.Node)),
 				),
 			),
 			trainJob: &trainer.TrainJob{
@@ -198,8 +198,9 @@ func TestRunEnforceMLPolicyPlugins(t *testing.T) {
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{
 						{
-							Name:  constants.DatasetInitializer,
-							Count: ptr.To[int32](1),
+							Name:     constants.DatasetInitializer,
+							Ancestor: ptr.To(constants.DatasetInitializer),
+							Count:    ptr.To[int32](1),
 							Containers: []runtime.Container{
 								{
 									Name: constants.DatasetInitializer,
@@ -207,8 +208,9 @@ func TestRunEnforceMLPolicyPlugins(t *testing.T) {
 							},
 						},
 						{
-							Name:  constants.ModelInitializer,
-							Count: ptr.To[int32](1),
+							Name:     constants.ModelInitializer,
+							Ancestor: ptr.To(constants.ModelInitializer),
+							Count:    ptr.To[int32](1),
 							Containers: []runtime.Container{
 								{
 									Name: constants.ModelInitializer,
@@ -216,10 +218,11 @@ func TestRunEnforceMLPolicyPlugins(t *testing.T) {
 							},
 						},
 						{
-							Name:  constants.JobTrainerNode,
-							Count: ptr.To[int32](10),
+							Name:     constants.Node,
+							Ancestor: ptr.To(constants.AncestorTrainer),
+							Count:    ptr.To[int32](10),
 							Containers: []runtime.Container{{
-								Name: constants.ContainerTrainer,
+								Name: constants.Node,
 							}},
 						},
 					},
@@ -233,18 +236,18 @@ func TestRunEnforceMLPolicyPlugins(t *testing.T) {
 				runtime.WithMLPolicySource(
 					testingutil.MakeMLPolicyWrapper().Obj(),
 				),
-				runtime.WithPodSet(constants.DatasetInitializer, 1, corev1.PodSpec{}, corev1ac.PodSpec().
+				runtime.WithPodSet(constants.DatasetInitializer, ptr.To(constants.DatasetInitializer), 1, corev1.PodSpec{}, corev1ac.PodSpec().
 					WithContainers(
 						corev1ac.Container().WithName(constants.DatasetInitializer),
 					),
 				),
-				runtime.WithPodSet(constants.ModelInitializer, 1, corev1.PodSpec{}, corev1ac.PodSpec().
+				runtime.WithPodSet(constants.ModelInitializer, ptr.To(constants.ModelInitializer), 1, corev1.PodSpec{}, corev1ac.PodSpec().
 					WithContainers(
 						corev1ac.Container().WithName(constants.ModelInitializer),
 					),
 				),
-				runtime.WithPodSet(constants.JobTrainerNode, 10, corev1.PodSpec{}, corev1ac.PodSpec().
-					WithContainers(corev1ac.Container().WithName(constants.ContainerTrainer)),
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 10, corev1.PodSpec{}, corev1ac.PodSpec().
+					WithContainers(corev1ac.Container().WithName(constants.Node)),
 				),
 			),
 			trainJob: &trainer.TrainJob{
@@ -261,8 +264,9 @@ func TestRunEnforceMLPolicyPlugins(t *testing.T) {
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{
 						{
-							Name:  constants.DatasetInitializer,
-							Count: ptr.To[int32](1),
+							Name:     constants.DatasetInitializer,
+							Ancestor: ptr.To(constants.DatasetInitializer),
+							Count:    ptr.To[int32](1),
 							Containers: []runtime.Container{
 								{
 									Name: constants.DatasetInitializer,
@@ -270,8 +274,9 @@ func TestRunEnforceMLPolicyPlugins(t *testing.T) {
 							},
 						},
 						{
-							Name:  constants.ModelInitializer,
-							Count: ptr.To[int32](1),
+							Name:     constants.ModelInitializer,
+							Ancestor: ptr.To(constants.ModelInitializer),
+							Count:    ptr.To[int32](1),
 							Containers: []runtime.Container{
 								{
 									Name: constants.ModelInitializer,
@@ -279,10 +284,11 @@ func TestRunEnforceMLPolicyPlugins(t *testing.T) {
 							},
 						},
 						{
-							Name:  constants.JobTrainerNode,
-							Count: ptr.To[int32](30),
+							Name:     constants.Node,
+							Ancestor: ptr.To(constants.AncestorTrainer),
+							Count:    ptr.To[int32](30),
 							Containers: []runtime.Container{{
-								Name: constants.ContainerTrainer,
+								Name: constants.Node,
 							}},
 						},
 					},
@@ -440,8 +446,9 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
 		wantError       error
 	}{
 		"succeeded to build PodGroup and JobSet with NumNodes from TrainJob": {
-			registry:        fwkplugins.NewRegistry(),
-			trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").DeepCopy(),
+			registry: fwkplugins.NewRegistry(),
+			trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").
+				Obj(),
 			runtimeInfo: &runtime.Info{
 				RuntimePolicy: runtime.RuntimePolicy{
 					PodGroupPolicy: &trainer.PodGroupPolicy{
@@ -455,8 +462,9 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{
 						{
-							Name:  constants.DatasetInitializer,
-							Count: ptr.To[int32](1),
+							Name:     constants.DatasetInitializer,
+							Ancestor: ptr.To(constants.DatasetInitializer),
+							Count:    ptr.To[int32](1),
 							SinglePodRequests: corev1.ResourceList{
 								corev1.ResourceCPU:    resource.MustParse("1"),
 								corev1.ResourceMemory: resource.MustParse("4Gi"),
@@ -479,8 +487,9 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
 							},
 						},
 						{
-							Name:  constants.ModelInitializer,
-							Count: ptr.To[int32](1),
+							Name:     constants.ModelInitializer,
+							Ancestor: ptr.To(constants.ModelInitializer),
+							Count:    ptr.To[int32](1),
 							SinglePodRequests: corev1.ResourceList{
 								corev1.ResourceCPU:    resource.MustParse("1"),
 								corev1.ResourceMemory: resource.MustParse("4Gi"),
@@ -503,7 +512,9 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
 							},
 						},
 						{
-							Name: constants.JobTrainerNode,
+							Name:     constants.Node,
+							Ancestor: ptr.To(constants.AncestorTrainer),
+							Count:    ptr.To[int32](1),
 							SinglePodRequests: corev1.ResourceList{
 								corev1.ResourceCPU:    resource.MustParse("1"),
 								corev1.ResourceMemory: resource.MustParse("4Gi"),
@@ -525,9 +536,11 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
 							jobsetv1alpha2ac.ReplicatedJob().
 								WithName(constants.DatasetInitializer).
 								WithTemplate(batchv1ac.JobTemplateSpec().
+									WithLabels(map[string]string{
+										constants.LabelTrainJobAncestor: constants.DatasetInitializer,
+									}).
 									WithSpec(batchv1ac.JobSpec().
 										WithTemplate(corev1ac.PodTemplateSpec().
-											WithLabels(map[string]string{constants.LabelTrainJobAncestor: constants.DatasetInitializer}).
 											WithSpec(corev1ac.PodSpec().
 												WithContainers(
 													corev1ac.Container().
@@ -551,9 +564,11 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
 							jobsetv1alpha2ac.ReplicatedJob().
 								WithName(constants.ModelInitializer).
 								WithTemplate(batchv1ac.JobTemplateSpec().
+									WithLabels(map[string]string{
+										constants.LabelTrainJobAncestor: constants.ModelInitializer,
+									}).
 									WithSpec(batchv1ac.JobSpec().
 										WithTemplate(corev1ac.PodTemplateSpec().
-											WithLabels(map[string]string{constants.LabelTrainJobAncestor: constants.ModelInitializer}).
 											WithSpec(corev1ac.PodSpec().
 												WithContainers(
 													corev1ac.Container().
@@ -575,14 +590,17 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
 									),
 								),
 							jobsetv1alpha2ac.ReplicatedJob().
-								WithName(constants.JobTrainerNode).
+								WithName(constants.Node).
 								WithTemplate(batchv1ac.JobTemplateSpec().
+									WithLabels(map[string]string{
+										constants.LabelTrainJobAncestor: constants.AncestorTrainer,
+									}).
 									WithSpec(batchv1ac.JobSpec().
 										WithTemplate(corev1ac.PodTemplateSpec().
 											WithSpec(corev1ac.PodSpec().
 												WithContainers(
 													corev1ac.Container().
-														WithName(constants.ContainerTrainer).
+														WithName(constants.Node).
 														WithVolumeMounts(
 															corev1ac.VolumeMount().
 																WithName(jobsetplgconsts.VolumeNameInitializer).
@@ -636,10 +654,12 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
 								WithReplicas(1).
 								WithName(constants.DatasetInitializer).
 								WithTemplate(batchv1ac.JobTemplateSpec().
+									WithLabels(map[string]string{
+										constants.LabelTrainJobAncestor: constants.DatasetInitializer,
+									}).
 									WithSpec(batchv1ac.JobSpec().
 										WithTemplate(corev1ac.PodTemplateSpec().
 											WithLabels(map[string]string{
-												constants.LabelTrainJobAncestor:        constants.DatasetInitializer,
 												schedulerpluginsv1alpha1.PodGroupLabel: "test-job",
 											}).
 											WithSpec(corev1ac.PodSpec().
@@ -666,10 +686,12 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
 								WithReplicas(1).
 								WithName(constants.ModelInitializer).
 								WithTemplate(batchv1ac.JobTemplateSpec().
+									WithLabels(map[string]string{
+										constants.LabelTrainJobAncestor: constants.ModelInitializer,
+									}).
 									WithSpec(batchv1ac.JobSpec().
 										WithTemplate(corev1ac.PodTemplateSpec().
 											WithLabels(map[string]string{
-												constants.LabelTrainJobAncestor:        constants.ModelInitializer,
 												schedulerpluginsv1alpha1.PodGroupLabel: "test-job",
 											}).
 											WithSpec(corev1ac.PodSpec().
@@ -693,9 +715,12 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
 									),
 								),
 							jobsetv1alpha2ac.ReplicatedJob().
-								WithName(constants.JobTrainerNode).
+								WithName(constants.Node).
 								WithReplicas(1).
 								WithTemplate(batchv1ac.JobTemplateSpec().
+									WithLabels(map[string]string{
+										constants.LabelTrainJobAncestor: constants.AncestorTrainer,
+									}).
 									WithSpec(batchv1ac.JobSpec().
 										WithParallelism(100).
 										WithCompletions(100).
@@ -706,7 +731,7 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
 											WithSpec(corev1ac.PodSpec().
 												WithContainers(
 													corev1ac.Container().
-														WithName(constants.ContainerTrainer).
+														WithName(constants.Node).
 														WithImage("test:trainjob").
 														WithCommand("trainjob").
 														WithArgs("trainjob").
@@ -737,8 +762,9 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
 						),
 					PodSets: []runtime.PodSet{
 						{
-							Name:  constants.DatasetInitializer,
-							Count: ptr.To[int32](1),
+							Name:     constants.DatasetInitializer,
+							Ancestor: ptr.To(constants.DatasetInitializer),
+							Count:    ptr.To[int32](1),
 							SinglePodRequests: corev1.ResourceList{
 								corev1.ResourceCPU:    resource.MustParse("1"),
 								corev1.ResourceMemory: resource.MustParse("4Gi"),
@@ -761,8 +787,9 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
 							},
 						},
 						{
-							Name:  constants.ModelInitializer,
-							Count: ptr.To[int32](1),
+							Name:     constants.ModelInitializer,
+							Ancestor: ptr.To(constants.ModelInitializer),
+							Count:    ptr.To[int32](1),
 							SinglePodRequests: corev1.ResourceList{
 								corev1.ResourceCPU:    resource.MustParse("1"),
 								corev1.ResourceMemory: resource.MustParse("4Gi"),
@@ -785,8 +812,9 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
 							},
 						},
 						{
-							Name:  constants.JobTrainerNode,
-							Count: ptr.To[int32](100),
+							Name:     constants.Node,
+							Ancestor: ptr.To(constants.AncestorTrainer),
+							Count:    ptr.To[int32](100),
 							SinglePodRequests: corev1.ResourceList{
 								corev1.ResourceCPU:    resource.MustParse("1"),
 								corev1.ResourceMemory: resource.MustParse("4Gi"),
@@ -821,9 +849,9 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
 				testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
 					ControllerReference(trainer.SchemeGroupVersion.WithKind("TrainJob"), "test-job", "uid").
 					PodLabel(schedulerpluginsv1alpha1.PodGroupLabel, "test-job").
-					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.JobTrainerNode).
+					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node).
 					NumNodes(100).
-					Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, corev1.ResourceList{
+					Container(constants.Node, constants.Node, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, corev1.ResourceList{
 						corev1.ResourceCPU:    resource.MustParse("1"),
 						corev1.ResourceMemory: resource.MustParse("4Gi"),
 					}).
@@ -1046,7 +1074,7 @@ func TestPodNetworkPlugins(t *testing.T) {
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{
 						{
-							Name:       constants.JobTrainerNode,
+							Name:       constants.Node,
 							Count:      ptr.To[int32](2),
 							Containers: make([]runtime.Container, 1),
 						},
@@ -1054,7 +1082,7 @@ func TestPodNetworkPlugins(t *testing.T) {
 					ObjApply: jobsetv1alpha2ac.JobSetSpec().
 						WithReplicatedJobs(
 							jobsetv1alpha2ac.ReplicatedJob().
-								WithName(constants.JobTrainerNode).
+								WithName(constants.Node).
 								WithTemplate(batchv1ac.JobTemplateSpec().
 									WithSpec(batchv1ac.JobSpec().
 										WithParallelism(1).
@@ -1063,7 +1091,7 @@ func TestPodNetworkPlugins(t *testing.T) {
 											WithSpec(corev1ac.PodSpec().
 												WithContainers(
 													corev1ac.Container().
-														WithName(constants.ContainerTrainer),
+														WithName(constants.Node),
 												),
 											),
 										),
@@ -1079,19 +1107,19 @@ func TestPodNetworkPlugins(t *testing.T) {
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{
 						{
-							Name:       constants.JobTrainerNode,
+							Name:       constants.Node,
 							Containers: make([]runtime.Container, 1),
 							Count:      ptr.To[int32](2),
 							Endpoints: func(yield func(string) bool) {
-								yield("test-job-trainer-node-0-0.test-job")
-								yield("test-job-trainer-node-0-1.test-job")
+								yield("test-job-node-0-0.test-job")
+								yield("test-job-node-0-1.test-job")
 							},
 						},
 					},
 					ObjApply: jobsetv1alpha2ac.JobSetSpec().
 						WithReplicatedJobs(
 							jobsetv1alpha2ac.ReplicatedJob().
-								WithName(constants.JobTrainerNode).
+								WithName(constants.Node).
 								WithTemplate(batchv1ac.JobTemplateSpec().
 									WithSpec(batchv1ac.JobSpec().
 										WithParallelism(1).
@@ -1100,7 +1128,7 @@ func TestPodNetworkPlugins(t *testing.T) {
 											WithSpec(corev1ac.PodSpec().
 												WithContainers(
 													corev1ac.Container().
-														WithName(constants.ContainerTrainer),
+														WithName(constants.Node),
 												),
 											),
 										),
diff --git a/pkg/runtime/framework/plugins/jobset/builder.go b/pkg/runtime/framework/plugins/jobset/builder.go
index d6ed50bb21..eb6fb578e8 100644
--- a/pkg/runtime/framework/plugins/jobset/builder.go
+++ b/pkg/runtime/framework/plugins/jobset/builder.go
@@ -41,11 +41,12 @@ func NewBuilder(jobSet *jobsetv1alpha2ac.JobSetApplyConfiguration) *Builder {
 // Initializer updates JobSet values for the initializer Job.
 func (b *Builder) Initializer(trainJob *trainer.TrainJob) *Builder {
 	for i, rJob := range b.Spec.ReplicatedJobs {
-		if rJob.Template.Spec.Template.ObjectMetaApplyConfiguration == nil {
+		jobMetadata := rJob.Template.ObjectMetaApplyConfiguration
+		if jobMetadata == nil || jobMetadata.Labels == nil {
 			continue
 		}
 		// Update values for the Dataset Initializer Job.
-		if ancestor, ok := rJob.Template.Spec.Template.Labels[constants.LabelTrainJobAncestor]; ok && ancestor == constants.DatasetInitializer {
+		if ancestor, ok := jobMetadata.Labels[constants.LabelTrainJobAncestor]; ok && ancestor == constants.DatasetInitializer {
 			// TODO: Support multiple replicas ('.template.spec.replicatedJobs[*].replicas') for replicated Jobs.
 			// REF: https://github.com/kubeflow/trainer/issues/2318
 			b.Spec.ReplicatedJobs[i].Replicas = ptr.To[int32](1)
@@ -71,7 +72,7 @@ func (b *Builder) Initializer(trainJob *trainer.TrainJob) *Builder {
 			}
 		}
 		// Update values for the Model Initializer Job.
-		if ancestor, ok := rJob.Template.Spec.Template.Labels[constants.LabelTrainJobAncestor]; ok && ancestor == constants.ModelInitializer {
+		if ancestor, ok := jobMetadata.Labels[constants.LabelTrainJobAncestor]; ok && ancestor == constants.ModelInitializer {
 			// TODO: Support multiple replicas ('.template.spec.replicatedJobs[*].replicas') for replicated Jobs.
 			// REF: https://github.com/kubeflow/trainer/issues/2318
 			b.Spec.ReplicatedJobs[i].Replicas = ptr.To[int32](1)
@@ -101,7 +102,7 @@ func (b *Builder) Initializer(trainJob *trainer.TrainJob) *Builder {
 }
 
 // Launcher updates JobSet values for the launcher Job.
-func (b *Builder) Launcher() *Builder {
+func (b *Builder) Launcher(info *runtime.Info, trainJob *trainer.TrainJob) *Builder {
 	for i, rJob := range b.Spec.ReplicatedJobs {
 		if *rJob.Name == constants.JobLauncher {
 			// TODO: Support multiple replicas ('.template.spec.replicatedJobs[*].replicas') for replicated Jobs.
@@ -115,17 +116,21 @@ func (b *Builder) Launcher() *Builder {
 // Trainer updates JobSet values for the trainer Job.
 func (b *Builder) Trainer(info *runtime.Info, trainJob *trainer.TrainJob) *Builder {
 	for i, rJob := range b.Spec.ReplicatedJobs {
-		if *rJob.Name == constants.JobTrainerNode {
+		jobMetadata := rJob.Template.ObjectMetaApplyConfiguration
+		if jobMetadata == nil || jobMetadata.Labels == nil {
+			continue
+		}
+		if ancestor, ok := jobMetadata.Labels[constants.LabelTrainJobAncestor]; ok && ancestor == constants.AncestorTrainer {
 			// TODO: Support multiple replicas ('.template.spec.replicatedJobs[*].replicas') for replicated Jobs.
 			// REF: https://github.com/kubeflow/trainer/issues/2318
 			b.Spec.ReplicatedJobs[i].Replicas = ptr.To[int32](1)
 			// Update the Parallelism and Completions values for the Trainer Job.
-			b.Spec.ReplicatedJobs[i].Template.Spec.Parallelism = info.TemplateSpec.PodSets[i].Count
-			b.Spec.ReplicatedJobs[i].Template.Spec.Completions = info.TemplateSpec.PodSets[i].Count
+			b.Spec.ReplicatedJobs[i].Template.Spec.Parallelism = info.FindPodSetByAncestor(constants.AncestorTrainer).Count
+			b.Spec.ReplicatedJobs[i].Template.Spec.Completions = info.FindPodSetByAncestor(constants.AncestorTrainer).Count
 
 			// Update values for the Trainer container.
 			for j, container := range rJob.Template.Spec.Template.Spec.Containers {
-				if *container.Name == constants.ContainerTrainer {
+				if *container.Name == constants.Node {
 					// Update values from the TrainJob trainer.
 					if jobTrainer := trainJob.Spec.Trainer; jobTrainer != nil {
 						if image := jobTrainer.Image; image != nil {
diff --git a/pkg/runtime/framework/plugins/jobset/jobset.go b/pkg/runtime/framework/plugins/jobset/jobset.go
index 9674cccdb1..6165255528 100644
--- a/pkg/runtime/framework/plugins/jobset/jobset.go
+++ b/pkg/runtime/framework/plugins/jobset/jobset.go
@@ -146,7 +146,7 @@ func (j *JobSet) IdentifyPodNetwork(info *runtime.Info, trainJob *trainer.TrainJ
 		// TODO: Support multiple replicas for replicated Jobs.
 		// REF: https://github.com/kubeflow/trainer/issues/2318
 		podCount := info.TemplateSpec.PodSets[rJobIdx].Count
-		rJobReplicas := 1
+		rJobReplicas := constants.DefaultJobReplicas
 		info.TemplateSpec.PodSets[rJobIdx].Endpoints = func(yield func(string) bool) {
 			for podIdx := range ptr.Deref(podCount, 1) {
 				endpoint := fmt.Sprintf("%s-%s-%d-%d.%s", trainJob.Name, *rJob.Name, rJobReplicas-1, podIdx, subDomain)
@@ -194,7 +194,7 @@ func (j *JobSet) Build(ctx context.Context, info *runtime.Info, trainJob *traine
 	// TODO: Once we remove deprecated runtime.Info.Trainer, we should remove JobSet Builder with DeprecatedTrainer().
 	jobSet := jobSetBuilder.
 		Initializer(trainJob).
-		Launcher().
+		Launcher(info, trainJob).
 		Trainer(info, trainJob).
 		PodLabels(info.Scheduler.PodLabels).
 		Suspend(trainJob.Spec.Suspend).
diff --git a/pkg/runtime/framework/plugins/jobset/jobset_test.go b/pkg/runtime/framework/plugins/jobset/jobset_test.go
index 1a56fba9ce..d7d475b7ad 100644
--- a/pkg/runtime/framework/plugins/jobset/jobset_test.go
+++ b/pkg/runtime/framework/plugins/jobset/jobset_test.go
@@ -74,7 +74,7 @@ func TestJobSet(t *testing.T) {
 				},
 			},
 		},
-		"trainer numNodes is respected rather than parallelism when replicatedJob name is trainer-node": {
+		"trainer numNodes is respected rather than parallelism when replicatedJob name is node": {
 			trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "trainJob").
 				Obj(),
 			info: &runtime.Info{
@@ -90,7 +90,7 @@ func TestJobSet(t *testing.T) {
 							Containers: make([]runtime.Container, 1),
 						},
 						{
-							Name:       constants.JobTrainerNode,
+							Name:       constants.Node,
 							Count:      ptr.To[int32](2),
 							Containers: make([]runtime.Container, 1),
 						},
@@ -113,14 +113,14 @@ func TestJobSet(t *testing.T) {
 									),
 								),
 							jobsetv1alpha2ac.ReplicatedJob().
-								WithName(constants.JobTrainerNode).
+								WithName(constants.Node).
 								WithTemplate(batchv1ac.JobTemplateSpec().
 									WithSpec(batchv1ac.JobSpec().
 										WithParallelism(2).
 										WithTemplate(corev1ac.PodTemplateSpec().
 											WithSpec(corev1ac.PodSpec().
 												WithContainers(
-													corev1ac.Container().WithName(constants.ContainerTrainer),
+													corev1ac.Container().WithName(constants.Node),
 												),
 											),
 										),
@@ -145,12 +145,12 @@ func TestJobSet(t *testing.T) {
 							},
 						},
 						{
-							Name:       constants.JobTrainerNode,
+							Name:       constants.Node,
 							Count:      ptr.To[int32](2),
 							Containers: make([]runtime.Container, 1),
 							Endpoints: func(yield func(string) bool) {
-								yield("trainJob-trainer-node-0-0.trainJob")
-								yield("trainJob-trainer-node-0-1.trainJob")
+								yield("trainJob-node-0-0.trainJob")
+								yield("trainJob-node-0-1.trainJob")
 							},
 						},
 					},
@@ -172,14 +172,14 @@ func TestJobSet(t *testing.T) {
 									),
 								),
 							jobsetv1alpha2ac.ReplicatedJob().
-								WithName(constants.JobTrainerNode).
+								WithName(constants.Node).
 								WithTemplate(batchv1ac.JobTemplateSpec().
 									WithSpec(batchv1ac.JobSpec().
 										WithParallelism(2).
 										WithTemplate(corev1ac.PodTemplateSpec().
 											WithSpec(corev1ac.PodSpec().
 												WithContainers(
-													corev1ac.Container().WithName(constants.ContainerTrainer),
+													corev1ac.Container().WithName(constants.Node),
 												),
 											),
 										),
@@ -203,7 +203,7 @@ func TestJobSet(t *testing.T) {
 							Containers: make([]runtime.Container, 1),
 						},
 						{
-							Name:       constants.JobTrainerNode,
+							Name:       constants.Node,
 							Containers: make([]runtime.Container, 1),
 						},
 					},
@@ -226,14 +226,14 @@ func TestJobSet(t *testing.T) {
 									),
 								),
 							jobsetv1alpha2ac.ReplicatedJob().
-								WithName(constants.JobTrainerNode).
+								WithName(constants.Node).
 								WithTemplate(batchv1ac.JobTemplateSpec().
 									WithSpec(batchv1ac.JobSpec().
 										WithParallelism(1).
 										WithTemplate(corev1ac.PodTemplateSpec().
 											WithSpec(corev1ac.PodSpec().
 												WithContainers(
-													corev1ac.Container().WithName(constants.ContainerTrainer),
+													corev1ac.Container().WithName(constants.Node),
 												),
 											),
 										),
@@ -256,10 +256,10 @@ func TestJobSet(t *testing.T) {
 							},
 						},
 						{
-							Name:       constants.JobTrainerNode,
+							Name:       constants.Node,
 							Containers: make([]runtime.Container, 1),
 							Endpoints: func(yield func(string) bool) {
-								yield("trainJob-trainer-node-0-0.kubeflow.org")
+								yield("trainJob-node-0-0.kubeflow.org")
 							},
 						},
 					},
@@ -282,14 +282,14 @@ func TestJobSet(t *testing.T) {
 									),
 								),
 							jobsetv1alpha2ac.ReplicatedJob().
-								WithName(constants.JobTrainerNode).
+								WithName(constants.Node).
 								WithTemplate(batchv1ac.JobTemplateSpec().
 									WithSpec(batchv1ac.JobSpec().
 										WithParallelism(1).
 										WithTemplate(corev1ac.PodTemplateSpec().
 											WithSpec(corev1ac.PodSpec().
 												WithContainers(
-													corev1ac.Container().WithName(constants.ContainerTrainer),
+													corev1ac.Container().WithName(constants.Node),
 												),
 											),
 										),
diff --git a/pkg/runtime/framework/plugins/mpi/mpi.go b/pkg/runtime/framework/plugins/mpi/mpi.go
index 7d7927c705..f222fc32da 100644
--- a/pkg/runtime/framework/plugins/mpi/mpi.go
+++ b/pkg/runtime/framework/plugins/mpi/mpi.go
@@ -79,7 +79,7 @@ func (m *MPI) Name() string {
 // TODO (andreyvelich): We should validate that envs from different plugins don't conflict with each other.
 // Ref: https://github.com/kubeflow/trainer/pull/2308#discussion_r1823229940
 
-func (m *MPI) Validate(runtimeJobTemplate client.Object, runtimeInfo *runtime.Info, oldJobObj, newJobObj *trainer.TrainJob) (admission.Warnings, field.ErrorList) {
+func (m *MPI) Validate(_ client.Object, runtimeInfo *runtime.Info, _, newJobObj *trainer.TrainJob) (admission.Warnings, field.ErrorList) {
 	var allErrs field.ErrorList
 	if runtimeInfo == nil || runtimeInfo.RuntimePolicy.MLPolicySource == nil || runtimeInfo.RuntimePolicy.MLPolicySource.MPI == nil {
 		return nil, allErrs
@@ -103,7 +103,16 @@ func (m *MPI) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob) er
 
 	// TrainJob contains the actual information for the Trainer.
 	if trainJob.Spec.Trainer != nil && trainJob.Spec.Trainer.NumNodes != nil {
-		info.FindPodSetByName(constants.JobTrainerNode).Count = trainJob.Spec.Trainer.NumNodes
+		if node := info.FindPodSetByName(constants.Node); node != nil && node.Count != nil {
+			if ptr.Deref(info.RuntimePolicy.MLPolicySource.MPI.RunLauncherAsNode, false) {
+				// TODO: We should implement more strong validations for the MPIRuntime with runLauncherAsNode.
+				// REF: https://github.com/kubeflow/trainer/issues/2550
+				// When runLauncherAsNode is enabled, 1 nodes should be allocated to launcher.
+				*node.Count = max(*trainJob.Spec.Trainer.NumNodes-1, 1)
+			} else {
+				*node.Count = *trainJob.Spec.Trainer.NumNodes
+			}
+		}
 	}
 
 	if trainJob.Spec.Trainer != nil && trainJob.Spec.Trainer.NumProcPerNode != nil {
@@ -112,7 +121,7 @@ func (m *MPI) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob) er
 
 	// Add Secret and ConfigMap volumes to the Info object
 	for psIdx, ps := range info.TemplateSpec.PodSets {
-		if ps.Name != constants.JobTrainerNode && ps.Name != constants.JobLauncher {
+		if ps.Name != constants.Node && ps.Name != constants.JobLauncher {
 			continue
 		}
 		apply.UpsertVolumes(
@@ -155,7 +164,7 @@ func (m *MPI) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob) er
 			)
 		}
 		for cIdx, container := range ps.Containers {
-			if container.Name != constants.JobLauncher && container.Name != constants.JobTrainerNode {
+			if container.Name != constants.JobLauncher && container.Name != constants.Node {
 				continue
 			}
 			apply.UpsertVolumeMounts(
@@ -274,7 +283,7 @@ func (m *MPI) buildHostFileConfigMap(info *runtime.Info, trainJob *trainer.Train
 	runLauncherAsNode := ptr.Deref(info.RuntimePolicy.MLPolicySource.MPI.RunLauncherAsNode, false)
 	slots := ptr.Deref(info.RuntimePolicy.MLPolicySource.MPI.NumProcPerNode, 1)
 	for _, ps := range info.TemplateSpec.PodSets {
-		if !isTrainerNode(runLauncherAsNode, ps) {
+		if !isNode(runLauncherAsNode, ps) {
 			continue
 		}
 		switch *info.RuntimePolicy.MLPolicySource.MPI.MPIImplementation {
@@ -297,6 +306,6 @@ func (m *MPI) buildHostFileConfigMap(info *runtime.Info, trainJob *trainer.Train
 			WithBlockOwnerDeletion(true))
 }
 
-func isTrainerNode(runLauncherAsNode bool, ps runtime.PodSet) bool {
-	return (runLauncherAsNode && ps.Name == constants.JobLauncher) || ps.Name == constants.JobTrainerNode
+func isNode(runLauncherAsNode bool, ps runtime.PodSet) bool {
+	return (runLauncherAsNode && ps.Name == constants.JobLauncher) || ps.Name == constants.Node
 }
diff --git a/pkg/runtime/framework/plugins/mpi/mpi_test.go b/pkg/runtime/framework/plugins/mpi/mpi_test.go
index 7f5f02dc79..b7910a91d1 100644
--- a/pkg/runtime/framework/plugins/mpi/mpi_test.go
+++ b/pkg/runtime/framework/plugins/mpi/mpi_test.go
@@ -47,6 +47,7 @@ func TestMPI(t *testing.T) {
 		cmpopts.SortSlices(func(a, b apiruntime.Object) int {
 			return cmp.Compare(a.GetObjectKind().GroupVersionKind().String(), b.GetObjectKind().GroupVersionKind().String())
 		}),
+		cmpopts.SortSlices(func(a, b corev1.EnvVar) int { return cmp.Compare(a.Name, b.Name) }),
 		gocmp.Comparer(utiltesting.MPISecretDataComparer),
 	}
 	errorGetSSHAuthSecretFromAPI := errors.New("failed to get SSH Auth Secret from API during Build")
@@ -88,16 +89,19 @@ func TestMPI(t *testing.T) {
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{
 						{
-							Name: constants.JobLauncher,
+							Name:  constants.JobLauncher,
+							Count: ptr.To[int32](1),
 							Endpoints: func(yield func(string) bool) {
 								yield("trainJob-launcher-0-0.trainJob")
 							},
 						},
 						{
-							Name: constants.JobTrainerNode,
+							Name:     constants.Node,
+							Ancestor: ptr.To(constants.AncestorTrainer),
+							Count:    ptr.To[int32](1),
 							Endpoints: func(yield func(string) bool) {
-								yield("trainJob-trainer-node-1-0.trainJob")
-								yield("trainJob-trainer-node-1-1.trainJob")
+								yield("trainJob-node-1-0.trainJob")
+								yield("trainJob-node-1-1.trainJob")
 							},
 						},
 					},
@@ -124,7 +128,8 @@ func TestMPI(t *testing.T) {
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{
 						{
-							Name: constants.JobLauncher,
+							Name:  constants.JobLauncher,
+							Count: ptr.To[int32](1),
 							Volumes: []corev1ac.VolumeApplyConfiguration{
 								*corev1ac.Volume().
 									WithName(constants.MPISSHAuthVolumeName).
@@ -159,8 +164,9 @@ func TestMPI(t *testing.T) {
 							},
 						},
 						{
-							Name:  constants.JobTrainerNode,
-							Count: ptr.To[int32](2),
+							Name:     constants.Node,
+							Ancestor: ptr.To(constants.AncestorTrainer),
+							Count:    ptr.To[int32](2),
 							Volumes: []corev1ac.VolumeApplyConfiguration{
 								*corev1ac.Volume().
 									WithName(constants.MPISSHAuthVolumeName).
@@ -180,8 +186,8 @@ func TestMPI(t *testing.T) {
 									),
 							},
 							Endpoints: func(yield func(string) bool) {
-								yield("trainJob-trainer-node-1-0.trainJob")
-								yield("trainJob-trainer-node-1-1.trainJob")
+								yield("trainJob-node-1-0.trainJob")
+								yield("trainJob-node-1-1.trainJob")
 							},
 						},
 					},
@@ -205,8 +211,8 @@ func TestMPI(t *testing.T) {
 					Obj(),
 				utiltesting.MakeConfigMapWrapper(fmt.Sprintf("trainJob%s", constants.MPIHostfileConfigMapSuffix), metav1.NamespaceDefault).
 					WithData(map[string]string{
-						constants.MPIHostfileName: `trainJob-trainer-node-1-0.trainJob slots=1
-trainJob-trainer-node-1-1.trainJob slots=1
+						constants.MPIHostfileName: `trainJob-node-1-0.trainJob slots=1
+trainJob-node-1-1.trainJob slots=1
 `,
 					}).
 					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "trainJob", "trainJob").
@@ -220,15 +226,18 @@ trainJob-trainer-node-1-1.trainJob slots=1
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{
 						{
-							Name: constants.JobLauncher,
+							Name:  constants.JobLauncher,
+							Count: ptr.To[int32](1),
 							Endpoints: func(yield func(string) bool) {
 								yield("trainJob-launcher-0-0.trainJob")
 							},
 						},
 						{
-							Name: constants.JobTrainerNode,
+							Name:     constants.Node,
+							Ancestor: ptr.To(constants.AncestorTrainer),
+							Count:    ptr.To[int32](1),
 							Endpoints: func(yield func(string) bool) {
-								yield("trainJob-trainer-node-1-0.trainJob")
+								yield("trainJob-node-1-0.trainJob")
 							},
 						},
 					},
@@ -256,7 +265,8 @@ trainJob-trainer-node-1-1.trainJob slots=1
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{
 						{
-							Name: constants.JobLauncher,
+							Name:  constants.JobLauncher,
+							Count: ptr.To[int32](1),
 							Volumes: []corev1ac.VolumeApplyConfiguration{
 								*corev1ac.Volume().
 									WithName(constants.MPISSHAuthVolumeName).
@@ -291,8 +301,9 @@ trainJob-trainer-node-1-1.trainJob slots=1
 							},
 						},
 						{
-							Name:  constants.JobTrainerNode,
-							Count: ptr.To[int32](1),
+							Name:     constants.Node,
+							Ancestor: ptr.To(constants.AncestorTrainer),
+							Count:    ptr.To[int32](1),
 							Volumes: []corev1ac.VolumeApplyConfiguration{
 								*corev1ac.Volume().
 									WithName(constants.MPISSHAuthVolumeName).
@@ -312,7 +323,7 @@ trainJob-trainer-node-1-1.trainJob slots=1
 									),
 							},
 							Endpoints: func(yield func(string) bool) {
-								yield("trainJob-trainer-node-1-0.trainJob")
+								yield("trainJob-node-1-0.trainJob")
 							},
 						},
 					},
@@ -336,7 +347,7 @@ trainJob-trainer-node-1-1.trainJob slots=1
 					Obj(),
 				utiltesting.MakeConfigMapWrapper(fmt.Sprintf("trainJob%s", constants.MPIHostfileConfigMapSuffix), metav1.NamespaceDefault).
 					WithData(map[string]string{
-						constants.MPIHostfileName: `trainJob-trainer-node-1-0.trainJob slots=2
+						constants.MPIHostfileName: `trainJob-node-1-0.trainJob slots=2
 `,
 					}).
 					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "trainJob", "trainJob").
@@ -355,16 +366,25 @@ trainJob-trainer-node-1-1.trainJob slots=1
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{
 						{
-							Name: constants.JobLauncher,
+							Name:  constants.JobLauncher,
+							Count: ptr.To[int32](1),
 							Endpoints: func(yield func(string) bool) {
 								yield("trainJob-launcher-0-0.trainJob")
 							},
+							Containers: []runtime.Container{{
+								Name: constants.ContainerLauncher,
+							}},
 						},
 						{
-							Name: constants.JobTrainerNode,
+							Name:     constants.Node,
+							Ancestor: ptr.To(constants.AncestorTrainer),
+							Count:    ptr.To[int32](1),
 							Endpoints: func(yield func(string) bool) {
-								yield("trainJob-trainer-node-1-0.trainJob")
+								yield("trainJob-node-1-0.trainJob")
 							},
+							Containers: []runtime.Container{{
+								Name: constants.Node,
+							}},
 						},
 					},
 				},
@@ -374,7 +394,7 @@ trainJob-trainer-node-1-1.trainJob slots=1
 				UID("trainJob").
 				Trainer(
 					utiltesting.MakeTrainJobTrainerWrapper().
-						NumNodes(1).
+						NumNodes(10).
 						Obj()).
 				Obj(),
 			wantInfo: &runtime.Info{
@@ -388,7 +408,33 @@ trainJob-trainer-node-1-1.trainJob slots=1
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{
 						{
-							Name: constants.JobLauncher,
+							Name:  constants.JobLauncher,
+							Count: ptr.To[int32](1),
+							Containers: []runtime.Container{{
+								Name: constants.ContainerLauncher,
+								VolumeMounts: []corev1ac.VolumeMountApplyConfiguration{
+									*corev1ac.VolumeMount().
+										WithName(constants.MPISSHAuthVolumeName).
+										WithMountPath("/root/.ssh"),
+									*corev1ac.VolumeMount().
+										WithName(constants.MPIHostfileVolumeName).
+										WithMountPath("/etc/mpi"),
+								},
+								Env: []corev1ac.EnvVarApplyConfiguration{
+									*corev1ac.EnvVar().
+										WithName(constants.OpenMPIEnvHostFileLocation).
+										WithValue(fmt.Sprintf("%s/%s", constants.MPIHostfileDir, constants.MPIHostfileName)),
+									*corev1ac.EnvVar().
+										WithName(constants.OpenMPIEnvKeepFQDNHostNames).
+										WithValue("true"),
+									*corev1ac.EnvVar().
+										WithName(constants.OpenMPIEnvDefaultSlots).
+										WithValue("1"),
+									*corev1ac.EnvVar().
+										WithName(constants.OpenMPIEnvKeyRSHArgs).
+										WithValue(constants.OpenMPIEnvDefaultValueRSHArgs),
+								},
+							}},
 							Volumes: []corev1ac.VolumeApplyConfiguration{
 								*corev1ac.Volume().
 									WithName(constants.MPISSHAuthVolumeName).
@@ -423,8 +469,17 @@ trainJob-trainer-node-1-1.trainJob slots=1
 							},
 						},
 						{
-							Name:  constants.JobTrainerNode,
-							Count: ptr.To[int32](1),
+							Name:     constants.Node,
+							Ancestor: ptr.To(constants.AncestorTrainer),
+							Count:    ptr.To[int32](9),
+							Containers: []runtime.Container{{
+								Name: constants.Node,
+								VolumeMounts: []corev1ac.VolumeMountApplyConfiguration{
+									*corev1ac.VolumeMount().
+										WithName(constants.MPISSHAuthVolumeName).
+										WithMountPath("/root/.ssh"),
+								},
+							}},
 							Volumes: []corev1ac.VolumeApplyConfiguration{
 								*corev1ac.Volume().
 									WithName(constants.MPISSHAuthVolumeName).
@@ -444,7 +499,7 @@ trainJob-trainer-node-1-1.trainJob slots=1
 									),
 							},
 							Endpoints: func(yield func(string) bool) {
-								yield("trainJob-trainer-node-1-0.trainJob")
+								yield("trainJob-node-1-0.trainJob")
 							},
 						},
 					},
@@ -464,7 +519,7 @@ trainJob-trainer-node-1-1.trainJob slots=1
 				utiltesting.MakeConfigMapWrapper(fmt.Sprintf("trainJob%s", constants.MPIHostfileConfigMapSuffix), metav1.NamespaceDefault).
 					WithData(map[string]string{
 						constants.MPIHostfileName: `trainJob-launcher-0-0.trainJob slots=1
-trainJob-trainer-node-1-0.trainJob slots=1
+trainJob-node-1-0.trainJob slots=1
 `,
 					}).
 					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "trainJob", "trainJob").
diff --git a/pkg/runtime/framework/plugins/plainml/plainml.go b/pkg/runtime/framework/plugins/plainml/plainml.go
index ab17752cc0..d60108e89f 100644
--- a/pkg/runtime/framework/plugins/plainml/plainml.go
+++ b/pkg/runtime/framework/plugins/plainml/plainml.go
@@ -50,13 +50,15 @@ func (p *PlainML) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob
 
 	// TrainJob contains the actual information for the number of nodes.
 	if trainJob.Spec.Trainer != nil && trainJob.Spec.Trainer.NumNodes != nil {
-		info.FindPodSetByName(constants.JobTrainerNode).Count = trainJob.Spec.Trainer.NumNodes
+		if trainerPS := info.FindPodSetByAncestor(constants.AncestorTrainer); trainerPS != nil && trainerPS.Count != nil {
+			*trainerPS.Count = *trainJob.Spec.Trainer.NumNodes
+		}
 	}
 
 	// Add envs from the TrainJob.
 	var trainerContainer *runtime.Container
 	if trainJob.Spec.Trainer != nil {
-		if trainerContainer = info.FindContainerByPodSetContainerName(constants.JobTrainerNode, constants.ContainerTrainer); trainerContainer != nil {
+		if trainerContainer = info.FindContainerByPodSetAncestorContainerName(constants.AncestorTrainer, constants.Node); trainerContainer != nil {
 			apply.UpsertEnvVars(&trainerContainer.Env, apply.EnvVars(trainJob.Spec.Trainer.Env...)...)
 		}
 	}
diff --git a/pkg/runtime/framework/plugins/plainml/plainml_test.go b/pkg/runtime/framework/plugins/plainml/plainml_test.go
index b026323ec0..44ac32c7b4 100644
--- a/pkg/runtime/framework/plugins/plainml/plainml_test.go
+++ b/pkg/runtime/framework/plugins/plainml/plainml_test.go
@@ -89,8 +89,8 @@ func TestPlainML(t *testing.T) {
 				runtime.WithMLPolicySource(
 					utiltesting.MakeMLPolicyWrapper().Obj(),
 				),
-				runtime.WithPodSet(constants.JobTrainerNode, 100, corev1.PodSpec{}, corev1ac.PodSpec().
-					WithContainers(corev1ac.Container().WithName(constants.ContainerTrainer)),
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 100, corev1.PodSpec{}, corev1ac.PodSpec().
+					WithContainers(corev1ac.Container().WithName(constants.Node)),
 				),
 			),
 			wantInfo: &runtime.Info{
@@ -101,11 +101,12 @@ func TestPlainML(t *testing.T) {
 				},
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{{
-						Name:              constants.JobTrainerNode,
+						Name:              constants.Node,
+						Ancestor:          ptr.To(constants.AncestorTrainer),
 						Count:             ptr.To[int32](200),
 						SinglePodRequests: make(corev1.ResourceList),
 						Containers: []runtime.Container{{
-							Name: constants.ContainerTrainer,
+							Name: constants.Node,
 						}},
 					}},
 				},
@@ -129,9 +130,9 @@ func TestPlainML(t *testing.T) {
 				runtime.WithMLPolicySource(
 					utiltesting.MakeMLPolicyWrapper().Obj(),
 				),
-				runtime.WithPodSet(constants.JobTrainerNode, 100, corev1.PodSpec{}, corev1ac.PodSpec().
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 100, corev1.PodSpec{}, corev1ac.PodSpec().
 					WithContainers(corev1ac.Container().
-						WithName(constants.ContainerTrainer).
+						WithName(constants.Node).
 						WithEnv(corev1ac.EnvVar().
 							WithName("CONFLICT").
 							WithValue("FROM_TRAINER"),
@@ -147,11 +148,12 @@ func TestPlainML(t *testing.T) {
 				},
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{{
-						Name:              constants.JobTrainerNode,
+						Name:              constants.Node,
+						Ancestor:          ptr.To(constants.AncestorTrainer),
 						Count:             ptr.To[int32](1),
 						SinglePodRequests: make(corev1.ResourceList),
 						Containers: []runtime.Container{{
-							Name: constants.ContainerTrainer,
+							Name: constants.Node,
 							Env: []corev1ac.EnvVarApplyConfiguration{
 								*corev1ac.EnvVar().WithName("CONFLICT").WithValue("FROM_TRAINER"),
 							},
diff --git a/pkg/runtime/framework/plugins/torch/torch.go b/pkg/runtime/framework/plugins/torch/torch.go
index dbd62c512f..7950a83de4 100644
--- a/pkg/runtime/framework/plugins/torch/torch.go
+++ b/pkg/runtime/framework/plugins/torch/torch.go
@@ -52,7 +52,7 @@ func (t *Torch) Name() string {
 	return Name
 }
 
-func (t *Torch) Validate(runtimeJobTemplate client.Object, runtimeInfo *runtime.Info, oldObj, newObj *trainer.TrainJob) (admission.Warnings, field.ErrorList) {
+func (t *Torch) Validate(_ client.Object, runtimeInfo *runtime.Info, _, newObj *trainer.TrainJob) (admission.Warnings, field.ErrorList) {
 	var allErrs field.ErrorList
 	if runtimeInfo == nil || runtimeInfo.RuntimePolicy.MLPolicySource == nil || runtimeInfo.RuntimePolicy.MLPolicySource.Torch == nil || newObj.Spec.Trainer.NumProcPerNode == nil {
 		return nil, allErrs
@@ -93,9 +93,9 @@ func (t *Torch) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob)
 	}
 
 	// TrainJob contains the actual information for the Trainer.
-	numNodes := info.FindPodSetByName(constants.JobTrainerNode).Count
-	if trainJob.Spec.Trainer != nil && trainJob.Spec.Trainer.NumNodes != nil {
-		*numNodes = *trainJob.Spec.Trainer.NumNodes
+	trainerPS := info.FindPodSetByAncestor(constants.AncestorTrainer)
+	if trainerPS != nil && trainerPS.Count != nil && trainJob.Spec.Trainer != nil && trainJob.Spec.Trainer.NumNodes != nil {
+		*trainerPS.Count = *trainJob.Spec.Trainer.NumNodes
 	}
 
 	numProcPerNode := ptr.Deref(info.RuntimePolicy.MLPolicySource.Torch.NumProcPerNode, intstr.FromString("auto"))
@@ -142,7 +142,7 @@ func (t *Torch) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob)
 	// Ref: https://github.com/kubeflow/trainer/pull/2308#discussion_r1823229940
 	var trainerContainer *runtime.Container
 	if trainJob.Spec.Trainer != nil {
-		if trainerContainer = info.FindContainerByPodSetContainerName(constants.JobTrainerNode, constants.ContainerTrainer); trainerContainer != nil {
+		if trainerContainer = info.FindContainerByPodSetAncestorContainerName(constants.AncestorTrainer, constants.Node); trainerContainer != nil {
 			apply.UpsertEnvVars(&trainerContainer.Env, apply.EnvVars(trainJob.Spec.Trainer.Env...)...)
 		}
 	}
@@ -150,7 +150,7 @@ func (t *Torch) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob)
 		apply.UpsertEnvVar(&trainerContainer.Env,
 			*corev1ac.EnvVar().
 				WithName(constants.TorchEnvNumNodes).
-				WithValue(fmt.Sprintf("%d", ptr.Deref(numNodes, 1))),
+				WithValue(fmt.Sprintf("%d", ptr.Deref(ptr.Deref(trainerPS, runtime.PodSet{}).Count, 1))),
 			*corev1ac.EnvVar().
 				WithName(constants.TorchEnvNumProcPerNode).
 				WithValue(numProcPerNode.String()),
@@ -161,7 +161,7 @@ func (t *Torch) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob)
 						WithFieldPath(constants.JobCompletionIndexFieldPath))),
 			*corev1ac.EnvVar().
 				WithName(constants.TorchEnvMasterAddr).
-				WithValue(fmt.Sprintf("%s-%s-0-0.%s", trainJob.Name, constants.JobTrainerNode, trainJob.Name)),
+				WithValue(fmt.Sprintf("%s-%s-0-0.%s", trainJob.Name, constants.Node, trainJob.Name)),
 			*corev1ac.EnvVar().
 				WithName(constants.TorchEnvMasterPort).
 				WithValue(fmt.Sprintf("%d", constants.ContainerTrainerPort)),
diff --git a/pkg/runtime/framework/plugins/torch/torch_test.go b/pkg/runtime/framework/plugins/torch/torch_test.go
index bd47f4b326..f082ce8dc7 100644
--- a/pkg/runtime/framework/plugins/torch/torch_test.go
+++ b/pkg/runtime/framework/plugins/torch/torch_test.go
@@ -76,9 +76,9 @@ func TestTorch(t *testing.T) {
 						).
 						Obj(),
 				),
-				runtime.WithPodSet(constants.JobTrainerNode, 1, corev1.PodSpec{}, corev1ac.PodSpec().
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 1, corev1.PodSpec{}, corev1ac.PodSpec().
 					WithContainers(
-						corev1ac.Container().WithName(constants.ContainerTrainer),
+						corev1ac.Container().WithName(constants.Node),
 					),
 				),
 			),
@@ -98,11 +98,12 @@ func TestTorch(t *testing.T) {
 				},
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{{
-						Name:              constants.JobTrainerNode,
+						Name:              constants.Node,
+						Ancestor:          ptr.To(constants.AncestorTrainer),
 						Count:             ptr.To[int32](2),
 						SinglePodRequests: make(corev1.ResourceList),
 						Containers: []runtime.Container{{
-							Name: constants.ContainerTrainer,
+							Name: constants.Node,
 							Ports: []corev1ac.ContainerPortApplyConfiguration{{
 								ContainerPort: ptr.To[int32](constants.ContainerTrainerPort),
 							}},
@@ -125,7 +126,7 @@ func TestTorch(t *testing.T) {
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterAddr),
-									Value: ptr.To("trainJob-trainer-node-0-0.trainJob"),
+									Value: ptr.To("trainJob-node-0-0.trainJob"),
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterPort),
@@ -158,8 +159,8 @@ func TestTorch(t *testing.T) {
 						).
 						Obj(),
 				),
-				runtime.WithPodSet(constants.JobTrainerNode, 1, corev1.PodSpec{}, corev1ac.PodSpec().
-					WithContainers(corev1ac.Container().WithName(constants.ContainerTrainer)),
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 1, corev1.PodSpec{}, corev1ac.PodSpec().
+					WithContainers(corev1ac.Container().WithName(constants.Node)),
 				),
 			),
 			wantInfo: &runtime.Info{
@@ -172,11 +173,12 @@ func TestTorch(t *testing.T) {
 				},
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{{
-						Name:              constants.JobTrainerNode,
+						Name:              constants.Node,
+						Ancestor:          ptr.To(constants.AncestorTrainer),
 						Count:             ptr.To[int32](1),
 						SinglePodRequests: make(corev1.ResourceList),
 						Containers: []runtime.Container{{
-							Name: constants.ContainerTrainer,
+							Name: constants.Node,
 							Ports: []corev1ac.ContainerPortApplyConfiguration{{
 								ContainerPort: ptr.To[int32](constants.ContainerTrainerPort),
 							}},
@@ -199,7 +201,7 @@ func TestTorch(t *testing.T) {
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterAddr),
-									Value: ptr.To("test-job-trainer-node-0-0.test-job"),
+									Value: ptr.To("test-job-node-0-0.test-job"),
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterPort),
@@ -230,8 +232,8 @@ func TestTorch(t *testing.T) {
 						).
 						Obj(),
 				),
-				runtime.WithPodSet(constants.JobTrainerNode, 1, corev1.PodSpec{}, corev1ac.PodSpec().
-					WithContainers(corev1ac.Container().WithName(constants.ContainerTrainer)),
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 1, corev1.PodSpec{}, corev1ac.PodSpec().
+					WithContainers(corev1ac.Container().WithName(constants.Node)),
 				),
 			),
 			wantInfo: &runtime.Info{
@@ -244,11 +246,12 @@ func TestTorch(t *testing.T) {
 				},
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{{
-						Name:              constants.JobTrainerNode,
+						Name:              constants.Node,
+						Ancestor:          ptr.To(constants.AncestorTrainer),
 						Count:             ptr.To[int32](1),
 						SinglePodRequests: make(corev1.ResourceList),
 						Containers: []runtime.Container{{
-							Name: constants.ContainerTrainer,
+							Name: constants.Node,
 							Ports: []corev1ac.ContainerPortApplyConfiguration{{
 								ContainerPort: ptr.To[int32](constants.ContainerTrainerPort),
 							}},
@@ -271,7 +274,7 @@ func TestTorch(t *testing.T) {
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterAddr),
-									Value: ptr.To("test-job-trainer-node-0-0.test-job"),
+									Value: ptr.To("test-job-node-0-0.test-job"),
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterPort),
@@ -304,8 +307,8 @@ func TestTorch(t *testing.T) {
 						).
 						Obj(),
 				),
-				runtime.WithPodSet(constants.JobTrainerNode, 1, corev1.PodSpec{}, corev1ac.PodSpec().
-					WithContainers(corev1ac.Container().WithName(constants.ContainerTrainer)),
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 1, corev1.PodSpec{}, corev1ac.PodSpec().
+					WithContainers(corev1ac.Container().WithName(constants.Node)),
 				),
 			),
 			wantInfo: &runtime.Info{
@@ -318,11 +321,12 @@ func TestTorch(t *testing.T) {
 				},
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{{
-						Name:              constants.JobTrainerNode,
+						Name:              constants.Node,
+						Ancestor:          ptr.To(constants.AncestorTrainer),
 						Count:             ptr.To[int32](1),
 						SinglePodRequests: make(corev1.ResourceList),
 						Containers: []runtime.Container{{
-							Name: constants.ContainerTrainer,
+							Name: constants.Node,
 							Ports: []corev1ac.ContainerPortApplyConfiguration{{
 								ContainerPort: ptr.To[int32](constants.ContainerTrainerPort),
 							}},
@@ -345,7 +349,7 @@ func TestTorch(t *testing.T) {
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterAddr),
-									Value: ptr.To("test-job-trainer-node-0-0.test-job"),
+									Value: ptr.To("test-job-node-0-0.test-job"),
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterPort),
@@ -378,8 +382,8 @@ func TestTorch(t *testing.T) {
 						).
 						Obj(),
 				),
-				runtime.WithPodSet(constants.JobTrainerNode, 1, corev1.PodSpec{}, corev1ac.PodSpec().
-					WithContainers(corev1ac.Container().WithName(constants.ContainerTrainer)),
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 1, corev1.PodSpec{}, corev1ac.PodSpec().
+					WithContainers(corev1ac.Container().WithName(constants.Node)),
 				),
 			),
 			wantInfo: &runtime.Info{
@@ -392,11 +396,12 @@ func TestTorch(t *testing.T) {
 				},
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{{
-						Name:              constants.JobTrainerNode,
+						Name:              constants.Node,
+						Ancestor:          ptr.To(constants.AncestorTrainer),
 						Count:             ptr.To[int32](1),
 						SinglePodRequests: make(corev1.ResourceList),
 						Containers: []runtime.Container{{
-							Name: constants.ContainerTrainer,
+							Name: constants.Node,
 							Ports: []corev1ac.ContainerPortApplyConfiguration{{
 								ContainerPort: ptr.To[int32](constants.ContainerTrainerPort),
 							}},
@@ -419,7 +424,7 @@ func TestTorch(t *testing.T) {
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterAddr),
-									Value: ptr.To("test-job-trainer-node-0-0.test-job"),
+									Value: ptr.To("test-job-node-0-0.test-job"),
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterPort),
@@ -452,8 +457,8 @@ func TestTorch(t *testing.T) {
 						).
 						Obj(),
 				),
-				runtime.WithPodSet(constants.JobTrainerNode, 1, corev1.PodSpec{}, corev1ac.PodSpec().
-					WithContainers(corev1ac.Container().WithName(constants.ContainerTrainer)),
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 1, corev1.PodSpec{}, corev1ac.PodSpec().
+					WithContainers(corev1ac.Container().WithName(constants.Node)),
 				),
 			),
 			wantInfo: &runtime.Info{
@@ -466,11 +471,12 @@ func TestTorch(t *testing.T) {
 				},
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{{
-						Name:              constants.JobTrainerNode,
+						Name:              constants.Node,
+						Ancestor:          ptr.To(constants.AncestorTrainer),
 						Count:             ptr.To[int32](1),
 						SinglePodRequests: make(corev1.ResourceList),
 						Containers: []runtime.Container{{
-							Name: constants.ContainerTrainer,
+							Name: constants.Node,
 							Ports: []corev1ac.ContainerPortApplyConfiguration{{
 								ContainerPort: ptr.To[int32](constants.ContainerTrainerPort),
 							}},
@@ -493,7 +499,7 @@ func TestTorch(t *testing.T) {
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterAddr),
-									Value: ptr.To("test-job-trainer-node-0-0.test-job"),
+									Value: ptr.To("test-job-node-0-0.test-job"),
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterPort),
@@ -526,8 +532,8 @@ func TestTorch(t *testing.T) {
 						).
 						Obj(),
 				),
-				runtime.WithPodSet(constants.JobTrainerNode, 1, corev1.PodSpec{}, corev1ac.PodSpec().
-					WithContainers(corev1ac.Container().WithName(constants.ContainerTrainer)),
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 1, corev1.PodSpec{}, corev1ac.PodSpec().
+					WithContainers(corev1ac.Container().WithName(constants.Node)),
 				),
 			),
 			wantInfo: &runtime.Info{
@@ -540,11 +546,12 @@ func TestTorch(t *testing.T) {
 				},
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{{
-						Name:              constants.JobTrainerNode,
+						Name:              constants.Node,
+						Ancestor:          ptr.To(constants.AncestorTrainer),
 						Count:             ptr.To[int32](1),
 						SinglePodRequests: make(corev1.ResourceList),
 						Containers: []runtime.Container{{
-							Name: constants.ContainerTrainer,
+							Name: constants.Node,
 							Ports: []corev1ac.ContainerPortApplyConfiguration{{
 								ContainerPort: ptr.To[int32](constants.ContainerTrainerPort),
 							}},
@@ -567,7 +574,7 @@ func TestTorch(t *testing.T) {
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterAddr),
-									Value: ptr.To("test-job-trainer-node-0-0.test-job"),
+									Value: ptr.To("test-job-node-0-0.test-job"),
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterPort),
@@ -600,8 +607,8 @@ func TestTorch(t *testing.T) {
 						).
 						Obj(),
 				),
-				runtime.WithPodSet(constants.JobTrainerNode, 1, corev1.PodSpec{}, corev1ac.PodSpec().
-					WithContainers(corev1ac.Container().WithName(constants.ContainerTrainer)),
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 1, corev1.PodSpec{}, corev1ac.PodSpec().
+					WithContainers(corev1ac.Container().WithName(constants.Node)),
 				),
 			),
 			wantInfo: &runtime.Info{
@@ -614,11 +621,12 @@ func TestTorch(t *testing.T) {
 				},
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{{
-						Name:              constants.JobTrainerNode,
+						Name:              constants.Node,
+						Ancestor:          ptr.To(constants.AncestorTrainer),
 						Count:             ptr.To[int32](1),
 						SinglePodRequests: make(corev1.ResourceList),
 						Containers: []runtime.Container{{
-							Name: constants.ContainerTrainer,
+							Name: constants.Node,
 							Ports: []corev1ac.ContainerPortApplyConfiguration{{
 								ContainerPort: ptr.To[int32](constants.ContainerTrainerPort),
 							}},
@@ -641,7 +649,7 @@ func TestTorch(t *testing.T) {
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterAddr),
-									Value: ptr.To("test-job-trainer-node-0-0.test-job"),
+									Value: ptr.To("test-job-node-0-0.test-job"),
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterPort),
@@ -674,8 +682,8 @@ func TestTorch(t *testing.T) {
 						).
 						Obj(),
 				),
-				runtime.WithPodSet(constants.JobTrainerNode, 1, corev1.PodSpec{}, corev1ac.PodSpec().
-					WithContainers(corev1ac.Container().WithName(constants.ContainerTrainer)),
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 1, corev1.PodSpec{}, corev1ac.PodSpec().
+					WithContainers(corev1ac.Container().WithName(constants.Node)),
 				),
 			),
 			wantInfo: &runtime.Info{
@@ -688,11 +696,12 @@ func TestTorch(t *testing.T) {
 				},
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{{
-						Name:              constants.JobTrainerNode,
+						Name:              constants.Node,
+						Ancestor:          ptr.To(constants.AncestorTrainer),
 						Count:             ptr.To[int32](1),
 						SinglePodRequests: make(corev1.ResourceList),
 						Containers: []runtime.Container{{
-							Name: constants.ContainerTrainer,
+							Name: constants.Node,
 							Ports: []corev1ac.ContainerPortApplyConfiguration{{
 								ContainerPort: ptr.To[int32](constants.ContainerTrainerPort),
 							}},
@@ -715,7 +724,7 @@ func TestTorch(t *testing.T) {
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterAddr),
-									Value: ptr.To("test-job-trainer-node-0-0.test-job"),
+									Value: ptr.To("test-job-node-0-0.test-job"),
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterPort),
@@ -748,8 +757,8 @@ func TestTorch(t *testing.T) {
 						).
 						Obj(),
 				),
-				runtime.WithPodSet(constants.JobTrainerNode, 1, corev1.PodSpec{}, corev1ac.PodSpec().
-					WithContainers(corev1ac.Container().WithName(constants.ContainerTrainer)),
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 1, corev1.PodSpec{}, corev1ac.PodSpec().
+					WithContainers(corev1ac.Container().WithName(constants.Node)),
 				),
 			),
 			wantInfo: &runtime.Info{
@@ -762,11 +771,12 @@ func TestTorch(t *testing.T) {
 				},
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{{
-						Name:              constants.JobTrainerNode,
+						Name:              constants.Node,
+						Ancestor:          ptr.To(constants.AncestorTrainer),
 						Count:             ptr.To[int32](1),
 						SinglePodRequests: make(corev1.ResourceList),
 						Containers: []runtime.Container{{
-							Name: constants.ContainerTrainer,
+							Name: constants.Node,
 							Ports: []corev1ac.ContainerPortApplyConfiguration{{
 								ContainerPort: ptr.To[int32](constants.ContainerTrainerPort),
 							}},
@@ -789,7 +799,7 @@ func TestTorch(t *testing.T) {
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterAddr),
-									Value: ptr.To("test-job-trainer-node-0-0.test-job"),
+									Value: ptr.To("test-job-node-0-0.test-job"),
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterPort),
@@ -822,8 +832,8 @@ func TestTorch(t *testing.T) {
 						).
 						Obj(),
 				),
-				runtime.WithPodSet(constants.JobTrainerNode, 1, corev1.PodSpec{}, corev1ac.PodSpec().
-					WithContainers(corev1ac.Container().WithName(constants.ContainerTrainer)),
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 1, corev1.PodSpec{}, corev1ac.PodSpec().
+					WithContainers(corev1ac.Container().WithName(constants.Node)),
 				),
 			),
 			wantInfo: &runtime.Info{
@@ -836,11 +846,12 @@ func TestTorch(t *testing.T) {
 				},
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{{
-						Name:              constants.JobTrainerNode,
+						Name:              constants.Node,
+						Ancestor:          ptr.To(constants.AncestorTrainer),
 						Count:             ptr.To[int32](1),
 						SinglePodRequests: make(corev1.ResourceList),
 						Containers: []runtime.Container{{
-							Name: constants.ContainerTrainer,
+							Name: constants.Node,
 							Ports: []corev1ac.ContainerPortApplyConfiguration{{
 								ContainerPort: ptr.To[int32](constants.ContainerTrainerPort),
 							}},
@@ -863,7 +874,7 @@ func TestTorch(t *testing.T) {
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterAddr),
-									Value: ptr.To("cpu-job-trainer-node-0-0.cpu-job"),
+									Value: ptr.To("cpu-job-node-0-0.cpu-job"),
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterPort),
@@ -897,8 +908,8 @@ func TestTorch(t *testing.T) {
 						).
 						Obj(),
 				),
-				runtime.WithPodSet(constants.JobTrainerNode, 1, corev1.PodSpec{}, corev1ac.PodSpec().
-					WithContainers(corev1ac.Container().WithName(constants.ContainerTrainer)),
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 1, corev1.PodSpec{}, corev1ac.PodSpec().
+					WithContainers(corev1ac.Container().WithName(constants.Node)),
 				),
 			),
 			wantInfo: &runtime.Info{
@@ -911,11 +922,12 @@ func TestTorch(t *testing.T) {
 				},
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{{
-						Name:              constants.JobTrainerNode,
+						Name:              constants.Node,
+						Ancestor:          ptr.To(constants.AncestorTrainer),
 						Count:             ptr.To[int32](1),
 						SinglePodRequests: make(corev1.ResourceList),
 						Containers: []runtime.Container{{
-							Name: constants.ContainerTrainer,
+							Name: constants.Node,
 							Ports: []corev1ac.ContainerPortApplyConfiguration{{
 								ContainerPort: ptr.To[int32](constants.ContainerTrainerPort),
 							}},
@@ -938,7 +950,7 @@ func TestTorch(t *testing.T) {
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterAddr),
-									Value: ptr.To("cpu-gpu-job-trainer-node-0-0.cpu-gpu-job"),
+									Value: ptr.To("cpu-gpu-job-node-0-0.cpu-gpu-job"),
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterPort),
@@ -971,8 +983,8 @@ func TestTorch(t *testing.T) {
 						).
 						Obj(),
 				),
-				runtime.WithPodSet(constants.JobTrainerNode, 1, corev1.PodSpec{}, corev1ac.PodSpec().
-					WithContainers(corev1ac.Container().WithName(constants.ContainerTrainer)),
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 1, corev1.PodSpec{}, corev1ac.PodSpec().
+					WithContainers(corev1ac.Container().WithName(constants.Node)),
 				),
 			),
 			wantInfo: &runtime.Info{
@@ -985,11 +997,12 @@ func TestTorch(t *testing.T) {
 				},
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{{
-						Name:              constants.JobTrainerNode,
+						Name:              constants.Node,
+						Ancestor:          ptr.To(constants.AncestorTrainer),
 						Count:             ptr.To[int32](1),
 						SinglePodRequests: make(corev1.ResourceList),
 						Containers: []runtime.Container{{
-							Name: constants.ContainerTrainer,
+							Name: constants.Node,
 							Ports: []corev1ac.ContainerPortApplyConfiguration{{
 								ContainerPort: ptr.To[int32](constants.ContainerTrainerPort),
 							}},
@@ -1012,7 +1025,7 @@ func TestTorch(t *testing.T) {
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterAddr),
-									Value: ptr.To("cpu-frac-job-trainer-node-0-0.cpu-frac-job"),
+									Value: ptr.To("cpu-frac-job-node-0-0.cpu-frac-job"),
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterPort),
@@ -1052,8 +1065,8 @@ func TestTorch(t *testing.T) {
 					"app": "pytorch-training",
 					"env": "production",
 				}),
-				runtime.WithPodSet(constants.JobTrainerNode, 2, corev1.PodSpec{}, corev1ac.PodSpec().
-					WithContainers(corev1ac.Container().WithName(constants.ContainerTrainer)),
+				runtime.WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 2, corev1.PodSpec{}, corev1ac.PodSpec().
+					WithContainers(corev1ac.Container().WithName(constants.Node)),
 				),
 			),
 			wantInfo: &runtime.Info{
@@ -1069,11 +1082,12 @@ func TestTorch(t *testing.T) {
 				},
 				TemplateSpec: runtime.TemplateSpec{
 					PodSets: []runtime.PodSet{{
-						Name:              constants.JobTrainerNode,
+						Name:              constants.Node,
+						Ancestor:          ptr.To(constants.AncestorTrainer),
 						Count:             ptr.To[int32](4),
 						SinglePodRequests: make(corev1.ResourceList),
 						Containers: []runtime.Container{{
-							Name: constants.ContainerTrainer,
+							Name: constants.Node,
 							Ports: []corev1ac.ContainerPortApplyConfiguration{{
 								ContainerPort: ptr.To[int32](constants.ContainerTrainerPort),
 							}},
@@ -1096,7 +1110,7 @@ func TestTorch(t *testing.T) {
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterAddr),
-									Value: ptr.To("gpu-job-trainer-node-0-0.gpu-job"),
+									Value: ptr.To("gpu-job-node-0-0.gpu-job"),
 								},
 								{
 									Name:  ptr.To(constants.TorchEnvMasterPort),
diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go
index 5277e7d452..9ec15730f2 100644
--- a/pkg/runtime/runtime.go
+++ b/pkg/runtime/runtime.go
@@ -64,7 +64,10 @@ type TemplateSpec struct {
 type PodSet struct {
 	// PodSet name is the name to identify PodSpec.
 	// This typically has the name stored in each PodSpec.
-	Name           string
+	Name string
+	// Ancestor is built by `trainer.kubeflow.org/trainjob-ancestor-step` label value
+	// in Runtime CRDs.
+	Ancestor       *string
 	Count          *int32
 	InitContainers []Container
 	Containers     []Container
@@ -133,11 +136,12 @@ func WithTemplateSpecObjApply(objApply any) InfoOption {
 // WithPodSet construct Info.TemplateSpec.PodSet from PodSpec.
 // The third argument, 'typedPodSpec' is used only to calculate requested resources.
 func WithPodSet(
-	psName string, count int32, typedPodSpec corev1.PodSpec, podSpecApply *corev1ac.PodSpecApplyConfiguration,
+	psName string, ancestor *string, count int32, typedPodSpec corev1.PodSpec, podSpecApply *corev1ac.PodSpecApplyConfiguration,
 ) InfoOption {
 	return func(o *InfoOptions) {
 		ps := PodSet{
 			Name:              psName,
+			Ancestor:          ancestor,
 			Count:             ptr.To(max(count, 1)),
 			Volumes:           podSpecApply.Volumes,
 			SinglePodRequests: resourcehelpers.PodRequests(&corev1.Pod{Spec: typedPodSpec}, resourcehelpers.PodResourcesOptions{}),
@@ -203,25 +207,28 @@ func TemplateSpecApply[A any](info *Info) (*A, bool) {
 	return spec, ok
 }
 
-// FindContainerByPodSetContainerName finds runtime.Container from Info.TemplateSpec.PodSet by PodSet and Container name.
-func (i *Info) FindContainerByPodSetContainerName(psName, containerName string) *Container {
-	ps := i.FindPodSetByName(psName)
+// FindContainerByPodSetAncestorContainerName finds runtime.Container from Info.TemplateSpec.PodSet by PodSet Ancestor and Container name.
+func (i *Info) FindContainerByPodSetAncestorContainerName(psAncestor, containerName string) *Container {
+	ps := i.FindPodSetByAncestor(psAncestor)
 	if ps == nil {
 		return nil
 	}
-	for containerIdx, container := range ps.Containers {
-		if container.Name == containerName {
-			return &ps.Containers[containerIdx]
-		}
+	if idx := slices.IndexFunc(ps.Containers, func(c Container) bool { return c.Name == containerName }); idx != -1 {
+		return &ps.Containers[idx]
+	}
+	return nil
+}
+
+func (i *Info) FindPodSetByAncestor(ancestor string) *PodSet {
+	if idx := slices.IndexFunc(i.TemplateSpec.PodSets, func(ps PodSet) bool { return ptr.Equal(ps.Ancestor, &ancestor) }); idx != -1 {
+		return &i.TemplateSpec.PodSets[idx]
 	}
 	return nil
 }
 
 func (i *Info) FindPodSetByName(psName string) *PodSet {
-	for psIdx, ps := range i.TemplateSpec.PodSets {
-		if ps.Name == psName {
-			return &i.TemplateSpec.PodSets[psIdx]
-		}
+	if idx := slices.IndexFunc(i.TemplateSpec.PodSets, func(ps PodSet) bool { return ps.Name == psName }); idx != -1 {
+		return &i.TemplateSpec.PodSets[idx]
 	}
 	return nil
 }
diff --git a/pkg/runtime/runtime_test.go b/pkg/runtime/runtime_test.go
index 32f16464b7..63911d0d25 100644
--- a/pkg/runtime/runtime_test.go
+++ b/pkg/runtime/runtime_test.go
@@ -45,7 +45,7 @@ func TestNewInfo(t *testing.T) {
 				WithAnnotations(map[string]string{
 					"annotationKey": "annotationValue",
 				}),
-				WithPodSet(constants.DatasetInitializer, 1, corev1.PodSpec{
+				WithPodSet(constants.DatasetInitializer, ptr.To(constants.DatasetInitializer), 1, corev1.PodSpec{
 					Containers: []corev1.Container{{
 						Name: constants.DatasetInitializer,
 						Resources: corev1.ResourceRequirements{
@@ -82,7 +82,7 @@ func TestNewInfo(t *testing.T) {
 								})),
 					),
 				),
-				WithPodSet(constants.ModelInitializer, 1, corev1.PodSpec{
+				WithPodSet(constants.ModelInitializer, ptr.To(constants.ModelInitializer), 1, corev1.PodSpec{
 					Containers: []corev1.Container{{
 						Name: constants.ModelInitializer,
 						Resources: corev1.ResourceRequirements{
@@ -119,9 +119,9 @@ func TestNewInfo(t *testing.T) {
 								})),
 					),
 				),
-				WithPodSet(constants.JobTrainerNode, 10, corev1.PodSpec{
+				WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 10, corev1.PodSpec{
 					Containers: []corev1.Container{{
-						Name: constants.ContainerTrainer,
+						Name: constants.Node,
 						Resources: corev1.ResourceRequirements{
 							Requests: corev1.ResourceList{
 								corev1.ResourceCPU: resource.MustParse("15"),
@@ -140,7 +140,7 @@ func TestNewInfo(t *testing.T) {
 				}, corev1ac.PodSpec().
 					WithContainers(
 						corev1ac.Container().
-							WithName(constants.ContainerTrainer).
+							WithName(constants.Node).
 							WithResources(corev1ac.ResourceRequirements().
 								WithRequests(corev1.ResourceList{
 									corev1.ResourceCPU: resource.MustParse("25"),
@@ -228,7 +228,7 @@ func TestNewInfo(t *testing.T) {
 									),
 								),
 							jobsetv1alpha2ac.ReplicatedJob().
-								WithName(constants.JobTrainerNode).
+								WithName(constants.Node).
 								WithTemplate(batchv1ac.JobTemplateSpec().
 									WithLabels(nil).
 									WithSpec(batchv1ac.JobSpec().
@@ -237,7 +237,7 @@ func TestNewInfo(t *testing.T) {
 											WithSpec(corev1ac.PodSpec().
 												WithContainers(
 													corev1ac.Container().
-														WithName(constants.ContainerTrainer).
+														WithName(constants.Node).
 														WithVolumeMounts(
 															corev1ac.VolumeMount().
 																WithName(jobsetplgconsts.VolumeNameInitializer).
@@ -272,8 +272,9 @@ func TestNewInfo(t *testing.T) {
 				TemplateSpec: TemplateSpec{
 					PodSets: []PodSet{
 						{
-							Name:  constants.DatasetInitializer,
-							Count: ptr.To[int32](1),
+							Name:     constants.DatasetInitializer,
+							Ancestor: ptr.To(constants.DatasetInitializer),
+							Count:    ptr.To[int32](1),
 							InitContainers: []Container{{
 								Name: "setup-initializer",
 							}},
@@ -285,8 +286,9 @@ func TestNewInfo(t *testing.T) {
 							},
 						},
 						{
-							Name:  constants.ModelInitializer,
-							Count: ptr.To[int32](1),
+							Name:     constants.ModelInitializer,
+							Ancestor: ptr.To(constants.ModelInitializer),
+							Count:    ptr.To[int32](1),
 							InitContainers: []Container{{
 								Name: "setup-initializer",
 							}},
@@ -298,10 +300,11 @@ func TestNewInfo(t *testing.T) {
 							},
 						},
 						{
-							Name:  constants.JobTrainerNode,
-							Count: ptr.To[int32](10),
+							Name:     constants.Node,
+							Ancestor: ptr.To(constants.AncestorTrainer),
+							Count:    ptr.To[int32](10),
 							Containers: []Container{{
-								Name: constants.ContainerTrainer,
+								Name: constants.Node,
 								Env: []corev1ac.EnvVarApplyConfiguration{{
 									Name:  ptr.To("TEST"),
 									Value: ptr.To("TEST"),
@@ -384,7 +387,7 @@ func TestNewInfo(t *testing.T) {
 									),
 								),
 							jobsetv1alpha2ac.ReplicatedJob().
-								WithName(constants.JobTrainerNode).
+								WithName(constants.Node).
 								WithTemplate(batchv1ac.JobTemplateSpec().
 									WithLabels(nil).
 									WithSpec(batchv1ac.JobSpec().
@@ -393,7 +396,7 @@ func TestNewInfo(t *testing.T) {
 											WithSpec(corev1ac.PodSpec().
 												WithContainers(
 													corev1ac.Container().
-														WithName(constants.ContainerTrainer).
+														WithName(constants.Node).
 														WithVolumeMounts(
 															corev1ac.VolumeMount().
 																WithName(jobsetplgconsts.VolumeNameInitializer).
@@ -439,10 +442,10 @@ func TestNewInfo(t *testing.T) {
 	}
 }
 
-func TestFindContainerByPodSetContainerName(t *testing.T) {
+func TestFindContainerByPodSetAncestorContainerName(t *testing.T) {
 	cases := map[string]struct {
 		info          *Info
-		psName        string
+		psAncestor    string
 		containerName string
 		wantContainer *Container
 	}{
@@ -451,7 +454,7 @@ func TestFindContainerByPodSetContainerName(t *testing.T) {
 				TemplateSpec: TemplateSpec{
 					PodSets: []PodSet{
 						{
-							Name: "alpha",
+							Ancestor: ptr.To("alpha"),
 							Containers: []Container{
 								{
 									Name: "one",
@@ -462,13 +465,13 @@ func TestFindContainerByPodSetContainerName(t *testing.T) {
 							},
 						},
 						{
-							Name:       "beta",
+							Ancestor:   ptr.To("beta"),
 							Containers: []Container{{Name: "one"}},
 						},
 					},
 				},
 			},
-			psName:        "alpha",
+			psAncestor:    "alpha",
 			containerName: "one",
 			wantContainer: &Container{
 				Name: "one",
@@ -478,12 +481,12 @@ func TestFindContainerByPodSetContainerName(t *testing.T) {
 			info: &Info{
 				TemplateSpec: TemplateSpec{
 					PodSets: []PodSet{{
-						Name:       "alpha",
+						Ancestor:   ptr.To("alpha"),
 						Containers: []Container{{Name: "one"}},
 					}},
 				},
 			},
-			psName:        "alpha",
+			psAncestor:    "alpha",
 			containerName: "two",
 			wantContainer: nil,
 		},
@@ -491,19 +494,19 @@ func TestFindContainerByPodSetContainerName(t *testing.T) {
 			info: &Info{
 				TemplateSpec: TemplateSpec{
 					PodSets: []PodSet{{
-						Name:       "alpha",
+						Ancestor:   ptr.To("alpha"),
 						Containers: []Container{{Name: "one"}},
 					}},
 				},
 			},
-			psName:        "beta",
+			psAncestor:    "beta",
 			containerName: "one",
 			wantContainer: nil,
 		},
 	}
 	for name, tc := range cases {
 		t.Run(name, func(t *testing.T) {
-			got := tc.info.FindContainerByPodSetContainerName(tc.psName, tc.containerName)
+			got := tc.info.FindContainerByPodSetAncestorContainerName(tc.psAncestor, tc.containerName)
 			if diff := cmp.Diff(tc.wantContainer, got); len(diff) != 0 {
 				t.Errorf("Unexpected Container (-want,+got):\n%s", diff)
 			}
@@ -511,11 +514,11 @@ func TestFindContainerByPodSetContainerName(t *testing.T) {
 	}
 }
 
-func TestFindContainerByName(t *testing.T) {
+func TestFindPodSetByName(t *testing.T) {
 	cases := map[string]struct {
-		info   *Info
-		psName string
-		want   *PodSet
+		info       *Info
+		psName     string
+		wantPodSet *PodSet
 	}{
 		"PodSet exists": {
 			info: &Info{
@@ -531,10 +534,58 @@ func TestFindContainerByName(t *testing.T) {
 				},
 			},
 			psName: "alpha",
-			want: &PodSet{
+			wantPodSet: &PodSet{
 				Name: "alpha",
 			},
 		},
+		"PodSet does not exist": {
+			info: &Info{
+				TemplateSpec: TemplateSpec{
+					PodSets: []PodSet{
+						{
+							Name: "alpha",
+						},
+					},
+				},
+			},
+			psName:     "beta",
+			wantPodSet: nil,
+		},
+	}
+	for name, tc := range cases {
+		t.Run(name, func(t *testing.T) {
+			got := tc.info.FindPodSetByName(tc.psName)
+			if diff := cmp.Diff(tc.wantPodSet, got); len(diff) != 0 {
+				t.Errorf("Unexpected PodSet (-want,+got):\n%s", diff)
+			}
+		})
+	}
+}
+
+func TestFindContainerByName(t *testing.T) {
+	cases := map[string]struct {
+		info       *Info
+		psAncestor string
+		want       *PodSet
+	}{
+		"PodSet exists": {
+			info: &Info{
+				TemplateSpec: TemplateSpec{
+					PodSets: []PodSet{
+						{
+							Ancestor: ptr.To("alpha"),
+						},
+						{
+							Ancestor: ptr.To("beta"),
+						},
+					},
+				},
+			},
+			psAncestor: "alpha",
+			want: &PodSet{
+				Ancestor: ptr.To("alpha"),
+			},
+		},
 		"PodSet does not exist": {
 			info: &Info{
 				TemplateSpec: TemplateSpec{
@@ -543,13 +594,13 @@ func TestFindContainerByName(t *testing.T) {
 					}},
 				},
 			},
-			psName: "beta",
-			want:   nil,
+			psAncestor: "beta",
+			want:       nil,
 		},
 	}
 	for name, tc := range cases {
 		t.Run(name, func(t *testing.T) {
-			got := tc.info.FindPodSetByName(tc.psName)
+			got := tc.info.FindPodSetByAncestor(tc.psAncestor)
 			if diff := cmp.Diff(tc.want, got); len(diff) != 0 {
 				t.Errorf("Unexpected PodSet (-want,+got):\n%s", diff)
 			}
diff --git a/pkg/util/testing/wrapper.go b/pkg/util/testing/wrapper.go
index 6f7fad27ef..67be0d3fab 100644
--- a/pkg/util/testing/wrapper.go
+++ b/pkg/util/testing/wrapper.go
@@ -54,13 +54,13 @@ func MakeJobSetWrapper(namespace, name string) *JobSetWrapper {
 					{
 						Name: constants.DatasetInitializer,
 						Template: batchv1.JobTemplateSpec{
+							ObjectMeta: metav1.ObjectMeta{
+								Labels: map[string]string{
+									constants.LabelTrainJobAncestor: constants.DatasetInitializer,
+								},
+							},
 							Spec: batchv1.JobSpec{
 								Template: corev1.PodTemplateSpec{
-									ObjectMeta: metav1.ObjectMeta{
-										Labels: map[string]string{
-											constants.LabelTrainJobAncestor: constants.DatasetInitializer,
-										},
-									},
 									Spec: corev1.PodSpec{
 										Containers: []corev1.Container{
 											{
@@ -87,13 +87,13 @@ func MakeJobSetWrapper(namespace, name string) *JobSetWrapper {
 					{
 						Name: constants.ModelInitializer,
 						Template: batchv1.JobTemplateSpec{
+							ObjectMeta: metav1.ObjectMeta{
+								Labels: map[string]string{
+									constants.LabelTrainJobAncestor: constants.ModelInitializer,
+								},
+							},
 							Spec: batchv1.JobSpec{
 								Template: corev1.PodTemplateSpec{
-									ObjectMeta: metav1.ObjectMeta{
-										Labels: map[string]string{
-											constants.LabelTrainJobAncestor: constants.ModelInitializer,
-										},
-									},
 									Spec: corev1.PodSpec{
 										Containers: []corev1.Container{
 											{
@@ -118,14 +118,19 @@ func MakeJobSetWrapper(namespace, name string) *JobSetWrapper {
 						},
 					},
 					{
-						Name: constants.JobTrainerNode,
+						Name: constants.Node,
 						Template: batchv1.JobTemplateSpec{
+							ObjectMeta: metav1.ObjectMeta{
+								Labels: map[string]string{
+									constants.LabelTrainJobAncestor: constants.AncestorTrainer,
+								},
+							},
 							Spec: batchv1.JobSpec{
 								Template: corev1.PodTemplateSpec{
 									Spec: corev1.PodSpec{
 										Containers: []corev1.Container{
 											{
-												Name: constants.ContainerTrainer,
+												Name: constants.Node,
 												VolumeMounts: []corev1.VolumeMount{
 													{
 														Name:      jobsetplgconsts.VolumeNameInitializer,
@@ -168,7 +173,7 @@ func (j *JobSetWrapper) Replicas(replicas int32, rJobNames ...string) *JobSetWra
 
 func (j *JobSetWrapper) NumNodes(numNodes int32) *JobSetWrapper {
 	for i, rJob := range j.Spec.ReplicatedJobs {
-		if rJob.Name == constants.JobTrainerNode {
+		if rJob.Name == constants.Node {
 			j.Spec.ReplicatedJobs[i].Template.Spec.Parallelism = &numNodes
 			j.Spec.ReplicatedJobs[i].Template.Spec.Completions = &numNodes
 		}
@@ -196,7 +201,7 @@ func (j *JobSetWrapper) Completions(c int32, rJobNames ...string) *JobSetWrapper
 
 func (j *JobSetWrapper) LauncherReplica() *JobSetWrapper {
 	for i, rJob := range j.Spec.ReplicatedJobs {
-		if rJob.Name == constants.JobTrainerNode {
+		if rJob.Name == constants.Node {
 			j.Spec.ReplicatedJobs = append(j.Spec.ReplicatedJobs, jobsetv1alpha2.ReplicatedJob{})
 			copy(j.Spec.ReplicatedJobs[i+1:], j.Spec.ReplicatedJobs[i:])
 			j.Spec.ReplicatedJobs[i] = jobsetv1alpha2.ReplicatedJob{
@@ -238,9 +243,9 @@ func (j *JobSetWrapper) Container(rJobName, containerName, image string, command
 
 func (j *JobSetWrapper) ContainerTrainerPorts(ports []corev1.ContainerPort) *JobSetWrapper {
 	for i, rJob := range j.Spec.ReplicatedJobs {
-		if rJob.Name == constants.JobTrainerNode {
+		if rJob.Name == constants.Node {
 			for k, container := range rJob.Template.Spec.Template.Spec.Containers {
-				if container.Name == constants.ContainerTrainer {
+				if container.Name == constants.Node {
 					j.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.Containers[k].Ports = ports
 				}
 			}
@@ -594,13 +599,13 @@ func MakeTrainingRuntimeWrapper(namespace, name string) *TrainingRuntimeWrapper
 							{
 								Name: constants.DatasetInitializer,
 								Template: batchv1.JobTemplateSpec{
+									ObjectMeta: metav1.ObjectMeta{
+										Labels: map[string]string{
+											constants.LabelTrainJobAncestor: constants.DatasetInitializer,
+										},
+									},
 									Spec: batchv1.JobSpec{
 										Template: corev1.PodTemplateSpec{
-											ObjectMeta: metav1.ObjectMeta{
-												Labels: map[string]string{
-													constants.LabelTrainJobAncestor: constants.DatasetInitializer,
-												},
-											},
 											Spec: corev1.PodSpec{
 												Containers: []corev1.Container{
 													{
@@ -627,13 +632,13 @@ func MakeTrainingRuntimeWrapper(namespace, name string) *TrainingRuntimeWrapper
 							{
 								Name: constants.ModelInitializer,
 								Template: batchv1.JobTemplateSpec{
+									ObjectMeta: metav1.ObjectMeta{
+										Labels: map[string]string{
+											constants.LabelTrainJobAncestor: constants.ModelInitializer,
+										},
+									},
 									Spec: batchv1.JobSpec{
 										Template: corev1.PodTemplateSpec{
-											ObjectMeta: metav1.ObjectMeta{
-												Labels: map[string]string{
-													constants.LabelTrainJobAncestor: constants.ModelInitializer,
-												},
-											},
 											Spec: corev1.PodSpec{
 												Containers: []corev1.Container{
 													{
@@ -658,14 +663,19 @@ func MakeTrainingRuntimeWrapper(namespace, name string) *TrainingRuntimeWrapper
 								},
 							},
 							{
-								Name: constants.JobTrainerNode,
+								Name: constants.Node,
 								Template: batchv1.JobTemplateSpec{
+									ObjectMeta: metav1.ObjectMeta{
+										Labels: map[string]string{
+											constants.LabelTrainJobAncestor: constants.AncestorTrainer,
+										},
+									},
 									Spec: batchv1.JobSpec{
 										Template: corev1.PodTemplateSpec{
 											Spec: corev1.PodSpec{
 												Containers: []corev1.Container{
 													{
-														Name: constants.ContainerTrainer,
+														Name: constants.Node,
 														VolumeMounts: []corev1.VolumeMount{
 															{
 																Name:      jobsetplgconsts.VolumeNameInitializer,
@@ -799,14 +809,14 @@ func MakeClusterTrainingRuntimeWrapper(name string) *ClusterTrainingRuntimeWrapp
 								},
 							},
 							{
-								Name: constants.JobTrainerNode,
+								Name: constants.Node,
 								Template: batchv1.JobTemplateSpec{
 									Spec: batchv1.JobSpec{
 										Template: corev1.PodTemplateSpec{
 											Spec: corev1.PodSpec{
 												Containers: []corev1.Container{
 													{
-														Name: constants.ContainerTrainer,
+														Name: constants.Node,
 														VolumeMounts: []corev1.VolumeMount{
 															{
 																Name:      jobsetplgconsts.VolumeNameInitializer,
@@ -871,7 +881,7 @@ func (s *TrainingRuntimeSpecWrapper) WithMLPolicy(mlPolicy *trainer.MLPolicy) *T
 
 func (s *TrainingRuntimeSpecWrapper) LauncherReplica() *TrainingRuntimeSpecWrapper {
 	for i, rJob := range s.Template.Spec.ReplicatedJobs {
-		if rJob.Name == constants.JobTrainerNode {
+		if rJob.Name == constants.Node {
 			s.Template.Spec.ReplicatedJobs = append(s.Template.Spec.ReplicatedJobs, jobsetv1alpha2.ReplicatedJob{})
 			copy(s.Template.Spec.ReplicatedJobs[i+1:], s.Template.Spec.ReplicatedJobs[i:])
 			s.Template.Spec.ReplicatedJobs[i] = jobsetv1alpha2.ReplicatedJob{
diff --git a/pkg/webhooks/trainingruntime_webhook_test.go b/pkg/webhooks/trainingruntime_webhook_test.go
index 8d8d4ff899..944528abc7 100644
--- a/pkg/webhooks/trainingruntime_webhook_test.go
+++ b/pkg/webhooks/trainingruntime_webhook_test.go
@@ -35,12 +35,12 @@ func TestValidateReplicatedJobs(t *testing.T) {
 	}{
 		"valid replicatedJobs": {
 			rJobs: testingutil.MakeJobSetWrapper("ns", "valid").
-				Replicas(1, constants.JobTrainerNode, constants.DatasetInitializer, constants.ModelInitializer).
+				Replicas(1, constants.Node, constants.DatasetInitializer, constants.ModelInitializer).
 				Obj().Spec.ReplicatedJobs,
 		},
 		"invalid replicas": {
 			rJobs: testingutil.MakeJobSetWrapper("ns", "valid").
-				Replicas(2, constants.JobTrainerNode, constants.DatasetInitializer, constants.ModelInitializer).
+				Replicas(2, constants.Node, constants.DatasetInitializer, constants.ModelInitializer).
 				Obj().Spec.ReplicatedJobs,
 			wantError: field.ErrorList{
 				field.Invalid(field.NewPath("spec").Child("template").Child("spec").Child("replicatedJobs").Index(0).Child("replicas"),
diff --git a/test/integration/controller/trainjob_controller_test.go b/test/integration/controller/trainjob_controller_test.go
index b767cafd2b..d381e80b60 100644
--- a/test/integration/controller/trainjob_controller_test.go
+++ b/test/integration/controller/trainjob_controller_test.go
@@ -119,7 +119,7 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 						PodGroupPolicyCoscheduling(&trainer.CoschedulingPodGroupPolicySource{ScheduleTimeoutSeconds: ptr.To[int32](100)}).
 						Container(constants.ModelInitializer, constants.ModelInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 						Container(constants.DatasetInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
-						Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+						Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 						Obj()).
 				Obj()
 		})
@@ -144,7 +144,7 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 							Label("testingKey", "testingVal").
 							Annotation("testingKey", "testingVal").
 							PodLabel(schedulerpluginsv1alpha1.PodGroupLabel, trainJobKey.Name).
-							Replicas(1, constants.JobTrainerNode, constants.DatasetInitializer, constants.ModelInitializer).
+							Replicas(1, constants.Node, constants.DatasetInitializer, constants.ModelInitializer).
 							Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
 							Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
 							NumNodes(100).
@@ -166,7 +166,7 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 									},
 								}...,
 							).
-							Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
+							Container(constants.Node, constants.Node, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
 							Obj(),
 						util.IgnoreObjectMetadata))
 					pg := &schedulerpluginsv1alpha1.PodGroup{}
@@ -219,7 +219,7 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 							Label("testingKey", "testingVal").
 							Annotation("testingKey", "testingVal").
 							PodLabel(schedulerpluginsv1alpha1.PodGroupLabel, trainJobKey.Name).
-							Replicas(1, constants.JobTrainerNode, constants.DatasetInitializer, constants.ModelInitializer).
+							Replicas(1, constants.Node, constants.DatasetInitializer, constants.ModelInitializer).
 							Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
 							Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
 							NumNodes(100).
@@ -241,7 +241,7 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 									},
 								}...,
 							).
-							Container(constants.JobTrainerNode, constants.ContainerTrainer, updatedImageName, []string{"trainjob"}, []string{"trainjob"}, resRequests).
+							Container(constants.Node, constants.Node, updatedImageName, []string{"trainjob"}, []string{"trainjob"}, resRequests).
 							Obj(),
 						util.IgnoreObjectMetadata))
 					pg := &schedulerpluginsv1alpha1.PodGroup{}
@@ -283,7 +283,7 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 					jobSet := &jobsetv1alpha2.JobSet{}
 					g.Expect(k8sClient.Get(ctx, trainJobKey, jobSet)).Should(gomega.Succeed())
 					for _, rJob := range jobSet.Spec.ReplicatedJobs {
-						if rJob.Name == constants.JobTrainerNode {
+						if rJob.Name == constants.Node {
 							g.Expect(rJob.Template.Spec.Template.Spec.Containers[0].Image).Should(gomega.Equal(updatedImageName))
 						}
 					}
@@ -304,7 +304,7 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 					g.Expect(jobSet.Spec.Suspend).ShouldNot(gomega.BeNil())
 					g.Expect(*jobSet.Spec.Suspend).Should(gomega.BeTrue())
 					for _, rJob := range jobSet.Spec.ReplicatedJobs {
-						if rJob.Name == constants.JobTrainerNode {
+						if rJob.Name == constants.Node {
 							g.Expect(rJob.Template.Spec.Template.Spec.Containers[0].Image).Should(gomega.Equal(originImageName))
 						}
 					}
@@ -337,7 +337,7 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 									).
 									Obj(),
 							).
-							Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+							Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
 							Obj()).
 					Obj()
 				gomega.Expect(k8sClient.Create(ctx, trainingRuntime)).Should(gomega.Succeed())
@@ -354,13 +354,13 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 						testingutil.MakeJobSetWrapper(ns.Name, trainJobKey.Name).
 							ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), trainJobKey.Name, string(trainJob.UID)).
 							Suspend(false).
-							Replicas(1, constants.JobTrainerNode, constants.DatasetInitializer, constants.ModelInitializer).
+							Replicas(1, constants.Node, constants.DatasetInitializer, constants.ModelInitializer).
 							Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
 							Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
 							NumNodes(100).
-							Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
+							Container(constants.Node, constants.Node, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
 							ContainerTrainerPorts([]corev1.ContainerPort{{ContainerPort: constants.ContainerTrainerPort, Protocol: "TCP"}}).
-							Env(constants.JobTrainerNode, constants.ContainerTrainer,
+							Env(constants.Node, constants.Node,
 								[]corev1.EnvVar{
 									{
 										Name:  "TRAIN_JOB",
@@ -384,7 +384,7 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 									},
 									{
 										Name:  constants.TorchEnvMasterAddr,
-										Value: fmt.Sprintf("alpha-%s-0-0.alpha", constants.JobTrainerNode),
+										Value: fmt.Sprintf("alpha-%s-0-0.alpha", constants.Node),
 									},
 									{
 										Name:  constants.TorchEnvMasterPort,
@@ -616,7 +616,7 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 									).
 									Obj(),
 							).
-							Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
+							Container(constants.Node, constants.Node, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
 							Obj()).
 					Obj()
 				gomega.Expect(k8sClient.Create(ctx, trainingRuntime)).Should(gomega.Succeed())
@@ -634,11 +634,11 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 							ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), trainJobKey.Name, string(trainJob.UID)).
 							Suspend(false).
 							LauncherReplica().
-							Replicas(1, constants.JobTrainerNode, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
+							Replicas(1, constants.Node, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
 							Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
 							Completions(1, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
 							NumNodes(2).
-							Container(constants.JobTrainerNode, constants.ContainerTrainer, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
+							Container(constants.Node, constants.Node, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
 							Volumes(constants.JobLauncher,
 								corev1.Volume{
 									Name: constants.MPISSHAuthVolumeName,
@@ -678,7 +678,7 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 									},
 								},
 							).
-							Volumes(constants.JobTrainerNode,
+							Volumes(constants.Node,
 								corev1.Volume{
 									Name: constants.MPISSHAuthVolumeName,
 									VolumeSource: corev1.VolumeSource{
@@ -706,6 +706,9 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 								corev1.VolumeMount{Name: constants.MPISSHAuthVolumeName, MountPath: "/root/.ssh"},
 								corev1.VolumeMount{Name: constants.MPIHostfileVolumeName, MountPath: constants.MPIHostfileDir},
 							).
+							VolumeMounts(constants.Node, constants.Node,
+								corev1.VolumeMount{Name: constants.MPISSHAuthVolumeName, MountPath: "/root/.ssh"},
+							).
 							Env(constants.JobLauncher, constants.ContainerLauncher,
 								corev1.EnvVar{
 									Name:  constants.OpenMPIEnvHostFileLocation,
@@ -724,7 +727,7 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 									Value: constants.OpenMPIEnvDefaultValueRSHArgs,
 								},
 							).
-							Env(constants.JobTrainerNode, constants.ContainerTrainer,
+							Env(constants.Node, constants.Node,
 								corev1.EnvVar{
 									Name:  "TRAIN_JOB",
 									Value: "value",
@@ -741,8 +744,8 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 					g.Expect(cm).Should(gomega.BeComparableTo(
 						testingutil.MakeConfigMapWrapper(cmKey.Name, cmKey.Namespace).
 							WithData(map[string]string{
-								constants.MPIHostfileName: `alpha-trainer-node-0-0.alpha slots=8
-alpha-trainer-node-0-1.alpha slots=8
+								constants.MPIHostfileName: `alpha-node-0-0.alpha slots=8
+alpha-node-0-1.alpha slots=8
 `,
 							}).
 							ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), trainJobKey.Name, string(trainJob.UID)).
diff --git a/test/integration/webhooks/trainingruntime_webhook_test.go b/test/integration/webhooks/trainingruntime_webhook_test.go
index bf03bab0c9..c714bc304f 100644
--- a/test/integration/webhooks/trainingruntime_webhook_test.go
+++ b/test/integration/webhooks/trainingruntime_webhook_test.go
@@ -189,7 +189,7 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
 						},
 					}
 					runtime.Spec.Template.Spec = testingutil.MakeJobSetWrapper(ns.Name, "runtime").
-						Replicas(1, constants.JobTrainerNode, constants.DatasetInitializer, constants.ModelInitializer).
+						Replicas(1, constants.Node, constants.DatasetInitializer, constants.ModelInitializer).
 						Obj().
 						Spec
 					return runtime
@@ -209,7 +209,7 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
 							).
 							JobSetSpec(
 								testingutil.MakeJobSetWrapper(ns.Name, "jobset").
-									Replicas(1, constants.JobTrainerNode, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
+									Replicas(1, constants.Node, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
 									Obj().
 									Spec,
 							).
@@ -231,7 +231,7 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
 							).
 							JobSetSpec(
 								testingutil.MakeJobSetWrapper(ns.Name, "jobset").
-									Replicas(1, constants.JobTrainerNode, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
+									Replicas(1, constants.Node, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
 									Obj().
 									Spec,
 							).
@@ -254,7 +254,7 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
 							).
 							JobSetSpec(
 								testingutil.MakeJobSetWrapper(ns.Name, "jobset").
-									Replicas(1, constants.JobTrainerNode, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
+									Replicas(1, constants.Node, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
 									Obj().
 									Spec,
 							).
@@ -276,7 +276,7 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
 							).
 							JobSetSpec(
 								testingutil.MakeJobSetWrapper(ns.Name, "jobset").
-									Replicas(1, constants.JobTrainerNode, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
+									Replicas(1, constants.Node, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
 									Obj().
 									Spec,
 							).
@@ -299,7 +299,7 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
 							).
 							JobSetSpec(
 								testingutil.MakeJobSetWrapper(ns.Name, "jobset").
-									Replicas(1, constants.JobTrainerNode, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
+									Replicas(1, constants.Node, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
 									Obj().
 									Spec,
 							).
@@ -321,7 +321,7 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
 							).
 							JobSetSpec(
 								testingutil.MakeJobSetWrapper(ns.Name, "jobset").
-									Replicas(1, constants.JobTrainerNode, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
+									Replicas(1, constants.Node, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
 									Obj().
 									Spec,
 							).
@@ -344,7 +344,7 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
 							).
 							JobSetSpec(
 								testingutil.MakeJobSetWrapper(ns.Name, "jobset").
-									Replicas(1, constants.JobTrainerNode, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
+									Replicas(1, constants.Node, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
 									Obj().
 									Spec,
 							).
@@ -366,7 +366,7 @@ var _ = ginkgo.Describe("TrainingRuntime marker validations and defaulting", gin
 							).
 							JobSetSpec(
 								testingutil.MakeJobSetWrapper(ns.Name, "jobset").
-									Replicas(1, constants.JobTrainerNode, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
+									Replicas(1, constants.Node, constants.DatasetInitializer, constants.ModelInitializer, constants.JobLauncher).
 									Obj().
 									Spec,
 							).


Reviews:
user: login: andreyvelich
body: Thank you for these updates @tenzen-y!

state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T13:53:45Z
user: login: andreyvelich
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T15:23:26Z
user: login: andreyvelich
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T15:29:20Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T15:48:22Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T15:49:15Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T15:50:15Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T15:51:11Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T15:52:10Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T15:54:32Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T15:55:11Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T15:55:43Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T15:57:03Z
user: login: andreyvelich
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T15:59:16Z
user: login: andreyvelich
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T16:00:41Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T16:02:34Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-20T05:41:03Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-20T07:30:21Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-20T07:30:28Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-20T07:30:37Z
user: login: andreyvelich
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-20T13:33:28Z
user: login: andreyvelich
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-20T13:36:44Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-20T13:42:31Z
user: login: tenzen-y
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-20T13:42:35Z
user: login: andreyvelich
body: Thanks for this great contribution @tenzen-y!
/lgtm
/approve

state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-20T13:46:10Z

Review comments:
diff_hunk: @@ -12,36 +12,56 @@ spec:
       numProcPerNode: 1
       mpiImplementation: OpenMPI
       sshAuthMountPath: /home/mpiuser/.ssh
+      runLauncherAsNode: true
   template:
     spec:
       network:
-        publishNotReadyAddresses: false
+        publishNotReadyAddresses: true
path: manifests/base/runtimes/pretraining/mpi_distributed.yaml
user: login: andreyvelich
body: Why did we change this setting ?
created_at: 2025-03-19T11:59:26Z
author_association: MEMBER
diff_hunk: @@ -12,36 +12,56 @@ spec:
       numProcPerNode: 1
       mpiImplementation: OpenMPI
       sshAuthMountPath: /home/mpiuser/.ssh
+      runLauncherAsNode: true
path: manifests/base/runtimes/pretraining/mpi_distributed.yaml
user: login: andreyvelich
body: Could we also change our defaulter webhook to set `runLauncherAsNode: true` by default ?
created_at: 2025-03-19T12:00:03Z
author_association: MEMBER
diff_hunk: @@ -12,36 +12,56 @@ spec:
       numProcPerNode: 1
       mpiImplementation: OpenMPI
       sshAuthMountPath: /home/mpiuser/.ssh
+      runLauncherAsNode: true
   template:
     spec:
       network:
-        publishNotReadyAddresses: false
+        publishNotReadyAddresses: true
       successPolicy:
         operator: All
         targetReplicatedJobs:
           - launcher
       replicatedJobs:
         - name: launcher
           template:
+            metadata:
+              labels:
+                trainer.kubeflow.org/trainjob-ancestor-step: trainer
             spec:
               template:
                 spec:
+                  initContainers:
path: manifests/base/runtimes/pretraining/mpi_distributed.yaml
user: login: andreyvelich
body: Actually, I was able to run distributed MLX when launcher runs as a node without OpenSSH server on the Launcher.
Have you tried it for PI example ?

created_at: 2025-03-19T12:01:27Z
author_association: MEMBER
diff_hunk: @@ -12,36 +12,56 @@ spec:
       numProcPerNode: 1
       mpiImplementation: OpenMPI
       sshAuthMountPath: /home/mpiuser/.ssh
+      runLauncherAsNode: true
   template:
     spec:
       network:
-        publishNotReadyAddresses: false
+        publishNotReadyAddresses: true
       successPolicy:
         operator: All
         targetReplicatedJobs:
           - launcher
       replicatedJobs:
         - name: launcher
           template:
+            metadata:
+              labels:
+                trainer.kubeflow.org/trainjob-ancestor-step: trainer
             spec:
               template:
                 spec:
+                  initContainers:
+                  - name: node
+                    image: mpioperator/mpi-pi:openmpi
+                    securityContext:
+                      runAsUser: 1000
+                    command:
+                      - /usr/sbin/sshd
+                    args:
+                      - -De
+                      - -f
+                      - /home/mpiuser/.sshd_config
+                    restartPolicy: Always
path: manifests/base/runtimes/pretraining/mpi_distributed.yaml
user: login: andreyvelich
body: I like the sidecar approach here!
I prefer it over: https://github.com/kubeflow/mpi-operator/pull/612#issuecomment-1950254976
cc @kuizhiqing  
created_at: 2025-03-19T12:25:38Z
author_association: MEMBER
diff_hunk: @@ -115,17 +116,23 @@ func (b *Builder) Launcher() *Builder {
 // Trainer updates JobSet values for the trainer Job.
 func (b *Builder) Trainer(info *runtime.Info, trainJob *trainer.TrainJob) *Builder {
 	for i, rJob := range b.Spec.ReplicatedJobs {
-		if *rJob.Name == constants.JobTrainerNode {
+		jobMetadata := rJob.Template.ObjectMetaApplyConfiguration
+		if jobMetadata == nil || jobMetadata.Labels == nil {
+			continue
+		}
+		if ancestor, ok := jobMetadata.Labels[constants.LabelTrainJobAncestor]; ok && ancestor == constants.AncestorTrainer {
 			// TODO: Support multiple replicas ('.template.spec.replicatedJobs[*].replicas') for replicated Jobs.
 			// REF: https://github.com/kubeflow/trainer/issues/2318
 			b.Spec.ReplicatedJobs[i].Replicas = ptr.To[int32](1)
 			// Update the Parallelism and Completions values for the Trainer Job.
-			b.Spec.ReplicatedJobs[i].Template.Spec.Parallelism = info.TemplateSpec.PodSets[i].Count
-			b.Spec.ReplicatedJobs[i].Template.Spec.Completions = info.TemplateSpec.PodSets[i].Count
+			b.Spec.ReplicatedJobs[i].Template.Spec.Parallelism = info.FindPodSetByAncestor(constants.AncestorTrainer).Count
+			b.Spec.ReplicatedJobs[i].Template.Spec.Completions = info.FindPodSetByAncestor(constants.AncestorTrainer).Count
 
 			// Update values for the Trainer container.
 			for j, container := range rJob.Template.Spec.Template.Spec.Containers {
-				if *container.Name == constants.ContainerTrainer {
+				// TODO: This `Launcher` should be handled more correctly.
+				// Currently, we assume the 'launcher' container name is used for MPI Launcher.
path: pkg/runtime/framework/plugins/jobset/builder.go
user: login: andreyvelich
body: Can we always use `node` container name for the launcher ReplicatedJob ?

created_at: 2025-03-19T13:42:45Z
author_association: MEMBER
diff_hunk: @@ -103,7 +103,14 @@ func (m *MPI) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob) er
 
 	// TrainJob contains the actual information for the Trainer.
 	if trainJob.Spec.Trainer != nil && trainJob.Spec.Trainer.NumNodes != nil {
-		info.FindPodSetByName(constants.JobTrainerNode).Count = trainJob.Spec.Trainer.NumNodes
+		if node := info.FindPodSetByName(constants.JobNode); node != nil && node.Count != nil {
+			if ptr.Deref(info.RuntimePolicy.MLPolicySource.MPI.RunLauncherAsNode, false) {
+				// When runLauncherAsNode is enabled, 1 nodes should be allocated to launcher.
+				*node.Count = max(*trainJob.Spec.Trainer.NumNodes-1, 1)
path: pkg/runtime/framework/plugins/mpi/mpi.go
user: login: andreyvelich
body: We might need to add validation in the future that if NumNodes >= 2 and RunLauncherAsNode = true, the `node` and `launcher` ReplicatedJob must be presented in the runtime.
@tenzen-y We can do it in the followup PRs, please can you add TODO statement ?

created_at: 2025-03-19T13:46:55Z
author_association: MEMBER
diff_hunk: @@ -154,50 +161,64 @@ func (m *MPI) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob) er
 				}...,
 			)
 		}
-		for cIdx, container := range ps.Containers {
-			if container.Name != constants.JobLauncher && container.Name != constants.JobTrainerNode {
-				continue
-			}
+		containers, err := enforceMPIParametersToContainers(info, ps.Name, psIdx, info.TemplateSpec.PodSets[psIdx].Containers...)
+		if err != nil {
+			return err
+		}
+		info.TemplateSpec.PodSets[psIdx].Containers = containers
+		initContainers, err := enforceMPIParametersToContainers(info, ps.Name, psIdx, info.TemplateSpec.PodSets[psIdx].InitContainers...)
path: pkg/runtime/framework/plugins/mpi/mpi.go
user: login: andreyvelich
body: Why do we need to enforce MPI params to the InitContainers ?

created_at: 2025-03-19T13:49:10Z
author_association: MEMBER
diff_hunk: @@ -50,13 +50,15 @@ func (p *PlainML) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob
 
 	// TrainJob contains the actual information for the number of nodes.
 	if trainJob.Spec.Trainer != nil && trainJob.Spec.Trainer.NumNodes != nil {
-		info.FindPodSetByName(constants.JobTrainerNode).Count = trainJob.Spec.Trainer.NumNodes
+		if trainerPS := info.FindPodSetByAncestor(constants.AncestorTrainer); trainerPS != nil && trainerPS.Count != nil {
path: pkg/runtime/framework/plugins/plainml/plainml.go
user: login: andreyvelich
body: Do we need to check trainerPS.Count here ?

created_at: 2025-03-19T13:50:15Z
author_association: MEMBER
diff_hunk: @@ -12,36 +12,56 @@ spec:
       numProcPerNode: 1
       mpiImplementation: OpenMPI
       sshAuthMountPath: /home/mpiuser/.ssh
+      runLauncherAsNode: true
   template:
     spec:
       network:
-        publishNotReadyAddresses: false
+        publishNotReadyAddresses: true
       successPolicy:
         operator: All
         targetReplicatedJobs:
           - launcher
       replicatedJobs:
         - name: launcher
           template:
+            metadata:
+              labels:
+                trainer.kubeflow.org/trainjob-ancestor-step: trainer
path: manifests/base/runtimes/pretraining/mpi_distributed.yaml
user: login: andreyvelich
body: Please can we propagate this label all to way to ReplicatedJob Pods ?
```
rJob.spec.template.metadata.labels
```

We will use this label to detect appropriate pods in the SDK.

created_at: 2025-03-19T15:23:23Z
author_association: MEMBER
diff_hunk: @@ -12,36 +12,56 @@ spec:
       numProcPerNode: 1
       mpiImplementation: OpenMPI
       sshAuthMountPath: /home/mpiuser/.ssh
+      runLauncherAsNode: true
   template:
     spec:
       network:
-        publishNotReadyAddresses: false
+        publishNotReadyAddresses: true
       successPolicy:
         operator: All
         targetReplicatedJobs:
           - launcher
       replicatedJobs:
         - name: launcher
           template:
+            metadata:
+              labels:
+                trainer.kubeflow.org/trainjob-ancestor-step: trainer
path: manifests/base/runtimes/pretraining/mpi_distributed.yaml
user: login: andreyvelich
body: Actually, sorry that is not required for now: https://github.com/kubeflow/trainer/blob/796645aee7c0be0aa0f4709eed0cac1fd1e38df5/sdk/kubeflow/trainer/api/trainer_client.py#L491-L499

We still need to rely on ReplicatedJob name to list Pods for now, since we need to get Launcher + Node pods for the MPI.

created_at: 2025-03-19T15:29:20Z
author_association: MEMBER
diff_hunk: @@ -12,36 +12,56 @@ spec:
       numProcPerNode: 1
       mpiImplementation: OpenMPI
       sshAuthMountPath: /home/mpiuser/.ssh
+      runLauncherAsNode: true
path: manifests/base/runtimes/pretraining/mpi_distributed.yaml
user: login: tenzen-y
body: That is pretty aggressive.
It would be better to leave the selection on users.

created_at: 2025-03-19T15:48:22Z
author_association: MEMBER
diff_hunk: @@ -12,36 +12,56 @@ spec:
       numProcPerNode: 1
       mpiImplementation: OpenMPI
       sshAuthMountPath: /home/mpiuser/.ssh
+      runLauncherAsNode: true
path: manifests/base/runtimes/pretraining/mpi_distributed.yaml
user: login: tenzen-y
body: Enforcing something does not fit on templating mechanism.

created_at: 2025-03-19T15:49:15Z
author_association: MEMBER
diff_hunk: @@ -12,36 +12,56 @@ spec:
       numProcPerNode: 1
       mpiImplementation: OpenMPI
       sshAuthMountPath: /home/mpiuser/.ssh
+      runLauncherAsNode: true
   template:
     spec:
       network:
-        publishNotReadyAddresses: false
+        publishNotReadyAddresses: true
path: manifests/base/runtimes/pretraining/mpi_distributed.yaml
user: login: tenzen-y
body: Because there is a chicken and egg problem between ready for Launcher or starting mpirun within Launcher.

created_at: 2025-03-19T15:50:14Z
author_association: MEMBER
diff_hunk: @@ -12,36 +12,56 @@ spec:
       numProcPerNode: 1
       mpiImplementation: OpenMPI
       sshAuthMountPath: /home/mpiuser/.ssh
+      runLauncherAsNode: true
   template:
     spec:
       network:
-        publishNotReadyAddresses: false
+        publishNotReadyAddresses: true
       successPolicy:
         operator: All
         targetReplicatedJobs:
           - launcher
       replicatedJobs:
         - name: launcher
           template:
+            metadata:
+              labels:
+                trainer.kubeflow.org/trainjob-ancestor-step: trainer
             spec:
               template:
                 spec:
+                  initContainers:
path: manifests/base/runtimes/pretraining/mpi_distributed.yaml
user: login: tenzen-y
body: As I checked the solution without ssh server for launcher, it works well.
So, I will remove these sidecar.

created_at: 2025-03-19T15:51:11Z
author_association: MEMBER
diff_hunk: @@ -115,17 +116,23 @@ func (b *Builder) Launcher() *Builder {
 // Trainer updates JobSet values for the trainer Job.
 func (b *Builder) Trainer(info *runtime.Info, trainJob *trainer.TrainJob) *Builder {
 	for i, rJob := range b.Spec.ReplicatedJobs {
-		if *rJob.Name == constants.JobTrainerNode {
+		jobMetadata := rJob.Template.ObjectMetaApplyConfiguration
+		if jobMetadata == nil || jobMetadata.Labels == nil {
+			continue
+		}
+		if ancestor, ok := jobMetadata.Labels[constants.LabelTrainJobAncestor]; ok && ancestor == constants.AncestorTrainer {
 			// TODO: Support multiple replicas ('.template.spec.replicatedJobs[*].replicas') for replicated Jobs.
 			// REF: https://github.com/kubeflow/trainer/issues/2318
 			b.Spec.ReplicatedJobs[i].Replicas = ptr.To[int32](1)
 			// Update the Parallelism and Completions values for the Trainer Job.
-			b.Spec.ReplicatedJobs[i].Template.Spec.Parallelism = info.TemplateSpec.PodSets[i].Count
-			b.Spec.ReplicatedJobs[i].Template.Spec.Completions = info.TemplateSpec.PodSets[i].Count
+			b.Spec.ReplicatedJobs[i].Template.Spec.Parallelism = info.FindPodSetByAncestor(constants.AncestorTrainer).Count
+			b.Spec.ReplicatedJobs[i].Template.Spec.Completions = info.FindPodSetByAncestor(constants.AncestorTrainer).Count
 
 			// Update values for the Trainer container.
 			for j, container := range rJob.Template.Spec.Template.Spec.Containers {
-				if *container.Name == constants.ContainerTrainer {
+				// TODO: This `Launcher` should be handled more correctly.
+				// Currently, we assume the 'launcher' container name is used for MPI Launcher.
path: pkg/runtime/framework/plugins/jobset/builder.go
user: login: tenzen-y
body: Using same container or initContainer name within the same Pod is not allowed in Kubernetes, which means Pod creation failed
created_at: 2025-03-19T15:52:10Z
author_association: MEMBER
diff_hunk: @@ -50,13 +50,15 @@ func (p *PlainML) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob
 
 	// TrainJob contains the actual information for the number of nodes.
 	if trainJob.Spec.Trainer != nil && trainJob.Spec.Trainer.NumNodes != nil {
-		info.FindPodSetByName(constants.JobTrainerNode).Count = trainJob.Spec.Trainer.NumNodes
+		if trainerPS := info.FindPodSetByAncestor(constants.AncestorTrainer); trainerPS != nil && trainerPS.Count != nil {
path: pkg/runtime/framework/plugins/plainml/plainml.go
user: login: tenzen-y
body: Avoiding null pointer.
created_at: 2025-03-19T15:54:31Z
author_association: MEMBER
diff_hunk: @@ -154,50 +161,64 @@ func (m *MPI) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob) er
 				}...,
 			)
 		}
-		for cIdx, container := range ps.Containers {
-			if container.Name != constants.JobLauncher && container.Name != constants.JobTrainerNode {
-				continue
-			}
+		containers, err := enforceMPIParametersToContainers(info, ps.Name, psIdx, info.TemplateSpec.PodSets[psIdx].Containers...)
+		if err != nil {
+			return err
+		}
+		info.TemplateSpec.PodSets[psIdx].Containers = containers
+		initContainers, err := enforceMPIParametersToContainers(info, ps.Name, psIdx, info.TemplateSpec.PodSets[psIdx].InitContainers...)
path: pkg/runtime/framework/plugins/mpi/mpi.go
user: login: tenzen-y
body: For sidecar container. The sidecar container requires ssh configurations.
created_at: 2025-03-19T15:55:11Z
author_association: MEMBER
diff_hunk: @@ -154,50 +161,64 @@ func (m *MPI) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob) er
 				}...,
 			)
 		}
-		for cIdx, container := range ps.Containers {
-			if container.Name != constants.JobLauncher && container.Name != constants.JobTrainerNode {
-				continue
-			}
+		containers, err := enforceMPIParametersToContainers(info, ps.Name, psIdx, info.TemplateSpec.PodSets[psIdx].Containers...)
+		if err != nil {
+			return err
+		}
+		info.TemplateSpec.PodSets[psIdx].Containers = containers
+		initContainers, err := enforceMPIParametersToContainers(info, ps.Name, psIdx, info.TemplateSpec.PodSets[psIdx].InitContainers...)
path: pkg/runtime/framework/plugins/mpi/mpi.go
user: login: tenzen-y
body: However, I decided to remove all initContainer impls from this PR since we do not need to launch ssh server for Launcher.

created_at: 2025-03-19T15:55:42Z
author_association: MEMBER
diff_hunk: @@ -103,7 +103,14 @@ func (m *MPI) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob) er
 
 	// TrainJob contains the actual information for the Trainer.
 	if trainJob.Spec.Trainer != nil && trainJob.Spec.Trainer.NumNodes != nil {
-		info.FindPodSetByName(constants.JobTrainerNode).Count = trainJob.Spec.Trainer.NumNodes
+		if node := info.FindPodSetByName(constants.JobNode); node != nil && node.Count != nil {
+			if ptr.Deref(info.RuntimePolicy.MLPolicySource.MPI.RunLauncherAsNode, false) {
+				// When runLauncherAsNode is enabled, 1 nodes should be allocated to launcher.
+				*node.Count = max(*trainJob.Spec.Trainer.NumNodes-1, 1)
path: pkg/runtime/framework/plugins/mpi/mpi.go
user: login: tenzen-y
body: Sure
created_at: 2025-03-19T15:57:03Z
author_association: MEMBER
diff_hunk: @@ -12,36 +12,56 @@ spec:
       numProcPerNode: 1
       mpiImplementation: OpenMPI
       sshAuthMountPath: /home/mpiuser/.ssh
+      runLauncherAsNode: true
path: manifests/base/runtimes/pretraining/mpi_distributed.yaml
user: login: andreyvelich
body: Sounds good.

created_at: 2025-03-19T15:59:16Z
author_association: MEMBER
diff_hunk: @@ -115,17 +116,23 @@ func (b *Builder) Launcher() *Builder {
 // Trainer updates JobSet values for the trainer Job.
 func (b *Builder) Trainer(info *runtime.Info, trainJob *trainer.TrainJob) *Builder {
 	for i, rJob := range b.Spec.ReplicatedJobs {
-		if *rJob.Name == constants.JobTrainerNode {
+		jobMetadata := rJob.Template.ObjectMetaApplyConfiguration
+		if jobMetadata == nil || jobMetadata.Labels == nil {
+			continue
+		}
+		if ancestor, ok := jobMetadata.Labels[constants.LabelTrainJobAncestor]; ok && ancestor == constants.AncestorTrainer {
 			// TODO: Support multiple replicas ('.template.spec.replicatedJobs[*].replicas') for replicated Jobs.
 			// REF: https://github.com/kubeflow/trainer/issues/2318
 			b.Spec.ReplicatedJobs[i].Replicas = ptr.To[int32](1)
 			// Update the Parallelism and Completions values for the Trainer Job.
-			b.Spec.ReplicatedJobs[i].Template.Spec.Parallelism = info.TemplateSpec.PodSets[i].Count
-			b.Spec.ReplicatedJobs[i].Template.Spec.Completions = info.TemplateSpec.PodSets[i].Count
+			b.Spec.ReplicatedJobs[i].Template.Spec.Parallelism = info.FindPodSetByAncestor(constants.AncestorTrainer).Count
+			b.Spec.ReplicatedJobs[i].Template.Spec.Completions = info.FindPodSetByAncestor(constants.AncestorTrainer).Count
 
 			// Update values for the Trainer container.
 			for j, container := range rJob.Template.Spec.Template.Spec.Containers {
-				if *container.Name == constants.ContainerTrainer {
+				// TODO: This `Launcher` should be handled more correctly.
+				// Currently, we assume the 'launcher' container name is used for MPI Launcher.
path: pkg/runtime/framework/plugins/jobset/builder.go
user: login: andreyvelich
body: But since we are going to remove the `InitContainer` accroding to this: https://github.com/kubeflow/trainer/pull/2548#discussion_r2003167928, we can just keep single Container: `node`  in the ReplicatedJob.name = `launcher`.
Isn't ?
created_at: 2025-03-19T16:00:41Z
author_association: MEMBER
diff_hunk: @@ -115,17 +116,23 @@ func (b *Builder) Launcher() *Builder {
 // Trainer updates JobSet values for the trainer Job.
 func (b *Builder) Trainer(info *runtime.Info, trainJob *trainer.TrainJob) *Builder {
 	for i, rJob := range b.Spec.ReplicatedJobs {
-		if *rJob.Name == constants.JobTrainerNode {
+		jobMetadata := rJob.Template.ObjectMetaApplyConfiguration
+		if jobMetadata == nil || jobMetadata.Labels == nil {
+			continue
+		}
+		if ancestor, ok := jobMetadata.Labels[constants.LabelTrainJobAncestor]; ok && ancestor == constants.AncestorTrainer {
 			// TODO: Support multiple replicas ('.template.spec.replicatedJobs[*].replicas') for replicated Jobs.
 			// REF: https://github.com/kubeflow/trainer/issues/2318
 			b.Spec.ReplicatedJobs[i].Replicas = ptr.To[int32](1)
 			// Update the Parallelism and Completions values for the Trainer Job.
-			b.Spec.ReplicatedJobs[i].Template.Spec.Parallelism = info.TemplateSpec.PodSets[i].Count
-			b.Spec.ReplicatedJobs[i].Template.Spec.Completions = info.TemplateSpec.PodSets[i].Count
+			b.Spec.ReplicatedJobs[i].Template.Spec.Parallelism = info.FindPodSetByAncestor(constants.AncestorTrainer).Count
+			b.Spec.ReplicatedJobs[i].Template.Spec.Completions = info.FindPodSetByAncestor(constants.AncestorTrainer).Count
 
 			// Update values for the Trainer container.
 			for j, container := range rJob.Template.Spec.Template.Spec.Containers {
-				if *container.Name == constants.ContainerTrainer {
+				// TODO: This `Launcher` should be handled more correctly.
+				// Currently, we assume the 'launcher' container name is used for MPI Launcher.
path: pkg/runtime/framework/plugins/jobset/builder.go
user: login: tenzen-y
body: Yes, removal initContainer implementations allow us to remove this 
created_at: 2025-03-19T16:02:34Z
author_association: MEMBER
diff_hunk: @@ -103,7 +103,14 @@ func (m *MPI) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob) er
 
 	// TrainJob contains the actual information for the Trainer.
 	if trainJob.Spec.Trainer != nil && trainJob.Spec.Trainer.NumNodes != nil {
-		info.FindPodSetByName(constants.JobTrainerNode).Count = trainJob.Spec.Trainer.NumNodes
+		if node := info.FindPodSetByName(constants.JobNode); node != nil && node.Count != nil {
+			if ptr.Deref(info.RuntimePolicy.MLPolicySource.MPI.RunLauncherAsNode, false) {
+				// When runLauncherAsNode is enabled, 1 nodes should be allocated to launcher.
+				*node.Count = max(*trainJob.Spec.Trainer.NumNodes-1, 1)
path: pkg/runtime/framework/plugins/mpi/mpi.go
user: login: tenzen-y
body: Done: https://github.com/kubeflow/trainer/issues/2550
created_at: 2025-03-20T05:41:03Z
author_association: MEMBER
diff_hunk: @@ -12,36 +12,56 @@ spec:
       numProcPerNode: 1
       mpiImplementation: OpenMPI
       sshAuthMountPath: /home/mpiuser/.ssh
+      runLauncherAsNode: true
   template:
     spec:
       network:
-        publishNotReadyAddresses: false
+        publishNotReadyAddresses: true
       successPolicy:
         operator: All
         targetReplicatedJobs:
           - launcher
       replicatedJobs:
         - name: launcher
           template:
+            metadata:
+              labels:
+                trainer.kubeflow.org/trainjob-ancestor-step: trainer
             spec:
               template:
                 spec:
+                  initContainers:
path: manifests/base/runtimes/pretraining/mpi_distributed.yaml
user: login: tenzen-y
body: Done.
created_at: 2025-03-20T07:30:21Z
author_association: MEMBER
diff_hunk: @@ -115,17 +116,23 @@ func (b *Builder) Launcher() *Builder {
 // Trainer updates JobSet values for the trainer Job.
 func (b *Builder) Trainer(info *runtime.Info, trainJob *trainer.TrainJob) *Builder {
 	for i, rJob := range b.Spec.ReplicatedJobs {
-		if *rJob.Name == constants.JobTrainerNode {
+		jobMetadata := rJob.Template.ObjectMetaApplyConfiguration
+		if jobMetadata == nil || jobMetadata.Labels == nil {
+			continue
+		}
+		if ancestor, ok := jobMetadata.Labels[constants.LabelTrainJobAncestor]; ok && ancestor == constants.AncestorTrainer {
 			// TODO: Support multiple replicas ('.template.spec.replicatedJobs[*].replicas') for replicated Jobs.
 			// REF: https://github.com/kubeflow/trainer/issues/2318
 			b.Spec.ReplicatedJobs[i].Replicas = ptr.To[int32](1)
 			// Update the Parallelism and Completions values for the Trainer Job.
-			b.Spec.ReplicatedJobs[i].Template.Spec.Parallelism = info.TemplateSpec.PodSets[i].Count
-			b.Spec.ReplicatedJobs[i].Template.Spec.Completions = info.TemplateSpec.PodSets[i].Count
+			b.Spec.ReplicatedJobs[i].Template.Spec.Parallelism = info.FindPodSetByAncestor(constants.AncestorTrainer).Count
+			b.Spec.ReplicatedJobs[i].Template.Spec.Completions = info.FindPodSetByAncestor(constants.AncestorTrainer).Count
 
 			// Update values for the Trainer container.
 			for j, container := range rJob.Template.Spec.Template.Spec.Containers {
-				if *container.Name == constants.ContainerTrainer {
+				// TODO: This `Launcher` should be handled more correctly.
+				// Currently, we assume the 'launcher' container name is used for MPI Launcher.
path: pkg/runtime/framework/plugins/jobset/builder.go
user: login: tenzen-y
body: Done.
created_at: 2025-03-20T07:30:28Z
author_association: MEMBER
diff_hunk: @@ -154,50 +161,64 @@ func (m *MPI) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob) er
 				}...,
 			)
 		}
-		for cIdx, container := range ps.Containers {
-			if container.Name != constants.JobLauncher && container.Name != constants.JobTrainerNode {
-				continue
-			}
+		containers, err := enforceMPIParametersToContainers(info, ps.Name, psIdx, info.TemplateSpec.PodSets[psIdx].Containers...)
+		if err != nil {
+			return err
+		}
+		info.TemplateSpec.PodSets[psIdx].Containers = containers
+		initContainers, err := enforceMPIParametersToContainers(info, ps.Name, psIdx, info.TemplateSpec.PodSets[psIdx].InitContainers...)
path: pkg/runtime/framework/plugins/mpi/mpi.go
user: login: tenzen-y
body: Done.
created_at: 2025-03-20T07:30:37Z
author_association: MEMBER
diff_hunk: @@ -12,36 +12,40 @@ spec:
       numProcPerNode: 1
       mpiImplementation: OpenMPI
       sshAuthMountPath: /home/mpiuser/.ssh
+      runLauncherAsNode: true
   template:
     spec:
       network:
-        publishNotReadyAddresses: false
+        publishNotReadyAddresses: true
       successPolicy:
         operator: All
         targetReplicatedJobs:
           - launcher
       replicatedJobs:
         - name: launcher
           template:
+            metadata:
+              labels:
+                trainer.kubeflow.org/trainjob-ancestor-step: trainer
             spec:
               template:
                 spec:
                   containers:
-                    - name: launcher
-                      image: mpioperator/mpi-pi:openmpi
-                      securityContext:
-                        runAsUser: 1000
-                      command:
-                        - mpirun
-                      args:
-                        - /home/mpiuser/pi
-        - name: trainer-node
+                  - name: launcher
path: manifests/base/runtimes/pretraining/mpi_distributed.yaml
user: login: andreyvelich
body: Since we run Launcher as Node can we keep the container name: `node` ?
```suggestion
                  - name: node
```
created_at: 2025-03-20T13:33:24Z
author_association: MEMBER
diff_hunk: @@ -34,11 +34,15 @@ const (
 	// ModelMountPath is the volumeMount path for model.
 	ModelMountPath string = "/workspace/model"
 
-	// JobTrainerNode is the Job name for the trainer node.
-	JobTrainerNode string = "trainer-node"
+	// AncestorTrainer is the ancestor name for Trainer, which is mostly used for the value of
+	// 'trainer.kubeflow.org/trainjob-ancestor-step'.
+	AncestorTrainer string = "trainer"
 
-	// ContainerTrainer is the container name for the trainer.
-	ContainerTrainer string = "trainer"
+	// JobNode is the Job name for the trainer node.
+	JobNode string = "node"
+
+	// ContainerNode is the container name for the node.
+	ContainerNode string = "node"
path: pkg/constants/constants.go
user: login: andreyvelich
body: Similar to initializer constants, shell we keep single value to reduce number of constants we maintain ?
https://github.com/kubeflow/trainer/blob/ac6bc1f098642529031dcbb1e4fb0d774a52a976/pkg/constants/constants.go#L32
```suggestion
	// Node is the name of the Job and container for the trainer node
	Node string = "node"
```
created_at: 2025-03-20T13:35:22Z
author_association: MEMBER
diff_hunk: @@ -12,36 +12,40 @@ spec:
       numProcPerNode: 1
       mpiImplementation: OpenMPI
       sshAuthMountPath: /home/mpiuser/.ssh
+      runLauncherAsNode: true
   template:
     spec:
       network:
-        publishNotReadyAddresses: false
+        publishNotReadyAddresses: true
       successPolicy:
         operator: All
         targetReplicatedJobs:
           - launcher
       replicatedJobs:
         - name: launcher
           template:
+            metadata:
+              labels:
+                trainer.kubeflow.org/trainjob-ancestor-step: trainer
             spec:
               template:
                 spec:
                   containers:
-                    - name: launcher
-                      image: mpioperator/mpi-pi:openmpi
-                      securityContext:
-                        runAsUser: 1000
-                      command:
-                        - mpirun
-                      args:
-                        - /home/mpiuser/pi
-        - name: trainer-node
+                  - name: launcher
path: manifests/base/runtimes/pretraining/mpi_distributed.yaml
user: login: tenzen-y
body: That's right, good catch.
Thank you
created_at: 2025-03-20T13:42:31Z
author_association: MEMBER
diff_hunk: @@ -34,11 +34,15 @@ const (
 	// ModelMountPath is the volumeMount path for model.
 	ModelMountPath string = "/workspace/model"
 
-	// JobTrainerNode is the Job name for the trainer node.
-	JobTrainerNode string = "trainer-node"
+	// AncestorTrainer is the ancestor name for Trainer, which is mostly used for the value of
+	// 'trainer.kubeflow.org/trainjob-ancestor-step'.
+	AncestorTrainer string = "trainer"
 
-	// ContainerTrainer is the container name for the trainer.
-	ContainerTrainer string = "trainer"
+	// JobNode is the Job name for the trainer node.
+	JobNode string = "node"
+
+	// ContainerNode is the container name for the node.
+	ContainerNode string = "node"
path: pkg/constants/constants.go
user: login: tenzen-y
body: Sure.
created_at: 2025-03-20T13:42:35Z
author_association: MEMBER

PR comments:
user: login: coveralls
created_at: 2025-03-19T11:58:35Z
body: 
[![Coverage Status](https://coveralls.io/builds/72827343/badge)](https://coveralls.io/builds/72827343)

coverage: 64.313% (+1.6%) from 62.757%
when pulling **ac6bc1f098642529031dcbb1e4fb0d774a52a976 on tenzen-y:treat-ancestor-label-as-identifier**
into **5c89faa2c0f75fc23b0835286222f709cb9d30cf on kubeflow:master**.

author_association: NONE
user: login: tenzen-y
created_at: 2025-03-20T07:44:29Z
body: @andreyvelich I addressed all comments, PTAL thank you
author_association: MEMBER
user: login: google-oss-prow[bot]
created_at: 2025-03-20T13:46:17Z
body: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/kubeflow/trainer/pull/2548#pullrequestreview-2702809743" title="Approved">andreyvelich</a>*

The full list of commands accepted by this bot can be found [here](https://oss.gprow.dev/command-help?repo=kubeflow%2Ftrainer).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/kubeflow/trainer/blob/master/OWNERS)~~ [andreyvelich]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
author_association: NONE