Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[YUNIKORN-2504] Support canonical labels and align metadata retrievin order in shim #64

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions pkg/cache/placeholder.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
8 changes: 4 additions & 4 deletions pkg/cache/placeholder_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
43 changes: 33 additions & 10 deletions pkg/common/utils/utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down Expand Up @@ -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
Expand Down
141 changes: 88 additions & 53 deletions pkg/common/utils/utils_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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{
Expand All @@ -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 {
Expand Down Expand Up @@ -892,6 +907,7 @@ func TestGetUserFromPodAnnotation(t *testing.T) {
}

func TestGetQueueNameFromPod(t *testing.T) {
queueInCanonicalLabel := "sandboxCanonicalLabel"
queueInLabel := "sandboxLabel"
queueInAnnotation := "sandboxAnnotation"
testCases := []struct {
Expand All @@ -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},
Expand All @@ -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{},
},
Expand Down