minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Feature] Add process-local in-memory KV indexer and Router integration (#33370)

Browse Source 
Co-authored-by: Wu, Yutong <yutong.wu@amd.com>
Co-authored-by: TianDi101 <ditian12@amd.com>
Co-authored-by: Zhangheng <hzh0425@apache.org>
This commit is contained in:
wuyl1
2026-08-20 10:45:35 +08:00
committed by GitHub
co-authored by Wu, Yutong TianDi101 Zhangheng
parent 238ba40c27
commit 360d10d6bc
42 changed files with 6603 additions and 174 deletions
Download Patch File Download Diff File Expand all files Collapse all files
experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py
+153 -86
View File
@@ -1,44 +1,28 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
# SPDX-License-Identifier: Apache-2.0
"""Content-based cross-router routing test for cache-aware-zmq.
"""Content-based routing test for both cache-aware-zmq index backends.
Two routers + two SGLang workers + one shared model. Each router runs an
independent ``cache_aware_zmq`` policy whose ``KvEventIndex`` subscribes
to **both** workers' KV publishers.
Two SGLang workers publish KV events to two routers at once: one runs the
local ``KvEventIndex`` (SUB straight to the workers) and one runs against an
external KV Indexer fed by a ``kv-indexer-bridge`` per worker. Every
subscriber attaches before the single warmup, so one pair of disjoint
prefixes exercises both index backends without a second model load.
The test warms each worker with a DIFFERENT prefix DIRECTLY (bypassing
both routers), then sends those prefixes through each router and
asserts that routing follows the prefix CONTENT: ``PREFIX_X`` lands on
the worker holding X, ``PREFIX_Y`` lands on the worker holding Y, on
both routers.
# Why content-based, not convergence
An earlier version of this test asserted that both routers converged on
the *same dominant worker* after a one-prefix warmup. That property
sounds like it pins the ZMQ-fan-out contract, but it doesn't: when the
KV-event path is broken (subscribers never opened, e.g. a worker's
``/server_info`` lacks the ``kv_events`` block), ``cache_aware_zmq``
silently degrades to **min-load** — which, with sequential requests
holding ``active_load`` at zero, picks the same worker deterministically
on every call within a router. Both routers' min-load picks happened to
agree often enough (about half the time, modulo HashSet seed) to make
the convergence assertion pass even when no event ever flowed.
Content-based routing is uniquely sensitive to the KV-event path. Two
disjoint prefixes warmed on two different workers can only be routed
correctly if the router knows *which worker holds which content* — the
only mechanism that supplies that information is the ``BlockStored``
event stream. Under min-load fallback, both prefixes route to the same
default worker on each router, so the ``PREFIX_Y → worker_y`` assertion
fails regardless of which worker min-load defaults to.
Assert on content, not on convergence: a broken event path degrades
``cache_aware_zmq`` to content-blind min-load, which routes both prefixes to
one worker and so fails at least one assertion below.
"""
from __future__ import annotations
import os
import re
import socket
import subprocess
import time
from contextlib import contextmanager
from pathlib import Path
import httpx
import pytest
@@ -79,6 +63,69 @@ _REQ_TOTAL_RE = re.compile(
_LABEL_RE = re.compile(r'(\w+)="([^"]*)"')
def _open_port() -> int:
with socket.socket() as sock:
sock.bind(("127.0.0.1", 0))
return sock.getsockname()[1]
@contextmanager
def _run(binary: Path, env: dict[str, str], log_path: Path):
with log_path.open("w") as log:
process = subprocess.Popen(
[str(binary)],
env={**os.environ, **env},
stdout=log,
stderr=subprocess.STDOUT,
start_new_session=True,
)
try:
yield process
finally:
process.terminate()
try:
process.wait(timeout=10)
except subprocess.TimeoutExpired:
process.kill()
process.wait(timeout=5)
def _wait_for_indexer(process: subprocess.Popen, port: int, log_path: Path) -> None:
deadline = time.time() + 10
while time.time() < deadline:
if process.poll() is not None:
raise RuntimeError(
f"KV Indexer exited during startup:\n{log_path.read_text()}"
)
try:
with socket.create_connection(("127.0.0.1", port), timeout=0.2):
return
except OSError:
time.sleep(0.1)
raise RuntimeError("timed out waiting for KV Indexer")
def _wait_for_bridge(process: subprocess.Popen, log_path: Path) -> None:
deadline = time.time() + 10
while time.time() < deadline:
output = log_path.read_text(errors="replace")
if "bridge session established" in output:
# ZMQ connect is asynchronous; let the subscription reach the PUB.
time.sleep(0.5)
return
if process.poll() is not None:
raise RuntimeError(f"KV Indexer Bridge exited during startup:\n{output}")
time.sleep(0.1)
raise RuntimeError(f"timed out waiting for KV Indexer Bridge:\n{output}")
def _dump_logs(logs: dict[str, Path]) -> None:
"""Print the tail of each Indexer/Bridge log so a routing failure is debuggable."""
for name, path in logs.items():
tail = path.read_text(errors="replace")[-4000:] if path.exists() else "<no log>"
print(f"\n----- {name} -----\n{tail}")
def _success_counts_by_worker(router_url: str) -> dict[str, int]:
"""Scrape ``/metrics`` and return ``{worker_url: success_count}``."""
r = httpx.get(f"{router_url}/metrics", timeout=5.0)
@@ -173,26 +220,21 @@ def _route_through(router_url: str, model_id: str, prompt: str) -> str:
@pytest.mark.real_gpu
@pytest.mark.slow
def test_two_routers_route_by_prefix_content(
router_binary, # noqa: ARG001 — fixture forces release-binary presence
def test_routers_route_by_prefix_content(
router_binary,
gpu_allocator,
tmp_path,
):
"""Each router must route by prefix CONTENT, agreeing across routers.
With each worker direct-warmed by a different disjoint prefix, the
only way a router can route ``PREFIX_X → worker_x`` AND
``PREFIX_Y → worker_y`` is by consulting a HashTree populated from
the BlockStored events the workers emit. Min-load fallback (the
failure mode when no SUB socket opened) is content-blind and would
route both prefixes to whichever worker its tiebreaker prefers.
"""
"""Both the local ZMQ index and the external Indexer must route by content."""
spec = get_model_spec("qwen3-0.6b")
gpus = gpu_allocator.acquire(2)
# Workers run with the model's REAL chat template (no override): the engine
# caches chat-templated tokens, and the router renders the same template
# (loaded from the model's tokenizer_config.json) before hashing. This
# exercises the production chat-template tokenization path, which aligns
# router query hashes with the engine's templated blocks.
indexer_port = _open_port()
indexer_endpoint = f"http://127.0.0.1:{indexer_port}"
indexer_binary = router_binary.parent / "kv-indexer-server"
bridge_binary = router_binary.parent / "kv-indexer-bridge"
logs = {
name: tmp_path / f"{name}.log" for name in ("indexer", "bridge-x", "bridge-y")
}
try:
with (
spawn_worker(
@@ -205,54 +247,79 @@ def test_two_routers_route_by_prefix_content(
gpu_ids=[gpus[1]],
enable_kv_events=True,
) as worker_y,
Gateway() as router_a,
Gateway() as router_b,
_run(
indexer_binary,
{"KV_INDEXER_LISTEN_ADDR": f"127.0.0.1:{indexer_port}"},
logs["indexer"],
) as indexer,
):
_wait_for_indexer(indexer, indexer_port, logs["indexer"])
worker_urls = [worker_x.url, worker_y.url]
for gw in (router_a, router_b):
gw.start_regular(
def bridge_env(worker, worker_id: str) -> dict[str, str]:
assert worker.kv_events_endpoint is not None
return {
"KV_INDEXER_WORKER_ID": worker_id,
"KV_INDEXER_WORKER_ADDRESS": worker.url,
"KV_INDEXER_ENDPOINT": indexer_endpoint,
"SGLANG_KV_EVENT_ENDPOINT": worker.kv_events_endpoint.replace(
"*", "127.0.0.1"
),
"SGLANG_KV_EVENT_TOPIC": "kv",
}
with (
_run(
bridge_binary, bridge_env(worker_x, "worker-x"), logs["bridge-x"]
) as bridge_x,
_run(
bridge_binary, bridge_env(worker_y, "worker-y"), logs["bridge-y"]
) as bridge_y,
Gateway() as local,
Gateway() as external,
):
local.start_regular(
model_id=spec["model"],
tokenizer_path=spec["model"],
worker_urls=worker_urls,
policy="cache_aware_zmq",
timeout=120.0,
)
# 1. Direct-warm each worker with its own prefix. Must happen
# AFTER both routers have started — ZMQ PUB/SUB doesn't
# replay messages emitted before SUB attaches, so any
# BlockStored event predating subscription is lost and
# the HashTree never sees it.
_direct_warm(worker_x.url, spec["model"], PREFIX_X)
_direct_warm(worker_y.url, spec["model"], PREFIX_Y)
# 2. Drain the SUB mpsc + pump-apply path. Sub-second under
# loopback ZMQ; 2 s leaves comfortable headroom.
time.sleep(2.0)
# 3. Content-routing assertion (×4): each prefix must land
# on the worker that holds it, on either router.
#
# The four assertions below are independently strong:
# min-load fallback routes both prefixes on a given
# router to a single default worker, so for ANY broken-
# fan-out scenario at least one of the four fails.
for router, label in ((router_a, "A"), (router_b, "B")):
landed = _route_through(router.base_url, spec["model"], PREFIX_X)
assert landed == worker_x.url, (
f"router {label}: PREFIX_X must route to worker_x "
f"({worker_x.url}); landed on {landed}. "
f"Likely cause: HashTree is empty — KV-event "
f"subscriber never opened, or BlockStored events "
f"never reached the pump."
)
landed = _route_through(router.base_url, spec["model"], PREFIX_Y)
assert landed == worker_y.url, (
f"router {label}: PREFIX_Y must route to worker_y "
f"({worker_y.url}); landed on {landed}. "
f"Likely cause: HashTree is empty — KV-event "
f"subscriber never opened, or BlockStored events "
f"never reached the pump."
external.start_regular(
model_id=spec["model"],
tokenizer_path=spec["model"],
worker_urls=worker_urls,
policy="cache_aware_zmq",
kv_indexer_endpoint=indexer_endpoint,
timeout=120.0,
)
_wait_for_bridge(bridge_x, logs["bridge-x"])
_wait_for_bridge(bridge_y, logs["bridge-y"])
_direct_warm(worker_x.url, spec["model"], PREFIX_X)
_direct_warm(worker_y.url, spec["model"], PREFIX_Y)
time.sleep(2.0)
try:
for router, label in (
(local, "local-index"),
(external, "external-indexer"),
):
landed = _route_through(
router.base_url, spec["model"], PREFIX_X
)
assert (
landed == worker_x.url
), f"router {label}: PREFIX_X must route to {worker_x.url}; landed on {landed}"
landed = _route_through(
router.base_url, spec["model"], PREFIX_Y
)
assert (
landed == worker_y.url
), f"router {label}: PREFIX_Y must route to {worker_y.url}; landed on {landed}"
except Exception:
_dump_logs(logs)
raise
finally:
gpu_allocator.release(gpus)
experimental/sgl-router/tests/e2e/infra/gateway.py
+6
View File
@@ -172,6 +172,7 @@ class Gateway:
tokenizer_path: str,
worker_urls: list[str],
policy: str = "round_robin",
kv_indexer_endpoint: str | None = None,
timeout: float = 60.0,
) -> None:
"""Start the router in regular (non-PD) mode.
@@ -186,6 +187,7 @@ class Gateway:
metadata are learned from ``/server_info``.
policy: Policy kind — ``round_robin``, ``random``, ``power_of_two``,
or ``cache_aware_zmq``.
kv_indexer_endpoint: Optional external KV Indexer gRPC endpoint.
timeout: How long to wait for ``/readyz`` before giving up.
"""
self._launch(
@@ -194,6 +196,7 @@ class Gateway:
tokenizer_path=tokenizer_path,
urls=list(worker_urls),
policy=policy,
kv_indexer_endpoint=kv_indexer_endpoint,
),
timeout=timeout,
)
@@ -293,6 +296,7 @@ class Gateway:
tokenizer_path: str,
urls: list[str],
policy: str,
kv_indexer_endpoint: str | None = None,
) -> list[str]:
resolved_tokenizer = _resolve_tokenizer_path(tokenizer_path)
@@ -317,6 +321,8 @@ class Gateway:
"--stale-request-timeout-secs",
str(self.stale_request_timeout_secs),
]
if kv_indexer_endpoint is not None:
args += ["--kv-indexer-endpoint", kv_indexer_endpoint]
# `--worker-urls` is multi-valued; keep it last so clap doesn't
# absorb a following flag as a URL.
args += ["--worker-urls", *urls]
experimental/sgl-router/tests/e2e/k8s_integration/Dockerfile.router
+5 -1
View File
@@ -10,7 +10,11 @@ FROM rust:1.90-bookworm AS builder
# `channel = "1.90"`.
ENV RUSTUP_TOOLCHAIN=1.90.0
# libssl-dev + pkg-config ship with rust:1.90-bookworm already; no apt-get needed.
# libssl-dev + pkg-config ship with rust:1.90-bookworm already; protoc does not,
# and the Indexer's build script needs it to compile the KV-indexer protos.
RUN apt-get update \
&& apt-get install -y --no-install-recommends protobuf-compiler \
&& rm -rf /var/lib/apt/lists/*
WORKDIR /build