add labels and annotations from servingruntime

Signed-off-by: Lize Cai <[email protected]>
kserve · Oct 27, 2022 · b8f17fa · b8f17fa
1 parent f129168
commit b8f17fa
Show file tree

Hide file tree

Showing 3 changed files with 107 additions and 23 deletions.
diff --git a/controllers/modelmesh/runtime.go b/controllers/modelmesh/runtime.go
@@ -360,32 +360,15 @@ func (m *Deployment) addPassThroughPodFieldsToDeployment(deployment *appsv1.Depl
 }
 
 func (m *Deployment) configureRuntimePodSpecAnnotations(deployment *appsv1.Deployment) error {
-
-	if deployment.Spec.Template.Annotations == nil {
-		deployment.Spec.Template.Annotations = m.AnnotationsMap
-		return nil
-	}
-
-	// apply user configmap annotations
-	for key, value := range m.AnnotationsMap {
-		// set labels for pods created by deployment
-		deployment.Spec.Template.Annotations[key] = value
-	}
-
+	// merging the annotations
+	// priority: ServingRuntimePodSpec > AnnotationsMap > whatever in the deployment template
+	deployment.Spec.Template.Annotations = Union(deployment.Spec.Template.Annotations, m.AnnotationsMap, m.SRSpec.ServingRuntimePodSpec.Annotations)
 	return nil
 }
 
 func (m *Deployment) configureRuntimePodSpecLabels(deployment *appsv1.Deployment) error {
-
-	if deployment.Spec.Template.Labels == nil {
-		deployment.Spec.Template.Labels = m.LabelsMap
-		return nil
-	}
-
-	for key, value := range m.LabelsMap {
-		// set labels for pods created by deployment
-		deployment.Spec.Template.Labels[key] = value
-	}
-
+	// merging the labels
+	// priority: ServingRuntimePodSpec > LabelsMap > whatever in the deployment template
+	deployment.Spec.Template.Labels = Union(deployment.Spec.Template.Labels, m.LabelsMap, m.SRSpec.ServingRuntimePodSpec.Labels)
 	return nil
 }
diff --git a/controllers/modelmesh/runtime_test.go b/controllers/modelmesh/runtime_test.go
@@ -375,6 +375,52 @@ func TestConfigureRuntimeAnnotations(t *testing.T) {
 		assert.Nil(t, err)
 		assert.Empty(t, deploy.Spec.Template.Annotations)
 	})
+
+	t.Run("success-set-annotations-from-servingruntime-spec", func(t *testing.T) {
+		deploy := &appsv1.Deployment{}
+		sr := &kserveapi.ServingRuntime{
+			Spec: kserveapi.ServingRuntimeSpec{
+				ServingRuntimePodSpec: kserveapi.ServingRuntimePodSpec{
+					Annotations: map[string]string{
+						"foo":            "bar",
+						"network-policy": "allow-egress",
+					},
+				},
+			},
+		}
+
+		m := Deployment{Owner: sr, AnnotationsMap: map[string]string{}, SRSpec: &sr.Spec}
+
+		err := m.configureRuntimePodSpecAnnotations(deploy)
+		assert.Nil(t, err)
+		assert.Equal(t, deploy.Spec.Template.Annotations["foo"], "bar")
+		assert.Equal(t, deploy.Spec.Template.Annotations["network-policy"], "allow-egress")
+	})
+
+	t.Run("success-overwrite-annotations-from-servingruntime-spec", func(t *testing.T) {
+		deploy := &appsv1.Deployment{}
+		// annotations from user config
+		annotationsData := map[string]string{
+			"foo":            "bar",
+			"network-policy": "allow-egress",
+		}
+		sr := &kserveapi.ServingRuntime{
+			Spec: kserveapi.ServingRuntimeSpec{
+				ServingRuntimePodSpec: kserveapi.ServingRuntimePodSpec{
+					Annotations: map[string]string{
+						"network-policy": "overwritten-by-servingruntime",
+					},
+				},
+			},
+		}
+
+		m := Deployment{Owner: sr, AnnotationsMap: annotationsData, SRSpec: &sr.Spec}
+
+		err := m.configureRuntimePodSpecAnnotations(deploy)
+		assert.Nil(t, err)
+		assert.Equal(t, deploy.Spec.Template.Annotations["foo"], "bar")
+		assert.Equal(t, deploy.Spec.Template.Annotations["network-policy"], "overwritten-by-servingruntime")
+	})
 }
 
 func TestConfigureRuntimeLabels(t *testing.T) {
@@ -407,4 +453,54 @@ func TestConfigureRuntimeLabels(t *testing.T) {
 		assert.Nil(t, err)
 		assert.Empty(t, deploy.Spec.Template.Labels)
 	})
+
+	t.Run("success-set-labels-from-servingruntime-spec", func(t *testing.T) {
+		deploy := &appsv1.Deployment{}
+		sr := &kserveapi.ServingRuntime{
+			Spec: kserveapi.ServingRuntimeSpec{
+				ServingRuntimePodSpec: kserveapi.ServingRuntimePodSpec{
+					Labels: map[string]string{
+						"foo":            "bar",
+						"network-policy": "allow-egress",
+						"cp4s-internet":  "allow",
+					},
+				},
+			},
+		}
+
+		m := Deployment{Owner: sr, LabelsMap: map[string]string{}, SRSpec: &sr.Spec}
+
+		err := m.configureRuntimePodSpecLabels(deploy)
+		assert.Nil(t, err)
+		assert.Equal(t, deploy.Spec.Template.Labels["foo"], "bar")
+		assert.Equal(t, deploy.Spec.Template.Labels["network-policy"], "allow-egress")
+		assert.Equal(t, deploy.Spec.Template.Labels["cp4s-internet"], "allow")
+	})
+
+	t.Run("success-overwrite-labels-from-servingruntime-spec", func(t *testing.T) {
+		deploy := &appsv1.Deployment{}
+		// labels from user config
+		labelData := map[string]string{
+			"foo":            "bar",
+			"network-policy": "allow-egress",
+			"cp4s-internet":  "allow",
+		}
+		sr := &kserveapi.ServingRuntime{
+			Spec: kserveapi.ServingRuntimeSpec{
+				ServingRuntimePodSpec: kserveapi.ServingRuntimePodSpec{
+					Labels: map[string]string{
+						"network-policy": "overwritten-by-servingruntime",
+					},
+				},
+			},
+		}
+
+		m := Deployment{Owner: sr, LabelsMap: labelData, SRSpec: &sr.Spec}
+
+		err := m.configureRuntimePodSpecLabels(deploy)
+		assert.Nil(t, err)
+		assert.Equal(t, deploy.Spec.Template.Labels["foo"], "bar")
+		assert.Equal(t, deploy.Spec.Template.Labels["network-policy"], "overwritten-by-servingruntime")
+		assert.Equal(t, deploy.Spec.Template.Labels["cp4s-internet"], "allow")
+	})
 }
diff --git a/controllers/modelmesh/util.go b/controllers/modelmesh/util.go
@@ -20,12 +20,17 @@ import (
 	"strings"
 
 	kserveapi "github.com/kserve/kserve/pkg/apis/serving/v1alpha1"
+	kserveutils "github.com/kserve/kserve/pkg/utils"
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"sigs.k8s.io/yaml"
 )
 
+var (
+	Union = kserveutils.Union
+)
+
 // Find a container by name in the given deployment, returns (-1, nil) if not found
 func findContainer(name string, deployment *appsv1.Deployment) (index int, container *corev1.Container) {
 	for i := range deployment.Spec.Template.Spec.Containers {