Addressed comments

Author: ahg-g

pkg/ext-proc/backend/datastore.go

@@ -61,11 +61,10 @@ func WithPodListerFactory(factory PodListerFactory) K8sDatastoreOption {
 type PodLister struct {
 	Lister         listersv1.PodLister
 	sharedInformer informers.SharedInformerFactory
-	ctx            context.Context
 }
 
-func (l *PodLister) list(selector labels.Selector) ([]*corev1.Pod, error) {
-	return l.Lister.List(selector)
+func (l *PodLister) listEverything() ([]*corev1.Pod, error) {
+	return l.Lister.List(labels.Everything())
 
 }
 
@@ -97,8 +96,9 @@ func (ds *K8sDatastore) setInferencePool(pool *v1alpha1.InferencePool) {
 		// Create a new informer with the new selector.
 		ds.podLister = ds.podListerFactory(ds.inferencePool)
 		if ds.podLister != nil && ds.podLister.sharedInformer != nil {
-			ds.podLister.sharedInformer.Start(ds.podLister.ctx.Done())
-			ds.podLister.sharedInformer.WaitForCacheSync(ds.podLister.ctx.Done())
+			ctx := context.Background()
+			ds.podLister.sharedInformer.Start(ctx.Done())
+			ds.podLister.sharedInformer.WaitForCacheSync(ctx.Done())
 		}
 	}
 }
@@ -123,7 +123,7 @@ func (ds *K8sDatastore) createPodLister(pool *v1alpha1.InferencePool) *PodLister
 	}
 
 	newPodInformer := func(cs clientset.Interface, resyncPeriod time.Duration) cache.SharedIndexInformer {
-		informer := informersv1.NewFilteredPodInformer(cs, pool.Namespace, 0, nil, func(options *metav1.ListOptions) {
+		informer := informersv1.NewFilteredPodInformer(cs, pool.Namespace, resyncPeriod, cache.Indexers{}, func(options *metav1.ListOptions) {
 			options.LabelSelector = labels.SelectorFromSet(selectorSet).String()
 		})
 		err := informer.SetTransform(func(obj interface{}) (interface{}, error) {
@@ -140,30 +140,30 @@ func (ds *K8sDatastore) createPodLister(pool *v1alpha1.InferencePool) *PodLister
 		}
 		return informer
 	}
-	sharedInformer := informers.NewSharedInformerFactory(ds.client, 0)
+	// 0 means we disable resyncing, it is not really useful to resync every hour (the controller-runtime default),
+	// if things go wrong in the watch, no one will wait for an hour for things to get fixed.
+	// As precedence, kube-scheduler also disables this since it is expensive to list all pods from the api-server regularly.
+	resyncPeriod := time.Duration(0)
+	sharedInformer := informers.NewSharedInformerFactory(ds.client, resyncPeriod)
 	sharedInformer.InformerFor(&v1.Pod{}, newPodInformer)
 
 	return &PodLister{
 		Lister:         sharedInformer.Core().V1().Pods().Lister(),
 		sharedInformer: sharedInformer,
-		ctx:            context.Background(),
 	}
 }
 
-func (ds *K8sDatastore) getPods() []*corev1.Pod {
+func (ds *K8sDatastore) getPods() ([]*corev1.Pod, error) {
 	ds.poolMu.RLock()
 	defer ds.poolMu.RUnlock()
-	if ds.podLister == nil {
-		klog.V(logutil.DEFAULT).Info("InferencePool not yet initialized")
-		return []*corev1.Pod{}
+	if !ds.HasSynced() {
+		return nil, errors.New("InferencePool is not initialized in datastore")
 	}
-
-	pods, err := ds.podLister.list(labels.Everything())
+	pods, err := ds.podLister.listEverything()
 	if err != nil {
-		klog.Errorf("Failed to list pods for pool %s/%s: %v", ds.inferencePool.Namespace, ds.inferencePool.Name, err)
-		return []*corev1.Pod{}
+		return nil, err
 	}
-	return pods
+	return pods, nil
 }
 
 func (s *K8sDatastore) FetchModelData(modelName string) (returnModel *v1alpha1.InferenceModel) {

pkg/ext-proc/backend/provider.go

@@ -104,14 +104,18 @@ func (p *Provider) Init(refreshPodsInterval, refreshMetricsInterval time.Duratio
 // refreshPodsOnce lists pods and updates keys in the podMetrics map.
 // Note this function doesn't update the PodMetrics value, it's done separately.
 func (p *Provider) refreshPodsOnce() {
-	pool, err := p.datastore.getInferencePool()
+	pods, err := p.datastore.getPods()
 	if err != nil {
-		klog.V(logutil.DEFAULT).Infof("Pool not ready: %v", err)
+		klog.V(logutil.DEFAULT).Infof("Couldn't list pods: %v", err)
 		p.podMetrics.Clear()
 		return
 	}
-
-	pods := p.datastore.getPods()
+	pool, _ := p.datastore.getInferencePool()
+	// revision is used to track which entries we need to remove in the next iteration that removes
+	// metrics for pods that don't exist anymore. Otherwise we have to build a map of the listed pods,
+	// which is not efficient. Revision can be any random id as long as it is different from the last
+	// refresh, so it should be very reliable (as reliable as the probability of randomly picking two
+	// different numbers from range 0 - maxInt).
 	revision := rand.Int()
 	ready := 0
 	for _, pod := range pods {

pkg/ext-proc/server/runserver.go

@@ -34,7 +34,7 @@ type ExtProcServerRunner struct {
 // Default values for CLI flags in main
 const (
 	DefaultGrpcPort               = 9002                             // default for --grpcPort
-	DefaultTargetEndpointKey      = "x-gateway-destination-endpoint" // default for --targetPodHeader
+	DefaultTargetEndpointKey      = "x-gateway-destination-endpoint" // default for --targetEndpointKey
 	DefaultPoolName               = ""                               // required but no default
 	DefaultPoolNamespace          = "default"                        // default for --poolNamespace
 	DefaultRefreshPodsInterval    = 10 * time.Second                 // default for --refreshPodsInterval
@@ -130,8 +130,7 @@ func (r *ExtProcServerRunner) StartManager() {
 	}
 	// Start the controller manager. Blocking and will return when shutdown is complete.
 	klog.Infof("Starting controller manager")
-	mgr := r.Manager
-	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
+	if err := r.Manager.Start(ctrl.SetupSignalHandler()); err != nil {
 		klog.Fatalf("Error starting controller manager: %v", err)
 	}
 	klog.Info("Controller manager shutting down")

pkg/ext-proc/test/utils.go

@@ -24,7 +24,7 @@ func StartExtProc(port int, refreshPodsInterval, refreshMetricsInterval time.Dur
 		pms[pod.Pod.Name] = pod
 	}
 	pmc := &backend.FakePodMetricsClient{Res: pms}
-	pp := backend.NewProvider(pmc, backend.NewK8sDataStore()) // backend.WithPods(pods)
+	pp := backend.NewProvider(pmc, backend.NewK8sDataStore())
 	if err := pp.Init(refreshPodsInterval, refreshMetricsInterval); err != nil {
 		klog.Fatalf("failed to initialize: %v", err)
 	}

test/integration/hermetic_test.go

@@ -206,44 +206,38 @@ func TestKubeInferenceModelRequest(t *testing.T) {
 		},
 	}
 
-	pods := []*backend.PodMetrics{
+	metrics := []*backend.Metrics{
 		{
-			Metrics: backend.Metrics{
-				WaitingQueueSize:    0,
-				KVCacheUsagePercent: 0.2,
-				ActiveModels: map[string]int{
-					"foo": 1,
-					"bar": 1,
-				},
+			WaitingQueueSize:    0,
+			KVCacheUsagePercent: 0.2,
+			ActiveModels: map[string]int{
+				"foo": 1,
+				"bar": 1,
 			},
 		},
 		{
-			Metrics: backend.Metrics{
-				WaitingQueueSize:    0,
-				KVCacheUsagePercent: 0.1,
-				ActiveModels: map[string]int{
-					"foo":            1,
-					"sql-lora-1fdg2": 1,
-				},
+			WaitingQueueSize:    0,
+			KVCacheUsagePercent: 0.1,
+			ActiveModels: map[string]int{
+				"foo":            1,
+				"sql-lora-1fdg2": 1,
 			},
 		},
 		{
-			Metrics: backend.Metrics{
-				WaitingQueueSize:    10,
-				KVCacheUsagePercent: 0.2,
-				ActiveModels: map[string]int{
-					"foo": 1,
-				},
+			WaitingQueueSize:    10,
+			KVCacheUsagePercent: 0.2,
+			ActiveModels: map[string]int{
+				"foo": 1,
 			},
 		},
 	}
 
 	// Set up global k8sclient and extproc server runner with test environment config
-	BeforeSuit(pods)
+	podMetrics := BeforeSuit(metrics)
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			client, cleanup := setUpHermeticServer(t, pods)
+			client, cleanup := setUpHermeticServer(t, podMetrics)
 			t.Cleanup(cleanup)
 			want := &extProcPb.ProcessingResponse{
 				Response: &extProcPb.ProcessingResponse_RequestBody{
@@ -374,7 +368,7 @@ func setUpHermeticServer(t *testing.T, pods []*backend.PodMetrics) (client extPr
 }
 
 // Sets up a test environment and returns the runner struct
-func BeforeSuit(metrics []*backend.PodMetrics) {
+func BeforeSuit(metrics []*backend.Metrics) []*backend.PodMetrics {
 	// Set up mock k8s API Client
 	testEnv = &envtest.Environment{
 		CRDDirectoryPaths:     []string{filepath.Join("..", "..", "config", "crd", "bases")},
@@ -396,17 +390,21 @@ func BeforeSuit(metrics []*backend.PodMetrics) {
 		log.Fatalf("No error, but returned kubernetes client is nil, cfg: %v", cfg)
 	}
 
+	var podMetrics []*backend.PodMetrics
 	fakeLister := &testingutil.FakePodLister{
 		PodsList: []*corev1.Pod{},
 	}
 	for i, m := range metrics {
 		podName := "pod-" + strconv.Itoa(i)
 		pod := testingutil.MakePod(podName).SetReady().SetPodIP(podName).Obj()
 		fakeLister.PodsList = append(fakeLister.PodsList, pod)
-		m.Pod = backend.Pod{
-			Name:    pod.Name,
-			Address: pod.Status.PodIP + ":8000",
-		}
+		podMetrics = append(podMetrics, &backend.PodMetrics{
+			Pod: backend.Pod{
+				Name:    pod.Name,
+				Address: pod.Status.PodIP + ":8000",
+			},
+			Metrics: *m,
+		})
 	}
 
 	serverRunner = runserver.NewDefaultExtProcServerRunner()
@@ -431,6 +429,7 @@ func BeforeSuit(metrics []*backend.PodMetrics) {
 
 	// Wait the reconcilers to populate the datastore.
 	time.Sleep(5 * time.Second)
+	return podMetrics
 }
 
 func sendRequest(t *testing.T, client extProcPb.ExternalProcessor_ProcessClient, req *extProcPb.ProcessingRequest) (*extProcPb.ProcessingResponse, error) {