Repository: armadaproject/armada
Pull request #2635 (2023-07-03 16:07:49)

Add gang scheduling node uniformity constraints
I.e., schedule gangs across a set of nodes for which a user-specified label has the same value, whatever that value may be. Used to ensure, e.g., that all gang jobs land on nodes in the same rack.

┆Issue is synchronized with this [Jira Task](https://gr-oss.atlassian.net/browse/BATCH-424) by [Unito](https://www.unito.io)


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/armadaproject/armada/pull/2635?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject) Report
Patch coverage has no change and project coverage change: **`+18.30`** :tada:
> Comparison is base [(`24fb052`)](https://app.codecov.io/gh/armadaproject/armada/commit/24fb0529ac2a123b5f32e8cfa47419293842effa?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject) 16.07% compared to head [(`cc36b77`)](https://app.codecov.io/gh/armadaproject/armada/pull/2635?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject) 34.38%.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##           master    #2635       +/-   ##
===========================================
+ Coverage   16.07%   34.38%   +18.30%     
===========================================
  Files         159      112       -47     
  Lines       12490     3467     -9023     
  Branches      470      470               
===========================================
- Hits         2008     1192      -816     
+ Misses      10320     2192     -8128     
+ Partials      162       83       -79     
```

| Flag | Coverage Δ | |
|---|---|---|
| unittests | `34.38% <ø> (+18.30%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject#carryforward-flags-in-the-pull-request-comment) to find out more.

[see 47 files with indirect coverage changes](https://app.codecov.io/gh/armadaproject/armada/pull/2635/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/armadaproject/armada/pull/2635?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=armadaproject).


Patch:
diff --git a/internal/armada/configuration/constants.go b/internal/armada/configuration/constants.go
index 992ab6c9..51669f2a 100644
--- a/internal/armada/configuration/constants.go
+++ b/internal/armada/configuration/constants.go
@@ -7,18 +7,16 @@ const (
 	// GangCardinalityAnnotation All jobs in a gang must specify the total number of jobs in the gang via this annotation.
 	// The cardinality should be expressed as an integer, e.g., "3".
 	GangCardinalityAnnotation = "armadaproject.io/gangCardinality"
+	// The jobs that make up a gang may be constrained to be scheduled across a set of uniform nodes.
+	// Specifically, if provided, all gang jobs are scheduled onto nodes for which the value of the provided label is equal.
+	// Used to ensure, e.g., that all gang jobs are scheduled onto the same cluster or rack.
+	GangNodeUniformityLabelAnnotation = "armadaproject.io/gangNodeUniformityLabel"
 	// Armada normally tries to re-schedule jobs for which a pod fails to start.
 	// Pods for which this annotation has value "true" are not retried.
 	// Instead, the job the pod is part of fails immediately.
 	FailFastAnnotation = "armadaproject.io/failFast"
 )
 
-var ArmadaManagedAnnotations = []string{
-	GangIdAnnotation,
-	GangCardinalityAnnotation,
-	FailFastAnnotation,
-}
-
 var ReturnLeaseRequestTrackedAnnotations = map[string]struct{}{
 	FailFastAnnotation: {},
 }
diff --git a/internal/armada/configuration/types.go b/internal/armada/configuration/types.go
index 80f4ed4e..23fea24b 100644
--- a/internal/armada/configuration/types.go
+++ b/internal/armada/configuration/types.go
@@ -153,6 +153,8 @@ type SchedulingConfig struct {
 	//
 	// Applies only to the new scheduler.
 	IndexedTaints []string
+	// Default value of GangNodeUniformityLabelAnnotation if none is provided.
+	DefaultGangNodeUniformityLabel string
 	// Kubernetes pods may specify a termination grace period.
 	// When Pods are cancelled/preempted etc., they are first sent a SIGTERM.
 	// If a pod has not exited within its termination grace period,
diff --git a/internal/armada/server/applydefaults.go b/internal/armada/server/applydefaults.go
index 6ffbd42a..510110c7 100644
--- a/internal/armada/server/applydefaults.go
+++ b/internal/armada/server/applydefaults.go
@@ -10,6 +10,21 @@ import (
 	armadaresource "github.com/armadaproject/armada/internal/common/resource"
 )
 
+func applyDefaultsToAnnotations(annotations map[string]string, config configuration.SchedulingConfig) {
+	if annotations == nil {
+		return
+	}
+	applyDefaultNodeUniformityLabelAnnotation(annotations, config)
+}
+
+func applyDefaultNodeUniformityLabelAnnotation(annotations map[string]string, config configuration.SchedulingConfig) {
+	if _, ok := annotations[configuration.GangIdAnnotation]; ok {
+		if _, ok := annotations[configuration.GangNodeUniformityLabelAnnotation]; !ok {
+			annotations[configuration.GangNodeUniformityLabelAnnotation] = config.DefaultGangNodeUniformityLabel
+		}
+	}
+}
+
 func applyDefaultsToPodSpec(spec *v1.PodSpec, config configuration.SchedulingConfig) {
 	if spec == nil {
 		return
diff --git a/internal/armada/server/submit.go b/internal/armada/server/submit.go
index 0acd6969..ca444ff3 100644
--- a/internal/armada/server/submit.go
+++ b/internal/armada/server/submit.go
@@ -794,6 +794,7 @@ func (server *SubmitServer) createJobsObjects(request *api.JobSubmitRequest, own
 			namespace = "default"
 		}
 		fillContainerRequestsAndLimits(podSpec.Containers)
+		applyDefaultsToAnnotations(item.Annotations, *server.schedulingConfig)
 		applyDefaultsToPodSpec(podSpec, *server.schedulingConfig)
 		if err := validation.ValidatePodSpec(podSpec, server.schedulingConfig); err != nil {
 			return nil, errors.Errorf("[createJobs] error validating the %d-th job of job set %s: %v", i, request.JobSetId, err)
diff --git a/internal/common/validation/job.go b/internal/common/validation/job.go
index fcc25f1b..5b50f6d8 100644
--- a/internal/common/validation/job.go
+++ b/internal/common/validation/job.go
@@ -12,8 +12,7 @@ import (
 )
 
 func ValidateApiJobs(jobs []*api.Job, config configuration.SchedulingConfig) error {
-	err := validateGangs(jobs)
-	if err != nil {
+	if err := validateGangs(jobs); err != nil {
 		return err
 	}
 	for _, job := range jobs {
@@ -26,13 +25,15 @@ func ValidateApiJobs(jobs []*api.Job, config configuration.SchedulingConfig) err
 
 func validateGangs(jobs []*api.Job) error {
 	gangDetailsByGangId := make(map[string]struct {
-		actualCardinality         int
-		expectedCardinality       int
-		expectedPriorityClassName string
+		actualCardinality           int
+		expectedCardinality         int
+		expectedPriorityClassName   string
+		expectedNodeUniformityLabel string
 	})
 	for i, job := range jobs {
 		annotations := job.Annotations
 		gangId, gangCardinality, isGangJob, err := scheduler.GangIdAndCardinalityFromAnnotations(annotations)
+		nodeUniformityLabel := annotations[configuration.GangNodeUniformityLabelAnnotation]
 		if err != nil {
 			return errors.WithMessagef(err, "%d-th job with id %s in gang %s", i, job.Id, gangId)
 		}
@@ -56,6 +57,12 @@ func validateGangs(jobs []*api.Job) error {
 					i, job.Id, gangId, details.expectedPriorityClassName, podSpec.PriorityClassName,
 				)
 			}
+			if nodeUniformityLabel != details.expectedNodeUniformityLabel {
+				return errors.Errorf(
+					"inconsistent nodeUniformityLabel for %d-th job with id %s in gang %s: expected %s but got %s",
+					i, job.Id, gangId, details.expectedNodeUniformityLabel, nodeUniformityLabel,
+				)
+			}
 			details.actualCardinality++
 			gangDetailsByGangId[gangId] = details
 		} else {
@@ -64,6 +71,7 @@ func validateGangs(jobs []*api.Job) error {
 			if podSpec != nil {
 				details.expectedPriorityClassName = podSpec.PriorityClassName
 			}
+			details.expectedNodeUniformityLabel = nodeUniformityLabel
 			gangDetailsByGangId[gangId] = details
 		}
 	}
diff --git a/internal/scheduler/configuration/configuration.go b/internal/scheduler/configuration/configuration.go
index c4caa350..2234e77c 100644
--- a/internal/scheduler/configuration/configuration.go
+++ b/internal/scheduler/configuration/configuration.go
@@ -13,8 +13,8 @@ const (
 	// IsEvictedAnnotation is set on evicted jobs; the scheduler uses it to differentiate between
 	// already-running and queued jobs.
 	IsEvictedAnnotation = "armadaproject.io/isEvicted"
-	// NodeIdLabel is set on evicted jobs, so that the scheduler only tries to schedule them on the
-	// nodes that they are already running on; nodedb is responsible for labelling its Node objects.
+	// NodeIdLabel maps to a unique id associated with each node.
+	// This label is automatically added to nodes within the NodeDb.
 	NodeIdLabel = "armadaproject.io/nodeId"
 )
 
diff --git a/internal/scheduler/context/context.go b/internal/scheduler/context/context.go
index f28756fa..3fd61bfd 100644
--- a/internal/scheduler/context/context.go
+++ b/internal/scheduler/context/context.go
@@ -567,6 +567,7 @@ type GangSchedulingContext struct {
 	JobSchedulingContexts []*JobSchedulingContext
 	TotalResourceRequests schedulerobjects.ResourceList
 	AllJobsEvicted        bool
+	NodeUniformityLabel   string
 }
 
 func NewGangSchedulingContext(jctxs []*JobSchedulingContext) *GangSchedulingContext {
@@ -574,9 +575,13 @@ func NewGangSchedulingContext(jctxs []*JobSchedulingContext) *GangSchedulingCont
 	// (which we enforce at job submission).
 	queue := ""
 	priorityClassName := ""
+	nodeUniformityLabel := ""
 	if len(jctxs) > 0 {
 		queue = jctxs[0].Job.GetQueue()
 		priorityClassName = jctxs[0].Job.GetPriorityClassName()
+		if jctxs[0].PodRequirements != nil {
+			nodeUniformityLabel = jctxs[0].PodRequirements.Annotations[configuration.GangNodeUniformityLabelAnnotation]
+		}
 	}
 	allJobsEvicted := true
 	totalResourceRequests := schedulerobjects.NewResourceList(4)
@@ -591,6 +596,7 @@ func NewGangSchedulingContext(jctxs []*JobSchedulingContext) *GangSchedulingCont
 		JobSchedulingContexts: jctxs,
 		TotalResourceRequests: totalResourceRequests,
 		AllJobsEvicted:        allJobsEvicted,
+		NodeUniformityLabel:   nodeUniformityLabel,
 	}
 }
 
@@ -661,6 +667,9 @@ type PodSchedulingContext struct {
 	NodeId string
 	// Score indicates how well the pod fits on the selected node.
 	Score int
+	// Priority class priority at which this pod was scheduled.
+	// Only set if NodeId is.
+	ScheduledAtPriority int32
 	// Node types on which this pod could be scheduled.
 	MatchingNodeTypes []*schedulerobjects.NodeType
 	// Total number of nodes in the cluster when trying to schedule.
diff --git a/internal/scheduler/gang_scheduler.go b/internal/scheduler/gang_scheduler.go
index 34bdf4f3..f1a39e31 100644
--- a/internal/scheduler/gang_scheduler.go
+++ b/internal/scheduler/gang_scheduler.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/hashicorp/go-memdb"
+
 	"github.com/armadaproject/armada/internal/common/util"
 	schedulerconstraints "github.com/armadaproject/armada/internal/scheduler/constraints"
 	schedulercontext "github.com/armadaproject/armada/internal/scheduler/context"
@@ -103,27 +105,124 @@ func (sch *GangScheduler) Schedule(ctx context.Context, gctx *schedulercontext.G
 			return
 		}
 	}
-	if ok, unschedulableReason, err = sch.trySchedule(ctx, gctx); err != nil || ok {
+	return sch.trySchedule(ctx, gctx)
+}
+
+func (sch *GangScheduler) trySchedule(ctx context.Context, gctx *schedulercontext.GangSchedulingContext) (ok bool, unschedulableReason string, err error) {
+	// If no node uniformity constraint, try scheduling across all nodes.
+	if gctx.NodeUniformityLabel == "" {
+		return sch.tryScheduleGang(ctx, gctx)
+	}
+
+	// Otherwise try scheduling such that all nodes onto which a gang job lands have the same value for gctx.NodeUniformityLabel.
+	// We do this by making a separate scheduling attempt for each unique value of gctx.NodeUniformityLabel.
+	nodeUniformityLabelValues, ok := sch.nodeDb.IndexedNodeLabelValues(gctx.NodeUniformityLabel)
+	if !ok {
+		ok = false
+		unschedulableReason = fmt.Sprintf("uniformity label %s is not indexed", gctx.NodeUniformityLabel)
+		return
+	}
+	if len(nodeUniformityLabelValues) == 0 {
+		ok = false
+		unschedulableReason = fmt.Sprintf("no nodes with uniformity label %s", gctx.NodeUniformityLabel)
 		return
 	}
-	return
+
+	// Try all possible values of nodeUniformityLabel one at a time to find the best fit.
+	bestValue := ""
+	var minMeanScheduledAtPriority float64
+	var i int
+	for value := range nodeUniformityLabelValues {
+		i++
+		if value == "" {
+			continue
+		}
+		addNodeSelectorToGctx(gctx, gctx.NodeUniformityLabel, value)
+		txn := sch.nodeDb.Txn(true)
+		if ok, unschedulableReason, err = sch.tryScheduleGangWithTxn(ctx, txn, gctx); err != nil {
+			txn.Abort()
+			return
+		} else if ok {
+			meanScheduledAtPriority, ok := meanScheduledAtPriorityFromGctx(gctx)
+			if !ok {
+				txn.Abort()
+				continue
+			}
+			if meanScheduledAtPriority == float64(nodedb.MinPriority) {
+				// Best possible; no need to keep looking.
+				txn.Commit()
+				return true, "", nil
+			}
+			if bestValue == "" || meanScheduledAtPriority <= minMeanScheduledAtPriority {
+				if i == len(nodeUniformityLabelValues) {
+					// Minimal meanScheduledAtPriority and no more options; commit and return.
+					txn.Commit()
+					return true, "", nil
+				}
+				// Record the best value seen so far.
+				bestValue = value
+				minMeanScheduledAtPriority = meanScheduledAtPriority
+			}
+		}
+		txn.Abort()
+	}
+	if bestValue == "" {
+		ok = false
+		unschedulableReason = "at least one job in the gang does not fit on any node"
+		return
+	}
+	addNodeSelectorToGctx(gctx, gctx.NodeUniformityLabel, bestValue)
+	return sch.tryScheduleGang(ctx, gctx)
 }
 
-func (sch *GangScheduler) trySchedule(ctx context.Context, gctx *schedulercontext.GangSchedulingContext) (bool, string, error) {
-	ok, err := sch.nodeDb.ScheduleMany(gctx.JobSchedulingContexts)
-	if err != nil {
-		return false, "", err
+func (sch *GangScheduler) tryScheduleGang(ctx context.Context, gctx *schedulercontext.GangSchedulingContext) (ok bool, unschedulableReason string, err error) {
+	txn := sch.nodeDb.Txn(true)
+	defer txn.Abort()
+	ok, unschedulableReason, err = sch.tryScheduleGangWithTxn(ctx, txn, gctx)
+	if ok && err == nil {
+		txn.Commit()
 	}
-	if !ok {
-		unschedulableReason := ""
+	return
+}
+
+func (sch *GangScheduler) tryScheduleGangWithTxn(ctx context.Context, txn *memdb.Txn, gctx *schedulercontext.GangSchedulingContext) (ok bool, unschedulableReason string, err error) {
+	if ok, err = sch.nodeDb.ScheduleManyWithTxn(txn, gctx.JobSchedulingContexts); err != nil {
+		return
+	} else if !ok {
+		for _, jctx := range gctx.JobSchedulingContexts {
+			if jctx.PodSchedulingContext != nil {
+				// Clear any node bindings on failure to schedule.
+				jctx.PodSchedulingContext.NodeId = ""
+			}
+		}
 		if len(gctx.JobSchedulingContexts) > 1 {
 			unschedulableReason = "at least one job in the gang does not fit on any node"
 		} else {
 			unschedulableReason = "job does not fit on any node"
 		}
-		return false, unschedulableReason, nil
+		return
+	}
+	return
+}
+
+func addNodeSelectorToGctx(gctx *schedulercontext.GangSchedulingContext, nodeSelectorKey, nodeSelectorValue string) {
+	for _, jctx := range gctx.JobSchedulingContexts {
+		if jctx.PodRequirements.NodeSelector == nil {
+			jctx.PodRequirements.NodeSelector = make(map[string]string)
+		}
+		jctx.PodRequirements.NodeSelector[nodeSelectorKey] = nodeSelectorValue
+	}
+}
+
+func meanScheduledAtPriorityFromGctx(gctx *schedulercontext.GangSchedulingContext) (float64, bool) {
+	var sum int32
+	for _, jctx := range gctx.JobSchedulingContexts {
+		if jctx.PodSchedulingContext == nil {
+			return 0, false
+		}
+		sum += jctx.PodSchedulingContext.ScheduledAtPriority
 	}
-	return true, "", nil
+	return float64(sum) / float64(len(gctx.JobSchedulingContexts)), true
 }
 
 func requestsAreLargeEnough(totalResourceRequests, minRequest schedulerobjects.ResourceList) (bool, string) {
diff --git a/internal/scheduler/nodedb/nodedb.go b/internal/scheduler/nodedb/nodedb.go
index 9fa19a11..1dbbb71b 100644
--- a/internal/scheduler/nodedb/nodedb.go
+++ b/internal/scheduler/nodedb/nodedb.go
@@ -27,9 +27,15 @@ import (
 	"github.com/armadaproject/armada/pkg/api"
 )
 
-// evictedPriority is the priority class priority resources consumed by evicted jobs are accounted for at.
-// This helps avoid scheduling new jobs onto nodes that make it impossible to re-schedule evicted jobs.
-const evictedPriority int32 = -1
+const (
+	// evictedPriority is the priority class priority resources consumed by evicted jobs are accounted for at.
+	// This helps avoid scheduling new jobs onto nodes that make it impossible to re-schedule evicted jobs.
+	evictedPriority int32 = -1
+	// MinPriority is the smallest possible priority class priority within the NodeDb.
+	MinPriority int32 = evictedPriority
+)
+
+var empty struct{}
 
 type Node struct {
 	Id   string
@@ -125,6 +131,11 @@ func (nodeDb *NodeDb) create(node *schedulerobjects.Node) (*Node, error) {
 	}
 
 	nodeDb.mu.Lock()
+	for key := range nodeDb.indexedNodeLabels {
+		if value, ok := labels[key]; ok {
+			nodeDb.indexedNodeLabelValues[key][value] = empty
+		}
+	}
 	nodeDb.numNodes++
 	nodeDb.numNodesByNodeType[nodeType.Id]++
 	nodeDb.totalResources.Add(totalResources)
@@ -237,6 +248,8 @@ type NodeDb struct {
 	// Mutex for the remaining fields of this struct, which are mutated after initialization.
 	mu sync.Mutex
 
+	// Map from indexed label names to the set of values that label takes across all nodes in the NodeDb.
+	indexedNodeLabelValues map[string]map[string]struct{}
 	// Total number of nodes in the db.
 	numNodes int
 	// Number of nodes in the db by node type.
@@ -293,6 +306,10 @@ func NewNodeDb(
 			Message: "there must be at least one indexed resource",
 		})
 	}
+	indexedNodeLabelValues := make(map[string]map[string]struct{}, len(indexedNodeLabels))
+	for _, key := range indexedNodeLabels {
+		indexedNodeLabelValues[key] = make(map[string]struct{})
+	}
 	mapFromSlice := func(vs []string) map[string]interface{} {
 		rv := make(map[string]interface{})
 		for _, v := range vs {
@@ -311,13 +328,14 @@ func NewNodeDb(
 			indexedResources,
 			func(v configuration.IndexedResource) int64 { return v.Resolution.MilliValue() },
 		),
-		indexNameByPriority: indexNameByPriority,
-		indexedTaints:       mapFromSlice(indexedTaints),
-		indexedNodeLabels:   mapFromSlice(indexedNodeLabels),
-		nodeTypes:           make(map[uint64]*schedulerobjects.NodeType),
-		numNodesByNodeType:  make(map[uint64]int),
-		totalResources:      schedulerobjects.ResourceList{Resources: make(map[string]resource.Quantity)},
-		db:                  db,
+		indexNameByPriority:    indexNameByPriority,
+		indexedTaints:          mapFromSlice(indexedTaints),
+		indexedNodeLabels:      mapFromSlice(indexedNodeLabels),
+		indexedNodeLabelValues: indexedNodeLabelValues,
+		nodeTypes:              make(map[uint64]*schedulerobjects.NodeType),
+		numNodesByNodeType:     make(map[uint64]int),
+		totalResources:         schedulerobjects.ResourceList{Resources: make(map[string]resource.Quantity)},
+		db:                     db,
 		// Set the initial capacity (somewhat arbitrarily) to 128 reasons.
 		podRequirementsNotMetReasonStringCache: make(map[uint64]string, 128),
 	}, nil
@@ -342,6 +360,12 @@ func (nodeDb *NodeDb) String() string {
 	return sb.String()
 }
 
+// IndexedNodeLabelValues returns the set of possible values for a given indexed label across all nodes in the NodeDb.
+func (nodeDb *NodeDb) IndexedNodeLabelValues(label string) (map[string]struct{}, bool) {
+	values, ok := nodeDb.indexedNodeLabelValues[label]
+	return values, ok
+}
+
 func (nodeDb *NodeDb) NumNodes() int {
 	nodeDb.mu.Lock()
 	defer nodeDb.mu.Unlock()
@@ -433,15 +457,6 @@ func (nodeDb *NodeDb) ScheduleMany(jctxs []*schedulercontext.JobSchedulingContex
 	if ok && err == nil {
 		// All pods can be scheduled; commit the transaction.
 		txn.Commit()
-	} else {
-		// On failure, clear the node binding.
-		for _, jctx := range jctxs {
-			pctx := jctx.PodSchedulingContext
-			if pctx == nil {
-				continue
-			}
-			pctx.NodeId = ""
-		}
 	}
 	return ok, err
 }
@@ -643,6 +658,7 @@ func (nodeDb *NodeDb) selectNodeForPodWithIt(
 	if selectedNode != nil {
 		pctx.NodeId = selectedNode.Id
 		pctx.Score = selectedNodeScore
+		pctx.ScheduledAtPriority = priority
 	}
 	return selectedNode, nil
 }
diff --git a/internal/scheduler/submitcheck.go b/internal/scheduler/submitcheck.go
index 716d0d16..77acce6a 100644
--- a/internal/scheduler/submitcheck.go
+++ b/internal/scheduler/submitcheck.go
@@ -213,6 +213,8 @@ func (srv *SubmitChecker) getSchedulingResult(jctxs []*schedulercontext.JobSched
 	for id, executor := range executorById {
 		nodeDb := executor.nodeDb
 		txn := nodeDb.Txn(true)
+		// TODO: This doesn't account for per-queue limits or the NodeUniformityLabel.
+		// We should create a GangScheduler for this instead.
 		ok, err := nodeDb.ScheduleManyWithTxn(txn, jctxs)
 		txn.Abort()
 


Test patch:
diff --git a/internal/armada/server/applydefaults_test.go b/internal/armada/server/applydefaults_test.go
index 036c642d..62293a45 100644
--- a/internal/armada/server/applydefaults_test.go
+++ b/internal/armada/server/applydefaults_test.go
@@ -11,6 +11,53 @@ import (
 	"github.com/armadaproject/armada/internal/armada/configuration"
 )
 
+func TestApplyDefaultsToAnnotations(t *testing.T) {
+	tests := map[string]struct {
+		Config      configuration.SchedulingConfig
+		Annotations map[string]string
+		Expected    map[string]string
+	}{
+		"no change": {
+			Annotations: make(map[string]string),
+			Expected:    make(map[string]string),
+		},
+		"DefaultNodeUniformityLabelAnnotation no change for non-gang jobs": {
+			Config: configuration.SchedulingConfig{
+				DefaultGangNodeUniformityLabel: "foo",
+			},
+			Annotations: make(map[string]string),
+			Expected:    make(map[string]string),
+		},
+		"DefaultNodeUniformityLabelAnnotation empty default": {
+			Annotations: map[string]string{
+				configuration.GangIdAnnotation: "bar",
+			},
+			Expected: map[string]string{
+				configuration.GangIdAnnotation:                  "bar",
+				configuration.GangNodeUniformityLabelAnnotation: "",
+			},
+		},
+		"DefaultNodeUniformityLabelAnnotation": {
+			Config: configuration.SchedulingConfig{
+				DefaultGangNodeUniformityLabel: "foo",
+			},
+			Annotations: map[string]string{
+				configuration.GangIdAnnotation: "bar",
+			},
+			Expected: map[string]string{
+				configuration.GangIdAnnotation:                  "bar",
+				configuration.GangNodeUniformityLabelAnnotation: "foo",
+			},
+		},
+	}
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			applyDefaultsToAnnotations(tc.Annotations, tc.Config)
+			assert.Equal(t, tc.Expected, tc.Annotations)
+		})
+	}
+}
+
 func TestApplyDefaultsToPodSpec(t *testing.T) {
 	tests := map[string]struct {
 		Config   configuration.SchedulingConfig
diff --git a/internal/common/validation/job_test.go b/internal/common/validation/job_test.go
index 960b74ba..08b54a9c 100644
--- a/internal/common/validation/job_test.go
+++ b/internal/common/validation/job_test.go
@@ -325,7 +325,8 @@ func TestValidateGangs(t *testing.T) {
 				},
 				{
 					Annotations: map[string]string{
-						configuration.GangIdAnnotation: "bar",
+						configuration.GangIdAnnotation:          "bar",
+						configuration.GangCardinalityAnnotation: strconv.Itoa(2),
 					},
 					PodSpec: &v1.PodSpec{
 						PriorityClassName: "zab",
@@ -334,6 +335,27 @@ func TestValidateGangs(t *testing.T) {
 			},
 			ExpectSuccess: false,
 		},
+		"inconsistent NodeUniformityLabel": {
+			Jobs: []*api.Job{
+				{
+					Annotations: map[string]string{
+						configuration.GangIdAnnotation:                  "bar",
+						configuration.GangCardinalityAnnotation:         strconv.Itoa(2),
+						configuration.GangNodeUniformityLabelAnnotation: "foo",
+					},
+					PodSpec: &v1.PodSpec{},
+				},
+				{
+					Annotations: map[string]string{
+						configuration.GangIdAnnotation:                  "bar",
+						configuration.GangCardinalityAnnotation:         strconv.Itoa(2),
+						configuration.GangNodeUniformityLabelAnnotation: "bar",
+					},
+					PodSpec: &v1.PodSpec{},
+				},
+			},
+			ExpectSuccess: false,
+		},
 	}
 	for name, tc := range tests {
 		t.Run(name, func(t *testing.T) {
diff --git a/internal/scheduler/gang_scheduler_test.go b/internal/scheduler/gang_scheduler_test.go
index 2b98f665..ea7153b0 100644
--- a/internal/scheduler/gang_scheduler_test.go
+++ b/internal/scheduler/gang_scheduler_test.go
@@ -210,15 +210,105 @@ func TestGangScheduler(t *testing.T) {
 			},
 			ExpectedScheduledIndices: testfixtures.IntRange(0, 0),
 		},
+		"NodeUniformityLabel set but not indexed": {
+			SchedulingConfig: testfixtures.TestSchedulingConfig(),
+			Nodes: testfixtures.WithLabelsNodes(
+				map[string]string{"foo": "foov"},
+				testfixtures.N32CpuNodes(1, testfixtures.TestPriorities),
+			),
+			Gangs: [][]*jobdb.Job{
+				testfixtures.WithNodeUniformityLabelAnnotationJobs(
+					"foo",
+					testfixtures.N16Cpu128GiJobs("A", testfixtures.PriorityClass0, 1),
+				),
+			},
+			ExpectedScheduledIndices: nil,
+		},
+		"NodeUniformityLabel not set": {
+			SchedulingConfig: testfixtures.WithIndexedNodeLabelsConfig(
+				[]string{"foo", "bar"},
+				testfixtures.TestSchedulingConfig(),
+			),
+			Nodes: testfixtures.N32CpuNodes(1, testfixtures.TestPriorities),
+			Gangs: [][]*jobdb.Job{
+				testfixtures.WithNodeUniformityLabelAnnotationJobs(
+					"foo",
+					testfixtures.N16Cpu128GiJobs("A", testfixtures.PriorityClass0, 1),
+				),
+			},
+			ExpectedScheduledIndices: nil,
+		},
+		"NodeUniformityLabel insufficient capacity": {
+			SchedulingConfig: testfixtures.WithIndexedNodeLabelsConfig(
+				[]string{"foo", "bar"},
+				testfixtures.TestSchedulingConfig(),
+			),
+			Nodes: armadaslices.Concatenate(
+				testfixtures.WithLabelsNodes(
+					map[string]string{"foo": "foov1"},
+					testfixtures.N32CpuNodes(1, testfixtures.TestPriorities),
+				),
+				testfixtures.WithLabelsNodes(
+					map[string]string{"foo": "foov2"},
+					testfixtures.N32CpuNodes(1, testfixtures.TestPriorities),
+				),
+			),
+			Gangs: [][]*jobdb.Job{
+				testfixtures.WithNodeUniformityLabelAnnotationJobs(
+					"foo",
+					testfixtures.N16Cpu128GiJobs("A", testfixtures.PriorityClass0, 3),
+				),
+			},
+			ExpectedScheduledIndices: nil,
+		},
+		"NodeUniformityLabel": {
+			SchedulingConfig: testfixtures.WithIndexedNodeLabelsConfig(
+				[]string{"foo", "bar"},
+				testfixtures.TestSchedulingConfig(),
+			),
+			Nodes: armadaslices.Concatenate(
+				testfixtures.WithLabelsNodes(
+					map[string]string{"foo": "foov1"},
+					testfixtures.WithUsedResourcesNodes(
+						0,
+						schedulerobjects.ResourceList{Resources: map[string]resource.Quantity{"cpu": resource.MustParse("1")}},
+						testfixtures.N32CpuNodes(2, testfixtures.TestPriorities),
+					),
+				),
+				testfixtures.WithLabelsNodes(
+					map[string]string{"foo": "foov2"},
+					testfixtures.N32CpuNodes(2, testfixtures.TestPriorities),
+				),
+				testfixtures.WithLabelsNodes(
+					map[string]string{"foo": "foov3"},
+					testfixtures.WithUsedResourcesNodes(
+						0,
+						schedulerobjects.ResourceList{Resources: map[string]resource.Quantity{"cpu": resource.MustParse("1")}},
+						testfixtures.N32CpuNodes(2, testfixtures.TestPriorities),
+					),
+				),
+			),
+			Gangs: [][]*jobdb.Job{
+				testfixtures.WithNodeUniformityLabelAnnotationJobs(
+					"foo",
+					testfixtures.N16Cpu128GiJobs("A", testfixtures.PriorityClass0, 4),
+				),
+			},
+			ExpectedScheduledIndices: []int{0},
+		},
 	}
 	for name, tc := range tests {
 		t.Run(name, func(t *testing.T) {
+			nodesById := make(map[string]*schedulerobjects.Node, len(tc.Nodes))
+			for _, node := range tc.Nodes {
+				nodesById[node.Id] = node
+			}
 			nodeDb, err := nodedb.NewNodeDb(
 				testfixtures.TestPriorityClasses,
 				testfixtures.TestMaxExtraNodesToConsider,
 				tc.SchedulingConfig.IndexedResources,
 				testfixtures.TestIndexedTaints,
-				testfixtures.TestIndexedNodeLabels,
+				tc.SchedulingConfig.IndexedNodeLabels,
 			)
 			require.NoError(t, err)
 			txn := nodeDb.Txn(true)
@@ -267,6 +357,23 @@ func TestGangScheduler(t *testing.T) {
 				if ok {
 					require.Empty(t, reason)
 					actualScheduledIndices = append(actualScheduledIndices, i)
+
+					// If there's a node uniformity constraint, check that it's met.
+					if gctx.NodeUniformityLabel != "" {
+						nodeUniformityLabelValues := make(map[string]bool)
+						for _, jctx := range jctxs {
+							require.NotNil(t, jctx.PodSchedulingContext)
+							node := nodesById[jctx.PodSchedulingContext.NodeId]
+							require.NotNil(t, node)
+							value, ok := node.Labels[gctx.NodeUniformityLabel]
+							require.True(t, ok, "gang job scheduled onto node with missing nodeUniformityLabel")
+							nodeUniformityLabelValues[value] = true
+						}
+						require.Equal(
+							t, 1, len(nodeUniformityLabelValues),
+							"node uniformity constraint not met: %s", nodeUniformityLabelValues,
+						)
+					}
 				} else {
 					require.NotEmpty(t, reason)
 				}
diff --git a/internal/scheduler/testfixtures/testfixtures.go b/internal/scheduler/testfixtures/testfixtures.go
index 79b3ffb1..6396103c 100644
--- a/internal/scheduler/testfixtures/testfixtures.go
+++ b/internal/scheduler/testfixtures/testfixtures.go
@@ -94,7 +94,8 @@ func TestSchedulingConfig() configuration.SchedulingConfig {
 			NodeEvictionProbability:                 1.0,
 			NodeOversubscriptionEvictionProbability: 1.0,
 		},
-		IndexedResources: TestResources,
+		IndexedResources:  TestResources,
+		IndexedNodeLabels: TestIndexedNodeLabels,
 		DominantResourceFairnessResourcesToConsider: TestResourceNames,
 		ExecutorTimeout:                  15 * time.Minute,
 		MaxUnacknowledgedJobsPerExecutor: math.MaxInt,
@@ -226,6 +227,17 @@ func WithNodeSelectorPodReq(selector map[string]string, req *schedulerobjects.Po
 	return req
 }
 
+func WithNodeUniformityLabelAnnotationJobs(label string, jobs []*jobdb.Job) []*jobdb.Job {
+	for _, job := range jobs {
+		req := job.PodRequirements()
+		if req.Annotations == nil {
+			req.Annotations = make(map[string]string)
+		}
+		req.Annotations[configuration.GangNodeUniformityLabelAnnotation] = label
+	}
+	return jobs
+}
+
 func WithNodeAffinityJobs(nodeSelectorTerms []v1.NodeSelectorTerm, jobs []*jobdb.Job) []*jobdb.Job {
 	for _, job := range jobs {
 		req := job.PodRequirements()