[Metrics] Add average kv cache and waiting queue size metrics for

inference pool

Author: JeffLuoo

pkg/ext-proc/backend/provider.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"go.uber.org/multierr"
+	"inference.networking.x-k8s.io/gateway-api-inference-extension/pkg/ext-proc/metrics"
 	logutil "inference.networking.x-k8s.io/gateway-api-inference-extension/pkg/ext-proc/util/logging"
 	corev1 "k8s.io/api/core/v1"
 	klog "k8s.io/klog/v2"
@@ -89,6 +90,15 @@ func (p *Provider) Init(refreshPodsInterval, refreshMetricsInterval time.Duratio
 		}
 	}()
 
+	// Periodically flush prometheus metrics for inference pool
+	flushPrometheusMetricsInterval := 10 * time.Second
+	go func() {
+		for {
+			time.Sleep(flushPrometheusMetricsInterval)
+			p.flushPrometheusMetricsOnce()
+		}
+	}()
+
 	// Periodically print out the pods and metrics for DEBUGGING.
 	if klog.V(logutil.DEBUG).Enabled() {
 		go func() {
@@ -217,3 +227,40 @@ func (p *Provider) refreshMetricsOnce() error {
 	}
 	return errs
 }
+
+func (p *Provider) flushPrometheusMetricsOnce() {
+	klog.V(logutil.DEFAULT).Infof("Flushing Prometheus Metrics")
+
+	pool, _ := p.datastore.getInferencePool()
+	if pool == nil {
+		// No inference pool or not initialize.
+		return
+	}
+
+	var kvCacheTotal float64
+	var queueTotal int
+
+	podTotalCount := 0
+	pods, err := p.datastore.getPods()
+	if err != nil {
+		klog.V(logutil.DEFAULT).Infof("Couldn't list pods: %v", err)
+		return
+	}
+
+	for _, pod := range pods {
+		if !podIsReady(pod) {
+			continue
+		}
+		podTotalCount++
+		if val, ok := p.podMetrics.Load(pod.Name); ok {
+			pm := val.(*PodMetrics)
+			kvCacheTotal += pm.KVCacheUsagePercent
+			queueTotal += pm.WaitingQueueSize
+		}
+	}
+
+	if podTotalCount != 0 {
+		metrics.RecordInferencePoolAvgKVCache(pool.Name, kvCacheTotal/float64(podTotalCount))
+		metrics.RecordInferencePoolAvgQueueSize(pool.Name, float64(queueTotal/podTotalCount))
+	}
+}

pkg/ext-proc/metrics/metrics.go

@@ -11,9 +11,11 @@ import (
 
 const (
 	InferenceModelComponent = "inference_model"
+	InferencePoolComponent  = "inference_pool"
 )
 
 var (
+	// Inference Model Metrics
 	requestCounter = compbasemetrics.NewCounterVec(
 		&compbasemetrics.CounterOpts{
 			Subsystem:      InferenceModelComponent,
@@ -88,6 +90,27 @@ var (
 		},
 		[]string{"model_name", "target_model_name"},
 	)
+
+	// Inference Pool Metrics
+	inferencePoolAvgKVCache = compbasemetrics.NewGaugeVec(
+		&compbasemetrics.GaugeOpts{
+			Subsystem:      InferencePoolComponent,
+			Name:           "average_kv_cache_utilization",
+			Help:           "The average kv cache utilization for an inference server pool.",
+			StabilityLevel: compbasemetrics.ALPHA,
+		},
+		[]string{"name"},
+	)
+
+	inferencePoolAvgQueueSize = compbasemetrics.NewGaugeVec(
+		&compbasemetrics.GaugeOpts{
+			Subsystem:      InferencePoolComponent,
+			Name:           "average_queue_size",
+			Help:           " The average number of requests pending in the model server queue.",
+			StabilityLevel: compbasemetrics.ALPHA,
+		},
+		[]string{"name"},
+	)
 )
 
 var registerMetrics sync.Once
@@ -101,6 +124,9 @@ func Register() {
 		legacyregistry.MustRegister(responseSizes)
 		legacyregistry.MustRegister(inputTokens)
 		legacyregistry.MustRegister(outputTokens)
+
+		legacyregistry.MustRegister(inferencePoolAvgKVCache)
+		legacyregistry.MustRegister(inferencePoolAvgQueueSize)
 	})
 }
 
@@ -143,3 +169,11 @@ func RecordOutputTokens(modelName, targetModelName string, size int) {
 		outputTokens.WithLabelValues(modelName, targetModelName).Observe(float64(size))
 	}
 }
+
+func RecordInferencePoolAvgKVCache(name string, utilization float64) {
+	inferencePoolAvgKVCache.WithLabelValues(name).Set(utilization)
+}
+
+func RecordInferencePoolAvgQueueSize(name string, utilization float64) {
+	inferencePoolAvgQueueSize.WithLabelValues(name).Set(utilization)
+}

pkg/ext-proc/metrics/metrics_test.go

@@ -15,6 +15,8 @@ const RequestSizesMetric = InferenceModelComponent + "_request_sizes"
 const ResponseSizesMetric = InferenceModelComponent + "_response_sizes"
 const InputTokensMetric = InferenceModelComponent + "_input_tokens"
 const OutputTokensMetric = InferenceModelComponent + "_output_tokens"
+const KVCacheAvgUsageMetric = InferencePoolComponent + "_average_kv_cache_utilization"
+const QueueAvgSizeMetric = InferencePoolComponent + "_average_queue_size"
 
 func TestRecordRequestCounterandSizes(t *testing.T) {
 	type requests struct {
@@ -257,3 +259,53 @@ func TestRecordResponseMetrics(t *testing.T) {
 		})
 	}
 }
+
+func TestInferencePoolMetrics(t *testing.T) {
+	scenarios := []struct {
+		name         string
+		poolName     string
+		kvCacheAvg   float64
+		queueSizeAvg float64
+	}{
+		{
+			name:         "basic test",
+			poolName:     "p1",
+			kvCacheAvg:   0.3,
+			queueSizeAvg: 0.4,
+		},
+	}
+	Register()
+	for _, scenario := range scenarios {
+		t.Run(scenario.name, func(t *testing.T) {
+
+			RecordInferencePoolAvgKVCache(scenario.poolName, scenario.kvCacheAvg)
+			RecordInferencePoolAvgQueueSize(scenario.poolName, scenario.queueSizeAvg)
+
+			wantKVCache, err := os.Open("testdata/kv_cache_avg_metrics")
+			defer func() {
+				if err := wantKVCache.Close(); err != nil {
+					t.Error(err)
+				}
+			}()
+			if err != nil {
+				t.Fatal(err)
+			}
+			if err := testutil.GatherAndCompare(legacyregistry.DefaultGatherer, wantKVCache, KVCacheAvgUsageMetric); err != nil {
+				t.Error(err)
+			}
+
+			wantQueueSize, err := os.Open("testdata/queue_avg_size_metrics")
+			defer func() {
+				if err := wantQueueSize.Close(); err != nil {
+					t.Error(err)
+				}
+			}()
+			if err != nil {
+				t.Fatal(err)
+			}
+			if err := testutil.GatherAndCompare(legacyregistry.DefaultGatherer, wantQueueSize, QueueAvgSizeMetric); err != nil {
+				t.Error(err)
+			}
+		})
+	}
+}

pkg/ext-proc/metrics/testdata/kv_cache_avg_metrics

@@ -0,0 +1,3 @@
+# HELP inference_pool_average_kv_cache_utilization [ALPHA] The average kv cache utilization for an inference server pool.
+# TYPE inference_pool_average_kv_cache_utilization gauge
+inference_pool_average_kv_cache_utilization{name="p1"} 0.3

pkg/ext-proc/metrics/testdata/queue_avg_size_metrics

@@ -0,0 +1,3 @@
+# HELP inference_pool_average_queue_size [ALPHA]  The average number of requests pending in the model server queue.
+# TYPE inference_pool_average_queue_size gauge
+inference_pool_average_queue_size{name="p1"} 0.4