Affiliate ready state with leader election with a flag ha-enable-leader-election (#1337)

* feat: Affiliate ready state with leader election

* Fix manifest file for e2e test with leader election after rebase

* Update cmd/epp/runner/health.go

* Add Helm configuration for leader election

* Update config/charts/inferencepool/templates/leader-election-rbac.yaml

* docs: Add leader election section in helm chart README

---------

Co-authored-by: Kellen Swain <[email protected]>

Author: yangligt2

cmd/epp/runner/health.go

@@ -18,6 +18,8 @@ package runner
 
 import (
 	"context"
+	"fmt"
+	"sync/atomic"
 
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	"github.com/go-logr/logr"
@@ -30,37 +32,83 @@ import (
 )
 
 type healthServer struct {
-	logger    logr.Logger
-	datastore datastore.Datastore
+	logger                logr.Logger
+	datastore             datastore.Datastore
+	isLeader              *atomic.Bool
+	leaderElectionEnabled bool
 }
 
+const (
+	LivenessCheckService  = "liveness"
+	ReadinessCheckService = "readiness"
+)
+
 func (s *healthServer) Check(ctx context.Context, in *healthPb.HealthCheckRequest) (*healthPb.HealthCheckResponse, error) {
-	// TODO: we're accepting ANY service name for now as a temporary hack in alignment with
-	// upstream issues. See https://github.com/kubernetes-sigs/gateway-api-inference-extension/pull/788
-	// if in.Service != extProcPb.ExternalProcessor_ServiceDesc.ServiceName {
-	// 	s.logger.V(logutil.DEFAULT).Info("gRPC health check requested unknown service", "available-services", []string{extProcPb.ExternalProcessor_ServiceDesc.ServiceName}, "requested-service", in.Service)
-	// 	return &healthPb.HealthCheckResponse{Status: healthPb.HealthCheckResponse_SERVICE_UNKNOWN}, nil
-	// }
-
-	if !s.datastore.PoolHasSynced() {
-		s.logger.V(logutil.DEFAULT).Info("gRPC health check not serving", "service", in.Service)
+	isLive := s.datastore.PoolHasSynced()
+
+	// If leader election is disabled, use current logic: all checks are based on whether the pool has synced.
+	if !s.leaderElectionEnabled {
+		if !isLive {
+			s.logger.V(logutil.DEFAULT).Info("gRPC health check not serving (leader election disabled)", "service", in.Service)
+			return &healthPb.HealthCheckResponse{Status: healthPb.HealthCheckResponse_NOT_SERVING}, nil
+		}
+		s.logger.V(logutil.TRACE).Info("gRPC health check serving (leader election disabled)", "service", in.Service)
+		return &healthPb.HealthCheckResponse{Status: healthPb.HealthCheckResponse_SERVING}, nil
+	}
+
+	// When leader election is enabled, differentiate between liveness and readiness.
+	// The service name in the request determines which check to perform.
+	var checkName string
+	var isPassing bool
+
+	switch in.Service {
+	case ReadinessCheckService:
+		checkName = "readiness"
+		isPassing = isLive && s.isLeader.Load()
+	case LivenessCheckService, "": // Default to liveness check if service is empty
+		checkName = "liveness"
+		// Any pod that is running and can respond to this gRPC check is considered "live".
+		// The datastore sync status should not affect liveness, only readiness.
+		// This is to prevent the non-leader node from continuous restarts
+		isPassing = true
+	case extProcPb.ExternalProcessor_ServiceDesc.ServiceName:
+		// The main service is considered ready only on the leader.
+		checkName = "ext_proc"
+		isPassing = isLive && s.isLeader.Load()
+	default:
+		s.logger.V(logutil.DEFAULT).Info("gRPC health check requested unknown service", "available-services", []string{LivenessCheckService, ReadinessCheckService, extProcPb.ExternalProcessor_ServiceDesc.ServiceName}, "requested-service", in.Service)
+		return &healthPb.HealthCheckResponse{Status: healthPb.HealthCheckResponse_SERVICE_UNKNOWN}, nil
+	}
+
+	if !isPassing {
+		s.logger.V(logutil.DEFAULT).Info(fmt.Sprintf("gRPC %s check not serving", checkName), "service", in.Service, "isLive", isLive, "isLeader", s.isLeader.Load())
 		return &healthPb.HealthCheckResponse{Status: healthPb.HealthCheckResponse_NOT_SERVING}, nil
 	}
-	s.logger.V(logutil.TRACE).Info("gRPC health check serving", "service", in.Service)
+
+	s.logger.V(logutil.TRACE).Info(fmt.Sprintf("gRPC %s check serving", checkName), "service", in.Service)
 	return &healthPb.HealthCheckResponse{Status: healthPb.HealthCheckResponse_SERVING}, nil
 }
 
 func (s *healthServer) List(ctx context.Context, _ *healthPb.HealthListRequest) (*healthPb.HealthListResponse, error) {
-	// currently only the ext_proc service is provided
-	serviceHealthResponse, err := s.Check(ctx, &healthPb.HealthCheckRequest{Service: extProcPb.ExternalProcessor_ServiceDesc.ServiceName})
-	if err != nil {
-		return nil, err
+	statuses := make(map[string]*healthPb.HealthCheckResponse)
+
+	services := []string{extProcPb.ExternalProcessor_ServiceDesc.ServiceName}
+	if s.leaderElectionEnabled {
+		services = append(services, LivenessCheckService, ReadinessCheckService)
+	}
+
+	for _, service := range services {
+		resp, err := s.Check(ctx, &healthPb.HealthCheckRequest{Service: service})
+		if err != nil {
+			// Check can return an error for unknown services, but here we are iterating known services.
+			// If another error occurs, we should probably return it.
+			return nil, err
+		}
+		statuses[service] = resp
 	}
 
 	return &healthPb.HealthListResponse{
-		Statuses: map[string]*healthPb.HealthCheckResponse{
-			extProcPb.ExternalProcessor_ServiceDesc.ServiceName: serviceHealthResponse,
-		},
+		Statuses: statuses,
 	}, nil
 }
 

cmd/epp/runner/runner.go

@@ -25,6 +25,7 @@ import (
 	"net/http"
 	"net/http/pprof"
 	"os"
+	"sync/atomic"
 
 	"github.com/go-logr/logr"
 	"github.com/prometheus/client_golang/prometheus"
@@ -151,6 +152,10 @@ var (
 	modelServerMetricsPath                    = flag.String("model-server-metrics-path", "/metrics", "Path to scrape metrics from pods")
 	modelServerMetricsScheme                  = flag.String("model-server-metrics-scheme", "http", "Scheme to scrape metrics from pods")
 	modelServerMetricsHttpsInsecureSkipVerify = flag.Bool("model-server-metrics-https-insecure-skip-verify", true, "When using 'https' scheme for 'model-server-metrics-scheme', configure 'InsecureSkipVerify' (default to true)")
+	haEnableLeaderElection                    = flag.Bool(
+		"ha-enable-leader-election",
+		false,
+		"Enables leader election for high availability. When enabled, readiness probes will only pass on the leader.")
 
 	setupLog = ctrl.Log.WithName("setup")
 )
@@ -190,8 +195,9 @@ func bindEnvToFlags() {
 		"POOL_NAME":                                       "pool-name",
 		"POOL_NAMESPACE":                                  "pool-namespace",
 		// durations & bools work too; flag.Set expects the *string* form
-		"REFRESH_METRICS_INTERVAL": "refresh-metrics-interval",
-		"SECURE_SERVING":           "secure-serving",
+		"REFRESH_METRICS_INTERVAL":  "refresh-metrics-interval",
+		"SECURE_SERVING":            "secure-serving",
+		"HA_ENABLE_LEADER_ELECTION": "ha-enable-leader-election",
 	} {
 		if v := os.Getenv(env); v != "" {
 			// ignore error; Parse() will catch invalid values later
@@ -299,12 +305,28 @@ func (r *Runner) Run(ctx context.Context) error {
 		NamespacedName: poolNamespacedName,
 		GroupKind:      poolGroupKind,
 	}
-	mgr, err := runserver.NewDefaultManager(poolGKNN, cfg, metricsServerOptions)
+
+	isLeader := &atomic.Bool{}
+	isLeader.Store(false)
+
+	mgr, err := runserver.NewDefaultManager(poolGKNN, cfg, metricsServerOptions, *haEnableLeaderElection)
 	if err != nil {
 		setupLog.Error(err, "Failed to create controller manager")
 		return err
 	}
 
+	if *haEnableLeaderElection {
+		setupLog.Info("Leader election enabled")
+		go func() {
+			<-mgr.Elected()
+			isLeader.Store(true)
+			setupLog.Info("This instance is now the leader!")
+		}()
+	} else {
+		// If leader election is disabled, all instances are "leaders" for readiness purposes.
+		isLeader.Store(true)
+	}
+
 	if *enablePprof {
 		setupLog.Info("Enabling pprof handlers")
 		err = setupPprofHandlers(mgr)
@@ -356,7 +378,7 @@ func (r *Runner) Run(ctx context.Context) error {
 
 	// --- Add Runnables to Manager ---
 	// Register health server.
-	if err := registerHealthServer(mgr, ctrl.Log.WithName("health"), datastore, *grpcHealthPort); err != nil {
+	if err := registerHealthServer(mgr, ctrl.Log.WithName("health"), datastore, *grpcHealthPort, isLeader, *haEnableLeaderElection); err != nil {
 		return err
 	}
 
@@ -452,11 +474,13 @@ func registerExtProcServer(mgr manager.Manager, runner *runserver.ExtProcServerR
 }
 
 // registerHealthServer adds the Health gRPC server as a Runnable to the given manager.
-func registerHealthServer(mgr manager.Manager, logger logr.Logger, ds datastore.Datastore, port int) error {
+func registerHealthServer(mgr manager.Manager, logger logr.Logger, ds datastore.Datastore, port int, isLeader *atomic.Bool, leaderElectionEnabled bool) error {
 	srv := grpc.NewServer()
 	healthPb.RegisterHealthServer(srv, &healthServer{
-		logger:    logger,
-		datastore: ds,
+		logger:                logger,
+		datastore:             ds,
+		isLeader:              isLeader,
+		leaderElectionEnabled: leaderElectionEnabled,
 	})
 	if err := mgr.Add(
 		runnable.NoLeaderElection(runnable.GRPCServer("health", srv, port))); err != nil {

config/charts/inferencepool/README.md

@@ -79,6 +79,24 @@ $ helm install triton-llama3-8b-instruct \
   oci://us-central1-docker.pkg.dev/k8s-staging-images/gateway-api-inference-extension/charts/inferencepool --version v0
 ```
 
+### Install with High Availability (HA)
+
+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.enableLeaderElection` to `true` and increase the number of replicas in your `values.yaml` file:
+
+```yaml
+inferenceExtension:
+  replicas: 3
+  enableLeaderElection: true
+```
+
+Then apply it with:
+
+```txt
+helm install vllm-llama3-8b-instruct ./config/charts/inferencepool -f values.yaml \
+```
+
 ## Uninstall
 
 Run the following command to uninstall the chart:
@@ -107,6 +125,8 @@ The following table list the configurable parameters of the chart.
 | `inferenceExtension.extraServicePorts`      | List of additional service ports to expose. Defaults to `[]`.                                                         |
 | `inferenceExtension.logVerbosity`           | Logging verbosity level for the endpoint picker. Defaults to `"3"`.                                                   |
 | `provider.name`                             | Name of the Inference Gateway implementation being used. Possible values: `gke`. Defaults to `none`.                   |
+| `inferenceExtension.enableLeaderElection`   | Enable leader election for high availability. When enabled, only one EPP pod (the leader) will be ready to serve traffic. It is recommended to set `inferenceExtension.replicas` to a value greater than 1 when this is set to `true`. Defaults to `false`. |
+
 
 ## Notes
 

config/charts/inferencepool/templates/epp-deployment.yaml

@@ -44,6 +44,9 @@ spec:
         - "--model-server-metrics-path={{ .Values.inferenceExtension.modelServerMetricsPath }}"
         - "--model-server-metrics-scheme={{ .Values.inferenceExtension.modelServerMetricsScheme }}"
         - "--model-server-metrics-https-insecure-skip-verify={{ .Values.inferenceExtension.modelServerMetricsHttpsInsecureSkipVerify }}"
+        {{- if .Values.inferenceExtension.enableLeaderElection }}
+        - "--ha-enable-leader-election"
+        {{- end }}
         {{- if eq (.Values.inferencePool.modelServerType | default "vllm") "triton-tensorrt-llm" }}
         - --total-queued-requests-metric
         - "nv_trt_llm_request_metrics{request_type=waiting}"
@@ -63,15 +66,27 @@ spec:
         {{- toYaml . | nindent 8 }}
         {{- end }}
         livenessProbe:
+          {{- if .Values.inferenceExtension.enableLeaderElection }}
+          grpc:
+            port: 9003
+            service: liveness
+          {{- else }}
           grpc:
             port: 9003
             service: inference-extension
+          {{- end }}
           initialDelaySeconds: 5
           periodSeconds: 10
         readinessProbe:
+          {{- if .Values.inferenceExtension.enableLeaderElection }}
+          grpc:
+            port: 9003
+            service: readiness
+          {{- else }}
           grpc:
             port: 9003
             service: inference-extension
+          {{- end }}
           initialDelaySeconds: 5
           periodSeconds: 10
         {{- with .Values.inferenceExtension.env }}

config/charts/inferencepool/templates/leader-election-rbac.yaml

@@ -0,0 +1,30 @@
+{{- if .Values.inferenceExtension.enableLeaderElection }}
+---
+kind: Role
+apiVersion: rbac.authorization.k8s.io/v1
+metadata:
+  name: {{ include "gateway-api-inference-extension.name" . }}-leader-election
+  namespace: {{ .Release.Namespace }}
+  labels:
+    {{- include "gateway-api-inference-extension.labels" . | nindent 4 }}
+rules:
+- apiGroups: [ "coordination.k8s.io" ]
+  resources: [ "leases" ]
+  verbs: [ "get", "list", "watch", "create", "update", "patch", "delete" ]
+- apiGroups: [ "" ]
+  resources: [ "events" ]
+  verbs: [ "create", "patch" ]
+---
+kind: RoleBinding
+apiVersion: rbac.authorization.k8s.io/v1
+metadata:
+  name: {{ include "gateway-api-inference-extension.name" . }}-leader-election-binding
+  namespace: {{ .Release.Namespace }}
+subjects:
+- kind: ServiceAccount
+  name: {{ include "gateway-api-inference-extension.name" . }}
+roleRef:
+  apiGroup: rbac.authorization.k8s.io
+  kind: Role
+  name: {{ include "gateway-api-inference-extension.name" . }}-leader-election
+{{- end }}

config/charts/inferencepool/values.yaml

@@ -34,6 +34,9 @@ inferenceExtension:
   extraContainerPorts: []
   # Define additional service ports
   extraServicePorts: []
+  # Enable leader election for high availability. When enabled, it is recommended to set replicas > 1.
+  # Only the leader pod will be ready to serve traffic.
+  enableLeaderElection: false
 
 inferencePool:
   targetPortNumber: 8000

pkg/epp/server/controller_manager.go

@@ -85,12 +85,23 @@ func defaultManagerOptions(gknn common.GKNN, metricsServerOptions metricsserver.
 }
 
 // NewDefaultManager creates a new controller manager with default configuration.
-func NewDefaultManager(gknn common.GKNN, restConfig *rest.Config, metricsServerOptions metricsserver.Options) (ctrl.Manager, error) {
+func NewDefaultManager(gknn common.GKNN, restConfig *rest.Config, metricsServerOptions metricsserver.Options, leaderElectionEnabled bool) (ctrl.Manager, error) {
 	opt, err := defaultManagerOptions(gknn, metricsServerOptions)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create controller manager options: %v", err)
 	}
+
+	if leaderElectionEnabled {
+		opt.LeaderElection = true
+		opt.LeaderElectionResourceLock = "leases"
+		// The lease name needs to be unique per EPP deployment.
+		opt.LeaderElectionID = fmt.Sprintf("epp-%s-%s.gateway-api-inference-extension.sigs.k8s.io", gknn.Namespace, gknn.Name)
+		opt.LeaderElectionNamespace = gknn.Namespace
+		opt.LeaderElectionReleaseOnCancel = true
+	}
+
 	manager, err := ctrl.NewManager(restConfig, opt)
+
 	if err != nil {
 		return nil, fmt.Errorf("failed to create controller manager: %v", err)
 	}

test/e2e/epp/README.md

@@ -45,6 +45,19 @@ Follow these steps to run the end-to-end tests:
      export E2E_MANIFEST_PATH=[config/manifests/vllm/gpu-deployment.yaml|config/manifests/vllm/cpu-deployment.yaml]
      ```
 
+   - **Enable leader election tests**: By default, the e2e test runs the EPP server as a single replica.
+     To test the high-availability (HA) mode with leader election (3 replicas), set the following environment variable:
+
+     ```sh
+     export E2E_LEADER_ELECTION_ENABLED=true
+     ```
+
+   - **Pause before cleanup**: To pause the test run before cleaning up resources, set the `E2E_PAUSE_ON_EXIT` environment variable.
+     This is useful for debugging the state of the cluster after the test has run.
+
+     - To pause indefinitely, set it to `true`: `export E2E_PAUSE_ON_EXIT=true`
+     - To pause for a specific duration, provide a duration string: `export E2E_PAUSE_ON_EXIT=10m`
+
 1. **Run the Tests**: Run the `test-e2e` target:
 
    ```sh