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

Co-authored-by: Wu, Yutong <yutong.wu@amd.com> Co-authored-by: TianDi101 <ditian12@amd.com> Co-authored-by: Zhangheng <hzh0425@apache.org>
2026-08-20 10:45:35 +08:00
parent 238ba40c27
commit 360d10d6bc
42 changed files with 6603 additions and 174 deletions
@@ -109,6 +109,7 @@ async fn zmq_indexer_routes_to_publishing_worker_e2e() {
            cache_threshold: 0.0,
            balance_abs_threshold: 32,
            balance_rel_threshold: 1.1,
+            kv_indexer_endpoint: None,
        },
        kv_index.tree(),
        Arc::clone(&tokenizers),
@@ -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)
@@ -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]
@@ -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

@@ -12,18 +12,10 @@
 //!   doesn't render tool schemas, so its ids would diverge from the engine).
 //! * A request with multimodal (array) content → `input_ids` omitted (a text
 //!   tokenizer can't represent image content).
-//!
-//! The model id contains `deepseek-v4` so the tokenizer registry auto-attaches
-//! the built-in V4 chat encoder — the engine-equivalent path — without a
-//! template fixture.

 use axum::body::Body;
 use axum::http::{Request, StatusCode};
 use serde_json::{json, Value};
-use sgl_router::config::{
-    ActiveLoadConfig, CacheAwareConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig,
-    PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
-};
 use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
 use sgl_router::policies::factory::build_registry;
 use sgl_router::policies::kv_events::{BlockSizeOracle, HashTree};
@@ -36,33 +28,9 @@ use std::sync::Arc;
 use std::time::Duration;
 use tower::ServiceExt;

+use crate::common::cache_aware_fixture::{config, MODEL};
 use crate::common::mock_worker::MockWorker;

-const MODEL: &str = "deepseek-v4-tiny";
-
-fn config() -> Config {
-    Config {
-        server: ServerConfig {
-            host: "0".into(),
-            port: 0,
-        },
-        observability: ObservabilityConfig::default(),
-        model: ModelConfig {
-            id: MODEL.into(),
-            tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
-            policy: PolicyKind::CacheAwareZmq,
-            circuit_breaker: None,
-            cache_aware: Some(CacheAwareConfig::default()),
-            sticky: None,
-        },
-        discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
-            urls: vec!["http://placeholder:0".into()],
-        }),
-        proxy: ProxyConfig::default(),
-        active_load: ActiveLoadConfig::default(),
-    }
-}
-
 fn build_ctx(url: String) -> Arc<AppContext> {
    let cfg = config();
    let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
@@ -0,0 +1,39 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
+// SPDX-License-Identifier: Apache-2.0
+
+//! Shared router config for the cache-aware proxy tests.
+//!
+//! The model id contains `deepseek-v4` so the tokenizer registry auto-attaches the
+//! built-in V4 chat encoder — the engine-equivalent path — with no template fixture.
+
+use sgl_router::config::{
+    ActiveLoadConfig, CacheAwareConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig,
+    PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
+};
+
+pub const MODEL: &str = "deepseek-v4-tiny";
+
+/// A single-model `cache_aware_zmq` router. Discovery is a placeholder because
+/// every caller installs its own `WorkerRegistry`.
+pub fn config() -> Config {
+    Config {
+        server: ServerConfig {
+            host: "0".into(),
+            port: 0,
+        },
+        observability: ObservabilityConfig::default(),
+        model: ModelConfig {
+            id: MODEL.into(),
+            tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
+            policy: PolicyKind::CacheAwareZmq,
+            circuit_breaker: None,
+            cache_aware: Some(CacheAwareConfig::default()),
+            sticky: None,
+        },
+        discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
+            urls: vec!["http://placeholder:0".into()],
+        }),
+        proxy: ProxyConfig::default(),
+        active_load: ActiveLoadConfig::default(),
+    }
+}
@@ -3,5 +3,6 @@

 //! Shared test harness re-exports.

+pub mod cache_aware_fixture;
 pub mod mock_worker;
 pub mod streaming;
@@ -0,0 +1,135 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
+// SPDX-License-Identifier: Apache-2.0
+
+//! Full HTTP routing path backed by a real in-memory Indexer gRPC server.
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use axum::body::Body;
+use axum::http::{Request, StatusCode};
+use serde_json::json;
+use sgl_kv_indexer::pb::kv_indexer_client::KvIndexerClient;
+use sgl_kv_indexer::pb::{
+    ApplyExternalKvBatchRequest, ExternalKvAction, ExternalKvActionType, TierType,
+};
+use sgl_kv_indexer::{
+    server_builder, GrpcPrefixIndex, InMemoryKvIndexerBackend, KvIndexerService, PrefixIndexConfig,
+};
+use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
+use sgl_router::policies::factory::build_registry;
+use sgl_router::policies::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree};
+use sgl_router::policies::request_tokens_for;
+use sgl_router::proxy::Proxy;
+use sgl_router::server::app::build_router;
+use sgl_router::server::app_context::AppContext;
+use sgl_router::tokenizer::TokenizerRegistry;
+use sgl_router::workers::WorkerRegistry;
+use tokio_stream::wrappers::TcpListenerStream;
+use tower::ServiceExt;
+
+use crate::common::cache_aware_fixture::{config, MODEL};
+use crate::common::mock_worker::MockWorker;
+
+#[tokio::test]
+async fn external_indexer_routes_to_the_cached_worker() {
+    let cached = MockWorker::start(vec![]).await;
+    let uncached = MockWorker::start(vec![]).await;
+    let cfg = config();
+    let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
+    let body = json!({
+        "model": MODEL,
+        "messages": [{"role": "user", "content": "hello there friend"}],
+    });
+    let tokens = request_tokens_for(&tokenizers, &ModelId(MODEL.into()), &body)
+        .expect("test prompt tokenizes");
+    let hashes = compute_block_hashes(&tokens.ids, 1);
+    assert!(!hashes.is_empty());
+
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let endpoint = format!("http://{}", listener.local_addr().unwrap());
+    let server = tokio::spawn(async move {
+        server_builder()
+            .add_service(KvIndexerService::new(InMemoryKvIndexerBackend::new()).into_server())
+            .serve_with_incoming(TcpListenerStream::new(listener))
+            .await
+            .unwrap();
+    });
+
+    let mut indexer = KvIndexerClient::connect(endpoint.clone()).await.unwrap();
+    indexer
+        .apply_external_kv_batch(ApplyExternalKvBatchRequest {
+            worker_id: "cached-worker".into(),
+            seq: 1,
+            actions: vec![ExternalKvAction {
+                r#type: ExternalKvActionType::ActionReport as i32,
+                tier: TierType::TierHbm as i32,
+                hashes: hashes.clone(),
+                component_masks: Vec::new(),
+                block_sizes: Vec::new(),
+            }],
+            worker_address: cached.url.clone(),
+            cache_spec: None,
+        })
+        .await
+        .unwrap();
+
+    let registry = Arc::new(WorkerRegistry::default());
+    for url in [&cached.url, &uncached.url] {
+        registry
+            .add(WorkerSpec {
+                id: WorkerId(url.clone()),
+                url: url.clone(),
+                mode: WorkerMode::Plain,
+                model_ids: vec![ModelId(MODEL.into())],
+                bootstrap_port: None,
+            })
+            .unwrap();
+    }
+    let oracle = BlockSizeOracle::new();
+    oracle.try_set(1).unwrap();
+    let policies = Arc::new(
+        build_registry(
+            &cfg,
+            Arc::new(HashTree::new()),
+            Arc::clone(&tokenizers),
+            Arc::clone(&oracle),
+        )
+        .unwrap(),
+    );
+    let mut ctx = AppContext::new(
+        cfg,
+        tokenizers,
+        Arc::new(Proxy::new(Duration::from_secs(5)).unwrap()),
+        registry,
+        policies,
+    );
+    ctx.prefix_index = Some(Arc::new(
+        GrpcPrefixIndex::new(PrefixIndexConfig {
+            endpoint,
+            query_deadline: Duration::from_secs(1),
+            max_inflight: 4,
+        })
+        .unwrap(),
+    ));
+    ctx.block_size_oracle = oracle;
+
+    let app = build_router(Arc::new(ctx));
+    let response = app
+        .oneshot(
+            Request::builder()
+                .method("POST")
+                .uri("/v1/chat/completions")
+                .header("content-type", "application/json")
+                .body(Body::from(serde_json::to_vec(&body).unwrap()))
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    assert_eq!(response.status(), StatusCode::OK);
+    assert!(cached.captured.lock().unwrap().last_body.is_some());
+    assert!(uncached.captured.lock().unwrap().last_body.is_none());
+
+    server.abort();
+}
@@ -12,6 +12,7 @@ mod common;

 mod cache_aware_input_ids;
 mod chat_routing;
+mod external_indexer_routing;
 mod failover;
 mod graceful_shutdown;
 mod header_forwarding;