Add per-pool queue limit overrides

armadaproject · Jun 22, 2023 · 9698ea6 · 9698ea6
1 parent 11443f7
commit 9698ea6
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 8 deletions.
diff --git a/internal/armada/configuration/types.go b/internal/armada/configuration/types.go
@@ -240,6 +240,9 @@ type PriorityClass struct {
 	// jobs of this priority class are not scheduled if doing so would cause the total resources assigned
 	// to jobs of priority 10 or lower from the same queue to exceed 30% of the total.
 	MaximumResourceFractionPerQueue map[string]float64
+	// Per-pool override of MaximumResourceFractionPerQueue.
+	// If missing for a particular pool, MaximumResourceFractionPerQueue is used instead for that pool.
+	MaximumResourceFractionPerQueueByPool map[string]map[string]float64
 }
 
 func (p PreemptionConfig) PriorityByPriorityClassName() map[string]int32 {

diff --git a/internal/armada/server/submit_test.go b/internal/armada/server/submit_test.go
@@ -1637,7 +1637,7 @@ func withSubmitServerAndRepos(action func(s *SubmitServer, jobRepo repository.Jo
 		MaxPodSpecSizeBytes: 65535,
 		Preemption: configuration.PreemptionConfig{
 			DefaultPriorityClass: "high",
-			PriorityClasses:      map[string]configuration.PriorityClass{"high": {0, false, nil}},
+			PriorityClasses:      map[string]configuration.PriorityClass{"high": {0, false, nil, nil}},
 		},
 		MinTerminationGracePeriod: time.Duration(30 * time.Second),
 		MaxTerminationGracePeriod: time.Duration(300 * time.Second),

diff --git a/internal/scheduler/adapters/adapters_test.go b/internal/scheduler/adapters/adapters_test.go
@@ -18,10 +18,10 @@ import (
 
 var (
 	priorityByPriorityClassName = map[string]configuration.PriorityClass{
-		"priority-0": {0, true, nil},
-		"priority-1": {1, true, nil},
-		"priority-2": {2, true, nil},
-		"priority-3": {3, false, nil},
+		"priority-0": {0, true, nil, nil},
+		"priority-1": {1, true, nil, nil},
+		"priority-2": {2, true, nil, nil},
+		"priority-3": {3, false, nil, nil},
 	}
 
 	priority int32 = 1

diff --git a/internal/scheduler/constraints/constraints.go b/internal/scheduler/constraints/constraints.go
@@ -68,15 +68,21 @@ func SchedulingConstraintsFromSchedulingConfig(
 ) SchedulingConstraints {
 	priorityClassSchedulingConstraintsByPriorityClassName := make(map[string]PriorityClassSchedulingConstraints, len(config.Preemption.PriorityClasses))
 	for name, priorityClass := range config.Preemption.PriorityClasses {
+		maximumResourceFractionPerQueue := priorityClass.MaximumResourceFractionPerQueue
+		if m, ok := priorityClass.MaximumResourceFractionPerQueueByPool[pool]; ok {
+			// Use pool-specific config is available.
+			maximumResourceFractionPerQueue = m
+		}
 		priorityClassSchedulingConstraintsByPriorityClassName[name] = PriorityClassSchedulingConstraints{
 			PriorityClassName:                  name,
 			PriorityClassPriority:              priorityClass.Priority,
-			MaximumCumulativeResourcesPerQueue: absoluteFromRelativeLimits(totalResources, priorityClass.MaximumResourceFractionPerQueue),
+			MaximumCumulativeResourcesPerQueue: absoluteFromRelativeLimits(totalResources, maximumResourceFractionPerQueue),
 		}
 	}
 	maximumResourceFractionToSchedule := config.MaximumResourceFractionToSchedule
-	if limit, ok := config.MaximumResourceFractionToScheduleByPool[pool]; ok {
-		maximumResourceFractionToSchedule = limit
+	if m, ok := config.MaximumResourceFractionToScheduleByPool[pool]; ok {
+		// Use pool-specific config is available.
+		maximumResourceFractionToSchedule = m
 	}
 	return SchedulingConstraints{
 		MaximumJobsToSchedule:      config.MaximumJobsToSchedule,