Cleanup dead code, ensure all preparedata plugins execute

Author: rahulgurnani

pkg/epp/requestcontrol/director.go

@@ -24,7 +24,6 @@ import (
 	"math/rand"
 	"net"
 	"strings"
-	"sync"
 	"time"
 
 	"sigs.k8s.io/controller-runtime/pkg/log"
@@ -354,30 +353,6 @@ func (d *Director) runPreRequestPlugins(ctx context.Context, request *scheduling
 	}
 }
 
-// prepareData executes the PrepareRequestData plugins with retries and timeout.
-func prepareData(plugin PrepareDataPlugin, ctx context.Context, request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) {
-	currentTimeout := prepareDataTimeout
-	for i := 0; i <= prepareDataMaxRetries; i++ {
-		done := make(chan struct{})
-		go func() {
-			defer close(done)
-			plugin.PrepareRequestData(ctx, request, pods)
-		}()
-
-		select {
-		case <-done:
-			// Plugin executed successfully
-			return
-		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
-			}
-		}
-	}
-}
-
 // executePluginsAsDAG executes PrepareData plugins as a DAG based on their dependencies asynchronously.
 // So, a plugin is executed only after all its dependencies have been executed.
 // If there is a cycle or other error in the DAG, it returns an error.
@@ -388,9 +363,9 @@ func (d *Director) executePluginsAsDAG(ctx context.Context, request *schedulingt
 	if err != nil {
 		return err
 	}
-
 	// Execute the DAG
-	// Initialize channels and nameToNode map
+
+	// Channels to signal plugin execution completion.
 	pluginExecuted := map[string]chan struct{}{}
 	nameToNode := map[string]PrepareDataPlugin{}
 	for _, plugin := range plugins {
@@ -407,33 +382,25 @@ func (d *Director) executePluginsAsDAG(ctx context.Context, request *schedulingt
 			}
 			nameToNode[pluginName].PrepareRequestData(ctx, request, pods)
 			// Signal that the plugin has been executed.
-			<-pluginExecuted[pluginName]
+			close(pluginExecuted[pluginName])
 		}()
 	}
+	for pluginName := range dag {
+		<-pluginExecuted[pluginName]
+	}
 	return nil
 }
 
 func (d *Director) runPrepareDataPlugins(ctx context.Context,
 	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.
 	err := d.executePluginsAsDAG(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
 	}
-	loggerDebug := log.FromContext(ctx).V(logutil.DEBUG)
-	// 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.prepareDataPlugins {
-		loggerDebug.Info("Running PrepareData plugin", "plugin", plugin.TypedName())
-		wg.Add(1)
-		go func(p PrepareDataPlugin) {
-			defer wg.Done()
-			prepareData(p, ctx, request, pods)
-		}(plugin)
-	}
-	wg.Wait()
 	return nil
 }
 

pkg/epp/requestcontrol/graph_util.go

@@ -53,16 +53,6 @@ func buildDAG(plugins []PrepareDataPlugin) map[string][]string {
 // prepareDataGraph builds a DAG of data preparation plugins and checks for cycles.
 // If there is a cycle, it returns an error.
 func prepareDataGraph(plugins []PrepareDataPlugin) (map[string][]string, error) {
-	nameToNode := map[string]PrepareDataPlugin{}
-
-	for _, node := range nameToNode {
-		nameToNode[node.TypedName().String()] = node
-	}
-	// Channels to signal plugin execution completion.
-	pluginExecuted := map[string]chan struct{}{}
-	for _, plugin := range plugins {
-		pluginExecuted[plugin.TypedName().String()] = make(chan struct{})
-	}
 	dag := buildDAG(plugins)
 
 	// Check for cycles in the DAG.