Execute prepare data plugins sequentially with retries and timeout. Also added more tests and some refactoring

Author: rahulgurnani

pkg/epp/backend/metrics/types.go

@@ -75,8 +75,4 @@ func (f *PodMetricsFactory) ReleaseEndpoint(ep PodMetrics) {
 	}
 }
 
-func (f *PodMetricsFactory) GetAttributes() *datalayer.Attributes {
-	return datalayer.NewAttributes()
-}
-
 type PodMetrics = datalayer.Endpoint

pkg/epp/requestcontrol/director.go

@@ -24,7 +24,6 @@ import (
 	"math/rand"
 	"net"
 	"strings"
-	"sync"
 	"time"
 
 	"sigs.k8s.io/controller-runtime/pkg/log"
@@ -175,7 +174,9 @@ func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestCo
 
 	// Prepare per request data by running PrepareData plugins.
 	// NOTE: Failure in prepare data plugins does not block the request processing.
-	d.runPrepareDataPlugins(ctx, reqCtx.SchedulingRequest, snapshotOfCandidatePods)
+	if d.runPrepareDataPlugins(ctx, reqCtx.SchedulingRequest, snapshotOfCandidatePods) != nil {
+		return reqCtx, errutil.Error{Code: errutil.Internal, Msg: "failed to prepare request data"}
+	}
 
 	// Run admit request plugins
 	if !d.withAdmissionPlugins(ctx, reqCtx.SchedulingRequest, snapshotOfCandidatePods) {
@@ -352,45 +353,57 @@ func (d *Director) runPreRequestPlugins(ctx context.Context, request *scheduling
 	}
 }
 
-// prepareData executes the PrepareRequestData plugins with retries and timeout.
-func prepareData(plugin DataProducer, ctx context.Context, request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) {
+// executePlugins executes PrepareDataPlugins sequentially.
+// TODO: Change to DAG execution in the following PRs.
+func (d *Director) executePlugins(ctx context.Context, request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod, plugins []PrepareDataPlugin) error {
+	for _, plugin := range plugins {
+		err := prepareDataWithRetriesAndTimeout(plugin, ctx, request, pods)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// prepareDataWithRetriesAndTimeout executes the PrepareRequestData plugins with retries and timeout.
+func prepareDataWithRetriesAndTimeout(plugin PrepareDataPlugin, ctx context.Context, request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) error {
 	currentTimeout := prepareDataTimeout
 	for i := 0; i <= prepareDataMaxRetries; i++ {
-		done := make(chan struct{})
+		errCh := make(chan error, 1)
 		go func() {
-			defer close(done)
-			plugin.PrepareRequestData(ctx, request, pods)
+			errCh <- plugin.PrepareRequestData(ctx, request, pods)
 		}()
 
 		select {
-		case <-done:
-			// Plugin executed successfully
-			return
+		case <-ctx.Done():
+			return ctx.Err()
+		case err := <-errCh:
+			if err != nil {
+				log.FromContext(ctx).V(logutil.DEBUG).Info("PrepareData plugin failed, retrying...", "plugin", plugin.TypedName(), "retry", i+1, "error", err)
+				continue
+			}
+			return nil // Success
 		case <-time.After(currentTimeout):
 			log.FromContext(ctx).V(logutil.DEBUG).Info("PrepareData plugin timed out, retrying...", "plugin", plugin.TypedName(), "retry", i+1, "timeout", currentTimeout)
 			if i == prepareDataMaxRetries {
-				log.FromContext(ctx).Error(nil, "PrepareData plugin failed after multiple retries", "plugin", plugin.TypedName())
-				return
+				return fmt.Errorf("PrepareData plugin %s failed after %d retries", plugin.TypedName().String(), prepareDataMaxRetries)
 			}
 		}
 	}
+	return nil
 }
 
 func (d *Director) runPrepareDataPlugins(ctx context.Context,
-	request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) {
-	loggerDebug := log.FromContext(ctx).V(logutil.DEBUG)
+	request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) error {
 	// Parallelly execute PrepareData for all the plugins. Some plugins might take time to prepare data e.g. latency predictor.
 	// Failure in any prepareData doesn't block the request processing.
-	var wg sync.WaitGroup
-	for _, plugin := range d.requestControlPlugins.dataProducerPlugins {
-		loggerDebug.Info("Running PrepareData plugin", "plugin", plugin.TypedName())
-		wg.Add(1)
-		go func(p DataProducer) {
-			defer wg.Done()
-			prepareData(p, ctx, request, pods)
-		}(plugin)
+	err := d.executePlugins(ctx, request, pods, d.requestControlPlugins.prepareDataPlugins)
+	if err != nil {
+		log.FromContext(ctx).Error(err, "failed to execute PrepareData plugins as DAG, falling back to parallel execution")
+		return err
 	}
-	wg.Wait()
+
+	return nil
 }
 
 func (d *Director) withAdmissionPlugins(ctx context.Context,

pkg/epp/requestcontrol/director_test.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"maps"
 	"testing"
 	"time"
 
@@ -105,38 +106,23 @@ func (ds *mockDatastore) PodList(predicate func(backendmetrics.PodMetrics) bool)
 	return res
 }
 
-type mockDataProducerPlugin struct {
-	tn plugins.TypedName
-}
-
-func newMockDataProducerPlugin(name string) *mockDataProducerPlugin {
-	return &mockDataProducerPlugin{
-		tn: plugins.TypedName{Type: "mock-prepare-request-data", Name: name},
+func newMockPrepareDataPlugin(name string) *mockPrepareDataPlugin {
+	return &mockPrepareDataPlugin{
+		name:     name,
+		produces: map[string]any{mockProducedDataKey: 0},
+		consumes: map[string]any{},
 	}
 }
 
-func (m *mockDataProducerPlugin) TypedName() plugins.TypedName {
-	return m.tn
-}
-
-func (m *mockDataProducerPlugin) Produces() map[string]any {
-	// Produces data of type int, 0 denotes it is int.
-	return map[string]any{mockProducedDataKey: 0}
-}
-
-func (m *mockDataProducerPlugin) PrepareRequestData(ctx context.Context, request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) {
-	pods[0].Put(mockProducedDataKey, mockProducedDataType{value: 42})
-}
-
 type mockAdmissionPlugin struct {
-	tn plugins.TypedName
-	// TODO: Replace this will admission control.
-	admitRequestCalled bool
+	tn          plugins.TypedName
+	denialError error
 }
 
-func newMockAdmissionPlugin(name string) *mockAdmissionPlugin {
+func newMockAdmissionPlugin(name string, denialError error) *mockAdmissionPlugin {
 	return &mockAdmissionPlugin{
-		tn: plugins.TypedName{Type: "mock-admit-data", Name: name},
+		tn:          plugins.TypedName{Type: "mock-admit-data", Name: name},
+		denialError: denialError,
 	}
 }
 
@@ -145,8 +131,7 @@ func (m *mockAdmissionPlugin) TypedName() plugins.TypedName {
 }
 
 func (m *mockAdmissionPlugin) AdmitRequest(ctx context.Context, request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) error {
-	m.admitRequestCalled = true
-	return nil
+	return m.denialError
 }
 
 type mockProducedDataType struct {
@@ -279,9 +264,8 @@ 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
-		admitRequestCalled      bool
-		dataProducerPlugin      *mockDataProducerPlugin
-		admissionPlugin         *mockAdmissionPlugin
+		admitRequestDenialError error                    // Expected denial error from admission plugin
+		prepareDataPlugin       *mockPrepareDataPlugin
 	}{
 		{
 			name: "successful completions request",
@@ -364,7 +348,7 @@ func TestDirector_HandleRequest(t *testing.T) {
 			},
 			wantMutatedBodyModel: model,
 			targetModelName:      model,
-			dataProducerPlugin:   newMockDataProducerPlugin("test-plugin"),
+			prepareDataPlugin:    newMockPrepareDataPlugin("test-plugin"),
 		},
 		{
 			name: "successful chat completions request with admit request plugins",
@@ -391,10 +375,29 @@ func TestDirector_HandleRequest(t *testing.T) {
 				},
 				TargetEndpoint: "192.168.1.100:8000,192.168.2.100:8000,192.168.4.100:8000",
 			},
-			wantMutatedBodyModel: model,
-			targetModelName:      model,
-			admitRequestCalled:   true,
-			admissionPlugin:      newMockAdmissionPlugin("test-plugin"),
+			wantMutatedBodyModel:    model,
+			targetModelName:         model,
+			admitRequestDenialError: nil,
+		},
+		{
+			name: "denied request by admit request plugin",
+			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
+			},
+			wantMutatedBodyModel:    model,
+			targetModelName:         model,
+			admitRequestDenialError: errors.New("denied by admit plugin"),
+			wantErrCode:             errutil.Internal,
 		},
 		{
 			name: "successful chat completions request with multiple messages",
@@ -546,12 +549,10 @@ func TestDirector_HandleRequest(t *testing.T) {
 				test.schedulerMockSetup(mockSched)
 			}
 			config := NewConfig()
-			if test.dataProducerPlugin != nil {
-				config = config.WithDataProducers(test.dataProducerPlugin)
-			}
-			if test.admissionPlugin != nil {
-				config = config.WithAdmissionPlugins(test.admissionPlugin)
+			if test.prepareDataPlugin != nil {
+				config = config.WithPrepareDataPlugins(test.prepareDataPlugin)
 			}
+			config = config.WithAdmissionPlugins(newMockAdmissionPlugin("test-admit-plugin", test.admitRequestDenialError))
 			director := NewDirectorWithConfig(ds, mockSched, test.mockAdmissionController, config)
 
 			reqCtx := &handlers.RequestContext{
@@ -566,9 +567,7 @@ func TestDirector_HandleRequest(t *testing.T) {
 				TargetModelName: test.targetModelName,
 			}
 			// Deep copy the body map.
-			for k, v := range test.reqBodyMap {
-				reqCtx.Request.Body[k] = v
-			}
+			maps.Copy(reqCtx.Request.Body, test.reqBodyMap)
 
 			returnedReqCtx, err := director.HandleRequest(ctx, reqCtx)
 
@@ -596,9 +595,6 @@ func TestDirector_HandleRequest(t *testing.T) {
 				assert.Equal(t, test.wantMutatedBodyModel, returnedReqCtx.Request.Body["model"],
 					"Mutated reqCtx.Request.Body model mismatch")
 			}
-			if test.admissionPlugin != nil {
-				assert.True(t, test.admissionPlugin.admitRequestCalled, "AdmitRequest not called")
-			}
 		})
 	}
 }

pkg/epp/requestcontrol/plugins.go

@@ -59,10 +59,11 @@ type ResponseComplete interface {
 }
 
 // PrepareRequestData is called by the director before scheduling requests.
-// DataProducer plugin is implemented by data producers which produce data from different sources.
-type DataProducer interface {
+// PrepareDataPlugin plugin is implemented by data producers which produce data from different sources.
+type PrepareDataPlugin interface {
 	plugins.ProducerPlugin
-	PrepareRequestData(ctx context.Context, request *types.LLMRequest, pods []types.Pod)
+	plugins.ConsumerPlugin
+	PrepareRequestData(ctx context.Context, request *types.LLMRequest, pods []types.Pod) error
 }
 
 // AdmissionPlugin is called by the director after the prepare data phase and before scheduling.

pkg/epp/requestcontrol/request_control_config.go

@@ -24,7 +24,7 @@ import (
 func NewConfig() *Config {
 	return &Config{
 		admissionPlugins:         []AdmissionPlugin{},
-		dataProducerPlugins:      []DataProducer{},
+		prepareDataPlugins:       []PrepareDataPlugin{},
 		preRequestPlugins:        []PreRequest{},
 		responseReceivedPlugins:  []ResponseReceived{},
 		responseStreamingPlugins: []ResponseStreaming{},
@@ -35,7 +35,7 @@ func NewConfig() *Config {
 // Config provides a configuration for the requestcontrol plugins.
 type Config struct {
 	admissionPlugins         []AdmissionPlugin
-	dataProducerPlugins      []DataProducer
+	prepareDataPlugins       []PrepareDataPlugin
 	preRequestPlugins        []PreRequest
 	responseReceivedPlugins  []ResponseReceived
 	responseStreamingPlugins []ResponseStreaming
@@ -70,9 +70,9 @@ func (c *Config) WithResponseCompletePlugins(plugins ...ResponseComplete) *Confi
 	return c
 }
 
-// WithDataProducers sets the given plugins as the PrepareData plugins.
-func (c *Config) WithDataProducers(plugins ...DataProducer) *Config {
-	c.dataProducerPlugins = plugins
+// WithPrepareDataPlugins sets the given plugins as the PrepareData plugins.
+func (c *Config) WithPrepareDataPlugins(plugins ...PrepareDataPlugin) *Config {
+	c.prepareDataPlugins = plugins
 	return c
 }
 
@@ -85,7 +85,6 @@ func (c *Config) WithAdmissionPlugins(plugins ...AdmissionPlugin) *Config {
 // 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.
-
 func (c *Config) AddPlugins(pluginObjects ...plugins.Plugin) {
 	for _, plugin := range pluginObjects {
 		if preRequestPlugin, ok := plugin.(PreRequest); ok {
@@ -100,5 +99,8 @@ func (c *Config) AddPlugins(pluginObjects ...plugins.Plugin) {
 		if responseCompletePlugin, ok := plugin.(ResponseComplete); ok {
 			c.responseCompletePlugins = append(c.responseCompletePlugins, responseCompletePlugin)
 		}
+		if prepareDataPlugin, ok := plugin.(PrepareDataPlugin); ok {
+			c.prepareDataPlugins = append(c.prepareDataPlugins, prepareDataPlugin)
+		}
 	}
 }

pkg/epp/scheduling/framework/plugins/multi/prefix/plugin.go

@@ -209,7 +209,6 @@ func (p *Plugin) Score(ctx context.Context, cycleState *types.CycleState, reques
 	log.FromContext(ctx).V(logutil.TRACE).Info("prefix cached state", "cached-servers", state.PrefixCacheServers, "hashes", state.PrefixHashes)
 	// calculate the scores of pods
 	scores := make(map[types.Pod]float64, len(pods))
-
 	total := len(state.PrefixHashes)
 	podScoreFunc := func(pod types.Pod) float64 {
 		if total == 0 {