Scheduler: efficient resource implementation (armadaproject#3561)

config/armada/config.yaml

@@ -39,6 +39,15 @@ eventsApiRedis:
 # This config must be consistent with the scheduling config used by the scheduler.
 # You may want to insert the scheduling config used for the scheduler automatically, e.g., using PyYAML, to guarantee consistency.
 scheduling:
+  supportedResourceTypes:
+    - name: memory
+      resolution: "1"
+    - name: cpu
+      resolution: "1m"
+    - name: ephemeral-storage
+      resolution: "1"
+    - name: nvidia.com/gpu
+      resolution: "1"
   executorTimeout: "60m"
   executorUpdateFrequency: "1m"
   priorityClasses:

config/scheduler/config.yaml

@@ -79,6 +79,15 @@ grpc:
     enabled: false
 # You may want to configure indexedNodeLabels and indexedTaints to speed up scheduling.
 scheduling:
+  supportedResourceTypes:
+    - name: memory
+      resolution: "1"
+    - name: cpu
+      resolution: "1m"
+    - name: ephemeral-storage
+      resolution: "1"
+    - name: nvidia.com/gpu
+      resolution: "1"
   disableScheduling: false
   enableAssertions: false
   nodeEvictionProbability: 1.0

go.mod

@@ -92,6 +92,7 @@ require (
 	golang.org/x/time v0.3.0
 	gonum.org/v1/gonum v0.14.0
 	google.golang.org/genproto/googleapis/api v0.0.0-20230525234035-dd9d682886f9
+	gopkg.in/inf.v0 v0.9.1
 )
 
 require (
@@ -204,7 +205,6 @@ require (
 	google.golang.org/appengine v1.6.7 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20230525234030-28d5490b6b19 // indirect
 	google.golang.org/protobuf v1.31.0 // indirect
-	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/ini.v1 v1.67.0 // indirect
 	gopkg.in/square/go-jose.v2 v2.6.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect

internal/armada/configuration/types.go

@@ -224,10 +224,19 @@ type SchedulingConfig struct {
 	// Hence, a larger MaxExtraNodesToConsider would reduce the expected number of preemptions.
 	// TODO(albin): Remove. It's unused.
 	MaxExtraNodesToConsider uint
+	// Resource types (e.g. memory or nvidia.com/gpu) that the scheduler keeps track of.
+	// Resource types not on this list will be ignored if seen on a node, and any jobs requesting them will fail.
+	SupportedResourceTypes []ResourceType
 	// Resources, e.g., "cpu", "memory", and "nvidia.com/gpu", for which the scheduler creates indexes for efficient lookup.
 	// This list must contain at least one resource. Adding more than one resource is not required, but may speed up scheduling.
 	// Ideally, this list contains all resources that frequently constrain which nodes a job can be scheduled onto.
-	IndexedResources []IndexedResource
+	//
+	// In particular, the allocatable resources on each node are rounded to a multiple of the resolution.
+	// Lower resolution speeds up scheduling by improving node lookup speed but may prevent scheduling jobs,
+	// since the allocatable resources may be rounded down to be a multiple of the resolution.
+	//
+	// See NodeDb docs for more details.
+	IndexedResources []ResourceType
 	// Node labels that the scheduler creates indexes for efficient lookup of.
 	// Should include node labels frequently used by node selectors on submitted jobs.
 	//
@@ -295,16 +304,11 @@ func SchedulingConfigValidation(sl validator.StructLevel) {
 	}
 }
 
-// IndexedResource represents a resource the scheduler indexes for efficient lookup.
-type IndexedResource struct {
+// ResourceType represents a resource the scheduler indexes for efficient lookup.
+type ResourceType struct {
 	// Resource name, e.g., "cpu", "memory", or "nvidia.com/gpu".
 	Name string
 	// Resolution with which Armada tracks this resource; larger values indicate lower resolution.
-	// In particular, the allocatable resources on each node are rounded to a multiple of the resolution.
-	// Lower resolution speeds up scheduling by improving node lookup speed but may prevent scheduling jobs,
-	// since the allocatable resources may be rounded down to be a multiple of the resolution.
-	//
-	// See NodeDb docs for more details.
 	Resolution resource.Quantity
 }
 

internal/armada/server.go

@@ -32,6 +32,7 @@ import (
 	"github.com/armadaproject/armada/internal/common/pulsarutils"
 	"github.com/armadaproject/armada/internal/scheduler"
 	schedulerdb "github.com/armadaproject/armada/internal/scheduler/database"
+	"github.com/armadaproject/armada/internal/scheduler/internaltypes"
 	"github.com/armadaproject/armada/internal/scheduler/reports"
 	"github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
 	"github.com/armadaproject/armada/pkg/api"
@@ -144,10 +145,17 @@ func Serve(ctx *armadacontext.Context, config *configuration.ArmadaConfig, healt
 
 	// Executor Repositories for pulsar scheduler
 	pulsarExecutorRepo := schedulerdb.NewRedisExecutorRepository(db, "pulsar")
+
+	resourceListFactory, err := internaltypes.MakeResourceListFactory(config.Scheduling.SupportedResourceTypes)
+	if err != nil {
+		return errors.WithMessage(err, "Error with the .scheduling.supportedResourceTypes field in config")
+	}
+	ctx.Infof("Supported resource types: %s", resourceListFactory.SummaryString())
 	submitChecker := scheduler.NewSubmitChecker(
 		30*time.Minute,
 		config.Scheduling,
 		pulsarExecutorRepo,
+		resourceListFactory,
 	)
 	services = append(services, func() error {
 		return submitChecker.Run(ctx)

internal/armada/submit/submit.go

@@ -454,7 +454,6 @@ func (s *Server) ReprioritizeJobs(grpcCtx context.Context, req *api.JobRepriorit
 	}
 
 	err = s.publisher.PublishMessages(ctx, sequence)
-
 	if err != nil {
 		log.WithError(err).Error("failed send to Pulsar")
 		return nil, status.Error(codes.Internal, "Failed to send message")

internal/executor/job/processors/preempt_runs.go

@@ -96,7 +96,6 @@ func (j *RunPreemptedProcessor) reportPodPreempted(run *job.RunState, pod *v1.Po
 		domain.JobPreemptedAnnotation: time.Now().String(),
 		string(v1.PodFailed):          time.Now().String(),
 	})
-
 	if err != nil {
 		return fmt.Errorf("failed to annotate pod as preempted - %s", err)
 	}

internal/pulsartest/app.go

@@ -37,7 +37,6 @@ func New(params Params, cmdType string) (*App, error) {
 			Name:  producerName,
 			Topic: params.Pulsar.JobsetEventsTopic,
 		})
-
 		if err != nil {
 			return nil, errors.Wrapf(err, "error creating pulsar producer %s", producerName)
 		}
@@ -47,7 +46,6 @@ func New(params Params, cmdType string) (*App, error) {
 			Topic:          params.Pulsar.JobsetEventsTopic,
 			StartMessageID: pulsar.EarliestMessageID(),
 		})
-
 		if err != nil {
 			return nil, errors.Wrapf(err, "error creating pulsar reader")
 		}

internal/scheduler/gang_scheduler_test.go

@@ -212,7 +212,7 @@ func TestGangScheduler(t *testing.T) {
 		},
 		"resolution has no impact on jobs of size a multiple of the resolution": {
 			SchedulingConfig: testfixtures.WithIndexedResourcesConfig(
-				[]configuration.IndexedResource{
+				[]configuration.ResourceType{
 					{Name: "cpu", Resolution: resource.MustParse("16")},
 					{Name: "memory", Resolution: resource.MustParse("128Mi")},
 				},
@@ -233,7 +233,7 @@ func TestGangScheduler(t *testing.T) {
 		},
 		"jobs of size not a multiple of the resolution blocks scheduling new jobs": {
 			SchedulingConfig: testfixtures.WithIndexedResourcesConfig(
-				[]configuration.IndexedResource{
+				[]configuration.ResourceType{
 					{Name: "cpu", Resolution: resource.MustParse("17")},
 					{Name: "memory", Resolution: resource.MustParse("128Mi")},
 				},
@@ -252,7 +252,7 @@ func TestGangScheduler(t *testing.T) {
 		},
 		"consider all nodes in the bucket": {
 			SchedulingConfig: testfixtures.WithIndexedResourcesConfig(
-				[]configuration.IndexedResource{
+				[]configuration.ResourceType{
 					{Name: "cpu", Resolution: resource.MustParse("1")},
 					{Name: "memory", Resolution: resource.MustParse("1Mi")},
 					{Name: "gpu", Resolution: resource.MustParse("1")},
@@ -593,6 +593,7 @@ func TestGangScheduler(t *testing.T) {
 				tc.SchedulingConfig.IndexedNodeLabels,
 				tc.SchedulingConfig.WellKnownNodeTypes,
 				stringinterner.New(1024),
+				testfixtures.TestResourceListFactory,
 			)
 			require.NoError(t, err)
 			txn := nodeDb.Txn(true)

internal/scheduler/internaltypes/quantity_util.go

@@ -0,0 +1,18 @@
+package internaltypes
+
+import (
+	"k8s.io/apimachinery/pkg/api/resource"
+)
+
+func QuantityToInt64RoundUp(q resource.Quantity, scale resource.Scale) int64 {
+	return q.ScaledValue(scale)
+}
+
+func QuantityToInt64RoundDown(q resource.Quantity, scale resource.Scale) int64 {
+	result := q.ScaledValue(scale)
+	q2 := resource.NewScaledQuantity(result, scale)
+	if q2.Cmp(q) > 0 {
+		result--
+	}
+	return result
+}

internal/scheduler/internaltypes/quantity_util_test.go

@@ -0,0 +1,69 @@
+package internaltypes
+
+import (
+	"testing"
+
+	"gopkg.in/inf.v0"
+
+	"github.com/stretchr/testify/assert"
+	"k8s.io/apimachinery/pkg/api/resource"
+)
+
+type quantityTest struct {
+	q                    resource.Quantity
+	expectedIntRoundDown int64
+	expectedIntRoundUp   int64
+}
+
+func TestQuantityToInt64_WithScaleMillis(t *testing.T) {
+	tests := []quantityTest{
+		{resource.MustParse("0"), 0, 0},
+		{resource.MustParse("1"), 1000, 1000},
+		{resource.MustParse("1m"), 1, 1},
+		{resource.MustParse("50m"), 50, 50},
+		{resource.MustParse("1Mi"), 1024 * 1024 * 1000, 1024 * 1024 * 1000},
+		{resource.MustParse("1e3"), 1000 * 1000, 1000 * 1000},
+		{*resource.NewMilliQuantity(1, resource.DecimalExponent), 1, 1},
+		{*resource.NewDecimalQuantity(*inf.NewDec(1, inf.Scale(0)), resource.DecimalExponent), 1000, 1000},
+		{resource.MustParse("1n"), 0, 1},
+		{resource.MustParse("100n"), 0, 1},
+		{resource.MustParse("999999999n"), 999, 1000},
+		{resource.MustParse("1000000000n"), 1000, 1000},
+		{resource.MustParse("1000000001n"), 1000, 1001},
+		{resource.MustParse("12.3m"), 12, 13},
+		{resource.MustParse("0.99m"), 0, 1},
+		{resource.MustParse("1.001m"), 1, 2},
+	}
+
+	for _, test := range tests {
+		assert.Equal(t, test.expectedIntRoundDown, QuantityToInt64RoundDown(test.q, resource.Milli), test.q)
+		assert.Equal(t, test.expectedIntRoundUp, QuantityToInt64RoundUp(test.q, resource.Milli), test.q)
+	}
+}
+
+func TestQuantityToInt64_WithUnitScale(t *testing.T) {
+	const tebi = 1024 * 1024 * 1024 * 1024
+	tests := []quantityTest{
+		{resource.MustParse("0"), 0, 0},
+		{resource.MustParse("1"), 1, 1},
+		{resource.MustParse("1m"), 0, 1},
+		{resource.MustParse("50m"), 0, 1},
+		{resource.MustParse("1Mi"), 1024 * 1024, 1024 * 1024},
+		{resource.MustParse("1.5Mi"), 1536 * 1024, 1536 * 1024},
+		{resource.MustParse("4Ti"), 4 * tebi, 4 * tebi},
+		{resource.MustParse("100000Ti"), 100000 * tebi, 100000 * tebi},
+		{resource.MustParse("1e3"), 1000, 1000},
+		{*resource.NewMilliQuantity(1, resource.DecimalExponent), 0, 1},
+		{*resource.NewDecimalQuantity(*inf.NewDec(1, inf.Scale(0)), resource.DecimalExponent), 1, 1},
+		{resource.MustParse("1n"), 0, 1},
+		{resource.MustParse("100n"), 0, 1},
+		{resource.MustParse("999999999n"), 0, 1},
+		{resource.MustParse("1000000000n"), 1, 1},
+		{resource.MustParse("1000000001n"), 1, 2},
+	}
+
+	for _, test := range tests {
+		assert.Equal(t, test.expectedIntRoundDown, QuantityToInt64RoundDown(test.q, resource.Scale(0)), test.q)
+		assert.Equal(t, test.expectedIntRoundUp, QuantityToInt64RoundUp(test.q, resource.Scale(0)), test.q)
+	}
+}

internal/scheduler/internaltypes/resource_list.go

@@ -0,0 +1,16 @@
+package internaltypes
+
+import "fmt"
+
+type ResourceList struct {
+	resources []int64
+	factory   *ResourceListFactory
+}
+
+func (rl *ResourceList) GetByName(name string) (int64, error) {
+	index, ok := rl.factory.nameToIndex[name]
+	if !ok {
+		return 0, fmt.Errorf("resource type %s not found", name)
+	}
+	return rl.resources[index], nil
+}

internal/scheduler/internaltypes/resource_list_factory.go

@@ -0,0 +1,93 @@
+package internaltypes
+
+import (
+	"fmt"
+	"math"
+
+	"github.com/pkg/errors"
+
+	k8sResource "k8s.io/apimachinery/pkg/api/resource"
+
+	"github.com/armadaproject/armada/internal/armada/configuration"
+	"github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
+)
+
+type ResourceListFactory struct {
+	nameToIndex map[string]int
+	indexToName []string
+	scales      []k8sResource.Scale
+}
+
+func MakeResourceListFactory(supportedResourceTypes []configuration.ResourceType) (*ResourceListFactory, error) {
+	if len(supportedResourceTypes) == 0 {
+		return nil, errors.New("no resource types configured")
+	}
+	indexToName := make([]string, len(supportedResourceTypes))
+	nameToIndex := make(map[string]int, len(supportedResourceTypes))
+	scales := make([]k8sResource.Scale, len(supportedResourceTypes))
+	for i, t := range supportedResourceTypes {
+		if dup, exists := nameToIndex[t.Name]; exists {
+			return nil, fmt.Errorf("duplicate resource type name %q", dup)
+		}
+		nameToIndex[t.Name] = i
+		indexToName[i] = t.Name
+		scales[i] = resolutionToScale(t.Resolution)
+	}
+	return &ResourceListFactory{
+		indexToName: indexToName,
+		nameToIndex: nameToIndex,
+		scales:      scales,
+	}, nil
+}
+
+// Convert resolution to a k8sResource.Scale
+// e.g.
+// 1     ->  0
+// 0.001 -> -3
+// 1000  ->  3
+func resolutionToScale(resolution k8sResource.Quantity) k8sResource.Scale {
+	if resolution.Sign() < 1 {
+		return k8sResource.Milli
+	}
+	return k8sResource.Scale(math.Floor(math.Log10(resolution.AsApproximateFloat64())))
+}
+
+// Ignore unknown resources, round down.
+func (factory *ResourceListFactory) FromNodeProto(resources schedulerobjects.ResourceList) ResourceList {
+	result := make([]int64, len(factory.indexToName))
+	for k, v := range resources.Resources {
+		index, ok := factory.nameToIndex[k]
+		if ok {
+			result[index] = QuantityToInt64RoundDown(v, factory.scales[index])
+		}
+	}
+	return ResourceList{resources: result, factory: factory}
+}
+
+// Fail on unknown resources, round up.
+func (factory *ResourceListFactory) FromJobProto(resources schedulerobjects.ResourceList) (ResourceList, error) {
+	result := make([]int64, len(factory.indexToName))
+	for k, v := range resources.Resources {
+		index, ok := factory.nameToIndex[k]
+		if ok {
+			result[index] = QuantityToInt64RoundUp(v, factory.scales[index])
+		} else {
+			return ResourceList{}, fmt.Errorf("unknown resource type %q", k)
+		}
+	}
+	return ResourceList{resources: result, factory: factory}, nil
+}
+
+func (factory *ResourceListFactory) SummaryString() string {
+	result := ""
+	for i, name := range factory.indexToName {
+		if i > 0 {
+			result += " "
+		}
+		scale := factory.scales[i]
+		resolution := k8sResource.NewScaledQuantity(1, scale)
+		maxValue := k8sResource.NewScaledQuantity(math.MaxInt64, scale)
+		result += fmt.Sprintf("%s (scale %v, resolution %v, maxValue %f)", name, scale, resolution, maxValue.AsApproximateFloat64())
+	}
+	return result
+}