[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"
@@ -3892,6 +3893,129 @@ 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),
+	}
+	preemptions1 := []mockCommon.Preemption{
+		mockCommon.NewPreemption(true, "ask-2", nodeID1, []string{"ask-3", "ask-1"}, 0, 0),
+	}
+	wrongNodes := make(map[string]int, 1)
+	wrongNodes[nodeID2] = 10
+	rightNodes := make(map[string]int, 1)
+	rightNodes[nodeID1] = 10
+
+	type resInfo struct {
+		res        *resources.Resource
+		createTime time.Time
+	}
+	res1 := resources.NewResourceFromMap(map[string]resources.Quantity{"first": 15})
+	res2 := resources.NewResourceFromMap(map[string]resources.Quantity{"first": 15})
+	res3 := resources.NewResourceFromMap(map[string]resources.Quantity{"first": 5})
+
+	ask1Info := make(map[string]*resInfo)
+	ask1Info["ask-1"] = &resInfo{res1, time.Now().Add(-30 * time.Second)}
+
+	ask2Info := make(map[string]*resInfo)
+	ask2Info["ask-1"] = &resInfo{res3, time.Now().Add(-60 * time.Second)}
+	ask2Info["ask-3"] = &resInfo{res2, time.Now().Add(-30 * time.Second)}
+
+	tests := []struct {
+		name            string
+		asks            map[string]*resInfo
+		mockPlugin      *mockCommon.PreemptionPredicatePlugin
+		result          bool
+		mockPluginError error
+		expectedVictims []string
+	}{
+		{"nil plugin, so no predicate checks", ask1Info, nil, true, nil, []string{"ask-1"}},
+		{"prefilter fails", ask1Info, mockCommon.NewPreemptionPredicatePlugin(preemptions, nil, true, false), false, errors.New("fail"), nil},
+		{"prefilter passes but none of the node from iterator is available in feasible nodes", ask1Info, mockCommon.NewPreemptionPredicatePlugin(preemptions, wrongNodes, false, false), false, nil, nil},
+		{"prefilter pass with expected feasible nodes, filter fails", ask1Info, mockCommon.NewPreemptionPredicatePlugin(preemptions, rightNodes, false, true), false, errors.New("fail"), nil},
+		{"both prefilter and filter passes with correct feasible nodes", ask1Info, mockCommon.NewPreemptionPredicatePlugin(preemptions, rightNodes, false, false), true, nil, []string{"ask-1"}},
+		{"both prefilter and filter passes with correct feasible nodes but predicate checks failed for 1 victim", ask2Info, mockCommon.NewPreemptionPredicatePlugin(preemptions1, rightNodes, false, false), true, nil, []string{"ask-3"}},
+		{"both prefilter and filter passes with empty feasible nodes", ask1Info, mockCommon.NewPreemptionPredicatePlugin(preemptions, nil, false, false), true, nil, []string{"ask-1"}},
+	}
+	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)
+
+			headRoom := resources.NewResourceFromMap(map[string]resources.Quantity{"first": 50})
+			preemptionAttemptsRemaining := 1
+
+			// add an ask
+			askInfo := make(map[string]*Allocation)
+			for k, v := range tt.asks {
+				ask := newAllocationAsk(k, "app-1", v.res)
+				ask.createTime = v.createTime
+				err = app.AddAllocationAsk(ask)
+				assert.NilError(t, err, "could not add ask "+k)
+				askInfo[k] = ask
+
+				result := app.tryAllocate(headRoom, true, 30*time.Second, &preemptionAttemptsRemaining, iterator, iterator, getNode)
+				assert.Equal(t, result.ResultType, Allocated, "could not allocate ask")
+				assert.Equal(t, result.Request.allocationKey, k, "unexpected allocation key")
+			}
+
+			// add ask2 with required node
+			ask2 := newAllocationAsk("ask-2", "app-1", res1)
+			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 {
+				for _, expVictim := range tt.expectedVictims {
+					assert.Assert(t, askInfo[expVictim].IsPreempted(), "ask1 has not been preempted")
+				}
+				assert.Assert(t, ask2.HasTriggeredPreemption(), "ask2 has not triggered preemption")
+			} else {
+				for _, expVictim := range tt.expectedVictims {
+					assert.Assert(t, !askInfo[expVictim].IsPreempted(), "ask1 has not been 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,13 +21,10 @@ package objects
 import (
 	"fmt"
 	"strings"
-	"sync"
 
 	"go.uber.org/zap"
 
 	"github.com/apache/yunikorn-core/pkg/log"
-	"github.com/apache/yunikorn-scheduler-interface/lib/go/api"
-	"github.com/apache/yunikorn-scheduler-interface/lib/go/si"
 )
 
 type predicateCheckResult struct {
@@ -114,44 +111,6 @@ 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()
-	result := &predicateCheckResult{
-		allocationKey:   args.AllocationKey,
-		nodeID:          args.NodeID,
-		success:         false,
-		index:           -1,
-		predicateErrors: make(map[string]int32),
-	}
-	if len(args.PreemptAllocationKeys) == 0 {
-		// normal check; there are sufficient resources to run on this node
-		if err := plugin.Predicates(&si.PredicatesArgs{
-			AllocationKey: args.AllocationKey,
-			NodeID:        args.NodeID,
-			Allocate:      true,
-		}); err == nil {
-			result.success = true
-			result.index = -1
-		} else {
-			log.Log(log.SchedPreemption).Debug("Normal predicate check failed",
-				zap.String("AllocationKey", args.AllocationKey),
-				zap.String("NodeID", args.NodeID),
-				zap.Error(err))
-			result.predicateErrors[err.Error()]++
-		}
-	} else if response := plugin.PreemptionPredicates(args); response != nil {
-		// preemption check; at least one allocation will need preemption
-		result.success = response.GetSuccess()
-		if result.success {
-			result.index = int(response.GetIndex())
-		} else {
-			result.predicateErrors = response.GetErrorMessage()
-		}
-	}
-	ch <- result
-}
-
 func (p *predicateCheckResult) String() string {
 	if p.nodeID == "" {
 		return ""

pkg/scheduler/objects/preemption.go

@@ -407,7 +407,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 <- runPredicateChecks(plugin, args)
+			}()
 		}
 		// wait for completion and close channel
 		go func() {

pkg/scheduler/objects/preemption_test.go

@@ -96,6 +96,10 @@ func creatApp1WithTwoDifferentAllocations(
 	} else {
 		alloc2 = newAllocationWithKey("alloc2", appID1, nodeID1, resources.NewResourceFromMap(app2Rec))
 		alloc2.createTime = ask2.createTime
+		err := app1.AddAllocation(alloc2)
+		if err != nil {
+			return nil, nil, err
+		}
 		if !node1.TryAddAllocation(alloc2) {
 			return nil, nil, fmt.Errorf("node alloc2 failed")
 		}
@@ -133,15 +137,23 @@ 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.applications {
+	for _, v := range queue.GetCopyOfApps() {
 		for _, a := range v.allocations {
-			v.RemoveAllocationAsk(a.allocationKey)
+			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)
@@ -363,7 +375,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),
 	}
 	wrongNodes := make(map[string]int, 1)
 	wrongNodes[nodeID2] = 10
@@ -399,8 +411,9 @@ func TestTryPreemption(t *testing.T) {
 				assert.NilError(t, tt.mockPlugin.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")
+				childQ1.DecPreemptingResource(alloc1.GetAllocatedResource())
 				assert.Equal(t, len(ask3.GetAllocationLog()), 0)
 			} else {
 				assert.Assert(t, result == nil, "no result")
@@ -438,7 +451,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)
@@ -449,7 +462,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(preemptions, nil, false, false)
 	plugins.RegisterSchedulerPlugin(plugin)
@@ -460,12 +473,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)
@@ -493,7 +506,7 @@ func TestTryAllocateDoesNotWaitForPreemptionRMReply(t *testing.T) {
 	ask.allowPreemptOther = true
 
 	plugin := mock.NewPreemptionPredicatePlugin([]mock.Preemption{
-		mock.NewPreemption(true, "alloc3", nodeID1, []string{"alloc1"}, 0, 0),
+		mock.NewPreemption(true, "alloc3", nodeID1, []string{"alloc2"}, 0, 0),
 	}, nil, false, false)
 	plugins.RegisterSchedulerPlugin(plugin)
 	defer plugins.UnregisterSchedulerPlugins()
@@ -753,6 +766,7 @@ func TestTryPreemptionOnQueue(t *testing.T) {
 				assert.Check(t, ok == true, "either node1 or node2 chosen")
 				assert.Check(t, alloc1.IsPreempted() || alloc2.IsPreempted(), "either alloc1 or alloc2 preempted, but not both")
 				assert.Check(t, !alloc1.IsPreempted() || !alloc2.IsPreempted(), "either alloc1 or alloc2 not preempted, but not both")
+				childQ1.DecPreemptingResource(alloc1.GetAllocatedResource())
 				assert.Equal(t, len(ask3.GetAllocationLog()), 0)
 			} else {
 				assert.Assert(t, result == nil, "no result")

pkg/scheduler/objects/preemption_utilities.go

@@ -22,7 +22,12 @@ import (
 	"sort"
 	"time"
 
+	"go.uber.org/zap"
+
 	"github.com/apache/yunikorn-core/pkg/common/resources"
+	"github.com/apache/yunikorn-core/pkg/log"
+	"github.com/apache/yunikorn-scheduler-interface/lib/go/api"
+	"github.com/apache/yunikorn-scheduler-interface/lib/go/si"
 )
 
 var (
@@ -141,3 +146,42 @@ func scoreAllocationBasedOnAsk(allocation *Allocation, ask *resources.Resource)
 	score += allocation.GetAllocatedResource().TypeMatching(ask)
 	return score
 }
+
+func runPredicateChecks(plugin api.ResourceManagerCallback, args *si.PreemptionPredicatesArgs) *predicateCheckResult {
+	result := &predicateCheckResult{
+		allocationKey:   args.AllocationKey,
+		nodeID:          args.NodeID,
+		success:         false,
+		index:           -1,
+		predicateErrors: make(map[string]int32),
+	}
+	if len(args.PreemptAllocationKeys) == 0 {
+		// normal check; there are sufficient resources to run on this node
+		if err := plugin.Predicates(&si.PredicatesArgs{
+			AllocationKey: args.AllocationKey,
+			NodeID:        args.NodeID,
+			Allocate:      true,
+		}); err == nil {
+			result.success = true
+			result.index = -1
+		} else {
+			log.Log(log.SchedPreemption).Debug("Normal predicate check failed",
+				zap.String("AllocationKey", args.AllocationKey),
+				zap.String("NodeID", args.NodeID),
+				zap.Error(err))
+			result.predicateErrors[err.Error()]++
+		}
+	} else if response := plugin.PreemptionPredicates(args); response != nil {
+		// preemption check; at least one allocation will need preemption
+		result.success = response.GetSuccess()
+		if result.success {
+			result.index = int(response.GetIndex())
+		} else {
+			log.Log(log.SchedPreemption).Debug("Preemption predicate check failed",
+				zap.String("AllocationKey", args.AllocationKey),
+				zap.String("NodeID", args.NodeID))
+			result.predicateErrors = response.GetErrorMessage()
+		}
+	}
+	return result
+}