Support watcher pod rbac in kvcache controller (#1071)

* Support watcher pod rbac in kv cache controller
* Update hpkv key in redis
* Grant pods/exec role to controller-manager
* Add prometheus scrape configuration in controller and watcher

---------

Signed-off-by: Jiaxin Shan <[email protected]>

Author: Jeffwan

cmd/kvcache-watcher/main.go

@@ -56,7 +56,7 @@ import (
 const KVCacheLabelKeyIdentifier = "kvcache.orchestration.aibrix.ai/name"
 const KVCacheLabelKeyRole = "kvcache.orchestration.aibrix.ai/role"
 const KVCacheLabelValueRoleCache = "cache"
-const HPKVRedisNodeMemberKey = "hpkv_nodes"
+const HPKVRedisNodeMemberKey = "hpkv_cluster_metadata"
 const InfiniStoreRedisNodeMemberKey = "kvcache_nodes"
 
 const networkStatusAnnotation = "k8s.volcengine.com/network-status"

config/manager/manager.yaml

@@ -67,6 +67,10 @@ metadata:
     control-plane: controller-manager
     app.kubernetes.io/name: aibrix
     app.kubernetes.io/managed-by: kustomize
+  annotations:
+    prometheus.io/scrape: "true"
+    prometheus.io/port: "8080"
+    prometheus.io/path: "/metrics"
   name: controller-manager-metrics-service
   namespace: system
 spec:

config/rbac/controller-manager/role.yaml

@@ -12,10 +12,37 @@ rules:
   - create
   - patch
   - update
+- apiGroups:
+  - ""
+  resources:
+  - secrets
+  verbs:
+  - get
+  - list
+  - update
+  - watch
+- apiGroups:
+  - admissionregistration.k8s.io
+  resources:
+  - mutatingwebhookconfigurations
+  - validatingwebhookconfigurations
+  verbs:
+  - get
+  - list
+  - update
+  - watch
+- apiGroups:
+  - apiextensions.k8s.io
+  resources:
+  - customresourcedefinitions
+  verbs:
+  - get
+  - list
 - apiGroups:
   - apps
   resources:
   - deployments
+  - statefulsets
   verbs:
   - create
   - delete
@@ -28,30 +55,11 @@ rules:
   - apps
   resources:
   - deployments/status
+  - statefulsets/status
   verbs:
   - get
   - patch
   - update
-- apiGroups:
-    - apps
-  resources:
-    - statefulsets
-  verbs:
-    - create
-    - delete
-    - get
-    - list
-    - patch
-    - update
-    - watch
-- apiGroups:
-    - apps
-  resources:
-    - statefulsets/status
-  verbs:
-    - get
-    - patch
-    - update
 - apiGroups:
   - autoscaling
   resources:
@@ -106,7 +114,9 @@ rules:
 - apiGroups:
   - ""
   resources:
+  - pods/exec
   - pods/status
+  - services
   verbs:
   - create
   - delete
@@ -118,14 +128,12 @@ rules:
 - apiGroups:
   - ""
   resources:
-  - services
+  - serviceaccounts
   verbs:
   - create
   - delete
   - get
   - list
-  - patch
-  - update
   - watch
 - apiGroups:
   - ""
@@ -159,17 +167,6 @@ rules:
   - gateway.networking.k8s.io
   resources:
   - httproutes
-  verbs:
-  - create
-  - delete
-  - get
-  - list
-  - patch
-  - update
-  - watch
-- apiGroups:
-  - gateway.networking.k8s.io
-  resources:
   - referencegrants
   verbs:
   - create
@@ -209,57 +206,7 @@ rules:
   - orchestration.aibrix.ai
   resources:
   - kvcaches
-  verbs:
-  - create
-  - delete
-  - get
-  - list
-  - patch
-  - update
-  - watch
-- apiGroups:
-  - orchestration.aibrix.ai
-  resources:
-  - kvcaches/finalizers
-  verbs:
-  - update
-- apiGroups:
-  - orchestration.aibrix.ai
-  resources:
-  - kvcaches/status
-  verbs:
-  - get
-  - patch
-  - update
-- apiGroups:
-  - orchestration.aibrix.ai
-  resources:
   - rayclusterfleets
-  verbs:
-  - create
-  - delete
-  - get
-  - list
-  - patch
-  - update
-  - watch
-- apiGroups:
-  - orchestration.aibrix.ai
-  resources:
-  - rayclusterfleets/finalizers
-  verbs:
-  - update
-- apiGroups:
-  - orchestration.aibrix.ai
-  resources:
-  - rayclusterfleets/status
-  verbs:
-  - get
-  - patch
-  - update
-- apiGroups:
-  - orchestration.aibrix.ai
-  resources:
   - rayclusterreplicasets
   verbs:
   - create
@@ -272,12 +219,16 @@ rules:
 - apiGroups:
   - orchestration.aibrix.ai
   resources:
+  - kvcaches/finalizers
+  - rayclusterfleets/finalizers
   - rayclusterreplicasets/finalizers
   verbs:
   - update
 - apiGroups:
   - orchestration.aibrix.ai
   resources:
+  - kvcaches/status
+  - rayclusterfleets/status
   - rayclusterreplicasets/status
   verbs:
   - get
@@ -310,28 +261,23 @@ rules:
   - patch
   - update
 - apiGroups:
-  - ""
+  - rbac.authorization.k8s.io
   resources:
-  - secrets
+  - rolebindings
   verbs:
+  - create
+  - delete
   - get
   - list
-  - update
   - watch
 - apiGroups:
-  - admissionregistration.k8s.io
+  - rbac.authorization.k8s.io
   resources:
-  - mutatingwebhookconfigurations
-  - validatingwebhookconfigurations
+  - roles
   verbs:
+  - create
+  - delete
   - get
   - list
   - update
   - watch
-- apiGroups:
-  - apiextensions.k8s.io
-  resources:
-  - customresourcedefinitions
-  verbs:
-  - get
-  - list

config/standalone/kv-cache-controller/patch.yaml

@@ -12,6 +12,6 @@ spec:
             - --leader-elect
             - --leader-election-id=aibrix-kv-cache-controller
             - --health-probe-bind-address=:8081
-            - --metrics-bind-address=0
+            - --metrics-bind-address=:8080
             - --controllers=kv-cache-controller
             - --disable-webhook

pkg/controller/kvcache/backends/common.go

@@ -22,6 +22,7 @@ import (
 	orchestrationv1alpha1 "github.com/vllm-project/aibrix/api/orchestration/v1alpha1"
 	"github.com/vllm-project/aibrix/pkg/constants"
 	corev1 "k8s.io/api/core/v1"
+	rbacv1 "k8s.io/api/rbac/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/intstr"
 )
@@ -113,3 +114,75 @@ func buildRedisService(kvCache *orchestrationv1alpha1.KVCache) *corev1.Service {
 
 	return svc
 }
+
+// buildServiceAccount creates a new ServiceAccount for Distributed kv cache solution.
+func buildServiceAccount(kvCache *orchestrationv1alpha1.KVCache) *corev1.ServiceAccount {
+	sa := &corev1.ServiceAccount{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      kvCache.Name,
+			Namespace: kvCache.Namespace,
+			Labels: map[string]string{
+				constants.KVCacheLabelKeyIdentifier: kvCache.Name,
+				constants.KVCacheLabelKeyRole:       constants.KVCacheLabelValueRoleCache,
+			},
+		},
+	}
+
+	return sa
+}
+
+// buildRole creates a new Role for a KVCache resource.
+func buildRole(kvCache *orchestrationv1alpha1.KVCache) *rbacv1.Role {
+	role := &rbacv1.Role{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      kvCache.Name,
+			Namespace: kvCache.Namespace,
+			Labels: map[string]string{
+				constants.KVCacheLabelKeyIdentifier: kvCache.Name,
+				constants.KVCacheLabelKeyRole:       constants.KVCacheLabelValueRoleCache,
+			},
+		},
+		Rules: []rbacv1.PolicyRule{
+			{
+				APIGroups: []string{""},
+				Resources: []string{"pods"},
+				Verbs:     []string{"get", "list", "watch"},
+			},
+			{
+				APIGroups: []string{""},
+				Resources: []string{"pods/exec"},
+				Verbs:     []string{"create"},
+			},
+		},
+	}
+
+	return role
+}
+
+// buildRoleBinding creates rolebinding for a kvCache object
+func buildRoleBinding(kvCache *orchestrationv1alpha1.KVCache) *rbacv1.RoleBinding {
+	rb := &rbacv1.RoleBinding{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      kvCache.Name,
+			Namespace: kvCache.Namespace,
+			Labels: map[string]string{
+				constants.KVCacheLabelKeyIdentifier: kvCache.Name,
+				constants.KVCacheLabelKeyRole:       constants.KVCacheLabelValueRoleCache,
+			},
+		},
+		Subjects: []rbacv1.Subject{
+			{
+				Kind:      rbacv1.ServiceAccountKind,
+				Name:      kvCache.Name,
+				Namespace: kvCache.Namespace,
+			},
+		},
+		RoleRef: rbacv1.RoleRef{
+			APIGroup: rbacv1.GroupName,
+			Kind:     "Role",
+			Name:     kvCache.Name,
+		},
+	}
+
+	return rb
+}

pkg/controller/kvcache/backends/distributed.go

@@ -60,6 +60,18 @@ func (r *DistributedReconciler) Reconcile(ctx context.Context, kvCache *orchestr
 		return reconcile.Result{}, err
 	}
 
+	if err := r.reconcileWatcherPodServiceAccount(ctx, r.Backend.BuildWatcherPodServiceAccount(kvCache)); err != nil {
+		return reconcile.Result{}, err
+	}
+
+	if err := r.reconcileWatcherPodRole(ctx, r.Backend.BuildWatcherPodRole(kvCache)); err != nil {
+		return reconcile.Result{}, err
+	}
+
+	if err := r.reconcileWatcherPodRoleBinding(ctx, r.Backend.BuildWatcherPodRoleBinding(kvCache)); err != nil {
+		return reconcile.Result{}, err
+	}
+
 	// Handle infinistore kvCache Deployment
 	if err := r.ReconcileStatefulsetObject(ctx, r.Backend.BuildCacheStatefulSet(kvCache)); err != nil {
 		return ctrl.Result{}, err

pkg/controller/kvcache/backends/distributed_test.go

@@ -22,6 +22,7 @@ import (
 
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
+	rbacv1 "k8s.io/api/rbac/v1"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/vllm-project/aibrix/api/orchestration/v1alpha1"
@@ -181,6 +182,9 @@ type mockBackend struct {
 	watcher *corev1.Pod
 	svc     *corev1.Service
 	sts     *appsv1.StatefulSet
+	sa      *corev1.ServiceAccount
+	role    *rbacv1.Role
+	rb      *rbacv1.RoleBinding
 }
 
 func (m mockBackend) Name() string {
@@ -199,6 +203,18 @@ func (m mockBackend) BuildMetadataService(*v1alpha1.KVCache) *corev1.Service {
 	return m.svc
 }
 
+func (m mockBackend) BuildWatcherPodServiceAccount(*v1alpha1.KVCache) *corev1.ServiceAccount {
+	return m.sa
+}
+
+func (m mockBackend) BuildWatcherPodRole(*v1alpha1.KVCache) *rbacv1.Role {
+	return m.role
+}
+
+func (m mockBackend) BuildWatcherPodRoleBinding(*v1alpha1.KVCache) *rbacv1.RoleBinding {
+	return m.rb
+}
+
 func (m mockBackend) BuildWatcherPod(*v1alpha1.KVCache) *corev1.Pod {
 	return m.watcher
 }