Skip to content

Commit ed00058

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

File tree

7 files changed

+234
-105
lines changed

7 files changed

+234
-105
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: 2 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -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.
173+
// getKVCacheBackendFromMetadata returns the backend based on labels and annotations.
174174
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-
}
175+
return kv.Annotations[constants.KVCacheLabelKeyBackend]
206176
}

‎pkg/controller/kvcache/kvcache_controller_test.go

Lines changed: 0 additions & 73 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 {
@@ -92,44 +60,3 @@ func Test_getKVCacheBackendFromMetadata(t *testing.T) {
9260
})
9361
}
9462
}
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)
132-
assert.Equal(t, tc.expected, result)
133-
})
134-
}
135-
}

‎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+
}

‎pkg/webhook/kvcache_webhook.go

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
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 webhook
18+
19+
import (
20+
"context"
21+
"fmt"
22+
23+
"github.com/vllm-project/aibrix/pkg/constants"
24+
"github.com/vllm-project/aibrix/pkg/utils"
25+
26+
"k8s.io/apimachinery/pkg/runtime"
27+
ctrl "sigs.k8s.io/controller-runtime"
28+
logf "sigs.k8s.io/controller-runtime/pkg/log"
29+
"sigs.k8s.io/controller-runtime/pkg/webhook"
30+
"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
31+
32+
orchestrationv1alpha1 "github.com/vllm-project/aibrix/api/orchestration/v1alpha1"
33+
)
34+
35+
// nolint:unused
36+
// log is for logging in this package.
37+
var kvcachelog = logf.Log.WithName("kvcache-resource")
38+
39+
// SetupKVCacheWebhookWithManager registers the webhook for KVCache in the manager.
40+
func SetupKVCacheWebhookWithManager(mgr ctrl.Manager) error {
41+
return ctrl.NewWebhookManagedBy(mgr).For(&orchestrationv1alpha1.KVCache{}).
42+
WithValidator(&KVCacheCustomValidator{}).
43+
WithDefaulter(&KVCacheCustomDefaulter{}).
44+
Complete()
45+
}
46+
47+
// +kubebuilder:webhook:path=/mutate-orchestration-aibrix-ai-v1alpha1-kvcache,mutating=true,failurePolicy=fail,sideEffects=None,groups=orchestration.aibrix.ai,resources=kvcaches,verbs=create;update,versions=v1alpha1,name=mkvcache-v1alpha1.kb.io,admissionReviewVersions=v1
48+
49+
// KVCacheCustomDefaulter struct is responsible for setting default values on the custom resource of the
50+
// Kind KVCache when those are created or updated.
51+
//
52+
// NOTE: The +kubebuilder:object:generate=false marker prevents controller-gen from generating DeepCopy methods,
53+
// as it is used only for temporary operations and does not need to be deeply copied.
54+
type KVCacheCustomDefaulter struct {
55+
}
56+
57+
var _ webhook.CustomDefaulter = &KVCacheCustomDefaulter{}
58+
59+
// Default implements webhook.CustomDefaulter so a webhook will be registered for the Kind KVCache.
60+
func (d *KVCacheCustomDefaulter) Default(_ context.Context, obj runtime.Object) error {
61+
kvcache, ok := obj.(*orchestrationv1alpha1.KVCache)
62+
if !ok {
63+
return fmt.Errorf("expected a KVCache object but got %T", obj)
64+
}
65+
66+
if kvcache.Annotations == nil {
67+
kvcache.Annotations = map[string]string{}
68+
}
69+
70+
backend := kvcache.Annotations[constants.KVCacheLabelKeyBackend]
71+
if backend == "" {
72+
kvcache.Annotations[constants.KVCacheLabelKeyBackend] = constants.KVCacheBackendDefault
73+
}
74+
75+
mode := kvcache.Annotations[constants.KVCacheAnnotationMode]
76+
if mode == "" {
77+
kvcache.Annotations[constants.KVCacheAnnotationMode] = constants.KVCacheBackendDefault
78+
}
79+
return nil
80+
}
81+
82+
// TODO(user): change verbs to "verbs=create;update;delete" if you want to enable deletion validation.
83+
// NOTE: The 'path' attribute must follow a specific pattern and should not be modified directly here.
84+
// Modifying the path for an invalid path can cause API server errors; failing to locate the webhook.
85+
// +kubebuilder:webhook:path=/validate-orchestration-aibrix-ai-v1alpha1-kvcache,mutating=true,failurePolicy=fail,sideEffects=None,groups=orchestration.aibrix.ai,resources=kvcaches,verbs=create;update,versions=v1alpha1,name=vkvcache-v1alpha1.kb.io,admissionReviewVersions=v1
86+
87+
// KVCacheCustomValidator struct is responsible for validating the KVCache resource
88+
// when it is created, updated, or deleted.
89+
//
90+
// NOTE: The +kubebuilder:object:generate=false marker prevents controller-gen from generating DeepCopy methods,
91+
// as this struct is used only for temporary operations and does not need to be deeply copied.
92+
type KVCacheCustomValidator struct {
93+
}
94+
95+
var _ webhook.CustomValidator = &KVCacheCustomValidator{}
96+
97+
// ValidateCreate implements webhook.CustomValidator so a webhook will be registered for the type KVCache.
98+
func (v *KVCacheCustomValidator) ValidateCreate(_ context.Context, obj runtime.Object) (admission.Warnings, error) {
99+
kvcache, ok := obj.(*orchestrationv1alpha1.KVCache)
100+
if !ok {
101+
return nil, fmt.Errorf("expected a KVCache object but got %T", obj)
102+
}
103+
kvcachelog.Info("Validation for KVCache upon creation", "name", kvcache.GetName())
104+
err := utils.ValidateKVCacheBackend(kvcache)
105+
if err != nil {
106+
return nil, err
107+
}
108+
return nil, nil
109+
}
110+
111+
// ValidateUpdate implements webhook.CustomValidator so a webhook will be registered for the type KVCache.
112+
func (v *KVCacheCustomValidator) ValidateUpdate(_ context.Context, oldObj, newObj runtime.Object) (admission.Warnings, error) {
113+
return nil, nil
114+
}
115+
116+
// ValidateDelete implements webhook.CustomValidator so a webhook will be registered for the type KVCache.
117+
func (v *KVCacheCustomValidator) ValidateDelete(ctx context.Context, obj runtime.Object) (admission.Warnings, error) {
118+
return nil, nil
119+
}

0 commit comments

Comments
 (0)