feat: Add the manager field to the podTemplateOverride object (#3020)

Signed-off-by: Tomas Tormo <[email protected]>

api/openapi-spec/swagger.json

@@ -14084,6 +14084,10 @@
           "targetJobs"
         ],
         "properties": {
+          "manager": {
+            "description": "manager indicates which controller created this PodTemplateOverride object. It can be used by external controllers or admission webhooks to track ownership and avoid conflicts with other controllers. For example, Kueue sets this field to \"kueue.k8s.io/manager\". Defaults to `trainer.kubeflow.org/unknown`",
+            "type": "string"
+          },
           "metadata": {
             "description": "metadata overrides the Pod template metadata. These values will be merged with the TrainingRuntime's Pod template metadata.",
             "allOf": [

api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_pod_template_override.py

@@ -17,7 +17,7 @@ import pprint
 import re  # noqa: F401
 import json
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, StrictStr
 from typing import Any, ClassVar, Dict, List, Optional
 from kubeflow_trainer_api.models.io_k8s_apimachinery_pkg_apis_meta_v1_object_meta import IoK8sApimachineryPkgApisMetaV1ObjectMeta
 from kubeflow_trainer_api.models.trainer_v1alpha1_pod_template_override_target_job import TrainerV1alpha1PodTemplateOverrideTargetJob
@@ -29,10 +29,11 @@ class TrainerV1alpha1PodTemplateOverride(BaseModel):
     """
     PodTemplateOverride represents a custom PodTemplateSpec override that will be applied to the TrainJob's training runtime.
     """ # noqa: E501
+    manager: Optional[StrictStr] = Field(default=None, description="manager indicates which controller created this PodTemplateOverride object. It can be used by external controllers or admission webhooks to track ownership and avoid conflicts with other controllers. For example, Kueue sets this field to \"kueue.k8s.io/manager\". Defaults to `trainer.kubeflow.org/unknown`")
     metadata: Optional[IoK8sApimachineryPkgApisMetaV1ObjectMeta] = Field(default=None, description="metadata overrides the Pod template metadata. These values will be merged with the TrainingRuntime's Pod template metadata.")
     spec: Optional[TrainerV1alpha1PodTemplateSpecOverride] = Field(default=None, description="spec overrides the Pod template spec. These values will be merged with the TrainingRuntime's Pod template spec.")
     target_jobs: List[TrainerV1alpha1PodTemplateOverrideTargetJob] = Field(description="targetJobs is the list of replicated jobs in the training runtime template to apply the overrides.", alias="targetJobs")
-    __properties: ClassVar[List[str]] = ["metadata", "spec", "targetJobs"]
+    __properties: ClassVar[List[str]] = ["manager", "metadata", "spec", "targetJobs"]
 
     model_config = ConfigDict(
         populate_by_name=True,
@@ -98,6 +99,7 @@ class TrainerV1alpha1PodTemplateOverride(BaseModel):
             return cls.model_validate(obj)
 
         _obj = cls.model_validate({
+            "manager": obj.get("manager"),
             "metadata": IoK8sApimachineryPkgApisMetaV1ObjectMeta.from_dict(obj["metadata"]) if obj.get("metadata") is not None else None,
             "spec": TrainerV1alpha1PodTemplateSpecOverride.from_dict(obj["spec"]) if obj.get("spec") is not None else None,
             "targetJobs": [TrainerV1alpha1PodTemplateOverrideTargetJob.from_dict(_item) for _item in obj["targetJobs"]] if obj.get("targetJobs") is not None else None

charts/kubeflow-trainer/crds/trainer.kubeflow.org_trainjobs.yaml

@@ -465,6 +465,15 @@ spec:
                   description: PodTemplateOverride represents a custom PodTemplateSpec
                     override that will be applied to the TrainJob's training runtime.
                   properties:
+                    manager:
+                      default: trainer.kubeflow.org/unknown
+                      description: |-
+                        manager indicates which controller created this PodTemplateOverride object.
+                        It can be used by external controllers or admission webhooks to track ownership
+                        and avoid conflicts with other controllers. For example, Kueue sets this field
+                        to "kueue.k8s.io/manager".
+                        Defaults to `trainer.kubeflow.org/unknown`
+                      type: string
                     metadata:
                       description: |-
                         metadata overrides the Pod template metadata.

manifests/base/crds/trainer.kubeflow.org_trainjobs.yaml

@@ -465,6 +465,15 @@ spec:
                   description: PodTemplateOverride represents a custom PodTemplateSpec
                     override that will be applied to the TrainJob's training runtime.
                   properties:
+                    manager:
+                      default: trainer.kubeflow.org/unknown
+                      description: |-
+                        manager indicates which controller created this PodTemplateOverride object.
+                        It can be used by external controllers or admission webhooks to track ownership
+                        and avoid conflicts with other controllers. For example, Kueue sets this field
+                        to "kueue.k8s.io/manager".
+                        Defaults to `trainer.kubeflow.org/unknown`
+                      type: string
                     metadata:
                       description: |-
                         metadata overrides the Pod template metadata.

pkg/apis/trainer/v1alpha1/trainjob_types.go

@@ -269,6 +269,15 @@ type PodTemplateOverride struct {
 	// +required
 	TargetJobs []PodTemplateOverrideTargetJob `json:"targetJobs,omitempty"`
 
+	// manager indicates which controller created this PodTemplateOverride object.
+	// It can be used by external controllers or admission webhooks to track ownership
+	// and avoid conflicts with other controllers. For example, Kueue sets this field
+	// to "kueue.k8s.io/manager".
+	// Defaults to `trainer.kubeflow.org/unknown`
+	// +kubebuilder:default="trainer.kubeflow.org/unknown"
+	// +optional
+	Manager *string `json:"manager,omitempty"`
+
 	// metadata overrides the Pod template metadata.
 	// These values will be merged with the TrainingRuntime's Pod template metadata.
 	// +optional

pkg/apis/trainer/v1alpha1/zz_generated.deepcopy.go

@@ -451,6 +451,11 @@ func (in *PodTemplateOverride) DeepCopyInto(out *PodTemplateOverride) {
 		*out = make([]PodTemplateOverrideTargetJob, len(*in))
 		copy(*out, *in)
 	}
+	if in.Manager != nil {
+		in, out := &in.Manager, &out.Manager
+		*out = new(string)
+		**out = **in
+	}
 	if in.Metadata != nil {
 		in, out := &in.Metadata, &out.Metadata
 		*out = new(metav1.ObjectMeta)

pkg/apis/trainer/v1alpha1/zz_generated.openapi.go

@@ -1032,6 +1032,13 @@ func schema_pkg_apis_trainer_v1alpha1_PodTemplateOverride(ref common.ReferenceCa
 							},
 						},
 					},
+					"manager": {
+						SchemaProps: spec.SchemaProps{
+							Description: "manager indicates which controller created this PodTemplateOverride object. It can be used by external controllers or admission webhooks to track ownership and avoid conflicts with other controllers. For example, Kueue sets this field to \"kueue.k8s.io/manager\". Defaults to `trainer.kubeflow.org/unknown`",
+							Type:        []string{"string"},
+							Format:      "",
+						},
+					},
 					"metadata": {
 						SchemaProps: spec.SchemaProps{
 							Description: "metadata overrides the Pod template metadata. These values will be merged with the TrainingRuntime's Pod template metadata.",

pkg/client/applyconfiguration/trainer/v1alpha1/podtemplateoverride.go

@@ -27,6 +27,12 @@ import (
 type PodTemplateOverrideApplyConfiguration struct {
 	// targetJobs is the list of replicated jobs in the training runtime template to apply the overrides.
 	TargetJobs []PodTemplateOverrideTargetJobApplyConfiguration `json:"targetJobs,omitempty"`
+	// manager indicates which controller created this PodTemplateOverride object.
+	// It can be used by external controllers or admission webhooks to track ownership
+	// and avoid conflicts with other controllers. For example, Kueue sets this field
+	// to "kueue.k8s.io/manager".
+	// Defaults to `trainer.kubeflow.org/unknown`
+	Manager *string `json:"manager,omitempty"`
 	// metadata overrides the Pod template metadata.
 	// These values will be merged with the TrainingRuntime's Pod template metadata.
 	Metadata *v1.ObjectMetaApplyConfiguration `json:"metadata,omitempty"`
@@ -54,6 +60,14 @@ func (b *PodTemplateOverrideApplyConfiguration) WithTargetJobs(values ...*PodTem
 	return b
 }
 
+// WithManager sets the Manager field in the declarative configuration to the given value
+// and returns the receiver, so that objects can be built by chaining "With" function invocations.
+// If called multiple times, the Manager field is set to the value of the last call.
+func (b *PodTemplateOverrideApplyConfiguration) WithManager(value string) *PodTemplateOverrideApplyConfiguration {
+	b.Manager = &value
+	return b
+}
+
 // WithMetadata sets the Metadata field in the declarative configuration to the given value
 // and returns the receiver, so that objects can be built by chaining "With" function invocations.
 // If called multiple times, the Metadata field is set to the value of the last call.

pkg/runtime/core/trainingruntime_test.go

@@ -237,6 +237,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 				).
 				PodTemplateOverrides([]trainer.PodTemplateOverride{
 					{
+						Manager:    ptr.To("manager-1"),
 						TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: constants.DatasetInitializer}},
 						Spec: &trainer.PodTemplateSpecOverride{
 							InitContainers: []trainer.ContainerOverride{
@@ -322,6 +323,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 						},
 					},
 					{
+						Manager:    ptr.To("manager-2"),
 						TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: constants.Node}},
 						Spec: &trainer.PodTemplateSpecOverride{
 							ServiceAccountName: ptr.To("override-sa"),
@@ -615,12 +617,14 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 				).
 				PodTemplateOverrides([]trainer.PodTemplateOverride{
 					{
+						Manager:    ptr.To("manager-1"),
 						TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: constants.DatasetInitializer}},
 						Spec: &trainer.PodTemplateSpecOverride{
 							Affinity: nil,
 						},
 					},
 					{
+						Manager:    ptr.To("manager-2"),
 						TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: constants.Node}},
 						Spec: &trainer.PodTemplateSpecOverride{
 							Affinity: nil,
@@ -731,6 +735,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 				).
 				PodTemplateOverrides([]trainer.PodTemplateOverride{
 					{
+						Manager:    ptr.To("manager-1"),
 						TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: constants.DatasetInitializer}},
 						Metadata: &metav1.ObjectMeta{
 							Labels:      map[string]string{"k1": "v1"},
@@ -738,6 +743,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 						},
 					},
 					{
+						Manager:    ptr.To("manager-2"),
 						TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: constants.Node}},
 						Metadata: &metav1.ObjectMeta{
 							Labels:      map[string]string{"k2": "v2"},
@@ -775,6 +781,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 				RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime").
 				PodTemplateOverrides([]trainer.PodTemplateOverride{
 					{
+						Manager:    ptr.To("manager-1"),
 						TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: constants.Node}},
 						Spec: &trainer.PodTemplateSpecOverride{
 							NodeSelector: map[string]string{
@@ -783,6 +790,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 						},
 					},
 					{
+						Manager:    ptr.To("manager-1"),
 						TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: constants.Node}},
 						Spec: &trainer.PodTemplateSpecOverride{
 							ServiceAccountName: ptr.To("test-sa"),
@@ -816,6 +824,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 				RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime").
 				PodTemplateOverrides([]trainer.PodTemplateOverride{
 					{
+						Manager:    ptr.To("manager-1"),
 						TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: constants.Node}},
 						Spec: &trainer.PodTemplateSpecOverride{
 							NodeSelector: map[string]string{
@@ -824,6 +833,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 						},
 					},
 					{
+						Manager:    ptr.To("manager-1"),
 						TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: constants.Node}},
 						Spec: &trainer.PodTemplateSpecOverride{
 							NodeSelector: map[string]string{

pkg/runtime/framework/plugins/jobset/jobset.go

@@ -29,6 +29,7 @@ import (
 	apiruntime "k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/util/sets"
+	"k8s.io/apimachinery/pkg/util/validation"
 	"k8s.io/apimachinery/pkg/util/validation/field"
 	metav1ac "k8s.io/client-go/applyconfigurations/meta/v1"
 	"k8s.io/utils/ptr"
@@ -124,6 +125,9 @@ func (j *JobSet) Validate(ctx context.Context, info *runtime.Info, oldObj, newOb
 
 	// TODO (andreyvelich): Validate Volumes, VolumeMounts, and Tolerations.
 	for _, podTemplateOverride := range newObj.Spec.PodTemplateOverrides {
+		if podTemplateOverride.Manager != nil {
+			allErrs = append(allErrs, validation.IsDomainPrefixedPath(podTemplateOverridePath.Child("manager"), *podTemplateOverride.Manager)...)
+		}
 		for _, targetJob := range podTemplateOverride.TargetJobs {
 			containers, ok := rJobContainerNames[targetJob.Name]
 			if !ok {

test/e2e/e2e_test.go

@@ -236,4 +236,41 @@ 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() {
+			trainJob := testingutil.MakeTrainJobWrapper(ns.Name, "e2e-test").
+				RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind), torchRuntime).
+				PodTemplateOverrides([]trainer.PodTemplateOverride{
+					{
+						TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: constants.Node}},
+						Spec: &trainer.PodTemplateSpecOverride{
+							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"),
+						},
+					},
+				}).
+				Obj()
+
+			ginkgo.By("Create a TrainJob with PodTemplateOverrides", 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() {
+				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"))
+				}, util.Timeout, util.Interval).Should(gomega.Succeed())
+			})
+		})
+	})
 })

test/integration/webhooks/trainjob_test.go

@@ -372,6 +372,32 @@ var _ = ginkgo.Describe("TrainJob marker validations and defaulting", ginkgo.Ord
 						Obj()
 				},
 				gomega.Succeed()),
+			ginkgo.Entry("Should succeed to create trainJob with valid PodTemplateOverrides.Manager",
+				func() *trainer.TrainJob {
+					return testingutil.MakeTrainJobWrapper(ns.Name, "valid-podtemplateoverrides-manager").
+						RuntimeRef(trainer.GroupVersion.WithKind(trainer.TrainingRuntimeKind), "testing").
+						PodTemplateOverrides([]trainer.PodTemplateOverride{
+							{
+								TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: "node"}},
+								Manager:    ptr.To("kueue.k8s.io/manager"),
+							},
+						}).
+						Obj()
+				},
+				gomega.Succeed()),
+			ginkgo.Entry("Should fail to create trainJob with invalid PodTemplateOverrides.Manager",
+				func() *trainer.TrainJob {
+					return testingutil.MakeTrainJobWrapper(ns.Name, "invalid-podtemplateoverrides-manager").
+						RuntimeRef(trainer.GroupVersion.WithKind(trainer.TrainingRuntimeKind), "testing").
+						PodTemplateOverrides([]trainer.PodTemplateOverride{
+							{
+								TargetJobs: []trainer.PodTemplateOverrideTargetJob{{Name: "node"}},
+								Manager:    ptr.To("kueue.k8s.io/"),
+							},
+						}).
+						Obj()
+				},
+				testingutil.BeForbiddenError()),
 		)
 		ginkgo.DescribeTable("Defaulting TrainJob on creation", func(trainJob func() *trainer.TrainJob, wantTrainJob func() *trainer.TrainJob) {
 			created := trainJob()
@@ -392,6 +418,31 @@ 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").