Natasha/add queue metrics (#320)

* added metrics to queue controller

* added reset queue metrics when needed

* added tests + fixes

* small typo fix

* little fix for tests

* added some tests to improve coverage + fix the tests found

* fix tests

* CR fixes - remove unnecessary labels

* CR fixes - DefaultMetricsNamespace const in common consts

Author: natasharomm

cmd/queuecontroller/app/app.go

@@ -21,8 +21,10 @@ import (
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/healthz"
 	"sigs.k8s.io/controller-runtime/pkg/log/zap"
+	metricsserver "sigs.k8s.io/controller-runtime/pkg/metrics/server"
 
 	"github.com/NVIDIA/KAI-scheduler/pkg/queuecontroller/controllers"
+	"github.com/NVIDIA/KAI-scheduler/pkg/queuecontroller/metrics"
 	// +kubebuilder:scaffold:imports
 )
 
@@ -47,9 +49,14 @@ func Run() error {
 
 	initLogger()
 
+	metrics.InitMetrics(opts.MetricsNamespace, opts.QueueLabelToMetricLabel.Get(), opts.QueueLabelToDefaultMetricValue.Get())
+
 	var err error
 	options := ctrl.Options{
-		Scheme:           scheme,
+		Scheme: scheme,
+		Metrics: metricsserver.Options{
+			BindAddress: opts.MetricsAddress,
+		},
 		LeaderElection:   opts.EnableLeaderElection,
 		LeaderElectionID: "42ece193.run.ai",
 	}

cmd/queuecontroller/app/options.go

@@ -5,14 +5,31 @@ package app
 
 import (
 	"flag"
+
+	"github.com/NVIDIA/KAI-scheduler/pkg/common/constants"
+	kaiflags "github.com/NVIDIA/KAI-scheduler/pkg/common/flags"
+)
+
+const (
+	defaultSchedulingQueueLabelKey = "kai.scheduler/queue"
+	defaultMetricsAddress          = ":8080"
 )
 
 type Options struct {
 	EnableLeaderElection    bool
 	SchedulingQueueLabelKey string
+
+	MetricsAddress                 string
+	MetricsNamespace               string
+	QueueLabelToMetricLabel        kaiflags.StringMapFlag
+	QueueLabelToDefaultMetricValue kaiflags.StringMapFlag
 }
 
 func (o *Options) AddFlags(fs *flag.FlagSet) {
 	fs.BoolVar(&o.EnableLeaderElection, "leader-elect", false, "Enable leader election for controller manager. Enabling this will ensure there is only one active controller manager.")
-	fs.StringVar(&o.SchedulingQueueLabelKey, "queue-label-key", "kai.scheduler/queue", "Scheduling queue label key name")
+	fs.StringVar(&o.SchedulingQueueLabelKey, "queue-label-key", defaultSchedulingQueueLabelKey, "Scheduling queue label key name.")
+	fs.StringVar(&o.MetricsAddress, "metrics-listen-address", defaultMetricsAddress, "The address the metrics endpoint binds to.")
+	fs.StringVar(&o.MetricsNamespace, "metrics-namespace", constants.DefaultMetricsNamespace, "Metrics namespace.")
+	fs.Var(&o.QueueLabelToMetricLabel, "queue-label-to-metric-label", "Map of queue label keys to metric label keys, e.g. 'foo=bar,baz=qux'.")
+	fs.Var(&o.QueueLabelToDefaultMetricValue, "queue-label-to-default-metric-value", "Map of queue label keys to default metric values, in case the label doesn't exist on the queue, e.g. 'foo=1,baz=0'.")
 }

cmd/scheduler/app/options/options.go

@@ -25,13 +25,13 @@ import (
 	"github.com/spf13/pflag"
 	utilfeature "k8s.io/apiserver/pkg/util/feature"
 
+	"github.com/NVIDIA/KAI-scheduler/pkg/common/constants"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/log"
 )
 
 const (
 	defaultSchedulerName               = "kai-scheduler"
 	defaultResourceReservationAppLabel = "kai-resource-reservation"
-	defaultMetricsNamespace            = "kai"
 	defaultNamespace                   = "kai-scheduler"
 	defaultSchedulerPeriod             = time.Second
 	defaultStalenessGracePeriod        = 60 * time.Second
@@ -99,7 +99,7 @@ func (s *ServerOption) AddFlags(fs *pflag.FlagSet) {
 	// kai-scheduler will ignore pods with scheduler names other than specified with the option
 	fs.StringVar(&s.SchedulerName, "scheduler-name", defaultSchedulerName, "The scheduler name in pod spec that handled by this scheduler")
 	fs.StringVar(&s.Namspace, "namespace", defaultNamespace, "Scheduler service namespace")
-	fs.StringVar(&s.MetricsNamespace, "metrics-namespace", defaultMetricsNamespace, "Metrics namespace")
+	fs.StringVar(&s.MetricsNamespace, "metrics-namespace", constants.DefaultMetricsNamespace, "Metrics namespace")
 	fs.StringVar(&s.ResourceReservationAppLabel, "resource-reservation-app-label", defaultResourceReservationAppLabel, "App label value of resource reservation pods")
 	fs.BoolVar(&s.RestrictSchedulingNodes, "restrict-node-scheduling", false, "kai-scheduler will allocate jobs only to restricted nodes")
 	fs.StringVar(&s.NodePoolLabelKey, "nodepool-label-key", defaultNodePoolLabelKey, "The label key by which to filter scheduling nodepool")

cmd/scheduler/app/options/options_test.go

@@ -24,6 +24,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/NVIDIA/KAI-scheduler/pkg/common/constants"
+
 	"github.com/spf13/pflag"
 	"k8s.io/apimachinery/pkg/util/diff"
 	utilfeature "k8s.io/apiserver/pkg/util/feature"
@@ -45,7 +47,7 @@ func TestAddFlags(t *testing.T) {
 	expected := &ServerOption{
 		SchedulerName:                     defaultSchedulerName,
 		Namspace:                          defaultNamespace,
-		MetricsNamespace:                  defaultMetricsNamespace,
+		MetricsNamespace:                  constants.DefaultMetricsNamespace,
 		ResourceReservationAppLabel:       defaultResourceReservationAppLabel,
 		SchedulePeriod:                    5 * time.Minute,
 		PrintVersion:                      true,

pkg/common/constants/constants.go

@@ -10,6 +10,7 @@ const (
 	UnlimitedResourceQuantity = float64(-1)
 	DefaultQueuePriority      = 100
 	DefaultNodePoolName       = "default"
+	DefaultMetricsNamespace   = "kai"
 
 	// Pod Groups
 	PodGrouperWarning   = "PodGrouperWarning"

pkg/common/flags/flags_test.go

@@ -0,0 +1,69 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package flags
+
+import (
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+	"golang.org/x/exp/slices"
+	"strings"
+	"testing"
+)
+
+func TestCommonFlags(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "common flags tests")
+}
+
+var _ = Describe("StringMapFlag", func() {
+	It("parses an empty string as an empty map", func() {
+		var m StringMapFlag
+		Expect(m.Set("")).To(Succeed())
+		Expect(m.Get()).To(BeEmpty())
+	})
+
+	It("parses a single key=value pair", func() {
+		var m StringMapFlag
+		Expect(m.Set("foo=bar")).To(Succeed())
+		Expect(m.Get()).To(HaveKeyWithValue("foo", "bar"))
+		Expect(len(m.Get())).To(Equal(1))
+	})
+
+	It("parses multiple key=value pairs", func() {
+		var m StringMapFlag
+		Expect(m.Set("foo=bar,baz=qux")).To(Succeed())
+		Expect(m.Get()).To(HaveKeyWithValue("foo", "bar"))
+		Expect(m.Get()).To(HaveKeyWithValue("baz", "qux"))
+		Expect(len(m.Get())).To(Equal(2))
+	})
+
+	It("overwrites duplicate keys with the last value", func() {
+		var m StringMapFlag
+		Expect(m.Set("foo=bar,foo=baz")).To(Succeed())
+		Expect(m.Get()).To(HaveKeyWithValue("foo", "baz"))
+		Expect(len(m.Get())).To(Equal(1))
+	})
+
+	It("returns an error for invalid input", func() {
+		var m StringMapFlag
+		err := m.Set("foo,bar=baz")
+		Expect(err).To(HaveOccurred())
+	})
+
+	It("String() returns the correct string representation", func() {
+		var m StringMapFlag
+		err := m.Set("foo=bar,baz=qux")
+		Expect(err).ToNot(HaveOccurred())
+		str := m.String()
+		// Accept either order
+		Expect([]string{str, reversePairs(str)}).To(ContainElement("foo=bar,baz=qux"))
+	})
+})
+
+// Helper to reverse the order of pairs in a comma-separated string
+func reversePairs(s string) string {
+	pairs := strings.Split(s, ",")
+	slices.Reverse(pairs)
+	return strings.Join(pairs, ",")
+}

pkg/common/flags/string_map.go

@@ -0,0 +1,46 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package flags
+
+import (
+	"fmt"
+	"strings"
+)
+
+type StringMapFlag map[string]string
+
+func (m *StringMapFlag) Get() map[string]string {
+	if *m == nil {
+		return make(map[string]string)
+	}
+	return *m
+}
+
+func (m *StringMapFlag) String() string {
+	// Convert map to string: key1=val1,key2=val2
+	var pairs []string
+	for k, v := range *m {
+		pairs = append(pairs, fmt.Sprintf("%s=%s", k, v))
+	}
+	return strings.Join(pairs, ",")
+}
+
+func (m *StringMapFlag) Set(value string) error {
+	// Convert string to map: key1=val1,key2=val2
+	if *m == nil {
+		*m = make(map[string]string)
+	}
+	if value == "" {
+		return nil
+	}
+	pairs := strings.Split(value, ",")
+	for _, pair := range pairs {
+		kv := strings.SplitN(pair, "=", 2)
+		if len(kv) != 2 {
+			return fmt.Errorf("invalid map item: %q", pair)
+		}
+		(*m)[kv[0]] = kv[1]
+	}
+	return nil
+}

pkg/queuecontroller/controllers/queue_controller.go

@@ -20,6 +20,7 @@ import (
 	"github.com/NVIDIA/KAI-scheduler/pkg/queuecontroller/common"
 	"github.com/NVIDIA/KAI-scheduler/pkg/queuecontroller/controllers/childqueues_updater"
 	"github.com/NVIDIA/KAI-scheduler/pkg/queuecontroller/controllers/resource_updater"
+	"github.com/NVIDIA/KAI-scheduler/pkg/queuecontroller/metrics"
 )
 
 // QueueReconciler reconciles a Queue object
@@ -54,7 +55,12 @@ func (r *QueueReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl
 	queue := &v2.Queue{}
 	err := r.Get(ctx, req.NamespacedName, queue)
 	if err != nil {
-		return ctrl.Result{}, client.IgnoreNotFound(err)
+		ignoreNotFoundErr := client.IgnoreNotFound(err)
+		if ignoreNotFoundErr == nil {
+			// If the queue is not found, reset its metrics
+			metrics.ResetQueueMetrics(req.Name)
+		}
+		return ctrl.Result{}, ignoreNotFoundErr
 	}
 	originalQueue := queue.DeepCopy()
 
@@ -73,6 +79,8 @@ func (r *QueueReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl
 		return ctrl.Result{}, fmt.Errorf("failed to patch status for queue %s, error: %v", queue.Name, err)
 	}
 
+	metrics.SetQueueMetrics(queue)
+
 	return ctrl.Result{}, err
 }
 

pkg/queuecontroller/controllers/suite_test.go

@@ -38,8 +38,12 @@ import (
 	logf "sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/controller-runtime/pkg/log/zap"
 
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/testutil"
+
 	v2 "github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2"
 	"github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2alpha2"
+	"github.com/NVIDIA/KAI-scheduler/pkg/queuecontroller/metrics"
 )
 
 const (
@@ -106,6 +110,11 @@ var _ = Describe("QueueController", Ordered, func() {
 	BeforeAll(func() {
 		ctx, cancel = context.WithCancel(context.Background())
 
+		metrics.InitMetrics("testns",
+			map[string]string{"priority": "queue_priority", "some-other-label": "some_other_label"},
+			map[string]string{"priority": "normal"},
+		)
+
 		var err error
 		mgr, err = ctrl.NewManager(cfg, ctrl.Options{
 			Scheme: scheme.Scheme,
@@ -134,6 +143,12 @@ var _ = Describe("QueueController", Ordered, func() {
 	})
 
 	Context("When managing child queues", func() {
+		AfterAll(func() {
+			deleteQueue(ctx, k8sClient, "child-queue-1")
+			deleteQueue(ctx, k8sClient, "child-queue-2")
+			deleteQueue(ctx, k8sClient, "parent-queue")
+		})
+
 		It("Should update parent queue's childQueues field", func() {
 			parentQueue := &v2.Queue{
 				ObjectMeta: metav1.ObjectMeta{
@@ -175,6 +190,10 @@ var _ = Describe("QueueController", Ordered, func() {
 	})
 
 	Context("When managing pod groups", func() {
+		AfterAll(func() {
+			deleteQueue(ctx, k8sClient, "resource-queue")
+		})
+
 		It("Should update queue status with pod group resources", func() {
 			queue := &v2.Queue{
 				ObjectMeta: metav1.ObjectMeta{
@@ -283,4 +302,67 @@ var _ = Describe("QueueController", Ordered, func() {
 			}, timeout, interval).Should(BeTrue())
 		})
 	})
+
+	Context("When setting queue metrics", func() {
+		AfterAll(func() {
+			deleteQueue(ctx, k8sClient, "test-queue")
+		})
+
+		It("should set metrics for a queue with resources", func() {
+			queue := &v2.Queue{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:   "test-queue",
+					Labels: map[string]string{"priority": "high"},
+				},
+				Spec: v2.QueueSpec{
+					Resources: &v2.QueueResources{
+						GPU:    v2.QueueResource{Quota: 2},
+						CPU:    v2.QueueResource{Quota: 2000},
+						Memory: v2.QueueResource{Quota: 4},
+					},
+				},
+			}
+			Expect(k8sClient.Create(ctx, queue)).Should(Succeed())
+
+			labels := []string{"test-queue", "high", ""}
+
+			Eventually(func(q gomega.Gomega) {
+				expectMetricValue(q, metrics.GetQueueInfoMetric(), labels, 1)
+				expectMetricValue(q, metrics.GetQueueDeservedGPUsMetric(), labels, 2)
+				expectMetricValue(q, metrics.GetQueueQuotaCPUMetric(), labels, 2)
+				expectMetricValue(q, metrics.GetQueueQuotaMemoryMetric(), labels, 4000000)
+			}, timeout, interval).Should(Succeed())
+
+			// delete the queue and expect metrics to be deleted
+			Expect(k8sClient.Delete(ctx, queue)).Should(Succeed())
+
+			Eventually(func(q gomega.Gomega) {
+				gathered := testutil.CollectAndCount(metrics.GetQueueInfoMetric())
+				q.Expect(gathered).To(Equal(0))
+				gathered = testutil.CollectAndCount(metrics.GetQueueDeservedGPUsMetric())
+				q.Expect(gathered).To(Equal(0))
+				gathered = testutil.CollectAndCount(metrics.GetQueueQuotaCPUMetric())
+				q.Expect(gathered).To(Equal(0))
+				gathered = testutil.CollectAndCount(metrics.GetQueueQuotaMemoryMetric())
+				q.Expect(gathered).To(Equal(0))
+			}, timeout, interval).Should(Succeed())
+		})
+	})
 })
+
+func expectMetricValue(q gomega.Gomega, gauge *prometheus.GaugeVec, labels []string, expected float64) {
+	metricGauge, err := gauge.GetMetricWithLabelValues(labels...)
+	q.Expect(err).To(BeNil())
+	q.Expect(metricGauge).ToNot(BeNil())
+	q.Expect(testutil.ToFloat64(metricGauge)).To(BeEquivalentTo(expected))
+}
+
+func deleteQueue(ctx context.Context, k8sClient client.Client, queueName string) {
+	queueObj := &v2.Queue{}
+	err := k8sClient.Get(ctx, types.NamespacedName{Name: queueName}, queueObj)
+	if err != nil {
+		return
+	}
+
+	_ = k8sClient.Delete(ctx, queueObj)
+}