add metrics for model rewrite. (#1938)

Author: zetxqx

pkg/epp/datastore/datastore.go

@@ -63,7 +63,7 @@ type Datastore interface {
 	// InferenceModelRewrite operations
 	ModelRewriteSet(infModelRewrite *v1alpha2.InferenceModelRewrite)
 	ModelRewriteDelete(namespacedName types.NamespacedName)
-	ModelRewriteGet(modelName string) *v1alpha2.InferenceModelRewriteRule
+	ModelRewriteGet(modelName string) (*v1alpha2.InferenceModelRewriteRule, string)
 	ModelRewriteGetAll() []*v1alpha2.InferenceModelRewrite
 
 	// PodList lists pods matching the given predicate.
@@ -225,7 +225,7 @@ func (ds *datastore) ModelRewriteDelete(namespacedName types.NamespacedName) {
 	ds.modelRewrites.delete(namespacedName)
 }
 
-func (ds *datastore) ModelRewriteGet(modelName string) *v1alpha2.InferenceModelRewriteRule {
+func (ds *datastore) ModelRewriteGet(modelName string) (*v1alpha2.InferenceModelRewriteRule, string) {
 	ds.mu.RLock()
 	defer ds.mu.RUnlock()
 	return ds.modelRewrites.getRule(modelName)

pkg/epp/datastore/modelrewritestore.go

@@ -120,17 +120,18 @@ func (ms *modelRewriteStore) deleteInternal(n string) {
 
 // getRule returns the single, highest-precedence rule for a given model name.
 // It prioritizes exact matches over generic ones, and among those, the oldest rule wins.
-func (ms *modelRewriteStore) getRule(modelName string) *v1alpha2.InferenceModelRewriteRule {
+// It also returns the name of the InferenceModelRewrite resource that provided the rule.
+func (ms *modelRewriteStore) getRule(modelName string) (*v1alpha2.InferenceModelRewriteRule, string) {
 	// Exact matches have the highest precedence.
 	if rulesWithMd, ok := ms.rulesByExactModelMatch[modelName]; ok && len(rulesWithMd) > 0 {
-		return &rulesWithMd[0].rule // The list is pre-sorted, so the first element is the oldest.
+		return &rulesWithMd[0].rule, rulesWithMd[0].parentName() // The list is pre-sorted, so the first element is the oldest.
 	}
 
 	// If no exact match, fall back to the oldest generic rule.
 	if len(ms.genericRules) > 0 {
-		return &ms.genericRules[0].rule // The list is pre-sorted.
+		return &ms.genericRules[0].rule, ms.genericRules[0].parentName() // The list is pre-sorted.
 	}
-	return nil
+	return nil, ""
 }
 
 // getAll returns a slice of all InferenceModelRewrite objects currently in the store.

pkg/epp/datastore/modelrewritestore_test.go

@@ -73,55 +73,63 @@ func TestModelRewriteStore(t *testing.T) {
 		op           func(store *modelRewriteStore)
 		modelToGet   string
 		wantRule     *v1alpha2.InferenceModelRewriteRule
+		wantName     string
 		wantGetAll   []*v1alpha2.InferenceModelRewrite
 	}{
 		{
 			name:         "Simple exact match",
 			initialState: []*v1alpha2.InferenceModelRewrite{rewriteOld},
 			modelToGet:   "model1",
 			wantRule:     &ruleModel1V1,
+			wantName:     rewriteOld.Name,
 			wantGetAll:   []*v1alpha2.InferenceModelRewrite{rewriteOld},
 		},
 		{
 			name:         "Simple generic match",
 			initialState: []*v1alpha2.InferenceModelRewrite{rewriteGenericOld},
 			modelToGet:   "model2", // A different model to test generic fallback
 			wantRule:     &ruleGeneric,
+			wantName:     rewriteGenericOld.Name,
 			wantGetAll:   []*v1alpha2.InferenceModelRewrite{rewriteGenericOld},
 		},
 		{
 			name:         "No match",
 			initialState: []*v1alpha2.InferenceModelRewrite{rewriteOld},
 			modelToGet:   "model2",
 			wantRule:     nil,
+			wantName:     "",
 			wantGetAll:   []*v1alpha2.InferenceModelRewrite{rewriteOld},
 		},
 		{
 			name:         "Precedence: Exact match wins over generic",
 			initialState: []*v1alpha2.InferenceModelRewrite{rewriteOld, rewriteGenericOld},
 			modelToGet:   "model1",
 			wantRule:     &ruleModel1V1,
+			wantName:     rewriteOld.Name,
 			wantGetAll:   []*v1alpha2.InferenceModelRewrite{rewriteOld, rewriteGenericOld},
 		},
 		{
 			name:         "Precedence: Fallback to generic when no exact match",
 			initialState: []*v1alpha2.InferenceModelRewrite{rewriteOld, rewriteGenericOld},
 			modelToGet:   "model2",
 			wantRule:     &ruleGeneric,
+			wantName:     rewriteGenericOld.Name,
 			wantGetAll:   []*v1alpha2.InferenceModelRewrite{rewriteOld, rewriteGenericOld},
 		},
 		{
 			name:         "Precedence: Oldest exact match wins",
 			initialState: []*v1alpha2.InferenceModelRewrite{rewriteNew, rewriteOld},
 			modelToGet:   "model1",
 			wantRule:     &ruleModel1V1,
+			wantName:     rewriteOld.Name,
 			wantGetAll:   []*v1alpha2.InferenceModelRewrite{rewriteNew, rewriteOld},
 		},
 		{
 			name:         "Precedence: Oldest generic match wins",
 			initialState: []*v1alpha2.InferenceModelRewrite{rewriteGenericNew, rewriteGenericOld},
 			modelToGet:   "any-model",
 			wantRule:     &ruleGeneric,
+			wantName:     rewriteGenericOld.Name,
 			wantGetAll:   []*v1alpha2.InferenceModelRewrite{rewriteGenericNew, rewriteGenericOld},
 		},
 		{
@@ -132,6 +140,7 @@ func TestModelRewriteStore(t *testing.T) {
 			},
 			modelToGet: "model1",
 			wantRule:   &ruleGeneric, // Falls back to generic
+			wantName:   rewriteGenericOld.Name,
 			wantGetAll: []*v1alpha2.InferenceModelRewrite{rewriteGenericOld},
 		},
 		{
@@ -142,6 +151,7 @@ func TestModelRewriteStore(t *testing.T) {
 			},
 			modelToGet: "model1",
 			wantRule:   &ruleModel1V1,
+			wantName:   rewriteOld.Name,
 			wantGetAll: []*v1alpha2.InferenceModelRewrite{rewriteOld},
 		},
 		{
@@ -152,6 +162,7 @@ func TestModelRewriteStore(t *testing.T) {
 			},
 			modelToGet: "model1",
 			wantRule:   &ruleModel1V2,
+			wantName:   rewriteUpdated.Name,
 			wantGetAll: []*v1alpha2.InferenceModelRewrite{rewriteUpdated},
 		},
 	}
@@ -167,10 +178,13 @@ func TestModelRewriteStore(t *testing.T) {
 				tc.op(store)
 			}
 
-			gotRule := store.getRule(tc.modelToGet)
+			gotRule, gotName := store.getRule(tc.modelToGet)
 			if diff := cmp.Diff(tc.wantRule, gotRule); diff != "" {
 				t.Errorf("GetRule() mismatch (-want +got):\n%s", diff)
 			}
+			if gotName != tc.wantName {
+				t.Errorf("GetRule() returned incorrect name: got %s, want %s", gotName, tc.wantName)
+			}
 
 			if tc.wantGetAll != nil {
 				gotAll := store.getAll()

pkg/epp/metrics/metrics.go

@@ -373,6 +373,16 @@ var (
 		},
 		[]string{"fairness_id", "priority"},
 	)
+
+	// Inference Model Rewrite Metrics
+	inferenceModelRewriteDecisionsTotal = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Subsystem: InferenceExtension,
+			Name:      "model_rewrite_decisions_total",
+			Help:      metricsutil.HelpMsgWithStability("Total number of inference model rewrite decisions.", compbasemetrics.ALPHA),
+		},
+		[]string{"model_rewrite_name", "model_name", "target_model"},
+	)
 )
 
 var registerMetrics sync.Once
@@ -416,6 +426,7 @@ func Register(customCollectors ...prometheus.Collector) {
 		metrics.Registry.MustRegister(PrefixCacheHitLength)
 		metrics.Registry.MustRegister(flowControlRequestQueueDuration)
 		metrics.Registry.MustRegister(flowControlQueueSize)
+		metrics.Registry.MustRegister(inferenceModelRewriteDecisionsTotal)
 		for _, collector := range customCollectors {
 			metrics.Registry.MustRegister(collector)
 		}
@@ -460,6 +471,7 @@ func Reset() {
 	PrefixCacheHitLength.Reset()
 	flowControlRequestQueueDuration.Reset()
 	flowControlQueueSize.Reset()
+	inferenceModelRewriteDecisionsTotal.Reset()
 }
 
 // RecordRequstCounter records the number of requests.
@@ -737,3 +749,8 @@ func SetTTFTSLOThreshold(modelName, targetModelName string, threshold float64) {
 func SetTPOTSLOThreshold(modelName, targetModelName string, threshold float64) {
 	inferenceGauges.With(prometheus.Labels{"model_name": modelName, "target_model_name": targetModelName, "type": "tpot_slo_threshold"}).Set(threshold)
 }
+
+// RecordInferenceModelRewriteDecision records the routing decision for InferenceModelRewrite.
+func RecordInferenceModelRewriteDecision(modelRewriteName, modelName, targetModel string) {
+	inferenceModelRewriteDecisionsTotal.WithLabelValues(modelRewriteName, modelName, targetModel).Inc()
+}

pkg/epp/metrics/metrics_test.go

@@ -887,3 +887,36 @@ func TestFlowControlQueueSizeMetric(t *testing.T) {
 	require.NoError(t, err, "Failed to get gauge value for non-existent user-c/100")
 	require.Equal(t, 0.0, val, "Gauge value for non-existent labels should be 0")
 }
+
+func TestInferenceModelRewriteDecisionsTotalMetric(t *testing.T) {
+	Reset()
+
+	RecordInferenceModelRewriteDecision("rewrite-rule-1", "model-a", "model-b")
+	RecordInferenceModelRewriteDecision("rewrite-rule-1", "model-a", "model-b")
+	RecordInferenceModelRewriteDecision("rewrite-rule-2", "model-c", "model-d")
+
+	testCases := []struct {
+		name        string
+		labels      prometheus.Labels
+		expectCount float64
+	}{
+		{
+			name:        "rewrite-rule-1, model-a -> model-b",
+			labels:      prometheus.Labels{"model_rewrite_name": "rewrite-rule-1", "model_name": "model-a", "target_model": "model-b"},
+			expectCount: 2,
+		},
+		{
+			name:        "rewrite-rule-2, model-c -> model-d",
+			labels:      prometheus.Labels{"model_rewrite_name": "rewrite-rule-2", "model_name": "model-c", "target_model": "model-d"},
+			expectCount: 1,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			val, err := testutil.GetCounterMetricValue(inferenceModelRewriteDecisionsTotal.With(tc.labels))
+			require.NoError(t, err, "Failed to get counter value for labels %v", tc.labels)
+			require.Equal(t, tc.expectCount, val, "Counter value mismatch for labels %v", tc.labels)
+		})
+	}
+}

pkg/epp/requestcontrol/director.go

@@ -51,7 +51,8 @@ type Datastore interface {
 	PoolGet() (*datalayer.EndpointPool, error)
 	ObjectiveGet(objectiveName string) *v1alpha2.InferenceObjective
 	PodList(predicate func(backendmetrics.PodMetrics) bool) []backendmetrics.PodMetrics
-	ModelRewriteGet(modelName string) *v1alpha2.InferenceModelRewriteRule
+	// ModelRewriteGet returns the rewrite rule for a given model name and the name of the InferenceModelRewrite object.
+	ModelRewriteGet(modelName string) (*v1alpha2.InferenceModelRewriteRule, string)
 }
 
 // Scheduler defines the interface required by the Director for scheduling.
@@ -195,11 +196,12 @@ func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestCo
 }
 
 func (d *Director) applyWeightedModelRewrite(reqCtx *handlers.RequestContext) {
-	rewriteRule := d.datastore.ModelRewriteGet(reqCtx.IncomingModelName)
+	rewriteRule, modelRewriteName := d.datastore.ModelRewriteGet(reqCtx.IncomingModelName)
 	if rewriteRule == nil {
 		return
 	}
 	reqCtx.TargetModelName = d.selectWeightedModel(rewriteRule.Targets)
+	metrics.RecordInferenceModelRewriteDecision(modelRewriteName, reqCtx.IncomingModelName, reqCtx.TargetModelName)
 }
 
 func (d *Director) selectWeightedModel(models []v1alpha2.TargetModel) string {

pkg/epp/requestcontrol/director_test.go

@@ -169,7 +169,7 @@ func (m mockProducedDataType) Clone() datalayer.Cloneable {
 	return mockProducedDataType{value: m.value}
 }
 
-func (ds *mockDatastore) ModelRewriteGet(modelName string) *v1alpha2.InferenceModelRewriteRule {
+func (ds *mockDatastore) ModelRewriteGet(modelName string) (*v1alpha2.InferenceModelRewriteRule, string) {
 	// This mock implementation simulates the precedence logic for simplicity.
 	// It finds the oldest rewrite that has a rule matching the modelName.
 	var matchingRewrites []*v1alpha2.InferenceModelRewrite
@@ -185,7 +185,7 @@ func (ds *mockDatastore) ModelRewriteGet(modelName string) *v1alpha2.InferenceMo
 	}
 
 	if len(matchingRewrites) == 0 {
-		return nil
+		return nil, ""
 	}
 
 	// Sort by timestamp to find the oldest.
@@ -194,7 +194,7 @@ func (ds *mockDatastore) ModelRewriteGet(modelName string) *v1alpha2.InferenceMo
 	})
 
 	// Return the first rule from the oldest rewrite.
-	return &matchingRewrites[0].Spec.Rules[0]
+	return &matchingRewrites[0].Spec.Rules[0], matchingRewrites[0].Name
 }
 
 func TestDirector_HandleRequest(t *testing.T) {