[YUNIKORN-1697] [shim] Make namespace annotation to support max appli…

…cations update. (#885) fix lint Address comments Fix golint Closes: #885 Signed-off-by: qzhu <[email protected]>
apache · Aug 2, 2024 · ecb651e · ecb651e
1 parent 138d53a
commit ecb651e
Show file tree

Hide file tree

Showing 5 changed files with 86 additions and 0 deletions.
diff --git a/pkg/cache/context.go b/pkg/cache/context.go
@@ -934,6 +934,8 @@ func (ctx *Context) notifyTaskComplete(appID, taskID string) {
 // adds the following tags to the request based on annotations (if exist):
 //   - namespace.resourcequota
 //   - namespace.parentqueue
+//   - namespace.resourceguaranteed
+//   - namespace.resourcemaxapps
 func (ctx *Context) updateApplicationTags(request *AddApplicationRequest, namespace string) {
 	namespaceObj := ctx.getNamespaceObject(namespace)
 	if namespaceObj == nil {
@@ -955,6 +957,12 @@ func (ctx *Context) updateApplicationTags(request *AddApplicationRequest, namesp
 		}
 	}
 
+	// add maxApps resource info as an app tag
+	maxApps := utils.GetNamespaceMaxAppsFromAnnotation(namespaceObj)
+	if maxApps != "" {
+		request.Metadata.Tags[siCommon.AppTagNamespaceResourceMaxApps] = maxApps
+	}
+
 	// add parent queue info as an app tag
 	parentQueue := utils.GetNameSpaceAnnotationValue(namespaceObj, constants.AnnotationParentQueue)
 	if parentQueue != "" {

diff --git a/pkg/cache/context_test.go b/pkg/cache/context_test.go
@@ -1522,6 +1522,7 @@ func TestAddApplicationsWithTags(t *testing.T) {
 				constants.NamespaceQuota:                 "{\"cpu\": \"1\", \"memory\": \"256M\", \"nvidia.com/gpu\": \"1\"}",
 				constants.DomainYuniKorn + "parentqueue": "root.test",
 				constants.NamespaceGuaranteed:            "{\"cpu\": \"1\", \"memory\": \"256M\", \"nvidia.com/gpu\": \"1\"}",
+				constants.NamespaceMaxApps:               "1000",
 			},
 		},
 	}
@@ -1610,6 +1611,12 @@ func TestAddApplicationsWithTags(t *testing.T) {
 		t.Fatalf("resource parsing failed")
 	}
 
+	maxApps, ok := request.Metadata.Tags[siCommon.AppTagNamespaceResourceMaxApps]
+	if !ok {
+		t.Fatalf("max apps tag is not updated from the namespace")
+	}
+	assert.Equal(t, maxApps, "1000")
+
 	parentQueue, ok := request.Metadata.Tags[constants.AppTagNamespaceParentQueue]
 	if !ok {
 		t.Fatalf("parent queue tag is not updated from the namespace")

diff --git a/pkg/common/constants/constants.go b/pkg/common/constants/constants.go
@@ -97,6 +97,9 @@ const NamespaceQuota = DomainYuniKorn + "namespace.quota"
 // NamespaceGuaranteed Namespace Guaranteed
 const NamespaceGuaranteed = DomainYuniKorn + "namespace.guaranteed"
 
+// NamespaceMaxApps Namespace Max Apps
+const NamespaceMaxApps = DomainYuniKorn + "namespace.maxApps"
+
 // AnnotationAllowPreemption set on PriorityClass, opt out of preemption for pods with this priority class
 const AnnotationAllowPreemption = DomainYuniKorn + "allow-preemption"
 

diff --git a/pkg/common/utils/utils.go b/pkg/common/utils/utils.go
@@ -214,6 +214,27 @@ func GetNamespaceGuaranteedFromAnnotation(namespaceObj *v1.Namespace) *si.Resour
 	return nil
 }
 
+// get namespace max apps from namespace annotation
+func GetNamespaceMaxAppsFromAnnotation(namespaceObj *v1.Namespace) string {
+	if maxApps := GetNameSpaceAnnotationValue(namespaceObj, constants.NamespaceMaxApps); maxApps != "" {
+		numMaxApp, err := strconv.Atoi(maxApps)
+		if err != nil {
+			log.Log(log.ShimUtils).Warn("Unable to process namespace.maxApps annotation",
+				zap.String("namespace", namespaceObj.Name),
+				zap.String("namespace.maxApps is", maxApps))
+			return ""
+		}
+		if numMaxApp < 0 {
+			log.Log(log.ShimUtils).Warn("Invalid value for namespace.maxApps annotation",
+				zap.String("namespace", namespaceObj.Name),
+				zap.String("namespace.maxApps is", maxApps))
+			return ""
+		}
+		return maxApps
+	}
+	return ""
+}
+
 func GetNamespaceQuotaFromAnnotation(namespaceObj *v1.Namespace) *si.Resource {
 	// retrieve resource quota info from annotations
 	cpuQuota := GetNameSpaceAnnotationValue(namespaceObj, constants.CPUQuota)

diff --git a/pkg/common/utils/utils_test.go b/pkg/common/utils/utils_test.go
@@ -331,6 +331,53 @@ func TestGetNamespaceGuaranteedFromAnnotation(t *testing.T) {
 	}
 }
 
+func TestGetNamespaceMaxAppsFromAnnotation(t *testing.T) {
+	testCases := []struct {
+		namespace      *v1.Namespace
+		expectedMaxApp string
+	}{
+		{&v1.Namespace{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      "test",
+				Namespace: "test",
+			},
+		}, ""},
+		{&v1.Namespace{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      "test",
+				Namespace: "test",
+				Annotations: map[string]string{
+					constants.NamespaceMaxApps: "5",
+				},
+			},
+		}, "5"},
+		{&v1.Namespace{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      "test",
+				Namespace: "test",
+				Annotations: map[string]string{
+					constants.NamespaceMaxApps: "-5",
+				},
+			},
+		}, ""},
+		{&v1.Namespace{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      "test",
+				Namespace: "test",
+				Annotations: map[string]string{
+					constants.NamespaceMaxApps: "error",
+				},
+			},
+		}, ""},
+	}
+	for _, tc := range testCases {
+		t.Run(fmt.Sprintf("namespace: %v", tc.namespace), func(t *testing.T) {
+			maxApp := GetNamespaceMaxAppsFromAnnotation(tc.namespace)
+			assert.Equal(t, maxApp, tc.expectedMaxApp)
+		})
+	}
+}
+
 func TestGetNamespaceQuotaFromAnnotationUsingNewAndOldAnnotations(t *testing.T) {
 	testCases := []struct {
 		namespace        *v1.Namespace