add integration test, rename datasoter interface, change split to targets.

Author: zetxqx

apix/v1alpha2/inferencemodelrewrite_types.go

@@ -92,7 +92,7 @@ type InferenceModelRewriteRule struct {
 	// +optional
 	// +kubebuilder:validation:MinItems=1
 	//
-	Targets []TargetModel `json:"split,omitempty"`
+	Targets []TargetModel `json:"targets,omitempty"`
 }
 
 // TargetModel defines a weighted model destination for traffic distribution.

client-go/applyconfiguration/apix/v1alpha2/inferencemodelrewriterule.go

@@ -108,7 +108,7 @@ spec:
                         - model
                         type: object
                       type: array
-                    split:
+                    targets:
                       items:
                         description: TargetModel defines a weighted model destination
                           for traffic distribution.

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

@@ -54,17 +54,20 @@ func (c *InferenceModelRewriteReconciler) Reconcile(ctx context.Context, req ctr
 		notFound = true
 	}
 
-	if notFound || !infModelRewrite.DeletionTimestamp.IsZero() || infModelRewrite.Spec.PoolRef == nil ||
+	isDeleted := !infModelRewrite.DeletionTimestamp.IsZero()
+	isPooRefUnmatch := infModelRewrite.Spec.PoolRef == nil ||
 		infModelRewrite.Spec.PoolRef.Name != v1alpha2.ObjectName(c.PoolGKNN.Name) ||
-		(infModelRewrite.Spec.PoolRef.Group != v1alpha2.Group(c.PoolGKNN.Group) && infModelRewrite.Spec.PoolRef.Group != "inference.networking.x-k8s.io") {
+		infModelRewrite.Spec.PoolRef.Group != v1alpha2.Group(c.PoolGKNN.Group)
+
+	if notFound || isDeleted || isPooRefUnmatch {
 		// InferenceModelRewrite object got deleted or changed the referenced pool.
-		c.Datastore.RewriteDelete(req.NamespacedName)
+		c.Datastore.ModelRewriteDelete(req.NamespacedName)
 		return ctrl.Result{}, nil
 	}
 
 	// Add or update if the InferenceModelRewrite instance has a creation timestamp older than the existing entry of the model.
 	logger = logger.WithValues("poolRef", infModelRewrite.Spec.PoolRef)
-	c.Datastore.RewriteSet(infModelRewrite)
+	c.Datastore.ModelRewriteSet(infModelRewrite)
 	logger.Info("Added/Updated InferenceModelRewrite")
 
 	return ctrl.Result{}, nil

pkg/epp/controller/inferencemodelrewrite_reconciler.go

@@ -185,7 +185,7 @@ func TestInferenceModelRewriteReconciler(t *testing.T) {
 			pmf := backendmetrics.NewPodMetricsFactory(&backendmetrics.FakePodMetricsClient{}, time.Second)
 			ds := datastore.NewDatastore(t.Context(), pmf, 0)
 			for _, r := range test.rewritesInStore {
-				ds.RewriteSet(r)
+				ds.ModelRewriteSet(r)
 			}
 			endpointPool := poolutil.InferencePoolToEndpointPool(poolForRewrite)
 			_ = ds.PoolSet(context.Background(), fakeClient, endpointPool)
@@ -210,8 +210,8 @@ func TestInferenceModelRewriteReconciler(t *testing.T) {
 				t.Errorf("Unexpected result diff (+got/-want): %s", diff)
 			}
 
-			if len(test.wantRewrites) != len(ds.RewriteGetAll()) {
-				t.Errorf("Unexpected number of rewrites; want: %d, got:%d", len(test.wantRewrites), len(ds.RewriteGetAll()))
+			if len(test.wantRewrites) != len(ds.ModelRewriteGetAll()) {
+				t.Errorf("Unexpected number of rewrites; want: %d, got:%d", len(test.wantRewrites), len(ds.ModelRewriteGetAll()))
 			}
 
 			if diff := diffStoreRewrites(ds, test.wantRewrites); diff != "" {
@@ -226,7 +226,7 @@ func diffStoreRewrites(ds datastore.Datastore, wantRewrites []*v1alpha2.Inferenc
 		wantRewrites = []*v1alpha2.InferenceModelRewrite{}
 	}
 
-	gotRewrites := ds.RewriteGetAll()
+	gotRewrites := ds.ModelRewriteGetAll()
 	if diff := cmp.Diff(wantRewrites, gotRewrites); diff != "" {
 		return "rewrites:" + diff
 	}

pkg/epp/controller/inferencemodelrewrite_reconciler_test.go

@@ -60,10 +60,10 @@ type Datastore interface {
 	ObjectiveGetAll() []*v1alpha2.InferenceObjective
 
 	// InferenceModelRewrite operations
-	RewriteSet(infModelRewrite *v1alpha2.InferenceModelRewrite)
-	RewriteDelete(namespacedName types.NamespacedName)
-	RewriteGet(modelName string) *v1alpha2.InferenceModelRewriteRule
-	RewriteGetAll() []*v1alpha2.InferenceModelRewrite
+	ModelRewriteSet(infModelRewrite *v1alpha2.InferenceModelRewrite)
+	ModelRewriteDelete(namespacedName types.NamespacedName)
+	ModelRewriteGet(modelName string) *v1alpha2.InferenceModelRewriteRule
+	ModelRewriteGetAll() []*v1alpha2.InferenceModelRewrite
 
 	// PodList lists pods matching the given predicate.
 	PodList(predicate func(backendmetrics.PodMetrics) bool) []backendmetrics.PodMetrics
@@ -81,7 +81,7 @@ func NewDatastore(parentCtx context.Context, epFactory datalayer.EndpointFactory
 		pool:                   nil,
 		mu:                     sync.RWMutex{},
 		objectives:             make(map[string]*v1alpha2.InferenceObjective),
-		rewrites:               newModelRewriteStore(),
+		modelRewrites:          newModelRewriteStore(),
 		pods:                   &sync.Map{},
 		modelServerMetricsPort: modelServerMetricsPort,
 		epf:                    epFactory,
@@ -103,8 +103,8 @@ type datastore struct {
 	pool *datalayer.EndpointPool
 	// key: InferenceObjective name, value: *InferenceObjective
 	objectives map[string]*v1alpha2.InferenceObjective
-	// rewrites store for InferenceModelRewrite objects.
-	rewrites *modelRewriteStore
+	// modelRewrites store for InferenceModelRewrite objects.
+	modelRewrites *modelRewriteStore
 	// key: types.NamespacedName, value: backendmetrics.PodMetrics
 	pods *sync.Map
 	// modelServerMetricsPort metrics port from EPP command line argument
@@ -118,7 +118,7 @@ func (ds *datastore) Clear() {
 	defer ds.mu.Unlock()
 	ds.pool = nil
 	ds.objectives = make(map[string]*v1alpha2.InferenceObjective)
-	ds.rewrites = newModelRewriteStore()
+	ds.modelRewrites = newModelRewriteStore()
 	// stop all pods go routines before clearing the pods map.
 	ds.pods.Range(func(_, v any) bool {
 		ds.epf.ReleaseEndpoint(v.(backendmetrics.PodMetrics))
@@ -210,28 +210,28 @@ func (ds *datastore) ObjectiveGetAll() []*v1alpha2.InferenceObjective {
 	return res
 }
 
-func (ds *datastore) RewriteSet(infModelRewrite *v1alpha2.InferenceModelRewrite) {
+func (ds *datastore) ModelRewriteSet(infModelRewrite *v1alpha2.InferenceModelRewrite) {
 	ds.mu.Lock()
 	defer ds.mu.Unlock()
-	ds.rewrites.set(infModelRewrite)
+	ds.modelRewrites.set(infModelRewrite)
 }
 
-func (ds *datastore) RewriteDelete(namespacedName types.NamespacedName) {
+func (ds *datastore) ModelRewriteDelete(namespacedName types.NamespacedName) {
 	ds.mu.Lock()
 	defer ds.mu.Unlock()
-	ds.rewrites.delete(namespacedName)
+	ds.modelRewrites.delete(namespacedName)
 }
 
-func (ds *datastore) RewriteGet(modelName string) *v1alpha2.InferenceModelRewriteRule {
+func (ds *datastore) ModelRewriteGet(modelName string) *v1alpha2.InferenceModelRewriteRule {
 	ds.mu.RLock()
 	defer ds.mu.RUnlock()
-	return ds.rewrites.getRule(modelName)
+	return ds.modelRewrites.getRule(modelName)
 }
 
-func (ds *datastore) RewriteGetAll() []*v1alpha2.InferenceModelRewrite {
+func (ds *datastore) ModelRewriteGetAll() []*v1alpha2.InferenceModelRewrite {
 	ds.mu.RLock()
 	defer ds.mu.RUnlock()
-	return ds.rewrites.getAll()
+	return ds.modelRewrites.getAll()
 }
 
 // /// Pods/endpoints APIs ///

pkg/epp/datastore/datastore.go

@@ -50,7 +50,7 @@ type Datastore interface {
 	PoolGet() (*datalayer.EndpointPool, error)
 	ObjectiveGet(objectiveName string) *v1alpha2.InferenceObjective
 	PodList(predicate func(backendmetrics.PodMetrics) bool) []backendmetrics.PodMetrics
-	RewriteGet(modelName string) *v1alpha2.InferenceModelRewriteRule
+	ModelRewriteGet(modelName string) *v1alpha2.InferenceModelRewriteRule
 }
 
 // Scheduler defines the interface required by the Director for scheduling.
@@ -194,7 +194,7 @@ func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestCo
 }
 
 func (d *Director) applyWeightedModelRewrite(reqCtx *handlers.RequestContext) {
-	rewriteRule := d.datastore.RewriteGet(reqCtx.IncomingModelName)
+	rewriteRule := d.datastore.ModelRewriteGet(reqCtx.IncomingModelName)
 	if rewriteRule == nil {
 		return
 	}

pkg/epp/requestcontrol/director.go

@@ -169,7 +169,7 @@ func (m mockProducedDataType) Clone() datalayer.Cloneable {
 	return mockProducedDataType{value: m.value}
 }
 
-func (ds *mockDatastore) RewriteGet(modelName string) *v1alpha2.InferenceModelRewriteRule {
+func (ds *mockDatastore) ModelRewriteGet(modelName string) *v1alpha2.InferenceModelRewriteRule {
 	// 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
@@ -268,7 +268,7 @@ func TestDirector_HandleRequest(t *testing.T) {
 	ds.ObjectiveSet(ioFoodReview)
 	ds.ObjectiveSet(ioFoodReviewResolve)
 	ds.ObjectiveSet(ioFoodReviewSheddable)
-	ds.RewriteSet(rewrite)
+	ds.ModelRewriteSet(rewrite)
 
 	scheme := runtime.NewScheme()
 	_ = clientgoscheme.AddToScheme(scheme)

pkg/epp/requestcontrol/director_test.go

@@ -90,6 +90,8 @@ const (
 	// Model Names
 	modelMyModel         = "my-model"
 	modelMyModelTarget   = "my-model-12345"
+	modelToBeWritten     = "model-to-be-rewritten"
+	modelAfterRewrite    = "rewritten-model"
 	modelSQLLora         = "sql-lora"
 	modelSQLLoraTarget   = "sql-lora-1fdg2"
 	modelSheddable       = "sql-lora-sheddable"
@@ -981,6 +983,42 @@ func TestFullDuplexStreamed_KubeInferenceObjectiveRequest(t *testing.T) {
 				},
 			},
 		},
+		{
+			name:     "rewrite request model",
+			requests: integrationutils.GenerateStreamedRequestSet(logger, "test-rewrite", modelToBeWritten, modelToBeWritten, nil),
+			// Pod 0 will be picked.
+			// Expected flow:
+			// 1. Request asks for "model-to-be-rewritte"
+			// 2. Rewrite rule transforms "model-to-be-rewritten" -> "rewritten-model"
+			// 3. EPP sends request to backend with model "rewritten-model"
+			pods: newPodStates(
+				podState{index: 0, queueSize: 0, kvCacheUsage: 0.1, activeModels: []string{"foo", "rewritten-model"}},
+			),
+			wantMetrics: map[string]string{
+				"inference_objective_request_total": inferenceObjectiveRequestTotal([]label{
+					{"model_name", modelToBeWritten},
+					{"target_model_name", modelAfterRewrite},
+				}),
+			},
+			wantErr: false,
+			wantResponses: integrationutils.NewRequestBufferedResponse(
+				"192.168.1.1:8000",
+				// Note: The prompt remains "test-rewrite", but the model in the JSON body is updated to the *rewritten target* model.
+				fmt.Sprintf(`{"max_tokens":100,"model":%q,"prompt":"test-rewrite","temperature":0}`, modelAfterRewrite),
+				&configPb.HeaderValueOption{
+					Header: &configPb.HeaderValue{
+						Key:      "hi",
+						RawValue: []byte("mom"),
+					},
+				},
+				&configPb.HeaderValueOption{
+					Header: &configPb.HeaderValue{
+						Key:      requtil.RequestIdHeaderKey,
+						RawValue: []byte("test-request-id"),
+					},
+				},
+			),
+		},
 	}
 
 	for _, test := range tests {
@@ -1247,6 +1285,7 @@ func BeforeSuite() func() {
 		_ = testEnv.Stop()
 		_ = k8sClient.DeleteAllOf(context.Background(), &v1.InferencePool{})
 		_ = k8sClient.DeleteAllOf(context.Background(), &v1alpha2.InferenceObjective{})
+		_ = k8sClient.DeleteAllOf(context.Background(), &v1alpha2.InferenceModelRewrite{})
 	}
 }
 
@@ -1299,6 +1338,11 @@ func managerTestOptions(namespace, name string, metricsServerOptions metricsserv
 						namespace: {},
 					},
 				},
+				&v1alpha2.InferenceModelRewrite{}: {
+					Namespaces: map[string]cache.Config{
+						namespace: {},
+					},
+				},
 			},
 		},
 		Controller: crconfig.Controller{

test/integration/epp/hermetic_test.go

@@ -62,3 +62,18 @@ spec:
   priority: 2
   poolRef:
     name: vllm-llama3-8b-instruct-pool
+---
+apiVersion: inference.networking.x-k8s.io/v1alpha2
+kind: InferenceModelRewrite
+metadata:
+  name: rewrite-test
+  namespace: default
+spec:
+  poolRef:
+    name: vllm-llama3-8b-instruct-pool
+  rules:
+  - matches:
+    - model:
+        value: model-to-be-rewritten
+    targets:
+    - modelRewrite: rewritten-model