[FEATURE]: metrics server support for gateway plugins & dashboard (vllm-project#1211)

Signed-off-by: Venkat Raman <vraman2811@gmail.com>
Signed-off-by: Modi Tamam <modi.tamam@gmail.com>

Author: Venkat2811

cmd/plugins/main.go

@@ -1,5 +1,5 @@
 /*
-Copyright 2024 The Aibrix Team.
+Copyright 2025 The Aibrix Team.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -18,7 +18,6 @@ package main
 
 import (
 	"flag"
-	"fmt"
 	"net"
 	"net/http"
 	"os"
@@ -41,16 +40,17 @@ import (
 )
 
 var (
-	grpc_port int
+	grpcAddr    string
+	metricsAddr string
 )
 
 func main() {
-	flag.IntVar(&grpc_port, "port", 50052, "gRPC port")
+	flag.StringVar(&grpcAddr, "grpc-bind-address", ":50052", "The address the gRPC server binds to.")
+	flag.StringVar(&metricsAddr, "metrics-bind-address", ":8080", "The address the metric endpoint binds to.")
 	klog.InitFlags(flag.CommandLine)
 	defer klog.Flush()
 	flag.Parse()
 
-	// Connect to Redis
 	redisClient := utils.GetRedisClient()
 	defer func() {
 		if err := redisClient.Close(); err != nil {
@@ -79,14 +79,12 @@ func main() {
 
 	cache.InitForGateway(config, stopCh, redisClient)
 
-	// Connect to K8s cluster
 	k8sClient, err := kubernetes.NewForConfig(config)
 	if err != nil {
 		klog.Fatalf("Error creating kubernetes client: %v", err)
 	}
 
-	// grpc server init
-	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", grpc_port))
+	lis, err := net.Listen("tcp", grpcAddr)
 	if err != nil {
 		klog.Fatalf("failed to listen: %v", err)
 	}
@@ -95,27 +93,34 @@ func main() {
 		klog.Fatalf("Error on creating gateway k8s client: %v", err)
 	}
 
+	gatewayServer := gateway.NewServer(redisClient, k8sClient, gatewayK8sClient)
+
+	if err := gatewayServer.StartMetricsServer(metricsAddr); err != nil {
+		klog.Fatalf("Failed to start metrics server: %v", err)
+	}
+	klog.Infof("Started metrics server on %s", metricsAddr)
+
 	s := grpc.NewServer()
-	extProcPb.RegisterExternalProcessorServer(s, gateway.NewServer(redisClient, k8sClient, gatewayK8sClient))
+	extProcPb.RegisterExternalProcessorServer(s, gatewayServer)
 
 	healthCheck := health.NewServer()
 	healthPb.RegisterHealthServer(s, healthCheck)
 	healthCheck.SetServingStatus("gateway-plugin", healthPb.HealthCheckResponse_SERVING)
 
-	klog.Info("starting gRPC server on port :50052")
+	klog.Info("starting gRPC server on " + grpcAddr)
 
 	go func() {
 		if err := http.ListenAndServe("localhost:6060", nil); err != nil {
 			klog.Fatalf("failed to setup profiling: %v", err)
 		}
 	}()
 
-	// shutdown
 	var gracefulStop = make(chan os.Signal, 1)
 	signal.Notify(gracefulStop, syscall.SIGINT, syscall.SIGTERM)
 	go func() {
 		sig := <-gracefulStop
 		klog.Warningf("signal received: %v, initiating graceful shutdown...", sig)
+		gatewayServer.Shutdown()
 		s.GracefulStop()
 		os.Exit(0)
 	}()

config/gateway/gateway-plugin/gateway-plugin.yaml

@@ -3,6 +3,12 @@ kind: Service
 metadata:
   name: gateway-plugins
   namespace: aibrix-system
+  labels:
+    app: gateway-plugins
+  annotations:
+    prometheus.io/scrape: "true"
+    prometheus.io/port: "8080"
+    prometheus.io/path: "/metrics"
 spec:
   selector:
     app: gateway-plugins
@@ -15,6 +21,10 @@ spec:
       protocol: TCP
       port: 6060
       targetPort: 6060
+    - name: metrics
+      protocol: TCP
+      port: 8080
+      targetPort: 8080
 ---
 apiVersion: apps/v1
 kind: Deployment
@@ -70,6 +80,8 @@ spec:
               containerPort: 50052
             - name: profiling
               containerPort: 6060
+            - name: metrics
+              containerPort: 8080
           resources:
             limits:
               cpu: 1
@@ -186,6 +198,6 @@ spec:
       processingMode:
         request:
           body: Buffered
-        response: 
+        response:
           body: Streamed
-      messageTimeout: 5s
+      messageTimeout: 5s

config/prometheus/monitor.yaml

@@ -7,7 +7,7 @@ metadata:
     app.kubernetes.io/name: aibrix
     app.kubernetes.io/managed-by: kustomize
   name: controller-manager-metrics-monitor
-  namespace: system
+  namespace: aibrix-system
 spec:
   endpoints:
     - path: /metrics

config/test/gateway/kustomization.yaml

@@ -1,5 +1,5 @@
-apiVersion: kustomize.config.k8s.io/v1beta1
-kind: Kustomization
+apiVersion: kustomize.config.k8s.io/v1alpha1
+kind: Component
 
 patches:
 - path: vtc-test-env-patch.yaml

config/test/gateway/vtc-test-env-patch.yaml

@@ -2,6 +2,7 @@ apiVersion: apps/v1
 kind: Deployment
 metadata:
   name: aibrix-gateway-plugins
+  namespace: aibrix-system
 spec:
   template:
     spec:

config/test/kustomization.yaml

@@ -3,3 +3,10 @@ kind: Kustomization
 
 resources:
 - ../default
+
+patches:
+- path: gateway/vtc-test-env-patch.yaml
+  target:
+    kind: Deployment
+    name: aibrix-gateway-plugins
+    namespace: aibrix-system