Fix trainjob controller not setting the trainer podset count value correctly

Author: kaisoz

pkg/controller/jobs/trainjob/trainjob_controller.go

@@ -177,6 +177,15 @@ func getChildJobSet(ctx context.Context, t *TrainJob) (*jobsetapi.JobSet, error)
 	if !ok {
 		return nil, err
 	}
+
+	// Jobset replicaJob parallelism/completions are set outside of the jobset builder
+	for psIdx, ps := range info.TemplateSpec.PodSets {
+		if ps.Count != nil {
+			jobSetSpec.ReplicatedJobs[psIdx].Template.Spec.Parallelism = ps.Count
+			jobSetSpec.ReplicatedJobs[psIdx].Template.Spec.Completions = ps.Count
+		}
+	}
+
 	jobsetApply := kftrainerjobset.NewBuilder(jobsetapplyapi.JobSet(t.Name, t.Namespace).
 		WithSpec(jobSetSpec)).Initializer(trainJob).Trainer(info, trainJob).PodLabels(info.Scheduler.PodLabels).Build()
 

pkg/controller/jobs/trainjob/trainjob_controller_test.go

@@ -79,6 +79,9 @@ func TestRunWithPodsetsInfo(t *testing.T) {
 	testJobset := testingjobset.MakeJobSet("", "").ReplicatedJobs(
 		testingjobset.ReplicatedJobRequirements{
 			Name: "node",
+			Labels: map[string]string{
+				"trainer.kubeflow.org/trainjob-ancestor-step": "trainer",
+			},
 		}).Obj()
 	testCtr := testingtrainjob.MakeClusterTrainingRuntime("test", testJobset.Spec)
 
@@ -361,6 +364,15 @@ func TestReconciler(t *testing.T) {
 			Replicas:    1,
 			Parallelism: 1,
 			Completions: 1,
+			Labels: map[string]string{
+				"trainer.kubeflow.org/trainjob-ancestor-step": "trainer",
+			},
+		},
+		testingjobset.ReplicatedJobRequirements{
+			Name:        "foo",
+			Replicas:    1,
+			Parallelism: 1,
+			Completions: 1,
 		}).Obj()
 	testCtr := testingtrainjob.MakeClusterTrainingRuntime("test", testJobset.Spec)
 
@@ -386,6 +398,35 @@ func TestReconciler(t *testing.T) {
 							SubGroupIndexLabel(ptr.To(jobsetapi.JobIndexKey)).
 							SubGroupCount(ptr.To[int32](1)).
 							Obj(),
+						*utiltestingapi.MakePodSet("foo", 1).
+							PodIndexLabel(ptr.To("batch.kubernetes.io/job-completion-index")).
+							SubGroupIndexLabel(ptr.To(jobsetapi.JobIndexKey)).
+							SubGroupCount(ptr.To[int32](1)).
+							Obj(),
+					).
+					Obj(),
+			},
+		},
+		"podset count for the trainer job is set to .Spec.Trainer.NumNodes": {
+			reconcilerOptions: []jobframework.Option{
+				jobframework.WithManageJobsWithoutQueueName(true),
+				jobframework.WithManagedJobsNamespaceSelector(labels.Everything()),
+			},
+			trainJob:     testTrainJob.Clone().TrainerNumNodes(2).Obj(),
+			wantTrainJob: testTrainJob.Clone().TrainerNumNodes(2).Obj(),
+			wantWorkloads: []kueue.Workload{
+				*utiltestingapi.MakeWorkload(testTrainJob.Name, testTrainJob.Namespace).
+					PodSets(
+						*utiltestingapi.MakePodSet("node", 2).
+							PodIndexLabel(ptr.To("batch.kubernetes.io/job-completion-index")).
+							SubGroupIndexLabel(ptr.To(jobsetapi.JobIndexKey)).
+							SubGroupCount(ptr.To[int32](1)).
+							Obj(),
+						*utiltestingapi.MakePodSet("foo", 1).
+							PodIndexLabel(ptr.To("batch.kubernetes.io/job-completion-index")).
+							SubGroupIndexLabel(ptr.To(jobsetapi.JobIndexKey)).
+							SubGroupCount(ptr.To[int32](1)).
+							Obj(),
 					).
 					Obj(),
 			},

pkg/util/testingjobs/trainjob/wrappers.go

@@ -79,6 +79,12 @@ func (t *TrainJobWrapper) TrainerImage(image string, cmd, args []string) *TrainJ
 	return t
 }
 
+// TrainerNumNodes sets a the number of nodes that will be used in the Trainer job
+func (t *TrainJobWrapper) TrainerNumNodes(numNodes int32) *TrainJobWrapper {
+	t.Spec.Trainer.NumNodes = ptr.To(numNodes)
+	return t
+}
+
 // Label sets a Trainjob annotation key and value
 func (t *TrainJobWrapper) Annotation(key, value string) *TrainJobWrapper {
 	if t.Annotations == nil {

test/integration/singlecluster/controller/jobs/trainjob/trainjob_controller_test.go

@@ -92,6 +92,13 @@ var _ = ginkgo.Describe("Trainjob controller", ginkgo.Ordered, ginkgo.ContinueOn
 				testingjobset.ReplicatedJobRequirements{
 					Name:     "node",
 					Replicas: 1,
+					Labels: map[string]string{
+						"trainer.kubeflow.org/trainjob-ancestor-step": "trainer",
+					},
+				},
+				testingjobset.ReplicatedJobRequirements{
+					Name:     "foo",
+					Replicas: 1,
 				}).
 				Obj()
 			testCtr = testingtrainjob.MakeClusterTrainingRuntime("test", testJobSet.Spec)
@@ -126,6 +133,7 @@ var _ = ginkgo.Describe("Trainjob controller", ginkgo.Ordered, ginkgo.ContinueOn
 					Name:     "test",
 					Kind:     ptr.To("ClusterTrainingRuntime"),
 				}).
+					TrainerNumNodes(2).
 					Suspend(false).
 					Queue("local-queue").
 					Obj()
@@ -136,11 +144,14 @@ var _ = ginkgo.Describe("Trainjob controller", ginkgo.Ordered, ginkgo.ContinueOn
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
 			})
 
-			ginkgo.By("checking the workload is created", func() {
+			ginkgo.By("checking the workload is created with the correct values", func() {
 				wlLookupKey = types.NamespacedName{Name: workloadtrainjob.GetWorkloadNameForTrainJob(createdTrainJob.Name, createdTrainJob.UID), Namespace: ns.Name}
 				gomega.Eventually(func(g gomega.Gomega) {
 					g.Expect(k8sClient.Get(ctx, wlLookupKey, createdWorkload)).Should(gomega.Succeed())
 					g.Expect(createdWorkload.Spec.QueueName).Should(gomega.Equal(kueue.LocalQueueName("local-queue")))
+					g.Expect(createdWorkload.Spec.PodSets).Should(gomega.HaveLen(2))
+					g.Expect(createdWorkload.Spec.PodSets[0].Count).Should(gomega.Equal(int32(2)))
+					g.Expect(createdWorkload.Spec.PodSets[1].Count).Should(gomega.Equal(int32(1)))
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
 			})
 
@@ -160,6 +171,12 @@ var _ = ginkgo.Describe("Trainjob controller", ginkgo.Ordered, ginkgo.ContinueOn
 							corev1.ResourceCPU: kueue.ResourceFlavorReference(onDemandFlavor.Name),
 						},
 					},
+					kueue.PodSetAssignment{
+						Name: createdWorkload.Spec.PodSets[1].Name,
+						Flavors: map[corev1.ResourceName]kueue.ResourceFlavorReference{
+							corev1.ResourceCPU: kueue.ResourceFlavorReference(onDemandFlavor.Name),
+						},
+					},
 				).Obj()
 				util.SetQuotaReservation(ctx, k8sClient, wlLookupKey, admission)
 				util.SyncAdmittedConditionForWorkloads(ctx, k8sClient, createdWorkload)
@@ -234,6 +251,12 @@ var _ = ginkgo.Describe("Trainjob controller", ginkgo.Ordered, ginkgo.ContinueOn
 							corev1.ResourceCPU: kueue.ResourceFlavorReference(onDemandFlavor.Name),
 						},
 					},
+					kueue.PodSetAssignment{
+						Name: createdWorkload.Spec.PodSets[1].Name,
+						Flavors: map[corev1.ResourceName]kueue.ResourceFlavorReference{
+							corev1.ResourceCPU: kueue.ResourceFlavorReference(onDemandFlavor.Name),
+						},
+					},
 				).Obj()
 				util.SetQuotaReservation(ctx, k8sClient, wlLookupKey, admission)
 				util.SyncAdmittedConditionForWorkloads(ctx, k8sClient, createdWorkload)