feat(api): Replace PodTemplateOverrides with RuntimePatches API (#3309)

* feat(api): Replace PodTemplateOverrides with RuntimePatches API Signed-off-by: Andrey Velichkevich <[email protected]> * Update goal with RuntimePatches API Signed-off-by: Andrey Velichkevich <[email protected]> * Add diagram for RuntimePatches Signed-off-by: Andrey Velichkevich <[email protected]> * feat(api): Replace PodTemplateOverrides with RuntimePatches API Signed-off-by: Antonin Stefanutti <[email protected]> * Fix kube lint errors Signed-off-by: Antonin Stefanutti <[email protected]> * Update Python models Signed-off-by: Antonin Stefanutti <[email protected]> * Enable maxlength Kube-API-Linter rule Signed-off-by: Antonin Stefanutti <[email protected]> * Restore manger field IsDomainPrefixedPath validation Signed-off-by: Antonin Stefanutti <[email protected]> * Fix integration tests Signed-off-by: Antonin Stefanutti <[email protected]> * Apply StrategicMergePatch to batchv1.JobTemplateSpec Signed-off-by: Antonin Stefanutti <[email protected]> --------- Signed-off-by: Andrey Velichkevich <[email protected]> Signed-off-by: Antonin Stefanutti <[email protected]> Co-authored-by: Andrey Velichkevich <[email protected]>
2026-09-24 00:31:34 +00:00 · 2026-03-12 16:29:44 +00:00
parent f105c84c58
commit 37c8a07a8a
52 changed files with 14122 additions and 9852 deletions
@@ -298,38 +298,67 @@ var _ = ginkgo.Describe("TrainJob e2e", func() {
 		})
 	})

-	ginkgo.When("Creating a TrainJob with PodTemplateOverrides", func() {
-		ginkgo.It("should allow the user to set the manager field and default to 'Unknown' if empty", func() {
+	ginkgo.When("Creating a TrainJob with RuntimePatches", func() {
+		ginkgo.It("should preserve user-provided manager fields", func() {
 			trainJob := testingutil.MakeTrainJobWrapper(ns.Name, "e2e-test").
 				RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), torchRuntime).
-				PodTemplateOverrides([]trainer.PodTemplateOverride{
+				RuntimePatches([]trainer.RuntimePatch{
 					{
-						TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: constants.Node}},
-						Spec: &trainer.PodTemplateSpecOverride{
-							ServiceAccountName: ptr.To("test-sa-1"),
+						Manager: "test.io/manager-one",
+						TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
+							Template: &trainer.JobSetTemplatePatch{
+								Spec: &trainer.JobSetSpecPatch{
+									ReplicatedJobs: []trainer.ReplicatedJobPatch{{
+										Name: constants.Node,
+										Template: &trainer.JobTemplatePatch{
+											Spec: &trainer.JobSpecPatch{
+												Template: &trainer.PodTemplatePatch{
+													Spec: &trainer.PodSpecPatch{
+														ServiceAccountName: ptr.To("test-sa-1"),
+													},
+												},
+											},
+										},
+									}},
+								},
+							},
 						},
 					},
 					{
-						Manager:    ptr.To("kueue.k8s.io/manager"),
-						TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: constants.Node}},
-						Spec: &trainer.PodTemplateSpecOverride{
-							ServiceAccountName: ptr.To("test-sa-2"),
+						Manager: "kueue.k8s.io/manager",
+						TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
+							Template: &trainer.JobSetTemplatePatch{
+								Spec: &trainer.JobSetSpecPatch{
+									ReplicatedJobs: []trainer.ReplicatedJobPatch{{
+										Name: constants.Node,
+										Template: &trainer.JobTemplatePatch{
+											Spec: &trainer.JobSpecPatch{
+												Template: &trainer.PodTemplatePatch{
+													Spec: &trainer.PodSpecPatch{
+														ServiceAccountName: ptr.To("test-sa-2"),
+													},
+												},
+											},
+										},
+									}},
+								},
+							},
 						},
 					},
 				}).
 				Obj()

-			ginkgo.By("Create a TrainJob with PodTemplateOverrides", func() {
+			ginkgo.By("Create a TrainJob with RuntimePatches", func() {
 				gomega.Expect(k8sClient.Create(ctx, trainJob)).Should(gomega.Succeed())
 			})

-			ginkgo.By("Verify manager field preserves user input and defaults to 'Unknown' if empty", func() {
+			ginkgo.By("Verify manager fields are preserved", func() {
 				gomega.Eventually(func(g gomega.Gomega) {
 					gotTrainJob := &trainer.TrainJob{}
 					g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(trainJob), gotTrainJob)).Should(gomega.Succeed())
-					g.Expect(gotTrainJob.Spec.PodTemplateOverrides).Should(gomega.HaveLen(2))
-					g.Expect(*gotTrainJob.Spec.PodTemplateOverrides[0].Manager).To(gomega.Equal("trainer.kubeflow.org/unknown"))
-					g.Expect(*gotTrainJob.Spec.PodTemplateOverrides[1].Manager).To(gomega.Equal("kueue.k8s.io/manager"))
+					g.Expect(gotTrainJob.Spec.RuntimePatches).Should(gomega.HaveLen(2))
+					g.Expect(gotTrainJob.Spec.RuntimePatches[0].Manager).To(gomega.Equal("test.io/manager-one"))
+					g.Expect(gotTrainJob.Spec.RuntimePatches[1].Manager).To(gomega.Equal("kueue.k8s.io/manager"))
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
 			})
 		})
@@ -85,8 +85,17 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 			trainJob = testingutil.MakeTrainJobWrapper(ns.Name, "alpha").
 				Suspend(true).
 				RuntimeRef(trainer.GroupVersion.WithKind(trainer.TrainingRuntimeKind), "alpha").
-				SpecLabel("testingKey", "testingVal").
-				SpecAnnotation("testingKey", "testingVal").
+				RuntimePatches([]trainer.RuntimePatch{{
+					Manager: "test.io/manager",
+					TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
+						Template: &trainer.JobSetTemplatePatch{
+							Metadata: &metav1.ObjectMeta{
+								Labels:      map[string]string{"testingKey": "testingVal"},
+								Annotations: map[string]string{"testingKey": "testingVal"},
+							},
+						},
+					},
+				}}).
 				Trainer(
 					testingutil.MakeTrainJobTrainerWrapper().
 						Container("test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
@@ -201,22 +210,35 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {

 				ginkgo.By("Updating suspended TrainJob node selector")
 				updatedSelector := map[string]string{"updated": "selector"}
-				podTemplateOverrides := []trainer.PodTemplateOverride{
-					{
-						TargetJobs: []trainer.PodTemplateOverrideTargetJob{
-							{
-								Name: "node",
+				runtimePatches := []trainer.RuntimePatch{{
+					Manager: "test.io/manager",
+					TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
+						Template: &trainer.JobSetTemplatePatch{
+							Metadata: &metav1.ObjectMeta{
+								Labels:      map[string]string{"testingKey": "testingVal"},
+								Annotations: map[string]string{"testingKey": "testingVal"},
+							},
+							Spec: &trainer.JobSetSpecPatch{
+								ReplicatedJobs: []trainer.ReplicatedJobPatch{{
+									Name: "node",
+									Template: &trainer.JobTemplatePatch{
+										Spec: &trainer.JobSpecPatch{
+											Template: &trainer.PodTemplatePatch{
+												Spec: &trainer.PodSpecPatch{
+													NodeSelector: updatedSelector,
+												},
+											},
+										},
+									},
+								}},
 							},
 						},
-						Spec: &trainer.PodTemplateSpecOverride{
-							NodeSelector: updatedSelector,
-						},
 					},
-				}
+				}}

 				gomega.Eventually(func(g gomega.Gomega) {
 					g.Expect(k8sClient.Get(ctx, trainJobKey, trainJob)).Should(gomega.Succeed())
-					trainJob.Spec.PodTemplateOverrides = podTemplateOverrides
+					trainJob.Spec.RuntimePatches = runtimePatches
 					g.Expect(k8sClient.Update(ctx, trainJob)).Should(gomega.Succeed())
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())

@@ -773,7 +795,7 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
 			})

-			ginkgo.It("Should succeed to create TrainJob with PodTemplateOverrides", func() {
+			ginkgo.It("Should succeed to create TrainJob with RuntimePatches", func() {
 				ginkgo.By("Creating Torch TrainingRuntime and TrainJob")
 				trainJob = testingutil.MakeTrainJobWrapper(ns.Name, "alpha").
 					RuntimeRef(trainer.GroupVersion.WithKind(trainer.TrainingRuntimeKind), "alpha").
@@ -782,34 +804,49 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 							Container("test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
 							Env([]corev1.EnvVar{{Name: "TRAIN_JOB", Value: "value"}}...).
 							Obj()).
-					PodTemplateOverrides([]trainer.PodTemplateOverride{
+					RuntimePatches([]trainer.RuntimePatch{
 						{
-							TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: constants.Node}},
-							Metadata: &metav1.ObjectMeta{
-								Labels: map[string]string{
-									"override-label-key": "override-label-value",
-									"custom-label":       "custom-value",
-								},
-								Annotations: map[string]string{
-									"override-annotation-key": "override-annotation-value",
-									"custom-annotation":       "custom-annotation-value",
-								},
-							},
-							Spec: &trainer.PodTemplateSpecOverride{
-								ServiceAccountName: ptr.To("override-sa"),
-								InitContainers: []trainer.ContainerOverride{
-									{
-										Name: "override-init-container",
-										Env: []corev1.EnvVar{
-											{
-												Name:  "INIT_ENV",
-												Value: "override_init",
+							Manager: "test.io/manager",
+							TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
+								Template: &trainer.JobSetTemplatePatch{
+									Spec: &trainer.JobSetSpecPatch{
+										ReplicatedJobs: []trainer.ReplicatedJobPatch{{
+											Name: constants.Node,
+											Template: &trainer.JobTemplatePatch{
+												Spec: &trainer.JobSpecPatch{
+													Template: &trainer.PodTemplatePatch{
+														Metadata: &metav1.ObjectMeta{
+															Labels: map[string]string{
+																"override-label-key": "override-label-value",
+																"custom-label":       "custom-value",
+															},
+															Annotations: map[string]string{
+																"override-annotation-key": "override-annotation-value",
+																"custom-annotation":       "custom-annotation-value",
+															},
+														},
+														Spec: &trainer.PodSpecPatch{
+															ServiceAccountName: ptr.To("override-sa"),
+															InitContainers: []trainer.ContainerPatch{
+																{
+																	Name: "override-init-container",
+																	Env: []corev1.EnvVar{
+																		{
+																			Name:  "INIT_ENV",
+																			Value: "override_init",
+																		},
+																		{
+																			Name:  "NEW_VALUE",
+																			Value: "from_overrides",
+																		},
+																	},
+																},
+															},
+														},
+													},
+												},
 											},
-											{
-												Name:  "NEW_VALUE",
-												Value: "from_overrides",
-											},
-										},
+										}},
 									},
 								},
 							},
@@ -330,48 +330,76 @@ var _ = ginkgo.Describe("TrainJob marker validations and defaulting", ginkgo.Ord
 						Obj()
 				},
 				testingutil.BeInvalidError()),
-			ginkgo.Entry("Should succeed to create trainJob with podTemplateOverrides containing duplicate targetJob",
+			ginkgo.Entry("Should succeed to create trainJob with runtimePatches containing two replicated job patches",
 				func() *trainer.TrainJob {
-					return testingutil.MakeTrainJobWrapper(ns.Name, "duplicated-podspecoverrides-target-jobs").
+					return testingutil.MakeTrainJobWrapper(ns.Name, "two-rjob-patches").
 						RuntimeRef(trainer.GroupVersion.WithKind(trainer.TrainingRuntimeKind), "testing").
-						PodTemplateOverrides([]trainer.PodTemplateOverride{
+						RuntimePatches([]trainer.RuntimePatch{
 							{
-								TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: "node"}},
-								Spec: &trainer.PodTemplateSpecOverride{
-									ServiceAccountName: ptr.To("custom-sa"),
+								Manager: "acme.io/manager-one",
+								TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
+									Template: &trainer.JobSetTemplatePatch{
+										Spec: &trainer.JobSetSpecPatch{
+											ReplicatedJobs: []trainer.ReplicatedJobPatch{{
+												Name: "node",
+												Template: &trainer.JobTemplatePatch{
+													Spec: &trainer.JobSpecPatch{
+														Template: &trainer.PodTemplatePatch{
+															Spec: &trainer.PodSpecPatch{
+																ServiceAccountName: ptr.To("custom-sa"),
+															},
+														},
+													},
+												},
+											}},
+										},
+									},
 								},
 							},
 							{
-								TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: "node"}},
-								Spec: &trainer.PodTemplateSpecOverride{
-									ServiceAccountName: ptr.To("custom-sa-two"),
+								Manager: "acme.io/manager-two",
+								TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
+									Template: &trainer.JobSetTemplatePatch{
+										Spec: &trainer.JobSetSpecPatch{
+											ReplicatedJobs: []trainer.ReplicatedJobPatch{{
+												Name: "node",
+												Template: &trainer.JobTemplatePatch{
+													Spec: &trainer.JobSpecPatch{
+														Template: &trainer.PodTemplatePatch{
+															Spec: &trainer.PodSpecPatch{
+																ServiceAccountName: ptr.To("custom-sa-two"),
+															},
+														},
+													},
+												},
+											}},
+										},
+									},
 								},
 							},
 						}).
 						Obj()
 				},
 				gomega.Succeed()),
-			ginkgo.Entry("Should succeed to create trainJob with valid PodTemplateOverrides.Manager",
+			ginkgo.Entry("Should succeed to create trainJob with valid RuntimePatches manager",
 				func() *trainer.TrainJob {
-					return testingutil.MakeTrainJobWrapper(ns.Name, "valid-podtemplateoverrides-manager").
+					return testingutil.MakeTrainJobWrapper(ns.Name, "valid-runtimepatches-manager").
 						RuntimeRef(trainer.GroupVersion.WithKind(trainer.TrainingRuntimeKind), "testing").
-						PodTemplateOverrides([]trainer.PodTemplateOverride{
+						RuntimePatches([]trainer.RuntimePatch{
 							{
-								TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: "node"}},
-								Manager:    ptr.To("kueue.k8s.io/manager"),
+								Manager: "kueue.k8s.io/manager",
 							},
 						}).
 						Obj()
 				},
 				gomega.Succeed()),
-			ginkgo.Entry("Should fail to create trainJob with invalid PodTemplateOverrides.Manager",
+			ginkgo.Entry("Should fail to create trainJob with invalid RuntimePatches manager",
 				func() *trainer.TrainJob {
-					return testingutil.MakeTrainJobWrapper(ns.Name, "invalid-podtemplateoverrides-manager").
+					return testingutil.MakeTrainJobWrapper(ns.Name, "invalid-runtimepatches-manager").
 						RuntimeRef(trainer.GroupVersion.WithKind(trainer.TrainingRuntimeKind), "testing").
-						PodTemplateOverrides([]trainer.PodTemplateOverride{
+						RuntimePatches([]trainer.RuntimePatch{
 							{
-								TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: "node"}},
-								Manager:    ptr.To("kueue.k8s.io/"),
+								Manager: "kueue.k8s.io/",
 							},
 						}).
 						Obj()
@@ -397,31 +425,6 @@ var _ = ginkgo.Describe("TrainJob marker validations and defaulting", ginkgo.Ord
 						Suspend(false).
 						Obj()
 				}),
-			ginkgo.Entry("Should succeed to default PodTemplateOverrides.Manager='trainer.kubeflow.org/unknown'",
-				func() *trainer.TrainJob {
-					return testingutil.MakeTrainJobWrapper(ns.Name, "null-podtemplateoverrides-manager").
-						RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "testing").
-						PodTemplateOverrides([]trainer.PodTemplateOverride{
-							{
-								TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: "node"}},
-							},
-						}).
-						Suspend(true).
-						Obj()
-				},
-				func() *trainer.TrainJob {
-					return testingutil.MakeTrainJobWrapper(ns.Name, "null-podtemplateoverrides-manager").
-						ManagedBy("trainer.kubeflow.org/trainjob-controller").
-						RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "testing").
-						PodTemplateOverrides([]trainer.PodTemplateOverride{
-							{
-								TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: "node"}},
-								Manager:    ptr.To("trainer.kubeflow.org/unknown"),
-							},
-						}).
-						Suspend(true).
-						Obj()
-				}),
 			ginkgo.Entry("Should succeed to default managedBy=trainer.kubeflow.org/trainjob-controller",
 				func() *trainer.TrainJob {
 					return testingutil.MakeTrainJobWrapper(ns.Name, "null-managed-by").
@@ -527,29 +530,38 @@ var _ = ginkgo.Describe("TrainJob marker validations and defaulting", ginkgo.Ord
 					return job
 				},
 				testingutil.BeInvalidError()),
-			ginkgo.Entry("Should succeed to update podTemplateOverride when suspend is true",
+			ginkgo.Entry("Should succeed to update runtimePatches when suspend is true",
 				func() *trainer.TrainJob {
-					return testingutil.MakeTrainJobWrapper(ns.Name, "valid-trainer").
+					return testingutil.MakeTrainJobWrapper(ns.Name, "valid-runtimepatches").
 						RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "testing").
 						Suspend(true).
-						PodTemplateOverrides(
-							[]trainer.PodTemplateOverride{
-								{
-									TargetJobs: []trainer.PodTemplateOverrideTargetJob{
-										{
-											Name: "node",
+						RuntimePatches([]trainer.RuntimePatch{
+							{
+								Manager: "test.io/manager",
+								TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
+									Template: &trainer.JobSetTemplatePatch{
+										Spec: &trainer.JobSetSpecPatch{
+											ReplicatedJobs: []trainer.ReplicatedJobPatch{{
+												Name: "node",
+												Template: &trainer.JobTemplatePatch{
+													Spec: &trainer.JobSpecPatch{
+														Template: &trainer.PodTemplatePatch{
+															Spec: &trainer.PodSpecPatch{
+																NodeSelector: map[string]string{"test": "test"},
+															},
+														},
+													},
+												},
+											}},
 										},
 									},
-									Spec: &trainer.PodTemplateSpecOverride{
-										NodeSelector: map[string]string{"test": "test"},
-									},
 								},
 							},
-						).
+						}).
 						Obj()
 				},
 				func(job *trainer.TrainJob) *trainer.TrainJob {
-					job.Spec.PodTemplateOverrides[0].Spec.NodeSelector = map[string]string{"allow": "update"}
+					job.Spec.RuntimePatches[0].TrainingRuntimeSpec.Template.Spec.ReplicatedJobs[0].Template.Spec.Template.Spec.NodeSelector = map[string]string{"allow": "update"}
 					return job
 				},
 				gomega.Succeed()),