Avoid copying nodes where possible

internal/armada/server/lease.go

@@ -378,24 +378,8 @@ func (q *AggregatedQueueServer) getJobs(ctx context.Context, req *api.StreamingL
 		}
 
 		// Bind pods to nodes, thus ensuring resources are marked as allocated on the node.
-		skipNode := false
-		for _, job := range jobs {
-			node, err = nodedb.BindJobToNode(
-				q.schedulingConfig.Preemption.PriorityClasses,
-				job,
-				node,
-			)
-			if err != nil {
-				logging.WithStacktrace(log, err).Warnf(
-					"skipping node %s from executor %s: failed to bind job %s to node",
-					nodeInfo.GetName(), req.GetClusterId(), job.Id,
-				)
-				skipNode = true
-				break
-			}
-		}
-		if skipNode {
-			continue
+		if node, err = nodedb.BindJobsToNode(q.schedulingConfig.Preemption.PriorityClasses, jobs, node); err != nil {
+			return nil, err
 		}
 
 		// Record which node each job is scheduled on. Necessary for gang preemption.

internal/scheduler/context/context.go

@@ -638,7 +638,7 @@ func (jctx *JobSchedulingContext) IsSuccessful() bool {
 	return jctx.UnschedulableReason == ""
 }
 
-func JobSchedulingContextsFromJobs[T interfaces.LegacySchedulerJob](priorityClasses map[string]configuration.PriorityClass, jobs []T) []*JobSchedulingContext {
+func JobSchedulingContextsFromJobs[J interfaces.LegacySchedulerJob](priorityClasses map[string]configuration.PriorityClass, jobs []J) []*JobSchedulingContext {
 	jctxs := make([]*JobSchedulingContext, len(jobs))
 	timestamp := time.Now()
 	for i, job := range jobs {

internal/scheduler/nodedb/nodedb.go

@@ -486,13 +486,31 @@ func (nodeDb *NodeDb) selectNodeForPodWithIt(
 	return selectedNode, nil
 }
 
+// BindJobsToNode returns a copy of node with all elements of jobs bound to it.
+func BindJobsToNode[J interfaces.LegacySchedulerJob](priorityClasses map[string]configuration.PriorityClass, jobs []J, node *schedulerobjects.Node) (*schedulerobjects.Node, error) {
+	node = node.DeepCopy()
+	for _, job := range jobs {
+		if err := bindJobToNodeInPlace(priorityClasses, job, node); err != nil {
+			return nil, err
+		}
+	}
+	return node, nil
+}
+
 // BindJobToNode returns a copy of node with job bound to it.
 func BindJobToNode(priorityClasses map[string]configuration.PriorityClass, job interfaces.LegacySchedulerJob, node *schedulerobjects.Node) (*schedulerobjects.Node, error) {
+	node = node.DeepCopy()
+	if err := bindJobToNodeInPlace(priorityClasses, job, node); err != nil {
+		return nil, err
+	}
+	return node, nil
+}
+
+// bindJobToNodeInPlace is like BindJobToNode, but doesn't make a copy of node.
+func bindJobToNodeInPlace(priorityClasses map[string]configuration.PriorityClass, job interfaces.LegacySchedulerJob, node *schedulerobjects.Node) error {
 	jobId := job.GetId()
 	requests := job.GetResourceRequirements().Requests
 
-	node = node.DeepCopy()
-
 	_, isEvicted := node.EvictedJobRunIds[jobId]
 	delete(node.EvictedJobRunIds, jobId)
 
@@ -501,7 +519,7 @@ func BindJobToNode(priorityClasses map[string]configuration.PriorityClass, job i
 			node.AllocatedByJobId = make(map[string]schedulerobjects.ResourceList)
 		}
 		if allocatedToJob, ok := node.AllocatedByJobId[jobId]; ok {
-			return nil, errors.Errorf("job %s already has resources allocated on node %s", jobId, node.Id)
+			return errors.Errorf("job %s already has resources allocated on node %s", jobId, node.Id)
 		} else {
 			allocatedToJob.AddV1ResourceList(requests)
 			node.AllocatedByJobId[jobId] = allocatedToJob
@@ -523,60 +541,89 @@ func BindJobToNode(priorityClasses map[string]configuration.PriorityClass, job i
 		allocatable.MarkAllocatableV1ResourceList(evictedPriority, requests)
 	}
 
-	return node, nil
+	return nil
 }
 
-// EvictJobFromNode returns a copy of node with job evicted from it. Specifically:
+// EvictJobsFromNode returns a copy of node with all elements of jobs for which jobFilter returns
+// true evicted from it, together with a slice containing exactly those jobs.
 //
-// - The job is marked as evicted on the node.
-// - AllocatedByJobId and AllocatedByQueue are not updated.
-// - Resources requested by the evicted pod are marked as allocated at priority evictedPriority.
-func EvictJobFromNode(priorityClasses map[string]configuration.PriorityClass, job interfaces.LegacySchedulerJob, node *schedulerobjects.Node) (*schedulerobjects.Node, error) {
-	jobId := job.GetId()
-	queue := job.GetQueue()
-	requests := job.GetResourceRequirements().Requests
-
+// Specifically:
+//
+//   - The jobs that jobFilter returns true for are marked as evicted on the node.
+//   - Within AllocatableByPriorityAndResource, the resources allocated to these jobs are moved from
+//     the jobs' priorities to evictedPriority; they are not subtracted from AllocatedByJobId and
+//     AllocatedByQueue.
+func EvictJobsFromNode(
+	priorityClasses map[string]configuration.PriorityClass,
+	jobFilter func(interfaces.LegacySchedulerJob) bool,
+	jobs []interfaces.LegacySchedulerJob,
+	node *schedulerobjects.Node,
+) ([]interfaces.LegacySchedulerJob, *schedulerobjects.Node, error) {
+	evicted := make([]interfaces.LegacySchedulerJob, 0)
 	node = node.DeepCopy()
-
-	// Ensure we track allocated resources at evictedPriority.
-	if _, ok := node.AllocatableByPriorityAndResource[evictedPriority]; !ok {
-		pMin := int32(math.MaxInt32)
-		ok := false
-		for p := range node.AllocatableByPriorityAndResource {
-			if p < pMin {
-				pMin = p
-				ok = true
-			}
+	if jobFilter == nil {
+		return evicted, node, nil
+	}
+	for _, job := range jobs {
+		if !jobFilter(job) {
+			continue
 		}
-		if ok {
-			node.AllocatableByPriorityAndResource[evictedPriority] = node.AllocatableByPriorityAndResource[pMin].DeepCopy()
+		evicted = append(evicted, job)
+		if err := evictJobFromNodeInPlace(priorityClasses, job, node); err != nil {
+			return nil, nil, err
 		}
 	}
+	return evicted, node, nil
+}
 
+// evictJobFromNodeInPlace is the in-place operation backing EvictJobsFromNode.
+func evictJobFromNodeInPlace(priorityClasses map[string]configuration.PriorityClass, job interfaces.LegacySchedulerJob, node *schedulerobjects.Node) error {
+	jobId := job.GetId()
 	if _, ok := node.AllocatedByJobId[jobId]; !ok {
-		return nil, errors.Errorf("job %s has no resources allocated on node %s", jobId, node.Id)
+		return errors.Errorf("job %s has no resources allocated on node %s", jobId, node.Id)
 	}
+
+	queue := job.GetQueue()
 	if _, ok := node.AllocatedByQueue[queue]; !ok {
-		return nil, errors.Errorf("queue %s has no resources allocated on node %s", queue, node.Id)
+		return errors.Errorf("queue %s has no resources allocated on node %s", queue, node.Id)
 	}
+
 	if node.EvictedJobRunIds == nil {
 		node.EvictedJobRunIds = make(map[string]bool)
 	}
 	if _, ok := node.EvictedJobRunIds[jobId]; ok {
-		// TODO: We're using run ids instead of job ids for now.
-		return nil, errors.Errorf("job %s is already evicted from node %s", jobId, node.Id)
-	} else {
-		node.EvictedJobRunIds[jobId] = true
+		return errors.Errorf("job %s is already evicted from node %s", jobId, node.Id)
 	}
+	node.EvictedJobRunIds[jobId] = true
 
+	if _, ok := node.AllocatableByPriorityAndResource[evictedPriority]; !ok {
+		minimumPriority := int32(math.MaxInt32)
+		foundPriority := false
+		for p := range node.AllocatableByPriorityAndResource {
+			if p < minimumPriority {
+				minimumPriority = p
+				foundPriority = true
+			}
+		}
+		if foundPriority {
+			node.AllocatableByPriorityAndResource[evictedPriority] = node.AllocatableByPriorityAndResource[minimumPriority].DeepCopy()
+		} else {
+			// We do not expect to hit this branch; however, if we do, then we need
+			// to make sure that evictedPriority is in this map so that the call to
+			// MarkAllocatedV1ResourceList() below knows about it.
+			node.AllocatableByPriorityAndResource[evictedPriority] = schedulerobjects.ResourceList{}
+		}
+	}
 	allocatable := schedulerobjects.AllocatableByPriorityAndResourceType(node.AllocatableByPriorityAndResource)
 	priority := priorityClasses[job.GetPriorityClassName()].Priority
+	requests := job.GetResourceRequirements().Requests
 	allocatable.MarkAllocatableV1ResourceList(priority, requests)
 	allocatable.MarkAllocatedV1ResourceList(evictedPriority, requests)
-	return node, nil
+
+	return nil
 }
 
-// UnbindJobsFromNode returns a node with all reqs unbound from it.
+// UnbindJobsFromNode returns a node with all elements of jobs unbound from it.
 func UnbindJobsFromNode(priorityClasses map[string]configuration.PriorityClass, jobs []interfaces.LegacySchedulerJob, node *schedulerobjects.Node) (*schedulerobjects.Node, error) {
 	node = node.DeepCopy()
 	for _, job := range jobs {
@@ -587,7 +634,7 @@ func UnbindJobsFromNode(priorityClasses map[string]configuration.PriorityClass,
 	return node, nil
 }
 
-// UnbindJobFromNode returns a copy of node with req unbound from it.
+// UnbindJobFromNode returns a copy of node with job unbound from it.
 func UnbindJobFromNode(priorityClasses map[string]configuration.PriorityClass, job interfaces.LegacySchedulerJob, node *schedulerobjects.Node) (*schedulerobjects.Node, error) {
 	node = node.DeepCopy()
 	if err := unbindJobFromNodeInPlace(priorityClasses, job, node); err != nil {
@@ -596,7 +643,7 @@ func UnbindJobFromNode(priorityClasses map[string]configuration.PriorityClass, j
 	return node, nil
 }
 
-// unbindPodFromNodeInPlace is like UnbindJobFromNode, but doesn't make a copy of the node.
+// unbindPodFromNodeInPlace is like UnbindJobFromNode, but doesn't make a copy of node.
 func unbindJobFromNodeInPlace(priorityClasses map[string]configuration.PriorityClass, job interfaces.LegacySchedulerJob, node *schedulerobjects.Node) error {
 	jobId := job.GetId()
 	requests := job.GetResourceRequirements().Requests

internal/scheduler/nodedb/nodedb_test.go

@@ -116,6 +116,7 @@ func TestSelectNodeForPod_NodeIdLabel_Failure(t *testing.T) {
 }
 
 func TestNodeBindingEvictionUnbinding(t *testing.T) {
+	jobFilter := func(job interfaces.LegacySchedulerJob) bool { return true }
 	node := testfixtures.Test8GpuNode(append(testfixtures.TestPriorities, evictedPriority))
 	job := testfixtures.Test1GpuJob("A", testfixtures.PriorityClass0)
 	request := schedulerobjects.ResourceListFromV1ResourceList(job.GetResourceRequirements().Requests)
@@ -130,16 +131,17 @@ func TestNodeBindingEvictionUnbinding(t *testing.T) {
 	unboundMultipleNode, err := UnbindJobsFromNode(testfixtures.TestPriorityClasses, []interfaces.LegacySchedulerJob{job}, boundNode)
 	require.NoError(t, err)
 
-	evictedNode, err := EvictJobFromNode(testfixtures.TestPriorityClasses, job, boundNode)
+	evictedJobs, evictedNode, err := EvictJobsFromNode(testfixtures.TestPriorityClasses, jobFilter, []interfaces.LegacySchedulerJob{job}, boundNode)
 	require.NoError(t, err)
+	assert.Equal(t, []interfaces.LegacySchedulerJob{job}, evictedJobs)
 
 	evictedUnboundNode, err := UnbindJobFromNode(testfixtures.TestPriorityClasses, job, evictedNode)
 	require.NoError(t, err)
 
 	evictedBoundNode, err := BindJobToNode(testfixtures.TestPriorityClasses, job, evictedNode)
 	require.NoError(t, err)
 
-	_, err = EvictJobFromNode(testfixtures.TestPriorityClasses, job, node)
+	_, _, err = EvictJobsFromNode(testfixtures.TestPriorityClasses, jobFilter, []interfaces.LegacySchedulerJob{job}, node)
 	require.Error(t, err)
 
 	_, err = UnbindJobFromNode(testfixtures.TestPriorityClasses, job, node)
@@ -148,7 +150,7 @@ func TestNodeBindingEvictionUnbinding(t *testing.T) {
 	_, err = BindJobToNode(testfixtures.TestPriorityClasses, job, boundNode)
 	require.Error(t, err)
 
-	_, err = EvictJobFromNode(testfixtures.TestPriorityClasses, job, evictedNode)
+	_, _, err = EvictJobsFromNode(testfixtures.TestPriorityClasses, jobFilter, []interfaces.LegacySchedulerJob{job}, evictedNode)
 	require.Error(t, err)
 
 	assertNodeAccountingEqual(t, node, unboundNode)
@@ -255,6 +257,26 @@ func assertNodeAccountingEqual(t *testing.T, node1, node2 *schedulerobjects.Node
 	return rv
 }
 
+func TestEviction(t *testing.T) {
+	jobFilter := func(job interfaces.LegacySchedulerJob) bool {
+		priorityClassName := job.GetPriorityClassName()
+		priorityClass := testfixtures.TestPriorityClasses[priorityClassName]
+		return priorityClass.Preemptible
+	}
+	node := testfixtures.Test32CpuNode(testfixtures.TestPriorities)
+	job0 := testfixtures.Test1CpuJob("queue-alice", testfixtures.PriorityClass0)
+	job3 := testfixtures.Test1CpuJob("queue-alice", testfixtures.PriorityClass3)
+	var err error
+	node, err = BindJobToNode(testfixtures.TestPriorityClasses, job0, node)
+	require.NoError(t, err)
+	node, err = BindJobToNode(testfixtures.TestPriorityClasses, job3, node)
+	require.NoError(t, err)
+
+	evictedJobs, _, err := EvictJobsFromNode(testfixtures.TestPriorityClasses, jobFilter, []interfaces.LegacySchedulerJob{job0, job3}, node)
+	require.NoError(t, err)
+	assert.Equal(t, []interfaces.LegacySchedulerJob{job0}, evictedJobs)
+}
+
 func TestScheduleIndividually(t *testing.T) {
 	tests := map[string]struct {
 		Nodes         []*schedulerobjects.Node

internal/scheduler/preempting_queue_scheduler.go

@@ -16,7 +16,6 @@ import (
 	"github.com/armadaproject/armada/internal/armada/configuration"
 	armadamaps "github.com/armadaproject/armada/internal/common/maps"
 	armadaslices "github.com/armadaproject/armada/internal/common/slices"
-	"github.com/armadaproject/armada/internal/common/util"
 	schedulerconfig "github.com/armadaproject/armada/internal/scheduler/configuration"
 	schedulerconstraints "github.com/armadaproject/armada/internal/scheduler/constraints"
 	schedulercontext "github.com/armadaproject/armada/internal/scheduler/context"
@@ -791,43 +790,45 @@ func (evi *Evictor) Evict(ctx context.Context, it nodedb.NodeIterator) (*Evictor
 	evictedJobsById := make(map[string]interfaces.LegacySchedulerJob)
 	affectedNodesById := make(map[string]*schedulerobjects.Node)
 	nodeIdByJobId := make(map[string]string)
+	result := &EvictorResult{
+		EvictedJobsById:   evictedJobsById,
+		AffectedNodesById: affectedNodesById,
+		NodeIdByJobId:     nodeIdByJobId,
+	}
+	if evi.jobFilter == nil || evi.nodeFilter == nil {
+		return result, nil
+	}
+	jobFilter := func(job interfaces.LegacySchedulerJob) bool { return evi.jobFilter(ctx, job) }
 	for node := it.NextNode(); node != nil; node = it.NextNode() {
-		if evi.nodeFilter != nil && !evi.nodeFilter(ctx, node) {
+		if !evi.nodeFilter(ctx, node) {
 			continue
 		}
-		jobIds := util.Filter(
-			maps.Keys(node.AllocatedByJobId),
-			func(jobId string) bool {
-				_, ok := node.EvictedJobRunIds[jobId]
-				return !ok
-			},
-		)
+		jobIds := make([]string, 0)
+		for jobId := range node.AllocatedByJobId {
+			if _, ok := node.EvictedJobRunIds[jobId]; !ok {
+				jobIds = append(jobIds, jobId)
+			}
+		}
 		jobs, err := evi.jobRepo.GetExistingJobsByIds(jobIds)
 		if err != nil {
 			return nil, err
 		}
-		for _, job := range jobs {
-			if evi.jobFilter != nil && !evi.jobFilter(ctx, job) {
-				continue
-			}
-			node, err = nodedb.EvictJobFromNode(evi.priorityClasses, job, node)
-			if err != nil {
-				return nil, err
-			}
+		evictedJobs, node, err := nodedb.EvictJobsFromNode(evi.priorityClasses, jobFilter, jobs, node)
+		if err != nil {
+			return nil, err
+		}
+		for _, job := range evictedJobs {
+			evictedJobsById[job.GetId()] = job
+			nodeIdByJobId[job.GetId()] = node.Id
 			if evi.postEvictFunc != nil {
 				evi.postEvictFunc(ctx, job, node)
 			}
-
-			evictedJobsById[job.GetId()] = job
-			nodeIdByJobId[job.GetId()] = node.Id
 		}
-		affectedNodesById[node.Id] = node
+		if len(evictedJobs) > 0 {
+			affectedNodesById[node.Id] = node
+		}
 	}
-	return &EvictorResult{
-		EvictedJobsById:   evictedJobsById,
-		AffectedNodesById: affectedNodesById,
-		NodeIdByJobId:     nodeIdByJobId,
-	}, nil
+	return result, nil
 }
 
 // TODO: This is only necessary for jobs not scheduled in this cycle.

internal/scheduler/preempting_queue_scheduler_test.go

@@ -1277,7 +1277,7 @@ func TestPreemptingQueueScheduler(t *testing.T) {
 	}
 	for name, tc := range tests {
 		t.Run(name, func(t *testing.T) {
-			nodeDb, err := CreateNodeDb(tc.Nodes)
+			nodeDb, err := NewNodeDb(tc.Nodes)
 			require.NoError(t, err)
 
 			// Repo. for storing jobs to be queued.
@@ -1603,7 +1603,7 @@ func BenchmarkPreemptingQueueScheduler(b *testing.B) {
 				priorityFactorByQueue[queue] = float64(rand.Intn(tc.MaxPriorityFactor-tc.MinPriorityFactor+1) + tc.MinPriorityFactor)
 			}
 
-			nodeDb, err := CreateNodeDb(tc.Nodes)
+			nodeDb, err := NewNodeDb(tc.Nodes)
 			require.NoError(b, err)
 			jobRepo := NewInMemoryJobRepository(testfixtures.TestPriorityClasses)
 
@@ -1669,7 +1669,7 @@ func BenchmarkPreemptingQueueScheduler(b *testing.B) {
 				require.NoError(b, err)
 				nodesById[nodeId] = node
 			}
-			nodeDb, err = CreateNodeDb(maps.Values(nodesById))
+			nodeDb, err = NewNodeDb(maps.Values(nodesById))
 			require.NoError(b, err)
 
 			allocatedByQueueAndPriorityClass := sctx.AllocatedByQueueAndPriority()