Fixed the provider test and covered the pool deletion events.

Author: ahg-g

Diff for: pkg/ext-proc/backend/datastore.go

@@ -35,10 +35,13 @@ type Datastore interface {
 	PodUpdateMetricsIfExist(pm *PodMetrics)
 	PodGet(namespacedName types.NamespacedName) (*PodMetrics, bool)
 	PodDelete(namespacedName types.NamespacedName)
-	PodFlush(ctx context.Context, ctrlClient client.Client)
+	PodFlushAll(ctx context.Context, ctrlClient client.Client)
 	PodGetAll() []*PodMetrics
-	PodRange(f func(key, value any) bool)
 	PodDeleteAll() // This is only for testing.
+	PodRange(f func(key, value any) bool)
+
+	// Clears the store state, happens when the pool gets deleted.
+	Clear()
 }
 
 func NewDatastore() Datastore {
@@ -59,6 +62,14 @@ type datastore struct {
 	pods *sync.Map
 }
 
+func (ds *datastore) Clear() {
+	ds.poolMu.Lock()
+	defer ds.poolMu.Unlock()
+	ds.pool = nil
+	ds.models.Clear()
+	ds.pods.Clear()
+}
+
 // /// InferencePool APIs ///
 func (ds *datastore) PoolSet(pool *v1alpha1.InferencePool) {
 	ds.poolMu.Lock()
@@ -158,7 +169,7 @@ func (ds *datastore) PodAddIfNotExist(pod *corev1.Pod) bool {
 	return false
 }
 
-func (ds *datastore) PodFlush(ctx context.Context, ctrlClient client.Client) {
+func (ds *datastore) PodFlushAll(ctx context.Context, ctrlClient client.Client) {
 	// Pool must exist to invoke this function.
 	pool, _ := ds.PoolGet()
 	podList := &corev1.PodList{}

Diff for: pkg/ext-proc/backend/inferencepool_reconciler.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"reflect"
 
+	"k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/client-go/tools/record"
@@ -36,12 +37,19 @@ func (c *InferencePoolReconciler) Reconcile(ctx context.Context, req ctrl.Reques
 	loggerDefault.Info("Reconciling InferencePool", "name", req.NamespacedName)
 
 	serverPool := &v1alpha1.InferencePool{}
+
 	if err := c.Get(ctx, req.NamespacedName, serverPool); err != nil {
+		if errors.IsNotFound(err) {
+			loggerDefault.Info("InferencePool not found. Clearing the datastore", "name", req.NamespacedName)
+			c.Datastore.Clear()
+			return ctrl.Result{}, nil
+		}
 		loggerDefault.Error(err, "Unable to get InferencePool", "name", req.NamespacedName)
 		return ctrl.Result{}, err
-
-		// TODO: Handle InferencePool deletions. Need to flush the datastore.
-		// TODO: Handle port updates, podMetrics should not be storing that as part of the address.
+	} else if !serverPool.DeletionTimestamp.IsZero() {
+		loggerDefault.Info("InferencePool is marked for deletion. Clearing the datastore", "name", req.NamespacedName)
+		c.Datastore.Clear()
+		return ctrl.Result{}, nil
 	}
 
 	c.updateDatastore(ctx, serverPool)
@@ -55,7 +63,7 @@ func (c *InferencePoolReconciler) updateDatastore(ctx context.Context, newPool *
 	c.Datastore.PoolSet(newPool)
 	if oldPool == nil || !reflect.DeepEqual(newPool.Spec.Selector, oldPool.Spec.Selector) {
 		logger.V(logutil.DEFAULT).Info("Updating inference pool endpoints", "target", klog.KMetadata(&newPool.ObjectMeta))
-		c.Datastore.PodFlush(ctx, c.Client)
+		c.Datastore.PodFlushAll(ctx, c.Client)
 	}
 }
 

Diff for: pkg/ext-proc/backend/inferencepool_reconciler_test.go

@@ -26,7 +26,10 @@ var (
 			Name:      "pool1",
 			Namespace: "pool1-ns",
 		},
-		Spec: v1alpha1.InferencePoolSpec{Selector: selector_v1},
+		Spec: v1alpha1.InferencePoolSpec{
+			Selector:         selector_v1,
+			TargetPortNumber: 8080,
+		},
 	}
 	pool2 = &v1alpha1.InferencePool{
 		ObjectMeta: metav1.ObjectMeta{
@@ -48,6 +51,9 @@ var (
 )
 
 func TestReconcile_InferencePoolReconciler(t *testing.T) {
+	// The best practice is to use table-driven tests, however in this scaenario it seems
+	// more logical to do a single test with steps that depend on each other.
+
 	// Set up the scheme.
 	scheme := runtime.NewScheme()
 	_ = clientgoscheme.AddToScheme(scheme)
@@ -63,7 +69,7 @@ func TestReconcile_InferencePoolReconciler(t *testing.T) {
 		WithObjects(initialObjects...).
 		Build()
 
-		// Create a request for the existing resource.
+	// Create a request for the existing resource.
 	namespacedName := types.NamespacedName{Name: pool1.Name, Namespace: pool1.Namespace}
 	req := ctrl.Request{NamespacedName: namespacedName}
 	ctx := context.Background()
@@ -103,6 +109,35 @@ func TestReconcile_InferencePoolReconciler(t *testing.T) {
 	if diff := diffPool(datastore, newPool1, []string{"pod5"}); diff != "" {
 		t.Errorf("Unexpected diff (+got/-want): %s", diff)
 	}
+
+	// Step 4: update the pool port
+	if err := fakeClient.Get(ctx, req.NamespacedName, newPool1); err != nil {
+		t.Errorf("Unexpected pool get error: %v", err)
+	}
+	newPool1.Spec.TargetPortNumber = 9090
+	if err := fakeClient.Update(ctx, newPool1, &client.UpdateOptions{}); err != nil {
+		t.Errorf("Unexpected pool update error: %v", err)
+	}
+	if _, err := inferencePoolReconciler.Reconcile(ctx, req); err != nil {
+		t.Errorf("Unexpected InferencePool reconcile error: %v", err)
+	}
+	if diff := diffPool(datastore, newPool1, []string{"pod5"}); diff != "" {
+		t.Errorf("Unexpected diff (+got/-want): %s", diff)
+	}
+
+	// Step 5: 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)
+	}
+	if err := fakeClient.Delete(ctx, newPool1, &client.DeleteOptions{}); err != nil {
+		t.Errorf("Unexpected pool delete error: %v", err)
+	}
+	if _, err := inferencePoolReconciler.Reconcile(ctx, req); err != nil {
+		t.Errorf("Unexpected InferencePool reconcile error: %v", err)
+	}
+	if diff := diffPool(datastore, nil, []string{}); diff != "" {
+		t.Errorf("Unexpected diff (+got/-want): %s", diff)
+	}
 }
 
 func diffPool(datastore Datastore, wantPool *v1alpha1.InferencePool, wantPods []string) string {

Diff for: pkg/ext-proc/backend/provider_test.go

@@ -1,13 +1,16 @@
 package backend
 
 import (
+	"context"
+	"errors"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
+	"github.com/stretchr/testify/assert"
 	"k8s.io/apimachinery/pkg/types"
-	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/ext-proc/util/logging"
 )
 
 var (
@@ -42,20 +45,15 @@ var (
 )
 
 func TestProvider(t *testing.T) {
-	logger := logutil.NewTestLogger()
-
 	tests := []struct {
 		name      string
 		pmc       PodMetricsClient
 		datastore Datastore
 		want      []*PodMetrics
 	}{
 		{
-			name: "Fetch metrics error",
+			name: "Probing metrics success",
 			pmc: &FakePodMetricsClient{
-				// Err: map[string]error{
-				// 	pod2.Name: errors.New("injected error"),
-				// },
 				Res: map[types.NamespacedName]*PodMetrics{
 					pod1.NamespacedName: pod1,
 					pod2.NamespacedName: pod2,
@@ -67,42 +65,72 @@ func TestProvider(t *testing.T) {
 			want: []*PodMetrics{
 				pod1,
 				pod2,
-				// // Failed to fetch pod2 metrics so it remains the default values.
-				// {
-				// 	Name: "pod2",
-				// 	Metrics: Metrics{
-				// 		WaitingQueueSize:    0,
-				// 		KVCacheUsagePercent: 0,
-				// 		MaxActiveModels:     0,
-				// 		ActiveModels:        map[string]int{},
-				// 	},
-				// },
+			},
+		},
+		{
+			name: "Only pods in the datastore are probed",
+			pmc: &FakePodMetricsClient{
+				Res: map[types.NamespacedName]*PodMetrics{
+					pod1.NamespacedName: pod1,
+					pod2.NamespacedName: pod2,
+				},
+			},
+			datastore: &datastore{
+				pods: populateMap(pod1),
+			},
+			want: []*PodMetrics{
+				pod1,
+			},
+		},
+		{
+			name: "Probing metrics error",
+			pmc: &FakePodMetricsClient{
+				Err: map[types.NamespacedName]error{
+					pod2.NamespacedName: errors.New("injected error"),
+				},
+				Res: map[types.NamespacedName]*PodMetrics{
+					pod1.NamespacedName: pod1,
+				},
+			},
+			datastore: &datastore{
+				pods: populateMap(pod1, pod2),
+			},
+			want: []*PodMetrics{
+				pod1,
+				// Failed to fetch pod2 metrics so it remains the default values.
+				{
+					NamespacedName: pod2.NamespacedName,
+					Metrics: Metrics{
+						WaitingQueueSize:    0,
+						KVCacheUsagePercent: 0,
+						MaxActiveModels:     0,
+					},
+				},
 			},
 		},
 	}
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			p := NewProvider(test.pmc, test.datastore)
-			// if err := p.refreshMetricsOnce(logger); err != nil {
-			// 	t.Fatalf("Unexpected error: %v", err)
-			// }
-			_ = p.refreshMetricsOnce(logger)
-			metrics := test.datastore.PodGetAll()
-			lessFunc := func(a, b *PodMetrics) bool {
-				return a.String() < b.String()
-			}
-			if diff := cmp.Diff(test.want, metrics, cmpopts.SortSlices(lessFunc)); diff != "" {
-				t.Errorf("Unexpected output (-want +got): %v", diff)
-			}
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			_ = p.Init(ctx, time.Millisecond, time.Millisecond)
+			assert.EventuallyWithT(t, func(t *assert.CollectT) {
+				metrics := test.datastore.PodGetAll()
+				diff := cmp.Diff(test.want, metrics, cmpopts.SortSlices(func(a, b *PodMetrics) bool {
+					return a.String() < b.String()
+				}))
+				assert.Equal(t, "", diff, "Unexpected diff (+got/-want)")
+			}, 5*time.Second, time.Millisecond)
 		})
 	}
 }
 
 func populateMap(pods ...*PodMetrics) *sync.Map {
 	newMap := &sync.Map{}
 	for _, pod := range pods {
-		newMap.Store(pod.NamespacedName, pod)
+		newMap.Store(pod.NamespacedName, &PodMetrics{NamespacedName: pod.NamespacedName})
 	}
 	return newMap
 }

Diff for: test/integration/hermetic_test.go

@@ -17,6 +17,7 @@ import (
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	envoyTypePb "github.com/envoyproxy/go-control-plane/envoy/type/v3"
 	"github.com/google/go-cmp/cmp"
+	"github.com/stretchr/testify/assert"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/protobuf/testing/protocmp"
@@ -320,7 +321,7 @@ func TestKubeInferenceModelRequest(t *testing.T) {
 	}
 
 	// Set up global k8sclient and extproc server runner with test environment config
-	cleanup := BeforeSuit()
+	cleanup := BeforeSuit(t)
 	defer cleanup()
 
 	for _, test := range tests {
@@ -409,7 +410,7 @@ func setUpHermeticServer(podMetrics []*backend.PodMetrics) (client extProcPb.Ext
 }
 
 // Sets up a test environment and returns the runner struct
-func BeforeSuit() func() {
+func BeforeSuit(t *testing.T) func() {
 	// Set up mock k8s API Client
 	testEnv = &envtest.Environment{
 		CRDDirectoryPaths:     []string{filepath.Join("..", "..", "config", "crd", "bases")},
@@ -488,26 +489,16 @@ func BeforeSuit() func() {
 		}
 	}
 
-	if !blockUntilPoolSyncs(serverRunner.Datastore) {
-		logutil.Fatal(logger, nil, "Timeout waiting for the pool and models to sync")
-	}
+	assert.EventuallyWithT(t, func(t *assert.CollectT) {
+		synced := serverRunner.Datastore.PoolHasSynced() && serverRunner.Datastore.ModelGet("my-model") != nil
+		assert.True(t, synced, "Timeout waiting for the pool and models to sync")
+	}, 10*time.Second, 10*time.Millisecond)
 
 	return func() {
 		_ = testEnv.Stop()
 	}
 }
 
-func blockUntilPoolSyncs(datastore backend.Datastore) bool {
-	// We really need to move those tests to gingo so we can use Eventually...
-	for i := 1; i < 10; i++ {
-		if datastore.PoolHasSynced() && datastore.ModelGet("my-model") != nil {
-			return true
-		}
-		time.Sleep(1 * time.Second)
-	}
-	return false
-}
-
 func sendRequest(t *testing.T, client extProcPb.ExternalProcessor_ProcessClient, req *extProcPb.ProcessingRequest) (*extProcPb.ProcessingResponse, error) {
 	t.Logf("Sending request: %v", req)
 	if err := client.Send(req); err != nil {