refactor(sidecar): encapsulate code better in order to share protocol implementation between different connectors

.gitignore

@@ -39,3 +39,4 @@ vendor
 
 # Build output
 /build
+

cmd/pd-sidecar/main.go

@@ -21,30 +21,21 @@ import (
 	"net/url"
 	"os"
 	"strconv"
-	"strings"
 
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/controller-runtime/pkg/log/zap"
 
 	"github.com/llm-d/llm-d-inference-scheduler/pkg/sidecar/proxy"
+	"github.com/llm-d/llm-d-inference-scheduler/pkg/sidecar/proxy/runners/types"
 	"github.com/llm-d/llm-d-inference-scheduler/pkg/sidecar/version"
 )
 
-var (
-	// supportedConnectors defines all valid P/D connector types
-	supportedConnectors = []string{
-		proxy.ConnectorNIXLV2,
-		proxy.ConnectorSharedStorage,
-		proxy.ConnectorSGLang,
-	}
-)
-
 func main() {
 	port := flag.String("port", "8000", "the port the sidecar is listening on")
 	vLLMPort := flag.String("vllm-port", "8001", "the port vLLM is listening on")
 	vLLMDataParallelSize := flag.Int("data-parallel-size", 1, "the vLLM DATA-PARALLEL-SIZE value")
-	connector := flag.String("connector", proxy.ConnectorNIXLV2, "the P/D connector being used. Supported: "+strings.Join(supportedConnectors, ", "))
+	connectorStr := flag.String("connector", types.ConnectorNIXLV2.String(), "the P/D connector being used. Supported: "+types.AllConnectorStrings())
 	prefillerUseTLS := flag.Bool("prefiller-use-tls", false, "whether to use TLS when sending requests to prefillers")
 	decoderUseTLS := flag.Bool("decoder-use-tls", false, "whether to use TLS when sending requests to the decoder")
 	prefillerInsecureSkipVerify := flag.Bool("prefiller-tls-insecure-skip-verify", false, "configures the proxy to skip TLS verification for requests to prefiller")
@@ -72,16 +63,9 @@ func main() {
 
 	logger.Info("Proxy starting", "Built on", version.BuildRef, "From Git SHA", version.CommitSHA)
 
-	// Validate connector
-	isValidConnector := false
-	for _, validConnector := range supportedConnectors {
-		if *connector == validConnector {
-			isValidConnector = true
-			break
-		}
-	}
-	if !isValidConnector {
-		logger.Info("Error: --connector must be one of: " + strings.Join(supportedConnectors, ", "))
+	connector := types.Connector(*connectorStr)
+	if err := connector.Validate(); err != nil {
+		logger.Error(err, "invalid connector")
 		return
 	}
 	logger.Info("p/d connector validated", "connector", connector)
@@ -127,7 +111,7 @@ func main() {
 	}
 
 	config := proxy.Config{
-		Connector:                   *connector,
+		Connector:                   connector,
 		PrefillerUseTLS:             *prefillerUseTLS,
 		PrefillerInsecureSkipVerify: *prefillerInsecureSkipVerify,
 		DecoderInsecureSkipVerify:   *decoderInsecureSkipVerify,

pkg/sidecar/proxy/allowlist.go

@@ -309,15 +309,15 @@ func (av *AllowlistValidator) createPodInformer(poolName string, selector labels
 				Group:    "",
 				Version:  "v1",
 				Resource: "pods",
-			}).Namespace(av.namespace).List(context.TODO(), options)
+			}).Namespace(av.namespace).List(context.TODO(), options) // FIXME: use stored context from Start() to respect shutdown
 		},
 		WatchFunc: func(options metav1.ListOptions) (watch.Interface, error) {
 			options.LabelSelector = selector.String()
 			return av.dynamicClient.Resource(schema.GroupVersionResource{
 				Group:    "",
 				Version:  "v1",
 				Resource: "pods",
-			}).Namespace(av.namespace).Watch(context.TODO(), options)
+			}).Namespace(av.namespace).Watch(context.TODO(), options) // FIXME: use stored context from Start() to respect shutdown
 		},
 	}
 
@@ -363,7 +363,12 @@ func (av *AllowlistValidator) onPodDelete(obj interface{}) {
 	av.rebuildAllowlist()
 }
 
-// rebuildAllowlist rebuilds the entire allowlist from current pod state
+// rebuildAllowlist rebuilds the entire allowlist from current pod state.
+// NOTE: This method acquires allowedTargetsMu then podInformersMu (RLock).
+// There is a brief window where allowedTargets may be inconsistent with
+// podInformers if a pod informer is being added/removed concurrently.
+// This is acceptable because the allowlist is rebuilt on every pod event
+// and will converge quickly.
 func (av *AllowlistValidator) rebuildAllowlist() {
 	av.allowedTargetsMu.Lock()
 	defer av.allowedTargetsMu.Unlock()

pkg/sidecar/proxy/chat_completions.go

@@ -58,7 +58,7 @@ func (s *Server) chatCompletionsHandler(w http.ResponseWriter, r *http.Request)
 		s.logger.V(4).Info("skip disaggregated prefill")
 
 		if !s.forwardDataParallel || !s.dataParallelHandler(w, r) {
-			s.decoderProxy.ServeHTTP(w, r)
+			s.pdProxyManager.decoderProxy.ServeHTTP(w, r)
 		}
 		return
 	}
@@ -75,5 +75,5 @@ func (s *Server) chatCompletionsHandler(w http.ResponseWriter, r *http.Request)
 	}
 
 	s.logger.V(4).Info("SSRF protection: prefill target allowed", "target", prefillHostPort)
-	s.runConnectorProtocol(w, r, prefillHostPort)
+	s.protocolRunner.Run(w, r, prefillHostPort, s.logger)
 }

pkg/sidecar/proxy/chat_completions_test.go

@@ -22,7 +22,9 @@ import (
 	"net/http/httptest"
 	"testing"
 
+	"github.com/go-logr/logr"
 	"github.com/llm-d/llm-d-inference-scheduler/pkg/common"
+	"github.com/llm-d/llm-d-inference-scheduler/pkg/sidecar/proxy/runners/types/mock"
 )
 
 func TestServer_chatCompletionsHandler(t *testing.T) {
@@ -119,12 +121,15 @@ func TestServer_chatCompletionsHandler(t *testing.T) {
 				s.prefillSamplerFn = func(n int) int { return i % n }
 				// verify the hostPort value
 				var hostPort string
-				s.runConnectorProtocol = func(_ http.ResponseWriter, _ *http.Request, selectedHostPort string) { hostPort = selectedHostPort }
+				s.protocolRunner = &mock.ProtocolRunnerMock{
+					RunFunc: func(_ http.ResponseWriter, _ *http.Request, prefillPodHost string, _ logr.Logger) {
+						hostPort = prefillPodHost
+					},
+				}
 				var passthrough bool
-				s.decoderProxy = http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {
+				s.pdProxyManager.decoderProxy = http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {
 					passthrough = true
 				})
-				s.dataParallelProxies = make(map[string]http.Handler)
 				recorder := httptest.NewRecorder()
 				recorder.Code = 0
 				s.chatCompletionsHandler(recorder, tt.r)