[Feature] Add process-local in-memory KV indexer and Router integration (#33370)
Co-authored-by: Wu, Yutong <yutong.wu@amd.com> Co-authored-by: TianDi101 <ditian12@amd.com> Co-authored-by: Zhangheng <hzh0425@apache.org>
This commit is contained in:
co-authored by
Wu, Yutong
TianDi101
Zhangheng
parent
238ba40c27
commit
360d10d6bc
@@ -108,6 +108,9 @@ jobs:
|
|||||||
steps:
|
steps:
|
||||||
- uses: actions/checkout@v4
|
- uses: actions/checkout@v4
|
||||||
|
|
||||||
|
- name: Install protoc
|
||||||
|
run: sudo bash scripts/ci/utils/install_protoc.sh
|
||||||
|
|
||||||
- name: Verify dynamo-* SHA pin
|
- name: Verify dynamo-* SHA pin
|
||||||
run: |
|
run: |
|
||||||
set -euo pipefail
|
set -euo pipefail
|
||||||
@@ -160,17 +163,17 @@ jobs:
|
|||||||
|
|
||||||
- name: cargo check
|
- name: cargo check
|
||||||
working-directory: experimental/sgl-router
|
working-directory: experimental/sgl-router
|
||||||
run: cargo check --all-targets
|
run: cargo check --workspace --all-targets
|
||||||
|
|
||||||
- name: cargo clippy
|
- name: cargo clippy
|
||||||
working-directory: experimental/sgl-router
|
working-directory: experimental/sgl-router
|
||||||
run: cargo clippy --all-targets -- -D warnings
|
run: cargo clippy --workspace --all-targets -- -D warnings
|
||||||
|
|
||||||
- name: cargo fmt
|
- name: cargo fmt
|
||||||
working-directory: experimental/sgl-router
|
working-directory: experimental/sgl-router
|
||||||
run: |
|
run: |
|
||||||
rustup toolchain install nightly --profile minimal --component rustfmt
|
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.
|
# cargo-deny 0.19.6+ required for CVSS 4.0 parsing.
|
||||||
- name: Install cargo-deny
|
- name: Install cargo-deny
|
||||||
@@ -196,6 +199,9 @@ jobs:
|
|||||||
steps:
|
steps:
|
||||||
- uses: actions/checkout@v4
|
- uses: actions/checkout@v4
|
||||||
|
|
||||||
|
- name: Install protoc
|
||||||
|
run: sudo bash scripts/ci/utils/install_protoc.sh
|
||||||
|
|
||||||
- name: Install OS deps (libssl, pkg-config)
|
- name: Install OS deps (libssl, pkg-config)
|
||||||
run: |
|
run: |
|
||||||
sudo apt-get update
|
sudo apt-get update
|
||||||
@@ -246,7 +252,7 @@ jobs:
|
|||||||
# the Qwen3-0.6B tokenizer.json. Filter is matched against the
|
# the Qwen3-0.6B tokenizer.json. Filter is matched against the
|
||||||
# full test path `tokenizer::parity::parity_matrix`; substring
|
# full test path `tokenizer::parity::parity_matrix`; substring
|
||||||
# `parity_matrix` is unique to that test.
|
# `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
|
# Regenerate the cross-impl block-hash parity fixture and fail if it
|
||||||
# differs from the committed file. The Python script replicates
|
# differs from the committed file. The Python script replicates
|
||||||
@@ -366,7 +372,7 @@ jobs:
|
|||||||
working-directory: experimental/sgl-router
|
working-directory: experimental/sgl-router
|
||||||
run: |
|
run: |
|
||||||
source "$HOME/.cargo/env"
|
source "$HOME/.cargo/env"
|
||||||
cargo build --release
|
cargo build --release --workspace
|
||||||
|
|
||||||
# Install SGLang from the local checkout in editable mode (no PyPI
|
# Install SGLang from the local checkout in editable mode (no PyPI
|
||||||
# version pin) — mirrors `pr-test-rust.yml` so the router e2e runs
|
# version pin) — mirrors `pr-test-rust.yml` so the router e2e runs
|
||||||
|
|||||||
@@ -140,7 +140,7 @@ repos:
|
|||||||
pass_filenames: false
|
pass_filenames: false
|
||||||
- id: rustfmt-sgl-router
|
- id: rustfmt-sgl-router
|
||||||
name: rustfmt experimental/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
|
language: system
|
||||||
files: ^experimental/sgl-router/.*\.rs$
|
files: ^experimental/sgl-router/.*\.rs$
|
||||||
pass_filenames: false
|
pass_filenames: false
|
||||||
|
|||||||
@@ -34,21 +34,28 @@ ARG DEBIAN_VERSION=bookworm
|
|||||||
######################## STAGE 1 — chef recipe ##########################
|
######################## STAGE 1 — chef recipe ##########################
|
||||||
FROM rust:${RUST_VERSION}-${DEBIAN_VERSION} AS chef
|
FROM rust:${RUST_VERSION}-${DEBIAN_VERSION} AS chef
|
||||||
RUN cargo install cargo-chef --locked --version ^0.1
|
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/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
|
# Stub a minimal src tree so cargo can resolve the workspace, generate
|
||||||
# the lockfile (gitignored upstream), then prepare the chef recipe.
|
# 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 "" > 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 generate-lockfile \
|
||||||
&& cargo chef prepare --recipe-path recipe.json \
|
&& cargo chef prepare --recipe-path recipe.json \
|
||||||
&& rm -rf src
|
&& rm -rf src sgl-kv-indexer/src
|
||||||
|
|
||||||
######################## STAGE 2 — builder ##############################
|
######################## STAGE 2 — builder ##############################
|
||||||
FROM rust:${RUST_VERSION}-${DEBIAN_VERSION} AS builder
|
FROM rust:${RUST_VERSION}-${DEBIAN_VERSION} AS builder
|
||||||
RUN cargo install cargo-chef --locked --version ^0.1
|
RUN apt-get update \
|
||||||
WORKDIR /work
|
&& 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
|
# `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
|
# 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.
|
# it statically, keeping the runtime self-contained.
|
||||||
ENV PCRE2_SYS_STATIC=1
|
ENV PCRE2_SYS_STATIC=1
|
||||||
|
|
||||||
COPY --from=chef /work/recipe.json ./recipe.json
|
COPY --from=chef /work/sgl-router/recipe.json ./recipe.json
|
||||||
COPY --from=chef /work/Cargo.lock ./Cargo.lock
|
COPY --from=chef /work/sgl-router/Cargo.lock ./Cargo.lock
|
||||||
COPY experimental/sgl-router/rust-toolchain.toml ./
|
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
|
# Cook (compile + cache) the dep graph from the recipe. The recipe carries every
|
||||||
# inputs are recipe.json + the toolchain — code changes in src/ do NOT
|
# workspace member's manifest, so chef recreates the Indexer's source stubs itself.
|
||||||
# invalidate it.
|
|
||||||
RUN cargo chef cook --release --recipe-path recipe.json
|
RUN cargo chef cook --release --recipe-path recipe.json
|
||||||
|
|
||||||
# Now bring in the real sources and the manifest they need.
|
# Now bring in the real sources and the manifest they need.
|
||||||
COPY experimental/sgl-router/Cargo.toml ./
|
COPY experimental/sgl-router/Cargo.toml ./
|
||||||
COPY experimental/sgl-router/src ./src
|
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
|
# --locked is intentionally omitted: the lockfile is generated in-container
|
||||||
# (gitignored upstream) and `cargo chef cook` may have mutated it during the
|
# (gitignored upstream) and `cargo chef cook` may have mutated it during the
|
||||||
# dep-cook step, so a strict --locked check would spuriously fail.
|
# 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
|
&& strip target/release/sgl-router
|
||||||
|
|
||||||
######################## STAGE 3 — runtime ##############################
|
######################## STAGE 3 — runtime ##############################
|
||||||
FROM gcr.io/distroless/cc-debian12:nonroot AS 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 <host-path>:/etc/sgl-router`.
|
# Default config path; mount your own via `-v <host-path>:/etc/sgl-router`.
|
||||||
ENV SGL_ROUTER_CONFIG=/etc/sgl-router/sgl-router.yaml
|
ENV SGL_ROUTER_CONFIG=/etc/sgl-router/sgl-router.yaml
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[workspace]
|
[workspace]
|
||||||
resolver = "2"
|
resolver = "2"
|
||||||
members = ["."]
|
members = [".", "sgl-kv-indexer"]
|
||||||
|
|
||||||
[package]
|
[package]
|
||||||
name = "sgl-router"
|
name = "sgl-router"
|
||||||
@@ -33,6 +33,7 @@ axum = { version = "0.8", features = ["macros", "tracing"] }
|
|||||||
tower = { version = "0.5", features = ["full"] }
|
tower = { version = "0.5", features = ["full"] }
|
||||||
tower-http = { version = "0.6", features = ["trace", "compression-gzip", "cors", "timeout", "request-id"] }
|
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 }
|
reqwest = { version = "0.12", features = ["stream", "json", "rustls-tls"], default-features = false }
|
||||||
|
sgl-kv-indexer = { path = "sgl-kv-indexer" }
|
||||||
|
|
||||||
# Serialization
|
# Serialization
|
||||||
serde = { version = "1", features = ["derive"] }
|
serde = { version = "1", features = ["derive"] }
|
||||||
@@ -62,7 +63,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
|
|||||||
futures = "0.3"
|
futures = "0.3"
|
||||||
bytes = "1"
|
bytes = "1"
|
||||||
rand = "0.8"
|
rand = "0.8"
|
||||||
tokio-stream = "0.1"
|
tokio-stream = { version = "0.1", features = ["net"] }
|
||||||
dashmap = "6"
|
dashmap = "6"
|
||||||
kube = { version = "0.96", features = ["runtime", "derive"] }
|
kube = { version = "0.96", features = ["runtime", "derive"] }
|
||||||
k8s-openapi = { version = "0.23", features = ["v1_31"] }
|
k8s-openapi = { version = "0.23", features = ["v1_31"] }
|
||||||
@@ -86,6 +87,7 @@ serde_json = "1"
|
|||||||
tempfile = "3"
|
tempfile = "3"
|
||||||
tower = { version = "0.5", features = ["util"] }
|
tower = { version = "0.5", features = ["util"] }
|
||||||
tokio = { version = "1.42", features = ["test-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
|
# Low-level msgpack encoder used to hand-construct wire bytes in
|
||||||
# kv_events golden-bytes tests (decode-only path uses rmp-serde).
|
# kv_events golden-bytes tests (decode-only path uses rmp-serde).
|
||||||
rmp = "0.8"
|
rmp = "0.8"
|
||||||
|
|||||||
@@ -50,6 +50,26 @@ Omit `--service-discovery-namespace` to watch all namespaces (requires
|
|||||||
cluster-wide RBAC). For prefill/decode disaggregation, replace `--selector`
|
cluster-wide RBAC). For prefill/decode disaggregation, replace `--selector`
|
||||||
with `--prefill-selector` and `--decode-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
|
## License
|
||||||
|
|
||||||
Apache-2.0.
|
Apache-2.0.
|
||||||
|
|||||||
@@ -0,0 +1,3 @@
|
|||||||
|
target/
|
||||||
|
*.rs.bk
|
||||||
|
.DS_Store
|
||||||
@@ -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"
|
||||||
@@ -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=<model-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 <model> \
|
||||||
|
--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 <model-id> \
|
||||||
|
--tokenizer-path <huggingface-repo-or-tokenizer> \
|
||||||
|
--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
|
||||||
|
```
|
||||||
@@ -0,0 +1,10 @@
|
|||||||
|
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||||
|
// SPDX-License-Identifier: Apache-2.0
|
||||||
|
|
||||||
|
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||||
|
tonic_prost_build::configure()
|
||||||
|
.build_client(true)
|
||||||
|
.build_server(true)
|
||||||
|
.compile_protos(&["proto/kv_indexer.proto"], &["proto"])?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
@@ -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);
|
||||||
|
}
|
||||||
@@ -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<u64> {
|
||||||
|
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<Request<()>, 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::<Arrival>(), 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<Duration> {
|
||||||
|
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<Instant>) -> 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<u64> = (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());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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<dyn std::error::Error>> {
|
||||||
|
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(())
|
||||||
|
}
|
||||||
@@ -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<dyn std::error::Error + Send + Sync>> {
|
||||||
|
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::<SocketAddr>()?;
|
||||||
|
let prefix_query_max_inflight = prefix_query_max_inflight_from_env()?;
|
||||||
|
|
||||||
|
let backend: Arc<dyn KvIndexerBackend> = 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<usize> {
|
||||||
|
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<usize> {
|
||||||
|
let value = raw.parse::<usize>().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());
|
||||||
|
}
|
||||||
|
}
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -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<PrefixMatch>,
|
||||||
|
/// 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<String>) -> 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<i64>) -> Result<PrefixOutcome, PrefixIndexError>;
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 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<Self, InvalidEndpoint> {
|
||||||
|
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<SemaphorePermit<'_>, PrefixIndexError> {
|
||||||
|
self.prefix_query_semaphore
|
||||||
|
.try_acquire()
|
||||||
|
.map_err(|_| PrefixIndexError::Overloaded)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn truncate_prefix_query(hashes: &mut Vec<i64>) -> Option<usize> {
|
||||||
|
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<i64>) -> Result<PrefixOutcome, PrefixIndexError> {
|
||||||
|
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<Endpoint, InvalidEndpoint> {
|
||||||
|
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());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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;
|
||||||
@@ -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<WorkerCacheSpec>,
|
||||||
|
/// Reverse index used by CLEAR_ALL_AT_TIER.
|
||||||
|
holdings: HashMap<i32, HashSet<i64>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Default)]
|
||||||
|
struct State {
|
||||||
|
blocks: HashMap<i64, BlockRecord>,
|
||||||
|
workers: HashMap<String, WorkerRecord>,
|
||||||
|
hit_counts: HashMap<i64, u64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
struct WorkerView {
|
||||||
|
worker_id: String,
|
||||||
|
address: String,
|
||||||
|
spec: Option<WorkerCacheSpec>,
|
||||||
|
hashes_by_tier: BTreeMap<i32, Vec<(i64, u32, u32)>>,
|
||||||
|
blocks: Vec<Option<BlockComponents>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[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<State>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl InMemoryKvIndexerBackend {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self::default()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read_state(&self) -> Result<RwLockReadGuard<'_, State>, Status> {
|
||||||
|
self.state
|
||||||
|
.read()
|
||||||
|
.map_err(|_| Status::internal("in-memory backend lock poisoned"))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn write_state(&self) -> Result<RwLockWriteGuard<'_, State>, Status> {
|
||||||
|
self.state
|
||||||
|
.write()
|
||||||
|
.map_err(|_| Status::internal("in-memory backend lock poisoned"))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply(
|
||||||
|
&self,
|
||||||
|
req: ApplyExternalKvBatchRequest,
|
||||||
|
) -> Result<ApplyExternalKvBatchResponse, Status> {
|
||||||
|
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<i64, (u32, u32)> = 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<MatchExternalKvResponse, Status> {
|
||||||
|
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<WorkerView>, Vec<i64>) {
|
||||||
|
// 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<WorkerPrefixInput> {
|
||||||
|
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<MatchExternalKvPrefixResponse, Status> {
|
||||||
|
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<PrefixCandidate> = 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<String, usize> = candidates
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(index, candidate)| (candidate.worker_id.clone(), index))
|
||||||
|
.collect();
|
||||||
|
let mut present = vec![false; candidates.len()];
|
||||||
|
let mut block_views: Vec<BlockComponents> = (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<GetExternalKvHitCountsResponse, Status> {
|
||||||
|
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<i64> {
|
||||||
|
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<ApplyExternalKvBatchResponse, Status> {
|
||||||
|
self.apply(request)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn match_external_kv(
|
||||||
|
&self,
|
||||||
|
request: MatchExternalKvRequest,
|
||||||
|
) -> Result<MatchExternalKvResponse, Status> {
|
||||||
|
self.do_match(request)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn collect_worker_prefix_inputs(
|
||||||
|
&self,
|
||||||
|
hashes: &[i64],
|
||||||
|
) -> Result<Vec<WorkerPrefixInput>, Status> {
|
||||||
|
let state = self.read_state()?;
|
||||||
|
Ok(Self::collect_prefix_inputs_locked(&state, hashes))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn match_external_kv_prefix(
|
||||||
|
&self,
|
||||||
|
request: MatchExternalKvPrefixRequest,
|
||||||
|
) -> Result<MatchExternalKvPrefixResponse, Status> {
|
||||||
|
self.do_match_prefix(request)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_external_kv_hit_counts(
|
||||||
|
&self,
|
||||||
|
request: GetExternalKvHitCountsRequest,
|
||||||
|
) -> Result<GetExternalKvHitCountsResponse, Status> {
|
||||||
|
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();
|
||||||
|
}
|
||||||
|
}
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -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"),
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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()
|
||||||
|
}
|
||||||
@@ -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<ExternalKvAction>,
|
||||||
|
) -> ApplyExternalKvBatchRequest {
|
||||||
|
ApplyExternalKvBatchRequest {
|
||||||
|
worker_id: worker.to_string(),
|
||||||
|
seq,
|
||||||
|
actions,
|
||||||
|
worker_address: address.to_string(),
|
||||||
|
cache_spec: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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()
|
||||||
|
}
|
||||||
@@ -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<tonic::transport::Channel> {
|
||||||
|
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<AtomicUsize>,
|
||||||
|
release: Arc<Semaphore>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tonic::async_trait]
|
||||||
|
impl KvIndexerBackend for BlockingPrefixBackend {
|
||||||
|
async fn apply_external_kv_batch(
|
||||||
|
&self,
|
||||||
|
_request: ApplyExternalKvBatchRequest,
|
||||||
|
) -> Result<ApplyExternalKvBatchResponse, Status> {
|
||||||
|
Ok(ApplyExternalKvBatchResponse::default())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn match_external_kv(
|
||||||
|
&self,
|
||||||
|
_request: MatchExternalKvRequest,
|
||||||
|
) -> Result<MatchExternalKvResponse, Status> {
|
||||||
|
Ok(MatchExternalKvResponse::default())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn match_external_kv_prefix(
|
||||||
|
&self,
|
||||||
|
_request: MatchExternalKvPrefixRequest,
|
||||||
|
) -> Result<MatchExternalKvPrefixResponse, Status> {
|
||||||
|
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<GetExternalKvHitCountsResponse, Status> {
|
||||||
|
Ok(GetExternalKvHitCountsResponse::default())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn start_blocking_backend(
|
||||||
|
backend: BlockingPrefixBackend,
|
||||||
|
) -> KvIndexerClient<tonic::transport::Channel> {
|
||||||
|
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<tonic::transport::Channel> {
|
||||||
|
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<i64> = (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::<i64>() + 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<Mutex<Vec<String>>>) {
|
||||||
|
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:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
@@ -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<i32> {
|
||||||
|
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<i64> = (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<InMemoryKvIndexerBackend>);
|
||||||
|
|
||||||
|
#[tonic::async_trait]
|
||||||
|
impl KvIndexerBackend for DefaultViaMemory {
|
||||||
|
async fn apply_external_kv_batch(
|
||||||
|
&self,
|
||||||
|
request: ApplyExternalKvBatchRequest,
|
||||||
|
) -> Result<ApplyExternalKvBatchResponse, Status> {
|
||||||
|
self.0.apply_external_kv_batch(request).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn match_external_kv(
|
||||||
|
&self,
|
||||||
|
request: MatchExternalKvRequest,
|
||||||
|
) -> Result<MatchExternalKvResponse, Status> {
|
||||||
|
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<Vec<WorkerPrefixInput>, Status> {
|
||||||
|
self.0.collect_worker_prefix_inputs(hashes).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_external_kv_hit_counts(
|
||||||
|
&self,
|
||||||
|
request: GetExternalKvHitCountsRequest,
|
||||||
|
) -> Result<GetExternalKvHitCountsResponse, Status> {
|
||||||
|
self.0.get_external_kv_hit_counts(request).await
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn shared_state_pair() -> (Arc<InMemoryKvIndexerBackend>, 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<sgl_kv_indexer::pb::ExternalKvAction>,
|
||||||
|
) -> 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);
|
||||||
|
}
|
||||||
@@ -12,10 +12,14 @@ use std::num::NonZeroU32;
|
|||||||
use crate::config::{
|
use crate::config::{
|
||||||
default_cb_cool_down, default_proxy_request_timeout_secs, default_stale_request_timeout_secs,
|
default_cb_cool_down, default_proxy_request_timeout_secs, default_stale_request_timeout_secs,
|
||||||
resolve_mode, ActiveLoadConfig, CacheAwareConfig, CircuitBreakerConfig, Config,
|
resolve_mode, ActiveLoadConfig, CacheAwareConfig, CircuitBreakerConfig, Config,
|
||||||
DiscoveryBackend, K8sDiscoveryConfig, LogFormat, ModelConfig, ObservabilityConfig, PolicyKind,
|
DiscoveryBackend, K8sDiscoveryConfig, KvIndexerEndpointConfig, LogFormat, ModelConfig,
|
||||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig,
|
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.
|
/// `sgl-router` — slim KV-aware OpenAI-compatible router for SGLang workers.
|
||||||
///
|
///
|
||||||
/// Discovery is mutually exclusive: pass `--worker-urls` for a static
|
/// 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.
|
/// Multiplicative load spread gating the absolute balance check.
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
pub balance_rel_threshold: Option<f32>,
|
pub balance_rel_threshold: Option<f32>,
|
||||||
|
/// 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<String>,
|
||||||
|
/// KV Indexer query timeout in milliseconds. Requires
|
||||||
|
/// `--kv-indexer-endpoint`; defaults to 100.
|
||||||
|
#[arg(long)]
|
||||||
|
pub kv_indexer_query_timeout_ms: Option<u64>,
|
||||||
|
/// 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<usize>,
|
||||||
|
|
||||||
// ---- sticky-session policy (only used by `--policy sticky`) ----
|
// ---- sticky-session policy (only used by `--policy sticky`) ----
|
||||||
/// Request header carrying the routing key for sticky-session routing.
|
/// 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()
|
let tuned_cache_aware = self.cache_threshold.is_some()
|
||||||
|| self.balance_abs_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 {
|
if tuned_cache_aware && self.policy != PolicyKind::CacheAwareZmq {
|
||||||
return Err(anyhow!(
|
return Err(anyhow!(
|
||||||
"--cache-threshold / --balance-abs-threshold / --balance-rel-threshold \
|
"cache-aware tuning flags require --policy cache_aware_zmq"
|
||||||
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()
|
let tuned_sticky = self.routing_key_header.is_some()
|
||||||
|| self.sticky_fallback_policy.is_some()
|
|| self.sticky_fallback_policy.is_some()
|
||||||
|| self.sticky_idle_secs.is_some()
|
|| self.sticky_idle_secs.is_some()
|
||||||
@@ -234,6 +271,12 @@ impl Cli {
|
|||||||
// Only build a CacheAwareConfig when the operator tuned at least
|
// Only build a CacheAwareConfig when the operator tuned at least
|
||||||
// one knob; otherwise leave it None so the policy uses its own
|
// one knob; otherwise leave it None so the policy uses its own
|
||||||
// defaults. Unset knobs fall back to the per-field defaults.
|
// 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 cache_aware = if tuned_cache_aware {
|
||||||
let d = CacheAwareConfig::default();
|
let d = CacheAwareConfig::default();
|
||||||
Some(CacheAwareConfig {
|
Some(CacheAwareConfig {
|
||||||
@@ -244,6 +287,11 @@ impl Cli {
|
|||||||
balance_rel_threshold: self
|
balance_rel_threshold: self
|
||||||
.balance_rel_threshold
|
.balance_rel_threshold
|
||||||
.unwrap_or(d.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 {
|
} else {
|
||||||
None
|
None
|
||||||
@@ -744,6 +792,109 @@ mod tests {
|
|||||||
assert_eq!(ca.balance_abs_threshold, 32);
|
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]
|
#[test]
|
||||||
fn no_cache_aware_flags_leaves_none() {
|
fn no_cache_aware_flags_leaves_none() {
|
||||||
let c = into_config_owned(with_model(&[
|
let c = into_config_owned(with_model(&[
|
||||||
|
|||||||
@@ -146,7 +146,7 @@ pub struct ModelConfig {
|
|||||||
pub tokenizer_path: String,
|
pub tokenizer_path: String,
|
||||||
pub policy: PolicyKind,
|
pub policy: PolicyKind,
|
||||||
pub circuit_breaker: Option<CircuitBreakerConfig>,
|
pub circuit_breaker: Option<CircuitBreakerConfig>,
|
||||||
/// 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 = "cache_aware_zmq"`. `None` falls back to defaults at
|
||||||
/// policy construction time.
|
/// policy construction time.
|
||||||
pub cache_aware: Option<CacheAwareConfig>,
|
pub cache_aware: Option<CacheAwareConfig>,
|
||||||
@@ -157,8 +157,16 @@ pub struct ModelConfig {
|
|||||||
pub sticky: Option<StickyConfig>,
|
pub sticky: Option<StickyConfig>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Per-model cache-aware-ZMQ tuning.
|
/// External KV Indexer client settings.
|
||||||
#[derive(Debug, Clone, Copy)]
|
#[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 {
|
pub struct CacheAwareConfig {
|
||||||
/// Lower bound on `matched_blocks / total_blocks` for the tree match
|
/// Lower bound on `matched_blocks / total_blocks` for the tree match
|
||||||
/// to win the selection. Below this, the policy falls back to
|
/// 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
|
/// that the absolute check is gated on. Default 1.1 — 10 % relative
|
||||||
/// difference triggers re-balancing.
|
/// difference triggers re-balancing.
|
||||||
pub balance_rel_threshold: f32,
|
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<KvIndexerEndpointConfig>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for CacheAwareConfig {
|
impl Default for CacheAwareConfig {
|
||||||
@@ -182,6 +193,7 @@ impl Default for CacheAwareConfig {
|
|||||||
cache_threshold: default_cache_threshold(),
|
cache_threshold: default_cache_threshold(),
|
||||||
balance_abs_threshold: default_balance_abs(),
|
balance_abs_threshold: default_balance_abs(),
|
||||||
balance_rel_threshold: default_balance_rel(),
|
balance_rel_threshold: default_balance_rel(),
|
||||||
|
kv_indexer_endpoint: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -98,19 +98,43 @@ async fn main() -> Result<()> {
|
|||||||
);
|
);
|
||||||
|
|
||||||
let registry = Arc::new(sgl_router::workers::WorkerRegistry::default());
|
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
|
// Build the KV-event index up front so the cache-aware-zmq policy can
|
||||||
// share its `HashTree` handle + `BlockSizeOracle`. When no model uses
|
// share its `HashTree` handle + `BlockSizeOracle`. An external Indexer makes
|
||||||
// `cache_aware_zmq`, the index is still constructed (cheap) but no
|
// the local tree irrelevant to routing, so only discover hash metadata rather
|
||||||
// subscribers are ever added.
|
// than duplicating every KV event.
|
||||||
let block_size_oracle = sgl_router::policies::kv_events::BlockSizeOracle::new();
|
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(
|
let kv_event_http = reqwest::Client::builder()
|
||||||
reqwest::Client::builder()
|
|
||||||
.timeout(std::time::Duration::from_secs(2))
|
.timeout(std::time::Duration::from_secs(2))
|
||||||
.build()
|
.build()
|
||||||
.expect("default http client builds"),
|
.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),
|
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(
|
let policies = Arc::new(
|
||||||
sgl_router::policies::factory::build_registry(
|
sgl_router::policies::factory::build_registry(
|
||||||
&cfg,
|
&cfg,
|
||||||
@@ -163,16 +187,17 @@ async fn main() -> Result<()> {
|
|||||||
.context("build proxy client")?,
|
.context("build proxy client")?,
|
||||||
);
|
);
|
||||||
|
|
||||||
let ctx = Arc::new(
|
let mut app_ctx = sgl_router::server::app_context::AppContext::with_active_load(
|
||||||
sgl_router::server::app_context::AppContext::with_active_load(
|
|
||||||
cfg.clone(),
|
cfg.clone(),
|
||||||
tokenizers,
|
tokenizers,
|
||||||
proxy,
|
proxy,
|
||||||
registry,
|
registry,
|
||||||
policies,
|
policies,
|
||||||
active_load,
|
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();
|
ctx.mark_ready();
|
||||||
|
|
||||||
let app = sgl_router::server::app::build_router(ctx.clone());
|
let app = sgl_router::server::app::build_router(ctx.clone());
|
||||||
|
|||||||
@@ -130,6 +130,46 @@ impl CacheAwareZmqPolicy {
|
|||||||
let rel_threshold = (min_load as f32 * self.config.balance_rel_threshold) as usize;
|
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
|
abs_diff > self.config.balance_abs_threshold && max_load > rel_threshold
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn select_external(
|
||||||
|
&self,
|
||||||
|
workers: &[Arc<Worker>],
|
||||||
|
ctx: &SelectionContext<'_>,
|
||||||
|
signal: &crate::policies::ExternalPrefixSignal,
|
||||||
|
) -> Option<Arc<Worker>> {
|
||||||
|
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 {
|
impl Policy for CacheAwareZmqPolicy {
|
||||||
@@ -144,6 +184,14 @@ impl Policy for CacheAwareZmqPolicy {
|
|||||||
return Self::pick_min_load(workers);
|
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
|
// 2. Routing tokens. Prefer the ids computed once at ingress; fall
|
||||||
// back to tokenizing the body here so the policy stays usable for
|
// back to tokenizing the body here so the policy stays usable for
|
||||||
// callers that don't pre-tokenize (e.g. unit tests). In production
|
// callers that don't pre-tokenize (e.g. unit tests). In production
|
||||||
@@ -263,6 +311,7 @@ mod tests {
|
|||||||
cache_threshold: 0.5,
|
cache_threshold: 0.5,
|
||||||
balance_abs_threshold: 32,
|
balance_abs_threshold: 32,
|
||||||
balance_rel_threshold: 1.1,
|
balance_rel_threshold: 1.1,
|
||||||
|
kv_indexer_endpoint: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -350,6 +399,70 @@ mod tests {
|
|||||||
assert_eq!(chosen.url, "http://w1:30000");
|
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
|
/// Tree contains w0's prefix; cache-aware selection picks w0 even
|
||||||
/// though w1 has lower load (the load skew is below the imbalance
|
/// though w1 has lower load (the load skew is below the imbalance
|
||||||
/// threshold, so cache wins).
|
/// threshold, so cache wins).
|
||||||
@@ -377,6 +490,7 @@ mod tests {
|
|||||||
cache_threshold: 0.0, // any match counts
|
cache_threshold: 0.0, // any match counts
|
||||||
balance_abs_threshold: 32,
|
balance_abs_threshold: 32,
|
||||||
balance_rel_threshold: 1.1,
|
balance_rel_threshold: 1.1,
|
||||||
|
kv_indexer_endpoint: None,
|
||||||
},
|
},
|
||||||
tree,
|
tree,
|
||||||
registry,
|
registry,
|
||||||
@@ -414,6 +528,7 @@ mod tests {
|
|||||||
cache_threshold: 0.0,
|
cache_threshold: 0.0,
|
||||||
balance_abs_threshold: 32,
|
balance_abs_threshold: 32,
|
||||||
balance_rel_threshold: 1.1,
|
balance_rel_threshold: 1.1,
|
||||||
|
kv_indexer_endpoint: None,
|
||||||
},
|
},
|
||||||
tree,
|
tree,
|
||||||
registry,
|
registry,
|
||||||
@@ -459,6 +574,7 @@ mod tests {
|
|||||||
cache_threshold: 0.0,
|
cache_threshold: 0.0,
|
||||||
balance_abs_threshold: 32,
|
balance_abs_threshold: 32,
|
||||||
balance_rel_threshold: 1.1,
|
balance_rel_threshold: 1.1,
|
||||||
|
kv_indexer_endpoint: None,
|
||||||
},
|
},
|
||||||
tree,
|
tree,
|
||||||
toks,
|
toks,
|
||||||
@@ -510,6 +626,7 @@ mod tests {
|
|||||||
cache_threshold: 1.0, // match_rate <= 1.0 always -> always fall back
|
cache_threshold: 1.0, // match_rate <= 1.0 always -> always fall back
|
||||||
balance_abs_threshold: 32,
|
balance_abs_threshold: 32,
|
||||||
balance_rel_threshold: 1.1,
|
balance_rel_threshold: 1.1,
|
||||||
|
kv_indexer_endpoint: None,
|
||||||
},
|
},
|
||||||
tree,
|
tree,
|
||||||
toks,
|
toks,
|
||||||
@@ -593,6 +710,7 @@ mod tests {
|
|||||||
cache_threshold: 0.0,
|
cache_threshold: 0.0,
|
||||||
balance_abs_threshold: 32,
|
balance_abs_threshold: 32,
|
||||||
balance_rel_threshold: 1.1,
|
balance_rel_threshold: 1.1,
|
||||||
|
kv_indexer_endpoint: None,
|
||||||
},
|
},
|
||||||
tree,
|
tree,
|
||||||
Arc::clone(®istry),
|
Arc::clone(®istry),
|
||||||
@@ -632,6 +750,7 @@ mod tests {
|
|||||||
cache_threshold: 0.0,
|
cache_threshold: 0.0,
|
||||||
balance_abs_threshold: 32,
|
balance_abs_threshold: 32,
|
||||||
balance_rel_threshold: 1.1,
|
balance_rel_threshold: 1.1,
|
||||||
|
kv_indexer_endpoint: None,
|
||||||
},
|
},
|
||||||
tree,
|
tree,
|
||||||
Arc::clone(®istry),
|
Arc::clone(®istry),
|
||||||
@@ -690,6 +809,7 @@ mod tests {
|
|||||||
cache_threshold: 0.0,
|
cache_threshold: 0.0,
|
||||||
balance_abs_threshold: 32,
|
balance_abs_threshold: 32,
|
||||||
balance_rel_threshold: 1.1,
|
balance_rel_threshold: 1.1,
|
||||||
|
kv_indexer_endpoint: None,
|
||||||
},
|
},
|
||||||
tree,
|
tree,
|
||||||
registry,
|
registry,
|
||||||
@@ -762,6 +882,7 @@ mod tests {
|
|||||||
cache_threshold: 0.0,
|
cache_threshold: 0.0,
|
||||||
balance_abs_threshold: 32,
|
balance_abs_threshold: 32,
|
||||||
balance_rel_threshold: 1.1,
|
balance_rel_threshold: 1.1,
|
||||||
|
kv_indexer_endpoint: None,
|
||||||
},
|
},
|
||||||
tree,
|
tree,
|
||||||
registry,
|
registry,
|
||||||
@@ -800,6 +921,7 @@ mod tests {
|
|||||||
cache_threshold: 0.0,
|
cache_threshold: 0.0,
|
||||||
balance_abs_threshold: 32,
|
balance_abs_threshold: 32,
|
||||||
balance_rel_threshold: 1.1,
|
balance_rel_threshold: 1.1,
|
||||||
|
kv_indexer_endpoint: None,
|
||||||
},
|
},
|
||||||
tree,
|
tree,
|
||||||
registry,
|
registry,
|
||||||
@@ -903,6 +1025,7 @@ mod tests {
|
|||||||
cache_threshold: 0.0,
|
cache_threshold: 0.0,
|
||||||
balance_abs_threshold: 32,
|
balance_abs_threshold: 32,
|
||||||
balance_rel_threshold: 1.1,
|
balance_rel_threshold: 1.1,
|
||||||
|
kv_indexer_endpoint: None,
|
||||||
},
|
},
|
||||||
tree,
|
tree,
|
||||||
registry,
|
registry,
|
||||||
@@ -938,6 +1061,7 @@ mod tests {
|
|||||||
cache_threshold: 0.0, // would normally always match
|
cache_threshold: 0.0, // would normally always match
|
||||||
balance_abs_threshold: 5,
|
balance_abs_threshold: 5,
|
||||||
balance_rel_threshold: 2.0,
|
balance_rel_threshold: 2.0,
|
||||||
|
kv_indexer_endpoint: None,
|
||||||
},
|
},
|
||||||
tree,
|
tree,
|
||||||
registry,
|
registry,
|
||||||
@@ -1063,6 +1187,7 @@ mod tests {
|
|||||||
cache_threshold: 0.99,
|
cache_threshold: 0.99,
|
||||||
balance_abs_threshold: 32,
|
balance_abs_threshold: 32,
|
||||||
balance_rel_threshold: 1.1,
|
balance_rel_threshold: 1.1,
|
||||||
|
kv_indexer_endpoint: None,
|
||||||
},
|
},
|
||||||
tree,
|
tree,
|
||||||
tokenizer_registry_with_tiny(),
|
tokenizer_registry_with_tiny(),
|
||||||
@@ -1146,6 +1271,7 @@ mod tests {
|
|||||||
cache_threshold: 0.0,
|
cache_threshold: 0.0,
|
||||||
balance_abs_threshold: 32,
|
balance_abs_threshold: 32,
|
||||||
balance_rel_threshold: 1.1,
|
balance_rel_threshold: 1.1,
|
||||||
|
kv_indexer_endpoint: None,
|
||||||
},
|
},
|
||||||
tree.clone(),
|
tree.clone(),
|
||||||
registry,
|
registry,
|
||||||
@@ -1243,6 +1369,7 @@ mod tests {
|
|||||||
cache_threshold: 0.0,
|
cache_threshold: 0.0,
|
||||||
balance_abs_threshold: 32,
|
balance_abs_threshold: 32,
|
||||||
balance_rel_threshold: 1.1,
|
balance_rel_threshold: 1.1,
|
||||||
|
kv_indexer_endpoint: None,
|
||||||
},
|
},
|
||||||
tree,
|
tree,
|
||||||
registry,
|
registry,
|
||||||
|
|||||||
@@ -65,7 +65,7 @@ pub fn build_policy(
|
|||||||
PolicyKind::PowerOfTwo => Arc::new(PowerOfTwoChoicesPolicy::new()),
|
PolicyKind::PowerOfTwo => Arc::new(PowerOfTwoChoicesPolicy::new()),
|
||||||
PolicyKind::LoadBased => Arc::new(LoadBasedPolicy::new()),
|
PolicyKind::LoadBased => Arc::new(LoadBasedPolicy::new()),
|
||||||
PolicyKind::CacheAwareZmq => {
|
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(
|
Arc::new(CacheAwareZmqPolicy::new(
|
||||||
cache_cfg,
|
cache_cfg,
|
||||||
tree,
|
tree,
|
||||||
|
|||||||
@@ -65,6 +65,7 @@ struct WorkerEntry {
|
|||||||
/// routing path entirely.
|
/// routing path entirely.
|
||||||
pub struct KvEventIndex {
|
pub struct KvEventIndex {
|
||||||
tree: Arc<HashTree>,
|
tree: Arc<HashTree>,
|
||||||
|
maintain_tree: bool,
|
||||||
subscribers: Arc<KvEventSubscriberRegistry>,
|
subscribers: Arc<KvEventSubscriberRegistry>,
|
||||||
pump: Mutex<Option<JoinHandle<()>>>,
|
pump: Mutex<Option<JoinHandle<()>>>,
|
||||||
pump_cancel: CancellationToken,
|
pump_cancel: CancellationToken,
|
||||||
@@ -114,6 +115,24 @@ impl KvEventIndex {
|
|||||||
pub fn new_with_http_and_oracle(
|
pub fn new_with_http_and_oracle(
|
||||||
http: reqwest::Client,
|
http: reqwest::Client,
|
||||||
block_size_oracle: Arc<BlockSizeOracle>,
|
block_size_oracle: Arc<BlockSizeOracle>,
|
||||||
|
) -> Arc<Self> {
|
||||||
|
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<BlockSizeOracle>,
|
||||||
|
) -> Arc<Self> {
|
||||||
|
Self::new_with_mode(http, block_size_oracle, false)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn new_with_mode(
|
||||||
|
http: reqwest::Client,
|
||||||
|
block_size_oracle: Arc<BlockSizeOracle>,
|
||||||
|
maintain_tree: bool,
|
||||||
) -> Arc<Self> {
|
) -> Arc<Self> {
|
||||||
let tree = Arc::new(HashTree::new());
|
let tree = Arc::new(HashTree::new());
|
||||||
let (tx, rx) = mpsc::channel::<WorkerEvent>(EVENT_CHANNEL_BUFFER);
|
let (tx, rx) = mpsc::channel::<WorkerEvent>(EVENT_CHANNEL_BUFFER);
|
||||||
@@ -130,6 +149,7 @@ impl KvEventIndex {
|
|||||||
));
|
));
|
||||||
Arc::new(Self {
|
Arc::new(Self {
|
||||||
tree,
|
tree,
|
||||||
|
maintain_tree,
|
||||||
subscribers,
|
subscribers,
|
||||||
pump: Mutex::new(Some(pump)),
|
pump: Mutex::new(Some(pump)),
|
||||||
pump_cancel,
|
pump_cancel,
|
||||||
@@ -208,6 +228,15 @@ impl KvEventIndex {
|
|||||||
// hash KV blocks over token bigrams, so the policy must use the bigram
|
// 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.
|
// hasher for its query hashes to match the worker's stored hashes.
|
||||||
self.block_size_oracle.set_bigram(cfg.is_bigram);
|
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!(
|
info!(
|
||||||
worker_url = %worker_url,
|
worker_url = %worker_url,
|
||||||
dp_size = cfg.dp_size,
|
dp_size = cfg.dp_size,
|
||||||
@@ -710,4 +739,28 @@ mod tests {
|
|||||||
);
|
);
|
||||||
index.shutdown().await;
|
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;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -34,6 +34,13 @@ pub struct RequestTokens {
|
|||||||
pub engine_equivalent: bool,
|
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 —
|
/// Produce the routing tokens — and whether they are engine-equivalent —
|
||||||
/// from an already-parsed request body, using the shared tokenizer registry.
|
/// from an already-parsed request body, using the shared tokenizer registry.
|
||||||
///
|
///
|
||||||
@@ -181,6 +188,7 @@ pub struct SelectionContext<'a> {
|
|||||||
request_body: Option<&'a [u8]>,
|
request_body: Option<&'a [u8]>,
|
||||||
routing_key: Option<&'a str>,
|
routing_key: Option<&'a str>,
|
||||||
request_tokens: Option<&'a [u32]>,
|
request_tokens: Option<&'a [u32]>,
|
||||||
|
external_prefix: Option<&'a ExternalPrefixSignal>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<'a> SelectionContext<'a> {
|
impl<'a> SelectionContext<'a> {
|
||||||
@@ -190,6 +198,7 @@ impl<'a> SelectionContext<'a> {
|
|||||||
request_body,
|
request_body,
|
||||||
routing_key: None,
|
routing_key: None,
|
||||||
request_tokens: None,
|
request_tokens: None,
|
||||||
|
external_prefix: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -203,6 +212,7 @@ impl<'a> SelectionContext<'a> {
|
|||||||
request_body,
|
request_body,
|
||||||
routing_key,
|
routing_key,
|
||||||
request_tokens: None,
|
request_tokens: None,
|
||||||
|
external_prefix: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -214,6 +224,14 @@ impl<'a> SelectionContext<'a> {
|
|||||||
self
|
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 {
|
pub fn model(&self) -> &ModelId {
|
||||||
self.model
|
self.model
|
||||||
}
|
}
|
||||||
@@ -231,6 +249,10 @@ impl<'a> SelectionContext<'a> {
|
|||||||
pub fn request_tokens(&self) -> Option<&[u32]> {
|
pub fn request_tokens(&self) -> Option<&[u32]> {
|
||||||
self.request_tokens
|
self.request_tokens
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn external_prefix(&self) -> Option<&ExternalPrefixSignal> {
|
||||||
|
self.external_prefix
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub trait Policy: Send + Sync + std::fmt::Debug {
|
pub trait Policy: Send + Sync + std::fmt::Debug {
|
||||||
|
|||||||
@@ -4,6 +4,7 @@
|
|||||||
use crate::config::Config;
|
use crate::config::Config;
|
||||||
|
|
||||||
use crate::policies::active_load::ActiveLoadRegistry;
|
use crate::policies::active_load::ActiveLoadRegistry;
|
||||||
|
use crate::policies::kv_events::BlockSizeOracle;
|
||||||
use crate::policies::PolicyRegistry;
|
use crate::policies::PolicyRegistry;
|
||||||
use crate::proxy::Proxy;
|
use crate::proxy::Proxy;
|
||||||
use crate::server::metrics::MetricsRegistry;
|
use crate::server::metrics::MetricsRegistry;
|
||||||
@@ -12,7 +13,6 @@ use crate::workers::WorkerRegistry;
|
|||||||
use std::sync::atomic::{AtomicBool, Ordering};
|
use std::sync::atomic::{AtomicBool, Ordering};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
#[derive(Debug)]
|
|
||||||
pub struct AppContext {
|
pub struct AppContext {
|
||||||
pub config: Config,
|
pub config: Config,
|
||||||
pub tokenizers: Arc<TokenizerRegistry>,
|
pub tokenizers: Arc<TokenizerRegistry>,
|
||||||
@@ -30,6 +30,8 @@ pub struct AppContext {
|
|||||||
/// (active_load gauge + stale_requests_total), and PD resolver
|
/// (active_load gauge + stale_requests_total), and PD resolver
|
||||||
/// (decode_affinity_total).
|
/// (decode_affinity_total).
|
||||||
pub metrics: Arc<MetricsRegistry>,
|
pub metrics: Arc<MetricsRegistry>,
|
||||||
|
pub prefix_index: Option<Arc<sgl_kv_indexer::GrpcPrefixIndex>>,
|
||||||
|
pub block_size_oracle: Arc<BlockSizeOracle>,
|
||||||
ready: AtomicBool,
|
ready: AtomicBool,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -82,6 +84,8 @@ impl AppContext {
|
|||||||
policies,
|
policies,
|
||||||
active_load,
|
active_load,
|
||||||
metrics,
|
metrics,
|
||||||
|
prefix_index: None,
|
||||||
|
block_size_oracle: BlockSizeOracle::new(),
|
||||||
ready: AtomicBool::new(false),
|
ready: AtomicBool::new(false),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -127,6 +131,8 @@ impl AppContext {
|
|||||||
policies: Arc::new(PolicyRegistry::default()),
|
policies: Arc::new(PolicyRegistry::default()),
|
||||||
active_load: ActiveLoadRegistry::with_defaults(),
|
active_load: ActiveLoadRegistry::with_defaults(),
|
||||||
metrics: MetricsRegistry::new(),
|
metrics: MetricsRegistry::new(),
|
||||||
|
prefix_index: None,
|
||||||
|
block_size_oracle: BlockSizeOracle::new(),
|
||||||
ready: AtomicBool::new(false),
|
ready: AtomicBool::new(false),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,8 +2,9 @@
|
|||||||
// SPDX-License-Identifier: Apache-2.0
|
// SPDX-License-Identifier: Apache-2.0
|
||||||
|
|
||||||
use crate::discovery::{ModelId, WorkerMode};
|
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::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::app_context::AppContext;
|
||||||
use crate::server::error::ApiError;
|
use crate::server::error::ApiError;
|
||||||
use crate::server::metrics::{
|
use crate::server::metrics::{
|
||||||
@@ -16,6 +17,7 @@ use axum::http::{HeaderMap, HeaderName, HeaderValue, Response};
|
|||||||
use bytes::Bytes;
|
use bytes::Bytes;
|
||||||
use serde::de::IgnoredAny;
|
use serde::de::IgnoredAny;
|
||||||
use serde::Deserialize;
|
use serde::Deserialize;
|
||||||
|
use sgl_kv_indexer::PrefixIndex;
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
@@ -167,6 +169,34 @@ pub async fn chat_completions(
|
|||||||
let request_tokens = request_value
|
let request_tokens = request_value
|
||||||
.as_ref()
|
.as_ref()
|
||||||
.and_then(|v| request_tokens_for(&ctx.tokenizers, &model_id, v));
|
.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,
|
// Sticky-session routing key. When the sticky policy is configured,
|
||||||
// read the routing key from the operator-chosen header into the
|
// 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())
|
.and_then(|v| v.to_str().ok())
|
||||||
.filter(|s| !s.is_empty());
|
.filter(|s| !s.is_empty());
|
||||||
let selection_ctx = SelectionContext::with_routing_key(&model_id, Some(&body), routing_key)
|
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 =
|
let worker =
|
||||||
policy
|
policy
|
||||||
.select(&workers, &selection_ctx)
|
.select(&workers, &selection_ctx)
|
||||||
@@ -620,6 +651,42 @@ pub async fn chat_completions(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn resolve_prefix_query(
|
||||||
|
result: Result<sgl_kv_indexer::PrefixOutcome, sgl_kv_indexer::PrefixIndexError>,
|
||||||
|
model: &str,
|
||||||
|
) -> Result<sgl_kv_indexer::PrefixOutcome, ApiError> {
|
||||||
|
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
|
/// Estimate prefill-token count from the raw request body for use as
|
||||||
/// the active-load `prefill_load` counter. Returns 1 at minimum so
|
/// the active-load `prefill_load` counter. Returns 1 at minimum so
|
||||||
/// a registered request always shows up as "load > 0" — under-counting
|
/// a registered request always shows up as "load > 0" — under-counting
|
||||||
@@ -911,6 +978,38 @@ fn parse_probe(body: &Bytes) -> Result<RequestProbe, ApiError> {
|
|||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
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
|
/// `generate_room_id` MUST return values in `[0, i64::MAX]`. The
|
||||||
/// SGLang prefill stores `bootstrap_room` as `torch.int64`; a u64
|
/// SGLang prefill stores `bootstrap_room` as `torch.int64`; a u64
|
||||||
/// with the top bit set would wrap negative on the engine side.
|
/// with the top bit set would wrap negative on the engine side.
|
||||||
|
|||||||
@@ -109,6 +109,7 @@ async fn zmq_indexer_routes_to_publishing_worker_e2e() {
|
|||||||
cache_threshold: 0.0,
|
cache_threshold: 0.0,
|
||||||
balance_abs_threshold: 32,
|
balance_abs_threshold: 32,
|
||||||
balance_rel_threshold: 1.1,
|
balance_rel_threshold: 1.1,
|
||||||
|
kv_indexer_endpoint: None,
|
||||||
},
|
},
|
||||||
kv_index.tree(),
|
kv_index.tree(),
|
||||||
Arc::clone(&tokenizers),
|
Arc::clone(&tokenizers),
|
||||||
|
|||||||
+147
-80
@@ -1,44 +1,28 @@
|
|||||||
# SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
# SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||||
# SPDX-License-Identifier: Apache-2.0
|
# 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
|
Two SGLang workers publish KV events to two routers at once: one runs the
|
||||||
independent ``cache_aware_zmq`` policy whose ``KvEventIndex`` subscribes
|
local ``KvEventIndex`` (SUB straight to the workers) and one runs against an
|
||||||
to **both** workers' KV publishers.
|
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
|
Assert on content, not on convergence: a broken event path degrades
|
||||||
both routers), then sends those prefixes through each router and
|
``cache_aware_zmq`` to content-blind min-load, which routes both prefixes to
|
||||||
asserts that routing follows the prefix CONTENT: ``PREFIX_X`` lands on
|
one worker and so fails at least one assertion below.
|
||||||
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.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
import re
|
import re
|
||||||
|
import socket
|
||||||
|
import subprocess
|
||||||
import time
|
import time
|
||||||
|
from contextlib import contextmanager
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
import httpx
|
import httpx
|
||||||
import pytest
|
import pytest
|
||||||
@@ -79,6 +63,69 @@ _REQ_TOTAL_RE = re.compile(
|
|||||||
_LABEL_RE = re.compile(r'(\w+)="([^"]*)"')
|
_LABEL_RE = re.compile(r'(\w+)="([^"]*)"')
|
||||||
|
|
||||||
|
|
||||||
|
def _open_port() -> int:
|
||||||
|
with socket.socket() as sock:
|
||||||
|
sock.bind(("127.0.0.1", 0))
|
||||||
|
return sock.getsockname()[1]
|
||||||
|
|
||||||
|
|
||||||
|
@contextmanager
|
||||||
|
def _run(binary: Path, env: dict[str, str], log_path: Path):
|
||||||
|
with log_path.open("w") as log:
|
||||||
|
process = subprocess.Popen(
|
||||||
|
[str(binary)],
|
||||||
|
env={**os.environ, **env},
|
||||||
|
stdout=log,
|
||||||
|
stderr=subprocess.STDOUT,
|
||||||
|
start_new_session=True,
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
yield process
|
||||||
|
finally:
|
||||||
|
process.terminate()
|
||||||
|
try:
|
||||||
|
process.wait(timeout=10)
|
||||||
|
except subprocess.TimeoutExpired:
|
||||||
|
process.kill()
|
||||||
|
process.wait(timeout=5)
|
||||||
|
|
||||||
|
|
||||||
|
def _wait_for_indexer(process: subprocess.Popen, port: int, log_path: Path) -> None:
|
||||||
|
deadline = time.time() + 10
|
||||||
|
while time.time() < deadline:
|
||||||
|
if process.poll() is not None:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"KV Indexer exited during startup:\n{log_path.read_text()}"
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
with socket.create_connection(("127.0.0.1", port), timeout=0.2):
|
||||||
|
return
|
||||||
|
except OSError:
|
||||||
|
time.sleep(0.1)
|
||||||
|
raise RuntimeError("timed out waiting for KV Indexer")
|
||||||
|
|
||||||
|
|
||||||
|
def _wait_for_bridge(process: subprocess.Popen, log_path: Path) -> None:
|
||||||
|
deadline = time.time() + 10
|
||||||
|
while time.time() < deadline:
|
||||||
|
output = log_path.read_text(errors="replace")
|
||||||
|
if "bridge session established" in output:
|
||||||
|
# ZMQ connect is asynchronous; let the subscription reach the PUB.
|
||||||
|
time.sleep(0.5)
|
||||||
|
return
|
||||||
|
if process.poll() is not None:
|
||||||
|
raise RuntimeError(f"KV Indexer Bridge exited during startup:\n{output}")
|
||||||
|
time.sleep(0.1)
|
||||||
|
raise RuntimeError(f"timed out waiting for KV Indexer Bridge:\n{output}")
|
||||||
|
|
||||||
|
|
||||||
|
def _dump_logs(logs: dict[str, Path]) -> None:
|
||||||
|
"""Print the tail of each Indexer/Bridge log so a routing failure is debuggable."""
|
||||||
|
for name, path in logs.items():
|
||||||
|
tail = path.read_text(errors="replace")[-4000:] if path.exists() else "<no log>"
|
||||||
|
print(f"\n----- {name} -----\n{tail}")
|
||||||
|
|
||||||
|
|
||||||
def _success_counts_by_worker(router_url: str) -> dict[str, int]:
|
def _success_counts_by_worker(router_url: str) -> dict[str, int]:
|
||||||
"""Scrape ``/metrics`` and return ``{worker_url: success_count}``."""
|
"""Scrape ``/metrics`` and return ``{worker_url: success_count}``."""
|
||||||
r = httpx.get(f"{router_url}/metrics", timeout=5.0)
|
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.real_gpu
|
||||||
@pytest.mark.slow
|
@pytest.mark.slow
|
||||||
def test_two_routers_route_by_prefix_content(
|
def test_routers_route_by_prefix_content(
|
||||||
router_binary, # noqa: ARG001 — fixture forces release-binary presence
|
router_binary,
|
||||||
gpu_allocator,
|
gpu_allocator,
|
||||||
|
tmp_path,
|
||||||
):
|
):
|
||||||
"""Each router must route by prefix CONTENT, agreeing across routers.
|
"""Both the local ZMQ index and the external Indexer must route by content."""
|
||||||
|
|
||||||
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.
|
|
||||||
"""
|
|
||||||
spec = get_model_spec("qwen3-0.6b")
|
spec = get_model_spec("qwen3-0.6b")
|
||||||
gpus = gpu_allocator.acquire(2)
|
gpus = gpu_allocator.acquire(2)
|
||||||
# Workers run with the model's REAL chat template (no override): the engine
|
indexer_port = _open_port()
|
||||||
# caches chat-templated tokens, and the router renders the same template
|
indexer_endpoint = f"http://127.0.0.1:{indexer_port}"
|
||||||
# (loaded from the model's tokenizer_config.json) before hashing. This
|
indexer_binary = router_binary.parent / "kv-indexer-server"
|
||||||
# exercises the production chat-template tokenization path, which aligns
|
bridge_binary = router_binary.parent / "kv-indexer-bridge"
|
||||||
# router query hashes with the engine's templated blocks.
|
logs = {
|
||||||
|
name: tmp_path / f"{name}.log" for name in ("indexer", "bridge-x", "bridge-y")
|
||||||
|
}
|
||||||
try:
|
try:
|
||||||
with (
|
with (
|
||||||
spawn_worker(
|
spawn_worker(
|
||||||
@@ -205,54 +247,79 @@ def test_two_routers_route_by_prefix_content(
|
|||||||
gpu_ids=[gpus[1]],
|
gpu_ids=[gpus[1]],
|
||||||
enable_kv_events=True,
|
enable_kv_events=True,
|
||||||
) as worker_y,
|
) as worker_y,
|
||||||
Gateway() as router_a,
|
_run(
|
||||||
Gateway() as router_b,
|
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]
|
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"],
|
model_id=spec["model"],
|
||||||
tokenizer_path=spec["model"],
|
tokenizer_path=spec["model"],
|
||||||
worker_urls=worker_urls,
|
worker_urls=worker_urls,
|
||||||
policy="cache_aware_zmq",
|
policy="cache_aware_zmq",
|
||||||
timeout=120.0,
|
timeout=120.0,
|
||||||
)
|
)
|
||||||
|
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"])
|
||||||
|
|
||||||
# 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_x.url, spec["model"], PREFIX_X)
|
||||||
_direct_warm(worker_y.url, spec["model"], PREFIX_Y)
|
_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)
|
time.sleep(2.0)
|
||||||
|
|
||||||
# 3. Content-routing assertion (×4): each prefix must land
|
try:
|
||||||
# on the worker that holds it, on either router.
|
for router, label in (
|
||||||
#
|
(local, "local-index"),
|
||||||
# The four assertions below are independently strong:
|
(external, "external-indexer"),
|
||||||
# min-load fallback routes both prefixes on a given
|
):
|
||||||
# router to a single default worker, so for ANY broken-
|
landed = _route_through(
|
||||||
# fan-out scenario at least one of the four fails.
|
router.base_url, spec["model"], PREFIX_X
|
||||||
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 (
|
||||||
assert landed == worker_y.url, (
|
landed == worker_x.url
|
||||||
f"router {label}: PREFIX_Y must route to worker_y "
|
), f"router {label}: PREFIX_X must route to {worker_x.url}; landed on {landed}"
|
||||||
f"({worker_y.url}); landed on {landed}. "
|
landed = _route_through(
|
||||||
f"Likely cause: HashTree is empty — KV-event "
|
router.base_url, spec["model"], PREFIX_Y
|
||||||
f"subscriber never opened, or BlockStored events "
|
|
||||||
f"never reached the pump."
|
|
||||||
)
|
)
|
||||||
|
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:
|
finally:
|
||||||
gpu_allocator.release(gpus)
|
gpu_allocator.release(gpus)
|
||||||
|
|||||||
@@ -172,6 +172,7 @@ class Gateway:
|
|||||||
tokenizer_path: str,
|
tokenizer_path: str,
|
||||||
worker_urls: list[str],
|
worker_urls: list[str],
|
||||||
policy: str = "round_robin",
|
policy: str = "round_robin",
|
||||||
|
kv_indexer_endpoint: str | None = None,
|
||||||
timeout: float = 60.0,
|
timeout: float = 60.0,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Start the router in regular (non-PD) mode.
|
"""Start the router in regular (non-PD) mode.
|
||||||
@@ -186,6 +187,7 @@ class Gateway:
|
|||||||
metadata are learned from ``/server_info``.
|
metadata are learned from ``/server_info``.
|
||||||
policy: Policy kind — ``round_robin``, ``random``, ``power_of_two``,
|
policy: Policy kind — ``round_robin``, ``random``, ``power_of_two``,
|
||||||
or ``cache_aware_zmq``.
|
or ``cache_aware_zmq``.
|
||||||
|
kv_indexer_endpoint: Optional external KV Indexer gRPC endpoint.
|
||||||
timeout: How long to wait for ``/readyz`` before giving up.
|
timeout: How long to wait for ``/readyz`` before giving up.
|
||||||
"""
|
"""
|
||||||
self._launch(
|
self._launch(
|
||||||
@@ -194,6 +196,7 @@ class Gateway:
|
|||||||
tokenizer_path=tokenizer_path,
|
tokenizer_path=tokenizer_path,
|
||||||
urls=list(worker_urls),
|
urls=list(worker_urls),
|
||||||
policy=policy,
|
policy=policy,
|
||||||
|
kv_indexer_endpoint=kv_indexer_endpoint,
|
||||||
),
|
),
|
||||||
timeout=timeout,
|
timeout=timeout,
|
||||||
)
|
)
|
||||||
@@ -293,6 +296,7 @@ class Gateway:
|
|||||||
tokenizer_path: str,
|
tokenizer_path: str,
|
||||||
urls: list[str],
|
urls: list[str],
|
||||||
policy: str,
|
policy: str,
|
||||||
|
kv_indexer_endpoint: str | None = None,
|
||||||
) -> list[str]:
|
) -> list[str]:
|
||||||
resolved_tokenizer = _resolve_tokenizer_path(tokenizer_path)
|
resolved_tokenizer = _resolve_tokenizer_path(tokenizer_path)
|
||||||
|
|
||||||
@@ -317,6 +321,8 @@ class Gateway:
|
|||||||
"--stale-request-timeout-secs",
|
"--stale-request-timeout-secs",
|
||||||
str(self.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
|
# `--worker-urls` is multi-valued; keep it last so clap doesn't
|
||||||
# absorb a following flag as a URL.
|
# absorb a following flag as a URL.
|
||||||
args += ["--worker-urls", *urls]
|
args += ["--worker-urls", *urls]
|
||||||
|
|||||||
@@ -10,7 +10,11 @@ FROM rust:1.90-bookworm AS builder
|
|||||||
# `channel = "1.90"`.
|
# `channel = "1.90"`.
|
||||||
ENV RUSTUP_TOOLCHAIN=1.90.0
|
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
|
WORKDIR /build
|
||||||
|
|
||||||
|
|||||||
@@ -12,18 +12,10 @@
|
|||||||
//! doesn't render tool schemas, so its ids would diverge from the engine).
|
//! doesn't render tool schemas, so its ids would diverge from the engine).
|
||||||
//! * A request with multimodal (array) content → `input_ids` omitted (a text
|
//! * A request with multimodal (array) content → `input_ids` omitted (a text
|
||||||
//! tokenizer can't represent image content).
|
//! 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::body::Body;
|
||||||
use axum::http::{Request, StatusCode};
|
use axum::http::{Request, StatusCode};
|
||||||
use serde_json::{json, Value};
|
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::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||||
use sgl_router::policies::factory::build_registry;
|
use sgl_router::policies::factory::build_registry;
|
||||||
use sgl_router::policies::kv_events::{BlockSizeOracle, HashTree};
|
use sgl_router::policies::kv_events::{BlockSizeOracle, HashTree};
|
||||||
@@ -36,33 +28,9 @@ use std::sync::Arc;
|
|||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
use tower::ServiceExt;
|
use tower::ServiceExt;
|
||||||
|
|
||||||
|
use crate::common::cache_aware_fixture::{config, MODEL};
|
||||||
use crate::common::mock_worker::MockWorker;
|
use crate::common::mock_worker::MockWorker;
|
||||||
|
|
||||||
const MODEL: &str = "deepseek-v4-tiny";
|
|
||||||
|
|
||||||
fn config() -> Config {
|
|
||||||
Config {
|
|
||||||
server: ServerConfig {
|
|
||||||
host: "0".into(),
|
|
||||||
port: 0,
|
|
||||||
},
|
|
||||||
observability: ObservabilityConfig::default(),
|
|
||||||
model: ModelConfig {
|
|
||||||
id: MODEL.into(),
|
|
||||||
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
|
|
||||||
policy: PolicyKind::CacheAwareZmq,
|
|
||||||
circuit_breaker: None,
|
|
||||||
cache_aware: Some(CacheAwareConfig::default()),
|
|
||||||
sticky: None,
|
|
||||||
},
|
|
||||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
|
||||||
urls: vec!["http://placeholder:0".into()],
|
|
||||||
}),
|
|
||||||
proxy: ProxyConfig::default(),
|
|
||||||
active_load: ActiveLoadConfig::default(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn build_ctx(url: String) -> Arc<AppContext> {
|
fn build_ctx(url: String) -> Arc<AppContext> {
|
||||||
let cfg = config();
|
let cfg = config();
|
||||||
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
|
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
|
||||||
|
|||||||
@@ -0,0 +1,39 @@
|
|||||||
|
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||||
|
// SPDX-License-Identifier: Apache-2.0
|
||||||
|
|
||||||
|
//! Shared router config for the cache-aware proxy tests.
|
||||||
|
//!
|
||||||
|
//! The model id contains `deepseek-v4` so the tokenizer registry auto-attaches the
|
||||||
|
//! built-in V4 chat encoder — the engine-equivalent path — with no template fixture.
|
||||||
|
|
||||||
|
use sgl_router::config::{
|
||||||
|
ActiveLoadConfig, CacheAwareConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig,
|
||||||
|
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
||||||
|
};
|
||||||
|
|
||||||
|
pub const MODEL: &str = "deepseek-v4-tiny";
|
||||||
|
|
||||||
|
/// A single-model `cache_aware_zmq` router. Discovery is a placeholder because
|
||||||
|
/// every caller installs its own `WorkerRegistry`.
|
||||||
|
pub fn config() -> Config {
|
||||||
|
Config {
|
||||||
|
server: ServerConfig {
|
||||||
|
host: "0".into(),
|
||||||
|
port: 0,
|
||||||
|
},
|
||||||
|
observability: ObservabilityConfig::default(),
|
||||||
|
model: ModelConfig {
|
||||||
|
id: MODEL.into(),
|
||||||
|
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
|
||||||
|
policy: PolicyKind::CacheAwareZmq,
|
||||||
|
circuit_breaker: None,
|
||||||
|
cache_aware: Some(CacheAwareConfig::default()),
|
||||||
|
sticky: None,
|
||||||
|
},
|
||||||
|
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||||
|
urls: vec!["http://placeholder:0".into()],
|
||||||
|
}),
|
||||||
|
proxy: ProxyConfig::default(),
|
||||||
|
active_load: ActiveLoadConfig::default(),
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -3,5 +3,6 @@
|
|||||||
|
|
||||||
//! Shared test harness re-exports.
|
//! Shared test harness re-exports.
|
||||||
|
|
||||||
|
pub mod cache_aware_fixture;
|
||||||
pub mod mock_worker;
|
pub mod mock_worker;
|
||||||
pub mod streaming;
|
pub mod streaming;
|
||||||
|
|||||||
@@ -0,0 +1,135 @@
|
|||||||
|
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||||
|
// SPDX-License-Identifier: Apache-2.0
|
||||||
|
|
||||||
|
//! Full HTTP routing path backed by a real in-memory Indexer gRPC server.
|
||||||
|
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use axum::body::Body;
|
||||||
|
use axum::http::{Request, StatusCode};
|
||||||
|
use serde_json::json;
|
||||||
|
use sgl_kv_indexer::pb::kv_indexer_client::KvIndexerClient;
|
||||||
|
use sgl_kv_indexer::pb::{
|
||||||
|
ApplyExternalKvBatchRequest, ExternalKvAction, ExternalKvActionType, TierType,
|
||||||
|
};
|
||||||
|
use sgl_kv_indexer::{
|
||||||
|
server_builder, GrpcPrefixIndex, InMemoryKvIndexerBackend, KvIndexerService, PrefixIndexConfig,
|
||||||
|
};
|
||||||
|
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||||
|
use sgl_router::policies::factory::build_registry;
|
||||||
|
use sgl_router::policies::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree};
|
||||||
|
use sgl_router::policies::request_tokens_for;
|
||||||
|
use sgl_router::proxy::Proxy;
|
||||||
|
use sgl_router::server::app::build_router;
|
||||||
|
use sgl_router::server::app_context::AppContext;
|
||||||
|
use sgl_router::tokenizer::TokenizerRegistry;
|
||||||
|
use sgl_router::workers::WorkerRegistry;
|
||||||
|
use tokio_stream::wrappers::TcpListenerStream;
|
||||||
|
use tower::ServiceExt;
|
||||||
|
|
||||||
|
use crate::common::cache_aware_fixture::{config, MODEL};
|
||||||
|
use crate::common::mock_worker::MockWorker;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn external_indexer_routes_to_the_cached_worker() {
|
||||||
|
let cached = MockWorker::start(vec![]).await;
|
||||||
|
let uncached = MockWorker::start(vec![]).await;
|
||||||
|
let cfg = config();
|
||||||
|
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
|
||||||
|
let body = json!({
|
||||||
|
"model": MODEL,
|
||||||
|
"messages": [{"role": "user", "content": "hello there friend"}],
|
||||||
|
});
|
||||||
|
let tokens = request_tokens_for(&tokenizers, &ModelId(MODEL.into()), &body)
|
||||||
|
.expect("test prompt tokenizes");
|
||||||
|
let hashes = compute_block_hashes(&tokens.ids, 1);
|
||||||
|
assert!(!hashes.is_empty());
|
||||||
|
|
||||||
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||||
|
let endpoint = format!("http://{}", listener.local_addr().unwrap());
|
||||||
|
let server = tokio::spawn(async move {
|
||||||
|
server_builder()
|
||||||
|
.add_service(KvIndexerService::new(InMemoryKvIndexerBackend::new()).into_server())
|
||||||
|
.serve_with_incoming(TcpListenerStream::new(listener))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
});
|
||||||
|
|
||||||
|
let mut indexer = KvIndexerClient::connect(endpoint.clone()).await.unwrap();
|
||||||
|
indexer
|
||||||
|
.apply_external_kv_batch(ApplyExternalKvBatchRequest {
|
||||||
|
worker_id: "cached-worker".into(),
|
||||||
|
seq: 1,
|
||||||
|
actions: vec![ExternalKvAction {
|
||||||
|
r#type: ExternalKvActionType::ActionReport as i32,
|
||||||
|
tier: TierType::TierHbm as i32,
|
||||||
|
hashes: hashes.clone(),
|
||||||
|
component_masks: Vec::new(),
|
||||||
|
block_sizes: Vec::new(),
|
||||||
|
}],
|
||||||
|
worker_address: cached.url.clone(),
|
||||||
|
cache_spec: None,
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let registry = Arc::new(WorkerRegistry::default());
|
||||||
|
for url in [&cached.url, &uncached.url] {
|
||||||
|
registry
|
||||||
|
.add(WorkerSpec {
|
||||||
|
id: WorkerId(url.clone()),
|
||||||
|
url: url.clone(),
|
||||||
|
mode: WorkerMode::Plain,
|
||||||
|
model_ids: vec![ModelId(MODEL.into())],
|
||||||
|
bootstrap_port: None,
|
||||||
|
})
|
||||||
|
.unwrap();
|
||||||
|
}
|
||||||
|
let oracle = BlockSizeOracle::new();
|
||||||
|
oracle.try_set(1).unwrap();
|
||||||
|
let policies = Arc::new(
|
||||||
|
build_registry(
|
||||||
|
&cfg,
|
||||||
|
Arc::new(HashTree::new()),
|
||||||
|
Arc::clone(&tokenizers),
|
||||||
|
Arc::clone(&oracle),
|
||||||
|
)
|
||||||
|
.unwrap(),
|
||||||
|
);
|
||||||
|
let mut ctx = AppContext::new(
|
||||||
|
cfg,
|
||||||
|
tokenizers,
|
||||||
|
Arc::new(Proxy::new(Duration::from_secs(5)).unwrap()),
|
||||||
|
registry,
|
||||||
|
policies,
|
||||||
|
);
|
||||||
|
ctx.prefix_index = Some(Arc::new(
|
||||||
|
GrpcPrefixIndex::new(PrefixIndexConfig {
|
||||||
|
endpoint,
|
||||||
|
query_deadline: Duration::from_secs(1),
|
||||||
|
max_inflight: 4,
|
||||||
|
})
|
||||||
|
.unwrap(),
|
||||||
|
));
|
||||||
|
ctx.block_size_oracle = oracle;
|
||||||
|
|
||||||
|
let app = build_router(Arc::new(ctx));
|
||||||
|
let response = app
|
||||||
|
.oneshot(
|
||||||
|
Request::builder()
|
||||||
|
.method("POST")
|
||||||
|
.uri("/v1/chat/completions")
|
||||||
|
.header("content-type", "application/json")
|
||||||
|
.body(Body::from(serde_json::to_vec(&body).unwrap()))
|
||||||
|
.unwrap(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(response.status(), StatusCode::OK);
|
||||||
|
assert!(cached.captured.lock().unwrap().last_body.is_some());
|
||||||
|
assert!(uncached.captured.lock().unwrap().last_body.is_none());
|
||||||
|
|
||||||
|
server.abort();
|
||||||
|
}
|
||||||
@@ -12,6 +12,7 @@ mod common;
|
|||||||
|
|
||||||
mod cache_aware_input_ids;
|
mod cache_aware_input_ids;
|
||||||
mod chat_routing;
|
mod chat_routing;
|
||||||
|
mod external_indexer_routing;
|
||||||
mod failover;
|
mod failover;
|
||||||
mod graceful_shutdown;
|
mod graceful_shutdown;
|
||||||
mod header_forwarding;
|
mod header_forwarding;
|
||||||
|
|||||||
Reference in New Issue
Block a user