Address comments

Author: liu-cong

pkg/epp/backend/metrics/fake.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"sync"
 
+	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/types"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/gateway-api-inference-extension/api/v1alpha2"
@@ -39,8 +40,8 @@ func (fpm *FakePodMetrics) GetPod() *Pod {
 func (fpm *FakePodMetrics) GetMetrics() *Metrics {
 	return fpm.Metrics
 }
-func (fpm *FakePodMetrics) UpdatePod(pod *Pod) {
-	fpm.Pod = pod
+func (fpm *FakePodMetrics) UpdatePod(pod *corev1.Pod) {
+	fpm.Pod = toInternalPod(pod)
 }
 func (fpm *FakePodMetrics) StopRefreshLoop() {} // noop
 

pkg/epp/backend/metrics/logger.go

@@ -20,7 +20,10 @@ import (
 	"context"
 	"time"
 
+	"github.com/go-logr/logr"
 	"sigs.k8s.io/controller-runtime/pkg/log"
+	"sigs.k8s.io/gateway-api-inference-extension/api/v1alpha2"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
@@ -31,12 +34,30 @@ const (
 )
 
 type Datastore interface {
+	PoolGet() (*v1alpha2.InferencePool, error)
+	// PodMetrics operations
+	// PodGetAll returns all pods and metrics, including fresh and stale.
+	PodGetAll() []PodMetrics
 	PodList(func(PodMetrics) bool) []PodMetrics
 }
 
-func PrintMetricsForDebugging(ctx context.Context, datastore Datastore) {
+func LogMetricsPeriodically(ctx context.Context, datastore Datastore, refreshPrometheusMetricsInterval time.Duration) {
 	logger := log.FromContext(ctx)
 
+	// Periodically flush prometheus metrics for inference pool
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				logger.V(logutil.DEFAULT).Info("Shutting down prometheus metrics thread")
+				return
+			default:
+				time.Sleep(refreshPrometheusMetricsInterval)
+				flushPrometheusMetricsOnce(logger, datastore)
+			}
+		}
+	}()
+
 	// Periodically print out the pods and metrics for DEBUGGING.
 	if logger := logger.V(logutil.DEBUG); logger.Enabled() {
 		go func() {
@@ -59,3 +80,30 @@ func PrintMetricsForDebugging(ctx context.Context, datastore Datastore) {
 		}()
 	}
 }
+
+func flushPrometheusMetricsOnce(logger logr.Logger, datastore Datastore) {
+	pool, err := datastore.PoolGet()
+	if err != nil {
+		// No inference pool or not initialize.
+		logger.V(logutil.VERBOSE).Info("pool is not initialized, skipping flushing metrics")
+		return
+	}
+
+	var kvCacheTotal float64
+	var queueTotal int
+
+	podMetrics := datastore.PodGetAll()
+	logger.V(logutil.VERBOSE).Info("Flushing Prometheus Metrics", "ReadyPods", len(podMetrics))
+	if len(podMetrics) == 0 {
+		return
+	}
+
+	for _, pod := range podMetrics {
+		kvCacheTotal += pod.GetMetrics().KVCacheUsagePercent
+		queueTotal += pod.GetMetrics().WaitingQueueSize
+	}
+
+	podTotalCount := len(podMetrics)
+	metrics.RecordInferencePoolAvgKVCache(pool.Name, kvCacheTotal/float64(podTotalCount))
+	metrics.RecordInferencePoolAvgQueueSize(pool.Name, float64(queueTotal/podTotalCount))
+}

pkg/epp/backend/metrics/pod_metrics.go

@@ -24,6 +24,8 @@ import (
 	"unsafe"
 
 	"github.com/go-logr/logr"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/types"
 
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
@@ -58,8 +60,18 @@ func (pm *podMetrics) GetMetrics() *Metrics {
 	return (*Metrics)(atomic.LoadPointer(&pm.metrics))
 }
 
-func (pm *podMetrics) UpdatePod(pod *Pod) {
-	atomic.StorePointer(&pm.pod, unsafe.Pointer(pod))
+func (pm *podMetrics) UpdatePod(in *corev1.Pod) {
+	atomic.StorePointer(&pm.pod, unsafe.Pointer(toInternalPod(in)))
+}
+
+func toInternalPod(in *corev1.Pod) *Pod {
+	return &Pod{
+		NamespacedName: types.NamespacedName{
+			Name:      in.Name,
+			Namespace: in.Namespace,
+		},
+		Address: in.Status.PodIP,
+	}
 }
 
 // start starts a goroutine exactly once to periodically update metrics. The goroutine will be

pkg/epp/backend/metrics/pod_metrics_test.go

@@ -23,16 +23,19 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/stretchr/testify/assert"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
 )
 
 var (
-	pod1 = &Pod{
-		NamespacedName: types.NamespacedName{
-			Name: "pod1",
+	pod1 = &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "pod1",
+			Namespace: "default",
 		},
 	}
-	metrics = &Metrics{
+	initial = &Metrics{
 		WaitingQueueSize:    0,
 		KVCacheUsagePercent: 0.2,
 		MaxActiveModels:     2,
@@ -59,22 +62,19 @@ func TestMetricsRefresh(t *testing.T) {
 	// The refresher is initialized with empty metrics.
 	pm := pmf.NewPodMetrics(ctx, pod1, 8000)
 
+	namespacedName := types.NamespacedName{Name: pod1.Name, Namespace: pod1.Namespace}
 	// Use SetRes to simulate an update of metrics from the pod.
 	// Verify that the metrics are updated.
-	pmc.SetRes(map[types.NamespacedName]*Metrics{
-		pod1.NamespacedName: metrics,
-	})
+	pmc.SetRes(map[types.NamespacedName]*Metrics{namespacedName: initial})
 	condition := func(collect *assert.CollectT) {
-		assert.True(collect, cmp.Equal(pm.GetMetrics(), metrics, cmpopts.IgnoreFields(Metrics{}, "UpdateTime")))
+		assert.True(collect, cmp.Equal(pm.GetMetrics(), initial, cmpopts.IgnoreFields(Metrics{}, "UpdateTime")))
 	}
 	assert.EventuallyWithT(t, condition, time.Second, time.Millisecond)
 
 	// Stop the loop, and simulate metric update again, this time the PodMetrics won't get the
 	// new update.
 	pm.StopRefreshLoop()
-	pmc.SetRes(map[types.NamespacedName]*Metrics{
-		pod1.NamespacedName: updated,
-	})
+	pmc.SetRes(map[types.NamespacedName]*Metrics{namespacedName: updated})
 	// Still expect the same condition (no metrics update).
 	assert.EventuallyWithT(t, condition, time.Second, time.Millisecond)
 }

pkg/epp/backend/metrics/types.go

@@ -24,6 +24,7 @@ import (
 	"time"
 	"unsafe"
 
+	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/types"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 )
@@ -40,9 +41,9 @@ type PodMetricsFactory struct {
 	refreshMetricsInterval time.Duration
 }
 
-func (f *PodMetricsFactory) NewPodMetrics(parentCtx context.Context, pod *Pod, port int32) PodMetrics {
+func (f *PodMetricsFactory) NewPodMetrics(parentCtx context.Context, in *corev1.Pod, port int32) PodMetrics {
 	pm := &podMetrics{
-		pod:        unsafe.Pointer(pod),
+		pod:        unsafe.Pointer(toInternalPod(in)),
 		metrics:    unsafe.Pointer(newMetrics()),
 		pmc:        f.pmc,
 		targetPort: port,
@@ -59,7 +60,7 @@ func (f *PodMetricsFactory) NewPodMetrics(parentCtx context.Context, pod *Pod, p
 type PodMetrics interface {
 	GetPod() *Pod
 	GetMetrics() *Metrics
-	UpdatePod(*Pod)
+	UpdatePod(*corev1.Pod)
 	StopRefreshLoop()
 }