[Misc] feature: use kvcache webhook

PROJECT

@@ -54,4 +54,8 @@ resources:
   kind: KVCache
   path: github.com/vllm-project/aibrix/api/orchestration/v1alpha1
   version: v1alpha1
+  webhooks:
+    defaulting: true
+    validation: true
+    webhookVersion: v1
 version: "3"

cmd/controllers/main.go

@@ -316,6 +316,10 @@ func setupControllers(mgr ctrl.Manager, runtimeConfig config.RuntimeConfig, cert
 			setupLog.Error(err, "unable to setup webhook", "webhook", "ModelAdapter")
 			os.Exit(1)
 		}
+		if err := apiwebhook.SetupKVCacheWebhookWithManager(mgr); err != nil {
+			setupLog.Error(err, "unable to setup webhook", "webhook", "KVCache")
+			os.Exit(1)
+		}
 	}
 
 	// Kind controller registration is encapsulated inside the pkg/controller/controller.go

config/webhook/manifests.yaml

@@ -4,6 +4,26 @@ kind: MutatingWebhookConfiguration
 metadata:
   name: mutating-webhook-configuration
 webhooks:
+- admissionReviewVersions:
+  - v1
+  clientConfig:
+    service:
+      name: webhook-service
+      namespace: system
+      path: /mutate-orchestration-aibrix-ai-v1alpha1-kvcache
+  failurePolicy: Fail
+  name: mkvcache-v1alpha1.kb.io
+  rules:
+  - apiGroups:
+    - orchestration.aibrix.ai
+    apiVersions:
+    - v1alpha1
+    operations:
+    - CREATE
+    - UPDATE
+    resources:
+    - kvcaches
+  sideEffects: None
 - admissionReviewVersions:
   - v1
   clientConfig:
@@ -24,6 +44,26 @@ webhooks:
     resources:
     - modeladapters
   sideEffects: None
+- admissionReviewVersions:
+  - v1
+  clientConfig:
+    service:
+      name: webhook-service
+      namespace: system
+      path: /validate-orchestration-aibrix-ai-v1alpha1-kvcache
+  failurePolicy: Fail
+  name: vkvcache-v1alpha1.kb.io
+  rules:
+  - apiGroups:
+    - orchestration.aibrix.ai
+    apiVersions:
+    - v1alpha1
+    operations:
+    - CREATE
+    - UPDATE
+    resources:
+    - kvcaches
+  sideEffects: None
 ---
 apiVersion: admissionregistration.k8s.io/v1
 kind: ValidatingWebhookConfiguration

pkg/controller/kvcache/kvcache_controller.go

@@ -153,14 +153,14 @@ func (r *KVCacheReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ct
 		return reconcile.Result{}, err
 	}
 
-	backend := getKVCacheBackendFromMetadata(kvCache)
-	handler, ok := r.Backends[backend]
+	backend := getKVCacheBackendFromAnnotations(kvCache)
+	backendHandler, ok := r.Backends[backend]
 	if !ok {
 		klog.Warningf("unsupported backend %s", backend)
 		return ctrl.Result{}, fmt.Errorf("unsupported backend: %s", backend)
 	}
 
-	return handler.Reconcile(ctx, kvCache)
+	return backendHandler.Reconcile(ctx, kvCache)
 }
 
 // SetupWithManager sets up the controller with the Manager.
@@ -170,37 +170,14 @@ func (r *KVCacheReconciler) SetupWithManager(mgr ctrl.Manager) error {
 		Complete(r)
 }
 
-// getKVCacheBackendFromMetadata returns the backend based on labels and annotations with fallback logic.
-func getKVCacheBackendFromMetadata(kv *orchestrationv1alpha1.KVCache) string {
+// getKVCacheBackendFromAnnotations returns the backend based on labels and annotations.
+func getKVCacheBackendFromAnnotations(kv *orchestrationv1alpha1.KVCache) string {
 	backend := kv.Annotations[constants.KVCacheLabelKeyBackend]
-	if backend != "" {
-		if isValidKVCacheBackend(backend) {
-			return backend
-		}
-
-		// TODO: Move validation logic to webhook.
-		// invalid value provided, fall back to default backend
-		return constants.KVCacheBackendDefault
-	}
-
-	// provide the compatibility for distributed, centralized mode.
-	mode := kv.Annotations[constants.KVCacheAnnotationMode]
-	switch mode {
-	case "distributed":
-		return constants.KVCacheBackendInfinistore
-	case "centralized":
-		return constants.KVCacheBackendVineyard
-	default:
+	if backend == "" {
+		// In some case where webhooks are disabled (e.g., via --disable-webhook),
+		// mutating webhooks cannot inject default values into the spec. Using annotations
+		// ensures that users can still explicitly specify the backend, even when no webhook is active.
 		return constants.KVCacheBackendDefault
 	}
-}
-
-// isValidKVCacheBackend returns true if the backend is one of the supported backends.
-func isValidKVCacheBackend(b string) bool {
-	switch b {
-	case constants.KVCacheBackendVineyard, constants.KVCacheBackendHPKV, constants.KVCacheBackendInfinistore:
-		return true
-	default:
-		return false
-	}
+	return backend
 }

pkg/controller/kvcache/kvcache_controller_ginkgo_test.go

@@ -55,7 +55,7 @@ var _ = Describe("KVCache Controller", func() {
 					ObjectMeta: metav1.ObjectMeta{
 						Name:      resourceName,
 						Namespace: "default",
-						Labels: map[string]string{
+						Annotations: map[string]string{
 							constants.KVCacheLabelKeyBackend: constants.KVCacheBackendVineyard,
 						},
 					},

pkg/controller/kvcache/kvcache_controller_test.go

@@ -46,38 +46,6 @@ func Test_getKVCacheBackendFromMetadata(t *testing.T) {
 			},
 			expected: constants.KVCacheBackendInfinistore,
 		},
-		{
-			name: "invalid backend annotation falls back to default",
-			annotations: map[string]string{
-				constants.KVCacheLabelKeyBackend: "unknown-backend",
-			},
-			expected: constants.KVCacheBackendDefault,
-		},
-		{
-			name: "no annotation, distributed mode via annotation",
-			annotations: map[string]string{
-				constants.KVCacheAnnotationMode: "distributed",
-			},
-			expected: constants.KVCacheBackendInfinistore,
-		},
-		{
-			name: "no annotation, centralized mode via annotation",
-			annotations: map[string]string{
-				constants.KVCacheAnnotationMode: "centralized",
-			},
-			expected: constants.KVCacheBackendVineyard,
-		},
-		{
-			name: "no annotation, unknown mode falls back to default",
-			annotations: map[string]string{
-				constants.KVCacheAnnotationMode: "invalid-mode",
-			},
-			expected: constants.KVCacheBackendDefault,
-		},
-		{
-			name:     "no annotation or annotation, falls back to default",
-			expected: constants.KVCacheBackendDefault,
-		},
 	}
 
 	for _, tc := range testCases {
@@ -87,48 +55,7 @@ func Test_getKVCacheBackendFromMetadata(t *testing.T) {
 					Annotations: tc.annotations,
 				},
 			}
-			result := getKVCacheBackendFromMetadata(kv)
-			assert.Equal(t, tc.expected, result)
-		})
-	}
-}
-
-func Test_isValidKVCacheBackend(t *testing.T) {
-	testCases := []struct {
-		name     string
-		input    string
-		expected bool
-	}{
-		{
-			name:     "valid vineyard backend",
-			input:    constants.KVCacheBackendVineyard,
-			expected: true,
-		},
-		{
-			name:     "valid infinistore backend",
-			input:    constants.KVCacheBackendInfinistore,
-			expected: true,
-		},
-		{
-			name:     "valid hpkv backend",
-			input:    constants.KVCacheBackendHPKV,
-			expected: true,
-		},
-		{
-			name:     "invalid backend",
-			input:    "not-a-valid-backend",
-			expected: false,
-		},
-		{
-			name:     "empty backend",
-			input:    "",
-			expected: false,
-		},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.name, func(t *testing.T) {
-			result := isValidKVCacheBackend(tc.input)
+			result := getKVCacheBackendFromAnnotations(kv)
 			assert.Equal(t, tc.expected, result)
 		})
 	}

pkg/utils/kvcache.go

@@ -0,0 +1,65 @@
+/*
+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.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package utils
+
+import (
+	"fmt"
+
+	orchestrationv1alpha1 "github.com/vllm-project/aibrix/api/orchestration/v1alpha1"
+	"github.com/vllm-project/aibrix/pkg/constants"
+)
+
+// ValidateKVCacheBackend validates that the backend specified in annotations is valid.
+func ValidateKVCacheBackend(kv *orchestrationv1alpha1.KVCache) error {
+	backend := getKVCacheBackendFromMetadata(kv)
+	if !isValidKVCacheBackend(backend) {
+		return fmt.Errorf("invalid backend %q specified, supported backends are: %s, %s, %s",
+			backend,
+			constants.KVCacheBackendVineyard,
+			constants.KVCacheBackendHPKV,
+			constants.KVCacheBackendInfinistore)
+	}
+	return nil
+}
+
+// getKVCacheBackendFromMetadata returns the backend based on labels and annotations with fallback logic.
+func getKVCacheBackendFromMetadata(kv *orchestrationv1alpha1.KVCache) string {
+	backend := kv.Annotations[constants.KVCacheLabelKeyBackend]
+	if backend != "" {
+		return backend
+	}
+
+	mode := kv.Annotations[constants.KVCacheAnnotationMode]
+	switch mode {
+	case "distributed":
+		return constants.KVCacheBackendInfinistore
+	case "centralized":
+		return constants.KVCacheBackendVineyard
+	default:
+		return constants.KVCacheBackendDefault
+	}
+}
+
+// isValidKVCacheBackend returns true if the backend is one of the supported backends.
+func isValidKVCacheBackend(b string) bool {
+	switch b {
+	case constants.KVCacheBackendVineyard, constants.KVCacheBackendHPKV, constants.KVCacheBackendInfinistore:
+		return true
+	default:
+		return false
+	}
+}