Add Trial Labels During Pod Mutation

pkg/controller.v1beta1/consts/const.go

@@ -68,6 +68,8 @@ const (
 	LabelExperimentName = "katib.kubeflow.org/experiment"
 	// LabelSuggestionName is the label of suggestion name.
 	LabelSuggestionName = "katib.kubeflow.org/suggestion"
+	// LabelTrialName is the label of trial name.
+	LabelTrialName = "katib.kubeflow.org/trial"
 	// LabelDeploymentName is the label of deployment name.
 	LabelDeploymentName = "katib.kubeflow.org/deployment"
 

pkg/webhook/v1beta1/pod/inject_webhook.go

@@ -94,7 +94,7 @@ func (s *SidecarInjector) Handle(ctx context.Context, req admission.Request) adm
 	// Do mutation
 	mutatedPod, err := s.Mutate(pod, namespace)
 	if err != nil {
-		log.Error(err, "Failed to inject metrics collector")
+		log.Error(err, "Failed to mutate Trial's pod")
 		return admission.Errored(http.StatusBadRequest, err)
 	}
 
@@ -124,17 +124,6 @@ func (s *SidecarInjector) MutationRequired(pod *v1.Pod, ns string) (bool, error)
 		return false, err
 	}
 
-	// If PrimaryPodLabel is not set we mutate all pods which are related to Trial job
-	// Otherwise mutate pod only with appropriate labels
-	if trial.Spec.PrimaryPodLabels != nil {
-		if !isPrimaryPod(pod.Labels, trial.Spec.PrimaryPodLabels) {
-			return false, nil
-		}
-	}
-
-	if trial.Spec.MetricsCollector.Collector.Kind == common.NoneCollector {
-		return false, nil
-	}
 	return true, nil
 }
 
@@ -155,6 +144,21 @@ func (s *SidecarInjector) Mutate(pod *v1.Pod, namespace string) (*v1.Pod, error)
 		return nil, err
 	}
 
+	// Add Katib Trial labels to the Pod metadata.
+	mutatePodMetadata(mutatedPod, trial)
+
+	// Do the following mutation only for the Primary pod.
+	// If PrimaryPodLabel is not set we mutate all pods which are related to Trial job.
+	// Otherwise, mutate pod only with the appropriate labels.
+	if trial.Spec.PrimaryPodLabels != nil && !isPrimaryPod(pod.Labels, trial.Spec.PrimaryPodLabels) {
+		return mutatedPod, nil
+	}
+
+	// If Metrics Collector in None, skip the mutation.
+	if trial.Spec.MetricsCollector.Collector.Kind == common.NoneCollector {
+		return mutatedPod, nil
+	}
+
 	// Create metrics sidecar container spec
 	injectContainer, err := s.getMetricsCollectorContainer(trial, pod)
 	if err != nil {

pkg/webhook/v1beta1/pod/inject_webhook_test.go

@@ -1019,3 +1019,49 @@ func TestIsPrimaryPod(t *testing.T) {
 		}
 	}
 }
+
+func TestMutatePodMetadata(t *testing.T) {
+	mutatedPodLabels := map[string]string{
+		"custom-pod-label":    "custom-value",
+		"katib-experiment":    "katib-value",
+		consts.LabelTrialName: "test-trial",
+	}
+
+	testCases := []struct {
+		pod             *v1.Pod
+		trial           *trialsv1beta1.Trial
+		mutatedPod      *v1.Pod
+		testDescription string
+	}{
+		{
+			pod: &v1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: map[string]string{
+						"custom-pod-label": "custom-value",
+					},
+				},
+			},
+			trial: &trialsv1beta1.Trial{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "test-trial",
+					Labels: map[string]string{
+						"katib-experiment": "katib-value",
+					},
+				},
+			},
+			mutatedPod: &v1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: mutatedPodLabels,
+				},
+			},
+			testDescription: "Mutated Pod should contain label from the origin Pod and Trial",
+		},
+	}
+
+	for _, tc := range testCases {
+		mutatePodMetadata(tc.pod, tc.trial)
+		if !reflect.DeepEqual(tc.mutatedPod, tc.pod) {
+			t.Errorf("Case %v. Expected Pod %v, got %v", tc.testDescription, tc.mutatedPod, tc.pod)
+		}
+	}
+}

pkg/webhook/v1beta1/pod/utils.go

@@ -31,6 +31,7 @@ import (
 
 	common "github.com/kubeflow/katib/pkg/apis/controller/common/v1beta1"
 	trialsv1beta1 "github.com/kubeflow/katib/pkg/apis/controller/trials/v1beta1"
+	"github.com/kubeflow/katib/pkg/controller.v1beta1/consts"
 	mccommon "github.com/kubeflow/katib/pkg/metricscollector/v1beta1/common"
 )
 
@@ -260,6 +261,26 @@ func mutateMetricsCollectorVolume(pod *v1.Pod, mountPath, sidecarContainerName,
 	return nil
 }
 
+func mutatePodMetadata(pod *v1.Pod, trial *trialsv1beta1.Trial) {
+	podLabels := map[string]string{}
+
+	// Get labels from the created pod.
+	if pod.Labels != nil {
+		podLabels = pod.Labels
+	}
+
+	// Get labels from Trial.
+	for k, v := range trial.Labels {
+		podLabels[k] = v
+	}
+
+	// Add Trial name label.
+	podLabels[consts.LabelTrialName] = trial.GetName()
+
+	// Append label to the Pod metadata.
+	pod.Labels = podLabels
+}
+
 func getSidecarContainerName(cKind common.CollectorKind) string {
 	if cKind == common.StdOutCollector || cKind == common.FileCollector {
 		return mccommon.MetricLoggerCollectorContainerName