Repository: kubernetes-sigs/gateway-api-inference-extension

Issue:
Rollout is stuck when leader election is enabled
<!-- Please use this template while reporting a bug and provide as much info as possible. Not doing so may result in your bug not being addressed in a timely manner. Thanks!

If the matter is security related, please disclose it privately via https://kubernetes.io/security/
-->

**What happened**:

When leader election is enabled with 3 EPP replicas, updating EPP deployment gets stuck. This is because when leader election is enabled, there is at most 1 ready replica at any given time. Killing the old leader to give the lease to a new leader would violate the default `maxUnavailable=25%` in k8s deployment rolling update.

**What you expected to happen**:

We would need the `Recreate` update policy to allow the old leader to be killed and let the new leader take over. 

**How to reproduce it (as minimally and precisely as possible)**:

**Anything else we need to know?**:

**Environment**:
- Kubernetes version (use `kubectl version`):
- Inference extension version (use `git describe --tags --dirty --always`):
- Cloud provider or hardware configuration:
- Install tools:
- Others:


Discussion:

Solution patch:
diff --git a/config/charts/inferencepool/README.md b/config/charts/inferencepool/README.md
index b6629d2..9a8be09 100644
--- a/config/charts/inferencepool/README.md
+++ b/config/charts/inferencepool/README.md
@@ -103,19 +103,30 @@ $ helm install triton-llama3-8b-instruct \
 
 To deploy the EndpointPicker in a high-availability (HA) active-passive configuration, you can enable leader election. When enabled, the EPP deployment will have multiple replicas, but only one "leader" replica will be active and ready to process traffic at any given time. If the leader pod fails, another pod will be elected as the new leader, ensuring service continuity.
 
-To enable HA, set `inferenceExtension.flags.has-enable-leader-election` to `true` and increase the number of replicas in your `values.yaml` file:
+To enable HA, set `inferenceExtension.enableLeaderElection` to `true`.
 
-```yaml
-inferenceExtension:
-  replicas: 3
-  has-enable-leader-election: true
-```
+* Via `--set` flag:
 
-Then apply it with:
+  ```txt
+  helm install vllm-llama3-8b-instruct \
+  --set inferencePool.modelServers.matchLabels.app=vllm-llama3-8b-instruct \
+  --set inferenceExtension.enableLeaderElection=true \
+  --set provider=[none|gke] \
+  oci://us-central1-docker.pkg.dev/k8s-staging-images/gateway-api-inference-extension/charts/inferencepool --version v0
+  ```
 
-```txt
-helm install vllm-llama3-8b-instruct ./config/charts/inferencepool -f values.yaml
-```
+* Via `values.yaml`:
+
+  ```yaml
+  inferenceExtension:
+    enableLeaderElection: true
+  ```
+
+  Then apply it with:
+
+  ```txt
+  helm install vllm-llama3-8b-instruct ./config/charts/inferencepool -f values.yaml
+  ```
 
 ### Install with Monitoring
 
@@ -171,8 +182,7 @@ The following table list the configurable parameters of the chart.
 | `inferenceExtension.extraServicePorts`      | List of additional service ports to expose. Defaults to `[]`.                                                          |
 | `inferenceExtension.flags`                  | List of flags which are passed through to endpoint picker. Example flags, enable-pprof, grpc-port etc. Refer [runner.go](https://github.com/kubernetes-sigs/gateway-api-inference-extension/blob/main/cmd/epp/runner/runner.go) for complete list.                                                            |
 | `inferenceExtension.affinity`               | Affinity for the endpoint picker. Defaults to `{}`.                                                                    |
-| `inferenceExtension.tolerations`            | Tolerations for the endpoint picker. Defaults to `[]`.                                                                 |
-| `inferenceExtension.flags.has-enable-leader-election` | Enable leader election for high availability. When enabled, only one EPP pod (the leader) will be ready to serve traffic.       |
+| `inferenceExtension.tolerations`            | Tolerations for the endpoint picker. Defaults to `[]`.                                                                 |   |
 | `inferenceExtension.monitoring.interval`   | Metrics scraping interval for monitoring. Defaults to `10s`.                                                           |
 | `inferenceExtension.monitoring.secret.name` | Name of the service account token secret for metrics authentication. Defaults to `inference-gateway-sa-metrics-reader-secret`. |
 | `inferenceExtension.monitoring.prometheus.enabled` | Enable Prometheus ServiceMonitor creation for EPP metrics collection. Defaults to `false`.                      |
diff --git a/config/charts/inferencepool/templates/epp-deployment.yaml b/config/charts/inferencepool/templates/epp-deployment.yaml
index abdc05a..d37ba3e 100644
--- a/config/charts/inferencepool/templates/epp-deployment.yaml
+++ b/config/charts/inferencepool/templates/epp-deployment.yaml
@@ -6,7 +6,19 @@ metadata:
   labels:
     {{- include "gateway-api-inference-extension.labels" . | nindent 4 }}
 spec:
-  replicas: {{ .Values.inferenceExtension.replicas | default 1 }}
+  {{- if .Values.inferenceExtension.enableLeaderElection }}
+  replicas: 3
+  {{- else }}
+  replicas: 1
+  {{- end }}
+  strategy:
+    # The current recommended EPP deployment pattern is to have a single active replica. This ensures 
+    # optimal performance of the stateful operations such prefix cache aware scorer.
+    # The Recreate strategy the old replica is killed immediately, and allow the new replica(s) to 
+    # quickly take over. This is particularly important in the high availability set up with leader
+    # election, as the rolling update strategy would prevent the old leader being killed because 
+    # otherwise the maxUnavailable would be 100%.
+    type: Recreate
   selector:
     matchLabels:
       {{- include "gateway-api-inference-extension.selectorLabels" . | nindent 6 }}
@@ -33,10 +45,6 @@ spec:
         - "json"
         - --config-file
         - "/config/{{ .Values.inferenceExtension.pluginsConfigFile }}"
-        {{- range .Values.inferenceExtension.flags }}
-        - "--{{ .name }}"
-        - "{{ .value }}"
-        {{- end }}
         {{- if eq (.Values.inferencePool.modelServerType | default "vllm") "triton-tensorrt-llm" }}
         - --total-queued-requests-metric
         - "nv_trt_llm_request_metrics{request_type=waiting}"
@@ -45,6 +53,14 @@ spec:
         - --lora-info-metric
         - "" # Set an empty metric to disable LoRA metric scraping as they are not supported by Triton yet.
         {{- end }}
+        {{- if .Values.inferenceExtension.enableLeaderElection }}
+        - --ha-enable-leader-election
+        {{- end }}
+        # Pass additional flags via the inferenceExtension.flags field in values.yaml.
+        {{- range .Values.inferenceExtension.flags }}
+        - "--{{ .name }}"
+        - "{{ .value }}"
+        {{- end }}
         ports:
         - name: grpc
           containerPort: 9002
@@ -77,8 +93,8 @@ spec:
             port: 9003
             service: inference-extension
           {{- end }}
-          initialDelaySeconds: 5
-          periodSeconds: 10
+          periodSeconds: 2
+
         env:
         - name: NAMESPACE
           valueFrom:
diff --git a/config/charts/inferencepool/templates/gke.yaml b/config/charts/inferencepool/templates/gke.yaml
index 59e186a..a0118f5 100644
--- a/config/charts/inferencepool/templates/gke.yaml
+++ b/config/charts/inferencepool/templates/gke.yaml
@@ -13,6 +13,10 @@ spec:
     kind: InferencePool
     name: {{ .Release.Name }}
   default:
+    # Set a more aggressive health check than the default 5s for faster switch
+    # over during EPP rollout.
+    timeoutSec: 2
+    checkIntervalSec: 2
     config:
       type: HTTP
       httpHealthCheck:
diff --git a/config/manifests/gateway/gke/gcp-backend-policy.yaml b/config/manifests/gateway/gke/gcp-backend-policy.yaml
deleted file mode 100644
index 9367865..0000000
--- a/config/manifests/gateway/gke/gcp-backend-policy.yaml
+++ /dev/null
@@ -1,13 +0,0 @@
-apiVersion: networking.gke.io/v1
-kind: GCPBackendPolicy
-metadata:
-  name: inferencepool-backend-policy
-spec:
-  targetRef:
-    group: "inference.networking.k8s.io"
-    kind: InferencePool
-    name: vllm-llama3-8b-instruct
-  default:
-    timeoutSec: 300
-    logging:
-      enabled: true
diff --git a/config/manifests/gateway/gke/healthcheck.yaml b/config/manifests/gateway/gke/healthcheck.yaml
deleted file mode 100644
index c9abb69..0000000
--- a/config/manifests/gateway/gke/healthcheck.yaml
+++ /dev/null
@@ -1,16 +0,0 @@
-kind: HealthCheckPolicy
-apiVersion: networking.gke.io/v1
-metadata:
-  name: health-check-policy
-  namespace: default
-spec:
-  targetRef:
-    group: "inference.networking.k8s.io"
-    kind: InferencePool
-    name: vllm-llama3-8b-instruct
-  default:
-    config:
-      type: HTTP
-      httpHealthCheck:
-          requestPath: /health
-          port: 8000
diff --git a/pkg/epp/scheduling/framework/plugins/multi/prefix/plugin.go b/pkg/epp/scheduling/framework/plugins/multi/prefix/plugin.go
index 40e8806..6bc81a4 100644
--- a/pkg/epp/scheduling/framework/plugins/multi/prefix/plugin.go
+++ b/pkg/epp/scheduling/framework/plugins/multi/prefix/plugin.go
@@ -37,7 +37,7 @@ import (
 
 const (
 	// vLLM default token block size is 16, and a good guess of average characters per token is 4.
-	DefaultHashBlockSize = 64
+	DefaultBlockSize = 64
 	// The maximum number of blocks to match. Two long requests with the same prefix up to this
 	// limit will be indistinguishable.
 	// This parameter provides a trade-off between cache size, prefix matching speed and matching
@@ -58,15 +58,15 @@ const (
 )
 
 var DefaultConfig = Config{
-	HashBlockSize:          DefaultHashBlockSize,
+	BlockSize:              DefaultBlockSize,
 	MaxPrefixBlocksToMatch: DefaultMaxPrefixBlocks,
 	LRUCapacityPerServer:   DefaultLRUCapacityPerServer,
 }
 
 type Config struct {
-	// The input prompt is broken into sizes of HashBlockSize to calculate block hashes . Requests
+	// The input prompt is broken into sizes of BlockSize to calculate block hashes . Requests
 	// with length shorter than the block size will be ignored.
-	HashBlockSize int `json:"hashBlockSize"`
+	BlockSize int `json:"blockSize"`
 	// MaxPrefixBlocksToMatch is the maximum number of prefix blocks to match. Input beyond this limit will
 	// be ignored.
 	MaxPrefixBlocksToMatch int `json:"maxPrefixBlocksToMatch"`
@@ -133,7 +133,7 @@ var (
 // PrefixCachePluginFactory defines the factory function for Prefix plugin.
 func PrefixCachePluginFactory(name string, rawParameters json.RawMessage, handle plugins.Handle) (plugins.Plugin, error) {
 	parameters := Config{
-		HashBlockSize:          DefaultHashBlockSize,
+		BlockSize:              DefaultBlockSize,
 		MaxPrefixBlocksToMatch: DefaultMaxPrefixBlocks,
 		LRUCapacityPerServer:   DefaultLRUCapacityPerServer,
 	}
@@ -180,7 +180,7 @@ func (p *Plugin) WithName(name string) *Plugin {
 // Score returns the scoring result for the given list of pods based on context.
 func (p *Plugin) Score(ctx context.Context, cycleState *types.CycleState, request *types.LLMRequest, pods []types.Pod) map[types.Pod]float64 {
 	// pre score step, hashing prompt and find longest prefix match.
-	hashes := hashPrompt(ctx, request, p.config.HashBlockSize, p.config.MaxPrefixBlocksToMatch)
+	hashes := hashPrompt(ctx, request, p.config.BlockSize, p.config.MaxPrefixBlocksToMatch)
 	state := &SchedulingContextState{
 		PrefixHashes:       hashes,
 		PrefixCacheServers: p.matchLongestPrefix(ctx, hashes),
@@ -231,7 +231,7 @@ func (p *Plugin) PreRequest(ctx context.Context, request *types.LLMRequest, sche
 
 	total := len(state.PrefixHashes)
 	matchLen := state.PrefixCacheServers[ServerID(targetPod.NamespacedName)]
-	metrics.RecordPrefixCacheMatch(matchLen*p.config.HashBlockSize, total*p.config.HashBlockSize)
+	metrics.RecordPrefixCacheMatch(matchLen*p.config.BlockSize, total*p.config.BlockSize)
 }
 
 // matchLongestPrefix returns a map of servers and length of prefix that each server caches.
diff --git a/site-src/guides/epp-configuration/config-text.md b/site-src/guides/epp-configuration/config-text.md
index cdb3714..382edba 100644
--- a/site-src/guides/epp-configuration/config-text.md
+++ b/site-src/guides/epp-configuration/config-text.md
@@ -85,7 +85,7 @@ kind: EndpointPickerConfig
 plugins:
 - type: prefix-cache-scorer
   parameters:
-    hashBlockSize: 5
+    blockSize: 5
     maxPrefixBlocksToMatch: 256
     lruCapacityPerServer: 31250
 schedulingProfiles:
@@ -152,7 +152,7 @@ spec:
           plugins:
           - type: prefix-cache-scorer
             parameters:
-              hashBlockSize: 5
+              blockSize: 5
               maxPrefixBlocksToMatch: 256
               lruCapacityPerServer: 31250
           schedulingProfiles:
@@ -171,7 +171,7 @@ kind: EndpointPickerConfig
 plugins:
 - type: prefix-cache-scorer
   parameters:
-    hashBlockSize: 5
+    blockSize: 5
     maxPrefixBlocksToMatch: 256
     lruCapacityPerServer: 31250
 - type: single-profile-handler
@@ -201,7 +201,7 @@ Scores pods based on the amount of the prompt is believed to be in the pod's KvC
 
 - *Type*: prefix-cache-scorer
 - *Parameters*:
-  - `hashBlockSize` specified the size of the blocks to break up the input prompt when
+  - `blockSize` specified the size of the blocks to break up the input prompt when
     calculating the block hashes. If not specified defaults to `64`
   - `maxPrefixBlocksToMatch` specifies the maximum number of prefix blocks to match. If
    not specified defaults to `256`
diff --git a/site-src/guides/epp-configuration/prefix-aware.md b/site-src/guides/epp-configuration/prefix-aware.md
index 9c1074b..88573c4 100644
--- a/site-src/guides/epp-configuration/prefix-aware.md
+++ b/site-src/guides/epp-configuration/prefix-aware.md
@@ -14,7 +14,7 @@ Like any other plugins, the prefix cache aware plugin can be enabled/disabled vi
 
 The prefix cache plugin exposes the following advanced configuration parameters:
 
-* `hashBlockSize`: The plugin matches prefixes in the unit of blocks. This is the size
+* `blockSize`: The plugin matches prefixes in the unit of blocks. This is the size
 of each block in number of bytes. vLLM default block size is 16 tokens. Assume 4 characters per token, the default
 is set to 64 in EPP. The default is recommended unless performance is critical for use cases with
 extremely long inputs.



Test patch:
diff --git a/pkg/epp/config/loader/configloader_test.go b/pkg/epp/config/loader/configloader_test.go
index ff7b652..b0b0741 100644
--- a/pkg/epp/config/loader/configloader_test.go
+++ b/pkg/epp/config/loader/configloader_test.go
@@ -73,7 +73,7 @@ func TestLoadRawConfiguration(t *testing.T) {
 			},
 			{
 				Type:       test2Type,
-				Parameters: json.RawMessage("{\"hashBlockSize\":32}"),
+				Parameters: json.RawMessage("{\"blockSize\":32}"),
 			},
 			{
 				Name: "testPicker",
@@ -175,7 +175,7 @@ func TestLoadRawConfigurationWithDefaults(t *testing.T) {
 			{
 				Name:       test2Type,
 				Type:       test2Type,
-				Parameters: json.RawMessage("{\"hashBlockSize\":32}"),
+				Parameters: json.RawMessage("{\"blockSize\":32}"),
 			},
 			{
 				Name: "testPicker",
@@ -464,7 +464,7 @@ plugins:
   type: test-profile-handler
 - type: test-two
   parameters:
-    hashBlockSize: 32
+    blockSize: 32
 - name: testPicker
   type: test-picker
 schedulingProfiles:
@@ -767,7 +767,7 @@ plugins:
 - name: prefixCacheScorer
   type: prefix-cache-scorer
   parameters:
-    hashBlockSize: 32
+    blockSize: 32
 - name: maxScorePicker
   type: max-score-picker
 - name: profileHandler
@@ -792,7 +792,7 @@ plugins:
 - name: prefixCacheScorer
   type: prefix-cache-scorer
   parameters:
-    hashBlockSize: 32
+    blockSize: 32
 schedulingProfiles:
 - name: default
   plugins:
@@ -826,7 +826,7 @@ plugins:
 - name: prefixCacheScorer
   type: prefix-cache-scorer
   parameters:
-    hashBlockSize: asdf
+    blockSize: asdf
 schedulingProfiles:
 - name: default
   plugins:
diff --git a/pkg/epp/scheduling/framework/plugins/multi/prefix/plugin_test.go b/pkg/epp/scheduling/framework/plugins/multi/prefix/plugin_test.go
index df2eae6..54a00ab 100644
--- a/pkg/epp/scheduling/framework/plugins/multi/prefix/plugin_test.go
+++ b/pkg/epp/scheduling/framework/plugins/multi/prefix/plugin_test.go
@@ -35,7 +35,7 @@ import (
 
 func TestPrefixPluginCompletion(t *testing.T) {
 	config := Config{
-		HashBlockSize:          4,
+		BlockSize:              4,
 		MaxPrefixBlocksToMatch: DefaultMaxPrefixBlocks,
 		LRUCapacityPerServer:   DefaultLRUCapacityPerServer,
 	}
@@ -201,7 +201,7 @@ func TestPrefixPluginCompletion(t *testing.T) {
 
 func TestPrefixPluginChatCompletions(t *testing.T) {
 	config := Config{
-		HashBlockSize:          4,
+		BlockSize:              4,
 		MaxPrefixBlocksToMatch: DefaultMaxPrefixBlocks,
 		LRUCapacityPerServer:   DefaultLRUCapacityPerServer,
 	}
@@ -235,7 +235,7 @@ func TestPrefixPluginChatCompletions(t *testing.T) {
 
 func TestPrefixPluginChatCompletionsGrowth(t *testing.T) {
 	config := Config{
-		HashBlockSize:          8, // Use larger block size for more predictable JSON marshaling
+		BlockSize:              8, // Use larger block size for more predictable JSON marshaling
 		MaxPrefixBlocksToMatch: DefaultMaxPrefixBlocks,
 		LRUCapacityPerServer:   DefaultLRUCapacityPerServer,
 	}
@@ -349,7 +349,7 @@ func BenchmarkPrefixPluginStress(b *testing.B) {
 	blockSize := 4
 	maxPrefixBlocks := 50000
 	config := Config{
-		HashBlockSize:          blockSize,
+		BlockSize:              blockSize,
 		MaxPrefixBlocksToMatch: maxPrefixBlocks,
 		LRUCapacityPerServer:   DefaultLRUCapacityPerServer,
 	}
@@ -418,7 +418,7 @@ func BenchmarkPrefixPluginChatCompletionsStress(b *testing.B) {
 	blockSize := 8
 	maxPrefixBlocks := 50000
 	config := Config{
-		HashBlockSize:          blockSize,
+		BlockSize:              blockSize,
 		MaxPrefixBlocksToMatch: maxPrefixBlocks,
 		LRUCapacityPerServer:   DefaultLRUCapacityPerServer,
 	}
diff --git a/test/testdata/configloader_1_test.yaml b/test/testdata/configloader_1_test.yaml
index f1f167e..db75a42 100644
--- a/test/testdata/configloader_1_test.yaml
+++ b/test/testdata/configloader_1_test.yaml
@@ -9,7 +9,7 @@ plugins:
   type: test-profile-handler
 - type: test-two
   parameters:
-    hashBlockSize: 32
+    blockSize: 32
 - name: testPicker
   type: test-picker