[YUNIKORN-3228] Assess the impact of restart on scheduled preemption and propose a solution (#1068)

Allow existing allocations coming out of recovery process, VPA for existing running pods and ask transitioned into allocations as usual even if this put up the usage over the queue max resources. Either already scheduled quota preemption or configured quota preemption but completed would bring down the usage once the delay expires.

Closes: #1068

Signed-off-by: mani <manirajv06@gmail.com>

Author: manirajv06

pkg/scheduler/objects/application.go

@@ -685,7 +685,7 @@ func (sa *Application) AddAllocationAsk(ask *Allocation) error {
 
 // UpdateAllocationResources updates the app, queue, and user tracker with deltas for an allocation.
 // If an existing allocation cannot be found or alloc is invalid, an error is returned.
-func (sa *Application) UpdateAllocationResources(alloc *Allocation) error {
+func (sa *Application) UpdateAllocationResources(alloc *Allocation, isQuotaPreemptionEnabled bool) error {
 	sa.Lock()
 	defer sa.Unlock()
 	if alloc == nil {
@@ -711,7 +711,7 @@ func (sa *Application) UpdateAllocationResources(alloc *Allocation) error {
 		// update allocated resources
 		sa.allocatedResource = resources.Add(sa.allocatedResource, delta)
 		sa.allocatedResource.Prune()
-		sa.queue.IncAllocatedResource(delta)
+		sa.queue.IncAllocatedResource(delta, isQuotaPreemptionEnabled)
 
 		// update user usage
 		sa.incUserResourceUsage(delta)

pkg/scheduler/objects/application_test.go

@@ -1054,22 +1054,22 @@ func TestUpdateAllocationResourcePending(t *testing.T) {
 	assert.Check(t, resources.Equals(res, queue.GetPendingResource()), "resources not on queue")
 
 	// check nil alloc update
-	err = app.UpdateAllocationResources(nil)
+	err = app.UpdateAllocationResources(nil, false)
 	assert.Check(t, err != nil, "error not returned on nil alloc")
 
 	// check zero alloc update
 	zero := newAllocationWithKey(alloc, appID1, nodeID1, resources.NewResource())
-	err = app.UpdateAllocationResources(zero)
+	err = app.UpdateAllocationResources(zero, false)
 	assert.Check(t, err != nil, "error not returned on zero alloc")
 
 	// check missing alloc update
 	missing := newAllocationWithKey("missing", appID1, nodeID1, res)
-	err = app.UpdateAllocationResources(missing)
+	err = app.UpdateAllocationResources(missing, false)
 	assert.Check(t, err != nil, "error not returned on missing alloc")
 
 	// check zero delta
 	same := newAllocationWithKey(alloc, appID1, nodeID1, res)
-	err = app.UpdateAllocationResources(same)
+	err = app.UpdateAllocationResources(same, false)
 	assert.NilError(t, err, "error returned on same alloc size")
 	assert.Check(t, resources.Equals(res, app.GetPendingResource()), "resources not on app")
 	assert.Check(t, resources.Equals(res, queue.GetPendingResource()), "resources not on queue")
@@ -1078,7 +1078,7 @@ func TestUpdateAllocationResourcePending(t *testing.T) {
 	res2, err := resources.NewResourceFromConf(map[string]string{"first": "3"})
 	assert.NilError(t, err, "failed to create resource with error")
 	inc := newAllocationWithKey(alloc, appID1, nodeID1, res2)
-	err = app.UpdateAllocationResources(inc)
+	err = app.UpdateAllocationResources(inc, false)
 	assert.NilError(t, err, "error returned on incremented alloc")
 	assert.Check(t, resources.Equals(res2, app.GetPendingResource()), "resources not updated on app")
 	assert.Check(t, resources.Equals(res2, queue.GetPendingResource()), "resources not updated on queue")
@@ -1096,30 +1096,30 @@ func TestUpdateAllocationResourceAllocated(t *testing.T) {
 	res, err := resources.NewResourceFromConf(map[string]string{"first": "2"})
 	assert.NilError(t, err, "failed to create resource with error")
 	alloc1 := newAllocationWithKey(alloc, appID1, nodeID1, res)
-	queue.IncAllocatedResource(res)
+	queue.IncAllocatedResource(res, false)
 	app.RecoverAllocationAsk(alloc1)
 	app.AddAllocation(alloc1)
 	assert.Check(t, resources.Equals(res, queue.GetAllocatedResource()), "resources not on queue")
 
 	// check nil alloc update
-	err = app.UpdateAllocationResources(nil)
+	err = app.UpdateAllocationResources(nil, false)
 	assert.Check(t, err != nil, "error not returned on nil alloc")
 
 	// check zero alloc update
 	zero := newAllocationWithKey(alloc, appID1, nodeID1, resources.NewResource())
-	err = app.UpdateAllocationResources(zero)
+	err = app.UpdateAllocationResources(zero, false)
 	assert.Check(t, err != nil, "error not returned on zero alloc")
 
 	// check missing alloc update
 	missing := newAllocationWithKey("missing", appID1, nodeID1, res)
-	err = app.UpdateAllocationResources(missing)
+	err = app.UpdateAllocationResources(missing, false)
 	assert.Check(t, err != nil, "error not returned on missing alloc")
 	assert.Check(t, resources.Equals(res, app.GetAllocatedResource()), "resources not on app")
 	assert.Check(t, resources.Equals(res, queue.GetAllocatedResource()), "resources not on queue")
 
 	// check zero delta
 	same := newAllocationWithKey(alloc, appID1, nodeID1, res)
-	err = app.UpdateAllocationResources(same)
+	err = app.UpdateAllocationResources(same, false)
 	assert.NilError(t, err, "error returned on same alloc size")
 	assert.Check(t, resources.Equals(res, app.GetAllocatedResource()), "resources not on app")
 	assert.Check(t, resources.Equals(res, queue.GetAllocatedResource()), "resources not on queue")
@@ -1128,7 +1128,7 @@ func TestUpdateAllocationResourceAllocated(t *testing.T) {
 	res2, err := resources.NewResourceFromConf(map[string]string{"first": "3"})
 	assert.NilError(t, err, "failed to create resource with error")
 	inc := newAllocationWithKey(alloc, appID1, nodeID1, res2)
-	err = app.UpdateAllocationResources(inc)
+	err = app.UpdateAllocationResources(inc, false)
 	assert.NilError(t, err, "error returned on incremented alloc")
 	assert.Check(t, resources.Equals(res2, app.GetAllocatedResource()), "resources not updated on app")
 	assert.Check(t, resources.Equals(res2, queue.GetAllocatedResource()), "resources not updated on queue")

pkg/scheduler/objects/preemption_utilities_test.go

@@ -149,7 +149,7 @@ func assignAllocationsToQueue(allocations []*Allocation, queue *Queue) {
 			app = queue.applications[allocation.applicationID]
 		}
 		app.AddAllocation(allocation)
-		queue.IncAllocatedResource(allocation.GetAllocatedResource())
+		queue.IncAllocatedResource(allocation.GetAllocatedResource(), false)
 	}
 }
 

pkg/scheduler/objects/queue.go

@@ -496,6 +496,15 @@ func (sq *Queue) setPreemptionTime(oldMaxResource *resources.Resource, oldDelay
 	}
 }
 
+// ResetPreemptionTime reset the quota preemption variables
+// Only for testing
+func (sq *Queue) ResetPreemptionTime() {
+	sq.Lock()
+	defer sq.Unlock()
+	sq.quotaPreemptionStartTime = time.Time{}
+	sq.quotaPreemptionDelay = 0
+}
+
 // shouldTriggerPreemption returns true if quota preemption should be triggered based on the settings and the
 // current time.
 func (sq *Queue) shouldTriggerPreemption() bool {
@@ -1259,20 +1268,39 @@ func (sq *Queue) TryIncAllocatedResource(alloc *resources.Resource) error {
 }
 
 // IncAllocatedResource increments the allocated resources for this queue (recursively). No queue limits are checked.
-func (sq *Queue) IncAllocatedResource(alloc *resources.Resource) {
+// In case any quota preemption already scheduled or any quota preemption delay is configured, set the same delay again
+// so that any usage overflow over the max quota caused by alloc would be brought down when the delay expires.
+func (sq *Queue) IncAllocatedResource(alloc *resources.Resource, isQuotaPreemptionEnabled bool) {
 	// fall through if nil
 	if sq == nil {
 		return
 	}
 
 	// update parent
-	sq.parent.IncAllocatedResource(alloc)
+	sq.parent.IncAllocatedResource(alloc, isQuotaPreemptionEnabled)
 
 	// update this queue
 	sq.Lock()
 	defer sq.Unlock()
 	sq.allocatedResource = resources.Add(sq.allocatedResource, alloc)
 	sq.updateAllocatedResourceMetrics()
+
+	// Should apply quota preemption based on the config?
+	if !isQuotaPreemptionEnabled ||
+		!sq.quotaPreemptionStartTime.IsZero() ||
+		!sq.isManaged ||
+		sq.quotaPreemptionDelay == 0 ||
+		resources.IsZero(sq.maxResource) ||
+		sq.maxResource.StrictlyGreaterThanOrEqualsOnlyExisting(sq.allocatedResource) {
+		return
+	}
+	// Override the earlier set quota preemption time with the configured delay.
+	// Delay clock ticking from now.
+	sq.quotaPreemptionStartTime = time.Now().Add(sq.quotaPreemptionDelay)
+	log.Log(log.SchedQueue).Info("Overridden quota preemption time",
+		zap.String("queue", sq.QueuePath),
+		zap.Duration("delay", sq.quotaPreemptionDelay),
+		zap.Time("quotaPreemptionStartTime", sq.quotaPreemptionStartTime))
 }
 
 // allocatedResFits adds the passed in resource to the allocatedResource of the queue and checks if it still fits in the

pkg/scheduler/objects/queue_test.go

@@ -668,8 +668,8 @@ func TestHeadroom(t *testing.T) {
 	var res *resources.Resource
 	res, err = resources.NewResourceFromConf(map[string]string{"first": "5", "second": "3"})
 	assert.NilError(t, err, "failed to create resource")
-	leaf1.IncAllocatedResource(res)
-	leaf2.IncAllocatedResource(res)
+	leaf1.IncAllocatedResource(res, false)
+	leaf2.IncAllocatedResource(res, false)
 
 	// headRoom root should be this (max 20-10 - alloc 10-6)
 	res, err = resources.NewResourceFromConf(map[string]string{"first": "10", "second": "4"})
@@ -726,10 +726,10 @@ func TestHeadroomMerge(t *testing.T) {
 	var res *resources.Resource
 	res, err = resources.NewResourceFromConf(map[string]string{"first": "5", "second": "5", "third": "5"})
 	assert.NilError(t, err, "failed to create resource")
-	leaf1.IncAllocatedResource(res)
+	leaf1.IncAllocatedResource(res, false)
 	res, err = resources.NewResourceFromConf(map[string]string{"third": "5", "fourth": "5"})
 	assert.NilError(t, err, "failed to create resource")
-	leaf2.IncAllocatedResource(res)
+	leaf2.IncAllocatedResource(res, false)
 
 	// root headroom should be nil
 	headRoom := root.getHeadRoom()
@@ -796,8 +796,8 @@ func TestMaxHeadroomNoMax(t *testing.T) {
 	var res *resources.Resource
 	res, err = resources.NewResourceFromConf(map[string]string{"first": "5", "second": "3"})
 	assert.NilError(t, err, "failed to create resource")
-	leaf1.IncAllocatedResource(res)
-	leaf2.IncAllocatedResource(res)
+	leaf1.IncAllocatedResource(res, false)
+	leaf2.IncAllocatedResource(res, false)
 
 	headRoom = root.getMaxHeadRoom()
 	assert.Assert(t, headRoom == nil, "headRoom of root should be nil because no max set for all queues")
@@ -834,8 +834,8 @@ func TestMaxHeadroomMax(t *testing.T) {
 	var res *resources.Resource
 	res, err = resources.NewResourceFromConf(map[string]string{"first": "5", "second": "3"})
 	assert.NilError(t, err, "failed to create resource")
-	leaf1.IncAllocatedResource(res)
-	leaf2.IncAllocatedResource(res)
+	leaf1.IncAllocatedResource(res, false)
+	leaf2.IncAllocatedResource(res, false)
 
 	// root headroom should be nil
 	headRoom := root.getMaxHeadRoom()
@@ -1475,7 +1475,7 @@ func TestOutStandingRequestMultipleChildrenWithMax(t *testing.T) {
 	leaf1App := newApplication("app-leaf1", "default", "root.parent.leaf1")
 	leaf1App.SetQueue(leaf1)
 	leaf1.AddApplication(leaf1App)
-	leaf1.IncAllocatedResource(allocatedRes)
+	leaf1.IncAllocatedResource(allocatedRes, false)
 	// use priority = 1000 for this ask to force ordering of queues when sorting
 	askLeaf1 := newAllocationAskAll("ask-leaf1", "app-leaf1", "", askRes, false, 1000)
 	askLeaf1.SetSchedulingAttempted(true)
@@ -1493,7 +1493,7 @@ func TestOutStandingRequestMultipleChildrenWithMax(t *testing.T) {
 	err = leaf2App.AddAllocationAsk(ask2Leaf2)
 	assert.NilError(t, err, "could not add ask")
 	leaf2.AddApplication(leaf2App)
-	leaf2.IncAllocatedResource(allocatedRes)
+	leaf2.IncAllocatedResource(allocatedRes, false)
 
 	outstanding := root.GetOutstandingRequests()
 	assert.Equal(t, 2, len(outstanding), "expected 2 outstanding requests to be collected")
@@ -2374,59 +2374,60 @@ func TestQuotaPreemptionSettings(t *testing.T) {
 	assert.NilError(t, err, "failed to create basic queue: %v", err)
 
 	testCases := []struct {
-		name       string
-		maxRes     map[string]string
-		conf       configs.QueueConfig
-		oldDelay   time.Duration
-		timeChange bool
+		name                       string
+		maxRes                     map[string]string
+		conf                       configs.QueueConfig
+		oldDelay                   time.Duration
+		timeChange                 bool
+		shouldApplyQuotaPreemption bool
 	}{
 		{"clearing max",
 			map[string]string{"memory": "500"},
 			configs.QueueConfig{
 				Resources: configs.Resources{
 					Max: nil,
 				},
-			}, 0, false},
+			}, 0, false, false},
 		{"clearing max with delay",
 			map[string]string{"memory": "500"},
 			configs.QueueConfig{
 				Resources: configs.Resources{
 					Max: nil,
 				},
 				Properties: map[string]string{configs.QuotaPreemptionDelay: "50ms"},
-			}, 0, false},
+			}, 0, false, false},
 		{"incorrect delay",
 			map[string]string{"memory": "500"},
 			configs.QueueConfig{
 				Resources: configs.Resources{
 					Max: nil,
 				},
 				Properties: map[string]string{configs.QuotaPreemptionDelay: "-50s"},
-			}, 0, false},
+			}, 0, false, false},
 		{"increase max with delay",
 			map[string]string{"memory": "500"},
 			configs.QueueConfig{
 				Resources: configs.Resources{
 					Max: map[string]string{"memory": "1000"},
 				},
 				Properties: map[string]string{configs.QuotaPreemptionDelay: "50ms"},
-			}, 50 * time.Millisecond, false},
+			}, 50 * time.Millisecond, false, false},
 		{"decrease max with delay",
 			map[string]string{"memory": "500"},
 			configs.QueueConfig{
 				Resources: configs.Resources{
 					Max: map[string]string{"memory": "100"},
 				},
 				Properties: map[string]string{configs.QuotaPreemptionDelay: "50ms"},
-			}, 50 * time.Millisecond, true},
+			}, 50 * time.Millisecond, true, true},
 		{"delay changed from 0 no max change",
 			map[string]string{"memory": "500"},
 			configs.QueueConfig{
 				Resources: configs.Resources{
 					Max: map[string]string{"memory": "500"},
 				},
 				Properties: map[string]string{configs.QuotaPreemptionDelay: "50ms"},
-			}, 0, true},
+			}, 0, true, true},
 	}
 
 	var oldMax *resources.Resource
@@ -3266,3 +3267,54 @@ func TestQueue_setPreemptionTime(t *testing.T) {
 		})
 	}
 }
+
+func TestOverrideAndResetPreemptionTime(t *testing.T) {
+	root, e := createRootQueue(nil)
+	assert.NilError(t, e, "failed to create basic root queue")
+	tests := []struct {
+		name               string
+		oldMaxResource     *resources.Resource
+		maxRes             map[string]string
+		currentUsage       *resources.Resource
+		delay              time.Duration
+		timeChange         bool
+		overrideTimeChange bool
+	}{
+		{"usage lesser than max res, so preemption time is not set. try to override", resources.NewResourceFromMap(map[string]resources.Quantity{"test": 100}), map[string]string{"test": "150"}, resources.NewResourceFromMap(map[string]resources.Quantity{"test": 80}), 5, false, false},
+		{"setting preemption time first, try to override", resources.NewResourceFromMap(map[string]resources.Quantity{"test": 100}), map[string]string{"test": "50"}, resources.NewResourceFromMap(map[string]resources.Quantity{"test": 80}), 10, true, false},
+		{"preemption time did not set first time, try to override", resources.NewResourceFromMap(map[string]resources.Quantity{"test": 100}), map[string]string{"test": "150"}, resources.NewResourceFromMap(map[string]resources.Quantity{"test": 180}), 10, false, true},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			queue, err := createManagedQueue(root, "test", false, tt.maxRes)
+			assert.NilError(t, err, "queue creation failed unexpectedly")
+			queue.quotaPreemptionDelay = tt.delay
+			queue.allocatedResource = tt.currentUsage
+			before := queue.quotaPreemptionStartTime
+			queue.setPreemptionTime(tt.oldMaxResource, tt.delay)
+			after := queue.quotaPreemptionStartTime
+			if tt.timeChange {
+				assert.Assert(t, !before.Equal(after), "time change is expected")
+				assert.Assert(t, !queue.quotaPreemptionStartTime.IsZero(), "schedule is expected")
+			}
+			queue.IncAllocatedResource(resources.NewResourceFromMap(map[string]resources.Quantity{"test": 1}), true)
+			afterOverride := queue.quotaPreemptionStartTime
+			if tt.overrideTimeChange {
+				assert.Assert(t, !after.Equal(afterOverride), "time change is not expected")
+				assert.Assert(t, !queue.quotaPreemptionStartTime.IsZero(), "schedule is expected")
+			}
+			if tt.timeChange && !tt.overrideTimeChange {
+				assert.Assert(t, after.Equal(afterOverride), "time change is not expected")
+				assert.Assert(t, !queue.quotaPreemptionStartTime.IsZero(), "schedule is expected")
+			}
+			if tt.timeChange || tt.overrideTimeChange {
+				queue.ResetPreemptionTime()
+				afterReset := queue.quotaPreemptionStartTime
+				assert.Assert(t, !afterOverride.Equal(afterReset), "time change is expected")
+				assert.Assert(t, afterReset.IsZero(), "time change is expected")
+				assert.Assert(t, queue.quotaPreemptionDelay == 0, "time change is expected")
+				assert.Assert(t, queue.quotaPreemptionStartTime.IsZero(), "schedule is not expected")
+			}
+		})
+	}
+}

pkg/scheduler/objects/quota_preemptor_test.go

@@ -65,7 +65,7 @@ func TestQuotaChangeGetPreemptableResource(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			tc.queue.parent.guaranteedResource = tc.parentGuaranteed
 			tc.queue.maxResource = tc.maxResource
-			tc.queue.IncAllocatedResource(tc.usedResource)
+			tc.queue.IncAllocatedResource(tc.usedResource, false)
 			preemptor := NewQuotaPreemptor(tc.queue)
 			preemptor.setPreemptableResources()
 			assert.Equal(t, resources.Equals(preemptor.preemptableResource, tc.preemptable), true)