Skip to content

add kvcache aware plugin e2e test #1222

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
FAUST-BENCHOU wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

add kvcache aware plugin e2e test #1222

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fec1d59
add kvcache-aware plugin e2e test
FAUST-BENCHOU Jun 16, 2026
8cc3718
upd
FAUST-BENCHOU Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion test/e2e/router/e2e_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ func TestMain(m *testing.M) {
os.Exit(1)
}

if err := plugincontext.SetupPluginComponents(testCtx.KubeClient, testCtx.KthenaClient, testNamespace); err != nil {
if err := plugincontext.SetupPluginComponents(testCtx.KubeClient, testCtx.KthenaClient, testNamespace, kthenaNamespace); err != nil {
fmt.Printf("Failed to setup plugin components: %v\n", err)
_ = testCtx.CleanupCommonComponents()
_ = testCtx.DeleteTestNamespace()
Expand Down
117 changes: 117 additions & 0 deletions test/e2e/router/plugins_helpers.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,14 @@ limitations under the License.
package router

import (
"context"
"encoding/json"
"fmt"
"net/http"
"testing"
"time"

"github.com/redis/go-redis/v9"
"github.com/stretchr/testify/require"
backendmetrics "github.com/volcano-sh/kthena/pkg/kthena-router/backend/metrics"
"github.com/volcano-sh/kthena/pkg/kthena-router/backend/vllm"
Expand All @@ -42,6 +44,12 @@ const (
gpuCacheUsageLoadWaitTimeout = 90 * time.Second
gpuCacheUsageLoadConcurrency = 2
gpuCacheUsageLoadMaxTokens = 256

kvCacheRedisWaitTimeout = 90 * time.Second
kvCacheWarmupRequests = 30
kvCacheE2EMaxTokens = 8
redisServerAppLabel = "app.kubernetes.io/component=redis-server"
kvCacheMatrixKeyPrefix = "matrix:kv:block:"
)

func listReadyMockPods(t *testing.T, kube kubernetes.Interface, namespace string) []corev1.Pod {
Expand Down Expand Up @@ -334,4 +342,113 @@ const (
enabled:
- name: gpu-usage
weight: 1`

schedulerOnlyKVCacheAware = `scheduler:
pluginConfig:
- name: kvcache-aware
args:
blockSizeToHash: 8
maxBlocksToMatch: 128
plugins:
Filter:
enabled: []
Score:
enabled:
- name: kvcache-aware
weight: 1`
)

func setupRedisClient(t *testing.T, kube kubernetes.Interface, namespace string) (*redis.Client, func()) {
t.Helper()
pods := utils.ListReadyPodsByLabel(t, kube, namespace, redisServerAppLabel)
require.NotEmpty(t, pods, "no ready redis pods in namespace %s", namespace)

localPort := utils.AllocateLocalPort(t)
pf, err := utils.SetupPortForwardToPod(namespace, pods[0].Name, localPort, "6379")
require.NoError(t, err, "port-forward to redis pod %s", pods[0].Name)

addr := fmt.Sprintf("127.0.0.1:%s", localPort)
client := redis.NewClient(&redis.Options{Addr: addr})
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
require.NoError(t, client.Ping(ctx).Err(), "redis ping via port-forward")

return client, func() {
_ = client.Close()
pf.Close()
}
}
Comment on lines +361 to +380

@gemini-code-assist gemini-code-assist Bot Jun 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Using require assertions (which call t.FailNow()) inside a helper function that is executed within require.Eventually defeats the purpose of retrying. If any transient error occurs (such as a temporary port-forwarding failure or a Redis ping timeout while the container is starting up), the test will fail immediately instead of waiting and retrying. Refactor setupRedisClient to return an error so that the caller can handle it gracefully inside Eventually.

func setupRedisClient(t *testing.T, kube kubernetes.Interface, namespace string) (*redis.Client, func(), error) {
	t.Helper()
	pods := utils.ListReadyPodsByLabel(t, kube, namespace, redisServerAppLabel)
	if len(pods) == 0 {
		return nil, nil, fmt.Errorf("no ready redis pods in namespace %s", namespace)
	}

	localPort := utils.AllocateLocalPort(t)
	pf, err := utils.SetupPortForwardToPod(namespace, pods[0].Name, localPort, "6379")
	if err != nil {
		return nil, nil, fmt.Errorf("port-forward to redis pod %s: %w", pods[0].Name, err)
	}

	addr := fmt.Sprintf("127.0.0.1:%s", localPort)
	client := redis.NewClient(&redis.Options{Addr: addr})
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := client.Ping(ctx).Err(); err != nil {
		_ = client.Close()
		pf.Close()
		return nil, nil, fmt.Errorf("redis ping via port-forward: %w", err)
	}

	return client, func() {
		_ = client.Close()
		pf.Close()
	}, nil
}


func logMockPodContainerTail(t *testing.T, kube kubernetes.Interface, pod corev1.Pod, container string, tailLines int64) {
t.Helper()
raw, err := kube.CoreV1().Pods(pod.Namespace).GetLogs(pod.Name, &corev1.PodLogOptions{
Container: container,
TailLines: &tailLines,
}).Do(context.Background()).Raw()
Comment thread
FAUST-BENCHOU marked this conversation as resolved.
Comment thread
FAUST-BENCHOU marked this conversation as resolved.
Comment thread
FAUST-BENCHOU marked this conversation as resolved.
if err != nil {
t.Logf("kvcache-aware: failed to read %s logs from pod %s: %v", container, pod.Name, err)
return
}
t.Logf("kvcache-aware: pod %s container %s (tail %d lines):\n%s", pod.Name, container, tailLines, string(raw))
}

func waitForKVCachePodInRedis(t *testing.T, kube kubernetes.Interface, redisNamespace string, pod corev1.Pod, modelName string) {
t.Helper()
podIdentifier := fmt.Sprintf("%s.%s", pod.Name, pod.Namespace)
keyPattern := fmt.Sprintf("%s%s@*", kvCacheMatrixKeyPrefix, modelName)

deadline := time.Now().Add(kvCacheRedisWaitTimeout)
poll := 0
for time.Now().Before(deadline) {
poll++
client, closeRedis := setupRedisClient(t, kube, redisNamespace)

ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
keys, err := client.Keys(ctx, keyPattern).Result()
cancel()
Comment on lines +404 to +408
Comment on lines +404 to +408
closeRedis()
Comment thread
FAUST-BENCHOU marked this conversation as resolved.
Comment thread
FAUST-BENCHOU marked this conversation as resolved.

if err != nil {
if poll%5 == 0 {
t.Logf("kvcache-aware: redis poll #%d keys lookup failed: %v", poll, err)
}
time.Sleep(2 * time.Second)
continue
}
Comment thread
FAUST-BENCHOU marked this conversation as resolved.
if len(keys) == 0 {
if poll%5 == 0 {
t.Logf("kvcache-aware: redis poll #%d no keys matching %q", poll, keyPattern)
}
time.Sleep(2 * time.Second)
continue
}

for _, key := range keys {
client, closeRedis := setupRedisClient(t, kube, redisNamespace)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
fields, err := client.HKeys(ctx, key).Result()
cancel()
closeRedis()
Comment thread
FAUST-BENCHOU marked this conversation as resolved.
Comment thread
FAUST-BENCHOU marked this conversation as resolved.
Comment on lines +426 to +431
if err != nil {
continue
}
Comment thread
FAUST-BENCHOU marked this conversation as resolved.
for _, field := range fields {
if field == podIdentifier {
t.Logf("kvcache-aware redis ready: key=%s pod=%s (poll #%d, %d keys)", key, podIdentifier, poll, len(keys))
logMockPodContainerTail(t, kube, pod, "zmq-bridge", 30)
logMockPodContainerTail(t, kube, pod, "runtime", 60)
return
}
}
}

if poll%5 == 0 {
t.Logf("kvcache-aware: redis poll #%d found %d keys but pod %s not listed yet", poll, len(keys), podIdentifier)
}
time.Sleep(2 * time.Second)
}

logMockPodContainerTail(t, kube, pod, "zmq-bridge", 60)
logMockPodContainerTail(t, kube, pod, "runtime", 120)
t.Fatalf("redis did not contain kv block mappings for pod %s model %q (pattern %q)", podIdentifier, modelName, keyPattern)
}
74 changes: 66 additions & 8 deletions test/e2e/router/plugins_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ package router
import (
"context"
"fmt"
"strings"
"testing"
"time"

Expand Down Expand Up @@ -58,7 +59,7 @@ func TestSchedulerPluginPrefixCache(t *testing.T) {
}
}
t.Logf("prefix-cache: dominant pod %d/%d (of %d log lines)", maxCount, 200, routed)
require.GreaterOrEqual(t, routed, 200/2, "expected access logs for routed requests")
require.GreaterOrEqual(t, routed, 190, "expected access logs for at least 190 routed requests")
require.GreaterOrEqual(t, float64(maxCount)/float64(routed), 0.9)

waitForSchedulerPluginInMetrics(t, metricsURL, plugins.PrefixCachePluginName, "score")
Expand Down Expand Up @@ -95,7 +96,7 @@ func TestSchedulerPluginLeastRequest(t *testing.T) {
idleCount := utils.CountSelectedPodInRouterLogs(t, testCtx.KubeClient, kthenaNamespace, idlePod.Name, since)
routed := busyCount + idleCount
t.Logf("least-request: busy pool %d, idle pod %s %d (of %d log lines)", busyCount, idlePod.Name, idleCount, routed)
require.GreaterOrEqual(t, routed, 200/2, "expected access logs for routed requests")
require.GreaterOrEqual(t, routed, 190, "expected access logs for at least 190 routed requests")
require.Greater(t, idleCount, busyCount, "least-request should prefer the idle pod over saturated pods")
require.GreaterOrEqual(t, float64(idleCount)/float64(routed), 0.9,
"least-request should route at least 90%% to the idle pod")
Expand Down Expand Up @@ -140,7 +141,7 @@ func TestSchedulerPluginGPUCacheUsage(t *testing.T) {
idleCount := utils.CountSelectedPodInRouterLogs(t, testCtx.KubeClient, kthenaNamespace, idlePod.Name, since)
routed := busyCount + idleCount
t.Logf("gpu-usage: busy pool %d, idle pod %s %d (of %d log lines)", busyCount, idlePod.Name, idleCount, routed)
require.GreaterOrEqual(t, routed, 200/2, "expected access logs for routed requests")
require.GreaterOrEqual(t, routed, 190, "expected access logs for at least 190 routed requests")
require.Greater(t, idleCount, busyCount, "gpu-usage should prefer the idle pod over kv-cache-hot pods")
require.GreaterOrEqual(t, float64(idleCount)/float64(routed), 0.9,
"gpu-usage should route at least 90%% to the idle pod")
Expand All @@ -150,6 +151,62 @@ func TestSchedulerPluginGPUCacheUsage(t *testing.T) {
waitForSchedulerPluginInMetrics(t, metricsURL, plugins.GPUCacheUsagePluginName, "score")
}

// TestSchedulerPluginKVCacheAware verifies the full kvcache-aware chain:
// sim completions -> native ZMQ -> zmq-bridge -> runtime -> Redis -> router plugin -> routing.
func TestSchedulerPluginKVCacheAware(t *testing.T) {
ctx := context.Background()
t.Cleanup(ensureRedis(t, testCtx.KubeClient, kthenaNamespace))

chatURL, metricsURL, restoreCfg := utils.ApplySchedulerConfig(
t, testCtx.KubeClient, testCtx.KthenaClient, kthenaNamespace, testNamespace,
schedulerOnlyKVCacheAware, plugincontext.ModelServerName, plugincontext.ModelName)
t.Cleanup(restoreCfg)

route := utils.CreateModelRouteFromFile(t, ctx, testCtx.KthenaClient, plugincontext.TestDataDir, testNamespace, "ModelRoute-plugins.yaml")
model := route.Spec.ModelName
utils.WaitForChatModelReady(t, chatURL, model, []utils.ChatMessage{utils.NewChatMessage("user", "ready")}, 90*time.Second)
// Must fit sim kv-cache-size=8 blocks (block-size=8): prompt blocks + max_tokens blocks <= 8.

@YaoZengzeng YaoZengzeng Jun 18, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Isn't it >= 8?

@FAUST-BENCHOU FAUST-BENCHOU Jun 18, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no , <= 8 refers to the maximum kv-cache size of the sim (kv-cache-size=8), not a requirement of at least 8 blocks. We keep the prompt and max_tokens very small to avoid the prompt blocks + output blocks exceeding 8, which would result in a 500 error.

@FAUST-BENCHOU FAUST-BENCHOU Jun 18, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe need to update comment

prompt := "kthena-kvcache-e2e " + strings.Repeat("cache-block-token ", 8)

Comment on lines +168 to +170
pods := listReadyMockPods(t, testCtx.KubeClient, testNamespace)
require.Len(t, pods, pluginMockReplicaCount, "kvcache-aware test needs %d mock pods", pluginMockReplicaCount)
warmedPod := pods[0]

t.Logf("kvcache-aware: warming pod %s with %d direct chat requests (max_tokens=%d)",
warmedPod.Name, kvCacheWarmupRequests, kvCacheE2EMaxTokens)

// Warm one pod only: chat requests populate KV cache, runtime writes Redis.
utils.DirectChatToPod(t, warmedPod, model, prompt, kvCacheWarmupRequests, kvCacheE2EMaxTokens)
logMockPodContainerTail(t, testCtx.KubeClient, warmedPod, "zmq-bridge", 20)
waitForKVCachePodInRedis(t, testCtx.KubeClient, kthenaNamespace, warmedPod, model)

since := metav1.NewTime(time.Now())
Comment thread
FAUST-BENCHOU marked this conversation as resolved.
utils.SendRouterChatRequests(t, chatURL, model, prompt, 200)
time.Sleep(2 * time.Second)
Comment thread
FAUST-BENCHOU marked this conversation as resolved.

warmedCount := 0
otherCount := 0
for _, pod := range pods {
c := utils.CountSelectedPodInRouterLogs(t, testCtx.KubeClient, kthenaNamespace, pod.Name, since)
t.Logf("kvcache-aware: pod %s selected %d/%d", pod.Name, c, 200)
if pod.Name == warmedPod.Name {
warmedCount = c
} else {
otherCount += c
}
}
routed := warmedCount + otherCount
t.Logf("kvcache-aware: warmed pod %s %d, other pods %d (of %d log lines)", warmedPod.Name, warmedCount, otherCount, routed)
require.GreaterOrEqual(t, routed, 190, "expected access logs for at least 190 routed requests")
require.Greater(t, warmedCount, otherCount, "kvcache-aware should prefer the pod with runtime-populated redis blocks")
require.GreaterOrEqual(t, float64(warmedCount)/float64(routed), 0.9,
"kvcache-aware should route at least 90%% to the warmed pod")
require.LessOrEqual(t, float64(otherCount)/float64(routed), 0.1,
"kvcache-aware should route at most 10%% to pods without redis block mappings")

waitForSchedulerPluginInMetrics(t, metricsURL, plugins.KVCacheAwarePluginName, "score")
}

// TestSchedulerPluginLeastLatency verifies least-latency prefers the intrinsically faster
// backend when both pools are idle and scored by observed TTFT/TPOT only.
func TestSchedulerPluginLeastLatency(t *testing.T) {
Expand All @@ -170,13 +227,14 @@ func TestSchedulerPluginLeastLatency(t *testing.T) {
// Prime both pools after the scheduler-specific router restart. Histogram deltas need
// an initial scrape baseline, so send a small baseline batch first, then the measured
// batch that should make the fast pool strictly lower-latency than the slow pool.
const latencyPrimeMaxTokens = 32
latencyPods := append(append([]corev1.Pod{}, fastPods...), slowPods...)
for _, pod := range latencyPods {
utils.DirectChatToPod(t, pod, model, "kthena-router-plugin-e2e-fixed-prompt-latency-baseline-prime", 2)
utils.DirectChatToPod(t, pod, model, "kthena-router-plugin-e2e-fixed-prompt-latency-baseline-prime", 2, latencyPrimeMaxTokens)
}
time.Sleep(2 * time.Second)
for _, pod := range latencyPods {
utils.DirectChatToPod(t, pod, model, "kthena-router-plugin-e2e-fixed-prompt-latency-measured-prime", 8)
utils.DirectChatToPod(t, pod, model, "kthena-router-plugin-e2e-fixed-prompt-latency-measured-prime", 8, latencyPrimeMaxTokens)
}
waitForLeastLatencyMetricsSeparation(t, testCtx.KubeClient, kthenaNamespace, fastPods, slowPods)

Expand All @@ -188,7 +246,7 @@ func TestSchedulerPluginLeastLatency(t *testing.T) {
slowCount := utils.CountSelectedPodsInRouterLogs(t, testCtx.KubeClient, kthenaNamespace, since, slowPods)
routed := fastCount + slowCount
t.Logf("least-latency: fast pool %d, slow pool %d (of %d log lines)", fastCount, slowCount, routed)
require.GreaterOrEqual(t, routed, 200/2, "expected access logs for routed requests")
require.GreaterOrEqual(t, routed, 190, "expected access logs for at least 190 routed requests")
require.Greater(t, fastCount, slowCount, "least-latency should prefer the faster backend when both pools are idle")
require.GreaterOrEqual(t, float64(fastCount)/float64(routed), 0.9,
"least-latency should route at least 90%% to the fast pool")
Expand Down Expand Up @@ -230,7 +288,7 @@ func TestSchedulerPluginLoraAffinity(t *testing.T) {
}
routed := loadedCount + otherCount
t.Logf("lora-affinity: loaded pod %s %d, other pods %d (of %d log lines)", loadedPod.Name, loadedCount, otherCount, routed)
require.GreaterOrEqual(t, routed, 200/2, "expected access logs for routed requests")
require.GreaterOrEqual(t, routed, 190, "expected access logs for at least 190 routed requests")
require.Equal(t, 0, otherCount, "lora-affinity filter should not route to pods without the adapter")
require.GreaterOrEqual(t, float64(loadedCount)/float64(routed), 0.9,
"lora-affinity should route at least 90%% to the pod that loaded the adapter")
Expand Down Expand Up @@ -261,7 +319,7 @@ func TestSchedulerPluginRandom(t *testing.T) {
routed += c
t.Logf("random: pod %s selected %d/%d", pod.Name, c, 200)
}
require.GreaterOrEqual(t, routed, 200/2, "expected access logs for routed requests")
require.GreaterOrEqual(t, routed, 190, "expected access logs for at least 190 routed requests")

// Each pod should receive roughly 1/3 of traffic (±10% absolute ratio).
const randomMaxRatioDeviation = 0.10
Expand Down
32 changes: 31 additions & 1 deletion test/e2e/router/router-plugins/context/context.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import (
networkingv1alpha1 "github.com/volcano-sh/kthena/pkg/apis/networking/v1alpha1"
"github.com/volcano-sh/kthena/test/e2e/utils"
appsv1 "k8s.io/api/apps/v1"
corev1 "k8s.io/api/core/v1"
apierrors "k8s.io/apimachinery/pkg/api/errors"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/client-go/kubernetes"
Expand All @@ -38,12 +39,41 @@ const (
TestDataDir = "test/e2e/router/router-plugins/testdata"
SlowMockDeploymentName = "router-plugin-mock-slow"
SlowMockAppLabel = "router-plugin-mock-slow"

BridgeConfigName = "router-plugin-mock-bridge"
RuntimeEnvConfigName = "router-plugin-mock-runtime-env"
)

// SetupPluginComponents deploys fast/slow plugin mocks and ModelServers shared by plugin e2e tests.
func SetupPluginComponents(kubeClient *kubernetes.Clientset, kthenaClient *clientset.Clientset, namespace string) error {
func SetupPluginComponents(kubeClient *kubernetes.Clientset, kthenaClient *clientset.Clientset, namespace, kthenaNamespace string) error {

@FAUST-BENCHOU FAUST-BENCHOU Jun 16, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Write REDIS_HOST into the ConfigMap before deploying SetupPluginComponents . I tried patching mid-deployment, but the patch triggered Pod restarts, interrupted warmup and kv cache writes.

ctx := stdcontext.Background()

bridgeConfigMap := &corev1.ConfigMap{
ObjectMeta: metav1.ObjectMeta{
Name: BridgeConfigName,
Namespace: namespace,
},
Data: map[string]string{
"zmq-bridge.py": utils.ZMQBridgePy,
},
Comment thread
FAUST-BENCHOU marked this conversation as resolved.
}
if _, err := kubeClient.CoreV1().ConfigMaps(namespace).Create(ctx, bridgeConfigMap, metav1.CreateOptions{}); err != nil && !apierrors.IsAlreadyExists(err) {
return fmt.Errorf("create zmq-bridge configmap: %w", err)
}
Comment thread
FAUST-BENCHOU marked this conversation as resolved.

runtimeEnvConfigMap := &corev1.ConfigMap{
ObjectMeta: metav1.ObjectMeta{
Name: RuntimeEnvConfigName,
Namespace: namespace,
},
Data: map[string]string{
"REDIS_HOST": fmt.Sprintf("redis-server.%s.svc.cluster.local", kthenaNamespace),
},
}
if _, err := kubeClient.CoreV1().ConfigMaps(namespace).Create(ctx, runtimeEnvConfigMap, metav1.CreateOptions{}); err != nil && !apierrors.IsAlreadyExists(err) {
return fmt.Errorf("create runtime env configmap: %w", err)
}
Comment thread
FAUST-BENCHOU marked this conversation as resolved.
Comment thread
FAUST-BENCHOU marked this conversation as resolved.

deployment := utils.LoadYAMLFromFile[appsv1.Deployment](filepath.Join(TestDataDir, "LLM-Mock-plugins.yaml"))
deployment.Namespace = namespace
if _, err := kubeClient.AppsV1().Deployments(namespace).Create(ctx, deployment, metav1.CreateOptions{}); err != nil && !apierrors.IsAlreadyExists(err) {
Expand Down
Loading
Loading