diff --git a/.github/workflows/pr-test-sgl-router.yml b/.github/workflows/pr-test-sgl-router.yml index 585ff2749..45f17cac0 100644 --- a/.github/workflows/pr-test-sgl-router.yml +++ b/.github/workflows/pr-test-sgl-router.yml @@ -108,6 +108,9 @@ jobs: steps: - uses: actions/checkout@v4 + - name: Install protoc + run: sudo bash scripts/ci/utils/install_protoc.sh + - name: Verify dynamo-* SHA pin run: | set -euo pipefail @@ -160,17 +163,17 @@ jobs: - name: cargo check working-directory: experimental/sgl-router - run: cargo check --all-targets + run: cargo check --workspace --all-targets - name: cargo clippy working-directory: experimental/sgl-router - run: cargo clippy --all-targets -- -D warnings + run: cargo clippy --workspace --all-targets -- -D warnings - name: cargo fmt working-directory: experimental/sgl-router run: | rustup toolchain install nightly --profile minimal --component rustfmt - cargo +nightly fmt -- --check + cargo +nightly fmt --all -- --check # cargo-deny 0.19.6+ required for CVSS 4.0 parsing. - name: Install cargo-deny @@ -196,6 +199,9 @@ jobs: steps: - uses: actions/checkout@v4 + - name: Install protoc + run: sudo bash scripts/ci/utils/install_protoc.sh + - name: Install OS deps (libssl, pkg-config) run: | sudo apt-get update @@ -246,7 +252,7 @@ jobs: # the Qwen3-0.6B tokenizer.json. Filter is matched against the # full test path `tokenizer::parity::parity_matrix`; substring # `parity_matrix` is unique to that test. - run: cargo test --release -- --skip parity_matrix + run: cargo test --release --workspace -- --skip parity_matrix # Regenerate the cross-impl block-hash parity fixture and fail if it # differs from the committed file. The Python script replicates @@ -366,7 +372,7 @@ jobs: working-directory: experimental/sgl-router run: | source "$HOME/.cargo/env" - cargo build --release + cargo build --release --workspace # Install SGLang from the local checkout in editable mode (no PyPI # version pin) — mirrors `pr-test-rust.yml` so the router e2e runs diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c2a339279..f7a14d03e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -140,7 +140,7 @@ repos: pass_filenames: false - id: rustfmt-sgl-router name: rustfmt experimental/sgl-router - entry: bash -c 'cd experimental/sgl-router && cargo fmt -- --check' + entry: bash -c 'cd experimental/sgl-router && cargo fmt --all -- --check' language: system files: ^experimental/sgl-router/.*\.rs$ pass_filenames: false diff --git a/docker/sgl-router.Dockerfile b/docker/sgl-router.Dockerfile index c8b3c0d26..67471a52c 100644 --- a/docker/sgl-router.Dockerfile +++ b/docker/sgl-router.Dockerfile @@ -34,21 +34,28 @@ ARG DEBIAN_VERSION=bookworm ######################## STAGE 1 — chef recipe ########################## FROM rust:${RUST_VERSION}-${DEBIAN_VERSION} AS chef RUN cargo install cargo-chef --locked --version ^0.1 -WORKDIR /work +WORKDIR /work/sgl-router COPY experimental/sgl-router/Cargo.toml ./ -COPY experimental/sgl-router/rust-toolchain.toml ./ +COPY experimental/sgl-router/sgl-kv-indexer/Cargo.toml sgl-kv-indexer/Cargo.toml # Stub a minimal src tree so cargo can resolve the workspace, generate # the lockfile (gitignored upstream), then prepare the chef recipe. -RUN mkdir -p src && echo "fn main() {}" > src/main.rs \ +RUN mkdir -p src sgl-kv-indexer/src/bin \ + && echo "fn main() {}" > src/main.rs \ && echo "" > src/lib.rs \ + && echo "" > sgl-kv-indexer/src/lib.rs \ + && echo "fn main() {}" > sgl-kv-indexer/src/bin/kv-indexer-server.rs \ + && echo "fn main() {}" > sgl-kv-indexer/src/bin/kv-indexer-bridge.rs \ && cargo generate-lockfile \ && cargo chef prepare --recipe-path recipe.json \ - && rm -rf src + && rm -rf src sgl-kv-indexer/src ######################## STAGE 2 — builder ############################## FROM rust:${RUST_VERSION}-${DEBIAN_VERSION} AS builder -RUN cargo install cargo-chef --locked --version ^0.1 -WORKDIR /work +RUN apt-get update \ + && apt-get install -y --no-install-recommends protobuf-compiler \ + && rm -rf /var/lib/apt/lists/* \ + && cargo install cargo-chef --locked --version ^0.1 +WORKDIR /work/sgl-router # `dynamo-tokenizers` pulls in `pcre2-sys`, whose build.rs links the SYSTEM # libpcre2-8 whenever pkg-config finds it (it does here — the rust:bookworm @@ -58,29 +65,33 @@ WORKDIR /work # it statically, keeping the runtime self-contained. ENV PCRE2_SYS_STATIC=1 -COPY --from=chef /work/recipe.json ./recipe.json -COPY --from=chef /work/Cargo.lock ./Cargo.lock -COPY experimental/sgl-router/rust-toolchain.toml ./ +COPY --from=chef /work/sgl-router/recipe.json ./recipe.json +COPY --from=chef /work/sgl-router/Cargo.lock ./Cargo.lock +COPY experimental/sgl-router/sgl-kv-indexer/Cargo.toml sgl-kv-indexer/Cargo.toml -# Cook (compile + cache) the dep graph from the recipe. This layer's -# inputs are recipe.json + the toolchain — code changes in src/ do NOT -# invalidate it. +# Cook (compile + cache) the dep graph from the recipe. The recipe carries every +# workspace member's manifest, so chef recreates the Indexer's source stubs itself. RUN cargo chef cook --release --recipe-path recipe.json # Now bring in the real sources and the manifest they need. COPY experimental/sgl-router/Cargo.toml ./ COPY experimental/sgl-router/src ./src +COPY experimental/sgl-router/sgl-kv-indexer/Cargo.toml sgl-kv-indexer/Cargo.toml +COPY experimental/sgl-router/sgl-kv-indexer/build.rs sgl-kv-indexer/build.rs +COPY experimental/sgl-router/sgl-kv-indexer/proto sgl-kv-indexer/proto +COPY experimental/sgl-router/sgl-kv-indexer/src sgl-kv-indexer/src # --locked is intentionally omitted: the lockfile is generated in-container # (gitignored upstream) and `cargo chef cook` may have mutated it during the # dep-cook step, so a strict --locked check would spuriously fail. -RUN cargo build --release --bin sgl-router \ +RUN touch sgl-kv-indexer/build.rs \ + && cargo build --release --bin sgl-router \ && strip target/release/sgl-router ######################## STAGE 3 — runtime ############################## FROM gcr.io/distroless/cc-debian12:nonroot AS runtime -COPY --from=builder /work/target/release/sgl-router /usr/local/bin/sgl-router +COPY --from=builder /work/sgl-router/target/release/sgl-router /usr/local/bin/sgl-router # Default config path; mount your own via `-v :/etc/sgl-router`. ENV SGL_ROUTER_CONFIG=/etc/sgl-router/sgl-router.yaml diff --git a/experimental/sgl-router/Cargo.toml b/experimental/sgl-router/Cargo.toml index 80416b29d..74906599f 100644 --- a/experimental/sgl-router/Cargo.toml +++ b/experimental/sgl-router/Cargo.toml @@ -1,6 +1,6 @@ [workspace] resolver = "2" -members = ["."] +members = [".", "sgl-kv-indexer"] [package] name = "sgl-router" @@ -33,6 +33,7 @@ axum = { version = "0.8", features = ["macros", "tracing"] } tower = { version = "0.5", features = ["full"] } tower-http = { version = "0.6", features = ["trace", "compression-gzip", "cors", "timeout", "request-id"] } reqwest = { version = "0.12", features = ["stream", "json", "rustls-tls"], default-features = false } +sgl-kv-indexer = { path = "sgl-kv-indexer" } # Serialization serde = { version = "1", features = ["derive"] } @@ -62,7 +63,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] } futures = "0.3" bytes = "1" rand = "0.8" -tokio-stream = "0.1" +tokio-stream = { version = "0.1", features = ["net"] } dashmap = "6" kube = { version = "0.96", features = ["runtime", "derive"] } k8s-openapi = { version = "0.23", features = ["v1_31"] } @@ -86,6 +87,7 @@ serde_json = "1" tempfile = "3" tower = { version = "0.5", features = ["util"] } tokio = { version = "1.42", features = ["test-util"] } +tonic = { version = "0.14.6", features = ["transport"] } # Low-level msgpack encoder used to hand-construct wire bytes in # kv_events golden-bytes tests (decode-only path uses rmp-serde). rmp = "0.8" diff --git a/experimental/sgl-router/README.md b/experimental/sgl-router/README.md index fa1e05b01..6b6a661e0 100644 --- a/experimental/sgl-router/README.md +++ b/experimental/sgl-router/README.md @@ -50,6 +50,26 @@ Omit `--service-discovery-namespace` to watch all namespaces (requires cluster-wide RBAC). For prefill/decode disaggregation, replace `--selector` with `--prefill-selector` and `--decode-selector`. +External KV indexer as the cache-aware signal source: + +```bash +sgl-router \ + --model-id qwen3 \ + --tokenizer-path /models/qwen3/tokenizer.json \ + --worker-urls http://10.0.0.1:30000 http://10.0.0.2:30000 \ + --policy cache_aware_zmq \ + --kv-indexer-endpoint http://10.0.0.10:50051 \ + --kv-indexer-query-timeout-ms 100 \ + --kv-indexer-query-max-inflight 32 +``` + +The existing cache-aware policy and thresholds are reused. When configured, the +Indexer replaces the Router-local radix tree as the cache signal. A successful +query with no usable match selects by minimum active load; connection failures, +timeouts, local admission rejection, and server rejection fail the Router +request with `503` rather than silently switching signals. The timeout and local +concurrency bound default to 100ms and 32 respectively. + ## License Apache-2.0. diff --git a/experimental/sgl-router/sgl-kv-indexer/.gitignore b/experimental/sgl-router/sgl-kv-indexer/.gitignore new file mode 100644 index 000000000..99e6c6109 --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/.gitignore @@ -0,0 +1,3 @@ +target/ +*.rs.bk +.DS_Store diff --git a/experimental/sgl-router/sgl-kv-indexer/Cargo.toml b/experimental/sgl-router/sgl-kv-indexer/Cargo.toml new file mode 100644 index 000000000..4ba5cd421 --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/Cargo.toml @@ -0,0 +1,35 @@ +[package] +name = "sgl-kv-indexer" +version = "0.1.0" +edition = "2021" +license = "Apache-2.0" +description = "In-memory KV cache indexer with a SGLang KV-event bridge" +publish = false # experimental crate; not published to crates.io + +[lib] +name = "sgl_kv_indexer" +path = "src/lib.rs" + +[[bin]] +name = "kv-indexer-server" +path = "src/bin/kv-indexer-server.rs" + +[[bin]] +name = "kv-indexer-bridge" +path = "src/bin/kv-indexer-bridge.rs" + +[dependencies] +bytes = "1" +prost = "0.14" +rmpv = "1.3.1" +tokio = { version = "1", features = ["macros", "rt-multi-thread", "signal", "sync", "time"] } +tonic = { version = "0.14.6", features = ["transport"] } +tonic-prost = "0.14.6" +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] } +# Must mirror the router's feature set: Cargo unifies features across the +# workspace, so enabling more transports here also enlarges the router binary. +zeromq = { version = "0.6", default-features = false, features = ["tokio-runtime", "tcp-transport"] } + +[build-dependencies] +tonic-prost-build = "0.14.6" diff --git a/experimental/sgl-router/sgl-kv-indexer/README.md b/experimental/sgl-router/sgl-kv-indexer/README.md new file mode 100644 index 000000000..a62653770 --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/README.md @@ -0,0 +1,256 @@ +# SGL KV Indexer (in-memory build) + +`sgl-kv-indexer` is an experimental metadata service for SGLang KV-cache +blocks. It records which worker and storage tier currently holds each +content-addressed block, allowing a router to query likely cache hits without +moving KV data itself. + +This build deliberately uses one process-local in-memory index. It has no +external storage dependency, but it is soft-state: restarting the Indexer loses +all placement metadata. + +## Architecture + +```text +SGLang worker ── ZMQ PUB ──> bridge ── gRPC ──> in-memory indexer +``` + +- SGLang publishes `BlockStored`, `BlockRemoved`, and `AllBlocksCleared` events. +- One `kv-indexer-bridge` follows each independent worker/rank event stream. +- One `kv-indexer-server` applies event batches and serves match queries. +- Placement, worker metadata, reverse holdings, and hit counters live in that + server process behind a single read/write lock. + +Each apply RPC is ordered and atomic within the process, and a query sees a +consistent snapshot. The bridge splits larger worker event batches into +ordered RPCs of at most 16,384 hashes and 256 actions while keeping per-hash +metadata aligned. If a later RPC fails, earlier chunks may already be applied; +there is no rollback or replay. There is no persistence, replication, or state +sharing between Indexer servers. + +## Operational contract + +Run exactly one Indexer server for a deployment. Multiple bridge processes and +workers may report to it, but active-active Indexer servers have independent +state and must not be treated as replicas. + +This build has no sequence gate, incarnation fencing, replay recovery, worker +liveness TTL, or restart recovery: + +- An Indexer restart starts with an empty index. +- A worker death is not detected; its last placements remain until revoked or + until the Indexer restarts. +- Events published while a bridge is disconnected are not replayed. +- A publisher sequence gap is logged and otherwise ignored. +- Redelivered or reordered batches are applied again in arrival order. + +Individual report, revoke, and clear mutations are idempotent. A future +Snapshot plus event-replay mechanism is required before production high +availability can rebuild state safely after restart or event loss. + +## In-memory data model + +The server stores: + +- block hash → token count and `(worker, tier) → component mask` +- worker → router-facing address, `WorkerCacheSpec`, and reverse holdings by tier +- block hash → cumulative hit count + +Component masks are `FULL=1`, `SWA=2`, `MAMBA=4`, and `0` for legacy +whole-block events. `REPORT` replaces the component snapshot for one +`(worker, tier, block)` placement. `REVOKE` removes that placement, and +`CLEAR_ALL_AT_TIER` removes every placement for the worker at that tier. + +## Build + +```bash +cd experimental/sgl-router +cargo build --release -p sgl-kv-indexer +``` + +This produces: + +- `target/release/kv-indexer-server` +- `target/release/kv-indexer-bridge` + +## End-to-end quickstart + +Component-aware routing requires an SGLang engine build that supports +`component_types`. Start each command in its own terminal. + +1. Start the single Indexer server: + +```bash +KV_INDEXER_LISTEN_ADDR=127.0.0.1:50051 \ + cargo run --release --bin kv-indexer-server +``` + +`KV_INDEXER_LISTEN_ADDR` defaults to `[::1]:50051`. +`KV_INDEXER_PREFIX_QUERY_MAX_INFLIGHT` sets the maximum number of prefix +queries executing concurrently and defaults to `32`. Requests above the limit +are rejected immediately with gRPC `RESOURCE_EXHAUSTED`. +There is no backend or storage configuration. + +2. Start one bridge per worker event stream. This FULL+SWA example uses the +worker URL registered with the Router: + +```bash +KV_INDEXER_WORKER_ID=worker-0 \ +KV_INDEXER_WORKER_ADDRESS=http://127.0.0.1:30000 \ +KV_INDEXER_ENDPOINT=http://127.0.0.1:50051 \ +SGLANG_KV_EVENT_ENDPOINT=tcp://127.0.0.1:5567 \ +SGLANG_KV_EVENT_TOPIC=kv-events \ +KV_INDEXER_CACHE_COMPONENTS=full,swa \ +KV_INDEXER_SWA_WINDOW_TOKENS= \ +KV_INDEXER_FULL_TIERS=HBM \ +KV_INDEXER_SWA_TIERS=HBM \ + cargo run --release --bin kv-indexer-bridge +``` + +For FULL+MAMBA, use: + +```bash +KV_INDEXER_CACHE_COMPONENTS=full,mamba +KV_INDEXER_FULL_TIERS=HBM +KV_INDEXER_MAMBA_TIERS=HBM +``` + +3. Start the matching SGLang worker: + +```bash +python -m sglang.launch_server \ + --model-path \ + --port 30000 \ + --kv-events-config \ + '{"publisher":"zmq","endpoint":"tcp://*:5567","topic":"kv-events"}' \ + --enable-kv-events-component-types +``` + +4. Start the Router with the Indexer as the authoritative cache signal: + +```bash +sgl-router \ + --model-id \ + --tokenizer-path \ + --worker-urls http://127.0.0.1:30000 \ + --policy cache_aware_zmq \ + --kv-indexer-endpoint http://127.0.0.1:50051 \ + --kv-indexer-query-timeout-ms 100 \ + --kv-indexer-query-max-inflight 32 +``` + +For multiple workers, repeat steps 2–3 with unique worker IDs and ports. +`KV_INDEXER_WORKER_ADDRESS` must exactly match the corresponding Router URL. + +The bridge sends its `WorkerCacheSpec` with every batch. Omitting +`KV_INDEXER_CACHE_COMPONENTS` clears any previously stored spec and uses legacy +whole-block matching. + +## API + +The protobuf service in `proto/kv_indexer.proto` provides: + +- `ApplyExternalKvBatch`: ordered placement reports, revocations, and clears. + The request `seq` is carried for observability only. +- `MatchExternalKv`: workers and tiers holding requested block hashes. +- `MatchExternalKvPrefix`: per-worker longest contiguous reusable prefix. +- `GetExternalKvHitCounts`: per-block hit counters. + +There is no gRPC health service in this build. + +## Prefix routing semantics + +For a legacy worker, `matched_prefix_blocks` is the largest `n` such that it +holds every block in `hashes[0..n)` without a gap. For a component-aware worker: + +- FULL must be contiguous on every matched block. +- SWA must cover the trailing `swa_window_tokens` at the candidate boundary, or + form an unbroken run from the prompt head. +- MAMBA must be present on the candidate boundary block. + +Component placements without a worker spec fail closed. Workers with an empty +router-facing address are excluded. + +The Indexer returns every candidate sorted by prefix length; it does not choose +a worker. When configured, it replaces the Router's local radix tree as the +cache signal: the Router intersects Indexer results with its healthy candidates, +and a successful query with no usable match selects by minimum active load. +Indexer connection failures, timeouts, overload, and a prompt too long to fit one +gRPC message fall back to that same minimum-active-load selection, so an +unreachable Indexer costs cache affinity rather than availability; a rejected RPC +still fails the Router request with `503`, because it means the two sides +disagree on the request contract. An +endpoint the Router could never dial is rejected at startup instead of failing +every query later. The local radix tree is used only when no Indexer endpoint is +configured. The per-query deadline defaults to 100ms and can be changed with +`--kv-indexer-query-timeout-ms`. The Router-side admission bound defaults to 32 +concurrent calls and can be changed with `--kv-indexer-query-max-inflight`. + +Prefix queries carry no Indexer-imposed block cap beyond the caller's +`max_blocks` ceiling — unlike applies and `MatchExternalKv`, which reject above +16,384 hashes. The in-memory backend scans the request in one pass over a single +consistent snapshot, holding O(1) matching state per candidate worker and +considering only workers that hold the first block, so request length costs time +but not memory. Block hashes use packed `sfixed64` encoding, and the server +accepts decoded gRPC messages up to 8 MiB (roughly one million hashes). +Server work is bounded by `max_blocks` when the caller supplies one and by that +transport limit. The Router's per-query deadline bounds how long it waits for an +answer, but does not cancel a synchronous scan already in progress. A first-block +miss returns immediately with `blocks_read=1`. + +Message decoding happens before a request reaches the service, so +`KV_INDEXER_PREFIX_QUERY_MAX_INFLIGHT` bounds the scan but not the bytes a peer +makes the server buffer. That is bounded instead by the HTTP/2 stream limit: each +connection is capped at 64 concurrent streams, bounding that connection to +64 × 8 MiB of undecoded requests. For a query past the 8 MiB ceiling, the Router +sends only the leading hashes that fit and still divides the returned prefix by +the full request's block count. This preserves a useful lower-bound cache signal +without overstating the match rate. If an Indexer has a lower ceiling and returns +gRPC `OUT_OF_RANGE`, the Router falls back to minimum active load. + +Long scans hold the read lock throughout. Operators serving very long prompts +should set `max_blocks` instead of relying on the message-size limit. + +## Overload behavior and observability + +The Router and server apply separate admission bounds. The Router rejects a +query locally when its `--kv-indexer-query-max-inflight` permits are exhausted; +the server returns gRPC `RESOURCE_EXHAUSTED` when +`KV_INDEXER_PREFIX_QUERY_MAX_INFLIGHT` is exhausted. Both leave the request +routed by minimum active load, logged at `WARN` on the Router. + +Every Router query publishes its timeout through the gRPC `grpc-timeout` header. +The server timestamps arrival and returns `DEADLINE_EXCEEDED` before backend work +when queueing has already consumed that budget. Apply/event RPCs are never shed, +because dropping one would permanently diverge the soft-state index. + +Deadline shedding is logged at `INFO`; server admission rejection is logged at +`WARN`. Each rejection class reports totals 1, 2, 4, 8, and so on, making the +first overload visible at the default log level without log volume growing +linearly with sustained overload. + +## Bridge configuration + +Required or commonly used bridge variables: + +- `KV_INDEXER_WORKER_ID`: unique ID for the worker event stream +- `KV_INDEXER_WORKER_ADDRESS`: Router-facing worker URL +- `KV_INDEXER_ENDPOINT`: Indexer endpoint, default `http://[::1]:50051` +- `SGLANG_KV_EVENT_ENDPOINT`: worker PUB endpoint +- `SGLANG_KV_EVENT_TOPIC`: ZMQ subscription topic +- `KV_INDEXER_CLEAR_TIERS`: tiers affected by clear, default `HBM,DRAM,SSD` +- `KV_INDEXER_CACHE_COMPONENTS`: optional `full,swa` or `full,mamba` +- `KV_INDEXER_SWA_WINDOW_TOKENS`: required when SWA is configured +- `KV_INDEXER_FULL_TIERS`, `KV_INDEXER_SWA_TIERS`, + `KV_INDEXER_MAMBA_TIERS`: servable component tiers +- `KV_INDEXER_CACHE_SPEC_VERSION`: component-rule version, default `1` + +## Tests + +No external service is needed: + +```bash +cargo fmt --all -- --check +cargo clippy --all-targets -- -D warnings +cargo test +``` diff --git a/experimental/sgl-router/sgl-kv-indexer/build.rs b/experimental/sgl-router/sgl-kv-indexer/build.rs new file mode 100644 index 000000000..cc37b84f4 --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/build.rs @@ -0,0 +1,10 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +fn main() -> Result<(), Box> { + tonic_prost_build::configure() + .build_client(true) + .build_server(true) + .compile_protos(&["proto/kv_indexer.proto"], &["proto"])?; + Ok(()) +} diff --git a/experimental/sgl-router/sgl-kv-indexer/proto/kv_indexer.proto b/experimental/sgl-router/sgl-kv-indexer/proto/kv_indexer.proto new file mode 100644 index 000000000..f9a7cbf1b --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/proto/kv_indexer.proto @@ -0,0 +1,185 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +syntax = "proto3"; + +package kv_indexer.v1; + +// Storage tier for externally managed KV cache blocks. +enum TierType { + TIER_UNKNOWN = 0; + TIER_HBM = 1; + TIER_DRAM = 2; + TIER_SSD = 3; +} + +// KV components are a fixed set; each component's rule is a property of its +// type, not a per-worker choice: +// FULL (bit 0, value 1) -- resident on every prefix block. +// SWA (bit 1, value 2) -- resident contiguously over at least +// `swa_window_tokens` tokens ending at the boundary, +// or in an unbroken run from the head. +// MAMBA (bit 2, value 4) -- resident on the boundary block only. +// Carried as a bitmask; FULL is always present on a stored block. + +// A worker's versioned cache spec: which components its cache holds and the +// parameters the fixed rules need. Absent means legacy / full-only behaviour. +message WorkerCacheSpec { + uint32 version = 1; + // Bitmask of components this worker's cache holds (FULL|SWA|MAMBA); selects + // which fixed rules gate the prefix. + uint32 components = 2; + // Sliding-window size in tokens; used only when SWA is present. + uint32 swa_window_tokens = 3; + // Per-component servable tiers, each a bitmask of (1 << TierType), intersected + // with the indexer's servable set (V1: HBM+DRAM) at query time. + uint32 full_tier_mask = 4; + uint32 swa_tier_mask = 5; + uint32 mamba_tier_mask = 6; +} + +message MatchExternalKvRequest { + // Block-level hashes matched against placement metadata. This RPC returns + // placement matches only; it does not compute the longest reusable prefix. + repeated sfixed64 hashes = 1 [packed = true]; + + // When true, only hashes that are actually matched should increase hit + // counters. Diagnostic callers should leave this false. + bool count_as_hit = 2; +} + +// The kind of mutation a single ExternalKvAction carries. A whole SGLang +// KVEventBatch is applied in one call while preserving exact action order. +enum ExternalKvActionType { + ACTION_UNKNOWN = 0; + ACTION_REPORT = 1; + ACTION_REVOKE = 2; + ACTION_CLEAR_ALL_AT_TIER = 3; +} + +message ExternalKvAction { + ExternalKvActionType type = 1; + TierType tier = 2; + // Non-empty for REPORT/REVOKE; ignored for CLEAR_ALL_AT_TIER. + repeated sfixed64 hashes = 3 [packed = true]; + + // REPORT only. Per-hash resident component bitmask at this tier, a REPLACE + // snapshot index-aligned with `hashes`. Empty means every hash is a legacy + // whole-block store (no component detail). REVOKE/CLEAR ignore it. + repeated uint32 component_masks = 4; + + // REPORT only. Per-hash token count (block_size), index-aligned with `hashes`, + // used to accumulate SWA trailing windows. Empty when not supplied (legacy). + repeated uint32 block_sizes = 5; +} + +message ApplyExternalKvBatchRequest { + string worker_id = 1; + // The SGLang batch sequence number, monotonic per worker. Observability only: + // every batch is applied, with no deduplication, fencing, or checkpointing. + uint64 seq = 2; + repeated ExternalKvAction actions = 3; + // The worker's KV-transfer address, used to populate MatchExternalKvResponse. + // Supplied by the bridge; may be empty, and then matches carry no address. + string worker_address = 4; + + // Field 5 held the worker incarnation token used for restart fencing. + reserved 5; + + // The worker's component cache spec, carried on every batch so it self-heals + // across reconnects and indexer restarts. Absent means legacy / full-only. + WorkerCacheSpec cache_spec = 6; +} + +// Deliberately empty: applies are unconditional, so there is nothing to report +// back. Kept as a message so the RPC signature stays stable. +message ApplyExternalKvBatchResponse { + // Fields 1-3 held the durable seq checkpoint and duplicate flag. + reserved 1, 2, 3; +} + +message TierHashes { + TierType tier = 1; + repeated sfixed64 hashes = 2 [packed = true]; + // Diagnostic snapshots aligned with `hashes`. Empty only for backends that + // cannot expose component detail. + repeated uint32 component_masks = 3; + repeated uint32 block_sizes = 4; +} + +message ExternalKvNodeMatch { + string worker_id = 1; + string address = 2; + repeated TierHashes hashes_by_tier = 3; +} + +message MatchExternalKvResponse { + repeated ExternalKvNodeMatch matches = 1; +} + +message MatchExternalKvPrefixRequest { + // Block hashes in prompt order. hashes[0] MUST be the request's first block: + // prefix matching starts there and stops at the first block a worker is + // missing, so a misordered list silently truncates every match. + repeated sfixed64 hashes = 1 [packed = true]; + + // Caller-supplied ceiling on how many leading blocks to consider. 0 means no + // caller ceiling: the server considers every block sent. Unlike the mutating + // RPCs, a prefix query is not rejected for length; its scan holds O(1) state + // per candidate worker, so length costs time rather than memory. + uint32 max_blocks = 2; +} + +message ExternalKvPrefixMatch { + // The worker's router-facing routing identity, NOT its KV-transfer address. + // The router intersects this byte-for-byte with the worker URLs it registered, + // so a mismatch silently disables cache-aware routing. Workers with an empty + // address are excluded from this response entirely. + string worker_address = 1; + + // Largest n such that this worker holds hashes[0..n) contiguously. + uint32 matched_prefix_blocks = 2; + + // Opaque worker id, carried for the caller's logs only; not a routing key. + string worker_id = 3; +} + +message MatchExternalKvPrefixResponse { + // Matches sorted by matched_prefix_blocks, descending. + repeated ExternalKvPrefixMatch matches = 1; + + // The longest contiguous prefix held by any single worker (0 when no match). + uint32 best_prefix_blocks = 2; + + // How many blocks the server actually read placement for, so early termination + // and truncation are observable. NOT part of the prefix semantics: backends may + // report different values for the same matches. + uint32 blocks_read = 3; +} + +message HitCountEntry { + sfixed64 hash = 1; + uint64 hit_count_total = 2; +} + +message GetExternalKvHitCountsRequest { + repeated sfixed64 hashes = 1 [packed = true]; +} + +message GetExternalKvHitCountsResponse { + repeated HitCountEntry entries = 1; +} + +service KVIndexer { + // The sole mutation API: applies an ordered SGLang KVEventBatch. + rpc ApplyExternalKvBatch(ApplyExternalKvBatchRequest) returns (ApplyExternalKvBatchResponse); + + rpc MatchExternalKv(MatchExternalKvRequest) returns (MatchExternalKvResponse); + + // Answers, per worker, how long a contiguous prefix of the request it holds. + // The indexer never picks a worker: it cannot see the router's health, load, or + // pool split, so the final choice stays with the router. + rpc MatchExternalKvPrefix(MatchExternalKvPrefixRequest) returns (MatchExternalKvPrefixResponse); + + rpc GetExternalKvHitCounts(GetExternalKvHitCountsRequest) returns (GetExternalKvHitCountsResponse); +} diff --git a/experimental/sgl-router/sgl-kv-indexer/src/admission.rs b/experimental/sgl-router/sgl-kv-indexer/src/admission.rs new file mode 100644 index 000000000..e5c896a00 --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/src/admission.rs @@ -0,0 +1,170 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Deadline-based load shedding for the query path. +//! +//! A query that already waited out its caller's whole deadline can no longer be +//! answered usefully, so serving it only delays the rest of the backlog. The +//! budget is the caller's own `grpc-timeout`: no server-side threshold to tune, +//! and a caller that declared no deadline is never shed. + +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, Instant}; + +use tonic::metadata::MetadataMap; +use tonic::{Extensions, Request, Status}; + +/// When the request's headers were read off the connection. +#[derive(Clone, Copy)] +struct Arrival(Instant); + +/// Counts rejections of one kind and reports on doubling totals, so the first +/// rejection is visible immediately and a sustained overload cannot flood the log. +pub(crate) struct RejectionLog(AtomicU64); + +impl RejectionLog { + pub(crate) const fn new() -> Self { + Self(AtomicU64::new(0)) + } + + /// Records one rejection, returning the running total when it should be + /// logged and `None` when it should be absorbed. + pub(crate) fn record(&self) -> Option { + let total = self.0.fetch_add(1, Ordering::Relaxed) + 1; + total.is_power_of_two().then_some(total) + } +} + +static DEADLINE_SHED_LOG: RejectionLog = RejectionLog::new(); + +/// Timestamps a request's arrival so the query path can measure how long it then +/// waited. Runs before the per-request task is spawned, so the stamp precedes +/// any scheduling delay. Without this interceptor nothing is ever shed. +pub fn stamp_arrival(mut request: Request<()>) -> Result, Status> { + request.extensions_mut().insert(Arrival(Instant::now())); + Ok(request) +} + +/// Rejects a query that spent its caller's entire deadline waiting to be served. +/// +/// Never applied to the apply path: dropping a KV event would leave the index +/// permanently diverged from the worker that reported it. +pub(crate) fn reject_if_deadline_passed( + metadata: &MetadataMap, + extensions: &Extensions, +) -> Result<(), Status> { + let (Some(arrival), Some(budget)) = (extensions.get::(), caller_deadline(metadata)) + else { + return Ok(()); + }; + let waited = arrival.0.elapsed(); + if waited < budget { + return Ok(()); + } + if let Some(shed_total) = DEADLINE_SHED_LOG.record() { + tracing::info!( + shed_total, + waited_ms = waited.as_millis(), + budget_ms = budget.as_millis(), + "shedding prefix query whose caller deadline already passed" + ); + } + Err(Status::deadline_exceeded( + "prefix query waited longer than its caller deadline", + )) +} + +/// The budget the caller declared in `grpc-timeout`, per the gRPC wire spec (up +/// to 8 digits followed by a unit). `None` for an absent or unparsable value, +/// which leaves the request unshed. Measured against the wait since arrival, so +/// transit time is ignored and shedding can only be late, never early. +fn caller_deadline(metadata: &MetadataMap) -> Option { + let raw = metadata.get("grpc-timeout")?.to_str().ok()?; + let unit = *raw.as_bytes().last()?; + let value: u64 = raw.get(..raw.len() - 1)?.parse().ok()?; + match unit { + b'H' => value.checked_mul(60 * 60).map(Duration::from_secs), + b'M' => value.checked_mul(60).map(Duration::from_secs), + b'S' => Some(Duration::from_secs(value)), + b'm' => Some(Duration::from_millis(value)), + b'u' => Some(Duration::from_micros(value)), + b'n' => Some(Duration::from_nanos(value)), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn metadata(timeout: Option<&str>) -> MetadataMap { + let mut metadata = MetadataMap::new(); + if let Some(timeout) = timeout { + metadata.insert("grpc-timeout", timeout.parse().unwrap()); + } + metadata + } + + fn extensions(arrival: Option) -> Extensions { + let mut extensions = Extensions::new(); + if let Some(arrival) = arrival { + extensions.insert(Arrival(arrival)); + } + extensions + } + + #[test] + fn parses_every_wire_unit() { + for (raw, expected) in [ + ("1H", Duration::from_secs(3600)), + ("2M", Duration::from_secs(120)), + ("3S", Duration::from_secs(3)), + ("100m", Duration::from_millis(100)), + ("250u", Duration::from_micros(250)), + ("400n", Duration::from_nanos(400)), + ] { + assert_eq!(caller_deadline(&metadata(Some(raw))), Some(expected)); + } + } + + #[test] + fn unparsable_deadline_is_ignored() { + for raw in ["", "m", "100", "100x", "abcm", "99999999999999999999H"] { + assert_eq!(caller_deadline(&metadata(Some(raw))), None, "{raw}"); + } + assert_eq!(caller_deadline(&metadata(None)), None); + } + + #[test] + fn sheds_only_once_the_caller_deadline_has_passed() { + let now = Instant::now(); + let waited_past = now - Duration::from_millis(150); + let within = now - Duration::from_millis(10); + + assert_eq!( + reject_if_deadline_passed(&metadata(Some("100m")), &extensions(Some(waited_past))) + .unwrap_err() + .code(), + tonic::Code::DeadlineExceeded + ); + assert!( + reject_if_deadline_passed(&metadata(Some("100m")), &extensions(Some(within))).is_ok() + ); + } + + #[test] + fn rejections_are_reported_on_doubling_counts() { + let log = RejectionLog::new(); + let reported: Vec = (0..16).filter_map(|_| log.record()).collect(); + assert_eq!(reported, vec![1, 2, 4, 8, 16]); + } + + #[test] + fn missing_deadline_or_arrival_never_sheds() { + let long_wait = Instant::now() - Duration::from_secs(60); + // A caller that declared no deadline keeps the pre-existing behaviour. + assert!(reject_if_deadline_passed(&metadata(None), &extensions(Some(long_wait))).is_ok()); + // No interceptor installed: nothing to measure the wait against. + assert!(reject_if_deadline_passed(&metadata(Some("1m")), &extensions(None)).is_ok()); + } +} diff --git a/experimental/sgl-router/sgl-kv-indexer/src/bin/kv-indexer-bridge.rs b/experimental/sgl-router/sgl-kv-indexer/src/bin/kv-indexer-bridge.rs new file mode 100644 index 000000000..da09104af --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/src/bin/kv-indexer-bridge.rs @@ -0,0 +1,18 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use sgl_kv_indexer::bridge::{run_bridge_until, BridgeConfig}; +use sgl_kv_indexer::shutdown_signal; + +#[tokio::main] +async fn main() -> Result<(), Box> { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()), + ) + .init(); + + let config = BridgeConfig::from_env()?; + run_bridge_until(config, shutdown_signal()).await?; + Ok(()) +} diff --git a/experimental/sgl-router/sgl-kv-indexer/src/bin/kv-indexer-server.rs b/experimental/sgl-router/sgl-kv-indexer/src/bin/kv-indexer-server.rs new file mode 100644 index 000000000..0d150c561 --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/src/bin/kv-indexer-server.rs @@ -0,0 +1,94 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use std::net::SocketAddr; +use std::sync::Arc; +use std::{env, io}; + +use sgl_kv_indexer::{ + server_builder, shutdown_signal, stamp_arrival, InMemoryKvIndexerBackend, KvIndexerBackend, + KvIndexerService, DEFAULT_PREFIX_QUERY_MAX_INFLIGHT, MAX_CONCURRENT_STREAMS, +}; +use tonic::service::interceptor::InterceptedService; +use tracing::info; + +const PREFIX_QUERY_MAX_INFLIGHT_ENV: &str = "KV_INDEXER_PREFIX_QUERY_MAX_INFLIGHT"; + +#[tokio::main] +async fn main() -> Result<(), Box> { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()), + ) + .init(); + + let addr = std::env::var("KV_INDEXER_LISTEN_ADDR") + .unwrap_or_else(|_| "[::1]:50051".to_string()) + .parse::()?; + let prefix_query_max_inflight = prefix_query_max_inflight_from_env()?; + + let backend: Arc = Arc::new(InMemoryKvIndexerBackend::new()); + // The interceptor timestamps each request before its own task is queued, + // which is what lets the query path shed work whose deadline expired. + let service = InterceptedService::new( + KvIndexerService::with_prefix_query_max_inflight(backend, prefix_query_max_inflight) + .into_server(), + stamp_arrival, + ); + + info!( + %addr, + prefix_query_max_inflight, + max_concurrent_streams = MAX_CONCURRENT_STREAMS, + "starting single-server in-memory SGLang KV Indexer" + ); + server_builder() + .add_service(service) + .serve_with_shutdown(addr, shutdown_signal()) + .await?; + + Ok(()) +} + +fn prefix_query_max_inflight_from_env() -> io::Result { + match env::var(PREFIX_QUERY_MAX_INFLIGHT_ENV) { + Ok(raw) => parse_prefix_query_max_inflight(&raw), + Err(env::VarError::NotPresent) => Ok(DEFAULT_PREFIX_QUERY_MAX_INFLIGHT), + Err(env::VarError::NotUnicode(_)) => Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("{PREFIX_QUERY_MAX_INFLIGHT_ENV} must be valid UTF-8"), + )), + } +} + +fn parse_prefix_query_max_inflight(raw: &str) -> io::Result { + let value = raw.parse::().map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + format!("{PREFIX_QUERY_MAX_INFLIGHT_ENV} must be a positive integer, got {raw:?}"), + ) + })?; + if value == 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("{PREFIX_QUERY_MAX_INFLIGHT_ENV} must be greater than zero"), + )); + } + Ok(value) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_positive_prefix_query_limit() { + assert_eq!(parse_prefix_query_max_inflight("64").unwrap(), 64); + } + + #[test] + fn rejects_invalid_prefix_query_limit() { + assert!(parse_prefix_query_max_inflight("0").is_err()); + assert!(parse_prefix_query_max_inflight("many").is_err()); + } +} diff --git a/experimental/sgl-router/sgl-kv-indexer/src/bridge.rs b/experimental/sgl-router/sgl-kv-indexer/src/bridge.rs new file mode 100644 index 000000000..17eea10a2 --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/src/bridge.rs @@ -0,0 +1,1424 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! SGLang KV event bridge. +//! +//! Subscribes to a worker's ZMQ KV-event stream, decodes each batch, and +//! forwards it to the indexer over gRPC. +//! +//! It keeps a reconnect supervisor but does not recover data: no sequence +//! tracking, no replay of missed batches, no incarnation token, no liveness +//! heartbeat. A sequence gap is logged and ignored, and events produced while +//! the bridge is disconnected are lost. + +use std::io::Cursor; +use std::time::Duration; + +use rmpv::decode::value::read_value; +use rmpv::Value; +use tonic::transport::{Channel, Endpoint}; +use tonic::{Code, Status}; +use tracing::{debug, info, warn}; +use zeromq::{Socket, SocketRecv, SubSocket}; + +use crate::pb::kv_indexer_client::KvIndexerClient; +use crate::pb::{ + ApplyExternalKvBatchRequest, ExternalKvAction, ExternalKvActionType, TierType, WorkerCacheSpec, +}; +use crate::service::{component_bit, COMPONENT_SWA, MAX_ACTIONS_PER_BATCH, MAX_HASHES_PER_REQUEST}; + +/// Backoff bounds for the reconnect supervisor loop. +const RECONNECT_MIN_DELAY: Duration = Duration::from_millis(500); +const RECONNECT_MAX_DELAY: Duration = Duration::from_secs(10); +const GRPC_CONNECT_TIMEOUT: Duration = Duration::from_secs(5); +const GRPC_REQUEST_TIMEOUT: Duration = Duration::from_secs(10); + +#[derive(Debug, Clone)] +pub struct BridgeConfig { + pub worker_id: String, + /// The worker's KV-transfer address, forwarded on every apply batch so the + /// indexer can answer MatchExternalKv with an address. Empty if unset. + pub worker_address: String, + pub event_endpoint: String, + pub event_topic: String, + pub indexer_endpoint: String, + pub clear_tiers: Vec, + /// The worker's component cache spec, forwarded on every apply batch. `None` + /// for a legacy / full-only worker that reports no component metadata. + pub cache_spec: Option, +} + +impl BridgeConfig { + pub fn from_env() -> Result { + let worker_id = std::env::var("KV_INDEXER_WORKER_ID") + .map_err(|_| BridgeError::Config("KV_INDEXER_WORKER_ID is required".to_string()))?; + let worker_address = std::env::var("KV_INDEXER_WORKER_ADDRESS").unwrap_or_default(); + let event_endpoint = std::env::var("SGLANG_KV_EVENT_ENDPOINT") + .unwrap_or_else(|_| "tcp://127.0.0.1:5557".to_string()); + // Match SGLang's upstream ZMQ publisher default. Deployments that use a + // non-empty topic must configure the same value on both sides. + let event_topic = std::env::var("SGLANG_KV_EVENT_TOPIC").unwrap_or_default(); + let indexer_endpoint = std::env::var("KV_INDEXER_ENDPOINT") + .unwrap_or_else(|_| "http://[::1]:50051".to_string()); + let clear_tiers = parse_clear_tiers( + &std::env::var("KV_INDEXER_CLEAR_TIERS").unwrap_or_else(|_| "HBM,DRAM,SSD".to_string()), + )?; + let cache_spec = cache_spec_from_env()?; + + Ok(Self { + worker_id, + worker_address, + event_endpoint, + event_topic, + indexer_endpoint, + clear_tiers, + cache_spec, + }) + } +} + +#[derive(Debug)] +pub enum BridgeError { + Config(String), + Decode(String), + Rpc(tonic::Status), + PermanentRpc(tonic::Status), + Transport(tonic::transport::Error), + Zmq(zeromq::ZmqError), +} + +impl std::fmt::Display for BridgeError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + BridgeError::Config(message) => write!(f, "bridge config error: {message}"), + BridgeError::Decode(message) => write!(f, "bridge decode error: {message}"), + BridgeError::Rpc(status) => write!(f, "indexer rpc error: {status}"), + BridgeError::PermanentRpc(status) => { + write!(f, "permanent indexer rpc error: {status}") + } + BridgeError::Transport(error) => write!(f, "indexer transport error: {error}"), + BridgeError::Zmq(error) => write!(f, "zmq error: {error}"), + } + } +} + +impl std::error::Error for BridgeError {} + +impl BridgeError { + fn is_permanent(&self) -> bool { + matches!(self, BridgeError::Config(_) | BridgeError::PermanentRpc(_)) + } +} + +impl From for BridgeError { + fn from(error: zeromq::ZmqError) -> Self { + BridgeError::Zmq(error) + } +} + +impl From for BridgeError { + fn from(error: tonic::transport::Error) -> Self { + BridgeError::Transport(error) + } +} + +fn classify_rpc(status: Status) -> BridgeError { + match status.code() { + Code::InvalidArgument + | Code::FailedPrecondition + | Code::NotFound + | Code::AlreadyExists + | Code::OutOfRange + | Code::Unauthenticated + | Code::PermissionDenied + | Code::Unimplemented + | Code::DataLoss => BridgeError::PermanentRpc(status), + // RESOURCE_EXHAUSTED is the indexer shedding load or refusing an + // oversized batch. Reconnecting loses that batch's events, which costs + // routing accuracy; exiting loses every later batch too. + _ => BridgeError::Rpc(status), + } +} + +/// A single indexer mutation, kept in the exact order it appeared in the event +/// batch so mutations on the same hash are never reordered. +/// +/// `Report` carries per-hash component metadata in arrays index-aligned with +/// `hashes`: `masks[i]` is `None` for a legacy whole-block store, and +/// `block_sizes[i]` is the reported token count, `None` when none was supplied. +#[derive(Debug, PartialEq, Eq)] +enum Action { + Report { + tier: i32, + hashes: Vec, + masks: Vec>, + block_sizes: Vec>, + }, + Revoke { + tier: i32, + hashes: Vec, + }, + ClearAll, +} + +#[derive(Debug, Default)] +struct EventActions { + actions: Vec, +} + +impl EventActions { + /// Append a store for the block hashes of one `BlockStored`, coalescing only + /// with an immediately-preceding store to the same tier and never across a + /// revoke/clear, so the final per-hash state is preserved. All hashes here + /// share the event's component mask and block size. + fn report(&mut self, tier: i32, hashes: Vec, mask: Option, block_size: Option) { + if hashes.is_empty() { + return; + } + let n = hashes.len(); + if let Some(Action::Report { + tier: last_tier, + hashes: last, + masks, + block_sizes, + }) = self.actions.last_mut() + { + if *last_tier == tier { + last.extend(hashes); + masks.extend(std::iter::repeat_n(mask, n)); + block_sizes.extend(std::iter::repeat_n(block_size, n)); + return; + } + } + self.actions.push(Action::Report { + tier, + hashes, + masks: vec![mask; n], + block_sizes: vec![block_size; n], + }); + } + + fn revoke(&mut self, tier: i32, hashes: Vec) { + if hashes.is_empty() { + return; + } + if let Some(Action::Revoke { + tier: last_tier, + hashes: last, + }) = self.actions.last_mut() + { + if *last_tier == tier { + last.extend(hashes); + return; + } + } + self.actions.push(Action::Revoke { tier, hashes }); + } + + fn clear_all(&mut self) { + self.actions.push(Action::ClearAll); + } +} + +pub async fn run_bridge(config: BridgeConfig) -> Result<(), BridgeError> { + run_bridge_until(config, std::future::pending()).await +} + +/// [`run_bridge`], but returns as soon as `shutdown` resolves. An in-flight apply +/// is dropped: acknowledgements are not tracked, so that batch is lost. +pub async fn run_bridge_until(config: BridgeConfig, shutdown: F) -> Result<(), BridgeError> +where + F: std::future::Future, +{ + tokio::select! { + result = supervise(config) => result, + () = shutdown => { + info!("bridge stopped by shutdown signal"); + Ok(()) + } + } +} + +async fn supervise(config: BridgeConfig) -> Result<(), BridgeError> { + info!( + worker_id = %config.worker_id, + event_endpoint = %config.event_endpoint, + event_topic = %config.event_topic, + indexer_endpoint = %config.indexer_endpoint, + "starting SGLang KV event bridge" + ); + + // Supervisor loop: (re)connect to both the indexer and the ZMQ publisher, + // run until a connection-level error, then back off and retry. Decode-level + // problems are handled inside the session and never tear down the bridge. + // Reconnecting recovers the connection only: events published while + // disconnected are lost. + let mut delay = RECONNECT_MIN_DELAY; + loop { + match connect(&config).await { + Ok((client, subscriber)) => { + delay = RECONNECT_MIN_DELAY; + match run_session(&config, client, subscriber).await { + Ok(()) => { + info!("bridge shut down cleanly"); + return Ok(()); + } + Err(error) => { + if error.is_permanent() { + return Err(error); + } + warn!(%error, retry_in = ?delay, "bridge session lost; reconnecting"); + } + } + } + Err(error) => { + if error.is_permanent() { + return Err(error); + } + warn!(%error, retry_in = ?delay, "bridge connect failed; retrying"); + } + } + + tokio::time::sleep(delay).await; + delay = (delay * 2).min(RECONNECT_MAX_DELAY); + } +} + +async fn connect( + config: &BridgeConfig, +) -> Result<(KvIndexerClient, SubSocket), BridgeError> { + let channel = Endpoint::from_shared(config.indexer_endpoint.clone())? + .connect_timeout(GRPC_CONNECT_TIMEOUT) + .timeout(GRPC_REQUEST_TIMEOUT) + .connect() + .await?; + let client = KvIndexerClient::new(channel); + let mut subscriber = SubSocket::new(); + subscriber.subscribe(&config.event_topic).await?; + subscriber.connect(&config.event_endpoint).await?; + info!("bridge session established"); + Ok((client, subscriber)) +} + +/// Runs a single connected session. Returns `Ok(())` only on a clean shutdown +/// (ctrl-c); any connection-level error is propagated so the supervisor can +/// reconnect. +async fn run_session( + config: &BridgeConfig, + mut client: KvIndexerClient, + mut subscriber: SubSocket, +) -> Result<(), BridgeError> { + // Tracked only to log a discontinuity. Nothing acts on it. + let mut last_seq: Option = None; + + loop { + let message = tokio::select! { + result = subscriber.recv() => result?, + _ = tokio::signal::ctrl_c() => { + info!("received ctrl-c; shutting down bridge"); + return Ok(()); + } + }; + + let (seq, payload) = match parse_zmq_frames(&message.into_vec()) { + Ok((seq, payload)) => (seq, payload.to_vec()), + Err(error) => { + warn!(%error, "skipping malformed ZMQ message"); + continue; + } + }; + + if let Some(previous) = last_seq { + if seq != previous.wrapping_add(1) { + warn!( + previous, + actual = seq, + "SGLang KV event sequence is not contiguous; this build does not recover the gap" + ); + } + } + last_seq = Some(seq); + + forward_raw_batch(config, &mut client, seq, &payload).await?; + } +} +fn parse_zmq_frames(frames: &[bytes::Bytes]) -> Result<(u64, &[u8]), BridgeError> { + match frames.len() { + 2 => Ok((decode_seq(&frames[0])?, frames[1].as_ref())), + 3 => Ok((decode_seq(&frames[1])?, frames[2].as_ref())), + n => Err(BridgeError::Decode(format!( + "expected 2 or 3 ZMQ frames, got {n}" + ))), + } +} + +fn decode_seq(bytes: &[u8]) -> Result { + let seq_bytes: [u8; 8] = bytes + .try_into() + .map_err(|_| BridgeError::Decode("sequence frame must be 8 bytes".to_string()))?; + Ok(u64::from_be_bytes(seq_bytes)) +} + +/// Decodes one raw batch and forwards it. A batch that cannot be decoded, or +/// that carries no supported mutation, is skipped without an RPC. +async fn forward_raw_batch( + config: &BridgeConfig, + client: &mut KvIndexerClient, + seq: u64, + payload: &[u8], +) -> Result<(), BridgeError> { + let actions = match decode_event_batch(payload) { + Ok(actions) => actions, + Err(error) => { + warn!(seq, %error, "skipping undecodable event batch"); + return Ok(()); + } + }; + + let request = build_apply_request(config, seq, actions); + // One ZMQ batch can hold more mutations than a single apply RPC admits, so + // send the parts in order. A later failure leaves an applied prefix, which + // beats rejecting and losing the whole event batch. + for request in split_apply_request(request) { + client + .apply_external_kv_batch(request) + .await + .map_err(classify_rpc)?; + } + Ok(()) +} + +/// Maps a decoded `EventActions` into a single `ApplyExternalKvBatchRequest`, +/// preserving per-action order. A `ClearAll` is expanded in place into one +/// `CLEAR_ALL_AT_TIER` action per configured clear tier. +fn build_apply_request( + config: &BridgeConfig, + seq: u64, + events: EventActions, +) -> ApplyExternalKvBatchRequest { + let mut actions = Vec::with_capacity(events.actions.len()); + for action in events.actions { + match action { + Action::Report { + tier, + hashes, + masks, + block_sizes, + } => actions.push(ExternalKvAction { + r#type: ExternalKvActionType::ActionReport as i32, + tier, + hashes, + // Emit the per-hash arrays only when some hash carries + // component data; a fully-legacy report leaves them empty so + // the backend keeps the whole-block fast path. + component_masks: encode_component_masks(&masks), + block_sizes: encode_block_sizes(&block_sizes), + }), + Action::Revoke { tier, hashes } => actions.push(ExternalKvAction { + r#type: ExternalKvActionType::ActionRevoke as i32, + tier, + hashes, + component_masks: Vec::new(), + block_sizes: Vec::new(), + }), + Action::ClearAll => { + for tier in &config.clear_tiers { + actions.push(ExternalKvAction { + r#type: ExternalKvActionType::ActionClearAllAtTier as i32, + tier: *tier, + hashes: Vec::new(), + component_masks: Vec::new(), + block_sizes: Vec::new(), + }); + } + } + } + } + + ApplyExternalKvBatchRequest { + worker_id: config.worker_id.clone(), + seq, + actions, + worker_address: config.worker_address.clone(), + cache_spec: config.cache_spec, + } +} + +/// Splits one decoded ZMQ batch into apply RPCs within the service's action and +/// hash bounds, preserving action order and per-hash index alignment. +/// +/// Every part reuses the source `seq`, which is safe because `seq` is +/// observability only (see the proto): applies are never deduplicated or fenced, +/// so a repeated `seq` cannot get a part dropped as stale. +fn split_apply_request(request: ApplyExternalKvBatchRequest) -> Vec { + let mut template = request; + let actions = std::mem::take(&mut template.actions); + let mut batches = Vec::new(); + let mut current = Vec::new(); + let mut current_hashes = 0usize; + + let flush = |actions: &mut Vec, + batches: &mut Vec| { + if actions.is_empty() { + return; + } + batches.push(ApplyExternalKvBatchRequest { + actions: std::mem::take(actions), + ..template.clone() + }); + }; + + for action in actions.into_iter().flat_map(split_action) { + let action_hashes = action.hashes.len(); + if !current.is_empty() + && (current.len() == MAX_ACTIONS_PER_BATCH + || current_hashes + action_hashes > MAX_HASHES_PER_REQUEST) + { + flush(&mut current, &mut batches); + current_hashes = 0; + } + current_hashes += action_hashes; + current.push(action); + } + flush(&mut current, &mut batches); + batches +} + +fn split_action(action: ExternalKvAction) -> Vec { + if action.hashes.len() <= MAX_HASHES_PER_REQUEST { + return vec![action]; + } + + (0..action.hashes.len()) + .step_by(MAX_HASHES_PER_REQUEST) + .map(|start| { + let end = (start + MAX_HASHES_PER_REQUEST).min(action.hashes.len()); + ExternalKvAction { + r#type: action.r#type, + tier: action.tier, + hashes: action.hashes[start..end].to_vec(), + component_masks: slice_or_empty(&action.component_masks, start, end), + block_sizes: slice_or_empty(&action.block_sizes, start, end), + } + }) + .collect() +} + +/// Slices a per-hash array alongside its `hashes` slice. Empty is the legacy +/// "field absent" signal; non-empty arrays are aligned by `build_apply_request`. +fn slice_or_empty(values: &[T], start: usize, end: usize) -> Vec { + if values.is_empty() { + return Vec::new(); + } + values[start..end].to_vec() +} + +/// Maps per-hash component sets to the wire form: an empty vector (the legacy +/// signal) when no hash carries components, otherwise one mask per hash. +fn encode_component_masks(masks: &[Option]) -> Vec { + if masks.iter().all(Option::is_none) { + return Vec::new(); + } + masks.iter().map(|mask| mask.unwrap_or_default()).collect() +} + +/// Maps per-hash block sizes to the wire form. Returns an empty vector when no +/// hash carries a size, otherwise one entry per hash (`0` for a legacy hash). +fn encode_block_sizes(block_sizes: &[Option]) -> Vec { + if block_sizes.iter().all(Option::is_none) { + return Vec::new(); + } + block_sizes + .iter() + .map(|size| size.unwrap_or_default()) + .collect() +} + +fn decode_event_batch(payload: &[u8]) -> Result { + decode_event_batch_impl(payload, true) +} + +fn decode_event_batch_impl( + payload: &[u8], + log_event_errors: bool, +) -> Result { + let mut cursor = Cursor::new(payload); + let value = read_value(&mut cursor).map_err(|error| BridgeError::Decode(error.to_string()))?; + let batch = expect_array(&value, "KVEventBatch")?; + if batch.len() < 2 { + return Err(BridgeError::Decode( + "KVEventBatch must contain timestamp and events".to_string(), + )); + } + + let events = expect_array(&batch[1], "KVEventBatch.events")?; + let mut actions = EventActions::default(); + for (event_index, event) in events.iter().enumerate() { + if let Err(error) = decode_event(event, &mut actions) { + if log_event_errors { + warn!( + event_index, + %error, + "skipping one undecodable SGLang KV event; preserving valid siblings" + ); + } + } + } + Ok(actions) +} + +fn decode_event(event: &Value, actions: &mut EventActions) -> Result<(), BridgeError> { + let event = expect_array(event, "KV event")?; + let event_type = expect_str( + event + .first() + .ok_or_else(|| BridgeError::Decode("KV event is empty".to_string()))?, + "KV event tag", + )?; + + match event_type { + "BlockStored" => { + // At least 7 fields (the legacy schema); an 8th `component_types` + // slot appears with `--enable-kv-events-component-types`. Both + // shapes are accepted. + if event.len() < 7 { + return Err(BridgeError::Decode( + "BlockStored must have at least 7 array fields".to_string(), + )); + } + let tier = medium_to_tier(expect_optional_str(&event[6], "BlockStored.medium")?)?; + // `component_types` is the trailing slot: a list of component labels + // folded into a bitmask, or nil/absent for a legacy whole-block store. + let mask = match event.get(7) { + Some(value) => decode_component_mask(value)?, + None => None, + }; + // The token count is only carried alongside component-aware stores, + // where the query path needs it to accumulate trailing windows. + let block_size = match mask { + Some(_) => Some(decode_block_size(&event[4])?), + None => None, + }; + actions.report(tier, decode_hashes(&event[1])?, mask, block_size); + } + "BlockRemoved" => { + if event.len() < 3 { + return Err(BridgeError::Decode( + "BlockRemoved must have 3 array fields".to_string(), + )); + } + let tier = medium_to_tier(expect_optional_str(&event[2], "BlockRemoved.medium")?)?; + actions.revoke(tier, decode_hashes(&event[1])?); + } + "AllBlocksCleared" => { + actions.clear_all(); + } + other => { + debug!(event_type = other, "ignoring unsupported SGLang KV event"); + } + } + Ok(()) +} + +fn decode_hashes(value: &Value) -> Result, BridgeError> { + expect_array(value, "block_hashes")? + .iter() + .map(|value| { + if let Some(value) = value.as_i64() { + return Ok(value); + } + // SGLang folds the unsigned top 64 bits of the SHA-256 into the + // signed range by subtracting 2^64 (`hash_str_to_int64`), which is + // two's complement, so a producer that serialises the unsigned half + // instead is carrying identical bits. Reinterpreting recovers the + // hash the router queries for; refusing the value would instead skip + // the whole event and lose every placement it carried. + if let Some(value) = value.as_u64() { + return Ok(value as i64); + } + Err(BridgeError::Decode( + "block hash must be an integer".to_string(), + )) + }) + .collect() +} + +/// Decodes the optional `component_types` slot of a `BlockStored` into a component +/// bitmask. `nil` maps to `None`, a legacy whole-block store; an array of labels +/// folds into a bitmask, and labels this build does not model are ignored. +fn decode_component_mask(value: &Value) -> Result, BridgeError> { + if matches!(value, Value::Nil) { + return Ok(None); + } + let mut mask = 0u32; + for item in expect_array(value, "BlockStored.component_types")? { + let name = item + .as_str() + .ok_or_else(|| BridgeError::Decode("component type must be a string".to_string()))?; + if let Some(bit) = component_bit(name) { + mask |= bit; + } + } + Ok(Some(mask)) +} + +/// Decodes the `block_size` (token count) slot of a `BlockStored`. +fn decode_block_size(value: &Value) -> Result { + let raw = value + .as_u64() + .or_else(|| value.as_i64().and_then(|v| u64::try_from(v).ok())) + .ok_or_else(|| { + BridgeError::Decode("block_size must be a non-negative integer".to_string()) + })?; + u32::try_from(raw).map_err(|_| BridgeError::Decode("block_size exceeds u32".to_string())) +} + +fn medium_to_tier(medium: Option<&str>) -> Result { + match medium { + Some("GPU") => Ok(TierType::TierHbm as i32), + Some("CPU_PINNED") => Ok(TierType::TierDram as i32), + Some("DISK") => Ok(TierType::TierSsd as i32), + Some("EXTERNAL") => Err(BridgeError::Decode( + "EXTERNAL medium does not map to a local indexer tier".to_string(), + )), + Some(other) => Err(BridgeError::Decode(format!( + "unsupported SGLang storage medium: {other}" + ))), + None => Err(BridgeError::Decode( + "SGLang storage medium is missing".to_string(), + )), + } +} + +/// Builds the worker's [`WorkerCacheSpec`] from the environment, `None` for a +/// legacy / full-only worker. Rules are fixed, so the config only declares which +/// components are present, the SWA window, and their servable tiers: +/// +/// ```text +/// KV_INDEXER_CACHE_COMPONENTS = full,swa (present components; FULL implied) +/// KV_INDEXER_SWA_WINDOW_TOKENS = 4096 (required when swa present) +/// KV_INDEXER_FULL_TIERS = HBM,DRAM (optional, default HBM,DRAM) +/// KV_INDEXER_SWA_TIERS = HBM +/// KV_INDEXER_MAMBA_TIERS = HBM,DRAM +/// KV_INDEXER_CACHE_SPEC_VERSION = 1 (optional, default 1) +/// ``` +fn cache_spec_from_env() -> Result, BridgeError> { + let Some(list) = env_nonempty("KV_INDEXER_CACHE_COMPONENTS") else { + return Ok(None); + }; + let mut components = 0u32; + for name in list.split(',').map(str::trim).filter(|s| !s.is_empty()) { + let bit = component_bit(name) + .ok_or_else(|| BridgeError::Config(format!("unknown cache component: {name}")))?; + components |= bit; + } + // FULL is the base component and always present on a stored block. + components |= crate::service::COMPONENT_FULL; + + let swa_window_tokens = match env_nonempty("KV_INDEXER_SWA_WINDOW_TOKENS") { + Some(v) => v.parse::().map_err(|_| { + BridgeError::Config(format!("KV_INDEXER_SWA_WINDOW_TOKENS is not a u32: {v}")) + })?, + None => 0, + }; + if components & COMPONENT_SWA != 0 && swa_window_tokens == 0 { + return Err(BridgeError::Config( + "swa component requires KV_INDEXER_SWA_WINDOW_TOKENS".to_string(), + )); + } + + let version = match env_nonempty("KV_INDEXER_CACHE_SPEC_VERSION") { + Some(v) => v + .parse::() + .map_err(|_| BridgeError::Config(format!("cache spec version is not a u32: {v}")))?, + None => 1, + }; + + Ok(Some(WorkerCacheSpec { + version, + components, + swa_window_tokens, + full_tier_mask: env_tier_mask("KV_INDEXER_FULL_TIERS")?, + swa_tier_mask: env_tier_mask("KV_INDEXER_SWA_TIERS")?, + mamba_tier_mask: env_tier_mask("KV_INDEXER_MAMBA_TIERS")?, + })) +} + +fn env_nonempty(key: &str) -> Option { + std::env::var(key) + .ok() + .map(|v| v.trim().to_string()) + .filter(|v| !v.is_empty()) +} + +/// Parses a `KV_INDEXER_*_TIERS` list into a `1 << TierType` bitmask, defaulting +/// to HBM+DRAM when unset. +fn env_tier_mask(key: &str) -> Result { + let list = env_nonempty(key).unwrap_or_else(|| "HBM,DRAM".to_string()); + let mut mask = 0u32; + for tier in list.split(',').map(str::trim).filter(|s| !s.is_empty()) { + mask |= 1u32 << tier_name_to_type(tier)?; + } + Ok(mask) +} + +fn tier_name_to_type(name: &str) -> Result { + match name { + "HBM" | "GPU" => Ok(TierType::TierHbm as i32), + "DRAM" | "CPU" | "CPU_PINNED" => Ok(TierType::TierDram as i32), + "SSD" | "DISK" => Ok(TierType::TierSsd as i32), + other => Err(BridgeError::Config(format!( + "cache spec has unsupported tier: {other}" + ))), + } +} + +fn parse_clear_tiers(value: &str) -> Result, BridgeError> { + value + .split(',') + .map(str::trim) + .filter(|part| !part.is_empty()) + .map(|part| match part { + "HBM" | "GPU" => Ok(TierType::TierHbm as i32), + "DRAM" | "CPU" | "CPU_PINNED" => Ok(TierType::TierDram as i32), + "SSD" | "DISK" => Ok(TierType::TierSsd as i32), + other => Err(BridgeError::Config(format!( + "unsupported clear tier: {other}" + ))), + }) + .collect() +} + +fn expect_array<'a>(value: &'a Value, field: &str) -> Result<&'a [Value], BridgeError> { + value + .as_array() + .map(Vec::as_slice) + .ok_or_else(|| BridgeError::Decode(format!("{field} must be an array"))) +} + +fn expect_str<'a>(value: &'a Value, field: &str) -> Result<&'a str, BridgeError> { + value + .as_str() + .ok_or_else(|| BridgeError::Decode(format!("{field} must be a string"))) +} + +fn expect_optional_str<'a>(value: &'a Value, field: &str) -> Result, BridgeError> { + if matches!(value, Value::Nil) { + return Ok(None); + } + expect_str(value, field).map(Some) +} + +#[cfg(test)] +mod tests { + use super::*; + use bytes::Bytes; + use rmpv::Value; + + fn hbm() -> i32 { + TierType::TierHbm as i32 + } + fn dram() -> i32 { + TierType::TierDram as i32 + } + fn ssd() -> i32 { + TierType::TierSsd as i32 + } + + fn encode(value: &Value) -> Vec { + let mut buf = Vec::new(); + rmpv::encode::write_value(&mut buf, value).unwrap(); + buf + } + + fn ints(values: &[i64]) -> Value { + Value::Array(values.iter().map(|v| Value::from(*v)).collect()) + } + + fn stored(hashes: &[i64], medium: &str) -> Value { + Value::Array(vec![ + Value::String("BlockStored".into()), + ints(hashes), + Value::Nil, // parent_block_hash + ints(&[1]), // token_ids + Value::from(1_i64), // block_size + Value::Nil, // lora_id + Value::String(medium.into()), + ]) + } + + /// A component-aware `BlockStored` (8-element schema): trailing + /// `component_types` slot plus a concrete `block_size` token count. + fn stored_c(hashes: &[i64], medium: &str, block_size: i64, components: Value) -> Value { + Value::Array(vec![ + Value::String("BlockStored".into()), + ints(hashes), + Value::Nil, // parent_block_hash + ints(&[1]), // token_ids + Value::from(block_size), + Value::Nil, // lora_id + Value::String(medium.into()), + components, // component_types (Nil or array of strings) + ]) + } + + fn strv(items: &[&str]) -> Value { + Value::Array(items.iter().map(|s| Value::String((*s).into())).collect()) + } + + /// Legacy (whole-block) report action expectation. + fn rep(tier: i32, hashes: &[&str]) -> Action { + Action::Report { + tier, + hashes: hashes.iter().map(|h| h.parse().unwrap()).collect(), + masks: vec![None; hashes.len()], + block_sizes: vec![None; hashes.len()], + } + } + + fn rev(tier: i32, hashes: &[&str]) -> Action { + Action::Revoke { + tier, + hashes: hashes.iter().map(|h| h.parse().unwrap()).collect(), + } + } + + fn removed(hashes: &[i64], medium: &str) -> Value { + Value::Array(vec![ + Value::String("BlockRemoved".into()), + ints(hashes), + Value::String(medium.into()), + ]) + } + + fn cleared() -> Value { + Value::Array(vec![Value::String("AllBlocksCleared".into())]) + } + + /// Wrap events in a 3-element batch [ts, events, attn_dp_rank]. + fn batch(events: Vec) -> Vec { + encode(&Value::Array(vec![ + Value::from(1.0_f64), + Value::Array(events), + Value::from(0_i64), + ])) + } + + fn actions_of(events: Vec) -> Vec { + decode_event_batch(&batch(events)).unwrap().actions + } + + fn golden_bytes(hex: &str) -> Vec { + assert_eq!(hex.len() % 2, 0); + hex.as_bytes() + .chunks_exact(2) + .map(|pair| { + let high = (pair[0] as char).to_digit(16).unwrap(); + let low = (pair[1] as char).to_digit(16).unwrap(); + ((high << 4) | low) as u8 + }) + .collect() + } + + fn test_config(clear_tiers: Vec) -> BridgeConfig { + BridgeConfig { + worker_id: "worker-1".to_string(), + worker_address: "127.0.0.1:9000".to_string(), + event_endpoint: "tcp://127.0.0.1:5557".to_string(), + event_topic: "kv-events".to_string(), + indexer_endpoint: "http://[::1]:50051".to_string(), + clear_tiers, + cache_spec: None, + } + } + + /// Build the apply-batch request the bridge would send for a set of events. + fn request_of( + config: &BridgeConfig, + seq: u64, + events: Vec, + ) -> ApplyExternalKvBatchRequest { + build_apply_request(config, seq, decode_event_batch(&batch(events)).unwrap()) + } + + fn report(tier: i32, hashes: &[&str]) -> ExternalKvAction { + ExternalKvAction { + r#type: ExternalKvActionType::ActionReport as i32, + tier, + hashes: hashes.iter().map(|h| h.parse().unwrap()).collect(), + component_masks: Vec::new(), + block_sizes: Vec::new(), + } + } + + fn revoke(tier: i32, hashes: &[&str]) -> ExternalKvAction { + ExternalKvAction { + r#type: ExternalKvActionType::ActionRevoke as i32, + tier, + hashes: hashes.iter().map(|h| h.parse().unwrap()).collect(), + component_masks: Vec::new(), + block_sizes: Vec::new(), + } + } + + fn clear_at(tier: i32) -> ExternalKvAction { + ExternalKvAction { + r#type: ExternalKvActionType::ActionClearAllAtTier as i32, + tier, + hashes: Vec::new(), + component_masks: Vec::new(), + block_sizes: Vec::new(), + } + } + + #[test] + fn request_carries_worker_id_and_seq() { + let config = test_config(vec![hbm()]); + let request = request_of(&config, 42, vec![stored(&[1], "GPU")]); + assert_eq!(request.worker_id, "worker-1"); + assert_eq!(request.seq, 42); + } + + #[test] + fn request_carries_worker_address() { + let config = test_config(vec![hbm()]); + let request = request_of(&config, 0, vec![stored(&[1], "GPU")]); + assert_eq!(request.worker_address, "127.0.0.1:9000"); + } + + #[test] + fn oversized_report_is_split_with_aligned_metadata() { + let count = MAX_HASHES_PER_REQUEST + 1; + let request = ApplyExternalKvBatchRequest { + worker_id: "worker-1".into(), + seq: 42, + actions: vec![ExternalKvAction { + r#type: ExternalKvActionType::ActionReport as i32, + tier: hbm(), + hashes: (0..count).map(|index| index as i64).collect(), + component_masks: (0..count as u32).collect(), + block_sizes: (0..count as u32).map(|index| index + 1).collect(), + }], + worker_address: "http://worker-1".into(), + cache_spec: None, + }; + + let batches = split_apply_request(request); + + assert_eq!(batches.len(), 2); + assert_eq!(batches[0].actions[0].hashes.len(), MAX_HASHES_PER_REQUEST); + assert_eq!( + batches[1].actions[0].hashes, + vec![MAX_HASHES_PER_REQUEST as i64] + ); + assert_eq!( + batches[1].actions[0].component_masks, + vec![MAX_HASHES_PER_REQUEST as u32] + ); + assert_eq!( + batches[1].actions[0].block_sizes, + vec![MAX_HASHES_PER_REQUEST as u32 + 1] + ); + assert!(batches.iter().all(|batch| batch.seq == 42)); + } + + #[test] + fn too_many_clear_actions_are_split_in_order() { + let request = ApplyExternalKvBatchRequest { + worker_id: "worker-1".into(), + seq: 7, + actions: (0..=MAX_ACTIONS_PER_BATCH) + .map(|index| clear_at(if index % 2 == 0 { hbm() } else { dram() })) + .collect(), + worker_address: "http://worker-1".into(), + cache_spec: None, + }; + + let batches = split_apply_request(request); + + assert_eq!(batches.len(), 2); + assert_eq!(batches[0].actions.len(), MAX_ACTIONS_PER_BATCH); + assert_eq!(batches[1].actions, vec![clear_at(hbm())]); + } + + #[test] + fn report_and_revoke_map_to_actions_in_order() { + let config = test_config(vec![hbm()]); + let request = request_of( + &config, + 0, + vec![removed(&[9], "GPU"), stored(&[9], "CPU_PINNED")], + ); + assert_eq!( + request.actions, + vec![revoke(hbm(), &["9"]), report(dram(), &["9"])] + ); + } + + #[test] + fn clear_all_expands_to_one_action_per_clear_tier_in_place() { + let config = test_config(vec![hbm(), dram(), ssd()]); + let request = request_of( + &config, + 7, + vec![stored(&[1], "GPU"), cleared(), stored(&[2], "GPU")], + ); + assert_eq!( + request.actions, + vec![ + report(hbm(), &["1"]), + clear_at(hbm()), + clear_at(dram()), + clear_at(ssd()), + report(hbm(), &["2"]), + ] + ); + } + + #[test] + fn batch_with_only_ignored_events_has_no_actions() { + let config = test_config(vec![hbm()]); + let events = vec![Value::Array(vec![Value::String("BlockUpdated".into())])]; + assert!(request_of(&config, 0, events).actions.is_empty()); + } + + #[test] + fn block_stored_maps_to_report_on_tier() { + assert_eq!( + actions_of(vec![stored(&[123], "GPU")]), + vec![rep(hbm(), &["123"])] + ); + } + + #[test] + fn mediums_map_to_expected_tiers() { + assert_eq!( + actions_of(vec![stored(&[1], "CPU_PINNED")]), + vec![rep(dram(), &["1"])] + ); + assert_eq!( + actions_of(vec![removed(&[2], "DISK")]), + vec![rev(ssd(), &["2"])] + ); + } + + #[test] + fn bad_event_does_not_drop_valid_siblings() { + assert_eq!( + actions_of(vec![ + stored(&[1], "GPU"), + stored(&[2], "EXTERNAL"), + removed(&[3], "DISK"), + ]), + vec![rep(hbm(), &["1"]), rev(ssd(), &["3"])] + ); + } + + #[test] + fn permanent_rpc_codes_are_not_retried() { + for code in [ + Code::InvalidArgument, + Code::FailedPrecondition, + Code::PermissionDenied, + ] { + assert!(classify_rpc(Status::new(code, "bad batch")).is_permanent()); + } + assert!(!classify_rpc(Status::unavailable("retry")).is_permanent()); + assert!(!classify_rpc(Status::deadline_exceeded("retry")).is_permanent()); + } + + /// Indexer backpressure must not take the bridge down: the router treats the + /// same code as recoverable, and a rejected batch is worth less than the + /// entire event stream. + #[test] + fn shed_batches_keep_the_bridge_alive() { + assert!(!classify_rpc(Status::resource_exhausted("batch too large")).is_permanent()); + } + + // --- ordering regressions --- + + #[test] + fn remove_then_store_same_hash_keeps_order() { + // Net state must be "stored"; reordering to report-then-revoke would drop it. + assert_eq!( + actions_of(vec![removed(&[9], "GPU"), stored(&[9], "GPU")]), + vec![rev(hbm(), &["9"]), rep(hbm(), &["9"])] + ); + } + + #[test] + fn clear_then_store_keeps_order() { + assert_eq!( + actions_of(vec![cleared(), stored(&[7], "GPU")]), + vec![Action::ClearAll, rep(hbm(), &["7"])] + ); + } + + #[test] + fn store_then_clear_keeps_order() { + assert_eq!( + actions_of(vec![stored(&[7], "GPU"), cleared()]), + vec![rep(hbm(), &["7"]), Action::ClearAll] + ); + } + + // --- coalescing rules --- + + #[test] + fn adjacent_same_tier_stores_coalesce() { + assert_eq!( + actions_of(vec![stored(&[1], "GPU"), stored(&[2], "GPU")]), + vec![rep(hbm(), &["1", "2"])] + ); + } + + #[test] + fn different_tier_stores_do_not_coalesce() { + assert_eq!( + actions_of(vec![stored(&[1], "GPU"), stored(&[2], "CPU_PINNED")]), + vec![rep(hbm(), &["1"]), rep(dram(), &["2"])] + ); + } + + #[test] + fn store_then_remove_same_tier_do_not_merge() { + assert_eq!( + actions_of(vec![stored(&[1], "GPU"), removed(&[1], "GPU")]), + vec![rep(hbm(), &["1"]), rev(hbm(), &["1"])] + ); + } + + #[test] + fn unknown_event_tag_is_ignored() { + let events = vec![Value::Array(vec![Value::String("BlockUpdated".into())])]; + assert!(actions_of(events).is_empty()); + } + + #[test] + fn two_element_batch_without_dp_rank_decodes() { + let payload = encode(&Value::Array(vec![ + Value::from(1.0_f64), + Value::Array(vec![stored(&[5], "GPU")]), + ])); + assert_eq!( + decode_event_batch(&payload).unwrap().actions, + vec![rep(hbm(), &["5"])] + ); + } + + #[test] + fn python_msgspec_mixed_batch_golden_decodes() { + // Generated by msgspec.msgpack.Encoder from the authoritative Python + // KVEventBatch schema in sglang.srt.disaggregation.kv_events. + let payload = golden_bytes(concat!( + "93cb405edd2f1a9fbe779397ab426c6f636b53746f72656492", + "cf0000011f71fb04cbd2c521974f2a940a141e280407a3475055", + "93ac426c6f636b52656d6f7665649264ccc8a44449534b", + "91b0416c6c426c6f636b73436c656172656402" + )); + assert_eq!( + decode_event_batch(&payload).unwrap().actions, + vec![ + rep(hbm(), &["1234567890123", "-987654321"]), + rev(ssd(), &["100", "200"]), + Action::ClearAll, + ] + ); + } + + #[test] + fn python_msgspec_bigram_tokens_golden_decodes() { + // token_ids contains Python tuples as nested msgpack arrays; the + // bridge ignores payload shape and indexes the published hashes. + let payload = golden_bytes(concat!( + "93cb3ff80000000000009197ab426c6f636b53746f726564916f", + "c092920a1492141e02c0a347505503" + )); + assert_eq!( + decode_event_batch(&payload).unwrap().actions, + vec![rep(hbm(), &["111"])] + ); + } + + #[test] + fn python_msgspec_nil_medium_golden_is_safely_skipped() { + // The Python schema permits medium=None; such events map to no + // Indexer tier, so they are isolated rather than given a placement. + let payload = golden_bytes(concat!( + "93cb00000000000000009297ab426c6f636b53746f7265649101", + "c092050602c0c093ac426c6f636b52656d6f7665649102c0c0" + )); + assert!(decode_event_batch(&payload).unwrap().actions.is_empty()); + } + + #[test] + fn negative_hashes_remain_signed_integers() { + assert_eq!( + actions_of(vec![stored(&[-1905904552702706914], "GPU")]), + vec![rep(hbm(), &["-1905904552702706914"])] + ); + } + + // --- error / mapping units --- + + #[test] + fn external_medium_event_is_skipped() { + assert!(decode_event_batch(&batch(vec![stored(&[1], "EXTERNAL")])) + .unwrap() + .actions + .is_empty()); + } + + #[test] + fn unknown_medium_event_is_skipped() { + assert!(decode_event_batch(&batch(vec![stored(&[1], "TAPE")])) + .unwrap() + .actions + .is_empty()); + } + + #[test] + fn medium_to_tier_mapping() { + assert_eq!(medium_to_tier(Some("GPU")).unwrap(), hbm()); + assert_eq!(medium_to_tier(Some("CPU_PINNED")).unwrap(), dram()); + assert_eq!(medium_to_tier(Some("DISK")).unwrap(), ssd()); + assert!(medium_to_tier(Some("EXTERNAL")).is_err()); + assert!(medium_to_tier(None).is_err()); + } + + #[test] + fn parse_clear_tiers_defaults_and_aliases() { + assert_eq!( + parse_clear_tiers("HBM,DRAM,SSD").unwrap(), + vec![hbm(), dram(), ssd()] + ); + assert_eq!( + parse_clear_tiers(" GPU , CPU_PINNED , DISK ").unwrap(), + vec![hbm(), dram(), ssd()] + ); + assert!(parse_clear_tiers("HBM,NVME").is_err()); + } + + /// An event that serialises a hash as unsigned must decode to the value the + /// router queries for. Anything else would file the block under a hash no + /// query can reach, which reads as a silent cache miss rather than an error. + #[test] + fn decode_hashes_reinterprets_unsigned_as_the_same_bits() { + let value = Value::Array(vec![ + Value::from(1_i64), + Value::from(-2_i64), + Value::from(i64::MAX as u64), + Value::from(u64::MAX), + Value::from(1_u64 << 63), + ]); + assert_eq!( + decode_hashes(&value).unwrap(), + vec![1, -2, i64::MAX, -1, i64::MIN] + ); + assert!(decode_hashes(&Value::Array(vec![Value::String("x".into())])).is_err()); + } + + // --- frame / sequence parsing --- + + #[test] + fn parse_zmq_frames_two_and_three() { + let seq = 42_u64; + let two = [ + Bytes::copy_from_slice(&seq.to_be_bytes()), + Bytes::from_static(b"p"), + ]; + assert_eq!(parse_zmq_frames(&two).unwrap().0, seq); + let three = [ + Bytes::from_static(b"kv-events"), + Bytes::copy_from_slice(&seq.to_be_bytes()), + Bytes::from_static(b"p"), + ]; + assert_eq!(parse_zmq_frames(&three).unwrap().0, seq); + let one = [Bytes::from_static(b"p")]; + assert!(parse_zmq_frames(&one).is_err()); + } + + #[test] + fn seq_decoders_are_big_endian() { + assert_eq!(decode_seq(&5_u64.to_be_bytes()).unwrap(), 5); + assert!(decode_seq(&[0_u8; 4]).is_err()); + } + + // --- component-aware decoding --- + + #[test] + fn component_types_list_decodes_into_report() { + assert_eq!( + actions_of(vec![stored_c(&[1], "GPU", 64, strv(&["full", "swa"]))]), + vec![Action::Report { + tier: hbm(), + hashes: vec![1], + masks: vec![Some( + crate::service::COMPONENT_FULL | crate::service::COMPONENT_SWA + )], + block_sizes: vec![Some(64)], + }] + ); + } + + #[test] + fn component_types_nil_decodes_as_legacy() { + // An 8-element BlockStored whose trailing slot is nil is exactly the + // legacy whole-block store: no components, no size. + assert_eq!( + actions_of(vec![stored_c(&[1], "GPU", 64, Value::Nil)]), + vec![rep(hbm(), &["1"])] + ); + } + + #[test] + fn component_aware_report_carries_aligned_wire_arrays() { + let config = test_config(vec![hbm()]); + let request = request_of( + &config, + 0, + vec![ + stored_c(&[1], "GPU", 64, strv(&["full", "swa"])), + stored_c(&[2], "GPU", 32, strv(&["full"])), + ], + ); + assert_eq!(request.actions.len(), 1); + let action = &request.actions[0]; + assert_eq!(action.hashes, vec![1, 2]); + assert_eq!( + action.component_masks, + vec![ + crate::service::COMPONENT_FULL | crate::service::COMPONENT_SWA, + crate::service::COMPONENT_FULL, + ] + ); + assert_eq!(action.block_sizes, vec![64, 32]); + } + + #[test] + fn cache_spec_forwarded_on_request() { + let mut config = test_config(vec![hbm()]); + config.cache_spec = Some(WorkerCacheSpec { + version: 1, + components: crate::service::COMPONENT_FULL, + swa_window_tokens: 0, + full_tier_mask: 1 << hbm(), + swa_tier_mask: 0, + mamba_tier_mask: 0, + }); + let request = request_of(&config, 0, vec![stored(&[1], "GPU")]); + assert_eq!(request.cache_spec, config.cache_spec); + } + + // --- cache spec config helpers --- + + #[test] + fn config_helpers_map_tiers_and_components() { + assert_eq!(tier_name_to_type("HBM").unwrap(), hbm()); + assert_eq!(tier_name_to_type("CPU_PINNED").unwrap(), dram()); + assert!(tier_name_to_type("NVME").is_err()); + assert_eq!(component_bit("full"), Some(crate::service::COMPONENT_FULL)); + assert_eq!(component_bit("swa"), Some(COMPONENT_SWA)); + assert_eq!(component_bit("bogus"), None); + } +} diff --git a/experimental/sgl-router/sgl-kv-indexer/src/client.rs b/experimental/sgl-router/sgl-kv-indexer/src/client.rs new file mode 100644 index 000000000..42ccaac3a --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/src/client.rs @@ -0,0 +1,382 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Router-facing client for the prefix-match query. +//! +//! A successful query distinguishes a real match from an empty result. Transport +//! failures, deadlines, and server rejections stay distinct errors so the caller +//! chooses between degrading and failing the request, instead of silently using +//! a different signal. + +use std::time::Duration; + +use tokio::sync::{Semaphore, SemaphorePermit}; +use tonic::transport::{Channel, Endpoint}; + +use crate::pb::kv_indexer_client::KvIndexerClient; +use crate::pb::MatchExternalKvPrefixRequest; +use crate::service::MAX_GRPC_DECODING_MESSAGE_SIZE; + +/// Default per-query deadline. Indexer failures are request failures, so this +/// absorbs normal cross-host jitter without stalling a request indefinitely. +pub const DEFAULT_QUERY_DEADLINE: Duration = Duration::from_millis(100); +/// Default process-local bound on prefix-query RPCs issued by one client. +pub const DEFAULT_QUERY_MAX_INFLIGHT: usize = 32; +// Leave room for the packed field tag, length prefix, and future scalar fields. +const PREFIX_QUERY_ENCODING_HEADROOM: usize = 16; +const MAX_PREFIX_HASHES_PER_QUERY: usize = + (MAX_GRPC_DECODING_MESSAGE_SIZE - PREFIX_QUERY_ENCODING_HEADROOM) / 8; + +/// One worker's contiguous prefix hit. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PrefixMatch { + /// Router-facing routing identity; intersect byte-for-byte with registered + /// worker URLs. Never empty (the indexer drops unroutable workers). + pub address: String, + /// Length of the contiguous request prefix this worker holds. + pub matched_prefix_blocks: u32, + /// Opaque worker id, for the caller's logs only. + pub worker_id: String, +} + +/// A failed prefix query. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum PrefixIndexError { + /// The endpoint could not be reached. + Unreachable, + /// The query exceeded its deadline. + Timeout, + /// The client or Indexer shed the query because its in-flight limit was hit. + Overloaded, + /// The query exceeded the Indexer's gRPC message-size limit, so no worker's + /// prefix was scanned. Bounded by prompt length, not by load: retrying the + /// same prompt cannot succeed. + QueryTooLarge, + /// The server rejected the request. + Rejected(tonic::Code), +} + +impl std::fmt::Display for PrefixIndexError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Unreachable => f.write_str("KV Indexer is unreachable"), + Self::Timeout => f.write_str("KV Indexer query timed out"), + Self::Overloaded => f.write_str("KV Indexer is overloaded"), + Self::QueryTooLarge => { + f.write_str("KV Indexer query exceeded the gRPC message-size limit") + } + Self::Rejected(code) => write!(f, "KV Indexer rejected the query with {code}"), + } + } +} + +impl std::error::Error for PrefixIndexError {} + +/// A configured endpoint that is not a usable gRPC target. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct InvalidEndpoint { + endpoint: String, + reason: &'static str, +} + +impl std::fmt::Display for InvalidEndpoint { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "invalid KV Indexer endpoint `{}`: {}", + self.endpoint, self.reason + ) + } +} + +impl std::error::Error for InvalidEndpoint {} + +/// Result of a successful prefix query. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum PrefixOutcome { + Matched { + /// Sorted by `matched_prefix_blocks`, descending. + matches: Vec, + /// Longest contiguous prefix held by any single worker. + best_prefix_blocks: u32, + }, + /// No worker holds a prefix (or the request had no hashes). + Empty, +} + +/// Client configuration. +#[derive(Debug, Clone)] +pub struct PrefixIndexConfig { + /// gRPC endpoint of the indexer, e.g. `http://10.0.0.1:50051`. + pub endpoint: String, + /// Per-query deadline. + pub query_deadline: Duration, + /// Maximum prefix-query RPCs issued concurrently by this client. + pub max_inflight: usize, +} + +impl PrefixIndexConfig { + /// Config with the default query deadline ([`DEFAULT_QUERY_DEADLINE`]). + pub fn new(endpoint: impl Into) -> Self { + Self { + endpoint: endpoint.into(), + query_deadline: DEFAULT_QUERY_DEADLINE, + max_inflight: DEFAULT_QUERY_MAX_INFLIGHT, + } + } +} + +/// The prefix-match query the router links against. +#[tonic::async_trait] +pub trait PrefixIndex: Send + Sync { + /// Queries the longest contiguous prefix each worker holds for `hashes` + /// (prompt order, `hashes[0]` first). + async fn match_prefix(&self, hashes: Vec) -> Result; +} + +/// tonic-backed [`PrefixIndex`] with a lazily-established connection. +pub struct GrpcPrefixIndex { + channel: Channel, + deadline: Duration, + prefix_query_semaphore: Semaphore, +} + +impl GrpcPrefixIndex { + /// Rejects an unusable endpoint instead of building a client that can only + /// fail, so a misconfigured address stops startup rather than silently + /// costing every request its cache affinity. + pub fn new(config: PrefixIndexConfig) -> Result { + assert!( + config.max_inflight > 0, + "prefix query max inflight must be greater than zero" + ); + Ok(Self { + channel: parse_endpoint(&config.endpoint)?.connect_lazy(), + deadline: config.query_deadline, + prefix_query_semaphore: Semaphore::new(config.max_inflight), + }) + } + + fn try_acquire_prefix_query(&self) -> Result, PrefixIndexError> { + self.prefix_query_semaphore + .try_acquire() + .map_err(|_| PrefixIndexError::Overloaded) + } +} + +fn truncate_prefix_query(hashes: &mut Vec) -> Option { + let total = hashes.len(); + hashes.truncate(MAX_PREFIX_HASHES_PER_QUERY); + (hashes.len() != total).then_some(total) +} + +#[tonic::async_trait] +impl PrefixIndex for GrpcPrefixIndex { + async fn match_prefix(&self, mut hashes: Vec) -> Result { + if hashes.is_empty() { + return Ok(PrefixOutcome::Empty); + } + + if let Some(total_hashes) = truncate_prefix_query(&mut hashes) { + tracing::warn!( + total_hashes, + queried_hashes = hashes.len(), + "KV Indexer query truncated to the gRPC message-size limit" + ); + } + + let _permit = self.try_acquire_prefix_query()?; + + let mut client = KvIndexerClient::new(self.channel.clone()); + let mut request = tonic::Request::new(MatchExternalKvPrefixRequest { + hashes, + // The policy retains the full query length as its denominator, so a + // transport-limited prefix cannot turn a partial scan into a perfect + // hit. + max_blocks: 0, + }); + // On the wire so the indexer can drop a query this caller already stopped + // waiting for. The local timeout below stays the hard stop, since it also + // covers a stall before the channel applies its own deadline. + request.set_timeout(self.deadline); + + match tokio::time::timeout(self.deadline, client.match_external_kv_prefix(request)).await { + Err(_) => Err(PrefixIndexError::Timeout), + Ok(Err(status)) => Err(classify(status.code())), + Ok(Ok(response)) => { + let response = response.into_inner(); + if response.matches.is_empty() { + return Ok(PrefixOutcome::Empty); + } + let matches = response + .matches + .into_iter() + .map(|m| PrefixMatch { + address: m.worker_address, + matched_prefix_blocks: m.matched_prefix_blocks, + worker_id: m.worker_id, + }) + .collect(); + Ok(PrefixOutcome::Matched { + matches, + best_prefix_blocks: response.best_prefix_blocks, + }) + } + } + } +} + +/// Validates the endpoint the operator configured. tonic itself only checks URI +/// syntax, which accepts a host:port with no scheme and then fails on every +/// connect, so the scheme and host are checked here. +fn parse_endpoint(endpoint: &str) -> Result { + let reject = |reason: &'static str| InvalidEndpoint { + endpoint: endpoint.to_string(), + reason, + }; + let parsed = + Endpoint::from_shared(endpoint.to_string()).map_err(|_| reject("not a valid URI"))?; + // A `unix:` endpoint is fully specified by its socket path. + if endpoint.starts_with("unix:") { + return Ok(parsed); + } + match parsed.uri().scheme_str() { + None => Err(reject("missing scheme, expected http:// or https://")), + Some("http" | "https") => { + if parsed.uri().host().unwrap_or_default().is_empty() { + Err(reject("missing host")) + } else { + Ok(parsed) + } + } + Some(_) => Err(reject("unsupported scheme, expected http:// or https://")), + } +} + +fn classify(code: tonic::Code) -> PrefixIndexError { + match code { + tonic::Code::Unavailable => PrefixIndexError::Unreachable, + // The indexer sheds an expired query as DEADLINE_EXCEEDED, while tonic + // reports its own enforcement of the same `grpc-timeout` as CANCELLED. + // This client cancels a query for no other reason. + tonic::Code::DeadlineExceeded | tonic::Code::Cancelled => PrefixIndexError::Timeout, + tonic::Code::ResourceExhausted => PrefixIndexError::Overloaded, + // The indexer's decoder refuses a message past its size limit with + // OUT_OF_RANGE. A prompt too long to carry is not a disagreement about + // the request contract, so it stays separable from `Rejected` and the + // caller can degrade instead of failing the request. + tonic::Code::OutOfRange => PrefixIndexError::QueryTooLarge, + _ => PrefixIndexError::Rejected(code), + } +} + +#[cfg(test)] +mod tests { + use prost::Message; + + use super::*; + + #[test] + fn classifies_resource_exhausted_as_overload() { + assert_eq!( + classify(tonic::Code::ResourceExhausted), + PrefixIndexError::Overloaded + ); + } + + /// An over-limit query must stay distinguishable from a contract rejection: + /// the caller degrades on the former and fails the request on the latter. + #[test] + fn classifies_over_limit_message_as_too_large() { + assert_eq!( + classify(tonic::Code::OutOfRange), + PrefixIndexError::QueryTooLarge + ); + assert_eq!( + classify(tonic::Code::InvalidArgument), + PrefixIndexError::Rejected(tonic::Code::InvalidArgument) + ); + } + + #[test] + fn oversized_query_keeps_a_prefix_within_the_transport_limit() { + let total = MAX_PREFIX_HASHES_PER_QUERY + 1; + let mut hashes = vec![-1; total]; + + assert_eq!(truncate_prefix_query(&mut hashes), Some(total)); + assert_eq!(hashes.len(), MAX_PREFIX_HASHES_PER_QUERY); + assert!( + MatchExternalKvPrefixRequest { + hashes, + max_blocks: 0, + } + .encoded_len() + <= MAX_GRPC_DECODING_MESSAGE_SIZE + ); + } + + #[test] + fn classifies_both_deadline_signals_as_timeout() { + for code in [tonic::Code::DeadlineExceeded, tonic::Code::Cancelled] { + assert_eq!(classify(code), PrefixIndexError::Timeout); + } + } + + #[test] + fn accepts_endpoints_the_client_can_actually_dial() { + for endpoint in [ + "http://10.0.0.1:50051", + "https://indexer.svc:443", + "unix:/tmp/i", + ] { + assert!( + parse_endpoint(endpoint).is_ok(), + "{endpoint} should be accepted" + ); + } + } + + /// A host:port with no scheme parses as a URI but can never connect, which + /// is the misconfiguration that otherwise only shows up under traffic. + #[test] + fn rejects_endpoints_that_could_only_fail_at_query_time() { + for endpoint in [ + "10.0.0.1:50051", + "indexer.svc", + "grpc://10.0.0.1:50051", + "http://", + ] { + let error = parse_endpoint(endpoint) + .expect_err(&format!("{endpoint} should be rejected")) + .to_string(); + assert!( + error.contains(endpoint), + "error should name the endpoint: {error}" + ); + } + } + + #[test] + fn construction_fails_on_an_invalid_endpoint() { + assert!(GrpcPrefixIndex::new(PrefixIndexConfig::new("10.0.0.1:50051")).is_err()); + } + + #[tokio::test] + async fn local_admission_rejects_without_queueing() { + let index = GrpcPrefixIndex::new(PrefixIndexConfig { + endpoint: "http://127.0.0.1:1".to_string(), + query_deadline: DEFAULT_QUERY_DEADLINE, + max_inflight: 1, + }) + .expect("valid endpoint"); + + let permit = index.try_acquire_prefix_query().unwrap(); + assert_eq!( + index.try_acquire_prefix_query().unwrap_err(), + PrefixIndexError::Overloaded + ); + + drop(permit); + assert!(index.try_acquire_prefix_query().is_ok()); + } +} diff --git a/experimental/sgl-router/sgl-kv-indexer/src/lib.rs b/experimental/sgl-router/sgl-kv-indexer/src/lib.rs new file mode 100644 index 000000000..7d8244e43 --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/src/lib.rs @@ -0,0 +1,34 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! SGLang KV Indexer: a gRPC service that tracks externally-managed KV cache +//! block placements (as reported by inference engines such as SGLang HiCache) +//! and answers placement-match queries for KV-aware routing. + +pub mod bridge; +pub mod client; + +pub mod pb { + tonic::include_proto!("kv_indexer.v1"); +} + +mod admission; +mod memory_backend; +mod service; +mod shutdown; + +pub use admission::stamp_arrival; +pub use client::{ + GrpcPrefixIndex, InvalidEndpoint, PrefixIndex, PrefixIndexConfig, PrefixIndexError, + PrefixMatch, PrefixOutcome, DEFAULT_QUERY_MAX_INFLIGHT, +}; +pub use memory_backend::InMemoryKvIndexerBackend; +pub use service::{ + component_bit, server_builder, BlockComponents, KvIndexerBackend, KvIndexerService, + WorkerPrefixInput, COMPONENT_FULL, COMPONENT_MAMBA, COMPONENT_SWA, + DEFAULT_PREFIX_QUERY_MAX_INFLIGHT, MAX_CONCURRENT_STREAMS, MAX_GRPC_DECODING_MESSAGE_SIZE, +}; +pub use shutdown::shutdown_signal; +/// Re-exported because [`PrefixIndexError::Rejected`] carries it, so callers can +/// match on a rejection without depending on tonic. +pub use tonic::Code as RpcCode; diff --git a/experimental/sgl-router/sgl-kv-indexer/src/memory_backend.rs b/experimental/sgl-router/sgl-kv-indexer/src/memory_backend.rs new file mode 100644 index 000000000..bf1144f59 --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/src/memory_backend.rs @@ -0,0 +1,500 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Process-local storage backend for the KV Indexer. +//! +//! The complete placement view lives behind one [`RwLock`], making an apply batch +//! atomic and every query a consistent snapshot. The state is soft: not shared +//! with another server, and lost when the process exits. + +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard}; + +use tonic::Status; + +use crate::pb::{ + ApplyExternalKvBatchRequest, ApplyExternalKvBatchResponse, ExternalKvActionType, + ExternalKvNodeMatch, GetExternalKvHitCountsRequest, GetExternalKvHitCountsResponse, + HitCountEntry, MatchExternalKvPrefixRequest, MatchExternalKvPrefixResponse, + MatchExternalKvRequest, MatchExternalKvResponse, TierHashes, WorkerCacheSpec, +}; +use crate::service::{assemble_prefix_response, prefix_limit, WorkerPrefixScanner}; +use crate::{BlockComponents, KvIndexerBackend, WorkerPrefixInput}; + +#[derive(Debug, Default)] +struct BlockRecord { + /// Shared block token count. A zero value means legacy/unspecified. + token_count: u32, + /// Resident component snapshot for each `(worker, tier)`. + placements: HashMap<(String, i32), u32>, +} + +#[derive(Debug, Default)] +struct WorkerRecord { + address: String, + spec: Option, + /// Reverse index used by CLEAR_ALL_AT_TIER. + holdings: HashMap>, +} + +#[derive(Debug, Default)] +struct State { + blocks: HashMap, + workers: HashMap, + hit_counts: HashMap, +} + +struct WorkerView { + worker_id: String, + address: String, + spec: Option, + hashes_by_tier: BTreeMap>, + blocks: Vec>, +} + +#[derive(Debug)] +struct PrefixCandidate { + worker_id: String, + address: String, + scanner: WorkerPrefixScanner, +} + +/// Single-process, soft-state KV placement index. +#[derive(Debug, Default)] +pub struct InMemoryKvIndexerBackend { + state: RwLock, +} + +impl InMemoryKvIndexerBackend { + pub fn new() -> Self { + Self::default() + } + + fn read_state(&self) -> Result, Status> { + self.state + .read() + .map_err(|_| Status::internal("in-memory backend lock poisoned")) + } + + fn write_state(&self) -> Result, Status> { + self.state + .write() + .map_err(|_| Status::internal("in-memory backend lock poisoned")) + } + + fn apply( + &self, + req: ApplyExternalKvBatchRequest, + ) -> Result { + let mut state = self.write_state()?; + let worker_id = req.worker_id; + + // Address and spec are snapshots carried on every batch. Empty address + // makes the worker unroutable; absent spec returns it to legacy mode. + { + let worker = state.workers.entry(worker_id.clone()).or_default(); + worker.address = req.worker_address; + worker.spec = req.cache_spec; + } + + for action in req.actions { + match ExternalKvActionType::try_from(action.r#type) { + Ok(ExternalKvActionType::ActionReport) => { + let has_masks = !action.component_masks.is_empty(); + let has_sizes = !action.block_sizes.is_empty(); + + // REPORT is a REPLACE snapshot. Keep the final occurrence + // when a coalesced action repeats one hash. + let mut last_by_hash: HashMap = HashMap::new(); + for (index, hash) in action.hashes.into_iter().enumerate() { + let mask = if has_masks { + action.component_masks[index] + } else { + 0 + }; + let token_count = if has_sizes { + action.block_sizes[index] + } else { + 0 + }; + last_by_hash.insert(hash, (mask, token_count)); + } + + for (hash, (mask, token_count)) in last_by_hash { + let block = state.blocks.entry(hash).or_default(); + block + .placements + .insert((worker_id.clone(), action.tier), mask); + // A legacy report carries no size, so 0 means + // "unknown" and must not erase a known count. + if token_count > 0 { + block.token_count = token_count; + } + state + .workers + .entry(worker_id.clone()) + .or_default() + .holdings + .entry(action.tier) + .or_default() + .insert(hash); + } + } + Ok(ExternalKvActionType::ActionRevoke) => { + for hash in action.hashes { + revoke_one(&mut state, &worker_id, &hash, action.tier); + } + } + Ok(ExternalKvActionType::ActionClearAllAtTier) => { + let hashes = state + .workers + .get(&worker_id) + .and_then(|worker| worker.holdings.get(&action.tier)) + .cloned() + .unwrap_or_default(); + for hash in hashes { + revoke_one(&mut state, &worker_id, &hash, action.tier); + } + } + Ok(ExternalKvActionType::ActionUnknown) | Err(_) => { + return Err(Status::invalid_argument("unsupported action type")); + } + } + } + + Ok(ApplyExternalKvBatchResponse {}) + } + + fn do_match(&self, req: MatchExternalKvRequest) -> Result { + let hashes = dedup_preserve_order(&req.hashes); + let workers = if req.count_as_hit { + let mut state = self.write_state()?; + let (workers, matched_hashes) = Self::collect_worker_views(&state, &hashes, false); + for hash in matched_hashes { + let count = state.hit_counts.entry(hash).or_default(); + *count = count.saturating_add(1); + } + workers + } else { + let state = self.read_state()?; + Self::collect_worker_views(&state, &hashes, false).0 + }; + let matches = workers + .into_iter() + .map(|worker| ExternalKvNodeMatch { + worker_id: worker.worker_id, + address: worker.address, + hashes_by_tier: worker + .hashes_by_tier + .into_iter() + .map(|(tier, placements)| TierHashes { + tier, + hashes: placements.iter().map(|(hash, _, _)| *hash).collect(), + component_masks: placements.iter().map(|(_, mask, _)| *mask).collect(), + block_sizes: placements + .into_iter() + .map(|(_, _, block_size)| block_size) + .collect(), + }) + .collect(), + }) + .collect(); + + Ok(MatchExternalKvResponse { matches }) + } + + fn collect_worker_views( + state: &State, + hashes: &[i64], + with_blocks: bool, + ) -> (Vec, Vec) { + // Keyed by a borrow of the stored worker id, so it is copied once per + // worker in the result rather than once per scanned placement. + let mut worker_order: Vec<&str> = Vec::new(); + let mut by_worker: HashMap<&str, WorkerView> = HashMap::new(); + let mut matched_hashes = Vec::new(); + + for (index, hash) in hashes.iter().enumerate() { + let Some(block) = state.blocks.get(hash) else { + continue; + }; + if block.placements.is_empty() { + continue; + } + matched_hashes.push(*hash); + for ((worker, tier), mask) in &block.placements { + let view = by_worker.entry(worker.as_str()).or_insert_with(|| { + worker_order.push(worker.as_str()); + let metadata = state.workers.get(worker); + WorkerView { + worker_id: worker.clone(), + address: metadata + .map(|worker| worker.address.clone()) + .unwrap_or_default(), + spec: metadata.and_then(|worker| worker.spec), + hashes_by_tier: BTreeMap::new(), + blocks: if with_blocks { + vec![None; hashes.len()] + } else { + Vec::new() + }, + } + }); + if with_blocks { + let components = view.blocks[index].get_or_insert_with(|| BlockComponents { + token_count: block.token_count, + tier_masks: Vec::new(), + }); + components.tier_masks.push((*tier, *mask)); + } else { + view.hashes_by_tier.entry(*tier).or_default().push(( + *hash, + *mask, + block.token_count, + )); + } + } + } + + let workers = worker_order + .into_iter() + .filter_map(|worker| by_worker.remove(worker)) + .collect(); + (workers, matched_hashes) + } + + fn collect_prefix_inputs_locked(state: &State, hashes: &[i64]) -> Vec { + Self::collect_worker_views(state, hashes, true) + .0 + .into_iter() + .map(|worker| WorkerPrefixInput { + worker_id: worker.worker_id, + address: worker.address, + spec: worker.spec, + blocks: worker.blocks, + }) + .collect() + } + + fn do_match_prefix( + &self, + req: MatchExternalKvPrefixRequest, + ) -> Result { + let limit = prefix_limit(req.hashes.len(), req.max_blocks); + let hashes = &req.hashes[..limit]; + if hashes.is_empty() { + return Ok(MatchExternalKvPrefixResponse::default()); + } + + // Only workers holding block zero can own a non-empty prefix, so the + // candidate set is fixed up front and each candidate reuses one scanner and + // one block view. Allocation is O(first-block holders) whatever the request + // length, which is why there is no scan cap: length costs time, not memory. + let state = self.read_state()?; + let Some(first) = state + .blocks + .get(&hashes[0]) + .filter(|block| !block.placements.is_empty()) + else { + return Ok(MatchExternalKvPrefixResponse { + matches: Vec::new(), + best_prefix_blocks: 0, + blocks_read: 1, + }); + }; + let mut seen = HashSet::new(); + let mut candidates: Vec = first + .placements + .keys() + .filter(|(worker, _)| seen.insert(worker.as_str())) + .map(|(worker, _)| { + let metadata = state.workers.get(worker); + PrefixCandidate { + address: metadata + .map(|worker| worker.address.clone()) + .unwrap_or_default(), + scanner: WorkerPrefixScanner::new( + metadata.and_then(|worker| worker.spec.as_ref()), + ), + worker_id: worker.clone(), + } + }) + .collect(); + let candidate_by_id: HashMap = candidates + .iter() + .enumerate() + .map(|(index, candidate)| (candidate.worker_id.clone(), index)) + .collect(); + let mut present = vec![false; candidates.len()]; + let mut block_views: Vec = (0..candidates.len()) + .map(|_| BlockComponents { + token_count: 0, + tier_masks: Vec::new(), + }) + .collect(); + + for hash in hashes { + present.fill(false); + for block in &mut block_views { + block.token_count = 0; + block.tier_masks.clear(); + } + if let Some(block) = state.blocks.get(hash) { + for ((worker, tier), mask) in &block.placements { + let Some(&index) = candidate_by_id.get(worker) else { + continue; + }; + present[index] = true; + block_views[index].token_count = block.token_count; + block_views[index].tier_masks.push((*tier, *mask)); + } + } + for (index, candidate) in candidates.iter_mut().enumerate() { + candidate + .scanner + .push(present[index].then_some(&block_views[index])); + } + } + + let entries = candidates + .into_iter() + .filter_map(|candidate| { + let prefix = candidate.scanner.prefix(); + (!candidate.address.is_empty() && prefix > 0).then_some(( + candidate.worker_id, + candidate.address, + prefix, + )) + }) + .collect(); + Ok(assemble_prefix_response(entries, limit as u32)) + } + + fn do_hit_counts( + &self, + req: GetExternalKvHitCountsRequest, + ) -> Result { + let state = self.read_state()?; + let entries = dedup_preserve_order(&req.hashes) + .into_iter() + .filter_map(|hash| { + state + .hit_counts + .get(&hash) + .copied() + .map(|hit_count_total| HitCountEntry { + hash, + hit_count_total, + }) + }) + .collect(); + Ok(GetExternalKvHitCountsResponse { entries }) + } +} + +fn revoke_one(state: &mut State, worker_id: &str, hash: &i64, tier: i32) { + let mut remove_block = false; + if let Some(block) = state.blocks.get_mut(hash) { + block.placements.remove(&(worker_id.to_string(), tier)); + remove_block = block.placements.is_empty(); + } + + if let Some(worker) = state.workers.get_mut(worker_id) { + if let Some(hashes) = worker.holdings.get_mut(&tier) { + hashes.remove(hash); + if hashes.is_empty() { + worker.holdings.remove(&tier); + } + } + } + + if remove_block { + state.blocks.remove(hash); + state.hit_counts.remove(hash); + } +} + +fn dedup_preserve_order(hashes: &[i64]) -> Vec { + let mut seen = HashSet::new(); + hashes + .iter() + .filter(|hash| seen.insert(**hash)) + .copied() + .collect() +} + +#[tonic::async_trait] +impl KvIndexerBackend for InMemoryKvIndexerBackend { + async fn apply_external_kv_batch( + &self, + request: ApplyExternalKvBatchRequest, + ) -> Result { + self.apply(request) + } + + async fn match_external_kv( + &self, + request: MatchExternalKvRequest, + ) -> Result { + self.do_match(request) + } + + async fn collect_worker_prefix_inputs( + &self, + hashes: &[i64], + ) -> Result, Status> { + let state = self.read_state()?; + Ok(Self::collect_prefix_inputs_locked(&state, hashes)) + } + + async fn match_external_kv_prefix( + &self, + request: MatchExternalKvPrefixRequest, + ) -> Result { + self.do_match_prefix(request) + } + + async fn get_external_kv_hit_counts( + &self, + request: GetExternalKvHitCountsRequest, + ) -> Result { + self.do_hit_counts(request) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{mpsc, Arc}; + use std::time::Duration; + + #[test] + fn dedup_keeps_first_seen_order() { + let hashes = vec![1, -2, 1, 3]; + assert_eq!(dedup_preserve_order(&hashes), vec![1, -2, 3]); + } + + #[test] + fn match_without_hit_count_uses_shared_lock() { + let backend = Arc::new(InMemoryKvIndexerBackend::new()); + let read_guard = backend.read_state().unwrap(); + let query_backend = Arc::clone(&backend); + let (result_tx, result_rx) = mpsc::channel(); + + let query = std::thread::spawn(move || { + result_tx + .send(query_backend.do_match(MatchExternalKvRequest { + hashes: vec![-1], + count_as_hit: false, + })) + .unwrap(); + }); + + result_rx + .recv_timeout(Duration::from_secs(1)) + .expect("read-only match should not wait for another reader") + .unwrap(); + drop(read_guard); + query.join().unwrap(); + } +} diff --git a/experimental/sgl-router/sgl-kv-indexer/src/service.rs b/experimental/sgl-router/sgl-kv-indexer/src/service.rs new file mode 100644 index 000000000..91f807e19 --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/src/service.rs @@ -0,0 +1,1036 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::HashMap; + +use tokio::sync::Semaphore; +use tonic::transport::Server; +use tonic::{Request, Response, Status}; + +use crate::admission::{reject_if_deadline_passed, RejectionLog}; +use crate::pb::kv_indexer_server::{KvIndexer, KvIndexerServer}; +use crate::pb::{ + ApplyExternalKvBatchRequest, ApplyExternalKvBatchResponse, ExternalKvAction, + ExternalKvActionType, ExternalKvPrefixMatch, GetExternalKvHitCountsRequest, + GetExternalKvHitCountsResponse, MatchExternalKvPrefixRequest, MatchExternalKvPrefixResponse, + MatchExternalKvRequest, MatchExternalKvResponse, TierType, WorkerCacheSpec, +}; + +/// Protocol-level resource bounds, enforced before a backend sees the request so +/// no caller can make it allocate work proportional to an unbounded field. The +/// prefix query is exempt from the hash bound; see [`validate_hashes`]. +pub(crate) const MAX_HASHES_PER_REQUEST: usize = 16_384; +pub(crate) const MAX_ACTIONS_PER_BATCH: usize = 256; +pub const DEFAULT_PREFIX_QUERY_MAX_INFLIGHT: usize = 32; +/// Maximum encoded gRPC request size accepted by the Indexer server. With +/// packed `sfixed64` hashes this holds roughly one million blocks. +pub const MAX_GRPC_DECODING_MESSAGE_SIZE: usize = 8 * 1024 * 1024; +/// Per-connection bound on concurrently served HTTP/2 streams. Decoding happens +/// in tonic's codec before a method body runs, so `prefix_query_max_inflight` +/// bounds only the scan, not the bytes a peer makes the server buffer — left +/// unset, one connection can hold an unbounded number of +/// [`MAX_GRPC_DECODING_MESSAGE_SIZE`] messages at once. Sized well above the +/// router's own default of 32 in-flight queries so it never throttles a healthy +/// caller. +pub const MAX_CONCURRENT_STREAMS: u32 = 64; + +static OVERLOAD_LOG: RejectionLog = RejectionLog::new(); + +/// Storage backend for the indexer. Every mutation flows through +/// `apply_external_kv_batch`, preserving one ordered write path. +/// +/// Async so a backend that does IO fits without reshaping the trait, and +/// dyn-safe so the server can hold it as `Arc`. +#[tonic::async_trait] +pub trait KvIndexerBackend: Send + Sync + 'static { + /// Applies a whole SGLang KVEventBatch. The actions are pre-validated and + /// must be applied in order. Applies are unconditional: the request `seq` is + /// informational only and a redelivered batch is applied again. + async fn apply_external_kv_batch( + &self, + request: ApplyExternalKvBatchRequest, + ) -> Result; + + async fn match_external_kv( + &self, + request: MatchExternalKvRequest, + ) -> Result; + + /// Collects the per-worker, per-block component placement needed to compute a + /// prefix, aligned with `hashes`. + /// + /// The default implementation is component-blind: every held block becomes a + /// legacy whole-block placement. Component-aware backends override it to + /// attach each worker's `WorkerCacheSpec` and the resident component set. + async fn collect_worker_prefix_inputs( + &self, + hashes: &[i64], + ) -> Result, Status> { + let matched = self + .match_external_kv(MatchExternalKvRequest { + hashes: hashes.to_vec(), + count_as_hit: false, + }) + .await?; + Ok(legacy_inputs_from_match(hashes, &matched)) + } + + /// Answers, per worker, the longest reusable request prefix it holds. + /// + /// This default implementation *is* the written definition of the prefix + /// semantics, so a backend that overrides it for performance must stay + /// field-for-field identical except for `blocks_read`, which is + /// observability rather than semantics. + /// + /// The result is a safe lower bound: every required component's rule is + /// applied, so an accurate index can only under-report, never over-report. + async fn match_external_kv_prefix( + &self, + request: MatchExternalKvPrefixRequest, + ) -> Result { + let limit = prefix_limit(request.hashes.len(), request.max_blocks); + let hashes: Vec = request.hashes.into_iter().take(limit).collect(); + if hashes.is_empty() { + return Ok(MatchExternalKvPrefixResponse::default()); + } + // The default path reads placement for every considered block. + let inputs = self.collect_worker_prefix_inputs(&hashes).await?; + Ok(compute_prefix_response(&inputs, hashes.len() as u32)) + } + + async fn get_external_kv_hit_counts( + &self, + request: GetExternalKvHitCountsRequest, + ) -> Result; +} + +/// Blanket impl so the server can hold the selected backend as +/// `Arc` and still satisfy `KvIndexerService`. +#[tonic::async_trait] +impl KvIndexerBackend for std::sync::Arc { + async fn apply_external_kv_batch( + &self, + request: ApplyExternalKvBatchRequest, + ) -> Result { + (**self).apply_external_kv_batch(request).await + } + + async fn match_external_kv( + &self, + request: MatchExternalKvRequest, + ) -> Result { + (**self).match_external_kv(request).await + } + + async fn collect_worker_prefix_inputs( + &self, + hashes: &[i64], + ) -> Result, Status> { + (**self).collect_worker_prefix_inputs(hashes).await + } + + async fn match_external_kv_prefix( + &self, + request: MatchExternalKvPrefixRequest, + ) -> Result { + (**self).match_external_kv_prefix(request).await + } + + async fn get_external_kv_hit_counts( + &self, + request: GetExternalKvHitCountsRequest, + ) -> Result { + (**self).get_external_kv_hit_counts(request).await + } +} + +#[derive(Debug)] +pub struct KvIndexerService { + backend: B, + prefix_query_semaphore: Semaphore, +} + +impl KvIndexerService +where + B: KvIndexerBackend, +{ + pub fn new(backend: B) -> Self { + Self::with_prefix_query_max_inflight(backend, DEFAULT_PREFIX_QUERY_MAX_INFLIGHT) + } + + pub fn with_prefix_query_max_inflight(backend: B, max_inflight: usize) -> Self { + assert!( + max_inflight > 0, + "prefix query max inflight must be greater than zero" + ); + Self { + backend, + prefix_query_semaphore: Semaphore::new(max_inflight), + } + } + + /// Wraps the service in its generated server with the decoding limit a + /// full-length prefix query needs. Constructing the server any other way + /// silently reinstates tonic's 4 MiB default, so production and tests that + /// exercise large requests go through here. + pub fn into_server(self) -> KvIndexerServer { + KvIndexerServer::new(self).max_decoding_message_size(MAX_GRPC_DECODING_MESSAGE_SIZE) + } +} + +/// A transport builder carrying the Indexer's stream bound. Pairs with +/// [`KvIndexerService::into_server`]: that sets the per-message ceiling, this +/// bounds how many messages can be in flight against it at once. +pub fn server_builder() -> Server { + Server::builder().max_concurrent_streams(MAX_CONCURRENT_STREAMS) +} + +#[tonic::async_trait] +impl KvIndexer for KvIndexerService +where + B: KvIndexerBackend, +{ + async fn match_external_kv( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + validate_hashes_bounded(&request.hashes)?; + let response = self.backend.match_external_kv(request).await?; + Ok(Response::new(response)) + } + + async fn match_external_kv_prefix( + &self, + request: Request, + ) -> Result, Status> { + let (metadata, extensions, request) = request.into_parts(); + // Before any work: an expired query must not spend the capacity the rest + // of the backlog needs to drain. + reject_if_deadline_passed(&metadata, &extensions)?; + validate_hashes(&request.hashes)?; + // Caps concurrent prefix queries; excess is rejected, never queued. + let _permit = self.prefix_query_semaphore.try_acquire().map_err(|_| { + if let Some(rejected_total) = OVERLOAD_LOG.record() { + tracing::warn!( + rejected_total, + "rejecting prefix query: too many in-flight prefix queries" + ); + } + Status::resource_exhausted("too many in-flight prefix queries") + })?; + let response = self.backend.match_external_kv_prefix(request).await?; + Ok(Response::new(response)) + } + + async fn get_external_kv_hit_counts( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + validate_hashes_bounded(&request.hashes)?; + let response = self.backend.get_external_kv_hit_counts(request).await?; + Ok(Response::new(response)) + } + + async fn apply_external_kv_batch( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + validate_worker_id(&request.worker_id)?; + validate_actions(&request.actions)?; + let response = self.backend.apply_external_kv_batch(request).await?; + Ok(Response::new(response)) + } +} + +fn validate_worker_id(worker_id: &str) -> Result<(), Status> { + if worker_id.is_empty() { + return Err(Status::invalid_argument("worker_id must not be empty")); + } + Ok(()) +} + +/// Well-formedness plus the per-request hash ceiling, for the RPCs that mutate +/// state or build a per-hash response. +fn validate_hashes_bounded(hashes: &[i64]) -> Result<(), Status> { + validate_hashes(hashes)?; + if hashes.len() > MAX_HASHES_PER_REQUEST { + return Err(Status::resource_exhausted(format!( + "request contains {} hashes; maximum is {MAX_HASHES_PER_REQUEST}", + hashes.len() + ))); + } + Ok(()) +} + +/// Well-formedness only: no hash ceiling. A prefix scan uses O(1) state per +/// candidate worker, and truncating it would silently understate a worker's +/// reusable prefix. Length is bounded by `max_blocks` and the transport limit, +/// not by the caller's deadline, which cannot cancel a scan already under way. +fn validate_hashes(hashes: &[i64]) -> Result<(), Status> { + if hashes.is_empty() { + return Err(Status::invalid_argument("hashes must not be empty")); + } + Ok(()) +} + +fn validate_tier(tier: i32) -> Result<(), Status> { + match tier { + 1..=3 => Ok(()), + 0 => Err(Status::invalid_argument("tier must not be TIER_UNKNOWN")), + _ => Err(Status::invalid_argument("tier is not supported")), + } +} + +fn validate_actions(actions: &[ExternalKvAction]) -> Result<(), Status> { + // An empty actions list is a no-op that only refreshes the worker's recorded + // address. Non-empty batches still have every action validated below. + if actions.len() > MAX_ACTIONS_PER_BATCH { + return Err(Status::resource_exhausted(format!( + "batch contains {} actions; maximum is {MAX_ACTIONS_PER_BATCH}", + actions.len() + ))); + } + let total_hashes: usize = actions.iter().map(|action| action.hashes.len()).sum(); + if total_hashes > MAX_HASHES_PER_REQUEST { + return Err(Status::resource_exhausted(format!( + "batch contains {total_hashes} hashes; maximum is {MAX_HASHES_PER_REQUEST}" + ))); + } + for action in actions { + validate_tier(action.tier)?; + match ExternalKvActionType::try_from(action.r#type) { + Ok(ExternalKvActionType::ActionReport) | Ok(ExternalKvActionType::ActionRevoke) => { + validate_hashes_bounded(&action.hashes)?; + } + // CLEAR_ALL_AT_TIER carries only a tier; hashes are ignored. + Ok(ExternalKvActionType::ActionClearAllAtTier) => {} + Ok(ExternalKvActionType::ActionUnknown) | Err(_) => { + return Err(Status::invalid_argument("action type is not supported")); + } + } + // The per-hash arrays are either absent (legacy) or index-aligned with + // `hashes`; a partial array is a malformed batch, not a silent legacy hash. + validate_aligned( + action.component_masks.len(), + action.hashes.len(), + "component_masks", + )?; + validate_aligned(action.block_sizes.len(), action.hashes.len(), "block_sizes")?; + } + Ok(()) +} + +/// A per-hash side array must be empty (legacy) or exactly as long as `hashes`. +fn validate_aligned(array_len: usize, hashes_len: usize, field: &str) -> Result<(), Status> { + if array_len != 0 && array_len != hashes_len { + return Err(Status::invalid_argument(format!( + "{field} has {array_len} entries but must be empty or match {hashes_len} hashes" + ))); + } + Ok(()) +} + +/// Number of leading blocks to consider for a prefix query: bounded by the +/// request length and, when the caller set one, by `max_blocks` (0 disables the +/// caller ceiling). +pub(crate) fn prefix_limit(len: usize, max_blocks: u32) -> usize { + if max_blocks == 0 { + len + } else { + len.min(max_blocks as usize) + } +} + +/// KV component bits. Each component's rule is a property of its type, so the +/// indexer applies fixed semantics rather than a per-worker rule binding. +pub const COMPONENT_FULL: u32 = 1 << 0; +pub const COMPONENT_SWA: u32 = 1 << 1; +pub const COMPONENT_MAMBA: u32 = 1 << 2; + +/// On-wire component label to its bit; `None` for a label this build does not +/// model (ignored, so an unknown future component never counts). +pub fn component_bit(name: &str) -> Option { + match name { + "full" => Some(COMPONENT_FULL), + "swa" => Some(COMPONENT_SWA), + "mamba" => Some(COMPONENT_MAMBA), + _ => None, + } +} + +/// Servable tiers as a `1 << TierType` bitmask. V1: HBM + DRAM, SSD excluded. +const SERVABLE_TIER_MASK: u32 = + (1 << (TierType::TierHbm as u32)) | (1 << (TierType::TierDram as u32)); + +/// Highest `WorkerCacheSpec.version` this build interprets; a higher (future) +/// version fails closed. Version 0 (proto default) is accepted as current. +const SUPPORTED_SPEC_VERSION: u32 = 1; + +/// Whether `tier` is set in a `1 << TierType` bitmask. +fn tier_in_mask(mask: u32, tier: i32) -> bool { + tier >= 0 && mask & (1u32 << tier) != 0 +} + +/// One block's placement at one worker: token count plus, per tier held, the +/// resident component bitmask (mask `0` = legacy whole-block, held with no detail). +#[derive(Debug, Clone)] +pub struct BlockComponents { + pub token_count: u32, + pub tier_masks: Vec<(i32, u32)>, +} + +/// One candidate worker for the rule engine: routing identity, optional spec, and +/// per-query-block placement (`None` where the worker does not hold the block). +#[derive(Debug, Clone)] +pub struct WorkerPrefixInput { + pub worker_id: String, + pub address: String, + pub spec: Option, + pub blocks: Vec>, +} + +/// Builds component-blind (legacy) prefix inputs from a `MatchExternalKv` result: +/// each held block becomes a whole-block placement (mask `0`, no size, no spec). +pub(crate) fn legacy_inputs_from_match( + hashes: &[i64], + matched: &MatchExternalKvResponse, +) -> Vec { + matched + .matches + .iter() + .map(|node| { + let mut tiers_by_hash: HashMap> = HashMap::new(); + for tier in &node.hashes_by_tier { + for hash in &tier.hashes { + tiers_by_hash.entry(*hash).or_default().push(tier.tier); + } + } + let blocks = hashes + .iter() + .map(|hash| { + tiers_by_hash.get(hash).map(|tiers| BlockComponents { + token_count: 0, + tier_masks: tiers.iter().map(|tier| (*tier, 0u32)).collect(), + }) + }) + .collect(); + WorkerPrefixInput { + worker_id: node.worker_id.clone(), + address: node.address.clone(), + spec: None, + blocks, + } + }) + .collect() +} + +/// Runs the component-aware rule engine over each worker and assembles the +/// response. Every backend feeds this same engine, so fast paths cannot drift. +pub(crate) fn compute_prefix_response( + inputs: &[WorkerPrefixInput], + blocks_read: u32, +) -> MatchExternalKvPrefixResponse { + let entries = inputs + .iter() + .filter_map(|worker| { + // An empty address is unroutable (see the proto worker_address contract). + if worker.address.is_empty() { + return None; + } + let prefix = compute_worker_prefix(worker.spec.as_ref(), &worker.blocks); + (prefix > 0).then(|| (worker.worker_id.clone(), worker.address.clone(), prefix)) + }) + .collect(); + assemble_prefix_response(entries, blocks_read) +} + +/// The reusable prefix length for one worker: a safe lower bound on what it can +/// serve. Returns 0 (the worker is excluded) when a component-aware store lacks a +/// spec or the spec carries an unusable rule. +pub(crate) fn compute_worker_prefix( + spec: Option<&WorkerCacheSpec>, + blocks: &[Option], +) -> u32 { + let mut scanner = WorkerPrefixScanner::new(spec); + for block in blocks { + scanner.push(block.as_ref()); + } + scanner.prefix() +} + +/// Incremental form of the component rule engine: one forward pass, one block at +/// a time, O(1) state. Lets a backend answer a prefix query without materializing +/// a `workers × request_blocks` placement array. +#[derive(Debug)] +pub(crate) struct WorkerPrefixScanner { + processed: u32, + state: PrefixScanState, +} + +#[derive(Debug)] +enum PrefixScanState { + /// A worker reporting no components: the count of leading blocks it holds, + /// unless some block carries a component mask, which fails the whole result + /// closed. + Legacy { + prefix: u32, + /// False once a gap appears, after which `prefix` is final. + contiguous: bool, + saw_components: bool, + }, + /// An unusable spec. Fails closed no matter what blocks arrive. + Invalid, + /// The largest boundary `N` where every required component's rule holds: + /// * FULL (always required) — present on every block `0..N`. + /// * SWA (if present) — an unbroken run ending at `N-1` covering + /// `swa_window_tokens`, or reaching the head. + /// * MAMBA (if present) — present on block `N-1`. + ComponentAware { + /// Cleared once FULL is missing, which freezes `best`. + active: bool, + best: u32, + spec: WorkerCacheSpec, + /// Contiguous SWA tokens ending at the block just processed. + swa_run: u64, + swa_head_broken: bool, + }, +} + +impl WorkerPrefixScanner { + pub(crate) fn new(spec: Option<&WorkerCacheSpec>) -> Self { + let state = match spec { + // No spec to interpret components with: legacy until a block proves + // otherwise. + None => PrefixScanState::Legacy { + prefix: 0, + contiguous: true, + saw_components: false, + }, + Some(spec) if spec.components == 0 || spec.version > SUPPORTED_SPEC_VERSION => { + PrefixScanState::Invalid + } + Some(spec) if spec.components & COMPONENT_SWA != 0 && spec.swa_window_tokens == 0 => { + PrefixScanState::Invalid + } + Some(spec) => PrefixScanState::ComponentAware { + active: true, + best: 0, + spec: *spec, + swa_run: 0, + swa_head_broken: false, + }, + }; + Self { + processed: 0, + state, + } + } + + pub(crate) fn push(&mut self, block: Option<&BlockComponents>) { + self.processed = self.processed.saturating_add(1); + match &mut self.state { + PrefixScanState::Legacy { + prefix, + contiguous, + saw_components, + } => { + // Runs past the gap too: a mask on any later block still fails the + // whole result closed. + *saw_components |= + block.is_some_and(|block| block.tier_masks.iter().any(|(_, mask)| *mask != 0)); + if *contiguous { + match block { + Some(_) => *prefix = self.processed, + None => *contiguous = false, + } + } + } + PrefixScanState::Invalid => {} + PrefixScanState::ComponentAware { + active, + best, + spec, + swa_run, + swa_head_broken, + } => { + if !*active { + return; + } + // FULL gates contiguity, so a block missing it settles this worker. + if !component_available(block, COMPONENT_FULL, spec.full_tier_mask) { + *active = false; + return; + } + let mut boundary_ok = true; + if spec.components & COMPONENT_SWA != 0 { + if component_available(block, COMPONENT_SWA, spec.swa_tier_mask) { + *swa_run += block.map(|block| block.token_count as u64).unwrap_or(0); + // Reaching the head counts as valid, matching the unified + // cache's accumulator seeded at infinity. + boundary_ok &= + !*swa_head_broken || *swa_run >= spec.swa_window_tokens as u64; + } else { + *swa_run = 0; + *swa_head_broken = true; + boundary_ok = false; // boundary block itself must carry SWA + } + } + if spec.components & COMPONENT_MAMBA != 0 { + boundary_ok &= + component_available(block, COMPONENT_MAMBA, spec.mamba_tier_mask); + } + if boundary_ok { + *best = self.processed; + } + } + } + } + + pub(crate) fn prefix(&self) -> u32 { + match &self.state { + PrefixScanState::Legacy { + prefix, + saw_components, + .. + } => { + if *saw_components { + 0 + } else { + *prefix + } + } + PrefixScanState::Invalid => 0, + PrefixScanState::ComponentAware { best, .. } => *best, + } + } +} + +/// Whether `component` (a single bit) is resident on `block` at some tier that is +/// both declared servable for that component (`spec_tier_mask`) and servable by +/// the indexer (`SERVABLE_TIER_MASK`). +fn component_available( + block: Option<&BlockComponents>, + component: u32, + spec_tier_mask: u32, +) -> bool { + let Some(block) = block else { + return false; + }; + block.tier_masks.iter().any(|(tier, mask)| { + mask & component != 0 + && tier_in_mask(SERVABLE_TIER_MASK, *tier) + && tier_in_mask(spec_tier_mask, *tier) + }) +} + +/// Sorts `(worker_id, address, prefix)` entries by prefix descending and builds +/// the response, so `best_prefix_blocks` and the order come from one place. +pub(crate) fn assemble_prefix_response( + mut entries: Vec<(String, String, u32)>, + blocks_read: u32, +) -> MatchExternalKvPrefixResponse { + entries.sort_by_key(|entry| std::cmp::Reverse(entry.2)); + let best_prefix_blocks = entries.first().map(|entry| entry.2).unwrap_or(0); + let matches = entries + .into_iter() + .map( + |(worker_id, worker_address, matched_prefix_blocks)| ExternalKvPrefixMatch { + worker_address, + matched_prefix_blocks, + worker_id, + }, + ) + .collect(); + MatchExternalKvPrefixResponse { + matches, + best_prefix_blocks, + blocks_read, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + use tokio::time::Duration; + + #[derive(Clone)] + struct BlockingPrefixBackend { + entered: Arc, + release: Arc, + } + + #[tonic::async_trait] + impl KvIndexerBackend for BlockingPrefixBackend { + async fn apply_external_kv_batch( + &self, + _request: ApplyExternalKvBatchRequest, + ) -> Result { + Ok(ApplyExternalKvBatchResponse::default()) + } + + async fn match_external_kv( + &self, + _request: MatchExternalKvRequest, + ) -> Result { + Ok(MatchExternalKvResponse::default()) + } + + async fn match_external_kv_prefix( + &self, + _request: MatchExternalKvPrefixRequest, + ) -> Result { + self.entered.fetch_add(1, Ordering::SeqCst); + let _permit = self + .release + .acquire() + .await + .expect("release semaphore closed"); + Ok(MatchExternalKvPrefixResponse::default()) + } + + async fn get_external_kv_hit_counts( + &self, + _request: GetExternalKvHitCountsRequest, + ) -> Result { + Ok(GetExternalKvHitCountsResponse::default()) + } + } + + /// Runs the arrival stamp, waits out `queued_for` to stand in for the time a + /// dispatched request spends waiting in the runtime, then serves it. + async fn serve_after_queueing( + service: &KvIndexerService, + caller_deadline: Duration, + queued_for: Duration, + ) -> Result, Status> { + let mut arriving = Request::new(()); + arriving.set_timeout(caller_deadline); + let (metadata, extensions, ()) = crate::admission::stamp_arrival(arriving) + .expect("arrival stamp never rejects") + .into_parts(); + + tokio::time::sleep(queued_for).await; + + let request = Request::from_parts( + metadata, + extensions, + MatchExternalKvPrefixRequest { + hashes: vec![-1], + max_blocks: 0, + }, + ); + KvIndexer::match_external_kv_prefix(service, request).await + } + + fn non_blocking_backend(entered: &Arc) -> BlockingPrefixBackend { + BlockingPrefixBackend { + entered: Arc::clone(entered), + release: Arc::new(Semaphore::new(Semaphore::MAX_PERMITS)), + } + } + + #[tokio::test] + async fn query_that_outlived_its_caller_never_reaches_the_backend() { + let entered = Arc::new(AtomicUsize::new(0)); + let service = KvIndexerService::new(non_blocking_backend(&entered)); + + let status = serve_after_queueing( + &service, + Duration::from_millis(20), + Duration::from_millis(60), + ) + .await + .unwrap_err(); + + assert_eq!(status.code(), tonic::Code::DeadlineExceeded); + assert_eq!(entered.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn query_still_inside_its_deadline_is_served() { + let entered = Arc::new(AtomicUsize::new(0)); + let service = KvIndexerService::new(non_blocking_backend(&entered)); + + serve_after_queueing(&service, Duration::from_secs(30), Duration::from_millis(10)) + .await + .expect("a query within its deadline must still be answered"); + + assert_eq!(entered.load(Ordering::SeqCst), 1); + } + + fn hbm() -> i32 { + crate::pb::TierType::TierHbm as i32 + } + + fn action(r#type: ExternalKvActionType, tier: i32, hashes: &[&str]) -> ExternalKvAction { + ExternalKvAction { + r#type: r#type as i32, + tier, + hashes: hashes.iter().map(|h| h.parse().unwrap()).collect(), + component_masks: Vec::new(), + block_sizes: Vec::new(), + } + } + + #[test] + fn validate_actions_allows_empty_batch() { + // An empty batch carries no mutation but is not an error. + assert!(validate_actions(&[]).is_ok()); + } + + #[test] + fn validate_actions_rejects_unknown_type() { + let actions = [action(ExternalKvActionType::ActionUnknown, hbm(), &["1"])]; + assert!(validate_actions(&actions).is_err()); + } + + #[test] + fn validate_actions_rejects_bad_tier() { + let actions = [action(ExternalKvActionType::ActionReport, 0, &["1"])]; + assert!(validate_actions(&actions).is_err()); + } + + #[test] + fn validate_actions_rejects_misaligned_side_arrays() { + let base = action(ExternalKvActionType::ActionReport, hbm(), &["1", "2"]); + assert!(validate_actions(std::slice::from_ref(&base)).is_ok()); + let mut aligned = base.clone(); + aligned.component_masks = vec![COMPONENT_FULL, COMPONENT_FULL]; + aligned.block_sizes = vec![16, 16]; + assert!(validate_actions(&[aligned]).is_ok()); + // A short component_masks array is a malformed batch, not a silent legacy. + let mut bad_masks = base.clone(); + bad_masks.component_masks = vec![COMPONENT_FULL]; + assert_eq!( + validate_actions(&[bad_masks]).unwrap_err().code(), + tonic::Code::InvalidArgument + ); + // A short block_sizes array is rejected too. + let mut bad_sizes = base; + bad_sizes.block_sizes = vec![16]; + assert_eq!( + validate_actions(&[bad_sizes]).unwrap_err().code(), + tonic::Code::InvalidArgument + ); + } + + #[test] + fn validate_actions_requires_hashes_for_report_and_revoke() { + assert!( + validate_actions(&[action(ExternalKvActionType::ActionReport, hbm(), &[])]).is_err() + ); + assert!( + validate_actions(&[action(ExternalKvActionType::ActionRevoke, hbm(), &[])]).is_err() + ); + } + + #[test] + fn validate_actions_allows_empty_hashes_for_clear_all_at_tier() { + let actions = [action( + ExternalKvActionType::ActionClearAllAtTier, + hbm(), + &[], + )]; + assert!(validate_actions(&actions).is_ok()); + } + + #[test] + fn validate_hashes_rejects_oversized_query() { + let hashes = vec![1; MAX_HASHES_PER_REQUEST + 1]; + let error = validate_hashes_bounded(&hashes).unwrap_err(); + assert_eq!(error.code(), tonic::Code::ResourceExhausted); + } + + /// Only the bounded variant rejects on length. + #[test] + fn validate_hashes_accepts_oversized_prefix_query() { + let hashes = vec![1; MAX_HASHES_PER_REQUEST + 1]; + assert!(validate_hashes(&hashes).is_ok()); + } + + #[test] + fn validate_actions_rejects_oversized_batch() { + let hashes = vec!["1"; MAX_HASHES_PER_REQUEST / 2 + 1]; + let actions = [ + action(ExternalKvActionType::ActionReport, hbm(), &hashes), + action(ExternalKvActionType::ActionReport, hbm(), &hashes), + ]; + let error = validate_actions(&actions).unwrap_err(); + assert_eq!(error.code(), tonic::Code::ResourceExhausted); + } + + #[test] + fn validate_actions_rejects_too_many_actions() { + let clear = action(ExternalKvActionType::ActionClearAllAtTier, hbm(), &[]); + let actions = vec![clear; MAX_ACTIONS_PER_BATCH + 1]; + let error = validate_actions(&actions).unwrap_err(); + assert_eq!(error.code(), tonic::Code::ResourceExhausted); + } + + #[test] + fn validate_worker_id_rejects_empty_value() { + assert!(validate_worker_id("").is_err()); + assert!(validate_worker_id("worker-1").is_ok()); + } + + // --- component-aware prefix rule engine --- + + fn dram() -> i32 { + crate::pb::TierType::TierDram as i32 + } + fn ssd() -> i32 { + crate::pb::TierType::TierSsd as i32 + } + + /// OR the tiers into a `1 << TierType` bitmask. + fn tmask(tiers: &[i32]) -> u32 { + tiers.iter().fold(0, |m, t| m | (1u32 << t)) + } + + /// A held block with `(tier, component bitmask)` placements and a token count. + fn blk(tiers: &[(i32, u32)], token_count: u32) -> Option { + Some(BlockComponents { + token_count, + tier_masks: tiers.to_vec(), + }) + } + + /// A legacy whole-block placement (mask 0) at HBM. + fn legacy_blk() -> Option { + blk(&[(hbm(), 0)], 0) + } + + fn spec( + components: u32, + swa_window_tokens: u32, + full_tiers: &[i32], + swa_tiers: &[i32], + mamba_tiers: &[i32], + ) -> WorkerCacheSpec { + WorkerCacheSpec { + version: 1, + components, + swa_window_tokens, + full_tier_mask: tmask(full_tiers), + swa_tier_mask: tmask(swa_tiers), + mamba_tier_mask: tmask(mamba_tiers), + } + } + + #[test] + fn legacy_no_spec_is_contiguous() { + let blocks = vec![legacy_blk(), legacy_blk(), legacy_blk(), None, legacy_blk()]; + assert_eq!(compute_worker_prefix(None, &blocks), 3); + } + + #[test] + fn component_report_without_spec_is_excluded() { + // A worker that reports components but declared no spec cannot be + // interpreted safely, so it contributes nothing (NoSignal-safe). + let blocks = vec![ + blk(&[(hbm(), COMPONENT_FULL)], 16), + blk(&[(hbm(), COMPONENT_FULL)], 16), + ]; + assert_eq!(compute_worker_prefix(None, &blocks), 0); + } + + #[test] + fn contiguous_full_stops_at_first_gap() { + let s = spec(COMPONENT_FULL, 0, &[hbm(), dram()], &[], &[]); + let blocks = vec![ + blk(&[(hbm(), COMPONENT_FULL)], 16), + blk(&[(dram(), COMPONENT_FULL)], 16), // full may live on a different servable tier + blk(&[(hbm(), COMPONENT_SWA)], 16), // no full here -> prefix stops + blk(&[(hbm(), COMPONENT_FULL)], 16), + ]; + assert_eq!(compute_worker_prefix(Some(&s), &blocks), 2); + } + + #[test] + fn ssd_only_is_not_servable_in_v1() { + let s = spec(COMPONENT_FULL, 0, &[hbm(), dram()], &[], &[]); + let blocks = vec![blk(&[(ssd(), COMPONENT_FULL)], 16)]; + assert_eq!(compute_worker_prefix(Some(&s), &blocks), 0); + } + + #[test] + fn trailing_window_requires_unbroken_window_before_boundary() { + // window = 100 tokens, 50 tokens per block: two contiguous swa blocks + // cover a window. full is present on every block. + let s = spec(COMPONENT_FULL | COMPONENT_SWA, 100, &[hbm()], &[hbm()], &[]); + let with_swa = || blk(&[(hbm(), COMPONENT_FULL | COMPONENT_SWA)], 50); + let no_swa = || blk(&[(hbm(), COMPONENT_FULL)], 50); + // swa present everywhere -> full length reusable. + let blocks = vec![with_swa(), with_swa(), with_swa(), with_swa(), with_swa()]; + assert_eq!(compute_worker_prefix(Some(&s), &blocks), 5); + // swa tombstoned at block index 3: the largest boundary whose trailing + // 100-token window is unbroken is N=3 (blocks 1..2 cover 100 tokens). + let holed = vec![with_swa(), with_swa(), with_swa(), no_swa(), with_swa()]; + assert_eq!(compute_worker_prefix(Some(&s), &holed), 3); + } + + #[test] + fn trailing_window_head_is_always_valid() { + // Fewer tokens than a window, but an unbroken run from the head is valid + // (matches the unified cache's window accumulator seeded at infinity). + let s = spec( + COMPONENT_FULL | COMPONENT_SWA, + 1000, + &[hbm()], + &[hbm()], + &[], + ); + let blocks = vec![blk(&[(hbm(), COMPONENT_FULL | COMPONENT_SWA)], 16); 2]; + assert_eq!(compute_worker_prefix(Some(&s), &blocks), 2); + } + + #[test] + fn exact_boundary_only_matches_at_a_checkpoint() { + // mamba lives only on the 4th block (a leaf checkpoint). full is on all. + let s = spec( + COMPONENT_FULL | COMPONENT_MAMBA, + 0, + &[hbm(), dram()], + &[], + &[hbm(), dram()], + ); + let blocks = vec![ + blk(&[(hbm(), COMPONENT_FULL)], 16), + blk(&[(hbm(), COMPONENT_FULL)], 16), + blk(&[(hbm(), COMPONENT_FULL)], 16), + blk(&[(hbm(), COMPONENT_FULL | COMPONENT_MAMBA)], 16), + ]; + assert_eq!(compute_worker_prefix(Some(&s), &blocks), 4); + // A shorter request that never reaches the checkpoint cannot reuse it. + assert_eq!(compute_worker_prefix(Some(&s), &blocks[..2]), 0); + } + + #[test] + fn unusable_specs_are_excluded() { + // Each of these declared specs is unusable and must fail closed: an empty + // component set, a future/unsupported version, and SWA without a window. + let blocks = vec![blk(&[(hbm(), COMPONENT_FULL | COMPONENT_SWA)], 16)]; + let empty = spec(0, 0, &[hbm()], &[], &[]); + let mut future = spec(COMPONENT_FULL, 0, &[hbm()], &[], &[]); + future.version = SUPPORTED_SPEC_VERSION + 1; + let swa_no_window = spec(COMPONENT_FULL | COMPONENT_SWA, 0, &[hbm()], &[hbm()], &[]); + for s in [empty, future, swa_no_window] { + assert_eq!(compute_worker_prefix(Some(&s), &blocks), 0); + } + } + + #[test] + fn missing_component_data_under_spec_excludes() { + // Spec requires full+swa but the worker reported legacy whole-block + // placement (mask 0), so full cannot be confirmed and it is excluded. + let s = spec(COMPONENT_FULL | COMPONENT_SWA, 100, &[hbm()], &[hbm()], &[]); + let blocks = vec![legacy_blk(), legacy_blk()]; + assert_eq!(compute_worker_prefix(Some(&s), &blocks), 0); + } +} diff --git a/experimental/sgl-router/sgl-kv-indexer/src/shutdown.rs b/experimental/sgl-router/sgl-kv-indexer/src/shutdown.rs new file mode 100644 index 000000000..bdc365e5d --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/src/shutdown.rs @@ -0,0 +1,42 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Termination signalling shared by the server and bridge binaries. + +use tracing::{info, warn}; + +/// Resolves once the process has been asked to terminate, on `SIGTERM` (which +/// container runtimes send before escalating to `SIGKILL`) or `SIGINT`. +/// +/// A handler that cannot be installed never resolves: firing immediately would +/// shut the process down at startup instead of leaving it running unsupervised. +pub async fn shutdown_signal() { + let interrupt = async { + if let Err(error) = tokio::signal::ctrl_c().await { + warn!(%error, "cannot handle SIGINT"); + std::future::pending::<()>().await; + } + }; + + #[cfg(unix)] + let terminate = async { + use tokio::signal::unix::{signal, SignalKind}; + match signal(SignalKind::terminate()) { + Ok(mut stream) => { + stream.recv().await; + } + Err(error) => { + warn!(%error, "cannot handle SIGTERM"); + std::future::pending::<()>().await; + } + } + }; + + #[cfg(not(unix))] + let terminate = std::future::pending::<()>(); + + tokio::select! { + () = interrupt => info!("received SIGINT; shutting down"), + () = terminate => info!("received SIGTERM; shutting down"), + } +} diff --git a/experimental/sgl-router/sgl-kv-indexer/tests/common/id.rs b/experimental/sgl-router/sgl-kv-indexer/tests/common/id.rs new file mode 100644 index 000000000..866969ede --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/tests/common/id.rs @@ -0,0 +1,11 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use std::time::{SystemTime, UNIX_EPOCH}; + +pub fn nanos() -> u128 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() +} diff --git a/experimental/sgl-router/sgl-kv-indexer/tests/common/kv.rs b/experimental/sgl-router/sgl-kv-indexer/tests/common/kv.rs new file mode 100644 index 000000000..4fa81e576 --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/tests/common/kv.rs @@ -0,0 +1,57 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use sgl_kv_indexer::pb::{ + ApplyExternalKvBatchRequest, ExternalKvAction, ExternalKvActionType, TierType, +}; + +pub fn hbm() -> i32 { + TierType::TierHbm as i32 +} + +pub fn dram() -> i32 { + TierType::TierDram as i32 +} + +pub fn action(kind: ExternalKvActionType, tier: i32, hashes: &[i64]) -> ExternalKvAction { + ExternalKvAction { + r#type: kind as i32, + tier, + hashes: hashes.to_vec(), + component_masks: Vec::new(), + block_sizes: Vec::new(), + } +} + +/// A component-aware REPORT action: each hash carries its component bitmask and +/// token count, index-aligned with `hashes`. +#[allow(dead_code)] // used by memory_integration, not grpc_contract +pub fn component_report( + tier: i32, + hashes: &[i64], + masks: &[u32], + block_sizes: &[u32], +) -> ExternalKvAction { + ExternalKvAction { + r#type: ExternalKvActionType::ActionReport as i32, + tier, + hashes: hashes.to_vec(), + component_masks: masks.to_vec(), + block_sizes: block_sizes.to_vec(), + } +} + +pub fn apply_request( + worker: &str, + address: &str, + seq: u64, + actions: Vec, +) -> ApplyExternalKvBatchRequest { + ApplyExternalKvBatchRequest { + worker_id: worker.to_string(), + seq, + actions, + worker_address: address.to_string(), + cache_spec: None, + } +} diff --git a/experimental/sgl-router/sgl-kv-indexer/tests/common/net.rs b/experimental/sgl-router/sgl-kv-indexer/tests/common/net.rs new file mode 100644 index 000000000..7c912ac49 --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/tests/common/net.rs @@ -0,0 +1,13 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use std::net::SocketAddr; + +/// Reserves an ephemeral loopback port. The listener is dropped immediately, +/// so callers that spawn a server should retain their connect-retry loop. +pub fn free_addr() -> SocketAddr { + std::net::TcpListener::bind("127.0.0.1:0") + .unwrap() + .local_addr() + .unwrap() +} diff --git a/experimental/sgl-router/sgl-kv-indexer/tests/grpc_contract.rs b/experimental/sgl-router/sgl-kv-indexer/tests/grpc_contract.rs new file mode 100644 index 000000000..b8fa8ba1c --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/tests/grpc_contract.rs @@ -0,0 +1,517 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! gRPC contract tests: exercise all four RPCs of the `KVIndexer` service +//! over the wire (real tonic server + client), not just the backend trait. + +#[path = "common/id.rs"] +mod test_id; +#[allow(dead_code)] +#[path = "common/kv.rs"] +mod test_kv; +#[path = "common/net.rs"] +mod test_net; + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use prost::Message; +use tokio::sync::Semaphore; +use tonic::transport::Server; +use tonic::{Code, Status}; + +use sgl_kv_indexer::pb::kv_indexer_client::KvIndexerClient; +use sgl_kv_indexer::pb::kv_indexer_server::KvIndexerServer; +use sgl_kv_indexer::pb::{ + ApplyExternalKvBatchRequest, ApplyExternalKvBatchResponse, ExternalKvAction, + ExternalKvActionType, GetExternalKvHitCountsRequest, GetExternalKvHitCountsResponse, + MatchExternalKvPrefixRequest, MatchExternalKvPrefixResponse, MatchExternalKvRequest, + MatchExternalKvResponse, +}; +use sgl_kv_indexer::{ + server_builder, GrpcPrefixIndex, InMemoryKvIndexerBackend, KvIndexerBackend, KvIndexerService, + PrefixIndex, PrefixIndexConfig, MAX_GRPC_DECODING_MESSAGE_SIZE, +}; +use test_id::nanos; +use test_kv::{action, apply_request, hbm}; +use test_net::free_addr; + +async fn start_backend( + backend: InMemoryKvIndexerBackend, +) -> KvIndexerClient { + let svc = KvIndexerService::new(backend).into_server(); + let addr = free_addr(); + tokio::spawn(async move { + server_builder() + .add_service(svc) + .serve(addr) + .await + .expect("server serve"); + }); + + let endpoint = format!("http://{addr}"); + for _ in 0..50 { + if let Ok(c) = KvIndexerClient::connect(endpoint.clone()).await { + return c; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + panic!("client failed to connect to {endpoint}"); +} + +#[derive(Clone)] +struct BlockingPrefixBackend { + entered: Arc, + release: Arc, +} + +#[tonic::async_trait] +impl KvIndexerBackend for BlockingPrefixBackend { + async fn apply_external_kv_batch( + &self, + _request: ApplyExternalKvBatchRequest, + ) -> Result { + Ok(ApplyExternalKvBatchResponse::default()) + } + + async fn match_external_kv( + &self, + _request: MatchExternalKvRequest, + ) -> Result { + Ok(MatchExternalKvResponse::default()) + } + + async fn match_external_kv_prefix( + &self, + _request: MatchExternalKvPrefixRequest, + ) -> Result { + self.entered.fetch_add(1, Ordering::SeqCst); + let _permit = self.release.acquire().await.expect("semaphore open"); + Ok(MatchExternalKvPrefixResponse::default()) + } + + async fn get_external_kv_hit_counts( + &self, + _request: GetExternalKvHitCountsRequest, + ) -> Result { + Ok(GetExternalKvHitCountsResponse::default()) + } +} + +async fn start_blocking_backend( + backend: BlockingPrefixBackend, +) -> KvIndexerClient { + let svc = KvIndexerService::with_prefix_query_max_inflight(backend, 2).into_server(); + let addr = free_addr(); + tokio::spawn(async move { + server_builder() + .add_service(svc) + .serve(addr) + .await + .expect("server serve"); + }); + + let endpoint = format!("http://{addr}"); + for _ in 0..50 { + if let Ok(client) = KvIndexerClient::connect(endpoint.clone()).await { + return client; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + panic!("client failed to connect to {endpoint}"); +} + +/// Starts a real gRPC server with isolated process-local state. +async fn start() -> KvIndexerClient { + start_backend(InMemoryKvIndexerBackend::new()).await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn prefix_limit_rejects_over_real_grpc_without_blocking_writes() { + let entered = Arc::new(AtomicUsize::new(0)); + let release = Arc::new(Semaphore::new(0)); + let backend = BlockingPrefixBackend { + entered: Arc::clone(&entered), + release: Arc::clone(&release), + }; + let client = start_blocking_backend(backend).await; + let request = || MatchExternalKvPrefixRequest { + hashes: vec![-1], + max_blocks: 0, + }; + + let mut first_client = client.clone(); + let first = tokio::spawn(async move { first_client.match_external_kv_prefix(request()).await }); + let mut second_client = client.clone(); + let second = + tokio::spawn(async move { second_client.match_external_kv_prefix(request()).await }); + + tokio::time::timeout(Duration::from_secs(1), async { + while entered.load(Ordering::SeqCst) != 2 { + tokio::task::yield_now().await; + } + }) + .await + .expect("two prefix queries should enter the backend"); + + let mut rejected_client = client.clone(); + let rejected = tokio::time::timeout( + Duration::from_secs(1), + rejected_client.match_external_kv_prefix(request()), + ) + .await + .expect("overload response should be immediate") + .expect_err("third prefix query should be rejected"); + assert_eq!(rejected.code(), Code::ResourceExhausted); + assert_eq!(entered.load(Ordering::SeqCst), 2); + + let mut write_client = client.clone(); + tokio::time::timeout( + Duration::from_secs(1), + write_client.apply_external_kv_batch(ApplyExternalKvBatchRequest { + worker_id: "worker".into(), + ..Default::default() + }), + ) + .await + .expect("writes should not share the prefix-query limit") + .expect("write should succeed"); + + release.add_permits(2); + first.await.expect("first task").expect("first response"); + second.await.expect("second task").expect("second response"); +} + +fn apply( + worker: &str, + addr: &str, + seq: u64, + action_type: ExternalKvActionType, + tier: i32, + hashes: &[i64], +) -> ApplyExternalKvBatchRequest { + apply_request(worker, addr, seq, vec![action(action_type, tier, hashes)]) +} + +fn apply_report( + worker: &str, + addr: &str, + seq: u64, + tier: i32, + hashes: &[i64], +) -> ApplyExternalKvBatchRequest { + apply( + worker, + addr, + seq, + ExternalKvActionType::ActionReport, + tier, + hashes, + ) +} + +#[tokio::test] +async fn multiple_workers_share_one_indexer_server() { + let mut indexer = start().await; + let suffix = nanos(); + let worker_0 = format!("worker-0-{suffix}"); + let worker_1 = format!("worker-1-{suffix}"); + let (hash_0, hash_1, shared_hash) = (1, 2, 3); + + indexer + .apply_external_kv_batch(apply_report( + &worker_0, + "10.0.0.1:9000", + 1, + hbm(), + &[hash_0, shared_hash], + )) + .await + .expect("apply worker-0"); + indexer + .apply_external_kv_batch(apply_report( + &worker_1, + "10.0.0.2:9000", + 1, + hbm(), + &[hash_1, shared_hash], + )) + .await + .expect("apply worker-1"); + + let response = indexer + .match_external_kv(MatchExternalKvRequest { + hashes: vec![hash_0, hash_1, shared_hash], + count_as_hit: false, + }) + .await + .expect("query indexer") + .into_inner(); + assert!(response + .matches + .iter() + .any(|entry| entry.worker_id == worker_0)); + assert!(response + .matches + .iter() + .any(|entry| entry.worker_id == worker_1)); + + // Keep one wire-level smoke check for hit counting; detailed counter + // semantics live in memory_integration.rs. + indexer + .match_external_kv(MatchExternalKvRequest { + hashes: vec![hash_0], + count_as_hit: true, + }) + .await + .expect("counting match over gRPC"); + let miss = 4; + let counts = indexer + .get_external_kv_hit_counts(GetExternalKvHitCountsRequest { + hashes: vec![hash_0, miss], + }) + .await + .expect("hit counts over gRPC") + .into_inner(); + let count = |hash: i64| { + counts + .entries + .iter() + .find(|entry| entry.hash == hash) + .map(|entry| entry.hit_count_total) + .unwrap_or(0) + }; + assert!(count(hash_0) >= 1, "matched hash should have a hit"); + assert_eq!(count(miss), 0, "unmatched hash must not be counted"); +} + +#[tokio::test] +async fn validation_errors_map_to_invalid_argument_over_grpc() { + let mut c = start().await; + + let err = c + .apply_external_kv_batch(apply_report("", "addr", 1, hbm(), &[1])) + .await + .expect_err("empty worker_id must be rejected"); + assert_eq!(err.code(), Code::InvalidArgument); + + // An action type outside the enum can only arrive over the wire; the + // in-process tests cover the mapped `ActionUnknown` variant instead. + let unmapped_action_type = ApplyExternalKvBatchRequest { + worker_id: "w".into(), + seq: 1, + worker_address: String::new(), + cache_spec: None, + actions: vec![ExternalKvAction { + r#type: 999, + tier: hbm(), + hashes: vec![1], + component_masks: Vec::new(), + block_sizes: Vec::new(), + }], + }; + let err = c + .apply_external_kv_batch(unmapped_action_type) + .await + .expect_err("unknown action type must be rejected"); + assert_eq!(err.code(), Code::InvalidArgument); +} + +#[tokio::test] +async fn match_prefix_over_grpc() { + let mut c = start().await; + let (w_long, w_short) = (format!("long-{}", nanos()), format!("short-{}", nanos())); + let (a, b, d) = (1, 2, 3); + + c.apply_external_kv_batch(apply_report(&w_long, "10.0.0.1:9000", 1, hbm(), &[a, b, d])) + .await + .expect("apply long"); + c.apply_external_kv_batch(apply_report(&w_short, "10.0.0.2:9000", 1, hbm(), &[a])) + .await + .expect("apply short"); + + let resp = c + .match_external_kv_prefix(MatchExternalKvPrefixRequest { + hashes: vec![a, b, d], + max_blocks: 0, + }) + .await + .expect("prefix ok") + .into_inner(); + + assert_eq!(resp.best_prefix_blocks, 3); + assert_eq!(resp.blocks_read, 3); + // Descending by prefix length: long (3) before short (1). + assert_eq!(resp.matches.len(), 2); + assert_eq!(resp.matches[0].worker_id, w_long); + assert_eq!(resp.matches[0].matched_prefix_blocks, 3); + assert_eq!(resp.matches[0].worker_address, "10.0.0.1:9000"); + assert_eq!(resp.matches[1].worker_id, w_short); + assert_eq!(resp.matches[1].matched_prefix_blocks, 1); +} + +#[tokio::test] +async fn prefix_query_scans_more_than_one_apply_chunk_over_grpc() { + const APPLY_CHUNK_SIZE: usize = 16_384; + + let mut indexer = start().await; + let hashes: Vec = (0..=APPLY_CHUNK_SIZE as i64).collect(); + for (seq, chunk) in hashes.chunks(APPLY_CHUNK_SIZE).enumerate() { + indexer + .apply_external_kv_batch(apply_report( + "large-prefix-worker", + "10.0.0.1:9000", + seq as u64, + hbm(), + chunk, + )) + .await + .expect("bounded apply chunk"); + } + + let response = indexer + .match_external_kv_prefix(MatchExternalKvPrefixRequest { + hashes, + max_blocks: 0, + }) + .await + .expect("prefix request larger than one apply chunk") + .into_inner(); + + assert_eq!(response.best_prefix_blocks as usize, APPLY_CHUNK_SIZE + 1); + assert_eq!(response.blocks_read as usize, APPLY_CHUNK_SIZE + 1); +} + +#[tokio::test] +async fn packed_signed_hash_query_can_exceed_tonics_default_receive_limit() { + const TONIC_DEFAULT_RECEIVE_LIMIT: usize = 4 * 1024 * 1024; + const HASH_COUNT: usize = 600_000; + + let mut indexer = start().await; + indexer + .apply_external_kv_batch(apply_report( + "large-wire-worker", + "10.0.0.1:9000", + 1, + hbm(), + &[-1], + )) + .await + .expect("store the signed first hash"); + + let mut hashes = Vec::with_capacity(HASH_COUNT); + hashes.push(-1); + hashes.extend((1..HASH_COUNT).map(|value| value as i64)); + let request = MatchExternalKvPrefixRequest { + hashes, + max_blocks: 0, + }; + assert!(request.encoded_len() > TONIC_DEFAULT_RECEIVE_LIMIT); + assert!(request.encoded_len() < MAX_GRPC_DECODING_MESSAGE_SIZE); + + let response = indexer + .match_external_kv_prefix(request) + .await + .expect("configured server accepts a packed request larger than 4 MiB") + .into_inner(); + assert_eq!(response.best_prefix_blocks, 1); +} + +/// Past the configured ceiling the server must answer OUT_OF_RANGE, because that +/// is the code the router maps to a degraded (cache-affinity-free) route rather +/// than to a failed request. A different code there would fail the request. +#[tokio::test] +async fn query_past_the_configured_limit_is_refused_as_out_of_range() { + let hash_count = MAX_GRPC_DECODING_MESSAGE_SIZE / std::mem::size_of::() + 1_024; + let request = MatchExternalKvPrefixRequest { + hashes: (0..hash_count).map(|value| value as i64).collect(), + max_blocks: 0, + }; + assert!(request.encoded_len() > MAX_GRPC_DECODING_MESSAGE_SIZE); + + let status = start() + .await + .match_external_kv_prefix(request) + .await + .expect_err("a request past the ceiling must be refused"); + assert_eq!(status.code(), Code::OutOfRange); +} + +/// Serves an empty backend behind an interceptor that records the `grpc-timeout` +/// of every request, and returns the router-facing client alongside the capture. +async fn start_recording_deadlines( + query_deadline: Duration, +) -> (GrpcPrefixIndex, Arc>>) { + let seen = Arc::new(Mutex::new(Vec::new())); + let recorder = Arc::clone(&seen); + let svc = KvIndexerServer::with_interceptor( + KvIndexerService::new(InMemoryKvIndexerBackend::new()), + move |request: tonic::Request<()>| { + if let Some(timeout) = request.metadata().get("grpc-timeout") { + recorder + .lock() + .expect("deadline recorder") + .push(timeout.to_str().expect("ascii timeout").to_string()); + } + Ok(request) + }, + ); + let addr = free_addr(); + tokio::spawn(async move { + Server::builder() + .add_service(svc) + .serve(addr) + .await + .expect("server serve"); + }); + + let endpoint = format!("http://{addr}"); + for _ in 0..50 { + if KvIndexerClient::connect(endpoint.clone()).await.is_ok() { + break; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + + let index = GrpcPrefixIndex::new(PrefixIndexConfig { + endpoint, + query_deadline, + max_inflight: sgl_kv_indexer::DEFAULT_QUERY_MAX_INFLIGHT, + }) + .expect("test endpoint is valid"); + (index, seen) +} + +/// The router-facing client must publish its deadline on the wire: that header is +/// the only thing letting the indexer shed a query whose caller gave up. +#[tokio::test] +async fn router_client_publishes_its_deadline_on_the_wire() { + let (index, seen) = start_recording_deadlines(Duration::from_millis(250)).await; + + index + .match_prefix(vec![1, 2, 3]) + .await + .expect("query reaches the indexer"); + + let seen = seen.lock().expect("deadline recorder").clone(); + assert_eq!( + seen.len(), + 1, + "exactly one query reached the server: {seen:?}" + ); + let raw = &seen[0]; + // Asserted structurally, not byte-for-byte: the wire spec lets the sender + // pick any unit that fits, so pinning tonic's choice would fail on a + // change that is still correct. + let (digits, unit) = raw.split_at(raw.len() - 1); + assert!( + matches!(unit, "H" | "M" | "S" | "m" | "u" | "n"), + "unit is one the wire spec defines: {raw:?}" + ); + let value: u64 = digits.parse().expect("timeout value is numeric"); + assert!( + value > 0, + "a budget of zero would shed every query: {raw:?}" + ); +} diff --git a/experimental/sgl-router/sgl-kv-indexer/tests/memory_integration.rs b/experimental/sgl-router/sgl-kv-indexer/tests/memory_integration.rs new file mode 100644 index 000000000..70f5966dd --- /dev/null +++ b/experimental/sgl-router/sgl-kv-indexer/tests/memory_integration.rs @@ -0,0 +1,886 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Integration tests for the process-local in-memory backend. + +#[path = "common/kv.rs"] +mod test_kv; + +use std::sync::Arc; + +use sgl_kv_indexer::pb::{ + ApplyExternalKvBatchRequest, ApplyExternalKvBatchResponse, ExternalKvActionType, + GetExternalKvHitCountsRequest, GetExternalKvHitCountsResponse, MatchExternalKvPrefixRequest, + MatchExternalKvPrefixResponse, MatchExternalKvRequest, MatchExternalKvResponse, + WorkerCacheSpec, +}; +use sgl_kv_indexer::{ + InMemoryKvIndexerBackend, KvIndexerBackend, WorkerPrefixInput, COMPONENT_FULL, COMPONENT_SWA, +}; +use test_kv::{action, apply_request as apply_req, component_report, dram, hbm}; +use tonic::Status; + +fn backend() -> InMemoryKvIndexerBackend { + InMemoryKvIndexerBackend::new() +} + +fn match_req(hs: &[i64], count_as_hit: bool) -> MatchExternalKvRequest { + MatchExternalKvRequest { + hashes: hs.to_vec(), + count_as_hit, + } +} + +/// Returns the tiers a worker holds a hash at, per the match response. +fn tiers_for(resp: &MatchExternalKvResponse, worker: &str, hash: i64) -> Vec { + let mut tiers = Vec::new(); + for m in &resp.matches { + if m.worker_id != worker { + continue; + } + for th in &m.hashes_by_tier { + if th.hashes.contains(&hash) { + tiers.push(th.tier); + } + } + } + tiers.sort_unstable(); + tiers +} + +macro_rules! itest { + ($name:ident, $b:ident, $body:block) => { + #[tokio::test] + async fn $name() { + let $b = backend(); + $body + } + }; +} + +itest!(report_then_match_returns_worker_and_address, b, { + b.apply_external_kv_batch(apply_req( + "w1", + "10.0.0.1:9000", + 1, + vec![action(ExternalKvActionType::ActionReport, hbm(), &[1, 2])], + )) + .await + .unwrap(); + + let resp = b + .match_external_kv(match_req(&[1, 2, 3], false)) + .await + .unwrap(); + assert_eq!(resp.matches.len(), 1); + let m = &resp.matches[0]; + assert_eq!(m.worker_id, "w1"); + assert_eq!(m.address, "10.0.0.1:9000"); + assert_eq!(tiers_for(&resp, "w1", 1), vec![hbm()]); + assert_eq!(tiers_for(&resp, "w1", 2), vec![hbm()]); + assert!(tiers_for(&resp, "w1", 3).is_empty()); +}); + +itest!(large_request_preserves_complete_ordered_results, b, { + // Exercise a large write and read while preserving complete ordered results. + let expected_hashes: Vec = (0..300).collect(); + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 1, + vec![action( + ExternalKvActionType::ActionReport, + hbm(), + &expected_hashes, + )], + )) + .await + .unwrap(); + + let resp = b + .match_external_kv(match_req(&expected_hashes, false)) + .await + .unwrap(); + let worker = resp + .matches + .iter() + .find(|m| m.worker_id == "w1") + .expect("worker must match"); + let tier = worker + .hashes_by_tier + .iter() + .find(|t| t.tier == hbm()) + .expect("HBM tier must match"); + assert_eq!(tier.hashes, expected_hashes); +}); + +itest!(duplicate_report_is_idempotent, b, { + for _ in 0..3 { + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 1, + vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])], + )) + .await + .unwrap(); + } + let resp = b.match_external_kv(match_req(&[1], false)).await.unwrap(); + assert_eq!(tiers_for(&resp, "w1", 1), vec![hbm()]); +}); + +itest!(identical_batch_replay_is_idempotent, b, { + // Stores, removes, then stores the same hash again; the net state is + // "stored". Re-delivering the identical batch must not change it. + let batch = apply_req( + "w1", + "a", + 7, + vec![ + action(ExternalKvActionType::ActionReport, hbm(), &[9]), + action(ExternalKvActionType::ActionRevoke, hbm(), &[9]), + action(ExternalKvActionType::ActionReport, hbm(), &[9]), + ], + ); + b.apply_external_kv_batch(batch.clone()).await.unwrap(); + let first = b.match_external_kv(match_req(&[9], false)).await.unwrap(); + b.apply_external_kv_batch(batch).await.unwrap(); + let second = b.match_external_kv(match_req(&[9], false)).await.unwrap(); + assert_eq!(tiers_for(&first, "w1", 9), vec![hbm()]); + assert_eq!(tiers_for(&second, "w1", 9), vec![hbm()]); +}); + +itest!(recomputed_full_node_restores_hbm_placement, b, { + // HiRadixCache lifecycle for an exact-match recomputation: + // BlockStored(GPU) -> BlockRemoved(GPU) -> BlockStored(GPU). + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 1, + vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])], + )) + .await + .unwrap(); + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 2, + vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[1])], + )) + .await + .unwrap(); + let evicted = b.match_external_kv(match_req(&[1], false)).await.unwrap(); + assert!(tiers_for(&evicted, "w1", 1).is_empty()); + + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 3, + vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])], + )) + .await + .unwrap(); + let restored = b.match_external_kv(match_req(&[1], false)).await.unwrap(); + assert_eq!(tiers_for(&restored, "w1", 1), vec![hbm()]); +}); + +itest!(recomputed_split_reports_only_materialized_hashes, b, { + // An evicted [prefix -> old suffix] is partially recomputed as + // [prefix -> new suffix]. The old suffix must remain absent. + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 1, + vec![action(ExternalKvActionType::ActionReport, hbm(), &[1, 2])], + )) + .await + .unwrap(); + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 2, + vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[1, 2])], + )) + .await + .unwrap(); + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 3, + vec![ + action(ExternalKvActionType::ActionReport, hbm(), &[1]), + action(ExternalKvActionType::ActionReport, hbm(), &[3]), + ], + )) + .await + .unwrap(); + + let result = b + .match_external_kv(match_req(&[1, 2, 3], false)) + .await + .unwrap(); + assert_eq!(tiers_for(&result, "w1", 1), vec![hbm()]); + assert!(tiers_for(&result, "w1", 2).is_empty()); + assert_eq!(tiers_for(&result, "w1", 3), vec![hbm()]); +}); + +itest!(recomputed_batch_replay_keeps_cpu_copy, b, { + // Re-materializing on GPU must not revoke the existing host backup, and + // re-delivering the same batch must leave both tiers unchanged. + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 1, + vec![action(ExternalKvActionType::ActionReport, dram(), &[1])], + )) + .await + .unwrap(); + let recomputed = apply_req( + "w1", + "a", + 2, + vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])], + ); + b.apply_external_kv_batch(recomputed.clone()).await.unwrap(); + b.apply_external_kv_batch(recomputed).await.unwrap(); + + let result = b.match_external_kv(match_req(&[1], false)).await.unwrap(); + assert_eq!(tiers_for(&result, "w1", 1), vec![hbm(), dram()]); +}); + +itest!(revoke_partial_tier_keeps_other_tier, b, { + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 1, + vec![ + action(ExternalKvActionType::ActionReport, hbm(), &[1]), + action(ExternalKvActionType::ActionReport, dram(), &[1]), + action(ExternalKvActionType::ActionRevoke, hbm(), &[1]), + ], + )) + .await + .unwrap(); + let resp = b.match_external_kv(match_req(&[1], false)).await.unwrap(); + assert_eq!(tiers_for(&resp, "w1", 1), vec![dram()]); +}); + +itest!(revoke_missing_hash_is_idempotent, b, { + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 1, + vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[404])], + )) + .await + .unwrap(); + let resp = b.match_external_kv(match_req(&[404], false)).await.unwrap(); + assert!(resp.matches.is_empty()); +}); + +itest!(multi_worker_multi_tier, b, { + b.apply_external_kv_batch(apply_req( + "w1", + "a1", + 1, + vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])], + )) + .await + .unwrap(); + b.apply_external_kv_batch(apply_req( + "w2", + "a2", + 1, + vec![action(ExternalKvActionType::ActionReport, dram(), &[1])], + )) + .await + .unwrap(); + let resp = b.match_external_kv(match_req(&[1], false)).await.unwrap(); + assert_eq!(resp.matches.len(), 2); + assert_eq!(tiers_for(&resp, "w1", 1), vec![hbm()]); + assert_eq!(tiers_for(&resp, "w2", 1), vec![dram()]); +}); + +itest!(clear_all_at_tier_removes_only_that_tier, b, { + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 1, + vec![ + action(ExternalKvActionType::ActionReport, hbm(), &[1, 2, 3]), + action(ExternalKvActionType::ActionReport, dram(), &[1]), + ], + )) + .await + .unwrap(); + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 2, + vec![action( + ExternalKvActionType::ActionClearAllAtTier, + hbm(), + &[], + )], + )) + .await + .unwrap(); + let resp = b + .match_external_kv(match_req(&[1, 2, 3], false)) + .await + .unwrap(); + assert_eq!(tiers_for(&resp, "w1", 1), vec![dram()]); + assert!(tiers_for(&resp, "w1", 2).is_empty()); + assert!(tiers_for(&resp, "w1", 3).is_empty()); +}); + +itest!( + count_as_hit_only_counts_matched_and_replay_does_not_double, + b, + { + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 1, + vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])], + )) + .await + .unwrap(); + + // Diagnostic match (count_as_hit=false) must not count. + b.match_external_kv(match_req(&[1, 2], false)) + .await + .unwrap(); + let counts = b + .get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1, 2] }) + .await + .unwrap(); + assert!(counts.entries.is_empty()); + + // Counting match: only the matched hash "1" is counted, "2" (a miss) is not. + b.match_external_kv(match_req(&[1, 2], true)).await.unwrap(); + let counts = b + .get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1, 2] }) + .await + .unwrap(); + assert_eq!(counts.entries.len(), 1); + assert_eq!(counts.entries[0].hash, 1); + assert_eq!(counts.entries[0].hit_count_total, 1); + + // Replaying the apply batch must not touch hit counts. + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 1, + vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])], + )) + .await + .unwrap(); + let counts = b + .get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1] }) + .await + .unwrap(); + assert_eq!(counts.entries[0].hit_count_total, 1); + } +); + +itest!(full_revoke_drops_hit_key, b, { + // Report a block, count a hit (creates the co-located :h key), then fully + // revoke it. The hit key must go with the placement, or a + // matched-then-evicted block leaks its counter forever. + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 1, + vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])], + )) + .await + .unwrap(); + + // Counting match creates the hit key with c=1. + b.match_external_kv(match_req(&[1], true)).await.unwrap(); + let counts = b + .get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1] }) + .await + .unwrap(); + assert_eq!(counts.entries.len(), 1); + assert_eq!(counts.entries[0].hit_count_total, 1); + + // Fully revoke the block: placement empties, so the hit key must go too. + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 2, + vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[1])], + )) + .await + .unwrap(); + + let resp = b.match_external_kv(match_req(&[1], false)).await.unwrap(); + assert!(resp.matches.is_empty()); + + // Hit key is gone too: a leaked :h would still report a count here. + let counts = b + .get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1] }) + .await + .unwrap(); + assert!( + counts.entries.is_empty(), + "hit key leaked after full revoke: {:?}", + counts.entries + ); +}); + +itest!(partial_revoke_keeps_hit_key, b, { + // Block present at two tiers; count a hit, then revoke only one tier. Placement + // is still non-empty, so the hit key must survive (guard against over-deletion). + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 1, + vec![ + action(ExternalKvActionType::ActionReport, hbm(), &[1]), + action(ExternalKvActionType::ActionReport, dram(), &[1]), + ], + )) + .await + .unwrap(); + + b.match_external_kv(match_req(&[1], true)).await.unwrap(); + + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 2, + vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[1])], + )) + .await + .unwrap(); + + let counts = b + .get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1] }) + .await + .unwrap(); + assert_eq!(counts.entries.len(), 1); + assert_eq!(counts.entries[0].hit_count_total, 1); +}); + +itest!(batch_action_order_is_preserved, b, { + // revoke-then-report on the same hash within one batch must net to "stored". + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 1, + vec![ + action(ExternalKvActionType::ActionRevoke, hbm(), &[5]), + action(ExternalKvActionType::ActionReport, hbm(), &[5]), + ], + )) + .await + .unwrap(); + let resp = b.match_external_kv(match_req(&[5], false)).await.unwrap(); + assert_eq!(tiers_for(&resp, "w1", 5), vec![hbm()]); + + // report-then-revoke on the same hash must net to "absent". + b.apply_external_kv_batch(apply_req( + "w1", + "a", + 2, + vec![ + action(ExternalKvActionType::ActionReport, hbm(), &[6]), + action(ExternalKvActionType::ActionRevoke, hbm(), &[6]), + ], + )) + .await + .unwrap(); + let resp = b.match_external_kv(match_req(&[6], false)).await.unwrap(); + assert!(tiers_for(&resp, "w1", 6).is_empty()); +}); + +// --- prefix query: backend override vs. the trait's default implementation --- +// +// The trait default is the written semantics and the backend override is a read +// optimization, so they must agree field-for-field on the parts that ARE the +// contract (per-worker prefix set and best_prefix_blocks). `blocks_read` is +// observability and legitimately differs, so it is not compared. + +/// Delegates every RPC to an in-memory backend EXCEPT `match_external_kv_prefix`, +/// which it leaves to the trait default — giving a reference answer computed from +/// the same state the optimized path reads. +struct DefaultViaMemory(Arc); + +#[tonic::async_trait] +impl KvIndexerBackend for DefaultViaMemory { + async fn apply_external_kv_batch( + &self, + request: ApplyExternalKvBatchRequest, + ) -> Result { + self.0.apply_external_kv_batch(request).await + } + + async fn match_external_kv( + &self, + request: MatchExternalKvRequest, + ) -> Result { + self.0.match_external_kv(request).await + } + + // Delegate the component-aware read to the same backend so the trait + // default computes over the same placement and specs the fast path sees. + async fn collect_worker_prefix_inputs( + &self, + hashes: &[i64], + ) -> Result, Status> { + self.0.collect_worker_prefix_inputs(hashes).await + } + + async fn get_external_kv_hit_counts( + &self, + request: GetExternalKvHitCountsRequest, + ) -> Result { + self.0.get_external_kv_hit_counts(request).await + } +} + +fn shared_state_pair() -> (Arc, DefaultViaMemory) { + let backend = Arc::new(InMemoryKvIndexerBackend::new()); + let reference = DefaultViaMemory(Arc::clone(&backend)); + (backend, reference) +} + +fn prefix_req(hs: &[i64]) -> MatchExternalKvPrefixRequest { + MatchExternalKvPrefixRequest { + hashes: hs.to_vec(), + max_blocks: 0, + } +} + +/// Sorted `(worker_id, matched_prefix_blocks)` — the semantic content of a +/// prefix response, independent of `blocks_read`. +fn prefix_pairs(resp: &MatchExternalKvPrefixResponse) -> Vec<(String, u32)> { + let mut pairs: Vec<(String, u32)> = resp + .matches + .iter() + .map(|m| (m.worker_id.clone(), m.matched_prefix_blocks)) + .collect(); + pairs.sort(); + pairs +} + +fn report(worker: &str, addr: &str, seq: u64, hs: &[i64]) -> ApplyExternalKvBatchRequest { + apply_req( + worker, + addr, + seq, + vec![action(ExternalKvActionType::ActionReport, hbm(), hs)], + ) +} + +#[tokio::test] +async fn prefix_fast_path_matches_default_impl() { + let (fast, reference) = shared_state_pair(); + + // Nested prefixes (hole-free), a diverging branch, and a hole. + fast.apply_external_kv_batch(report("w-long", "10.0.0.1:1", 1, &[1, 2, 3, 4])) + .await + .unwrap(); + fast.apply_external_kv_batch(report("w-short", "10.0.0.2:1", 1, &[1, 2])) + .await + .unwrap(); + // w-hole holds 1, 3, 4 but not 2: strict prefix must be 1. + fast.apply_external_kv_batch(report("w-hole", "10.0.0.3:1", 1, &[1, 3, 4])) + .await + .unwrap(); + // w-noaddr is unroutable and must be excluded by both paths. + fast.apply_external_kv_batch(report("w-noaddr", "", 1, &[1, 2])) + .await + .unwrap(); + + let query = [1, 2, 3, 4]; + let fast_resp = fast + .match_external_kv_prefix(prefix_req(&query)) + .await + .unwrap(); + let ref_resp = reference + .match_external_kv_prefix(prefix_req(&query)) + .await + .unwrap(); + + assert_eq!(prefix_pairs(&fast_resp), prefix_pairs(&ref_resp)); + assert_eq!(fast_resp.best_prefix_blocks, ref_resp.best_prefix_blocks); + assert_eq!(fast_resp.best_prefix_blocks, 4); + assert_eq!( + prefix_pairs(&fast_resp), + vec![ + ("w-hole".to_string(), 1), + ("w-long".to_string(), 4), + ("w-short".to_string(), 2), + ] + ); + assert!(fast_resp + .matches + .iter() + .all(|m| !m.worker_address.is_empty())); + // Descending order and first-block read are part of the response contract. + assert_eq!(fast_resp.matches[0].matched_prefix_blocks, 4); + assert!(fast_resp.blocks_read >= 1); +} + +#[tokio::test] +async fn prefix_first_block_miss_reads_one_block() { + let b = backend(); + // No worker holds the first queried block; the scan stops after one read. + b.apply_external_kv_batch(report("w1", "10.0.0.1:1", 1, &[2, 3])) + .await + .unwrap(); + let resp = b + .match_external_kv_prefix(prefix_req(&[1, 2, 3])) + .await + .unwrap(); + assert!(resp.matches.is_empty()); + assert_eq!(resp.best_prefix_blocks, 0); + assert_eq!(resp.blocks_read, 1); +} + +#[tokio::test] +async fn prefix_max_blocks_caps_the_scan() { + let b = backend(); + b.apply_external_kv_batch(report("w1", "10.0.0.1:1", 1, &[1, 2, 3, 4])) + .await + .unwrap(); + let resp = b + .match_external_kv_prefix(MatchExternalKvPrefixRequest { + hashes: vec![1, 2, 3, 4], + max_blocks: 2, + }) + .await + .unwrap(); + // Capped at 2 even though the worker holds all four. + assert_eq!(resp.best_prefix_blocks, 2); + assert_eq!(resp.blocks_read, 2); + assert_eq!(resp.matches.len(), 1); + assert_eq!(resp.matches[0].matched_prefix_blocks, 2); +} + +// --- component-aware placement & prefix ------------------------------------- + +/// A hybrid-SWA spec: full servable from HBM+DRAM, swa a 100-token trailing +/// window servable from HBM. +fn swa_spec() -> WorkerCacheSpec { + WorkerCacheSpec { + version: 1, + components: COMPONENT_FULL | COMPONENT_SWA, + swa_window_tokens: 100, + full_tier_mask: (1 << hbm()) | (1 << dram()), + swa_tier_mask: 1 << hbm(), + mamba_tier_mask: 0, + } +} + +fn apply_with_spec( + worker: &str, + addr: &str, + seq: u64, + spec: WorkerCacheSpec, + actions: Vec, +) -> ApplyExternalKvBatchRequest { + let mut req = apply_req(worker, addr, seq, actions); + req.cache_spec = Some(spec); + req +} + +#[tokio::test] +async fn component_prefix_matches_default_impl() { + let (fast, reference) = shared_state_pair(); + + // Four full blocks (50 tokens each); swa present on all but the 4th, so the + // largest boundary with an unbroken 100-token swa window is 3. + let report = component_report( + hbm(), + &[1, 2, 3, 4], + &[ + COMPONENT_FULL | COMPONENT_SWA, + COMPONENT_FULL | COMPONENT_SWA, + COMPONENT_FULL | COMPONENT_SWA, + COMPONENT_FULL, + ], + &[50, 50, 50, 50], + ); + fast.apply_external_kv_batch(apply_with_spec( + "w-swa", + "10.0.0.1:1", + 1, + swa_spec(), + vec![report], + )) + .await + .unwrap(); + + let query = [1, 2, 3, 4]; + let fast_resp = fast + .match_external_kv_prefix(prefix_req(&query)) + .await + .unwrap(); + let ref_resp = reference + .match_external_kv_prefix(prefix_req(&query)) + .await + .unwrap(); + + assert_eq!(prefix_pairs(&fast_resp), prefix_pairs(&ref_resp)); + assert_eq!(fast_resp.best_prefix_blocks, ref_resp.best_prefix_blocks); + assert_eq!(prefix_pairs(&fast_resp), vec![("w-swa".to_string(), 3)]); +} + +#[tokio::test] +async fn partial_eviction_replace_shrinks_component_set() { + let b = backend(); + // Store full+swa, then restate to full-only (partial swa eviction) via a + // REPLACE snapshot for the same (hash, tier). No BlockRemoved is involved. + b.apply_external_kv_batch(apply_with_spec( + "w1", + "10.0.0.1:1", + 1, + swa_spec(), + vec![component_report( + hbm(), + &[1, 2], + &[ + COMPONENT_FULL | COMPONENT_SWA, + COMPONENT_FULL | COMPONENT_SWA, + ], + &[80, 80], + )], + )) + .await + .unwrap(); + // Both blocks reusable (window 100 met by 2x80 tokens; head rule also holds). + let before = b + .match_external_kv_prefix(prefix_req(&[1, 2])) + .await + .unwrap(); + assert_eq!(before.best_prefix_blocks, 2); + + // Restate the second block to full only: swa gone there. + b.apply_external_kv_batch(apply_with_spec( + "w1", + "10.0.0.1:1", + 2, + swa_spec(), + vec![component_report(hbm(), &[2], &[COMPONENT_FULL], &[80])], + )) + .await + .unwrap(); + // That block has no swa now, and its trailing window (only 80 < 100) is not + // headed, so the largest valid boundary drops to 1. + let after = b + .match_external_kv_prefix(prefix_req(&[1, 2])) + .await + .unwrap(); + assert_eq!(after.best_prefix_blocks, 1); + + let snapshot = b + .match_external_kv(match_req(&[1, 2], false)) + .await + .unwrap(); + let tier = &snapshot.matches[0].hashes_by_tier[0]; + assert_eq!(tier.hashes, vec![1, 2]); + assert_eq!( + tier.component_masks, + vec![COMPONENT_FULL | COMPONENT_SWA, COMPONENT_FULL] + ); + assert_eq!(tier.block_sizes, vec![80, 80]); +} + +#[tokio::test] +async fn component_aware_worker_without_spec_is_excluded() { + let b = backend(); + // Report component-aware placement but never send a spec: the worker cannot + // be interpreted and must be excluded (NoSignal-safe), never over-reported. + b.apply_external_kv_batch(apply_req( + "w1", + "10.0.0.1:1", + 1, + vec![component_report( + hbm(), + &[1, 2], + &[COMPONENT_FULL, COMPONENT_FULL], + &[16, 16], + )], + )) + .await + .unwrap(); + let resp = b + .match_external_kv_prefix(prefix_req(&[1, 2])) + .await + .unwrap(); + assert!(resp.matches.is_empty()); + assert_eq!(resp.best_prefix_blocks, 0); +} + +#[tokio::test] +async fn duplicate_hash_in_one_report_keeps_last_snapshot() { + let b = backend(); + // A single REPORT action naming the same hash twice (a coalesced + // store+restate): the LAST snapshot must win deterministically, never a race. + b.apply_external_kv_batch(apply_with_spec( + "w1", + "10.0.0.1:1", + 1, + swa_spec(), + vec![component_report( + hbm(), + &[1, 1], + &[COMPONENT_FULL | COMPONENT_SWA, COMPONENT_FULL], + &[80, 80], + )], + )) + .await + .unwrap(); + // The hash ends as full-only (last snapshot); with swa required and a lone 80-token + // block that is not a full head window, the boundary requiring swa fails, + // so no reusable prefix. + let resp = b.match_external_kv_prefix(prefix_req(&[1])).await.unwrap(); + assert_eq!(resp.best_prefix_blocks, 0); +} + +#[tokio::test] +async fn absent_spec_batch_clears_stored_spec() { + let b = backend(); + // First a component-aware batch establishes a spec + a reusable block. + b.apply_external_kv_batch(apply_with_spec( + "w1", + "10.0.0.1:1", + 1, + swa_spec(), + vec![component_report( + hbm(), + &[1], + &[COMPONENT_FULL | COMPONENT_SWA], + &[200], + )], + )) + .await + .unwrap(); + assert_eq!( + b.match_external_kv_prefix(prefix_req(&[1])) + .await + .unwrap() + .best_prefix_blocks, + 1 + ); + + // A later batch with NO spec (worker reverted to legacy) must clear the old + // spec. The still-component-aware placement can then no longer be interpreted + // (component data but no spec) -> fail closed, never scored on stale rules. + b.apply_external_kv_batch(apply_req( + "w1", + "10.0.0.1:1", + 2, + vec![action(ExternalKvActionType::ActionReport, dram(), &[2])], + )) + .await + .unwrap(); + let resp = b.match_external_kv_prefix(prefix_req(&[1])).await.unwrap(); + assert!(resp.matches.is_empty()); + assert_eq!(resp.best_prefix_blocks, 0); +} diff --git a/experimental/sgl-router/src/config/cli.rs b/experimental/sgl-router/src/config/cli.rs index 4cb7be1ed..8d0e02b2d 100644 --- a/experimental/sgl-router/src/config/cli.rs +++ b/experimental/sgl-router/src/config/cli.rs @@ -12,10 +12,14 @@ use std::num::NonZeroU32; use crate::config::{ default_cb_cool_down, default_proxy_request_timeout_secs, default_stale_request_timeout_secs, resolve_mode, ActiveLoadConfig, CacheAwareConfig, CircuitBreakerConfig, Config, - DiscoveryBackend, K8sDiscoveryConfig, LogFormat, ModelConfig, ObservabilityConfig, PolicyKind, - ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig, + DiscoveryBackend, K8sDiscoveryConfig, KvIndexerEndpointConfig, LogFormat, ModelConfig, + ObservabilityConfig, PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, + StickyConfig, }; +const DEFAULT_KV_INDEXER_QUERY_TIMEOUT_MS: u64 = 100; +const DEFAULT_KV_INDEXER_QUERY_MAX_INFLIGHT: usize = sgl_kv_indexer::DEFAULT_QUERY_MAX_INFLIGHT; + /// `sgl-router` — slim KV-aware OpenAI-compatible router for SGLang workers. /// /// Discovery is mutually exclusive: pass `--worker-urls` for a static @@ -69,6 +73,18 @@ pub struct Cli { /// Multiplicative load spread gating the absolute balance check. #[arg(long)] pub balance_rel_threshold: Option, + /// External KV indexer gRPC endpoint used as the authoritative cache signal. + /// Needs an explicit scheme, e.g. `http://10.0.0.1:50051`. + #[arg(long)] + pub kv_indexer_endpoint: Option, + /// KV Indexer query timeout in milliseconds. Requires + /// `--kv-indexer-endpoint`; defaults to 100. + #[arg(long)] + pub kv_indexer_query_timeout_ms: Option, + /// Maximum concurrent KV Indexer queries issued by this Router. Requires + /// `--kv-indexer-endpoint`; defaults to 32. + #[arg(long)] + pub kv_indexer_query_max_inflight: Option, // ---- sticky-session policy (only used by `--policy sticky`) ---- /// Request header carrying the routing key for sticky-session routing. @@ -157,14 +173,35 @@ impl Cli { } let tuned_cache_aware = self.cache_threshold.is_some() || self.balance_abs_threshold.is_some() - || self.balance_rel_threshold.is_some(); + || self.balance_rel_threshold.is_some() + || self.kv_indexer_endpoint.is_some() + || self.kv_indexer_query_timeout_ms.is_some() + || self.kv_indexer_query_max_inflight.is_some(); if tuned_cache_aware && self.policy != PolicyKind::CacheAwareZmq { return Err(anyhow!( - "--cache-threshold / --balance-abs-threshold / --balance-rel-threshold \ - require --policy cache_aware_zmq" + "cache-aware tuning flags require --policy cache_aware_zmq" + )); + } + if self.kv_indexer_query_timeout_ms == Some(0) { + return Err(anyhow!( + "--kv-indexer-query-timeout-ms must be greater than zero" + )); + } + if self.kv_indexer_query_timeout_ms.is_some() && self.kv_indexer_endpoint.is_none() { + return Err(anyhow!( + "--kv-indexer-query-timeout-ms requires --kv-indexer-endpoint" + )); + } + if self.kv_indexer_query_max_inflight == Some(0) { + return Err(anyhow!( + "--kv-indexer-query-max-inflight must be greater than zero" + )); + } + if self.kv_indexer_query_max_inflight.is_some() && self.kv_indexer_endpoint.is_none() { + return Err(anyhow!( + "--kv-indexer-query-max-inflight requires --kv-indexer-endpoint" )); } - let tuned_sticky = self.routing_key_header.is_some() || self.sticky_fallback_policy.is_some() || self.sticky_idle_secs.is_some() @@ -234,6 +271,12 @@ impl Cli { // Only build a CacheAwareConfig when the operator tuned at least // one knob; otherwise leave it None so the policy uses its own // defaults. Unset knobs fall back to the per-field defaults. + let kv_indexer_query_timeout_ms = self + .kv_indexer_query_timeout_ms + .unwrap_or(DEFAULT_KV_INDEXER_QUERY_TIMEOUT_MS); + let kv_indexer_query_max_inflight = self + .kv_indexer_query_max_inflight + .unwrap_or(DEFAULT_KV_INDEXER_QUERY_MAX_INFLIGHT); let cache_aware = if tuned_cache_aware { let d = CacheAwareConfig::default(); Some(CacheAwareConfig { @@ -244,6 +287,11 @@ impl Cli { balance_rel_threshold: self .balance_rel_threshold .unwrap_or(d.balance_rel_threshold), + kv_indexer_endpoint: self.kv_indexer_endpoint.map(|url| KvIndexerEndpointConfig { + url, + query_timeout_ms: kv_indexer_query_timeout_ms, + query_max_inflight: kv_indexer_query_max_inflight, + }), }) } else { None @@ -744,6 +792,109 @@ mod tests { assert_eq!(ca.balance_abs_threshold, 32); } + #[test] + fn kv_indexer_reuses_cache_aware_policy_config() { + let c = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--policy", + "cache_aware_zmq", + "--kv-indexer-endpoint", + "http://indexer:50051", + "--kv-indexer-query-timeout-ms", + "75", + "--kv-indexer-query-max-inflight", + "17", + ])) + .unwrap(); + let cache = c.model.cache_aware.expect("cache-aware config"); + let indexer = cache.kv_indexer_endpoint.expect("Indexer config"); + assert_eq!(indexer.url, "http://indexer:50051"); + assert_eq!(indexer.query_timeout_ms, 75); + assert_eq!(indexer.query_max_inflight, 17); + } + + #[test] + fn kv_indexer_uses_safe_query_defaults() { + let c = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--policy", + "cache_aware_zmq", + "--kv-indexer-endpoint", + "http://indexer:50051", + ])) + .unwrap(); + let indexer = c + .model + .cache_aware + .expect("cache-aware config") + .kv_indexer_endpoint + .expect("Indexer config"); + assert_eq!(indexer.query_timeout_ms, 100); + assert_eq!(indexer.query_max_inflight, 32); + } + + #[test] + fn kv_indexer_requires_cache_aware_policy() { + let err = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--kv-indexer-endpoint", + "http://indexer:50051", + ])) + .unwrap_err() + .to_string(); + assert!(err.contains("require --policy cache_aware_zmq")); + } + + #[test] + fn kv_indexer_timeout_requires_endpoint() { + let err = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--policy", + "cache_aware_zmq", + "--kv-indexer-query-timeout-ms", + "75", + ])) + .unwrap_err() + .to_string(); + assert!(err.contains("requires --kv-indexer-endpoint")); + } + + #[test] + fn kv_indexer_max_inflight_requires_endpoint() { + let err = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--policy", + "cache_aware_zmq", + "--kv-indexer-query-max-inflight", + "17", + ])) + .unwrap_err() + .to_string(); + assert!(err.contains("requires --kv-indexer-endpoint")); + } + + #[test] + fn kv_indexer_max_inflight_must_be_positive() { + let err = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--policy", + "cache_aware_zmq", + "--kv-indexer-endpoint", + "http://indexer:50051", + "--kv-indexer-query-max-inflight", + "0", + ])) + .unwrap_err() + .to_string(); + assert!(err.contains("must be greater than zero")); + } + #[test] fn no_cache_aware_flags_leaves_none() { let c = into_config_owned(with_model(&[ diff --git a/experimental/sgl-router/src/config/types.rs b/experimental/sgl-router/src/config/types.rs index f4c414be7..519f78769 100644 --- a/experimental/sgl-router/src/config/types.rs +++ b/experimental/sgl-router/src/config/types.rs @@ -146,7 +146,7 @@ pub struct ModelConfig { pub tokenizer_path: String, pub policy: PolicyKind, pub circuit_breaker: Option, - /// Tuning for the cache-aware ZMQ policy. Ignored unless + /// Tuning for cache-aware routing. Ignored unless /// `policy = "cache_aware_zmq"`. `None` falls back to defaults at /// policy construction time. pub cache_aware: Option, @@ -157,8 +157,16 @@ pub struct ModelConfig { pub sticky: Option, } -/// Per-model cache-aware-ZMQ tuning. -#[derive(Debug, Clone, Copy)] +/// External KV Indexer client settings. +#[derive(Debug, Clone)] +pub struct KvIndexerEndpointConfig { + pub url: String, + pub query_timeout_ms: u64, + pub query_max_inflight: usize, +} + +/// Per-model cache-aware tuning. +#[derive(Debug, Clone)] pub struct CacheAwareConfig { /// Lower bound on `matched_blocks / total_blocks` for the tree match /// to win the selection. Below this, the policy falls back to @@ -174,6 +182,9 @@ pub struct CacheAwareConfig { /// that the absolute check is gated on. Default 1.1 — 10 % relative /// difference triggers re-balancing. pub balance_rel_threshold: f32, + /// Optional external KV Indexer client configuration. When configured, it + /// replaces the local ZMQ radix tree as the cache signal. + pub kv_indexer_endpoint: Option, } impl Default for CacheAwareConfig { @@ -182,6 +193,7 @@ impl Default for CacheAwareConfig { cache_threshold: default_cache_threshold(), balance_abs_threshold: default_balance_abs(), balance_rel_threshold: default_balance_rel(), + kv_indexer_endpoint: None, } } } diff --git a/experimental/sgl-router/src/main.rs b/experimental/sgl-router/src/main.rs index 08ea5472e..e1fb8b08f 100644 --- a/experimental/sgl-router/src/main.rs +++ b/experimental/sgl-router/src/main.rs @@ -98,19 +98,43 @@ async fn main() -> Result<()> { ); let registry = Arc::new(sgl_router::workers::WorkerRegistry::default()); + let prefix_index = cfg + .model + .cache_aware + .as_ref() + .and_then(|cache| cache.kv_indexer_endpoint.as_ref()) + .map(|indexer| { + let config = sgl_kv_indexer::PrefixIndexConfig { + endpoint: indexer.url.clone(), + query_deadline: std::time::Duration::from_millis(indexer.query_timeout_ms), + max_inflight: indexer.query_max_inflight, + }; + sgl_kv_indexer::GrpcPrefixIndex::new(config) + .map(Arc::new) + .context("configure KV Indexer client") + }) + .transpose()?; // Build the KV-event index up front so the cache-aware-zmq policy can - // share its `HashTree` handle + `BlockSizeOracle`. When no model uses - // `cache_aware_zmq`, the index is still constructed (cheap) but no - // subscribers are ever added. + // share its `HashTree` handle + `BlockSizeOracle`. An external Indexer makes + // the local tree irrelevant to routing, so only discover hash metadata rather + // than duplicating every KV event. let block_size_oracle = sgl_router::policies::kv_events::BlockSizeOracle::new(); - let kv_index = sgl_router::policies::kv_events::KvEventIndex::new_with_http_and_oracle( - reqwest::Client::builder() - .timeout(std::time::Duration::from_secs(2)) - .build() - .expect("default http client builds"), - Arc::clone(&block_size_oracle), - ); + let kv_event_http = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(2)) + .build() + .expect("default http client builds"); + let kv_index = if prefix_index.is_some() { + sgl_router::policies::kv_events::KvEventIndex::new_metadata_only_with_http_and_oracle( + kv_event_http, + Arc::clone(&block_size_oracle), + ) + } else { + sgl_router::policies::kv_events::KvEventIndex::new_with_http_and_oracle( + kv_event_http, + Arc::clone(&block_size_oracle), + ) + }; let policies = Arc::new( sgl_router::policies::factory::build_registry( &cfg, @@ -163,16 +187,17 @@ async fn main() -> Result<()> { .context("build proxy client")?, ); - let ctx = Arc::new( - sgl_router::server::app_context::AppContext::with_active_load( - cfg.clone(), - tokenizers, - proxy, - registry, - policies, - active_load, - ), + let mut app_ctx = sgl_router::server::app_context::AppContext::with_active_load( + cfg.clone(), + tokenizers, + proxy, + registry, + policies, + active_load, ); + app_ctx.prefix_index = prefix_index; + app_ctx.block_size_oracle = block_size_oracle; + let ctx = Arc::new(app_ctx); ctx.mark_ready(); let app = sgl_router::server::app::build_router(ctx.clone()); diff --git a/experimental/sgl-router/src/policies/cache_aware_zmq.rs b/experimental/sgl-router/src/policies/cache_aware_zmq.rs index 3ebf51183..f8430134b 100644 --- a/experimental/sgl-router/src/policies/cache_aware_zmq.rs +++ b/experimental/sgl-router/src/policies/cache_aware_zmq.rs @@ -130,6 +130,46 @@ impl CacheAwareZmqPolicy { let rel_threshold = (min_load as f32 * self.config.balance_rel_threshold) as usize; abs_diff > self.config.balance_abs_threshold && max_load > rel_threshold } + + fn select_external( + &self, + workers: &[Arc], + ctx: &SelectionContext<'_>, + signal: &crate::policies::ExternalPrefixSignal, + ) -> Option> { + let sgl_kv_indexer::PrefixOutcome::Matched { matches, .. } = &signal.outcome else { + return None; + }; + if signal.query_blocks == 0 { + return None; + } + + // The index may include unhealthy workers or workers in another pool. + let best_routable_blocks = matches + .iter() + .filter(|m| workers.iter().any(|worker| worker.url == m.address)) + .map(|m| m.matched_prefix_blocks) + .max() + .unwrap_or(0); + + let match_rate = best_routable_blocks as f32 / signal.query_blocks as f32; + if let Some(metrics) = self.metrics.get() { + metrics.observe_overlap_blocks(ctx.model().0.as_str(), best_routable_blocks as u64); + } + if match_rate <= self.config.cache_threshold { + return None; + } + + workers + .iter() + .filter(|worker| { + matches.iter().any(|m| { + m.matched_prefix_blocks == best_routable_blocks && m.address == worker.url + }) + }) + .min_by_key(|worker| worker.active_load()) + .cloned() + } } impl Policy for CacheAwareZmqPolicy { @@ -144,6 +184,14 @@ impl Policy for CacheAwareZmqPolicy { return Self::pick_min_load(workers); } + // An external signal is authoritative: an empty/unusable result + // degrades only to min-load and never consults the local radix tree. + if let Some(signal) = ctx.external_prefix() { + return self + .select_external(workers, ctx, signal) + .or_else(|| Self::pick_min_load(workers)); + } + // 2. Routing tokens. Prefer the ids computed once at ingress; fall // back to tokenizing the body here so the policy stays usable for // callers that don't pre-tokenize (e.g. unit tests). In production @@ -263,6 +311,7 @@ mod tests { cache_threshold: 0.5, balance_abs_threshold: 32, balance_rel_threshold: 1.1, + kv_indexer_endpoint: None, } } @@ -350,6 +399,70 @@ mod tests { assert_eq!(chosen.url, "http://w1:30000"); } + #[test] + fn external_prefix_signal_skips_unroutable_best_match() { + let mut config = cfg_default(); + config.cache_threshold = 0.0; + let policy = CacheAwareZmqPolicy::new( + config, + Arc::new(HashTree::new()), + tokenizer_registry_with_tiny(), + oracle_for_tests(4), + ); + let w0 = worker("http://w0:30000", "tiny"); + let w1 = worker("http://w1:30000", "tiny"); + let _load = w0.load_guard(); + let workers = vec![Arc::clone(&w0), Arc::clone(&w1)]; + let signal = crate::policies::ExternalPrefixSignal { + outcome: sgl_kv_indexer::PrefixOutcome::Matched { + matches: vec![ + sgl_kv_indexer::PrefixMatch { + address: "http://gone:30000".into(), + matched_prefix_blocks: 4, + worker_id: "gone".into(), + }, + sgl_kv_indexer::PrefixMatch { + address: w0.url.clone(), + matched_prefix_blocks: 2, + worker_id: "w0".into(), + }, + ], + best_prefix_blocks: 4, + }, + query_blocks: 4, + }; + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None).with_external_prefix(Some(&signal)); + let chosen = policy.select(&workers, &ctx).expect("must pick"); + assert_eq!(chosen.url, w0.url); + } + + #[test] + fn external_empty_result_uses_min_load_without_local_tree() { + let tree = Arc::new(HashTree::new()); + let registry = tokenizer_registry_with_tiny(); + let text = "hello world hello world hello world"; + let tok = registry.get("tiny").unwrap(); + let ids = adapter::encode(&tok, text).unwrap(); + let hashes = compute_block_hashes(&ids, 4); + tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes); + + let policy = CacheAwareZmqPolicy::new(cfg_default(), tree, registry, oracle_for_tests(4)); + let w0 = worker("http://w0:30000", "tiny"); + let w1 = worker("http://w1:30000", "tiny"); + let _load = w0.load_guard(); + let workers = vec![Arc::clone(&w0), Arc::clone(&w1)]; + let signal = crate::policies::ExternalPrefixSignal { + outcome: sgl_kv_indexer::PrefixOutcome::Empty, + query_blocks: 4, + }; + let model = ModelId("tiny".into()); + let body = serde_json::to_vec(&serde_json::json!({"prompt": text})).unwrap(); + let ctx = SelectionContext::new(&model, Some(&body)).with_external_prefix(Some(&signal)); + let chosen = policy.select(&workers, &ctx).expect("must pick"); + assert_eq!(chosen.url, w1.url); + } + /// Tree contains w0's prefix; cache-aware selection picks w0 even /// though w1 has lower load (the load skew is below the imbalance /// threshold, so cache wins). @@ -377,6 +490,7 @@ mod tests { cache_threshold: 0.0, // any match counts balance_abs_threshold: 32, balance_rel_threshold: 1.1, + kv_indexer_endpoint: None, }, tree, registry, @@ -414,6 +528,7 @@ mod tests { cache_threshold: 0.0, balance_abs_threshold: 32, balance_rel_threshold: 1.1, + kv_indexer_endpoint: None, }, tree, registry, @@ -459,6 +574,7 @@ mod tests { cache_threshold: 0.0, balance_abs_threshold: 32, balance_rel_threshold: 1.1, + kv_indexer_endpoint: None, }, tree, toks, @@ -510,6 +626,7 @@ mod tests { cache_threshold: 1.0, // match_rate <= 1.0 always -> always fall back balance_abs_threshold: 32, balance_rel_threshold: 1.1, + kv_indexer_endpoint: None, }, tree, toks, @@ -593,6 +710,7 @@ mod tests { cache_threshold: 0.0, balance_abs_threshold: 32, balance_rel_threshold: 1.1, + kv_indexer_endpoint: None, }, tree, Arc::clone(®istry), @@ -632,6 +750,7 @@ mod tests { cache_threshold: 0.0, balance_abs_threshold: 32, balance_rel_threshold: 1.1, + kv_indexer_endpoint: None, }, tree, Arc::clone(®istry), @@ -690,6 +809,7 @@ mod tests { cache_threshold: 0.0, balance_abs_threshold: 32, balance_rel_threshold: 1.1, + kv_indexer_endpoint: None, }, tree, registry, @@ -762,6 +882,7 @@ mod tests { cache_threshold: 0.0, balance_abs_threshold: 32, balance_rel_threshold: 1.1, + kv_indexer_endpoint: None, }, tree, registry, @@ -800,6 +921,7 @@ mod tests { cache_threshold: 0.0, balance_abs_threshold: 32, balance_rel_threshold: 1.1, + kv_indexer_endpoint: None, }, tree, registry, @@ -903,6 +1025,7 @@ mod tests { cache_threshold: 0.0, balance_abs_threshold: 32, balance_rel_threshold: 1.1, + kv_indexer_endpoint: None, }, tree, registry, @@ -938,6 +1061,7 @@ mod tests { cache_threshold: 0.0, // would normally always match balance_abs_threshold: 5, balance_rel_threshold: 2.0, + kv_indexer_endpoint: None, }, tree, registry, @@ -1063,6 +1187,7 @@ mod tests { cache_threshold: 0.99, balance_abs_threshold: 32, balance_rel_threshold: 1.1, + kv_indexer_endpoint: None, }, tree, tokenizer_registry_with_tiny(), @@ -1146,6 +1271,7 @@ mod tests { cache_threshold: 0.0, balance_abs_threshold: 32, balance_rel_threshold: 1.1, + kv_indexer_endpoint: None, }, tree.clone(), registry, @@ -1243,6 +1369,7 @@ mod tests { cache_threshold: 0.0, balance_abs_threshold: 32, balance_rel_threshold: 1.1, + kv_indexer_endpoint: None, }, tree, registry, diff --git a/experimental/sgl-router/src/policies/factory.rs b/experimental/sgl-router/src/policies/factory.rs index 498f892f9..f9d8ca663 100644 --- a/experimental/sgl-router/src/policies/factory.rs +++ b/experimental/sgl-router/src/policies/factory.rs @@ -65,7 +65,7 @@ pub fn build_policy( PolicyKind::PowerOfTwo => Arc::new(PowerOfTwoChoicesPolicy::new()), PolicyKind::LoadBased => Arc::new(LoadBasedPolicy::new()), PolicyKind::CacheAwareZmq => { - let cache_cfg = model.cache_aware.unwrap_or_default(); + let cache_cfg = model.cache_aware.clone().unwrap_or_default(); Arc::new(CacheAwareZmqPolicy::new( cache_cfg, tree, diff --git a/experimental/sgl-router/src/policies/kv_events/index.rs b/experimental/sgl-router/src/policies/kv_events/index.rs index de979b6f3..a91e1d205 100644 --- a/experimental/sgl-router/src/policies/kv_events/index.rs +++ b/experimental/sgl-router/src/policies/kv_events/index.rs @@ -65,6 +65,7 @@ struct WorkerEntry { /// routing path entirely. pub struct KvEventIndex { tree: Arc, + maintain_tree: bool, subscribers: Arc, pump: Mutex>>, pump_cancel: CancellationToken, @@ -114,6 +115,24 @@ impl KvEventIndex { pub fn new_with_http_and_oracle( http: reqwest::Client, block_size_oracle: Arc, + ) -> Arc { + Self::new_with_mode(http, block_size_oracle, true) + } + + /// Discovers worker hash metadata only: seeds the shared [`BlockSizeOracle`] + /// but neither subscribes to KV events nor maintains the local tree, because + /// an external Indexer is the routing signal. + pub fn new_metadata_only_with_http_and_oracle( + http: reqwest::Client, + block_size_oracle: Arc, + ) -> Arc { + Self::new_with_mode(http, block_size_oracle, false) + } + + fn new_with_mode( + http: reqwest::Client, + block_size_oracle: Arc, + maintain_tree: bool, ) -> Arc { let tree = Arc::new(HashTree::new()); let (tx, rx) = mpsc::channel::(EVENT_CHANNEL_BUFFER); @@ -130,6 +149,7 @@ impl KvEventIndex { )); Arc::new(Self { tree, + maintain_tree, subscribers, pump: Mutex::new(Some(pump)), pump_cancel, @@ -208,6 +228,15 @@ impl KvEventIndex { // hash KV blocks over token bigrams, so the policy must use the bigram // hasher for its query hashes to match the worker's stored hashes. self.block_size_oracle.set_bigram(cfg.is_bigram); + if !self.maintain_tree { + info!( + worker_url = %worker_url, + block_size = cfg.block_size, + is_bigram = cfg.is_bigram, + "kv-events: external Indexer configured; discovered hash metadata without subscribing" + ); + return; + } info!( worker_url = %worker_url, dp_size = cfg.dp_size, @@ -710,4 +739,28 @@ mod tests { ); index.shutdown().await; } + + #[tokio::test] + async fn metadata_only_mode_seeds_oracle_without_registering_subscribers() { + let oracle = BlockSizeOracle::new(); + let index = KvEventIndex::new_metadata_only_with_http_and_oracle( + reqwest::Client::new(), + Arc::clone(&oracle), + ); + let cfg = EventConfig { + host: "127.0.0.1".into(), + port_base: 30400, + topic: "kv-events".into(), + block_size: 64, + dp_size: 2, + is_bigram: true, + }; + + index.add_worker("http://127.0.0.1:30400", Some(cfg)).await; + + assert_eq!(oracle.get(), Some(64)); + assert!(oracle.is_bigram()); + assert_eq!(index.known_worker_count(), 0); + index.shutdown().await; + } } diff --git a/experimental/sgl-router/src/policies/mod.rs b/experimental/sgl-router/src/policies/mod.rs index 8b58861d3..ec2cd3c9b 100644 --- a/experimental/sgl-router/src/policies/mod.rs +++ b/experimental/sgl-router/src/policies/mod.rs @@ -34,6 +34,13 @@ pub struct RequestTokens { pub engine_equivalent: bool, } +/// External indexer answer prepared by the async ingress path for the +/// synchronous cache-aware policy. +pub struct ExternalPrefixSignal { + pub outcome: sgl_kv_indexer::PrefixOutcome, + pub query_blocks: usize, +} + /// Produce the routing tokens — and whether they are engine-equivalent — /// from an already-parsed request body, using the shared tokenizer registry. /// @@ -181,6 +188,7 @@ pub struct SelectionContext<'a> { request_body: Option<&'a [u8]>, routing_key: Option<&'a str>, request_tokens: Option<&'a [u32]>, + external_prefix: Option<&'a ExternalPrefixSignal>, } impl<'a> SelectionContext<'a> { @@ -190,6 +198,7 @@ impl<'a> SelectionContext<'a> { request_body, routing_key: None, request_tokens: None, + external_prefix: None, } } @@ -203,6 +212,7 @@ impl<'a> SelectionContext<'a> { request_body, routing_key, request_tokens: None, + external_prefix: None, } } @@ -214,6 +224,14 @@ impl<'a> SelectionContext<'a> { self } + pub fn with_external_prefix( + mut self, + external_prefix: Option<&'a ExternalPrefixSignal>, + ) -> Self { + self.external_prefix = external_prefix; + self + } + pub fn model(&self) -> &ModelId { self.model } @@ -231,6 +249,10 @@ impl<'a> SelectionContext<'a> { pub fn request_tokens(&self) -> Option<&[u32]> { self.request_tokens } + + pub fn external_prefix(&self) -> Option<&ExternalPrefixSignal> { + self.external_prefix + } } pub trait Policy: Send + Sync + std::fmt::Debug { diff --git a/experimental/sgl-router/src/server/app_context.rs b/experimental/sgl-router/src/server/app_context.rs index 04778aa04..a655913f0 100644 --- a/experimental/sgl-router/src/server/app_context.rs +++ b/experimental/sgl-router/src/server/app_context.rs @@ -4,6 +4,7 @@ use crate::config::Config; use crate::policies::active_load::ActiveLoadRegistry; +use crate::policies::kv_events::BlockSizeOracle; use crate::policies::PolicyRegistry; use crate::proxy::Proxy; use crate::server::metrics::MetricsRegistry; @@ -12,7 +13,6 @@ use crate::workers::WorkerRegistry; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; -#[derive(Debug)] pub struct AppContext { pub config: Config, pub tokenizers: Arc, @@ -30,6 +30,8 @@ pub struct AppContext { /// (active_load gauge + stale_requests_total), and PD resolver /// (decode_affinity_total). pub metrics: Arc, + pub prefix_index: Option>, + pub block_size_oracle: Arc, ready: AtomicBool, } @@ -82,6 +84,8 @@ impl AppContext { policies, active_load, metrics, + prefix_index: None, + block_size_oracle: BlockSizeOracle::new(), ready: AtomicBool::new(false), } } @@ -127,6 +131,8 @@ impl AppContext { policies: Arc::new(PolicyRegistry::default()), active_load: ActiveLoadRegistry::with_defaults(), metrics: MetricsRegistry::new(), + prefix_index: None, + block_size_oracle: BlockSizeOracle::new(), ready: AtomicBool::new(false), } } diff --git a/experimental/sgl-router/src/server/routes/chat.rs b/experimental/sgl-router/src/server/routes/chat.rs index bb27d2f11..1fce3ec26 100644 --- a/experimental/sgl-router/src/server/routes/chat.rs +++ b/experimental/sgl-router/src/server/routes/chat.rs @@ -2,8 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 use crate::discovery::{ModelId, WorkerMode}; +use crate::policies::kv_events::{compute_block_hashes, compute_block_hashes_bigram}; use crate::policies::registry::{PdPoolResolver, PdResolveError}; -use crate::policies::{request_tokens_for, RequestTokens, SelectionContext}; +use crate::policies::{request_tokens_for, ExternalPrefixSignal, RequestTokens, SelectionContext}; use crate::server::app_context::AppContext; use crate::server::error::ApiError; use crate::server::metrics::{ @@ -16,6 +17,7 @@ use axum::http::{HeaderMap, HeaderName, HeaderValue, Response}; use bytes::Bytes; use serde::de::IgnoredAny; use serde::Deserialize; +use sgl_kv_indexer::PrefixIndex; use std::collections::HashMap; use std::sync::Arc; @@ -167,6 +169,34 @@ pub async fn chat_completions( let request_tokens = request_value .as_ref() .and_then(|v| request_tokens_for(&ctx.tokenizers, &model_id, v)); + let external_prefix = match ( + ctx.prefix_index.as_ref(), + request_tokens.as_ref(), + ctx.block_size_oracle.get(), + ) { + (Some(index), Some(tokens), Some(block_size)) => { + let hashes = if ctx.block_size_oracle.is_bigram() { + compute_block_hashes_bigram(&tokens.ids, block_size as usize) + } else { + compute_block_hashes(&tokens.ids, block_size as usize) + }; + let query_blocks = hashes.len(); + let outcome = if hashes.is_empty() { + sgl_kv_indexer::PrefixOutcome::Empty + } else { + resolve_prefix_query(index.match_prefix(hashes).await, &model_str)? + }; + Some(ExternalPrefixSignal { + outcome, + query_blocks, + }) + } + (Some(_), _, _) => Some(ExternalPrefixSignal { + outcome: sgl_kv_indexer::PrefixOutcome::Empty, + query_blocks: 0, + }), + _ => None, + }; // Sticky-session routing key. When the sticky policy is configured, // read the routing key from the operator-chosen header into the @@ -181,7 +211,8 @@ pub async fn chat_completions( .and_then(|v| v.to_str().ok()) .filter(|s| !s.is_empty()); let selection_ctx = SelectionContext::with_routing_key(&model_id, Some(&body), routing_key) - .with_request_tokens(request_tokens.as_ref().map(|t| t.ids.as_slice())); + .with_request_tokens(request_tokens.as_ref().map(|t| t.ids.as_slice())) + .with_external_prefix(external_prefix.as_ref()); let worker = policy .select(&workers, &selection_ctx) @@ -620,6 +651,42 @@ pub async fn chat_completions( } } +fn resolve_prefix_query( + result: Result, + model: &str, +) -> Result { + use sgl_kv_indexer::PrefixIndexError; + match result { + Ok(outcome) => Ok(outcome), + // The prefix hit only improves worker choice, so an indexer that is + // shedding, slow, or down costs cache affinity — not availability. + Err( + error @ (PrefixIndexError::Overloaded + | PrefixIndexError::Timeout + | PrefixIndexError::Unreachable), + ) => { + tracing::warn!(%model, error = %error, "KV Indexer unavailable; falling back to min-load routing"); + Ok(sgl_kv_indexer::PrefixOutcome::Empty) + } + // A prompt too long to fit one gRPC message is still a prompt a worker + // can serve, so it costs cache affinity like the cases above. Logged + // separately because the remedy is operational — raise the indexer's + // message limit — rather than waiting for the indexer to recover. + Err(error @ PrefixIndexError::QueryTooLarge) => { + tracing::warn!(%model, error = %error, "prompt exceeds the KV Indexer query size limit; falling back to min-load routing"); + Ok(sgl_kv_indexer::PrefixOutcome::Empty) + } + // A rejection means the router and the indexer disagree on the request + // contract; degrading would hide that from every request. + Err(error) => { + tracing::warn!(%model, error = %error, "KV Indexer rejected the query"); + Err(ApiError::PolicySelectionFailed { + model: model.to_string(), + }) + } + } +} + /// Estimate prefill-token count from the raw request body for use as /// the active-load `prefill_load` counter. Returns 1 at minimum so /// a registered request always shows up as "load > 0" — under-counting @@ -911,6 +978,38 @@ fn parse_probe(body: &Bytes) -> Result { mod tests { use super::*; + /// An unavailable indexer must never fail a request that min-load routing + /// can still serve. `QueryTooLarge` belongs here too: a prompt that outgrows + /// the query's message limit loses cache affinity, not availability. + #[test] + fn unavailable_indexer_degrades_to_empty_prefix_signal() { + for error in [ + sgl_kv_indexer::PrefixIndexError::Overloaded, + sgl_kv_indexer::PrefixIndexError::Timeout, + sgl_kv_indexer::PrefixIndexError::Unreachable, + sgl_kv_indexer::PrefixIndexError::QueryTooLarge, + ] { + assert_eq!( + resolve_prefix_query(Err(error.clone()), "tiny").unwrap(), + sgl_kv_indexer::PrefixOutcome::Empty, + "{error} should degrade" + ); + } + } + + #[test] + fn rejected_indexer_query_still_fails_selection() { + assert!(matches!( + resolve_prefix_query( + Err(sgl_kv_indexer::PrefixIndexError::Rejected( + sgl_kv_indexer::RpcCode::InvalidArgument + )), + "tiny" + ), + Err(ApiError::PolicySelectionFailed { .. }) + )); + } + /// `generate_room_id` MUST return values in `[0, i64::MAX]`. The /// SGLang prefill stores `bootstrap_room` as `torch.int64`; a u64 /// with the top bit set would wrap negative on the engine side. diff --git a/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs b/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs index 7c6d81b82..2f4e7e6fe 100644 --- a/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs +++ b/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs @@ -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), diff --git a/experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py b/experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py index a46f44c79..5534b428d 100644 --- a/experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py +++ b/experimental/sgl-router/tests/e2e/chat_completions/test_two_router_convergence.py @@ -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 "" + 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) diff --git a/experimental/sgl-router/tests/e2e/infra/gateway.py b/experimental/sgl-router/tests/e2e/infra/gateway.py index b580975b4..af2cb2409 100644 --- a/experimental/sgl-router/tests/e2e/infra/gateway.py +++ b/experimental/sgl-router/tests/e2e/infra/gateway.py @@ -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] diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/Dockerfile.router b/experimental/sgl-router/tests/e2e/k8s_integration/Dockerfile.router index fe8a870f8..78a84b0c0 100644 --- a/experimental/sgl-router/tests/e2e/k8s_integration/Dockerfile.router +++ b/experimental/sgl-router/tests/e2e/k8s_integration/Dockerfile.router @@ -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 diff --git a/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs b/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs index 06f605812..d311a5213 100644 --- a/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs +++ b/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs @@ -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 { let cfg = config(); let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap()); diff --git a/experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs b/experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs new file mode 100644 index 000000000..456a4e659 --- /dev/null +++ b/experimental/sgl-router/tests/proxy/common/cache_aware_fixture.rs @@ -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(), + } +} diff --git a/experimental/sgl-router/tests/proxy/common/mod.rs b/experimental/sgl-router/tests/proxy/common/mod.rs index 9ad730c43..8c63c5722 100644 --- a/experimental/sgl-router/tests/proxy/common/mod.rs +++ b/experimental/sgl-router/tests/proxy/common/mod.rs @@ -3,5 +3,6 @@ //! Shared test harness re-exports. +pub mod cache_aware_fixture; pub mod mock_worker; pub mod streaming; diff --git a/experimental/sgl-router/tests/proxy/external_indexer_routing.rs b/experimental/sgl-router/tests/proxy/external_indexer_routing.rs new file mode 100644 index 000000000..54501c9d2 --- /dev/null +++ b/experimental/sgl-router/tests/proxy/external_indexer_routing.rs @@ -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(); +} diff --git a/experimental/sgl-router/tests/proxy/main.rs b/experimental/sgl-router/tests/proxy/main.rs index 00cde22bd..4b7d440ce 100644 --- a/experimental/sgl-router/tests/proxy/main.rs +++ b/experimental/sgl-router/tests/proxy/main.rs @@ -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;