feat: Affiliate ready state with leader election

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

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

test/e2e/epp/e2e_suite_test.go

@@ -85,8 +85,10 @@ const (
 	xInferObjectiveManifest = "../../../config/crd/bases/inference.networking.x-k8s.io_inferenceobjectives.yaml"
 	// inferPoolManifest is the manifest for the inference pool CRD with 'inference.networking.k8s.io' group.
 	inferPoolManifest = "../../../config/crd/bases/inference.networking.k8s.io_inferencepools.yaml"
-	// inferExtManifest is the manifest for the inference extension test resources.
-	inferExtManifest = "../../testdata/inferencepool-e2e.yaml"
+	// inferExtManifestDefault is the manifest for the default inference extension test resources (single replica).
+	inferExtManifestDefault = "../../testdata/inferencepool-e2e.yaml"
+	// inferExtManifestLeaderElection is the manifest for the inference extension test resources with leader election enabled (3 replicas).
+	inferExtManifestLeaderElection = "../../testdata/inferencepool-leader-election-e2e.yaml"
 	// envoyManifest is the manifest for the envoy proxy test resources.
 	envoyManifest = "../../testdata/envoy.yaml"
 	// metricsRbacManifest is the manifest for the rbac resources for testing metrics.
@@ -95,15 +97,18 @@ const (
 	modelServerManifestFilepathEnvVar = "MANIFEST_PATH"
 )
 
+const e2eLeaderElectionEnabledEnvVar = "E2E_LEADER_ELECTION_ENABLED"
+
 var (
 	ctx = context.Background()
 	cli client.Client
 	// Required for exec'ing in curl pod
-	kubeCli  *kubernetes.Clientset
-	scheme   = runtime.NewScheme()
-	cfg      = config.GetConfigOrDie()
-	nsName   string
-	e2eImage string
+	kubeCli               *kubernetes.Clientset
+	scheme                = runtime.NewScheme()
+	cfg                   = config.GetConfigOrDie()
+	nsName                string
+	e2eImage              string
+	leaderElectionEnabled bool
 )
 
 func TestAPIs(t *testing.T) {
@@ -121,6 +126,11 @@ var _ = ginkgo.BeforeSuite(func() {
 	e2eImage = os.Getenv("E2E_IMAGE")
 	gomega.Expect(e2eImage).NotTo(gomega.BeEmpty(), "E2E_IMAGE environment variable is not set")
 
+	if os.Getenv(e2eLeaderElectionEnabledEnvVar) == "true" {
+		leaderElectionEnabled = true
+		ginkgo.By("Leader election test mode enabled via " + e2eLeaderElectionEnabledEnvVar)
+	}
+
 	ginkgo.By("Setting up the test suite")
 	setupSuite()
 
@@ -146,7 +156,12 @@ func setupInfra() {
 	}
 
 	createCRDs(cli, crds)
-	createInferExt(cli, inferExtManifest)
+
+	inferExtManifestPath := inferExtManifestDefault
+	if leaderElectionEnabled {
+		inferExtManifestPath = inferExtManifestLeaderElection
+	}
+	createInferExt(cli, inferExtManifestPath)
 	createClient(cli, clientManifest)
 	createEnvoy(cli, envoyManifest)
 	createMetricsRbac(cli, metricsRbacManifest)
@@ -156,6 +171,20 @@ func setupInfra() {
 }
 
 var _ = ginkgo.AfterSuite(func() {
+	// If E2E_PAUSE_ON_EXIT is set, pause the test run before cleanup.
+	// This is useful for debugging the state of the cluster after the test has run.
+	if pauseStr := os.Getenv("E2E_PAUSE_ON_EXIT"); pauseStr != "" {
+		ginkgo.By("Pausing before cleanup as requested by E2E_PAUSE_ON_EXIT=" + pauseStr)
+		pauseDuration, err := time.ParseDuration(pauseStr)
+		if err != nil {
+			// If it's not a valid duration (e.g., "true"), just wait indefinitely.
+			ginkgo.By("Invalid duration, pausing indefinitely. Press Ctrl+C to stop the test runner when you are done.")
+			select {} // Block forever
+		}
+		ginkgo.By(fmt.Sprintf("Pausing for %v...", pauseDuration))
+		time.Sleep(pauseDuration)
+	}
+
 	ginkgo.By("Performing global cleanup")
 	cleanupResources()
 })
@@ -423,8 +452,12 @@ func createInferExt(k8sClient client.Client, filePath string) {
 		return k8sClient.Get(ctx, types.NamespacedName{Namespace: nsName, Name: inferExtName}, deploy)
 	}, existsTimeout, interval)
 
-	// Wait for the deployment to be available.
-	testutils.DeploymentAvailable(ctx, k8sClient, deploy, modelReadyTimeout, interval)
+	if leaderElectionEnabled {
+		// With leader election enabled, only 1 replica will be "Ready" at any given time (the leader).
+		testutils.DeploymentReadyReplicas(ctx, k8sClient, deploy, 1, modelReadyTimeout, interval)
+	} else {
+		testutils.DeploymentAvailable(ctx, k8sClient, deploy, modelReadyTimeout, interval)
+	}
 
 	// Wait for the service to exist.
 	testutils.EventuallyExists(ctx, func() error {