From 4305c07ec6f4a7032bafe20a0fab7f092a68bdec Mon Sep 17 00:00:00 2001 From: Yu-Lin Chen Date: Mon, 11 Sep 2023 11:21:26 +0800 Subject: [PATCH] [YUNIKORN-1351] alway prefer annotations above labels --- pkg/admission/admission_controller.go | 66 +++- pkg/admission/admission_controller_test.go | 360 ++++++++++++------- pkg/admission/metadata/usergroup.go | 5 +- pkg/admission/metadata/usergroup_test.go | 7 +- pkg/admission/util.go | 110 ++++-- pkg/admission/util_test.go | 394 +++++++++++++++------ pkg/cache/metadata.go | 3 +- pkg/cache/placeholder.go | 5 +- pkg/cache/placeholder_test.go | 20 +- pkg/common/constants/constants.go | 7 +- pkg/common/utils/utils.go | 16 +- pkg/common/utils/utils_test.go | 10 +- 12 files changed, 718 insertions(+), 285 deletions(-) diff --git a/pkg/admission/admission_controller.go b/pkg/admission/admission_controller.go index 3b058da6b..5c3b5f405 100644 --- a/pkg/admission/admission_controller.go +++ b/pkg/admission/admission_controller.go @@ -196,7 +196,7 @@ func (c *AdmissionController) processPod(req *admissionv1.AdmissionRequest, name patch = updateSchedulerName(patch) if c.shouldLabelNamespace(namespace) { - patch = c.updateLabels(namespace, &pod, patch) + patch = c.updateApplicationInfo(namespace, &pod, patch) patch = c.updatePreemptionInfo(&pod, patch) } else { patch = disableYuniKorn(namespace, &pod, patch) @@ -369,8 +369,8 @@ func (c *AdmissionController) checkUserInfoAnnotation(getAnnotation func() (stri func updateSchedulerName(patch []common.PatchOperation) []common.PatchOperation { log.Log(log.Admission).Info("updating scheduler name") return append(patch, common.PatchOperation{ - Op: "add", - Path: "/spec/schedulerName", + Op: constants.AddPatchOp, + Path: constants.SchedulerNamePatchPath, Value: constants.SchedulerName, }) } @@ -397,7 +397,7 @@ func (c *AdmissionController) updatePreemptionInfo(pod *v1.Pod, patch []common.P // check for an existing patch on annotations and update it for _, p := range patch { - if p.Op == "add" && p.Path == "/metadata/annotations" { + if p.Op == constants.AddPatchOp && p.Path == constants.AnnotationPatchPath { if annotations, ok := p.Value.(map[string]string); ok { annotations[constants.AnnotationAllowPreemption] = value return patch @@ -407,29 +407,61 @@ func (c *AdmissionController) updatePreemptionInfo(pod *v1.Pod, patch []common.P result := updatePodAnnotation(pod, constants.AnnotationAllowPreemption, value) patch = append(patch, common.PatchOperation{ - Op: "add", - Path: "/metadata/annotations", + Op: constants.AddPatchOp, + Path: constants.AnnotationPatchPath, Value: result, }) return patch } -func (c *AdmissionController) updateLabels(namespace string, pod *v1.Pod, patch []common.PatchOperation) []common.PatchOperation { - log.Log(log.Admission).Info("updating pod labels", +func (c *AdmissionController) updateApplicationInfo(namespace string, pod *v1.Pod, patch []common.PatchOperation) []common.PatchOperation { + log.Log(log.Admission).Info("updating pod application labels and annotations", zap.String("podName", pod.Name), zap.String("generateName", pod.GenerateName), zap.String("namespace", namespace), - zap.Any("labels", pod.Labels)) + zap.Any("labels", pod.Labels), + zap.Any("annotations", pod.Annotations)) - result := updatePodLabel(pod, namespace, c.conf.GetGenerateUniqueAppIds(), c.conf.GetDefaultQueueName()) + newLabels, newAnnotations := getNewApplicationInfo(pod, namespace, c.conf.GetGenerateUniqueAppIds(), c.conf.GetDefaultQueueName()) - patch = append(patch, common.PatchOperation{ - Op: "add", - Path: "/metadata/labels", - Value: result, - }) + patch = updatePatch(pod, newLabels, constants.LabelPatchPath, patch) + patch = updatePatch(pod, newAnnotations, constants.AnnotationPatchPath, patch) + + return patch +} + +func updatePatch(pod *v1.Pod, newValues map[string]string, patchPath string, patch []common.PatchOperation) []common.PatchOperation { + // if found an existing patch, add new values to it + for _, p := range patch { + if p.Op == constants.AddPatchOp && p.Path == patchPath { + if existingPatchValues, ok := p.Value.(map[string]string); ok { + for k, v := range newValues { + existingPatchValues[k] = v + } + return patch + } + } + } + // Add a new patch + if len(newValues) != 0 { + // combine new values with existing values in pod to create first patch + var existingPodValues map[string]string + if patchPath == constants.LabelPatchPath { + // newly add labels patch should include existing labels in pod + existingPodValues = pod.Labels + } else if patchPath == constants.AnnotationPatchPath { + // newly add annotations patch should include existing annotations in pod + existingPodValues = pod.Annotations + } + patchValue := utils.MergeMaps(existingPodValues, newValues) + patch = append(patch, common.PatchOperation{ + Op: constants.AddPatchOp, + Path: patchPath, + Value: patchValue, + }) + } return patch } @@ -442,8 +474,8 @@ func disableYuniKorn(namespace string, pod *v1.Pod, patch []common.PatchOperatio result := updatePodAnnotation(pod, constants.AnnotationIgnoreApplication, constants.True) patch = append(patch, common.PatchOperation{ - Op: "add", - Path: "/metadata/annotations", + Op: constants.AddPatchOp, + Path: constants.AnnotationPatchPath, Value: result, }) diff --git a/pkg/admission/admission_controller_test.go b/pkg/admission/admission_controller_test.go index 56f9b0b33..d1bf9a68e 100644 --- a/pkg/admission/admission_controller_test.go +++ b/pkg/admission/admission_controller_test.go @@ -51,9 +51,9 @@ const ( ) // nolint: funlen -func TestUpdateLabels(t *testing.T) { +func TestUpdateApplicationInfo(t *testing.T) { // verify when appId/queue are not given, - // we patch it correctly + // we patch it correctly. (Add appId/queueName/disableStateAware to labels and annotations) var patch []common.PatchOperation pod := &v1.Pod{ @@ -66,32 +66,30 @@ func TestUpdateLabels(t *testing.T) { Namespace: "default", UID: "7f5fd6c5d5", ResourceVersion: "10654", - Labels: map[string]string{ - "random": "random", - }, }, Spec: v1.PodSpec{}, Status: v1.PodStatus{}, } c := createAdmissionControllerForTest() - patch = c.updateLabels("default", pod, patch) - assert.Equal(t, len(patch), 1) - assert.Equal(t, patch[0].Op, "add") - assert.Equal(t, patch[0].Path, "/metadata/labels") - if updatedMap, ok := patch[0].Value.(map[string]string); ok { - assert.Equal(t, len(updatedMap), 4) - assert.Equal(t, updatedMap["random"], "random") - assert.Equal(t, updatedMap["queue"], "root.default") - assert.Equal(t, updatedMap["disableStateAware"], "true") - assert.Equal(t, strings.HasPrefix(updatedMap["applicationId"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("patch info content is not as expected") - } + patch = c.updateApplicationInfo("default", pod, patch) - // verify if applicationId is given in the labels, - // we won't modify it + labelsInPatch := fetchPatchValues(constants.AddPatchOp, constants.LabelPatchPath, patch) + annotationsInPatch := fetchPatchValues(constants.AddPatchOp, constants.AnnotationPatchPath, patch) + + assert.Equal(t, len(patch), 2) + assert.Equal(t, len(labelsInPatch), 3) + assert.Equal(t, labelsInPatch["queue"], "root.default") + assert.Equal(t, labelsInPatch["disableStateAware"], "true") + assert.Equal(t, strings.HasPrefix(labelsInPatch["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(annotationsInPatch), 3) + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, strings.HasPrefix(annotationsInPatch["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) + + // verify if applicationId is given in the pod's labels, + // we won't modify it in admission controller, and will patch the value to annotations patch = make([]common.PatchOperation, 0) pod = &v1.Pod{ @@ -105,29 +103,28 @@ func TestUpdateLabels(t *testing.T) { UID: "7f5fd6c5d5", ResourceVersion: "10654", Labels: map[string]string{ - "random": "random", "applicationId": "app-0001", }, }, Spec: v1.PodSpec{}, Status: v1.PodStatus{}, } - patch = c.updateLabels("default", pod, patch) - - assert.Equal(t, len(patch), 1) - assert.Equal(t, patch[0].Op, "add") - assert.Equal(t, patch[0].Path, "/metadata/labels") - if updatedMap, ok := patch[0].Value.(map[string]string); ok { - assert.Equal(t, len(updatedMap), 3) - assert.Equal(t, updatedMap["random"], "random") - assert.Equal(t, updatedMap["queue"], "root.default") - assert.Equal(t, updatedMap["applicationId"], "app-0001") - } else { - t.Fatal("patch info content is not as expected") - } - - // verify if queue is given in the labels, - // we won't modify it + patch = c.updateApplicationInfo("default", pod, patch) + + labelsInPatch = fetchPatchValues(constants.AddPatchOp, constants.LabelPatchPath, patch) + annotationsInPatch = fetchPatchValues(constants.AddPatchOp, constants.AnnotationPatchPath, patch) + + assert.Equal(t, len(patch), 2) + assert.Equal(t, len(labelsInPatch), 2) + assert.Equal(t, labelsInPatch["queue"], "root.default") + assert.Equal(t, labelsInPatch["applicationId"], "app-0001") + assert.Equal(t, len(patch), 2) + assert.Equal(t, len(annotationsInPatch), 2) + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/app-id"], "app-0001") + + // verify if queue is given in the pod's labels, + // we won't modify it in admission controller, and will patch the value to annotations patch = make([]common.PatchOperation, 0) pod = &v1.Pod{ @@ -141,31 +138,30 @@ func TestUpdateLabels(t *testing.T) { UID: "7f5fd6c5d5", ResourceVersion: "10654", Labels: map[string]string{ - "random": "random", - "queue": "root.abc", + "queue": "root.abc", }, }, Spec: v1.PodSpec{}, Status: v1.PodStatus{}, } - patch = c.updateLabels("default", pod, patch) + patch = c.updateApplicationInfo("default", pod, patch) - assert.Equal(t, len(patch), 1) - assert.Equal(t, patch[0].Op, "add") - assert.Equal(t, patch[0].Path, "/metadata/labels") - if updatedMap, ok := patch[0].Value.(map[string]string); ok { - assert.Equal(t, len(updatedMap), 4) - assert.Equal(t, updatedMap["random"], "random") - assert.Equal(t, updatedMap["queue"], "root.abc") - assert.Equal(t, updatedMap["disableStateAware"], "true") - assert.Equal(t, strings.HasPrefix(updatedMap["applicationId"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("patch info content is not as expected") - } + labelsInPatch = fetchPatchValues(constants.AddPatchOp, constants.LabelPatchPath, patch) + annotationsInPatch = fetchPatchValues(constants.AddPatchOp, constants.AnnotationPatchPath, patch) + + assert.Equal(t, len(patch), 2) + assert.Equal(t, len(labelsInPatch), 3) + assert.Equal(t, labelsInPatch["queue"], "root.abc") + assert.Equal(t, labelsInPatch["disableStateAware"], "true") + assert.Equal(t, strings.HasPrefix(labelsInPatch["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(annotationsInPatch), 3) + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/queue"], "root.abc") + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, strings.HasPrefix(annotationsInPatch["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) - // namespace might be empty - // labels might be empty + // verify if applicationId is given in the pod's annotations, + // we won't modify it in admission controller, and will patch the value to labels patch = make([]common.PatchOperation, 0) pod = &v1.Pod{ @@ -175,90 +171,192 @@ func TestUpdateLabels(t *testing.T) { }, ObjectMeta: metav1.ObjectMeta{ Name: "a-test-pod", + Namespace: "default", UID: "7f5fd6c5d5", ResourceVersion: "10654", + Annotations: map[string]string{ + "yunikorn.apache.org/app-id": "app-0001", + }, }, Spec: v1.PodSpec{}, Status: v1.PodStatus{}, } - - patch = c.updateLabels("default", pod, patch) - - assert.Equal(t, len(patch), 1) - assert.Equal(t, patch[0].Op, "add") - assert.Equal(t, patch[0].Path, "/metadata/labels") - if updatedMap, ok := patch[0].Value.(map[string]string); ok { - assert.Equal(t, len(updatedMap), 3) - assert.Equal(t, updatedMap["queue"], "root.default") - assert.Equal(t, updatedMap["disableStateAware"], "true") - assert.Equal(t, strings.HasPrefix(updatedMap["applicationId"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("patch info content is not as expected") - } - - // pod name might be empty, it can comes from generatedName + patch = c.updateApplicationInfo("default", pod, patch) + + labelsInPatch = fetchPatchValues(constants.AddPatchOp, constants.LabelPatchPath, patch) + annotationsInPatch = fetchPatchValues(constants.AddPatchOp, constants.AnnotationPatchPath, patch) + + assert.Equal(t, len(patch), 2) + assert.Equal(t, len(labelsInPatch), 2) + assert.Equal(t, labelsInPatch["queue"], "root.default") + assert.Equal(t, labelsInPatch["applicationId"], "app-0001") + assert.Equal(t, len(patch), 2) + assert.Equal(t, len(annotationsInPatch), 2) + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/app-id"], "app-0001") + + // if both annotation and labels patch existing + // we won't create new patch and will add new labels/annotation to the existing patch. patch = make([]common.PatchOperation, 0) - - pod = &v1.Pod{ - TypeMeta: metav1.TypeMeta{ - Kind: "Pod", - APIVersion: "v1", + patch = append(patch, common.PatchOperation{ + Op: constants.AddPatchOp, + Path: constants.LabelPatchPath, + Value: map[string]string{ + "existingLabelKey": "existingLabelValue", }, - ObjectMeta: metav1.ObjectMeta{ - GenerateName: "some-pod-", + }) + patch = append(patch, common.PatchOperation{ + Op: constants.AddPatchOp, + Path: constants.AnnotationPatchPath, + Value: map[string]string{ + "existingAnnotationKey": "existingAnnotationValue", }, - Spec: v1.PodSpec{}, - Status: v1.PodStatus{}, - } - - patch = c.updateLabels("default", pod, patch) + }) + pod = &v1.Pod{} + + patch = c.updateApplicationInfo("default", pod, patch) + + labelsInPatch = fetchPatchValues(constants.AddPatchOp, constants.LabelPatchPath, patch) + annotationsInPatch = fetchPatchValues(constants.AddPatchOp, constants.AnnotationPatchPath, patch) + + assert.Equal(t, len(labelsInPatch), 4) + assert.Equal(t, labelsInPatch["existingLabelKey"], "existingLabelValue") + assert.Equal(t, labelsInPatch["queue"], "root.default") + assert.Equal(t, labelsInPatch["disableStateAware"], "true") + assert.Equal(t, strings.HasPrefix(labelsInPatch["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(annotationsInPatch), 4) + assert.Equal(t, annotationsInPatch["existingAnnotationKey"], "existingAnnotationValue") + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, annotationsInPatch["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, strings.HasPrefix(annotationsInPatch["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) +} - assert.Equal(t, len(patch), 1) - assert.Equal(t, patch[0].Op, "add") - assert.Equal(t, patch[0].Path, "/metadata/labels") - if updatedMap, ok := patch[0].Value.(map[string]string); ok { - assert.Equal(t, len(updatedMap), 3) - assert.Equal(t, updatedMap["queue"], "root.default") - assert.Equal(t, updatedMap["disableStateAware"], "true") - assert.Equal(t, strings.HasPrefix(updatedMap["applicationId"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("patch info content is not as expected") +func TestUpdatePatch(t *testing.T) { + // dummy pod for following test cases + pod := &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{ + "label_key_in_pod": "label_value_in_pod", + }, + Annotations: map[string]string{ + "annotation_key_in_pod": "annotation_value_in_pod", + }, + }, } - // pod name and generate name could be both empty - patch = make([]common.PatchOperation, 0) - - pod = &v1.Pod{ - TypeMeta: metav1.TypeMeta{ - Kind: "Pod", - APIVersion: "v1", + testCases := []struct { + description string + pod *v1.Pod + patch []common.PatchOperation + patchPath string + newValues map[string]string + expectedPatchValues map[string]string + expectedPatchSize int + }{ + { + description: "Verify if no patch on labels in patch list, create new patch with pods's existing labels and put new labels into the patch", + pod: pod, + patch: make([]common.PatchOperation, 0), + patchPath: constants.LabelPatchPath, + newValues: map[string]string{ + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", + }, + expectedPatchValues: map[string]string{ + "label_key_in_pod": "label_value_in_pod", + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", + }, + expectedPatchSize: 1, + }, { + description: "Verify if no patch on annotations in patch list, create new patch with pods's existing annotation and put new annotation into the patch", + pod: pod, + patch: make([]common.PatchOperation, 0), + patchPath: constants.AnnotationPatchPath, + newValues: map[string]string{ + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", + }, + expectedPatchValues: map[string]string{ + "annotation_key_in_pod": "annotation_value_in_pod", + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", + }, + expectedPatchSize: 1, + }, { + description: "Verify if have patch on labels in patch list, won't crete new patch and will add new labels into the existing patch", + pod: pod, + patch: []common.PatchOperation{ + { + Op: constants.AddPatchOp, + Path: constants.LabelPatchPath, + Value: map[string]string{ + "label_key_in_patch": "label_value_in_patch", + }, + }, + }, + patchPath: constants.LabelPatchPath, + newValues: map[string]string{ + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", + }, + expectedPatchValues: map[string]string{ + "label_key_in_patch": "label_value_in_patch", + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", + }, + expectedPatchSize: 1, + }, { + description: "Verify if have patch on annotations in patch list, won't crete new patch and will add new annotations into the existing patch", + pod: pod, + patch: []common.PatchOperation{ + { + Op: constants.AddPatchOp, + Path: constants.AnnotationPatchPath, + Value: map[string]string{ + "annotation_key_in_patch": "annotation_value_in_patch", + }, + }, + }, + patchPath: constants.AnnotationPatchPath, + newValues: map[string]string{ + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", + }, + expectedPatchValues: map[string]string{ + "annotation_key_in_patch": "annotation_value_in_patch", + "new_key_1": "new_value_1", + "new_key_2": "new_value_2", + }, + expectedPatchSize: 1, }, - ObjectMeta: metav1.ObjectMeta{}, - Spec: v1.PodSpec{}, - Status: v1.PodStatus{}, } + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + patch := updatePatch(tc.pod, tc.newValues, tc.patchPath, tc.patch) + assert.DeepEqual(t, len(patch), tc.expectedPatchSize) + assert.DeepEqual(t, patch[0].Path, tc.patchPath) + assert.DeepEqual(t, patch[0].Value, tc.expectedPatchValues) + }) + } +} - patch = c.updateLabels("default", pod, patch) - - assert.Equal(t, len(patch), 1) - assert.Equal(t, patch[0].Op, "add") - assert.Equal(t, patch[0].Path, "/metadata/labels") - if updatedMap, ok := patch[0].Value.(map[string]string); ok { - assert.Equal(t, len(updatedMap), 3) - assert.Equal(t, updatedMap["queue"], "root.default") - assert.Equal(t, updatedMap["disableStateAware"], "true") - assert.Equal(t, strings.HasPrefix(updatedMap["applicationId"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("patch info content is not as expected") +func fetchPatchValues(op string, path string, patch []common.PatchOperation) map[string]string { + // fetch an existing patch by op and path + for _, p := range patch { + if p.Op == op && p.Path == path { + return p.Value.(map[string]string) + } } + return make(map[string]string) } func TestUpdateSchedulerName(t *testing.T) { var patch []common.PatchOperation patch = updateSchedulerName(patch) assert.Equal(t, len(patch), 1) - assert.Equal(t, patch[0].Op, "add") - assert.Equal(t, patch[0].Path, "/spec/schedulerName") + assert.Equal(t, patch[0].Op, constants.AddPatchOp) + assert.Equal(t, patch[0].Path, constants.SchedulerNamePatchPath) if name, ok := patch[0].Value.(string); ok { assert.Equal(t, name, constants.SchedulerName) } else { @@ -454,9 +552,9 @@ func TestMutate(t *testing.T) { resp = ac.mutate(req) assert.Check(t, resp.Allowed, "response not allowed for pod") assert.Equal(t, schedulerName(t, resp.Patch), "yunikorn", "yunikorn not set as scheduler for pod") - assert.Equal(t, labels(t, resp.Patch)["applicationId"], "yunikorn-default-autogen", "wrong applicationId label") - assert.Equal(t, labels(t, resp.Patch)["disableStateAware"], "true", "missing disableStateAware label") - assert.Equal(t, labels(t, resp.Patch)["queue"], "root.default", "incorrect queue name") + assert.Equal(t, annotations(t, resp.Patch)["yunikorn.apache.org/app-id"], "yunikorn-default-autogen", "wrong app-id annotation") + assert.Equal(t, annotations(t, resp.Patch)["yunikorn.apache.org/disable-state-aware"], "true", "missing disable-state-aware annotation") + assert.Equal(t, annotations(t, resp.Patch)["yunikorn.apache.org/queue"], "root.default", "incorrect queue name") // pod without applicationID pod = v1.Pod{ObjectMeta: metav1.ObjectMeta{ @@ -473,8 +571,8 @@ func TestMutate(t *testing.T) { resp = ac.mutate(req) assert.Check(t, resp.Allowed, "response not allowed for pod") assert.Equal(t, schedulerName(t, resp.Patch), "yunikorn", "yunikorn not set as scheduler for pod") - assert.Equal(t, labels(t, resp.Patch)["applicationId"], "yunikorn-test-ns-autogen", "wrong applicationId label") - assert.Equal(t, labels(t, resp.Patch)["disableStateAware"], "true", "missing disableStateAware label") + assert.Equal(t, annotations(t, resp.Patch)["yunikorn.apache.org/app-id"], "yunikorn-test-ns-autogen", "wrong app-id annotation") + assert.Equal(t, annotations(t, resp.Patch)["yunikorn.apache.org/disable-state-aware"], "true", "missing disable-state-aware annotation") // pod with applicationId pod.ObjectMeta.Labels = map[string]string{"applicationId": "test-app"} @@ -484,7 +582,7 @@ func TestMutate(t *testing.T) { resp = ac.mutate(req) assert.Check(t, resp.Allowed, "response not allowed for pod") assert.Equal(t, schedulerName(t, resp.Patch), "yunikorn", "yunikorn not set as scheduler for pod") - assert.Equal(t, labels(t, resp.Patch)["applicationId"], "test-app", "wrong applicationId label") + assert.Equal(t, annotations(t, resp.Patch)["yunikorn.apache.org/app-id"], "test-app", "wrong app-id annotation") // pod in bypassed namespace pod = v1.Pod{ObjectMeta: metav1.ObjectMeta{ @@ -517,7 +615,9 @@ func TestMutate(t *testing.T) { resp = ac.mutate(req) assert.Check(t, resp.Allowed, "response not allowed for nolabel pod") assert.Equal(t, schedulerName(t, resp.Patch), "yunikorn", "yunikorn not set as scheduler for nolabel pod") - assert.Equal(t, len(labels(t, resp.Patch)), 0, "non-empty labels for nolabel pod") + assert.Equal(t, annotations(t, resp.Patch)["yunikorn.apache.org/app-id"], nil, "non-empty app-id annotation for nolabel pod") + assert.Equal(t, annotations(t, resp.Patch)["yunikorn.apache.org/disable-state-aware"], nil, "non-empty disable-state-aware annotation for nolabel pod") + assert.Equal(t, annotations(t, resp.Patch)["yunikorn.apache.org/queue"], nil, "non-empty queue annotation for nolabel pod") // unknown object type pod = v1.Pod{ObjectMeta: metav1.ObjectMeta{ @@ -842,7 +942,7 @@ func parsePatch(t *testing.T, patch []byte) []common.PatchOperation { func schedulerName(t *testing.T, patch []byte) string { ops := parsePatch(t, patch) for _, op := range ops { - if op.Path == "/spec/schedulerName" { + if op.Path == constants.SchedulerNamePatchPath { return op.Value.(string) } } @@ -852,7 +952,17 @@ func schedulerName(t *testing.T, patch []byte) string { func labels(t *testing.T, patch []byte) map[string]interface{} { ops := parsePatch(t, patch) for _, op := range ops { - if op.Path == "/metadata/labels" { + if op.Path == constants.LabelPatchPath { + return op.Value.(map[string]interface{}) + } + } + return make(map[string]interface{}) +} + +func annotations(t *testing.T, patch []byte) map[string]interface{} { + ops := parsePatch(t, patch) + for _, op := range ops { + if op.Path == constants.AnnotationPatchPath { return op.Value.(map[string]interface{}) } } diff --git a/pkg/admission/metadata/usergroup.go b/pkg/admission/metadata/usergroup.go index a93ad3042..bd8e43a60 100644 --- a/pkg/admission/metadata/usergroup.go +++ b/pkg/admission/metadata/usergroup.go @@ -29,6 +29,7 @@ import ( "github.com/apache/yunikorn-k8shim/pkg/admission/common" "github.com/apache/yunikorn-k8shim/pkg/admission/conf" + "github.com/apache/yunikorn-k8shim/pkg/common/constants" "github.com/apache/yunikorn-k8shim/pkg/log" ) @@ -127,7 +128,7 @@ func (u *UserGroupAnnotationHandler) GetPatchForWorkload(req *admissionv1.Admiss } func (u *UserGroupAnnotationHandler) GetPatchForPod(annotations map[string]string, user string, groups []string) (*common.PatchOperation, error) { - patchOp, err := u.getPatchOperation(annotations, "/metadata/annotations", user, groups) + patchOp, err := u.getPatchOperation(annotations, constants.AnnotationPatchPath, user, groups) if err != nil { return nil, err } @@ -152,7 +153,7 @@ func (u *UserGroupAnnotationHandler) getPatchOperation(annotations map[string]st newAnnotations[common.UserInfoAnnotation] = string(jsonBytes) return &common.PatchOperation{ - Op: "add", + Op: constants.AddPatchOp, Path: path, Value: newAnnotations, }, nil diff --git a/pkg/admission/metadata/usergroup_test.go b/pkg/admission/metadata/usergroup_test.go index 5dc4a9125..89d6419dc 100644 --- a/pkg/admission/metadata/usergroup_test.go +++ b/pkg/admission/metadata/usergroup_test.go @@ -36,6 +36,7 @@ import ( "github.com/apache/yunikorn-k8shim/pkg/admission/common" "github.com/apache/yunikorn-k8shim/pkg/admission/conf" + "github.com/apache/yunikorn-k8shim/pkg/common/constants" ) const ( @@ -163,7 +164,7 @@ func TestGetPatchForWorkload(t *testing.T) { assert.NilError(t, err) assert.Equal(t, 1, len(patch)) patchOp := patch[0] - assert.Equal(t, patchOp.Op, "add") + assert.Equal(t, patchOp.Op, constants.AddPatchOp) assert.Equal(t, patchOp.Path, testCase.path) verifyUserGroupAnnotation(t, patchOp.Value) }) @@ -175,8 +176,8 @@ func TestGetPatchForPod(t *testing.T) { patchOp, err := ah.GetPatchForPod(annotation, "yunikorn", []string{"users", "dev"}) assert.NilError(t, err) assert.Assert(t, patchOp != nil) - assert.Equal(t, patchOp.Op, "add") - assert.Equal(t, patchOp.Path, "/metadata/annotations") + assert.Equal(t, patchOp.Op, constants.AddPatchOp) + assert.Equal(t, patchOp.Path, constants.AnnotationPatchPath) verifyUserGroupAnnotation(t, patchOp.Value) } diff --git a/pkg/admission/util.go b/pkg/admission/util.go index 5c9576482..bf5a1d42f 100644 --- a/pkg/admission/util.go +++ b/pkg/admission/util.go @@ -29,42 +29,86 @@ import ( "github.com/apache/yunikorn-k8shim/pkg/log" ) -func updatePodLabel(pod *v1.Pod, namespace string, generateUniqueAppIds bool, defaultQueueName string) map[string]string { - existingLabels := pod.Labels - result := make(map[string]string) - for k, v := range existingLabels { - result[k] = v +func getNewApplicationInfo(pod *v1.Pod, namespace string, generateUniqueAppIds bool, defaultQueueName string) (map[string]string, map[string]string) { + newLabels := make(map[string]string) + newAnnotations := make(map[string]string) + + appID, isAppIdFromLabel := getApplicationIDFromPod(pod) + disableStateAware, isDisableStateAwareFromLabel := getDisableStateAwareFromPod(pod) + + // for undefined configuration, am_conf will add 'root.default' as default queue name + // if a custom name is configured for default queue, it will be used instead of root.default + queueName, isQueueNameFromLabel := utils.GetQueueNameFromPod(pod, defaultQueueName) + + if isAppIdFromLabel || isDisableStateAwareFromLabel || isQueueNameFromLabel { + log.Log(log.Admission).Warn("Pod contains label for Yunikorn configuration. This is deprecated and will be ignored in a future release. Please migrate to annotation-based configuration.", + zap.Bool("applicationId from label", isAppIdFromLabel), + zap.Bool("disableStateAware from label", isDisableStateAwareFromLabel), + zap.Bool("queue name from label", isQueueNameFromLabel)) } - sparkAppID := utils.GetPodLabelValue(pod, constants.SparkLabelAppID) - appID := utils.GetPodLabelValue(pod, constants.LabelApplicationID) - if sparkAppID == "" && appID == "" { + if appID == "" { // if app id not exist, generate one // for each namespace, we group unnamed pods to one single app - if GenerateUniqueAppId is not set // if GenerateUniqueAppId: // application ID convention: ${NAMESPACE}-${POD_UID} // else // application ID convention: ${AUTO_GEN_PREFIX}-${NAMESPACE}-${AUTO_GEN_SUFFIX} - generatedID := utils.GenerateApplicationID(namespace, generateUniqueAppIds, string(pod.UID)) - result[constants.LabelApplicationID] = generatedID + appID = utils.GenerateApplicationID(namespace, generateUniqueAppIds, string(pod.UID)) // if we generate an app ID, disable state-aware scheduling for this app - if _, ok := existingLabels[constants.LabelDisableStateAware]; !ok { - result[constants.LabelDisableStateAware] = "true" + // skip it if disableStateAware has already been set in the pod + if disableStateAware == "" { + disableStateAware = "true" } } - // if existing label exist, it takes priority over everything else - if _, ok := existingLabels[constants.LabelQueueName]; !ok { - // if defaultQueueName is "", skip adding default queue name to the pod labels - if defaultQueueName != "" { - // for undefined configuration, am_conf will add 'root.default' to retain existing behavior - // if a custom name is configured for default queue, it will be used instead of root.default - result[constants.LabelQueueName] = defaultQueueName - } + // we're looking forward to deprecate the labels and move everything to annotations + // but for backforward compatibility, we still add to labels + newLabels = updateLabelIfNotPresentInPod(pod, newLabels, constants.LabelApplicationID, appID) + newAnnotations = updateAnnotationIfNotPresentInPod(pod, newAnnotations, constants.AnnotationApplicationID, appID) + + // skip adding empty disableStateAware to the pod + if disableStateAware != "" { + newLabels = updateLabelIfNotPresentInPod(pod, newLabels, constants.LabelDisableStateAware, disableStateAware) + newAnnotations = updateAnnotationIfNotPresentInPod(pod, newAnnotations, constants.AnnotationDisableStateAware, disableStateAware) } - return result + // skip adding empty queue name to the pod + if queueName != "" { + newLabels = updateLabelIfNotPresentInPod(pod, newLabels, constants.LabelQueueName, queueName) + newAnnotations = updateAnnotationIfNotPresentInPod(pod, newAnnotations, constants.AnnotationQueueName, queueName) + } + + return newLabels, newAnnotations +} + +func updateLabelIfNotPresentInPod(pod *v1.Pod, labelMap map[string]string, label string, value string) map[string]string { + existingLabels := pod.Labels + + if existingLabels == nil { + labelMap[label] = value + } + + // skip update label if it has already been set in the pod + if _, ok := existingLabels[label]; !ok { + labelMap[label] = value + } + return labelMap +} + +func updateAnnotationIfNotPresentInPod(pod *v1.Pod, annotationMap map[string]string, annotation string, value string) map[string]string { + existingAnnotations := pod.Annotations + + if existingAnnotations == nil { + annotationMap[annotation] = value + } + + // skip update annotation if it has already been set in the pod + if _, ok := existingAnnotations[annotation]; !ok { + annotationMap[annotation] = value + } + return annotationMap } func updatePodAnnotation(pod *v1.Pod, key string, value string) map[string]string { @@ -84,3 +128,27 @@ func convert2Namespace(obj interface{}) *v1.Namespace { log.Log(log.AdmissionUtils).Warn("cannot convert to *v1.Namespace", zap.Stringer("type", reflect.TypeOf(obj))) return nil } + +func getApplicationIDFromPod(pod *v1.Pod) (value string, isFromLabel bool) { + // if existing annotation exist, it takes priority over everything else + if value := utils.GetPodAnnotationValue(pod, constants.AnnotationApplicationID); value != "" { + return value, false + } else if value := utils.GetPodLabelValue(pod, constants.LabelApplicationID); value != "" { + return value, true + } + // application ID can be defined in Spark label + if value := utils.GetPodLabelValue(pod, constants.SparkLabelAppID); value != "" { + return value, true + } + return "", false +} + +func getDisableStateAwareFromPod(pod *v1.Pod) (value string, isFromLabel bool) { + // if existing annotation exist, it takes priority over everything else + if value := utils.GetPodAnnotationValue(pod, constants.AnnotationDisableStateAware); value != "" { + return value, false + } else if value := utils.GetPodLabelValue(pod, constants.LabelDisableStateAware); value != "" { + return value, true + } + return "", false +} diff --git a/pkg/admission/util_test.go b/pkg/admission/util_test.go index 1b768631c..52a567d2b 100644 --- a/pkg/admission/util_test.go +++ b/pkg/admission/util_test.go @@ -67,13 +67,22 @@ func createTestingPodWithMeta() *v1.Pod { return pod } -func createTestingPodWithAppId() *v1.Pod { +func createTestingPodWithLabelAppId() *v1.Pod { pod := createTestingPodWithMeta() pod.ObjectMeta.Labels["applicationId"] = "app-0001" return pod } +func createTestingPodWithAnnotationAppId() *v1.Pod { + pod := createTestingPodWithMeta() + pod.ObjectMeta.Annotations = map[string]string{ + "yunikorn.apache.org/app-id": "app-0001", + } + + return pod +} + func createTestingPodWithGenerateName() *v1.Pod { pod := createMinimalTestingPod() pod.ObjectMeta.GenerateName = "some-pod-" @@ -81,13 +90,31 @@ func createTestingPodWithGenerateName() *v1.Pod { return pod } -func createTestingPodWithQueue() *v1.Pod { +func createTestingPodWithLabelQueue() *v1.Pod { pod := createTestingPodWithMeta() pod.ObjectMeta.Labels["queue"] = "root.abc" return pod } +func createTestingPodWithAnnotationQueue() *v1.Pod { + pod := createTestingPodWithMeta() + pod.ObjectMeta.Annotations = map[string]string{ + "yunikorn.apache.org/queue": "root.abc", + } + + return pod +} + +func createTestingPodWithLabelEnableStateAware() *v1.Pod { + pod := createTestingPodWithMeta() + pod.ObjectMeta.Labels = map[string]string{ + constants.LabelDisableStateAware: "false", + } + + return pod +} + func createTestingPodNoNamespaceAndLabels() *v1.Pod { pod := createMinimalTestingPod() pod.ObjectMeta = @@ -99,131 +126,294 @@ func createTestingPodNoNamespaceAndLabels() *v1.Pod { return pod } -func TestUpdatePodLabelForAdmissionController(t *testing.T) { - // verify when appId/queue are not given, +// nolint: funlen +func TestGetNewApplicationInfo(t *testing.T) { + // verify when appId/queue are not given, will generate it pod := createTestingPodWithMeta() + newLabels, newAnnotations := getNewApplicationInfo(pod, "default", false, "root.default") - if result := updatePodLabel(pod, "default", false, "root.default"); result != nil { - assert.Equal(t, len(result), 4) - assert.Equal(t, result["random"], "random") - assert.Equal(t, result["queue"], "root.default") - assert.Equal(t, result["disableStateAware"], "true") - assert.Equal(t, strings.HasPrefix(result["applicationId"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("UpdatePodLabelForAdmissionController is not as expected") - } + assert.Equal(t, len(newLabels), 3) + assert.Equal(t, newLabels["queue"], "root.default") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Equal(t, strings.HasPrefix(newLabels["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(newAnnotations), 3) + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, strings.HasPrefix(newAnnotations["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) // verify if applicationId is given in the labels, - // we won't modify it - pod = createTestingPodWithAppId() - - if result := updatePodLabel(pod, "default", false, "root.default"); result != nil { - assert.Equal(t, len(result), 3) - assert.Equal(t, result["random"], "random") - assert.Equal(t, result["queue"], "root.default") - assert.Equal(t, result["applicationId"], "app-0001") - } else { - t.Fatal("UpdatePodLabelForAdmissionController is not as expected") - } + // will copy applicationId to annotation and won't change existing value + pod = createTestingPodWithLabelAppId() + newLabels, newAnnotations = getNewApplicationInfo(pod, "default", false, "root.default") + + assert.Equal(t, len(newLabels), 1) + assert.Equal(t, newLabels["queue"], "root.default") + assert.Equal(t, newLabels["disableStateAware"], "") + assert.Equal(t, newLabels["applicationId"], "") + assert.Equal(t, len(newAnnotations), 2) + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, newAnnotations["yunikorn.apache.org/app-id"], "app-0001") // verify if queue is given in the labels, - // we won't modify it - pod = createTestingPodWithQueue() - if result := updatePodLabel(pod, "default", false, "root.default"); result != nil { - assert.Equal(t, len(result), 4) - assert.Equal(t, result["random"], "random") - assert.Equal(t, result["queue"], "root.abc") - assert.Equal(t, result["disableStateAware"], "true") - assert.Equal(t, strings.HasPrefix(result["applicationId"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("UpdatePodLabelForAdmissionControllert is not as expected") - } - - // namespace might be empty - // labels might be empty - pod = createTestingPodNoNamespaceAndLabels() - - if result := updatePodLabel(pod, "default", false, "root.default"); result != nil { - assert.Equal(t, len(result), 3) - assert.Equal(t, result["queue"], "root.default") - assert.Equal(t, result["disableStateAware"], "true") - assert.Equal(t, strings.HasPrefix(result["applicationId"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("UpdatePodLabelForAdmissionController is not as expected") - } - - // pod name might be empty, it can comes from generatedName + // will copy queue to annotation and won't change existing value + pod = createTestingPodWithLabelQueue() + newLabels, newAnnotations = getNewApplicationInfo(pod, "default", false, "root.default") + + assert.Equal(t, len(newLabels), 2) + assert.Equal(t, newLabels["queue"], "") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Equal(t, strings.HasPrefix(newLabels["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(newAnnotations), 3) + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "root.abc") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, strings.HasPrefix(newAnnotations["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) + + // verify if applicationId is given in the annotation, + // will copy to label and won't change existing value + pod = createTestingPodWithAnnotationAppId() + newLabels, newAnnotations = getNewApplicationInfo(pod, "default", false, "root.default") + assert.Equal(t, len(newLabels), 2) + assert.Equal(t, newLabels["applicationId"], "app-0001") + assert.Equal(t, newLabels["queue"], "root.default") + assert.Equal(t, newLabels["disableStateAware"], "") + assert.Equal(t, len(newAnnotations), 1) + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, newAnnotations["yunikorn.apache.org/app-id"], "") + + // verify if queue is given in the annotation + // will copy to labels and won't change existing value + pod = createTestingPodWithAnnotationQueue() + newLabels, newAnnotations = getNewApplicationInfo(pod, "default", false, "root.default") + assert.Equal(t, len(newLabels), 3) + assert.Equal(t, newLabels["queue"], "root.abc") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Equal(t, strings.HasPrefix(newLabels["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(newAnnotations), 2) + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, strings.HasPrefix(newAnnotations["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) + + // verify if state aware label set in pod, + // won't change state aware setting and will copy to annotation + pod = createTestingPodWithLabelEnableStateAware() + newLabels, newAnnotations = getNewApplicationInfo(pod, "default", false, "root.default") + + assert.Equal(t, len(newLabels), 2) + assert.Equal(t, newLabels["queue"], "root.default") + assert.Equal(t, newLabels["disableStateAware"], "") + assert.Equal(t, strings.HasPrefix(newLabels["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(newAnnotations), 3) + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "false") + assert.Equal(t, strings.HasPrefix(newAnnotations["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) + + // verify empty pod info and default queue name is "" + // won't add queue name to pod + pod = createTestingPodWithMeta() + newLabels, newAnnotations = getNewApplicationInfo(pod, "default", false, "") + + assert.Equal(t, len(newLabels), 2) + assert.Equal(t, newLabels["queue"], "") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Equal(t, strings.HasPrefix(newLabels["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(newAnnotations), 2) + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, strings.HasPrefix(newAnnotations["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) + + // pod name might be empty, applicationId can come from generatedName pod = createTestingPodWithGenerateName() - if result := updatePodLabel(pod, "default", false, "root.default"); result != nil { - assert.Equal(t, len(result), 3) - assert.Equal(t, result["queue"], "root.default") - assert.Equal(t, result["disableStateAware"], "true") - assert.Equal(t, strings.HasPrefix(result["applicationId"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("UpdatePodLabelForAdmissionController is not as expected") - } + newLabels, newAnnotations = getNewApplicationInfo(pod, "default", false, "root.default") + assert.Equal(t, strings.HasPrefix(newLabels["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, strings.HasPrefix(newAnnotations["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) pod = createMinimalTestingPod() - if result := updatePodLabel(pod, "default", false, "root.default"); result != nil { - assert.Equal(t, len(result), 3) - assert.Equal(t, result["queue"], "root.default") - assert.Equal(t, result["disableStateAware"], "true") - assert.Equal(t, strings.HasPrefix(result["applicationId"], constants.AutoGenAppPrefix), true) - } else { - t.Fatal("UpdatePodLabelForAdmissionController is not as expected") - } + newLabels, newAnnotations = getNewApplicationInfo(pod, "default", false, "root.default") + assert.Equal(t, len(newLabels), 3) + assert.Equal(t, newLabels["queue"], "root.default") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Equal(t, strings.HasPrefix(newLabels["applicationId"], constants.AutoGenAppPrefix), true) + assert.Equal(t, len(newAnnotations), 3) + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "root.default") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, strings.HasPrefix(newAnnotations["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true) } func TestDefaultQueueName(t *testing.T) { defaultConf := createConfig() pod := createTestingPodWithMeta() - if result := updatePodLabel(pod, defaultConf.GetNamespace(), defaultConf.GetGenerateUniqueAppIds(), defaultConf.GetDefaultQueueName()); result != nil { - assert.Equal(t, len(result), 4) - assert.Equal(t, result["random"], "random") - assert.Equal(t, result["applicationId"], "yunikorn-default-autogen") - assert.Equal(t, result["disableStateAware"], "true") - assert.Equal(t, result["queue"], "root.default") - } else { - t.Fatal("UpdatePodLabelForAdmissionController is not as expected") - } + newLabels, newAnnotations := getNewApplicationInfo(pod, defaultConf.GetNamespace(), defaultConf.GetGenerateUniqueAppIds(), defaultConf.GetDefaultQueueName()) + + assert.Equal(t, len(newLabels), 3) + assert.Equal(t, newLabels["applicationId"], "yunikorn-default-autogen") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Equal(t, newLabels["queue"], "root.default") + assert.Equal(t, len(newAnnotations), 3) + assert.Equal(t, newAnnotations["yunikorn.apache.org/app-id"], "yunikorn-default-autogen") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "root.default") queueNameEmptyConf := createConfigWithOverrides(map[string]string{ conf.AMFilteringDefaultQueueName: "", }) - if result := updatePodLabel(pod, queueNameEmptyConf.GetNamespace(), queueNameEmptyConf.GetGenerateUniqueAppIds(), queueNameEmptyConf.GetDefaultQueueName()); result != nil { - assert.Equal(t, len(result), 3) - assert.Equal(t, result["random"], "random") - assert.Equal(t, result["applicationId"], "yunikorn-default-autogen") - assert.Equal(t, result["disableStateAware"], "true") - assert.Equal(t, result["queue"], "") - } else { - t.Fatal("UpdatePodLabelForAdmissionController is not as expected") - } + newLabels, newAnnotations = getNewApplicationInfo(pod, queueNameEmptyConf.GetNamespace(), queueNameEmptyConf.GetGenerateUniqueAppIds(), queueNameEmptyConf.GetDefaultQueueName()) + assert.Equal(t, len(newLabels), 2) + assert.Equal(t, newLabels["applicationId"], "yunikorn-default-autogen") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Equal(t, newLabels["queue"], "") + assert.Equal(t, len(newAnnotations), 2) + assert.Equal(t, newAnnotations["yunikorn.apache.org/app-id"], "yunikorn-default-autogen") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "") customQueueNameConf := createConfigWithOverrides(map[string]string{ conf.AMFilteringDefaultQueueName: "yunikorn", }) - if result := updatePodLabel(pod, customQueueNameConf.GetNamespace(), customQueueNameConf.GetGenerateUniqueAppIds(), customQueueNameConf.GetDefaultQueueName()); result != nil { - assert.Equal(t, len(result), 4) - assert.Equal(t, result["random"], "random") - assert.Equal(t, result["applicationId"], "yunikorn-default-autogen") - assert.Equal(t, result["disableStateAware"], "true") - assert.Assert(t, result["queue"] != "yunikorn") - } else { - t.Fatal("UpdatePodLabelForAdmissionController is not as expected") - } + newLabels, newAnnotations = getNewApplicationInfo(pod, customQueueNameConf.GetNamespace(), customQueueNameConf.GetGenerateUniqueAppIds(), customQueueNameConf.GetDefaultQueueName()) + assert.Equal(t, len(newLabels), 3) + assert.Equal(t, newLabels["applicationId"], "yunikorn-default-autogen") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Assert(t, newLabels["queue"] != "yunikorn") + assert.Equal(t, len(newAnnotations), 3) + assert.Equal(t, newAnnotations["yunikorn.apache.org/app-id"], "yunikorn-default-autogen") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Assert(t, newAnnotations["yunikorn.apache.org/queue"] != "yunikorn") customValidQueueNameConf := createConfigWithOverrides(map[string]string{ conf.AMFilteringDefaultQueueName: "root.yunikorn", }) - if result := updatePodLabel(pod, customValidQueueNameConf.GetNamespace(), - customValidQueueNameConf.GetGenerateUniqueAppIds(), customValidQueueNameConf.GetDefaultQueueName()); result != nil { - assert.Equal(t, len(result), 4) - assert.Equal(t, result["random"], "random") - assert.Equal(t, result["applicationId"], "yunikorn-default-autogen") - assert.Equal(t, result["disableStateAware"], "true") - assert.Equal(t, result["queue"], "root.yunikorn") - } else { - t.Fatal("UpdatePodLabelForAdmissionController is not as expected") + newLabels, newAnnotations = getNewApplicationInfo(pod, customValidQueueNameConf.GetNamespace(), + customValidQueueNameConf.GetGenerateUniqueAppIds(), customValidQueueNameConf.GetDefaultQueueName()) + assert.Equal(t, len(newLabels), 3) + assert.Equal(t, newLabels["applicationId"], "yunikorn-default-autogen") + assert.Equal(t, newLabels["disableStateAware"], "true") + assert.Equal(t, newLabels["queue"], "root.yunikorn") + assert.Equal(t, len(newAnnotations), 3) + assert.Equal(t, newAnnotations["yunikorn.apache.org/app-id"], "yunikorn-default-autogen") + assert.Equal(t, newAnnotations["yunikorn.apache.org/disable-state-aware"], "true") + assert.Equal(t, newAnnotations["yunikorn.apache.org/queue"], "root.yunikorn") +} + +func TestUpdateLabelIfNotPresentInPod(t *testing.T) { + pod := createTestingPodWithMeta() + pod.ObjectMeta.Labels = map[string]string{ + "exist_key": "exist_value", + } + + // Verify updating label that exist in pod + result := make(map[string]string) + result = updateLabelIfNotPresentInPod(pod, result, "exist_key", "new_value") + assert.Equal(t, len(result), 0) + assert.Equal(t, result["exist_key"], "") + + // Verify updating label that not exists in pod + result = make(map[string]string) + result = updateLabelIfNotPresentInPod(pod, result, "new_key", "new_value") + assert.Equal(t, len(result), 1) + assert.Equal(t, result["new_key"], "new_value") + + // Verify if no label in pod + pod.ObjectMeta.Labels = nil + result = make(map[string]string) + result = updateLabelIfNotPresentInPod(pod, result, "new_key", "new_value") + assert.Equal(t, len(result), 1) + assert.Equal(t, result["new_key"], "new_value") +} + +func TestUpdateAnnotationIfNotPresentInPod(t *testing.T) { + pod := createTestingPodWithMeta() + pod.ObjectMeta.Annotations = map[string]string{ + "exist_key": "exist_value", + } + + // Verify updating annotation that exist in pod + result := make(map[string]string) + result = updateAnnotationIfNotPresentInPod(pod, result, "exist_key", "new_value") + assert.Equal(t, len(result), 0) + assert.Equal(t, result["exist_key"], "") + + // Verify updating annotation that not exists in pod + result = make(map[string]string) + result = updateAnnotationIfNotPresentInPod(pod, result, "new_key", "new_value") + assert.Equal(t, len(result), 1) + assert.Equal(t, result["new_key"], "new_value") + + // Verify if no annotation in pod + pod.ObjectMeta.Labels = nil + result = make(map[string]string) + result = updateAnnotationIfNotPresentInPod(pod, result, "new_key", "new_value") + assert.Equal(t, len(result), 1) + assert.Equal(t, result["new_key"], "new_value") +} + +func TestGetApplicationIDFromPod(t *testing.T) { + appName := []string{"app00001", "app00002"} + + // test empty pod + pod := &v1.Pod{} + appId, isFromLabel := getApplicationIDFromPod(pod) + assert.Equal(t, appId, "") + assert.Equal(t, isFromLabel, false) + + // test annotation take precedence over label + pod = &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{ + constants.LabelApplicationID: appName[0], + }, + Annotations: map[string]string{ + constants.AnnotationApplicationID: appName[1], + }, + }, + } + appId, isFromLabel = getApplicationIDFromPod(pod) + assert.Equal(t, appId, appName[1]) + assert.Equal(t, isFromLabel, false) + + // test pod with label only + pod = &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{ + constants.LabelApplicationID: appName[0], + }, + }, + } + appId, isFromLabel = getApplicationIDFromPod(pod) + assert.Equal(t, appId, appName[0]) + assert.Equal(t, isFromLabel, true) +} + +func TestGetDisableStateAwareFromPod(t *testing.T) { + // test empty pod + pod := &v1.Pod{} + disableStateAware, isFromLabel := getDisableStateAwareFromPod(pod) + assert.Equal(t, disableStateAware, "") + assert.Equal(t, isFromLabel, false) + + // test annotation take precedence over label + pod = &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{ + constants.LabelDisableStateAware: "true", + }, + Annotations: map[string]string{ + constants.AnnotationDisableStateAware: "false", + }, + }, + } + disableStateAware, isFromLabel = getDisableStateAwareFromPod(pod) + assert.Equal(t, disableStateAware, "false") + assert.Equal(t, isFromLabel, false) + + // test pod with label only + pod = &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{ + constants.LabelDisableStateAware: "true", + }, + }, } + disableStateAware, isFromLabel = getDisableStateAwareFromPod(pod) + assert.Equal(t, disableStateAware, "true") + assert.Equal(t, isFromLabel, true) } diff --git a/pkg/cache/metadata.go b/pkg/cache/metadata.go index b3b038733..86b72ced5 100644 --- a/pkg/cache/metadata.go +++ b/pkg/cache/metadata.go @@ -68,6 +68,7 @@ func getAppMetadata(pod *v1.Pod) (ApplicationMetadata, bool) { zap.String("name", pod.Name)) return ApplicationMetadata{}, false } + queueName, _ := utils.GetQueueNameFromPod(pod, constants.ApplicationDefaultQueue) // tags will at least have namespace info // labels or annotations from the pod can be added when needed @@ -117,7 +118,7 @@ func getAppMetadata(pod *v1.Pod) (ApplicationMetadata, bool) { return ApplicationMetadata{ ApplicationID: appID, - QueueName: utils.GetQueueNameFromPod(pod), + QueueName: queueName, User: user, Groups: groups, Tags: tags, diff --git a/pkg/cache/placeholder.go b/pkg/cache/placeholder.go index 0d1345a03..def095e00 100644 --- a/pkg/cache/placeholder.go +++ b/pkg/cache/placeholder.go @@ -94,7 +94,10 @@ func newPlaceholder(placeholderName string, app *Application, taskGroup TaskGrou constants.LabelQueueName: app.GetQueue(), constants.LabelPlaceholderFlag: "true", }), - Annotations: annotations, + Annotations: utils.MergeMaps(annotations, map[string]string{ + constants.AnnotationApplicationID: app.GetApplicationID(), + constants.AnnotationQueueName: app.GetQueue(), + }), OwnerReferences: ownerRefs, }, Spec: v1.PodSpec{ diff --git a/pkg/cache/placeholder_test.go b/pkg/cache/placeholder_test.go index 8dee656bc..d3ce03dd6 100644 --- a/pkg/cache/placeholder_test.go +++ b/pkg/cache/placeholder_test.go @@ -109,11 +109,19 @@ func TestNewPlaceholder(t *testing.T) { assert.Equal(t, holder.pod.Name, "ph-name") assert.Equal(t, holder.pod.Namespace, namespace) assert.Equal(t, len(holder.pod.Labels), 5, "unexpected number of labels") + assert.Equal(t, len(holder.pod.Annotations), 7, "unexpected number of annotations") + assert.Equal(t, holder.pod.Labels["labelKey0"], "labelKeyValue0") + assert.Equal(t, holder.pod.Labels["labelKey1"], "labelKeyValue1") assert.Equal(t, holder.pod.Labels[constants.LabelApplicationID], appID) assert.Equal(t, holder.pod.Labels[constants.LabelQueueName], queue) - assert.Equal(t, holder.pod.Labels["placeholder"], "true") - assert.Equal(t, len(holder.pod.Annotations), 5, "unexpected number of annotations") + assert.Equal(t, holder.pod.Labels[constants.LabelPlaceholderFlag], "true") + assert.Equal(t, holder.pod.Annotations["annotationKey0"], "annotationValue0") + assert.Equal(t, holder.pod.Annotations["annotationKey1"], "annotationValue1") + assert.Equal(t, holder.pod.Annotations["annotationKey2"], "annotationValue2") assert.Equal(t, holder.pod.Annotations[constants.AnnotationTaskGroupName], app.taskGroups[0].Name) + assert.Equal(t, holder.pod.Annotations[constants.AnnotationApplicationID], appID) + assert.Equal(t, holder.pod.Annotations[constants.AnnotationQueueName], queue) + assert.Equal(t, holder.pod.Annotations[constants.AnnotationPlaceholderFlag], "true") assert.Equal(t, common.GetPodResource(holder.pod).Resources[siCommon.CPU].Value, int64(500)) assert.Equal(t, common.GetPodResource(holder.pod).Resources[siCommon.Memory].Value, int64(1024*1000*1000)) assert.Equal(t, common.GetPodResource(holder.pod).Resources["pods"].Value, int64(1)) @@ -149,12 +157,18 @@ func TestNewPlaceholderWithLabelsAndAnnotations(t *testing.T) { "queue": "root.default", }) - assert.Equal(t, len(holder.pod.Annotations), 6) + assert.Equal(t, len(holder.pod.Annotations), 8) assert.Equal(t, holder.pod.Annotations["annotationKey0"], "annotationValue0") assert.Equal(t, holder.pod.Annotations["annotationKey1"], "annotationValue1") assert.Equal(t, holder.pod.Annotations["annotationKey2"], "annotationValue2") + assert.Equal(t, holder.pod.Annotations[constants.AnnotationTaskGroupName], app.taskGroups[0].Name) + assert.Equal(t, holder.pod.Annotations[constants.AnnotationApplicationID], appID) + assert.Equal(t, holder.pod.Annotations[constants.AnnotationQueueName], queue) + assert.Equal(t, holder.pod.Annotations[constants.AnnotationPlaceholderFlag], "true") + var taskGroupsDef []TaskGroup err = json.Unmarshal([]byte(holder.pod.Annotations[siCommon.DomainYuniKorn+"task-groups"]), &taskGroupsDef) + assert.NilError(t, err, "taskGroupsDef unmarshal failed") var priority *int32 assert.Equal(t, priority, holder.pod.Spec.Priority) diff --git a/pkg/common/constants/constants.go b/pkg/common/constants/constants.go index ab4dbd1c8..492688877 100644 --- a/pkg/common/constants/constants.go +++ b/pkg/common/constants/constants.go @@ -42,6 +42,7 @@ const RootQueue = "root" const AnnotationQueueName = DomainYuniKorn + "queue" const AnnotationParentQueue = DomainYuniKorn + "parentqueue" const LabelDisableStateAware = "disableStateAware" +const AnnotationDisableStateAware = DomainYuniKorn + "disable-state-aware" const ApplicationDefaultQueue = "root.default" const DefaultPartition = "default" const AppTagNamespace = "namespace" @@ -109,9 +110,13 @@ const AnnotationGenerateAppID = DomainYuniKorn + "namespace.generateAppId" // false: do not do anything const AnnotationEnableYuniKorn = DomainYuniKorn + "namespace.enableYuniKorn" -// Admission Controller pod label update constants +// Admission Controller pod update constants const AutoGenAppPrefix = "yunikorn" const AutoGenAppSuffix = "autogen" +const AddPatchOp = "add" +const AnnotationPatchPath = "/metadata/annotations" +const LabelPatchPath = "/metadata/labels" +const SchedulerNamePatchPath = "/spec/schedulerName" // Compression Algorithms for schedulerConfig const GzipSuffix = "gz" diff --git a/pkg/common/utils/utils.go b/pkg/common/utils/utils.go index b0895ef83..3f559f9e9 100644 --- a/pkg/common/utils/utils.go +++ b/pkg/common/utils/utils.go @@ -102,14 +102,16 @@ func IsAssignedPod(pod *v1.Pod) bool { return len(pod.Spec.NodeName) != 0 } -func GetQueueNameFromPod(pod *v1.Pod) string { - queueName := constants.ApplicationDefaultQueue - if an := GetPodLabelValue(pod, constants.LabelQueueName); an != "" { - queueName = an - } else if qu := GetPodAnnotationValue(pod, constants.AnnotationQueueName); qu != "" { - queueName = qu +func GetQueueNameFromPod(pod *v1.Pod, defaultQueueName string) (value string, isFromLabel bool) { + queueName := defaultQueueName + isFromLabel = false + if value := GetPodAnnotationValue(pod, constants.AnnotationQueueName); value != "" { + queueName = value + } else if value := GetPodLabelValue(pod, constants.LabelQueueName); value != "" { + isFromLabel = true + queueName = value } - return queueName + return queueName, isFromLabel } // GenerateApplicationID generates an appID based on the namespace value diff --git a/pkg/common/utils/utils_test.go b/pkg/common/utils/utils_test.go index c7e8a45b4..3900cfcfc 100644 --- a/pkg/common/utils/utils_test.go +++ b/pkg/common/utils/utils_test.go @@ -812,6 +812,7 @@ func TestGetQueueNameFromPod(t *testing.T) { name string pod *v1.Pod expectedQueue string + isFromLabel bool }{ { name: "With queue label", @@ -821,6 +822,7 @@ func TestGetQueueNameFromPod(t *testing.T) { }, }, expectedQueue: queueInLabel, + isFromLabel: true, }, { name: "With queue annotation", @@ -830,6 +832,7 @@ func TestGetQueueNameFromPod(t *testing.T) { }, }, expectedQueue: queueInAnnotation, + isFromLabel: false, }, { name: "With queue label and annotation", @@ -839,7 +842,8 @@ func TestGetQueueNameFromPod(t *testing.T) { Annotations: map[string]string{constants.AnnotationQueueName: queueInAnnotation}, }, }, - expectedQueue: queueInLabel, + expectedQueue: queueInAnnotation, + isFromLabel: false, }, { name: "Without queue label and annotation", @@ -847,13 +851,15 @@ func TestGetQueueNameFromPod(t *testing.T) { ObjectMeta: metav1.ObjectMeta{}, }, expectedQueue: constants.ApplicationDefaultQueue, + isFromLabel: false, }, } for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { - queue := GetQueueNameFromPod(tc.pod) + queue, isFromLabel := GetQueueNameFromPod(tc.pod, constants.ApplicationDefaultQueue) assert.Equal(t, queue, tc.expectedQueue) + assert.Equal(t, isFromLabel, tc.isFromLabel) }) } }