From f5a7c9f0a101e1c38b66b2b046b851071b3437f4 Mon Sep 17 00:00:00 2001
From: Ce Gao <gaoce@caicloud.io>
Date: Mon, 6 Jan 2020 17:08:40 +0800
Subject: [PATCH] fix: Inherit labels and annotations from experiment

Signed-off-by: Ce Gao <gaoce@caicloud.io>
---
 pkg/controller.v1alpha3/consts/const.go       |  2 ++
 .../experiment/suggestion/suggestion.go       |  6 +++--
 .../suggestion/composer/composer.go           |  9 +++++---
 pkg/controller.v1alpha3/util/labels.go        | 22 ++++++++++++++-----
 4 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/pkg/controller.v1alpha3/consts/const.go b/pkg/controller.v1alpha3/consts/const.go
index 8b69dd24d74..e45c74d7a9a 100644
--- a/pkg/controller.v1alpha3/consts/const.go
+++ b/pkg/controller.v1alpha3/consts/const.go
@@ -18,6 +18,8 @@ const (
 	LabelExperimentName = "experiment"
 	// LabelSuggestionName is the label of suggestion name.
 	LabelSuggestionName = "suggestion"
+	// LabelDeploymentName is the label of deployment name.
+	LabelDeploymentName = "deployment"
 
 	// ContainerSuggestion is the container name in Suggestion.
 	ContainerSuggestion = "suggestion"
diff --git a/pkg/controller.v1alpha3/experiment/suggestion/suggestion.go b/pkg/controller.v1alpha3/experiment/suggestion/suggestion.go
index 2540719fa50..de95d2a350f 100644
--- a/pkg/controller.v1alpha3/experiment/suggestion/suggestion.go
+++ b/pkg/controller.v1alpha3/experiment/suggestion/suggestion.go
@@ -57,8 +57,10 @@ func (g *General) createSuggestion(instance *experimentsv1alpha3.Experiment, sug
 	logger := log.WithValues("experiment", types.NamespacedName{Name: instance.Name, Namespace: instance.Namespace})
 	suggestion := &suggestionsv1alpha3.Suggestion{
 		ObjectMeta: metav1.ObjectMeta{
-			Name:      instance.Name,
-			Namespace: instance.Namespace,
+			Name:        instance.Name,
+			Namespace:   instance.Namespace,
+			Labels:      instance.Labels,
+			Annotations: instance.Annotations,
 		},
 		Spec: suggestionsv1alpha3.SuggestionSpec{
 			AlgorithmName: instance.Spec.Algorithm.AlgorithmName,
diff --git a/pkg/controller.v1alpha3/suggestion/composer/composer.go b/pkg/controller.v1alpha3/suggestion/composer/composer.go
index 184c0a8e212..f81e0b32e46 100644
--- a/pkg/controller.v1alpha3/suggestion/composer/composer.go
+++ b/pkg/controller.v1alpha3/suggestion/composer/composer.go
@@ -54,8 +54,10 @@ func (g *General) DesiredDeployment(s *suggestionsv1alpha3.Suggestion) (*appsv1.
 	}
 	d := &appsv1.Deployment{
 		ObjectMeta: metav1.ObjectMeta{
-			Name:      util.GetAlgorithmDeploymentName(s),
-			Namespace: s.Namespace,
+			Name:        util.GetAlgorithmDeploymentName(s),
+			Namespace:   s.Namespace,
+			Labels:      s.Labels,
+			Annotations: s.Annotations,
 		},
 		Spec: appsv1.DeploymentSpec{
 			Selector: &metav1.LabelSelector{
@@ -63,7 +65,8 @@ func (g *General) DesiredDeployment(s *suggestionsv1alpha3.Suggestion) (*appsv1.
 			},
 			Template: corev1.PodTemplateSpec{
 				ObjectMeta: metav1.ObjectMeta{
-					Labels: util.SuggestionLabels(s),
+					Labels:      util.SuggestionLabels(s),
+					Annotations: s.Annotations,
 				},
 				Spec: corev1.PodSpec{
 					Containers: []corev1.Container{
diff --git a/pkg/controller.v1alpha3/util/labels.go b/pkg/controller.v1alpha3/util/labels.go
index d941063cf46..4b3e6904b33 100644
--- a/pkg/controller.v1alpha3/util/labels.go
+++ b/pkg/controller.v1alpha3/util/labels.go
@@ -6,16 +6,26 @@ import (
 	"github.com/kubeflow/katib/pkg/controller.v1alpha3/consts"
 )
 
+// SuggestionLabels returns the expected suggestion labels.
 func SuggestionLabels(instance *suggestionsv1alpha3.Suggestion) map[string]string {
-	return map[string]string{
-		"deployment":               GetAlgorithmDeploymentName(instance),
-		consts.LabelExperimentName: instance.Name,
-		consts.LabelSuggestionName: instance.Name,
+	res := make(map[string]string)
+	for k, v := range instance.Labels {
+		res[k] = v
 	}
+	res[consts.LabelDeploymentName] = GetAlgorithmDeploymentName(instance)
+	res[consts.LabelExperimentName] = instance.Name
+	res[consts.LabelSuggestionName] = instance.Name
+
+	return res
 }
 
+// TrialLabels returns the expected trial labels.
 func TrialLabels(instance *experimentsv1alpha3.Experiment) map[string]string {
-	return map[string]string{
-		consts.LabelExperimentName: instance.Name,
+	res := make(map[string]string)
+	for k, v := range instance.Labels {
+		res[k] = v
 	}
+	res[consts.LabelExperimentName] = instance.Name
+
+	return res
 }