Prefix aware scorer data structure fix

Signed-off-by: Ricardo Noriega <rnoriega@redhat.com>

Author: oglok

cmd/epp/main.go

@@ -23,6 +23,7 @@ import (
 	"net/http"
 	"os"
 	"strconv"
+	"time"
 
 	"github.com/go-logr/logr"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
@@ -42,6 +43,7 @@ import (
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
 	runserver "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/server"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling"
 )
 
 const (
@@ -162,7 +164,12 @@ func run() error {
 
 	pmf := backendmetrics.NewPodMetricsFactory(&backendmetrics.PodMetricsClientImpl{MetricMapping: mapping}, *refreshMetricsInterval)
 	// Setup runner.
-	datastore := datastore.NewDatastore(ctx, pmf)
+	datastore := datastore.NewDatastore(ctx, pmf, scheduling.NewPrefixStore(scheduling.PrefixStoreConfig{
+		MaxEntries:    1000,
+		MinPrefixLen:  1,
+		MaxPrefixLen:  100,
+		EntryTTL:      5 * time.Minute,
+	}))
 
 	serverRunner := &runserver.ExtProcServerRunner{
 		GrpcPort:                                 *grpcPort,

pkg/epp/datastore/datastore.go

@@ -32,6 +32,8 @@ import (
 	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 	podutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/pod"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
 const (
@@ -75,14 +77,15 @@ type Datastore interface {
 	Clear()
 }
 
-func NewDatastore(parentCtx context.Context, pmf *backendmetrics.PodMetricsFactory) Datastore {
+func NewDatastore(parentCtx context.Context, pmf *backendmetrics.PodMetricsFactory, prefixStore *scheduling.PrefixStore) Datastore {
 	store := &datastore{
 		parentCtx:       parentCtx,
 		poolAndModelsMu: sync.RWMutex{},
 		models:          make(map[string]*v1alpha2.InferenceModel),
 		pods:            &sync.Map{},
 		sessions:        &sync.Map{},
 		pmf:             pmf,
+		prefixStore:     prefixStore,
 	}
 
 	go store.cleanupSessions(sessionKeepAliveCheckFrequency, sessionKeepAliveTime, parentCtx)
@@ -103,6 +106,8 @@ type datastore struct {
 	// key: session id, value: *backendmetrics.Pod
 	sessions *sync.Map
 	pmf      *backendmetrics.PodMetricsFactory
+	// prefixStore is used to store and lookup model prefixes
+	prefixStore *scheduling.PrefixStore
 }
 
 func (ds *datastore) Clear() {
@@ -256,6 +261,19 @@ func (ds *datastore) PodUpdateOrAddIfNotExist(pod *corev1.Pod, pool *v1alpha2.In
 	}
 	// Update pod properties if anything changed.
 	pm.UpdatePod(pod)
+
+	// Update prefix store with pod's active models
+	if ds.prefixStore != nil {
+		metrics := pm.GetMetrics()
+		for model := range metrics.ActiveModels {
+			// Add a prefix for each active model
+			// The prefix is the model name itself
+			ds.prefixStore.AddPrefix(ds.parentCtx, model, namespacedName, model)
+		}
+	} else {
+		log.FromContext(ds.parentCtx).V(logging.DEBUG).Info("Prefix store is nil, skipping prefix updates", "pod", namespacedName)
+	}
+
 	return ok
 }
 
@@ -301,6 +319,7 @@ func (ds *datastore) PodDelete(namespacedName types.NamespacedName) {
 		pmr := v.(backendmetrics.PodMetrics)
 		pmr.StopRefreshLoop()
 	}
+	log.FromContext(ds.parentCtx).V(logging.DEBUG).Info("Pod removed or not added", "name", namespacedName)
 }
 
 type sessionInfo struct {

pkg/epp/handlers/request.go

@@ -51,6 +51,10 @@ func (s *StreamingServer) HandleRequestBody(
 
 	modelName := model
 
+	// Extract prompt from request body
+	prompt, _ := requestBodyMap["prompt"].(string) // We don't require prompt to be present
+	reqCtx.Prompt = prompt
+
 	// NOTE: The nil checking for the modelObject means that we DO allow passthrough currently.
 	// This might be a security risk in the future where adapters not registered in the InferenceModel
 	// are able to be requested by using their distinct name.
@@ -67,6 +71,7 @@ func (s *StreamingServer) HandleRequestBody(
 
 	llmReq := &schedulingtypes.LLMRequest{
 		Model:               model,
+		Prompt:             prompt,
 		ResolvedTargetModel: modelName,
 		Critical:            modelObj.Spec.Criticality != nil && *modelObj.Spec.Criticality == v1alpha2.Critical,
 		SessionID:           reqCtx.SessionID,

pkg/epp/handlers/response.go

@@ -25,6 +25,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
+	"k8s.io/apimachinery/pkg/types"
 )
 
 const (
@@ -62,6 +63,32 @@ func (s *StreamingServer) HandleResponseBody(
 	// will add the processing for streaming case.
 	reqCtx.ResponseComplete = true
 
+	// Add the prompt to the prefix store if we have a target pod and model
+	if reqCtx.TargetPod != "" && reqCtx.ResolvedTargetModel != "" {
+		// Get the prompt from the request context
+		prompt := reqCtx.Prompt
+		if prompt != "" {
+			// Convert TargetPod string to NamespacedName
+			parts := strings.Split(reqCtx.TargetPod, "/")
+			if len(parts) != 2 {
+				logger.Error(nil, "Invalid TargetPod format", "targetPod", reqCtx.TargetPod)
+				return reqCtx, nil
+			}
+			podName := types.NamespacedName{
+				Namespace: parts[0],
+				Name:      parts[1],
+			}
+
+			// Add the prefix to the store
+			err := s.scheduler.GetPrefixStore().AddPrefix(ctx, prompt, podName, reqCtx.ResolvedTargetModel)
+			if err != nil {
+				logger.Error(err, "Failed to add prefix to store", "prefix", prompt, "pod", reqCtx.TargetPod, "model", reqCtx.ResolvedTargetModel)
+			} else {
+				logger.Info("Added prefix to store", "prefix", prompt, "pod", reqCtx.TargetPod, "model", reqCtx.ResolvedTargetModel)
+			}
+		}
+	}
+
 	reqCtx.respBodyResp = &extProcPb.ProcessingResponse{
 		// The Endpoint Picker supports two approaches to communicating the target endpoint, as a request header
 		// and as an unstructure ext-proc response metadata key/value pair. This enables different integration

pkg/epp/handlers/server.go

@@ -40,6 +40,7 @@ import (
 	schedulingtypes "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
 	errutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/error"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling"
 )
 
 func NewStreamingServer(scheduler Scheduler, destinationEndpointHintMetadataNamespace, destinationEndpointHintKey string, datastore datastore.Datastore) *StreamingServer {
@@ -66,6 +67,7 @@ type StreamingServer struct {
 
 type Scheduler interface {
 	Schedule(ctx context.Context, b *schedulingtypes.LLMRequest) (targetPod schedulingtypes.Pod, err error)
+	GetPrefixStore() *scheduling.PrefixStore
 }
 
 // RequestContext stores context information during the life time of an HTTP request.
@@ -75,6 +77,7 @@ type RequestContext struct {
 	Model                     string
 	SessionID                 string
 	ResolvedTargetModel       string
+	Prompt                    string
 	RequestReceivedTimestamp  time.Time
 	ResponseCompleteTimestamp time.Time
 	RequestSize               int

pkg/epp/scheduling/prefix_aware_scorer.go

@@ -19,6 +19,7 @@ package scheduling
 import (
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
 // PrefixAwareScorer is a routing scorer that scores pods based on the longest prefix match
@@ -44,7 +45,9 @@ func (s *PrefixAwareScorer) ScoreTargets(ctx *types.Context, pods []*types.PodMe
 
 	// Get the prompt from the request
 	prompt := ctx.Req.Prompt
+
 	if prompt == "" {
+		logger.V(logging.DEBUG).Info("Empty prompt, returning zero scores for all pods")
 		// If no prompt, return zero scores for all pods
 		for i, pod := range pods {
 			scoredPods[i] = PodScore{
@@ -58,6 +61,7 @@ func (s *PrefixAwareScorer) ScoreTargets(ctx *types.Context, pods []*types.PodMe
 	// Find the best matching pod for the prompt
 	matchedPod, found := s.prefixStore.FindPodForPrefix(ctx, prompt, ctx.Req.ResolvedTargetModel)
 	if !found {
+		logger.V(logging.DEBUG).Info("No matching prefix found, returning zero scores for all pods")
 		// If no matching prefix found, return zero scores for all pods
 		for i, pod := range pods {
 			scoredPods[i] = PodScore{
@@ -71,12 +75,18 @@ func (s *PrefixAwareScorer) ScoreTargets(ctx *types.Context, pods []*types.PodMe
 	// Assign scores based on pod match
 	for i, pod := range pods {
 		if pod.NamespacedName == matchedPod {
-			logger.Info("Pod found for prefix", "prompt", prompt, "pod", pod.NamespacedName.String())
+			logger.V(logging.DEBUG).Info("Pod matched for prefix",
+				"prompt", prompt,
+				"pod", pod.NamespacedName.String(),
+				"score", s.weight)
 			scoredPods[i] = PodScore{
 				Score: s.weight, // Use the configured weight for the matching pod
 				Pod:   pod,
 			}
 		} else {
+			logger.V(logging.DEBUG).Info("Pod did not match",
+			"pod", pod.NamespacedName.String(),
+			"score", 0)
 			scoredPods[i] = PodScore{
 				Score: 0, // Zero score for non-matching pods
 				Pod:   pod,

pkg/epp/scheduling/prefix_store.go

@@ -9,6 +9,7 @@ import (
 	"k8s.io/apimachinery/pkg/types"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	errutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/error"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
 // PrefixEntry represents a single entry in the prefix store
@@ -35,6 +36,8 @@ type PrefixStore struct {
 
 // NewPrefixStore creates a new PrefixStore with the given configuration
 func NewPrefixStore(config PrefixStoreConfig) *PrefixStore {
+	logger := log.FromContext(context.Background())
+	logger.V(logging.DEBUG).Info("Creating new PrefixStore", "config", config)
 	return &PrefixStore{
 		tree:   radix.New(),
 		config: config,
@@ -56,13 +59,15 @@ func (ps *PrefixStore) AddPrefix(ctx context.Context, prefix string, pod types.N
 		}
 	}
 	if len(prefix) > ps.config.MaxPrefixLen {
+		logger.V(logging.DEBUG).Info("Truncating prefix", "originalLength", len(prefix), "maxLength", ps.config.MaxPrefixLen)
 		prefix = prefix[:ps.config.MaxPrefixLen]
 	}
 
 	// Check if we're updating an existing entry
 	if val, exists := ps.tree.Get(prefix); exists {
 		entry := val.(*PrefixEntry)
 		if entry.PodRef == pod && entry.ModelName == modelName {
+			logger.V(logging.DEBUG).Info("Updating existing entry", "prefix", prefix, "pod", pod.String())
 			entry.LastUsed = time.Now()
 			ps.tree.Insert(prefix, entry)
 			return nil
@@ -71,6 +76,7 @@ func (ps *PrefixStore) AddPrefix(ctx context.Context, prefix string, pod types.N
 
 	// Check total entries limit
 	if ps.tree.Len() >= ps.config.MaxEntries {
+		logger.V(logging.DEBUG).Info("Store at capacity, evicting oldest entry", "currentSize", ps.tree.Len(), "maxSize", ps.config.MaxEntries)
 		ps.evictOldest()
 	}
 
@@ -82,7 +88,7 @@ func (ps *PrefixStore) AddPrefix(ctx context.Context, prefix string, pod types.N
 	}
 	ps.tree.Insert(prefix, entry)
 
-	logger.Info("Added prefix entry", "prefix", prefix, "pod", pod.String(), "model", modelName)
+	logger.V(logging.DEBUG).Info("Successfully added new prefix entry", "prefix", prefix, "pod", pod.String(), "model", modelName, "totalEntries", ps.tree.Len())
 	return nil
 }
 
@@ -94,32 +100,36 @@ func (ps *PrefixStore) FindPodForPrefix(ctx context.Context, prefix string, mode
 	logger := log.FromContext(ctx)
 
 	if len(prefix) < ps.config.MinPrefixLen {
+		logger.V(logging.DEBUG).Info("Prefix too short", "prefix", prefix, "minLength", ps.config.MinPrefixLen)
 		return types.NamespacedName{}, false
 	}
 
 	if len(prefix) > ps.config.MaxPrefixLen {
+		logger.V(logging.DEBUG).Info("Truncating prefix", "originalLength", len(prefix), "maxLength", ps.config.MaxPrefixLen)
 		prefix = prefix[:ps.config.MaxPrefixLen]
 	}
 
 	// Use LongestPrefix to find the best match
 	matchedPrefix, val, found := ps.tree.LongestPrefix(prefix)
 	if !found {
+		logger.V(logging.DEBUG).Info("No matching prefix found", "prefix", prefix)
 		return types.NamespacedName{}, false
 	}
 
 	entry := val.(*PrefixEntry)
 
 	// Check if entry has expired or model doesn't match
-	if time.Since(entry.LastUsed) > ps.config.EntryTTL || entry.ModelName != modelName {
-		// Don't remove here to avoid write lock
+	if time.Since(entry.LastUsed) > ps.config.EntryTTL {
+		return types.NamespacedName{}, false
+	}
+	if entry.ModelName != modelName {
 		return types.NamespacedName{}, false
 	}
 
 	// Update LastUsed time for the matched entry
 	entry.LastUsed = time.Now()
 	ps.tree.Insert(matchedPrefix, entry)
 
-	logger.Info("Found pod for prefix", "prefix", prefix, "matchedPrefix", matchedPrefix, "pod", entry.PodRef.String(), "model", modelName)
 	return entry.PodRef, true
 }
 
@@ -169,7 +179,9 @@ func (ps *PrefixStore) cleanupExpired(ctx context.Context) {
 	}
 
 	if len(keysToDelete) > 0 {
-		logger.Info("Cleaned up expired entries", "count", len(keysToDelete))
+		logger.V(logging.DEBUG).Info("Cleaned up expired entries", "count", len(keysToDelete), "remainingEntries", ps.tree.Len())
+	} else {
+		logger.V(logging.DEBUG).Info("No expired entries found", "totalEntries", ps.tree.Len())
 	}
 }
 
@@ -179,14 +191,14 @@ func (ps *PrefixStore) RunMaintenance(ctx context.Context) {
 	ticker := time.NewTicker(ps.config.EntryTTL / 2)
 	defer ticker.Stop()
 
+	logger.V(logging.DEBUG).Info("Starting maintenance routine", "interval", ps.config.EntryTTL/2)
+
 	for {
 		select {
 		case <-ctx.Done():
-			logger.Info("Maintenance routine stopping")
 			return
 		case <-ticker.C:
 			ps.cleanupExpired(ctx)
-			logger.Info("Completed maintenance cycle")
 		}
 	}
 }

pkg/epp/scheduling/scheduler.go

@@ -137,7 +137,7 @@ type Scheduler struct {
 
 func NewScheduler(datastore Datastore) *Scheduler {
 	sMng := NewScorerMng()
-	sMng.addScorer(NewSessionAffinityScorer(1, datastore))
+	sMng.addScorer(NewSessionAffinityScorer(0, datastore))
 
 	// Initialize prefix store with configuration from environment variables
 	prefixStore := NewPrefixStore(PrefixStoreConfig{
@@ -193,3 +193,8 @@ func (s *Scheduler) Schedule(ctx context.Context, req *types.LLMRequest) (target
 
 	return selectedPod, nil
 }
+
+// GetPrefixStore returns the prefix store for this scheduler
+func (s *Scheduler) GetPrefixStore() *PrefixStore {
+	return s.prefixStore
+}

pkg/epp/scheduling/scorer.go

@@ -92,6 +92,7 @@ func (sm *ScorerMng) scoreTargets(ctx *types.Context, pods []*types.PodMetrics)
 		if isFirst {
 			maxScore = score
 			highestScoreTargets = []*types.PodMetrics{pod}
+			isFirst = false
 		} else {
 			if score > maxScore {
 				maxScore = score