Skip to content

Commit a51508a

Browse files
googs1025googs1025
committed
feature: use kvcache webhook
Signed-off-by: googs1025 <[email protected]>
1 parent 2734c70 commit a51508a

File tree

9 files changed

+369
-111
lines changed

9 files changed

+369
-111
lines changed

‎PROJECT

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,4 +54,8 @@ resources:
5454
kind: KVCache
5555
path: github.com/vllm-project/aibrix/api/orchestration/v1alpha1
5656
version: v1alpha1
57+
webhooks:
58+
defaulting: true
59+
validation: true
60+
webhookVersion: v1
5761
version: "3"

‎cmd/controllers/main.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -316,6 +316,10 @@ func setupControllers(mgr ctrl.Manager, runtimeConfig config.RuntimeConfig, cert
316316
setupLog.Error(err, "unable to setup webhook", "webhook", "ModelAdapter")
317317
os.Exit(1)
318318
}
319+
if err := apiwebhook.SetupKVCacheWebhookWithManager(mgr); err != nil {
320+
setupLog.Error(err, "unable to setup webhook", "webhook", "KVCache")
321+
os.Exit(1)
322+
}
319323
}
320324

321325
// Kind controller registration is encapsulated inside the pkg/controller/controller.go

‎config/webhook/manifests.yaml

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,26 @@ kind: MutatingWebhookConfiguration
44
metadata:
55
name: mutating-webhook-configuration
66
webhooks:
7+
- admissionReviewVersions:
8+
- v1
9+
clientConfig:
10+
service:
11+
name: webhook-service
12+
namespace: system
13+
path: /mutate-orchestration-aibrix-ai-v1alpha1-kvcache
14+
failurePolicy: Fail
15+
name: mkvcache-v1alpha1.kb.io
16+
rules:
17+
- apiGroups:
18+
- orchestration.aibrix.ai
19+
apiVersions:
20+
- v1alpha1
21+
operations:
22+
- CREATE
23+
- UPDATE
24+
resources:
25+
- kvcaches
26+
sideEffects: None
727
- admissionReviewVersions:
828
- v1
929
clientConfig:
@@ -30,6 +50,26 @@ kind: ValidatingWebhookConfiguration
3050
metadata:
3151
name: validating-webhook-configuration
3252
webhooks:
53+
- admissionReviewVersions:
54+
- v1
55+
clientConfig:
56+
service:
57+
name: webhook-service
58+
namespace: system
59+
path: /validate-orchestration-aibrix-ai-v1alpha1-kvcache
60+
failurePolicy: Fail
61+
name: vkvcache-v1alpha1.kb.io
62+
rules:
63+
- apiGroups:
64+
- orchestration.aibrix.ai
65+
apiVersions:
66+
- v1alpha1
67+
operations:
68+
- CREATE
69+
- UPDATE
70+
resources:
71+
- kvcaches
72+
sideEffects: None
3373
- admissionReviewVersions:
3474
- v1
3575
clientConfig:

‎pkg/controller/kvcache/kvcache_controller.go

Lines changed: 6 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -153,14 +153,14 @@ func (r *KVCacheReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ct
153153
return reconcile.Result{}, err
154154
}
155155

156-
backend := getKVCacheBackendFromMetadata(kvCache)
157-
handler, ok := r.Backends[backend]
156+
backend := getKVCacheBackendFromAnnotations(kvCache)
157+
backendHandler, ok := r.Backends[backend]
158158
if !ok {
159159
klog.Warningf("unsupported backend %s", backend)
160160
return ctrl.Result{}, fmt.Errorf("unsupported backend: %s", backend)
161161
}
162162

163-
return handler.Reconcile(ctx, kvCache)
163+
return backendHandler.Reconcile(ctx, kvCache)
164164
}
165165

166166
// SetupWithManager sets up the controller with the Manager.
@@ -170,37 +170,7 @@ func (r *KVCacheReconciler) SetupWithManager(mgr ctrl.Manager) error {
170170
Complete(r)
171171
}
172172

173-
// getKVCacheBackendFromMetadata returns the backend based on labels and annotations with fallback logic.
174-
func getKVCacheBackendFromMetadata(kv *orchestrationv1alpha1.KVCache) string {
175-
backend := kv.Annotations[constants.KVCacheLabelKeyBackend]
176-
if backend != "" {
177-
if isValidKVCacheBackend(backend) {
178-
return backend
179-
}
180-
181-
// TODO: Move validation logic to webhook.
182-
// invalid value provided, fall back to default backend
183-
return constants.KVCacheBackendDefault
184-
}
185-
186-
// provide the compatibility for distributed, centralized mode.
187-
mode := kv.Annotations[constants.KVCacheAnnotationMode]
188-
switch mode {
189-
case "distributed":
190-
return constants.KVCacheBackendInfinistore
191-
case "centralized":
192-
return constants.KVCacheBackendVineyard
193-
default:
194-
return constants.KVCacheBackendDefault
195-
}
196-
}
197-
198-
// isValidKVCacheBackend returns true if the backend is one of the supported backends.
199-
func isValidKVCacheBackend(b string) bool {
200-
switch b {
201-
case constants.KVCacheBackendVineyard, constants.KVCacheBackendHPKV, constants.KVCacheBackendInfinistore:
202-
return true
203-
default:
204-
return false
205-
}
173+
// getKVCacheBackendFromAnnotations returns the backend based on labels and annotations.
174+
func getKVCacheBackendFromAnnotations(kv *orchestrationv1alpha1.KVCache) string {
175+
return kv.Annotations[constants.KVCacheLabelKeyBackend]
206176
}

‎pkg/controller/kvcache/kvcache_controller_ginkgo_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ var _ = Describe("KVCache Controller", func() {
5555
ObjectMeta: metav1.ObjectMeta{
5656
Name: resourceName,
5757
Namespace: "default",
58-
Labels: map[string]string{
58+
Annotations: map[string]string{
5959
constants.KVCacheLabelKeyBackend: constants.KVCacheBackendVineyard,
6060
},
6161
},

‎pkg/controller/kvcache/kvcache_controller_test.go

Lines changed: 1 addition & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -46,38 +46,6 @@ func Test_getKVCacheBackendFromMetadata(t *testing.T) {
4646
},
4747
expected: constants.KVCacheBackendInfinistore,
4848
},
49-
{
50-
name: "invalid backend annotation falls back to default",
51-
annotations: map[string]string{
52-
constants.KVCacheLabelKeyBackend: "unknown-backend",
53-
},
54-
expected: constants.KVCacheBackendDefault,
55-
},
56-
{
57-
name: "no annotation, distributed mode via annotation",
58-
annotations: map[string]string{
59-
constants.KVCacheAnnotationMode: "distributed",
60-
},
61-
expected: constants.KVCacheBackendInfinistore,
62-
},
63-
{
64-
name: "no annotation, centralized mode via annotation",
65-
annotations: map[string]string{
66-
constants.KVCacheAnnotationMode: "centralized",
67-
},
68-
expected: constants.KVCacheBackendVineyard,
69-
},
70-
{
71-
name: "no annotation, unknown mode falls back to default",
72-
annotations: map[string]string{
73-
constants.KVCacheAnnotationMode: "invalid-mode",
74-
},
75-
expected: constants.KVCacheBackendDefault,
76-
},
77-
{
78-
name: "no annotation or annotation, falls back to default",
79-
expected: constants.KVCacheBackendDefault,
80-
},
8149
}
8250

8351
for _, tc := range testCases {
@@ -87,48 +55,7 @@ func Test_getKVCacheBackendFromMetadata(t *testing.T) {
8755
Annotations: tc.annotations,
8856
},
8957
}
90-
result := getKVCacheBackendFromMetadata(kv)
91-
assert.Equal(t, tc.expected, result)
92-
})
93-
}
94-
}
95-
96-
func Test_isValidKVCacheBackend(t *testing.T) {
97-
testCases := []struct {
98-
name string
99-
input string
100-
expected bool
101-
}{
102-
{
103-
name: "valid vineyard backend",
104-
input: constants.KVCacheBackendVineyard,
105-
expected: true,
106-
},
107-
{
108-
name: "valid infinistore backend",
109-
input: constants.KVCacheBackendInfinistore,
110-
expected: true,
111-
},
112-
{
113-
name: "valid hpkv backend",
114-
input: constants.KVCacheBackendHPKV,
115-
expected: true,
116-
},
117-
{
118-
name: "invalid backend",
119-
input: "not-a-valid-backend",
120-
expected: false,
121-
},
122-
{
123-
name: "empty backend",
124-
input: "",
125-
expected: false,
126-
},
127-
}
128-
129-
for _, tc := range testCases {
130-
t.Run(tc.name, func(t *testing.T) {
131-
result := isValidKVCacheBackend(tc.input)
58+
result := getKVCacheBackendFromAnnotations(kv)
13259
assert.Equal(t, tc.expected, result)
13360
})
13461
}

‎pkg/utils/kvcache.go

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
/*
2+
Copyright 2025 The Aibrix Team.
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
17+
package utils
18+
19+
import (
20+
"fmt"
21+
22+
orchestrationv1alpha1 "github.com/vllm-project/aibrix/api/orchestration/v1alpha1"
23+
"github.com/vllm-project/aibrix/pkg/constants"
24+
)
25+
26+
// ValidateKVCacheBackend validates that the backend specified in annotations is valid.
27+
func ValidateKVCacheBackend(kv *orchestrationv1alpha1.KVCache) error {
28+
backend := getKVCacheBackendFromMetadata(kv)
29+
if !isValidKVCacheBackend(backend) {
30+
return fmt.Errorf("invalid backend %q specified, supported backends are: %s, %s, %s",
31+
backend,
32+
constants.KVCacheBackendVineyard,
33+
constants.KVCacheBackendHPKV,
34+
constants.KVCacheBackendInfinistore)
35+
}
36+
return nil
37+
}
38+
39+
// getKVCacheBackendFromMetadata returns the backend based on labels and annotations with fallback logic.
40+
func getKVCacheBackendFromMetadata(kv *orchestrationv1alpha1.KVCache) string {
41+
backend := kv.Annotations[constants.KVCacheLabelKeyBackend]
42+
if backend != "" {
43+
return backend
44+
}
45+
46+
mode := kv.Annotations[constants.KVCacheAnnotationMode]
47+
switch mode {
48+
case "distributed":
49+
return constants.KVCacheBackendInfinistore
50+
case "centralized":
51+
return constants.KVCacheBackendVineyard
52+
default:
53+
return constants.KVCacheBackendDefault
54+
}
55+
}
56+
57+
// isValidKVCacheBackend returns true if the backend is one of the supported backends.
58+
func isValidKVCacheBackend(b string) bool {
59+
switch b {
60+
case constants.KVCacheBackendVineyard, constants.KVCacheBackendHPKV, constants.KVCacheBackendInfinistore:
61+
return true
62+
default:
63+
return false
64+
}
65+
}

0 commit comments

Comments
 (0)