Skip to content

Commit

Permalink
do not consider previously preempted allocations as potential victims
Browse files Browse the repository at this point in the history
  • Loading branch information
psantacl committed Dec 19, 2024
1 parent 149820f commit fee6a1f
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 1 deletion.
6 changes: 6 additions & 0 deletions pkg/scheduler/objects/allocation.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
5 changes: 5 additions & 0 deletions pkg/scheduler/objects/queue.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
10 changes: 9 additions & 1 deletion pkg/scheduler/objects/queue_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -1886,7 +1886,8 @@ func TestFindEligiblePreemptionVictims(t *testing.T) {

// 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
Expand Down Expand Up @@ -1974,6 +1975,13 @@ func TestFindEligiblePreemptionVictims(t *testing.T) {
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

Check failure on line 1978 in pkg/scheduler/objects/queue_test.go

View workflow job for this annotation

GitHub Actions / build

commentFormatting: put a space between `//` and comment text (gocritic)
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)
Expand Down

0 comments on commit fee6a1f

Please sign in to comment.