Make plugin execution sequential in topologically sorted order

Author: rahulgurnani

cmd/epp/runner/runner.go

@@ -459,8 +459,8 @@ func (r *Runner) parseConfigurationPhaseTwo(ctx context.Context, rawConfig *conf
 
 	// Add requestControl plugins
 	r.requestControlConfig.AddPlugins(handle.GetAllPlugins()...)
-	// Check prepare data plugins for cycles.
-	if _, err := r.requestControlConfig.PrepareDataPluginGraph(); err != nil {
+	// Sort prepare data plugins in DAG order (topological sort). Also check prepare data plugins for cycles.
+	if r.requestControlConfig.PrepareDataPluginGraph() != nil {
 		return nil, errors.New("failed to load the configuration - prepare data plugins have cyclic dependencies")
 	}
 

pkg/epp/requestcontrol/dag.go

@@ -16,7 +16,10 @@ limitations under the License.
 
 package requestcontrol
 
-import "errors"
+import (
+	"errors"
+	"slices"
+)
 
 // buildDAG builds a dependency graph among data preparation plugins based on their
 // produced and consumed data keys.
@@ -50,50 +53,76 @@ func buildDAG(plugins []PrepareDataPlugin) map[string][]string {
 	return dag
 }
 
-// prepareDataGraph builds a DAG of data preparation plugins and checks for cycles.
+// prepareDataGraph builds the dependency graph and returns the plugins ordered in topological order.
 // If there is a cycle, it returns an error.
-func prepareDataGraph(plugins []PrepareDataPlugin) (map[string][]string, error) {
+func prepareDataGraph(plugins []PrepareDataPlugin) (map[string][]string, []PrepareDataPlugin, error) {
 	dag := buildDAG(plugins)
-
-	// Check for cycles in the DAG.
-	if cycleExistsInDAG(dag) {
-		return nil, errors.New("cycle detected in data preparation plugin dependencies")
+	nameToPlugin := map[string]PrepareDataPlugin{}
+	for _, plugin := range plugins {
+		nameToPlugin[plugin.TypedName().String()] = plugin
+	}
+	sortedPlugins, err := topologicalSort(dag)
+	if err != nil {
+		return nil, nil, err
+	}
+	orderedPlugins := []PrepareDataPlugin{}
+	for _, pluginName := range sortedPlugins {
+		orderedPlugins = append(orderedPlugins, nameToPlugin[pluginName])
 	}
 
-	return dag, nil
+	return dag, orderedPlugins, err
 }
 
-// cycleExistsInDAG checks if there are cycles in the given directed graph represented as an adjacency list.
-func cycleExistsInDAG(dag map[string][]string) bool {
-	visited := make(map[string]bool)
-	recStack := make(map[string]bool)
+// TopologicalSort performs Kahn's Algorithm on a DAG.
+// It returns the sorted order or an error if a cycle is detected.
+func topologicalSort(graph map[string][]string) ([]string, error) {
+	// 1. Initialize in-degree map
+	inDegree := make(map[string]int)
 
-	var dfs func(string) bool
-	dfs = func(node string) bool {
-		if recStack[node] {
-			return true // Cycle detected
+	// Ensure all nodes are present in the inDegree map, even those with no dependencies
+	for u, neighbors := range graph {
+		if _, ok := inDegree[u]; !ok {
+			inDegree[u] = 0
 		}
-		if visited[node] {
-			return false
+		for _, v := range neighbors {
+			inDegree[v]++ // Increment in-degree for the destination node
 		}
-		visited[node] = true
-		recStack[node] = true
+	}
 
-		for _, neighbor := range dag[node] {
-			if dfs(neighbor) {
-				return true
-			}
+	// 2. Initialize the queue with nodes having 0 in-degree
+	var queue []string
+	for node, degree := range inDegree {
+		if degree == 0 {
+			queue = append(queue, node)
 		}
-		recStack[node] = false
-		return false
 	}
 
-	for pluginName := range dag {
-		if !visited[pluginName] {
-			if dfs(pluginName) {
-				return true
+	var result []string
+
+	// 3. Process the queue
+	for len(queue) > 0 {
+		// Dequeue
+		u := queue[0]
+		queue = queue[1:]
+
+		result = append(result, u)
+
+		// Decrease in-degree of neighbors
+		if neighbors, ok := graph[u]; ok {
+			for _, v := range neighbors {
+				inDegree[v]--
+				if inDegree[v] == 0 {
+					queue = append(queue, v)
+				}
 			}
 		}
 	}
-	return false
+
+	// 4. Check for cycles
+	// If the result size != total nodes, there is a cycle
+	if len(result) != len(inDegree) {
+		return nil, errors.New("cycle detected: graph is not a DAG")
+	}
+	slices.Reverse(result)
+	return result, nil
 }

pkg/epp/requestcontrol/dag_test.go

@@ -18,6 +18,7 @@ package requestcontrol
 
 import (
 	"context"
+	"maps"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -82,7 +83,7 @@ func TestPrepareDataGraph(t *testing.T) {
 			expectError: false,
 		},
 		{
-			name:    "Simple linear dependency (A -> B -> C)",
+			name:    "Simple linear dependency (C -> B -> A)",
 			plugins: []PrepareDataPlugin{pluginA, pluginB, pluginC},
 			expectedDAG: map[string][]string{
 				"A/mock": {},
@@ -92,7 +93,7 @@ func TestPrepareDataGraph(t *testing.T) {
 			expectError: false,
 		},
 		{
-			name:    "DAG with multiple dependencies (A -> B, A -> D)",
+			name:    "DAG with multiple dependencies (B -> A, D -> A, E independent)",
 			plugins: []PrepareDataPlugin{pluginA, pluginB, pluginD, pluginE},
 			expectedDAG: map[string][]string{
 				"A/mock": {},
@@ -108,17 +109,11 @@ func TestPrepareDataGraph(t *testing.T) {
 			expectedDAG: nil,
 			expectError: true,
 		},
-		{
-			name:        "Complex graph with a cycle",
-			plugins:     []PrepareDataPlugin{pluginA, pluginB, pluginX, pluginY},
-			expectedDAG: nil,
-			expectError: true,
-		},
 	}
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			dag, err := prepareDataGraph(tc.plugins)
+			dag, orderedPlugins, err := prepareDataGraph(tc.plugins)
 
 			if tc.expectError {
 				assert.Error(t, err)
@@ -129,9 +124,7 @@ func TestPrepareDataGraph(t *testing.T) {
 
 				// Normalize the slices in the maps for consistent comparison
 				normalizedDAG := make(map[string][]string)
-				for k, v := range dag {
-					normalizedDAG[k] = v
-				}
+				maps.Copy(normalizedDAG, dag)
 				normalizedExpectedDAG := make(map[string][]string)
 				for k, v := range tc.expectedDAG {
 					normalizedExpectedDAG[k] = v
@@ -140,7 +133,27 @@ func TestPrepareDataGraph(t *testing.T) {
 				if diff := cmp.Diff(normalizedExpectedDAG, normalizedDAG); diff != "" {
 					t.Errorf("prepareDataGraph() mismatch (-want +got):\n%s", diff)
 				}
+
+				orderedPluginNames := make([]string, len(orderedPlugins))
+				for i, p := range orderedPlugins {
+					orderedPluginNames[i] = p.TypedName().String()
+				}
+				assertTopologicalOrder(t, dag, orderedPlugins)
 			}
 		})
 	}
 }
+
+func assertTopologicalOrder(t *testing.T, dag map[string][]string, ordered []PrepareDataPlugin) {
+	t.Helper()
+	positions := make(map[string]int)
+	for i, p := range ordered {
+		positions[p.TypedName().String()] = i
+	}
+
+	for node, dependencies := range dag {
+		for _, dep := range dependencies {
+			assert.Less(t, positions[dep], positions[node], "Dependency %s should come before %s", dep, node)
+		}
+	}
+}

pkg/epp/requestcontrol/plugin_executor.go

@@ -19,7 +19,6 @@ package requestcontrol
 import (
 	"context"
 	"errors"
-	"fmt"
 	"time"
 
 	schedulingtypes "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
@@ -29,63 +28,9 @@ import (
 // So, a plugin is executed only after all its dependencies have been executed.
 // If there is a cycle or any plugin fails with error, it returns an error.
 func executePluginsAsDAG(plugins []PrepareDataPlugin, ctx context.Context, request *schedulingtypes.LLMRequest, pods []schedulingtypes.Pod) error {
-	// Build the DAG
-	// The error validation happens on startup when loading the config. So, here there should not be any error.
-	dag, err := prepareDataGraph(plugins)
-	if err != nil {
-		return err
-	}
-	// Create a readonly map of plugin name to plugin instance.
-	nameToNode := map[string]PrepareDataPlugin{}
 	for _, plugin := range plugins {
-		nameToNode[plugin.TypedName().String()] = plugin
-	}
-	// Execute the DAG
-
-	// Channels to signal plugin execution completion.
-	pluginExecuted := make(map[string]chan error)
-	// The capacity of the channel is equal to the number of dependents + 1 (for itself).
-	capacityMap := make(map[string]int)
-	for pluginName := range dag {
-		capacityMap[pluginName]++
-		for _, dep := range dag[pluginName] {
-			capacityMap[dep]++
-		}
-	}
-	for pluginName, capacity := range capacityMap {
-		pluginExecuted[pluginName] = make(chan error, capacity)
-	}
-	for pluginName, dependents := range dag {
-		// Execute plugins based on dependencies.
-		//  Wait for the dependencies to complete before executing a plugin.
-		go func(pN string, depds []string) {
-			for _, dep := range depds {
-				// Wait for the dependency plugin to signal completion.
-				if err, closed := <-pluginExecuted[dep]; closed {
-					if err != nil {
-						for range cap(pluginExecuted[pN]) {
-							// Notify all dependents about the failure.
-							pluginExecuted[pN] <- fmt.Errorf("dependency plugin %s failed: %w", dep, err)
-						}
-						// Do not execute this plugin as one of its dependencies failed.
-						return
-					}
-				}
-			}
-			res := nameToNode[pN].PrepareRequestData(ctx, request, pods)
-			for range cap(pluginExecuted[pN]) {
-				// Notify all dependents about the completion.
-				pluginExecuted[pN] <- res
-			}
-		}(pluginName, dependents)
-	}
-
-	// Check for errors in plugin execution.
-	// This will also ensure that all plugins have completed execution before returning.
-	for pluginName := range dag {
-		err := <-pluginExecuted[pluginName]
-		if err != nil {
-			return fmt.Errorf("prepare data plugin %s failed: %v", pluginName, err)
+		if err := plugin.PrepareRequestData(ctx, request, pods); err != nil {
+			return errors.New("prepare data plugin " + plugin.TypedName().String() + " failed: " + err.Error())
 		}
 	}
 	return nil

pkg/epp/requestcontrol/request_control_config.go

@@ -40,7 +40,6 @@ type Config struct {
 	responseReceivedPlugins  []ResponseReceived
 	responseStreamingPlugins []ResponseStreaming
 	responseCompletePlugins  []ResponseComplete
-	prepareDataPluginGraph   map[string][]string
 }
 
 // WithPreRequestPlugins sets the given plugins as the PreRequest plugins.
@@ -106,16 +105,14 @@ func (c *Config) AddPlugins(pluginObjects ...plugins.Plugin) {
 	}
 }
 
-// PrepareDataPluginGraph creates and returns the data dependency graph of PrepareData plugins.
+// PrepareDataPluginGraph creates data dependency graph and sorts the plugins in topological order.
 // If a cycle is detected, it returns an error.
-func (c *Config) PrepareDataPluginGraph() (map[string][]string, error) {
-	if c.prepareDataPluginGraph != nil {
-		return c.prepareDataPluginGraph, nil
-	}
-	graph, err := prepareDataGraph(c.prepareDataPlugins)
+func (c *Config) PrepareDataPluginGraph() error {
+	_, plugins, err := prepareDataGraph(c.prepareDataPlugins)
 	if err != nil {
-		return nil, err
+		return err
 	}
-	c.prepareDataPluginGraph = graph
-	return graph, nil
+	c.prepareDataPlugins = plugins
+
+	return nil
 }