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

Author: ahg-g

pkg/epp/controller/inferencemodel_reconciler.go

@@ -85,17 +85,15 @@ func (c *InferenceModelReconciler) Reconcile(ctx context.Context, req ctrl.Reque
 func (c *InferenceModelReconciler) handleModelDeleted(ctx context.Context, req types.NamespacedName) error {
 	logger := log.FromContext(ctx)
 
-	// We will lookup the modelName associated with this object to search for
-	// other instance referencing the same ModelName if exist to store the oldest in
+	// 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.ModelGetByObjName(req)
+	existing, exists := c.Datastore.ModelDelete(req)
 	if !exists {
 		// No entry exists in the first place, nothing to do.
 		return nil
 	}
-	// Delete the internal object, it may be replaced with another version below.
-	c.Datastore.ModelDelete(req)
 	logger.Info("InferenceModel removed from datastore", "poolRef", existing.Spec.PoolRef, "modelName", existing.Spec.ModelName)
 
 	// List all InferenceModels with a matching ModelName.

pkg/epp/datastore/datastore.go

@@ -46,9 +46,8 @@ type Datastore interface {
 
 	// InferenceModel operations
 	ModelSetIfOlder(infModel *v1alpha2.InferenceModel) bool
-	ModelGetByModelName(modelName string) (*v1alpha2.InferenceModel, bool)
-	ModelGetByObjName(namespacedName types.NamespacedName) (*v1alpha2.InferenceModel, bool)
-	ModelDelete(namespacedName types.NamespacedName)
+	ModelGet(modelName string) (*v1alpha2.InferenceModel, bool)
+	ModelDelete(namespacedName types.NamespacedName) (*v1alpha2.InferenceModel, bool)
 	ModelGetAll() []*v1alpha2.InferenceModel
 
 	// PodMetrics operations
@@ -67,11 +66,9 @@ type Datastore interface {
 
 func NewDatastore() Datastore {
 	store := &datastore{
-		poolMu:            sync.RWMutex{},
-		modelsMu:          sync.RWMutex{},
-		modelsByModelName: make(map[string]*v1alpha2.InferenceModel),
-		modelsByObjName:   make(map[types.NamespacedName]*v1alpha2.InferenceModel),
-		pods:              &sync.Map{},
+		poolAndModelsMu: sync.RWMutex{},
+		models:          make(map[string]*v1alpha2.InferenceModel),
+		pods:            &sync.Map{},
 	}
 	return store
 }
@@ -97,127 +94,102 @@ func NewFakeDatastore(pods []*PodMetrics, models []*v1alpha2.InferenceModel, poo
 }
 
 type datastore struct {
-	// poolMu is used to synchronize access to the inferencePool.
-	poolMu   sync.RWMutex
-	pool     *v1alpha2.InferencePool
-	modelsMu sync.RWMutex
-	// key: types.NamespacedName, value: *InferenceModel
-	modelsByObjName map[types.NamespacedName]*v1alpha2.InferenceModel
+	// poolAndModelsMu is used to synchronize access to pool and the models map.
+	poolAndModelsMu sync.RWMutex
+	pool            *v1alpha2.InferencePool
 	// key: InferenceModel.Spec.ModelName, value: *InferenceModel
-	modelsByModelName map[string]*v1alpha2.InferenceModel
+	models map[string]*v1alpha2.InferenceModel
 	// key: types.NamespacedName, value: *PodMetrics
 	pods *sync.Map
 }
 
 func (ds *datastore) Clear() {
-	ds.poolMu.Lock()
-	defer ds.poolMu.Unlock()
+	ds.poolAndModelsMu.Lock()
+	defer ds.poolAndModelsMu.Unlock()
 	ds.pool = nil
-	ds.modelsMu.Lock()
-	ds.modelsByModelName = make(map[string]*v1alpha2.InferenceModel)
-	ds.modelsByObjName = make(map[types.NamespacedName]*v1alpha2.InferenceModel)
-	ds.modelsMu.Unlock()
+	ds.models = make(map[string]*v1alpha2.InferenceModel)
 	ds.pods.Clear()
 }
 
 // /// InferencePool APIs ///
 func (ds *datastore) PoolSet(pool *v1alpha2.InferencePool) {
-	ds.poolMu.Lock()
-	defer ds.poolMu.Unlock()
+	ds.poolAndModelsMu.Lock()
+	defer ds.poolAndModelsMu.Unlock()
 	ds.pool = pool
 }
 
 func (ds *datastore) PoolGet() (*v1alpha2.InferencePool, error) {
-	ds.poolMu.RLock()
-	defer ds.poolMu.RUnlock()
+	ds.poolAndModelsMu.RLock()
+	defer ds.poolAndModelsMu.RUnlock()
 	if !ds.PoolHasSynced() {
 		return nil, errPoolNotSynced
 	}
 	return ds.pool, nil
 }
 
 func (ds *datastore) PoolHasSynced() bool {
-	ds.poolMu.RLock()
-	defer ds.poolMu.RUnlock()
+	ds.poolAndModelsMu.RLock()
+	defer ds.poolAndModelsMu.RUnlock()
 	return ds.pool != nil
 }
 
 func (ds *datastore) PoolLabelsMatch(podLabels map[string]string) bool {
-	ds.poolMu.RLock()
-	defer ds.poolMu.RUnlock()
+	ds.poolAndModelsMu.RLock()
+	defer ds.poolAndModelsMu.RUnlock()
 	poolSelector := selectorFromInferencePoolSelector(ds.pool.Spec.Selector)
 	podSet := labels.Set(podLabels)
 	return poolSelector.Matches(podSet)
 }
 
 // /// InferenceModel APIs ///
 func (ds *datastore) ModelSetIfOlder(infModel *v1alpha2.InferenceModel) bool {
-	ds.modelsMu.Lock()
-	defer ds.modelsMu.Unlock()
+	ds.poolAndModelsMu.Lock()
+	defer ds.poolAndModelsMu.Unlock()
 
 	// Check first if the existing model is older.
 	// One exception is if the incoming model object is the same, in which case, we should not
 	// check for creation timestamp since that means the object was re-created, and so we should override.
-	existing, exists := ds.modelsByModelName[infModel.Spec.ModelName]
+	existing, exists := ds.models[infModel.Spec.ModelName]
 	if exists {
 		diffObj := infModel.Name != existing.Name || infModel.Namespace != existing.Namespace
 		if diffObj && existing.ObjectMeta.CreationTimestamp.Before(&infModel.ObjectMeta.CreationTimestamp) {
 			return false
 		}
 	}
-
-	// Deleting the model first ensures that the two maps are always aligned.
-	namespacedName := types.NamespacedName{Name: infModel.Name, Namespace: infModel.Namespace}
-	ds.modelDeleteByObjName(namespacedName)
-	ds.modelDeleteByModelName(infModel.Spec.ModelName)
-	ds.modelsByModelName[infModel.Spec.ModelName] = infModel
-	ds.modelsByObjName[namespacedName] = infModel
+	// Delete the model
+	ds.modelDelete(types.NamespacedName{Name: infModel.Name, Namespace: infModel.Namespace})
+	ds.models[infModel.Spec.ModelName] = infModel
 	return true
 }
 
-func (ds *datastore) ModelGetByModelName(modelName string) (*v1alpha2.InferenceModel, bool) {
-	ds.modelsMu.RLock()
-	defer ds.modelsMu.RUnlock()
-	m, exists := ds.modelsByModelName[modelName]
-	return m, exists
-}
-
-func (ds *datastore) ModelGetByObjName(namespacedName types.NamespacedName) (*v1alpha2.InferenceModel, bool) {
-	ds.modelsMu.RLock()
-	defer ds.modelsMu.RUnlock()
-	m, exists := ds.modelsByObjName[namespacedName]
+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) {
-	ds.modelsMu.Lock()
-	defer ds.modelsMu.Unlock()
-	ds.modelDeleteByObjName(namespacedName)
-}
-
-func (ds *datastore) modelDeleteByObjName(namespacedName types.NamespacedName) {
-	infModel, ok := ds.modelsByObjName[namespacedName]
-	if !ok {
-		return
+func (ds *datastore) modelDelete(namespacedName types.NamespacedName) (*v1alpha2.InferenceModel, bool) {
+	for _, m := range ds.models {
+		if m.Name == namespacedName.Name && m.Namespace == namespacedName.Namespace {
+			delete(ds.models, m.Spec.ModelName)
+			return m, true
+		}
 	}
-	delete(ds.modelsByObjName, namespacedName)
-	delete(ds.modelsByModelName, infModel.Spec.ModelName)
+	return nil, false
 }
 
-func (ds *datastore) modelDeleteByModelName(modelName string) {
-	infModel, ok := ds.modelsByModelName[modelName]
-	if !ok {
-		return
-	}
-	delete(ds.modelsByObjName, types.NamespacedName{Name: infModel.Name, Namespace: infModel.Namespace})
-	delete(ds.modelsByModelName, modelName)
+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.modelsMu.RLock()
-	defer ds.modelsMu.RUnlock()
+	ds.poolAndModelsMu.RLock()
+	defer ds.poolAndModelsMu.RUnlock()
 	res := []*v1alpha2.InferenceModel{}
-	for _, v := range ds.modelsByObjName {
+	for _, v := range ds.models {
 		res = append(res, v)
 	}
 	return res

pkg/epp/datastore/datastore_test.go

@@ -185,32 +185,22 @@ func TestModel1(t *testing.T) {
 			name:           "Getting by model name, chat -> model2",
 			existingModels: []*v1alpha2.InferenceModel{model2chat, model1ts},
 			op: func(ds Datastore) bool {
-				gotChat, exists := ds.ModelGetByModelName(chatModel)
+				gotChat, exists := ds.ModelGet(chatModel)
 				return exists && cmp.Diff(model2chat, gotChat) == ""
 			},
 			wantOpResult: true,
 			wantModels:   []*v1alpha2.InferenceModel{model2chat, model1ts},
 		},
 		{
-			name:           "Getting by obj name, model1 -> tweet-summary",
+			name:           "Delete the model",
 			existingModels: []*v1alpha2.InferenceModel{model2chat, model1ts},
 			op: func(ds Datastore) bool {
-				got, exists := ds.ModelGetByObjName(types.NamespacedName{Name: model1ts.Name, Namespace: model1ts.Namespace})
-				return exists && cmp.Diff(model1ts, got) == ""
-			},
-			wantOpResult: true,
-			wantModels:   []*v1alpha2.InferenceModel{model2chat, model1ts},
-		},
-		{
-			name:           "Getting by model name, chat -> model2",
-			existingModels: []*v1alpha2.InferenceModel{model2chat, model1ts},
-			op: func(ds Datastore) bool {
-				ds.ModelDelete(types.NamespacedName{Name: model1ts.Name, Namespace: model1ts.Namespace})
-				_, exists := ds.ModelGetByModelName(tsModel)
-				return exists
+				_, existed := ds.ModelDelete(types.NamespacedName{Name: model1ts.Name, Namespace: model1ts.Namespace})
+				_, exists := ds.ModelGet(tsModel)
+				return existed && !exists
 
 			},
-			wantOpResult: false,
+			wantOpResult: true,
 			wantModels:   []*v1alpha2.InferenceModel{model2chat},
 		},
 	}
@@ -221,26 +211,15 @@ func TestModel1(t *testing.T) {
 			if gotOpResult != test.wantOpResult {
 				t.Errorf("Unexpected operation result, want: %v, got: %v", test.wantOpResult, gotOpResult)
 			}
-			if diff := diffModelMaps(ds.(*datastore), test.wantModels); diff != "" {
+
+			if diff := testutil.DiffModelLists(test.wantModels, ds.ModelGetAll()); diff != "" {
 				t.Errorf("Unexpected models diff: %s", diff)
 			}
 
 		})
 	}
 }
 
-func diffModelMaps(ds *datastore, want []*v1alpha2.InferenceModel) string {
-	byObjName := ds.ModelGetAll()
-	byModelName := []*v1alpha2.InferenceModel{}
-	for _, v := range ds.modelsByModelName {
-		byModelName = append(byModelName, v)
-	}
-	if diff := testutil.DiffModelLists(byObjName, byModelName); diff != "" {
-		return "Inconsistent maps diff: " + diff
-	}
-	return testutil.DiffModelLists(want, byObjName)
-}
-
 func TestRandomWeightedDraw(t *testing.T) {
 	logger := logutil.NewTestLogger()
 	tests := []struct {

pkg/epp/handlers/request.go

@@ -64,7 +64,7 @@ func (s *Server) HandleRequestBody(
 	// NOTE: The nil checking for the modelObject means that we DO allow passthrough currently.
 	// This might be a security risk in the future where adapters not registered in the InferenceModel
 	// are able to be requested by using their distinct name.
-	modelObj, exist := s.datastore.ModelGetByModelName(model)
+	modelObj, exist := s.datastore.ModelGet(model)
 	if !exist {
 		return nil, errutil.Error{Code: errutil.BadConfiguration, Msg: fmt.Sprintf("error finding a model object in InferenceModel for input %v", model)}
 	}

test/integration/hermetic_test.go

@@ -476,7 +476,7 @@ func BeforeSuit(t *testing.T) func() {
 	}
 
 	assert.EventuallyWithT(t, func(t *assert.CollectT) {
-		_, modelExist := serverRunner.Datastore.ModelGetByModelName("my-model")
+		_, modelExist := serverRunner.Datastore.ModelGet("my-model")
 		synced := serverRunner.Datastore.PoolHasSynced() && modelExist
 		assert.True(t, synced, "Timeout waiting for the pool and models to sync")
 	}, 10*time.Second, 10*time.Millisecond)