Make the node database operate on jobs instead of pods (#2612)

internal/armada/server/lease.go

@@ -370,11 +370,9 @@ 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.BindPodToNode(
-				scheduler.PodRequirementFromLegacySchedulerJob(
-					job,
-					q.schedulingConfig.Preemption.PriorityClasses,
-				),
+			node, err = nodedb.BindJobToNode(
+				q.schedulingConfig.Preemption.PriorityClasses,
+				job,
 				node,
 			)
 			if err != nil {

internal/scheduler/common.go

@@ -3,7 +3,6 @@ package scheduler
 import (
 	"fmt"
 	"strconv"
-	"time"
 
 	"github.com/pkg/errors"
 	"golang.org/x/exp/maps"
@@ -12,7 +11,6 @@ import (
 	armadamaps "github.com/armadaproject/armada/internal/common/maps"
 	armadaslices "github.com/armadaproject/armada/internal/common/slices"
 	schedulerconfig "github.com/armadaproject/armada/internal/scheduler/configuration"
-	schedulercontext "github.com/armadaproject/armada/internal/scheduler/context"
 	"github.com/armadaproject/armada/internal/scheduler/interfaces"
 	"github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
 )
@@ -111,27 +109,6 @@ func JobsSummary(jobs []interfaces.LegacySchedulerJob) string {
 	)
 }
 
-func jobSchedulingContextsFromJobs[T interfaces.LegacySchedulerJob](jobs []T, executorId string, priorityClasses map[string]configuration.PriorityClass) []*schedulercontext.JobSchedulingContext {
-	if jobs == nil {
-		return nil
-	}
-	if len(jobs) == 0 {
-		return make([]*schedulercontext.JobSchedulingContext, 0)
-	}
-	jctxs := make([]*schedulercontext.JobSchedulingContext, len(jobs))
-	timestamp := time.Now()
-	for i, job := range jobs {
-		jctxs[i] = &schedulercontext.JobSchedulingContext{
-			Created:    timestamp,
-			ExecutorId: executorId,
-			JobId:      job.GetId(),
-			Job:        job,
-			Req:        PodRequirementFromLegacySchedulerJob(job, priorityClasses),
-		}
-	}
-	return jctxs
-}
-
 func isEvictedJob(job interfaces.LegacySchedulerJob) bool {
 	return job.GetAnnotations()[schedulerconfig.IsEvictedAnnotation] == "true"
 }
@@ -168,40 +145,3 @@ func GangIdAndCardinalityFromAnnotations(annotations map[string]string) (string,
 	}
 	return gangId, gangCardinality, true, nil
 }
-
-func PodRequirementsFromLegacySchedulerJobs[S ~[]E, E interfaces.LegacySchedulerJob](jobs S, priorityClasses map[string]configuration.PriorityClass) []*schedulerobjects.PodRequirements {
-	rv := make([]*schedulerobjects.PodRequirements, len(jobs))
-	for i, job := range jobs {
-		rv[i] = PodRequirementFromLegacySchedulerJob(job, priorityClasses)
-	}
-	return rv
-}
-
-func PodRequirementFromLegacySchedulerJob[E interfaces.LegacySchedulerJob](job E, priorityClasses map[string]configuration.PriorityClass) *schedulerobjects.PodRequirements {
-	annotations := make(map[string]string, len(configuration.ArmadaManagedAnnotations)+len(schedulerconfig.ArmadaSchedulerManagedAnnotations))
-	for _, key := range configuration.ArmadaManagedAnnotations {
-		if value, ok := job.GetAnnotations()[key]; ok {
-			annotations[key] = value
-		}
-	}
-	for _, key := range schedulerconfig.ArmadaSchedulerManagedAnnotations {
-		if value, ok := job.GetAnnotations()[key]; ok {
-			annotations[key] = value
-		}
-	}
-	annotations[schedulerconfig.JobIdAnnotation] = job.GetId()
-	annotations[schedulerconfig.QueueAnnotation] = job.GetQueue()
-	info := job.GetJobSchedulingInfo(priorityClasses)
-	req := PodRequirementFromJobSchedulingInfo(info)
-	req.Annotations = annotations
-	return req
-}
-
-func PodRequirementFromJobSchedulingInfo(info *schedulerobjects.JobSchedulingInfo) *schedulerobjects.PodRequirements {
-	for _, oreq := range info.ObjectRequirements {
-		if preq := oreq.GetPodRequirements(); preq != nil {
-			return preq
-		}
-	}
-	return nil
-}

internal/scheduler/common_test.go

@@ -10,12 +10,11 @@ import (
 
 	"github.com/armadaproject/armada/internal/armada/configuration"
 	"github.com/armadaproject/armada/internal/common/util"
-	schedulerconfig "github.com/armadaproject/armada/internal/scheduler/configuration"
 	"github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
 	"github.com/armadaproject/armada/pkg/api"
 )
 
-func TestPodRequirementFromLegacySchedulerJob(t *testing.T) {
+func TestGetPodRequirements(t *testing.T) {
 	resourceLimit := v1.ResourceList{
 		"cpu":               resource.MustParse("1"),
 		"memory":            resource.MustParse("128Mi"),
@@ -64,13 +63,12 @@ func TestPodRequirementFromLegacySchedulerJob(t *testing.T) {
 		PreemptionPolicy:     string(v1.PreemptLowerPriority),
 		ResourceRequirements: requirements,
 		Annotations: map[string]string{
+			"something":                             "test",
 			configuration.GangIdAnnotation:          "gang-id",
 			configuration.GangCardinalityAnnotation: "1",
-			schedulerconfig.JobIdAnnotation:         j.Id,
-			schedulerconfig.QueueAnnotation:         j.Queue,
 		},
 	}
-	actual := PodRequirementFromLegacySchedulerJob(j, map[string]configuration.PriorityClass{"armada-default": {Priority: int32(1)}})
+	actual := j.GetPodRequirements(map[string]configuration.PriorityClass{"armada-default": {Priority: int32(1)}})
 	assert.Equal(t, expected, actual)
 }
 

internal/scheduler/configuration/configuration.go

@@ -10,23 +10,14 @@ import (
 )
 
 const (
-	// IsEvictedAnnotation, indicates a pod was evicted in this round and is currently running.
-	// Used by the scheduler to differentiate between pods from running and queued jobs.
+	// IsEvictedAnnotation is set on evicted jobs; the scheduler uses it to differentiate between
+	// already-running and queued jobs.
 	IsEvictedAnnotation = "armadaproject.io/isEvicted"
-	// JobIdAnnotation if set on a pod, indicates which job this pod is part of.
-	JobIdAnnotation = "armadaproject.io/jobId"
-	// QueueAnnotation if set on a pod, indicates which queue this pod is part of.
-	QueueAnnotation = "armadaproject.io/queue"
-	// IdNodeLabel is automatically added to nodes in the NodeDb.
+	// 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 = "armadaproject.io/nodeId"
 )
 
-var ArmadaSchedulerManagedAnnotations = []string{
-	IsEvictedAnnotation,
-	JobIdAnnotation,
-	QueueAnnotation,
-}
-
 type Configuration struct {
 	// Database configuration
 	Postgres configuration.PostgresConfig

internal/scheduler/context/context.go

@@ -228,12 +228,12 @@ func (sctx *SchedulingContext) AddJobSchedulingContext(jctx *JobSchedulingContex
 	}
 	if jctx.IsSuccessful() {
 		if evictedInThisRound {
-			sctx.EvictedResources.SubV1ResourceList(jctx.Req.ResourceRequirements.Requests)
-			sctx.EvictedResourcesByPriorityClass.SubV1ResourceList(jctx.Job.GetPriorityClassName(), jctx.Req.ResourceRequirements.Requests)
+			sctx.EvictedResources.SubV1ResourceList(jctx.PodRequirements.ResourceRequirements.Requests)
+			sctx.EvictedResourcesByPriorityClass.SubV1ResourceList(jctx.Job.GetPriorityClassName(), jctx.PodRequirements.ResourceRequirements.Requests)
 			sctx.NumEvictedJobs--
 		} else {
-			sctx.ScheduledResources.AddV1ResourceList(jctx.Req.ResourceRequirements.Requests)
-			sctx.ScheduledResourcesByPriorityClass.AddV1ResourceList(jctx.Job.GetPriorityClassName(), jctx.Req.ResourceRequirements.Requests)
+			sctx.ScheduledResources.AddV1ResourceList(jctx.PodRequirements.ResourceRequirements.Requests)
+			sctx.ScheduledResourcesByPriorityClass.AddV1ResourceList(jctx.Job.GetPriorityClassName(), jctx.PodRequirements.ResourceRequirements.Requests)
 			sctx.NumScheduledJobs++
 		}
 	}
@@ -440,23 +440,23 @@ func (qctx *QueueSchedulingContext) AddJobSchedulingContext(jctx *JobSchedulingC
 	}
 	_, evictedInThisRound := qctx.EvictedJobsById[jctx.JobId]
 	if jctx.IsSuccessful() {
-		if jctx.Req == nil {
+		if jctx.PodRequirements == nil {
 			return false, errors.Errorf("failed adding job %s to queue: job requirements are missing", jctx.JobId)
 		}
 
 		// Always update ResourcesByPriority.
 		// Since ResourcesByPriority is used to order queues by fraction of fair share.
-		qctx.Allocated.AddV1ResourceList(jctx.Req.ResourceRequirements.Requests)
-		qctx.AllocatedByPriorityClass.AddV1ResourceList(jctx.Job.GetPriorityClassName(), jctx.Req.ResourceRequirements.Requests)
+		qctx.Allocated.AddV1ResourceList(jctx.PodRequirements.ResourceRequirements.Requests)
+		qctx.AllocatedByPriorityClass.AddV1ResourceList(jctx.Job.GetPriorityClassName(), jctx.PodRequirements.ResourceRequirements.Requests)
 
 		// Only if the job is not evicted, update ScheduledResourcesByPriority.
 		// Since ScheduledResourcesByPriority is used to control per-round scheduling constraints.
 		if evictedInThisRound {
 			delete(qctx.EvictedJobsById, jctx.JobId)
-			qctx.EvictedResourcesByPriorityClass.SubV1ResourceList(jctx.Job.GetPriorityClassName(), jctx.Req.ResourceRequirements.Requests)
+			qctx.EvictedResourcesByPriorityClass.SubV1ResourceList(jctx.Job.GetPriorityClassName(), jctx.PodRequirements.ResourceRequirements.Requests)
 		} else {
 			qctx.SuccessfulJobSchedulingContexts[jctx.JobId] = jctx
-			qctx.ScheduledResourcesByPriorityClass.AddV1ResourceList(jctx.Job.GetPriorityClassName(), jctx.Req.ResourceRequirements.Requests)
+			qctx.ScheduledResourcesByPriorityClass.AddV1ResourceList(jctx.Job.GetPriorityClassName(), jctx.PodRequirements.ResourceRequirements.Requests)
 		}
 	} else {
 		qctx.UnsuccessfulJobSchedulingContexts[jctx.JobId] = jctx
@@ -531,7 +531,7 @@ func NewGangSchedulingContext(jctxs []*JobSchedulingContext) *GangSchedulingCont
 	totalResourceRequests := schedulerobjects.NewResourceList(4)
 	for _, jctx := range jctxs {
 		allJobsEvicted = allJobsEvicted && isEvictedJob(jctx.Job)
-		totalResourceRequests.AddV1ResourceList(jctx.Req.ResourceRequirements.Requests)
+		totalResourceRequests.AddV1ResourceList(jctx.PodRequirements.ResourceRequirements.Requests)
 	}
 	return &GangSchedulingContext{
 		Created:               time.Now(),
@@ -543,14 +543,6 @@ func NewGangSchedulingContext(jctxs []*JobSchedulingContext) *GangSchedulingCont
 	}
 }
 
-func (gctx GangSchedulingContext) PodRequirements() []*schedulerobjects.PodRequirements {
-	rv := make([]*schedulerobjects.PodRequirements, len(gctx.JobSchedulingContexts))
-	for i, jctx := range gctx.JobSchedulingContexts {
-		rv[i] = jctx.Req
-	}
-	return rv
-}
-
 func isEvictedJob(job interfaces.LegacySchedulerJob) bool {
 	return job.GetAnnotations()[schedulerconfig.IsEvictedAnnotation] == "true"
 }
@@ -560,17 +552,13 @@ func isEvictedJob(job interfaces.LegacySchedulerJob) bool {
 type JobSchedulingContext struct {
 	// Time at which this context was created.
 	Created time.Time
-	// Executor this job was attempted to be assigned to.
-	ExecutorId string
-	// Total number of nodes in the cluster when trying to schedule.
-	NumNodes int
 	// Id of the job this pod corresponds to.
 	JobId string
 	// Job spec.
 	Job interfaces.LegacySchedulerJob
 	// Scheduling requirements of this job.
 	// We currently require that each job contains exactly one pod spec.
-	Req *schedulerobjects.PodRequirements
+	PodRequirements *schedulerobjects.PodRequirements
 	// Reason for why the job could not be scheduled.
 	// Empty if the job was scheduled successfully.
 	UnschedulableReason string
@@ -583,7 +571,6 @@ func (jctx *JobSchedulingContext) String() string {
 	w := tabwriter.NewWriter(&sb, 1, 1, 1, ' ', 0)
 	fmt.Fprintf(w, "Time:\t%s\n", jctx.Created)
 	fmt.Fprintf(w, "Job ID:\t%s\n", jctx.JobId)
-	fmt.Fprintf(w, "Number of nodes in cluster:\t%d\n", jctx.NumNodes)
 	if jctx.UnschedulableReason != "" {
 		fmt.Fprintf(w, "UnschedulableReason:\t%s\n", jctx.UnschedulableReason)
 	} else {
@@ -600,6 +587,20 @@ func (jctx *JobSchedulingContext) IsSuccessful() bool {
 	return jctx.UnschedulableReason == ""
 }
 
+func JobSchedulingContextsFromJobs[T interfaces.LegacySchedulerJob](priorityClasses map[string]configuration.PriorityClass, jobs []T) []*JobSchedulingContext {
+	jctxs := make([]*JobSchedulingContext, len(jobs))
+	timestamp := time.Now()
+	for i, job := range jobs {
+		jctxs[i] = &JobSchedulingContext{
+			Created:         timestamp,
+			JobId:           job.GetId(),
+			Job:             job,
+			PodRequirements: job.GetPodRequirements(priorityClasses),
+		}
+	}
+	return jctxs
+}
+
 // PodSchedulingContext is returned by SelectAndBindNodeToPod and
 // contains detailed information on the scheduling decision made for this pod.
 type PodSchedulingContext struct {
@@ -626,6 +627,7 @@ func (pctx *PodSchedulingContext) String() string {
 	} else {
 		fmt.Fprint(w, "Node:\tnone\n")
 	}
+	fmt.Fprintf(w, "Number of nodes in cluster:\t%d\n", pctx.NumNodes)
 	if len(pctx.NumExcludedNodesByReason) == 0 {
 		fmt.Fprint(w, "Excluded nodes:\tnone\n")
 	} else {

internal/scheduler/context/context_test.go

@@ -84,10 +84,8 @@ func testNSmallCpuJobSchedulingContext(queue, priorityClassName string, n int) [
 func testSmallCpuJobSchedulingContext(queue, priorityClassName string) *JobSchedulingContext {
 	job := testfixtures.Test1CpuJob(queue, priorityClassName)
 	return &JobSchedulingContext{
-		ExecutorId: "executor",
-		NumNodes:   1,
-		JobId:      job.GetId(),
-		Job:        job,
-		Req:        job.GetJobSchedulingInfo(nil).ObjectRequirements[0].GetPodRequirements(),
+		JobId:           job.GetId(),
+		Job:             job,
+		PodRequirements: job.GetPodRequirements(testfixtures.TestPriorityClasses),
 	}
 }

internal/scheduler/gang_scheduler.go

@@ -4,8 +4,6 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/pkg/errors"
-
 	"github.com/armadaproject/armada/internal/common/util"
 	schedulerconstraints "github.com/armadaproject/armada/internal/scheduler/constraints"
 	schedulercontext "github.com/armadaproject/armada/internal/scheduler/context"
@@ -94,7 +92,7 @@ func (sch *GangScheduler) Schedule(ctx context.Context, gctx *schedulercontext.G
 		// Check that the job is large enough for this executor.
 		// This check needs to be here, since it relates to a specific job.
 		// Only perform limit checks for new jobs to avoid preempting jobs if, e.g., MinimumJobSize changes.
-		if ok, unschedulableReason = requestIsLargeEnough(gctx.TotalResourceRequests, sch.constraints.MinimumJobSize); !ok {
+		if ok, unschedulableReason = requestsAreLargeEnough(gctx.TotalResourceRequests, sch.constraints.MinimumJobSize); !ok {
 			return
 		}
 		if ok, unschedulableReason, err = sch.constraints.CheckPerQueueAndPriorityClassConstraints(
@@ -112,20 +110,10 @@ func (sch *GangScheduler) Schedule(ctx context.Context, gctx *schedulercontext.G
 }
 
 func (sch *GangScheduler) trySchedule(ctx context.Context, gctx *schedulercontext.GangSchedulingContext) (bool, string, error) {
-	pctxs, ok, err := sch.nodeDb.ScheduleMany(gctx.PodRequirements())
+	ok, err := sch.nodeDb.ScheduleMany(gctx.JobSchedulingContexts)
 	if err != nil {
 		return false, "", err
 	}
-	if len(pctxs) > len(gctx.JobSchedulingContexts) {
-		return false, "", errors.Errorf(
-			"received %d pod scheduling context(s), but gang has cardinality %d",
-			len(pctxs), len(gctx.JobSchedulingContexts),
-		)
-	}
-	for i, pctx := range pctxs {
-		gctx.JobSchedulingContexts[i].PodSchedulingContext = pctx
-		gctx.JobSchedulingContexts[i].NumNodes = pctx.NumNodes
-	}
 	if !ok {
 		unschedulableReason := ""
 		if len(gctx.JobSchedulingContexts) > 1 {
@@ -138,7 +126,7 @@ func (sch *GangScheduler) trySchedule(ctx context.Context, gctx *schedulercontex
 	return true, "", nil
 }
 
-func requestIsLargeEnough(totalResourceRequests, minRequest schedulerobjects.ResourceList) (bool, string) {
+func requestsAreLargeEnough(totalResourceRequests, minRequest schedulerobjects.ResourceList) (bool, string) {
 	if len(minRequest.Resources) == 0 {
 		return true, ""
 	}

internal/scheduler/gang_scheduler_test.go

@@ -256,7 +256,7 @@ func TestGangScheduler(t *testing.T) {
 
 			var actualScheduledIndices []int
 			for i, gang := range tc.Gangs {
-				jctxs := jobSchedulingContextsFromJobs(gang, "", testfixtures.TestPriorityClasses)
+				jctxs := schedulercontext.JobSchedulingContextsFromJobs(testfixtures.TestPriorityClasses, gang)
 				gctx := schedulercontext.NewGangSchedulingContext(jctxs)
 				ok, reason, err := sch.Schedule(context.Background(), gctx)
 				require.NoError(t, err)

internal/scheduler/interfaces/interfaces.go

@@ -17,7 +17,7 @@ type LegacySchedulerJob interface {
 	GetPerQueuePriority() uint32
 	GetSubmitTime() time.Time
 	GetAnnotations() map[string]string
-	GetJobSchedulingInfo(map[string]configuration.PriorityClass) *schedulerobjects.JobSchedulingInfo
+	GetPodRequirements(priorityClasses map[string]configuration.PriorityClass) *schedulerobjects.PodRequirements
 	GetPriorityClassName() string
 	GetNodeSelector() map[string]string
 	GetAffinity() *v1.Affinity