Repository: kubernetes-sigs/gateway-api-inference-extension

Issue:
flaky test in prefix plugin
<!-- Please use this template while reporting a bug and provide as much info as possible. Not doing so may result in your bug not being addressed in a timely manner. Thanks!

If the matter is security related, please disclose it privately via https://kubernetes.io/security/
-->

**What happened**:
flaky test in prefix plugin - I made PR #1528 which only updates few log lines and prefix unit test failed. 
more information is here:
https://prow.k8s.io/view/gs/kubernetes-ci-logs/pr-logs/pull/kubernetes-sigs_gateway-api-inference-extension/1528/pull-gateway-api-inference-extension-test-unit-main/1963475514790252544

**What you expected to happen**:
tests should always pass
**How to reproduce it (as minimally and precisely as possible)**:

**Anything else we need to know?**:

**Environment**:
- Kubernetes version (use `kubectl version`):
- Inference extension version (use `git describe --tags --dirty --always`):
- Cloud provider or hardware configuration:
- Install tools:
- Others:


Discussion:
@nirrozenbaum: cc: @liu-cong @kfswain @kfirtoledo @nirrozenbaum: another link to PR #1524 (which only updates quickstart guide) with this flake :
https://prow.k8s.io/view/gs/kubernetes-ci-logs/pr-logs/pull/kubernetes-sigs_gateway-api-inference-extension/1524/pull-gateway-api-inference-extension-test-unit-main/1963489858903085056@nirrozenbaum: from my short analysis - 
I think the introduction of the go routine in prefix plugin introduces race conditions in unit-tests
```go
go func() {
		p.indexer.Add(state.PrefixHashes, ServerID(targetPod.NamespacedName))
	}()
```

the prefix plugin isn't finishing to update the indexer and the test assumes the state is already saved.

Solution patch:
diff --git a/pkg/epp/scheduling/framework/plugins/multi/prefix/plugin.go b/pkg/epp/scheduling/framework/plugins/multi/prefix/plugin.go
index 7c6993d..49ec7fa 100644
--- a/pkg/epp/scheduling/framework/plugins/multi/prefix/plugin.go
+++ b/pkg/epp/scheduling/framework/plugins/multi/prefix/plugin.go
@@ -21,6 +21,7 @@ import (
 	"encoding/binary"
 	"encoding/json"
 	"fmt"
+	"sync"
 
 	"github.com/cespare/xxhash/v2"
 	k8stypes "k8s.io/apimachinery/pkg/types"
@@ -78,6 +79,7 @@ type Plugin struct {
 	config      Config
 	pluginState *plugins.PluginState
 	indexer     Indexer
+	wg          sync.WaitGroup
 }
 
 // podSet holds an pods servers that may have a specific prefix hash.
@@ -175,7 +177,6 @@ func (p *Plugin) WithName(name string) *Plugin {
 
 // Score returns the scoring result for the given list of pods based on context.
 func (p *Plugin) Score(ctx context.Context, cycleState *types.CycleState, request *types.LLMRequest, pods []types.Pod) map[types.Pod]float64 {
-	loggerTrace := log.FromContext(ctx).V(logutil.TRACE)
 	// pre score step, hashing prompt and find longest prefix match.
 	hashes := hashPrompt(ctx, request, p.config.HashBlockSize, p.config.MaxPrefixBlocksToMatch)
 	state := &SchedulingContextState{
@@ -185,7 +186,7 @@ func (p *Plugin) Score(ctx context.Context, cycleState *types.CycleState, reques
 
 	cycleState.Write(plugins.StateKey(p.TypedName().String()), state)
 	p.pluginState.Write(request.RequestId, plugins.StateKey(p.TypedName().String()), state)
-	loggerTrace.Info(fmt.Sprintf("cached servers: %+v", state.PrefixCacheServers), "hashes", state.PrefixHashes)
+	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))
 
@@ -219,8 +220,11 @@ func (p *Plugin) PreRequest(ctx context.Context, request *types.LLMRequest, sche
 	// This function is just adding data, it does not need to block other operations.
 	// TODO: look into making this entire function async, none of this needs to be done in-band
 	// The PR that introduces this change is meant as a cherrypick, so it was minimally invasive.
+	// WaitGroup is added to the Plugin struct to allow waiting in tests.
+	p.wg.Add(1)
 	go func() {
 		p.indexer.Add(state.PrefixHashes, ServerID(targetPod.NamespacedName))
+		p.wg.Done()
 	}()
 
 	total := len(state.PrefixHashes)
diff --git a/pkg/epp/scheduling/framework/plugins/picker/max_score_picker.go b/pkg/epp/scheduling/framework/plugins/picker/max_score_picker.go
index 90100d7..325f735 100644
--- a/pkg/epp/scheduling/framework/plugins/picker/max_score_picker.go
+++ b/pkg/epp/scheduling/framework/plugins/picker/max_score_picker.go
@@ -82,8 +82,8 @@ func (p *MaxScorePicker) TypedName() plugins.TypedName {
 
 // Pick selects the pod with the maximum score from the list of candidates.
 func (p *MaxScorePicker) Pick(ctx context.Context, cycleState *types.CycleState, scoredPods []*types.ScoredPod) *types.ProfileRunResult {
-	log.FromContext(ctx).V(logutil.DEBUG).Info("Selecting pods from candidates sorted by max score: ", "NumberOfPods", p.maxNumOfEndpoints,
-		"scoredPodsLength", len(scoredPods), "scoredPods", scoredPods)
+	log.FromContext(ctx).V(logutil.DEBUG).Info("Selecting pods from candidates sorted by max score", "max-num-of-endpoints", p.maxNumOfEndpoints,
+		"num-of-candidates", len(scoredPods), "scored-pods", scoredPods)
 
 	// TODO: merge this with the logic in RandomPicker
 	// Rand package is not safe for concurrent use, so we create a new instance.
diff --git a/pkg/epp/scheduling/framework/plugins/picker/random_picker.go b/pkg/epp/scheduling/framework/plugins/picker/random_picker.go
index b571ce3..87a1747 100644
--- a/pkg/epp/scheduling/framework/plugins/picker/random_picker.go
+++ b/pkg/epp/scheduling/framework/plugins/picker/random_picker.go
@@ -38,6 +38,7 @@ const (
 // compile-time type validation
 var _ framework.Picker = &RandomPicker{}
 
+// RandomPickerFactory defines the factory function for RandomPicker.
 func RandomPickerFactory(name string, rawParameters json.RawMessage, _ plugins.Handle) (plugins.Plugin, error) {
 	parameters := pickerParameters{MaxNumOfEndpoints: DefaultMaxNumOfEndpoints}
 	if rawParameters != nil {
@@ -80,8 +81,8 @@ func (p *RandomPicker) TypedName() plugins.TypedName {
 
 // Pick selects random pod(s) from the list of candidates.
 func (p *RandomPicker) Pick(ctx context.Context, _ *types.CycleState, scoredPods []*types.ScoredPod) *types.ProfileRunResult {
-	log.FromContext(ctx).V(logutil.DEBUG).Info(fmt.Sprintf("Selecting maximum '%d' pods from %d candidates randomly: %+v", p.maxNumOfEndpoints,
-		len(scoredPods), scoredPods))
+	log.FromContext(ctx).V(logutil.DEBUG).Info("Selecting pods from candidates randomly", "max-num-of-endpoints", p.maxNumOfEndpoints,
+		"num-of-candidates", len(scoredPods), "scored-pods", scoredPods)
 
 	// TODO: merge this with the logic in MaxScorePicker
 	// Rand package is not safe for concurrent use, so we create a new instance.
diff --git a/pkg/epp/scheduling/framework/plugins/picker/weighted_random_picker.go b/pkg/epp/scheduling/framework/plugins/picker/weighted_random_picker.go
index c12ab72..540ede4 100644
--- a/pkg/epp/scheduling/framework/plugins/picker/weighted_random_picker.go
+++ b/pkg/epp/scheduling/framework/plugins/picker/weighted_random_picker.go
@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"math"
 	"math/rand"
+	"slices"
 	"sort"
 	"time"
 
@@ -43,12 +44,12 @@ type weightedScoredPod struct {
 	key float64
 }
 
+// compile-time type validation
 var _ framework.Picker = &WeightedRandomPicker{}
 
+// WeightedRandomPickerFactory defines the factory function for WeightedRandomPicker.
 func WeightedRandomPickerFactory(name string, rawParameters json.RawMessage, _ plugins.Handle) (plugins.Plugin, error) {
-	parameters := pickerParameters{
-		MaxNumOfEndpoints: DefaultMaxNumOfEndpoints,
-	}
+	parameters := pickerParameters{MaxNumOfEndpoints: DefaultMaxNumOfEndpoints}
 	if rawParameters != nil {
 		if err := json.Unmarshal(rawParameters, &parameters); err != nil {
 			return nil, fmt.Errorf("failed to parse the parameters of the '%s' picker - %w", WeightedRandomPickerType, err)
@@ -58,9 +59,10 @@ func WeightedRandomPickerFactory(name string, rawParameters json.RawMessage, _ p
 	return NewWeightedRandomPicker(parameters.MaxNumOfEndpoints).WithName(name), nil
 }
 
+// NewWeightedRandomPicker initializes a new WeightedRandomPicker and returns its pointer.
 func NewWeightedRandomPicker(maxNumOfEndpoints int) *WeightedRandomPicker {
 	if maxNumOfEndpoints <= 0 {
-		maxNumOfEndpoints = DefaultMaxNumOfEndpoints
+		maxNumOfEndpoints = DefaultMaxNumOfEndpoints // on invalid configuration value, fallback to default value
 	}
 
 	return &WeightedRandomPicker{
@@ -70,81 +72,68 @@ func NewWeightedRandomPicker(maxNumOfEndpoints int) *WeightedRandomPicker {
 	}
 }
 
+// WeightedRandomPicker picks pod(s) from the list of candidates based on weighted random sampling using A-Res algorithm.
+// Reference: https://utopia.duth.gr/~pefraimi/research/data/2007EncOfAlg.pdf.
+//
+// The picker at its core is picking pods randomly, where the probability of the pod to get picked is derived
+// from its weighted score.
+// Algorithm:
+// - Uses A-Res (Algorithm for Reservoir Sampling): keyᵢ = Uᵢ^(1/wᵢ)
+// - Selects k items with largest keys for mathematically correct weighted sampling
+// - More efficient than traditional cumulative probability approach
+//
+// Key characteristics:
+// - Mathematically correct weighted random sampling
+// - Single pass algorithm with O(n + k log k) complexity
 type WeightedRandomPicker struct {
 	typedName         plugins.TypedName
 	maxNumOfEndpoints int
 	randomPicker      *RandomPicker // fallback for zero weights
 }
 
+// WithName sets the name of the picker.
 func (p *WeightedRandomPicker) WithName(name string) *WeightedRandomPicker {
 	p.typedName.Name = name
 	return p
 }
 
+// TypedName returns the type and name tuple of this plugin instance.
 func (p *WeightedRandomPicker) TypedName() plugins.TypedName {
 	return p.typedName
 }
 
-// WeightedRandomPicker performs weighted random sampling using A-Res algorithm.
-// Reference: https://utopia.duth.gr/~pefraimi/research/data/2007EncOfAlg.pdf
-// Algorithm:
-// - Uses A-Res (Algorithm for Reservoir Sampling): keyᵢ = Uᵢ^(1/wᵢ)
-// - Selects k items with largest keys for mathematically correct weighted sampling
-// - More efficient than traditional cumulative probability approach
-//
-// Key characteristics:
-// - Mathematically correct weighted random sampling
-// - Single pass algorithm with O(n + k log k) complexity
+// Pick selects the pod(s) randomly from the list of candidates, where the probability of the pod to get picked is derived
+// from its weighted score.
 func (p *WeightedRandomPicker) Pick(ctx context.Context, cycleState *types.CycleState, scoredPods []*types.ScoredPod) *types.ProfileRunResult {
-	log.FromContext(ctx).V(logutil.DEBUG).Info(fmt.Sprintf("Selecting maximum '%d' pods from %d candidates using weighted random sampling: %+v",
-		p.maxNumOfEndpoints, len(scoredPods), scoredPods))
-
-	// Check if all weights are zero or negative
-	allZeroWeights := true
-	for _, scoredPod := range scoredPods {
-		if scoredPod.Score > 0 {
-			allZeroWeights = false
-			break
-		}
-	}
-
-	// Delegate to RandomPicker for uniform selection when all weights are zero
-	if allZeroWeights {
-		log.FromContext(ctx).V(logutil.DEBUG).Info("All weights are zero, delegating to RandomPicker for uniform selection")
+	// Check if there is at least one pod with Score > 0, if not let random picker run
+	if slices.IndexFunc(scoredPods, func(scoredPod *types.ScoredPod) bool { return scoredPod.Score > 0 }) == -1 {
+		log.FromContext(ctx).V(logutil.DEBUG).Info("All scores are zero, delegating to RandomPicker for uniform selection")
 		return p.randomPicker.Pick(ctx, cycleState, scoredPods)
 	}
 
+	log.FromContext(ctx).V(logutil.DEBUG).Info("Selecting pods from candidates by random weighted picker", "max-num-of-endpoints", p.maxNumOfEndpoints,
+		"num-of-candidates", len(scoredPods), "scored-pods", scoredPods)
+
 	randomGenerator := rand.New(rand.NewSource(time.Now().UnixNano()))
 
 	// A-Res algorithm: keyᵢ = Uᵢ^(1/wᵢ)
-	weightedPods := make([]weightedScoredPod, 0, len(scoredPods))
-
-	for _, scoredPod := range scoredPods {
-		weight := float64(scoredPod.Score)
-
-		// Handle zero or negative weights
-		if weight <= 0 {
-			// Assign very small key for zero-weight pods (effectively excludes them)
-			weightedPods = append(weightedPods, weightedScoredPod{
-				ScoredPod: scoredPod,
-				key:       0,
-			})
+	weightedPods := make([]weightedScoredPod, len(scoredPods))
+
+	for i, scoredPod := range scoredPods {
+		// Handle zero score
+		if scoredPod.Score <= 0 {
+			// Assign key=0 for zero-score pods (effectively excludes them from selection)
+			weightedPods[i] = weightedScoredPod{ScoredPod: scoredPod, key: 0}
 			continue
 		}
 
-		// Generate random number U in (0,1)
+		// If we're here the scoredPod.Score > 0. Generate a random number U in (0,1)
 		u := randomGenerator.Float64()
 		if u == 0 {
 			u = 1e-10 // Avoid log(0)
 		}
 
-		// Calculate key = U^(1/weight)
-		key := math.Pow(u, 1.0/weight)
-
-		weightedPods = append(weightedPods, weightedScoredPod{
-			ScoredPod: scoredPod,
-			key:       key,
-		})
+		weightedPods[i] = weightedScoredPod{ScoredPod: scoredPod, key: math.Pow(u, 1.0/scoredPod.Score)} // key = U^(1/weight)
 	}
 
 	// Sort by key in descending order (largest keys first)
@@ -155,14 +144,9 @@ func (p *WeightedRandomPicker) Pick(ctx context.Context, cycleState *types.Cycle
 	// Select top k pods
 	selectedCount := min(p.maxNumOfEndpoints, len(weightedPods))
 
-	scoredPods = make([]*types.ScoredPod, selectedCount)
+	targetPods := make([]types.Pod, selectedCount)
 	for i := range selectedCount {
-		scoredPods[i] = weightedPods[i].ScoredPod
-	}
-
-	targetPods := make([]types.Pod, len(scoredPods))
-	for i, scoredPod := range scoredPods {
-		targetPods[i] = scoredPod
+		targetPods[i] = weightedPods[i].ScoredPod
 	}
 
 	return &types.ProfileRunResult{TargetPods: targetPods}



Test patch:
diff --git a/pkg/epp/scheduling/framework/plugins/multi/prefix/plugin_test.go b/pkg/epp/scheduling/framework/plugins/multi/prefix/plugin_test.go
index d6ec43c..3fbac2c 100644
--- a/pkg/epp/scheduling/framework/plugins/multi/prefix/plugin_test.go
+++ b/pkg/epp/scheduling/framework/plugins/multi/prefix/plugin_test.go
@@ -71,6 +71,7 @@ func TestPrefixPlugin(t *testing.T) {
 		},
 	}
 	plugin.PreRequest(context.Background(), req1, schedulingResult, 0)
+	plugin.wg.Wait()
 
 	// Second request doesn't share any prefix with first one. It should be added to the cache but
 	// the pod score should be 0.
@@ -98,6 +99,7 @@ func TestPrefixPlugin(t *testing.T) {
 		},
 	}
 	plugin.PreRequest(context.Background(), req2, schedulingResult, 0)
+	plugin.wg.Wait()
 
 	// Third request shares partial prefix with first one.
 	req3 := &types.LLMRequest{
@@ -123,6 +125,7 @@ func TestPrefixPlugin(t *testing.T) {
 		},
 	}
 	plugin.PreRequest(context.Background(), req3, schedulingResult, 0)
+	plugin.wg.Wait()
 
 	// 4th request is same as req3 except the model is different, still no match.
 	req4 := &types.LLMRequest{
@@ -148,6 +151,7 @@ func TestPrefixPlugin(t *testing.T) {
 		},
 	}
 	plugin.PreRequest(context.Background(), req4, schedulingResult, 0)
+	plugin.wg.Wait()
 
 	// 5th request shares partial prefix with 3rd one.
 	req5 := &types.LLMRequest{
@@ -173,6 +177,7 @@ func TestPrefixPlugin(t *testing.T) {
 		},
 	}
 	plugin.PreRequest(context.Background(), req5, schedulingResult, 0)
+	plugin.wg.Wait()
 }
 
 // TestPrefixPluginStress is a stress test for the prefix scoring plugin, using prompts of increasing length.
@@ -220,6 +225,7 @@ func BenchmarkPrefixPluginStress(b *testing.B) {
 			},
 		}
 		plugin.PreRequest(context.Background(), req, schedulingResult, 0)
+		plugin.wg.Wait()
 
 		// Second cycle: validate internal state
 		state, err := plugins.ReadPluginStateKey[*SchedulingContextState](plugin.pluginState, req.RequestId, plugins.StateKey(plugin.TypedName().String()))