diff --git a/pkg/cache/placeholder.go b/pkg/cache/placeholder.go index d259b5622..6208aa23a 100644 --- a/pkg/cache/placeholder.go +++ b/pkg/cache/placeholder.go @@ -91,8 +91,8 @@ func newPlaceholder(placeholderName string, app *Application, taskGroup TaskGrou Name: placeholderName, Namespace: app.tags[constants.AppTagNamespace], Labels: utils.MergeMaps(taskGroup.Labels, map[string]string{ - constants.LabelApplicationID: app.GetApplicationID(), - constants.LabelQueueName: app.GetQueue(), + constants.CanonicalLabelApplicationID: app.GetApplicationID(), + constants.CanonicalLabelQueueName: app.GetQueue(), }), Annotations: annotations, OwnerReferences: ownerRefs, diff --git a/pkg/cache/placeholder_test.go b/pkg/cache/placeholder_test.go index 35914af38..553ade1aa 100644 --- a/pkg/cache/placeholder_test.go +++ b/pkg/cache/placeholder_test.go @@ -125,10 +125,10 @@ func TestNewPlaceholder(t *testing.T) { assert.Equal(t, holder.pod.Name, "ph-name") assert.Equal(t, holder.pod.Namespace, namespace) assert.DeepEqual(t, holder.pod.Labels, map[string]string{ - constants.LabelApplicationID: appID, - constants.LabelQueueName: queue, - "labelKey0": "labelKeyValue0", - "labelKey1": "labelKeyValue1", + constants.CanonicalLabelApplicationID: appID, + constants.CanonicalLabelQueueName: queue, + "labelKey0": "labelKeyValue0", + "labelKey1": "labelKeyValue1", }) assert.Equal(t, len(holder.pod.Annotations), 6, "unexpected number of annotations") assert.Equal(t, holder.pod.Annotations[constants.AnnotationTaskGroupName], app.taskGroups[0].Name) diff --git a/pkg/common/utils/utils.go b/pkg/common/utils/utils.go index 281189972..515d4109c 100644 --- a/pkg/common/utils/utils.go +++ b/pkg/common/utils/utils.go @@ -109,12 +109,20 @@ func IsAssignedPod(pod *v1.Pod) bool { } func GetQueueNameFromPod(pod *v1.Pod) string { + // Queue name can be defined in multiple places + // The queue name is determined by the following order + // 1. Label: constants.CanonicalLabelQueueName + // 2. Annotation: constants.AnnotationQueueName + // 3. Label: constants.LabelQueueName queueName := "" - if an := GetPodLabelValue(pod, constants.LabelQueueName); an != "" { - queueName = an - } else if qu := GetPodAnnotationValue(pod, constants.AnnotationQueueName); qu != "" { - queueName = qu + if canonicalLabelQueueName := GetPodLabelValue(pod, constants.CanonicalLabelQueueName); canonicalLabelQueueName != "" { + queueName = canonicalLabelQueueName + } else if annotationQueueName := GetPodAnnotationValue(pod, constants.AnnotationQueueName); annotationQueueName != "" { + queueName = annotationQueueName + } else if labelQueueName := GetPodLabelValue(pod, constants.LabelQueueName); labelQueueName != "" { + queueName = labelQueueName } + return queueName } @@ -159,15 +167,30 @@ func GetApplicationIDFromPod(pod *v1.Pod) string { } } - // Application ID can be defined in annotation - appID := GetPodAnnotationValue(pod, constants.AnnotationApplicationID) + // Application ID can be defined in multiple places + // The application ID is determined by the following order. + // 1. Label: constants.CanonicalLabelApplicationID + // 2. Annotation: constants.AnnotationApplicationID + // 3. Label: constants.LabelApplicationID + // 4. Label: constants.SparkLabelAppID + + appID := GetPodLabelValue(pod, constants.CanonicalLabelApplicationID) + if appID == "" { - // Application ID can be defined in label - appID = GetPodLabelValue(pod, constants.LabelApplicationID) + appID = GetPodAnnotationValue(pod, constants.AnnotationApplicationID) } + if appID == "" { - // Spark can also define application ID - appID = GetPodLabelValue(pod, constants.SparkLabelAppID) + labelKeys := []string{ + constants.LabelApplicationID, + constants.SparkLabelAppID, + } + for _, label := range labelKeys { + appID = GetPodLabelValue(pod, label) + if appID != "" { + break + } + } } // If plugin mode, interpret missing Application ID as a non-YuniKorn pod diff --git a/pkg/common/utils/utils_test.go b/pkg/common/utils/utils_test.go index 0b20f1a95..4c2b64cda 100644 --- a/pkg/common/utils/utils_test.go +++ b/pkg/common/utils/utils_test.go @@ -587,10 +587,10 @@ func TestGetApplicationIDFromPod(t *testing.T) { defer SetPluginMode(false) defer func() { conf.GetSchedulerConf().GenerateUniqueAppIds = false }() - appIDInLabel := "labelAppID" + appIDInCanonicalLabel := "CanonicalLabelAppID" appIDInAnnotation := "annotationAppID" - appIDInSelector := "selectorAppID" - sparkIDInAnnotation := "sparkAnnotationAppID" + appIDInLabel := "labelAppID" + appIDInSelector := "sparkLabelAppID" testCases := []struct { name string pod *v1.Pod @@ -598,13 +598,55 @@ func TestGetApplicationIDFromPod(t *testing.T) { expectedAppIDPluginMode string generateUniqueAppIds bool }{ - {"AppID defined in label", &v1.Pod{ + {"AppID defined in canonical label", &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{constants.CanonicalLabelApplicationID: appIDInCanonicalLabel}, + }, + Spec: v1.PodSpec{SchedulerName: constants.SchedulerName}, + }, appIDInCanonicalLabel, appIDInCanonicalLabel, false}, + {"AppID defined in annotation", &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Annotations: map[string]string{constants.AnnotationApplicationID: appIDInAnnotation}, + }, + Spec: v1.PodSpec{SchedulerName: constants.SchedulerName}, + }, appIDInAnnotation, appIDInAnnotation, false}, + {"AppID defined in legacy label", &v1.Pod{ ObjectMeta: metav1.ObjectMeta{ Labels: map[string]string{constants.LabelApplicationID: appIDInLabel}, }, Spec: v1.PodSpec{SchedulerName: constants.SchedulerName}, }, appIDInLabel, appIDInLabel, false}, - {"No AppID defined", &v1.Pod{ + {"AppID defined in spark app selector label", &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{constants.SparkLabelAppID: appIDInSelector}, + }, + Spec: v1.PodSpec{SchedulerName: constants.SchedulerName}, + }, appIDInSelector, appIDInSelector, false}, + {"AppID defined in canonical label and annotation, canonical label win", &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Annotations: map[string]string{constants.AnnotationApplicationID: appIDInAnnotation}, + Labels: map[string]string{constants.CanonicalLabelApplicationID: appIDInCanonicalLabel}, + }, + Spec: v1.PodSpec{SchedulerName: constants.SchedulerName}, + }, appIDInCanonicalLabel, appIDInCanonicalLabel, false}, + {"AppID defined in annotation and legacy label, annotation win", &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Annotations: map[string]string{constants.AnnotationApplicationID: appIDInAnnotation}, + Labels: map[string]string{constants.LabelApplicationID: appIDInLabel}, + }, + Spec: v1.PodSpec{SchedulerName: constants.SchedulerName}, + }, appIDInAnnotation, appIDInAnnotation, false}, + {"Spark AppID defined in legacy label and spark app selector, legacy label win", &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{ + constants.LabelApplicationID: appIDInLabel, + constants.SparkLabelAppID: appIDInSelector, + }, + }, + Spec: v1.PodSpec{SchedulerName: constants.SchedulerName}, + }, appIDInLabel, appIDInLabel, false}, + {"No AppID defined", &v1.Pod{}, "", "", false}, + {"No AppID defined but not generateUnique", &v1.Pod{ ObjectMeta: metav1.ObjectMeta{ Namespace: "testns", UID: "podUid", @@ -618,7 +660,15 @@ func TestGetApplicationIDFromPod(t *testing.T) { }, Spec: v1.PodSpec{SchedulerName: constants.SchedulerName}, }, "testns-podUid", "", true}, - {"Unique autogen token found with generateUnique", &v1.Pod{ + {"Unique autogen token found with generateUnique in canonical AppId label", &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Namespace: "testns", + UID: "podUid", + Labels: map[string]string{constants.CanonicalLabelApplicationID: "testns-uniqueautogen"}, + }, + Spec: v1.PodSpec{SchedulerName: constants.SchedulerName}, + }, "testns-podUid", "testns-podUid", true}, + {"Unique autogen token found with generateUnique in legacy AppId labels", &v1.Pod{ ObjectMeta: metav1.ObjectMeta{ Namespace: "testns", UID: "podUid", @@ -626,18 +676,18 @@ func TestGetApplicationIDFromPod(t *testing.T) { }, Spec: v1.PodSpec{SchedulerName: constants.SchedulerName}, }, "testns-podUid", "testns-podUid", true}, - {"Non-yunikorn schedulerName", &v1.Pod{ + {"Non-yunikorn schedulerName with canonical AppId label", &v1.Pod{ ObjectMeta: metav1.ObjectMeta{ - Labels: map[string]string{constants.LabelApplicationID: appIDInLabel}, + Labels: map[string]string{constants.CanonicalLabelApplicationID: appIDInCanonicalLabel}, }, Spec: v1.PodSpec{SchedulerName: "default"}, }, "", "", false}, - {"AppID defined in annotation", &v1.Pod{ + {"Non-yunikorn schedulerName with legacy AppId label", &v1.Pod{ ObjectMeta: metav1.ObjectMeta{ - Annotations: map[string]string{constants.AnnotationApplicationID: appIDInAnnotation}, + Labels: map[string]string{constants.LabelApplicationID: appIDInLabel}, }, - Spec: v1.PodSpec{SchedulerName: constants.SchedulerName}, - }, appIDInAnnotation, appIDInAnnotation, false}, + Spec: v1.PodSpec{SchedulerName: "default"}, + }, "", "", false}, {"AppID defined but ignore-application set", &v1.Pod{ ObjectMeta: metav1.ObjectMeta{ Annotations: map[string]string{ @@ -656,41 +706,6 @@ func TestGetApplicationIDFromPod(t *testing.T) { }, Spec: v1.PodSpec{SchedulerName: constants.SchedulerName}, }, appIDInAnnotation, appIDInAnnotation, false}, - {"AppID defined in label and annotation", &v1.Pod{ - ObjectMeta: metav1.ObjectMeta{ - Annotations: map[string]string{constants.AnnotationApplicationID: appIDInAnnotation}, - Labels: map[string]string{constants.LabelApplicationID: appIDInLabel}, - }, - Spec: v1.PodSpec{SchedulerName: constants.SchedulerName}, - }, appIDInAnnotation, appIDInAnnotation, false}, - - {"Spark AppID defined in spark app selector", &v1.Pod{ - ObjectMeta: metav1.ObjectMeta{ - Labels: map[string]string{constants.SparkLabelAppID: appIDInSelector}, - }, - Spec: v1.PodSpec{SchedulerName: constants.SchedulerName}, - }, appIDInSelector, appIDInSelector, false}, - {"Spark AppID defined in spark app selector and annotation", &v1.Pod{ - ObjectMeta: metav1.ObjectMeta{ - Labels: map[string]string{constants.SparkLabelAppID: appIDInSelector}, - Annotations: map[string]string{constants.AnnotationApplicationID: sparkIDInAnnotation}, - }, - Spec: v1.PodSpec{SchedulerName: constants.SchedulerName}, - }, sparkIDInAnnotation, sparkIDInAnnotation, false}, - {"Spark AppID defined in spark app selector and annotation", &v1.Pod{ - ObjectMeta: metav1.ObjectMeta{ - Labels: map[string]string{constants.SparkLabelAppID: appIDInSelector, constants.LabelApplicationID: appIDInLabel}, - Annotations: map[string]string{constants.AnnotationApplicationID: sparkIDInAnnotation}, - }, - Spec: v1.PodSpec{SchedulerName: constants.SchedulerName}, - }, sparkIDInAnnotation, sparkIDInAnnotation, false}, - {"No AppID defined", &v1.Pod{}, "", "", false}, - {"Spark AppID defined in spark app selector and label", &v1.Pod{ - ObjectMeta: metav1.ObjectMeta{ - Labels: map[string]string{constants.SparkLabelAppID: appIDInSelector, constants.LabelApplicationID: appIDInLabel}, - }, - Spec: v1.PodSpec{SchedulerName: constants.SchedulerName}, - }, appIDInLabel, appIDInLabel, false}, } for _, tc := range testCases { @@ -892,6 +907,7 @@ func TestGetUserFromPodAnnotation(t *testing.T) { } func TestGetQueueNameFromPod(t *testing.T) { + queueInCanonicalLabel := "sandboxCanonicalLabel" queueInLabel := "sandboxLabel" queueInAnnotation := "sandboxAnnotation" testCases := []struct { @@ -900,7 +916,25 @@ func TestGetQueueNameFromPod(t *testing.T) { expectedQueue string }{ { - name: "With queue label", + name: "Queue defined in canonical label", + pod: &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{constants.CanonicalLabelQueueName: queueInCanonicalLabel}, + }, + }, + expectedQueue: queueInCanonicalLabel, + }, + { + name: "Queue defined in annotation", + pod: &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Annotations: map[string]string{constants.AnnotationQueueName: queueInAnnotation}, + }, + }, + expectedQueue: queueInAnnotation, + }, + { + name: "Queue defined in legacy label", pod: &v1.Pod{ ObjectMeta: metav1.ObjectMeta{ Labels: map[string]string{constants.LabelQueueName: queueInLabel}, @@ -909,26 +943,27 @@ func TestGetQueueNameFromPod(t *testing.T) { expectedQueue: queueInLabel, }, { - name: "With queue annotation", + name: "Queue defined in canonical label and annotation, canonical label win", pod: &v1.Pod{ ObjectMeta: metav1.ObjectMeta{ + Labels: map[string]string{constants.CanonicalLabelQueueName: queueInCanonicalLabel}, Annotations: map[string]string{constants.AnnotationQueueName: queueInAnnotation}, }, }, - expectedQueue: queueInAnnotation, + expectedQueue: queueInCanonicalLabel, }, { - name: "With queue label and annotation", + name: "Queue defined in annotation and legacy label, annotation win", pod: &v1.Pod{ ObjectMeta: metav1.ObjectMeta{ Labels: map[string]string{constants.LabelQueueName: queueInLabel}, Annotations: map[string]string{constants.AnnotationQueueName: queueInAnnotation}, }, }, - expectedQueue: queueInLabel, + expectedQueue: queueInAnnotation, }, { - name: "Without queue label and annotation", + name: "No queue defined", pod: &v1.Pod{ ObjectMeta: metav1.ObjectMeta{}, },