fix: refactor custom inferencepool metric

Signed-off-by: nayihz <[email protected]>

Author: nayihz

cmd/bbr/main.go

@@ -18,6 +18,7 @@ package main
 
 import (
 	"flag"
+	"fmt"
 	"os"
 
 	"github.com/go-logr/logr"
@@ -46,8 +47,9 @@ var (
 		"grpcHealthPort",
 		9005,
 		"The port used for gRPC liveness and readiness probes")
-	metricsAddr = flag.String("metrics-bind-address", ":9090", "The address the metric endpoint binds to.")
-	streaming   = flag.Bool(
+	metricsPort = flag.Int(
+		"metricsPort", 9090, "The metrics port")
+	streaming = flag.Bool(
 		"streaming", false, "Enables streaming support for Envoy full-duplex streaming mode")
 	logVerbosity = flag.Int("v", logging.DEFAULT, "number for the log level verbosity")
 
@@ -88,7 +90,7 @@ func run() error {
 	// - https://pkg.go.dev/sigs.k8s.io/[email protected]/pkg/metrics/server
 	// - https://book.kubebuilder.io/reference/metrics.html
 	metricsServerOptions := metricsserver.Options{
-		BindAddress:    *metricsAddr,
+		BindAddress:    fmt.Sprintf(":%d", *metricsPort),
 		FilterProvider: filters.WithAuthenticationAndAuthorization,
 	}
 	mgr, err := ctrl.NewManager(cfg, ctrl.Options{Metrics: metricsServerOptions})

cmd/epp/main.go

@@ -22,6 +22,7 @@ import (
 	"os"
 
 	"github.com/go-logr/logr"
+	"github.com/prometheus/client_golang/prometheus"
 	uberzap "go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 	"google.golang.org/grpc"
@@ -38,6 +39,7 @@ import (
 	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datastore"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics/collectors"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins/filter"
@@ -58,7 +60,8 @@ var (
 		"grpcHealthPort",
 		9003,
 		"The port used for gRPC liveness and readiness probes")
-	metricsAddr                = flag.String("metrics-bind-address", ":9090", "The address the metric endpoint binds to.")
+	metricsPort = flag.Int(
+		"metricsPort", 9090, "The metrics port")
 	destinationEndpointHintKey = flag.String(
 		"destinationEndpointHintKey",
 		runserver.DefaultDestinationEndpointHintKey,
@@ -154,14 +157,34 @@ func run() error {
 		return err
 	}
 
-	metrics.Register()
+	// Set up mapper for metric scraping.
+	mapping, err := backendmetrics.NewMetricMapping(
+		*totalQueuedRequestsMetric,
+		*kvCacheUsagePercentageMetric,
+		*loraInfoMetric,
+	)
+	if err != nil {
+		setupLog.Error(err, "Failed to create metric mapping from flags.")
+		return err
+	}
+	verifyMetricMapping(*mapping, setupLog)
+
+	pmf := backendmetrics.NewPodMetricsFactory(&backendmetrics.PodMetricsClientImpl{MetricMapping: mapping}, *refreshMetricsInterval)
+	// Setup runner.
+	ctx := ctrl.SetupSignalHandler()
+
+	datastore := datastore.NewDatastore(ctx, pmf)
+
+	customCollectors := []prometheus.Collector{collectors.NewInferencePoolMetricsCollector(datastore)}
+	metrics.Register(customCollectors...)
+	metrics.RecordInferenceExtensionInfo()
 	// Register metrics handler.
 	// Metrics endpoint is enabled in 'config/default/kustomization.yaml'. The Metrics options configure the server.
 	// More info:
 	// - https://pkg.go.dev/sigs.k8s.io/[email protected]/pkg/metrics/server
 	// - https://book.kubebuilder.io/reference/metrics.html
 	metricsServerOptions := metricsserver.Options{
-		BindAddress:    *metricsAddr,
+		BindAddress:    fmt.Sprintf(":%d", *metricsPort),
 		FilterProvider: filters.WithAuthenticationAndAuthorization,
 	}
 
@@ -175,24 +198,6 @@ func run() error {
 		return err
 	}
 
-	// Set up mapper for metric scraping.
-	mapping, err := backendmetrics.NewMetricMapping(
-		*totalQueuedRequestsMetric,
-		*kvCacheUsagePercentageMetric,
-		*loraInfoMetric,
-	)
-	if err != nil {
-		setupLog.Error(err, "Failed to create metric mapping from flags.")
-		return err
-	}
-	verifyMetricMapping(*mapping, setupLog)
-
-	pmf := backendmetrics.NewPodMetricsFactory(&backendmetrics.PodMetricsClientImpl{MetricMapping: mapping}, *refreshMetricsInterval)
-	// Setup runner.
-	ctx := ctrl.SetupSignalHandler()
-
-	datastore := datastore.NewDatastore(ctx, pmf)
-
 	scheduler := scheduling.NewScheduler(datastore)
 	if schedulerV2 == "true" {
 		queueScorerWeight := envutil.GetEnvInt("QUEUE_SCORE_WEIGHT", scorer.DefaultQueueScorerWeight, setupLog)

config/manifests/inferencepool-resources.yaml

@@ -107,6 +107,10 @@ rules:
   - subjectaccessreviews
   verbs:
   - create
+- nonResourceURLs:
+  - "/metrics"
+  verbs:
+  - get
 --- 
 kind: ClusterRoleBinding
 apiVersion: rbac.authorization.k8s.io/v1

pkg/bbr/metrics/metrics.go

@@ -17,12 +17,13 @@ limitations under the License.
 package metrics
 
 import (
-	"fmt"
 	"sync"
 
 	"github.com/prometheus/client_golang/prometheus"
 	compbasemetrics "k8s.io/component-base/metrics"
 	"sigs.k8s.io/controller-runtime/pkg/metrics"
+
+	metricsutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/metrics"
 )
 
 const component = "bbr"
@@ -32,23 +33,23 @@ var (
 		prometheus.CounterOpts{
 			Subsystem: component,
 			Name:      "success_total",
-			Help:      fmt.Sprintf("[%v] %v", compbasemetrics.ALPHA, "Count of successes pulling model name from body and injecting it in the request headers."),
+			Help:      metricsutil.HelpMsgWithStability("Count of successes pulling model name from body and injecting it in the request headers.", compbasemetrics.ALPHA),
 		},
 		[]string{},
 	)
 	modelNotInBodyCounter = prometheus.NewCounterVec(
 		prometheus.CounterOpts{
 			Subsystem: component,
 			Name:      "model_not_in_body_total",
-			Help:      fmt.Sprintf("[%v] %v", compbasemetrics.ALPHA, "Count of times the model was not present in the request body."),
+			Help:      metricsutil.HelpMsgWithStability("Count of times the model was not present in the request body.", compbasemetrics.ALPHA),
 		},
 		[]string{},
 	)
 	modelNotParsedCounter = prometheus.NewCounterVec(
 		prometheus.CounterOpts{
 			Subsystem: component,
 			Name:      "model_not_parsed_total",
-			Help:      fmt.Sprintf("[%v] %v", compbasemetrics.ALPHA, "Count of times the model was in the request body but we could not parse it."),
+			Help:      metricsutil.HelpMsgWithStability("Count of times the model was in the request body but we could not parse it.", compbasemetrics.ALPHA),
 		},
 		[]string{},
 	)

pkg/epp/metrics/collectors/inference_pool.go

@@ -17,47 +17,46 @@ limitations under the License.
 package collectors
 
 import (
-	"k8s.io/component-base/metrics"
+	"github.com/prometheus/client_golang/prometheus"
+	compbasemetrics "k8s.io/component-base/metrics"
+
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datastore"
+	metricsutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/metrics"
 )
 
 var (
-	descInferencePoolPerPodQueueSize = metrics.NewDesc(
+	descInferencePoolPerPodQueueSize = prometheus.NewDesc(
 		"inference_pool_per_pod_queue_size",
-		"The total number of requests pending in the model server queue for each underlying pod.",
+		metricsutil.HelpMsgWithStability("The total number of requests pending in the model server queue for each underlying pod.", compbasemetrics.ALPHA),
 		[]string{
 			"name",
 			"model_server_pod",
 		}, nil,
-		metrics.ALPHA,
-		"",
 	)
 )
 
 type inferencePoolMetricsCollector struct {
-	metrics.BaseStableCollector
-
 	ds datastore.Datastore
 }
 
 // Check if inferencePoolMetricsCollector implements necessary interface
-var _ metrics.StableCollector = &inferencePoolMetricsCollector{}
+var _ prometheus.Collector = &inferencePoolMetricsCollector{}
 
-// NewInferencePoolMetricsCollector implements the metrics.StableCollector interface and
+// NewInferencePoolMetricsCollector implements the prometheus.Collector interface and
 // exposes metrics about inference pool.
-func NewInferencePoolMetricsCollector(ds datastore.Datastore) metrics.StableCollector {
+func NewInferencePoolMetricsCollector(ds datastore.Datastore) prometheus.Collector {
 	return &inferencePoolMetricsCollector{
 		ds: ds,
 	}
 }
 
-// DescribeWithStability implements the metrics.StableCollector interface.
-func (c *inferencePoolMetricsCollector) DescribeWithStability(ch chan<- *metrics.Desc) {
+// DescribeWithStability implements the prometheus.Collector interface.
+func (c *inferencePoolMetricsCollector) Describe(ch chan<- *prometheus.Desc) {
 	ch <- descInferencePoolPerPodQueueSize
 }
 
-// CollectWithStability implements the metrics.StableCollector interface.
-func (c *inferencePoolMetricsCollector) CollectWithStability(ch chan<- metrics.Metric) {
+// CollectWithStability implements the prometheus.Collector interface.
+func (c *inferencePoolMetricsCollector) Collect(ch chan<- prometheus.Metric) {
 	pool, err := c.ds.PoolGet()
 	if err != nil {
 		return
@@ -69,9 +68,9 @@ func (c *inferencePoolMetricsCollector) CollectWithStability(ch chan<- metrics.M
 	}
 
 	for _, pod := range podMetrics {
-		ch <- metrics.NewLazyConstMetric(
+		ch <- prometheus.MustNewConstMetric(
 			descInferencePoolPerPodQueueSize,
-			metrics.GaugeValue,
+			prometheus.GaugeValue,
 			float64(pod.GetMetrics().WaitingQueueSize),
 			pool.Name,
 			pod.GetPod().NamespacedName.Name,

pkg/epp/metrics/collectors/inference_pool_test.go

@@ -55,7 +55,7 @@ func TestNoMetricsCollected(t *testing.T) {
 		ds: datastore,
 	}
 
-	if err := testutil.CustomCollectAndCompare(collector, strings.NewReader(""), ""); err != nil {
+	if err := testutil.CollectAndCompare(collector, strings.NewReader(""), ""); err != nil {
 		t.Fatal(err)
 	}
 }
@@ -90,7 +90,7 @@ func TestMetricsCollected(t *testing.T) {
 	collector := &inferencePoolMetricsCollector{
 		ds: ds,
 	}
-	err := testutil.CustomCollectAndCompare(collector, strings.NewReader(`
+	err := testutil.CollectAndCompare(collector, strings.NewReader(`
 		# HELP inference_pool_per_pod_queue_size [ALPHA] The total number of requests pending in the model server queue for each underlying pod.
 		# TYPE inference_pool_per_pod_queue_size gauge
 		inference_pool_per_pod_queue_size{model_server_pod="pod1",name="test-pool"} 100