[YUNIKORN-3330] Run predicate checks in Required node preemption

pkg/scheduler/objects/application_test.go

@@ -19,6 +19,7 @@
 package objects
 
 import (
+	"errors"
 	"fmt"
 	"math"
 	"strings"
@@ -3513,6 +3514,99 @@ func TestRequiredNodePreemption(t *testing.T) {
 	}
 }
 
+func TestRequiredNodePreemptionWithPredicates(t *testing.T) {
+	node := newNode(nodeID1, map[string]resources.Quantity{"first": 20})
+	iterator := getNodeIteratorFn(node)
+	getNode := func(nodeID string) *Node {
+		return node
+	}
+	appQueueMapping := NewAppQueueMapping()
+
+	// set queue
+	rootQ, err := createRootQueue(map[string]string{"first": "20"})
+	assert.NilError(t, err)
+	childQ, err := createManagedQueueWithAppQueueMapping(rootQ, "default", false, map[string]string{"first": "20"}, appQueueMapping)
+	assert.NilError(t, err)
+
+	// register predicate handler
+	preemptions := []mockCommon.Preemption{
+		mockCommon.NewPreemption(true, "ask-2", nodeID1, []string{"ask-1"}, 0, 0),
+	}
+	wrongNodes := make(map[string]int, 1)
+	wrongNodes[nodeID2] = 10
+	rightNodes := make(map[string]int, 1)
+	rightNodes[nodeID1] = 10
+	tests := []struct {
+		name            string
+		mockPlugin      *mockCommon.PreemptionPredicatePlugin
+		result          bool
+		mockPluginError error
+	}{
+		{"nil plugin, so no predicate checks", nil, true, nil},
+		{"prefilter fails", mockCommon.NewPreemptionPredicatePlugin(preemptions, nil, true, false), false, errors.New("fail")},
+		{"prefilter passes but none of the node from iterator is available in feasible nodes", mockCommon.NewPreemptionPredicatePlugin(preemptions, wrongNodes, false, false), false, nil},
+		{"prefilter pass with expected feasible nodes, filter fails", mockCommon.NewPreemptionPredicatePlugin(preemptions, rightNodes, false, true), false, errors.New("fail")},
+		{"both prefilter and filter passes with correct feasible nodes", mockCommon.NewPreemptionPredicatePlugin(preemptions, rightNodes, false, false), true, nil},
+		{"both prefilter and filter passes with empty feasible nodes", mockCommon.NewPreemptionPredicatePlugin(preemptions, nil, false, false), true, nil},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			app := newApplication(appID1, "default", "root.default")
+			app.SetQueue(childQ)
+			childQ.AddApplication(app)
+			appQueueMapping.AddAppQueueMapping(app.ApplicationID, childQ)
+
+			// add an ask
+			askRes := resources.NewResourceFromMap(map[string]resources.Quantity{"first": 15})
+			ask1 := newAllocationAsk("ask-1", "app-1", askRes)
+			err = app.AddAllocationAsk(ask1)
+			assert.NilError(t, err, "could not add ask-1")
+			preemptionAttemptsRemaining := 1
+
+			// allocate ask
+			headRoom := resources.NewResourceFromMap(map[string]resources.Quantity{"first": 50})
+			result := app.tryAllocate(headRoom, true, 30*time.Second, &preemptionAttemptsRemaining, iterator, iterator, getNode)
+			assert.Equal(t, result.ResultType, Allocated, "could not allocate ask-1")
+			assert.Equal(t, result.Request.allocationKey, "ask-1", "unexpected allocation key")
+
+			// add ask2 with required node
+			ask2 := newAllocationAsk("ask-2", "app-1", askRes)
+			ask2.requiredNode = nodeID1
+			err = app.AddAllocationAsk(ask2)
+			assert.NilError(t, err, "could not add ask-2")
+
+			// try to allocate ask2 with node being full - expect a reservation
+			result = app.tryAllocate(headRoom, true, 30*time.Second, &preemptionAttemptsRemaining, iterator, iterator, getNode)
+			assert.Equal(t, result.ResultType, Reserved, "allocation result is not reserved")
+			assert.Equal(t, result.Request.allocationKey, "ask-2", "unexpected allocation key")
+			err = app.Reserve(node, ask2)
+			assert.NilError(t, err, "reservation failed")
+
+			if tt.mockPlugin != nil {
+				plugins.RegisterSchedulerPlugin(tt.mockPlugin)
+			}
+			app.tryReservedAllocate(headRoom, iterator)
+			if tt.result {
+				assert.Assert(t, ask1.IsPreempted(), "ask1 has not been preempted")
+				assert.Assert(t, ask2.HasTriggeredPreemption(), "ask2 has not triggered preemption")
+			} else {
+				assert.Assert(t, !ask1.IsPreempted(), "ask1 preempted")
+				assert.Assert(t, !ask2.HasTriggeredPreemption(), "ask2 triggered preemption")
+			}
+			if tt.mockPluginError != nil {
+				assert.ErrorContains(t, tt.mockPlugin.GetPredicateError(), tt.mockPluginError.Error())
+			}
+			// reset
+			resetQ(t, childQ)
+			resetNode(node)
+			appQueueMapping.RemoveAppQueueMapping(app.ApplicationID)
+			if tt.mockPlugin != nil {
+				plugins.UnregisterSchedulerPlugins()
+			}
+		})
+	}
+}
+
 func TestRequiredNodePreemptionFailed(t *testing.T) {
 	// tests RequiredNode (DaemonSet) preemption where the victim pod has a high priority, hence preemption is not possible
 	app := newApplication(appID0, "default", "root.default")

pkg/scheduler/objects/predicates.go

@@ -21,7 +21,6 @@ package objects
 import (
 	"fmt"
 	"strings"
-	"sync"
 
 	"go.uber.org/zap"
 
@@ -113,9 +112,7 @@ func (pcr *predicateCheckResult) populateVictims(victimsByNode map[string][]*All
 	}
 }
 
-// preemptPredicateCheck performs a single predicate check and reports the resultType on a channel
-func preemptPredicateCheck(plugin api.ResourceManagerCallback, ch chan<- *predicateCheckResult, wg *sync.WaitGroup, args *si.PreemptionPredicatesArgs) {
-	defer wg.Done()
+func PredicateChecks(plugin api.ResourceManagerCallback, args *si.PreemptionPredicatesArgs) *predicateCheckResult {
 	result := &predicateCheckResult{
 		allocationKey: args.AllocationKey,
 		nodeID:        args.NodeID,
@@ -144,7 +141,7 @@ func preemptPredicateCheck(plugin api.ResourceManagerCallback, ch chan<- *predic
 			result.index = int(response.GetIndex())
 		}
 	}
-	ch <- result
+	return result
 }
 
 func (p *predicateCheckResult) String() string {

pkg/scheduler/objects/preemption.go

@@ -444,7 +444,10 @@ func (p *Preemptor) checkPreemptionPredicates(predicateChecks []*si.PreemptionPr
 			// add goroutine for checking preemption
 			wg.Add(1)
 			expected++
-			go preemptPredicateCheck(plugin, ch, &wg, args)
+			go func() {
+				defer wg.Done()
+				ch <- PredicateChecks(plugin, args)
+			}()
 		}
 		// wait for completion and close channel
 		go func() {

pkg/scheduler/objects/preemption_test.go

@@ -87,6 +87,7 @@ func creatApp1WithTwoDifferentAllocations(
 	} else {
 		alloc2 = newAllocationWithKey("alloc2", appID1, nodeID1, resources.NewResourceFromMap(app2Rec))
 		alloc2.createTime = ask2.createTime
+		app1.AddAllocation(alloc2)
 		if !node1.TryAddAllocation(alloc2) {
 			return nil, nil, fmt.Errorf("node alloc2 failed")
 		}
@@ -120,6 +121,32 @@ func creatApp2(
 	return app2, ask3, nil
 }
 
+func resetNode(node *Node) {
+	for _, v := range node.allocations {
+		node.RemoveAllocation(v.allocationKey)
+	}
+	node.reservations = make(map[string]*reservation)
+}
+
+func resetQ(t *testing.T, queue *Queue) {
+	for _, v := range queue.GetCopyOfApps() {
+		for _, a := range v.allocations {
+			v.RemoveAllocation(a.allocationKey, si.TerminationType_STOPPED_BY_RM)
+			err := queue.DecAllocatedResource(a.GetAllocatedResource())
+			assert.NilError(t, err)
+		}
+		for _, req := range v.requests {
+			v.RemoveAllocationAsk(req.allocationKey)
+		}
+		for _, r := range v.reservations {
+			v.unReserveInternal(r)
+		}
+		v.queue.UnReserve(v.ApplicationID, len(v.reservations))
+		queue.RemoveApplication(v)
+	}
+	queue.applications = make(map[string]*Application)
+}
+
 func TestCheckPreconditions(t *testing.T) {
 	node := newNode("node1", map[string]resources.Quantity{"first": 5})
 	iterator := getNodeIteratorFn(node)
@@ -345,7 +372,7 @@ func TestTryPreemption(t *testing.T) {
 
 	// register predicate handler
 	preemptions := []mock.Preemption{
-		mock.NewPreemption(true, "alloc3", nodeID1, []string{"alloc1"}, 0, 0),
+		mock.NewPreemption(true, "alloc3", nodeID1, []string{"alloc2"}, 0, 0),
 	}
 	plugin := mock.NewPreemptionPredicatePlugin(nil, nil, preemptions)
 	plugins.RegisterSchedulerPlugin(plugin)
@@ -356,9 +383,10 @@ func TestTryPreemption(t *testing.T) {
 	assert.NilError(t, plugin.GetPredicateError())
 	assert.Assert(t, ok, "no victims found")
 	assert.Equal(t, "alloc3", result.Request.GetAllocationKey(), "wrong alloc")
-	assert.Check(t, alloc1.IsPreempted(), "alloc1 not preempted")
-	assert.Check(t, !alloc2.IsPreempted(), "alloc2 preempted")
+	assert.Check(t, !alloc1.IsPreempted(), "alloc1 not preempted")
+	assert.Check(t, alloc2.IsPreempted(), "alloc2 preempted")
 	assert.Equal(t, len(ask3.GetAllocationLog()), 0)
+	childQ1.DecPreemptingResource(alloc1.GetAllocatedResource())
 }
 
 func TestTryPreemption_SendEvent(t *testing.T) {
@@ -379,7 +407,7 @@ func TestTryPreemption_SendEvent(t *testing.T) {
 
 	eventSystem := evtMock.NewEventSystem()
 	events := schedEvt.NewAskEvents(eventSystem)
-	alloc1.askEvents = events
+	alloc2.askEvents = events
 
 	app2, ask3, err := creatApp2(childQ2, map[string]resources.Quantity{"first": 5, "pods": 1}, "alloc3", appQueueMapping)
 	assert.NilError(t, err)
@@ -390,7 +418,7 @@ func TestTryPreemption_SendEvent(t *testing.T) {
 
 	// register predicate handler
 	preemptions := []mock.Preemption{
-		mock.NewPreemption(true, "alloc3", nodeID1, []string{"alloc1"}, 0, 0),
+		mock.NewPreemption(true, "alloc3", nodeID1, []string{"alloc2"}, 0, 0),
 	}
 	plugin := mock.NewPreemptionPredicatePlugin(nil, nil, preemptions)
 	plugins.RegisterSchedulerPlugin(plugin)
@@ -401,12 +429,12 @@ func TestTryPreemption_SendEvent(t *testing.T) {
 	assert.NilError(t, plugin.GetPredicateError())
 	assert.Assert(t, ok, "no victims found")
 	assert.Equal(t, "alloc3", result.Request.GetAllocationKey(), "wrong alloc")
-	assert.Check(t, alloc1.IsPreempted(), "alloc1 not preempted")
-	assert.Check(t, !alloc2.IsPreempted(), "alloc2 preempted")
+	assert.Check(t, !alloc1.IsPreempted(), "alloc1 not preempted")
+	assert.Check(t, alloc2.IsPreempted(), "alloc2 preempted")
 	assert.Equal(t, 1, len(eventSystem.Events))
 	event := eventSystem.Events[0]
-	assert.Equal(t, alloc1.applicationID, event.ReferenceID)
-	assert.Equal(t, alloc1.allocationKey, event.ObjectID)
+	assert.Equal(t, alloc2.applicationID, event.ReferenceID)
+	assert.Equal(t, alloc2.allocationKey, event.ObjectID)
 	assert.Equal(t, si.EventRecord_NONE, event.EventChangeType)
 	assert.Equal(t, si.EventRecord_DETAILS_NONE, event.EventChangeDetail)
 	assert.Equal(t, si.EventRecord_REQUEST, event.Type)
@@ -618,6 +646,7 @@ func TestTryPreemptionOnQueue(t *testing.T) {
 	assert.Equal(t, nodeID2, result.NodeID, "wrong node")
 	assert.Check(t, !alloc1.IsPreempted(), "alloc1 preempted")
 	assert.Check(t, alloc2.IsPreempted(), "alloc2 not preempted")
+	childQ1.DecPreemptingResource(alloc1.GetAllocatedResource())
 	assert.Equal(t, len(ask3.GetAllocationLog()), 0)
 }