Session affinity scorer (#117)

* 'session affinity scorer' partial implementation (without headers in response)

* Fix in filling request headers

* Encoded value of namespaced pod name is sent in response to client

* Support of session affinity scorer configuration via environment variables, is added

* Go file for session affinity scorer is renamed

* Redundant 'sessions' field is removed

* Redundant 'ScorerWithPostResponse' struct is removed

* - SessionID is renamed to sessionToken
- Map fetch is done instead of loop

* Session token name is changed to 'x-session-token'

* Minor fixes are made in README

* Small fix after merge

---------

Co-authored-by: Shmuel Kallner <[email protected]>

Author: dmitripikus

README.md

@@ -23,6 +23,12 @@ export ENABLE_LOAD_AWARE_SCORER=true
 export LOAD_AWARE_SCORER_WEIGHT=1.0
 ```
 
+To enable the SessionAwareScorer, the following environment variables must be configured:
+```
+export ENABLE_SESSION_AWARE_SCORER=true
+export SESSION_AWARE_SCORER_WEIGHT=1.0
+```
+
 To enable Prefill/Decode (PD) processing, the following environment variable must be configured:
 ```
 export PD_ENABLED=true

pkg/epp/handlers/request.go

@@ -159,7 +159,7 @@ func (s *StreamingServer) HandleRequestHeaders(ctx context.Context, reqCtx *Requ
 	}
 
 	for _, header := range req.RequestHeaders.Headers.Headers {
-		reqCtx.RequestHeaders[header.Key] = header.Value
+		reqCtx.RequestHeaders[header.Key] = string(header.RawValue)
 	}
 
 	return nil

pkg/epp/scheduling/local_config.go

@@ -27,14 +27,16 @@ import (
 )
 
 const (
-	kvCacheScorerEnablementEnvVar   = "ENABLE_KVCACHE_AWARE_SCORER"
-	loadAwareScorerEnablementEnvVar = "ENABLE_LOAD_AWARE_SCORER"
-	prefixScorerEnablementEnvVar    = "ENABLE_PREFIX_AWARE_SCORER"
-	pdFilterEnablementEnvVar        = "ENABLE_PD_FILTER"
-
-	kvCacheScorerWeightEnvVar   = "KVCACHE_AWARE_SCORER_WEIGHT"
-	loadAwareScorerWeightEnvVar = "LOAD_AWARE_SCORER_WEIGHT"
-	prefixScorerWeightEnvVar    = "PREFIX_AWARE_SCORER_WEIGHT"
+	kvCacheScorerEnablementEnvVar      = "ENABLE_KVCACHE_AWARE_SCORER"
+	loadAwareScorerEnablementEnvVar    = "ENABLE_LOAD_AWARE_SCORER"
+	prefixScorerEnablementEnvVar       = "ENABLE_PREFIX_AWARE_SCORER"
+	sessionAwareScorerEnablementEnvVar = "ENABLE_SESSION_AWARE_SCORER"
+	pdFilterEnablementEnvVar           = "ENABLE_PD_FILTER"
+
+	kvCacheScorerWeightEnvVar      = "KVCACHE_AWARE_SCORER_WEIGHT"
+	loadAwareScorerWeightEnvVar    = "LOAD_AWARE_SCORER_WEIGHT"
+	prefixScorerWeightEnvVar       = "PREFIX_AWARE_SCORER_WEIGHT"
+	sessionAwareScorerWeightEnvVar = "SESSION_AWARE_SCORER_WEIGHT"
 )
 
 func init() {
@@ -45,6 +47,7 @@ func setDefaultConfig() {
 	// since the default config is a global variable, we add this function to minimize rebase conflicts.
 	// this configuration is a temporary state, it should be better streamlined.
 	setLoadAwareScorer()
+	setSessionAwareScorer()
 	setKVCacheAwareScorer()
 	setPrefixScorer()
 
@@ -65,6 +68,23 @@ func setLoadAwareScorer() {
 	loggerDebug.Info("Initialized LoadAwareScorer", "weight", loadBasedScorerWeight)
 }
 
+func setSessionAwareScorer() {
+	ctx := context.Background()
+	loggerDebug := log.FromContext(ctx).WithName("scheduler_config").V(logutil.DEBUG)
+
+	if envutil.GetEnvString(sessionAwareScorerEnablementEnvVar, "false", loggerDebug) != "true" {
+		loggerDebug.Info("Skipping SessionAwareScorer creation as it is not enabled")
+		return
+	}
+
+	sessionBasedScorerWeight := envutil.GetEnvInt(sessionAwareScorerWeightEnvVar, 1, loggerDebug)
+	sessionAffinity := scorer.NewSessionAffinity()
+
+	defaultConfig.scorers[sessionAffinity] = sessionBasedScorerWeight
+	defaultConfig.postResponsePlugins = append(defaultConfig.postResponsePlugins, sessionAffinity)
+	loggerDebug.Info("Initialized SessionAwareScorer", "weight", sessionBasedScorerWeight)
+}
+
 func setKVCacheAwareScorer() {
 	ctx := context.Background()
 	loggerDebug := log.FromContext(ctx).WithName("scheduler_config").V(logutil.DEBUG)

pkg/epp/scheduling/plugins/scorer/session-affinity-scorer.go

@@ -0,0 +1,79 @@
+/*
+Copyright 2025 The Kubernetes Authors.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+	http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package scorer
+
+import (
+	"encoding/base64"
+	"time"
+
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
+)
+
+const (
+	sessionKeepAliveTime           = 60 * time.Minute  // How long should an idle session be kept alive
+	sessionKeepAliveCheckFrequency = 15 * time.Minute  // How often to check for overly idle sessions
+	sessionTokenHeader             = "x-session-token" // name of the session header in request
+)
+
+// sessionAffinity is a routing scorer that routes subsequent
+// requests in a session to the same pod as the first request in the
+// session was sent to, by giving that pod the specified weight and assigning
+// zero score to the rest of the targets
+type SessionAffinity struct {
+}
+
+func NewSessionAffinity() *SessionAffinity {
+	return &SessionAffinity{}
+}
+
+func (s *SessionAffinity) Name() string {
+	return "session affinity scorer"
+}
+
+func (s *SessionAffinity) Score(ctx *types.SchedulingContext, pods []types.Pod) map[types.Pod]float64 {
+	scoredPods := make(map[types.Pod]float64)
+
+	reqHeaders := ctx.Req.Headers
+
+	var sessionToken = ""
+	v, ok := reqHeaders[sessionTokenHeader]
+	if ok {
+		sessionToken = v
+	}
+
+	podName := ""
+	if sessionToken != "" {
+		decodedBytes, err := base64.StdEncoding.DecodeString(sessionToken)
+		if err != nil {
+			ctx.Logger.Error(err, "Error decoding")
+		} else {
+			podName = string(decodedBytes)
+		}
+	}
+	for _, pod := range pods {
+		if podName == "" {
+			scoredPods[pod] = 0.0
+		} else {
+			if pod.GetPod().NamespacedName.String() == podName {
+				scoredPods[pod] = 1.0
+			}
+		}
+	}
+
+	return scoredPods
+}
+
+func (s *SessionAffinity) PostResponse(ctx *types.SchedulingContext, pod types.Pod) {
+	ctx.MutatedHeaders[sessionTokenHeader] = base64.StdEncoding.EncodeToString([]byte(pod.GetPod().NamespacedName.String()))
+}