[YUNIKORN-2432] Add unit test coverage for UserTracker/GroupTracker/Q…

…ueueTracker.canRunApp() (apache#812) Closes: apache#812 Signed-off-by: Peter Bacsko <[email protected]>
pbacsko · Feb 26, 2024 · ce6adb3 · ce6adb3
1 parent 908d1cb
commit ce6adb3
Show file tree

Hide file tree

Showing 3 changed files with 93 additions and 7 deletions.
diff --git a/pkg/scheduler/ugm/group_tracker_test.go b/pkg/scheduler/ugm/group_tracker_test.go
@@ -265,6 +265,34 @@ func TestGTSetAndClearMaxLimits(t *testing.T) {
 	assert.Equal(t, si.EventRecord_UG_GROUP_LIMIT, eventSystem.Events[1].EventChangeDetail)
 }
 
+func TestGTCanRunApp(t *testing.T) {
+	manager := GetUserManager()
+	defer manager.ClearConfigLimits()
+	user := security.UserGroup{User: "test", Groups: []string{"test"}}
+	groupTracker := newGroupTracker(user.User, newUGMEvents(mock.NewEventSystemDisabled()))
+	manager.userWildCardLimitsConfig["root.parent"] = &LimitConfig{
+		maxResources: resources.NewResourceFromMap(map[string]resources.Quantity{
+			"cpu": 1000,
+		}),
+		maxApplications: 3,
+	} // manipulate map directly to avoid hierarchy creation
+
+	hierarchy1 := strings.Split(path1, configs.DOT)
+	assert.Assert(t, groupTracker.canRunApp(hierarchy1, TestApp1))
+	// make sure wildcard limits are not applied due to the tracker type
+	assert.Equal(t, uint64(0), groupTracker.queueTracker.childQueueTrackers["parent"].maxRunningApps)
+	assert.Assert(t, groupTracker.queueTracker.childQueueTrackers["parent"].maxResources == nil)
+	assert.Assert(t, !groupTracker.queueTracker.childQueueTrackers["parent"].useWildCard)
+
+	// maxApps limit hit
+	groupTracker.setLimits(path1, nil, 1)
+	groupTracker.increaseTrackedResource(path1, TestApp1, resources.NewResourceFromMap(map[string]resources.Quantity{
+		"cpu": 1000,
+	}), user.User)
+	assert.Assert(t, groupTracker.canRunApp(hierarchy1, TestApp1))
+	assert.Assert(t, !groupTracker.canRunApp(hierarchy1, TestApp2))
+}
+
 func getGroupResource(gt *GroupTracker) map[string]*resources.Resource {
 	resources := make(map[string]*resources.Resource)
 	usage := gt.GetGroupResourceUsageDAOInfo()

diff --git a/pkg/scheduler/ugm/queue_tracker_test.go b/pkg/scheduler/ugm/queue_tracker_test.go
@@ -249,23 +249,20 @@ func TestQTQuotaEnforcement(t *testing.T) {
 }
 
 func TestHeadroom(t *testing.T) {
+	GetUserManager()
 	var nilResource *resources.Resource
 	path := "root.parent.leaf"
 	hierarchy := strings.Split(path, configs.DOT)
 
-	// nothing exists make sure the hierarchy gets created
+	// validate that the hierarchy gets created
 	root := newRootQueueTracker(user)
-	root.childQueueTrackers["parent"] = newQueueTracker("root", "parent", user)
+	headroom := root.headroom(hierarchy, user)
+	assert.Equal(t, headroom, nilResource, "auto create: expected nil resource")
 	parent := root.childQueueTrackers["parent"]
 	assert.Assert(t, parent != nil, "parent queue tracker should have been created")
-	parent.childQueueTrackers["leaf"] = newQueueTracker("root.parent", "leaf", user)
 	leaf := parent.childQueueTrackers["leaf"]
 	assert.Assert(t, leaf != nil, "leaf queue tracker should have been created")
 
-	// auto created trackers no max resource set
-	headroom := root.headroom(hierarchy, none)
-	assert.Equal(t, headroom, nilResource, "auto create: expected nil resource")
-
 	// prep resources to set as usage and max
 	usage, err := resources.NewResourceFromConf(map[string]string{"mem": "10M", "vcore": "10"})
 	assert.NilError(t, err, "usage: new resource create returned error")
@@ -313,6 +310,38 @@ func TestHeadroom(t *testing.T) {
 	assert.Assert(t, resources.Equals(headroom, combined), "headroom should be same as combined")
 }
 
+func TestQTCanRunApp(t *testing.T) {
+	GetUserManager()
+
+	// validate that the hierarchy gets created
+	root := newRootQueueTracker(user)
+	hierarchy := strings.Split("root.parent.leaf", configs.DOT)
+	assert.Assert(t, root.canRunApp(hierarchy, TestApp1, user))
+	parent := root.childQueueTrackers["parent"]
+	assert.Assert(t, parent != nil, "parent queue tracker should have been created")
+	leaf := parent.childQueueTrackers["leaf"]
+	assert.Assert(t, leaf != nil, "leaf queue tracker should have been created")
+
+	// limit in the leaf queue
+	leaf.maxRunningApps = 2
+	leaf.runningApplications[TestApp1] = true
+	leaf.runningApplications[TestApp2] = true
+	parent.runningApplications[TestApp1] = true
+	parent.runningApplications[TestApp2] = true
+	root.runningApplications[TestApp1] = true
+	root.runningApplications[TestApp2] = true
+	assert.Assert(t, root.canRunApp(hierarchy, TestApp1, user))
+	assert.Assert(t, root.canRunApp(hierarchy, TestApp2, user))
+	assert.Assert(t, !root.canRunApp(hierarchy, TestApp3, user))
+
+	// limit in the parent queue
+	leaf.maxRunningApps = 0
+	parent.maxRunningApps = 2
+	assert.Assert(t, root.canRunApp(hierarchy, TestApp1, user))
+	assert.Assert(t, root.canRunApp(hierarchy, TestApp2, user))
+	assert.Assert(t, !root.canRunApp(hierarchy, TestApp3, user))
+}
+
 func TestNewQueueTracker(t *testing.T) {
 	manager := GetUserManager()
 	defer manager.ClearConfigLimits()

diff --git a/pkg/scheduler/ugm/user_tracker_test.go b/pkg/scheduler/ugm/user_tracker_test.go
@@ -302,6 +302,35 @@ func TestSetAndClearMaxLimits(t *testing.T) {
 	assert.Equal(t, si.EventRecord_UG_USER_LIMIT, eventSystem.Events[1].EventChangeDetail)
 }
 
+func TestUTCanRunApp(t *testing.T) {
+	manager := GetUserManager()
+	defer manager.ClearConfigLimits()
+	user := security.UserGroup{User: "test", Groups: []string{"test"}}
+	userTracker := newUserTracker(user.User, newUGMEvents(mock.NewEventSystemDisabled()))
+	maxRes := resources.NewResourceFromMap(map[string]resources.Quantity{
+		"cpu": 1000,
+	})
+	manager.userWildCardLimitsConfig["root.parent"] = &LimitConfig{
+		maxResources:    maxRes,
+		maxApplications: 3,
+	} // manipulate map directly to avoid hierarchy creation
+
+	hierarchy1 := strings.Split(path1, configs.DOT)
+	assert.Assert(t, userTracker.canRunApp(hierarchy1, TestApp1))
+	// make sure wildcard limits are applied
+	assert.Equal(t, uint64(3), userTracker.queueTracker.childQueueTrackers["parent"].maxRunningApps)
+	assert.Assert(t, resources.Equals(maxRes, userTracker.queueTracker.childQueueTrackers["parent"].maxResources))
+	assert.Assert(t, userTracker.queueTracker.childQueueTrackers["parent"].useWildCard)
+
+	// maxApps limit hit
+	userTracker.setLimits(path1, nil, 1, false, false)
+	userTracker.increaseTrackedResource(path1, TestApp1, resources.NewResourceFromMap(map[string]resources.Quantity{
+		"cpu": 1000,
+	}))
+	assert.Assert(t, userTracker.canRunApp(hierarchy1, TestApp1))
+	assert.Assert(t, !userTracker.canRunApp(hierarchy1, TestApp2))
+}
+
 func getUserResource(ut *UserTracker) map[string]*resources.Resource {
 	resources := make(map[string]*resources.Resource)
 	usage := ut.GetUserResourceUsageDAOInfo()