Skip to content

Commit

Permalink
[YUNIKORN-2636] Admission Controller ignores existing Queue/Applicati…
Browse files Browse the repository at this point in the history 
…onID annotations (#848)

AM also check app-id/queue annotations in updatePodLabel().

Closes: #848

Signed-off-by: Yu-Lin Chen <[email protected]>
  • Loading branch information
@chenyulin0719
chenyulin0719 committed May 28, 2024
1 parent 08a9b93 commit 58adfe9
Show file tree
Hide file tree
Showing 3 changed files with 74 additions and 68 deletions.
42 changes: 21 additions & 21 deletions pkg/admission/admission_controller_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -83,8 +83,8 @@ func TestUpdateLabels(t *testing.T) {
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, strings.HasPrefix(updatedMap["applicationId"], constants.AutoGenAppPrefix), true)
assert.Equal(t, updatedMap[constants.LabelQueueName], "root.default")
assert.Equal(t, strings.HasPrefix(updatedMap[constants.LabelApplicationID], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("patch info content is not as expected")
}
Expand All @@ -104,8 +104,8 @@ func TestUpdateLabels(t *testing.T) {
UID: "7f5fd6c5d5",
ResourceVersion: "10654",
Labels: map[string]string{
"random": "random",
"applicationId": "app-0001",
"random": "random",
constants.LabelApplicationID: "app-0001",
},
},
Spec: v1.PodSpec{},
Expand All @@ -119,8 +119,8 @@ func TestUpdateLabels(t *testing.T) {
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")
assert.Equal(t, updatedMap[constants.LabelQueueName], "root.default")
assert.Equal(t, updatedMap[constants.LabelApplicationID], "app-0001")
} else {
t.Fatal("patch info content is not as expected")
}
Expand All @@ -140,8 +140,8 @@ func TestUpdateLabels(t *testing.T) {
UID: "7f5fd6c5d5",
ResourceVersion: "10654",
Labels: map[string]string{
"random": "random",
"queue": "root.abc",
"random": "random",
constants.LabelQueueName: "root.abc",
},
},
Spec: v1.PodSpec{},
Expand All @@ -156,8 +156,8 @@ func TestUpdateLabels(t *testing.T) {
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.abc")
assert.Equal(t, strings.HasPrefix(updatedMap["applicationId"], constants.AutoGenAppPrefix), true)
assert.Equal(t, updatedMap[constants.LabelQueueName], "root.abc")
assert.Equal(t, strings.HasPrefix(updatedMap[constants.LabelApplicationID], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("patch info content is not as expected")
}
Expand Down Expand Up @@ -187,8 +187,8 @@ func TestUpdateLabels(t *testing.T) {
assert.Equal(t, patch[0].Path, "/metadata/labels")
if updatedMap, ok := patch[0].Value.(map[string]string); ok {
assert.Equal(t, len(updatedMap), 2)
assert.Equal(t, updatedMap["queue"], "root.default")
assert.Equal(t, strings.HasPrefix(updatedMap["applicationId"], constants.AutoGenAppPrefix), true)
assert.Equal(t, updatedMap[constants.LabelQueueName], "root.default")
assert.Equal(t, strings.HasPrefix(updatedMap[constants.LabelApplicationID], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("patch info content is not as expected")
}
Expand All @@ -215,8 +215,8 @@ func TestUpdateLabels(t *testing.T) {
assert.Equal(t, patch[0].Path, "/metadata/labels")
if updatedMap, ok := patch[0].Value.(map[string]string); ok {
assert.Equal(t, len(updatedMap), 2)
assert.Equal(t, updatedMap["queue"], "root.default")
assert.Equal(t, strings.HasPrefix(updatedMap["applicationId"], constants.AutoGenAppPrefix), true)
assert.Equal(t, updatedMap[constants.LabelQueueName], "root.default")
assert.Equal(t, strings.HasPrefix(updatedMap[constants.LabelApplicationID], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("patch info content is not as expected")
}
Expand All @@ -241,8 +241,8 @@ func TestUpdateLabels(t *testing.T) {
assert.Equal(t, patch[0].Path, "/metadata/labels")
if updatedMap, ok := patch[0].Value.(map[string]string); ok {
assert.Equal(t, len(updatedMap), 2)
assert.Equal(t, updatedMap["queue"], "root.default")
assert.Equal(t, strings.HasPrefix(updatedMap["applicationId"], constants.AutoGenAppPrefix), true)
assert.Equal(t, updatedMap[constants.LabelQueueName], "root.default")
assert.Equal(t, strings.HasPrefix(updatedMap[constants.LabelApplicationID], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("patch info content is not as expected")
}
Expand Down Expand Up @@ -449,8 +449,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-default-autogen", "wrong applicationId label")
assert.Equal(t, labels(t, resp.Patch)["queue"], "root.default", "incorrect queue name")
assert.Equal(t, labels(t, resp.Patch)[constants.LabelApplicationID], "yunikorn-default-autogen", "wrong applicationId label")
assert.Equal(t, labels(t, resp.Patch)[constants.LabelQueueName], "root.default", "incorrect queue name")

// pod without applicationID
pod = v1.Pod{ObjectMeta: metav1.ObjectMeta{
Expand All @@ -467,17 +467,17 @@ 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)[constants.LabelApplicationID], "yunikorn-test-ns-autogen", "wrong applicationId label")

// pod with applicationId
pod.ObjectMeta.Labels = map[string]string{"applicationId": "test-app"}
pod.ObjectMeta.Labels = map[string]string{constants.LabelApplicationID: "test-app"}
podJSON, err = json.Marshal(pod)
assert.NilError(t, err, "failed to marshal pod")
req.Object = runtime.RawExtension{Raw: podJSON}
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, labels(t, resp.Patch)[constants.LabelApplicationID], "test-app", "wrong applicationId label")

// pod in bypassed namespace
pod = v1.Pod{ObjectMeta: metav1.ObjectMeta{
Expand Down
17 changes: 9 additions & 8 deletions pkg/admission/util.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,15 +30,15 @@ import (
)

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 {
for k, v := range pod.Labels {
result[k] = v
}

sparkAppID := utils.GetPodLabelValue(pod, constants.SparkLabelAppID)
appID := utils.GetPodLabelValue(pod, constants.LabelApplicationID)
if sparkAppID == "" && appID == "" {
labelAppID := utils.GetPodLabelValue(pod, constants.LabelApplicationID)
annotationAppID := utils.GetPodAnnotationValue(pod, constants.AnnotationApplicationID)
if sparkAppID == "" && labelAppID == "" && annotationAppID == "" {
// 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:
Expand All @@ -49,8 +49,10 @@ func updatePodLabel(pod *v1.Pod, namespace string, generateUniqueAppIds bool, de
result[constants.LabelApplicationID] = generatedID
}

// if existing label exist, it takes priority over everything else
if _, ok := existingLabels[constants.LabelQueueName]; !ok {
labelQueueName := utils.GetPodLabelValue(pod, constants.LabelQueueName)
annotationQueueName := utils.GetPodAnnotationValue(pod, constants.AnnotationQueueName)
if labelQueueName == "" && annotationQueueName == "" {
// if queueName not exist, generate one
// 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
Expand All @@ -63,9 +65,8 @@ func updatePodLabel(pod *v1.Pod, namespace string, generateUniqueAppIds bool, de
}

func updatePodAnnotation(pod *v1.Pod, key string, value string) map[string]string {
existingAnnotations := pod.Annotations
result := make(map[string]string)
for k, v := range existingAnnotations {
for k, v := range pod.Annotations {
result[k] = v
}
result[key] = value
Expand Down
83 changes: 44 additions & 39 deletions pkg/admission/util_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,14 +62,16 @@ func createTestingPodWithMeta() *v1.Pod {
Labels: map[string]string{
"random": "random",
},
Annotations: map[string]string{},
}

return pod
}

func createTestingPodWithAppId() *v1.Pod {
func createTestingPodWithLabels(appId string, queue string) *v1.Pod {
pod := createTestingPodWithMeta()
pod.ObjectMeta.Labels["applicationId"] = "app-0001"
pod.ObjectMeta.Labels[constants.LabelApplicationID] = appId
pod.ObjectMeta.Labels[constants.LabelQueueName] = queue

return pod
}
Expand All @@ -81,9 +83,10 @@ func createTestingPodWithGenerateName() *v1.Pod {
return pod
}

func createTestingPodWithQueue() *v1.Pod {
func createTestingPodWithAnnotations(appId string, queue string) *v1.Pod {
pod := createTestingPodWithMeta()
pod.ObjectMeta.Labels["queue"] = "root.abc"
pod.ObjectMeta.Annotations[constants.AnnotationApplicationID] = appId
pod.ObjectMeta.Annotations[constants.AnnotationQueueName] = queue

return pod
}
Expand All @@ -100,70 +103,72 @@ func createTestingPodNoNamespaceAndLabels() *v1.Pod {
}

func TestUpdatePodLabelForAdmissionController(t *testing.T) {
dummyAppId := "app-0001"
dummyQueueName := "root.abc"
defaultQueueName := "root.default"

// verify when appId/queue are not given,
// we generate new appId/queue labels
pod := createTestingPodWithMeta()

if result := updatePodLabel(pod, "default", false, "root.default"); result != nil {
if result := updatePodLabel(pod, "default", false, defaultQueueName); result != nil {
assert.Equal(t, len(result), 3)
assert.Equal(t, result["random"], "random")
assert.Equal(t, result["queue"], "root.default")
assert.Equal(t, strings.HasPrefix(result["applicationId"], constants.AutoGenAppPrefix), true)
assert.Equal(t, strings.HasPrefix(result[constants.LabelApplicationID], constants.AutoGenAppPrefix), true)
assert.Equal(t, result[constants.LabelQueueName], defaultQueueName)
} else {
t.Fatal("UpdatePodLabelForAdmissionController is not as expected")
}

// verify if applicationId is given in the labels,
// verify if applicationId and queue is given in the labels,
// we won't modify it
pod = createTestingPodWithAppId()

if result := updatePodLabel(pod, "default", false, "root.default"); result != nil {
pod = createTestingPodWithLabels(dummyAppId, dummyQueueName)
if result := updatePodLabel(pod, "default", false, defaultQueueName); 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")
assert.Equal(t, result[constants.LabelApplicationID], dummyAppId)
assert.Equal(t, result[constants.LabelQueueName], dummyQueueName)
} else {
t.Fatal("UpdatePodLabelForAdmissionController is not as expected")
}

// 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), 3)
// verify if applicationId and queue is given in the annotations,
// we won't generate new labels
pod = createTestingPodWithAnnotations(dummyAppId, dummyQueueName)
if result := updatePodLabel(pod, "default", false, defaultQueueName); result != nil {
t.Log(result)
assert.Equal(t, len(result), 1)
assert.Equal(t, result["random"], "random")
assert.Equal(t, result["queue"], "root.abc")
assert.Equal(t, strings.HasPrefix(result["applicationId"], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("UpdatePodLabelForAdmissionControllert is not as expected")
t.Fatal("UpdatePodLabelForAdmissionController is not as expected")
}

// namespace might be empty
// labels might be empty
pod = createTestingPodNoNamespaceAndLabels()

if result := updatePodLabel(pod, "default", false, "root.default"); result != nil {
if result := updatePodLabel(pod, "default", false, defaultQueueName); result != nil {
assert.Equal(t, len(result), 2)
assert.Equal(t, result["queue"], "root.default")
assert.Equal(t, strings.HasPrefix(result["applicationId"], constants.AutoGenAppPrefix), true)
assert.Equal(t, result[constants.LabelQueueName], defaultQueueName)
assert.Equal(t, strings.HasPrefix(result[constants.LabelApplicationID], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("UpdatePodLabelForAdmissionController is not as expected")
}

// pod name might be empty, it can comes from generatedName
pod = createTestingPodWithGenerateName()
if result := updatePodLabel(pod, "default", false, "root.default"); result != nil {
if result := updatePodLabel(pod, "default", false, defaultQueueName); result != nil {
assert.Equal(t, len(result), 2)
assert.Equal(t, result["queue"], "root.default")
assert.Equal(t, strings.HasPrefix(result["applicationId"], constants.AutoGenAppPrefix), true)
assert.Equal(t, result[constants.LabelQueueName], defaultQueueName)
assert.Equal(t, strings.HasPrefix(result[constants.LabelApplicationID], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("UpdatePodLabelForAdmissionController is not as expected")
}

pod = createMinimalTestingPod()
if result := updatePodLabel(pod, "default", false, "root.default"); result != nil {
if result := updatePodLabel(pod, "default", false, defaultQueueName); result != nil {
assert.Equal(t, len(result), 2)
assert.Equal(t, result["queue"], "root.default")
assert.Equal(t, strings.HasPrefix(result["applicationId"], constants.AutoGenAppPrefix), true)
assert.Equal(t, result[constants.LabelQueueName], defaultQueueName)
assert.Equal(t, strings.HasPrefix(result[constants.LabelApplicationID], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("UpdatePodLabelForAdmissionController is not as expected")
}
Expand All @@ -175,8 +180,8 @@ func TestDefaultQueueName(t *testing.T) {
if result := updatePodLabel(pod, defaultConf.GetNamespace(), defaultConf.GetGenerateUniqueAppIds(), defaultConf.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["queue"], "root.default")
assert.Equal(t, result[constants.LabelApplicationID], "yunikorn-default-autogen")
assert.Equal(t, result[constants.LabelQueueName], "root.default")
} else {
t.Fatal("UpdatePodLabelForAdmissionController is not as expected")
}
Expand All @@ -187,8 +192,8 @@ func TestDefaultQueueName(t *testing.T) {
if result := updatePodLabel(pod, queueNameEmptyConf.GetNamespace(), queueNameEmptyConf.GetGenerateUniqueAppIds(), queueNameEmptyConf.GetDefaultQueueName()); result != nil {
assert.Equal(t, len(result), 2)
assert.Equal(t, result["random"], "random")
assert.Equal(t, result["applicationId"], "yunikorn-default-autogen")
assert.Equal(t, result["queue"], "")
assert.Equal(t, result[constants.LabelApplicationID], "yunikorn-default-autogen")
assert.Equal(t, result[constants.LabelQueueName], "")
} else {
t.Fatal("UpdatePodLabelForAdmissionController is not as expected")
}
Expand All @@ -199,8 +204,8 @@ func TestDefaultQueueName(t *testing.T) {
if result := updatePodLabel(pod, customQueueNameConf.GetNamespace(), customQueueNameConf.GetGenerateUniqueAppIds(), customQueueNameConf.GetDefaultQueueName()); result != nil {
assert.Equal(t, len(result), 3)
assert.Equal(t, result["random"], "random")
assert.Equal(t, result["applicationId"], "yunikorn-default-autogen")
assert.Assert(t, result["queue"] != "yunikorn")
assert.Equal(t, result[constants.LabelApplicationID], "yunikorn-default-autogen")
assert.Assert(t, result[constants.LabelQueueName] != "yunikorn")
} else {
t.Fatal("UpdatePodLabelForAdmissionController is not as expected")
}
Expand All @@ -212,8 +217,8 @@ func TestDefaultQueueName(t *testing.T) {
customValidQueueNameConf.GetGenerateUniqueAppIds(), customValidQueueNameConf.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["queue"], "root.yunikorn")
assert.Equal(t, result[constants.LabelApplicationID], "yunikorn-default-autogen")
assert.Equal(t, result[constants.LabelQueueName], "root.yunikorn")
} else {
t.Fatal("UpdatePodLabelForAdmissionController is not as expected")
}
Expand Down

0 comments on commit 58adfe9

Please sign in to comment.