Carry over ObjectMeta from pod template

Signed-off-by: Fabio Grätz <[email protected]>
flyteorg · Dec 10, 2022 · ffbe158 · ffbe158
1 parent 1c8f4de
commit ffbe158
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 6 deletions.
diff --git a/go/tasks/plugins/k8s/kfoperators/mpi/mpi.go b/go/tasks/plugins/k8s/kfoperators/mpi/mpi.go
@@ -71,12 +71,15 @@ func (mpiOperatorResourceHandler) BuildResource(ctx context.Context, taskCtx plu
 
 	podTemplate := flytek8s.DefaultPodTemplateStore.LoadOrDefault(taskCtx.TaskExecutionMetadata().GetNamespace())
 
+	objectMeta := metav1.ObjectMeta{}
+
 	if podTemplate != nil {
 		mergedPodSpec, err := flytek8s.MergePodSpecs(&podTemplate.Template.Spec, podSpec, kubeflowv1.MPIJobDefaultContainerName)
 		if err != nil {
 			return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "Unable to merge default pod template: [%v]", err.Error())
 		}
 		podSpec = mergedPodSpec
+		objectMeta = podTemplate.Template.ObjectMeta
 	}
 
 	// workersPodSpec is deepCopy of podSpec submitted by flyte
@@ -101,14 +104,16 @@ func (mpiOperatorResourceHandler) BuildResource(ctx context.Context, taskCtx plu
 			kubeflowv1.MPIJobReplicaTypeLauncher: {
 				Replicas: &launcherReplicas,
 				Template: v1.PodTemplateSpec{
-					Spec: *podSpec,
+					ObjectMeta: objectMeta,
+					Spec:       *podSpec,
 				},
 				RestartPolicy: commonKf.RestartPolicyNever,
 			},
 			kubeflowv1.MPIJobReplicaTypeWorker: {
 				Replicas: &workers,
 				Template: v1.PodTemplateSpec{
-					Spec: *workersPodSpec,
+					ObjectMeta: objectMeta,
+					Spec:       *workersPodSpec,
 				},
 				RestartPolicy: commonKf.RestartPolicyNever,
 			},

diff --git a/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch.go b/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch.go
@@ -70,12 +70,15 @@ func (pytorchOperatorResourceHandler) BuildResource(ctx context.Context, taskCtx
 
 	podTemplate := flytek8s.DefaultPodTemplateStore.LoadOrDefault(taskCtx.TaskExecutionMetadata().GetNamespace())
 
+	objectMeta := metav1.ObjectMeta{}
+
 	if podTemplate != nil {
 		mergedPodSpec, err := flytek8s.MergePodSpecs(&podTemplate.Template.Spec, podSpec, kubeflowv1.PytorchJobDefaultContainerName)
 		if err != nil {
 			return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "Unable to merge default pod template: [%v]", err.Error())
 		}
 		podSpec = mergedPodSpec
+		objectMeta = podTemplate.Template.ObjectMeta
 	}
 
 	workers := pytorchTaskExtraArgs.GetWorkers()
@@ -84,14 +87,16 @@ func (pytorchOperatorResourceHandler) BuildResource(ctx context.Context, taskCtx
 		PyTorchReplicaSpecs: map[commonOp.ReplicaType]*commonOp.ReplicaSpec{
 			kubeflowv1.PyTorchJobReplicaTypeMaster: {
 				Template: v1.PodTemplateSpec{
-					Spec: *podSpec,
+					ObjectMeta: objectMeta,
+					Spec:       *podSpec,
 				},
 				RestartPolicy: commonOp.RestartPolicyNever,
 			},
 			kubeflowv1.PyTorchJobReplicaTypeWorker: {
 				Replicas: &workers,
 				Template: v1.PodTemplateSpec{
-					Spec: *podSpec,
+					ObjectMeta: objectMeta,
+					Spec:       *podSpec,
 				},
 				RestartPolicy: commonOp.RestartPolicyNever,
 			},

diff --git a/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow.go b/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow.go
@@ -70,12 +70,15 @@ func (tensorflowOperatorResourceHandler) BuildResource(ctx context.Context, task
 
 	podTemplate := flytek8s.DefaultPodTemplateStore.LoadOrDefault(taskCtx.TaskExecutionMetadata().GetNamespace())
 
+	objectMeta := metav1.ObjectMeta{}
+
 	if podTemplate != nil {
 		mergedPodSpec, err := flytek8s.MergePodSpecs(&podTemplate.Template.Spec, podSpec, kubeflowv1.TFJobDefaultContainerName)
 		if err != nil {
 			return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "Unable to merge default pod template: [%v]", err.Error())
 		}
 		podSpec = mergedPodSpec
+		objectMeta = podTemplate.Template.ObjectMeta
 	}
 
 	workers := tensorflowTaskExtraArgs.GetWorkers()
@@ -87,14 +90,16 @@ func (tensorflowOperatorResourceHandler) BuildResource(ctx context.Context, task
 			kubeflowv1.TFJobReplicaTypePS: {
 				Replicas: &psReplicas,
 				Template: v1.PodTemplateSpec{
-					Spec: *podSpec,
+					ObjectMeta: objectMeta,
+					Spec:       *podSpec,
 				},
 				RestartPolicy: commonOp.RestartPolicyNever,
 			},
 			kubeflowv1.TFJobReplicaTypeChief: {
 				Replicas: &chiefReplicas,
 				Template: v1.PodTemplateSpec{
-					Spec: *podSpec,
+					ObjectMeta: objectMeta,
+					Spec:       *podSpec,
 				},
 				RestartPolicy: commonOp.RestartPolicyNever,
 			},