[Router] k8s e2e for cache-aware peer bootstrap (12/13)

experimental/sgl-router/tests/e2e/k8s_integration/Dockerfile.fake_kv_worker

@@ -0,0 +1,7 @@
+FROM python:3.12-slim
+WORKDIR /app
+# pyzmq for the PUB socket, msgspec for the engine's msgpack event encoding.
+RUN pip install --no-cache-dir fastapi uvicorn pyzmq msgspec
+COPY fake_kv_worker.py .
+EXPOSE 8000 5557
+CMD ["python", "fake_kv_worker.py"]

experimental/sgl-router/tests/e2e/k8s_integration/conftest.py

@@ -2,7 +2,8 @@
 
 These tests require:
   - A kind cluster named 'sgl-router-kind'
-  - The sgl-router:e2e and sgl-router-fake-worker:e2e images loaded into kind
+  - The sgl-router:e2e, sgl-router-fake-worker:e2e and
+    sgl-router-fake-kv-worker:e2e images loaded into kind
   - kubectl configured to use the kind-sgl-router-kind context
 
 Setup:  ./tests/e2e/k8s_integration/setup.sh

experimental/sgl-router/tests/e2e/k8s_integration/fake_kv_worker.py

@@ -0,0 +1,185 @@
+"""Fake SGLang worker that publishes KV-cache events, for kind E2E testing.
+
+Unlike ``fake_worker.py`` this one advertises a ``kv_events`` block on
+``/server_info`` and runs a real ZMQ PUB socket speaking SGLang's wire format
+(3-frame multipart: topic, big-endian i64 seq, msgpack ``KVEventBatch``), so the
+router's real subscriber path is exercised end to end.
+
+Events are emitted **only** when the test asks for them, via ``/control/store``.
+That is deliberate: a worker that emitted on a timer would make every
+view-comparison assertion a race against the next event. The test drives the
+event stream, then quiesces, then compares.
+
+Endpoints beyond the usual worker surface:
+
+  POST /control/store   {"chains": [[h1, h2, ...], ...], "dp_rank": 0}
+                        -> publish one BlockStored batch per chain; returns
+                           {"last_seq": N}, so the test can wait for routers to
+                           catch up to a known watermark instead of sleeping.
+"""
+
+from __future__ import annotations
+
+import os
+import threading
+import time
+from typing import Optional
+
+import msgspec
+import uvicorn
+import zmq
+from fastapi import FastAPI, Request
+
+app = FastAPI()
+
+MODEL_ID = os.environ.get("MODEL_ID", "tiny")
+BLOCK_SIZE = int(os.environ.get("BLOCK_SIZE", "4"))
+DP_SIZE = int(os.environ.get("DP_SIZE", "1"))
+KV_PORT_BASE = int(os.environ.get("KV_PORT_BASE", "5557"))
+POD_IP = os.environ.get("POD_IP", "0.0.0.0")
+
+
+# Mirrors of the engine's event structs (`EventBatch`, `KVCacheEvent`,
+# `BlockStored` in python/sglang/srt/disaggregation/kv_events.py),
+# declared with the same msgspec options so msgspec itself produces the engine's
+# encoding: the batch is an array `[ts, events, attn_dp_rank]`, and each event
+# is a tagged MAP whose `type` key carries the class name. The router rejects
+# the older tagged-array event shape outright, so hand-building events here
+# would test nothing but a decode failure.
+class _KVCacheEvent(msgspec.Struct, omit_defaults=True, gc=False, tag=True):
+    pass
+
+
+class BlockStored(_KVCacheEvent):
+    block_hashes: list[int]
+    parent_block_hash: Optional[int]
+    token_ids: list[int]
+    block_size: int
+    lora_id: Optional[int]
+    medium: Optional[str] = None
+
+
+class _EventBatch(msgspec.Struct, array_like=True, gc=False):
+    ts: float
+    events: list[BlockStored]
+    attn_dp_rank: Optional[int] = None
+
+
+class _Publisher:
+    """One PUB socket per DP rank, mirroring ``ZmqEventPublisher``.
+
+    The seq counter starts at 0 and is monotonic per rank, matching the
+    engine's ``itertools.count()``; the router resolves a rank whose stream it
+    received from batch 0 without a peer snapshot. ``last_seq`` is -1 until the
+    first publish.
+    A lock serialises publishes so concurrent control calls cannot interleave a
+    frame or hand out a duplicate seq.
+    """
+
+    def __init__(self, dp_rank: int) -> None:
+        self._sock = zmq.Context.instance().socket(zmq.PUB)
+        self._sock.set_hwm(100_000)
+        self._sock.bind(f"tcp://0.0.0.0:{KV_PORT_BASE + dp_rank}")
+        self._seq = -1
+        self._lock = threading.Lock()
+        self._encoder = msgspec.msgpack.Encoder()
+
+    def publish(self, events: list[BlockStored]) -> None:
+        batch = _EventBatch(ts=time.time(), events=events)
+        with self._lock:
+            self._seq += 1
+            self._sock.send_multipart(
+                (b"", self._seq.to_bytes(8, "big"), self._encoder.encode(batch))
+            )
+
+    @property
+    def last_seq(self) -> int:
+        with self._lock:
+            return self._seq
+
+
+# Bound before the HTTP server starts, so the sockets exist before the readiness
+# probe can pass and the router can discover this worker and connect its SUB.
+_publishers: dict[int, _Publisher] = {rank: _Publisher(rank) for rank in range(DP_SIZE)}
+
+
+@app.get("/health")
+async def health():
+    return {"status": "ok"}
+
+
+@app.get("/server_info")
+async def server_info():
+    return {
+        "served_model_name": MODEL_ID,
+        # No speculative decoding -> unigram block hashing.
+        "speculative_algorithm": None,
+        "kv_events": {
+            "publisher": "zmq",
+            # The router replaces a wildcard host with the host from the worker
+            # URL; report the pod IP explicitly so it is unambiguous.
+            "endpoint_host": POD_IP,
+            "endpoint_port_base": KV_PORT_BASE,
+            "topic": "",
+            "block_size": BLOCK_SIZE,
+            "dp_size": DP_SIZE,
+        },
+    }
+
+
+@app.get("/v1/models")
+async def models():
+    return {
+        "object": "list",
+        "data": [
+            {"id": MODEL_ID, "object": "model", "created": 0, "owned_by": "sglang"}
+        ],
+    }
+
+
+@app.post("/control/store")
+async def control_store(request: Request):
+    payload = await request.json()
+    dp_rank = int(payload.get("dp_rank", 0))
+    chains: list[list[int]] = payload["chains"]
+    for chain in chains:
+        # One BlockStored per block, parent-chained, exactly as
+        # `_record_store_event` emits them.
+        events = []
+        parent = None
+        for h in chain:
+            events.append(
+                BlockStored(
+                    block_hashes=[h],
+                    parent_block_hash=parent,
+                    token_ids=[],
+                    block_size=BLOCK_SIZE,
+                    lora_id=None,
+                )
+            )
+            parent = h
+        _publishers[dp_rank].publish(events)
+    return {"last_seq": _publishers[dp_rank].last_seq}
+
+
+@app.post("/v1/chat/completions")
+async def chat_completions(request: Request):
+    payload = await request.json()
+    messages = payload.get("messages", [])
+    last = messages[-1]["content"] if messages else ""
+    return {
+        "id": "chatcmpl-mock",
+        "object": "chat.completion",
+        "model": payload.get("model", MODEL_ID),
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": last},
+                "finish_reason": "stop",
+            }
+        ],
+    }
+
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)

experimental/sgl-router/tests/e2e/k8s_integration/manifests/kv-bootstrap.yaml

@@ -0,0 +1,165 @@
+# Fixture for the cache-aware peer-bootstrap tests.
+#
+# Deliberately a SEPARATE worker Service + router Deployment from the ones in
+# router.yaml / setup.sh, so these tests cannot perturb (or be perturbed by) the
+# discovery / lifecycle / reconciliation suites that share the namespace.
+---
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: fake-kv-worker
+  namespace: sgl-router-test
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: sglang-kv
+  template:
+    metadata:
+      labels:
+        app: sglang-kv
+    spec:
+      containers:
+        - name: worker
+          image: sgl-router-fake-kv-worker:e2e
+          imagePullPolicy: Never
+          env:
+            # The worker advertises this as `kv_events.endpoint_host`, so the
+            # router connects to a concrete address rather than a wildcard.
+            - name: POD_IP
+              valueFrom:
+                fieldRef:
+                  fieldPath: status.podIP
+            - name: MODEL_ID
+              value: "tiny"
+            - name: BLOCK_SIZE
+              value: "4"
+            - name: DP_SIZE
+              value: "1"
+          ports:
+            - containerPort: 8000
+              name: http
+            - containerPort: 5557
+              name: kv-events
+          readinessProbe:
+            httpGet:
+              path: /health
+              port: 8000
+            initialDelaySeconds: 1
+            periodSeconds: 2
+---
+apiVersion: v1
+kind: Service
+metadata:
+  name: fake-kv-worker
+  namespace: sgl-router-test
+  # Copied onto the generated EndpointSlice, which is what the router's
+  # `--selector app=sglang-kv` matches.
+  labels:
+    app: sglang-kv
+spec:
+  selector:
+    app: sglang-kv
+  ports:
+    - name: http
+      port: 8000
+      targetPort: 8000
+---
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: sgl-router-kv
+  namespace: sgl-router-test
+spec:
+  replicas: 3
+  selector:
+    matchLabels:
+      app: sgl-router-kv
+  # Surge-only so a rolling update always keeps the warm replicas serving while
+  # the new ones bootstrap — the configuration the peer-bootstrap design assumes.
+  strategy:
+    type: RollingUpdate
+    rollingUpdate:
+      maxSurge: 2
+      maxUnavailable: 0
+  template:
+    metadata:
+      labels:
+        app: sgl-router-kv
+    spec:
+      serviceAccountName: sgl-router
+      containers:
+        - name: router
+          image: sgl-router:e2e
+          imagePullPolicy: Never
+          env:
+            # Lets a replica exclude itself from its own peer list: POD_NAME
+            # matches the EndpointSlice endpoints' targetRef (covers every
+            # address of the pod); POD_IP is the fallback for endpoints with
+            # no targetRef.
+            - name: POD_NAME
+              valueFrom:
+                fieldRef:
+                  fieldPath: metadata.name
+            - name: POD_NAMESPACE
+              valueFrom:
+                fieldRef:
+                  fieldPath: metadata.namespace
+            - name: POD_IP
+              valueFrom:
+                fieldRef:
+                  fieldPath: status.podIP
+          args:
+            - "--host"
+            - "0.0.0.0"
+            - "--port"
+            - "8090"
+            - "--model-id"
+            - "tiny"
+            - "--tokenizer-path"
+            - "/etc/tokenizer/tiny.json"
+            - "--policy"
+            - "cache_aware"
+            - "--service-discovery"
+            - "--service-discovery-namespace"
+            - "sgl-router-test"
+            - "--selector"
+            - "app=sglang-kv"
+            # Matched against EndpointSlice labels. The EndpointSlice
+            # controller stamps `kubernetes.io/service-name` on every slice, so
+            # this finds sibling replicas through this Deployment's own Service.
+            - "--kv-peer-selector"
+            - "kubernetes.io/service-name=sgl-router-kv"
+            # Generous relative to a loopback fetch inside kind, so a slow CI
+            # node cannot turn a working bootstrap into a cold one and make the
+            # view assertions flaky.
+            - "--kv-bootstrap-timeout-ms"
+            - "20000"
+          ports:
+            - containerPort: 8090
+              name: http
+          readinessProbe:
+            httpGet:
+              path: /readyz
+              port: 8090
+            initialDelaySeconds: 2
+            periodSeconds: 2
+          livenessProbe:
+            httpGet:
+              path: /healthz
+              port: 8090
+            initialDelaySeconds: 5
+            periodSeconds: 10
+---
+apiVersion: v1
+kind: Service
+metadata:
+  name: sgl-router-kv
+  namespace: sgl-router-test
+spec:
+  selector:
+    app: sgl-router-kv
+  ports:
+    - name: http
+      port: 8090
+      targetPort: 8090