Add support for custom inference engine port (#1140)

Signed-off-by: Varun Gupta <[email protected]>

Author: varungup90

development/app/config/templates/deployment/deployment.yaml

@@ -19,6 +19,7 @@ spec:
       labels:
         adapter.model.aibrix.ai/enabled: "true"
         model.aibrix.ai/name: "llama2-7b"
+        model.aibrix.ai/port: "8000"
         app: "mock-llama2-7b"
     spec:
       serviceAccountName: mocked-app-sa

development/vllm/macos/deployment.yaml

@@ -19,6 +19,7 @@ spec:
       labels:
         adapter.model.aibrix.ai/enabled: "true"
         model.aibrix.ai/name: "facebook-opt-125m"
+        model.aibrix.ai/port: "8000"
         app: "mock-facebook-opt-125m"
     spec:
       serviceAccountName: mocked-app-sa

pkg/types/router_context.go

@@ -23,11 +23,10 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/vllm-project/aibrix/pkg/utils"
 	v1 "k8s.io/api/core/v1"
 )
 
-const podMetricPort = "8000"
-
 var nilPod = &v1.Pod{}
 
 // RoutingAlgorithm defines the routing algorithms
@@ -93,7 +92,7 @@ func (r *RoutingContext) HasRouted() bool {
 }
 
 func (r *RoutingContext) targetAddress(pod *v1.Pod) string {
-	return fmt.Sprintf("%v:%v", pod.Status.PodIP, podMetricPort)
+	return fmt.Sprintf("%v:%v", pod.Status.PodIP, utils.GetModelPortForPod(r.RequestID, pod))
 }
 
 func (r *RoutingContext) reset(ctx context.Context, algorithms RoutingAlgorithm, model string, message string, requestID string, user *string) {

pkg/utils/pod.go

@@ -19,6 +19,7 @@ package utils
 import (
 	"context"
 	"fmt"
+	"strconv"
 	"strings"
 
 	"k8s.io/klog/v2"
@@ -30,7 +31,9 @@ import (
 )
 
 const (
-	NAMESPACE = "aibrix-system"
+	NAMESPACE            = "aibrix-system"
+	modelPortIdentifier  = "model.aibrix.ai/port"
+	defaultPodMetricPort = 8000
 )
 
 // GeneratePodKey generates a key in the format "namespace/name" for a given pod.
@@ -238,3 +241,20 @@ func SelectRandomPod(pods []*v1.Pod, randomFn func(int) int) (*v1.Pod, error) {
 	randomPod := readyPods[randomFn(len(readyPods))]
 	return randomPod, nil
 }
+
+func GetModelPortForPod(requestID string, pod *v1.Pod) int64 {
+	value, ok := pod.Labels[modelPortIdentifier]
+	if !ok {
+		klog.Warningf("requestID: %v, pod: %v is missing port identifier label: %v, hence default to port: %v",
+			requestID, pod.Name, modelPortIdentifier, defaultPodMetricPort)
+		return defaultPodMetricPort
+	}
+
+	modelPort, err := strconv.ParseInt(value, 10, 32)
+	if err != nil {
+		klog.Warningf("requestID: %v, pod: %v has incorrect value: %v for port identifier label: %v, hence default to port: %v",
+			requestID, pod.Name, value, modelPortIdentifier, defaultPodMetricPort)
+		modelPort = defaultPodMetricPort
+	}
+	return modelPort
+}

pkg/utils/pod_test.go

@@ -18,10 +18,12 @@ package utils
 import (
 	"fmt"
 	"math/rand"
+	"testing"
 	"time"
 
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
+	"github.com/stretchr/testify/assert"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
@@ -81,3 +83,59 @@ var _ = Describe("Pod", func() {
 		Expect(modified).To(Equal(expected))
 	})
 })
+
+func TestModePortForPod(t *testing.T) {
+	testcases := []struct {
+		message      string
+		pod          *v1.Pod
+		expectedPort int64
+	}{
+		{
+			message: "read port from pod labels",
+			pod: &v1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "p1",
+					Labels: map[string]string{
+						"model.aibrix.ai/port": "9000",
+					},
+				},
+			},
+			expectedPort: 9000,
+		},
+		{
+			message: "incorrect model port label value",
+			pod: &v1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "p1",
+					Labels: map[string]string{
+						"model.aibrix.ai/port": "port",
+					},
+				},
+			},
+			expectedPort: 8000,
+		},
+		{
+			message: "return default port if not configured in pod labels",
+			pod: &v1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:   "p1",
+					Labels: map[string]string{},
+				},
+			},
+			expectedPort: 8000,
+		},
+		{
+			message: "return default port if no label is present",
+			pod: &v1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "p1",
+				},
+			},
+			expectedPort: 8000,
+		},
+	}
+
+	for _, tt := range testcases {
+		assert.Equal(t, tt.expectedPort, GetModelPortForPod("1", tt.pod), tt.message)
+	}
+}