Skip to content

Commit

Permalink
[YUNIKORN-2631] Support canonical labels for queue/applicationId in A…
Browse files Browse the repository at this point in the history 
…dmission Controller
  • Loading branch information
@chenyulin0719
chenyulin0719 committed May 17, 2024
1 parent 5f80f49 commit 88f9adf
Show file tree
Hide file tree
Showing 5 changed files with 108 additions and 60 deletions.
2 changes: 1 addition & 1 deletion pkg/admission/admission_controller.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -322,7 +322,7 @@ func (c *AdmissionController) processPodUpdate(req *admissionv1.AdmissionRequest

func (c *AdmissionController) shouldProcessAdmissionReview(namespace string, labels map[string]string) bool {
if c.shouldProcessNamespace(namespace) &&
(labels[constants.LabelApplicationID] != "" || labels[constants.SparkLabelAppID] != "" || c.shouldLabelNamespace(namespace)) {
(labels[constants.CanonicalLabelApplicationID] != "" || labels[constants.LabelApplicationID] != "" || labels[constants.SparkLabelAppID] != "" || c.shouldLabelNamespace(namespace)) {
return true
}

Expand Down
56 changes: 34 additions & 22 deletions pkg/admission/admission_controller_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,10 +81,12 @@ func TestUpdateLabels(t *testing.T) {
assert.Equal(t, patch[0].Op, "add")
assert.Equal(t, patch[0].Path, "/metadata/labels")
if updatedMap, ok := patch[0].Value.(map[string]string); ok {
assert.Equal(t, len(updatedMap), 3)
assert.Equal(t, len(updatedMap), 5)
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.CanonicalLabelQueueName], "root.default")
assert.Equal(t, updatedMap[constants.LabelQueueName], "root.default")
assert.Equal(t, strings.HasPrefix(updatedMap[constants.CanonicalLabelApplicationID], constants.AutoGenAppPrefix), true)
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 +106,8 @@ func TestUpdateLabels(t *testing.T) {
UID: "7f5fd6c5d5",
ResourceVersion: "10654",
Labels: map[string]string{
"random": "random",
"applicationId": "app-0001",
"random": "random",
constants.CanonicalLabelApplicationID: "app-0001",
},
},
Spec: v1.PodSpec{},
Expand All @@ -117,10 +119,12 @@ func TestUpdateLabels(t *testing.T) {
assert.Equal(t, patch[0].Op, "add")
assert.Equal(t, patch[0].Path, "/metadata/labels")
if updatedMap, ok := patch[0].Value.(map[string]string); ok {
assert.Equal(t, len(updatedMap), 3)
assert.Equal(t, len(updatedMap), 5)
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.CanonicalLabelQueueName], "root.default")
assert.Equal(t, updatedMap[constants.LabelQueueName], "root.default")
assert.Equal(t, updatedMap[constants.CanonicalLabelApplicationID], "app-0001")
assert.Equal(t, updatedMap[constants.LabelApplicationID], "app-0001")
} else {
t.Fatal("patch info content is not as expected")
}
Expand All @@ -140,8 +144,8 @@ func TestUpdateLabels(t *testing.T) {
UID: "7f5fd6c5d5",
ResourceVersion: "10654",
Labels: map[string]string{
"random": "random",
"queue": "root.abc",
"random": "random",
constants.CanonicalLabelQueueName: "root.abc",
},
},
Spec: v1.PodSpec{},
Expand All @@ -154,10 +158,12 @@ func TestUpdateLabels(t *testing.T) {
assert.Equal(t, patch[0].Op, "add")
assert.Equal(t, patch[0].Path, "/metadata/labels")
if updatedMap, ok := patch[0].Value.(map[string]string); ok {
assert.Equal(t, len(updatedMap), 3)
assert.Equal(t, len(updatedMap), 5)
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.CanonicalLabelQueueName], "root.abc")
assert.Equal(t, updatedMap[constants.LabelQueueName], "root.abc")
assert.Equal(t, strings.HasPrefix(updatedMap[constants.CanonicalLabelApplicationID], constants.AutoGenAppPrefix), true)
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 @@ -186,9 +192,11 @@ func TestUpdateLabels(t *testing.T) {
assert.Equal(t, patch[0].Op, "add")
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, len(updatedMap), 4)
assert.Equal(t, updatedMap[constants.CanonicalLabelQueueName], "root.default")
assert.Equal(t, updatedMap[constants.LabelQueueName], "root.default")
assert.Equal(t, strings.HasPrefix(updatedMap[constants.CanonicalLabelApplicationID], constants.AutoGenAppPrefix), true)
assert.Equal(t, strings.HasPrefix(updatedMap[constants.LabelApplicationID], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("patch info content is not as expected")
}
Expand All @@ -214,9 +222,11 @@ func TestUpdateLabels(t *testing.T) {
assert.Equal(t, patch[0].Op, "add")
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, len(updatedMap), 4)
assert.Equal(t, updatedMap[constants.CanonicalLabelQueueName], "root.default")
assert.Equal(t, updatedMap[constants.LabelQueueName], "root.default")
assert.Equal(t, strings.HasPrefix(updatedMap[constants.CanonicalLabelApplicationID], constants.AutoGenAppPrefix), true)
assert.Equal(t, strings.HasPrefix(updatedMap[constants.LabelApplicationID], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("patch info content is not as expected")
}
Expand All @@ -240,9 +250,11 @@ func TestUpdateLabels(t *testing.T) {
assert.Equal(t, patch[0].Op, "add")
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, len(updatedMap), 4)
assert.Equal(t, updatedMap[constants.CanonicalLabelQueueName], "root.default")
assert.Equal(t, updatedMap[constants.LabelQueueName], "root.default")
assert.Equal(t, strings.HasPrefix(updatedMap[constants.CanonicalLabelApplicationID], constants.AutoGenAppPrefix), true)
assert.Equal(t, strings.HasPrefix(updatedMap[constants.LabelApplicationID], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("patch info content is not as expected")
}
Expand Down
18 changes: 16 additions & 2 deletions pkg/admission/util.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,26 +37,40 @@ func updatePodLabel(pod *v1.Pod, namespace string, generateUniqueAppIds bool, de
}

sparkAppID := utils.GetPodLabelValue(pod, constants.SparkLabelAppID)
canonicalAppID := utils.GetPodLabelValue(pod, constants.CanonicalLabelApplicationID)
appID := utils.GetPodLabelValue(pod, constants.LabelApplicationID)
if sparkAppID == "" && appID == "" {
if canonicalAppID == "" && sparkAppID == "" && 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}-${POD_UID}
// else
// application ID convention: ${AUTO_GEN_PREFIX}-${NAMESPACE}-${AUTO_GEN_SUFFIX}
generatedID := utils.GenerateApplicationID(namespace, generateUniqueAppIds, string(pod.UID))

result[constants.CanonicalLabelApplicationID] = generatedID
// Deprecated: After 1.7.0, admission controller will only add canonical label if application ID was not set
result[constants.LabelApplicationID] = generatedID
} else if canonicalAppID != "" {
// Deprecated: Added in 1.6.0 for backward compatibility, in case the prior shim version can't handle canonical label
result[constants.LabelApplicationID] = canonicalAppID
}

// if existing label exist, it takes priority over everything else
if _, ok := existingLabels[constants.LabelQueueName]; !ok {
canonicalQueueName := utils.GetPodLabelValue(pod, constants.CanonicalLabelQueueName)
queueName := utils.GetPodLabelValue(pod, constants.LabelQueueName)
if canonicalQueueName == "" && 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.CanonicalLabelQueueName] = defaultQueueName
// Deprecated: After 1.7.0, admission controller will only add canonical label if queue was not set
result[constants.LabelQueueName] = defaultQueueName
}
} else if canonicalQueueName != "" {
// Deprecated: Added in 1.6.0 for backward compatibility, in case the prior shim version can't handle canonical label
result[constants.LabelQueueName] = canonicalQueueName
}

return result
Expand Down
90 changes: 55 additions & 35 deletions pkg/admission/util_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ func createTestingPodWithMeta() *v1.Pod {

func createTestingPodWithAppId() *v1.Pod {
pod := createTestingPodWithMeta()
pod.ObjectMeta.Labels["applicationId"] = "app-0001"
pod.ObjectMeta.Labels[constants.CanonicalLabelApplicationID] = "app-0001"

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

func createTestingPodWithQueue() *v1.Pod {
pod := createTestingPodWithMeta()
pod.ObjectMeta.Labels["queue"] = "root.abc"
pod.ObjectMeta.Labels[constants.CanonicalLabelQueueName] = "root.abc"

return pod
}
Expand All @@ -104,35 +104,41 @@ func TestUpdatePodLabelForAdmissionController(t *testing.T) {
pod := createTestingPodWithMeta()

if result := updatePodLabel(pod, "default", false, "root.default"); result != nil {
assert.Equal(t, len(result), 3)
assert.Equal(t, len(result), 5)
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, result[constants.CanonicalLabelQueueName], "root.default")
assert.Equal(t, result[constants.LabelQueueName], "root.default")
assert.Equal(t, strings.HasPrefix(result[constants.CanonicalLabelApplicationID], constants.AutoGenAppPrefix), true)
assert.Equal(t, strings.HasPrefix(result[constants.LabelApplicationID], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("UpdatePodLabelForAdmissionController is not as expected")
}

// verify if applicationId is given in the labels,
// we won't modify it
// verify if applicationId is given in the canonical labels,
// we won't modify the value and will add it to non-canonical label for backward compatibility
pod = createTestingPodWithAppId()

if result := updatePodLabel(pod, "default", false, "root.default"); result != nil {
assert.Equal(t, len(result), 3)
assert.Equal(t, len(result), 5)
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.CanonicalLabelQueueName], "root.default")
assert.Equal(t, result[constants.LabelQueueName], "root.default")
assert.Equal(t, result[constants.CanonicalLabelApplicationID], "app-0001")
assert.Equal(t, result[constants.LabelApplicationID], "app-0001")
} else {
t.Fatal("UpdatePodLabelForAdmissionController is not as expected")
}

// verify if queue is given in the labels,
// we won't modify it
// we won't modify the value and will add it to non-canonical label for backward compatibility
pod = createTestingPodWithQueue()
if result := updatePodLabel(pod, "default", false, "root.default"); result != nil {
assert.Equal(t, len(result), 3)
assert.Equal(t, len(result), 5)
assert.Equal(t, result["random"], "random")
assert.Equal(t, result["queue"], "root.abc")
assert.Equal(t, strings.HasPrefix(result["applicationId"], constants.AutoGenAppPrefix), true)
assert.Equal(t, result[constants.CanonicalLabelQueueName], "root.abc")
assert.Equal(t, result[constants.LabelQueueName], "root.abc")
assert.Equal(t, strings.HasPrefix(result[constants.CanonicalLabelApplicationID], constants.AutoGenAppPrefix), true)
assert.Equal(t, strings.HasPrefix(result[constants.LabelApplicationID], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("UpdatePodLabelForAdmissionControllert is not as expected")
}
Expand All @@ -142,28 +148,34 @@ func TestUpdatePodLabelForAdmissionController(t *testing.T) {
pod = createTestingPodNoNamespaceAndLabels()

if result := updatePodLabel(pod, "default", false, "root.default"); 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, len(result), 4)
assert.Equal(t, result[constants.CanonicalLabelQueueName], "root.default")
assert.Equal(t, result[constants.LabelQueueName], "root.default")
assert.Equal(t, strings.HasPrefix(result[constants.CanonicalLabelApplicationID], constants.AutoGenAppPrefix), true)
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 {
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, len(result), 4)
assert.Equal(t, result[constants.CanonicalLabelQueueName], "root.default")
assert.Equal(t, result[constants.LabelQueueName], "root.default")
assert.Equal(t, strings.HasPrefix(result[constants.CanonicalLabelApplicationID], constants.AutoGenAppPrefix), true)
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 {
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, len(result), 4)
assert.Equal(t, result[constants.CanonicalLabelQueueName], "root.default")
assert.Equal(t, result[constants.LabelQueueName], "root.default")
assert.Equal(t, strings.HasPrefix(result[constants.CanonicalLabelApplicationID], constants.AutoGenAppPrefix), true)
assert.Equal(t, strings.HasPrefix(result[constants.LabelApplicationID], constants.AutoGenAppPrefix), true)
} else {
t.Fatal("UpdatePodLabelForAdmissionController is not as expected")
}
Expand All @@ -173,10 +185,12 @@ func TestDefaultQueueName(t *testing.T) {
defaultConf := createConfig()
pod := createTestingPodWithMeta()
if result := updatePodLabel(pod, defaultConf.GetNamespace(), defaultConf.GetGenerateUniqueAppIds(), defaultConf.GetDefaultQueueName()); result != nil {
assert.Equal(t, len(result), 3)
assert.Equal(t, len(result), 5)
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.CanonicalLabelApplicationID], "yunikorn-default-autogen")
assert.Equal(t, result[constants.LabelApplicationID], "yunikorn-default-autogen")
assert.Equal(t, result[constants.CanonicalLabelQueueName], "root.default")
assert.Equal(t, result[constants.LabelQueueName], "root.default")
} else {
t.Fatal("UpdatePodLabelForAdmissionController is not as expected")
}
Expand All @@ -185,10 +199,12 @@ func TestDefaultQueueName(t *testing.T) {
conf.AMFilteringDefaultQueueName: "",
})
if result := updatePodLabel(pod, queueNameEmptyConf.GetNamespace(), queueNameEmptyConf.GetGenerateUniqueAppIds(), queueNameEmptyConf.GetDefaultQueueName()); result != nil {
assert.Equal(t, len(result), 2)
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"], "")
assert.Equal(t, result[constants.CanonicalLabelApplicationID], "yunikorn-default-autogen")
assert.Equal(t, result[constants.LabelApplicationID], "yunikorn-default-autogen")
assert.Equal(t, result[constants.CanonicalLabelQueueName], "")
assert.Equal(t, result[constants.LabelQueueName], "")
} else {
t.Fatal("UpdatePodLabelForAdmissionController is not as expected")
}
Expand All @@ -197,10 +213,12 @@ func TestDefaultQueueName(t *testing.T) {
conf.AMFilteringDefaultQueueName: "yunikorn",
})
if result := updatePodLabel(pod, customQueueNameConf.GetNamespace(), customQueueNameConf.GetGenerateUniqueAppIds(), customQueueNameConf.GetDefaultQueueName()); result != nil {
assert.Equal(t, len(result), 3)
assert.Equal(t, len(result), 5)
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.CanonicalLabelApplicationID], "yunikorn-default-autogen")
assert.Equal(t, result[constants.LabelApplicationID], "yunikorn-default-autogen")
assert.Assert(t, result[constants.CanonicalLabelQueueName] != "yunikorn")
assert.Assert(t, result[constants.LabelQueueName] != "yunikorn")
} else {
t.Fatal("UpdatePodLabelForAdmissionController is not as expected")
}
Expand All @@ -210,10 +228,12 @@ func TestDefaultQueueName(t *testing.T) {
})
if result := updatePodLabel(pod, customValidQueueNameConf.GetNamespace(),
customValidQueueNameConf.GetGenerateUniqueAppIds(), customValidQueueNameConf.GetDefaultQueueName()); result != nil {
assert.Equal(t, len(result), 3)
assert.Equal(t, len(result), 5)
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.CanonicalLabelApplicationID], "yunikorn-default-autogen")
assert.Equal(t, result[constants.LabelApplicationID], "yunikorn-default-autogen")
assert.Equal(t, result[constants.CanonicalLabelQueueName], "root.yunikorn")
assert.Equal(t, result[constants.LabelQueueName], "root.yunikorn")
} else {
t.Fatal("UpdatePodLabelForAdmissionController is not as expected")
}
Expand Down
2 changes: 2 additions & 0 deletions pkg/common/constants/constants.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,10 @@ const DomainYuniKornInternal = siCommon.DomainYuniKornInternal

// Application
const LabelApp = "app"
const CanonicalLabelApplicationID = DomainYuniKorn + "app-id"
const LabelApplicationID = "applicationId"
const AnnotationApplicationID = DomainYuniKorn + "app-id"
const CanonicalLabelQueueName = DomainYuniKorn + "queue"
const LabelQueueName = "queue"
const RootQueue = "root"
const AnnotationQueueName = DomainYuniKorn + "queue"
Expand Down

0 comments on commit 88f9adf

Please sign in to comment.