feat: rename PodGetAll to PodGetAllWithFreshMetrics and add unit test

Author: nayihz

conformance/testing-epp/sheduler_test.go

@@ -116,7 +116,7 @@ type fakeDataStore struct {
 	pods []*backendmetrics.FakePodMetrics
 }
 
-func (fds *fakeDataStore) PodGetAll() []backendmetrics.PodMetrics {
+func (fds *fakeDataStore) PodGetAllWithFreshMetrics() []backendmetrics.PodMetrics {
 	pm := make([]backendmetrics.PodMetrics, 0, len(fds.pods))
 	for _, pod := range fds.pods {
 		pm = append(pm, pod)

pkg/epp/backend/metrics/logger.go

@@ -38,8 +38,8 @@ const (
 type Datastore interface {
 	PoolGet() (*v1alpha2.InferencePool, error)
 	// PodMetrics operations
-	// PodGetAll returns all pods and metrics, including fresh and stale.
-	PodGetAll() []PodMetrics
+	// PodGetAllWithFreshMetrics returns all pods and metrics, only including fresh.
+	PodGetAllWithFreshMetrics() []PodMetrics
 	PodList(func(PodMetrics) bool) []PodMetrics
 }
 
@@ -97,7 +97,7 @@ func refreshPrometheusMetrics(logger logr.Logger, datastore Datastore) {
 	var kvCacheTotal float64
 	var queueTotal int
 
-	podMetrics := datastore.PodGetAll()
+	podMetrics := datastore.PodGetAllWithFreshMetrics()
 	logger.V(logutil.TRACE).Info("Refreshing Prometheus Metrics", "ReadyPods", len(podMetrics))
 	if len(podMetrics) == 0 {
 		return

pkg/epp/backend/metrics/pod_metrics_test.go

@@ -90,7 +90,7 @@ type fakeDataStore struct{}
 func (f *fakeDataStore) PoolGet() (*v1alpha2.InferencePool, error) {
 	return &v1alpha2.InferencePool{Spec: v1alpha2.InferencePoolSpec{TargetPortNumber: 8000}}, nil
 }
-func (f *fakeDataStore) PodGetAll() []PodMetrics {
+func (f *fakeDataStore) PodGetAllWithFreshMetrics() []PodMetrics {
 	// Not implemented.
 	return nil
 }

pkg/epp/controller/inferencepool_reconciler_test.go

@@ -165,7 +165,7 @@ func diffStore(datastore datastore.Datastore, params diffStoreParams) string {
 		return "pool:" + diff
 	}
 
-	// Default wantPods if not set because PodGetAll returns an empty slice when empty.
+	// Default wantPods if not set because PodGetAllWithFreshMetrics returns an empty slice when empty.
 	if params.wantPods == nil {
 		params.wantPods = []string{}
 	}

pkg/epp/datastore/datastore.go

@@ -62,8 +62,8 @@ type Datastore interface {
 	ModelGetAll() []*v1alpha2.InferenceModel
 
 	// PodMetrics operations
-	// PodGetAll returns all pods and metrics, including fresh and stale.
-	PodGetAll() []backendmetrics.PodMetrics
+	// PodGetAllWithFreshMetrics returns all pods and metrics, only including fresh.
+	PodGetAllWithFreshMetrics() []backendmetrics.PodMetrics
 	// PodList lists pods matching the given predicate.
 	PodList(predicate func(backendmetrics.PodMetrics) bool) []backendmetrics.PodMetrics
 	PodUpdateOrAddIfNotExist(pod *corev1.Pod) bool
@@ -246,7 +246,7 @@ func (ds *datastore) ModelGetAll() []*v1alpha2.InferenceModel {
 
 // /// Pods/endpoints APIs ///
 
-func (ds *datastore) PodGetAll() []backendmetrics.PodMetrics {
+func (ds *datastore) PodGetAllWithFreshMetrics() []backendmetrics.PodMetrics {
 	return ds.PodList(func(pm backendmetrics.PodMetrics) bool {
 		return time.Since(pm.GetMetrics().UpdateTime) <= ds.pmf.MetricsStalenessThreshold
 	})

pkg/epp/datastore/datastore_test.go

@@ -263,6 +263,12 @@ var (
 		},
 		WaitingModels: map[string]int{},
 	}
+	pod3 = &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name: "pod2",
+		},
+	}
+
 	pod1NamespacedName = types.NamespacedName{Name: pod1.Name, Namespace: pod1.Namespace}
 	pod2NamespacedName = types.NamespacedName{Name: pod2.Name, Namespace: pod2.Namespace}
 	inferencePool      = &v1alpha2.InferencePool{
@@ -312,17 +318,19 @@ func TestMetrics(t *testing.T) {
 				},
 			},
 			storePods: []*corev1.Pod{pod1, pod2},
-			want: []*backendmetrics.MetricsState{
-				pod1Metrics,
-				// Failed to fetch pod2 metrics so it remains the default values.
-				{
-					ActiveModels:        map[string]int{},
-					WaitingModels:       map[string]int{},
-					WaitingQueueSize:    0,
-					KVCacheUsagePercent: 0,
-					MaxActiveModels:     0,
+			want:      []*backendmetrics.MetricsState{pod1Metrics},
+		},
+		{
+			name: "Filter stale metrics",
+			pmc: &backendmetrics.FakePodMetricsClient{
+				Res: map[types.NamespacedName]*backendmetrics.MetricsState{
+					pod1NamespacedName: pod1Metrics,
+					pod2NamespacedName: pod2Metrics,
 				},
 			},
+			storePods: []*corev1.Pod{pod1, pod2, pod3},
+			want:      []*backendmetrics.MetricsState{pod1Metrics, pod2Metrics}, // pod3 metrics were stale and should not be included.
+
 		},
 	}
 
@@ -336,16 +344,15 @@ func TestMetrics(t *testing.T) {
 			fakeClient := fake.NewClientBuilder().
 				WithScheme(scheme).
 				Build()
-			pmf := backendmetrics.NewPodMetricsFactory(test.pmc, time.Millisecond, config.DefaultMetricsStalenessThreshold)
+			pmf := backendmetrics.NewPodMetricsFactory(test.pmc, time.Millisecond, time.Second*10)
 			ds := NewDatastore(ctx, pmf)
 			_ = ds.PoolSet(ctx, fakeClient, inferencePool)
 			for _, pod := range test.storePods {
 				ds.PodUpdateOrAddIfNotExist(pod)
 			}
+			time.Sleep(1 * time.Second) // Give some time for the metrics to be fetched.
 			assert.EventuallyWithT(t, func(t *assert.CollectT) {
-				got := ds.PodList(func(backendmetrics.PodMetrics) bool {
-					return true
-				})
+				got := ds.PodGetAllWithFreshMetrics()
 				metrics := []*backendmetrics.MetricsState{}
 				for _, one := range got {
 					metrics = append(metrics, one.GetMetrics())

pkg/epp/metrics/collectors/inference_pool.go

@@ -62,7 +62,7 @@ func (c *inferencePoolMetricsCollector) Collect(ch chan<- prometheus.Metric) {
 		return
 	}
 
-	podMetrics := c.ds.PodGetAll()
+	podMetrics := c.ds.PodGetAllWithFreshMetrics()
 	if len(podMetrics) == 0 {
 		return
 	}

pkg/epp/saturationdetector/saturationdetector.go

@@ -56,7 +56,7 @@ type Config struct {
 
 // Datastore provides an interface to access backend pod metrics.
 type Datastore interface {
-	PodGetAll() []backendmetrics.PodMetrics
+	PodGetAllWithFreshMetrics() []backendmetrics.PodMetrics
 }
 
 // Detector determines system saturation based on metrics from the Datastore.
@@ -100,7 +100,7 @@ func NewDetector(config *Config, datastore Datastore, logger logr.Logger) *Detec
 // (no capacity).
 func (d *Detector) IsSaturated(ctx context.Context) bool {
 	logger := log.FromContext(ctx).WithName(loggerName)
-	allPodsMetrics := d.datastore.PodGetAll()
+	allPodsMetrics := d.datastore.PodGetAllWithFreshMetrics()
 	if len(allPodsMetrics) == 0 {
 		logger.V(logutil.VERBOSE).Info("No pods found in datastore; system is considered SATURATED (no capacity).")
 		// If there are no pods, there is no capacity to serve requests.
@@ -121,13 +121,6 @@ func (d *Detector) IsSaturated(ctx context.Context) bool {
 			continue
 		}
 
-		// Check for metric staleness
-		// if time.Since(metrics.UpdateTime) > d.config.MetricsStalenessThreshold {
-		// 	logger.V(logutil.TRACE).Info("Pod metrics are stale, considered as not having good capacity",
-		// 		"pod", podNn, "updateTime", metrics.UpdateTime, "stalenessThreshold", d.config.MetricsStalenessThreshold)
-		// 	continue
-		// }
-
 		// Check queue depth
 		if metrics.WaitingQueueSize > d.config.QueueDepthThreshold {
 			logger.V(logutil.TRACE).Info("Pod WaitingQueueSize is above threshold, considered as not having good capacity",

pkg/epp/saturationdetector/saturationdetector_test.go

@@ -37,8 +37,8 @@ type mockDatastore struct {
 	pods []*backendmetrics.FakePodMetrics
 }
 
-// PodGetAll returns all pod metrics from the fake datastore.
-func (fds *mockDatastore) PodGetAll() []backendmetrics.PodMetrics {
+// PodGetAllWithFreshMetrics returns all pod metrics from the fake datastore.
+func (fds *mockDatastore) PodGetAllWithFreshMetrics() []backendmetrics.PodMetrics {
 	pm := make([]backendmetrics.PodMetrics, 0, len(fds.pods))
 	for _, pod := range fds.pods {
 		pm = append(pm, pod)

pkg/epp/scheduling/scheduler.go

@@ -35,7 +35,7 @@ import (
 )
 
 type Datastore interface {
-	PodGetAll() []backendmetrics.PodMetrics
+	PodGetAllWithFreshMetrics() []backendmetrics.PodMetrics
 }
 
 // NewScheduler returns a new scheduler with default scheduler plugins configuration.
@@ -106,7 +106,7 @@ func (s *Scheduler) Schedule(ctx context.Context, request *types.LLMRequest) (*t
 	// Snapshot pod metrics from the datastore to:
 	// 1. Reduce concurrent access to the datastore.
 	// 2. Ensure consistent data during the scheduling operation of a request between all scheduling cycles.
-	podsSnapshot := types.ToSchedulerPodMetrics(s.datastore.PodGetAll())
+	podsSnapshot := types.ToSchedulerPodMetrics(s.datastore.PodGetAllWithFreshMetrics())
 	loggerDebug.Info(fmt.Sprintf("Scheduling a request, Metrics: %+v", podsSnapshot))
 
 	profileRunResults := map[string]*types.ProfileRunResult{}