[YUNIKORN-3003] Don't consider already preempted allocations for preemption

pkg/scheduler/objects/allocation.go

@@ -354,6 +354,12 @@ func (a *Allocation) MarkPreempted() {
 	a.preempted = true
 }
 
+func (a *Allocation) UnmarkPreempted() {
+	a.Lock()
+	defer a.Unlock()
+	a.preempted = false
+}
+
 // IsPreempted returns whether the allocation has been marked for preemption or not.
 func (a *Allocation) IsPreempted() bool {
 	a.RLock()

pkg/scheduler/objects/queue.go

@@ -1885,6 +1885,11 @@ func (sq *Queue) findEligiblePreemptionVictims(results map[string]*QueuePreempti
 					continue
 				}
 
+				// skip allocs which have already been preempted
+				if alloc.IsPreempted() {
+					continue
+				}
+
 				// if we have encountered a fence then all tasks are eligible for preemption
 				// otherwise the task is a candidate if its priority is less than or equal to the ask priority
 				if fenced || int64(alloc.GetPriority()) <= askPriority {

pkg/scheduler/objects/queue_test.go

@@ -1886,7 +1886,8 @@
 
 	// verify no victims when no allocations exist
 	snapshot := leaf1.FindEligiblePreemptionVictims(leaf1.QueuePath, ask)
-	assert.Equal(t, 5, len(snapshot), "wrong snapshot count") // leaf1, parent1, root
+
+	assert.Equal(t, 5, len(snapshot), "wrong snapshot count") // root, root.parent1, root.parent1.leaf1, root.parent2, root.parent2.leaf2
 	assert.Equal(t, 0, len(victims(snapshot)), "found victims")
 
 	// add two lower-priority allocs in leaf2
@@ -1974,6 +1975,13 @@
 	assert.Equal(t, alloc3.allocationKey, victims(snapshot)[0].allocationKey, "wrong alloc")
 	alloc2.released = false
 
+	//alloc2 has already been preempted and should not be considered a valid victim
+	alloc2.MarkPreempted()
+	snapshot = leaf1.FindEligiblePreemptionVictims(leaf1.QueuePath, ask)
+	assert.Equal(t, 1, len(victims(snapshot)), "wrong victim count")
+	assert.Equal(t, alloc3.allocationKey, victims(snapshot)[0].allocationKey, "wrong alloc")
+	alloc2.UnmarkPreempted()
+
 	// setting priority offset on parent2 queue should remove leaf2 victims
 	parent2.priorityOffset = 1001
 	snapshot = leaf1.FindEligiblePreemptionVictims(leaf1.QueuePath, ask)