implementing pod cache flushing logic

Author: kfswain

pkg/ext-proc/backend/datastore.go

@@ -1,15 +1,18 @@
 package backend
 
 import (
+	"context"
 	"errors"
 	"math/rand"
+	"strconv"
 	"sync"
 
 	"inference.networking.x-k8s.io/gateway-api-inference-extension/api/v1alpha1"
 	logutil "inference.networking.x-k8s.io/gateway-api-inference-extension/pkg/ext-proc/util/logging"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/klog/v2"
+	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
 func NewK8sDataStore(options ...K8sDatastoreOption) *K8sDatastore {
@@ -114,12 +117,36 @@ func IsCritical(model *v1alpha1.InferenceModel) bool {
 }
 
 func (ds *K8sDatastore) LabelsMatch(podLabels map[string]string) bool {
-	poolSelector := labels.SelectorFromSet(whimsicalWondersOfAliasTypes(ds.inferencePool.Spec.Selector))
+	poolSelector := selectorFromInferencePoolSelector(ds.inferencePool.Spec.Selector)
 	podSet := labels.Set(podLabels)
 	return poolSelector.Matches(podSet)
 }
 
-func whimsicalWondersOfAliasTypes(labels map[v1alpha1.LabelKey]v1alpha1.LabelValue) map[string]string {
+func (ds *K8sDatastore) flushPodsAndRefetch(ctx context.Context, ctrlClient client.Client, newServerPool *v1alpha1.InferencePool) {
+	podList := &corev1.PodList{}
+	if err := ctrlClient.List(ctx, podList, selectorFromInferencePoolSelector(newServerPool.Spec.Selector).(client.MatchingLabelsSelector)); err != nil {
+		klog.Error(err, "error listing clients")
+	}
+	podMap := sync.Map{}
+
+	for _, k8sPod := range podList.Items {
+		pod := Pod{
+			Name:    k8sPod.Name,
+			Address: k8sPod.Status.PodIP + ":" + strconv.Itoa(int(newServerPool.Spec.TargetPortNumber)),
+		}
+		podMap.Store(pod, true)
+	}
+
+	// Clear is thread-safe, so if there are any in-flight accesses this should prevent odd data behavior. Then we replace after.
+	ds.pods.Clear()
+	ds.pods = &podMap
+}
+
+func selectorFromInferencePoolSelector(selector map[v1alpha1.LabelKey]v1alpha1.LabelValue) labels.Selector {
+	return labels.SelectorFromSet(stripLabelKeyAliasFromLabelMap(selector))
+}
+
+func stripLabelKeyAliasFromLabelMap(labels map[v1alpha1.LabelKey]v1alpha1.LabelValue) map[string]string {
 	outMap := make(map[string]string)
 	for k, v := range labels {
 		outMap[string(k)] = string(v)

pkg/ext-proc/backend/inferencepool_reconciler.go

@@ -2,6 +2,7 @@ package backend
 
 import (
 	"context"
+	"reflect"
 
 	"inference.networking.x-k8s.io/gateway-api-inference-extension/api/v1alpha1"
 	"k8s.io/apimachinery/pkg/runtime"
@@ -34,7 +35,9 @@ func (c *InferencePoolReconciler) Reconcile(ctx context.Context, req ctrl.Reques
 		klog.Error(err, "unable to get InferencePool")
 		return ctrl.Result{}, err
 	}
-
+	if !reflect.DeepEqual(serverPool.Spec.Selector, c.Datastore.inferencePool.Spec.Selector) {
+		c.Datastore.flushPodsAndRefetch(ctx, c.Client, serverPool)
+	}
 	c.updateDatastore(serverPool)
 
 	return ctrl.Result{}, nil