Implement error propagation for prepare data plugins. Also add retry timeout logic back since it got removed in previous commit.

Author: rahulgurnani

pkg/epp/requestcontrol/director.go

@@ -366,27 +366,68 @@ func (d *Director) executePluginsAsDAG(ctx context.Context, request *schedulingt
 	// Execute the DAG
 
 	// Channels to signal plugin execution completion.
-	pluginExecuted := map[string]chan struct{}{}
+	pluginExecuted := make(map[string]chan error)
 	nameToNode := map[string]PrepareDataPlugin{}
 	for _, plugin := range plugins {
-		pluginExecuted[plugin.TypedName().String()] = make(chan struct{})
+		pluginExecuted[plugin.TypedName().String()] = make(chan error)
 		nameToNode[plugin.TypedName().String()] = plugin
 	}
 
 	for pluginName, dependents := range dag {
 		// Execute plugins based on dependencies.
 		//  Wait for the dependencies to complete before executing a plugin.
-		go func() {
+		go func() error {
 			for _, dep := range dependents {
-				<-pluginExecuted[dep]
+				err, open := <-pluginExecuted[dep]
+				if !open {
+					continue
+				}
+				if err != nil {
+					// If a dependency failed, propagate the error and do not execute this plugin.
+					pluginExecuted[pluginName] <- fmt.Errorf("dependency plugin %s failed: %w", dep, err)
+					return err
+				}
 			}
-			nameToNode[pluginName].PrepareRequestData(ctx, request, pods)
 			// Signal that the plugin has been executed.
-			close(pluginExecuted[pluginName])
+			defer close(pluginExecuted[pluginName])
+
+			return prepareDataWithRetriesAndTimeout(nameToNode[pluginName], ctx, request, pods)
 		}()
 	}
 	for pluginName := range dag {
-		<-pluginExecuted[pluginName]
+		err, open := <-pluginExecuted[pluginName]
+		if !open {
+			continue
+		}
+		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++ {
+		errCh := make(chan error, 1)
+		go func() {
+			errCh <- plugin.PrepareRequestData(ctx, request, pods)
+		}()
+
+		select {
+		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 {
+				return fmt.Errorf("PrepareData plugin %s failed after %d retries", plugin.TypedName().String(), prepareDataMaxRetries)
+			}
+		}
 	}
 	return nil
 }

pkg/epp/requestcontrol/graph_util_test.go

@@ -44,8 +44,9 @@ func (m *mockPrepareDataPlugin) Consumes() map[string]any {
 	return m.consumes
 }
 
-func (m *mockPrepareDataPlugin) PrepareRequestData(ctx context.Context, request *types.LLMRequest, pods []types.Pod) {
+func (m *mockPrepareDataPlugin) PrepareRequestData(ctx context.Context, request *types.LLMRequest, pods []types.Pod) error {
 	pods[0].Put(mockProducedDataKey, mockProducedDataType{value: 42})
+	return nil
 }
 
 func TestPrepareDataGraph(t *testing.T) {

pkg/epp/requestcontrol/plugins.go

@@ -63,7 +63,7 @@ type ResponseComplete interface {
 type PrepareDataPlugin interface {
 	plugins.ProducerPlugin
 	plugins.ConsumerPlugin
-	PrepareRequestData(ctx context.Context, request *types.LLMRequest, pods []types.Pod)
+	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.