diff --git a/Cargo.lock b/Cargo.lock index 351821513b76..9d232619f727 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2592,7 +2592,10 @@ dependencies = [ "dynamo-truthy", "fastrand", "flume 0.12.0", + "futures", "indexmap 2.14.0", + "k8s-openapi", + "kube", "once_cell", "ordered-float 4.6.0", "parking_lot", diff --git a/container/build-indexer-image.sh b/container/build-indexer-image.sh new file mode 100755 index 000000000000..4b713baf5105 --- /dev/null +++ b/container/build-indexer-image.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +# Build the standalone KV-cache indexer image from COMMITTED dynamo source. +# +# ./container/build-indexer-image.sh [tag] +# +# Builds a multi-stage image (compile wheel -> slim runtime) using a git-archive +# of HEAD as the build context, so the multi-GB local target/ is never shipped +# to the docker daemon. Default tag: /dynamo-indexer:kvtest-. +# Push separately: docker push +set -euo pipefail + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +REGISTRY="${REGISTRY:-localhost:30500}" +SHA="$(git -C "$REPO_ROOT" rev-parse --short=10 HEAD)" +TAG="${1:-${REGISTRY}/dynamo-indexer:kvtest-${SHA}}" + +BUILD_DIR="$(mktemp -d)" +trap 'rm -rf "$BUILD_DIR"' EXIT + +echo "==> git archive ${SHA} -> build context (committed files only)" +git -C "$REPO_ROOT" archive --format=tar HEAD -o "$BUILD_DIR/dynamo-src.tar" +cp "$REPO_ROOT/container/indexer.Dockerfile" "$BUILD_DIR/Dockerfile" + +echo "==> docker build $TAG" +# --ulimit nofile: BuildKit RUN steps otherwise inherit a low fd limit, which a +# big parallel cargo build exhausts ("too many open files"). The Dockerfile also +# caps CARGO_BUILD_JOBS, but raise the ceiling here too. +DOCKER_BUILDKIT=1 docker build \ + --ulimit nofile=1048576:1048576 \ + -f "$BUILD_DIR/Dockerfile" -t "$TAG" "$BUILD_DIR" + +echo "==> done: $TAG" +echo " push with: docker push $TAG" diff --git a/container/indexer.Dockerfile b/container/indexer.Dockerfile new file mode 100644 index 000000000000..ced345c512c6 --- /dev/null +++ b/container/indexer.Dockerfile @@ -0,0 +1,93 @@ +# syntax=docker/dockerfile:1.7 +# ============================================================================= +# Standalone KV-cache indexer image -> `python -m dynamo.indexer` +# +# Serves DeepInfra's per-model kv-indexer services; one image, the flavor is a flag: +# kv-indexer:reality (no flags) real evictions, per-shard trees +# kv-indexer:routing --keep-evictions evictions parked, per-shard trees +# kv-indexer:h24 --h24 flat "stored by anyone" counterfactual +# Each discovers the model's engine pods, subscribes to their ZMQ KV events and +# serves /query /query_by_hash /workers /dump /health (+ /metrics). +# +# Self-contained: compiles the Rust python bindings (dynamo._core) from source, +# then installs ONLY that wheel into a slim runtime. Both stages are ubuntu:24.04, +# so the runtime libstdc++ is identical to what the extension was linked against +# -> no LD_PRELOAD shim needed (that hack was only for the mismatched conda env). +# +# The build context is a `git archive` of the dynamo source (committed files +# only), so the multi-GB local target/ never enters the build. The indexer is +# fully standalone: no etcd / NATS, just HTTP + ZMQ. See build-indexer-image.sh. +# ============================================================================= + +# ---------- builder: compile the wheel ------------------------------------- +FROM ubuntu:24.04 AS builder +ENV DEBIAN_FRONTEND=noninteractive +RUN apt-get update && apt-get install -y --no-install-recommends \ + ca-certificates curl git \ + build-essential pkg-config cmake \ + clang libclang-dev \ + protobuf-compiler libprotobuf-dev \ + libzmq3-dev \ + python3 python3-dev python3-venv \ + && rm -rf /var/lib/apt/lists/* + +# Rust: rustup with NO default toolchain. dynamo's rust-toolchain.toml pins +# the toolchain, which rustup auto-fetches on the first cargo run in /src. +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \ + | sh -s -- -y --no-modify-path --default-toolchain none --profile minimal +ENV PATH=/root/.cargo/bin:$PATH + +# maturin in an isolated venv (Ubuntu 24.04 system python is externally-managed). +RUN python3 -m venv /opt/maturin && /opt/maturin/bin/pip install --no-cache-dir maturin +ENV PATH=/opt/maturin/bin:$PATH + +# Source tarball (git archive), auto-extracted by ADD. rust-toolchain.toml lands +# at /src so rustup resolves the pinned toolchain for the whole workspace. +ADD dynamo-src.tar /src +# git archive stamps every file with the commit time, and cargo judges a path +# crate fresh by mtime against the shared /cargo-target cache. A sibling build +# (another branch) can therefore leave a newer artifact that cargo reuses with +# the wrong contents; stamping the sources with the build time prevents that. +RUN find /src -type f -exec touch {} + + +WORKDIR /src/lib/bindings/python +ENV CARGO_TARGET_DIR=/cargo-target +# Cap parallelism: this compiles on a shared ~288-core services node. The default +# (one rustc per core) both starves prod neighbours and exhausts the build's +# file-descriptor limit (hundreds of parallel rustc -> "too many open files"). +# 16 jobs is plenty and well-behaved; build-indexer-image.sh also raises nofile. +ENV CARGO_BUILD_JOBS=16 +# Release build: optimized for real indexer throughput (slower to compile than a +# debug build, but the radix-tree query/apply hot paths need the optimizations). +# kv-indexer-metrics adds the `dynamo.indexer` binary + Prometheus /metrics. +# nixl-sys' build script runs bindgen, so point it at libclang (resolved +# dynamically to survive llvm bumps). Cache mounts keep the crate registry + +# compiled artifacts warm across rebuilds. +RUN --mount=type=cache,target=/root/.cargo/registry \ + --mount=type=cache,target=/cargo-target \ + export LIBCLANG_PATH="$(dirname "$(find /usr/lib/llvm-* -name 'libclang.so*' 2>/dev/null | head -1)")" \ + && maturin build --release --features kv-indexer-metrics --out /wheels + +# ---------- runtime: install just the wheel -------------------------------- +FROM ubuntu:24.04 AS runtime +ENV DEBIAN_FRONTEND=noninteractive +RUN apt-get update && apt-get install -y --no-install-recommends \ + ca-certificates \ + libzmq5 \ + python3 python3-venv \ + && rm -rf /var/lib/apt/lists/* + +# venv with only the indexer wheel and its deps (pydantic, uvloop -> manylinux +# wheels from PyPI, no compiler needed here). +COPY --from=builder /wheels/ /wheels/ +RUN python3 -m venv /opt/venv \ + && /opt/venv/bin/pip install --no-cache-dir /wheels/ai_dynamo_runtime*.whl \ + && rm -rf /wheels +ENV PATH=/opt/venv/bin:$PATH + +# Smoke-test the import + CLI at build time so a broken wheel fails the build. +RUN python -m dynamo.indexer --help >/dev/null + +EXPOSE 8090 +ENTRYPOINT ["python", "-m", "dynamo.indexer"] +CMD ["--port", "8090"] diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md index 3483286a242c..b0624e653124 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md @@ -201,6 +201,10 @@ Returns: Register a ZMQ endpoint for an instance. Each call creates or reuses the indexer for the given `(model_name, routing_group)` pair. Registration is non-blocking: if the worker is not up yet, the listener is accepted in `pending` state and transitions to `active` once the initial ZMQ connection succeeds. +Repeating `/register` for an existing `instance_id` and `dp_rank` returns an error; +it does not replace the listener or reset its sequence watermark. For worker restarts, +see [Worker Restarts](#worker-restarts). + ```bash # Single model, default routing group curl -X POST http://localhost:8090/register \ @@ -490,6 +494,29 @@ recovers history still retained by the engine. Without a replay endpoint, or whe range is no longer retained, restart with a healthy indexer in `--peers` to recover its startup snapshot, or perform another full resynchronization before relying on the rebuilt worker state. +## Worker Restarts + +A ZMQ reconnect does not identify a new worker or cache-event publisher lifetime. +If the publisher restarts with its sequence reset, an existing listener retains its +watermark and discards batches at or below it. Previously indexed ownership can also +remain stale. + +On a worker or publisher restart, call `/unregister` and wait for the response before +calling `/register` with the current worker metadata. For a whole-worker restart, +omit `dp_rank` from `/unregister` to remove all ranks. Apply this lifecycle change to +every indexer replica; peer registration does not propagate worker lifecycle changes. +Serialize lifecycle operations for the same worker. + +This sequence requests removal of old cache ownership and starts a fresh sequence +watermark. It does not certify a complete cache view or provide an atomic boundary +against in-flight old events. Registration success and listener status `active` do +not establish that startup events were received or replayed. Recover missing history +before relying on cache overlap, subject to the limits in +[Gap Detection and Replay](#gap-detection-and-replay). + +The [standalone selector](standalone-selection.md#worker-restarts) has a different +registration API: updating a schedulable worker reconciles its indexer registration. + ## Limitations - **Standalone mode is ZMQ only**: Workers must publish KV events via ZMQ PUB sockets. diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md index 4a740248663e..9bcd9b9da3fe 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md @@ -127,6 +127,26 @@ both default to `"default"` when omitted. `GET /health` is process liveness. `GET /ready` returns `200` only after at least one worker is schedulable, otherwise `503` with lifecycle details. +### Worker Restarts + +When a worker or its cache-event publisher restarts, update its registration on every +selector replica. A reconnect at the same address does not reset the existing event +sequence watermark or invalidate old cache ownership. + +With KV events enabled, `POST /workers` for an existing schedulable worker drains it, +removes its indexer registration, and creates new listeners with fresh sequence +watermarks. Supply the complete current worker record because this endpoint replaces +the catalog record. `PATCH /workers/{worker_id}` also reconciles a schedulable worker; +even an unchanged update can rebuild its cache view, so avoid using registration +updates as periodic heartbeats. Serialize lifecycle operations for the same worker. + +This differs from the standalone indexer's `/register`, which rejects duplicate +registrations and requires `/unregister` first. See +[Indexer Worker Restarts](standalone-indexer.md#worker-restarts) for cleanup and +replay limits. Neither a successful catalog update nor `/ready` certifies a complete +cache view: missed startup events must still be recovered, and expired replay +history requires another source of complete state. + ## Selection API ### `POST /select` diff --git a/lib/bindings/python/Cargo.lock b/lib/bindings/python/Cargo.lock index cb812629cd5a..c47ee5620a1c 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -2133,7 +2133,10 @@ dependencies = [ "dynamo-truthy", "fastrand", "flume", + "futures", "indexmap 2.14.0", + "k8s-openapi", + "kube", "once_cell", "ordered-float 4.6.0", "parking_lot", diff --git a/lib/bindings/python/Cargo.toml b/lib/bindings/python/Cargo.toml index ff24d9e9218d..8a8a7d852799 100644 --- a/lib/bindings/python/Cargo.toml +++ b/lib/bindings/python/Cargo.toml @@ -33,7 +33,7 @@ custom-policy = [ "dynamo-kv-router/standalone-selection", ] media-ffmpeg = ["dynamo-llm/media-ffmpeg"] -kv-indexer = ["dep:clap", "dep:tracing-subscriber", "dynamo-kv-router/standalone-indexer"] +kv-indexer = ["dep:clap", "dep:tracing-subscriber", "dynamo-kv-router/standalone-indexer", "dynamo-kv-router/kube-discovery"] kv-indexer-metrics = ["kv-indexer", "dynamo-kv-router/metrics"] # Exposes the optional Relay stats/snapshot methods to Python diagnostic builds. ckf-diagnostics = ["dynamo-llm/ckf-diagnostics"] diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 91c86ce80878..e08938ef7086 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -228,6 +228,122 @@ struct KvIndexerCli { /// Use local timezone for access log timestamps (default: UTC) #[arg(long)] access_log_local_time: bool, + + /// Total timeout in seconds for one engine `GET /kv_recover` download + /// (connect + body). A full TreeDump of a large engine is tens of MB. + #[arg(long, default_value_t = indexer::kv_recover::DEFAULT_RECOVER_TIMEOUT_S)] + recover_timeout_secs: u64, + + /// Maximum concurrent `/kv_recover` downloads across all listeners of this + /// indexer. Bounds the load a fleet-wide (re)subscription puts on engines. + #[arg(long, default_value_t = indexer::kv_recover::DEFAULT_RECOVER_CONCURRENCY)] + recover_concurrency: usize, + + /// Park Removed events in a buffer (and drop Cleared) instead of applying + /// them, approximating infinite memory (the kv-indexer:routing flavor). + /// Buffered evictions older than --evict-retention-secs are replayed into + /// the tree once memory use crosses --evict-memory-threshold of the limit. + #[arg(long, default_value_t = false)] + keep_evictions: bool, + + /// Minimum age (seconds) a parked eviction must reach before the memory + /// sweep may apply it. Only meaningful with --keep-evictions. + #[arg(long, default_value_t = 1800)] + evict_retention_secs: u64, + + /// Fraction of the memory limit (0..1] above which the sweep replays aged + /// evictions. Only meaningful with --keep-evictions. + #[arg(long, default_value_t = 0.75)] + evict_memory_threshold: f64, + + /// Emit verbose audit logs on the `kv_audit` tracing target: one line per + /// query (block hashes + full response) and one per store/evict/clear + /// event ingested from the engines. Filter with RUST_LOG=kv_audit=info. + #[arg(long, default_value_t = false)] + enable_logging: bool, + + /// Run as the flat h24 counterfactual (kv-indexer:h24): one map per model + /// of every prefix stored by any worker within --h24-horizon-secs, + /// ignoring evictions. Exclusive with --keep-evictions. + #[arg(long, default_value_t = false)] + h24: bool, + + /// Retention horizon of the h24 expiry sweep, in seconds. + #[arg(long, default_value_t = 172800)] + h24_horizon_secs: u64, + + /// Kubernetes namespace to watch for engine pods. Together with + /// --watch-model-name (or --watch-label) this enables pod auto-discovery: + /// subscribe on Ready, unsubscribe on delete. + #[arg(long)] + watch_namespace: Option, + + /// Raw label selector for this model's engine pods, e.g. + /// "engine_hash=d4b7a85131172ca6". Overrides the selector derived from + /// --watch-model-name. + #[arg(long)] + watch_label: Option, + + /// ZMQ KV-event port the discovered engines publish on. + #[arg(long, default_value_t = 5557)] + watch_zmq_port: u16, + + /// HTTP port serving `GET /kv_recover` on the engines, used for gap + /// recovery at `http://:/kv_recover`. + #[arg(long)] + watch_recover_port: Option, + + /// Data-parallel ranks per engine pod (vLLM --data-parallel-size). Rank r + /// publishes on --watch-zmq-port + r and recovers on + /// --watch-recover-port + r. Leaving this at 1 on a DP engine indexes + /// only rank 0's cache. + #[arg(long, default_value_t = 1)] + watch_dp_size: u32, + + /// Model whose engine pods to discover, e.g. "openai/gpt-oss-120b". The + /// watch uses `di/model_name=` unless --watch-label is + /// given; discovered pods register under this name (default: + /// --model-name). + #[arg(long)] + watch_model_name: Option, +} + +#[cfg(feature = "kv-indexer")] +impl KvIndexerCli { + /// Pod-discovery config, `None` unless --watch-namespace is set. + fn kube_discovery(&self) -> anyhow::Result> { + let Some(namespace) = self.watch_namespace.clone() else { + anyhow::ensure!( + self.watch_label.is_none() && self.watch_model_name.is_none(), + "--watch-label/--watch-model-name require --watch-namespace" + ); + return Ok(None); + }; + let label_selector = match (&self.watch_label, &self.watch_model_name) { + (Some(raw), _) => raw.clone(), + (None, Some(model_name)) => indexer::discovery::model_name_label_selector(model_name), + (None, None) => { + anyhow::bail!("--watch-namespace requires --watch-model-name or --watch-label") + } + }; + let block_size = self + .block_size + .ok_or_else(|| anyhow::anyhow!("--block-size is required with --watch-namespace"))?; + anyhow::ensure!(self.watch_dp_size >= 1, "--watch-dp-size must be at least 1"); + Ok(Some(indexer::discovery::KubeDiscoveryConfig { + namespace, + label_selector, + zmq_port: self.watch_zmq_port, + recover_port: self.watch_recover_port, + dp_size: self.watch_dp_size, + model_name: self + .watch_model_name + .clone() + .unwrap_or_else(|| self.model_name.clone()), + routing_group: self.routing_group.clone(), + block_size, + })) + } } pub fn run_kv_indexer_cli(args: I) -> anyhow::Result<()> @@ -248,6 +364,8 @@ where anyhow::anyhow!("invalid --trace-id-header '{}': {e}", cli.trace_id_header) })?; + let kube_discovery = cli.kube_discovery()?; + init_standalone_logging(); let rt = tokio::runtime::Runtime::new()?; @@ -262,6 +380,19 @@ where access_log: cli.access_log, trace_id_header, access_log_local_time: cli.access_log_local_time, + kv_recover: indexer::kv_recover::KvRecoverSettings { + timeout_s: cli.recover_timeout_secs, + concurrency: cli.recover_concurrency, + }, + kube_discovery, + audit_log: cli.enable_logging, + h24_horizon_s: cli.h24.then_some(cli.h24_horizon_secs), + keep_evictions: cli.keep_evictions.then_some( + indexer::evictions::KeepEvictionsConfig { + retention_s: cli.evict_retention_secs, + memory_threshold: cli.evict_memory_threshold, + }, + ), })) } diff --git a/lib/kv-router/Cargo.toml b/lib/kv-router/Cargo.toml index b1d1f4642bb9..c3dc1cff2452 100644 --- a/lib/kv-router/Cargo.toml +++ b/lib/kv-router/Cargo.toml @@ -21,6 +21,8 @@ profile = [] standalone-indexer = ["dep:axum", "dep:serde_json", "dep:reqwest", "dep:zmq", "dep:chrono", "dep:tracing-appender"] standalone-slot-tracker = ["dep:axum", "dep:serde_json", "dep:zmq"] standalone-selection = ["standalone-indexer", "dep:once_cell"] +# Kubernetes pod auto-discovery for the standalone indexer (pulls in the kube client). +kube-discovery = ["standalone-indexer", "dep:kube", "dep:k8s-openapi", "dep:futures"] [dependencies] # repo @@ -76,6 +78,11 @@ chrono = { workspace = true, optional = true } reqwest = { version = "0.12.24", default-features = false, features = ["json"], optional = true } tracing-appender = { version = "0.2", optional = true } +# kube-discovery (optional) +kube = { workspace = true, optional = true } +k8s-openapi = { workspace = true, optional = true } +futures = { workspace = true, optional = true } + [dev-dependencies] criterion = "0.5" rstest = "0.18.2" diff --git a/lib/kv-router/src/indexer/README.md b/lib/kv-router/src/indexer/README.md index 5f3e070b271a..e1a9e756d5e2 100644 --- a/lib/kv-router/src/indexer/README.md +++ b/lib/kv-router/src/indexer/README.md @@ -114,9 +114,25 @@ RadixBlock ├── state.worker_cutoffs: HashMap ├── state.full_edge_workers: HashSet ├── children: HashMap +├── parent: Weak> └── internal: bool ``` +Children and worker lookups hold strong references; parent links are weak to avoid +reference cycles. A child is keyed by its first local block hash. Splits preserve +that key for the prefix and update the parent links of children moved to the suffix. +Cleanup detaches an unowned node only if the parent's slot still points to it, so +an old detached node cannot remove a replacement. + +This single-threaded tree cleans up eagerly during mutations. The shared helper in +[`cleanup.rs`](../cleanup.rs) uses `Arc>` and periodic concurrent sweeps; +it does not support this tree's `Rc>` handles. + +The existing mid-chain removal gap can leave descendant lookup entries alive; see +the `TODO(CORRECTNESS)` in `apply_removed`. Such detached nodes can also retain weak +parent links. A weak link retains the parent's allocation after its value is +dropped. Complete descendant and back-link cleanup remains separate follow-up work. + ### Visual Representation ``` @@ -151,9 +167,9 @@ RadixBlock **remove_blocks(worker, block_hashes)**: 1. For each hash, find node via `lookup[worker][hash]` -2. Remove worker from node's `workers` set -3. If `workers` empty, clear children (cascading cleanup) -4. Remove from `lookup[worker]` +2. Truncate the worker's coverage at the removed block +3. Clear child edges if no worker covers the full node; detach the node if it has no owners +4. Remove newly uncovered hashes from `lookup[worker]` **find_matches(local_hashes, early_exit)**: 1. Start at root with all workers as candidates diff --git a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/mod.rs b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/mod.rs index 19b61e046781..7ef4efc217f9 100644 --- a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/mod.rs +++ b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/mod.rs @@ -19,7 +19,7 @@ use super::{ MatchDetails, PreBoundEventCounters, SyncIndexer, WorkerLookupStats, WorkerTask, }; use crate::cleanup::{CleanupGuard, CleanupState}; -use crate::lookup_update::update_arc_lookup_for_keys; +use crate::lookup_update::{update_arc_lookup_for_keys, update_existing_arc_lookup_for_keys}; use crate::protocols::*; mod children; diff --git a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/repair.rs b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/repair.rs index 1576c6e12c0a..18c654a70068 100644 --- a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/repair.rs +++ b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/repair.rs @@ -61,7 +61,7 @@ impl ConcurrentRadixTreeCompressed { let mut changed_entries_total = 0u64; for (&worker, worker_lookup) in lookup.iter_mut() { - let _changed_entries = update_arc_lookup_for_keys( + let _changed_entries = update_existing_arc_lookup_for_keys( worker_lookup, resolved.lookup_hashes_for_worker_repair(worker, hash, direction), resolved, diff --git a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/tests.rs b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/tests.rs index f327dfbefbe2..948674449b96 100644 --- a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/tests.rs +++ b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/tests.rs @@ -1379,3 +1379,66 @@ fn remove_after_split_and_children_clear_scrubs_lookup() { apply_direct(&index, &mut l2, remove_event(2, 103, 0, w2_leaf)); assert_eq!(worker_lookup_len(&l2, w2), Some(0)); } + +#[test] +fn successful_repair_does_not_restore_scrubbed_other_worker_entries() { + let index = ConcurrentRadixTreeCompressed::new(); + let scrubbed_worker = worker(1); + let repairing_worker = worker(2); + let mut shared_lookup = direct_lookup(); + let mut splitter_lookup = direct_lookup(); + + apply_direct( + &index, + &mut shared_lookup, + make_store_event(1, &[1, 2, 3, 4]), + ); + apply_direct( + &index, + &mut shared_lookup, + make_store_event(2, &[1, 2, 3, 4]), + ); + + // A split from another event thread leaves both local lookups pointing at + // the old prefix node for the [3, 4] suffix. + apply_direct( + &index, + &mut splitter_lookup, + make_store_event(3, &[1, 2, 9]), + ); + let suffix_hashes = remove_hashes_with_parent(&[1, 2], &[3, 4]); + + // A resolve-miss remove scrubs lookup state but cannot update coverage on + // the node it failed to find. Model that boundary directly while keeping + // the resolved live node's coverage intact. + let scrubbed_lookup = shared_lookup.get_mut(&scrubbed_worker).unwrap(); + for hash in &suffix_hashes { + scrubbed_lookup.remove(hash); + } + assert_eq!(worker_lookup_len(&shared_lookup, scrubbed_worker), Some(2)); + assert_direct_score(&index, &[1, 2, 3, 4], scrubbed_worker, 4); + + // A different worker's stale lookup successfully resolves the suffix and + // repairs the event thread's shared lookup map. + let resolved = index + .resolve_lookup( + &mut shared_lookup, + repairing_worker, + suffix_hashes[0], + LookupRepairDirection::TowardTail, + ) + .expect("repairing worker should resolve the split suffix"); + + let repairing_lookup = shared_lookup.get(&repairing_worker).unwrap(); + for hash in &suffix_hashes { + assert!(Arc::ptr_eq(repairing_lookup.get(hash).unwrap(), &resolved)); + } + + let scrubbed_lookup = shared_lookup.get(&scrubbed_worker).unwrap(); + for hash in suffix_hashes { + assert!( + !scrubbed_lookup.contains_key(&hash), + "repair for another worker restored a scrubbed lookup entry" + ); + } +} diff --git a/lib/kv-router/src/indexer/lower_tier.rs b/lib/kv-router/src/indexer/lower_tier.rs index 8445072149f6..77fe9d7948dd 100644 --- a/lib/kv-router/src/indexer/lower_tier.rs +++ b/lib/kv-router/src/indexer/lower_tier.rs @@ -622,7 +622,7 @@ impl LowerTierIndexer { let indexed_owner = IndexedResidencyOwner::from_exact(owner); let remove_worker_entry = { let Some(owner_state) = worker_blocks.get_mut(&indexed_owner) else { - return Err(KvCacheEventError::BlockNotFound); + return Ok(()); }; if owner_state.owner != owner { return Err(KvCacheEventError::UnsupportedResidencyDomain); @@ -631,7 +631,7 @@ impl LowerTierIndexer { for block_hash in block_hashes { let Some(key) = worker_map.remove(block_hash) else { - return Err(KvCacheEventError::BlockNotFound); + continue; }; self.remove_owner_from_edge(key, indexed_owner); @@ -1488,11 +1488,6 @@ mod tests { self.index.remove_worker(&mut self.worker_blocks, worker_id); } - fn remove_worker_dp_rank(&mut self, worker_id: u64, dp_rank: u32) { - self.index - .remove_worker_dp_rank(&mut self.worker_blocks, worker_id, dp_rank); - } - fn root_workers(&self, local_hash: LocalBlockHash) -> Vec { self.index .root_workers(local_hash, &projection(WorkerWithDpRank::new(7, 0))) @@ -2014,6 +2009,55 @@ mod tests { ); } + #[test] + fn removal_batches_are_idempotent_and_preserve_other_owners() { + let mut index = TestLowerTierIndex::new(); + let removed = WorkerWithDpRank::new(1, 0); + let other_rank = WorkerWithDpRank::new(1, 1); + let other_worker = WorkerWithDpRank::new(2, 0); + for worker in [removed, other_rank, other_worker] { + index + .apply_event(store_event( + worker.worker_id, + worker.dp_rank, + 0, + None, + &[11, 12], + &[101, 102], + )) + .unwrap(); + } + + let batch = remove_event( + 1, + 1, + 0, + vec![ + ExternalSequenceBlockHash(999), + ExternalSequenceBlockHash(101), + ExternalSequenceBlockHash(101), + ExternalSequenceBlockHash(102), + ], + ); + index.apply_event(batch.clone()).unwrap(); + index.apply_event(batch).unwrap(); + index + .apply_event(remove_event(99, 2, 0, vec![ExternalSequenceBlockHash(101)])) + .unwrap(); + + let continuations = [removed, other_rank, other_worker] + .into_iter() + .map(|worker| (worker, LowerTierContinuation::from_root(0))) + .collect::>(); + let hits = index.query_contiguous_hits(&local_hashes(&[11, 12]), &continuations); + assert_eq!(hits.get(&removed), Some(&0)); + assert_eq!(hits.get(&other_rank), Some(&2)); + assert_eq!(hits.get(&other_worker), Some(&2)); + assert!(index.dump_events().iter().all(|event| { + event.worker_id != removed.worker_id || event.event.dp_rank != removed.dp_rank + })); + } + #[test] fn removing_one_owner_preserves_shared_edge_for_other_workers() { let mut index = TestLowerTierIndex::new(); @@ -2095,21 +2139,21 @@ mod tests { } #[test] - fn start_pos_past_end_returns_zero() { + fn exhausted_sequence_returns_zero() { let mut index = TestLowerTierIndex::new(); index .apply_event(store_event(23, 0, 0, Some(1100), &[91], &[901])) .unwrap(); - let query = local_hashes(&[91]); - let mut continuations = FxHashMap::default(); - continuations.insert( - WorkerWithDpRank::new(23, 0), - LowerTierContinuation::new(1, ExternalSequenceBlockHash(1100)), - ); - - let hits = index.query_contiguous_hits(&query, &continuations); - assert_eq!(hits.get(&WorkerWithDpRank::new(23, 0)), Some(&0)); + let worker = WorkerWithDpRank::new(23, 0); + for query in [local_hashes(&[91]), Vec::new()] { + let continuation = + LowerTierContinuation::new(query.len(), ExternalSequenceBlockHash(1100)); + let continuations = FxHashMap::from_iter([(worker, continuation)]); + let details = index.query_match_details(&query, &continuations); + assert_eq!(details.hits.get(&worker), Some(&0)); + assert_eq!(details.next_continuations.get(&worker), Some(&continuation)); + } } #[test] @@ -2201,9 +2245,8 @@ mod tests { .apply_event(store_event(41, 0, 0, Some(3000), &[1], &[301])) .unwrap(); index - .apply_event(store_event(41, 1, 1, Some(4000), &[2], &[401])) + .apply_event(store_event(41, 1, 1, Some(4000), &[1], &[401])) .unwrap(); - index.remove_worker(41); let mut continuations = FxHashMap::default(); continuations.insert( @@ -2215,37 +2258,16 @@ mod tests { LowerTierContinuation::new(0, ExternalSequenceBlockHash(4000)), ); + let before = index.query_contiguous_hits(&local_hashes(&[1]), &continuations); + assert_eq!(before.get(&WorkerWithDpRank::new(41, 0)), Some(&1)); + assert_eq!(before.get(&WorkerWithDpRank::new(41, 1)), Some(&1)); + index.remove_worker(41); + let hits = index.query_contiguous_hits(&local_hashes(&[1]), &continuations); assert_eq!(hits.get(&WorkerWithDpRank::new(41, 0)), Some(&0)); assert_eq!(hits.get(&WorkerWithDpRank::new(41, 1)), Some(&0)); } - #[test] - fn remove_worker_dp_rank_keeps_other_ranks() { - let mut index = TestLowerTierIndex::new(); - index - .apply_event(store_event(43, 0, 0, Some(5000), &[1], &[501])) - .unwrap(); - index - .apply_event(store_event(43, 1, 1, Some(6000), &[2], &[601])) - .unwrap(); - index.remove_worker_dp_rank(43, 0); - - let mut continuations = FxHashMap::default(); - continuations.insert( - WorkerWithDpRank::new(43, 0), - LowerTierContinuation::new(0, ExternalSequenceBlockHash(5000)), - ); - continuations.insert( - WorkerWithDpRank::new(43, 1), - LowerTierContinuation::new(0, ExternalSequenceBlockHash(6000)), - ); - - let hits = index.query_contiguous_hits(&local_hashes(&[2]), &continuations); - assert_eq!(hits.get(&WorkerWithDpRank::new(43, 0)), Some(&0)); - assert_eq!(hits.get(&WorkerWithDpRank::new(43, 1)), Some(&1)); - } - #[test] fn removing_parent_block_keeps_child_continuation_edge() { let mut index = TestLowerTierIndex::new(); @@ -2626,24 +2648,6 @@ mod tests { assert_eq!(details.hits.get(&WorkerWithDpRank::new(103, 0)), Some(&0),); } - /// Empty sequence — every worker should get 0 hits. - #[test] - fn empty_sequence_returns_zero_hits() { - let mut index = TestLowerTierIndex::new(); - index - .apply_event(store_event(111, 0, 0, None, &[1], &[101])) - .unwrap(); - - let mut continuations = FxHashMap::default(); - continuations.insert( - WorkerWithDpRank::new(111, 0), - LowerTierContinuation::from_root(0), - ); - - let details = index.query_match_details(&local_hashes(&[]), &continuations); - assert_eq!(details.hits.get(&WorkerWithDpRank::new(111, 0)), Some(&0)); - } - // --- dump_events tests --- /// Helper: replay dumped events into a fresh indexer and return it. diff --git a/lib/kv-router/src/indexer/radix_tree.rs b/lib/kv-router/src/indexer/radix_tree.rs index 60ee6a3145c7..49d1dc496b3f 100644 --- a/lib/kv-router/src/indexer/radix_tree.rs +++ b/lib/kv-router/src/indexer/radix_tree.rs @@ -3,7 +3,11 @@ //! Single-threaded compressed radix tree for KV cache routing. -use std::{cell::RefCell, collections::VecDeque, rc::Rc}; +use std::{ + cell::RefCell, + collections::VecDeque, + rc::{Rc, Weak}, +}; use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet}; @@ -18,6 +22,7 @@ type WorkerLookup = FxHashMap; pub(crate) struct RadixBlock { state: NodeState, children: FxHashMap, + parent: Weak>, /// Once a node has children it is never eligible for leaf extension again. internal: bool, } @@ -27,21 +32,50 @@ impl RadixBlock { Self { state: NodeState::empty(), children: FxHashMap::default(), + parent: Weak::new(), internal: true, } } - fn for_blocks(blocks: &[KvCacheStoredBlockData], worker: WorkerWithDpRank) -> Self { + fn for_blocks( + blocks: &[KvCacheStoredBlockData], + worker: WorkerWithDpRank, + parent: &SharedRadixBlock, + ) -> Self { Self { state: NodeState::for_blocks(blocks, worker), children: FxHashMap::default(), + parent: Rc::downgrade(parent), internal: false, } } - fn clear_children_if_unreachable(&mut self) { - if self.state.full_edge_workers.is_empty() { - self.children.clear(); + fn prune_unreachable(node: &SharedRadixBlock) { + let (parent, key) = { + let mut node_ref = node.borrow_mut(); + debug_assert!(!node_ref.state.edge.is_empty(), "cannot prune the root"); + if node_ref.state.full_edge_workers.is_empty() { + node_ref.children.clear(); + } + if node_ref.state.has_any_workers() { + return; + } + ( + node_ref.parent.upgrade(), + node_ref.state.edge.first().map(|&(key, _)| key), + ) + }; + if let (Some(parent), Some(key)) = (parent, key) { + let mut parent = parent.borrow_mut(); + // A detached node can outlive its old edge through a worker lookup. + // Do not remove a replacement that now uses the same local hash. + if parent + .children + .get(&key) + .is_some_and(|child| Rc::ptr_eq(child, node)) + { + parent.children.remove(&key); + } } } } @@ -339,7 +373,9 @@ impl RadixTree { break; } - let child = Rc::new(RefCell::new(RadixBlock::for_blocks(remaining, worker))); + let child = Rc::new(RefCell::new(RadixBlock::for_blocks( + remaining, worker, &parent, + ))); { let mut parent_ref = parent.borrow_mut(); parent_ref.internal = true; @@ -397,7 +433,7 @@ impl RadixTree { self.update_lookup_for_blocks(worker, &remaining[..match_len], &child); let tail = &remaining[match_len..]; - let tail_node = Rc::new(RefCell::new(RadixBlock::for_blocks(tail, worker))); + let tail_node = Rc::new(RefCell::new(RadixBlock::for_blocks(tail, worker, &child))); { let mut child_ref = child.borrow_mut(); child_ref.internal = true; @@ -511,8 +547,12 @@ impl RadixTree { full_edge_workers: suffix_full, }, children: std::mem::take(&mut node_ref.children), + parent: Rc::downgrade(node), internal: node_ref.internal, })); + for child in suffix.borrow().children.values() { + child.borrow_mut().parent = Rc::downgrade(&suffix); + } node_ref.children.insert(suffix_first_local, suffix.clone()); node_ref.internal = true; suffix @@ -532,6 +572,8 @@ impl RadixTree { lookup.insert(hash, suffix.clone()); } } + drop(suffix_ref); + RadixBlock::prune_unreachable(&suffix); } fn update_lookup_for_blocks( @@ -638,12 +680,11 @@ impl RadixTree { // worker. Otherwise stale descendants can be reused as store parents, // reactivated by restoring only the removed block, or emitted by dumps // without a valid worker-specific parent. - let outcome = node_ref + node_ref .state - .remove_worker_at_pos(worker, min_pos, min_hash); - node_ref.clear_children_if_unreachable(); - outcome + .remove_worker_at_pos(worker, min_pos, min_hash) }; + RadixBlock::prune_unreachable(&node); for stale_hash in outcome.stale_hashes { lookup.remove(&stale_hash); eagerly_removed.insert(stale_hash); @@ -670,9 +711,8 @@ impl RadixTree { if !seen.insert(Rc::as_ptr(&node)) { continue; } - let mut node_ref = node.borrow_mut(); - node_ref.state.drop_worker(worker); - node_ref.clear_children_if_unreachable(); + node.borrow_mut().state.drop_worker(worker); + RadixBlock::prune_unreachable(&node); } if keep_worker { self.lookup.insert(worker_key, FxHashMap::default()); @@ -694,9 +734,8 @@ impl RadixTree { if !seen.insert(Rc::as_ptr(&node)) { continue; } - let mut node_ref = node.borrow_mut(); - node_ref.state.drop_worker(worker); - node_ref.clear_children_if_unreachable(); + node.borrow_mut().state.drop_worker(worker); + RadixBlock::prune_unreachable(&node); } } @@ -833,6 +872,113 @@ mod tests { create_remove_event, create_store_event, make_store_event, snapshot_events, }; + #[test] + fn cache_churn_releases_evicted_branches() { + let mut tree = RadixTree::new(); + tree.apply_event(create_store_event(1, 0, vec![1, 2, 3], None)) + .unwrap(); + for hash in [4, 5, 6, 7, 8] { + for worker in [1, 2] { + tree.apply_event(create_store_event(worker, hash, vec![1, 2, hash], None)) + .unwrap(); + } + let branch = Rc::downgrade( + &tree.lookup[&WorkerWithDpRank::new(1, 0)][&ExternalSequenceBlockHash(hash * 100)], + ); + if hash == 4 { + // Splitting the parent moves its existing children to the suffix. + tree.apply_event(create_store_event(1, hash, vec![1, 99], None)) + .unwrap(); + } + tree.apply_event(create_remove_event(1, hash + 1, vec![hash])) + .unwrap(); + assert!( + branch.upgrade().is_some(), + "another worker still owns the branch" + ); + match hash % 4 { + 0 => tree + .apply_event(create_remove_event(2, hash + 1, vec![hash])) + .unwrap(), + 1 => tree.remove_worker(2), + 2 => tree.remove_worker_dp_rank(2, 0), + _ => tree.clear_all_blocks(2), + } + assert!( + branch.upgrade().is_none(), + "evicted branch is still retained" + ); + assert_eq!(tree.edge_lengths_for_test(), vec![1, 1, 1, 1]); + assert_eq!(tree.current_size(), if hash % 4 == 0 { 6 } else { 4 }); + } + } + + #[test] + fn removing_detached_node_preserves_replacement() { + let mut tree = RadixTree::new(); + for blocks in [vec![1, 2, 3, 4], vec![1, 2, 3, 9]] { + tree.apply_event(create_store_event(2, 0, blocks, None)) + .unwrap(); + } + tree.apply_event(create_store_event(1, 0, vec![1, 2, 3], None)) + .unwrap(); + let old_branch = Rc::downgrade( + &tree.lookup[&WorkerWithDpRank::new(2, 0)][&ExternalSequenceBlockHash(400)], + ); + for worker in [1, 2] { + tree.apply_event(create_remove_event(worker, 1, vec![3])) + .unwrap(); + } + // The existing descendant-invalidation gap leaves worker 2's old node + // in its lookup after the parent has disconnected it. + assert!(old_branch.upgrade().is_some()); + tree.apply_event(create_store_event( + 1, + 2, + vec![3], + Some(ExternalSequenceBlockHash(200)), + )) + .unwrap(); + tree.apply_event(create_store_event( + 1, + 3, + vec![4], + Some(ExternalSequenceBlockHash(300)), + )) + .unwrap(); + tree.apply_event(create_remove_event(2, 3, vec![4])) + .unwrap(); + + assert!(old_branch.upgrade().is_none()); + let scores = tree.find_matches((1..=4).map(LocalBlockHash).collect(), false); + assert_eq!(scores.scores.get(&WorkerWithDpRank::new(1, 0)), Some(&4)); + } + + #[test] + fn split_releases_unowned_suffix() { + let mut tree = RadixTree::new(); + for worker in [1, 2] { + tree.apply_event(create_store_event(worker, 0, vec![1, 2, 3], None)) + .unwrap(); + } + tree.apply_event(create_remove_event(1, 1, vec![2])) + .unwrap(); + tree.apply_event(create_remove_event(2, 1, vec![1])) + .unwrap(); + assert_eq!(tree.edge_lengths_for_test(), vec![3]); + + tree.apply_event(create_store_event( + 1, + 2, + vec![9], + Some(ExternalSequenceBlockHash(100)), + )) + .unwrap(); + + assert_eq!(tree.edge_lengths_for_test(), vec![1, 1]); + assert_eq!(tree.current_size(), 2); + } + #[test] fn rejects_self_referencing_store() { let mut tree = RadixTree::new(); diff --git a/lib/kv-router/src/lookup_update.rs b/lib/kv-router/src/lookup_update.rs index 593e5146996c..c345403e2f5d 100644 --- a/lib/kv-router/src/lookup_update.rs +++ b/lib/kv-router/src/lookup_update.rs @@ -9,10 +9,9 @@ use rustc_hash::FxHashMap; /// Update an `Arc` lookup map for keys that should point at `node`. /// -/// Duplicate store/repair paths often revisit keys that already resolve to the -/// same node. Skipping those replacements avoids unnecessary hash table writes -/// and `Arc` refcount churn while still repairing stale entries. Returns the -/// number of entries that were inserted or changed. +/// Duplicate store paths often revisit keys that already resolve to the same +/// node. Skipping those replacements avoids unnecessary hash table writes and +/// `Arc` refcount churn. Returns the number of entries inserted or changed. pub(crate) fn update_arc_lookup_for_keys( lookup: &mut FxHashMap>, keys: impl IntoIterator, @@ -39,3 +38,31 @@ where changed } + +/// Update existing `Arc` lookup entries for keys that should point at `node`. +/// +/// Unlike [`update_arc_lookup_for_keys`], this does not insert missing keys. +/// Lookup repair uses absence as meaningful state: a remove can scrub an entry +/// before another worker on the same event thread repairs a stale lookup. +/// Returns the number of existing entries that changed. +pub(crate) fn update_existing_arc_lookup_for_keys( + lookup: &mut FxHashMap>, + keys: impl IntoIterator, + node: &Arc, +) -> usize +where + K: Eq + Hash, +{ + let mut changed = 0; + + for key in keys { + if let Some(entry) = lookup.get_mut(&key) + && !Arc::ptr_eq(entry, node) + { + *entry = Arc::clone(node); + changed += 1; + } + } + + changed +} diff --git a/lib/kv-router/src/services/common/zmq.rs b/lib/kv-router/src/services/common/zmq.rs index 406be80b5103..8b0a4acc27e0 100644 --- a/lib/kv-router/src/services/common/zmq.rs +++ b/lib/kv-router/src/services/common/zmq.rs @@ -18,6 +18,12 @@ pub(crate) type MultipartMessage = Vec>; pub(crate) type SharedSocket = Arc>; const ZMQ_RCVTIMEOUT_MS: i32 = 100; +// Unbounded on purpose: libzmq 4.3.4 asserts (`_input_stopped`) when a +// heartbeating peer restarts a pipe that HWM-stopped input, so an indexer +// listener stalled in gap recovery must never let its queue fill +// (zeromq/libzmq#3596). +#[cfg(feature = "standalone-indexer")] +const ZMQ_INDEXER_RCVHWM: i32 = 0; const ZMQ_SNDTIMEOUT_MS: i32 = 0; const ZMQ_RECONNECT_IVL_MS: i32 = 100; const ZMQ_RECONNECT_IVL_MAX_MS: i32 = 5000; @@ -217,6 +223,8 @@ pub(crate) fn create_sub_socket(topic: &[u8]) -> Result { #[cfg(feature = "standalone-indexer")] pub(crate) fn connect_sub_socket(endpoint: &str) -> Result { let socket = create_sub_socket(b"")?; + // Must precede connect: HWM applies to pipes created afterwards. + socket.socket().set_rcvhwm(ZMQ_INDEXER_RCVHWM)?; socket.connect(endpoint)?; Ok(Arc::new(Mutex::new(socket))) } @@ -298,3 +306,15 @@ pub(crate) fn validate_endpoint(endpoint: &str) -> Result<()> { )), } } + +#[cfg(all(test, feature = "standalone-indexer"))] +mod tests { + use super::*; + + #[tokio::test] + async fn indexer_sub_socket_receive_queue_is_unbounded() { + let socket = connect_sub_socket("inproc://rcvhwm-test").unwrap(); + let guard = socket.lock().await; + assert_eq!(guard.socket().get_rcvhwm().unwrap(), 0); + } +} diff --git a/lib/kv-router/src/services/indexer/audit.rs b/lib/kv-router/src/services/indexer/audit.rs new file mode 100644 index 000000000000..e342470f9ab4 --- /dev/null +++ b/lib/kv-router/src/services/indexer/audit.rs @@ -0,0 +1,106 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! `--enable-logging`: one `kv_audit` line per query and per ingested engine +//! event, for offline analysis (filter with `RUST_LOG=kv_audit=info`). +//! +//! Event lines are written before the `--keep-evictions` filter, so they show +//! what the engine published. `source` is `live` (ZMQ SUB), `replay` (ZMQ +//! replay socket) or `recover` (`/kv_recover`). + +use crate::protocols::{KvCacheEventData, LocalBlockHash, RouterEvent}; + +fn now_unix_ms() -> u64 { + use std::time::{SystemTime, UNIX_EPOCH}; + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0) +} + +pub(super) fn log_event(event: &RouterEvent, seq: u64, source: &'static str) { + let ts_ms = now_unix_ms(); + let worker_id = event.worker_id; + let dp_rank = event.event.dp_rank; + let storage_tier = event.storage_tier; + let event_id = event.event.event_id; + match &event.event.data { + KvCacheEventData::Stored(data) => { + // tokens_hash is the local block hash (what /query computes); + // block_hash is the engine's chained sequence hash. + let token_block_hashes: Vec = + data.blocks.iter().map(|b| b.tokens_hash.0).collect(); + let sequence_block_hashes: Vec = + data.blocks.iter().map(|b| b.block_hash.0).collect(); + tracing::info!( + target: "kv_audit", + kind = "STORE", + ts_ms, + source, + seq, + worker_id, + dp_rank, + storage_tier = ?storage_tier, + event_id, + parent_hash = ?data.parent_hash.map(|h| h.0), + num_blocks = token_block_hashes.len(), + token_block_hashes = ?token_block_hashes, + sequence_block_hashes = ?sequence_block_hashes, + "kv_audit STORE" + ); + } + KvCacheEventData::Removed(data) => { + let sequence_block_hashes: Vec = data.block_hashes.iter().map(|h| h.0).collect(); + tracing::info!( + target: "kv_audit", + kind = "EVICT", + ts_ms, + source, + seq, + worker_id, + dp_rank, + storage_tier = ?storage_tier, + event_id, + num_blocks = sequence_block_hashes.len(), + sequence_block_hashes = ?sequence_block_hashes, + "kv_audit EVICT" + ); + } + KvCacheEventData::Cleared => { + tracing::info!( + target: "kv_audit", + kind = "CLEAR", + ts_ms, + source, + seq, + worker_id, + dp_rank, + storage_tier = ?storage_tier, + event_id, + "kv_audit CLEAR" + ); + } + } +} + +pub(super) fn log_query( + model_name: &str, + queried_groups: &[String], + status: u16, + block_hashes: &[LocalBlockHash], + response: &serde_json::Value, +) { + let block_hashes: Vec = block_hashes.iter().map(|h| h.0).collect(); + tracing::info!( + target: "kv_audit", + kind = "QUERY", + ts_ms = now_unix_ms(), + model_name, + fanout_groups = ?queried_groups, + status, + num_blocks = block_hashes.len(), + block_hashes = ?block_hashes, + response = %response, + "kv_audit QUERY" + ); +} diff --git a/lib/kv-router/src/services/indexer/backend.rs b/lib/kv-router/src/services/indexer/backend.rs index 65cb770596ee..0ffbd1886209 100644 --- a/lib/kv-router/src/services/indexer/backend.rs +++ b/lib/kv-router/src/services/indexer/backend.rs @@ -31,6 +31,9 @@ pub enum Indexer { primary: Arc>, lower_tier: LowerTierIndexers, }, + /// The flat counterfactual indexer (`--h24`): no workers, no evictions, + /// no tiers. See [`super::h24`]. + H24(Arc), } impl Indexer { @@ -41,6 +44,7 @@ impl Indexer { match self { Indexer::Single { primary, .. } => primary.apply_event(event).await, Indexer::Concurrent { primary, .. } => primary.apply_event(event).await, + Indexer::H24(h24) => h24_apply(h24, &event), } } @@ -49,6 +53,12 @@ impl Indexer { /// indexer otherwise. `Cleared` events fan out to their applicable physical /// indexes according to the event's reset scope. pub async fn apply_event_routed(&self, event: RouterEvent) -> Result<(), KvRouterError> { + // Tier-agnostic on purpose: in the counterfactual a store on ANY tier + // means the content existed. + if let Indexer::H24(h24) = self { + h24_apply(h24, &event); + return Ok(()); + } let targets_primary = match event.targets_primary() { Ok(targets_primary) => targets_primary, Err(_) => { @@ -56,6 +66,7 @@ impl Indexer { Self::Single { lower_tier, .. } | Self::Concurrent { lower_tier, .. } => { lower_tier.record_unsupported_residency_event(&event); } + Self::H24(_) => unreachable!("h24 events return above"), } return Ok(()); } @@ -122,6 +133,7 @@ impl Indexer { .await; } } + Indexer::H24(_) => unreachable!("h24 events return above"), } Ok(()) } @@ -146,6 +158,8 @@ impl Indexer { } primary.remove_worker(worker_id).await; } + // The counterfactual keeps content regardless of worker lifecycle. + Indexer::H24(_) => {} } } @@ -169,6 +183,7 @@ impl Indexer { } primary.remove_worker_dp_rank(worker_id, dp_rank).await; } + Indexer::H24(_) => {} } } @@ -182,6 +197,7 @@ impl Indexer { Indexer::Concurrent { primary, .. } => { primary.find_matches(hashes).await.map_err(Into::into) } + Indexer::H24(h24) => Ok(h24.find_matches(&hashes)), } } @@ -215,6 +231,13 @@ impl Indexer { lower_tier: lt, }) } + Indexer::H24(h24) => Ok(TieredMatchDetails { + device: MatchDetails { + overlap_scores: h24.find_matches(&sequence), + ..Default::default() + }, + lower_tier: Default::default(), + }), } } @@ -242,6 +265,12 @@ impl Indexer { primary.dump_events().await.map_err(anyhow::Error::from)?, lower_tier.entries(), ), + // Flat maps cannot reconstruct chains; /dump and peer recovery + // are unsupported in h24 mode. + Indexer::H24(_) => { + tracing::warn!("h24 indexer does not support dump_events; returning empty"); + return Ok(Vec::new()); + } }; let mut out = primary_events; @@ -530,7 +559,9 @@ mod tests { lower_tier.get_or_create(StorageTier::HostPinned), lower_tier.get_or_create(StorageTier::Disk), ), - Indexer::Concurrent { .. } => unreachable!("test creates the single indexer"), + Indexer::Concurrent { .. } | Indexer::H24(_) => { + unreachable!("test creates the single indexer") + } }; indexer @@ -598,6 +629,19 @@ mod tests { } } +/// Stores chain into the flat h24 maps; evictions and clears are the +/// counterfactual's no-op and are only counted. +fn h24_apply(h24: &super::h24::H24Indexer, event: &RouterEvent) { + match &event.event.data { + KvCacheEventData::Stored(data) => h24.apply_store(data), + KvCacheEventData::Removed(_) | KvCacheEventData::Cleared => h24.note_ignored_event(), + } +} + +pub fn create_h24_indexer() -> Indexer { + Indexer::H24(Arc::new(super::h24::H24Indexer::new())) +} + pub fn create_indexer(block_size: u32, num_threads: usize) -> Indexer { create_indexer_with_metrics( block_size, diff --git a/lib/kv-router/src/services/indexer/block_size.rs b/lib/kv-router/src/services/indexer/block_size.rs new file mode 100644 index 000000000000..784116653209 --- /dev/null +++ b/lib/kv-router/src/services/indexer/block_size.rs @@ -0,0 +1,171 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Block-size check on engine `BlockStored` events before they are indexed. +//! +//! The shared ZMQ normalizer silently drops every block whose size differs +//! from the configured one, so an indexer started with the wrong +//! `--block-size` stays empty while its event stream looks healthy. The +//! standalone indexer treats that as a fatal misconfiguration instead, except +//! for vLLM's partial-prefix entries (see [`StoredBlockSize::PartialPrefix`]). + +use crate::protocols::WorkerWithDpRank; +use crate::zmq_wire::RawKvEvent; + +/// Per-listener gate applied to every preprocessed engine event. +pub(super) struct BlockSizeGuard { + configured_block_size: u32, + partial_prefix_skips: u64, +} + +impl BlockSizeGuard { + pub(super) fn new(configured_block_size: u32) -> Self { + Self { + configured_block_size, + partial_prefix_skips: 0, + } + } + + /// Whether `raw` should be indexed. Partial-prefix stores are dropped + /// (first few logged). A block-size mismatch **exits the process**: every + /// later store would be dropped the same way, and a crash loop is the + /// only signal k8s surfaces for an indexer that looks alive but is empty. + pub(super) fn admit(&mut self, raw: &RawKvEvent, worker: WorkerWithDpRank) -> bool { + match classify_stored_block_size(raw, self.configured_block_size) { + StoredBlockSize::Indexable => true, + StoredBlockSize::PartialPrefix { event_block_size } => { + self.partial_prefix_skips += 1; + if self.partial_prefix_skips <= 3 { + tracing::warn!( + worker_id = worker.worker_id, + dp_rank = worker.dp_rank, + event_block_size, + configured_block_size = self.configured_block_size, + "Skipping sub-block BlockStored: a partial-prefix entry \ + (prefix_match_unit < block_size), unless the indexer's \ + --block-size is larger than the engine's -- then the \ + index stays empty; check configured vs event size" + ); + } + false + } + StoredBlockSize::Mismatch { event_block_size } => { + tracing::error!( + worker_id = worker.worker_id, + dp_rank = worker.dp_rank, + event_block_size, + configured_block_size = self.configured_block_size, + "Block size mismatch: every stored block would be dropped and the \ + index would stay empty. Fix the indexer's --block-size to match \ + the engine's --block-size. Exiting." + ); + std::process::exit(1); + } + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum StoredBlockSize { + /// Not a store, an empty store, or a store of whole configured blocks. + Indexable, + /// A store whose `block_size` is smaller than the configured cache block: + /// vLLM hashes prefixes every `prefix_match_unit` (hash_block_size) + /// tokens, which can be finer than the cache block, and publishes the + /// prompt tail inside a cache block as a sub-block store (32/64/96 tokens + /// for a 128-token block). It carries nothing the index keys on. + PartialPrefix { event_block_size: usize }, + /// A store of blocks larger than configured: the indexer's `--block-size` + /// does not match the engine's, and every stored block would be dropped. + Mismatch { event_block_size: usize }, +} + +pub(super) fn classify_stored_block_size( + raw: &RawKvEvent, + configured_block_size: u32, +) -> StoredBlockSize { + let RawKvEvent::BlockStored { + block_hashes, + block_size, + .. + } = raw + else { + return StoredBlockSize::Indexable; + }; + let event_block_size = *block_size; + let configured = configured_block_size as usize; + if block_hashes.is_empty() || event_block_size == configured { + StoredBlockSize::Indexable + } else if event_block_size > 0 && event_block_size < configured { + StoredBlockSize::PartialPrefix { event_block_size } + } else { + StoredBlockSize::Mismatch { event_block_size } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::zmq_wire::BlockHashValue; + + fn stored(block_size: usize, num_blocks: u64) -> RawKvEvent { + RawKvEvent::BlockStored { + block_hashes: (0..num_blocks).map(BlockHashValue::Unsigned).collect(), + parent_block_hash: None, + token_ids: vec![0; block_size * num_blocks as usize], + block_size, + medium: None, + lora_name: None, + cache_namespace: None, + block_mm_infos: None, + is_eagle: None, + group_idx: None, + kv_cache_spec_kind: None, + kv_cache_spec_sliding_window: None, + locality: None, + ownership: None, + } + } + + #[test] + fn whole_blocks_are_indexable() { + assert_eq!( + classify_stored_block_size(&stored(128, 2), 128), + StoredBlockSize::Indexable + ); + } + + #[test] + fn sub_block_store_is_a_partial_prefix_entry() { + for size in [32, 64, 96] { + assert_eq!( + classify_stored_block_size(&stored(size, 1), 128), + StoredBlockSize::PartialPrefix { + event_block_size: size + } + ); + } + } + + #[test] + fn larger_engine_block_is_a_mismatch() { + assert_eq!( + classify_stored_block_size(&stored(256, 1), 128), + StoredBlockSize::Mismatch { + event_block_size: 256 + } + ); + } + + #[test] + fn empty_store_and_non_store_events_are_indexable() { + assert_eq!( + classify_stored_block_size(&stored(256, 0), 128), + StoredBlockSize::Indexable + ); + assert_eq!( + classify_stored_block_size(&RawKvEvent::Ignored, 128), + StoredBlockSize::Indexable + ); + } +} diff --git a/lib/kv-router/src/services/indexer/deepapi_contract_tests.rs b/lib/kv-router/src/services/indexer/deepapi_contract_tests.rs new file mode 100644 index 000000000000..10198bc2041f --- /dev/null +++ b/lib/kv-router/src/services/indexer/deepapi_contract_tests.rs @@ -0,0 +1,246 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The HTTP contract DeepInfra's deepapi relies on (`query_kv_indexer`, +//! `kv_shards_from_indexer_response` in deepinfra/backend): it POSTs +//! `{"model_name", "tenant_id": "default", "block_hashes": [i64]}` to +//! `/query_by_hash` and reads `instances[*].{pod_name, longest_matched, gpu}`. + +use std::sync::Arc; + +use axum::body::Body; +use axum::http::{Request, StatusCode, header}; +use tower::ServiceExt; + +use crate::protocols::StorageTier; + +use super::backend::Indexer; +use super::backend::test_util::store_event; +use super::registry::{ListenerExtras, WorkerRegistry}; +use super::server::{AppState, create_router}; + +const BLOCK_SIZE: u32 = 4; + +async fn register_pod( + registry: &WorkerRegistry, + instance_id: u64, + port: u16, + model_name: &str, + routing_group: &str, + pod_name: &str, +) { + registry + .register_with_extras( + instance_id, + format!("tcp://127.0.0.1:{port}"), + 0, + model_name.to_string(), + routing_group.to_string(), + BLOCK_SIZE, + None, + ListenerExtras { + recover_endpoint: None, + pod_name: Some(pod_name.to_string()), + }, + ) + .await + .unwrap(); +} + +async fn store_blocks(indexer: &Indexer, worker_id: u64, blocks: &[u64]) { + indexer + .apply_event_routed(store_event( + worker_id, + 0, + 1, + &[], + blocks, + StorageTier::Device, + )) + .await + .unwrap(); + indexer.dump_events().await.expect("flush indexer"); +} + +fn router(registry: Arc) -> axum::Router { + create_router(Arc::new(AppState { + registry, + access_log_sink: None, + #[cfg(feature = "metrics")] + prom_registry: prometheus::Registry::new(), + })) +} + +async fn post_json( + app: &axum::Router, + uri: &str, + body: serde_json::Value, +) -> (StatusCode, serde_json::Value) { + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri(uri) + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + (status, serde_json::from_slice(&bytes).unwrap()) +} + +#[tokio::test] +async fn query_by_hash_names_each_matched_instance_by_pod() { + let registry = Arc::new(WorkerRegistry::new(1)); + registry.signal_ready(); + register_pod(®istry, 7, 15700, "m", "default", "vllm-m-abcd").await; + let indexer = registry + .get_indexer(&crate::identity::RoutingPartitionId::new("m", "default")) + .unwrap() + .indexer + .clone(); + store_blocks(&indexer, 7, &[11, 12]).await; + + let app = router(registry); + let (status, body) = post_json( + &app, + "/query_by_hash", + serde_json::json!({"model_name": "m", "tenant_id": "default", "block_hashes": [11, 12]}), + ) + .await; + + assert_eq!(status, StatusCode::OK, "{body}"); + let instance = &body["instances"]["7"]; + assert_eq!(instance["pod_name"], "vllm-m-abcd"); + assert_eq!(instance["longest_matched"], 2 * BLOCK_SIZE); + assert_eq!(instance["gpu"], 2 * BLOCK_SIZE); + + let workers = app + .oneshot( + Request::builder() + .uri("/workers") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let bytes = axum::body::to_bytes(workers.into_body(), usize::MAX) + .await + .unwrap(); + let workers: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(workers[0]["pod_name"], "vllm-m-abcd"); +} + +fn indexer_of(registry: &WorkerRegistry, model_name: &str, routing_group: &str) -> Indexer { + registry + .get_indexer(&crate::identity::RoutingPartitionId::new( + model_name, + routing_group, + )) + .unwrap() + .indexer + .clone() +} + +/// One model under two routing groups (two engine_hash regimes) plus an +/// unrelated model: deepapi's request names only the model (and +/// `tenant_id: "default"`, which matches no tree), and must be answered from +/// both groups of that model and nothing else. +#[tokio::test] +async fn query_by_hash_fans_out_across_routing_groups_of_the_model() { + let registry = Arc::new(WorkerRegistry::new(1)); + registry.signal_ready(); + register_pod(®istry, 7, 15710, "m", "hash-a", "pod-a").await; + register_pod(®istry, 8, 15711, "m", "hash-b", "pod-b").await; + register_pod(®istry, 9, 15712, "other", "hash-a", "pod-c").await; + store_blocks(&indexer_of(®istry, "m", "hash-a"), 7, &[11, 12]).await; + store_blocks(&indexer_of(®istry, "m", "hash-b"), 8, &[11]).await; + store_blocks(&indexer_of(®istry, "other", "hash-a"), 9, &[11, 12]).await; + + let app = router(registry); + let (status, body) = post_json( + &app, + "/query_by_hash", + serde_json::json!({"model_name": "m", "tenant_id": "default", "block_hashes": [11, 12]}), + ) + .await; + + assert_eq!(status, StatusCode::OK, "{body}"); + let instances = body["instances"].as_object().unwrap(); + assert_eq!(instances.len(), 2, "{body}"); + assert_eq!(instances["7"]["pod_name"], "pod-a"); + assert_eq!(instances["7"]["longest_matched"], 2 * BLOCK_SIZE); + assert_eq!(instances["8"]["pod_name"], "pod-b"); + assert_eq!(instances["8"]["longest_matched"], BLOCK_SIZE); + assert_eq!(body["scores"]["7"]["0"], 2 * BLOCK_SIZE); + + let (status, _) = post_json( + &app, + "/query_by_hash", + serde_json::json!({"model_name": "nonexistent", "block_hashes": [11]}), + ) + .await; + assert_eq!(status, StatusCode::NOT_FOUND); +} + +/// deepapi hashes prompts itself (`deepinfra/utils.kv_block_hashes`: XXH3-64 +/// of little-endian u32 tokens, seed 1337, lora seed 1337 + xxh3(name)). These +/// vectors were produced by that function; `/query` must hash identically or +/// every deepapi probe misses. +#[test] +fn token_hashing_matches_deepapi_golden_vectors() { + use crate::protocols::{BlockHashOptions, compute_block_hash_for_seq}; + let tokens: Vec = (1..=8).collect(); + let plain: Vec = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()) + .into_iter() + .map(|h| h.0) + .collect(); + assert_eq!(plain, [14643705804678351452, 16777012769546811212]); + + let lora: Vec = compute_block_hash_for_seq( + &tokens, + 4, + BlockHashOptions { + lora_name: Some("lora_7"), + ..Default::default() + }, + ) + .into_iter() + .map(|h| h.0) + .collect(); + assert_eq!(lora, [6090315765804930700, 3193399369686574168]); +} + +#[tokio::test] +async fn query_by_tokens_fans_out_across_routing_groups() { + use crate::protocols::{BlockHashOptions, compute_block_hash_for_seq}; + let registry = Arc::new(WorkerRegistry::new(1)); + registry.signal_ready(); + register_pod(®istry, 7, 15713, "m", "hash-a", "pod-a").await; + register_pod(®istry, 8, 15714, "m", "hash-b", "pod-b").await; + let token_ids: Vec = (0..8).collect(); + let hashes: Vec = + compute_block_hash_for_seq(&token_ids, BLOCK_SIZE, BlockHashOptions::default()) + .into_iter() + .map(|h| h.0) + .collect(); + store_blocks(&indexer_of(®istry, "m", "hash-a"), 7, &hashes).await; + store_blocks(&indexer_of(®istry, "m", "hash-b"), 8, &hashes[..1]).await; + + let (status, body) = post_json( + &router(registry), + "/query", + serde_json::json!({"model_name": "m", "token_ids": token_ids}), + ) + .await; + + assert_eq!(status, StatusCode::OK, "{body}"); + assert_eq!(body["instances"]["7"]["longest_matched"], 2 * BLOCK_SIZE); + assert_eq!(body["instances"]["8"]["longest_matched"], BLOCK_SIZE); +} diff --git a/lib/kv-router/src/services/indexer/discovery.rs b/lib/kv-router/src/services/indexer/discovery.rs new file mode 100644 index 000000000000..b34713dc7370 --- /dev/null +++ b/lib/kv-router/src/services/indexer/discovery.rs @@ -0,0 +1,88 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Configuration for Kubernetes pod auto-discovery. +//! +//! Plain data, always compiled, so the CLI can build it without the +//! `kube-discovery` feature; only the feature-gated `pod_watcher` consumes it. + +/// Label the backend stamps on every engine pod with the sanitized model name. +/// Stable across GPU configs and engine versions, unlike `engine_hash`, which +/// fragments per GPU config (B200 vs B300 replicas of one model differ). +pub const MODEL_NAME_LABEL: &str = "di/model_name"; + +/// The backend's `di/model_name` label value for `name`: lowercase, then every +/// char not in `[a-z0-9-]` becomes `-`. Must match the backend's +/// `re.sub('[^a-z0-9-]', '-', name.lower())` exactly, or the selector matches +/// zero pods. +pub fn sanitize_model_name_label(name: &str) -> String { + name.to_lowercase() + .chars() + .map(|c| { + if c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-' { + c + } else { + '-' + } + }) + .collect() +} + +/// Selector matching every engine pod of `model_name`, across engine_hash +/// variants. +pub fn model_name_label_selector(model_name: &str) -> String { + format!( + "{MODEL_NAME_LABEL}={}", + sanitize_model_name_label(model_name) + ) +} + +#[derive(Debug, Clone)] +pub struct KubeDiscoveryConfig { + /// Namespace to watch for engine pods. + pub namespace: String, + /// Selector picking out this model's engine pods: derived from the model + /// name, or supplied raw (e.g. `engine_hash=d4b7a85131172ca6`). + pub label_selector: String, + /// KV-event ZMQ port; data-parallel rank `r` publishes on `zmq_port + r`. + pub zmq_port: u16, + /// Port serving `GET /kv_recover`; rank `r` serves on `recover_port + r`. + pub recover_port: Option, + /// Data-parallel ranks per engine pod, each registered as a listener of + /// the pod's instance. + pub dp_size: u32, + /// Model name discovered pods are registered under. + pub model_name: String, + /// Routing group for pods without an `engine_hash` label. + pub routing_group: String, + /// KV cache block size of the discovered engines. + pub block_size: u32, +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The sanitization contract with the backend, verbatim. + #[test] + fn model_name_label_selector_matches_backend_contract() { + assert_eq!( + model_name_label_selector("openai/gpt-oss-120b"), + "di/model_name=openai-gpt-oss-120b" + ); + assert_eq!( + model_name_label_selector("meta-llama/Llama-3.3-70B"), + "di/model_name=meta-llama-llama-3-3-70b" + ); + assert_eq!( + model_name_label_selector("Qwen/Qwen2.5-72B-Instruct"), + "di/model_name=qwen-qwen2-5-72b-instruct" + ); + } + + #[test] + fn sanitize_replaces_every_disallowed_char() { + assert_eq!(sanitize_model_name_label("A_b.c/Dé9-"), "a-b-c-d-9-"); + assert_eq!(sanitize_model_name_label("gpt-oss-20b"), "gpt-oss-20b"); + } +} diff --git a/lib/kv-router/src/services/indexer/evictions.rs b/lib/kv-router/src/services/indexer/evictions.rs new file mode 100644 index 000000000000..940a2b998259 --- /dev/null +++ b/lib/kv-router/src/services/indexer/evictions.rs @@ -0,0 +1,482 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! `--keep-evictions` support: park engine eviction events in a per-listener +//! buffer instead of applying them, then replay the aged ones into the tree +//! through the normal `Removed` apply path once the pod's memory crosses a +//! threshold. The radix tree itself is untouched — from its point of view +//! evictions simply arrive late. +//! +//! Correctness hinges on one rule: a buffered eviction is only replayed if it +//! is still the block's *latest* eviction. The map holds the latest eviction +//! time per hash; the FIFO queue entries each carry the time they were added, +//! and an entry whose time no longer matches the map is dead (the block was +//! re-stored, or re-evicted later — the newer queue entry handles it). + +use std::collections::VecDeque; +use std::sync::{Arc, LazyLock}; +use std::time::{Duration, Instant}; + +use rustc_hash::FxHashMap; +use tokio_util::sync::CancellationToken; + +use crate::protocols::{ + ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheRemoveData, RouterEvent, + StorageTier, +}; + +use super::registry::WorkerRegistry; + +/// How often the cleanup loop re-checks memory pressure. +const SWEEP_INTERVAL: Duration = Duration::from_secs(15); + +/// Max hashes per synthesized `Removed` event, to bound message size on the +/// indexer's internal apply channels. +const REMOVE_CHUNK: usize = 4096; + +/// Monotonic seconds since process start. Immune to wall-clock jumps, which +/// matters because entry age decides what gets removed from the tree. +pub(super) fn monotonic_secs() -> u64 { + static START: LazyLock = LazyLock::new(Instant::now); + START.elapsed().as_secs() +} + +/// `--keep-evictions` settings. +#[derive(Debug, Clone, Copy)] +pub struct KeepEvictionsConfig { + /// Minimum age a parked eviction must reach before the sweep replays it. + pub retention_s: u64, + /// Fraction of the memory limit, in (0, 1], above which the sweep runs. + pub memory_threshold: f64, +} + +impl KeepEvictionsConfig { + pub fn validate(&self) -> anyhow::Result<()> { + anyhow::ensure!( + self.memory_threshold > 0.0 && self.memory_threshold <= 1.0, + "--evict-memory-threshold must be in (0, 1], got {}", + self.memory_threshold + ); + Ok(()) + } +} + +/// Buffer of eviction events not yet applied to the tree, owned by one +/// listener (one `(worker_id, dp_rank)` event stream). +#[derive(Default)] +pub struct PendingEvictions { + /// Latest eviction per sequence hash: when it happened and on which tier. + pending: FxHashMap, + /// Arrival-ordered `(hash, at_secs, tier)`. Only pushed at the back and + /// popped at the front, so timestamps are non-decreasing front to back. + queue: VecDeque<(u64, u64, StorageTier)>, +} + +impl PendingEvictions { + /// Record an eviction batch instead of applying it. + pub fn buffer(&mut self, hashes: &[ExternalSequenceBlockHash], tier: StorageTier) { + self.buffer_at(hashes, tier, monotonic_secs()); + } + + /// Timestamp-explicit form of [`Self::buffer`]. `at` must be + /// non-decreasing across calls (queue order is arrival order). + fn buffer_at(&mut self, hashes: &[ExternalSequenceBlockHash], tier: StorageTier, at: u64) { + for hash in hashes { + self.pending.insert(hash.0, (at, tier)); + self.queue.push_back((hash.0, at, tier)); + } + } + + /// Cancel pending evictions for blocks the engine just re-stored. Their + /// stale queue entries stay behind and are skipped at drain time. + pub fn cancel(&mut self, hashes: impl IntoIterator) { + for hash in hashes { + self.pending.remove(&hash); + } + } + + /// Drop all state (used when a recovery snapshot replaces the worker's + /// tree state wholesale). + pub fn clear(&mut self) { + self.pending.clear(); + self.queue.clear(); + } + + /// Number of evictions currently pending replay. + pub fn len(&self) -> usize { + self.pending.len() + } + + /// Raw queue length, including stale (cancelled/superseded) entries that + /// still occupy memory until a drain pops them. `queue_len - len` is the + /// dead weight from re-stored blocks. + pub fn queue_len(&self) -> usize { + self.queue.len() + } + + pub fn is_empty(&self) -> bool { + self.pending.is_empty() + } + + /// Pop every queue entry with `at <= cutoff_secs` and return the hashes + /// whose latest eviction that entry is, grouped by storage tier. Entries + /// that were cancelled by a re-store, or superseded by a later eviction + /// of the same hash, are skipped — not stopped at. + pub fn drain_older_than( + &mut self, + cutoff_secs: u64, + ) -> FxHashMap> { + let mut doomed: FxHashMap> = + FxHashMap::default(); + while let Some(&(hash, at, _)) = self.queue.front() { + if at > cutoff_secs { + break; + } + self.queue.pop_front(); + // Only replay if this entry is still the hash's latest eviction. + // The tier stored in `pending` is authoritative (same write as + // the timestamp), so read it from there. + if let Some(&(pending_at, tier)) = self.pending.get(&hash) + && pending_at == at + { + self.pending.remove(&hash); + doomed + .entry(tier) + .or_default() + .push(ExternalSequenceBlockHash(hash)); + } + } + doomed + } +} + +/// Fraction of the memory limit currently in use, best-effort: +/// cgroup v2, then cgroup v1, then RSS against total host memory. +/// `None` if no source could be read (never expected on Linux). +pub fn memory_usage_fraction() -> Option { + fn read_u64(path: &str) -> Option { + std::fs::read_to_string(path).ok()?.trim().parse().ok() + } + // "max" in cgroup v2 (or the v1 sentinel ~2^63) means "no limit set"; + // fall through to the host total in that case. + fn read_limit(path: &str) -> Option { + let raw = std::fs::read_to_string(path).ok()?; + let raw = raw.trim(); + if raw == "max" { + return None; + } + let value: u64 = raw.parse().ok()?; + (value < (1 << 62)).then_some(value) + } + + let usage = read_u64("/sys/fs/cgroup/memory.current") + .or_else(|| read_u64("/sys/fs/cgroup/memory/memory.usage_in_bytes")) + .or_else(vm_rss_bytes)?; + let limit = read_limit("/sys/fs/cgroup/memory.max") + .or_else(|| read_limit("/sys/fs/cgroup/memory/memory.limit_in_bytes")) + .or_else(host_total_bytes)?; + (limit > 0).then(|| usage as f64 / limit as f64) +} + +/// Resident set size from `/proc/self/status` (`VmRSS`, reported in kB). +fn vm_rss_bytes() -> Option { + let status = std::fs::read_to_string("/proc/self/status").ok()?; + parse_kb_line(&status, "VmRSS:") +} + +/// Total host memory from `/proc/meminfo` (`MemTotal`, reported in kB). +fn host_total_bytes() -> Option { + let meminfo = std::fs::read_to_string("/proc/meminfo").ok()?; + parse_kb_line(&meminfo, "MemTotal:") +} + +fn parse_kb_line(text: &str, key: &str) -> Option { + let line = text.lines().find(|l| l.starts_with(key))?; + let kb: u64 = line.split_whitespace().nth(1)?.parse().ok()?; + Some(kb * 1024) +} + +/// Spawn the background sweep: every [`SWEEP_INTERVAL`], if memory usage is +/// at or above `memory_threshold` of the limit, drain evictions older than +/// `retention_secs` from every listener's buffer and replay them into that +/// listener's indexer as ordinary `Removed` events. +pub fn spawn_cleanup_loop( + registry: Arc, + retention_secs: u64, + memory_threshold: f64, + cancel: CancellationToken, +) { + tokio::spawn(async move { + let mut ticker = tokio::time::interval(SWEEP_INTERVAL); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let mut tick_count: u64 = 0; + loop { + tokio::select! { + _ = cancel.cancelled() => return, + _ = ticker.tick() => {} + } + tick_count += 1; + + let usage = memory_usage_fraction(); + + // Once a minute, log buffer/memory state unconditionally — this is + // the only visibility into the buffers below the sweep threshold, + // and the pod's log retention is far too short to rely on + // startup-time lines. + if tick_count % 4 == 1 { + let (mut pending, mut queued) = (0usize, 0usize); + let records = registry.listener_records(); + for (_, _, buffer, _) in &records { + let buffer = buffer.lock(); + pending += buffer.len(); + queued += buffer.queue_len(); + } + tracing::info!( + memory_fraction = usage.map(|u| format!("{u:.3}")), + pending_evictions = pending, + queued_entries = queued, + listeners = records.len(), + "keep-evictions status" + ); + } + + let Some(usage) = usage else { + tracing::warn!("keep-evictions sweep: could not read memory usage; skipping"); + continue; + }; + if usage < memory_threshold { + continue; + } + + let cutoff = monotonic_secs().saturating_sub(retention_secs); + let mut replayed = 0usize; + let mut still_pending = 0usize; + for (worker_id, dp_rank, pending_evictions, indexer) in registry.listener_records() { + // Drain under the lock, apply outside it: applies await on the + // indexer channels and must not block the listener's event path. + let (doomed, remaining) = { + let mut buffer = pending_evictions.lock(); + let doomed = buffer.drain_older_than(cutoff); + (doomed, buffer.len()) + }; + still_pending += remaining; + for (tier, hashes) in doomed { + replayed += hashes.len(); + for chunk in hashes.chunks(REMOVE_CHUNK) { + let event = RouterEvent::with_storage_tier( + worker_id, + KvCacheEvent { + // Synthetic id: the standalone apply path does + // gap-tracking on ZMQ seq numbers, not on + // event_id, so 0 is safe here. + event_id: 0, + data: KvCacheEventData::Removed(KvCacheRemoveData { + block_hashes: chunk.to_vec(), + }), + dp_rank, + }, + tier, + ); + if let Err(error) = indexer.apply_event_routed(event).await { + tracing::warn!( + worker_id, + dp_rank, + %error, + "keep-evictions sweep: replayed eviction was rejected" + ); + } + } + } + } + tracing::info!( + usage = format!("{:.1}%", usage * 100.0), + threshold = format!("{:.1}%", memory_threshold * 100.0), + replayed, + still_pending, + retention_secs, + "keep-evictions sweep: replayed aged evictions into the tree" + ); + } + }); +} + +#[cfg(test)] +mod tests { + use super::*; + + fn hashes(values: &[u64]) -> Vec { + values + .iter() + .copied() + .map(ExternalSequenceBlockHash) + .collect() + } + + fn drained_flat(buffer: &mut PendingEvictions, cutoff: u64) -> Vec { + let mut out: Vec = buffer + .drain_older_than(cutoff) + .into_values() + .flatten() + .map(|h| h.0) + .collect(); + out.sort_unstable(); + out + } + + #[test] + fn drain_respects_cutoff() { + let mut buffer = PendingEvictions::default(); + buffer.buffer_at(&hashes(&[1, 2]), StorageTier::Device, 100); + buffer.buffer_at(&hashes(&[3]), StorageTier::Device, 200); + + // Entries newer than the cutoff stay. + assert!(drained_flat(&mut buffer, 99).is_empty()); + assert_eq!(buffer.len(), 3); + + // Cutoff between the batches: only the older batch drains. + assert_eq!(drained_flat(&mut buffer, 150), vec![1, 2]); + assert_eq!(buffer.len(), 1); + + // Cutoff past everything: the rest drains; a second drain is a no-op. + assert_eq!(drained_flat(&mut buffer, 500), vec![3]); + assert!(buffer.is_empty()); + assert!(drained_flat(&mut buffer, 500).is_empty()); + } + + #[test] + fn restore_cancels_pending_eviction() { + let mut buffer = PendingEvictions::default(); + buffer.buffer(&hashes(&[1, 2, 3]), StorageTier::Device); + buffer.cancel([2]); + + assert_eq!(buffer.len(), 2); + let drained = drained_flat(&mut buffer, monotonic_secs() + 1); + assert_eq!(drained, vec![1, 3], "re-stored block 2 must not be removed"); + // The stale queue entry for 2 must not stall the drain of 3 behind it. + assert!(buffer.is_empty()); + } + + #[test] + fn later_re_eviction_is_not_removed_early() { + let mut buffer = PendingEvictions::default(); + // Evict at t=100, re-store, evict again at t=2000. A drain whose + // cutoff covers only the first eviction must remove NOTHING: the + // block's current eviction is too fresh, and the stale first entry + // must not act on the map's newer state. + buffer.buffer_at(&hashes(&[7]), StorageTier::Device, 100); + buffer.cancel([7]); + buffer.buffer_at(&hashes(&[7]), StorageTier::Device, 2000); + + assert!( + drained_flat(&mut buffer, 1000).is_empty(), + "the fresh re-eviction at t=2000 must not be removed by the stale t=100 entry" + ); + assert_eq!(buffer.len(), 1); + + // Once the second eviction ages past the cutoff, it drains exactly once. + assert_eq!(drained_flat(&mut buffer, 2000), vec![7]); + assert!(buffer.is_empty()); + assert!(buffer.queue.is_empty(), "no leftover queue entries"); + } + + #[test] + fn duplicate_eviction_without_restore_drains_once() { + let mut buffer = PendingEvictions::default(); + // Same hash evicted twice with no store in between (shouldn't happen, + // but must not double-remove): only the latest entry replays. + buffer.buffer_at(&hashes(&[9]), StorageTier::Device, 100); + buffer.buffer_at(&hashes(&[9]), StorageTier::Device, 200); + + assert_eq!(buffer.len(), 1); + assert_eq!(drained_flat(&mut buffer, 500), vec![9]); + assert!(buffer.queue.is_empty()); + } + + #[test] + fn drain_groups_by_storage_tier() { + let mut buffer = PendingEvictions::default(); + buffer.buffer(&hashes(&[1]), StorageTier::Device); + buffer.buffer(&hashes(&[2]), StorageTier::HostPinned); + + let drained = buffer.drain_older_than(monotonic_secs() + 1); + assert_eq!(drained.get(&StorageTier::Device).map(Vec::len), Some(1)); + assert_eq!(drained.get(&StorageTier::HostPinned).map(Vec::len), Some(1)); + } + + #[test] + fn clear_drops_everything() { + let mut buffer = PendingEvictions::default(); + buffer.buffer(&hashes(&[1, 2]), StorageTier::Device); + buffer.clear(); + assert!(buffer.is_empty()); + assert!(drained_flat(&mut buffer, monotonic_secs() + 1).is_empty()); + } + + /// End-to-end shape of the sweep: a stored-then-evicted block stays + /// queryable while parked, and disappears once the drained eviction is + /// replayed through the normal `Removed` apply path. + #[tokio::test] + async fn drained_evictions_replay_into_the_tree() { + use crate::protocols::{LocalBlockHash, WorkerWithDpRank, compute_seq_hash_for_block}; + use crate::services::indexer::backend::{create_indexer, test_util::store_event}; + + let worker = WorkerWithDpRank::new(7, 0); + let indexer = create_indexer(4, 1); + indexer + .apply_event_routed(store_event( + worker.worker_id, + worker.dp_rank, + 0, + &[], + &[11, 12], + StorageTier::Device, + )) + .await + .unwrap(); + let _ = indexer.dump_events().await; // flush in-flight applies + + let sequence = vec![LocalBlockHash(11), LocalBlockHash(12)]; + let scores = indexer.find_matches(sequence.clone()).await.unwrap(); + assert_eq!(scores.scores.get(&worker).copied(), Some(2)); + + // The engine evicts the tail block; the buffer parks it, so the tree + // still matches both blocks. + let tail_seq_hash = ExternalSequenceBlockHash(compute_seq_hash_for_block(&sequence)[1]); + let mut buffer = PendingEvictions::default(); + buffer.buffer_at(&[tail_seq_hash], StorageTier::Device, 100); + let scores = indexer.find_matches(sequence.clone()).await.unwrap(); + assert_eq!(scores.scores.get(&worker).copied(), Some(2)); + + // The sweep drains the aged entry and replays it as a Removed event — + // exactly what spawn_cleanup_loop's body does. + for (tier, hashes) in buffer.drain_older_than(1000) { + let event = RouterEvent::with_storage_tier( + worker.worker_id, + KvCacheEvent { + event_id: 0, + data: KvCacheEventData::Removed(KvCacheRemoveData { + block_hashes: hashes, + }), + dp_rank: worker.dp_rank, + }, + tier, + ); + indexer.apply_event_routed(event).await.unwrap(); + } + let _ = indexer.dump_events().await; + + let scores = indexer.find_matches(sequence).await.unwrap(); + assert_eq!( + scores.scores.get(&worker).copied().unwrap_or(0), + 1, + "only the surviving head block should match after the replayed eviction" + ); + assert!(buffer.is_empty()); + } + + #[test] + fn memory_usage_fraction_reads_something_sane() { + let usage = memory_usage_fraction().expect("some memory source must be readable"); + assert!(usage > 0.0 && usage < 10.0, "got {usage}"); + } +} diff --git a/lib/kv-router/src/services/indexer/h24.rs b/lib/kv-router/src/services/indexer/h24.rs new file mode 100644 index 000000000000..30ca263a3fa1 --- /dev/null +++ b/lib/kv-router/src/services/indexer/h24.rs @@ -0,0 +1,469 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! `--h24` counterfactual indexer: answers "how many prefix blocks of this +//! query were stored by ANY worker within the retention horizon", ignoring +//! evictions and worker identity entirely. +//! +//! Instead of a radix tree it keeps two flat maps: +//! - `attach`: engine sequence hash -> our chained prefix hash, so store +//! events that name their parent by engine hash can keep chaining. +//! - `seen`: chained prefix hash -> (last_stored, last_touched) timestamps. +//! +//! The chain recurrence is the crate-wide [`compute_next_seq_hash`] (first +//! block's prefix hash equals its tokens hash), so `seen` keys are the same +//! `SequenceHash` values the rest of the codebase computes for a prefix. +//! Queries recompute the chain from the caller's tokens hashes and walk it to +//! the first miss; matched-block ages feed the retention-curve histograms. +//! +//! Matches are reported under a single synthetic worker ([`H24_WORKER_ID`]) +//! so the HTTP response shape stays identical for existing callers. + +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; + +use parking_lot::Mutex; +use rustc_hash::FxHashMap; +use tokio_util::sync::CancellationToken; + +use crate::protocols::{ + KvCacheStoreData, LocalBlockHash, OverlapScores, WorkerWithDpRank, compute_next_seq_hash, +}; + +use super::backend::Indexer; +use super::evictions::memory_usage_fraction; +use super::metrics; +use super::registry::WorkerRegistry; + +/// Synthetic worker id all h24 matches are reported under. Real worker ids +/// are xxh3 hashes of pod names, so 0 cannot collide with one. +pub const H24_WORKER_ID: u64 = 0; + +/// How often the expiry sweep runs. +const SWEEP_INTERVAL: Duration = Duration::from_secs(600); + +/// Memory fraction above which the sweep tightens the horizon to half. +const MEMORY_GUARD_THRESHOLD: f64 = 0.92; + +/// Wall-clock seconds since the Unix epoch (u32 is fine until 2106). +fn unix_secs() -> u32 { + use std::time::{SystemTime, UNIX_EPOCH}; + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs() as u32) + .unwrap_or(0) +} + +#[derive(Clone, Copy)] +struct AttachEntry { + prefix: u64, + last_seen: u32, +} + +#[derive(Clone, Copy)] +struct SeenEntry { + last_stored: u32, + last_touched: u32, +} + +#[derive(Default)] +struct Inner { + /// Engine sequence hash -> our chained prefix hash. Store-side plumbing + /// only: lets a store that names its parent by engine hash keep chaining. + attach: FxHashMap, + /// Chained prefix hash -> timestamps. One entry per unique prefix. + seen: FxHashMap, +} + +pub struct H24Indexer { + inner: Mutex, + parent_misses: AtomicU64, + ignored_events: AtomicU64, +} + +impl H24Indexer { + pub fn new() -> Self { + Self { + inner: Mutex::new(Inner::default()), + parent_misses: AtomicU64::new(0), + ignored_events: AtomicU64::new(0), + } + } + + /// Apply a store event: resolve the parent's prefix hash, chain through + /// the event's blocks, and stamp `seen`/`attach`. A missing parent drops + /// the whole event (every block chains off it) and is counted. + pub fn apply_store(&self, data: &KvCacheStoreData) { + self.apply_store_at(data, unix_secs()); + } + + fn apply_store_at(&self, data: &KvCacheStoreData, now: u32) { + if data.blocks.is_empty() { + return; + } + + let mut inner = self.inner.lock(); + + // None parent = sequence start: the first block's prefix hash equals + // its tokens hash (compute_seq_hash_for_block's recurrence). + let mut parent_prefix: Option = match data.parent_hash { + None => None, + Some(parent) => match inner.attach.get(&parent.0) { + Some(entry) => Some(entry.prefix), + None => { + self.parent_misses.fetch_add(1, Ordering::Relaxed); + metrics::h24_inc_parent_miss(); + return; + } + }, + }; + + for block in &data.blocks { + let prefix = match parent_prefix { + None => block.tokens_hash.0, + Some(parent) => compute_next_seq_hash(parent, block.tokens_hash), + }; + let entry = inner.seen.entry(prefix).or_insert(SeenEntry { + last_stored: now, + last_touched: now, + }); + entry.last_stored = now; + entry.last_touched = entry.last_touched.max(now); + inner.attach.insert( + block.block_hash.0, + AttachEntry { + prefix, + last_seen: now, + }, + ); + parent_prefix = Some(prefix); + } + } + + /// Longest-prefix match: rebuild the chain from the query's tokens hashes + /// and walk `seen` to the first miss. Matched blocks bump `last_touched` + /// and feed the age histograms; the score is reported under the synthetic + /// worker. + pub fn find_matches(&self, sequence: &[LocalBlockHash]) -> OverlapScores { + self.find_matches_at(sequence, unix_secs()) + } + + fn find_matches_at(&self, sequence: &[LocalBlockHash], now: u32) -> OverlapScores { + let mut scores = OverlapScores::new(); + if sequence.is_empty() { + return scores; + } + + let mut inner = self.inner.lock(); + let mut depth = 0u32; + let mut prefix: Option = None; + for hash in sequence { + let next = match prefix { + None => hash.0, + Some(parent) => compute_next_seq_hash(parent, *hash), + }; + let Some(entry) = inner.seen.get_mut(&next) else { + break; + }; + metrics::h24_observe_age("since_stored", now.saturating_sub(entry.last_stored)); + metrics::h24_observe_age("since_touched", now.saturating_sub(entry.last_touched)); + entry.last_touched = entry.last_touched.max(now); + depth += 1; + prefix = Some(next); + } + drop(inner); + + if depth > 0 { + metrics::h24_add_matched_blocks(depth as usize); + scores + .scores + .insert(WorkerWithDpRank::new(H24_WORKER_ID, 0), depth); + } + scores + } + + pub fn note_ignored_event(&self) { + self.ignored_events.fetch_add(1, Ordering::Relaxed); + } + + /// Drop entries not seen (stored or touched) since `cutoff_secs`. + /// Returns (seen_dropped, attach_dropped). + pub fn expire_older_than(&self, cutoff_secs: u32) -> (usize, usize) { + let mut inner = self.inner.lock(); + let seen_before = inner.seen.len(); + let attach_before = inner.attach.len(); + inner + .seen + .retain(|_, e| e.last_stored.max(e.last_touched) >= cutoff_secs); + inner.attach.retain(|_, e| e.last_seen >= cutoff_secs); + ( + seen_before - inner.seen.len(), + attach_before - inner.attach.len(), + ) + } + + /// (seen entries, attach entries). + pub fn sizes(&self) -> (usize, usize) { + let inner = self.inner.lock(); + (inner.seen.len(), inner.attach.len()) + } + + pub fn parent_misses(&self) -> u64 { + self.parent_misses.load(Ordering::Relaxed) + } + + pub fn ignored_events(&self) -> u64 { + self.ignored_events.load(Ordering::Relaxed) + } +} + +impl Default for H24Indexer { + fn default() -> Self { + Self::new() + } +} + +/// Background expiry sweep: every [`SWEEP_INTERVAL`], drop entries older than +/// `horizon_secs` from every h24 indexer in the registry, refresh the size +/// gauges, and tighten to half the horizon if memory is still above +/// [`MEMORY_GUARD_THRESHOLD`] afterwards. +pub fn spawn_expiry_loop( + registry: Arc, + horizon_secs: u64, + cancel: CancellationToken, +) { + tokio::spawn(async move { + let mut ticker = tokio::time::interval(SWEEP_INTERVAL); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + tokio::select! { + _ = cancel.cancelled() => return, + _ = ticker.tick() => {} + } + + let now = unix_secs(); + let cutoff = now.saturating_sub(horizon_secs as u32); + let mut dropped = (0usize, 0usize); + let mut sizes = (0usize, 0usize); + let indexers = registry.all_indexers_with_block_size(); + for (_, indexer, _) in &indexers { + let Indexer::H24(h24) = indexer else { + continue; + }; + let (seen_dropped, attach_dropped) = h24.expire_older_than(cutoff); + dropped.0 += seen_dropped; + dropped.1 += attach_dropped; + let (seen, attach) = h24.sizes(); + sizes.0 += seen; + sizes.1 += attach; + } + + // Memory guard: horizon is the target, memory is the hard limit. + let usage = memory_usage_fraction(); + if usage.is_some_and(|u| u >= MEMORY_GUARD_THRESHOLD) { + let tight_cutoff = now.saturating_sub((horizon_secs / 2) as u32); + let mut guard_dropped = 0usize; + for (_, indexer, _) in &indexers { + if let Indexer::H24(h24) = indexer { + let (s, a) = h24.expire_older_than(tight_cutoff); + guard_dropped += s + a; + } + } + tracing::warn!( + usage = ?usage, + horizon_secs, + guard_dropped, + "h24 memory guard: usage above threshold, expired at half horizon" + ); + } + + metrics::h24_set_sizes(sizes.0, sizes.1); + tracing::info!( + seen_blocks = sizes.0, + attach_entries = sizes.1, + seen_dropped = dropped.0, + attach_dropped = dropped.1, + horizon_secs, + memory_fraction = usage.map(|u| format!("{u:.3}")), + "h24 expiry sweep" + ); + } + }); +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocols::{ + ExternalSequenceBlockHash, KvCacheStoredBlockData, compute_seq_hash_for_block, + }; + + /// Store `local_hashes` chained after `prefix_hashes`, with engine ext + /// hashes derived from the crate chain recurrence offset by `ext_salt` + /// (different salts model different hash regimes storing the same tokens). + fn store_data(prefix_hashes: &[u64], local_hashes: &[u64], ext_salt: u64) -> KvCacheStoreData { + let prefix: Vec = + prefix_hashes.iter().copied().map(LocalBlockHash).collect(); + let full: Vec = prefix_hashes + .iter() + .chain(local_hashes.iter()) + .copied() + .map(LocalBlockHash) + .collect(); + let full_seq = compute_seq_hash_for_block(&full); + let parent_hash = compute_seq_hash_for_block(&prefix) + .last() + .map(|h| ExternalSequenceBlockHash(h.wrapping_add(ext_salt))); + let blocks = local_hashes + .iter() + .zip(full_seq[prefix_hashes.len()..].iter()) + .map(|(&local, &seq)| KvCacheStoredBlockData { + block_hash: ExternalSequenceBlockHash(seq.wrapping_add(ext_salt)), + tokens_hash: LocalBlockHash(local), + mm_extra_info: None, + }) + .collect(); + KvCacheStoreData { + parent_hash, + start_position: None, + blocks, + } + } + + fn seq(hashes: &[u64]) -> Vec { + hashes.iter().copied().map(LocalBlockHash).collect() + } + + fn depth(scores: &OverlapScores) -> u32 { + scores + .scores + .get(&WorkerWithDpRank::new(H24_WORKER_ID, 0)) + .copied() + .unwrap_or(0) + } + + #[test] + fn store_then_query_matches_prefix() { + let h24 = H24Indexer::new(); + h24.apply_store_at(&store_data(&[], &[1, 2, 3], 0), 100); + + assert_eq!(depth(&h24.find_matches_at(&seq(&[1, 2, 3]), 200)), 3); + assert_eq!(depth(&h24.find_matches_at(&seq(&[1, 2]), 200)), 2); + // Divergence after block 1 stops the walk. + assert_eq!(depth(&h24.find_matches_at(&seq(&[1, 9, 3]), 200)), 1); + // Different first block = no match at all. + assert_eq!(depth(&h24.find_matches_at(&seq(&[9, 2, 3]), 200)), 0); + assert_eq!(h24.sizes(), (3, 3)); + } + + #[test] + fn chained_store_attaches_via_engine_hash() { + let h24 = H24Indexer::new(); + h24.apply_store_at(&store_data(&[], &[1, 2], 0), 100); + // Second event chains off ext hash of block 2. + h24.apply_store_at(&store_data(&[1, 2], &[3], 0), 150); + + assert_eq!(depth(&h24.find_matches_at(&seq(&[1, 2, 3]), 200)), 3); + assert_eq!(h24.parent_misses(), 0); + } + + #[test] + fn same_tokens_different_ext_regime_dedup() { + let h24 = H24Indexer::new(); + // Two regimes (different engine ext hashes) store identical tokens. + h24.apply_store_at(&store_data(&[], &[1, 2], 0), 100); + h24.apply_store_at(&store_data(&[], &[1, 2], 777), 150); + + let (seen, attach) = h24.sizes(); + assert_eq!(seen, 2, "identical token chains must share seen entries"); + assert_eq!(attach, 4, "each regime keeps its own attach aliases"); + // Both regimes can chain further stores off their own ext hashes. + h24.apply_store_at(&store_data(&[1, 2], &[3], 777), 160); + assert_eq!(depth(&h24.find_matches_at(&seq(&[1, 2, 3]), 200)), 3); + assert_eq!(h24.parent_misses(), 0); + } + + #[test] + fn missing_parent_drops_event_and_counts() { + let h24 = H24Indexer::new(); + h24.apply_store_at(&store_data(&[1, 2], &[3], 0), 100); + + assert_eq!(h24.parent_misses(), 1); + assert_eq!(h24.sizes(), (0, 0)); + assert_eq!(depth(&h24.find_matches_at(&seq(&[1, 2, 3]), 200)), 0); + } + + #[test] + fn restore_refreshes_timestamps_and_expiry_respects_touch() { + let h24 = H24Indexer::new(); + h24.apply_store_at(&store_data(&[], &[1, 2], 0), 100); + + // A query at t=5000 touches block 1 only (divergent second block). + let _ = h24.find_matches_at(&seq(&[1, 9]), 5000); + + // Expire everything not seen since t=1000: block 2's entry dies + // (stored 100, never touched), block 1 survives via its touch. + let (seen_dropped, attach_dropped) = h24.expire_older_than(1000); + assert_eq!(seen_dropped, 1); + assert_eq!(attach_dropped, 2, "attach expires on store recency only"); + assert_eq!(depth(&h24.find_matches_at(&seq(&[1, 2]), 6000)), 1); + + // Re-store refreshes: nothing left to expire at the same cutoff. + h24.apply_store_at(&store_data(&[], &[1, 2], 0), 7000); + assert_eq!(h24.expire_older_than(1000), (0, 0)); + assert_eq!(depth(&h24.find_matches_at(&seq(&[1, 2]), 8000)), 2); + } + + /// Every routing group of a model shares ONE h24 map, and deregistration + /// with a pod's real routing group still finds it. + #[tokio::test] + async fn registry_collapses_routing_groups_into_one_h24_indexer() { + use crate::services::indexer::registry::RegistryOptions; + let registry = WorkerRegistry::new(1) + .with_options(RegistryOptions { + h24: true, + ..Default::default() + }) + .unwrap(); + registry.signal_ready(); + for (instance, group) in [(1u64, "regime-a"), (2u64, "regime-b")] { + registry + .register( + instance, + format!("tcp://127.0.0.1:1576{instance}"), + 0, + "m".to_string(), + group.to_string(), + 4, + None, + ) + .await + .unwrap(); + } + + let trees = registry.indexers_for_model("m"); + assert_eq!(trees.len(), 1, "h24 keeps one tree per model"); + assert_eq!(trees[0].0.routing_group, "h24"); + assert!(matches!(trees[0].1, Indexer::H24(_))); + + registry.deregister(1, "m", "regime-a").await.unwrap(); + registry.deregister(2, "m", "regime-b").await.unwrap(); + assert!(registry.list().is_empty()); + } + + #[test] + fn empty_query_and_empty_store_are_noops() { + let h24 = H24Indexer::new(); + h24.apply_store_at( + &KvCacheStoreData { + parent_hash: None, + start_position: None, + blocks: vec![], + }, + 100, + ); + assert_eq!(h24.sizes(), (0, 0)); + assert!(h24.find_matches_at(&[], 100).scores.is_empty()); + } +} diff --git a/lib/kv-router/src/services/indexer/kv_recover.rs b/lib/kv-router/src/services/indexer/kv_recover.rs new file mode 100644 index 000000000000..30d0f5230644 --- /dev/null +++ b/lib/kv-router/src/services/indexer/kv_recover.rs @@ -0,0 +1,410 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Per-worker gap recovery over HTTP: `GET /kv_recover`. +//! +//! An alternative to the ZMQ DEALER replay socket for engines that serve their +//! local indexer over HTTP. On a sequence gap the listener fetches a +//! [`WorkerKvQueryResponse`] for `[start, end)` and applies it per +//! [`plan_recovery`]: +//! +//! - `Events` — apply all, resume at `last_event_id`. +//! - `TreeDump` — drop this `(worker, dp_rank)`'s state, apply all, resume at +//! `last_event_id` (dump event ids are synthetic). +//! - anything else — no-op, the watermark stays where it was. +//! +//! Recovery is at-least-once: the resume point is always the response's +//! `last_event_id`, never a count of applied events. Events are applied under +//! the consumer-assigned identity; the worker's self-reported one is ignored. +//! +//! A large engine's TreeDump is tens of MB serialized inside the engine +//! process, and a freshly started indexer asks every engine at once, so +//! downloads share a process-wide gate, get a long total timeout, and are +//! retried with backoff before the gap is given up. + +use std::time::Duration; + +use anyhow::Context; +use tokio::sync::Semaphore; +use tokio_util::sync::CancellationToken; + +use crate::indexer::WorkerKvQueryResponse; +use crate::protocols::{ResetScope, RouterEvent, WorkerId}; + +pub const DEFAULT_RECOVER_TIMEOUT_S: u64 = 120; +pub const DEFAULT_RECOVER_CONCURRENCY: usize = 8; + +const RECOVER_CONNECT_TIMEOUT: Duration = Duration::from_secs(5); +const RECOVER_ATTEMPTS: u32 = 3; +const RECOVER_BACKOFF_BASE: Duration = Duration::from_secs(2); + +#[derive(Debug, Clone, Copy)] +pub struct KvRecoverSettings { + /// Total timeout (connect + body) for one download. + pub timeout_s: u64, + /// Downloads allowed in flight across every listener of this process. + pub concurrency: usize, +} + +impl Default for KvRecoverSettings { + fn default() -> Self { + Self { + timeout_s: DEFAULT_RECOVER_TIMEOUT_S, + concurrency: DEFAULT_RECOVER_CONCURRENCY, + } + } +} + +/// Shared by every listener of a registry: one HTTP client and one gate. +pub struct KvRecoverClient { + http: reqwest::Client, + gate: Semaphore, + backoff_base: Duration, +} + +impl KvRecoverClient { + pub fn new(settings: KvRecoverSettings) -> anyhow::Result { + Self::with_backoff(settings, RECOVER_BACKOFF_BASE) + } + + fn with_backoff(settings: KvRecoverSettings, backoff_base: Duration) -> anyhow::Result { + let http = reqwest::Client::builder() + .connect_timeout(RECOVER_CONNECT_TIMEOUT) + .timeout(Duration::from_secs(settings.timeout_s.max(1))) + .build() + .context("failed to build kv_recover HTTP client")?; + Ok(Self { + http, + gate: Semaphore::new(settings.concurrency.max(1)), + backoff_base, + }) + } + + /// Download the worker's recovery response for `[start_seq, end_seq)`, + /// retrying failed downloads. `None` when every attempt failed (logged as + /// `kv_recover request failed; giving up`) or `cancel` fired. + pub(super) async fn fetch( + &self, + recover_endpoint: &str, + start_seq: u64, + end_seq: u64, + worker_id: WorkerId, + dp_rank: u32, + cancel: &CancellationToken, + ) -> Option { + let url = format!("{}/kv_recover", recover_endpoint.trim_end_matches('/')); + for attempt in 0..RECOVER_ATTEMPTS { + let result = tokio::select! { + _ = cancel.cancelled() => return None, + result = self.fetch_once(&url, start_seq, end_seq) => result, + }; + let error = match result { + Ok(response) => return Some(response), + Err(error) => error, + }; + if attempt + 1 == RECOVER_ATTEMPTS { + tracing::error!( + worker_id, + dp_rank, + attempts = RECOVER_ATTEMPTS, + gap_size = end_seq.saturating_sub(start_seq), + error = %format!("{error:#}"), + "kv_recover request failed; giving up, batches lost" + ); + return None; + } + let delay = self.backoff(attempt); + tracing::warn!( + worker_id, + dp_rank, + attempt = attempt + 1, + retry_in_ms = delay.as_millis() as u64, + error = %format!("{error:#}"), + "kv_recover request failed; retrying" + ); + tokio::select! { + _ = cancel.cancelled() => return None, + _ = tokio::time::sleep(delay) => {} + } + } + None + } + + async fn fetch_once( + &self, + url: &str, + start_seq: u64, + end_seq: u64, + ) -> anyhow::Result { + let _permit = self + .gate + .acquire() + .await + .expect("kv_recover gate is never closed"); + let response = self + .http + .get(url) + .query(&[("start", start_seq), ("end", end_seq)]) + .send() + .await?; + anyhow::ensure!( + response.status().is_success(), + "kv_recover returned status {}", + response.status() + ); + Ok(response.json().await?) + } + + /// 2 s, 4 s, 8 s, ... capped at 16x the base. + fn backoff(&self, attempt: u32) -> Duration { + self.backoff_base * (1u32 << attempt.min(4)) + } +} + +/// What a listener must do with one recovery response. +#[derive(Debug, Default)] +pub(super) struct RecoveryPlan { + /// Drop this `(worker, dp_rank)`'s indexed state before applying. + pub reset_dp_rank: bool, + /// Events to apply, already rewritten to the consumer-assigned identity. + pub events: Vec, + /// New watermark; `None` leaves it unchanged. + pub resume_at: Option, +} + +pub(super) fn plan_recovery( + response: WorkerKvQueryResponse, + worker_id: WorkerId, + dp_rank: u32, +) -> RecoveryPlan { + let reassign = |mut events: Vec| { + for event in &mut events { + event.worker_id = worker_id; + event.event.dp_rank = dp_rank; + } + events + }; + match response { + WorkerKvQueryResponse::Events { + events, + last_event_id, + } => RecoveryPlan { + reset_dp_rank: false, + events: reassign(events), + resume_at: Some(last_event_id), + }, + WorkerKvQueryResponse::TreeDump { + events, + last_event_id, + reset_scope, + } => { + // A domain-scoped snapshot does not describe the whole rank, so it + // must not wipe it; applying it on top stays at-least-once. + let reset_dp_rank = matches!(reset_scope, ResetScope::All); + if !reset_dp_rank { + tracing::warn!( + worker_id, + dp_rank, + ?reset_scope, + "kv_recover TreeDump is domain-scoped; applying without a reset" + ); + } + RecoveryPlan { + reset_dp_rank, + events: reassign(events), + resume_at: Some(last_event_id), + } + } + WorkerKvQueryResponse::StateAgentRecovery { response, .. } => { + plan_recovery(*response, worker_id, dp_rank) + } + WorkerKvQueryResponse::TooNew { + newest_available, .. + } => { + tracing::debug!( + worker_id, + dp_rank, + newest_available, + "kv_recover returned TooNew; consumer is ahead, no-op" + ); + RecoveryPlan::default() + } + other => { + tracing::warn!( + worker_id, + dp_rank, + response = ?other, + "kv_recover returned no recoverable state; no-op" + ); + RecoveryPlan::default() + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocols::{KvCacheEventData, StorageTier}; + use crate::services::indexer::backend::test_util::store_event; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + /// A real `/kv_recover` Events response captured from a prod + /// zai-org/GLM-5.3 engine (2026-09-23), trimmed to two events. + const CAPTURED_EVENTS_RESPONSE: &str = r#"{"Events": {"events": [{"worker_id": 0, "storage_tier": "device", "event": {"event_id": 22636960, "data": {"removed": {"block_hashes": [18405780835976372970]}}, "dp_rank": 0}}, {"worker_id": 0, "storage_tier": "device", "event": {"event_id": 22636961, "data": {"stored": {"parent_hash": 270656378982115943, "start_position": null, "blocks": [{"block_hash": 2733357436102908004, "tokens_hash": 2825904518766229643, "mm_extra_info": null}, {"block_hash": 2451106649480522483, "tokens_hash": 10358785732907640880, "mm_extra_info": null}]}}, "dp_rank": 0}}], "last_event_id": 22637089}}"#; + + #[test] + fn captured_engine_response_plans_events_under_consumer_identity() { + let response: WorkerKvQueryResponse = + serde_json::from_str(CAPTURED_EVENTS_RESPONSE).unwrap(); + let plan = plan_recovery(response, 7, 1); + + assert!(!plan.reset_dp_rank); + assert_eq!(plan.resume_at, Some(22637089)); + assert_eq!(plan.events.len(), 2); + assert!( + plan.events + .iter() + .all(|e| e.worker_id == 7 && e.event.dp_rank == 1) + ); + let KvCacheEventData::Stored(stored) = &plan.events[1].event.data else { + panic!("second captured event is a store"); + }; + assert_eq!(stored.blocks.len(), 2); + } + + #[test] + fn legacy_tree_dump_resets_the_rank() { + let response: WorkerKvQueryResponse = serde_json::from_value(serde_json::json!({ + "TreeDump": { + "events": [store_event(999, 3, 0, &[], &[22], StorageTier::Device)], + "last_event_id": 10 + } + })) + .unwrap(); + let plan = plan_recovery(response, 7, 0); + + assert!(plan.reset_dp_rank); + assert_eq!(plan.resume_at, Some(10)); + assert_eq!(plan.events[0].worker_id, 7); + } + + #[test] + fn too_new_leaves_the_watermark_alone() { + let plan = plan_recovery( + WorkerKvQueryResponse::TooNew { + requested_start: Some(4), + requested_end: Some(9), + newest_available: 3, + }, + 7, + 0, + ); + assert!(!plan.reset_dp_rank); + assert!(plan.events.is_empty()); + assert_eq!(plan.resume_at, None); + } + + /// Serve `/kv_recover` on a local port: the first `failures` requests get + /// a 500, the rest the captured Events body. + async fn flaky_recover_server(failures: usize) -> (String, Arc) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let requests = Arc::new(AtomicUsize::new(0)); + let seen = requests.clone(); + tokio::spawn(async move { + loop { + let (mut stream, _) = listener.accept().await.unwrap(); + let n = seen.fetch_add(1, Ordering::SeqCst); + let mut buf = [0u8; 4096]; + let _ = stream.read(&mut buf).await; + let (status, body) = if n < failures { + ("500 Internal Server Error", "boom") + } else { + ("200 OK", CAPTURED_EVENTS_RESPONSE) + }; + let reply = format!( + "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + let _ = stream.write_all(reply.as_bytes()).await; + } + }); + (endpoint, requests) + } + + fn fast_client(timeout_s: u64) -> KvRecoverClient { + KvRecoverClient::with_backoff( + KvRecoverSettings { + timeout_s, + concurrency: 1, + }, + Duration::from_millis(10), + ) + .unwrap() + } + + #[tokio::test] + async fn failed_downloads_are_retried_until_one_succeeds() { + let (endpoint, requests) = flaky_recover_server(2).await; + let response = fast_client(5) + .fetch(&endpoint, 5, 9, 7, 0, &CancellationToken::new()) + .await; + + assert!(matches!( + response, + Some(WorkerKvQueryResponse::Events { .. }) + )); + assert_eq!(requests.load(Ordering::SeqCst), 3); + } + + #[tokio::test] + async fn gap_is_given_up_after_the_last_attempt() { + let (endpoint, requests) = flaky_recover_server(usize::MAX).await; + let response = fast_client(5) + .fetch(&endpoint, 5, 9, 7, 0, &CancellationToken::new()) + .await; + + assert!(response.is_none()); + assert_eq!(requests.load(Ordering::SeqCst), RECOVER_ATTEMPTS as usize); + } + + #[tokio::test] + async fn a_download_slower_than_the_timeout_fails_the_attempt() { + // Accepts and never answers. + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let _hold = tokio::spawn(async move { + let mut open = Vec::new(); + loop { + let (stream, _) = listener.accept().await.unwrap(); + open.push(stream); + } + }); + + let started = std::time::Instant::now(); + let response = fast_client(1) + .fetch(&endpoint, 5, 9, 7, 0, &CancellationToken::new()) + .await; + + assert!(response.is_none()); + let elapsed = started.elapsed(); + assert!( + elapsed >= Duration::from_secs(RECOVER_ATTEMPTS as u64) + && elapsed < Duration::from_secs(10), + "each attempt must end at the 1 s timeout, took {elapsed:?}" + ); + } + + #[test] + fn default_settings_match_the_documented_flags() { + let settings = KvRecoverSettings::default(); + assert_eq!(settings.timeout_s, 120); + assert_eq!(settings.concurrency, 8); + let client = KvRecoverClient::new(settings).unwrap(); + assert_eq!(client.backoff(0), Duration::from_secs(2)); + assert_eq!(client.backoff(1), Duration::from_secs(4)); + assert_eq!(client.backoff(9), Duration::from_secs(32)); + } +} diff --git a/lib/kv-router/src/services/indexer/listener.rs b/lib/kv-router/src/services/indexer/listener.rs index c844c4fbbd80..9d6202e2e47b 100644 --- a/lib/kv-router/src/services/indexer/listener.rs +++ b/lib/kv-router/src/services/indexer/listener.rs @@ -7,12 +7,15 @@ use std::sync::atomic::{AtomicU64, Ordering}; use tokio::sync::watch; use tokio_util::sync::CancellationToken; -use crate::protocols::{WorkerId, WorkerWithDpRank}; +use crate::protocols::{KvCacheEventData, RouterEvent, WorkerId, WorkerWithDpRank}; use crate::recovery::{CursorObservation, CursorState}; use crate::zmq_wire::{ZmqEventNormalizer, decode_event_batch}; +use super::audit; use super::backend::Indexer; -use super::registry::ListenerRecord; +use super::block_size::BlockSizeGuard; +use super::kv_recover::plan_recovery; +use super::registry::{ListenerRecord, RecoverTarget, SharedPendingEvictions}; use crate::services::common::zmq::{ MultipartMessage, SharedSocket, connect_dealer_socket, connect_sub_socket, recv_multipart, send_multipart, @@ -151,8 +154,12 @@ struct ListenerLoop { cancel: CancellationToken, live_socket: SharedSocket, replay_socket: Option, + recover: Option, + pending_evictions: Option, + audit_log: bool, watermark: Arc, normalizer: ZmqEventNormalizer, + block_size_guard: BlockSizeGuard, messages_processed: u64, } @@ -166,6 +173,9 @@ impl ListenerLoop { cancel: CancellationToken, live_socket: SharedSocket, replay_socket: Option, + recover: Option, + pending_evictions: Option, + audit_log: bool, watermark: Arc, ) -> Self { Self { @@ -175,12 +185,49 @@ impl ListenerLoop { cancel, live_socket, replay_socket, + recover, + pending_evictions, + audit_log, watermark, normalizer: ZmqEventNormalizer::new(block_size), + block_size_guard: BlockSizeGuard::new(block_size), messages_processed: 0, } } + /// Audit-log `event` (as published) and apply the `--keep-evictions` + /// filter. `true` when the event must not reach the tree. + fn withhold_from_tree(&self, event: &RouterEvent, seq: u64, source: &'static str) -> bool { + if self.audit_log { + audit::log_event(event, seq, source); + } + self.park_kept_eviction(event) + } + + /// `--keep-evictions`: a store cancels any parked eviction of its blocks + /// and applies; an eviction is parked instead of applied; a clear is + /// dropped. `true` when the event must not reach the tree. + fn park_kept_eviction(&self, event: &RouterEvent) -> bool { + let Some(pending) = self.pending_evictions.as_ref() else { + return false; + }; + match &event.event.data { + KvCacheEventData::Stored(data) => { + pending + .lock() + .cancel(data.blocks.iter().map(|b| b.block_hash.0)); + false + } + KvCacheEventData::Removed(data) => { + pending + .lock() + .buffer(&data.block_hashes, event.storage_tier); + true + } + KvCacheEventData::Cleared => true, + } + } + fn cursor(&self) -> CursorState { cursor_from_watermark(self.watermark.load(Ordering::Acquire)) } @@ -229,25 +276,24 @@ impl ListenerLoop { } } }; - if msg.len() != 4 { - tracing::warn!( - worker_id, - dp_rank, - "Unexpected replay frame count: {}", - msg.len() - ); - break; - } - - // vLLM replay responses include the DEALER identity delimiter plus the - // original PUB topic, so the payload shape is: - // [empty delimiter, topic, sequence number, encoded event batch]. - let payload = msg.get(3).expect("frame count checked above"); + // DEALER strips the ROUTER identity. vLLM includes the PUB topic; + // SGLang sends only the delimiter, sequence number, and payload. + let (seq_bytes, payload) = match msg.as_slice() { + [_, seq, payload] | [_, _, seq, payload] => (seq, payload), + _ => { + tracing::warn!( + worker_id, + dp_rank, + "Unexpected replay frame count: {}", + msg.len() + ); + break; + } + }; if payload.is_empty() { break; } - let seq_bytes = msg.get(2).expect("frame count checked above"); if seq_bytes.len() != 8 { tracing::warn!( worker_id, @@ -268,16 +314,25 @@ impl ListenerLoop { .data_parallel_rank .map_or(dp_rank, |rank| rank.cast_unsigned()); for raw_event in batch.events { - let Some(placement_event) = self.normalizer.normalize( - raw_event, - seq, - WorkerWithDpRank::new(worker_id, effective_dp_rank), - ) else { + let worker = WorkerWithDpRank::new(worker_id, effective_dp_rank); + let Some(raw_event) = self.normalizer.preprocess(raw_event, worker) else { + continue; + }; + if !self.block_size_guard.admit(&raw_event, worker) { + continue; + } + let Some(placement_event) = self + .normalizer + .normalize_preprocessed(raw_event, seq, worker) + else { continue; }; let router_event = placement_event .into_router_event() .expect("local worker placement must convert to router event"); + if self.withhold_from_tree(&router_event, seq, "replay") { + continue; + } indexer .apply_event_routed(router_event) .await @@ -298,6 +353,82 @@ impl ListenerLoop { Ok(replayed) } + /// Recover `[start_seq, end_seq)` over HTTP `/kv_recover` when the worker + /// serves it, else over the ZMQ replay socket. + async fn recover_gap(&mut self, start_seq: u64, end_seq: u64) -> Result { + match self.recover.clone() { + Some(target) => self.recover_over_http(&target, start_seq, end_seq).await, + None => self.replay_gap(start_seq, end_seq).await, + } + } + + async fn recover_over_http( + &mut self, + target: &RecoverTarget, + start_seq: u64, + end_seq: u64, + ) -> Result { + let (worker_id, dp_rank) = (self.worker_id, self.dp_rank); + tracing::info!( + worker_id, + dp_rank, + start_seq, + end_seq, + "Requesting recovery from worker via /kv_recover" + ); + let Some(response) = target + .client + .fetch( + &target.endpoint, + start_seq, + end_seq, + worker_id, + dp_rank, + &self.cancel, + ) + .await + else { + return Ok(0); + }; + + let plan = plan_recovery(response, worker_id, dp_rank); + if plan.reset_dp_rank { + // Parked evictions describe the state being replaced; replaying + // them onto the snapshot could remove blocks it says are live. + if let Some(pending) = self.pending_evictions.as_ref() { + pending.lock().clear(); + } + self.indexer.remove_worker_dp_rank(worker_id, dp_rank).await; + } + let mut applied = 0u64; + for event in plan.events { + if self.withhold_from_tree(&event, event.event.event_id, "recover") { + continue; + } + self.indexer + .apply_event_routed(event) + .await + .map_err(|error| { + format!( + "failed to apply recovered event for worker {worker_id} dp_rank {dp_rank}: {error}" + ) + })?; + applied += 1; + } + if let Some(last_event_id) = plan.resume_at { + self.watermark.store(last_event_id, Ordering::Release); + } + tracing::info!( + worker_id, + dp_rank, + applied, + reset = plan.reset_dp_rank, + resume_at = ?plan.resume_at, + "Recovery via /kv_recover complete" + ); + Ok(applied) + } + async fn handle_gap(&mut self, seq: u64) -> Result<(), String> { match self.cursor().observe(seq) { CursorObservation::Initial { got } if got > 0 => { @@ -308,7 +439,7 @@ impl ListenerLoop { got, "Gap detected: expected seq 0, got {got}" ); - self.replay_gap(0, got).await?; + self.recover_gap(0, got).await?; } CursorObservation::Gap { expected, got } => { tracing::warn!( @@ -318,7 +449,7 @@ impl ListenerLoop { got, "Gap detected: expected seq {expected}, got {got}" ); - self.replay_gap(expected, got).await?; + self.recover_gap(expected, got).await?; } CursorObservation::Initial { .. } | CursorObservation::Contiguous { .. } @@ -344,16 +475,25 @@ impl ListenerLoop { .data_parallel_rank .map_or(self.dp_rank, |rank| rank.cast_unsigned()); for raw_event in batch.events { - let Some(placement_event) = self.normalizer.normalize( - raw_event, - seq, - WorkerWithDpRank::new(self.worker_id, effective_dp_rank), - ) else { + let worker = WorkerWithDpRank::new(self.worker_id, effective_dp_rank); + let Some(raw_event) = self.normalizer.preprocess(raw_event, worker) else { + continue; + }; + if !self.block_size_guard.admit(&raw_event, worker) { + continue; + } + let Some(placement_event) = self + .normalizer + .normalize_preprocessed(raw_event, seq, worker) + else { continue; }; let router_event = placement_event .into_router_event() .expect("local worker placement must convert to router event"); + if self.withhold_from_tree(&router_event, seq, "live") { + continue; + } self.indexer .apply_event_routed(router_event) .await @@ -516,6 +656,9 @@ async fn run_listener( cancel, socket, replay_socket, + record.recover_target(), + record.pending_evictions(), + record.audit_log(), watermark, ) .run() @@ -555,3 +698,258 @@ async fn connect_replay_socket( } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocols::compute_block_hash_for_seq; + use crate::services::indexer::backend::create_indexer; + use std::time::Duration; + + #[rstest::rstest] + #[case::sglang(false)] + #[case::vllm(true)] + #[tokio::test] + async fn replay_accepts_engine_framing(#[case] include_topic: bool) { + let context = zmq::Context::new(); + let router = context.socket(zmq::ROUTER).unwrap(); + router.set_linger(0).unwrap(); + router.set_rcvtimeo(5000).unwrap(); + router.set_sndtimeo(5000).unwrap(); + router.bind("tcp://127.0.0.1:*").unwrap(); + let endpoint = router.get_last_endpoint().unwrap().unwrap(); + let replay_socket = connect_dealer_socket(&endpoint).unwrap(); + let publisher = context.socket(zmq::PUB).unwrap(); + publisher.bind("tcp://127.0.0.1:*").unwrap(); + let live_socket = + connect_sub_socket(&publisher.get_last_endpoint().unwrap().unwrap()).unwrap(); + + let server = tokio::task::spawn_blocking(move || { + let request = router.recv_multipart(0).unwrap(); + assert_eq!(&request[1..], &[vec![], 0_u64.to_be_bytes().to_vec()]); + for seq in 0_u64..=2 { + let mut frames = vec![request[0].clone(), vec![]]; + if include_topic { + frames.push(b"kv-events".to_vec()); + } + if seq == 2 { + frames.extend([u64::MAX.to_be_bytes().to_vec(), vec![]]); + } else { + let event = ( + "BlockStored", + vec![seq + 100], + if seq == 0 { None } else { Some(100_u64) }, + vec![1_u32 + seq as u32; 4], + 4_usize, + Option::::None, + "GPU", + ); + let payload = rmp_serde::to_vec(&(0.0_f64, vec![event], Some(0_i32))).unwrap(); + frames.extend([seq.to_be_bytes().to_vec(), payload]); + } + router.send_multipart(frames, 0).unwrap(); + } + // Keep the socket open until the client consumes the queued replies. + assert_eq!(router.recv_multipart(0).unwrap()[1], b"done"); + }); + let indexer = create_indexer(4, 1); + let watermark = Arc::new(AtomicU64::new(WATERMARK_UNSET)); + let cancel = CancellationToken::new(); + let mut listener = ListenerLoop::new( + 1, + 0, + 4, + indexer.clone(), + cancel.clone(), + live_socket, + Some(replay_socket.clone()), + None, + None, + false, + watermark.clone(), + ); + let replayed = tokio::time::timeout(Duration::from_secs(5), listener.replay_gap(0, 2)) + .await + .expect("replay completion marker") + .unwrap(); + send_multipart(&replay_socket, vec![b"done".to_vec()]) + .await + .unwrap(); + server.await.unwrap(); + assert_eq!(replayed, 2); + assert_eq!(watermark.load(Ordering::Acquire), 1); + indexer.dump_events().await.expect("flush indexer"); + let hashes = compute_block_hash_for_seq(&[1, 1, 1, 1, 2, 2, 2, 2], 4, Default::default()); + let matches = indexer.find_matches(hashes).await.unwrap(); + assert_eq!(matches.scores.get(&WorkerWithDpRank::new(1, 0)), Some(&2)); + cancel.cancel(); + } +} + +#[cfg(test)] +mod deepinfra_tests { + use super::*; + use crate::protocols::{ + BlockHashOptions, LocalBlockHash, StorageTier, compute_block_hash_for_seq, + }; + use crate::services::indexer::backend::create_indexer; + use crate::services::indexer::evictions::PendingEvictions; + + const BLOCK_SIZE: u32 = 4; + + fn listener(pending: Option) -> (ListenerLoop, Indexer) { + // connect() is lazy, so no publisher is needed. + let live_socket = connect_sub_socket("tcp://127.0.0.1:1").unwrap(); + let indexer = create_indexer(BLOCK_SIZE, 1); + let listener = ListenerLoop::new( + 7, + 0, + BLOCK_SIZE, + indexer.clone(), + CancellationToken::new(), + live_socket, + None, + None, + pending, + false, + Arc::new(AtomicU64::new(WATERMARK_UNSET)), + ); + (listener, indexer) + } + + /// A vLLM batch of `BlockStored` events, positional msgpack as on the wire. + fn stored_batch(tokens: &[u32], block_size: usize, medium: &str) -> Vec { + let hashes: Vec = (0..tokens.len() / block_size) + .map(|i| 1000 + i as u64) + .collect(); + let event = ( + "BlockStored", + hashes, + Option::::None, + tokens.to_vec(), + block_size, + Option::::None, + medium, + ); + rmp_serde::to_vec(&(0.0_f64, vec![event], Some(0_i32))).unwrap() + } + + fn removed_batch(hashes: &[u64]) -> Vec { + let event = ("BlockRemoved", hashes.to_vec(), "GPU"); + rmp_serde::to_vec(&(0.0_f64, vec![event], Some(0_i32))).unwrap() + } + + fn probe(tokens: &[u32]) -> Vec { + compute_block_hash_for_seq(tokens, BLOCK_SIZE, BlockHashOptions::default()) + } + + /// vLLM's CPU offload connector publishes `medium: "CPU"`. The fork this + /// replaces filed those on the device tree; they must land in HostPinned + /// so `/query` reports them as `cpu`, not `gpu` (upstream #10368). + #[tokio::test] + async fn cpu_medium_events_are_indexed_on_the_host_pinned_tier() { + let (mut listener, indexer) = listener(None); + let tokens: Vec = (1..=8).collect(); + listener + .apply_live_batch(0, &stored_batch(&tokens, 4, "CPU")) + .await + .unwrap(); + indexer.dump_events().await.unwrap(); + + let tiered = indexer.find_tiered_matches(probe(&tokens)).await.unwrap(); + let worker = WorkerWithDpRank::new(7, 0); + assert_eq!(tiered.device.overlap_scores.scores.get(&worker), None); + assert_eq!( + tiered + .lower_tier + .get(&StorageTier::HostPinned) + .and_then(|m| m.hits.get(&worker)) + .copied(), + Some(2) + ); + } + + #[tokio::test] + async fn partial_prefix_stores_are_skipped_not_fatal() { + let (mut listener, indexer) = listener(None); + let tokens: Vec = (1..=2).collect(); + listener + .apply_live_batch(0, &stored_batch(&tokens, 2, "GPU")) + .await + .unwrap(); + assert!(indexer.dump_events().await.unwrap().is_empty()); + } + + /// `--keep-evictions`: an eviction is parked, not applied, so the block + /// stays matchable; a re-store cancels the parked eviction. + #[tokio::test] + async fn kept_evictions_are_parked_and_cancelled_by_a_restore() { + let pending: SharedPendingEvictions = + Arc::new(parking_lot::Mutex::new(PendingEvictions::default())); + let (mut listener, indexer) = listener(Some(pending.clone())); + let tokens: Vec = (1..=4).collect(); + listener + .apply_live_batch(0, &stored_batch(&tokens, 4, "GPU")) + .await + .unwrap(); + listener + .apply_live_batch(1, &removed_batch(&[1000])) + .await + .unwrap(); + indexer.dump_events().await.unwrap(); + + let scores = indexer.find_matches(probe(&tokens)).await.unwrap(); + assert_eq!(scores.scores.get(&WorkerWithDpRank::new(7, 0)), Some(&1)); + assert_eq!(pending.lock().len(), 1); + + listener + .apply_live_batch(2, &stored_batch(&tokens, 4, "GPU")) + .await + .unwrap(); + assert_eq!(pending.lock().len(), 0); + } + + /// A gap on a listener with a recover endpoint is filled from + /// `/kv_recover`: a TreeDump replaces the rank's state and moves the + /// watermark to its `last_event_id`. + #[tokio::test] + async fn gap_is_recovered_from_kv_recover_tree_dump() { + use crate::services::indexer::backend::test_util::store_event; + use crate::services::indexer::kv_recover::{KvRecoverClient, KvRecoverSettings}; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let tokens: Vec = (1..=8).collect(); + let hashes: Vec = probe(&tokens).into_iter().map(|h| h.0).collect(); + let body = serde_json::json!({ + "TreeDump": { + "events": [store_event(999, 3, 0, &[], &hashes, StorageTier::Device)], + "last_event_id": 41 + } + }) + .to_string(); + let server = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", server.local_addr().unwrap()); + tokio::spawn(async move { + let (mut stream, _) = server.accept().await.unwrap(); + let mut buf = [0u8; 4096]; + let _ = stream.read(&mut buf).await; + let reply = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + stream.write_all(reply.as_bytes()).await.unwrap(); + }); + + let (mut listener, indexer) = listener(None); + listener.recover = Some(RecoverTarget { + endpoint, + client: Arc::new(KvRecoverClient::new(KvRecoverSettings::default()).unwrap()), + }); + listener.handle_gap(42).await.unwrap(); + indexer.dump_events().await.unwrap(); + + assert_eq!(listener.watermark.load(Ordering::Acquire), 41); + let scores = indexer.find_matches(probe(&tokens)).await.unwrap(); + assert_eq!(scores.scores.get(&WorkerWithDpRank::new(7, 0)), Some(&2)); + } +} diff --git a/lib/kv-router/src/services/indexer/metrics.rs b/lib/kv-router/src/services/indexer/metrics.rs index ef3950b14bb6..c8f9f63abf3f 100644 --- a/lib/kv-router/src/services/indexer/metrics.rs +++ b/lib/kv-router/src/services/indexer/metrics.rs @@ -10,7 +10,8 @@ use std::time::Instant; use axum::{extract::MatchedPath, http::Request, middleware::Next, response::Response}; #[cfg(feature = "metrics")] use prometheus::{ - HistogramVec, IntCounterVec, IntGauge, IntGaugeVec, Opts, exponential_buckets, histogram_opts, + HistogramVec, IntCounter, IntCounterVec, IntGauge, IntGaugeVec, Opts, exponential_buckets, + histogram_opts, }; #[cfg(feature = "metrics")] @@ -39,6 +40,15 @@ pub struct StandaloneIndexerMetrics { pub models: IntGauge, pub workers: IntGauge, pub listeners: IntGaugeVec, + /// h24 mode: age of each matched block at query time, labeled + /// kind = since_stored | since_touched. The bucket CDF over queried + /// blocks is the hit-rate-vs-retention curve. + pub h24_age_seconds: HistogramVec, + pub h24_queried_blocks: IntCounter, + pub h24_matched_blocks: IntCounter, + pub h24_parent_miss: IntCounter, + pub h24_seen_blocks: IntGauge, + pub h24_attach_entries: IntGauge, } #[cfg(feature = "metrics")] @@ -86,6 +96,43 @@ static METRICS: LazyLock = LazyLock::new(|| Standalone &["status"], ) .expect("valid gauge"), + h24_age_seconds: HistogramVec::new( + histogram_opts!( + format!("{METRICS_PREFIX}_h24_age_seconds"), + "h24: age of each matched block at query time", + vec![ + 60.0, 300.0, 900.0, 1800.0, 3600.0, 10800.0, 21600.0, 43200.0, 86400.0, 129600.0, + 172800.0 + ] + ), + &["kind"], + ) + .expect("valid histogram"), + h24_queried_blocks: IntCounter::new( + format!("{METRICS_PREFIX}_h24_queried_blocks_total"), + "h24: blocks in incoming queries (denominator of the retention curve)", + ) + .expect("valid counter"), + h24_matched_blocks: IntCounter::new( + format!("{METRICS_PREFIX}_h24_matched_blocks_total"), + "h24: matched prefix blocks across queries", + ) + .expect("valid counter"), + h24_parent_miss: IntCounter::new( + format!("{METRICS_PREFIX}_h24_parent_miss_total"), + "h24: store events dropped because the parent hash was unknown", + ) + .expect("valid counter"), + h24_seen_blocks: IntGauge::new( + format!("{METRICS_PREFIX}_h24_seen_blocks"), + "h24: unique prefix blocks currently tracked", + ) + .expect("valid gauge"), + h24_attach_entries: IntGauge::new( + format!("{METRICS_PREFIX}_h24_attach_entries"), + "h24: engine-hash attach aliases currently tracked", + ) + .expect("valid gauge"), }); #[cfg(feature = "metrics")] @@ -97,6 +144,12 @@ pub fn register(registry: &prometheus::Registry) -> Result<(), prometheus::Error registry.register(Box::new(m.models.clone()))?; registry.register(Box::new(m.workers.clone()))?; registry.register(Box::new(m.listeners.clone()))?; + registry.register(Box::new(m.h24_age_seconds.clone()))?; + registry.register(Box::new(m.h24_queried_blocks.clone()))?; + registry.register(Box::new(m.h24_matched_blocks.clone()))?; + registry.register(Box::new(m.h24_parent_miss.clone()))?; + registry.register(Box::new(m.h24_seen_blocks.clone()))?; + registry.register(Box::new(m.h24_attach_entries.clone()))?; Ok(()) } @@ -144,6 +197,50 @@ pub fn set_worker_state(models: usize, workers: usize, listener_counts: [i64; 4] #[cfg(not(feature = "metrics"))] pub fn set_worker_state(_models: usize, _workers: usize, _listener_counts: [i64; 4]) {} +#[cfg(feature = "metrics")] +pub fn h24_observe_age(kind: &str, age_s: u32) { + METRICS + .h24_age_seconds + .with_label_values(&[kind]) + .observe(age_s as f64); +} + +#[cfg(not(feature = "metrics"))] +pub fn h24_observe_age(_kind: &str, _age_s: u32) {} + +#[cfg(feature = "metrics")] +pub fn h24_add_queried_blocks(n: usize) { + METRICS.h24_queried_blocks.inc_by(n as u64); +} + +#[cfg(not(feature = "metrics"))] +pub fn h24_add_queried_blocks(_n: usize) {} + +#[cfg(feature = "metrics")] +pub fn h24_add_matched_blocks(n: usize) { + METRICS.h24_matched_blocks.inc_by(n as u64); +} + +#[cfg(not(feature = "metrics"))] +pub fn h24_add_matched_blocks(_n: usize) {} + +#[cfg(feature = "metrics")] +pub fn h24_inc_parent_miss() { + METRICS.h24_parent_miss.inc(); +} + +#[cfg(not(feature = "metrics"))] +pub fn h24_inc_parent_miss() {} + +#[cfg(feature = "metrics")] +pub fn h24_set_sizes(seen: usize, attach: usize) { + METRICS.h24_seen_blocks.set(seen as i64); + METRICS.h24_attach_entries.set(attach as i64); +} + +#[cfg(not(feature = "metrics"))] +pub fn h24_set_sizes(_seen: usize, _attach: usize) {} + #[cfg(all(test, feature = "metrics"))] mod tests { use super::*; @@ -167,6 +264,11 @@ mod tests { .errors_total .with_label_values(&["/query", "4xx"]) .inc(); + h24_observe_age("since_stored", 60); + h24_add_queried_blocks(4); + h24_add_matched_blocks(2); + h24_inc_parent_miss(); + h24_set_sizes(10, 12); let encoder = prometheus::TextEncoder::new(); let mut buf = Vec::new(); @@ -180,5 +282,11 @@ mod tests { assert!(output.contains("dynamo_kvindexer_workers 2")); assert!(output.contains("dynamo_kvindexer_listeners{status=\"pending\"} 1")); assert!(output.contains("dynamo_kvindexer_listeners{status=\"active\"} 1")); + assert!(output.contains("dynamo_kvindexer_h24_age_seconds")); + assert!(output.contains("dynamo_kvindexer_h24_queried_blocks_total")); + assert!(output.contains("dynamo_kvindexer_h24_matched_blocks_total")); + assert!(output.contains("dynamo_kvindexer_h24_parent_miss_total")); + assert!(output.contains("dynamo_kvindexer_h24_seen_blocks 10")); + assert!(output.contains("dynamo_kvindexer_h24_attach_entries 12")); } } diff --git a/lib/kv-router/src/services/indexer/mod.rs b/lib/kv-router/src/services/indexer/mod.rs index 532366d882b1..7a6496f28eb9 100644 --- a/lib/kv-router/src/services/indexer/mod.rs +++ b/lib/kv-router/src/services/indexer/mod.rs @@ -29,10 +29,21 @@ //! Tier counts are CUMULATIVE through each tier's walk — see the doc on the //! response struct in [`server`] for the exact semantics. +mod audit; pub mod backend; +mod block_size; +#[cfg(test)] +mod deepapi_contract_tests; +pub mod discovery; +pub mod evictions; +pub mod h24; +pub mod kv_recover; pub mod listener; pub mod logging; pub mod metrics; +mod model_query; +#[cfg(feature = "kube-discovery")] +mod pod_watcher; pub mod recovery; pub mod registry; pub mod server; @@ -47,8 +58,11 @@ use tokio_util::sync::CancellationToken; use crate::config::min_initial_workers_from_env; use crate::services::common::zmq::validate_endpoint as validate_zmq_endpoint; use axum::http::header::HeaderName; +use discovery::KubeDiscoveryConfig; +use evictions::KeepEvictionsConfig; +use kv_recover::KvRecoverSettings; use logging::AccessLogSink; -use registry::WorkerRegistry; +use registry::{RegistryOptions, WorkerRegistry}; use server::{AppState, create_router}; pub struct IndexerConfig { @@ -62,6 +76,17 @@ pub struct IndexerConfig { pub access_log: Option, pub trace_id_header: HeaderName, pub access_log_local_time: bool, + /// Timeout and concurrency of `/kv_recover` gap-recovery downloads. + pub kv_recover: KvRecoverSettings, + /// When set, watch Kubernetes and register/deregister engine pods. + pub kube_discovery: Option, + /// When set, park evictions and replay aged ones under memory pressure. + pub keep_evictions: Option, + /// Emit `kv_audit` lines for every query and ingested event. + pub audit_log: bool, + /// Run every indexer as the flat h24 counterfactual with this retention + /// horizon (seconds). Exclusive with `keep_evictions`. + pub h24_horizon_s: Option, } pub(super) fn validate_listener_endpoints( @@ -77,6 +102,23 @@ pub(super) fn validate_listener_endpoints( Ok(()) } +/// A `recover_endpoint` is the base URL queried at `/kv_recover`: an +/// absolute `http`/`https` URL with a host. +pub(super) fn validate_recover_endpoint(endpoint: &str) -> anyhow::Result<()> { + let url = reqwest::Url::parse(endpoint) + .map_err(|error| anyhow::anyhow!("invalid recover endpoint `{endpoint}`: {error}"))?; + anyhow::ensure!( + matches!(url.scheme(), "http" | "https"), + "invalid recover endpoint `{endpoint}`: scheme must be http or https, got `{}`", + url.scheme() + ); + anyhow::ensure!( + url.host().is_some(), + "invalid recover endpoint `{endpoint}`: missing host" + ); + Ok(()) +} + pub fn parse_workers(s: &str) -> anyhow::Result> { let mut workers = Vec::new(); @@ -141,7 +183,20 @@ pub async fn run_server(config: IndexerConfig) -> anyhow::Result<()> { "Starting standalone KV cache indexer (HTTP-only mode)" ); - let mut state = AppState::new_with_cancel_token(config.threads, cancel_token.clone())?; + let options = RegistryOptions { + kv_recover: config.kv_recover, + keep_evictions: config.keep_evictions.is_some(), + audit_log: config.audit_log, + h24: config.h24_horizon_s.is_some(), + }; + if let Some(horizon_s) = config.h24_horizon_s { + anyhow::ensure!( + config.keep_evictions.is_none(), + "--h24 and --keep-evictions are mutually exclusive" + ); + anyhow::ensure!(horizon_s > 0, "--h24-horizon-secs must be positive"); + } + let mut state = AppState::new_with_cancel_token(config.threads, cancel_token.clone(), options)?; state.access_log_sink = match config.access_log { Some(ref path) => { let s = AccessLogSink::new( @@ -237,6 +292,49 @@ async fn run_common( wait_for_min_initial_workers(registry, &cancel_token).await?; registry.signal_ready(); + if config.audit_log { + tracing::info!( + target: "kv_audit", + "kv_audit logging enabled: queries and store/evict/clear events will be logged" + ); + } + + if let Some(horizon_s) = config.h24_horizon_s { + h24::spawn_expiry_loop(state.registry.clone(), horizon_s, cancel_token.clone()); + tracing::info!( + horizon_s, + "h24 mode enabled: flat counterfactual indexer, evictions ignored" + ); + } + + if let Some(keep) = config.keep_evictions { + keep.validate()?; + evictions::spawn_cleanup_loop( + state.registry.clone(), + keep.retention_s, + keep.memory_threshold, + cancel_token.clone(), + ); + tracing::info!( + retention_s = keep.retention_s, + memory_threshold = keep.memory_threshold, + "keep-evictions enabled: parking eviction events, sweeping under memory pressure" + ); + } + + if let Some(kube_config) = config.kube_discovery.clone() { + #[cfg(feature = "kube-discovery")] + pod_watcher::spawn_pod_watcher(kube_config, registry.clone(), cancel_token.clone()); + #[cfg(not(feature = "kube-discovery"))] + { + let _ = kube_config; + anyhow::bail!( + "pod discovery is configured but this binary was built without the \ + `kube-discovery` feature" + ); + } + } + let app = create_router(state); let listener = TcpListener::bind(("0.0.0.0", config.port)).await?; tracing::info!("HTTP server listening on 0.0.0.0:{}", config.port); @@ -277,6 +375,14 @@ mod tests { validate_zmq_endpoint("ipc:///tmp/dynamo.sock").unwrap(); } + #[test] + fn recover_endpoint_must_be_an_http_url_with_a_host() { + validate_recover_endpoint("http://10.0.0.1:5558").unwrap(); + validate_recover_endpoint("https://engine.local").unwrap(); + assert!(validate_recover_endpoint("tcp://10.0.0.1:5558").is_err()); + assert!(validate_recover_endpoint("10.0.0.1:5558").is_err()); + } + #[test] fn test_validate_zmq_endpoint_rejects_invalid_values() { assert!(validate_zmq_endpoint("tcp://host").is_err()); diff --git a/lib/kv-router/src/services/indexer/model_query.rs b/lib/kv-router/src/services/indexer/model_query.rs new file mode 100644 index 000000000000..3da5e046a37d --- /dev/null +++ b/lib/kv-router/src/services/indexer/model_query.rs @@ -0,0 +1,103 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Model-level queries: one `/query` or `/query_by_hash` answered from every +//! routing group's tree of the model, merged into one response. +//! +//! The pod watcher keys each pod's tree by its `engine_hash`, so a model's +//! routing groups are hash regimes, not caller-visible partitions. A probe +//! hashes in one regime and scores 0 against the others, so the merged +//! response converges on the pods that can reuse the cache. A request that +//! names a `routing_group` reads only that tree (upstream behavior); the +//! legacy `tenant_id` is ignored. + +use std::collections::HashMap; + +use axum::http::StatusCode; + +use crate::identity::RoutingPartitionId; +use crate::protocols::{LocalBlockHash, WorkerId}; + +use super::backend::Indexer; +use super::server::{ScoreResponse, build_score_response}; + +pub(super) struct ModelQueryOutcome { + pub status: StatusCode, + pub body: serde_json::Value, + /// Routing groups whose tree answered. + pub queried_groups: Vec, +} + +/// Query every tree in `trees` (non-empty; callers 404 before calling) and +/// merge. `hashes_for(block_size)` supplies each tree's probe hashes. A tree +/// that fails is skipped while any tree answers; 500 only if all fail. +pub(super) async fn query_model( + trees: Vec<(RoutingPartitionId, Indexer, u32)>, + mut hashes_for: impl FnMut(u32) -> Vec, + pod_names: &HashMap, +) -> ModelQueryOutcome { + let mut merged: Option = None; + let mut last_error: Option = None; + let mut queried_groups = Vec::with_capacity(trees.len()); + let mut h24_queried_blocks: Option = None; + + for (key, indexer, block_size) in trees { + let hashes = hashes_for(block_size); + if matches!(indexer, Indexer::H24(_)) { + h24_queried_blocks = Some(h24_queried_blocks.unwrap_or(0).max(hashes.len())); + } + match indexer.find_tiered_matches(hashes).await { + Ok(tiered) => { + let response = build_score_response(&tiered, block_size, pod_names); + match merged.as_mut() { + Some(acc) => merge_score_responses(acc, response), + None => merged = Some(response), + } + queried_groups.push(key.routing_group); + } + Err(error) => { + tracing::warn!( + model_name = %key.model_name, + routing_group = %key.routing_group, + %error, + "Routing-group tree query failed; continuing fan-out" + ); + last_error = Some(error.to_string()); + } + } + } + + // Once per request, not per tree: the h24 retention curve's denominator. + if let Some(blocks) = h24_queried_blocks { + super::metrics::h24_add_queried_blocks(blocks); + } + + let (status, body) = match merged { + Some(response) => (StatusCode::OK, serde_json::json!(response)), + None => ( + StatusCode::INTERNAL_SERVER_ERROR, + serde_json::json!({ + "error": last_error.unwrap_or_else(|| "no trees queried".to_string()) + }), + ), + }; + ModelQueryOutcome { + status, + body, + queried_groups, + } +} + +/// Worker sets are disjoint across trees (a worker registers into exactly one +/// routing group), so `scores` and `instances` merge by union; +/// `frequencies[i]` counts workers holding block `i`, so it sums element-wise. +fn merge_score_responses(acc: &mut ScoreResponse, other: ScoreResponse) { + acc.scores.extend(other.scores); + acc.instances.extend(other.instances); + if acc.frequencies.len() < other.frequencies.len() { + acc.frequencies.resize(other.frequencies.len(), 0); + } + for (slot, count) in acc.frequencies.iter_mut().zip(other.frequencies) { + *slot += count; + } +} diff --git a/lib/kv-router/src/services/indexer/pod_watcher.rs b/lib/kv-router/src/services/indexer/pod_watcher.rs new file mode 100644 index 000000000000..25792f7b94f8 --- /dev/null +++ b/lib/kv-router/src/services/indexer/pod_watcher.rs @@ -0,0 +1,566 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Kubernetes pod auto-discovery for the standalone indexer. +//! +//! Watches engine pods in a namespace (filtered by a label selector) and keeps +//! the [`WorkerRegistry`] in sync with them: +//! +//! - a pod that becomes **Ready** is registered (subscribe to its ZMQ KV-event +//! port), and +//! - a pod that is **deleted** (or stops being Ready) is deregistered. +//! +//! Each pod is registered under `routing_group = `, so +//! engine generations that hash KV blocks differently never share a tree +//! during a migration; `/query` fans out over a model's routing groups. +//! +//! This module is the *only* place that imports the Kubernetes client. It +//! depends on [`WorkerRegistry`] and calls nothing but +//! [`WorkerRegistry::register_with_extras`] / [`WorkerRegistry::deregister`], +//! so the indexer core never learns about Kubernetes. +//! +//! ## How resilience is handled +//! +//! We use the high-level `kube::runtime::watcher` + `reflector` machinery. The +//! reflector maintains an in-memory `Store` of the current pods and the watcher +//! transparently reconnects and **re-lists on `410 Gone`** (the "I fell behind +//! the resourceVersion history" case) with exponential backoff. On every change +//! we recompute the desired set of Ready pods from the Store and reconcile it +//! against what we have registered — so a missed delete during a disconnect is +//! corrected on the next re-list rather than leaking a subscription. +//! +//! ## Required RBAC +//! +//! The indexer's ServiceAccount needs `get`, `list`, and `watch` on `pods` in +//! the target namespace: +//! +//! ```yaml +//! apiVersion: rbac.authorization.k8s.io/v1 +//! kind: Role +//! metadata: +//! name: kv-indexer-pod-watcher +//! namespace: +//! rules: +//! - apiGroups: [""] +//! resources: ["pods"] +//! verbs: ["get", "list", "watch"] +//! --- +//! apiVersion: rbac.authorization.k8s.io/v1 +//! kind: RoleBinding +//! metadata: +//! name: kv-indexer-pod-watcher +//! namespace: +//! subjects: +//! - kind: ServiceAccount +//! name: +//! namespace: +//! roleRef: +//! kind: Role +//! name: kv-indexer-pod-watcher +//! apiGroup: rbac.authorization.k8s.io +//! ``` + +use std::collections::HashMap; +use std::sync::Arc; + +use futures::StreamExt; +use k8s_openapi::api::core::v1::Pod; +use kube::runtime::{WatchStreamExt, reflector, watcher}; +use kube::{Api, Client}; +use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; +use xxhash_rust::xxh3; + +use crate::protocols::WorkerId; + +use super::discovery::KubeDiscoveryConfig; +use super::registry::{ListenerExtras, WorkerRegistry}; + +/// Label the backend stamps on every engine pod with a hash of the exact +/// deployment (image, env, args). Pods with different values may hash KV +/// blocks differently, so each value gets its own tree via `routing_group`. +const ENGINE_HASH_LABEL: &str = "engine_hash"; + +/// A pod we have an active subscription for, keyed by pod name. +struct Subscribed { + instance_id: WorkerId, + ip: String, + routing_group: String, +} + +/// A Ready engine pod as extracted from the watch store. +#[derive(Debug, PartialEq)] +struct DiscoveredPod { + name: String, + ip: String, + /// Value of the pod's `engine_hash` label, when present. + engine_hash: Option, +} + +impl DiscoveredPod { + /// The routing group this pod's tree lives under: its `engine_hash`, so + /// hash-incompatible engine generations never share a tree during + /// migrations. Falls back to the configured group for unlabeled pods. + fn routing_group(&self, fallback: &str) -> String { + self.engine_hash + .clone() + .unwrap_or_else(|| fallback.to_string()) + } +} + +/// Spawn the pod watcher as a detached background task. +/// +/// Returns immediately; the watch loop runs until `cancel` fires. Any fatal +/// error (e.g. the Kubernetes client cannot be created) is logged and the task +/// exits — it never panics or blocks the caller. +pub fn spawn_pod_watcher( + config: KubeDiscoveryConfig, + registry: Arc, + cancel: CancellationToken, +) { + tokio::spawn(async move { + tracing::info!( + namespace = %config.namespace, + label_selector = %config.label_selector, + zmq_port = config.zmq_port, + dp_size = config.dp_size, + model_name = %config.model_name, + block_size = config.block_size, + "Starting Kubernetes pod watcher" + ); + if let Err(error) = run_pod_watcher(config, registry, cancel).await { + tracing::error!(error = %error, "Pod watcher exited with error"); + } + }); +} + +async fn run_pod_watcher( + config: KubeDiscoveryConfig, + registry: Arc, + cancel: CancellationToken, +) -> anyhow::Result<()> { + // Reads in-cluster config (mounted ServiceAccount token) when running as a + // pod, or the local kubeconfig otherwise. + let client = Client::try_default().await?; + let api: Api = Api::namespaced(client, &config.namespace); + let watcher_config = watcher::Config::default().labels(&config.label_selector); + + // The reflector keeps `reader` populated with the current pod set; the + // watcher feeds it and handles reconnects / 410-relist with backoff. + let (reader, writer) = reflector::store(); + let notify = Arc::new(Notify::new()); + + // Drive the watch stream on its own task. Each change (or error) pokes the + // reconcile loop below; the store itself is updated by the reflector. + let stream_notify = notify.clone(); + let stream_cancel = cancel.clone(); + tokio::spawn(async move { + let stream = reflector(writer, watcher(api, watcher_config)) + .default_backoff() + .touched_objects() + .for_each(move |event| { + if let Err(error) = event { + tracing::warn!(error = %error, "Pod watch stream error (auto-retrying)"); + } + stream_notify.notify_one(); + futures::future::ready(()) + }); + tokio::select! { + _ = stream => {} + _ = stream_cancel.cancelled() => {} + } + }); + + let mut subscribed: HashMap = HashMap::new(); + + // Initial sync. `wait_until_ready` returns once the reflector has committed + // its first full LIST of pods, so the Store is readable. We then reconcile + // once — this is what subscribes to engines that already exist at startup + // (and on every restart), rather than only reacting to later changes. + tokio::select! { + _ = cancel.cancelled() => { + tracing::info!("Pod watcher shutting down before initial sync"); + return Ok(()); + } + result = reader.wait_until_ready() => { + result?; + } + } + reconcile(&reader, &mut subscribed, ®istry, &config).await; + tracing::info!(engines = subscribed.len(), "Initial pod sync complete"); + + // Steady state: reconcile on every subsequent change. + loop { + tokio::select! { + _ = cancel.cancelled() => { + tracing::info!("Pod watcher shutting down"); + return Ok(()); + } + _ = notify.notified() => { + reconcile(&reader, &mut subscribed, ®istry, &config).await; + } + } + } +} + +/// Compare the current set of Ready pods against our subscriptions and apply +/// the difference: register newly-ready pods, deregister vanished ones. +async fn reconcile( + store: &reflector::Store, + subscribed: &mut HashMap, + registry: &WorkerRegistry, + config: &KubeDiscoveryConfig, +) { + // Desired state: every Ready, non-terminating pod, keyed by name. + let mut desired: HashMap = HashMap::new(); + for pod in store.state() { + if let Some(discovered) = ready_pod_endpoint(pod.as_ref()) { + desired.insert(discovered.name.clone(), discovered); + } + } + + // Unsubscribe first: pods that are gone, no longer Ready, or whose + // engine_hash changed (the same pod name must re-register under its new + // routing group, and register() rejects a worker that switches keys in place). + let vanished: Vec = subscribed + .iter() + .filter(|(name, entry)| { + desired + .get(*name) + .is_none_or(|pod| pod.routing_group(&config.routing_group) != entry.routing_group) + }) + .map(|(name, _)| name.clone()) + .collect(); + for name in vanished { + let entry = subscribed + .remove(&name) + .expect("name came from subscribed keys"); + match registry + .deregister(entry.instance_id, &config.model_name, &entry.routing_group) + .await + { + Ok(()) => tracing::info!( + pod = %name, + ip = %entry.ip, + instance_id = entry.instance_id, + routing_group = %entry.routing_group, + "Unsubscribed from engine pod" + ), + // Already gone from the registry — an idempotent no-op. + Err(error) => tracing::debug!( + pod = %name, + error = %error, + "Deregister was a no-op" + ), + } + } + + // Subscribe to pods that are newly Ready (or regrouped above). + for (name, pod) in &desired { + if subscribed.contains_key(name) { + continue; // already subscribed — idempotent no-op + } + + let instance_id = instance_id_for(name); + let ip = &pod.ip; + let routing_group = pod.routing_group(&config.routing_group); + + let mut failed: Option<(u32, anyhow::Error)> = None; + for rank in rank_endpoints(ip, config) { + if let Err(error) = registry + .register_with_extras( + instance_id, + rank.endpoint, + rank.dp_rank, + config.model_name.clone(), + routing_group.clone(), + config.block_size, + None, + ListenerExtras { + recover_endpoint: rank.recover_endpoint, + pod_name: Some(name.clone()), + }, + ) + .await + { + failed = Some((rank.dp_rank, error)); + break; + } + } + + match failed { + None => { + tracing::info!( + pod = %name, + ip = %ip, + instance_id, + routing_group = %routing_group, + dp_size = config.dp_size, + "Subscribed to engine pod" + ); + subscribed.insert( + name.clone(), + Subscribed { + instance_id, + ip: ip.clone(), + routing_group, + }, + ); + } + Some((dp_rank, error)) => { + // Not recorded, so it is retried on the next reconcile. Drop + // any ranks that did register: register() rejects a rank that + // is already present, so a partial instance would never + // converge on retry. + tracing::warn!( + pod = %name, + ip = %ip, + dp_rank, + error = %error, + "Failed to register engine pod; will retry" + ); + if let Err(error) = registry + .deregister(instance_id, &config.model_name, &routing_group) + .await + { + tracing::debug!( + pod = %name, + error = %error, + "Cleanup deregister was a no-op" + ); + } + } + } + } +} + +/// One listener to register for a pod: its data-parallel rank and the +/// per-rank event and recovery endpoints. +#[derive(Debug, PartialEq)] +struct RankEndpoint { + dp_rank: u32, + endpoint: String, + recover_endpoint: Option, +} + +/// Endpoints for every data-parallel rank of a pod. vLLM offsets both the ZMQ +/// KV-event port and the `/kv_recover` port by the rank, so rank `r` of a pod +/// at `ip` publishes on `zmq_port + r` and recovers on `recover_port + r`. +fn rank_endpoints(ip: &str, config: &KubeDiscoveryConfig) -> Vec { + (0..config.dp_size.max(1)) + .map(|dp_rank| RankEndpoint { + dp_rank, + endpoint: format!("tcp://{ip}:{}", u32::from(config.zmq_port) + dp_rank), + recover_endpoint: config + .recover_port + .map(|port| format!("http://{ip}:{}", u32::from(port) + dp_rank)), + }) + .collect() +} + +/// Derive a stable `WorkerId` from the pod name so the same pod always maps to +/// the same registry entry across MODIFIED events. +fn instance_id_for(pod_name: &str) -> WorkerId { + xxh3::xxh3_64(pod_name.as_bytes()) +} + +/// Return the pod's name, IP, and `engine_hash` label if it is Ready, has an +/// IP, and is not terminating; otherwise `None`. +fn ready_pod_endpoint(pod: &Pod) -> Option { + // A pod with a deletion timestamp is shutting down — treat it as gone. + if pod.metadata.deletion_timestamp.is_some() { + return None; + } + let name = pod.metadata.name.clone()?; + let status = pod.status.as_ref()?; + let ip = status.pod_ip.clone()?; + + let ready = status.conditions.as_ref().is_some_and(|conditions| { + conditions + .iter() + .any(|c| c.type_ == "Ready" && c.status == "True") + }); + if !ready { + return None; + } + + let engine_hash = pod + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get(ENGINE_HASH_LABEL)) + .cloned(); + + Some(DiscoveredPod { + name, + ip, + engine_hash, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use k8s_openapi::api::core::v1::{PodCondition, PodStatus}; + use k8s_openapi::apimachinery::pkg::apis::meta::v1::{ObjectMeta, Time}; + use k8s_openapi::chrono::Utc; + + fn pod(name: &str, ip: Option<&str>, ready: Option<&str>, terminating: bool) -> Pod { + pod_with_labels(name, ip, ready, terminating, &[]) + } + + fn pod_with_labels( + name: &str, + ip: Option<&str>, + ready: Option<&str>, + terminating: bool, + labels: &[(&str, &str)], + ) -> Pod { + Pod { + metadata: ObjectMeta { + name: Some(name.to_string()), + deletion_timestamp: terminating.then(|| Time(Utc::now())), + labels: (!labels.is_empty()).then(|| { + labels + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect() + }), + ..Default::default() + }, + status: Some(PodStatus { + pod_ip: ip.map(String::from), + conditions: ready.map(|status| { + vec![PodCondition { + type_: "Ready".to_string(), + status: status.to_string(), + ..Default::default() + }] + }), + ..Default::default() + }), + ..Default::default() + } + } + + #[test] + fn ready_pod_with_ip_is_subscribable() { + let p = pod("engine-abc", Some("10.0.0.1"), Some("True"), false); + assert_eq!( + ready_pod_endpoint(&p), + Some(DiscoveredPod { + name: "engine-abc".to_string(), + ip: "10.0.0.1".to_string(), + engine_hash: None, + }) + ); + } + + #[test] + fn engine_hash_label_becomes_routing_group() { + let p = pod_with_labels( + "engine-abc", + Some("10.0.0.1"), + Some("True"), + false, + &[("engine_hash", "975d72d83374a79b"), ("engine", "pytrtllm")], + ); + let discovered = ready_pod_endpoint(&p).unwrap(); + assert_eq!(discovered.engine_hash.as_deref(), Some("975d72d83374a79b")); + assert_eq!(discovered.routing_group("default"), "975d72d83374a79b"); + } + + #[test] + fn missing_engine_hash_falls_back_to_configured_routing_group() { + let p = pod_with_labels( + "engine-abc", + Some("10.0.0.1"), + Some("True"), + false, + &[("engine", "pytrtllm")], + ); + let discovered = ready_pod_endpoint(&p).unwrap(); + assert_eq!(discovered.engine_hash, None); + assert_eq!(discovered.routing_group("default"), "default"); + } + + #[test] + fn not_ready_pod_is_skipped() { + let p = pod("engine-abc", Some("10.0.0.1"), Some("False"), false); + assert_eq!(ready_pod_endpoint(&p), None); + } + + #[test] + fn ready_pod_without_ip_is_skipped() { + let p = pod("engine-abc", None, Some("True"), false); + assert_eq!(ready_pod_endpoint(&p), None); + } + + #[test] + fn terminating_pod_is_skipped_even_if_ready() { + let p = pod("engine-abc", Some("10.0.0.1"), Some("True"), true); + assert_eq!(ready_pod_endpoint(&p), None); + } + + #[test] + fn instance_id_is_stable_and_name_specific() { + assert_eq!(instance_id_for("engine-abc"), instance_id_for("engine-abc")); + assert_ne!(instance_id_for("engine-abc"), instance_id_for("engine-xyz")); + } + + fn discovery(dp_size: u32, recover_port: Option) -> KubeDiscoveryConfig { + KubeDiscoveryConfig { + namespace: "deepinfra".to_string(), + label_selector: "di/model_name=m".to_string(), + zmq_port: 5557, + recover_port, + dp_size, + model_name: "m".to_string(), + routing_group: "default".to_string(), + block_size: 128, + } + } + + #[test] + fn single_rank_engine_registers_rank_zero_on_base_ports() { + assert_eq!( + rank_endpoints("10.0.0.1", &discovery(1, Some(5559))), + vec![RankEndpoint { + dp_rank: 0, + endpoint: "tcp://10.0.0.1:5557".to_string(), + recover_endpoint: Some("http://10.0.0.1:5559".to_string()), + }] + ); + } + + #[test] + fn data_parallel_ranks_offset_event_and_recovery_ports() { + assert_eq!( + rank_endpoints("10.0.0.1", &discovery(2, Some(5559))), + vec![ + RankEndpoint { + dp_rank: 0, + endpoint: "tcp://10.0.0.1:5557".to_string(), + recover_endpoint: Some("http://10.0.0.1:5559".to_string()), + }, + RankEndpoint { + dp_rank: 1, + endpoint: "tcp://10.0.0.1:5558".to_string(), + recover_endpoint: Some("http://10.0.0.1:5560".to_string()), + }, + ] + ); + } + + #[test] + fn no_recover_port_means_no_recovery_endpoint_on_any_rank() { + let ranks = rank_endpoints("10.0.0.1", &discovery(2, None)); + assert_eq!(ranks.len(), 2); + assert!(ranks.iter().all(|r| r.recover_endpoint.is_none())); + } + + #[test] + fn zero_dp_size_still_registers_rank_zero() { + assert_eq!(rank_endpoints("10.0.0.1", &discovery(0, None)).len(), 1); + } +} diff --git a/lib/kv-router/src/services/indexer/registry.rs b/lib/kv-router/src/services/indexer/registry.rs index ce96d84b5221..2beec5144f4d 100644 --- a/lib/kv-router/src/services/indexer/registry.rs +++ b/lib/kv-router/src/services/indexer/registry.rs @@ -18,7 +18,9 @@ use crate::identity::RoutingPartitionId; use crate::indexer::KvIndexerMetrics; use crate::protocols::WorkerId; -use super::backend::{Indexer, create_indexer_with_metrics}; +use super::backend::{Indexer, create_h24_indexer, create_indexer_with_metrics}; +use super::evictions::PendingEvictions; +use super::kv_recover::{KvRecoverClient, KvRecoverSettings}; use super::listener::spawn_zmq_listener; pub struct IndexerEntry { @@ -102,6 +104,10 @@ pub struct ListenerInfo { #[derive(Debug, Clone, Serialize)] pub struct WorkerInfo { instance_id: WorkerId, + /// Pod name this worker was registered under, when the registrant knows + /// it. `instance_id` stays the canonical key. + #[serde(skip_serializing_if = "Option::is_none")] + pod_name: Option, source: WorkerSource, status: ListenerStatus, model_name: String, @@ -141,9 +147,50 @@ struct ListenerRuntime { generation: u64, } +/// DeepInfra registry behavior chosen at startup. +#[derive(Debug, Clone, Copy, Default)] +pub struct RegistryOptions { + pub kv_recover: KvRecoverSettings, + /// Park evictions instead of applying them (`--keep-evictions`). + pub keep_evictions: bool, + /// Emit `kv_audit` lines for queries and events (`--enable-logging`). + pub audit_log: bool, + /// Every indexer is the flat h24 counterfactual, one per model (`--h24`). + pub h24: bool, +} + +/// Where and how a listener fetches `/kv_recover` on a gap. +#[derive(Clone)] +pub struct RecoverTarget { + pub endpoint: String, + pub client: Arc, +} + +/// Optional registration fields beyond upstream's `register` arguments. +#[derive(Debug, Clone, Default)] +pub struct ListenerExtras { + /// HTTP base URL of the worker's `GET /kv_recover`. Takes precedence over + /// the ZMQ `replay_endpoint` for gap recovery when both are set. + pub recover_endpoint: Option, + /// Pod name, surfaced in `/workers` and in `/query` instances. + pub pod_name: Option, +} + +/// Eviction events a `--keep-evictions` listener parks instead of applying. +pub type SharedPendingEvictions = Arc>; + +/// DeepInfra per-listener ingest behavior, fixed at registration. +struct ListenerIngest { + recover: Option, + /// `Some` only under `--keep-evictions`; drained by the evictions sweep. + pending_evictions: Option, + audit_log: bool, +} + pub struct ListenerRecord { endpoint: String, replay_endpoint: Option, + ingest: ListenerIngest, block_size: u32, indexer: Indexer, watermark: Arc, @@ -154,6 +201,7 @@ impl ListenerRecord { fn new( endpoint: String, replay_endpoint: Option, + ingest: ListenerIngest, block_size: u32, indexer: Indexer, watermark: Arc, @@ -161,6 +209,7 @@ impl ListenerRecord { Self { endpoint, replay_endpoint, + ingest, block_size, indexer, watermark, @@ -181,6 +230,18 @@ impl ListenerRecord { self.replay_endpoint.as_deref() } + pub(super) fn recover_target(&self) -> Option { + self.ingest.recover.clone() + } + + pub(super) fn pending_evictions(&self) -> Option { + self.ingest.pending_evictions.clone() + } + + pub(super) fn audit_log(&self) -> bool { + self.ingest.audit_log + } + pub(super) fn block_size(&self) -> u32 { self.block_size } @@ -305,12 +366,15 @@ impl ListenerRecord { pub struct WorkerEntry { key: RoutingPartitionId, + pod_name: Option, listeners: HashMap>, } pub struct WorkerRegistry { workers: DashMap, indexers: DashMap, + // Serialize indexer claims through worker publication with empty-indexer removal. + indexer_lifecycle: tokio::sync::Mutex<()>, peers: DashMap, watermarks: DashMap<(WorkerId, u32), Arc>, num_threads: usize, @@ -318,6 +382,11 @@ pub struct WorkerRegistry { ready_tx: watch::Sender, ready_rx: watch::Receiver, root_cancel_token: CancellationToken, + retain_empty_indexers: bool, + kv_recover: Arc, + keep_evictions: bool, + audit_log: bool, + h24: bool, } impl WorkerRegistry { @@ -357,6 +426,7 @@ impl WorkerRegistry { Self { workers: DashMap::new(), indexers: DashMap::new(), + indexer_lifecycle: tokio::sync::Mutex::new(()), peers: DashMap::new(), watermarks: DashMap::new(), num_threads, @@ -364,9 +434,62 @@ impl WorkerRegistry { ready_tx, ready_rx, root_cancel_token, + retain_empty_indexers: false, + kv_recover: Arc::new( + KvRecoverClient::new(KvRecoverSettings::default()) + .expect("default kv_recover client builds"), + ), + keep_evictions: false, + audit_log: false, + h24: false, + } + } + + /// Apply startup options; call before any registration. + pub fn with_options(mut self, options: RegistryOptions) -> Result { + self.kv_recover = Arc::new(KvRecoverClient::new(options.kv_recover)?); + self.keep_evictions = options.keep_evictions; + self.audit_log = options.audit_log; + self.h24 = options.h24; + Ok(self) + } + + fn new_indexer(&self, block_size: u32) -> Indexer { + if self.h24 { + create_h24_indexer() + } else { + create_indexer_with_metrics(block_size, self.num_threads, self.indexer_metrics.clone()) } } + /// h24 is worker-less and regime-agnostic ("stored by anyone"), so every + /// routing group of a model shares one flat map. Per-group maps would count + /// a block matched in N regimes N times against a once-per-request + /// denominator (live incident: 120% hit rate). Requires one block size + /// across the model's regimes. + fn partition_key(&self, model_name: String, routing_group: String) -> RoutingPartitionId { + if self.h24 { + RoutingPartitionId::new(model_name, "h24") + } else { + RoutingPartitionId::new(model_name, routing_group) + } + } + + pub fn h24_enabled(&self) -> bool { + self.h24 + } + + pub fn audit_log_enabled(&self) -> bool { + self.audit_log + } + + #[cfg(feature = "standalone-selection")] + pub(crate) fn with_retained_indexers(mut self) -> Self { + // Selection entries and their schedulers retain these indexers across worker updates. + self.retain_empty_indexers = true; + self + } + pub fn signal_ready(&self) { let _ = self.ready_tx.send(true); } @@ -413,7 +536,34 @@ impl WorkerRegistry { block_size: u32, replay_endpoint: Option, ) -> Result<()> { - let key = RoutingPartitionId::new(model_name, routing_group); + self.register_with_extras( + instance_id, + endpoint, + dp_rank, + model_name, + routing_group, + block_size, + replay_endpoint, + ListenerExtras::default(), + ) + .await + } + + /// [`Self::register`] plus the DeepInfra-only registration fields. + #[expect(clippy::too_many_arguments)] + pub async fn register_with_extras( + &self, + instance_id: WorkerId, + endpoint: String, + dp_rank: u32, + model_name: String, + routing_group: String, + block_size: u32, + replay_endpoint: Option, + extras: ListenerExtras, + ) -> Result<()> { + let key = self.partition_key(model_name.to_string(), routing_group.to_string()); + let registration = self.indexer_lifecycle.lock().await; if let Some(entry) = self.workers.get(&instance_id) { if entry.key != key { @@ -437,11 +587,7 @@ impl WorkerRegistry { "Creating new indexer" ); IndexerEntry { - indexer: create_indexer_with_metrics( - block_size, - self.num_threads, - self.indexer_metrics.clone(), - ), + indexer: self.new_indexer(block_size), block_size, } }); @@ -466,9 +612,21 @@ impl WorkerRegistry { .or_insert_with(|| Arc::new(AtomicU64::new(u64::MAX))) .clone(); + let recover = extras.recover_endpoint.map(|endpoint| RecoverTarget { + endpoint, + client: self.kv_recover.clone(), + }); + let pending_evictions = self + .keep_evictions + .then(|| Arc::new(Mutex::new(PendingEvictions::default()))); let record = Arc::new(ListenerRecord::new( endpoint, replay_endpoint, + ListenerIngest { + recover, + pending_evictions, + audit_log: self.audit_log, + }, bs, indexer, watermark, @@ -481,11 +639,16 @@ impl WorkerRegistry { .entry(instance_id) .or_insert_with(|| WorkerEntry { key: key.clone(), + pod_name: None, listeners: HashMap::new(), }); + if extras.pod_name.is_some() { + entry.pod_name = extras.pod_name; + } entry.listeners.insert(dp_rank, record.clone()); } + drop(registration); self.spawn_listener(instance_id, dp_rank, attempt, record); Ok(()) } @@ -496,7 +659,7 @@ impl WorkerRegistry { model_name: &str, routing_group: &str, ) -> Result<()> { - let key = RoutingPartitionId::new(model_name, routing_group); + let key = self.partition_key(model_name.to_string(), routing_group.to_string()); if let Some(entry) = self.workers.get(&instance_id) { if entry.key != key { @@ -521,10 +684,12 @@ impl WorkerRegistry { } } - if let Some(ie) = self.indexers.get(&key) { - ie.indexer.remove_worker(instance_id).await; + // Registration needs this map's write lock while removal can wait on indexer queues. + let indexer = self.indexers.get(&key).map(|entry| entry.indexer.clone()); + if let Some(indexer) = indexer { + indexer.remove_worker(instance_id).await; } - self.maybe_remove_indexer(&key); + self.maybe_remove_indexer(&key).await; Ok(()) } @@ -535,7 +700,7 @@ impl WorkerRegistry { model_name: &str, routing_group: &str, ) -> Result<()> { - let key = RoutingPartitionId::new(model_name, routing_group); + let key = self.partition_key(model_name.to_string(), routing_group.to_string()); let (record, remove_worker) = { let mut entry = self @@ -569,13 +734,17 @@ impl WorkerRegistry { .remove_if(&instance_id, |_, entry| entry.listeners.is_empty()) .is_some(); if actually_removed { - if let Some(ie) = self.indexers.get(&key) { - ie.indexer.remove_worker(instance_id).await; + let indexer = self.indexers.get(&key).map(|entry| entry.indexer.clone()); + if let Some(indexer) = indexer { + indexer.remove_worker(instance_id).await; } - self.maybe_remove_indexer(&key); + self.maybe_remove_indexer(&key).await; + } + } else { + let indexer = self.indexers.get(&key).map(|entry| entry.indexer.clone()); + if let Some(indexer) = indexer { + indexer.remove_worker_dp_rank(instance_id, dp_rank).await; } - } else if let Some(ie) = self.indexers.get(&key) { - ie.indexer.remove_worker_dp_rank(instance_id, dp_rank).await; } Ok(()) @@ -610,10 +779,11 @@ impl WorkerRegistry { } } - if let Some(ie) = self.indexers.get(&key) { - ie.indexer.remove_worker(instance_id).await; + let indexer = self.indexers.get(&key).map(|entry| entry.indexer.clone()); + if let Some(indexer) = indexer { + indexer.remove_worker(instance_id).await; } - self.maybe_remove_indexer(&key); + self.maybe_remove_indexer(&key).await; Ok(()) } @@ -715,6 +885,7 @@ impl WorkerRegistry { let status = ListenerStatus::aggregate(listeners.values().map(|info| info.status)); Some(WorkerInfo { instance_id: *entry.key(), + pod_name: worker.pod_name.clone(), source: WorkerSource::Zmq, status, model_name: key.model_name.clone(), @@ -734,7 +905,19 @@ impl WorkerRegistry { self.indexers.get(key) } + /// Known pod names by worker id, for enriching query responses. + pub fn pod_names(&self) -> HashMap { + self.workers + .iter() + .filter_map(|entry| { + let pod_name = entry.value().pod_name.clone()?; + Some((*entry.key(), pod_name)) + }) + .collect() + } + pub fn get_or_create_indexer(&self, key: RoutingPartitionId, block_size: u32) -> Indexer { + let key = self.partition_key(key.model_name, key.routing_group); let entry = self.indexers.entry(key.clone()).or_insert_with(|| { tracing::info!( model_name = %key.model_name, @@ -743,11 +926,7 @@ impl WorkerRegistry { "Creating indexer from recovery dump" ); IndexerEntry { - indexer: create_indexer_with_metrics( - block_size, - self.num_threads, - self.indexer_metrics.clone(), - ), + indexer: self.new_indexer(block_size), block_size, } }); @@ -763,6 +942,58 @@ impl WorkerRegistry { entry.indexer.clone() } + /// Every `--keep-evictions` listener as `(worker_id, dp_rank, buffer, + /// indexer)`, for the evictions sweep. + pub(super) fn listener_records(&self) -> Vec<(WorkerId, u32, SharedPendingEvictions, Indexer)> { + self.workers + .iter() + .flat_map(|entry| { + let worker_id = *entry.key(); + entry + .value() + .listeners + .iter() + .filter_map(|(dp_rank, record)| { + let buffer = record.pending_evictions()?; + Some((worker_id, *dp_rank, buffer, record.indexer())) + }) + .collect::>() + }) + .collect() + } + + /// The trees a query for `model_name` reads: the named routing group's, + /// or every group's when none is named. + pub fn query_trees( + &self, + model_name: &str, + routing_group: Option<&str>, + ) -> Vec<(RoutingPartitionId, Indexer, u32)> { + let Some(routing_group) = routing_group else { + return self.indexers_for_model(model_name); + }; + let key = self.partition_key(model_name.to_string(), routing_group.to_string()); + self.indexers + .get(&key) + .map(|entry| vec![(key.clone(), entry.indexer.clone(), entry.block_size)]) + .unwrap_or_default() + } + + /// Every routing group's tree of `model_name`, for model-level queries. + pub fn indexers_for_model(&self, model_name: &str) -> Vec<(RoutingPartitionId, Indexer, u32)> { + self.indexers + .iter() + .filter(|entry| entry.key().model_name == model_name) + .map(|entry| { + ( + entry.key().clone(), + entry.value().indexer.clone(), + entry.value().block_size, + ) + }) + .collect() + } + pub fn all_indexers_with_block_size(&self) -> Vec<(RoutingPartitionId, Indexer, u32)> { self.indexers .iter() @@ -807,7 +1038,12 @@ impl WorkerRegistry { ); } - fn maybe_remove_indexer(&self, key: &RoutingPartitionId) { + async fn maybe_remove_indexer(&self, key: &RoutingPartitionId) { + if self.retain_empty_indexers { + return; + } + + let _lifecycle = self.indexer_lifecycle.lock().await; if self.workers.iter().any(|entry| entry.value().key == *key) { return; } @@ -819,41 +1055,184 @@ impl WorkerRegistry { #[cfg(test)] mod tests { use super::*; + use crate::protocols::{LocalBlockHash, StorageTier, WorkerWithDpRank}; + use crate::services::indexer::backend::test_util::store_event; + use std::future::Future; use std::sync::atomic::Ordering; + use std::task::{Context, Poll, Waker}; + use std::time::Duration; fn test_registry() -> WorkerRegistry { WorkerRegistry::new(1) } + #[rstest::rstest] + #[case("worker")] + #[case("all_groups")] + #[case("last_rank")] #[tokio::test] - async fn deregister_removes_watermark() { + async fn registration_progresses_while_removal_is_backpressured(#[case] removal: &str) { let registry = test_registry(); - registry.signal_ready(); - + let key = RoutingPartitionId::new("test-model", "default"); registry .register( 1, - "tcp://127.0.0.1:15557".to_string(), + "tcp://127.0.0.1:15557".into(), 0, - "test-model".to_string(), - "default".to_string(), + "test-model".into(), + "default".into(), 1, None, ) .await .unwrap(); + let indexer = registry.indexers.get(&key).unwrap().indexer.clone(); + let Indexer::Single { primary, .. } = indexer else { + unreachable!(); + }; + let sender = primary.remove_worker_sender(); + // Reserve the entire queue so removal must yield without blocking the executor. + let permits = sender.reserve_many(sender.max_capacity()).await.unwrap(); + let removal = async { + match removal { + "worker" => registry.deregister(1, "test-model", "default").await, + "all_groups" => { + registry + .deregister_all_routing_groups(1, "test-model") + .await + } + "last_rank" => { + registry + .deregister_dp_rank(1, 0, "test-model", "default") + .await + } + _ => unreachable!(), + } + }; + tokio::pin!(removal); + assert!(matches!( + removal + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())), + Poll::Pending + )); + // A nonblocking probe makes the regression fail instead of deadlocking the test. + assert!(matches!( + registry.indexers.try_get_mut(&key), + dashmap::try_result::TryResult::Present(_) + )); + registry + .register( + 2, + "tcp://127.0.0.1:15558".into(), + 0, + "test-model".into(), + "default".into(), + 1, + None, + ) + .await + .unwrap(); + drop(permits); + tokio::time::timeout(Duration::from_secs(5), removal) + .await + .expect("removal should finish after queue capacity is released") + .unwrap(); + assert!(registry.workers.contains_key(&2)); + assert!(registry.indexers.contains_key(&key)); + registry.root_cancel_token.cancel(); + } - assert!(registry.watermarks.contains_key(&(1, 0))); - + #[rstest::rstest] + #[case("worker")] + #[case("all_groups")] + #[case("last_rank")] + #[tokio::test] + async fn empty_indexer_removal_waits_for_registration(#[case] removal: &str) { + let registry = test_registry(); + let key = RoutingPartitionId::new("test-model", "default"); registry - .deregister(1, "test-model", "default") + .register( + 1, + "tcp://127.0.0.1:15557".into(), + 0, + "test-model".into(), + "default".into(), + 1, + None, + ) .await .unwrap(); - assert!( - !registry.watermarks.contains_key(&(1, 0)), - "watermark should be removed after deregister" + // Queue registration ahead of pruning, without timing or thread scheduling assumptions. + let lifecycle = registry.indexer_lifecycle.lock().await; + let registration = registry.register( + 2, + "tcp://127.0.0.1:15558".into(), + 0, + "test-model".into(), + "default".into(), + 1, + None, ); + tokio::pin!(registration); + assert!(matches!( + registration + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())), + Poll::Pending + )); + + let removal = async { + match removal { + "worker" => registry.deregister(1, "test-model", "default").await, + "all_groups" => { + registry + .deregister_all_routing_groups(1, "test-model") + .await + } + "last_rank" => { + registry + .deregister_dp_rank(1, 0, "test-model", "default") + .await + } + _ => unreachable!(), + } + }; + tokio::pin!(removal); + assert!(matches!( + removal + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())), + Poll::Pending + )); + assert!(registry.workers.is_empty()); + assert!(!registry.watermarks.contains_key(&(1, 0))); + assert!(registry.get_indexer(&key).is_some()); + drop(lifecycle); + registration.await.unwrap(); + removal.await.unwrap(); + + // Events written by the new listener must remain visible through the query registry. + let listener_indexer = registry.workers.get(&2).unwrap().listeners[&0] + .indexer + .clone(); + listener_indexer + .apply_event_routed(store_event(2, 0, 0, &[], &[11], StorageTier::Device)) + .await + .unwrap(); + listener_indexer.dump_events().await.unwrap(); + let query_indexer = registry + .get_indexer(&key) + .expect("registered indexer") + .indexer + .clone(); + let scores = query_indexer + .find_matches(vec![LocalBlockHash(11)]) + .await + .unwrap(); + assert_eq!(scores.scores.get(&WorkerWithDpRank::new(2, 0)), Some(&1)); + registry.root_cancel_token.cancel(); } #[tokio::test] @@ -969,15 +1348,23 @@ mod tests { .unwrap(); // Simulate that the listener advanced the watermark. - if let Some(wm) = registry.watermarks.get(&(1, 0)) { - wm.store(42, Ordering::Release); - } + registry + .watermarks + .get(&(1, 0)) + .unwrap() + .store(42, Ordering::Release); registry .deregister(1, "test-model", "default") .await .unwrap(); + assert!( + registry + .get_indexer(&RoutingPartitionId::new("test-model", "default")) + .is_none() + ); + registry .register( 1, @@ -1002,37 +1389,6 @@ mod tests { ); } - #[tokio::test] - async fn deregister_all_routing_groups_removes_watermarks() { - let registry = test_registry(); - registry.signal_ready(); - - registry - .register( - 1, - "tcp://127.0.0.1:15562".to_string(), - 0, - "test-model".to_string(), - "default".to_string(), - 1, - None, - ) - .await - .unwrap(); - - assert!(registry.watermarks.contains_key(&(1, 0))); - - registry - .deregister_all_routing_groups(1, "test-model") - .await - .unwrap(); - - assert!( - !registry.watermarks.contains_key(&(1, 0)), - "watermark should be removed after deregister_all_routing_groups" - ); - } - // ── list_filtered tests ─────────────────────────────────────────────────── #[tokio::test] @@ -1119,6 +1475,7 @@ mod tests { 99u64, WorkerEntry { key, + pod_name: None, listeners: HashMap::new(), }, ); diff --git a/lib/kv-router/src/services/indexer/server.rs b/lib/kv-router/src/services/indexer/server.rs index b07ee2f07405..67bb0e284d57 100644 --- a/lib/kv-router/src/services/indexer/server.rs +++ b/lib/kv-router/src/services/indexer/server.rs @@ -14,15 +14,16 @@ use prometheus::Encoder; use serde::{Deserialize, Serialize}; use tokio_util::sync::CancellationToken; -use crate::identity::{RoutingPartitionId, default_routing_group}; +use crate::identity::default_routing_group; #[cfg(feature = "metrics")] use crate::indexer::KvIndexerMetrics; use crate::indexer::TieredMatchDetails; use crate::protocols::{BlockHashOptions, LocalBlockHash, WorkerId, compute_block_hash_for_seq}; use crate::services::overlap::{MooncakeOverlapSummary, build_mooncake_overlap_summaries}; -use super::backend::Indexer; -use super::registry::{ListenerControlError, WorkerRegistry}; +use super::audit; +use super::model_query::{ModelQueryOutcome, query_model}; +use super::registry::{ListenerControlError, ListenerExtras, RegistryOptions, WorkerRegistry}; /// We need to fit one million tokens as JSON text, this should do it. const QUERY_REQUEST_BODY_LIMIT_BYTES: usize = 8 * 1024 * 1024; @@ -46,12 +47,17 @@ pub struct AppState { impl AppState { pub fn new(indexer_threads: usize) -> anyhow::Result { - Self::new_with_cancel_token(indexer_threads, CancellationToken::new()) + Self::new_with_cancel_token( + indexer_threads, + CancellationToken::new(), + RegistryOptions::default(), + ) } pub(super) fn new_with_cancel_token( indexer_threads: usize, root_cancel_token: CancellationToken, + options: RegistryOptions, ) -> anyhow::Result { #[cfg(feature = "metrics")] { @@ -59,11 +65,14 @@ impl AppState { super::metrics::register(&prom_registry)?; let indexer_metrics = KvIndexerMetrics::new_registered(&prom_registry)?; Ok(Self { - registry: Arc::new(WorkerRegistry::new_with_indexer_metrics_and_cancel_token( - indexer_threads, - indexer_metrics, - root_cancel_token, - )), + registry: Arc::new( + WorkerRegistry::new_with_indexer_metrics_and_cancel_token( + indexer_threads, + indexer_metrics, + root_cancel_token, + ) + .with_options(options)?, + ), access_log_sink: None, prom_registry, }) @@ -71,10 +80,10 @@ impl AppState { #[cfg(not(feature = "metrics"))] Ok(Self { - registry: Arc::new(WorkerRegistry::new_with_cancel_token( - indexer_threads, - root_cancel_token, - )), + registry: Arc::new( + WorkerRegistry::new_with_cancel_token(indexer_threads, root_cancel_token) + .with_options(options)?, + ), access_log_sink: None, }) } @@ -94,6 +103,14 @@ struct RegisterRequest { dp_rank: Option, #[serde(default)] replay_endpoint: Option, + /// HTTP base URL of the worker's `GET /kv_recover`, e.g. + /// `http://10.0.0.1:5558`. Used for gap recovery instead of + /// `replay_endpoint` when set. + #[serde(default)] + recover_endpoint: Option, + /// Optional pod name, surfaced in `/workers` and `/query` instances. + #[serde(default)] + pod_name: Option, /// Optional per-tenant salt (Mooncake RFC #1403 `additionalsalt`). /// Currently accepted but not yet mixed into hashes — engines apply /// their own salt internally. Plumbed for forward compatibility. @@ -117,8 +134,10 @@ struct UnregisterRequest { struct QueryRequest { token_ids: Vec, model_name: String, - #[serde(default = "default_routing_group")] - routing_group: String, + /// Queries only this group's tree when given; absent, the query fans out + /// over every routing group of `model_name` (see `model_query`). + #[serde(default)] + routing_group: Option, #[serde(default = "default_routing_group", rename = "tenant_id")] _tenant_id: String, #[serde(default)] @@ -132,8 +151,10 @@ struct QueryRequest { struct QueryByHashRequest { block_hashes: Vec, model_name: String, - #[serde(default = "default_routing_group")] - routing_group: String, + /// Queries only this group's tree when given; absent, the query fans out + /// over every routing group of `model_name` (see `model_query`). + #[serde(default)] + routing_group: Option, #[serde(default = "default_routing_group", rename = "tenant_id")] _tenant_id: String, /// Invalid for `/query_by_hash`. Callers must precompute `block_hashes` with the intended @@ -150,11 +171,21 @@ struct QueryByHashRequest { /// RFC #1403 (kvcache-ai/Mooncake#1403): /// `{instance_id: {longest_matched, gpu, dp: {rank: count}, cpu, disk}}`. #[derive(Serialize)] -struct ScoreResponse { - scores: HashMap>, - frequencies: Vec, +pub(super) struct ScoreResponse { + pub(super) scores: HashMap>, + pub(super) frequencies: Vec, /// Per-instance tier breakdown (Mooncake RFC #1403 alignment). - instances: HashMap, + pub(super) instances: HashMap, +} + +/// One `instances` entry: the shared Mooncake summary plus the worker's pod +/// name when known (deepapi maps matches back to shards by pod name). +#[derive(Serialize)] +pub(super) struct InstanceMatch { + #[serde(skip_serializing_if = "Option::is_none")] + pod_name: Option, + #[serde(flatten)] + summary: MooncakeOverlapSummary, } async fn register( @@ -163,7 +194,11 @@ async fn register( ) -> Response { let model = req.model_name.clone(); if let Err(error) = - super::validate_listener_endpoints(&req.endpoint, req.replay_endpoint.as_deref()) + super::validate_listener_endpoints(&req.endpoint, req.replay_endpoint.as_deref()).and( + req.recover_endpoint + .as_deref() + .map_or(Ok(()), super::validate_recover_endpoint), + ) { let mut resp = ( StatusCode::BAD_REQUEST, @@ -176,7 +211,7 @@ async fn register( let resp = match state .registry - .register( + .register_with_extras( req.instance_id, req.endpoint, req.dp_rank.unwrap_or(0), @@ -184,6 +219,10 @@ async fn register( req.routing_group, req.block_size, req.replay_endpoint, + ListenerExtras { + recover_endpoint: req.recover_endpoint, + pod_name: req.pod_name, + }, ) .await { @@ -268,7 +307,11 @@ async fn list_workers( /// (Mooncake RFC #1403) shapes from a tiered match result. /// /// All token counts are scaled from blocks → tokens via `block_size`. -fn build_score_response(tiered: &TieredMatchDetails, block_size: u32) -> ScoreResponse { +pub(super) fn build_score_response( + tiered: &TieredMatchDetails, + block_size: u32, + pod_names: &HashMap, +) -> ScoreResponse { // Flat fields (unchanged) come from the device-tier overlap. let device = &tiered.device.overlap_scores; @@ -282,7 +325,13 @@ fn build_score_response(tiered: &TieredMatchDetails, block_size: u32) -> ScoreRe let instances = build_mooncake_overlap_summaries(tiered, block_size, []) .into_iter() - .map(|(worker_id, summary)| (worker_id.to_string(), summary)) + .map(|(worker_id, summary)| { + let instance = InstanceMatch { + pod_name: pod_names.get(&worker_id).cloned(), + summary, + }; + (worker_id.to_string(), instance) + }) .collect(); ScoreResponse { @@ -292,57 +341,67 @@ fn build_score_response(tiered: &TieredMatchDetails, block_size: u32) -> ScoreRe } } -/// Run a tiered query and serialize the result, returning the appropriate -/// HTTP status. Shared between `/query` and `/query_by_hash`. -async fn run_tiered_query( - indexer: &Indexer, - block_hashes: Vec, - block_size: u32, -) -> (StatusCode, Json) { - match indexer.find_tiered_matches(block_hashes).await { - Ok(tiered) => ( - StatusCode::OK, - Json(serde_json::json!(build_score_response(&tiered, block_size))), - ), - Err(e) => ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": e.to_string()})), - ), - } -} - async fn query(State(state): State>, Json(req): Json) -> Response { let model = req.model_name.clone(); - let key = RoutingPartitionId::new(req.model_name, req.routing_group); - let Some(ie) = state.registry.get_indexer(&key) else { - let mut resp = ( - StatusCode::NOT_FOUND, - Json(serde_json::json!({ - "error": format!( - "no indexer for model={} routing_group={}", - key.model_name, key.routing_group - ) - })), - ) - .into_response(); - resp.extensions_mut().insert(AccessLogModel(model)); - return resp; - }; - let block_size = ie.block_size; - let indexer = ie.indexer.clone(); - drop(ie); - - let block_hashes = compute_block_hash_for_seq( - &req.token_ids, - block_size, - BlockHashOptions { - lora_name: req.lora_name.as_deref(), - cache_namespace: req.cache_salt.as_deref(), - ..Default::default() + let trees = state + .registry + .query_trees(&req.model_name, req.routing_group.as_deref()); + if trees.is_empty() { + return model_not_found(model); + } + + // Trees of one model normally share a block size; hash once per size. + let mut hashes_by_block_size: HashMap> = HashMap::new(); + let outcome = query_model( + trees, + |block_size| { + hashes_by_block_size + .entry(block_size) + .or_insert_with(|| { + compute_block_hash_for_seq( + &req.token_ids, + block_size, + BlockHashOptions { + lora_name: req.lora_name.as_deref(), + cache_namespace: req.cache_salt.as_deref(), + ..Default::default() + }, + ) + }) + .clone() }, + &state.registry.pod_names(), + ) + .await; + if state.registry.audit_log_enabled() { + let probe = hashes_by_block_size + .values() + .next() + .cloned() + .unwrap_or_default(); + audit_query(&model, &outcome, &probe); + } + let mut resp = (outcome.status, Json(outcome.body)).into_response(); + resp.extensions_mut().insert(AccessLogModel(model)); + resp +} + +fn audit_query(model_name: &str, outcome: &ModelQueryOutcome, probe: &[LocalBlockHash]) { + audit::log_query( + model_name, + &outcome.queried_groups, + outcome.status.as_u16(), + probe, + &outcome.body, ); - let (status, json) = run_tiered_query(&indexer, block_hashes, block_size).await; - let mut resp = (status, json).into_response(); +} + +fn model_not_found(model: String) -> Response { + let mut resp = ( + StatusCode::NOT_FOUND, + Json(serde_json::json!({ "error": format!("no indexer for model={model}") })), + ) + .into_response(); resp.extensions_mut().insert(AccessLogModel(model)); resp } @@ -364,32 +423,23 @@ async fn query_by_hash( return resp; } - let key = RoutingPartitionId::new(req.model_name, req.routing_group); - let Some(ie) = state.registry.get_indexer(&key) else { - let mut resp = ( - StatusCode::NOT_FOUND, - Json(serde_json::json!({ - "error": format!( - "no indexer for model={} routing_group={}", - key.model_name, key.routing_group - ) - })), - ) - .into_response(); - resp.extensions_mut().insert(AccessLogModel(model)); - return resp; - }; - let block_size = ie.block_size; - let indexer = ie.indexer.clone(); - drop(ie); + let trees = state + .registry + .query_trees(&req.model_name, req.routing_group.as_deref()); + if trees.is_empty() { + return model_not_found(model); + } let block_hashes: Vec = req .block_hashes .iter() .map(|h| LocalBlockHash(*h as u64)) .collect(); - let (status, json) = run_tiered_query(&indexer, block_hashes, block_size).await; - let mut resp = (status, json).into_response(); + let outcome = query_model(trees, |_| block_hashes.clone(), &state.registry.pod_names()).await; + if state.registry.audit_log_enabled() { + audit_query(&model, &outcome, &block_hashes); + } + let mut resp = (outcome.status, Json(outcome.body)).into_response(); resp.extensions_mut().insert(AccessLogModel(model)); resp } diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index f4c45b1fa7dd..6ee9143eb200 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -203,10 +203,10 @@ impl SelectionCore { tracking_hash: Arc, ) -> Self { let cancel_token = cancel_token.child_token(); - let indexer_registry = Arc::new(WorkerRegistry::new_with_cancel_token( - indexer_threads, - cancel_token.clone(), - )); + let indexer_registry = Arc::new( + WorkerRegistry::new_with_cancel_token(indexer_threads, cancel_token.clone()) + .with_retained_indexers(), + ); if signal_indexer_ready { indexer_registry.signal_ready(); } @@ -1413,41 +1413,85 @@ mod tests { } } - #[test] - fn shutdown_keeps_parent_alive() { + #[tokio::test] + async fn shutdown_cancels_listeners_but_keeps_parent_alive() { let parent = CancellationToken::new(); let core = SelectionCore::new_local( - test_config(false), + test_config(true), 1, parent.clone(), SelectionCacheConfig::default(), ); - core.shutdown(); + let record = core + .upsert_worker(worker_with_kv_events(1)) + .await + .expect("worker upsert"); + assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable); + assert_eq!(core.indexer_registry.listener_cancelled(1, 0), Some(false)); + core.shutdown(); assert!(core.cancel_token.is_cancelled()); assert!(!parent.is_cancelled()); + assert_eq!(core.indexer_registry.listener_cancelled(1, 0), Some(true)); } + #[rstest::rstest] + #[case(1)] + #[case(2)] #[tokio::test] - async fn shutdown_cancels_listeners() { - let parent = CancellationToken::new(); - let core = SelectionCore::new_local( + async fn selection_sees_cache_after_last_worker_replacement(#[case] indexer_threads: usize) { + let core = SelectionCore::try_new_local( test_config(true), - 1, - parent, + indexer_threads, + CancellationToken::new(), SelectionCacheConfig::default(), - ); + ) + .expect("valid test config"); + let key = RoutingPartitionId::new("model", "default"); + let request = || { + let mut request = select_request(); + request.prompt.token_ids = None; + request.prompt.block_hashes = Some(vec![11]); + request.prompt.sequence_hashes = Some(vec![101]); + request.prompt.isl_tokens = Some(4); + request + }; - let record = core - .upsert_worker(worker_with_kv_events(1)) - .await - .expect("worker upsert"); - assert_eq!(record.lifecycle, WorkerLifecycle::Schedulable); - assert_eq!(core.indexer_registry.listener_cancelled(1, 0), Some(false)); + // Exercise an in-place update, then removing the last worker and adding a new one. + for worker_id in [1, 1, 2] { + if worker_id == 2 { + core.delete_worker(1).await.expect("delete last worker"); + } + core.upsert_worker(worker_with_kv_events(worker_id)) + .await + .expect("worker upsert"); + assert_eq!(core.select(request()).await.unwrap().overlap.gpu, 0); + // Write through the registry used by listeners, not the selector's saved reference. + let indexer = core + .indexer_registry + .get_indexer(&key) + .unwrap() + .indexer + .clone(); + indexer + .apply_event_routed(store_event( + worker_id, + 0, + 1, + &[], + &[11], + StorageTier::Device, + )) + .await + .unwrap(); + indexer.dump_events().await.expect("flush indexer"); + let selected = core.select(request()).await.expect("select cached worker"); + assert_eq!(selected.worker_id, worker_id); + assert_eq!(selected.overlap.gpu, 4); + } core.shutdown(); - assert_eq!(core.indexer_registry.listener_cancelled(1, 0), Some(true)); } #[tokio::test] diff --git a/lib/kv-router/tests/standalone_indexer_http.rs b/lib/kv-router/tests/standalone_indexer_http.rs index ba5cfc0546bc..bd1a2d8c8150 100644 --- a/lib/kv-router/tests/standalone_indexer_http.rs +++ b/lib/kv-router/tests/standalone_indexer_http.rs @@ -281,7 +281,7 @@ async fn query_by_hash_returns_per_instance_tier_breakdown() { } #[tokio::test] -async fn query_by_hash_isolates_routing_groups_and_ignores_legacy_tenant_id() { +async fn query_by_hash_reads_the_named_routing_group_or_fans_out() { const BLOCK_SIZE: u32 = 4; const MODEL: &str = "test-model"; let registry = registry_with_events( @@ -315,8 +315,10 @@ async fn query_by_hash_isolates_routing_groups_and_ignores_legacy_tenant_id() { .json() .await .expect("legacy tenant-only response"); + // DeepInfra: without a routing_group the query fans out over every group + // of the model (engine_hash regimes); tenant_id stays ignored. assert_eq!(legacy_only["scores"]["7"]["0"], BLOCK_SIZE); - assert!(legacy_only["scores"].get("8").is_none()); + assert_eq!(legacy_only["scores"]["8"]["0"], BLOCK_SIZE); let explicit_group: serde_json::Value = client .post(format!("{base_url}/query_by_hash"))