[YUNIKORN-2504] Support canonical labels and align metadata retrievin…

…g order in shim
chenyulin0719 · Aug 16, 2024 · 36b7e48 · 36b7e48
1 parent fabbce3
commit 36b7e48
Show file tree

Hide file tree

Showing 4 changed files with 127 additions and 69 deletions.
diff --git 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
@@ -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
@@ -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
@@ -587,24 +587,66 @@ 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
 		expectedAppID           string
 		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,26 +660,34 @@ 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",
 				Labels:    map[string]string{constants.LabelApplicationID: "testns-uniqueautogen"},
 			},
 			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{},
 			},