Update prefix cache match to have total and match length

Author: rahulgurnani

pkg/epp/datalayer/plugins/data_types.go

@@ -25,21 +25,28 @@ const (
 )
 
 type PrefixCacheMatchInfo struct {
-	matchPercentage float64
+	matchLength int
+	totalBlocks int
 }
 
-func NewPrefixCacheMatchInfo(matchPercentage float64) *PrefixCacheMatchInfo {
+func NewPrefixCacheMatchInfo(matchLen int, blockHashLen int) *PrefixCacheMatchInfo {
 	return &PrefixCacheMatchInfo{
-		matchPercentage: matchPercentage,
+		matchLength: matchLen,
+		totalBlocks: blockHashLen,
 	}
 }
 
-func (p *PrefixCacheMatchInfo) MatchPercentage() float64 {
-	return p.matchPercentage
+func (p *PrefixCacheMatchInfo) MatchLength() int {
+	return p.matchLength
+}
+
+func (p *PrefixCacheMatchInfo) TotalLength() int {
+	return p.totalBlocks
 }
 
 func (p *PrefixCacheMatchInfo) Clone() datalayer.Cloneable {
 	return &PrefixCacheMatchInfo{
-		matchPercentage: p.matchPercentage,
+		matchLength: p.matchLength,
+		totalBlocks: p.totalBlocks,
 	}
 }

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

@@ -215,8 +215,8 @@ func (p *Plugin) Consumes() map[string]any {
 	return map[string]any{}
 }
 
+// PrepareRequestData hashes prompt, finds longest prefix match and stores it in pod as attribute.
 func (p *Plugin) PrepareRequestData(ctx context.Context, request *types.LLMRequest, pods []types.Pod) error {
-	// pre score step, hashing prompt and find longest prefix match.
 	hashes := hashPrompt(ctx, request, getBlockSize(pods, p.config), p.config.MaxPrefixBlocksToMatch)
 	state := &SchedulingContextState{
 		PrefixHashes:       hashes,
@@ -228,15 +228,10 @@ func (p *Plugin) PrepareRequestData(ctx context.Context, request *types.LLMReque
 	}
 
 	total := len(state.PrefixHashes)
-	podScoreFunc := func(pod types.Pod) float64 {
-		if total == 0 {
-			return 0
-		}
-		matchLen := state.PrefixCacheServers[ServerID(pod.GetPod().NamespacedName)]
-		return float64(matchLen) / float64(total)
-	}
+
 	for _, pod := range pods {
-		pod.Put(dplugins.PrefixCacheMatchInfoKey, dplugins.NewPrefixCacheMatchInfo(podScoreFunc(pod)))
+		matchLen := state.PrefixCacheServers[ServerID(pod.GetPod().NamespacedName)]
+		pod.Put(dplugins.PrefixCacheMatchInfoKey, dplugins.NewPrefixCacheMatchInfo(matchLen, total))
 	}
 	return nil
 }

pkg/epp/scheduling/framework/plugins/scorer/prefix_cache_match_scorer.go

@@ -70,7 +70,7 @@ func (s *PrefixCacheScorer) WithName(name string) *PrefixCacheScorer {
 }
 
 func (s *PrefixCacheScorer) Score(_ context.Context, cycleState *types.CycleState, _ *types.LLMRequest, pods []types.Pod) map[types.Pod]float64 {
-	// calculate the scores of pods
+	// Calculate the scores of pods based on prefix cache match percent.
 	scores := make(map[types.Pod]float64, len(pods))
 
 	for _, pod := range pods {
@@ -79,7 +79,7 @@ func (s *PrefixCacheScorer) Score(_ context.Context, cycleState *types.CycleStat
 			scores[pod] = 0.0
 			continue
 		}
-		scores[pod] = matchPercent.(*dplugins.PrefixCacheMatchInfo).MatchPercentage()
+		scores[pod] = float64(matchPercent.(*dplugins.PrefixCacheMatchInfo).MatchLength()) / float64(matchPercent.(*dplugins.PrefixCacheMatchInfo).TotalLength()) * 100
 	}
 	return scores
 }

pkg/epp/scheduling/framework/plugins/scorer/prefix_cache_match_scorer_test.go

@@ -70,10 +70,10 @@ func (p *mockPod) Keys() []string {
 
 func TestPrefixCacheScorer_Score(t *testing.T) {
 	pod1 := newMockPod()
-	pod1.Put(dplugins.PrefixCacheMatchInfoKey, dplugins.NewPrefixCacheMatchInfo(50.0))
+	pod1.Put(dplugins.PrefixCacheMatchInfoKey, dplugins.NewPrefixCacheMatchInfo(5, 10))
 
 	pod2 := newMockPod()
-	pod2.Put(dplugins.PrefixCacheMatchInfoKey, dplugins.NewPrefixCacheMatchInfo(100.0))
+	pod2.Put(dplugins.PrefixCacheMatchInfoKey, dplugins.NewPrefixCacheMatchInfo(10, 10))
 
 	pod3 := newMockPod()