Skip to content

Commit

Permalink
[YUNIKORN-1351] alway prefer annotations above labels
Browse files Browse the repository at this point in the history
  • Loading branch information
@chenyulin0719
chenyulin0719 committed Sep 15, 2023
1 parent 0b3ea77 commit b16c973
Show file tree
Hide file tree
Showing 11 changed files with 355 additions and 160 deletions.
35 changes: 25 additions & 10 deletions pkg/admission/admission_controller.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -415,20 +415,35 @@ func (c *AdmissionController) updatePreemptionInfo(pod *v1.Pod, patch []common.P
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 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())
annotations := getAnnotationsForApplicationInfoUpdate(pod, namespace, c.conf.GetGenerateUniqueAppIds(), c.conf.GetDefaultQueueName())

patch = append(patch, common.PatchOperation{
Op: "add",
Path: "/metadata/labels",
Value: result,
})
// check for an existing patch on annotations and update it
for _, p := range patch {
if p.Op == "add" && p.Path == "/metadata/annotations" {
if existingAnnotations, ok := p.Value.(map[string]string); ok {
for k, v := range annotations {
existingAnnotations[k] = v
}
return patch
}
}
}

if len(annotations) != 0 {
patch = append(patch, common.PatchOperation{
Op: "add",
Path: "/metadata/annotations",
Value: annotations,
})
}

return patch
}
Expand Down
87 changes: 43 additions & 44 deletions pkg/admission/admission_controller_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,17 +75,16 @@ func TestUpdateLabels(t *testing.T) {
}

c := createAdmissionControllerForTest()
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")
assert.Equal(t, patch[0].Path, "/metadata/annotations")
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)
assert.Equal(t, len(updatedMap), 3)
assert.Equal(t, updatedMap["yunikorn.apache.org/queue"], "root.default")
assert.Equal(t, updatedMap["yunikorn.apache.org/disable-state-aware"], "true")
assert.Equal(t, strings.HasPrefix(updatedMap["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("patch info content is not as expected")
}
Expand All @@ -112,16 +111,15 @@ func TestUpdateLabels(t *testing.T) {
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")
assert.Equal(t, patch[0].Path, "/metadata/annotations") //nolint:gosec
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, len(updatedMap), 2)
assert.Equal(t, updatedMap["yunikorn.apache.org/queue"], "root.default")
assert.Equal(t, updatedMap["yunikorn.apache.org/app-id"], "app-0001")
} else {
t.Fatal("patch info content is not as expected")
}
Expand Down Expand Up @@ -149,17 +147,16 @@ func TestUpdateLabels(t *testing.T) {
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")
assert.Equal(t, patch[0].Path, "/metadata/annotations") //nolint:gosec
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)
assert.Equal(t, len(updatedMap), 3)
assert.Equal(t, updatedMap["yunikorn.apache.org/queue"], "root.abc")
assert.Equal(t, updatedMap["yunikorn.apache.org/disable-state-aware"], "true")
assert.Equal(t, strings.HasPrefix(updatedMap["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("patch info content is not as expected")
}
Expand All @@ -182,16 +179,16 @@ func TestUpdateLabels(t *testing.T) {
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")
assert.Equal(t, patch[0].Path, "/metadata/annotations") //nolint:gosec
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)
assert.Equal(t, updatedMap["yunikorn.apache.org/queue"], "root.default")
assert.Equal(t, updatedMap["yunikorn.apache.org/disable-state-aware"], "true")
assert.Equal(t, strings.HasPrefix(updatedMap["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("patch info content is not as expected")
}
Expand All @@ -211,16 +208,16 @@ func TestUpdateLabels(t *testing.T) {
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")
assert.Equal(t, patch[0].Path, "/metadata/annotations") //nolint:gosec
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)
assert.Equal(t, updatedMap["yunikorn.apache.org/queue"], "root.default")
assert.Equal(t, updatedMap["yunikorn.apache.org/disable-state-aware"], "true")
assert.Equal(t, strings.HasPrefix(updatedMap["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("patch info content is not as expected")
}
Expand All @@ -238,16 +235,16 @@ func TestUpdateLabels(t *testing.T) {
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")
assert.Equal(t, patch[0].Path, "/metadata/annotations") //nolint:gosec
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)
assert.Equal(t, updatedMap["yunikorn.apache.org/queue"], "root.default")
assert.Equal(t, updatedMap["yunikorn.apache.org/disable-state-aware"], "true")
assert.Equal(t, strings.HasPrefix(updatedMap["yunikorn.apache.org/app-id"], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("patch info content is not as expected")
}
Expand Down Expand Up @@ -454,9 +451,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{
Expand All @@ -473,8 +470,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"}
Expand All @@ -484,7 +481,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{
Expand Down Expand Up @@ -517,7 +514,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{
Expand Down Expand Up @@ -849,10 +848,10 @@ func schedulerName(t *testing.T, patch []byte) string {
return ""
}

func labels(t *testing.T, patch []byte) map[string]interface{} {
func annotations(t *testing.T, patch []byte) map[string]interface{} {
ops := parsePatch(t, patch)
for _, op := range ops {
if op.Path == "/metadata/labels" {
if op.Path == "/metadata/annotations" {
return op.Value.(map[string]interface{})
}
}
Expand Down
64 changes: 50 additions & 14 deletions pkg/admission/util.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,38 +31,50 @@ 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
func getAnnotationsForApplicationInfoUpdate(pod *v1.Pod, namespace string, generateUniqueAppIds bool, defaultQueueName string) map[string]string {
result := make(map[string]string)
for k, v := range existingLabels {
existingAnnotations := pod.Annotations
for k, v := range existingAnnotations {
result[k] = v
}

sparkAppID := utils.GetPodLabelValue(pod, constants.SparkLabelAppID)
appID := utils.GetPodLabelValue(pod, constants.LabelApplicationID)
if sparkAppID == "" && appID == "" {
appID := getApplicationIDFromPod(pod)
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}-${GENERATED_UUID}
// else
// application ID convention: ${AUTO_GEN_PREFIX}-${NAMESPACE}-${AUTO_GEN_SUFFIX}
generatedID := generateAppID(namespace, generateUniqueAppIds)
result[constants.LabelApplicationID] = generatedID
appID = generateAppID(namespace, generateUniqueAppIds)
result[constants.AnnotationApplicationID] = appID

// if we generate an app ID, disable state-aware scheduling for this app
if _, ok := existingLabels[constants.LabelDisableStateAware]; !ok {
result[constants.LabelDisableStateAware] = "true"
disableStateAware := "true"
if value := utils.GetPodLabelValue(pod, constants.LabelDisableStateAware); value != "" {
disableStateAware = value
}
result[constants.AnnotationDisableStateAware] = disableStateAware
}

// if app id not in pod annotation, add it
if value := utils.GetPodAnnotationValue(pod, constants.AnnotationIgnoreApplication); value == "" {
result[constants.AnnotationApplicationID] = appID
}

// if existing label exist, it takes priority over everything else
if _, ok := existingLabels[constants.LabelQueueName]; !ok {
queueName := getQueueNameFromPod(pod)
if queueName == "" {
// 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
queueName = defaultQueueName
}
}

// if queue name not in pod annotation, add it
if value := utils.GetPodAnnotationValue(pod, constants.AnnotationQueueName); value == "" {
if queueName != "" {
result[constants.AnnotationQueueName] = queueName
}
}

Expand Down Expand Up @@ -106,3 +118,27 @@ func generateAppID(namespace string, generateUniqueAppIds bool) string {

return fmt.Sprintf("%.63s", generatedID)
}

func getApplicationIDFromPod(pod *v1.Pod) string {
// if existing annotation exist, it takes priority over everything else
if value := utils.GetPodAnnotationValue(pod, constants.AnnotationApplicationID); value != "" {
return value
} else if value := utils.GetPodLabelValue(pod, constants.LabelApplicationID); value != "" {
return value
}
// application ID can be defined in Spark label
if value := utils.GetPodLabelValue(pod, constants.SparkLabelAppID); value != "" {
return value
}
return ""
}

func getQueueNameFromPod(pod *v1.Pod) string {
// if existing annotation exist, it takes priority over everything else
if value := utils.GetPodAnnotationValue(pod, constants.AnnotationQueueName); value != "" {
return value
} else if value := utils.GetPodLabelValue(pod, constants.LabelQueueName); value != "" {
return value
}
return ""
}
Loading

0 comments on commit b16c973

Please sign in to comment.