Add CohortSubtreeResourceUsage metric

Author: mszadkow

pkg/cache/scheduler/cache.go

@@ -141,6 +141,19 @@ type Cache struct {
 	customLabels *metrics.CustomLabels
 }
 
+func (c *Cache) GetDirectCohortNameForWorkload(wl *kueue.Workload) kueue.CohortReference {
+	if wl == nil || wl.Status.Admission == nil {
+		return ""
+	}
+	c.RLock()
+	defer c.RUnlock()
+	cq := c.hm.ClusterQueue(wl.Status.Admission.ClusterQueue)
+	if cq == nil || !cq.HasParent() {
+		return ""
+	}
+	return cq.Parent().Name
+}
+
 func New(client client.Client, options ...Option) *Cache {
 	cache := &Cache{
 		client:                 client,

pkg/cache/scheduler/cache_test.go

@@ -3753,3 +3753,64 @@ func TestGetWorkloadFromCache(t *testing.T) {
 		})
 	}
 }
+
+func TestGetDirectCohortNameForWorkload(t *testing.T) {
+	now := time.Now().Truncate(time.Second)
+
+	cases := map[string]struct {
+		clusterQueues []kueue.ClusterQueue
+		wl            *kueue.Workload
+		wantCohort    kueue.CohortReference
+	}{
+		"nil workload": {
+			wl:         nil,
+			wantCohort: "",
+		},
+		"workload without admission": {
+			wl:         utiltestingapi.MakeWorkload("wl", "ns").Obj(),
+			wantCohort: "",
+		},
+		"cluster queue not found": {
+			wl: utiltestingapi.MakeWorkload("wl", "ns").ReserveQuotaAt(&kueue.Admission{
+				ClusterQueue: "missing-cq",
+			}, now).Obj(),
+			wantCohort: "",
+		},
+		"cluster queue without cohort": {
+			clusterQueues: []kueue.ClusterQueue{
+				*utiltestingapi.MakeClusterQueue("cq").Obj(),
+			},
+			wl: utiltestingapi.MakeWorkload("wl", "ns").ReserveQuotaAt(&kueue.Admission{
+				ClusterQueue: "cq",
+			}, now).Obj(),
+			wantCohort: "",
+		},
+		"returns workload cohort": {
+			clusterQueues: []kueue.ClusterQueue{
+				*utiltestingapi.MakeClusterQueue("cq").Cohort("team-a").Obj(),
+			},
+			wl: utiltestingapi.MakeWorkload("wl", "ns").ReserveQuotaAt(&kueue.Admission{
+				ClusterQueue: "cq",
+			}, now).Obj(),
+			wantCohort: "team-a",
+		},
+	}
+
+	for name, tc := range cases {
+		t.Run(name, func(t *testing.T) {
+			ctx, _ := utiltesting.ContextWithLog(t)
+			cache := New(utiltesting.NewFakeClient())
+
+			for i := range tc.clusterQueues {
+				if err := cache.AddClusterQueue(ctx, &tc.clusterQueues[i]); err != nil {
+					t.Fatalf("failed adding clusterQueue %s: %v", tc.clusterQueues[i].Name, err)
+				}
+			}
+
+			got := cache.GetDirectCohortNameForWorkload(tc.wl)
+			if diff := cmp.Diff(tc.wantCohort, got); diff != "" {
+				t.Errorf("GetDirectCohortNameForWorkload() mismatch (-want,+got):\n%s", diff)
+			}
+		})
+	}
+}

pkg/cache/scheduler/cohort_metrics.go

@@ -17,7 +17,11 @@ limitations under the License.
 package scheduler
 
 import (
+	"maps"
+	"slices"
+
 	"github.com/go-logr/logr"
+	"k8s.io/apimachinery/pkg/util/sets"
 
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta2"
 	"sigs.k8s.io/kueue/pkg/cache/hierarchy"
@@ -28,7 +32,8 @@ import (
 type cohortMetricPoint struct {
 	cohortName     kueue.CohortReference
 	flavorResource resources.FlavorResource
-	qty            int64
+	quotaQty       int64
+	usageQty       int64
 }
 
 func (c *Cache) RecordCohortMetrics(log logr.Logger, cohortName kueue.CohortReference) {
@@ -73,21 +78,39 @@ func (c *Cache) collectCohortMetricPoints(cohortName kueue.CohortReference, simu
 	}
 
 	removedQuota := ch.resourceNode.SubtreeQuota
+
+	// Cache total subtree usage calls during this run
+	usageCache := make(map[*cohort]resources.FlavorResourceQuantities)
+	removedUsage := totalSubtreeUsageWithCache(ch, usageCache)
+
 	var points []cohortMetricPoint
 	for ancestor := range ch.PathSelfToRoot() {
 		quotas := ancestor.resourceNode.SubtreeQuota
+		usages := totalSubtreeUsageWithCache(ancestor, usageCache)
+
+		var baselineUsage resources.FlavorResourceQuantities
 		if simulateRemoval {
 			quotas = removedQuota
+			usages = removedUsage
+			baselineUsage = totalSubtreeUsageWithCache(ancestor, usageCache)
 		}
 
-		for flr, qty := range quotas {
+		keys := sets.New[resources.FlavorResource]()
+		keys.Insert(slices.Collect(maps.Keys(quotas))...)
+		keys.Insert(slices.Collect(maps.Keys(usages))...)
+
+		for fr := range keys {
+			quotaQty := quotas[fr]
+			usageQty := usages[fr]
 			if simulateRemoval {
-				qty = max(ancestor.resourceNode.SubtreeQuota[flr]-qty, 0)
+				quotaQty = max(ancestor.resourceNode.SubtreeQuota[fr]-quotaQty, 0)
+				usageQty = max(baselineUsage[fr]-usageQty, 0)
 			}
 			points = append(points, cohortMetricPoint{
 				cohortName:     ancestor.Name,
-				flavorResource: flr,
-				qty:            qty,
+				flavorResource: fr,
+				quotaQty:       quotaQty,
+				usageQty:       usageQty,
 			})
 		}
 	}
@@ -99,19 +122,43 @@ func (c *Cache) withCohortLogger(log logr.Logger, cohortName kueue.CohortReferen
 }
 
 func (c *Cache) applyCohortMetricPoint(p cohortMetricPoint) {
-	if p.qty <= 0 {
-		metrics.ClearCohortSubtreeQuota(
-			p.cohortName,
-			string(p.flavorResource.Flavor),
-			string(p.flavorResource.Resource),
-		)
-		return
+	flavor := p.flavorResource.Flavor
+	resource := p.flavorResource.Resource
+
+	if p.quotaQty <= 0 {
+		metrics.ClearCohortSubtreeQuota(p.cohortName, flavor, resource)
+	} else {
+		metrics.ReportCohortSubtreeQuota(p.cohortName, flavor, resource, p.quotaQty, c.roleTracker)
+	}
+
+	if p.usageQty <= 0 {
+		metrics.ClearCohortSubtreeResourceUsage(p.cohortName, flavor, resource)
+	} else {
+		metrics.ReportCohortSubtreeResourceUsage(p.cohortName, flavor, resource, p.usageQty, c.roleTracker)
+	}
+}
+
+func accumulateUsage(total, usage resources.FlavorResourceQuantities) {
+	for fr, qty := range usage {
+		total[fr] += qty
+	}
+}
+
+// totalSubtreeUsageWithCache returns the sum of *actual* resource usage in the subtree rooted at given cohort
+func totalSubtreeUsageWithCache(ch *cohort, cache map[*cohort]resources.FlavorResourceQuantities) resources.FlavorResourceQuantities {
+	if cached, found := cache[ch]; found {
+		return cached
 	}
-	metrics.ReportCohortSubtreeQuota(
-		p.cohortName,
-		string(p.flavorResource.Flavor),
-		string(p.flavorResource.Resource),
-		float64(p.qty),
-		c.roleTracker,
-	)
+
+	total := make(resources.FlavorResourceQuantities)
+	for _, cq := range ch.ChildCQs() {
+		accumulateUsage(total, cq.getResourceNode().Usage)
+	}
+
+	for _, child := range ch.ChildCohorts() {
+		accumulateUsage(total, totalSubtreeUsageWithCache(child, cache))
+	}
+
+	cache[ch] = total
+	return total
 }

pkg/controller/core/workload_controller.go

@@ -317,7 +317,7 @@ func (r *WorkloadReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 				return ctrl.Result{}, err
 			}
 		} else {
-			if !r.cache.AddOrUpdateWorkload(log, wl.DeepCopy()) {
+			if !r.addOrUpdateWorkload(log, wl.DeepCopy()) {
 				log.V(2).Info("ClusterQueue for workload didn't exist; ignored for now")
 			}
 		}
@@ -607,7 +607,7 @@ func (r *WorkloadReconciler) deleteWorkloadFromCaches(ctx context.Context, names
 	// by the scheduler, and leaving them blocks ClusterQueue finalizer removal.
 	// The operation is idempotent if the workload was never in the cache.
 	r.queues.QueueAssociatedInadmissibleWorkloadsAfter(ctx, wlRef, func() {
-		if err := r.cache.DeleteWorkload(log, wlRef); err != nil {
+		if err := r.deleteWorkload(log, wlRef); err != nil {
 			log.Error(err, "Failed to delete workload from cache")
 		}
 	})
@@ -977,7 +977,7 @@ func (r *WorkloadReconciler) Create(e event.TypedCreateEvent[*kueue.Workload]) b
 		}
 		return true
 	}
-	if !r.cache.AddOrUpdateWorkload(log, wlCopy) {
+	if !r.addOrUpdateWorkload(log, wlCopy) {
 		log.V(2).Info("ClusterQueue for workload didn't exist; ignored for now")
 	}
 	r.queues.QueueSecondPassIfNeeded(ctx, e.Object, 0)
@@ -1050,7 +1050,7 @@ func (r *WorkloadReconciler) Update(e event.TypedUpdateEvent[*kueue.Workload]) b
 			// Delete the workload from cache while holding the queues lock
 			// to guarantee that requeued workloads are taken into account before
 			// the next scheduling cycle.
-			if err := r.cache.DeleteWorkload(log, wlKey); err != nil && prevStatus == workload.StatusAdmitted {
+			if err := r.deleteWorkload(log, wlKey); err != nil && prevStatus == workload.StatusAdmitted {
 				log.Error(err, "Failed to delete workload from cache")
 			}
 		})
@@ -1065,7 +1065,7 @@ func (r *WorkloadReconciler) Update(e event.TypedUpdateEvent[*kueue.Workload]) b
 		}
 	case prevStatus == workload.StatusPending && (status == workload.StatusQuotaReserved || status == workload.StatusAdmitted):
 		r.queues.DeleteWorkload(log, wlKey)
-		if !r.cache.AddOrUpdateWorkload(log, wlCopy) {
+		if !r.addOrUpdateWorkload(log, wlCopy) {
 			log.V(2).Info("ClusterQueue for workload didn't exist; ignored for now")
 		}
 		if afs.Enabled(r.admissionFSConfig) && status == workload.StatusAdmitted && r.cache.ClusterQueueUsesAdmissionFairSharing(wlCopy.Status.Admission.ClusterQueue) {
@@ -1083,7 +1083,7 @@ func (r *WorkloadReconciler) Update(e event.TypedUpdateEvent[*kueue.Workload]) b
 			// Delete the workload from cache while holding the queues lock
 			// to guarantee that requeued workloads are taken into account before
 			// the next scheduling cycle.
-			if err := r.cache.DeleteWorkload(log, wlKey); err != nil {
+			if err := r.deleteWorkload(log, wlKey); err != nil {
 				log.Error(err, "Failed to delete workload from cache")
 			}
 			// Here we don't take the lock as it is already taken by the wrapping function.
@@ -1111,13 +1111,13 @@ func (r *WorkloadReconciler) Update(e event.TypedUpdateEvent[*kueue.Workload]) b
 			// Update the workload from cache while holding the queues lock
 			// to guarantee that requeued workloads are taken into account before
 			// the next scheduling cycle.
-			r.cache.AddOrUpdateWorkload(log, wlCopy)
+			r.addOrUpdateWorkload(log, wlCopy)
 		})
 
 	default:
 		// Workload update in the cache is handled here; however, some fields are immutable
 		// and are not supposed to actually change anything.
-		r.cache.AddOrUpdateWorkload(log, wlCopy)
+		r.addOrUpdateWorkload(log, wlCopy)
 	}
 	r.queues.QueueSecondPassIfNeeded(ctx, wlCopy, 0)
 	return true
@@ -1452,3 +1452,23 @@ func (h *draEventHandler) Generic(ctx context.Context, e event.TypedGenericEvent
 func (r *WorkloadReconciler) GetDRAReconcileChannel() chan<- event.TypedGenericEvent[*kueue.Workload] {
 	return r.draReconcileChannel
 }
+
+func (r *WorkloadReconciler) addOrUpdateWorkload(log logr.Logger, wl *kueue.Workload) bool {
+	updated := r.cache.AddOrUpdateWorkload(log, wl)
+	if wl != nil && workload.HasQuotaReservation(wl) {
+		r.cache.RecordCohortMetrics(log, r.cache.GetDirectCohortNameForWorkload(wl))
+	}
+
+	return updated
+}
+
+func (r *WorkloadReconciler) deleteWorkload(log logr.Logger, wlKey workload.Reference) error {
+	wl := r.cache.GetWorkloadFromCache(wlKey)
+	cohortName := r.cache.GetDirectCohortNameForWorkload(wl)
+	err := r.cache.DeleteWorkload(log, wlKey)
+	if err != nil {
+		return err
+	}
+	r.cache.RecordCohortMetrics(log, cohortName)
+	return nil
+}

pkg/metrics/metrics.go

@@ -20,6 +20,7 @@ import (
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
+	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"sigs.k8s.io/controller-runtime/pkg/metrics"
 
@@ -254,6 +255,10 @@ var (
 	// +metricsdoc:group=cohort
 	// +metricsdoc:labels=cohort="the name of the Cohort",flavor="the resource flavor name",resource="the resource name",replica_role="one of `leader`, `follower`, or `standalone`"
 	CohortSubtreeQuota *prometheus.GaugeVec
+
+	// +metricsdoc:group=cohort
+	// +metricsdoc:labels=cohort="the name of the Cohort",flavor="the resource flavor name",resource="the resource name",replica_role="one of `leader`, `follower`, or `standalone`"
+	CohortSubtreeResourceUsage *prometheus.GaugeVec
 )
 
 func trackGaugeVec(g *prometheus.GaugeVec) *prometheus.GaugeVec {
@@ -728,6 +733,14 @@ If the Cohort has a weight of zero and is borrowing, this will return NaN.`,
 			Help:      `Reports the cohort's nominal quota aggregated within the cohort's subtree. The values are reported per resource and flavor`,
 		}, []string{"cohort", "flavor", "resource", "replica_role"},
 	))
+
+	CohortSubtreeResourceUsage = trackGaugeVec(prometheus.NewGaugeVec(
+		prometheus.GaugeOpts{
+			Subsystem: constants.KueueName,
+			Name:      "cohort_subtree_resource_usage",
+			Help:      `Reports the cohort's resource usage aggregated within the cohort's subtree. The values are reported per resource and flavor`,
+		}, []string{"cohort", "flavor", "resource", "replica_role"},
+	))
 }
 
 func init() {
@@ -916,6 +929,7 @@ func ClearLocalQueueMetrics(lq LocalQueueReference) {
 func ClearCohortMetrics(cohortName string) {
 	CohortSubtreeQuota.DeletePartialMatch(prometheus.Labels{"cohort": cohortName})
 	CohortWeightedShare.DeletePartialMatch(prometheus.Labels{"cohort": cohortName})
+	CohortSubtreeResourceUsage.DeletePartialMatch(prometheus.Labels{"cohort": cohortName})
 }
 
 func ReportClusterQueueStatus(cqName kueue.ClusterQueueReference, cqStatus ClusterQueueStatus, customLabelValues []string, tracker *roletracker.RoleTracker) {
@@ -965,23 +979,35 @@ func ReportClusterQueueQuotas(cohort kueue.CohortReference, queue, flavor, resou
 	ClusterQueueResourceLendingLimit.WithLabelValues(labels...).Set(lending)
 }
 
-func ReportCohortSubtreeQuota(cohort kueue.CohortReference, flavor, resource string, quota float64, tracker *roletracker.RoleTracker) {
-	CohortSubtreeQuota.WithLabelValues(string(cohort), flavor, resource, roletracker.GetRole(tracker)).Set(quota)
+func ReportCohortSubtreeQuota(cohort kueue.CohortReference, flavor kueue.ResourceFlavorReference, resource corev1.ResourceName, quota int64, tracker *roletracker.RoleTracker) {
+	CohortSubtreeQuota.WithLabelValues(string(cohort), string(flavor), string(resource), roletracker.GetRole(tracker)).Set(float64(quota))
 }
 
-func ClearCohortSubtreeQuota(cohort kueue.CohortReference, flavor, resource string) {
-	lbls := prometheus.Labels{
-		"cohort": string(cohort),
-	}
+func cohortPartialMatchLabels(cohort kueue.CohortReference, flavor, resource string) prometheus.Labels {
+	lbls := prometheus.Labels{"cohort": string(cohort)}
 	if len(flavor) != 0 {
 		lbls["flavor"] = flavor
 	}
 	if len(resource) != 0 {
 		lbls["resource"] = resource
 	}
+	return lbls
+}
+
+func ClearCohortSubtreeQuota(cohort kueue.CohortReference, flavor kueue.ResourceFlavorReference, resource corev1.ResourceName) {
+	lbls := cohortPartialMatchLabels(cohort, string(flavor), string(resource))
 	CohortSubtreeQuota.DeletePartialMatch(lbls)
 }
 
+func ReportCohortSubtreeResourceUsage(cohort kueue.CohortReference, flavor kueue.ResourceFlavorReference, resource corev1.ResourceName, usage int64, tracker *roletracker.RoleTracker) {
+	CohortSubtreeResourceUsage.WithLabelValues(string(cohort), string(flavor), string(resource), roletracker.GetRole(tracker)).Set(float64(usage))
+}
+
+func ClearCohortSubtreeResourceUsage(cohort kueue.CohortReference, flavor kueue.ResourceFlavorReference, resource corev1.ResourceName) {
+	lbls := cohortPartialMatchLabels(cohort, string(flavor), string(resource))
+	CohortSubtreeResourceUsage.DeletePartialMatch(lbls)
+}
+
 func ReportClusterQueueResourceReservations(cohort kueue.CohortReference, queue, flavor, resource string, usage float64, customLabelValues []string, tracker *roletracker.RoleTracker) {
 	labels := append([]string{string(cohort), queue, flavor, resource, roletracker.GetRole(tracker)}, customLabelValues...)
 	ClusterQueueResourceReservations.WithLabelValues(labels...).Set(usage)
@@ -1146,6 +1172,7 @@ func Register() {
 		ClusterQueueWeightedShare,
 		CohortWeightedShare,
 		CohortSubtreeQuota,
+		CohortSubtreeResourceUsage,
 	)
 	if features.Enabled(features.LocalQueueMetrics) {
 		RegisterLQMetrics()