feat: Add GetActivePods functionto plugins handle and datastore

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

Author: kfirtoledo

cmd/epp/runner/runner.go

@@ -230,7 +230,7 @@ func (r *Runner) Run(ctx context.Context) error {
 		}
 	}
 
-	err = r.parsePluginsConfiguration(ctx)
+	err = r.parsePluginsConfiguration(ctx, datastore)
 	if err != nil {
 		setupLog.Error(err, "Failed to parse plugins configuration")
 		return err
@@ -306,7 +306,7 @@ func (r *Runner) registerInTreePlugins() {
 	plugins.Register(testfilter.HeaderBasedTestingFilterType, testfilter.HeaderBasedTestingFilterFactory)
 }
 
-func (r *Runner) parsePluginsConfiguration(ctx context.Context) error {
+func (r *Runner) parsePluginsConfiguration(ctx context.Context, ds datastore.Datastore) error {
 	if *configText == "" && *configFile == "" {
 		return nil // configuring through code, not through file
 	}
@@ -325,8 +325,9 @@ func (r *Runner) parsePluginsConfiguration(ctx context.Context) error {
 	}
 
 	r.registerInTreePlugins()
-	handle := plugins.NewEppHandle(ctx)
+	handle := plugins.NewEppHandle(ctx, ds.GetActivePods)
 	config, err := loader.LoadConfig(configBytes, handle, logger)
+
 	if err != nil {
 		return fmt.Errorf("failed to load the configuration - %w", err)
 	}

pkg/epp/datastore/datastore.go

@@ -63,6 +63,7 @@ type Datastore interface {
 	PodList(predicate func(backendmetrics.PodMetrics) bool) []backendmetrics.PodMetrics
 	PodUpdateOrAddIfNotExist(pod *corev1.Pod) bool
 	PodDelete(namespacedName types.NamespacedName)
+	GetActivePods() []types.NamespacedName
 
 	// Clears the store state, happens when the pool gets deleted.
 	Clear()
@@ -232,6 +233,16 @@ func (ds *datastore) PodUpdateOrAddIfNotExist(pod *corev1.Pod) bool {
 	return ok
 }
 
+// GetActivePods returns a list of all active pods.
+func (ds *datastore) GetActivePods() []types.NamespacedName {
+	var namespacedNames []types.NamespacedName
+	ds.pods.Range(func(k, _ any) bool {
+		namespacedNames = append(namespacedNames, k.(types.NamespacedName))
+		return true
+	})
+	return namespacedNames
+}
+
 func (ds *datastore) PodDelete(namespacedName types.NamespacedName) {
 	v, ok := ds.pods.LoadAndDelete(namespacedName)
 	if ok {

pkg/epp/plugins/handle.go

@@ -19,6 +19,8 @@ package plugins
 import (
 	"context"
 	"fmt"
+
+	"k8s.io/apimachinery/pkg/types"
 )
 
 // Handle provides plugins a set of standard data and tools to work with
@@ -27,6 +29,9 @@ type Handle interface {
 	Context() context.Context
 
 	HandlePlugins
+
+	// GetActivePods returns a list of all active pods
+	GetActivePods() []types.NamespacedName
 }
 
 // HandlePlugins defines a set of APIs to work with instantiated plugins
@@ -44,10 +49,14 @@ type HandlePlugins interface {
 	GetAllPluginsWithNames() map[string]Plugin
 }
 
+// GetActivePodsFunc is a function that returns a list of all active pods.
+type GetActivePodsFunc func() []types.NamespacedName
+
 // eppHandle is an implementation of the interface plugins.Handle
 type eppHandle struct {
 	ctx context.Context
 	HandlePlugins
+	getActivePods GetActivePodsFunc
 }
 
 // Context returns a context the plugins can use, if they need one
@@ -84,7 +93,12 @@ func (h *eppHandlePlugins) GetAllPluginsWithNames() map[string]Plugin {
 	return h.plugins
 }
 
-func NewEppHandle(ctx context.Context) Handle {
+// GetActivePods returns a function that returns a list of all active pods
+func (h *eppHandle) GetActivePods() []types.NamespacedName {
+	return h.getActivePods()
+}
+
+func NewEppHandle(ctx context.Context, getActivePods GetActivePodsFunc) Handle {
 	return &eppHandle{
 		ctx: ctx,
 		HandlePlugins: &eppHandlePlugins{

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

@@ -149,3 +149,23 @@ func (i *indexer) ReportLRUSize(interval time.Duration) {
 		i.mu.RUnlock()
 	}
 }
+
+// RemovePod removes a pod and its associated entries from the indexer.
+func (i *indexer) RemovePod(pod ServerID) {
+	i.mu.RLock()
+	lruCache, exists := i.podToLRU[pod]
+	i.mu.RUnlock()
+
+	if !exists {
+		return
+	}
+
+	// Remove all hashes associated with the pod from hashToPods (triggers eviction callbacks).
+	for _, hash := range lruCache.Keys() {
+		lruCache.Remove(hash)
+	}
+
+	i.mu.Lock()
+	delete(i.podToLRU, pod)
+	i.mu.Unlock()
+}

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

@@ -45,3 +45,63 @@ func TestIndexer_AddAndGet(t *testing.T) {
 	servers = i.Get(BlockHash(4))
 	assert.Empty(t, servers, "Cache should not contain non-existent hash")
 }
+
+func TestIndexer_RemovePodAndEviction(t *testing.T) {
+	const indexerSize = 10
+
+	i := newIndexer(indexerSize)
+
+	server1 := ServerID{Namespace: "default", Name: "server1"}
+	server2 := ServerID{Namespace: "default", Name: "server2"}
+
+	// Add indexerSize hashes to both servers
+	var hashes []BlockHash
+	for j := 0; j < indexerSize; j++ {
+		h := BlockHash(j)
+		hashes = append(hashes, h)
+		i.Add([]BlockHash{h}, server1)
+		i.Add([]BlockHash{h}, server2)
+	}
+
+	// Ensure all entries are added
+	assert.Equal(t, indexerSize, i.podToLRU[server1].Len(), "server1 should have 10 entries")
+	assert.Equal(t, indexerSize, i.podToLRU[server2].Len(), "server2 should have 10 entries")
+
+	// Ensure each hash in hashToPods maps to both server1 and server2
+	for _, h := range hashes {
+		pods := i.hashToPods[h]
+		assert.Len(t, pods, 2, "Each hash should be associated with exactly 2 pods")
+		assert.Contains(t, pods, server1, "hash should be associated with server1")
+		assert.Contains(t, pods, server2, "hash should be associated with server2")
+	}
+
+	// Add indexerSize hash to server1 → should evict BlockHash(0)
+	evictedHash := BlockHash(0)
+	newHash := BlockHash(indexerSize)
+	i.Add([]BlockHash{newHash}, server1)
+
+	// server1 LRU should still be at max capacity
+	assert.Equal(t, indexerSize, i.podToLRU[server1].Len(), "server1 LRU should maintain max size")
+
+	// BlockHash(0) should no longer have server1 in hashToPods
+	pods := i.Get(evictedHash)
+	assert.NotContains(t, pods, server1, "server1 should be evicted from hashToPods for hash 0")
+	assert.Contains(t, pods, server2, "server2 should still have hash 0")
+
+	// Remove server2
+	i.RemovePod(server2)
+
+	// hashToPods for hash 0 should now be empty
+	pods = i.Get(evictedHash)
+	assert.NotContains(t, pods, server2, "server2 should be removed from hash 0")
+	assert.Empty(t, pods, "hash 0 should have no pods after both eviction and removal")
+
+	// All remaining hashes should map only to server1
+	for hash, pods := range i.hashToPods {
+		assert.Len(t, pods, 1, "hash %v should have only 1 pod after server2 removal", hash)
+		assert.Contains(t, pods, server1, "hash %v should only contain server1", hash)
+	}
+
+	// Ensure hashToPods contains exactly indexerSize hashes (post-eviction and server2 removal)
+	assert.Len(t, i.hashToPods, indexerSize, "hashToPods should contain %d hashes after cleanup", indexerSize)
+}

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

@@ -21,6 +21,7 @@ import (
 	"encoding/binary"
 	"encoding/json"
 	"fmt"
+	"time"
 
 	"github.com/cespare/xxhash/v2"
 	k8stypes "k8s.io/apimachinery/pkg/types"
@@ -56,6 +57,11 @@ const (
 	PrefixCachePluginType = "prefix-cache-scorer"
 )
 
+const (
+	PodActiveCheckInterval = 1 * time.Minute
+	PodInactivityTimeout   = 5 * time.Minute
+)
+
 var DefaultConfig = Config{
 	HashBlockSize:          DefaultHashBlockSize,
 	MaxPrefixBlocksToMatch: DefaultMaxPrefixBlocks,
@@ -86,6 +92,7 @@ type podSet map[ServerID]struct{}
 type Indexer interface {
 	Get(hash BlockHash) podSet
 	Add(hashes []BlockHash, server ServerID)
+	RemovePod(server ServerID)
 }
 
 // BlockHash is a hash of the block of request body.
@@ -140,7 +147,9 @@ func PrefixCachePluginFactory(name string, rawParameters json.RawMessage, handle
 		}
 	}
 
-	return New(handle.Context(), parameters).WithName(name), nil
+	p := New(handle.Context(), parameters).WithName(name)
+	go p.StartPodActiveWatcher(handle.Context(), handle)
+	return p, nil
 }
 
 // New initializes a new prefix Plugin and returns its pointer.
@@ -245,6 +254,45 @@ func (p *Plugin) matchLongestPrefix(ctx context.Context, hashes []BlockHash) map
 	return res
 }
 
+// StartPodActiveWatcher starts a goroutine that watches for active pods.
+func (m *Plugin) StartPodActiveWatcher(ctx context.Context, handle plugins.Handle) {
+	logger := log.FromContext(ctx).V(logutil.VERBOSE)
+
+	ticker := time.NewTicker(PodActiveCheckInterval)
+	defer ticker.Stop()
+
+	podLastSeen := make(map[ServerID]time.Time)
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			now := time.Now()
+			activePods := handle.GetActivePods()
+
+			// Track active pods
+			activeSet := make(map[ServerID]struct{}, len(activePods))
+			for _, np := range activePods {
+				id := ServerID(np)
+				activeSet[id] = struct{}{}
+				podLastSeen[id] = now
+			}
+
+			// Remove stale pods
+			for pod, lastSeen := range podLastSeen {
+				if _, stillActive := activeSet[pod]; !stillActive {
+					if now.Sub(lastSeen) > PodInactivityTimeout {
+						m.indexer.RemovePod(pod)
+						delete(podLastSeen, pod)
+						logger.Info("Removed inactive pod from prefix cache", "pod", pod)
+					}
+				}
+			}
+		}
+	}
+}
+
 // hashPrompt divides the prompt into blocks and calculate the prefix cache for each block.
 // hash(0) is the hash of the model name, since different models generally don't share prefix cache.
 // For block i, hash(i) = hash(block i content, hash(i-1)).

test/utils/handle.go

@@ -19,6 +19,8 @@ package utils
 import (
 	"context"
 
+	k8stypes "k8s.io/apimachinery/pkg/types"
+
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/plugins"
 )
 
@@ -33,6 +35,10 @@ func (h *testHandle) Context() context.Context {
 	return h.ctx
 }
 
+func (h *testHandle) GetActivePods() []k8stypes.NamespacedName {
+	return []k8stypes.NamespacedName{}
+}
+
 type testHandlePlugins struct {
 	plugins map[string]plugins.Plugin
 }