Fix InferenceModel deletion logic (kubernetes-sigs#393)

* Currently the logic tracks the models by Spec.ModelName, since this is not guaranteed to be unique within the cluster, we could run into two issues:

1) If the model name changes on the same InferenceModel object, we don't delete the original model entry in the datastore. 2) We don't enforce the semantics that the modelName with the oldest creation timestamp is retained. While the api is assuming that this is enforced by another controller via the Ready condition, we don't have this controller yet, and so currently the behavior is unpredictable depending on InferenceModel events order.

To address the above, the PR makes changes to both the InferenceModel reconciler and the Model APIs in the datastore to ensure thread safe updates of the entries. In the store, the sync.Map was replaced with two maps to track the InferenceModel entries by both ModelName and InferenceModel object NamespacedName. This is needed to properly handle deletions when the object doesn't exist anymore (could be handled in other ways, but this seemed like a reasonable approach).

The PR increases the datastore pkg unit test coverage the Pool and Model APIs. We still need to followup with adding unit test coverage to the pods APIs, which is currently non-existent.

* Convert unit test to a table

* remove the dual map for the models store, and rely on linear search when looking up the model by object name

* Added ModelResync to handle a race condition

* Update pkg/epp/controller/inferencemodel_reconciler.go

Author: ahg-g

cmd/epp/main.go

@@ -149,6 +149,8 @@ func run() error {
 		return err
 	}
 
+	ctx := ctrl.SetupSignalHandler()
+
 	// Setup runner.
 	datastore := datastore.NewDatastore()
 	provider := backend.NewProvider(&vllm.PodMetricsClientImpl{}, datastore)
@@ -165,7 +167,7 @@ func run() error {
 		CertPath:                                 *certPath,
 		Provider:                                 provider,
 	}
-	if err := serverRunner.SetupWithManager(mgr); err != nil {
+	if err := serverRunner.SetupWithManager(ctx, mgr); err != nil {
 		setupLog.Error(err, "Failed to setup ext-proc controllers")
 		return err
 	}
@@ -188,7 +190,7 @@ func run() error {
 
 	// Start the manager. This blocks until a signal is received.
 	setupLog.Info("Controller manager starting")
-	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
+	if err := mgr.Start(ctx); err != nil {
 		setupLog.Error(err, "Error starting controller manager")
 		return err
 	}

pkg/epp/backend/provider_test.go

@@ -19,7 +19,6 @@ package backend
 import (
 	"context"
 	"errors"
-	"sync"
 	"testing"
 	"time"
 
@@ -37,6 +36,9 @@ var (
 				Name: "pod1",
 			},
 		},
+	}
+	pod1WithMetrics = &datastore.PodMetrics{
+		Pod: pod1.Pod,
 		Metrics: datastore.Metrics{
 			WaitingQueueSize:    0,
 			KVCacheUsagePercent: 0.2,
@@ -53,6 +55,9 @@ var (
 				Name: "pod2",
 			},
 		},
+	}
+	pod2WithMetrics = &datastore.PodMetrics{
+		Pod: pod2.Pod,
 		Metrics: datastore.Metrics{
 			WaitingQueueSize:    1,
 			KVCacheUsagePercent: 0.2,
@@ -69,35 +74,30 @@ func TestProvider(t *testing.T) {
 	tests := []struct {
 		name      string
 		pmc       PodMetricsClient
-		datastore datastore.Datastore
+		storePods []*datastore.PodMetrics
 		want      []*datastore.PodMetrics
 	}{
 		{
 			name: "Probing metrics success",
 			pmc: &FakePodMetricsClient{
 				Res: map[types.NamespacedName]*datastore.PodMetrics{
-					pod1.NamespacedName: pod1,
-					pod2.NamespacedName: pod2,
+					pod1.NamespacedName: pod1WithMetrics,
+					pod2.NamespacedName: pod2WithMetrics,
 				},
 			},
-			datastore: datastore.NewFakeDatastore(populateMap(pod1, pod2), nil, nil),
-			want: []*datastore.PodMetrics{
-				pod1,
-				pod2,
-			},
+			storePods: []*datastore.PodMetrics{pod1, pod2},
+			want:      []*datastore.PodMetrics{pod1WithMetrics, pod2WithMetrics},
 		},
 		{
 			name: "Only pods in the datastore are probed",
 			pmc: &FakePodMetricsClient{
 				Res: map[types.NamespacedName]*datastore.PodMetrics{
-					pod1.NamespacedName: pod1,
-					pod2.NamespacedName: pod2,
+					pod1.NamespacedName: pod1WithMetrics,
+					pod2.NamespacedName: pod2WithMetrics,
 				},
 			},
-			datastore: datastore.NewFakeDatastore(populateMap(pod1), nil, nil),
-			want: []*datastore.PodMetrics{
-				pod1,
-			},
+			storePods: []*datastore.PodMetrics{pod1},
+			want:      []*datastore.PodMetrics{pod1WithMetrics},
 		},
 		{
 			name: "Probing metrics error",
@@ -106,13 +106,12 @@ func TestProvider(t *testing.T) {
 					pod2.NamespacedName: errors.New("injected error"),
 				},
 				Res: map[types.NamespacedName]*datastore.PodMetrics{
-					pod1.NamespacedName: pod1,
+					pod1.NamespacedName: pod1WithMetrics,
 				},
 			},
-			datastore: datastore.NewFakeDatastore(populateMap(pod1, pod2), nil, nil),
-
+			storePods: []*datastore.PodMetrics{pod1, pod2},
 			want: []*datastore.PodMetrics{
-				pod1,
+				pod1WithMetrics,
 				// Failed to fetch pod2 metrics so it remains the default values.
 				{
 					Pod: datastore.Pod{NamespacedName: pod2.NamespacedName},
@@ -128,12 +127,13 @@ func TestProvider(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			p := NewProvider(test.pmc, test.datastore)
+			ds := datastore.NewFakeDatastore(test.storePods, nil, nil)
+			p := NewProvider(test.pmc, ds)
 			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()
+				metrics := ds.PodGetAll()
 				diff := cmp.Diff(test.want, metrics, cmpopts.SortSlices(func(a, b *datastore.PodMetrics) bool {
 					return a.String() < b.String()
 				}))
@@ -142,11 +142,3 @@ func TestProvider(t *testing.T) {
 		})
 	}
 }
-
-func populateMap(pods ...*datastore.PodMetrics) *sync.Map {
-	newMap := &sync.Map{}
-	for _, pod := range pods {
-		newMap.Store(pod.NamespacedName, &datastore.PodMetrics{Pod: datastore.Pod{NamespacedName: pod.NamespacedName, Address: pod.Address}})
-	}
-	return newMap
-}

pkg/epp/controller/inferencemodel_reconciler.go

@@ -18,8 +18,8 @@ package controller
 
 import (
 	"context"
+	"fmt"
 
-	"github.com/go-logr/logr"
 	"k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
@@ -43,44 +43,80 @@ type InferenceModelReconciler struct {
 }
 
 func (c *InferenceModelReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
-	logger := log.FromContext(ctx)
-	loggerDefault := logger.V(logutil.DEFAULT)
-	loggerDefault.Info("Reconciling InferenceModel", "name", req.NamespacedName)
+	if req.Namespace != c.PoolNamespacedName.Namespace {
+		return ctrl.Result{}, nil
+	}
+	logger := log.FromContext(ctx).V(logutil.DEFAULT).WithValues("inferenceModel", req.Name)
+	ctx = ctrl.LoggerInto(ctx, logger)
+
+	logger.Info("Reconciling InferenceModel")
 
 	infModel := &v1alpha2.InferenceModel{}
+	notFound := false
 	if err := c.Get(ctx, req.NamespacedName, infModel); err != nil {
-		if errors.IsNotFound(err) {
-			loggerDefault.Info("InferenceModel not found. Removing from datastore since object must be deleted", "name", req.NamespacedName)
-			c.Datastore.ModelDelete(infModel.Spec.ModelName)
-			return ctrl.Result{}, nil
+		if !errors.IsNotFound(err) {
+			logger.Error(err, "Unable to get InferenceModel")
+			return ctrl.Result{}, err
 		}
-		loggerDefault.Error(err, "Unable to get InferenceModel", "name", req.NamespacedName)
+		notFound = true
+	}
+
+	if notFound || !infModel.DeletionTimestamp.IsZero() || infModel.Spec.PoolRef.Name != c.PoolNamespacedName.Name {
+		// InferenceModel object got deleted or changed the referenced pool.
+		err := c.handleModelDeleted(ctx, req.NamespacedName)
 		return ctrl.Result{}, err
-	} else if !infModel.DeletionTimestamp.IsZero() {
-		loggerDefault.Info("InferenceModel is marked for deletion. Removing from datastore", "name", req.NamespacedName)
-		c.Datastore.ModelDelete(infModel.Spec.ModelName)
-		return ctrl.Result{}, nil
 	}
 
-	c.updateDatastore(logger, infModel)
+	// Add or update if the InferenceModel instance has a creation timestamp older than the existing entry of the model.
+	logger = logger.WithValues("poolRef", infModel.Spec.PoolRef).WithValues("modelName", infModel.Spec.ModelName)
+	if !c.Datastore.ModelSetIfOlder(infModel) {
+		logger.Info("Skipping InferenceModel, existing instance has older creation timestamp")
+
+	}
+	logger.Info("Added/Updated InferenceModel")
+
 	return ctrl.Result{}, nil
 }
 
-func (c *InferenceModelReconciler) updateDatastore(logger logr.Logger, infModel *v1alpha2.InferenceModel) {
-	loggerDefault := logger.V(logutil.DEFAULT)
+func (c *InferenceModelReconciler) handleModelDeleted(ctx context.Context, req types.NamespacedName) error {
+	logger := log.FromContext(ctx)
 
-	if infModel.Spec.PoolRef.Name == c.PoolNamespacedName.Name {
-		loggerDefault.Info("Updating datastore", "poolRef", infModel.Spec.PoolRef, "serverPoolName", c.PoolNamespacedName)
-		loggerDefault.Info("Adding/Updating InferenceModel", "modelName", infModel.Spec.ModelName)
-		c.Datastore.ModelSet(infModel)
-		return
+	// We will lookup and delete the modelName associated with this object, and search for
+	// other instances referencing the same modelName if exist, and store the oldest in
+	// its place. This ensures that the InferenceModel with the oldest creation
+	// timestamp is active.
+	existing, exists := c.Datastore.ModelDelete(req)
+	if !exists {
+		// No entry exists in the first place, nothing to do.
+		return nil
+	}
+	logger.Info("InferenceModel removed from datastore", "poolRef", existing.Spec.PoolRef, "modelName", existing.Spec.ModelName)
+
+	// TODO(#409): replace this backfill logic with one that is based on InferenceModel Ready conditions once those are set by an external controller.
+	updated, err := c.Datastore.ModelResync(ctx, c.Client, existing.Spec.ModelName)
+	if err != nil {
+		return err
+	}
+	if updated {
+		logger.Info("Model replaced.", "modelName", existing.Spec.ModelName)
 	}
-	loggerDefault.Info("Removing/Not adding InferenceModel", "modelName", infModel.Spec.ModelName)
-	// If we get here. The model is not relevant to this pool, remove.
-	c.Datastore.ModelDelete(infModel.Spec.ModelName)
+	return nil
 }
 
-func (c *InferenceModelReconciler) SetupWithManager(mgr ctrl.Manager) error {
+func indexInferenceModelsByModelName(obj client.Object) []string {
+	m, ok := obj.(*v1alpha2.InferenceModel)
+	if !ok {
+		return nil
+	}
+	return []string{m.Spec.ModelName}
+}
+
+func (c *InferenceModelReconciler) SetupWithManager(ctx context.Context, mgr ctrl.Manager) error {
+	// Create an index on ModelName for InferenceModel objects.
+	indexer := mgr.GetFieldIndexer()
+	if err := indexer.IndexField(ctx, &v1alpha2.InferenceModel{}, datastore.ModelNameIndexKey, indexInferenceModelsByModelName); err != nil {
+		return fmt.Errorf("setting index on ModelName for InferenceModel: %w", err)
+	}
 	return ctrl.NewControllerManagedBy(mgr).
 		For(&v1alpha2.InferenceModel{}).
 		WithEventFilter(predicate.Funcs{