Prefix Aware Scorer #48

oglok · 2025-04-23T12:33:57Z

This PR implements:

prefix store based in a radix tree storing a map of prefixes (key) and prefixEntry (value) that contains the target pod, last time used and model name.
prefix aware scorer
addition of new scorer to the scorer manager
unit tests for new store and scorer

NOTE: there is a runMaintenance function in the store that is not being used. Need to figure out where to call it from without blocking all requests.

…ill be the target for a request. Session affinity scorer added

- Rename SessionId to SessionID - Remove datastore from scoreTargets, add datastore to SessionAffinityScorer - Rename ScoredPod to PodScore

…orerManager

…f ScoreMng - If some specific scorer failed to score pods - just log the problem, skip it and continue to the next scorer

Signed-off-by: Shane Utt <shaneutt@linux.com>

Bumps [golang.org/x/net](https://github.com/golang/net) from 0.37.0 to 0.38.0. - [Commits](golang/net@v0.37.0...v0.38.0) --- updated-dependencies: - dependency-name: golang.org/x/net dependency-version: 0.38.0 dependency-type: indirect ... Signed-off-by: dependabot[bot] <support@github.com>

…es/golang.org/x/net-0.38.0 Bump golang.org/x/net from 0.37.0 to 0.38.0

Add scorers support in scheduler

…eployments First iteration of development deployments & environments

Signed-off-by: Shane Utt <shaneutt@linux.com>

…ilds fix: basic container image builds for linux

Signed-off-by: Shane Utt <shaneutt@linux.com>

Signed-off-by: Etai Lev Ran <elevran@gmail.com>

empty top level kustomization.yaml - make CICD happy

Fix kustomize envs

provide more defaults in Makefile

Signed-off-by: Shane Utt <shaneutt@linux.com>

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

Copilot

Pull Request Overview

This PR implements a new prefix store based on a radix tree and a prefix-aware scorer to route requests more efficiently. Key changes include:

Adding a PrefixStore and its configuration to scheduler configuration with environmental overrides.
Implementing a new PrefixAwareScorer and integrating it into the scheduler.
Providing unit tests for the new store and scorer components.

Reviewed Changes

Copilot reviewed 7 out of 8 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
pkg/epp/scheduling/types/types.go	Added a Prompt field to the LLMRequest structure.
pkg/epp/scheduling/scheduler.go	Introduced prefix store configuration and integrated the new scorer.
pkg/epp/scheduling/prefix_store.go	Implements the prefix store with entry addition, lookup, eviction, and maintenance.
pkg/epp/scheduling/prefix_store_test.go	Unit tests for basic prefix operations, constraints, TTL expiration, max entries, and maintenance routine.
pkg/epp/scheduling/prefix_aware_scorer.go	Defines the scorer which uses the prefix store to score pods.
pkg/epp/scheduling/prefix_aware_scorer_test.go	Tests behavior of the new scorer with various input scenarios.

Files not reviewed (1)

go.mod: Language not supported

Copilot · 2025-04-23T12:56:39Z

pkg/epp/scheduling/prefix_store.go

+	// Update LastUsed time for the matched entry
+	entry.LastUsed = time.Now()
+	ps.tree.Insert(matchedPrefix, entry)


Modifying 'entry.LastUsed' under a read lock (RLock) in FindPodForPrefix may lead to a race condition. Consider acquiring the write lock (Lock) when updating mutable state or restructuring the update to ensure thread safety.

Suggested change

// Update LastUsed time for the matched entry

entry.LastUsed = time.Now()

ps.tree.Insert(matchedPrefix, entry)

// Update LastUsed time for the matched entry

ps.mu.RUnlock() // Release the read lock before acquiring the write lock

ps.mu.Lock()

entry.LastUsed = time.Now()

ps.tree.Insert(matchedPrefix, entry)

ps.mu.Unlock()

rootfs · 2025-04-23T12:58:36Z

pkg/epp/scheduling/scheduler.go

@@ -43,6 +48,10 @@ const (
 	defaultQueueThresholdCritical = 5
 	defaultQueueingThresholdLoRA  = 128
 	defaultLoraAffinityThreshold  = 0.999
+	defaultPrefixStoreMaxEntries  = 1000
+	defaultPrefixStoreMinLen      = 3


this is probably too small :D

oops! testing defaults xD

rootfs · 2025-04-23T12:59:50Z

pkg/epp/scheduling/scheduler.go

+	defaultPrefixStoreMaxEntries  = 1000
+	defaultPrefixStoreMinLen      = 3
+	defaultPrefixStoreMaxLen      = 100
+	defaultPrefixStoreTTLHours    = 24


minute instead or hour?

I'm not really sure what should be a reasonable TTL for LLM inferencing.

vMaroon · 2025-04-23T14:21:06Z

The dev branch was rebased on upstream main, which includes the new and different Scorer interface. The Scorer pulled from there defines a Score function that works per-pod and not group of pods. We need to discuss whether to propose changing that but it makes sense to first adapt to it.

I'll open a KVCacheAwareScorer PR soon and we can sync.

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

elevran · 2025-04-27T12:47:05Z

pkg/epp/scheduling/prefix_aware_scorer.go

+// between the request's prompt and stored prefixes. The score is normalized between 0 and 1,
+// where 1 represents the longest matching prefix.
+type PrefixAwareScorer struct {
+	weight      float64


IIRC, weights are extracted to the layer above.

elevran · 2025-04-27T12:48:39Z

pkg/epp/scheduling/prefix_aware_scorer.go

+}
+
+// NewPrefixAwareScorer creates a new PrefixAwareScorer with the given weight and prefix store
+func NewPrefixAwareScorer(weight float64, prefixStore *PrefixStore) Scorer {


What is the rationale for receiving (instead of creating internally) the prefixStore?
Seems like this would be an internal implementation decision?
If so, PrefixStore is likely not an exported type.

elevran · 2025-04-27T12:51:26Z

pkg/epp/scheduling/prefix_aware_scorer.go

+	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 {


consider moving this up to the function initialization?
That way, you'd be assured that you start with 0s for all Pods and would not need to repeat it multiple times.

elevran · 2025-04-27T12:53:08Z

pkg/epp/scheduling/prefix_store.go

+
+// PrefixEntry represents a single entry in the prefix store
+type PrefixEntry struct {
+	PodRef    types.NamespacedName


q: is this how defined in PodMetrics.Pod or would a simple string suffice?

elevran · 2025-04-27T12:54:03Z

pkg/epp/scheduling/prefix_store.go

+)
+
+// PrefixEntry represents a single entry in the prefix store
+type PrefixEntry struct {


q: should the struct and its fields be exported?
On first glance seems that it should be internal use.

elevran · 2025-04-27T13:02:37Z

pkg/epp/scheduling/prefix_store.go

+	errutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/error"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
+)
+


It would be helpful to have a high level description of data structure and algorithms, and tie those to the processing steps of the LLM request (e.g., consulted in read only or also updated on prompt request, are responses handed, ...)

elevran · 2025-04-27T13:05:07Z

pkg/epp/scheduling/prefix_store.go

+		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)


will this allow multiple hits on a prefix (e.g., both pods A and B hold the prefix) or just one (last writer wins)?
I think we would want to have multiple (e.g., in case other scorers score these Pods differently).

elevran · 2025-04-27T13:07:04Z

pkg/epp/scheduling/prefix_store.go

+
+	// 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)


Q: would any entries "below" the evicted node become dangling following the eviction?

elevran · 2025-04-27T13:07:58Z

pkg/epp/scheduling/prefix_store.go

+	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]
+	}


consider extracting to an isValidPrompt function that can be called from multiple places

elevran · 2025-04-27T13:08:44Z

pkg/epp/scheduling/prefix_store.go

+	}
+
+	// Use LongestPrefix to find the best match
+	matchedPrefix, val, found := ps.tree.LongestPrefix(prefix)


The current implementation seems to score only the longest prefix and not any of the other pods?

elevran · 2025-04-27T13:11:49Z

@oglok would be great if you can rebase off latest dev for easier review. Most of the 78 modified files are coming in from different PRs and "polluting" your PR

elevran · 2025-04-27T13:21:21Z

pkg/epp/scheduling/session_affinity_scorer.go

Are there downsides to having this feature (including scorer, store and tests) in its own directory?
That would minimize the changes in the existing code base

elevran · 2025-04-27T13:22:30Z

pkg/epp/scheduling/session_affinity_scorer.go

+	}
+}
+
+// ScoreTargets does the actual scoring of the target pods by the session affinity.


It is not clear where/when the store is updated with new prompt requests and replies.
Is this currently handled?

oglok · 2025-04-28T10:02:21Z

Hey @elevran ! thanks for reviewing the PR. I'm working on the rebase, and rework my code because apparently the scorers interface merged upstream has changed. I'll get your comments in while doing that!

mayabar and others added 30 commits April 10, 2025 14:50

Add initial support for scorers, used as part of decision which pod w…

0c95c2a

…ill be the target for a request. Session affinity scorer added

Fixes in scores infrastructure & session aware scorer

bde57da

- Add cleanup for session->pod map

8f9785e

- Rename SessionId to SessionID - Remove datastore from scoreTargets, add datastore to SessionAffinityScorer - Rename ScoredPod to PodScore

Export score and pod for external implementations

52af66f

Rename session id header

90e23bc

Separate code of Scorer interface and scorer implementations + add sc…

c946acc

…orerManager

Remove vllmRequest from scoreTargets API since it exists in the context

137ca09

Support negative score weights

c42f72a

Fix fakeDataStore to be compatible with DataStore intereface

a05a573

- Check for nils in list of available pods in main scoring function o…

aca8e07

…f ScoreMng - If some specific scorer failed to score pods - just log the problem, skip it and continue to the next scorer

[version bump] Promote 0.0.2 to prod, bump dev to 0.0.3

bae2a66

chore: move openshift router deployment to extra

c398f4b

Signed-off-by: Shane Utt <shaneutt@linux.com>

feat: add deployment for sail operator

95e6bb1

Signed-off-by: Shane Utt <shaneutt@linux.com>

feat: add istio control-plane deployment

6d12b06

Signed-off-by: Shane Utt <shaneutt@linux.com>

feat: add vllm simulator deployment

8c4eb46

Signed-off-by: Shane Utt <shaneutt@linux.com>

feat: add inference-gateway deployment

58ef159

Signed-off-by: Shane Utt <shaneutt@linux.com>

feat: add kind environment deployment

f606e0d

Signed-off-by: Shane Utt <shaneutt@linux.com>

feat: kind dev env deployment script

c679724

Signed-off-by: Shane Utt <shaneutt@linux.com>

Merge pull request neuralmagic#5 from neuralmagic/dependabot/go_modul…

09e79e6

…es/golang.org/x/net-0.38.0 Bump golang.org/x/net from 0.37.0 to 0.38.0

Merge pull request neuralmagic#1 from mayabar/main

d8303a0

Add scorers support in scheduler

Merge pull request neuralmagic#4 from shaneutt/shaneutt/initial-dev-d…

dad8db2

…eployments First iteration of development deployments & environments

fix: basic container image builds for linux

f608526

Signed-off-by: Shane Utt <shaneutt@linux.com>

fix: lint fix

2dd2ee7

Signed-off-by: Shane Utt <shaneutt@linux.com>

Merge pull request neuralmagic#10 from shaneutt/shaneutt/fix-image-bu…

10be213

…ilds fix: basic container image builds for linux

fix: move openshift deployment to environments

a24a801

Signed-off-by: Shane Utt <shaneutt@linux.com>

fix: retarget kustomize deployments in Makefile

950e07b

Signed-off-by: Shane Utt <shaneutt@linux.com>

empty top level kustomization.yaml - make CICD happy

47bed9d

Signed-off-by: Etai Lev Ran <elevran@gmail.com>

Merge pull request neuralmagic#17 from elevran/deploy_kustomization_yaml

0c4e6c8

empty top level kustomization.yaml - make CICD happy

Merge pull request neuralmagic#15 from shaneutt/fix-kustomize-envs

e0dcba6

Fix kustomize envs

elevran and others added 15 commits April 22, 2025 23:49

Merge pull request neuralmagic#46 from shaneutt/defaults-in-namespace

47ce2f4

provide more defaults in Makefile

fix: broken vars in makefile target

b467a55

Signed-off-by: Shane Utt <shaneutt@linux.com>

testing deploy

f9b5889

fix: add REGISTRY_SECRET to tektok ocp deploy

064166f

Signed-off-by: Shane Utt <shaneutt@linux.com>

fix: remove ctx change in makefile for k8s dev env

4036bd0

Signed-off-by: Shane Utt <shaneutt@linux.com>

updating rbac

840f048

[version bump] Promote 0.0.3 to prod, bump dev to 0.0.4

b0d29ec

Add prefix store data structures

caf3b78

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

Add prefix store functionality

a153c34

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

Follow project pattern for errors

765db02

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

Prefix Aware Scorer

c778139

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

Integrate prefix aware scorer to score manager

b6d6487

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

Load config from env vars

dd0dfe2

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

Add unit tests for prefix store

bde07c5

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

Add unit tests for prefix aware scorer

4966409

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

oglok changed the title ~~WIP: Prefix Aware Scorer~~ Prefix Aware Scorer Apr 23, 2025

rootfs requested a review from Copilot April 23, 2025 12:55

Copilot AI reviewed Apr 23, 2025

View reviewed changes

rootfs reviewed Apr 23, 2025

View reviewed changes

mayabar force-pushed the dev branch from b0d29ec to a80bcfc Compare April 23, 2025 14:13

oglok force-pushed the prefix_scorer branch from 37ed68d to 7383242 Compare April 26, 2025 09:29

Prefix aware scorer data structure fix

037fadd

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

oglok force-pushed the prefix_scorer branch from 7383242 to 037fadd Compare April 26, 2025 09:30

elevran requested changes Apr 27, 2025

View reviewed changes

elevran reviewed Apr 27, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Prefix Aware Scorer #48

Prefix Aware Scorer #48

oglok commented Apr 23, 2025

Copilot AI left a comment

Copilot AI Apr 23, 2025

rootfs Apr 23, 2025

oglok Apr 23, 2025

rootfs Apr 23, 2025

oglok Apr 23, 2025

vMaroon commented Apr 23, 2025

elevran Apr 27, 2025

elevran Apr 27, 2025

elevran Apr 27, 2025

elevran Apr 27, 2025

elevran Apr 27, 2025

elevran Apr 27, 2025

elevran Apr 27, 2025

elevran Apr 27, 2025

elevran Apr 27, 2025

elevran Apr 27, 2025

elevran commented Apr 27, 2025

elevran Apr 27, 2025

elevran Apr 27, 2025 •

edited

Loading

oglok commented Apr 28, 2025

Prefix Aware Scorer #48

Are you sure you want to change the base?

Prefix Aware Scorer #48

Conversation

oglok commented Apr 23, 2025

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Copilot AI Apr 23, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

vMaroon commented Apr 23, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

elevran commented Apr 27, 2025

Choose a reason for hiding this comment

elevran Apr 27, 2025 • edited Loading

Choose a reason for hiding this comment

oglok commented Apr 28, 2025

elevran Apr 27, 2025 •

edited

Loading