Pluggable BBR framework (WiP)

Author: davidbreitgand

config/charts/body-based-routing/templates/_helpers.tpl

@@ -28,9 +28,6 @@ spec:
         - containerPort: {{ .Values.bbr.port }}
         # health check
         - containerPort: {{ .Values.bbr.healthCheckPort }}
-        envFrom:
-        - configMapRef:
-            name: {{ include "body-based-routing.fullname" . }}-plugins-config
 ---
 apiVersion: v1
 kind: Service

config/charts/body-based-routing/templates/bbr.yaml

@@ -19,18 +19,6 @@ provider:
 inferenceGateway:
   name: inference-gateway
 
-plugins:
-  requestChain:
-#    - type: ModelSelector
-#      name: semantic-model-selector
-    - type: MetadataExtractor 
-      name: simple-model-extractor
-
-  responseChain:
-#    - type: GuardRail
-#      name: pid-disclosure-blocker
-
-  metaDataKeys:
-    - model
-#  - temperature
-#  - max_tokens
+env:
+  - name: REQUEST_PLUGINS_CHAIN
+    value: "simple-model-extractor"

config/charts/body-based-routing/templates/configmap.yaml

@@ -1,5 +1,5 @@
 # Body-Based Routing
-This package provides an extension that can be deployed to write the `model`
+By deafult this package provides a plugable extension that can be to set the `model`
 HTTP body parameter as a header (`X-Gateway-Model-Name`) so as to enable routing capabilities on the
 model name.
 

config/charts/body-based-routing/values.yaml

@@ -20,6 +20,13 @@ import (
 	bbrplugins "sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/plugins"
 )
 
+const (
+	RequestPluginChain   = "REQUEST_PLUGINS_CHAIN"
+	ResponsePluginChain  = "RESPONSE_PLUGINS_CHAIN"
+	MaxPluginChainLength = 10  //arbitrary number that can be changed **later**; it's here to short-circuit long inputs
+	MaxNumberOfPlugins   = 100 //arbitrary number that can be changed **later**; can be used to optimize memory allocation for the registry **later**
+)
+
 // placeholder for Plugin constructors
 type PluginFactoryFunc func() bbrplugins.BBRPlugin //any no-argument function that returns bbrplugins.BBRPlugin can be assigned to this type (including a constructor function)
 
@@ -51,11 +58,3 @@ type PluginsChain interface {
 	Run(bodyBytes []byte, metaDataKeys []string, registry PluginRegistry) (map[string]string, []byte, error) //return potentially mutated body and all headers map safely merged
 	String() string
 }
-
-// ------------------------- Supported plugin interfaces -----------------------------------------
-// Extend this map when new interfaces and/or implementations are added
-var SupportedInterfaces = map[string][]string{
-	"ModelSelector":     {"semantic-model-selector"},                     //future
-	"GuardRail":         {"bad-words-blocker", "pid-disclosure-blocker"}, //future
-	"MetadataExtractor": {"simple-model-extractor"},
-}

pkg/bbr/README.md

@@ -18,7 +18,6 @@ package framework
 
 import (
 	"fmt"
-	"slices"
 
 	bbrplugins "sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/plugins"
 )
@@ -44,10 +43,10 @@ func NewPluginRegistry() PluginRegistry {
 
 // Register a plugin factory by type key (e.g., "ModelSelector", "MetadataExtractor")
 func (r *pluginRegistry) RegisterFactory(typeKey string, factory PluginFactoryFunc) error {
-	//validate whether this interface is supported
-	_, ok := SupportedInterfaces[typeKey]
-	if !ok {
-		err := fmt.Errorf("non-supported plugin interface type %s", typeKey)
+	//validate whether already registered
+	alreadyRegistered := r.ContainsFactory(typeKey)
+	if alreadyRegistered {
+		err := fmt.Errorf("factory fot plugin interface type %s is already registered", typeKey)
 		return err
 	}
 	r.pluginsFactory[typeKey] = factory
@@ -58,16 +57,13 @@ func (r *pluginRegistry) RegisterFactory(typeKey string, factory PluginFactoryFu
 // Register a plugin instance (created through the appropriate factory)
 func (r *pluginRegistry) RegisterPlugin(plugin bbrplugins.BBRPlugin) error {
 	//validate whether this interface is supported
-	implementations, ok := SupportedInterfaces[plugin.TypedName().Type]
-	if !ok {
-		err := fmt.Errorf("non-supported plugin interface type %s", plugin.TypedName().Type)
-		return err
-	}
-	//validate whether this implementation for this plugin interface is supported
-	if len(implementations) == 0 || !slices.Contains(implementations, plugin.TypedName().Name) {
-		err := fmt.Errorf("no implementation for plugin interface type %s", plugin.TypedName().Type)
+	alreadyRegistered := r.ContainsPlugin(plugin.TypedName().Type)
+
+	if alreadyRegistered {
+		err := fmt.Errorf("plugin implementing interface type %s is already registered", plugin.TypedName().Type)
 		return err
 	}
+
 	// validate that the factory for this plugin is registered: always register factory before the plugin
 	if _, ok := r.pluginsFactory[plugin.TypedName().Type]; !ok {
 		err := fmt.Errorf("no plugin factory registered for plugin interface type %s", plugin.TypedName().Type)

pkg/bbr/framework/interfaces.go

@@ -30,13 +30,19 @@ import (
 	crmetrics "sigs.k8s.io/controller-runtime/pkg/metrics"
 
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/metrics"
+	utils "sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/utils"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
 func TestHandleRequestBody(t *testing.T) {
 	metrics.Register()
 	ctx := logutil.NewTestLoggerIntoContext(context.Background())
 
+	//set environment variables expected by the code under test
+	//testing a minimal configuration
+	//request plugin chain always contains the default bbr plugin that extracts a model name and sets it on the X-Gateway-Model-Name header
+	t.Setenv("REQUEST_PLUGINS_CHAIN", "simple_model_extractor")
+
 	tests := []struct {
 		name      string
 		body      map[string]any
@@ -176,9 +182,15 @@ func TestHandleRequestBody(t *testing.T) {
 		},
 	}
 
+	//Initialize PluginRegistry and request/response PluginsChain instances based on the minimal configuration setting vi env vars
+	registry, requestChain, responseChain, metaDataKeys, err := utils.InitPlugins()
+	if err != nil {
+		t.Fatalf("processRequestBody(): %v", err)
+	}
+
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			server := &Server{streaming: test.streaming}
+			server := NewServer(test.streaming, registry, requestChain, responseChain, metaDataKeys)
 			bodyBytes, _ := json.Marshal(test.body)
 			resp, err := server.HandleRequestBody(ctx, bodyBytes)
 			if err != nil {

pkg/bbr/framework/registry.go

@@ -27,13 +27,18 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
 	bbrplugins "sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/plugins"
-	bbrutils "sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/utils"
+	utils "sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/utils"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
 func TestProcessRequestBody(t *testing.T) {
 	ctx := logutil.NewTestLoggerIntoContext(context.Background())
 
+	//set environment variables expected by the code under test
+	//testing a minimal configuration
+	//request plugin chain always contains the default bbr plugin that extracts a model name and sets it on the X-Gateway-Model-Name header
+	t.Setenv("REQUEST_PLUGINS_CHAIN", "simple_model_extractor")
+
 	cases := []struct {
 		desc      string
 		streaming bool
@@ -127,14 +132,15 @@ func TestProcessRequestBody(t *testing.T) {
 		},
 	}
 
-	//Initialize PluginRegistry and request/response PluginsChain instances
-	//Change to hermetic
-	registry, requestChain, responseChain, metaDataKeys, _ := bbrutils.InitPlugins()
+	//Initialize PluginRegistry and request/response PluginsChain instances based on the minimal configuration setting vi env vars
+	registry, requestChain, responseChain, metaDataKeys, err := utils.InitPlugins()
+	if err != nil {
+		t.Fatalf("processRequestBody(): %v", err)
+	}
 
 	for _, tc := range cases {
 		t.Run(tc.desc, func(t *testing.T) {
 			srv := NewServer(tc.streaming, registry, requestChain, responseChain, metaDataKeys)
-			//srv := NewServer(tc.streaming)
 			streamedBody := &streamedBody{}
 			for i, body := range tc.bodys {
 				got, err := srv.processRequestBody(context.Background(), body, streamedBody, log.FromContext(ctx))

pkg/bbr/handlers/request_test.go

@@ -20,13 +20,21 @@ import (
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/plugins"
 )
 
-// ---------------------------- Well-known custom headers ---------------------------------
-const ModelHeader = "X-Gateway-Model-Name"
-
 // ----------------------------------------------------------------------------------------
 // ------------------------------------ Defaults ------------------------------------------
-const DefaultPluginType = "MetadataExtractor"
-const DefaultPluginImplementation = "simple-model-selector"
+const (
+	//The deafult plugin will always be configured for request plugins chain
+	DefaultPluginType = "MetadataExtractor"
+	//Even though BBRPlugin is not a K8s resource yet, let's make its naming compliant with K8s resource naming
+	// Allows: lowercase letters, digits, hyphens, dots.
+	// Must start and end with a lowercase alphanumeric character.
+	// Middle characters group can contain lowercase alphanumerics, hyphens, and dots
+	// Middle and rightmost groups are optional
+	PluginNamePattern   = `^[a-z0-9]([-a-z0-9.]*[a-z0-9])?$`
+	MaxPluginNameLength = 253
+	//Well-known custom header set to a model name
+	ModelHeader = "X-Gateway-Model-Name"
+)
 
 // BBRPlugin defines the interface for plugins in the BBR framework should never mutate the body directly.
 type BBRPlugin interface {