fix: add safety check for LRUCapacityPerServer

Signed-off-by: Kfir Toledo <[email protected]>

Author: kfirtoledo

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

@@ -18,7 +18,6 @@ package prefix
 
 import (
 	"context"
-	"fmt"
 	"sync"
 	"time"
 
@@ -44,24 +43,22 @@ func newIndexer(maxLRUSize int) *indexer {
 		podToLRU:   make(map[ServerID]*lru.Cache[BlockHash, struct{}]),
 		maxLRUSize: maxLRUSize,
 	}
+
 	go ix.ReportLRUSize(time.Second)
 	return ix
 }
 
 // Add adds a list of prefix hashes to the cache, tied to the server.
-func (i *indexer) Add(hashes []BlockHash, pod ServerID) error {
+func (i *indexer) Add(hashes []BlockHash, pod ServerID) {
 	i.mu.Lock()
 	// Check if the LRU pod exist
 	lruForPod, exists := i.podToLRU[pod]
 	if !exists {
-		newLRU, err := lru.NewWithEvict[BlockHash, struct{}](i.maxLRUSize, i.makeEvictionFn(pod))
-		if err != nil {
-			i.mu.Unlock()
-			return fmt.Errorf("failed to create LRU for pod %s: %w", pod, err)
-		}
+		newLRU, _ := lru.NewWithEvict[BlockHash, struct{}](i.maxLRUSize, i.makeEvictionFn(pod))
 		i.podToLRU[pod] = newLRU
 		lruForPod = newLRU
 	}
+
 	i.mu.Unlock()
 
 	// Add to LRU (may evict)
@@ -79,8 +76,8 @@ func (i *indexer) Add(hashes []BlockHash, pod ServerID) error {
 		pods[pod] = struct{}{}
 		i.hashToPods[hash] = pods
 	}
+
 	i.mu.Unlock()
-	return nil
 }
 
 // Get returns a set of servers that have the given prefix hash cached.

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

@@ -27,21 +27,18 @@ func TestIndexer_AddAndGet(t *testing.T) {
 	hash1 := BlockHash(1)
 	server := ServerID{Namespace: "default", Name: "server1"}
 	// Add an entry to the cache
-	err := i.Add([]BlockHash{hash1}, server)
-	assert.NoError(t, err)
+	i.Add([]BlockHash{hash1}, server)
 
 	// Retrieve the entry
 	assert.Equal(t, 1, i.podToLRU[server].Len(), "Cache size should be 1 after adding an entry")
 	servers := i.Get(hash1)
 	assert.Contains(t, servers, server, "Cache should contain the added server")
 
 	// Add another entry to the cache, the cache size should be incremented to 2.
-	err = i.Add([]BlockHash{BlockHash(2)}, server)
-	assert.NoError(t, err)
+	i.Add([]BlockHash{BlockHash(2)}, server)
 	assert.Equal(t, 2, i.podToLRU[server].Len(), "Cache size should  be 2 after adding an entry")
 
 	// Add another entry to the cache, which should evict the first one due to max size.
-	err = i.Add([]BlockHash{BlockHash(3)}, server)
-	assert.NoError(t, err)
+	i.Add([]BlockHash{BlockHash(3)}, server)
 	assert.Equal(t, 2, i.podToLRU[server].Len(), "Cache size should still be 2 after adding an entry")
 }

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

@@ -72,7 +72,7 @@ type podSet map[ServerID]struct{}
 
 type Indexer interface {
 	Get(hash BlockHash) podSet
-	Add(hashes []BlockHash, server ServerID) error
+	Add(hashes []BlockHash, server ServerID)
 }
 
 // BlockHash is a hash of the block of request body.
@@ -115,9 +115,18 @@ var _ framework.PostCycle = &Plugin{}
 
 // New initializes a new prefix Plugin and returns its pointer.
 func New(config Config) *Plugin {
+	capacity := config.LRUCapacityPerServer
+	if capacity <= 0 {
+		capacity = DefaultLRUCapacityPerServer
+		log.FromContext(context.TODO()).V(logutil.DEFAULT).Info(
+			"LRUCapacityPerServer is not positive, using default value",
+			"defaultCapacity", DefaultLRUCapacityPerServer,
+		)
+	}
+
 	m := &Plugin{
 		Config:  config,
-		indexer: newIndexer(config.LRUCapacityPerServer),
+		indexer: newIndexer(capacity),
 	}
 	return m
 }
@@ -165,11 +174,9 @@ func (m *Plugin) PostCycle(ctx context.Context, cycleState *types.CycleState, re
 		log.FromContext(ctx).Error(err, "failed to read prefix plugin cycle state")
 		return
 	}
-	err = m.indexer.Add(state.PrefixHashes, ServerID(targetPod.NamespacedName))
-	if err != nil {
-		log.FromContext(ctx).Error(err, "failed to add prefix hashes to indexer for target pod", "pod", targetPod.NamespacedName)
-		return
-	}
+
+	m.indexer.Add(state.PrefixHashes, ServerID(targetPod.NamespacedName))
+
 	total := len(state.PrefixHashes)
 	matchLen := state.PrefixCacheServers[ServerID(targetPod.NamespacedName)]
 	metrics.RecordPrefixCacheMatch(matchLen*m.HashBlockSize, total*m.HashBlockSize)

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

@@ -154,7 +154,12 @@ func BenchmarkPrefixPluginStress(b *testing.B) {
 
 	plugin := New(config)
 	types.NewCycleState()
-	promptLen := []int{10, 20, 50, 100, 200, 500, 1000, 2000, 5000, 10000, 20000, 50000}
+	var promptLen []int
+	for i := 1; i <= 1024; i++ {
+		promptLen = append(promptLen, i)
+	}
+	promptLen = append(promptLen, 2048, 4096, 8192, 10000, 20000, 50000)
+
 	for _, i := range promptLen {
 		// Generate increasing-length random prompts
 		prompt := randomPrompt(4 + i)