Add Idealised Value Metric (#4414)

This Pr adds a new metric (idealised_scheduled_value) which represents
the "ideal" value of jobs that could be scheduled each scheduling cycle.

We calculate this by considering that all (pool) resources are present
on one giang node. This frees us from pesky constraints such as "we need
to binpack efficiently" and allows us to quickly generate an upper bound
for the maximum amount of value that could be scheduled.

---------

Signed-off-by: Chris Martin <chris@cmartinit.co.uk>
Co-authored-by: Christopher Martin <Chris.Martin@gresearch.co.uk>

Author: d80tb7

internal/scheduler/internaltypes/node.go

@@ -249,6 +249,10 @@ func (node *Node) GetLabels() map[string]string {
 	return deepCopyLabels(node.labels)
 }
 
+func (node *Node) GetRunningJobIds() []string {
+	return maps.Keys(node.AllocatedByJobId)
+}
+
 func (node *Node) GetLabelValue(key string) (string, bool) {
 	val, ok := node.labels[key]
 	return val, ok

internal/scheduler/internaltypes/node_factory.go

@@ -64,7 +64,6 @@ func (f *NodeFactory) CreateNodeAndType(
 	labels map[string]string,
 	totalResources ResourceList,
 	allocatableResources ResourceList,
-	unallocatableResources map[int32]ResourceList,
 	allocatableByPriority map[int32]ResourceList,
 ) *Node {
 	return CreateNodeAndType(

internal/scheduler/jobdb/job.go

@@ -425,10 +425,22 @@ func (job *Job) WithBidPrices(bids map[string]pricing.Bid) *Job {
 	return j
 }
 
+// GetBidPrice resolves the current bid price.
+// It considers running, non-preemptible jobs to have an effectively infinite price.
 func (job *Job) GetBidPrice(pool string) float64 {
 	if !job.queued && !job.priorityClass.Preemptible {
 		return pricing.NonPreemptibleRunningPrice
 	}
+	return job.getBidPrice(pool)
+}
+
+// GetRawBidPrice resolves the current bid price.
+// It considers running, non-preemptible jobs to have the simple job running price (i.e. ignores premtibility pricing).
+func (job *Job) GetRawBidPrice(pool string) float64 {
+	return job.getBidPrice(pool)
+}
+
+func (job *Job) getBidPrice(pool string) float64 {
 	bidPrice, present := job.bidPricesPool[pool]
 	if !present {
 		return 0

internal/scheduler/metrics/cycle_metrics.go

@@ -63,6 +63,8 @@ type perCycleMetrics struct {
 	nodeAllocatedResource        *prometheus.GaugeVec
 	indicativePrice              *prometheus.GaugeVec
 	indicativePriceSchedulable   *prometheus.GaugeVec
+	idealisedScheduledValue      *prometheus.GaugeVec
+	realisedScheduledValue       *prometheus.GaugeVec
 }
 
 func newPerCycleMetrics() *perCycleMetrics {
@@ -290,6 +292,22 @@ func newPerCycleMetrics() *perCycleMetrics {
 		poolAndShapeAndReasonLabels,
 	)
 
+	idealisedScheduledValue := prometheus.NewGaugeVec(
+		prometheus.GaugeOpts{
+			Name: prefix + "idealised_scheduled_value",
+			Help: "Idealised value scheduled per queue",
+		},
+		poolAndQueueLabels,
+	)
+
+	realisedScheduledValue := prometheus.NewGaugeVec(
+		prometheus.GaugeOpts{
+			Name: prefix + "realised_scheduled_value",
+			Help: "Realised value scheduled per queue",
+		},
+		poolAndQueueLabels,
+	)
+
 	return &perCycleMetrics{
 		consideredJobs:               consideredJobs,
 		fairShare:                    fairShare,
@@ -319,6 +337,8 @@ func newPerCycleMetrics() *perCycleMetrics {
 		nodeAllocatedResource:        nodeAllocatedResource,
 		indicativePrice:              indicativePrice,
 		indicativePriceSchedulable:   indicativePriceSchedulable,
+		idealisedScheduledValue:      idealisedScheduledValue,
+		realisedScheduledValue:       realisedScheduledValue,
 	}
 }
 
@@ -432,6 +452,8 @@ func (m *cycleMetrics) ReportSchedulerResult(ctx *armadacontext.Context, result
 			currentCycle.shortJobPenalty.WithLabelValues(pool, queue).Set(shortJobPenalty)
 			currentCycle.queueWeight.WithLabelValues(pool, queue).Set(queueContext.Weight)
 			currentCycle.rawQueueWeight.WithLabelValues(pool, queue).Set(queueContext.RawWeight)
+			currentCycle.idealisedScheduledValue.WithLabelValues(pool, queue).Set(queueContext.IdealisedValue)
+			currentCycle.realisedScheduledValue.WithLabelValues(pool, queue).Set(queueContext.RealisedValue)
 			for _, r := range queueContext.GetBillableResource().GetResources() {
 				currentCycle.billableResource.WithLabelValues(pool, queue, r.Name).Set(float64(r.RawValue))
 			}
@@ -563,6 +585,8 @@ func (m *cycleMetrics) describe(ch chan<- *prometheus.Desc) {
 		currentCycle.nodeAllocatedResource.Describe(ch)
 		currentCycle.indicativePrice.Describe(ch)
 		currentCycle.indicativePriceSchedulable.Describe(ch)
+		currentCycle.idealisedScheduledValue.Describe(ch)
+		currentCycle.realisedScheduledValue.Describe(ch)
 	}
 
 	m.reconciliationCycleTime.Describe(ch)
@@ -603,6 +627,8 @@ func (m *cycleMetrics) collect(ch chan<- prometheus.Metric) {
 		currentCycle.nodeAllocatedResource.Collect(ch)
 		currentCycle.indicativePrice.Collect(ch)
 		currentCycle.indicativePriceSchedulable.Collect(ch)
+		currentCycle.idealisedScheduledValue.Collect(ch)
+		currentCycle.realisedScheduledValue.Collect(ch)
 	}
 
 	m.reconciliationCycleTime.Collect(ch)

internal/scheduler/pricing/bid_price.go

@@ -3,6 +3,8 @@ package pricing
 import (
 	"time"
 
+	"github.com/armadaproject/armada/internal/scheduler/internaltypes"
+
 	"github.com/armadaproject/armada/internal/common/armadacontext"
 	"github.com/armadaproject/armada/internal/common/cache"
 	"github.com/armadaproject/armada/pkg/bidstore"
@@ -12,8 +14,8 @@ type BidPriceCache struct {
 	cache *cache.GenericCache[BidPriceSnapshot]
 }
 
-func NewBidPriceCache(client bidstore.BidRetrieverServiceClient, updateFrequency time.Duration) *BidPriceCache {
-	svc := NewExternalBidPriceService(client)
+func NewBidPriceCache(client bidstore.BidRetrieverServiceClient, rlf *internaltypes.ResourceListFactory, updateFrequency time.Duration) *BidPriceCache {
+	svc := NewExternalBidPriceService(client, rlf)
 	return &BidPriceCache{
 		cache: cache.NewGenericCache(svc.GetBidPrices, updateFrequency),
 	}

internal/scheduler/pricing/bid_price_service_test.go

@@ -6,19 +6,69 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"k8s.io/apimachinery/pkg/api/resource"
 
+	schedulerconfiguration "github.com/armadaproject/armada/internal/scheduler/configuration"
+	"github.com/armadaproject/armada/internal/scheduler/internaltypes"
 	"github.com/armadaproject/armada/pkg/bidstore"
 )
 
+var testResourceListFactory, _ = internaltypes.NewResourceListFactory(
+	[]schedulerconfiguration.ResourceType{
+		{Name: "memory", Resolution: resource.MustParse("1")},
+		{Name: "cpu", Resolution: resource.MustParse("1m")},
+		{Name: "nvidia.com/gpu", Resolution: resource.MustParse("1m")},
+	},
+	[]schedulerconfiguration.FloatingResourceConfig{},
+)
+
 func TestConvert(t *testing.T) {
 	tests := map[string]struct {
 		input    *bidstore.RetrieveBidsResponse
 		expected BidPriceSnapshot
 	}{
 		"empty input": {
 			input: &bidstore.RetrieveBidsResponse{},
+			expected: BidPriceSnapshot{
+				Bids:          make(map[PriceKey]map[string]Bid),
+				ResourceUnits: map[string]internaltypes.ResourceList{},
+			},
+		},
+		"resourceUnit": {
+			input: &bidstore.RetrieveBidsResponse{
+				PoolResourceUnits: map[string]*bidstore.ResourceShape{
+					"pool1": {
+						Resources: map[string]*resource.Quantity{
+							"cpu":    mustParsePtr("1"),
+							"memory": mustParsePtr("1Gi"),
+						},
+					},
+					"pool2": {
+						Resources: map[string]*resource.Quantity{
+							"cpu":            mustParsePtr("2"),
+							"memory":         mustParsePtr("2Gi"),
+							"nvidia.com/gpu": mustParsePtr("1"),
+						},
+					},
+				},
+			},
 			expected: BidPriceSnapshot{
 				Bids: make(map[PriceKey]map[string]Bid),
+				ResourceUnits: map[string]internaltypes.ResourceList{
+					"pool1": testResourceListFactory.FromNodeProto(
+						map[string]*resource.Quantity{
+							"cpu":    mustParsePtr("1"),
+							"memory": mustParsePtr("1Gi"),
+						},
+					),
+					"pool2": testResourceListFactory.FromNodeProto(
+						map[string]*resource.Quantity{
+							"cpu":            mustParsePtr("2"),
+							"memory":         mustParsePtr("2Gi"),
+							"nvidia.com/gpu": mustParsePtr("1"),
+						},
+					),
+				},
 			},
 		},
 		"single bid with both phases": {
@@ -55,6 +105,7 @@ func TestConvert(t *testing.T) {
 						},
 					},
 				},
+				ResourceUnits: map[string]internaltypes.ResourceList{},
 			},
 		},
 		"uses fallback if direct bid is missing": {
@@ -63,7 +114,7 @@ func TestConvert(t *testing.T) {
 					"queue1": {
 						PoolBids: map[string]*bidstore.PoolBids{
 							"pool1": {
-								FallbackBid: &bidstore.PriceBandBids{
+								FallbackBids: &bidstore.PriceBandBids{
 									PricingPhaseBids: []*bidstore.PricingPhaseBid{
 										CreateQueuedBid(5.0),
 										CreateRunningBid(10.0),
@@ -158,13 +209,14 @@ func TestConvert(t *testing.T) {
 						},
 					},
 				},
+				ResourceUnits: map[string]internaltypes.ResourceList{},
 			},
 		},
 	}
 
 	for name, tc := range tests {
 		t.Run(name, func(t *testing.T) {
-			actual := convert(tc.input)
+			actual := convert(tc.input, testResourceListFactory)
 
 			// Remove non-deterministic time field before comparison
 			require.False(t, actual.Timestamp.IsZero(), "timestamp should be set")
@@ -192,3 +244,8 @@ func CreateRunningBid(p float64) *bidstore.PricingPhaseBid {
 		},
 	}
 }
+
+func mustParsePtr(qty string) *resource.Quantity {
+	q := resource.MustParse(qty)
+	return &q
+}

internal/scheduler/pricing/bid_service.go

@@ -4,6 +4,7 @@ import (
 	"time"
 
 	"github.com/armadaproject/armada/internal/common/armadacontext"
+	"github.com/armadaproject/armada/internal/scheduler/internaltypes"
 	"github.com/armadaproject/armada/internal/scheduler/queue"
 	"github.com/armadaproject/armada/pkg/bidstore"
 )
@@ -59,11 +60,13 @@ func (b *LocalBidPriceService) GetBidPrices(ctx *armadacontext.Context) (BidPric
 
 type ExternalBidPriceService struct {
 	client bidstore.BidRetrieverServiceClient
+	rlf    *internaltypes.ResourceListFactory
 }
 
-func NewExternalBidPriceService(client bidstore.BidRetrieverServiceClient) *ExternalBidPriceService {
+func NewExternalBidPriceService(client bidstore.BidRetrieverServiceClient, rlf *internaltypes.ResourceListFactory) *ExternalBidPriceService {
 	return &ExternalBidPriceService{
 		client: client,
+		rlf:    rlf,
 	}
 }
 
@@ -72,13 +75,18 @@ func (b *ExternalBidPriceService) GetBidPrices(ctx *armadacontext.Context) (BidP
 	if err != nil {
 		return BidPriceSnapshot{}, err
 	}
-	return convert(resp), nil
+	return convert(resp, b.rlf), nil
 }
 
-func convert(resp *bidstore.RetrieveBidsResponse) BidPriceSnapshot {
+func convert(resp *bidstore.RetrieveBidsResponse, rlf *internaltypes.ResourceListFactory) BidPriceSnapshot {
 	snapshot := BidPriceSnapshot{
-		Timestamp: time.Now(),
-		Bids:      make(map[PriceKey]map[string]Bid),
+		Timestamp:     time.Now(),
+		Bids:          make(map[PriceKey]map[string]Bid),
+		ResourceUnits: make(map[string]internaltypes.ResourceList),
+	}
+
+	for pool, resourceShape := range resp.PoolResourceUnits {
+		snapshot.ResourceUnits[pool] = rlf.FromNodeProto(resourceShape.Resources)
 	}
 
 	for queue, qb := range resp.QueueBids {
@@ -88,7 +96,7 @@ func convert(resp *bidstore.RetrieveBidsResponse) BidPriceSnapshot {
 
 			for pool, poolBids := range qb.PoolBids {
 				bb, _ := poolBids.GetBidsForBand(band)
-				fallback := poolBids.GetFallbackBid()
+				fallback := poolBids.GetFallbackBids()
 
 				queued, hasQueued := getPrice(bb, fallback, bidstore.PricingPhase_PRICING_PHASE_QUEUEING)
 				running, hasRunning := getPrice(bb, fallback, bidstore.PricingPhase_PRICING_PHASE_RUNNING)

internal/scheduler/pricing/types.go

@@ -4,6 +4,7 @@ import (
 	"time"
 
 	"github.com/armadaproject/armada/internal/common/armadacontext"
+	"github.com/armadaproject/armada/internal/scheduler/internaltypes"
 	"github.com/armadaproject/armada/pkg/bidstore"
 )
 
@@ -19,8 +20,9 @@ type PriceKey struct {
 }
 
 type BidPriceSnapshot struct {
-	Timestamp time.Time
-	Bids      map[PriceKey]map[string]Bid
+	Timestamp     time.Time
+	Bids          map[PriceKey]map[string]Bid
+	ResourceUnits map[string]internaltypes.ResourceList
 }
 
 type Bid struct {

internal/scheduler/scheduler.go

@@ -325,14 +325,14 @@ func (s *Scheduler) cycle(ctx *armadacontext.Context, updateAll bool, leaderToke
 	// Schedule jobs.
 	if shouldSchedule {
 		start := time.Now()
-		err := s.updateJobPrices(ctx, txn)
+		resourceUnits, err := s.updateJobPrices(ctx, txn)
 		if err != nil {
 			return overallSchedulerResult, err
 		}
 		ctx.Logger().Infof("updating job prices in %s", time.Now().Sub(start))
 
 		var result *scheduling.SchedulerResult
-		result, err = s.schedulingAlgo.Schedule(ctx, txn)
+		result, err = s.schedulingAlgo.Schedule(ctx, resourceUnits, txn)
 		if err != nil {
 			return overallSchedulerResult, err
 		}
@@ -463,16 +463,16 @@ func (s *Scheduler) syncState(ctx *armadacontext.Context, initial bool) ([]*jobd
 
 // TODO - This is highly inefficient and will only work if market driven pools are small
 // We should rewrite how bids are stored in an efficient manner
-func (s *Scheduler) updateJobPrices(ctx *armadacontext.Context, txn *jobdb.Txn) error {
+func (s *Scheduler) updateJobPrices(ctx *armadacontext.Context, txn *jobdb.Txn) (map[string]internaltypes.ResourceList, error) {
 	if len(s.marketDrivenPools) == 0 {
-		return nil
+		return nil, nil
 	}
 
 	jobs := txn.GetAll()
 	jobsByQueue := map[string]map[bidstore.PriceBand][]*jobdb.Job{}
 	updatedBids, err := s.bidPriceProvider.GetBidPrices(ctx)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	hasMarketDrivenPool := func(j *jobdb.Job) bool {
@@ -512,7 +512,11 @@ func (s *Scheduler) updateJobPrices(ctx *armadacontext.Context, txn *jobdb.Txn)
 		}
 	}
 	ctx.Logger().Infof("updating the prices of %d jobs", len(updatedJobs))
-	return txn.Upsert(updatedJobs)
+	err = txn.Upsert(updatedJobs)
+	if err != nil {
+		return nil, err
+	}
+	return updatedBids.ResourceUnits, nil
 }
 
 func (s *Scheduler) createSchedulingInfoWithNodeAntiAffinityForAttemptedRuns(job *jobdb.Job) (*internaltypes.JobSchedulingInfo, error) {

internal/scheduler/scheduler_test.go

@@ -1763,7 +1763,7 @@ type testSchedulingAlgo struct {
 	persisted bool
 }
 
-func (t *testSchedulingAlgo) Schedule(_ *armadacontext.Context, txn *jobdb.Txn) (*scheduling.SchedulerResult, error) {
+func (t *testSchedulingAlgo) Schedule(_ *armadacontext.Context, _ map[string]internaltypes.ResourceList, txn *jobdb.Txn) (*scheduling.SchedulerResult, error) {
 	t.numberOfScheduleCalls++
 	if t.shouldError {
 		return nil, errors.New("error scheduling jobs")