Scheduler: refactor - more schedulerobjects to internaltypes (#3998)

Signed-off-by: Robert Smith <[email protected]>

internal/scheduler/internaltypes/resource_list.go

@@ -120,6 +120,19 @@ func (rl ResourceList) AllZero() bool {
 	return true
 }
 
+func (rl ResourceList) FloorAtZero() ResourceList {
+	if rl.IsEmpty() {
+		return rl
+	}
+	result := make([]int64, len(rl.resources))
+	for i, r := range rl.resources {
+		if r > 0 {
+			result[i] = r
+		}
+	}
+	return ResourceList{factory: rl.factory, resources: result}
+}
+
 func (rl ResourceList) HasNegativeValues() bool {
 	if rl.IsEmpty() {
 		return false

internal/scheduler/internaltypes/resource_list_test.go

@@ -116,6 +116,16 @@ func TestAllZero_HandlesEmptyCorrectly(t *testing.T) {
 	assert.True(t, empty.AllZero())
 }
 
+func TestFloorAtZero(t *testing.T) {
+	factory := testFactory()
+
+	assert.Equal(t, testResourceList(factory, "0", "1Ki"), testResourceList(factory, "-1", "1Ki").FloorAtZero())
+}
+
+func TestFloorAtZero_HandlesEmptyCorrectly(t *testing.T) {
+	assert.Equal(t, ResourceList{}, ResourceList{}.FloorAtZero())
+}
+
 func TestHasNegativeValues(t *testing.T) {
 	factory := testFactory()
 	assert.False(t, testResourceList(factory, "0", "0").HasNegativeValues())

internal/scheduler/jobdb/job.go

@@ -50,6 +50,8 @@ type Job struct {
 	queuedVersion int32
 	// Scheduling requirements of this job.
 	jobSchedulingInfo *schedulerobjects.JobSchedulingInfo
+	// All resource requirements, including floating resources, for this job
+	allResourceRequirements internaltypes.ResourceList
 	// Kubernetes (i.e. non-floating) resource requirements of this job
 	kubernetesResourceRequirements internaltypes.ResourceList
 	// Priority class of this job. Populated automatically on job creation.
@@ -508,6 +510,11 @@ func (job *Job) ResourceRequirements() v1.ResourceRequirements {
 	return v1.ResourceRequirements{}
 }
 
+// All resource requirements, including floating resources, for this job
+func (job *Job) AllResourceRequirements() internaltypes.ResourceList {
+	return job.allResourceRequirements
+}
+
 // Kubernetes (i.e. non-floating) resource requirements of this job
 func (job *Job) KubernetesResourceRequirements() internaltypes.ResourceList {
 	return job.kubernetesResourceRequirements
@@ -779,7 +786,7 @@ func (job *Job) Validated() bool {
 
 // Does this job request any floating resources?
 func (job *Job) RequestsFloatingResources() bool {
-	return !job.jobDb.getResourceRequirements(job.jobSchedulingInfo).OfType(internaltypes.Floating).AllZero()
+	return !job.AllResourceRequirements().OfType(internaltypes.Floating).AllZero()
 }
 
 // WithJobSchedulingInfo returns a copy of the job with the job scheduling info updated.
@@ -790,7 +797,9 @@ func (job *Job) WithJobSchedulingInfo(jobSchedulingInfo *schedulerobjects.JobSch
 
 	// Changing the scheduling info invalidates the scheduling key stored with the job.
 	j.schedulingKey = SchedulingKeyFromJob(j.jobDb.schedulingKeyGenerator, j)
-	j.kubernetesResourceRequirements = job.jobDb.getKubernetesResourceRequirements(jobSchedulingInfo)
+
+	j.allResourceRequirements = j.jobDb.getResourceRequirements(jobSchedulingInfo)
+	j.kubernetesResourceRequirements = j.allResourceRequirements.OfType(internaltypes.Kubernetes)
 
 	return j, nil
 }

internal/scheduler/jobdb/job_test.go

@@ -3,6 +3,9 @@ package jobdb
 import (
 	"testing"
 
+	v1 "k8s.io/api/core/v1"
+	k8sResource "k8s.io/apimachinery/pkg/api/resource"
+
 	"github.com/gogo/protobuf/proto"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
@@ -16,6 +19,12 @@ var jobSchedulingInfo = &schedulerobjects.JobSchedulingInfo{
 		{
 			Requirements: &schedulerobjects.ObjectRequirements_PodRequirements{
 				PodRequirements: &schedulerobjects.PodRequirements{
+					ResourceRequirements: v1.ResourceRequirements{
+						Requests: v1.ResourceList{
+							"cpu":                 k8sResource.MustParse("1"),
+							"storage-connections": k8sResource.MustParse("1"),
+						},
+					},
 					Annotations: map[string]string{
 						"foo": "bar",
 					},
@@ -339,6 +348,12 @@ func TestJob_TestWithJobSchedulingInfo(t *testing.T) {
 			{
 				Requirements: &schedulerobjects.ObjectRequirements_PodRequirements{
 					PodRequirements: &schedulerobjects.PodRequirements{
+						ResourceRequirements: v1.ResourceRequirements{
+							Requests: v1.ResourceList{
+								"cpu":                 k8sResource.MustParse("2"),
+								"storage-connections": k8sResource.MustParse("2"),
+							},
+						},
 						Annotations: map[string]string{
 							"fish": "chips",
 						},
@@ -350,6 +365,37 @@ func TestJob_TestWithJobSchedulingInfo(t *testing.T) {
 	newJob := JobWithJobSchedulingInfo(baseJob, newSchedInfo)
 	assert.Equal(t, jobSchedulingInfo, baseJob.JobSchedulingInfo())
 	assert.Equal(t, newSchedInfo, newJob.JobSchedulingInfo())
+
+	assert.Equal(t, int64(1000), baseJob.AllResourceRequirements().GetByNameZeroIfMissing("cpu"))
+	assert.Equal(t, int64(1), baseJob.AllResourceRequirements().GetByNameZeroIfMissing("storage-connections"))
+	assert.Equal(t, int64(1000), baseJob.KubernetesResourceRequirements().GetByNameZeroIfMissing("cpu"))
+	assert.Equal(t, int64(0), baseJob.KubernetesResourceRequirements().GetByNameZeroIfMissing("storage-connections"))
+
+	assert.Equal(t, int64(2000), newJob.AllResourceRequirements().GetByNameZeroIfMissing("cpu"))
+	assert.Equal(t, int64(2), newJob.AllResourceRequirements().GetByNameZeroIfMissing("storage-connections"))
+	assert.Equal(t, int64(2000), newJob.KubernetesResourceRequirements().GetByNameZeroIfMissing("cpu"))
+	assert.Equal(t, int64(0), newJob.KubernetesResourceRequirements().GetByNameZeroIfMissing("storage-connections"))
+}
+
+func TestRequestsFloatingResources(t *testing.T) {
+	noFloatingResourcesJob := JobWithJobSchedulingInfo(baseJob, &schedulerobjects.JobSchedulingInfo{
+		ObjectRequirements: []*schedulerobjects.ObjectRequirements{
+			{
+				Requirements: &schedulerobjects.ObjectRequirements_PodRequirements{
+					PodRequirements: &schedulerobjects.PodRequirements{
+						ResourceRequirements: v1.ResourceRequirements{
+							Requests: v1.ResourceList{
+								"cpu":                 k8sResource.MustParse("1"),
+								"storage-connections": k8sResource.MustParse("0"),
+							},
+						},
+					},
+				},
+			},
+		},
+	})
+	assert.True(t, baseJob.RequestsFloatingResources())
+	assert.False(t, noFloatingResourcesJob.RequestsFloatingResources())
 }
 
 func TestJobSchedulingInfoFieldsInitialised(t *testing.T) {
@@ -396,3 +442,13 @@ func TestJob_TestResolvedPools(t *testing.T) {
 	// Job with an active run
 	assert.Equal(t, []string{"testPool2"}, jobWithJobRunPool.ResolvedPools())
 }
+
+func TestJob_TestAllResourceRequirements(t *testing.T) {
+	assert.Equal(t, int64(1000), baseJob.AllResourceRequirements().GetByNameZeroIfMissing("cpu"))
+	assert.Equal(t, int64(1), baseJob.AllResourceRequirements().GetByNameZeroIfMissing("storage-connections"))
+}
+
+func TestJob_TestKubernetesResourceRequirements(t *testing.T) {
+	assert.Equal(t, int64(1000), baseJob.KubernetesResourceRequirements().GetByNameZeroIfMissing("cpu"))
+	assert.Equal(t, int64(0), baseJob.KubernetesResourceRequirements().GetByNameZeroIfMissing("storage-connections"))
+}

internal/scheduler/jobdb/jobdb.go

@@ -149,6 +149,8 @@ func (jobDb *JobDb) NewJob(
 		priorityClass = jobDb.defaultPriorityClass
 	}
 
+	rr := jobDb.getResourceRequirements(schedulingInfo)
+
 	job := &Job{
 		jobDb:                          jobDb,
 		id:                             jobId,
@@ -160,7 +162,8 @@ func (jobDb *JobDb) NewJob(
 		requestedPriority:              priority,
 		submittedTime:                  created,
 		jobSchedulingInfo:              jobDb.internJobSchedulingInfoStrings(schedulingInfo),
-		kubernetesResourceRequirements: jobDb.getKubernetesResourceRequirements(schedulingInfo),
+		allResourceRequirements:        rr,
+		kubernetesResourceRequirements: rr.OfType(internaltypes.Kubernetes),
 		priorityClass:                  priorityClass,
 		cancelRequested:                cancelRequested,
 		cancelByJobSetRequested:        cancelByJobSetRequested,
@@ -174,10 +177,6 @@ func (jobDb *JobDb) NewJob(
 	return job, nil
 }
 
-func (jobDb *JobDb) getKubernetesResourceRequirements(schedulingInfo *schedulerobjects.JobSchedulingInfo) internaltypes.ResourceList {
-	return jobDb.getResourceRequirements(schedulingInfo).OfType(internaltypes.Kubernetes)
-}
-
 func (jobDb *JobDb) getResourceRequirements(schedulingInfo *schedulerobjects.JobSchedulingInfo) internaltypes.ResourceList {
 	return jobDb.resourceListFactory.FromJobResourceListIgnoreUnknown(safeGetRequirements(schedulingInfo))
 }

internal/scheduler/jobdb/test_utils.go

@@ -13,7 +13,7 @@ var testResourceListFactory = makeTestResourceListFactory()
 func makeTestResourceListFactory() *internaltypes.ResourceListFactory {
 	result, _ := internaltypes.NewResourceListFactory(
 		getTestSupportedResourceTypes(),
-		nil,
+		getTestFloatingResourceTypes(),
 	)
 	return result
 }
@@ -27,6 +27,12 @@ func getTestSupportedResourceTypes() []schedulerconfiguration.ResourceType {
 	}
 }
 
+func getTestFloatingResourceTypes() []schedulerconfiguration.FloatingResourceConfig {
+	return []schedulerconfiguration.FloatingResourceConfig{
+		{Name: "storage-connections", Resolution: resource.MustParse("1")},
+	}
+}
+
 func WithJobDbJobPodRequirements(job *Job, reqs *schedulerobjects.PodRequirements) *Job {
 	return JobWithJobSchedulingInfo(job, &schedulerobjects.JobSchedulingInfo{
 		PriorityClassName: job.JobSchedulingInfo().PriorityClassName,

internal/scheduler/schedulerobjects/node.go

@@ -2,7 +2,6 @@ package schedulerobjects
 
 import (
 	"fmt"
-	"math"
 
 	"golang.org/x/exp/maps"
 	"golang.org/x/exp/slices"
@@ -45,15 +44,3 @@ func (node *Node) AvailableArmadaResource() ResourceList {
 	}
 	return tr
 }
-
-func (node *Node) MarkResourceUnallocatable(unallocatable ResourceList) {
-	currentAllocatable := node.UnallocatableResources[math.MaxInt32]
-	(&currentAllocatable).Add(unallocatable)
-	node.UnallocatableResources[math.MaxInt32] = currentAllocatable
-
-	for priority, allocatable := range node.AllocatableByPriorityAndResource {
-		allocatable.Sub(unallocatable)
-		allocatable.LimitToZero()
-		node.AllocatableByPriorityAndResource[priority] = allocatable
-	}
-}