Repository: vllm-project/aibrix
Pull request #1140: Add support for custom inference engine port
Author: login: varungup90
Created: 2025-05-27T23:22:43Z | Merged: 2025-05-30T03:27:01Z

## Pull Request Description
[Please provide a clear and concise description of your changes here]

## Related Issues
Resolves: #1111 

**Important: Before submitting, please complete the description above and review the checklist below.**

---

<details>
<summary><strong>Contribution Guidelines (Expand for Details)</strong></summary>

<p>We appreciate your contribution to aibrix! To ensure a smooth review process and maintain high code quality, please adhere to the following guidelines:</p>

<h3>Pull Request Title Format</h3>
<p>Your PR title should start with one of these prefixes to indicate the nature of the change:</p>
<ul>
    <li><code>[Bug]</code>: Corrections to existing functionality</li>
    <li><code>[CI]</code>: Changes to build process or CI pipeline</li>
    <li><code>[Docs]</code>: Updates or additions to documentation</li>
    <li><code>[API]</code>: Modifications to aibrix's API or interface</li>
    <li><code>[CLI]</code>: Changes or additions to the Command Line Interface</li>
    <li><code>[Misc]</code>: For changes not covered above (use sparingly)</li>
</ul>
<p><em>Note: For changes spanning multiple categories, use multiple prefixes in order of importance.</em></p>

<h3>Submission Checklist</h3>
<ul>
    <li>[ ] PR title includes appropriate prefix(es)</li>
    <li>[ ] Changes are clearly explained in the PR description</li>
    <li>[ ] New and existing tests pass successfully</li>
    <li>[ ] Code adheres to project style and best practices</li>
    <li>[ ] Documentation updated to reflect changes (if applicable)</li>
    <li>[ ] Thorough testing completed, no regressions introduced</li>
</ul>

<p>By submitting this PR, you confirm that you've read these guidelines and your changes align with the project's contribution standards.</p>

</details>


Diff:
diff --git a/development/app/config/templates/deployment/deployment.yaml b/development/app/config/templates/deployment/deployment.yaml
index 7431faef3..37a697d27 100644
--- a/development/app/config/templates/deployment/deployment.yaml
+++ b/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
diff --git a/development/vllm/macos/deployment.yaml b/development/vllm/macos/deployment.yaml
index 511db8121..960492254 100644
--- a/development/vllm/macos/deployment.yaml
+++ b/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
diff --git a/pkg/types/router_context.go b/pkg/types/router_context.go
index 24643397f..0d9b0dbf2 100644
--- a/pkg/types/router_context.go
+++ b/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) {
diff --git a/pkg/utils/pod.go b/pkg/utils/pod.go
index a5e4d6117..6e699b475 100644
--- a/pkg/utils/pod.go
+++ b/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
+}
diff --git a/pkg/utils/pod_test.go b/pkg/utils/pod_test.go
index 587fa9fcd..b5fd8b84a 100644
--- a/pkg/utils/pod_test.go
+++ b/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)
+	}
+}


Reviews:
user: login: googs1025
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-28T01:18:50Z
user: login: varungup90
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-28T03:30:55Z
user: login: Jeffwan
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-28T22:27:14Z
user: login: varungup90
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-29T17:46:06Z
user: login: varungup90
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-29T23:20:15Z

Review comments:
diff_hunk: @@ -238,3 +241,11 @@ func SelectRandomPod(pods []*v1.Pod, randomFn func(int) int) (*v1.Pod, error) {
 	randomPod := readyPods[randomFn(len(readyPods))]
 	return randomPod, nil
 }
+
+func GetModelPortForPod(pod *v1.Pod) int64 {
+	modelPort, err := strconv.ParseInt(pod.Labels[modelPortIdentifier], 10, 32)
path: pkg/utils/pod.go
user: login: googs1025
body: Should we add similar boundary conditions to make the code more robust?
🤔 
```go

   if pod == nil || pod.Labels == nil {
        return defaultPodMetricPort
    }

    portStr, exists := pod.Labels[modelPortIdentifier]
    if !exists {
        return defaultPodMetricPort
    }
```
created_at: 2025-05-28T01:18:49Z
author_association: COLLABORATOR
diff_hunk: @@ -238,3 +241,11 @@ func SelectRandomPod(pods []*v1.Pod, randomFn func(int) int) (*v1.Pod, error) {
 	randomPod := readyPods[randomFn(len(readyPods))]
 	return randomPod, nil
 }
+
+func GetModelPortForPod(pod *v1.Pod) int64 {
+	modelPort, err := strconv.ParseInt(pod.Labels[modelPortIdentifier], 10, 32)
path: pkg/utils/pod.go
user: login: varungup90
body: Current code already handles corner case scenarios, added more unit tests to cover them. Only missing is pod == nil check, which will not happen in routing use case.
created_at: 2025-05-28T03:30:55Z
author_association: COLLABORATOR
diff_hunk: @@ -238,3 +241,11 @@ func SelectRandomPod(pods []*v1.Pod, randomFn func(int) int) (*v1.Pod, error) {
 	randomPod := readyPods[randomFn(len(readyPods))]
 	return randomPod, nil
 }
+
+func GetModelPortForPod(pod *v1.Pod) int64 {
+	modelPort, err := strconv.ParseInt(pod.Labels[modelPortIdentifier], 10, 32)
path: pkg/utils/pod.go
user: login: Jeffwan
body: I think the current code can catch the error. even the label is not there, it return the value "". 

I recommend to add some warning logs here. If we do want to know whether it's missing label or invalid port value, we need to split the values.

labelValue, ok := pod.Labels[modelPortIdentifier]
if !ok {
    warn("label %q not found in pod", modelPortIdentifier)
}
created_at: 2025-05-28T22:27:10Z
author_association: COLLABORATOR
diff_hunk: @@ -238,3 +241,11 @@ func SelectRandomPod(pods []*v1.Pod, randomFn func(int) int) (*v1.Pod, error) {
 	randomPod := readyPods[randomFn(len(readyPods))]
 	return randomPod, nil
 }
+
+func GetModelPortForPod(pod *v1.Pod) int64 {
+	modelPort, err := strconv.ParseInt(pod.Labels[modelPortIdentifier], 10, 32)
path: pkg/utils/pod.go
user: login: varungup90
body: Problem with warning logs is that this method is invoked for each request, it will flood the logs. By default, 8000 will be picked up and if custom port is incorrectly set then for debugging purpose, targetpod ip & port are logged (with each request_id) and returned in response header as well. 

Other alternative is that we can add warning message for now as recommended by @googs1025 because its a new feature and help users to debug labels are correctly set. In future we can evaluate, whether to keep them or not.
created_at: 2025-05-29T17:46:06Z
author_association: COLLABORATOR
diff_hunk: @@ -238,3 +241,11 @@ func SelectRandomPod(pods []*v1.Pod, randomFn func(int) int) (*v1.Pod, error) {
 	randomPod := readyPods[randomFn(len(readyPods))]
 	return randomPod, nil
 }
+
+func GetModelPortForPod(pod *v1.Pod) int64 {
+	modelPort, err := strconv.ParseInt(pod.Labels[modelPortIdentifier], 10, 32)
path: pkg/utils/pod.go
user: login: varungup90
body: added a warn message if label is empty or incorrect value is set. Will re-evaluate later if need to keep the warn message. 
created_at: 2025-05-29T23:20:15Z
author_association: COLLABORATOR

PR comments:
user: login: nwangfw
created_at: 2025-05-27T23:41:45Z
body: Does the current approach support the multi-node case? In multi-node inference, we may need to dynamically assign multiple ports. For example, https://github.com/vllm-project/aibrix/blob/4acc1c062ca221d578b48f940e41e146ad7add44/samples/distributed/fleet-two-node.yaml#L59-L66
author_association: COLLABORATOR
user: login: varungup90
created_at: 2025-05-28T00:39:47Z
body: > Does the current approach support the multi-node case? In multi-node inference, we may need to dynamically assign multiple ports. For example, https://github.com/vllm-project/aibrix/blob/4acc1c062ca221d578b48f940e41e146ad7add44/samples/distributed/fleet-two-node.yaml#L59-L66

For inference, only one port is required.
author_association: COLLABORATOR
user: login: nwangfw
created_at: 2025-05-28T00:57:20Z
body: > Does the current approach support the multi-node case? In multi-node inference, we may need to dynamically assign multiple ports. For example,
> 
> https://github.com/vllm-project/aibrix/blob/4acc1c062ca221d578b48f940e41e146ad7add44/samples/distributed/fleet-two-node.yaml#L59-L66
In the multi-node inference scenario, Ray worker needs to find the master pod port.  Just want to throw it out. I feel that we may need to address it here, or at the kuberay operator side or both.
author_association: COLLABORATOR
user: login: googs1025
created_at: 2025-05-30T01:02:24Z
body: LGTM. 😄 
author_association: COLLABORATOR