Repository: armadaproject/armada
Pull request #3076 (2023-10-27 10:43:37)

Compute scheduling key once on creation for jobdb.Job
This results in an overall speedup of about 23% when running the simulator for a particular scenario.

Before this change:
```
Measure-Command { .\simulator.exe --clusters .\internal\scheduler\simulator\testdata\clusters\cpu_1_1_100.yaml --workloads '.\internal\scheduler\simulator\testdata\workloads\small_big\*.yaml' --configs .\internal\scheduler\simulator\testdata\configs\defaultSchedulingConfig.yaml --logInterval 10000 }
...
Seconds           : 21
Milliseconds      : 70
Ticks             : 210707754
TotalDays         : 0.000243874715277778
TotalHours        : 0.00585299316666667
TotalMinutes      : 0.35117959
TotalSeconds      : 21.0707754
TotalMilliseconds : 21070.7754
```

After this change:
```
Measure-Command { .\simulator.exe --clusters .\internal\scheduler\simulator\testdata\clusters\cpu_1_1_100.yaml --workloads '.\internal\scheduler\simulator\testdata\workloads\small_big\*.yaml' --configs .\internal\scheduler\simulator\testdata\configs\defaultSchedulingConfig.yaml --logInterval 10000 }
...
Seconds           : 17
Milliseconds      : 69
Ticks             : 170699021
TotalDays         : 0.000197568311342593
TotalHours        : 0.00474163947222222
TotalMinutes      : 0.284498368333333
TotalSeconds      : 17.0699021
TotalMilliseconds : 17069.9021
```


Patch:
diff --git a/internal/scheduler/context/context.go b/internal/scheduler/context/context.go
index 8a52b497..5f22f7d0 100644
--- a/internal/scheduler/context/context.go
+++ b/internal/scheduler/context/context.go
@@ -225,7 +225,6 @@ func (sctx *SchedulingContext) ReportString(verbosity int32) string {
 
 func (sctx *SchedulingContext) AddGangSchedulingContext(gctx *GangSchedulingContext) (bool, error) {
 	allJobsEvictedInThisRound := true
-	allJobsSuccessful := true
 	numberOfSuccessfulJobs := 0
 	for _, jctx := range gctx.JobSchedulingContexts {
 		evictedInThisRound, err := sctx.AddJobSchedulingContext(jctx)
@@ -234,7 +233,6 @@ func (sctx *SchedulingContext) AddGangSchedulingContext(gctx *GangSchedulingCont
 		}
 		allJobsEvictedInThisRound = allJobsEvictedInThisRound && evictedInThisRound
 		isSuccess := jctx.IsSuccessful()
-		allJobsSuccessful = allJobsSuccessful && isSuccess
 		if isSuccess {
 			numberOfSuccessfulJobs++
 		}
diff --git a/internal/scheduler/gang_scheduler.go b/internal/scheduler/gang_scheduler.go
index f81e6bca..ad8ee5f4 100644
--- a/internal/scheduler/gang_scheduler.go
+++ b/internal/scheduler/gang_scheduler.go
@@ -11,6 +11,7 @@ import (
 	schedulercontext "github.com/armadaproject/armada/internal/scheduler/context"
 	"github.com/armadaproject/armada/internal/scheduler/interfaces"
 	"github.com/armadaproject/armada/internal/scheduler/nodedb"
+	"github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
 )
 
 // GangScheduler schedules one gang at a time. GangScheduler is not aware of queues.
@@ -38,20 +39,42 @@ func (sch *GangScheduler) SkipUnsuccessfulSchedulingKeyCheck() {
 	sch.skipUnsuccessfulSchedulingKeyCheck = true
 }
 
-func (sch *GangScheduler) updateGangSchedulingContextOnFailure(gctx *schedulercontext.GangSchedulingContext, gangAddedToSchedulingContext bool, unschedulableReason string) (err error) {
+func (sch *GangScheduler) updateGangSchedulingContextOnSuccess(gctx *schedulercontext.GangSchedulingContext, gangAddedToSchedulingContext bool) error {
+	if !gangAddedToSchedulingContext {
+		// Nothing to do.
+		return nil
+	}
+
+	// Evict any jobs added to the context marked as unsuccessful.
+	// This is necessary to support min-max gang-scheduling,
+	// where the gang is scheduled successfully if at least min of its members scheduled successfully.
+	// Here, we evict the memebers of the gang that were not scheduled successfully.
+	for _, jctx := range gctx.JobSchedulingContexts {
+		if !jctx.IsSuccessful() {
+			if _, err := sch.schedulingContext.EvictJob(jctx.Job); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func (sch *GangScheduler) updateGangSchedulingContextOnFailure(gctx *schedulercontext.GangSchedulingContext, gangAddedToSchedulingContext bool, unschedulableReason string) error {
+	// If the job was added to the context, remove it first.
 	if gangAddedToSchedulingContext {
 		failedJobs := util.Map(gctx.JobSchedulingContexts, func(jctx *schedulercontext.JobSchedulingContext) interfaces.LegacySchedulerJob { return jctx.Job })
-		if _, err = sch.schedulingContext.EvictGang(failedJobs); err != nil {
-			return
+		if _, err := sch.schedulingContext.EvictGang(failedJobs); err != nil {
+			return err
 		}
 	}
 
+	// Ensure all jobs have an unschedulableReason.
+	// Adding jobs with an unschedulableReason to the context ensures they're correctly accounted for as failed.
 	for _, jctx := range gctx.JobSchedulingContexts {
 		jctx.UnschedulableReason = unschedulableReason
 	}
-
-	if _, err = sch.schedulingContext.AddGangSchedulingContext(gctx); err != nil {
-		return
+	if _, err := sch.schedulingContext.AddGangSchedulingContext(gctx); err != nil {
+		return err
 	}
 
 	// Register unfeasible scheduling keys.
@@ -60,29 +83,19 @@ func (sch *GangScheduler) updateGangSchedulingContextOnFailure(gctx *schedulerco
 	// Since a gang may be unschedulable even if all its members are individually schedulable.
 	if !sch.skipUnsuccessfulSchedulingKeyCheck && gctx.Cardinality() == 1 {
 		jctx := gctx.JobSchedulingContexts[0]
-		schedulingKey := sch.schedulingContext.SchedulingKeyFromLegacySchedulerJob(jctx.Job)
-		if _, ok := sch.schedulingContext.UnfeasibleSchedulingKeys[schedulingKey]; !ok {
-			// Keep the first jctx for each unfeasible schedulingKey.
-			sch.schedulingContext.UnfeasibleSchedulingKeys[schedulingKey] = jctx
+		schedulingKey, ok := jctx.Job.GetSchedulingKey()
+		if !ok {
+			schedulingKey = sch.schedulingContext.SchedulingKeyFromLegacySchedulerJob(jctx.Job)
 		}
-	}
-
-	return
-}
-
-func (sch *GangScheduler) updateGangSchedulingContextOnSuccess(gctx *schedulercontext.GangSchedulingContext, gangAddedToSchedulingContext bool) (err error) {
-	if gangAddedToSchedulingContext {
-		jobs := util.Map(gctx.JobSchedulingContexts, func(jctx *schedulercontext.JobSchedulingContext) interfaces.LegacySchedulerJob { return jctx.Job })
-		if _, err = sch.schedulingContext.EvictGang(jobs); err != nil {
-			return
+		if schedulingKey != schedulerobjects.EmptySchedulingKey {
+			if _, ok := sch.schedulingContext.UnfeasibleSchedulingKeys[schedulingKey]; !ok {
+				// Keep the first jctx for each unfeasible schedulingKey.
+				sch.schedulingContext.UnfeasibleSchedulingKeys[schedulingKey] = jctx
+			}
 		}
 	}
 
-	if _, err = sch.schedulingContext.AddGangSchedulingContext(gctx); err != nil {
-		return
-	}
-
-	return
+	return nil
 }
 
 func (sch *GangScheduler) Schedule(ctx *armadacontext.Context, gctx *schedulercontext.GangSchedulingContext) (ok bool, unschedulableReason string, err error) {
@@ -93,8 +106,7 @@ func (sch *GangScheduler) Schedule(ctx *armadacontext.Context, gctx *schedulerco
 		}
 	}
 
-	// This deferred function ensures unschedulable jobs are registered as such
-	// and sets sch.queueScheduledInPreviousCall.
+	// This deferred function ensures unschedulable jobs are registered as such.
 	gangAddedToSchedulingContext := false
 	defer func() {
 		// Do nothing if an error occurred.
@@ -115,8 +127,6 @@ func (sch *GangScheduler) Schedule(ctx *armadacontext.Context, gctx *schedulerco
 		} else {
 			err = sch.updateGangSchedulingContextOnFailure(gctx, gangAddedToSchedulingContext, unschedulableReason)
 		}
-
-		return
 	}()
 
 	if _, err = sch.schedulingContext.AddGangSchedulingContext(gctx); err != nil {
diff --git a/internal/scheduler/interfaces/interfaces.go b/internal/scheduler/interfaces/interfaces.go
index 121fae18..674282fe 100644
--- a/internal/scheduler/interfaces/interfaces.go
+++ b/internal/scheduler/interfaces/interfaces.go
@@ -24,4 +24,7 @@ type LegacySchedulerJob interface {
 	GetTolerations() []v1.Toleration
 	GetResourceRequirements() v1.ResourceRequirements
 	GetQueueTtlSeconds() int64
+	// Returns (schedulingKey, true) if the job has a scheduling key associated with it.
+	// Returns (emptySchedulingKey, false) otherwise, where emptySchedulingKey is the zero value of the SchedulingKey type.
+	GetSchedulingKey() (schedulerobjects.SchedulingKey, bool)
 }
diff --git a/internal/scheduler/jobdb/job.go b/internal/scheduler/jobdb/job.go
index 8980b55c..06c61167 100644
--- a/internal/scheduler/jobdb/job.go
+++ b/internal/scheduler/jobdb/job.go
@@ -29,6 +29,8 @@ type Job struct {
 	// Logical timestamp indicating the order in which jobs are submitted.
 	// Jobs with identical Queue and Priority are sorted by this.
 	created int64
+	// Hash of the scheduling requirements of the job.
+	schedulingKey schedulerobjects.SchedulingKey
 	// True if the job is currently queued.
 	// If this is set then the job will not be considered for scheduling.
 	queued bool
@@ -58,39 +60,6 @@ func EmptyJob(id string) *Job {
 	return &Job{id: id, runsById: map[uuid.UUID]*JobRun{}}
 }
 
-// NewJob creates a new scheduler job
-func NewJob(
-	jobId string,
-	jobset string,
-	queue string,
-	priority uint32,
-	schedulingInfo *schedulerobjects.JobSchedulingInfo,
-	queued bool,
-	queuedVersion int32,
-	cancelRequested bool,
-	cancelByJobsetRequested bool,
-	cancelled bool,
-	created int64,
-) *Job {
-	job := &Job{
-		id:                      jobId,
-		jobset:                  jobset,
-		queue:                   queue,
-		queued:                  queued,
-		queuedVersion:           queuedVersion,
-		priority:                priority,
-		requestedPriority:       priority,
-		jobSchedulingInfo:       schedulingInfo,
-		cancelRequested:         cancelRequested,
-		cancelByJobsetRequested: cancelByJobsetRequested,
-		cancelled:               cancelled,
-		created:                 created,
-		runsById:                map[uuid.UUID]*JobRun{},
-	}
-	job.ensureJobSchedulingInfoFieldsInitialised()
-	return job
-}
-
 func (job *Job) ensureJobSchedulingInfoFieldsInitialised() {
 	// Initialise the annotation and nodeSelector maps if nil.
 	// Since those need to be mutated in-place.
@@ -146,6 +115,13 @@ func (job *Job) Priority() uint32 {
 	return job.priority
 }
 
+// GetSchedulingKey returns the scheduling key associated with a job.
+// The second return value is always true since scheduling keys are computed at job creation time.
+// This is needed for compatibility with interfaces.LegacySchedulerJob.
+func (job *Job) GetSchedulingKey() (schedulerobjects.SchedulingKey, bool) {
+	return job.schedulingKey, true
+}
+
 // GetPerQueuePriority exists for compatibility with the LegacyJob interface.
 func (job *Job) GetPerQueuePriority() uint32 {
 	return job.priority
diff --git a/internal/scheduler/jobdb/jobdb.go b/internal/scheduler/jobdb/jobdb.go
index 45d02b5b..01b167cf 100644
--- a/internal/scheduler/jobdb/jobdb.go
+++ b/internal/scheduler/jobdb/jobdb.go
@@ -7,6 +7,8 @@ import (
 	"github.com/google/uuid"
 	"github.com/pkg/errors"
 	"golang.org/x/exp/maps"
+
+	"github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
 )
 
 var (
@@ -15,25 +17,69 @@ var (
 )
 
 type JobDb struct {
-	jobsById        *immutable.Map[string, *Job]
-	jobsByRunId     *immutable.Map[uuid.UUID, string]
-	jobsByQueue     map[string]immutable.SortedSet[*Job]
-	queuedJobsByTtl *immutable.SortedSet[*Job]
-	copyMutex       sync.Mutex
-	writerMutex     sync.Mutex
+	jobsById               *immutable.Map[string, *Job]
+	jobsByRunId            *immutable.Map[uuid.UUID, string]
+	jobsByQueue            map[string]immutable.SortedSet[*Job]
+	queuedJobsByTtl        *immutable.SortedSet[*Job]
+	schedulingKeyGenerator *schedulerobjects.SchedulingKeyGenerator
+	copyMutex              sync.Mutex
+	writerMutex            sync.Mutex
 }
 
 func NewJobDb() *JobDb {
+	return NewJobDbWithSchedulingKeyGenerator(schedulerobjects.NewSchedulingKeyGenerator())
+}
+
+func NewJobDbWithSchedulingKeyGenerator(skg *schedulerobjects.SchedulingKeyGenerator) *JobDb {
 	return &JobDb{
-		jobsById:        immutable.NewMap[string, *Job](nil),
-		jobsByRunId:     immutable.NewMap[uuid.UUID, string](&UUIDHasher{}),
-		jobsByQueue:     map[string]immutable.SortedSet[*Job]{},
-		queuedJobsByTtl: &emptyQueuedJobsByTtl,
-		copyMutex:       sync.Mutex{},
+		jobsById:               immutable.NewMap[string, *Job](nil),
+		jobsByRunId:            immutable.NewMap[uuid.UUID, string](&UUIDHasher{}),
+		jobsByQueue:            map[string]immutable.SortedSet[*Job]{},
+		queuedJobsByTtl:        &emptyQueuedJobsByTtl,
+		schedulingKeyGenerator: skg,
+	}
+}
+
+// NewJob creates a new scheduler job.
+// The new job is not automatically inserted into the jobDb; call jobDb.Upsert to upsert it.
+func (jobDb *JobDb) NewJob(
+	jobId string,
+	jobset string,
+	queue string,
+	priority uint32,
+	schedulingInfo *schedulerobjects.JobSchedulingInfo,
+	queued bool,
+	queuedVersion int32,
+	cancelRequested bool,
+	cancelByJobsetRequested bool,
+	cancelled bool,
+	created int64,
+) *Job {
+	var schedulingKey schedulerobjects.SchedulingKey
+	if preq := schedulingInfo.GetPodRequirements(); preq != nil {
+		schedulingKey = jobDb.schedulingKeyGenerator.KeyFromPodRequirements(preq)
+	}
+	job := &Job{
+		id:                      jobId,
+		queue:                   queue,
+		jobset:                  jobset,
+		priority:                priority,
+		queued:                  queued,
+		queuedVersion:           queuedVersion,
+		requestedPriority:       priority,
+		created:                 created,
+		schedulingKey:           schedulingKey,
+		jobSchedulingInfo:       schedulingInfo,
+		cancelRequested:         cancelRequested,
+		cancelByJobsetRequested: cancelByJobsetRequested,
+		cancelled:               cancelled,
+		runsById:                map[uuid.UUID]*JobRun{},
 	}
+	job.ensureJobSchedulingInfoFieldsInitialised()
+	return job
 }
 
-// Upsert will insert the given jobs if they don't already exist or update the if they do
+// Upsert will insert the given jobs if they don't already exist or update them if they do.
 func (jobDb *JobDb) Upsert(txn *Txn, jobs []*Job) error {
 	if err := jobDb.checkWritableTransaction(txn); err != nil {
 		return err
@@ -57,7 +103,7 @@ func (jobDb *JobDb) Upsert(txn *Txn, jobs []*Job) error {
 		}
 	}
 
-	// Now need to insert jobs, runs and queuedJobs.  This can be done in parallel
+	// Now need to insert jobs, runs and queuedJobs. This can be done in parallel.
 	wg := sync.WaitGroup{}
 	wg.Add(3)
 
diff --git a/internal/scheduler/preempting_queue_scheduler.go b/internal/scheduler/preempting_queue_scheduler.go
index 7be10464..73b195fa 100644
--- a/internal/scheduler/preempting_queue_scheduler.go
+++ b/internal/scheduler/preempting_queue_scheduler.go
@@ -503,7 +503,7 @@ func addEvictedJobsToNodeDb(ctx *armadacontext.Context, sctx *schedulercontext.S
 		if err != nil {
 			return err
 		}
-		gangItByQueue[qctx.Queue] = NewQueuedGangIterator(sctx, jobIt, 0)
+		gangItByQueue[qctx.Queue] = NewQueuedGangIterator(sctx, jobIt, 0, false)
 	}
 	qr := NewMinimalQueueRepositoryFromSchedulingContext(sctx)
 	candidateGangIterator, err := NewCandidateGangIterator(qr, sctx.FairnessCostProvider, gangItByQueue)
diff --git a/internal/scheduler/queue_scheduler.go b/internal/scheduler/queue_scheduler.go
index 95683904..96497bb9 100644
--- a/internal/scheduler/queue_scheduler.go
+++ b/internal/scheduler/queue_scheduler.go
@@ -43,7 +43,7 @@ func NewQueueScheduler(
 	}
 	gangIteratorsByQueue := make(map[string]*QueuedGangIterator)
 	for queue, it := range jobIteratorByQueue {
-		gangIteratorsByQueue[queue] = NewQueuedGangIterator(sctx, it, constraints.MaxQueueLookback)
+		gangIteratorsByQueue[queue] = NewQueuedGangIterator(sctx, it, constraints.MaxQueueLookback, true)
 	}
 	candidateGangIterator, err := NewCandidateGangIterator(sctx, sctx.FairnessCostProvider, gangIteratorsByQueue)
 	if err != nil {
@@ -148,17 +148,20 @@ type QueuedGangIterator struct {
 	jobsByGangId map[string][]interfaces.LegacySchedulerJob
 	// Maximum number of jobs to look at before giving up.
 	maxLookback uint
+	// If true, do not yield jobs known to be unschedulable.
+	skipKnownUnschedulableJobs bool
 	// Number of jobs we have seen so far.
 	jobsSeen uint
 	next     *schedulercontext.GangSchedulingContext
 }
 
-func NewQueuedGangIterator(sctx *schedulercontext.SchedulingContext, it JobIterator, maxLookback uint) *QueuedGangIterator {
+func NewQueuedGangIterator(sctx *schedulercontext.SchedulingContext, it JobIterator, maxLookback uint, skipKnownUnschedulableJobs bool) *QueuedGangIterator {
 	return &QueuedGangIterator{
-		schedulingContext:  sctx,
-		queuedJobsIterator: it,
-		maxLookback:        maxLookback,
-		jobsByGangId:       make(map[string][]interfaces.LegacySchedulerJob),
+		schedulingContext:          sctx,
+		queuedJobsIterator:         it,
+		maxLookback:                maxLookback,
+		skipKnownUnschedulableJobs: skipKnownUnschedulableJobs,
+		jobsByGangId:               make(map[string][]interfaces.LegacySchedulerJob),
 	}
 }
 
@@ -209,23 +212,28 @@ func (it *QueuedGangIterator) Peek() (*schedulercontext.GangSchedulingContext, e
 		}
 
 		// Skip this job if it's known to be unschedulable.
-		if len(it.schedulingContext.UnfeasibleSchedulingKeys) > 0 {
-			schedulingKey := it.schedulingContext.SchedulingKeyFromLegacySchedulerJob(job)
-			if unsuccessfulJctx, ok := it.schedulingContext.UnfeasibleSchedulingKeys[schedulingKey]; ok {
-				// TODO: For performance, we should avoid creating new objects and instead reference the existing one.
-				jctx := &schedulercontext.JobSchedulingContext{
-					Created:              time.Now(),
-					JobId:                job.GetId(),
-					Job:                  job,
-					UnschedulableReason:  unsuccessfulJctx.UnschedulableReason,
-					PodSchedulingContext: unsuccessfulJctx.PodSchedulingContext,
-					// TODO: Move this into gang scheduling context
-					GangMinCardinality: 1,
-				}
-				if _, err := it.schedulingContext.AddJobSchedulingContext(jctx); err != nil {
-					return nil, err
+		if it.skipKnownUnschedulableJobs && len(it.schedulingContext.UnfeasibleSchedulingKeys) > 0 {
+			schedulingKey, ok := job.GetSchedulingKey()
+			if !ok {
+				schedulingKey = it.schedulingContext.SchedulingKeyFromLegacySchedulerJob(job)
+			}
+			if schedulingKey != schedulerobjects.EmptySchedulingKey {
+				if unsuccessfulJctx, ok := it.schedulingContext.UnfeasibleSchedulingKeys[schedulingKey]; ok {
+					// TODO: For performance, we should avoid creating new objects and instead reference the existing one.
+					jctx := &schedulercontext.JobSchedulingContext{
+						Created:              time.Now(),
+						JobId:                job.GetId(),
+						Job:                  job,
+						UnschedulableReason:  unsuccessfulJctx.UnschedulableReason,
+						PodSchedulingContext: unsuccessfulJctx.PodSchedulingContext,
+						// TODO: Move this into gang scheduling context
+						GangMinCardinality: 1,
+					}
+					if _, err := it.schedulingContext.AddJobSchedulingContext(jctx); err != nil {
+						return nil, err
+					}
+					continue
 				}
-				continue
 			}
 		}
 
diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go
index 5b63bd78..6df625db 100644
--- a/internal/scheduler/scheduler.go
+++ b/internal/scheduler/scheduler.go
@@ -338,12 +338,10 @@ func (s *Scheduler) syncState(ctx *armadacontext.Context) ([]*jobdb.Job, error)
 	}
 
 	jobsToUpdate := maps.Values(jobsToUpdateById)
-	err = s.jobDb.Upsert(txn, jobsToUpdate)
-	if err != nil {
+	if err := s.jobDb.Upsert(txn, jobsToUpdate); err != nil {
 		return nil, err
 	}
-	err = s.jobDb.BatchDelete(txn, jobsToDelete)
-	if err != nil {
+	if err := s.jobDb.BatchDelete(txn, jobsToDelete); err != nil {
 		return nil, err
 	}
 	txn.Commit()
@@ -952,7 +950,7 @@ func (s *Scheduler) schedulerJobFromDatabaseJob(dbJob *database.Job) (*jobdb.Job
 		return nil, errors.Wrapf(err, "error unmarshalling scheduling info for job %s", dbJob.JobID)
 	}
 	s.internJobSchedulingInfoStrings(schedulingInfo)
-	return jobdb.NewJob(
+	return s.jobDb.NewJob(
 		dbJob.JobID,
 		s.stringInterner.Intern(dbJob.JobSet),
 		s.stringInterner.Intern(dbJob.Queue),
diff --git a/internal/scheduler/schedulerobjects/podutils.go b/internal/scheduler/schedulerobjects/podutils.go
index 9bd476ca..5e2f92b9 100644
--- a/internal/scheduler/schedulerobjects/podutils.go
+++ b/internal/scheduler/schedulerobjects/podutils.go
@@ -3,14 +3,19 @@ package schedulerobjects
 import (
 	"crypto/rand"
 	"encoding/binary"
+	"sync"
 
 	"github.com/minio/highwayhash"
 	"golang.org/x/exp/slices"
 	v1 "k8s.io/api/core/v1"
 )
 
+// SchedulingKey is a hash of the scheduling requirements of a job.
+// This key is used to efficiently mark jobs as unschedulable.
 type SchedulingKey [highwayhash.Size]byte
 
+var EmptySchedulingKey SchedulingKey
+
 func (req *PodRequirements) GetAffinityNodeSelector() *v1.NodeSelector {
 	affinity := req.Affinity
 	if affinity == nil {
@@ -25,11 +30,12 @@ func (req *PodRequirements) GetAffinityNodeSelector() *v1.NodeSelector {
 
 // SchedulingKeyGenerator is used to generate scheduling keys efficiently.
 // A scheduling key is the canonical hash of the scheduling requirements of a job.
-// All memory is allocated up-front and re-used. Not thread-safe.
+// All memory is allocated up-front and re-used. Thread-safe.
 type SchedulingKeyGenerator struct {
 	s      PodRequirementsSerialiser
 	key    []byte
 	buffer []byte
+	sync.Mutex
 }
 
 func NewSchedulingKeyGenerator() *SchedulingKeyGenerator {
@@ -45,6 +51,16 @@ func NewSchedulingKeyGenerator() *SchedulingKeyGenerator {
 	}
 }
 
+func (skg *SchedulingKeyGenerator) KeyFromPodRequirements(preq *PodRequirements) SchedulingKey {
+	return skg.Key(
+		preq.NodeSelector,
+		preq.Affinity,
+		preq.Tolerations,
+		preq.ResourceRequirements.Requests,
+		preq.Priority,
+	)
+}
+
 func (skg *SchedulingKeyGenerator) Key(
 	nodeSelector map[string]string,
 	affinity *v1.Affinity,
@@ -52,6 +68,8 @@ func (skg *SchedulingKeyGenerator) Key(
 	requests v1.ResourceList,
 	priority int32,
 ) SchedulingKey {
+	skg.Mutex.Lock()
+	defer skg.Mutex.Unlock()
 	skg.buffer = skg.buffer[0:0]
 	skg.buffer = skg.s.AppendRequirements(
 		skg.buffer,
diff --git a/internal/scheduler/simulator/simulator.go b/internal/scheduler/simulator/simulator.go
index 2fa165a7..53d65735 100644
--- a/internal/scheduler/simulator/simulator.go
+++ b/internal/scheduler/simulator/simulator.go
@@ -568,7 +568,7 @@ func (s *Simulator) handleEventSequence(ctx *armadacontext.Context, es *armadaev
 		switch eventType := event.GetEvent().(type) {
 		case *armadaevents.EventSequence_Event_SubmitJob:
 			s.shouldSchedule = true
-			ok, err = s.handleSubmitJob(txn, event.GetSubmitJob(), *event.Created, es)
+			ok, err = s.handleSubmitJob(s.jobDb, txn, event.GetSubmitJob(), *event.Created, es)
 		case *armadaevents.EventSequence_Event_JobRunLeased:
 			ok, err = s.handleJobRunLeased(txn, event.GetJobRunLeased())
 		case *armadaevents.EventSequence_Event_JobSucceeded:
@@ -616,12 +616,12 @@ func (s *Simulator) handleEventSequence(ctx *armadacontext.Context, es *armadaev
 	return nil
 }
 
-func (s *Simulator) handleSubmitJob(txn *jobdb.Txn, e *armadaevents.SubmitJob, time time.Time, eventSequence *armadaevents.EventSequence) (bool, error) {
+func (s *Simulator) handleSubmitJob(jobDb *jobdb.JobDb, txn *jobdb.Txn, e *armadaevents.SubmitJob, time time.Time, eventSequence *armadaevents.EventSequence) (bool, error) {
 	schedulingInfo, err := scheduleringester.SchedulingInfoFromSubmitJob(e, time, s.schedulingConfig.Preemption.PriorityClasses)
 	if err != nil {
 		return false, err
 	}
-	job := jobdb.NewJob(
+	job := jobDb.NewJob(
 		armadaevents.UlidFromProtoUuid(e.JobId).String(),
 		eventSequence.JobSetName,
 		eventSequence.Queue,
diff --git a/pkg/api/util.go b/pkg/api/util.go
index 7906d5cd..21a41717 100644
--- a/pkg/api/util.go
+++ b/pkg/api/util.go
@@ -250,6 +250,12 @@ func (job *Job) GetResourceRequirements() v1.ResourceRequirements {
 	}
 }
 
+// GetSchedulingKey returns the scheduling key associated with a job.
+// The second return value is always false since scheduling keys are not pre-computed for these jobs.
+func (job *Job) GetSchedulingKey() (schedulerobjects.SchedulingKey, bool) {
+	return schedulerobjects.SchedulingKey{}, false
+}
+
 func (job *Job) GetJobSet() string {
 	return job.JobSetId
 }


Test patch:
diff --git a/internal/scheduler/gang_scheduler_test.go b/internal/scheduler/gang_scheduler_test.go
index cce5b296..dcd79654 100644
--- a/internal/scheduler/gang_scheduler_test.go
+++ b/internal/scheduler/gang_scheduler_test.go
@@ -60,7 +60,10 @@ func TestGangScheduler(t *testing.T) {
 			SchedulingConfig: testfixtures.TestSchedulingConfig(),
 			Nodes:            testfixtures.N32CpuNodes(1, testfixtures.TestPriorities),
 			Gangs: [][]*jobdb.Job{
-				testfixtures.WithGangAnnotationsAndMinCardinalityJobs(testfixtures.N1Cpu4GiJobs("A", testfixtures.PriorityClass0, 40), 32),
+				testfixtures.WithGangAnnotationsAndMinCardinalityJobs(
+					32,
+					testfixtures.N1Cpu4GiJobs("A", testfixtures.PriorityClass0, 40),
+				),
 			},
 			ExpectedScheduledIndices: testfixtures.IntRange(0, 0),
 			ExpectedScheduledJobs:    []int{32},
@@ -69,7 +72,10 @@ func TestGangScheduler(t *testing.T) {
 			SchedulingConfig: testfixtures.TestSchedulingConfig(),
 			Nodes:            testfixtures.N32CpuNodes(1, testfixtures.TestPriorities),
 			Gangs: [][]*jobdb.Job{
-				testfixtures.WithGangAnnotationsAndMinCardinalityJobs(testfixtures.N1Cpu4GiJobs("A", testfixtures.PriorityClass0, 40), 33),
+				testfixtures.WithGangAnnotationsAndMinCardinalityJobs(
+					33,
+					testfixtures.N1Cpu4GiJobs("A", testfixtures.PriorityClass0, 40),
+				),
 			},
 			ExpectedScheduledIndices: nil,
 			ExpectedScheduledJobs:    []int{0},
@@ -88,7 +94,10 @@ func TestGangScheduler(t *testing.T) {
 			SchedulingConfig: testfixtures.TestSchedulingConfig(),
 			Nodes:            testfixtures.N32CpuNodes(1, testfixtures.TestPriorities),
 			Gangs: [][]*jobdb.Job{
-				testfixtures.WithGangAnnotationsAndMinCardinalityJobs(testfixtures.N1Cpu4GiJobs("A", testfixtures.PriorityClass0, 33), 32),
+				testfixtures.WithGangAnnotationsAndMinCardinalityJobs(
+					32,
+					testfixtures.N1Cpu4GiJobs("A", testfixtures.PriorityClass0, 33),
+				),
 				testfixtures.WithGangAnnotationsJobs(testfixtures.N1Cpu4GiJobs("A", testfixtures.PriorityClass0, 1)),
 			},
 			ExpectedScheduledIndices: testfixtures.IntRange(0, 0),
diff --git a/internal/scheduler/jobdb/job_test.go b/internal/scheduler/jobdb/job_test.go
index 2b3df48f..0d7ee875 100644
--- a/internal/scheduler/jobdb/job_test.go
+++ b/internal/scheduler/jobdb/job_test.go
@@ -24,7 +24,10 @@ var schedulingInfo = &schedulerobjects.JobSchedulingInfo{
 	},
 }
 
-var baseJob = NewJob(
+// Used for creating jobs.
+var jobDb = NewJobDb()
+
+var baseJob = jobDb.NewJob(
 	"test-job",
 	"test-jobset",
 	"test-queue",
@@ -35,7 +38,8 @@ var baseJob = NewJob(
 	false,
 	false,
 	false,
-	3)
+	3,
+)
 
 var baseRun = &JobRun{
 	id:        uuid.New(),
@@ -291,9 +295,9 @@ func TestJob_TestWithCreated(t *testing.T) {
 }
 
 func TestJob_DeepCopy(t *testing.T) {
-	original := NewJob("test-job", "test-jobset", "test-queue", 2, schedulingInfo, true, 0, false, false, false, 3)
+	original := jobDb.NewJob("test-job", "test-jobset", "test-queue", 2, schedulingInfo, true, 0, false, false, false, 3)
 	original = original.WithUpdatedRun(baseJobRun.DeepCopy())
-	expected := NewJob("test-job", "test-jobset", "test-queue", 2, schedulingInfo, true, 0, false, false, false, 3)
+	expected := jobDb.NewJob("test-job", "test-jobset", "test-queue", 2, schedulingInfo, true, 0, false, false, false, 3)
 	expected = expected.WithUpdatedRun(baseJobRun.DeepCopy())
 
 	result := original.DeepCopy()
@@ -345,7 +349,7 @@ func TestJobSchedulingInfoFieldsInitialised(t *testing.T) {
 	assert.Nil(t, infoWithNilFields.GetPodRequirements().NodeSelector)
 	assert.Nil(t, infoWithNilFields.GetPodRequirements().Annotations)
 
-	job := NewJob("test-job", "test-jobset", "test-queue", 2, infoWithNilFieldsCopy, true, 0, false, false, false, 3)
+	job := jobDb.NewJob("test-job", "test-jobset", "test-queue", 2, infoWithNilFieldsCopy, true, 0, false, false, false, 3)
 	assert.NotNil(t, job.GetNodeSelector())
 	assert.NotNil(t, job.GetAnnotations())
 
diff --git a/internal/scheduler/jobdb/jobdb_test.go b/internal/scheduler/jobdb/jobdb_test.go
index 1a1a6371..064db81a 100644
--- a/internal/scheduler/jobdb/jobdb_test.go
+++ b/internal/scheduler/jobdb/jobdb_test.go
@@ -10,6 +10,7 @@ import (
 	"golang.org/x/exp/slices"
 
 	"github.com/armadaproject/armada/internal/common/util"
+	"github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
 )
 
 func TestJobDb_TestUpsert(t *testing.T) {
@@ -198,6 +199,34 @@ func TestJobDb_TestBatchDelete(t *testing.T) {
 	require.Error(t, err)
 }
 
+func TestJobDb_SchedulingKey(t *testing.T) {
+	podRequirements := &schedulerobjects.PodRequirements{
+		NodeSelector: map[string]string{"foo": "bar"},
+		Priority:     2,
+	}
+	jobSchedulingInfo := &schedulerobjects.JobSchedulingInfo{
+		ObjectRequirements: []*schedulerobjects.ObjectRequirements{
+			{
+				Requirements: &schedulerobjects.ObjectRequirements_PodRequirements{
+					PodRequirements: podRequirements,
+				},
+			},
+		},
+	}
+	jobDb := NewJobDb()
+	job := jobDb.NewJob("jobId", "jobSet", "queue", 1, jobSchedulingInfo, false, 0, false, false, false, 2)
+	assert.Equal(t,
+		jobDb.schedulingKeyGenerator.Key(
+			podRequirements.NodeSelector,
+			podRequirements.Affinity,
+			podRequirements.Tolerations,
+			podRequirements.ResourceRequirements.Requests,
+			podRequirements.Priority,
+		),
+		job.schedulingKey,
+	)
+}
+
 func newJob() *Job {
 	return &Job{
 		id:                util.NewULID(),
diff --git a/internal/scheduler/nodedb/nodedb_test.go b/internal/scheduler/nodedb/nodedb_test.go
index afeeadec..163b4f51 100644
--- a/internal/scheduler/nodedb/nodedb_test.go
+++ b/internal/scheduler/nodedb/nodedb_test.go
@@ -498,7 +498,7 @@ func TestScheduleMany(t *testing.T) {
 		// Attempts to schedule 33 jobs with a minimum gang cardinality of 32 jobs. One fails, but the overall result is a success.
 		"simple success with min cardinality": {
 			Nodes:         testfixtures.N32CpuNodes(1, testfixtures.TestPriorities),
-			Jobs:          [][]*jobdb.Job{testfixtures.WithGangAnnotationsAndMinCardinalityJobs(testfixtures.N1Cpu4GiJobs("A", testfixtures.PriorityClass0, 33), 32)},
+			Jobs:          [][]*jobdb.Job{testfixtures.WithGangAnnotationsAndMinCardinalityJobs(32, testfixtures.N1Cpu4GiJobs("A", testfixtures.PriorityClass0, 33))},
 			ExpectSuccess: []bool{true},
 		},
 		// Attempts to schedule 33 jobs with a minimum gang cardinality of 33. The overall result fails.
diff --git a/internal/scheduler/preempting_queue_scheduler_test.go b/internal/scheduler/preempting_queue_scheduler_test.go
index 68dc698a..d9269d32 100644
--- a/internal/scheduler/preempting_queue_scheduler_test.go
+++ b/internal/scheduler/preempting_queue_scheduler_test.go
@@ -6,7 +6,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/exp/maps"
@@ -518,7 +517,10 @@ func TestPreemptingQueueScheduler(t *testing.T) {
 				{
 					// Schedule a gang across two nodes.
 					JobsByQueue: map[string][]*jobdb.Job{
-						"A": testfixtures.WithGangAnnotationsAndMinCardinalityJobs(testfixtures.N32Cpu256GiJobs("A", testfixtures.PriorityClass0, 2), 1),
+						"A": testfixtures.WithGangAnnotationsAndMinCardinalityJobs(
+							1,
+							testfixtures.N32Cpu256GiJobs("A", testfixtures.PriorityClass0, 2),
+						),
 					},
 					ExpectedScheduledIndices: map[string][]int{
 						"A": testfixtures.IntRange(0, 1),
@@ -1400,10 +1402,10 @@ func TestPreemptingQueueScheduler(t *testing.T) {
 			}
 
 			// Run the scheduler.
-			log := logrus.NewEntry(logrus.New())
+			ctx := armadacontext.Background()
 			for i, round := range tc.Rounds {
-				log = log.WithField("round", i)
-				log.Infof("starting scheduling round %d", i)
+				ctx.FieldLogger = ctx.WithField("round", i)
+				ctx.Infof("starting scheduling round %d", i)
 
 				// Reset the queues between rounds.
 				repo.jobsByQueue = make(map[string][]interfaces.LegacySchedulerJob)
@@ -1504,7 +1506,7 @@ func TestPreemptingQueueScheduler(t *testing.T) {
 				if tc.SchedulingConfig.EnableNewPreemptionStrategy {
 					sch.EnableNewPreemptionStrategy()
 				}
-				result, err := sch.Schedule(armadacontext.Background())
+				result, err := sch.Schedule(ctx)
 				require.NoError(t, err)
 				jobIdsByGangId = sch.jobIdsByGangId
 				gangIdByJobId = sch.gangIdByJobId
diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go
index f513d79a..473d89da 100644
--- a/internal/scheduler/scheduler_test.go
+++ b/internal/scheduler/scheduler_test.go
@@ -22,6 +22,7 @@ import (
 	"github.com/armadaproject/armada/internal/scheduler/jobdb"
 	"github.com/armadaproject/armada/internal/scheduler/kubernetesobjects/affinity"
 	"github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
+	"github.com/armadaproject/armada/internal/scheduler/testfixtures"
 	"github.com/armadaproject/armada/pkg/api"
 	"github.com/armadaproject/armada/pkg/armadaevents"
 )
@@ -105,7 +106,7 @@ var (
 	})
 )
 
-var queuedJob = jobdb.NewJob(
+var queuedJob = testfixtures.JobDb.NewJob(
 	util.NewULID(),
 	"testJobset",
 	"testQueue",
@@ -118,7 +119,7 @@ var queuedJob = jobdb.NewJob(
 	false,
 	1)
 
-var queuedJobWithExpiredTtl = jobdb.NewJob(
+var queuedJobWithExpiredTtl = testfixtures.JobDb.NewJob(
 	util.NewULID(),
 	"testJobset",
 	"testQueue",
@@ -131,7 +132,7 @@ var queuedJobWithExpiredTtl = jobdb.NewJob(
 	false,
 	1)
 
-var leasedJob = jobdb.NewJob(
+var leasedJob = testfixtures.JobDb.NewJob(
 	util.NewULID(),
 	"testJobset",
 	"testQueue",
@@ -153,7 +154,7 @@ var defaultJobRunError = &armadaevents.Error{
 	},
 }
 
-var leasedFailFastJob = jobdb.NewJob(
+var leasedFailFastJob = testfixtures.JobDb.NewJob(
 	util.NewULID(),
 	"testJobset",
 	"testQueue",
@@ -168,7 +169,7 @@ var leasedFailFastJob = jobdb.NewJob(
 
 var (
 	requeuedJobId = util.NewULID()
-	requeuedJob   = jobdb.NewJob(
+	requeuedJob   = testfixtures.JobDb.NewJob(
 		requeuedJobId,
 		"testJobset",
 		"testQueue",
@@ -1003,9 +1004,14 @@ func TestScheduler_TestSyncState(t *testing.T) {
 				1*time.Hour,
 				maxNumberOfAttempts,
 				nodeIdLabel,
-				schedulerMetrics)
+				schedulerMetrics,
+			)
 			require.NoError(t, err)
 
+			// The SchedulingKeyGenerator embedded in the jobDb has some randomness,
+			// which must be consistent within tests.
+			sched.jobDb = jobdb.NewJobDbWithSchedulingKeyGenerator(testfixtures.SchedulingKeyGenerator)
+
 			// insert initial jobs
 			txn := sched.jobDb.WriteTxn()
 			err = sched.jobDb.Upsert(txn, tc.initialJobs)
diff --git a/internal/scheduler/scheduling_algo_test.go b/internal/scheduler/scheduling_algo_test.go
index 894ee708..016396e5 100644
--- a/internal/scheduler/scheduling_algo_test.go
+++ b/internal/scheduler/scheduling_algo_test.go
@@ -277,10 +277,13 @@ func TestSchedule(t *testing.T) {
 			expectedScheduledIndices: []int{0, 1},
 		},
 		"gang scheduling successful with some jobs failing to schedule above min cardinality": {
-			schedulingConfig:         testfixtures.TestSchedulingConfig(),
-			executors:                []*schedulerobjects.Executor{testfixtures.Test1Node32CoreExecutor("executor1")},
-			queues:                   []*database.Queue{{Name: "A", Weight: 100}},
-			queuedJobs:               testfixtures.WithGangAnnotationsAndMinCardinalityJobs(testfixtures.N16Cpu128GiJobs("A", testfixtures.PriorityClass0, 10), 2),
+			schedulingConfig: testfixtures.TestSchedulingConfig(),
+			executors:        []*schedulerobjects.Executor{testfixtures.Test1Node32CoreExecutor("executor1")},
+			queues:           []*database.Queue{{Name: "A", Weight: 100}},
+			queuedJobs: testfixtures.WithGangAnnotationsAndMinCardinalityJobs(
+				2,
+				testfixtures.N16Cpu128GiJobs("A", testfixtures.PriorityClass0, 10),
+			),
 			expectedScheduledIndices: []int{0, 1},
 			expectedFailedJobCount:   8,
 		},
diff --git a/internal/scheduler/testfixtures/testfixtures.go b/internal/scheduler/testfixtures/testfixtures.go
index a4e45d8a..6efa1379 100644
--- a/internal/scheduler/testfixtures/testfixtures.go
+++ b/internal/scheduler/testfixtures/testfixtures.go
@@ -34,8 +34,11 @@ const (
 )
 
 var (
-	BaseTime, _         = time.Parse("2006-01-02T15:04:05.000Z", "2022-03-01T15:04:05.000Z")
-	TestPriorityClasses = map[string]types.PriorityClass{
+	// Used for job creation.
+	SchedulingKeyGenerator = schedulerobjects.NewSchedulingKeyGenerator()
+	JobDb                  = jobdb.NewJobDbWithSchedulingKeyGenerator(SchedulingKeyGenerator)
+	BaseTime, _            = time.Parse("2006-01-02T15:04:05.000Z", "2022-03-01T15:04:05.000Z")
+	TestPriorityClasses    = map[string]types.PriorityClass{
 		PriorityClass0:               {Priority: 0, Preemptible: true},
 		PriorityClass1:               {Priority: 1, Preemptible: true},
 		PriorityClass2:               {Priority: 2, Preemptible: true},
@@ -326,12 +329,16 @@ func WithGangAnnotationsJobs(jobs []*jobdb.Job) []*jobdb.Job {
 	)
 }
 
-func WithGangAnnotationsAndMinCardinalityJobs(jobs []*jobdb.Job, minimumCardinality int) []*jobdb.Job {
+func WithGangAnnotationsAndMinCardinalityJobs(minimumCardinality int, jobs []*jobdb.Job) []*jobdb.Job {
 	gangId := uuid.NewString()
 	gangCardinality := fmt.Sprintf("%d", len(jobs))
 	gangMinCardinality := fmt.Sprintf("%d", minimumCardinality)
 	return WithAnnotationsJobs(
-		map[string]string{configuration.GangIdAnnotation: gangId, configuration.GangCardinalityAnnotation: gangCardinality, configuration.GangMinimumCardinalityAnnotation: gangMinCardinality},
+		map[string]string{
+			configuration.GangIdAnnotation:                 gangId,
+			configuration.GangCardinalityAnnotation:        gangCardinality,
+			configuration.GangMinimumCardinalityAnnotation: gangMinCardinality,
+		},
 		jobs,
 	)
 }
@@ -399,7 +406,7 @@ func extractPriority(priorityClassName string) int32 {
 func TestJob(queue string, jobId ulid.ULID, priorityClassName string, req *schedulerobjects.PodRequirements) *jobdb.Job {
 	created := jobTimestamp.Add(1)
 	submitTime := time.Time{}.Add(time.Millisecond * time.Duration(created))
-	return jobdb.NewJob(
+	return JobDb.NewJob(
 		jobId.String(),
 		TestJobset,
 		queue,