Added ModelResync to handle a race condition

Author: ahg-g

pkg/epp/controller/inferencemodel_reconciler.go

@@ -34,10 +34,6 @@ import (
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
-const (
-	modelNameKey = "spec.modelName"
-)
-
 type InferenceModelReconciler struct {
 	client.Client
 	Scheme             *runtime.Scheme
@@ -96,35 +92,13 @@ func (c *InferenceModelReconciler) handleModelDeleted(ctx context.Context, req t
 	}
 	logger.Info("InferenceModel removed from datastore", "poolRef", existing.Spec.PoolRef, "modelName", existing.Spec.ModelName)
 
-	// List all InferenceModels with a matching ModelName.
-	var models v1alpha2.InferenceModelList
-	if err := c.List(ctx, &models, client.MatchingFields{modelNameKey: existing.Spec.ModelName}, client.InNamespace(c.PoolNamespacedName.Namespace)); err != nil {
-		return fmt.Errorf("listing models that match the modelName %s: %w", existing.Spec.ModelName, err)
+	updated, err := c.Datastore.ModelResync(ctx, c.Client, existing.Spec.ModelName)
+	if err != nil {
+		return err
 	}
-	if len(models.Items) == 0 {
-		// No other instances of InferenceModels with this ModelName exists.
-		return nil
+	if updated {
+		logger.Info("Model replaced.", "modelName", existing.Spec.ModelName)
 	}
-
-	var oldest *v1alpha2.InferenceModel
-	for i := range models.Items {
-		m := &models.Items[i]
-		if m.Spec.ModelName != existing.Spec.ModelName || // The index should filter those out, but just in case!
-			m.Spec.PoolRef.Name != c.PoolNamespacedName.Name || // We don't care about other pools, we could setup an index on this too!
-			m.Name == existing.Name { // We don't care about the same object, it could be in the list if it was only marked for deletion, but not yet deleted.
-			continue
-		}
-		if oldest == nil || m.ObjectMeta.CreationTimestamp.Before(&oldest.ObjectMeta.CreationTimestamp) {
-			oldest = m
-		}
-	}
-	if oldest != nil && c.Datastore.ModelSetIfOlder(oldest) {
-		logger.Info("InferenceModel replaced.",
-			"poolRef", oldest.Spec.PoolRef,
-			"modelName", oldest.Spec.ModelName,
-			"newInferenceModel", types.NamespacedName{Name: oldest.Name, Namespace: oldest.Namespace})
-	}
-
 	return nil
 }
 
@@ -139,7 +113,7 @@ func indexInferenceModelsByModelName(obj client.Object) []string {
 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{}, modelNameKey, indexInferenceModelsByModelName); err != nil {
+	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).

pkg/epp/controller/inferencemodel_reconciler_test.go

@@ -198,7 +198,7 @@ func TestInferenceModelReconciler(t *testing.T) {
 			fakeClient := fake.NewClientBuilder().
 				WithScheme(scheme).
 				WithObjects(initObjs...).
-				WithIndex(&v1alpha2.InferenceModel{}, modelNameKey, indexInferenceModelsByModelName).
+				WithIndex(&v1alpha2.InferenceModel{}, datastore.ModelNameIndexKey, indexInferenceModelsByModelName).
 				Build()
 
 			datastore := datastore.NewFakeDatastore(nil, test.modelsInStore, pool)

pkg/epp/datastore/datastore.go

@@ -19,6 +19,7 @@ package datastore
 import (
 	"context"
 	"errors"
+	"fmt"
 	"math/rand"
 	"sync"
 
@@ -32,6 +33,10 @@ import (
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
+const (
+	ModelNameIndexKey = "spec.modelName"
+)
+
 var (
 	errPoolNotSynced = errors.New("InferencePool is not initialized in data store")
 )
@@ -48,6 +53,7 @@ type Datastore interface {
 	ModelSetIfOlder(infModel *v1alpha2.InferenceModel) bool
 	ModelGet(modelName string) (*v1alpha2.InferenceModel, bool)
 	ModelDelete(namespacedName types.NamespacedName) (*v1alpha2.InferenceModel, bool)
+	ModelResync(ctx context.Context, ctrlClient client.Client, modelName string) (bool, error)
 	ModelGetAll() []*v1alpha2.InferenceModel
 
 	// PodMetrics operations
@@ -156,20 +162,53 @@ func (ds *datastore) ModelSetIfOlder(infModel *v1alpha2.InferenceModel) bool {
 			return false
 		}
 	}
-	// Delete the model
-	ds.modelDelete(types.NamespacedName{Name: infModel.Name, Namespace: infModel.Namespace})
+	// Set the model.
 	ds.models[infModel.Spec.ModelName] = infModel
 	return true
 }
 
+func (ds *datastore) ModelResync(ctx context.Context, c client.Client, modelName string) (bool, error) {
+	ds.poolAndModelsMu.Lock()
+	defer ds.poolAndModelsMu.Unlock()
+
+	var models v1alpha2.InferenceModelList
+	if err := c.List(ctx, &models, client.MatchingFields{ModelNameIndexKey: modelName}, client.InNamespace(ds.pool.Namespace)); err != nil {
+		return false, fmt.Errorf("listing models that match the modelName %s: %w", modelName, err)
+	}
+	if len(models.Items) == 0 {
+		// No other instances of InferenceModels with this ModelName exists.
+		return false, nil
+	}
+
+	var oldest *v1alpha2.InferenceModel
+	for i := range models.Items {
+		m := &models.Items[i]
+		if m.Spec.ModelName != modelName || // The index should filter those out, but just in case!
+			m.Spec.PoolRef.Name != ds.pool.Name || // We don't care about other pools, we could setup an index on this too!
+			!m.DeletionTimestamp.IsZero() { // ignore objects marked for deletion
+			continue
+		}
+		if oldest == nil || m.ObjectMeta.CreationTimestamp.Before(&oldest.ObjectMeta.CreationTimestamp) {
+			oldest = m
+		}
+	}
+	if oldest == nil {
+		return false, nil
+	}
+	ds.models[modelName] = oldest
+	return true, nil
+}
+
 func (ds *datastore) ModelGet(modelName string) (*v1alpha2.InferenceModel, bool) {
 	ds.poolAndModelsMu.RLock()
 	defer ds.poolAndModelsMu.RUnlock()
 	m, exists := ds.models[modelName]
 	return m, exists
 }
 
-func (ds *datastore) modelDelete(namespacedName types.NamespacedName) (*v1alpha2.InferenceModel, bool) {
+func (ds *datastore) ModelDelete(namespacedName types.NamespacedName) (*v1alpha2.InferenceModel, bool) {
+	ds.poolAndModelsMu.Lock()
+	defer ds.poolAndModelsMu.Unlock()
 	for _, m := range ds.models {
 		if m.Name == namespacedName.Name && m.Namespace == namespacedName.Namespace {
 			delete(ds.models, m.Spec.ModelName)
@@ -179,12 +218,6 @@ func (ds *datastore) modelDelete(namespacedName types.NamespacedName) (*v1alpha2
 	return nil, false
 }
 
-func (ds *datastore) ModelDelete(namespacedName types.NamespacedName) (*v1alpha2.InferenceModel, bool) {
-	ds.poolAndModelsMu.Lock()
-	defer ds.poolAndModelsMu.Unlock()
-	return ds.modelDelete(namespacedName)
-}
-
 func (ds *datastore) ModelGetAll() []*v1alpha2.InferenceModel {
 	ds.poolAndModelsMu.RLock()
 	defer ds.poolAndModelsMu.RUnlock()

pkg/epp/datastore/datastore_test.go

@@ -154,15 +154,6 @@ func TestModel(t *testing.T) {
 			wantOpResult: true,
 			wantModels:   []*v1alpha2.InferenceModel{model2ts},
 		},
-		{
-			name:           "Set model2 updated with a new modelName, should update modelName",
-			existingModels: []*v1alpha2.InferenceModel{model2ts},
-			op: func(ds Datastore) bool {
-				return ds.ModelSetIfOlder(model2chat)
-			},
-			wantOpResult: true,
-			wantModels:   []*v1alpha2.InferenceModel{model2chat},
-		},
 		{
 			name:           "Set model1 with the tweet-summary modelName, both models should exist",
 			existingModels: []*v1alpha2.InferenceModel{model2chat},