feat: initial implementation of Inference Extension

.golangci.yml

@@ -38,7 +38,6 @@ linters-settings:
     # Do not allow non-required aliases.
     no-extra-aliases: false
     alias:
-      # gateway-api
       - pkg: sigs.k8s.io/gateway-api/apis/v1
         alias: gwapiv1
       - pkg: sigs.k8s.io/gateway-api/apis/v1alpha2
@@ -51,7 +50,8 @@ linters-settings:
         alias: egv1a1
       - pkg: github.com/envoyproxy/ai-gateway/api/v1alpha1
         alias: aigv1a1
-      # kubernetes api
+      - pkg: sigs.k8s.io/gateway-api-inference-extension/api/v1alpha2
+        alias: gwaiev1a2
       - pkg: k8s.io/apimachinery/pkg/apis/meta/v1
         alias: metav1
       - pkg: k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1
@@ -64,6 +64,12 @@ linters-settings:
         alias: apierrors
       - pkg: github.com/envoyproxy/ai-gateway/internal/testing
         alias: internaltesting
+      - pkg: github.com/envoyproxy/go-control-plane/envoy/config/cluster/v3
+        alias: clusterv3
+      - pkg: github.com/envoyproxy/go-control-plane/envoy/config/route/v3
+        alias: routev3
+      - pkg: github.com/envoyproxy/gateway/proto/extension
+        alias: egextension
   gci:
     sections:
       # Captures all standard packages if they do not match another section.

api/v1alpha1/api.go

@@ -218,8 +218,30 @@ type AIGatewayRouteRule struct {
 	Matches []AIGatewayRouteRuleMatch `json:"matches,omitempty"`
 }
 
-// AIGatewayRouteRuleBackendRef is a reference to a AIServiceBackend with a weight.
+// AIGatewayRouteRuleBackendRefKind specifies the kind of the backend reference.
+type AIGatewayRouteRuleBackendRefKind string
+
+const (
+	// AIGatewayRouteRuleBackendRefAIServiceBackend is the kind of the AIServiceBackend.
+	AIGatewayRouteRuleBackendRefAIServiceBackend AIGatewayRouteRuleBackendRefKind = "AIServiceBackend"
+	// AIGatewayRouteRuleBackendRefInferencePool is the kind of the InferencePool in the Gateway API Inference Extension.
+	// https://github.com/kubernetes-sigs/gateway-api-inference-extension
+	AIGatewayRouteRuleBackendRefInferencePool AIGatewayRouteRuleBackendRefKind = "InferencePool"
+)
+
+// AIGatewayRouteRuleBackendRef is a reference to a backend with a weight.
 type AIGatewayRouteRuleBackendRef struct {
+	// Kind is the kind of the backend, which is either "AIServiceBackend" or "InferencePool" in Gateway API Inference Extension.
+	//
+	// When this references InferencePool, the selector of the InferencePool is used to select (multiple) AIServiceBackend(s)
+	// that can serve the same model sets that the InferencePool binds.
+	//
+	// Default is AIServiceBackend.
+	//
+	// +kubebuilder:validation:Enum=AIServiceBackend;InferencePool
+	// +kubebuilder:default=AIServiceBackend
+	Kind *AIGatewayRouteRuleBackendRefKind `json:"kind,omitempty"`
+
 	// Name is the name of the AIServiceBackend.
 	//
 	// +kubebuilder:validation:Required

cmd/aigw/testdata/translate_basic.out.yaml

@@ -62,13 +62,10 @@ spec:
         - headers:
             - name: x-ai-eg-selected-backend
               value: envoy-ai-gateway-basic-testupstream.default
-    - backendRefs:
-        - group: gateway.envoyproxy.io
-          kind: Backend
-          name: envoy-ai-gateway-basic-openai
-      matches:
+    - matches:
         - path:
             value: /
+      name: unreachable
 ---
 apiVersion: gateway.envoyproxy.io/v1alpha1
 kind: EnvoyExtensionPolicy

cmd/aigw/translate.go

@@ -183,7 +183,7 @@ func translateCustomResourceObjects(
 		WithStatusSubresource(&aigv1a1.AIGatewayRoute{}).
 		WithStatusSubresource(&aigv1a1.AIServiceBackend{}).
 		WithStatusSubresource(&aigv1a1.BackendSecurityPolicy{})
-	_ = controller.ApplyIndexing(ctx, func(_ context.Context, obj client.Object, field string, extractValue client.IndexerFunc) error {
+	_ = controller.ApplyIndexing(ctx, true, func(_ context.Context, obj client.Object, field string, extractValue client.IndexerFunc) error {
 		builder = builder.WithIndex(obj, field, extractValue)
 		return nil
 	}) // Error should never happen.

cmd/controller/main.go

@@ -12,7 +12,7 @@ import (
 	"net"
 	"os"
 
-	"github.com/envoyproxy/gateway/proto/extension"
+	egextension "github.com/envoyproxy/gateway/proto/extension"
 	"go.uber.org/zap/zapcore"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/health/grpc_health_v1"
@@ -31,6 +31,7 @@ func parseAndValidateFlags(args []string) (
 	enableLeaderElection bool,
 	logLevel zapcore.Level,
 	extensionServerPort string,
+	enableInfExt bool,
 	err error,
 ) {
 	fs := flag.NewFlagSet("AI Gateway Controller", flag.ContinueOnError)
@@ -60,6 +61,11 @@ func parseAndValidateFlags(args []string) (
 		":1063",
 		"gRPC port for the extension server",
 	)
+	enableInfExtPtr := fs.Bool(
+		"enableInferenceExtension",
+		false,
+		"Enable the Gateway Inference Extetension. When enabling this, the CRDs for the InferenceModel and InferencePool must be installed prior to starting the controller.",
+	)
 
 	if err = fs.Parse(args); err != nil {
 		err = fmt.Errorf("failed to parse flags: %w", err)
@@ -77,7 +83,7 @@ func parseAndValidateFlags(args []string) (
 		err = fmt.Errorf("invalid log level: %q", *logLevelPtr)
 		return
 	}
-	return *extProcLogLevelPtr, *extProcImagePtr, *enableLeaderElectionPtr, zapLogLevel, *extensionServerPortPtr, nil
+	return *extProcLogLevelPtr, *extProcImagePtr, *enableLeaderElectionPtr, zapLogLevel, *extensionServerPortPtr, *enableInfExtPtr, nil
 }
 
 func main() {
@@ -88,6 +94,7 @@ func main() {
 		flagEnableLeaderElection,
 		zapLogLevel,
 		flagExtensionServerPort,
+		enableInfExt,
 		err := parseAndValidateFlags(os.Args[1:])
 	if err != nil {
 		setupLog.Error(err, "failed to parse and validate flags")
@@ -110,8 +117,8 @@ func main() {
 
 	// Start the extension server running alongside the controller.
 	s := grpc.NewServer()
-	extSrv := extensionserver.New(setupLog)
-	extension.RegisterEnvoyGatewayExtensionServer(s, extSrv)
+	extSrv := extensionserver.New(ctrl.Log)
+	egextension.RegisterEnvoyGatewayExtensionServer(s, extSrv)
 	grpc_health_v1.RegisterHealthServer(s, extSrv)
 	go func() {
 		<-ctx.Done()
@@ -128,6 +135,7 @@ func main() {
 		ExtProcImage:         flagExtProcImage,
 		ExtProcLogLevel:      flagExtProcLogLevel,
 		EnableLeaderElection: flagEnableLeaderElection,
+		EnableInfExt:         enableInfExt,
 	}); err != nil {
 		setupLog.Error(err, "failed to start controller")
 	}

cmd/controller/main_test.go

@@ -13,12 +13,13 @@ import (
 
 func Test_parseAndValidateFlags(t *testing.T) {
 	t.Run("no flags", func(t *testing.T) {
-		extProcLogLevel, extProcImage, enableLeaderElection, logLevel, extensionServerPort, err := parseAndValidateFlags([]string{})
+		extProcLogLevel, extProcImage, enableLeaderElection, logLevel, extensionServerPort, enableInfExt, err := parseAndValidateFlags([]string{})
 		require.Equal(t, "info", extProcLogLevel)
 		require.Equal(t, "docker.io/envoyproxy/ai-gateway-extproc:latest", extProcImage)
 		require.True(t, enableLeaderElection)
 		require.Equal(t, "info", logLevel.String())
 		require.Equal(t, ":1063", extensionServerPort)
+		require.False(t, enableInfExt)
 		require.NoError(t, err)
 	})
 	t.Run("all flags", func(t *testing.T) {
@@ -36,13 +37,15 @@ func Test_parseAndValidateFlags(t *testing.T) {
 					tc.dash + "enableLeaderElection=false",
 					tc.dash + "logLevel=debug",
 					tc.dash + "port=:8080",
+					tc.dash + "enableInferenceExtension=true",
 				}
-				extProcLogLevel, extProcImage, enableLeaderElection, logLevel, extensionServerPort, err := parseAndValidateFlags(args)
+				extProcLogLevel, extProcImage, enableLeaderElection, logLevel, extensionServerPort, enableInfExt, err := parseAndValidateFlags(args)
 				require.Equal(t, "debug", extProcLogLevel)
 				require.Equal(t, "example.com/extproc:latest", extProcImage)
 				require.False(t, enableLeaderElection)
 				require.Equal(t, "debug", logLevel.String())
 				require.Equal(t, ":8080", extensionServerPort)
+				require.True(t, enableInfExt)
 				require.NoError(t, err)
 			})
 		}
@@ -66,7 +69,7 @@ func Test_parseAndValidateFlags(t *testing.T) {
 			},
 		} {
 			t.Run(tc.name, func(t *testing.T) {
-				_, _, _, _, _, err := parseAndValidateFlags(tc.flags)
+				_, _, _, _, _, _, err := parseAndValidateFlags(tc.flags)
 				require.ErrorContains(t, err, tc.expErr)
 			})
 		}

cmd/extproc/mainlib/main.go

@@ -188,7 +188,7 @@ func startMetricsServer(addr string, logger *slog.Logger) (*http.Server, metric.
 	}
 
 	go func() {
-		logger.Info("Starting metrics server", "address", addr)
+		logger.Info("starting metrics server", "address", addr)
 		if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
 			logger.Error("Metrics server failed", "error", err)
 		}

examples/inference_extension/README.md

@@ -0,0 +1,2 @@
+This example demonstrates how to use the [Inference Extension API](https://gateway-api-inference-extension.sigs.k8s.io/)in Envoy AI Gateway project.
+The feature can be used only when `--enableInferenceExtension` is set to `true` passed to the Envoy AI Gateway controller. See the helm values.yaml file for more details.

examples/inference_extension/inference_extension.yaml

@@ -0,0 +1,133 @@
+# Copyright Envoy AI Gateway Authors
+# SPDX-License-Identifier: Apache-2.0
+# The full text of the Apache license is available in the LICENSE file at
+# the root of the repo.
+
+apiVersion: gateway.networking.k8s.io/v1
+kind: GatewayClass
+metadata:
+  name: inference-extension-example
+spec:
+  controllerName: gateway.envoyproxy.io/gatewayclass-controller
+---
+apiVersion: gateway.networking.k8s.io/v1
+kind: Gateway
+metadata:
+  name: inference-extension-example
+  namespace: default
+spec:
+  gatewayClassName: inference-extension-example
+  listeners:
+    - name: http
+      protocol: HTTP
+      port: 80
+---
+apiVersion: aigateway.envoyproxy.io/v1alpha1
+kind: AIGatewayRoute
+metadata:
+  name: inference-extension-example
+  namespace: default
+spec:
+  schema:
+    name: OpenAI
+  targetRefs:
+    - name: inference-extension-example
+      kind: Gateway
+      group: gateway.networking.k8s.io
+  rules:
+    - matches:
+        - headers:
+            - type: Exact
+              name: x-target-inference-extension
+              value: "yes"
+      backendRefs:
+        - name: inference-extension-example-pool  # The name of the InferencePool that binds to the backend.
+          # Explicitly specify the kind of the backend to be InferenceExtension.
+          kind: InferencePool
+---
+apiVersion: inference.networking.x-k8s.io/v1alpha2
+kind: InferencePool
+metadata:
+  name: inference-extension-example-pool
+spec:
+  targetPortNumber: 8080
+  selector:
+    # Select multiple AIServiceBackend objects to bind to the InferencePool.
+    app: my-backend
+  extensionRef:
+    # Specify the static name "envoy-ai-gateway" to bind the InferencePool to the Envoy AI Gateway.
+    # This indicates that the InferencePool will be managed by the Envoy AI Gateway.
+    name: envoy-ai-gateway
+---
+apiVersion: inference.networking.x-k8s.io/v1alpha2
+kind: InferenceModel
+metadata:
+  name: inference-extension-example
+spec:
+  modelName: mistral:latest
+  criticality: Critical
+  poolRef:
+    # Bind the InferenceModel to the InferencePool.
+    name: inference-extension-example-pool
+---
+apiVersion: aigateway.envoyproxy.io/v1alpha1
+kind: AIServiceBackend
+metadata:
+  name: inference-extension-example-testupstream
+  namespace: default
+  labels:
+    # Indicate the backend is selected by the InferencePool.
+    app: my-backend
+spec:
+  schema:
+    name: OpenAI
+  backendRef:
+    name: inference-extension-example-testupstream
+    kind: Service
+    port: 8080
+---
+apiVersion: v1
+kind: Service
+metadata:
+  name: inference-extension-example-testupstream
+  namespace: default
+spec:
+  selector:
+    app: inference-extension-example-testupstream
+  ports:
+    - protocol: TCP
+      port: 8080
+      targetPort: 8080
+  # The headless service allows the IP addresses of the pods to be resolved via the Service DNS.
+  clusterIP: None
+---
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: inference-extension-example-testupstream
+  namespace: default
+spec:
+  replicas: 3
+  selector:
+    matchLabels:
+      app: inference-extension-example-testupstream
+  template:
+    metadata:
+      labels:
+        app: inference-extension-example-testupstream
+    spec:
+      containers:
+        - name: testupstream
+          image: docker.io/envoyproxy/ai-gateway-testupstream:latest
+          imagePullPolicy: IfNotPresent
+          ports:
+            - containerPort: 8080
+          env:
+            - name: TESTUPSTREAM_ID
+              value: test
+          readinessProbe:
+            httpGet:
+              path: /health
+              port: 8080
+            initialDelaySeconds: 1
+            periodSeconds: 1