fix metric scrape port not updated when inference pool target port updated

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

Author: Kuromesi

pkg/epp/backend/provider.go

@@ -94,7 +94,7 @@ func (p *Provider) Init(ctx context.Context, refreshMetricsInterval, refreshProm
 					return
 				default:
 					time.Sleep(5 * time.Second)
-					logger.Info("Current Pods and metrics gathered", "metrics", p.datastore.PodGetAll())
+					logger.Info("Current Pods and metrics gathered", "metrics", p.datastore.PodList())
 				}
 			}
 		}()
@@ -162,7 +162,7 @@ func (p *Provider) flushPrometheusMetricsOnce(logger logr.Logger) {
 	var kvCacheTotal float64
 	var queueTotal int
 
-	podMetrics := p.datastore.PodGetAll()
+	podMetrics := p.datastore.PodList()
 	if len(podMetrics) == 0 {
 		return
 	}

pkg/epp/backend/provider_test.go

@@ -116,6 +116,7 @@ func TestProvider(t *testing.T) {
 				{
 					Pod: datastore.Pod{NamespacedName: pod2.NamespacedName},
 					Metrics: datastore.Metrics{
+						ActiveModels:        map[string]int{},
 						WaitingQueueSize:    0,
 						KVCacheUsagePercent: 0,
 						MaxActiveModels:     0,
@@ -133,7 +134,7 @@ func TestProvider(t *testing.T) {
 			defer cancel()
 			_ = p.Init(ctx, time.Millisecond, time.Millisecond)
 			assert.EventuallyWithT(t, func(t *assert.CollectT) {
-				metrics := ds.PodGetAll()
+				metrics := ds.PodList()
 				diff := cmp.Diff(test.want, metrics, cmpopts.SortSlices(func(a, b *datastore.PodMetrics) bool {
 					return a.String() < b.String()
 				}))

pkg/epp/controller/inferencepool_reconciler.go

@@ -85,6 +85,13 @@ func (c *InferencePoolReconciler) updateDatastore(ctx context.Context, newPool *
 		//    the ones that may have existed already to the store.
 		c.Datastore.PodResyncAll(ctx, c.Client)
 	}
+	if oldPool != nil && newPool.Spec.TargetPortNumber != oldPool.Spec.TargetPortNumber {
+		logger.V(logutil.DEFAULT).Info("Updating scrape port of endpoints", "targetPortNumber", newPool.Spec.TargetPortNumber)
+		c.Datastore.PodUpdate(func(pod *datastore.PodMetrics) *datastore.PodMetrics {
+			pod.ScrapePort = newPool.Spec.TargetPortNumber
+			return pod
+		})
+	}
 }
 
 func (c *InferencePoolReconciler) SetupWithManager(mgr ctrl.Manager) error {

pkg/epp/controller/inferencepool_reconciler_test.go

@@ -134,6 +134,13 @@ func TestInferencePoolReconciler(t *testing.T) {
 		t.Errorf("Unexpected diff (+got/-want): %s", diff)
 	}
 
+	pods := datastore.PodList()
+	for _, pm := range pods {
+		if pm.ScrapePort != newPool1.Spec.TargetPortNumber {
+			t.Errorf("Unexpected pod scrape port: %d", pm.ScrapePort)
+		}
+	}
+
 	// Step 4: delete the pool to trigger a datastore clear
 	if err := fakeClient.Get(ctx, req.NamespacedName, newPool1); err != nil {
 		t.Errorf("Unexpected pool get error: %v", err)
@@ -166,7 +173,7 @@ func diffStore(datastore datastore.Datastore, params diffStoreParams) string {
 		params.wantPods = []string{}
 	}
 	gotPods := []string{}
-	for _, pm := range datastore.PodGetAll() {
+	for _, pm := range datastore.PodList() {
 		gotPods = append(gotPods, pm.NamespacedName.Name)
 	}
 	if diff := cmp.Diff(params.wantPods, gotPods, cmpopts.SortSlices(func(a, b string) bool { return a < b })); diff != "" {

pkg/epp/datastore/datastore.go

@@ -62,14 +62,19 @@ type Datastore interface {
 	PodGet(namespacedName types.NamespacedName) (*PodMetrics, bool)
 	PodDelete(namespacedName types.NamespacedName)
 	PodResyncAll(ctx context.Context, ctrlClient client.Client)
-	PodGetAll() []*PodMetrics
+	PodList(predicates ...PodListPredicate) []*PodMetrics
 	PodDeleteAll() // This is only for testing.
 	PodRange(f func(key, value any) bool)
+	PodUpdate(f PodUpdateFunc, predicates ...PodListPredicate)
 
 	// Clears the store state, happens when the pool gets deleted.
 	Clear()
 }
 
+type PodUpdateFunc func(pod *PodMetrics) *PodMetrics
+
+type PodListPredicate func(pod *PodMetrics) bool
+
 func NewDatastore() Datastore {
 	store := &datastore{
 		poolAndModelsMu: sync.RWMutex{},
@@ -231,8 +236,9 @@ func (ds *datastore) ModelGetAll() []*v1alpha2.InferenceModel {
 // /// Pods/endpoints APIs ///
 func (ds *datastore) PodUpdateMetricsIfExist(namespacedName types.NamespacedName, m *Metrics) bool {
 	if val, ok := ds.pods.Load(namespacedName); ok {
-		existing := val.(*PodMetrics)
+		existing := val.(*PodMetrics).Clone()
 		existing.Metrics = *m
+		ds.pods.Store(namespacedName, existing)
 		return true
 	}
 	return false
@@ -241,23 +247,35 @@ func (ds *datastore) PodUpdateMetricsIfExist(namespacedName types.NamespacedName
 func (ds *datastore) PodGet(namespacedName types.NamespacedName) (*PodMetrics, bool) {
 	val, ok := ds.pods.Load(namespacedName)
 	if ok {
-		return val.(*PodMetrics), true
+		return val.(*PodMetrics).Clone(), true
 	}
 	return nil, false
 }
 
-func (ds *datastore) PodGetAll() []*PodMetrics {
+// PodList returns a list of PodMetrics that match the given predicates.
+// If no predicates are provided, all PodMetrics will be returned.
+func (ds *datastore) PodList(predicates ...PodListPredicate) []*PodMetrics {
 	res := []*PodMetrics{}
-	fn := func(k, v any) bool {
-		res = append(res, v.(*PodMetrics))
+	fn := func(_, v any) bool {
+		pod := v.(*PodMetrics)
+		for _, p := range predicates {
+			if !p(pod) {
+				continue
+			}
+		}
+		res = append(res, pod)
 		return true
 	}
-	ds.pods.Range(fn)
+	ds.PodRange(fn)
 	return res
 }
 
 func (ds *datastore) PodRange(f func(key, value any) bool) {
-	ds.pods.Range(f)
+	withClone := func(key, value any) bool {
+		pod := value.(*PodMetrics)
+		return f(key, pod.Clone())
+	}
+	ds.pods.Range(withClone)
 }
 
 func (ds *datastore) PodDelete(namespacedName types.NamespacedName) {
@@ -280,15 +298,18 @@ func (ds *datastore) PodUpdateOrAddIfNotExist(pod *corev1.Pod) bool {
 			ActiveModels: make(map[string]int),
 		},
 	}
-	existing, ok := ds.pods.Load(new.NamespacedName)
-	if !ok {
-		ds.pods.Store(new.NamespacedName, new)
-		return true
-	}
+	_, ok := ds.pods.Load(new.NamespacedName)
+	ds.pods.Store(new.NamespacedName, new)
+	return ok
+}
 
-	// Update pod properties if anything changed.
-	existing.(*PodMetrics).Pod = new.Pod
-	return false
+func (ds *datastore) PodUpdate(f PodUpdateFunc, predicates ...PodListPredicate) {
+	pods := ds.PodList(predicates...)
+	for _, pod := range pods {
+		if new := f(pod); new != nil {
+			ds.pods.Store(new.NamespacedName, new)
+		}
+	}
 }
 
 func (ds *datastore) PodResyncAll(ctx context.Context, ctrlClient client.Client) {

pkg/epp/datastore/types.go

@@ -66,6 +66,7 @@ func (pm *PodMetrics) Clone() *PodMetrics {
 		},
 		Metrics: Metrics{
 			ActiveModels:            cm,
+			MaxActiveModels:         pm.MaxActiveModels,
 			RunningQueueSize:        pm.RunningQueueSize,
 			WaitingQueueSize:        pm.WaitingQueueSize,
 			KVCacheUsagePercent:     pm.KVCacheUsagePercent,

pkg/epp/scheduling/scheduler.go

@@ -124,7 +124,7 @@ type Scheduler struct {
 // Schedule finds the target pod based on metrics and the requested lora adapter.
 func (s *Scheduler) Schedule(ctx context.Context, req *LLMRequest) (targetPod datastore.PodMetrics, err error) {
 	logger := log.FromContext(ctx).WithValues("request", req)
-	podMetrics := s.datastore.PodGetAll()
+	podMetrics := s.datastore.PodList()
 	logger.V(logutil.VERBOSE).Info("Scheduling a request", "metrics", podMetrics)
 	pods, err := s.filter.Filter(logger, req, podMetrics)
 	if err != nil || len(pods) == 0 {