From dc3e9fb79270fb5174a44c3e8380a5b3039886cf Mon Sep 17 00:00:00 2001 From: Albin Severinson Date: Tue, 27 Jun 2023 11:53:10 +0100 Subject: [PATCH] Add tests --- .../scheduler/preempting_queue_scheduler.go | 3 +- .../preempting_queue_scheduler_test.go | 116 ++++++++++++++++++ .../scheduler/testfixtures/testfixtures.go | 11 +- 3 files changed, 126 insertions(+), 4 deletions(-) diff --git a/internal/scheduler/preempting_queue_scheduler.go b/internal/scheduler/preempting_queue_scheduler.go index a19aba06b63..636d8caf713 100644 --- a/internal/scheduler/preempting_queue_scheduler.go +++ b/internal/scheduler/preempting_queue_scheduler.go @@ -79,6 +79,7 @@ func NewPreemptingQueueScheduler( constraints: constraints, nodeEvictionProbability: nodeEvictionProbability, nodeOversubscriptionEvictionProbability: nodeOversubscriptionEvictionProbability, + protectedFractionOfFairShare: protectedFractionOfFairShare, jobRepo: jobRepo, nodeDb: nodeDb, nodeIdByJobId: maps.Clone(initialNodeIdByJobId), @@ -155,7 +156,7 @@ func (sch *PreemptingQueueScheduler) Schedule(ctx context.Context) (*SchedulerRe return false } if qctx, ok := sch.schedulingContext.QueueSchedulingContexts[job.GetQueue()]; ok { - if qctx.FractionOfFairShare() < sch.protectedFractionOfFairShare { + if qctx.FractionOfFairShare() <= sch.protectedFractionOfFairShare { return false } } diff --git a/internal/scheduler/preempting_queue_scheduler_test.go b/internal/scheduler/preempting_queue_scheduler_test.go index 17dee635be4..7ab6ae1d3fb 100644 --- a/internal/scheduler/preempting_queue_scheduler_test.go +++ b/internal/scheduler/preempting_queue_scheduler_test.go @@ -1136,6 +1136,122 @@ func TestPreemptingQueueScheduler(t *testing.T) { "B": 1, }, }, + "ProtectedFractionOfFairShare": { + SchedulingConfig: testfixtures.WithProtectedFractionOfFairShareConfig( + 1.0, + testfixtures.TestSchedulingConfig(), + ), + Nodes: testfixtures.N32CpuNodes(1, testfixtures.TestPriorities), + Rounds: []SchedulingRound{ + { + JobsByQueue: map[string][]*jobdb.Job{ + "A": testfixtures.N1CpuJobs("A", testfixtures.PriorityClass0, 10), + }, + ExpectedScheduledIndices: map[string][]int{ + "A": testfixtures.IntRange(0, 9), + }, + }, + { + JobsByQueue: map[string][]*jobdb.Job{ + "B": testfixtures.N1CpuJobs("B", testfixtures.PriorityClass3, 22), + }, + ExpectedScheduledIndices: map[string][]int{ + "B": testfixtures.IntRange(0, 21), + }, + }, + { + JobsByQueue: map[string][]*jobdb.Job{ + "C": testfixtures.N1CpuJobs("C", testfixtures.PriorityClass0, 1), + }, + }, + {}, // Empty round to make sure nothing changes. + }, + PriorityFactorByQueue: map[string]float64{ + "A": 1, + "B": 1, + "C": 1, + }, + }, + "ProtectedFractionOfFairShare at limit": { + SchedulingConfig: testfixtures.WithProtectedFractionOfFairShareConfig( + 0.5, + testfixtures.TestSchedulingConfig(), + ), + Nodes: testfixtures.N32CpuNodes(1, testfixtures.TestPriorities), + Rounds: []SchedulingRound{ + { + JobsByQueue: map[string][]*jobdb.Job{ + "A": testfixtures.N1CpuJobs("A", testfixtures.PriorityClass0, 8), + }, + ExpectedScheduledIndices: map[string][]int{ + "A": testfixtures.IntRange(0, 7), + }, + }, + { + JobsByQueue: map[string][]*jobdb.Job{ + "B": testfixtures.N1CpuJobs("B", testfixtures.PriorityClass3, 24), + }, + ExpectedScheduledIndices: map[string][]int{ + "B": testfixtures.IntRange(0, 23), + }, + }, + { + JobsByQueue: map[string][]*jobdb.Job{ + "C": testfixtures.N1CpuJobs("C", testfixtures.PriorityClass0, 1), + }, + }, + {}, // Empty round to make sure nothing changes. + }, + PriorityFactorByQueue: map[string]float64{ + "A": 0.5, + "B": 1, + "C": 1, + }, + }, + "ProtectedFractionOfFairShare above limit": { + SchedulingConfig: testfixtures.WithProtectedFractionOfFairShareConfig( + 0.5, + testfixtures.TestSchedulingConfig(), + ), + Nodes: testfixtures.N32CpuNodes(1, testfixtures.TestPriorities), + Rounds: []SchedulingRound{ + { + JobsByQueue: map[string][]*jobdb.Job{ + "A": testfixtures.N1CpuJobs("A", testfixtures.PriorityClass0, 9), + }, + ExpectedScheduledIndices: map[string][]int{ + "A": testfixtures.IntRange(0, 8), + }, + }, + { + JobsByQueue: map[string][]*jobdb.Job{ + "B": testfixtures.N1CpuJobs("B", testfixtures.PriorityClass3, 23), + }, + ExpectedScheduledIndices: map[string][]int{ + "B": testfixtures.IntRange(0, 22), + }, + }, + { + JobsByQueue: map[string][]*jobdb.Job{ + "C": testfixtures.N1CpuJobs("C", testfixtures.PriorityClass0, 1), + }, + ExpectedScheduledIndices: map[string][]int{ + "C": testfixtures.IntRange(0, 0), + }, + ExpectedPreemptedIndices: map[string]map[int][]int{ + "A": { + 0: testfixtures.IntRange(8, 8), + }, + }, + }, + {}, // Empty round to make sure nothing changes. + }, + PriorityFactorByQueue: map[string]float64{ + "A": 1, + "B": 1, + "C": 1, + }, + }, } for name, tc := range tests { t.Run(name, func(t *testing.T) { diff --git a/internal/scheduler/testfixtures/testfixtures.go b/internal/scheduler/testfixtures/testfixtures.go index f00ced1b4a0..07ccae0f531 100644 --- a/internal/scheduler/testfixtures/testfixtures.go +++ b/internal/scheduler/testfixtures/testfixtures.go @@ -88,7 +88,7 @@ func ContextWithDefaultLogger(ctx context.Context) context.Context { func TestSchedulingConfig() configuration.SchedulingConfig { return configuration.SchedulingConfig{ - ResourceScarcity: map[string]float64{"cpu": 1, "memory": 0}, + ResourceScarcity: map[string]float64{"cpu": 1}, Preemption: configuration.PreemptionConfig{ PriorityClasses: maps.Clone(TestPriorityClasses), DefaultPriorityClass: TestDefaultPriorityClass, @@ -101,8 +101,13 @@ func TestSchedulingConfig() configuration.SchedulingConfig { } } -func WithMaxUnacknowledgedJobsPerExecutor(i uint, config configuration.SchedulingConfig) configuration.SchedulingConfig { - config.MaxUnacknowledgedJobsPerExecutor = i +func WithMaxUnacknowledgedJobsPerExecutorConfig(v uint, config configuration.SchedulingConfig) configuration.SchedulingConfig { + config.MaxUnacknowledgedJobsPerExecutor = v + return config +} + +func WithProtectedFractionOfFairShareConfig(v float64, config configuration.SchedulingConfig) configuration.SchedulingConfig { + config.Preemption.ProtectedFractionOfFairShare = v return config }