Add unit tests and comments

Author: rahulgurnani

pkg/epp/requestcontrol/director.go

@@ -91,7 +91,7 @@ type Director struct {
 	defaultPriority int
 }
 
-// getInferenceObjective creates inferenceObjective based on reqCtx.
+// getInferenceObjective fetches the inferenceObjective from the datastore otherwise creates a new one based on reqCtx.
 func (d *Director) getInferenceObjective(logger logr.Logger, reqCtx *handlers.RequestContext) *v1alpha2.InferenceObjective {
 	infObjective := d.datastore.ObjectiveGet(reqCtx.ObjectiveKey)
 	if infObjective == nil {
@@ -109,7 +109,7 @@ func (d *Director) getInferenceObjective(logger logr.Logger, reqCtx *handlers.Re
 }
 
 // resolveTargetModel is a helper that resolves targetModel
-// and updates the reqCtx and ctx.
+// and updates the reqCtx.
 func (d *Director) resolveTargetModel(reqCtx *handlers.RequestContext) error {
 	requestBodyMap := reqCtx.Request.Body
 	var ok bool
@@ -168,18 +168,19 @@ func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestCo
 		logger.V(logutil.DEFAULT).Info("Request rejected by admission control", "error", err)
 		return reqCtx, err
 	}
-	copyOfCandidatePods := d.toSchedulerPodMetrics(candidatePods)
+	snapshotOfCandidatePods := d.toSchedulerPodMetrics(candidatePods)
 
 	// Prepare per request data
-	d.runPrepareDataPlugins(ctx, reqCtx.SchedulingRequest, copyOfCandidatePods)
+	// TODO(rahulgurnani): Add retries and timeout in the preparedata step.
+	d.runPrepareDataPlugins(ctx, reqCtx.SchedulingRequest, snapshotOfCandidatePods)
 
 	// Run admit request plugins
-	if !d.runAdmitRequestPlugins(ctx, reqCtx.SchedulingRequest, copyOfCandidatePods) {
+	if !d.runAdmitRequestPlugins(ctx, reqCtx.SchedulingRequest, snapshotOfCandidatePods) {
 		logger.V(logutil.DEFAULT).Info("Request cannot be admitted")
 		return reqCtx, errutil.Error{Code: errutil.Internal, Msg: "request cannot be admitted"}
 	}
 
-	result, err := d.scheduler.Schedule(ctx, reqCtx.SchedulingRequest, copyOfCandidatePods)
+	result, err := d.scheduler.Schedule(ctx, reqCtx.SchedulingRequest, snapshotOfCandidatePods)
 	if err != nil {
 		return reqCtx, errutil.Error{Code: errutil.InferencePoolResourceExhausted, Msg: fmt.Errorf("failed to find target pod: %w", err).Error()}
 	}

pkg/epp/requestcontrol/director_test.go

@@ -93,6 +93,45 @@ func (ds *mockDatastore) PodList(predicate func(backendmetrics.PodMetrics) bool)
 	return res
 }
 
+type mockPrepareRequestDataPlugin struct {
+	tn                plugins.TypedName
+	prepareDataCalled bool
+}
+
+func newmockPrepareRequestDataPlugin(name string) *mockPrepareRequestDataPlugin {
+	return &mockPrepareRequestDataPlugin{
+		tn: plugins.TypedName{Type: "mock-prepare-request-data", Name: name},
+	}
+}
+
+func (m *mockPrepareRequestDataPlugin) TypedName() plugins.TypedName {
+	return m.tn
+}
+
+func (m *mockPrepareRequestDataPlugin) PrepareData(ctx context.Context, request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) {
+	m.prepareDataCalled = true
+}
+
+type mockAdmitRequestPlugins struct {
+	tn                 plugins.TypedName
+	admitRequestCalled bool
+}
+
+func newmockAdmitRequestPlugins(name string) *mockAdmitRequestPlugins {
+	return &mockAdmitRequestPlugins{
+		tn: plugins.TypedName{Type: "mock-admit-data", Name: name},
+	}
+}
+
+func (m *mockAdmitRequestPlugins) TypedName() plugins.TypedName {
+	return m.tn
+}
+
+func (m *mockAdmitRequestPlugins) Admit(ctx context.Context, request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) bool {
+	m.admitRequestCalled = true
+	return true
+}
+
 func TestDirector_HandleRequest(t *testing.T) {
 	ctx := logutil.NewTestLoggerIntoContext(context.Background())
 
@@ -211,6 +250,10 @@ func TestDirector_HandleRequest(t *testing.T) {
 		wantReqCtx              *handlers.RequestContext // Fields to check in the returned RequestContext
 		wantMutatedBodyModel    string                   // Expected model in reqCtx.Request.Body after PostDispatch
 		targetModelName         string                   // Expected model name after target model resolution
+		prepareDataCalled       bool
+		admitRequestCalled      bool
+		prepareDataPlugins      *mockPrepareRequestDataPlugin
+		admitRequestPlugins     *mockAdmitRequestPlugins
 	}{
 		{
 			name: "successful completions request",
@@ -265,6 +308,66 @@ func TestDirector_HandleRequest(t *testing.T) {
 			wantMutatedBodyModel: model,
 			targetModelName:      model,
 		},
+		{
+			name: "successful chat completions request with prepare data plugins",
+			reqBodyMap: map[string]any{
+				"model": model,
+				"messages": []any{
+					map[string]any{
+						"role":    "user",
+						"content": "critical prompt",
+					},
+				},
+			},
+			mockAdmissionController: &mockAdmissionController{admitErr: nil},
+			schedulerMockSetup: func(m *mockScheduler) {
+				m.scheduleResults = defaultSuccessfulScheduleResults
+			},
+			wantReqCtx: &handlers.RequestContext{
+				TargetModelName: model,
+				TargetPod: &backend.Pod{
+					NamespacedName: types.NamespacedName{Namespace: "default", Name: "pod1"},
+					Address:        "192.168.1.100",
+					Port:           "8000",
+					MetricsHost:    "192.168.1.100:8000",
+				},
+				TargetEndpoint: "192.168.1.100:8000,192.168.2.100:8000,192.168.4.100:8000",
+			},
+			wantMutatedBodyModel: model,
+			targetModelName:      model,
+			prepareDataCalled:    true,
+			prepareDataPlugins:   newmockPrepareRequestDataPlugin("test-plugin"),
+		},
+		{
+			name: "successful chat completions request with admit request plugins",
+			reqBodyMap: map[string]any{
+				"model": model,
+				"messages": []any{
+					map[string]any{
+						"role":    "user",
+						"content": "critical prompt",
+					},
+				},
+			},
+			mockAdmissionController: &mockAdmissionController{admitErr: nil},
+			schedulerMockSetup: func(m *mockScheduler) {
+				m.scheduleResults = defaultSuccessfulScheduleResults
+			},
+			wantReqCtx: &handlers.RequestContext{
+				TargetModelName: model,
+				TargetPod: &backend.Pod{
+					NamespacedName: types.NamespacedName{Namespace: "default", Name: "pod1"},
+					Address:        "192.168.1.100",
+					Port:           "8000",
+					MetricsHost:    "192.168.1.100:8000",
+				},
+				TargetEndpoint: "192.168.1.100:8000,192.168.2.100:8000,192.168.4.100:8000",
+			},
+			wantMutatedBodyModel: model,
+			targetModelName:      model,
+			admitRequestCalled:   true,
+			admitRequestPlugins:  newmockAdmitRequestPlugins("test-plugin"),
+		},
 		{
 			name: "successful chat completions request with multiple messages",
 			reqBodyMap: map[string]any{
@@ -414,7 +517,14 @@ func TestDirector_HandleRequest(t *testing.T) {
 			if test.schedulerMockSetup != nil {
 				test.schedulerMockSetup(mockSched)
 			}
-			director := NewDirectorWithConfig(ds, mockSched, test.mockAdmissionController, NewConfig())
+			config := NewConfig()
+			if test.prepareDataPlugins != nil {
+				config = config.WithPrepareDataPlugins(test.prepareDataPlugins)
+			}
+			if test.admitRequestPlugins != nil {
+				config = config.WithAdmitRequestPlugins(test.admitRequestPlugins)
+			}
+			director := NewDirectorWithConfig(ds, mockSched, test.mockAdmissionController, config)
 
 			reqCtx := &handlers.RequestContext{
 				Request: &handlers.Request{
@@ -458,6 +568,12 @@ func TestDirector_HandleRequest(t *testing.T) {
 				assert.Equal(t, test.wantMutatedBodyModel, returnedReqCtx.Request.Body["model"],
 					"Mutated reqCtx.Request.Body model mismatch")
 			}
+			if test.admitRequestPlugins != nil {
+				assert.True(t, test.admitRequestPlugins.admitRequestCalled, "AdmitRequestPlugins not called")
+			}
+			if test.prepareDataPlugins != nil {
+				assert.True(t, test.prepareDataPlugins.prepareDataCalled, "PrepareDataPlugins not called")
+			}
 		})
 	}
 }

pkg/epp/requestcontrol/request_control_config.go

@@ -70,6 +70,18 @@ func (c *Config) WithResponseCompletePlugins(plugins ...ResponseComplete) *Confi
 	return c
 }
 
+// WithPrepareDataPlugins sets the given plugins as the PrepareData plugins.
+func (c *Config) WithPrepareDataPlugins(plugins ...PrepareData) *Config {
+	c.prepareDataPlugins = plugins
+	return c
+}
+
+// WithAdmitRequestPlugins sets the given plugins as the AdmitRequest plugins.
+func (c *Config) WithAdmitRequestPlugins(plugins ...AdmitRequest) *Config {
+	c.admitRequestPlugins = plugins
+	return c
+}
+
 // AddPlugins adds the given plugins to the Config.
 // The type of each plugin is checked and added to the corresponding list of plugins in the Config.
 // If a plugin implements multiple plugin interfaces, it will be added to each corresponding list.