diff --git a/test/e2e/router/e2e_test.go b/test/e2e/router/e2e_test.go index ee7a3ddff9..a29b7a0ae9 100644 --- a/test/e2e/router/e2e_test.go +++ b/test/e2e/router/e2e_test.go @@ -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() diff --git a/test/e2e/router/plugins_helpers.go b/test/e2e/router/plugins_helpers.go index 880ccdec65..bf2f7b6f4a 100644 --- a/test/e2e/router/plugins_helpers.go +++ b/test/e2e/router/plugins_helpers.go @@ -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" @@ -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 { @@ -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() + } +} + +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() + 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() + closeRedis() + + 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 + } + 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() + if err != nil { + continue + } + 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) +} diff --git a/test/e2e/router/plugins_test.go b/test/e2e/router/plugins_test.go index d1b6eb95de..f98d482745 100644 --- a/test/e2e/router/plugins_test.go +++ b/test/e2e/router/plugins_test.go @@ -19,6 +19,7 @@ package router import ( "context" "fmt" + "strings" "testing" "time" @@ -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") @@ -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") @@ -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") @@ -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. + prompt := "kthena-kvcache-e2e " + strings.Repeat("cache-block-token ", 8) + + 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()) + utils.SendRouterChatRequests(t, chatURL, model, prompt, 200) + time.Sleep(2 * time.Second) + + 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) { @@ -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) @@ -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") @@ -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") @@ -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 diff --git a/test/e2e/router/router-plugins/context/context.go b/test/e2e/router/router-plugins/context/context.go index e1503e9c66..5d0ab5c8ee 100644 --- a/test/e2e/router/router-plugins/context/context.go +++ b/test/e2e/router/router-plugins/context/context.go @@ -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" @@ -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 { ctx := stdcontext.Background() + bridgeConfigMap := &corev1.ConfigMap{ + ObjectMeta: metav1.ObjectMeta{ + Name: BridgeConfigName, + Namespace: namespace, + }, + Data: map[string]string{ + "zmq-bridge.py": utils.ZMQBridgePy, + }, + } + 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) + } + + 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) + } + 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) { diff --git a/test/e2e/router/router-plugins/testdata/LLM-Mock-plugins.yaml b/test/e2e/router/router-plugins/testdata/LLM-Mock-plugins.yaml index 91eae3db70..b10a28e6ac 100644 --- a/test/e2e/router/router-plugins/testdata/LLM-Mock-plugins.yaml +++ b/test/e2e/router/router-plugins/testdata/LLM-Mock-plugins.yaml @@ -1,5 +1,6 @@ # Mock vLLM backends for router scheduler plugin e2e tests. -# KV-cache simulation is enabled for gpu-usage; other plugin tests only use their own scheduler plugins. +# Includes sim native ZMQ + zmq-bridge + runtime for kvcache-aware full-chain tests. +# ConfigMap router-plugin-mock-bridge (from scripts/zmq-bridge.py) is created by SetupPluginComponents. apiVersion: apps/v1 kind: Deployment metadata: @@ -15,7 +16,31 @@ spec: app: router-plugin-mock router-plugin-pool: latency spec: + volumes: + - name: bridge-script + configMap: + name: router-plugin-mock-bridge containers: + - name: zmq-bridge + image: ghcr.io/volcano-sh/runtime:latest + imagePullPolicy: IfNotPresent + command: ["python3", "/config/zmq-bridge.py"] + env: + - name: SIM_SUB_BIND + value: "tcp://127.0.0.1:5556" + - name: RUNTIME_PUB_BIND + value: "tcp://127.0.0.1:5557" + volumeMounts: + - name: bridge-script + mountPath: /config + readinessProbe: + exec: + command: + - python3 + - -c + - import socket; s=socket.create_connection(('127.0.0.1',5557),1); s.close() + initialDelaySeconds: 2 + periodSeconds: 2 - name: llm-engine image: ghcr.io/llm-d/llm-d-inference-sim:latest imagePullPolicy: IfNotPresent @@ -36,11 +61,56 @@ spec: - --max-num-seqs=2 - --time-to-first-token=20ms - --inter-token-latency=5ms + - --zmq-endpoint=tcp://127.0.0.1:5556 ports: - containerPort: 8000 readinessProbe: httpGet: path: /health port: 8000 - initialDelaySeconds: 2 + initialDelaySeconds: 3 periodSeconds: 2 + - name: runtime + image: ghcr.io/volcano-sh/runtime:latest + imagePullPolicy: IfNotPresent + args: + - --port + - "8900" + - --engine + - vllm + - --engine-base-url + - http://127.0.0.1:8000 + - --engine-metrics-path + - /metrics + - --pod + - $(POD_NAME).$(NAMESPACE) + - --model + - router-plugin-model + env: + - name: POD_NAME + valueFrom: + fieldRef: + fieldPath: metadata.name + - name: NAMESPACE + valueFrom: + fieldRef: + fieldPath: metadata.namespace + - name: VLLM_USE_V1 + value: "1" + - name: VLLM_ZMQ_ENDPOINT + value: "tcp://127.0.0.1:5557" + - name: REDIS_HOST + valueFrom: + configMapKeyRef: + name: router-plugin-mock-runtime-env + key: REDIS_HOST + - name: REDIS_PORT + value: "6379" + ports: + - containerPort: 8900 + readinessProbe: + httpGet: + path: /health + port: 8900 + initialDelaySeconds: 5 + periodSeconds: 5 diff --git a/test/e2e/utils/assets/zmq-bridge.py b/test/e2e/utils/assets/zmq-bridge.py new file mode 100644 index 0000000000..46effddad2 --- /dev/null +++ b/test/e2e/utils/assets/zmq-bridge.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +# Copyright The Volcano Authors. +# +# E2E helper: bridge llm-d-inference-sim ZMQ (kv@IP@model) to Kthena Runtime (kv-events). + +import logging +import os +import signal +import sys + +import msgpack +import zmq + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger("zmq-bridge") + +SIM_SUB_BIND = os.getenv("SIM_SUB_BIND", "tcp://127.0.0.1:5556") +RUNTIME_PUB_BIND = os.getenv("RUNTIME_PUB_BIND", "tcp://127.0.0.1:5557") +SIM_TOPIC_PREFIX = os.getenv("SIM_TOPIC_PREFIX", "kv@") +RUNTIME_TOPIC = os.getenv("RUNTIME_TOPIC", "kv-events") + +running = True + + +def _shutdown(signum, _frame) -> None: + global running + logger.info("received signal %s, shutting down", signum) + running = False + + +def main() -> int: + signal.signal(signal.SIGTERM, _shutdown) + signal.signal(signal.SIGINT, _shutdown) + + ctx = zmq.Context() + sub = ctx.socket(zmq.SUB) + sub.bind(SIM_SUB_BIND) + sub.setsockopt_string(zmq.SUBSCRIBE, SIM_TOPIC_PREFIX) + sub.setsockopt(zmq.RCVTIMEO, 1000) + + pub = ctx.socket(zmq.PUB) + pub.bind(RUNTIME_PUB_BIND) + + logger.info( + "zmq-bridge listening sim=%s (prefix=%s) runtime=%s (topic=%s)", + SIM_SUB_BIND, + SIM_TOPIC_PREFIX, + RUNTIME_PUB_BIND, + RUNTIME_TOPIC, + ) + + forwarded = 0 + while running: + try: + parts = sub.recv_multipart() + except zmq.Again: + continue + except zmq.ZMQError as exc: + if running: + logger.error("recv failed: %s", exc) + break + + if len(parts) < 3: + logger.warning("ignored message with %d parts", len(parts)) + continue + + topic = parts[0].decode("utf-8", errors="replace") + if not topic.startswith(SIM_TOPIC_PREFIX): + continue + + payload = parts[2] + + # Best-effort introspection: help debug whether sim payload contains token ids. + if forwarded == 0: + try: + obj = msgpack.unpackb(payload, raw=False) + token_ids_present = False + if isinstance(obj, dict): + events = obj.get("events") + if isinstance(events, list) and events: + first = events[0] + if isinstance(first, dict) and "token_ids" in first: + token_ids_present = True + logger.info("first payload decoded type=%s token_ids_present=%s", type(obj).__name__, token_ids_present) + except Exception as exc: + logger.info("first payload msgpack decode failed: %s", exc) + + pub.send_multipart([RUNTIME_TOPIC.encode("utf-8"), b"", payload]) + forwarded += 1 + if forwarded == 1 or forwarded % 10 == 0: + logger.info("forwarded %d batches (latest topic=%s, %d bytes)", forwarded, topic, len(payload)) + + sub.close() + pub.close() + ctx.term() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) + diff --git a/test/e2e/utils/chat.go b/test/e2e/utils/chat.go index bdfdf33653..970288af69 100644 --- a/test/e2e/utils/chat.go +++ b/test/e2e/utils/chat.go @@ -276,7 +276,7 @@ func SendRouterChatRequests(t *testing.T, routerChatURL, modelName, prompt strin } // DirectChatToPod sends count streaming chat requests directly to a pod via port-forward. -func DirectChatToPod(t *testing.T, pod corev1.Pod, model, prompt string, count int) { +func DirectChatToPod(t *testing.T, pod corev1.Pod, model, prompt string, count, maxTokens int) { t.Helper() localPort := AllocateLocalPort(t) pf, err := SetupPortForwardToPod(pod.Namespace, pod.Name, localPort, "8000") @@ -284,12 +284,13 @@ func DirectChatToPod(t *testing.T, pod corev1.Pod, model, prompt string, count i defer pf.Close() url := fmt.Sprintf("http://127.0.0.1:%s/v1/chat/completions", localPort) - body, _ := json.Marshal(map[string]interface{}{ + body, err := json.Marshal(map[string]interface{}{ "model": model, "messages": []map[string]string{{"role": "user", "content": prompt}}, - "max_tokens": 32, + "max_tokens": maxTokens, "stream": true, }) + require.NoError(t, err) client := &http.Client{Timeout: 30 * time.Second} for i := 0; i < count; i++ { req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(body)) @@ -301,6 +302,7 @@ func DirectChatToPod(t *testing.T, pod corev1.Pod, model, prompt string, count i resp.Body.Close() require.Equal(t, http.StatusOK, resp.StatusCode) } + t.Logf("sent %d streaming chat warmup requests to pod %s (max_tokens=%d)", count, pod.Name, maxTokens) } // StartSustainedLongRequestsToPod keeps concurrent long requests on one pod via port-forward. diff --git a/test/e2e/utils/embedded_assets.go b/test/e2e/utils/embedded_assets.go new file mode 100644 index 0000000000..e39cd0f369 --- /dev/null +++ b/test/e2e/utils/embedded_assets.go @@ -0,0 +1,24 @@ +/* +Copyright The Volcano Authors. + +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 _ "embed" + +// ZMQBridgePy is the embedded helper script used by router plugin e2e tests. +// +//go:embed assets/zmq-bridge.py +var ZMQBridgePy string