From ce6adb3b81491aa55febe235227683225abcbaad Mon Sep 17 00:00:00 2001 From: Peter Bacsko Date: Mon, 26 Feb 2024 16:40:04 +0100 Subject: [PATCH] [YUNIKORN-2432] Add unit test coverage for UserTracker/GroupTracker/QueueTracker.canRunApp() (#812) Closes: #812 Signed-off-by: Peter Bacsko --- pkg/scheduler/ugm/group_tracker_test.go | 28 ++++++++++++++++ pkg/scheduler/ugm/queue_tracker_test.go | 43 +++++++++++++++++++++---- pkg/scheduler/ugm/user_tracker_test.go | 29 +++++++++++++++++ 3 files changed, 93 insertions(+), 7 deletions(-) diff --git a/pkg/scheduler/ugm/group_tracker_test.go b/pkg/scheduler/ugm/group_tracker_test.go index 08dca63fb..3b0e33bc5 100644 --- 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 index 929b77fd0..23a92c952 100644 --- 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 index 28db413fc..d07c1e219 100644 --- 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()