more efficient rewrite fetching per request.

Author: zetxqx

apix/v1alpha2/inferencemodelrewrite_types.go

@@ -57,20 +57,25 @@ type InferenceModelRewriteSpec struct {
 	// If multiple InferenceModelRewrite resources target the same
 	// InferencePool, the controller will merge them based on precedence.
 	//
-	// **Timestamp Wins:** If two rules from different rewrites all matches,
-	// the rule from the *oldest*
-	// InferenceModelRewrite resource (determined by
-	// metadata.creationTimestamp) will be used.
+	// Across all rules specified on applicable rewrites, precedence MUST be
+	// given to the match having an "Exact" model match over a generic match
+	// (a rule with an empty `matches` array).
+	//
+	// If ties still exist across multiple InferenceModelRewrite resources (e.g.
+	// two rewrites both have an exact match for the same model), matching
+	// precedence MUST be determined by the oldest resource based on
+	// creation timestamp.
+	//
+	// If ties still exist within a single InferenceModelRewrite resource, the
+	// FIRST matching rule (in list order) is used.
 	// +required
 	Rules []InferenceModelRewriteRule `json:"rules"`
 }
 
 // InferenceModelRewriteRule defines the match criteria and corresponding action.
-//
-// A specific model name can only be matched by one rule across all
-// rules attached to the same InferencePool. If multiple rules attempt
-// to match the same model name, the oldest rule (by creationTimestamp)
-// will be the only one considered valid.
+// For details on how precedence is determined across multiple rules and
+// InferenceModelRewrite resources, see the "Precedence and Conflict Resolution"
+// section in InferenceModelRewriteSpec.
 type InferenceModelRewriteRule struct {
 	// Matches defines the criteria for matching a request.
 	// If multiple match criteria are specified, a request matches if

config/crd/bases/inference.networking.x-k8s.io_inferencemodelrewrites.yaml

@@ -74,11 +74,9 @@ spec:
                 items:
                   description: |-
                     InferenceModelRewriteRule defines the match criteria and corresponding action.
-
-                    A specific model name can only be matched by one rule across all
-                    rules attached to the same InferencePool. If multiple rules attempt
-                    to match the same model name, the oldest rule (by creationTimestamp)
-                    will be the only one considered valid.
+                    For details on how precedence is determined across multiple rules and
+                    InferenceModelRewrite resources, see the "Precedence and Conflict Resolution"
+                    section in InferenceModelRewriteSpec.
                   properties:
                     matches:
                       items:

docs/proposals/1816-inferenceomodelrewrite/README.md

@@ -64,20 +64,25 @@ type InferenceModelRewriteSpec struct {
 	// If multiple InferenceModelRewrite resources target the same
 	// InferencePool, the controller will merge them based on precedence.
 	//
-	// **Timestamp Wins:** If two rules from different rewrite all matches,
-	// the rule from the *oldest*
-	// InferenceModelRewrite resource (determined by
-	// metadata.creationTimestamp) will be used.
+	// Across all rules specified on applicable rewrites, precedence MUST be
+	// given to the match having an "Exact" model match over a generic match
+	// (a rule with an empty `matches` array).
+	//
+	// If ties still exist across multiple InferenceModelRewrite resources (e.g.
+	// two rewrites both have an exact match for the same model), matching
+	// precedence MUST be determined by the oldest resource based on
+	// creation timestamp.
+	//
+	// If ties still exist within a single InferenceModelRewrite resource, the
+	// FIRST matching rule (in list order) is used.
 	// +required
 	Rules []InferenceModelRewriteRule `json:"rules"`
 }
 
 // InferenceModelRewriteRule defines the match criteria and corresponding action.
-//
-// A specific model name can only be matched by one rule across all
-// rewrites attached to the same InferencePool. If multiple rules attempt
-// to match the same model name, the oldest rule (by creationTimestamp)
-// will be the only one considered valid.
+// For details on how precedence is determined across multiple rules and
+// InferenceModelRewrite resources, see the "Precedence and Conflict Resolution"
+// section in InferenceModelRewriteSpec.
 type InferenceModelRewriteRule struct {
 	// Matches defines the criteria for matching a request.
 	// If multiple match criteria are specified, a request matches if

pkg/epp/controller/inferencemodelrewrite_reconciler_test.go

@@ -41,79 +41,59 @@ import (
 
 var (
 	poolForRewrite = utiltest.MakeInferencePool("test-pool1").Namespace("ns1").ObjRef()
-	rewrite1       = makeInferenceModelRewrite("rewrite1").
-			Namespace(poolForRewrite.Namespace).
-			PoolName(poolForRewrite.Name).
-			CreationTimestamp(metav1.Unix(1000, 0)).
-			ObjRef()
-	rewrite1Pool2 = makeInferenceModelRewrite(rewrite1.Name).
-			Namespace(rewrite1.Namespace).
-			PoolName("test-pool2").
-			CreationTimestamp(metav1.Unix(1001, 0)).
-			ObjRef()
-	rewrite1Updated = makeInferenceModelRewrite(rewrite1.Name).
-			Namespace(rewrite1.Namespace).
-			PoolName(poolForRewrite.Name).
-			CreationTimestamp(metav1.Unix(1003, 0)).
-			Rules([]v1alpha2.InferenceModelRewriteRule{{}}).
-			ObjRef()
-	rewrite1Deleted = makeInferenceModelRewrite(rewrite1.Name).
-			Namespace(rewrite1.Namespace).
-			PoolName(poolForRewrite.Name).
-			CreationTimestamp(metav1.Unix(1004, 0)).
-			DeletionTimestamp().
-			ObjRef()
-	rewrite2 = makeInferenceModelRewrite("rewrite2").
-			Namespace(poolForRewrite.Namespace).
-			PoolName(poolForRewrite.Name).
-			CreationTimestamp(metav1.Unix(1000, 0)).
-			ObjRef()
-)
-
-type inferenceModelRewriteBuilder struct {
-	*v1alpha2.InferenceModelRewrite
-}
-
-func makeInferenceModelRewrite(name string) *inferenceModelRewriteBuilder {
-	return &inferenceModelRewriteBuilder{
-		&v1alpha2.InferenceModelRewrite{
-			ObjectMeta: metav1.ObjectMeta{
-				Name: name,
-			},
+	rewrite1       = &v1alpha2.InferenceModelRewrite{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:              "rewrite1",
+			Namespace:         poolForRewrite.Namespace,
+			CreationTimestamp: metav1.Unix(1000, 0),
+		},
+		Spec: v1alpha2.InferenceModelRewriteSpec{
+			PoolRef: &v1alpha2.PoolObjectReference{Name: v1alpha2.ObjectName(poolForRewrite.Name)},
 		},
 	}
-}
-
-func (b *inferenceModelRewriteBuilder) Namespace(ns string) *inferenceModelRewriteBuilder {
-	b.ObjectMeta.Namespace = ns
-	return b
-}
-
-func (b *inferenceModelRewriteBuilder) PoolName(name string) *inferenceModelRewriteBuilder {
-	b.Spec.PoolRef = &v1alpha2.PoolObjectReference{}
-	b.Spec.PoolRef.Name = v1alpha2.ObjectName(name)
-	return b
-}
-
-func (b *inferenceModelRewriteBuilder) CreationTimestamp(t metav1.Time) *inferenceModelRewriteBuilder {
-	b.ObjectMeta.CreationTimestamp = t
-	return b
-}
-
-func (b *inferenceModelRewriteBuilder) DeletionTimestamp() *inferenceModelRewriteBuilder {
-	now := metav1.Now()
-	b.ObjectMeta.DeletionTimestamp = &now
-	return b
-}
-
-func (b *inferenceModelRewriteBuilder) Rules(rules []v1alpha2.InferenceModelRewriteRule) *inferenceModelRewriteBuilder {
-	b.Spec.Rules = rules
-	return b
-}
-
-func (b *inferenceModelRewriteBuilder) ObjRef() *v1alpha2.InferenceModelRewrite {
-	return b.InferenceModelRewrite
-}
+	rewrite1Pool2 = &v1alpha2.InferenceModelRewrite{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:              rewrite1.Name,
+			Namespace:         rewrite1.Namespace,
+			CreationTimestamp: metav1.Unix(1001, 0),
+		},
+		Spec: v1alpha2.InferenceModelRewriteSpec{
+			PoolRef: &v1alpha2.PoolObjectReference{Name: "test-pool2"},
+		},
+	}
+	rewrite1Updated = &v1alpha2.InferenceModelRewrite{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:              rewrite1.Name,
+			Namespace:         rewrite1.Namespace,
+			CreationTimestamp: metav1.Unix(1003, 0),
+		},
+		Spec: v1alpha2.InferenceModelRewriteSpec{
+			PoolRef: &v1alpha2.PoolObjectReference{Name: v1alpha2.ObjectName(poolForRewrite.Name)},
+			Rules:   []v1alpha2.InferenceModelRewriteRule{{}},
+		},
+	}
+	rewrite1Deleted = &v1alpha2.InferenceModelRewrite{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:              rewrite1.Name,
+			Namespace:         rewrite1.Namespace,
+			CreationTimestamp: metav1.Unix(1004, 0),
+			DeletionTimestamp: &metav1.Time{Time: time.Now()},
+		},
+		Spec: v1alpha2.InferenceModelRewriteSpec{
+			PoolRef: &v1alpha2.PoolObjectReference{Name: v1alpha2.ObjectName(poolForRewrite.Name)},
+		},
+	}
+	rewrite2 = &v1alpha2.InferenceModelRewrite{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:              "rewrite2",
+			Namespace:         poolForRewrite.Namespace,
+			CreationTimestamp: metav1.Unix(1001, 0),
+		},
+		Spec: v1alpha2.InferenceModelRewriteSpec{
+			PoolRef: &v1alpha2.PoolObjectReference{Name: v1alpha2.ObjectName(poolForRewrite.Name)},
+		},
+	}
+)
 
 func TestInferenceModelRewriteReconciler(t *testing.T) {
 	tests := []struct {

pkg/epp/datastore/datastore.go

@@ -63,6 +63,7 @@ type Datastore interface {
 	// InferenceModelRewrite operations
 	RewriteSet(infModelRewrite *v1alpha2.InferenceModelRewrite)
 	RewriteDelete(namespacedName types.NamespacedName)
+	RewriteGet(modelName string) *v1alpha2.InferenceModelRewriteRule
 	RewriteGetAll() []*v1alpha2.InferenceModelRewrite
 
 	// PodList lists pods matching the given predicate.
@@ -77,9 +78,9 @@ type Datastore interface {
 func NewDatastore(parentCtx context.Context, epFactory datalayer.EndpointFactory, modelServerMetricsPort int32) Datastore {
 	store := &datastore{
 		parentCtx:              parentCtx,
-		poolAndObjectivesMu:    sync.RWMutex{},
+		mu:                     sync.RWMutex{},
 		objectives:             make(map[string]*v1alpha2.InferenceObjective),
-		rewrites:               make(map[types.NamespacedName]*v1alpha2.InferenceModelRewrite),
+		rewrites:               NewModelRewriteStore(),
 		pods:                   &sync.Map{},
 		modelServerMetricsPort: modelServerMetricsPort,
 		epf:                    epFactory,
@@ -90,13 +91,13 @@ func NewDatastore(parentCtx context.Context, epFactory datalayer.EndpointFactory
 type datastore struct {
 	// parentCtx controls the lifecycle of the background metrics goroutines that spawn up by the datastore.
 	parentCtx context.Context
-	// poolAndObjectivesMu is used to synchronize access to pool and the objectives map.
-	poolAndObjectivesMu sync.RWMutex
-	pool                *v1.InferencePool
+	// mu is used to synchronize access to pool, objectives, and rewrites.
+	mu   sync.RWMutex
+	pool *v1.InferencePool
 	// key: InferenceObjective name, value: *InferenceObjective
 	objectives map[string]*v1alpha2.InferenceObjective
-	// key: types.NamespacedName, value: *v1alpha2.InferenceModelRewrite
-	rewrites map[types.NamespacedName]*v1alpha2.InferenceModelRewrite
+	// rewrites store for InferenceModelRewrite objects.
+	rewrites *ModelRewriteStore
 	// key: types.NamespacedName, value: backendmetrics.PodMetrics
 	pods *sync.Map
 	// modelServerMetricsPort metrics port from EPP command line argument
@@ -106,11 +107,11 @@ type datastore struct {
 }
 
 func (ds *datastore) Clear() {
-	ds.poolAndObjectivesMu.Lock()
-	defer ds.poolAndObjectivesMu.Unlock()
+	ds.mu.Lock()
+	defer ds.mu.Unlock()
 	ds.pool = nil
 	ds.objectives = make(map[string]*v1alpha2.InferenceObjective)
-	ds.rewrites = make(map[types.NamespacedName]*v1alpha2.InferenceModelRewrite)
+	ds.rewrites = NewModelRewriteStore()
 	// stop all pods go routines before clearing the pods map.
 	ds.pods.Range(func(_, v any) bool {
 		ds.epf.ReleaseEndpoint(v.(backendmetrics.PodMetrics))
@@ -126,8 +127,8 @@ func (ds *datastore) PoolSet(ctx context.Context, reader client.Reader, pool *v1
 		return nil
 	}
 	logger := log.FromContext(ctx)
-	ds.poolAndObjectivesMu.Lock()
-	defer ds.poolAndObjectivesMu.Unlock()
+	ds.mu.Lock()
+	defer ds.mu.Unlock()
 
 	oldPool := ds.pool
 	ds.pool = pool
@@ -148,23 +149,23 @@ func (ds *datastore) PoolSet(ctx context.Context, reader client.Reader, pool *v1
 }
 
 func (ds *datastore) PoolGet() (*v1.InferencePool, error) {
-	ds.poolAndObjectivesMu.RLock()
-	defer ds.poolAndObjectivesMu.RUnlock()
+	ds.mu.RLock()
+	defer ds.mu.RUnlock()
 	if !ds.PoolHasSynced() {
 		return nil, errPoolNotSynced
 	}
 	return ds.pool, nil
 }
 
 func (ds *datastore) PoolHasSynced() bool {
-	ds.poolAndObjectivesMu.RLock()
-	defer ds.poolAndObjectivesMu.RUnlock()
+	ds.mu.RLock()
+	defer ds.mu.RUnlock()
 	return ds.pool != nil
 }
 
 func (ds *datastore) PoolLabelsMatch(podLabels map[string]string) bool {
-	ds.poolAndObjectivesMu.RLock()
-	defer ds.poolAndObjectivesMu.RUnlock()
+	ds.mu.RLock()
+	defer ds.mu.RUnlock()
 	if ds.pool == nil {
 		return false
 	}
@@ -173,59 +174,57 @@ func (ds *datastore) PoolLabelsMatch(podLabels map[string]string) bool {
 	return poolSelector.Matches(podSet)
 }
 
+// /// InferenceObjective APIs ///
 func (ds *datastore) ObjectiveSet(infObjective *v1alpha2.InferenceObjective) {
-	ds.poolAndObjectivesMu.Lock()
-	defer ds.poolAndObjectivesMu.Unlock()
-	// Set the objective.
+	ds.mu.Lock()
+	defer ds.mu.Unlock()
 	ds.objectives[infObjective.Name] = infObjective
 }
 
 func (ds *datastore) ObjectiveGet(objectiveName string) *v1alpha2.InferenceObjective {
-	ds.poolAndObjectivesMu.RLock()
-	defer ds.poolAndObjectivesMu.RUnlock()
-	iObj, ok := ds.objectives[objectiveName]
-	if !ok {
-		return nil
-	}
-	return iObj
+	ds.mu.RLock()
+	defer ds.mu.RUnlock()
+	return ds.objectives[objectiveName]
 }
 
 func (ds *datastore) ObjectiveDelete(namespacedName types.NamespacedName) {
-	ds.poolAndObjectivesMu.Lock()
-	defer ds.poolAndObjectivesMu.Unlock()
+	ds.mu.Lock()
+	defer ds.mu.Unlock()
 	delete(ds.objectives, namespacedName.Name)
 }
 
 func (ds *datastore) ObjectiveGetAll() []*v1alpha2.InferenceObjective {
-	ds.poolAndObjectivesMu.RLock()
-	defer ds.poolAndObjectivesMu.RUnlock()
-	res := []*v1alpha2.InferenceObjective{}
+	ds.mu.RLock()
+	defer ds.mu.RUnlock()
+	res := make([]*v1alpha2.InferenceObjective, 0, len(ds.objectives))
 	for _, v := range ds.objectives {
 		res = append(res, v)
 	}
 	return res
 }
 
 func (ds *datastore) RewriteSet(infModelRewrite *v1alpha2.InferenceModelRewrite) {
-	ds.poolAndObjectivesMu.Lock()
-	defer ds.poolAndObjectivesMu.Unlock()
-	ds.rewrites[types.NamespacedName{Name: infModelRewrite.Name, Namespace: infModelRewrite.Namespace}] = infModelRewrite
+	ds.mu.Lock()
+	defer ds.mu.Unlock()
+	ds.rewrites.Set(infModelRewrite)
 }
 
 func (ds *datastore) RewriteDelete(namespacedName types.NamespacedName) {
-	ds.poolAndObjectivesMu.Lock()
-	defer ds.poolAndObjectivesMu.Unlock()
-	delete(ds.rewrites, namespacedName)
+	ds.mu.Lock()
+	defer ds.mu.Unlock()
+	ds.rewrites.Delete(namespacedName)
+}
+
+func (ds *datastore) RewriteGet(modelName string) *v1alpha2.InferenceModelRewriteRule {
+	ds.mu.RLock()
+	defer ds.mu.RUnlock()
+	return ds.rewrites.GetRule(modelName)
 }
 
 func (ds *datastore) RewriteGetAll() []*v1alpha2.InferenceModelRewrite {
-	ds.poolAndObjectivesMu.RLock()
-	defer ds.poolAndObjectivesMu.RUnlock()
-	res := []*v1alpha2.InferenceModelRewrite{}
-	for _, v := range ds.rewrites {
-		res = append(res, v)
-	}
-	return res
+	ds.mu.RLock()
+	defer ds.mu.RUnlock()
+	return ds.rewrites.GetAll()
 }
 
 // /// Pods/endpoints APIs ///