From 5ce23b1527725cf623a3fc7feb94dc6d3d5fafaa Mon Sep 17 00:00:00 2001 From: Dong Xu Date: Tue, 3 Nov 2026 09:51:50 -0800 Subject: [PATCH 1/8] Pin the vLLM quick start to a Dynamo revision that imports on the pinned vLLM nightly (#160) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem - `Dockerfile.quick-start-vllm` pins Dynamo v1.6.0 (`701da25`) with vLLM nightly `edad8b03`. Every worker built from it exits at import (`cannot import name 'get_engine_zmq_addresses'`), so the quick start cannot start a worker (#159). - The build-time check imports `dynamo.vllm` but not the worker entry point, so the mismatch passes the build. ## Fix - `DYNAMO_REF` → `318ab039` (Dynamo main, 2026-11-02). Dynamo dropped the renamed import on 2026-10-23 (ai-dynamo/dynamo#16757); v1.6.0 and v1.6.1 both still have it. `VLLM_IMAGE` and `KVCR_VLLM_REF` are unchanged. - The compatibility check now also runs `import dynamo.vllm.main`, so a Dynamo/vLLM pair that cannot start a worker fails the build instead of the first `python -m dynamo.vllm`. ## Validation - Image built from this Dockerfile unchanged otherwise (`ghcr.io/luckyq/kvcr-phase1:qs-218c7816b34b`, 7 min); the new import check passes. - 2×H100 (dynamo-aws-dev-02), Qwen3-8B, Guard mode, Phase 1 harness TC-01/02/03: 3/3 PASS (`source_write == remote_deliver`, 0 token mismatches). - Same Dynamo revision with the full matrix earlier today: dense 6/6, 2P+1D 6/6, Qwen3.5-4B hybrid 6/6. Fixes #159. Signed-off-by: doxu --- Dockerfile.quick-start-vllm | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/Dockerfile.quick-start-vllm b/Dockerfile.quick-start-vllm index e812b9a8..8619a64d 100644 --- a/Dockerfile.quick-start-vllm +++ b/Dockerfile.quick-start-vllm @@ -37,7 +37,7 @@ RUN sed -i 's/^from kvcr\.control_channels import ZmqPeerControlChannel$/from kv # Pin a landed Dynamo revision with the versioned KV hint protocol and # per-rank KVCR control ports. FROM build-base AS dynamo-builder -ARG DYNAMO_REF=701da25292ecca8f64d770cc0b17d389200019a1 +ARG DYNAMO_REF=318ab0390d8e3e1fa7517ea6a1ebf46a7854c236 RUN git clone --filter=blob:none https://github.com/ai-dynamo/dynamo.git /src/dynamo \ && git -C /src/dynamo fetch --depth 1 origin "${DYNAMO_REF}" \ && git -C /src/dynamo checkout --detach "${DYNAMO_REF}" @@ -97,6 +97,9 @@ import inspect from types import SimpleNamespace import dynamo.vllm +# The worker entry point: fails here, not at `python -m dynamo.vllm`, when the +# pinned Dynamo and vLLM disagree on an internal API. +import dynamo.vllm.main import kvcr import vllm from dynamo.common.constants import KV_HINT_TRANSFER_CAPABILITY_KEY From 4e0f7fb5f44a51fe25958126503fffd767275566 Mon Sep 17 00:00:00 2001 From: Dong Xu Date: Tue, 3 Nov 2026 17:35:29 -0800 Subject: [PATCH 2/8] Pin the vLLM quick start to the adapter that implements the renamed KeyAdapter API (#162) Signed-off-by: doxu --- Dockerfile.quick-start-vllm | 4 ++-- docs/dev-guide.md | 2 +- docs/quick-start-vllm.md | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/Dockerfile.quick-start-vllm b/Dockerfile.quick-start-vllm index 8619a64d..cf50c659 100644 --- a/Dockerfile.quick-start-vllm +++ b/Dockerfile.quick-start-vllm @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -ARG VLLM_IMAGE=vllm/vllm-openai:nightly-edad8b03c5285a1c12e7e4f83c1af709b319a01b@sha256:2e735ffeef26457f8a3ff44a62ae182870b10af37adc06ecbb08baca152365c6 +ARG VLLM_IMAGE=vllm/vllm-openai:nightly-40941f2c4ba96fba0a7e282a5355e89cd74d9ac5@sha256:298b8765981d9dabda0da1504487bcdca876d5c18e8533455b3fda65ae964ba2 # Keep the vLLM image, adapter source, and Dynamo revision compatible. FROM rust:1.93-bookworm AS build-base @@ -14,7 +14,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ # Pin the merged vLLM KVCR integration to an exact commit SHA. FROM build-base AS vllm-kvcr-source ARG KVCR_VLLM_REPO=https://github.com/vllm-project/vllm.git -ARG KVCR_VLLM_REF=fed7de909f894fbee0c45130376aed2e2c9bf0d3 +ARG KVCR_VLLM_REF=a5db4c9ff58a0ec4dc20594ba5cea9f776dcd693 RUN git init /src/vllm-kvcr \ && git -C /src/vllm-kvcr remote add origin "${KVCR_VLLM_REPO}" \ && git -C /src/vllm-kvcr fetch --depth 1 origin "${KVCR_VLLM_REF}" \ diff --git a/docs/dev-guide.md b/docs/dev-guide.md index cf340387..4c21045b 100644 --- a/docs/dev-guide.md +++ b/docs/dev-guide.md @@ -308,7 +308,7 @@ Validate that counters increase on both success and failure paths, byte counts a Complete standalone setup and validation first. Use a separate shared environment for Dynamo, vLLM, and KVCR to keep integration dependencies out of the standalone development loop. -Select a vLLM revision that includes the merged KVCR 0.2.0 secondary-tier adapter (`"type": "kvcr"`), starting with [`fed7de9`](https://github.com/vllm-project/vllm/commit/fed7de909f894fbee0c45130376aed2e2c9bf0d3) on `main`. Pair it with a landed Dynamo revision such as [`701da25`](https://github.com/ai-dynamo/dynamo/commit/701da25292ecca8f64d770cc0b17d389200019a1), which supports the versioned KV hint contract and one KVCR control port per local data-parallel rank. The [vLLM quick start](quick-start-vllm.md) records a pinned combination for its container. +Select a vLLM revision that includes the merged KVCR 0.2.0 secondary-tier adapter (`"type": "kvcr"`), starting with [`a5db4c9`](https://github.com/vllm-project/vllm/commit/a5db4c9ff58a0ec4dc20594ba5cea9f776dcd693) on `main` (the first revision implementing the `KeyAdapter` API of KVCR #161). Pair it with a landed Dynamo revision such as [`318ab03`](https://github.com/ai-dynamo/dynamo/commit/318ab0390d8e3e1fa7517ea6a1ebf46a7854c236), which supports the versioned KV hint contract and one KVCR control port per local data-parallel rank. The [vLLM quick start](quick-start-vllm.md) records a pinned combination for its container. ### Build the integration environment diff --git a/docs/quick-start-vllm.md b/docs/quick-start-vllm.md index 93229ad8..f4505e0a 100644 --- a/docs/quick-start-vllm.md +++ b/docs/quick-start-vllm.md @@ -32,7 +32,7 @@ The repository includes [Dockerfile.quick-start-vllm](../Dockerfile.quick-start- ```bash export KVCR_VLLM_REPO=https://github.com/vllm-project/vllm.git -export KVCR_VLLM_REF=fed7de909f894fbee0c45130376aed2e2c9bf0d3 +export KVCR_VLLM_REF=a5db4c9ff58a0ec4dc20594ba5cea9f776dcd693 DOCKER_BUILDKIT=1 docker build \ --build-arg KVCR_VLLM_REPO="$KVCR_VLLM_REPO" \ From f737e27e61075132150ff35e84c7d397871a1848 Mon Sep 17 00:00:00 2001 From: linhu-nv <141609318+linhu-nv@users.noreply.github.com> Date: Wed, 4 Nov 2026 13:15:43 +0800 Subject: [PATCH 3/8] docs: add SGLang direct linker quick start (#129) ## Summary - Add an experimental SGLang Direct Linker quick start and Dockerfile. - Use KVCR-owned inventory and Dynamo router hints for peer KV reuse. - Scope the example to dense Qwen3-32B BF16 with TP1 workers. The recipe pins Dynamo [#17655](https://github.com/ai-dynamo/dynamo/pull/17655) (`fe81cfd`) and SGLang [#45650](https://github.com/sgl-project/sglang/pull/45650) (`5b65e08`); KVCR is installed from the build checkout. ## Validation - Built the Dockerfile from this PR and ran its import/version canary. - Two B200 nodes, one TP1 worker per node: ten 116K-token requests restored 115,968 tokens via router-generated hints and RDMA, with matching output IDs. Mooncake linker `local_first` and `random` ran against the same inputs. - The cross-node test used a harness and a chroot in retained GPU pods; the guide's same-host deployment was not run as written. Signed-off-by: linhu-nv Co-authored-by: Ziqi Fan --- Dockerfile.quick-start-sglang | 70 ++++++++++---- README.md | 14 +-- docs/dev-guide.md | 2 +- docs/quick-start-sglang.md | 175 +++++++++++++++++++++++++++------- 4 files changed, 198 insertions(+), 63 deletions(-) diff --git a/Dockerfile.quick-start-sglang b/Dockerfile.quick-start-sglang index 4eca401b..ff48d51b 100644 --- a/Dockerfile.quick-start-sglang +++ b/Dockerfile.quick-start-sglang @@ -2,6 +2,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# QS_MODE selects the SGLang integration in docs/quick-start-sglang.md: hicache or +# linker. They are separate SGLang PR branches, so each mode pins its own SGLang +# and Dynamo revisions. +ARG QS_MODE=hicache ARG SGLANG_IMAGE=lmsysorg/sglang:v0.5.21-cu130-runtime@sha256:797614a377ead47cf89ccaba810685b56a1bf0d6f0c45f84af2e8f42c198e104 FROM rust:1.96.1-bookworm@sha256:d99f7b31f49909348dc59b51f3c95d1efded1701ffb222f095aaab7de3c4abd8 AS sglang-dynamo-builder @@ -9,12 +13,17 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ cmake git libclang-dev libhwloc-dev libprotobuf-dev libssl-dev libudev-dev \ pkg-config protobuf-compiler python3-dev python3-pip python3-venv \ && rm -rf /var/lib/apt/lists/* -ARG SGLANG_DYNAMO_REF=83ea37e849e883c27d7ed848e0f329961e1fb173 -RUN git init /src/dynamo \ +ARG QS_MODE +RUN case "${QS_MODE}" in \ + hicache) ref=83ea37e849e883c27d7ed848e0f329961e1fb173 ;; \ + linker) ref=fe81cfdccedc23b4eabef3a250b538ff11ba102c ;; \ + *) echo "QS_MODE must be hicache or linker" >&2; exit 1 ;; \ + esac \ + && git init /src/dynamo \ && git -C /src/dynamo remote add origin https://github.com/ai-dynamo/dynamo.git \ - && git -C /src/dynamo fetch --depth 1 origin "${SGLANG_DYNAMO_REF}" \ + && git -C /src/dynamo fetch --depth 1 origin "${ref}" \ && git -C /src/dynamo checkout --detach FETCH_HEAD \ - && test "$(git -C /src/dynamo rev-parse HEAD)" = "${SGLANG_DYNAMO_REF}" + && test "$(git -C /src/dynamo rev-parse HEAD)" = "${ref}" RUN python3 -m venv /venv \ && /venv/bin/pip install --no-cache-dir 'maturin[patchelf]==1.9.6' ENV CARGO_BUILD_JOBS=4 @@ -24,16 +33,20 @@ RUN /venv/bin/pip wheel --no-deps /src/dynamo --wheel-dir /wheels FROM ${SGLANG_IMAGE} AS sglang USER root -ARG SGLANG_REF=658755604ace9f476b84f98511f76b815606072a +ARG QS_MODE RUN test "$(uname -m)" = x86_64 \ && apt-get update && apt-get install -y --no-install-recommends \ git curl libatomic1 libhwloc15 libnuma1 libssl3t64 libudev1 \ && rm -rf /var/lib/apt/lists/* -RUN git init /opt/sglang-src \ +RUN case "${QS_MODE}" in \ + hicache) ref=24a3a79ed3ed3ee5811f78ba2c4cbd5ceaac27d1 ;; \ + linker) ref=002d6d8d1df91a58695a8868ce2c2db590809f75 ;; \ + esac \ + && git init /opt/sglang-src \ && git -C /opt/sglang-src remote add origin https://github.com/ziqifan617/sglang.git \ - && git -C /opt/sglang-src fetch --depth 1 origin "${SGLANG_REF}" \ + && git -C /opt/sglang-src fetch --depth 1 origin "${ref}" \ && git -C /opt/sglang-src checkout --detach FETCH_HEAD \ - && test "$(git -C /opt/sglang-src rev-parse HEAD)" = "${SGLANG_REF}" + && test "$(git -C /opt/sglang-src rev-parse HEAD)" = "${ref}" # Install KVCR from the checkout used as the Docker build context. COPY pyproject.toml README.md LICENSE /opt/kvcr-src/ COPY src /opt/kvcr-src/src @@ -80,9 +93,10 @@ ENV DYN_REQUEST_PLANE=tcp DYN_EVENT_PLANE=zmq DYN_DISCOVERY_BACKEND=file \ DYN_FILE_KV=/tmp/dynamo-sglang-quick-start DYN_SYSTEM_HOST=127.0.0.1 \ DYN_RUNTIME_NUM_WORKER_THREADS=4 DYN_RUNTIME_MAX_BLOCKING_THREADS=8 \ TOKENIZERS_PARALLELISM=false PYTHONHASHSEED=0 RUST_LOG=info \ - UCX_TLS=tcp,sm,self,cuda_copy OMP_NUM_THREADS=4 + UCX_TLS=tcp,sm,self,cuda_copy OMP_NUM_THREADS=4 SGLANG_JIT_DEEPGEMM_PRECOMPILE=0 # Import/version canary only; model/peer-transfer acceptance is in the guide. RUN python3 - <<'PY' +import os from importlib import metadata from pathlib import Path import tomllib @@ -90,11 +104,7 @@ import dynamo.sglang import torch from dynamo.common.constants import KV_HINT_TRANSFER_CAPABILITY_KEY from dynamo.llm import WorkerType -from dynamo.sglang import kv_hints -from dynamo.sglang.state_agent import StateAgentLifecycle from kvcr import ZmqPeerControlChannel -from sglang.srt.mem_cache.hicache_storage import HiCacheStorage -from sglang.srt.mem_cache.storage.kvcr.kvcr_store import KVCRStore for line in Path('/opt/sglang-constraints.txt').read_text().splitlines(): name, expected = line.split('==') @@ -103,15 +113,41 @@ for line in Path('/opt/sglang-constraints.txt').read_text().splitlines(): print(f'{name}=={actual}') assert torch.version.cuda == '13.0' assert KV_HINT_TRANSFER_CAPABILITY_KEY == 'router_hint' -assert callable(KVCRStore.batch_get_v1) -assert callable(HiCacheStorage.get_kv_hint_source) -assert callable(StateAgentLifecycle.start) assert callable(ZmqPeerControlChannel) +if os.environ['QS_MODE'] == 'hicache': + from dynamo.sglang import kv_hints + from dynamo.sglang.state_agent import StateAgentLifecycle + from sglang.srt.mem_cache.hicache_storage import HiCacheStorage + from sglang.srt.mem_cache.storage.kvcr.kvcr_store import KVCRStore + + assert callable(KVCRStore.batch_get_v1) + assert callable(HiCacheStorage.get_kv_hint_source) + assert callable(StateAgentLifecycle.start) +else: + from dynamo.sglang.kv_hints import is_supported_kvcr_linker_topology + from dynamo.sglang.state_agent import state_agent_descriptors + from kvcr.hints.fetch import KV_FETCH + from kvcr.hints.parser import parse_hints + from sglang.srt.managers.kv_hints import decode_kv_hints_envelope + from sglang.srt.mem_cache.storage.kvcr.kvcr_direct_linker import KVCRDirectLinker + + assert callable(is_supported_kvcr_linker_topology) and callable(state_agent_descriptors) + assert KVCRDirectLinker.local_storage_ownership == 'kvcr' + hint = decode_kv_hints_envelope({ + 'protocol_version': '0.1', 'message_id': 'canary', + 'actions': [{'action_id': 'fetch', 'action_type': 'kv.fetch', + 'action_version': '1.0', 'payload': { + 'source_control_endpoint': 'tcp://127.0.0.1:19500', + 'block_hashes': [1], + }}], + }) + parsed, rejected = parse_hints(hint, {'kv.fetch': KV_FETCH}) + assert len(parsed) == 1 and not rejected assert metadata.version('ai-dynamo') == '1.7.0' assert metadata.version('ai-dynamo-runtime') == '1.7.0' kvcr_project = tomllib.loads(Path('/opt/kvcr-src/pyproject.toml').read_text())['project'] assert metadata.version('kvcr') == kvcr_project['version'] -print('SGLang HiCache import/version checks passed; GPU verification still required') +print(f"SGLang {os.environ['QS_MODE']} import/version checks passed; GPU verification still required") PY RUN python3 -m pip freeze --all > /opt/sglang-installed-packages.txt WORKDIR /workspace diff --git a/README.md b/README.md index 5067483b..782887eb 100644 --- a/README.md +++ b/README.md @@ -2,11 +2,6 @@ KV Cache Runner - Data Plane Architecture

-> [!NOTE] -> **Initial release (experimental).** KVCR supports cross-node DRAM sharing, local disk caching, and resiliency for vLLM. Future releases will add SGLang and TRT-LLM integrations and object storage support. -> -> Feedback and contributions are welcome - see [CONTRIBUTING.md](CONTRIBUTING.md). - # KV Cache Runner Modern AI workloads require KV cache infrastructure that delivers speed, scale, and resilience. @@ -63,9 +58,10 @@ A hint envelope may carry several hints. KVCR parses and dispatches each one ind - vLLM - supported -- SGLang - - [[RFC] KVCR as a HiCacheStorage backend for peer-to-peer KV reuse #32903](https://github.com/sgl-project/sglang/issues/32903) - +- SGLang - supported + - HiCache storage backend: [#45647](https://github.com/sgl-project/sglang/pull/45647), [RFC #32903](https://github.com/sgl-project/sglang/issues/32903) + - Direct linker (experimental, TP1 for now): [#45650](https://github.com/sgl-project/sglang/pull/45650) + ### Routers @@ -79,7 +75,7 @@ A hint envelope may carry several hints. KVCR parses and dispatches each one ind The [vLLM quick start](docs/quick-start-vllm.md) is a public preview that builds and runs KVCR with vLLM, Dynamo, and NIXL using pinned source revisions and a compatible base image. See the [developer guide](docs/dev-guide.md#integrate-with-vllm-and-dynamo-optional) for source installation and verification. -The [SGLang HiCache quick start](docs/quick-start-sglang.md) is a preview for building and running KVCR with SGLang and Dynamo. +The [SGLang quick start](docs/quick-start-sglang.md) is a preview for building and running KVCR with SGLang and Dynamo, through either the HiCache storage backend or the experimental direct linker. ### Development diff --git a/docs/dev-guide.md b/docs/dev-guide.md index 4c21045b..2ca3a5d2 100644 --- a/docs/dev-guide.md +++ b/docs/dev-guide.md @@ -2,7 +2,7 @@ This guide is for developers who are comfortable with Python, Linux, and inference engines such as vLLM, but are new to KV Cache Runner (KVCR). It covers the local development path from environment setup through standalone validation. -To try the integrated stack without modifying source, follow the [vLLM quick start](quick-start-vllm.md) or [SGLang HiCache quick start](quick-start-sglang.md). +To try the integrated stack without modifying source, follow the [vLLM quick start](quick-start-vllm.md) or [SGLang quick start](quick-start-sglang.md). --- diff --git a/docs/quick-start-sglang.md b/docs/quick-start-sglang.md index 9f0b7686..99b91758 100644 --- a/docs/quick-start-sglang.md +++ b/docs/quick-start-sglang.md @@ -3,23 +3,26 @@ SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All SPDX-License-Identifier: Apache-2.0 --> -# KVCR Quick Start with SGLang HiCache +# KVCR Quick Start with SGLang -Build and run Dynamo, SGLang and KVCR without editing their serving code. SGLang stages KV through HiCache host memory; KVCR stores it in a separate DRAM pool. A Dynamo state agent indexes KVCR-owned blocks so the router can generate source hints for peer delivery. +Build and run Dynamo, SGLang and KVCR without editing their serving code. A Dynamo state agent indexes KVCR-owned blocks so the router can generate source hints for peer delivery. This guide covers two SGLang integrations, selected with `QS_MODE`: + +- `hicache`: SGLang stages KV through HiCache host memory; KVCR stores it in a separate DRAM pool ([SGLang #45647](https://github.com/sgl-project/sglang/pull/45647)). +- `linker` (experimental): SGLang's external cache linker hands GPU KV pages to KVCR, which keeps them in DRAM and delivers them layer by layer into another worker's device pools. There is no HiCache host tier ([SGLang #45650](https://github.com/sgl-project/sglang/pull/45650)). > [!IMPORTANT] -> This integration is a preview built from pinned source revisions, including the open SGLang PR below. +> Both modes are previews built from pinned source revisions, including the open SGLang PRs above. The linker supports TP1 workers for now and is tested with the dense Qwen3-32B BF16 model; TP>1, hybrid pools and other models are not validated yet. For vLLM, use the [vLLM quick start](quick-start-vllm.md). For implementation and API development, use the [developer guide](dev-guide.md). ## Prerequisites -- A Linux x86-64 host with two available NVIDIA GPUs compatible with the pinned SGLang/CUDA 13 image and enough free GPU memory for this example. -- NVIDIA driver 580.126.20 or newer for the pinned CUDA 13 runtime. +- A Linux x86-64 host with two available NVIDIA Hopper or Blackwell GPUs (tested on H20 and B200). +- An NVIDIA R580 driver with CUDA 13.0 support (tested with 580.95.05 and 580.126.20). - Docker with [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html), Git, and public access to Docker Hub, GitHub, PyPI, `pypi.nvidia.com`, `flashinfer.ai`, crates.io and Hugging Face. These public indexes need no login. -- Allow 32 GiB host RAM and 150 GiB disk for this build and model cache. +- Allow 32 GiB host RAM and 150 GiB disk for `hicache`, or 128 GiB and 250 GiB for `linker`. -This example uses two independent TP1 workers on one host, Qwen3-0.6B BF16, `trtllm_mha`, 64-token pages, and the Python radix tree. Each worker has 2 GiB of HiCache staging and a separate 1 GiB KVCR DRAM pool. The same-host transport is UCX TCP/shared memory with CUDA-copy support, not an RDMA performance test. +This example uses two independent TP1 workers on one host, `trtllm_mha`, 64-token pages, and the Python radix tree. `hicache` runs Qwen3-0.6B BF16 with 2 GiB of HiCache staging and a separate 1 GiB KVCR DRAM pool per worker. `linker` runs Qwen3-32B BF16 with a 4 GiB KVCR DRAM pool per worker. The same-host transport is UCX TCP/shared memory with CUDA-copy support, not an RDMA performance test. ## 1. Build the integration image @@ -27,16 +30,17 @@ Clone the repository and build the SGLang image from [Dockerfile.quick-start-sgl ```bash set -euo pipefail +export QS_MODE=hicache # or linker git clone https://github.com/ai-dynamo/kvcr.git kvcr-sgl-quick-start cd kvcr-sgl-quick-start git rev-parse HEAD -export QS_TAG="kvcr-sgl-quick-start:$(date -u +%Y%m%dT%H%M%SZ)-$$" +export QS_TAG="kvcr-sgl-quick-start:$QS_MODE-$(date -u +%Y%m%dT%H%M%SZ)-$$" DOCKER_BUILDKIT=1 docker build --platform linux/amd64 \ - --file Dockerfile.quick-start-sglang --tag "$QS_TAG" . + --file Dockerfile.quick-start-sglang --build-arg QS_MODE="$QS_MODE" --tag "$QS_TAG" . export QS_IMAGE_ID="$(docker image inspect "$QS_TAG" --format '{{.Id}}')" ``` -The build uses pinned [SGLang #45647](https://github.com/sgl-project/sglang/pull/45647) and Dynamo revisions from the Dockerfile, and installs KVCR from the local checkout used as the build context (`.`), including any local changes. Continue only if the final import/version checks pass. +The Dockerfile pins the SGLang and Dynamo revisions for each mode; mismatched SGLang and KVCR revisions silently drop peer hints. KVCR is installed from the local checkout used as the build context (`.`), including any local changes. Continue only if the final import/version checks pass. ## 2. Start the container @@ -44,16 +48,20 @@ Select two available GPUs; this example uses `0,1`. No host ports are published. ```bash nvidia-smi --query-compute-apps=gpu_uuid,pid,process_name --format=csv -export QS_NAME="kvcr-sgl-$(date -u +%Y%m%dT%H%M%SZ)-$$" +case "$QS_MODE" in + hicache) export QS_MODEL=Qwen/Qwen3-0.6B QS_MODEL_REV=c1899de289a04d12100db370d81485cdf75e47ca ;; + linker) export QS_MODEL=Qwen/Qwen3-32B QS_MODEL_REV=9216db5781bf21249d130ec9da846c4624c16137 ;; +esac +export QS_NAME="kvcr-sgl-$QS_MODE-$(date -u +%Y%m%dT%H%M%SZ)-$$" docker volume create "$QS_NAME-models" export QS_CID="$(docker run --detach --init --name "$QS_NAME" \ --gpus '"device=0,1"' --shm-size 8g \ --ulimit memlock=-1 --ulimit stack=67108864 \ - --env MODEL=Qwen/Qwen3-0.6B --env MODEL_PATH=/model \ + --env QS_MODE="$QS_MODE" --env MODEL="$QS_MODEL" --env MODEL_PATH=/model \ --volume "$QS_NAME-models:/model" \ "$QS_IMAGE_ID" sleep infinity)" docker exec "$QS_CID" python3 -c \ - 'from huggingface_hub import snapshot_download; snapshot_download("Qwen/Qwen3-0.6B", revision="c1899de289a04d12100db370d81485cdf75e47ca", local_dir="/model")' + "from huggingface_hub import snapshot_download; snapshot_download('$QS_MODEL', revision='$QS_MODEL_REV', local_dir='/model')" printf 'Container: %s\n' "$QS_CID" ``` @@ -82,7 +90,7 @@ env -u NATS_SERVER DYN_SYSTEM_PORT=24003 python3 -m dynamo.frontend \ The KV router supplies source hints; a round-robin router does not demonstrate router-guided peer reuse. -## 5. Start two HiCache workers +## 5. Start two workers Run this block in **each** worker terminal, setting `RANK=0` in the first and `RANK=1` in the second. Keep both running throughout verification. @@ -90,44 +98,58 @@ Run this block in **each** worker terminal, setting `RANK=0` in the first and `R set -euo pipefail export RANK=0 # Use 1 in the second terminal. export PROMETHEUS_MULTIPROC_DIR="/tmp/sgl-metrics-$RANK" -mkdir "$PROMETHEUS_MULTIPROC_DIR" +mkdir -p "$PROMETHEUS_MULTIPROC_DIR" KVCR_CONFIG="$(python3 - <<'PY' import ipaddress, json, os, socket rank = int(os.environ['RANK']) host = os.environ.get('KVCR_ADVERTISE_HOST') or socket.gethostbyname(socket.gethostname()) assert not ipaddress.ip_address(host).is_loopback, 'Set KVCR_ADVERTISE_HOST to a reachable container address' -print(json.dumps({ - 'control_advertise_host': host, - 'control_port': 24100 + rank * 32, 'nixl_listen_port': 24200 + rank * 32, - 'local_dram_bytes': 1 << 30, - 'publish_inventory': True, - 'enable_remote_hint': True, 'hicache_storage_pass_prefix_keys': True, - 'align_storage_prefix': True, 'combine_page_reads': True, - 'combine_page_writes': True, 'diagnostic_telemetry': True, - 'verify_transfer_bytes': True, - 'startup_timeout_seconds': 300, +config = { + 'publish_inventory': True, 'startup_timeout_seconds': 300, 'operation_timeout_ms': 20000, 'abandon_timeout_ms': 60000, -})) +} +if os.environ['QS_MODE'] == 'hicache': + config.update({ + 'control_advertise_host': host, + 'control_port': 24100 + rank * 32, 'nixl_listen_port': 24200 + rank * 32, + 'local_dram_bytes': 1 << 30, + 'enable_remote_hint': True, 'hicache_storage_pass_prefix_keys': True, + 'align_storage_prefix': True, 'combine_page_reads': True, + 'combine_page_writes': True, 'diagnostic_telemetry': True, + 'verify_transfer_bytes': True, + }) +else: + config.update({ + 'kvcr_service_socket_path': None, 'dram_size_gb': 4, 'advertise_host': host, + 'control_port': 24100 + rank * 32, 'nixl_port': 24200 + rank * 32, + 'enable_telemetry': True, + }) +print(json.dumps(config)) PY )" +if [ "$QS_MODE" = hicache ]; then + MODE_ARGS=(--mem-fraction-static 0.5 --enable-hierarchical-cache --hicache-size 2 + --hicache-mem-layout page_first --hicache-write-policy write_through + --hicache-storage-backend kvcr --hicache-storage-prefetch-policy wait_complete) +else + MODE_ARGS=(--mem-fraction-static 0.8 --kv-cache-dtype bfloat16 + --enable-unified-cache-external-linker --unified-cache-external-linker-backend kvcr) +fi env -u NATS_SERVER CUDA_VISIBLE_DEVICES="$RANK" DYN_SYSTEM_PORT="$((24001 + RANK))" \ python3 -m dynamo.sglang --endpoint dyn://kvcrqs.backend.generate \ --model-path "$MODEL_PATH" --served-model-name "$MODEL" --tp-size 1 \ --host 127.0.0.1 --port "$((24010 + RANK))" \ --dtype bfloat16 --attention-backend trtllm_mha \ --context-length 8192 --max-total-tokens 8192 --chunked-prefill-size 8192 \ - --max-running-requests 1 --page-size 64 --mem-fraction-static 0.5 \ + --max-running-requests 1 --page-size 64 "${MODE_ARGS[@]}" \ --disable-cuda-graph --disable-piecewise-cuda-graph \ --incremental-streaming-output --enable-metrics \ - --enable-hierarchical-cache --hicache-size 2 \ - --hicache-mem-layout page_first --hicache-write-policy write_through \ - --hicache-storage-backend kvcr --hicache-storage-prefetch-policy wait_complete \ --hicache-storage-backend-extra-config "$KVCR_CONFIG" \ --kv-events-config "{\"publisher\":\"zmq\",\"format\":\"dynamo\",\"topic\":\"\",\"endpoint\":\"tcp://*:$((24020 + RANK))\"}" \ 2>&1 | tee "/tmp/worker-$RANK.log" ``` -Control, NIXL, metrics and event ports must be distinct. The ownership descriptor requires a reachable non-loopback control address, even in this same-container example. Both terminals derive the container IP automatically; set the same `KVCR_ADVERTISE_HOST` in both if hostname resolution is unsuitable. KV events bind `*` inside that private network namespace. +`publish_inventory` enables KVCR-owned event publication. Control, NIXL, metrics and event ports must be distinct. The ownership descriptor requires a reachable non-loopback control address, even in this same-container example. Both terminals derive the container IP automatically; set the same `KVCR_ADVERTISE_HOST` in both if hostname resolution is unsuitable. KV events bind `*` inside that private network namespace. ## 6. Verify peer reuse @@ -144,6 +166,8 @@ print(status) PY ``` +### HiCache + This check sends a fresh 2048-token prefix to worker 0, then requests the identical prefix on worker 1. Do not probe this prefix before the test. The first request can compile additional attention kernels after readiness. ```bash @@ -225,7 +249,86 @@ print('Completed HiCache prefetched tokens:', total) PY ``` -A cache-hit counter alone is not proof of transfer. If any check fails, preserve the logs and investigate; retry with fresh workers and a new prefix, not a warm destination. This smoke test does not qualify general accuracy, performance, cross-host RDMA, TP>1 or hybrid pools. +### Direct linker + +This check sends a fresh prompt of about 4,300 tokens with a random code in the middle to worker 0, then the identical prompt to worker 1, then again to worker 0. It pins only the destination; Dynamo still supplies the source hint. Worker 1 must reuse every complete page before the last prompt token, every output must contain the code, and worker 1's output must match worker 0's warm replay, which uses the same cached prefix (a cold prefill can differ in later greedy tokens). The first request on each worker JIT-compiles kernels and can take several minutes. Do not probe this prompt before the test. + +```bash +python3 - <<'PY' +import json, os, random, re, time, urllib.request, uuid +from pathlib import Path +from transformers import AutoTokenizer + +workers = [] +for rank in (0, 1): + log = Path(f'/tmp/worker-{rank}.log').read_text() + log = re.sub(r'\x1b\[[0-9;]*m', '', log) + ids = set(re.findall(r'Backend admission gate created[^\n]*worker_id=([0-9a-f]+)', log)) + assert len(ids) == 1 and 'Successfully registered LLM with runtime config' in log + workers.append(int(ids.pop(), 16)) +assert workers[0] != workers[1] +frontend = Path('/tmp/frontend.log').read_text() +assert all(f'Adding worker WorkerWithDpRank {{ worker_id: {w}, dp_rank: 0 }}' in frontend for w in workers) +tokenizer = AutoTokenizer.from_pretrained(os.environ['MODEL_PATH']) +code = str(random.randrange(100000, 1000000)) +filler = 'The capital of France is Paris. ' * 300 +prompt = (f'Document {uuid.uuid4().hex}. {filler}The access code for vault 7 is {code}. ' + f'{filler}Question: what is the access code for vault 7? Answer: The access code for vault 7 is') +tokens = tokenizer.encode(prompt, add_special_tokens=False) +expected = (len(tokens) - 1) // 64 * 64 +outputs = [] +for rank, worker in enumerate(workers + workers[:1]): + if rank: + time.sleep(15) # Deposit, KV events and router indexing are asynchronous. + payload = dict(model=os.environ['MODEL'], prompt=tokens, temperature=0, + max_tokens=16, stream=True, stream_options={'include_usage': True}, + nvext={'backend_instance_id': worker, 'dp_rank': 0, + 'extra_fields': ['worker_id', 'engine_data', 'completion_token_ids']}) + request = urllib.request.Request('http://127.0.0.1:24000/v1/completions', + data=json.dumps(payload).encode(), headers={'Content-Type': 'application/json'}) + with urllib.request.urlopen(request, timeout=900) as response: + raw = response.read().decode() + Path(f'/tmp/response-{rank}.sse').write_text(raw) + events = [json.loads(line[5:]) for line in raw.splitlines() + if line.startswith('data:') and line[5:].strip() != '[DONE]'] + assert not any('error' in event for event in events) + ext = [event.get('nvext', {}) for event in events] + ids = [token for item in ext for token in item.get('completion_token_ids', [])] + reported = [item['worker_id'] for item in ext if 'worker_id' in item] + assert reported and all(x['prefill_worker_id'] == x['decode_worker_id'] == worker for x in reported) + cache = [item.get('engine_data', {}).get('kv_cache_hit', {}) for item in ext] + hits = [item['reused_tokens'] for item in cache if 'reused_tokens' in item] + assert hits and max(hits) == (expected if rank else 0), (rank, hits, expected) + text = tokenizer.decode(ids) + assert code in text, (rank, code, text) + outputs.append(ids) + print(f'request {rank} (worker {rank % 2}): cached_tokens={max(hits)}/{len(tokens)}, output={text!r}') +assert outputs[1] == outputs[2], 'Restored output differs from the warm replay on worker 0' +print('Response checks passed; verify transfer telemetry below before declaring success.') +PY +``` + +Then confirm a physical peer transfer from the `KVCR telemetry` lines. Each line covers one interval; sum them once: + +```bash +python3 - <<'PY' +import json, re +from collections import Counter +from pathlib import Path +for rank in (0, 1): + totals = Counter() + for line in Path(f'/tmp/worker-{rank}.log').read_text().splitlines(): + if (match := re.search(r'KVCR telemetry rank=\d+ interval_s=\S+ stats=(\{.*\})', line)): + totals.update({k: v for k, v in json.loads(match.group(1)).items() + if k.endswith(('.counter', '.count'))}) + print(f'worker {rank}:', json.dumps(dict(sorted(totals.items())), indent=2)) +PY +``` + +- Destination (worker 1): `kvcr_transfer_blocks[remote_deliver].counter` equals reused pages times layers (`cached_tokens / 64 * 64` for this 64-layer model, so it equals `cached_tokens`) and there is no `failed` entry. +- Source (worker 0): nonzero `kvcr_transfer_bytes[source_write].counter` and no `failed` entry. + +A cache-hit counter alone is not proof of transfer. If any check fails, preserve the logs and investigate; retry with fresh workers and a new prompt, not a warm destination. This smoke test does not qualify general accuracy, performance, cross-host RDMA, TP>1 or hybrid pools. ## Cleanup and troubleshooting @@ -236,8 +339,8 @@ export QS_EVIDENCE="$(mktemp -d /tmp/kvcr-sgl-evidence.XXXXXX)" for log in state-agent frontend worker-0 worker-1; do docker cp "$QS_CID:/tmp/$log.log" "$QS_EVIDENCE/" done -for rank in 0 1; do - docker cp "$QS_CID:/tmp/response-$rank.sse" "$QS_EVIDENCE/" +for rank in 0 1 2; do # response-2 exists only in linker mode + docker cp "$QS_CID:/tmp/response-$rank.sse" "$QS_EVIDENCE/" || true done docker cp "$QS_CID:/opt/sglang-installed-packages.txt" "$QS_EVIDENCE/" docker stop --time 90 "$QS_CID" @@ -248,5 +351,5 @@ Keep the model volume and any underlying GPU allocation unless you deliberately - **Build/import failure:** use the pinned source and kernel versions together and check the final import/version output. - **Startup failure:** check worker logs, GPU architecture, driver, memory and unique ports. Model readiness is not a transfer pass. -- **Cache miss:** check source deposit, KV events, router discovery and hints. -- **Transfer failure:** retain `UCX_TLS=tcp,sm,self,cuda_copy`; missing CUDA memory support can prevent NIXL startup even for host-staged data. +- **Cache miss:** check source deposit, KV events, router discovery and hints. In linker mode, `KVCR remote hints disabled` in a worker log means its layout cannot serve peers. +- **Transfer failure:** retain `UCX_TLS=tcp,sm,self,cuda_copy`; missing CUDA memory support can prevent NIXL startup even for host-staged data. In linker mode, failed loads are not recomputed and can stop the worker; size `dram_size_gb` for the working set. From dd5df02c076dfee4bb2cdc39bef11164f9dcaec5 Mon Sep 17 00:00:00 2001 From: Moein Khazraee <33970824+mkhazraee@users.noreply.github.com> Date: Tue, 3 Nov 2026 21:36:37 -0800 Subject: [PATCH 4/8] docs: clarify router support in README (#163) Signed-off-by: Moein Khazraee --- README.md | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 782887eb..aa72f6f2 100644 --- a/README.md +++ b/README.md @@ -33,7 +33,7 @@ For more details, see the [design document](docs/design_overview.md). ## Framework support -KVCR is vendor and framework agnostic, with no dependency on a specific request router. The list below summarizes framework integrations that are complete or in progress. We plan to extend support to additional frameworks and routers in the future. +KVCR is vendor and framework agnostic and can integrate with any request router that supports the KV Hint Protocol. The list below summarizes framework integrations that are complete or in progress. We plan to extend support to additional frameworks and routers in the future. ### KV Hint Protocol @@ -46,7 +46,7 @@ KVCR is vendor and framework agnostic, with no dependency on a specific request - SGLang - supported - Additional features tracked in [RFC #36224](https://github.com/sgl-project/sglang/issues/36224). -- Dynamo KV Router - supported +- [Dynamo KV Router](https://docs.nvidia.com/dynamo/dev/knowledge-base/modular-components/router/overview) - supported A hint envelope may carry several hints. KVCR parses and dispatches each one independently and records a per-hint outcome. @@ -59,16 +59,9 @@ A hint envelope may carry several hints. KVCR parses and dispatches each one ind - vLLM - supported - SGLang - supported - - HiCache storage backend: [#45647](https://github.com/sgl-project/sglang/pull/45647), [RFC #32903](https://github.com/sgl-project/sglang/issues/32903) + - HiCache storage backend: [#45647](https://github.com/sgl-project/sglang/pull/45647) - Direct linker (experimental, TP1 for now): [#45650](https://github.com/sgl-project/sglang/pull/45650) - -### Routers - -- [Dynamo KV Router](https://docs.nvidia.com/dynamo/dev/knowledge-base/modular-components/router/overview) -- [sgl-router](https://github.com/sgl-project/sglang/tree/main/experimental/sgl-router) -- [llmd-router](https://llm-d.ai/docs/dev/architecture/core/router) - ## Using KVCR ### Quick start From e6df92961f4c75bd18ac5eb50d7710c9b7abed34 Mon Sep 17 00:00:00 2001 From: Weiyang Wang Date: Fri, 6 Nov 2026 02:20:27 +0200 Subject: [PATCH 5/8] fix: make request-scoped kv.fetch advice copy-only Keep kv.fetch as source advice for later framework fetch/deliver calls. Remove obsolete payload handling and clarify that hint hashes differ from full BlockKey inventory and transfer identities. Validation: hint parser/fetch suites, Ruff, and whitespace checks. Signed-off-by: Weiyang Wang --- docs/design_overview.md | 12 ++++-------- src/kvcr/hints/fetch.py | 5 ----- tests/unit/test_hints_parser.py | 5 ----- 3 files changed, 4 insertions(+), 18 deletions(-) diff --git a/docs/design_overview.md b/docs/design_overview.md index f005a2d9..1d814e75 100644 --- a/docs/design_overview.md +++ b/docs/design_overview.md @@ -253,17 +253,13 @@ The router maintains an eventually-consistent global prefix inventory for cache- Framework-owned memory is reported through the framework's existing router event path. If needed, `tier_enter`/`tier_exit` reporting or `inventory_snapshot()` can be added to the KVCR API later. -KVCR-owned inventory events use the same `BlockKey` values as fetch and `submit_hint` in the normal case. KVCR deliberately reports only the affected keys, tier, and whether the residency was removed. KVCR delivers each event to `inventory_sink`; the framework or local indexer owns any buffering or batching for publication. Each event may contain multiple keys and is tier-specific because one block may be present in several tiers at once. +KVCR-owned inventory events use the same full `BlockKey` values as fetch. Request-scoped hints instead carry hashes interpreted by the configured key adapter. KVCR deliberately reports only the affected keys, tier, and whether the residency was removed. KVCR delivers each event to `inventory_sink`; the framework or local indexer owns any buffering or batching for publication. Each event may contain multiple keys and is tier-specific because one block may be present in several tiers at once. ### Request-scoped hints -For data shared among engines on the same node, the router can use one KVCR as the source in hints sent to the other local KVCRs. A destination may retain a local copy when replication is worthwhile. Otherwise, the router includes `no_retain` in `submit_hint`. The destination still satisfies any committed framework request; `no_retain` only advises it not to keep an additional KVCR-owned copy after the framework no longer needs it. +A `kv.fetch` hint identifies a source control endpoint and advertised block hashes for one request. `submit_hint` stores this advice and may establish the peer connection, but starts no block transfer or remote pinning. The framework supplies full `BlockKey` values through later `fetch()` or `deliver()` calls with the same request ID; the configured key adapter matches those keys against the hint hashes. Both calls copy without requesting source removal. -KVCR-to-KVCR movement is destination-initiated. A router hint gives the destination the source node and KVCR control endpoint. `mode=copy` retains the source copy, while `mode=move` permits source eviction only after successful completion. Timeout or cancellation keeps the source copy. When the hint has no source control endpoint, the destination uses its local storage information or checks object storage when enabled. The router does not need object-store inventory for locality decisions, but may track object presence to issue source hints. - -`submit_hint` may establish the peer connection, but does not start data movement or remote pinning yet. Each hint in an envelope is parsed and handled independently; proactive fetching, staging, or pinning may be added later as new hint types. `discard_hint` normally reports that the relevant request is over, and also permits an integration to discard the request-scoped hint early. A route-time `submit_hint` may carry a whole `BlockKeyList` from one source. In the future, if necessary, multi-source assembly can be added by splitting the list into multiple hinted operations. - -The router may choose not to send a hint based on its internal cost model. Router hints may later be extended to request proactive copy or move for cache rebalancing, which may require periodic utilization reports from KVCR. The router also considers sending hints for matching TP layouts, meaning prefill-to-prefill, decode-to-decode, and aggregated deployments, which cover many use cases. A TP-independent host representation further removes that constraint. +Failed or conflicting advice does not prevent local cache hits. `discard_hint` clears request advice without cancelling already-started operations. A `kv.fetch` without a request ID does not execute an independent operation. Source-less hints and `no_retain` are unsupported. ### Peer control and liveness @@ -313,7 +309,7 @@ At initialization, an integration selects one built-in policy or supplies an ext Eviction scores are finite floats. LRU favors evicting later sequence keys at comparable recency. Custom policies may use ordering position or ignore it. -Policy decides whether to stage, retain, move, or drop data within KVCR-owned DRAM and downstream storage. Fetch and `no_evict` deposit claims are hard mechanism constraints: while they exist, policy cannot evict the entry. `no_retain` is an advisory router retention hint. If it conflicts with a framework `no_evict` claim, `no_evict` wins until the framework calls `release` with the returned handle; policy may then honor or override `no_retain`. +Policy decides whether to stage, retain, move, or drop data within KVCR-owned DRAM and downstream storage. Fetch and `no_evict` deposit claims prevent eviction while held. `no_retain` is not currently supported. The policy methods are invoked as follows: diff --git a/src/kvcr/hints/fetch.py b/src/kvcr/hints/fetch.py index 18330ea1..6f982d3d 100644 --- a/src/kvcr/hints/fetch.py +++ b/src/kvcr/hints/fetch.py @@ -38,14 +38,9 @@ class _KVFetchHint(msgspec.Struct, frozen=True): def _parse_kv_fetch( action_id: str | None, payload: Mapping[str, object] ) -> _KVFetchHint: - # TODO: Add mode and no_retain when KVCR starts consuming them. if "no_retain" in payload: raise ValueError("no_retain is not currently supported") - mode = payload.get("mode", "copy") - if mode != "copy": - raise ValueError("only copy mode is currently supported") - source = payload.get("source_control_endpoint") if source is None: raise ValueError("source-less hints are not currently supported") diff --git a/tests/unit/test_hints_parser.py b/tests/unit/test_hints_parser.py index df2fa2e9..c95f3f48 100644 --- a/tests/unit/test_hints_parser.py +++ b/tests/unit/test_hints_parser.py @@ -178,11 +178,6 @@ def parse(action_id, payload): {"source_control_endpoint": "tcp://source:1", "block_hashes": [True]}, {"source_control_endpoint": "tcp://source:1", "block_hashes": [-1]}, {"source_control_endpoint": "tcp://source:1", "block_hashes": [1 << 64]}, - { - "source_control_endpoint": "tcp://source:1", - "block_hashes": [1], - "mode": "move", - }, { "source_control_endpoint": "tcp://source:1", "block_hashes": [1], From 76fb194fdc3e289b86966b5a705f06b64fc0376f Mon Sep 17 00:00:00 2001 From: Weiyang Wang Date: Thu, 5 Nov 2026 02:31:24 +0200 Subject: [PATCH 6/8] feat: resolve router hashes through key adapters Add metadata-only peer resolution through the adapter's reverse lookup, with complete keys, layouts, and per-hash failures. Reject self-addressed resolution by process incarnation. The integration owns the reverse index and ambiguity policy; lookup acquires no source-memory claims. Include the full client lifecycle: validate peer, hash set, reply rows, and deadline, then deliver completion on the owner thread through a callback. Neither endpoint depends on hint execution or move state. Extend the existing remote-source/target suites for lookup failures, owner-independent replies, malformed/late responses, and exactly-once completion. Validation: full unit suite, Ruff, and whitespace checks. Assisted-by: OpenAI Codex Assisted-by: Anthropic Claude Signed-off-by: Weiyang Wang --- docs/hash_management.md | 7 ++ src/kvcr/backends/remote_fw_dram.py | 103 ++++++++++++++++++++++++- src/kvcr/backends/remote_ops.py | 107 +++++++++++++++++++++++++- tests/unit/test_kvcr_remote_source.py | 85 +++++++++++++++++++- tests/unit/test_kvcr_remote_target.py | 78 +++++++++++++++++++ 5 files changed, 377 insertions(+), 3 deletions(-) create mode 100644 docs/hash_management.md diff --git a/docs/hash_management.md b/docs/hash_management.md new file mode 100644 index 00000000..08451e65 --- /dev/null +++ b/docs/hash_management.md @@ -0,0 +1,7 @@ +# Hash-addressed cache operations + +The source's `KeyAdapter.decode_from_router` maps each uint64 hash to a full `BlockKey`: `KeyError` reports `not_found`, `ValueError` reports `ambiguous_hash`, and any other exception fails only that hash with `resolution_failed`. The integration owns the reverse map and collision policy; KVCR adds no index. Lookups run on progress and must not block. Full keys remain the identities used for storage and transfers. + +Resolution exchanges keys, ordered piece labels, and sizes in one metadata round trip without pinning memory. Self-addressed resolution fails with `self_source`. + +An adapter may maintain reverse mappings from inventory events. The vLLM adapter requires KV events enabled; colliding full keys reject lookup until unique again, including hashes shared across cache groups. Guard recovery does not replay inventory: reverse mappings must be restored separately. Peers need compatible key conventions and layouts; KVCR does not transform them. diff --git a/src/kvcr/backends/remote_fw_dram.py b/src/kvcr/backends/remote_fw_dram.py index 935ee0d8..2c6a2e29 100644 --- a/src/kvcr/backends/remote_fw_dram.py +++ b/src/kvcr/backends/remote_fw_dram.py @@ -43,12 +43,17 @@ OpHandle, ) from .dangling_ops import _DanglingOps, _SourceWriteStatus +from .local_dram_pools import _LocalDramState from .remote_ops import ( _CONTROL_DECODER, _ControlMessage, _decode_notif, + _HashResolutionOp, _NvtxState, _PeerMessage, + _ResolvedHash, + _ResolveHashesReply, + _ResolveHashesRequest, _SourceWriteOp, _SourceWriteState, _StartWrite, @@ -375,6 +380,8 @@ def poll_main(self, items: Collection[object]) -> None: ) elif isinstance(item, _ProgressUpdate): self._apply_progress_update(item) + elif isinstance(item, _HashResolutionOp): + item.on_complete(item) else: raise TypeError(f"unsupported KVCR main item: {type(item)!r}") if not self._closed: @@ -567,6 +574,12 @@ def _process_control_messages(self, progress: _KVCRProgress) -> bool: self._dangling_ops.handle_probe(progress, payload) elif isinstance(payload, _WriteProbeAck): self._dangling_ops.handle_probe_ack(progress, payload) + elif isinstance(payload, _ResolveHashesRequest): + self._handle_resolve_hashes(progress, payload) + elif isinstance(payload, _ResolveHashesReply): + op = progress._in_flight_ops.get(("resolve_hashes", payload.op_handle)) + if isinstance(op, _HashResolutionOp): + op.receive(payload) handled = True if handled_at is not None: self._record_progress_duration( @@ -574,6 +587,86 @@ def _process_control_messages(self, progress: _KVCRProgress) -> bool: ) return handled + def _handle_resolve_hashes( + self, progress: _KVCRProgress, request: _ResolveHashesRequest + ) -> None: + if not request.sender_control_endpoint: + return + error, rows = None, () + if ( + not math.isfinite(request.remaining_timeout_ms) + or request.remaining_timeout_ms <= 0 + or not request.hashes + or len(set(request.hashes)) != len(request.hashes) + or any(not 0 <= value < 1 << 64 for value in request.hashes) + ): + error = "invalid_request" + elif request.sender_incarnation == self._dangling_ops.incarnation: + # The requester is this process: a destination cannot resolve + # hashes against its own cache, so refuse before any lookup. + error = "self_source" + elif progress._stop_requested: + error = "closed" + else: + try: + rows = self._resolve_hashes(request.hashes) + except Exception: + logger.warning("Hash metadata lookup failed", exc_info=True) + error = "resolution_failed" + self._send_control( + progress, + request.sender_control_endpoint, + _ResolveHashesReply( + op_handle=request.op_handle, + results=rows, + error=error, + # Preserve the address the requester used, including aliases. + sender_control_endpoint=request.source_control_endpoint, + ), + ) + + def _resolve_hashes(self, hashes: Collection[int]) -> tuple[_ResolvedHash, ...]: + """Resolve advisory identities; the later transfer claims source memory.""" + adapter = self._key_adapter + if adapter is None: + return tuple( + _ResolvedHash(value, reason="key_adapter_unavailable") + for value in hashes + ) + core = self._kvcr + rows = [] + sizes_by_layout: dict[tuple[str, ...], tuple[int, ...]] = {} + lookup_failed = False + for block_hash in hashes: + try: + key = adapter.decode_from_router(block_hash) + except KeyError: + rows.append(_ResolvedHash(block_hash, reason="not_found")) + continue + except ValueError: + rows.append(_ResolvedHash(block_hash, reason="ambiguous_hash")) + continue + except Exception: + # An adapter error fails this hash, not the whole request. + if not lookup_failed: + lookup_failed = True + logger.warning("Hash metadata lookup failed", exc_info=True) + rows.append(_ResolvedHash(block_hash, reason="resolution_failed")) + continue + record = core._block_record_map.get(key) + local = record.local_dram if record is not None else None + layout = local.key_layout if local is not None else None + if local is None or local.state is _LocalDramState.DISCARDING or not layout: + rows.append(_ResolvedHash(block_hash, reason="source_not_transferable")) + continue + sizes = sizes_by_layout.get(layout) + if sizes is None: + sizes = sizes_by_layout[layout] = tuple( + core._block_sizes[label.partition(":")[0]] for label in layout + ) + rows.append(_ResolvedHash(block_hash, key, layout, sizes)) + return tuple(rows) + def _send_control( self, progress: _KVCRProgress, @@ -593,7 +686,15 @@ def _send_control( if not isinstance(payload, _TargetMetadataAck): payload.source_control_endpoint = endpoint includes_metadata = ( - not isinstance(payload, (_TargetMetadataAck, _WriteRefused)) + not isinstance( + payload, + ( + _TargetMetadataAck, + _WriteRefused, + _ResolveHashesRequest, + _ResolveHashesReply, + ), + ) and endpoint not in self._metadata_acked_sources and ( isinstance(payload, (_TargetMetadata, _StartWrite)) diff --git a/src/kvcr/backends/remote_ops.py b/src/kvcr/backends/remote_ops.py index d31ab01e..62760fa8 100644 --- a/src/kvcr/backends/remote_ops.py +++ b/src/kvcr/backends/remote_ops.py @@ -3,7 +3,7 @@ """Remote transfer operations and wire-format helpers.""" import logging -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass, field, replace from enum import Enum, auto from itertools import chain @@ -39,6 +39,30 @@ class _StartWrite(_PeerMessage, tag="start_write"): source_incarnation: str | None = None +class _ResolvedHash(msgspec.Struct, frozen=True): + block_hash: int + key: BlockKey | None = None + layout: tuple[str, ...] = () + sizes: tuple[int, ...] = () + reason: str | None = None + + +class _ResolveHashesRequest(_PeerMessage, tag="resolve_hashes"): + """Wire request for full keys and layouts matching advertised hashes.""" + + op_handle: OpHandle + hashes: tuple[int, ...] + remaining_timeout_ms: float + + +class _ResolveHashesReply(_PeerMessage, tag="hashes_resolved"): + """Wire reply carrying resolution results, without block data.""" + + op_handle: OpHandle + results: tuple[_ResolvedHash, ...] + error: str | None = None + + class _TargetMetadata(_PeerMessage, tag="target_metadata"): pass @@ -65,6 +89,8 @@ class _WriteProbeAck(_PeerMessage, tag="write_probe_ack"): _ControlMessage = ( _StartWrite + | _ResolveHashesRequest + | _ResolveHashesReply | _TargetMetadata | _TargetMetadataAck | _WriteRefused @@ -74,6 +100,85 @@ class _WriteProbeAck(_PeerMessage, tag="write_probe_ack"): _CONTROL_DECODER = msgspec.msgpack.Decoder(_ControlMessage) +@dataclass +class _HashResolutionOp(_ProgressOp): + """Local state for a pending resolution exchange; never sent over the wire.""" + + endpoint: str + hashes: tuple[int, ...] + _backend: "_RemoteFWDram" = field(repr=False, compare=False) + # Delivered by owner polling; consumers can enqueue their continuation. + on_complete: Callable[["_HashResolutionOp"], None] = field( + repr=False, compare=False + ) + deadline: float = field(init=False) + response: tuple[_ResolvedHash, ...] | None = None + failure: str | None = None + sent: bool = False + + def __post_init__(self) -> None: + self.deadline = self._backend._kvcr._operation_deadline() + + def progress(self, progress: _KVCRProgress, _event): + now = self._backend._kvcr._clock() + if self.response is not None or self.failure is not None: + return True, True + if progress._stop_requested or now >= self.deadline: + self.failure = "closed" if progress._stop_requested else "timeout" + return True, True + if not self.sent: + self.sent = self._backend._send_control( + progress, + self.endpoint, + _ResolveHashesRequest( + op_handle=self.op_id[1], + hashes=self.hashes, + remaining_timeout_ms=(self.deadline - now) * 1000, + ), + ) + if not self.sent: + self.failure = "control_send_failed" + return True, True + return False, True + return False, False + + def receive(self, reply: _ResolveHashesReply) -> None: + if ( + not self.sent + or reply.sender_control_endpoint != self.endpoint + or self.response is not None + or self.failure is not None + ): + return + if self._backend._kvcr._clock() >= self.deadline: + self.failure = "timeout" + return + if reply.error is not None: + self.failure = ( + reply.error if reply.error and not reply.results else "invalid_response" + ) + return + reply_hashes = {row.block_hash for row in reply.results} + if len(reply.results) != len(self.hashes) or reply_hashes != set(self.hashes): + self.failure = "invalid_response" + return + for row in reply.results: + if row.key is None: + valid = bool(row.reason) and not row.layout and not row.sizes + else: + valid = ( + bool(row.key) + and row.reason is None + and bool(row.layout) + and len(row.layout) == len(row.sizes) + and all(size > 0 for size in row.sizes) + ) + if not valid: + self.failure = "invalid_response" + return + self.response = reply.results + + class _TargetPullState(Enum): START_WRITE = auto() WAITING_WRITE_DONE = auto() diff --git a/tests/unit/test_kvcr_remote_source.py b/tests/unit/test_kvcr_remote_source.py index 93b73bd7..9334862f 100644 --- a/tests/unit/test_kvcr_remote_source.py +++ b/tests/unit/test_kvcr_remote_source.py @@ -18,6 +18,7 @@ FakeTelemetryStats, PendingPrimaryPinning, _buffer_region, + _ConstantHashAdapter, _decode_control_message, _decode_notif, _has_outstanding_operations, @@ -33,7 +34,11 @@ from kvcr import DURATION_METRIC, TRANSFER_BLOCKS_METRIC, TRANSFER_BYTES_METRIC from kvcr.backends.local_dram_pools import _LocalDramResidency, _LocalDramState from kvcr.backends.remote_fw_dram import _RemoteFWDram -from kvcr.backends.remote_ops import _CONTROL_DECODER, _TargetMetadata +from kvcr.backends.remote_ops import ( + _CONTROL_DECODER, + _ResolveHashesRequest, + _TargetMetadata, +) from kvcr.backends.remote_sources import _FwMemResidency, _RemoteSources, _SourcePinOp from kvcr.config import KVCRConfig, LocalDramOptions from kvcr.runtime.core import _BlockRecord, _KVCRCore @@ -42,6 +47,84 @@ from kvcr.types import BlockKey, PinHandle, PinRequestId, RegionDescriptor +@pytest.mark.parametrize( + "case", ["ready", "ambiguous", "stale", "adapter_error", "self", "expired"] +) +def test_hash_resolution_replies_without_owner_poll_or_state_lock(case): + """Source resolution returns advisory metadata without pinning or copying.""" + key = BlockKey(b"full-key-with-group") + adapter = _ConstantHashAdapter() + adapter.decode_from_router = Mock( + side_effect=[ + { + "ambiguous": ValueError("collision"), + "adapter_error": RuntimeError("bug"), + }.get(case, key), + KeyError(999), + ] + ) + agent, pinning = FakeNixlAgent(), FakePrimaryPinning() + agent.state = "DONE" + control = FakeBytesControl("tcp://source:1") + primary, local = ctypes.create_string_buffer(16), ctypes.create_string_buffer(32) + source = _new_kvcr( + agent, + pinning, + control, + name="source", + key_adapter=adapter, + framework_regions=[_buffer_region(primary, 16)], + local_dram=LocalDramOptions([("", ctypes.addressof(local), 32)]), + ) + if case != "stale": + deposited = source.deposit({key: [_mem_descriptor(end_point_name="source")]}) + assert dict(_poll_until(source, bool))[deposited][key].success + backend = source._core._remote_fw_dram + before = len(agent.xfers) + request = _ResolveHashesRequest( + op_handle=7, + hashes=(123, 999), + remaining_timeout_ms=0 if case == "expired" else 1000, + sender_control_endpoint="tcp://target:1", + sender_incarnation=( + backend._dangling_ops.incarnation if case == "self" else "peer-generation" + ), + ) + with source._core._state_lock: + control.incoming.append(msgspec.msgpack.encode(request)) + _wait_until(lambda: bool(control.sent)) + response = _decode_control_message(control.sent[-1][1]) + assert response["type"] == "hashes_resolved" and response["op_handle"] == 7 + assert len(agent.xfers) == before and not pinning.searches + assert "target_regions" not in response and "target_agent_metadata" not in response + if case in ("self", "expired"): + assert response["error"] == ( + "self_source" if case == "self" else "invalid_request" + ) + adapter.decode_from_router.assert_not_called() + return + first, unknown = response["results"] + assert unknown["block_hash"] == 999 and unknown["reason"] == "not_found" + if case == "ready": + assert first == { + "block_hash": 123, + "key": key, + "layout": [""], + "sizes": [16], + "reason": None, + } + else: + assert first["key"] is None + assert ( + first["reason"] + == { + "ambiguous": "ambiguous_hash", + "stale": "source_not_transferable", + "adapter_error": "resolution_failed", + }[case] + ) + + @pytest.mark.parametrize("level", [logging.DEBUG, logging.INFO]) @pytest.mark.parametrize("custom_argument", [False, True]) @pytest.mark.parametrize( diff --git a/tests/unit/test_kvcr_remote_target.py b/tests/unit/test_kvcr_remote_target.py index f03d6435..a3324e6d 100644 --- a/tests/unit/test_kvcr_remote_target.py +++ b/tests/unit/test_kvcr_remote_target.py @@ -34,6 +34,11 @@ ) from kvcr import TRANSFER_BLOCKS_METRIC, TRANSFER_BYTES_METRIC +from kvcr.backends.remote_ops import ( + _HashResolutionOp, + _ResolvedHash, + _ResolveHashesReply, +) from kvcr.config import KVCRConfig, LocalDramOptions, RemoteFWDramOptions from kvcr.types import ( BlockKey, @@ -50,6 +55,79 @@ def _make_block_key(block_hash: bytes, group_idx: int) -> BlockKey: return BlockKey(block_hash + group_idx.to_bytes(4, "big", signed=False)) +@pytest.mark.parametrize( + "case", + [ + "ready", + "wrong_peer", + "wrong_hash", + "duplicate", + "source_error", + "timeout", + "missing_reason", + "invalid_sizes", + ], +) +def test_hash_resolution_client_validates_reply_and_completes_on_owner(case): + """Resolution is usable without independent hints; callbacks run exactly once.""" + agent, control = FakeNixlAgent(), FakeBytesControl("tcp://target:1") + target = _new_kvcr(agent, FakePrimaryPinning(), control) + core = target._core + core._clock = lambda: 0.0 + completed = [] + operation = _HashResolutionOp( + op_id=("resolve_hashes", 7), + keys=frozenset(), + endpoint="tcp://source:1", + hashes=(123,), + _backend=core._remote_fw_dram, + on_complete=completed.append, + ) + core._progress.submit(operation) + _wait_until(lambda: operation.sent) + request = _decode_control_message(control.sent[-1][1]) + assert request["type"] == "resolve_hashes" and request["hashes"] == [123] + assert "target_regions" not in request and "target_agent_metadata" not in request + row = _ResolvedHash(999 if case == "wrong_hash" else 123, b"full-key", ("",), (16,)) + if case == "missing_reason": + row = _ResolvedHash(123) + elif case == "invalid_sizes": + row = _ResolvedHash(123, b"full-key", ("",), ()) + reply = _ResolveHashesReply( + op_handle=7, + results=() + if case == "source_error" + else ((row, row) if case == "duplicate" else (row,)), + error="resolution_failed" if case == "source_error" else None, + sender_control_endpoint=operation.endpoint, + ) + if case == "wrong_peer": + wrong = msgspec.structs.replace(reply, sender_control_endpoint="tcp://other:1") + operation.receive(wrong) + assert operation.response is None and operation.failure is None + if case == "timeout": + core._clock = lambda: operation.deadline + 1 + control.incoming.append(msgspec.msgpack.encode(reply)) + _wait_until(lambda: operation.op_id not in core._progress._in_flight_ops) + assert not completed # Progress completion alone cannot invoke the owner callback. + assert not _poll_until(target, lambda _: bool(completed)) + assert completed == [operation] + assert operation.failure == { + "wrong_hash": "invalid_response", + "duplicate": "invalid_response", + "missing_reason": "invalid_response", + "invalid_sizes": "invalid_response", + "source_error": "resolution_failed", + "timeout": "timeout", + }.get(case) + if operation.failure is None: + assert operation.response == (row,) + control.incoming.append(msgspec.msgpack.encode(reply)) + _wait_until(lambda: not control.incoming) + assert not target.poll_completed() and completed == [operation] + assert not agent.xfers + + class _CopyingWriteAgent(FakeNixlAgent): def transfer(self, handle): self.transfers.append(handle) From ff1cefd0a79dd0b588615e85572cba3f23e05e9d Mon Sep 17 00:00:00 2001 From: Weiyang Wang Date: Fri, 6 Nov 2026 02:20:27 +0200 Subject: [PATCH 7/8] feat: add explicit-source transfers and source-owned moves Add internal full-key fetches from a named peer, with optional source removal. Keep the source through deferred admission and require managed source G2, without acquiring or borrowing framework pins. The public framework API and ordinary copying remain unchanged. Withdraw eligible source copies before transfer; successful writes leave removal pending until readers drain, while failure reopens admission. Honor no_evict, invalidate cached destination hits, and publish actual residency changes. Hash lookup excludes keys pending removal. Extend existing remote tests for exact bytes, cache hits, protection, partial failure, source validation, busy producers, readers, and pin exclusion. Validation: full unit suite, Ruff, and whitespace checks. Signed-off-by: Weiyang Wang --- docs/design_overview.md | 2 +- src/kvcr/backends/local_dram.py | 78 +++++++- src/kvcr/backends/local_dram_pools.py | 101 ++++++++-- src/kvcr/backends/remote_fw_dram.py | 106 ++++++++--- src/kvcr/backends/remote_ops.py | 19 ++ src/kvcr/backends/remote_sources.py | 63 +++++- src/kvcr/runtime/core.py | 89 ++++++++- src/kvcr/runtime/invalidation.py | 163 ++++++++++++++++ src/kvcr/service/recovery_journal.py | 2 + tests/unit/_kvcr_test_utils.py | 7 + tests/unit/test_kvcr_remote_source.py | 8 +- tests/unit/test_kvcr_remote_target.py | 263 ++++++++++++++++++++++++++ 12 files changed, 838 insertions(+), 63 deletions(-) create mode 100644 src/kvcr/runtime/invalidation.py diff --git a/docs/design_overview.md b/docs/design_overview.md index 1d814e75..27d27076 100644 --- a/docs/design_overview.md +++ b/docs/design_overview.md @@ -263,7 +263,7 @@ Failed or conflicting advice does not prevent local cache hits. `discard_hint` c ### Peer control and liveness -The router is never on the data path: KVCR instances execute transfers peer-to-peer via NIXL. Source endpoints remain request-scoped rather than becoming part of global block state, so this coordination adds no KVCR-side peer inventory. +The router is never on the data path: KVCR instances execute transfers peer-to-peer via NIXL. Source endpoints belong to request advice or an independent operation rather than global block state, so this coordination adds no KVCR-side peer inventory. Peer transfers use a separate control channel for connection metadata, acknowledgements, and transfer control; payload bytes move directly through NIXL and never traverse the router or control channel. Peer-protocol versioning and compatibility checks may be added as needed. diff --git a/src/kvcr/backends/local_dram.py b/src/kvcr/backends/local_dram.py index fa9c26e5..e1bb0d39 100644 --- a/src/kvcr/backends/local_dram.py +++ b/src/kvcr/backends/local_dram.py @@ -44,6 +44,43 @@ _Clock = Callable[[], float] +@dataclass +class _NamedSource: + """A fetch bound to one peer. For a move, complete local hits collect in + cached_keys until one invalidation message asks the source to drop them. + """ + + source: str + layout: tuple[str, ...] + keys: set[BlockKey] + move: bool = False + cached_keys: list[BlockKey] = field(default_factory=list) + + def invalidate_cached(self, core: "_KVCRCore") -> None: + """Send one invalidation for the complete local hits found so far.""" + from .remote_ops import _Invalidate + + if not self.move or not self.cached_keys: + return + keys = tuple(self.cached_keys) + self.cached_keys.clear() + try: + core._progress.submit( + _Invalidate( + keys=keys, + key_layout=self.layout, + source_control_endpoint=self.source, + ) + ) + except Exception: + # Removal is best effort; the destination result stands. + logger.warning( + "KVCR cached-key invalidation could not reach %s", + self.source, + exc_info=True, + ) + + @dataclass class _PendingResidencyOp(_Op): deadline: float @@ -59,6 +96,8 @@ class _PendingResidencyOp(_Op): key_layout: tuple[str, ...] | None = None fill_sources: Mapping[BlockKey, CacheTier] = field(default_factory=dict) claimed_keys: set[BlockKey] = field(default_factory=set) + # Explicit peer source retained through deferred fills; see _NamedSource. + named_source: "_NamedSource | None" = None @dataclass @@ -247,9 +286,23 @@ def deposit( key_layout=key_layout, ) self._pending_residency_ops[op.op_id] = op - self._kvcr._add_block_dependencies(op, new_operation=True) + self._admit(op) self._start_deposits(op, blocks) + def _admit(self, op: _PendingResidencyOp) -> None: + """Decide admission once: keys pending removal fail now and never hold a + dependency, so work on other keys cannot delay their removal. Accepted + keys may finish even if removal becomes pending later. + """ + rejected = self._kvcr._invalidation_handler.rejected_entries(op.keys) + if rejected: + op.results.update(rejected) + self._kvcr._add_block_dependencies( + op, + new_operation=True, + keys=op.keys - op.results.keys() if rejected else None, + ) + def _start_deposits( self, op: _PendingResidencyOp, keys: Collection[BlockKey] ) -> None: @@ -394,6 +447,7 @@ def fetch( hints: object | None, layout: list[str], key_layout: tuple[str, ...] | None = None, + named_source: "_NamedSource | None" = None, ) -> dict[BlockKey, list[_TransferRef]]: ordered_keys = tuple(dict.fromkeys(keys)) if not ordered_keys: @@ -409,9 +463,10 @@ def fetch( key_layout=key_layout, fill_sources=sources, hints=hints, + named_source=named_source, ) self._pending_residency_ops[op.op_id] = op - self._kvcr._add_block_dependencies(op, new_operation=True) + self._admit(op) destinations = {} for key in ordered_keys: refs = self._prepare_fetch(op, key) @@ -423,6 +478,8 @@ def fetch( def _prepare_fetch( self, op: _PendingResidencyOp, key: BlockKey ) -> list[_TransferRef] | None: + if key in op.results: + return None layout = op.requested_layout record = self._kvcr._block_record(key) residency = record.local_dram @@ -463,12 +520,25 @@ def _prepare_fetch( ready = [residency.slots[index] for index in indices] self._kvcr._record_access((key,)) op.results[key] = self._pools._new_public_claim(key, residency, slots=ready) + named = op.named_source + if named is not None and named.move and key in named.keys: + named.keys.remove(key) + named.cached_keys.append(key) else: if key not in op.claimed_keys: self._pools._acquire_claim(key, residency) op.claimed_keys.add(key) if source is None or residency.remote_indices is not None: return None + named = op.named_source + if ( + named is not None + and named.move + and key in named.keys + and residency.pending_deposits + ): + # Source removal requires every piece outside this write to be ready. + return None # One remote fill owns these destinations until its native terminal. missing = [ index @@ -783,6 +853,7 @@ def _resolve_fetch_waiter( {key: refs}, residency_op.request_id, residency_op.deadline, + named_source=residency_op.named_source, ) return key in residency_op.results @@ -898,6 +969,9 @@ def _finish_delivery_copy(self, copy: _LocalCopyOp) -> None: self._finish_deliver_if_ready(op) def _finish_residency_if_ready(self, op: _PendingResidencyOp) -> None: + if op.named_source is not None: + # A move's complete local hits ask the source to drop its copies now. + op.named_source.invalidate_cached(self._kvcr) if ( len(op.results) != len(op.keys) or self._pending_residency_ops.get(op.op_id) is not op diff --git a/src/kvcr/backends/local_dram_pools.py b/src/kvcr/backends/local_dram_pools.py index ed10bc8f..8a626c26 100644 --- a/src/kvcr/backends/local_dram_pools.py +++ b/src/kvcr/backends/local_dram_pools.py @@ -15,7 +15,7 @@ from ..config import LocalDramOptions from ..runtime.policy_runtime import _EvictionQueue -from ..runtime.progress import _TransferRef +from ..runtime.progress import _OpId, _TransferRef from ..types import ( BlockKey, CacheTier, @@ -86,6 +86,9 @@ class _LocalDramResidency(msgspec.Struct, gc=False, eq=False): refs: dict[int, _TransferRef] | tuple[_TransferRef, ...] = msgspec.field( default_factory=dict ) + # no_evict deposit claims, included in claim_count. Invalidation honours + # these guarantees but waits for the other, temporary claims to drain. + retention_claims: int = 0 indices_by_label: dict[str, int | None] | None = None _layout: tuple[str, ...] | None = None _pool_counts: tuple[tuple[str, int], ...] | None = None @@ -233,8 +236,9 @@ def __init__( self._evictable_slots: Counter[str] = Counter() self._unscored: set[BlockKey] = set() self._capacity_eviction_key: BlockKey | None = None + # Handle -> (key, residency, retention); retention marks a no_evict claim. self._public_claims: dict[ - ReleaseHandle, tuple[BlockKey, _LocalDramResidency] + ReleaseHandle, tuple[BlockKey, _LocalDramResidency, bool] ] = {} self._next_release_handle = 1 # A no-op until something attaches: the tiers publish residency @@ -269,7 +273,12 @@ def _publish_source( residency: _LocalDramResidency, refs: tuple[_TransferRef, ...] | None = None, ) -> None: - if self._kvcr.framework_control is None or residency.pending_labels: + # Finishing an existing fill must not reopen a source being drained. + if ( + self._kvcr.framework_control is None + or residency.pending_labels + or self._kvcr._block_record_map[key].pending_invalidation + ): return if refs is None: refs = ( @@ -381,6 +390,8 @@ def _drain_source_completions(self) -> None: state.active.pop(key) for key, residency in zip(keys, sources): record = self._kvcr._block_record_map.get(key) + if record is not None and record.pending_invalidation: + self._kvcr._invalidation_handler.notify(key) if ( record is not None and record.local_dram is residency @@ -493,7 +504,9 @@ def release(self, handles: Collection[ReleaseHandle]) -> list[ReleaseResult]: if claim is None: results.append((handle, False)) continue - key, residency = claim + key, residency, retention = claim + if retention: + residency.retention_claims -= 1 self._release_claim(key, residency) results.append((handle, True)) self._update_capacity_pressure() @@ -504,6 +517,7 @@ def acquire_sources( keys: Collection[BlockKey], *, allow_filling: bool = False, + op_id: _OpId | None = None, ) -> dict[BlockKey, _LocalDramResidency]: sources: dict[BlockKey, _LocalDramResidency] = {} for key in keys: @@ -511,8 +525,11 @@ def acquire_sources( continue record = self._kvcr._block_record_map.get(key) residency = record.local_dram if record is not None else None - if residency is None or ( - residency.state is not _LocalDramState.READY + if ( + residency is None + or record.pending_invalidation + and op_id not in (record.in_flight_ops or ()) + or residency.state is not _LocalDramState.READY and not (allow_filling and residency.state is _LocalDramState.FILLING) ): continue @@ -586,7 +603,11 @@ def _new_public_claim( self._acquire_claim(key, residency) handle = ReleaseHandle(self._next_release_handle) self._next_release_handle += 1 - self._public_claims[handle] = (key, residency) + # Fetch passes readable slots; a protected deposit claims only retention. + retention = slots is None + self._public_claims[handle] = (key, residency, retention) + if retention: + residency.retention_claims += 1 return OpEntryResult( OpEntryStatus.SUCCESS, [ @@ -617,6 +638,8 @@ def _release_claim(self, key: BlockKey, residency: _LocalDramResidency) -> None: ): raise RuntimeError(f"invalid local DRAM claim for {key!r}") residency.claim_count -= 1 + if record.pending_invalidation: + self._kvcr._invalidation_handler.notify(key) if residency.claim_count == 0: self._after_last_claim(key, residency) @@ -634,28 +657,64 @@ def _after_last_claim(self, key: BlockKey, residency: _LocalDramResidency) -> bo state.index.pop(key, None) record.local_dram = None if not retiring: - self._make_evictable(key) + if not record.pending_invalidation: + self._make_evictable(key) + # Pending removal owns cleanup; avoid scoring a copy about to drop. return False if residency.state is _LocalDramState.READY: - if self._kvcr._key_history_enabled: - self._kvcr._log_key_history( - "remove", (key,), CacheTier.LOCAL_G2, "retire_on_release" - ) - self._residency_observer(key, record) - self._free(residency.slots) - self.abandon_capacity_eviction(key) - self._kvcr._on_remove( - self._kvcr._block_meta(key, record, self._size_bytes(residency)) - ) - if not residency.pending_labels: - self._kvcr._publish_inventory((key,), CacheTier.LOCAL_G2, removed=True) + self._reclaim_withdrawn(key, record, residency, "retire_on_release") else: self._free(residency.slots) - self._kvcr._prune_block_record(key) + self._kvcr._prune_block_record(key) if not self._kvcr._closed and not self._draining_source_completions: self._resume_capacity_waiters() return True + def _drop(self, key: BlockKey) -> None: + """Reclaim G2 slots after source withdrawal and all readers/operations drain. + Unlike reader release, this removes the copy rather than retaining it. + """ + core = self._kvcr + record = core._block_record_map[key] + residency = record.local_dram + if residency is None or residency.claim_count or record.in_flight_ops: + raise RuntimeError(f"local DRAM entry is not ready to drop {key!r}") + self._remove_evictable(key, residency) + state = self._source_state + # The owner lock excludes index writers; an admitted fill may reinsert None. + if key in state.index: + with state.lock: + state.index.pop(key, None) + record.local_dram = None + self._reclaim_withdrawn(key, record, residency, "invalidate") + + def _reclaim_withdrawn( + self, + key: BlockKey, + record: "_BlockRecord", + residency: _LocalDramResidency, + reason: str, + ) -> None: + """Free a READY copy already withdrawn from the source index and report + its removal.""" + core = self._kvcr + if core._key_history_enabled: + core._log_key_history("remove", (key,), CacheTier.LOCAL_G2, reason) + self._residency_observer(key, record) + self._free(residency.slots) + self.abandon_capacity_eviction(key) + core._prune_block_record(key) + core._on_remove(core._block_meta(key, record, self._size_bytes(residency))) + # A callback's redeposit must not be followed by a stale removal event. + current = core._block_record_map.get(key) + replacement = current.local_dram if current is not None else None + if not residency.pending_labels and not ( + replacement is not None + and replacement.state is _LocalDramState.READY + and not replacement.pending_labels + ): + core._publish_inventory((key,), CacheTier.LOCAL_G2, removed=True) + def _allocate_slots( self, layout: list[str], protected: frozenset[BlockKey], deadline: float ) -> tuple[list[tuple[str, int]] | None, list[BlockKey], bool]: diff --git a/src/kvcr/backends/remote_fw_dram.py b/src/kvcr/backends/remote_fw_dram.py index 2c6a2e29..08c2e1c4 100644 --- a/src/kvcr/backends/remote_fw_dram.py +++ b/src/kvcr/backends/remote_fw_dram.py @@ -49,6 +49,7 @@ _ControlMessage, _decode_notif, _HashResolutionOp, + _Invalidate, _NvtxState, _PeerMessage, _ResolvedHash, @@ -70,6 +71,7 @@ if TYPE_CHECKING: from ..runtime.core import _KVCRCore + from .local_dram import _NamedSource @dataclass(frozen=True) @@ -229,26 +231,37 @@ def _start_target_pull( op_handle: OpHandle, *, local_fill: bool, + named_source: "_NamedSource | None" = None, ) -> bool: kvcr = self._kvcr started_at = kvcr._timer() scope = "remote_fetch" if local_fill else "remote_deliver" - current_hint = ( - self._request_hints.get(request_id) if request_id is not None else None - ) - if current_hint is not None and ( - current_hint.failed or current_hint.missing_keys.issuperset(blocks) - ): - kvcr._record_duration(scope, started_at, "failed") - return False - - if current_hint is None or current_hint.source is None: - kvcr._record_duration(scope, started_at, "failed") - return False - keys = tuple(key for key in blocks if key not in current_hint.missing_keys) - kvcr._record_duration("hint_wait", current_hint.submitted_at, "complete") - if request_id is not None: - self._request_hints[request_id] = replace(current_hint, submitted_at=None) + current_hint = None + if named_source is not None: + source, keys = named_source.source, tuple(blocks) + else: + current_hint = self._request_hints.get(request_id) + if ( + current_hint is None + or current_hint.source is None + or current_hint.failed + or current_hint.missing_keys.issuperset(blocks) + ): + kvcr._record_duration(scope, started_at, "failed") + return False + source = current_hint.source + keys = ( + tuple(key for key in blocks if key not in current_hint.missing_keys) + if current_hint.missing_keys + else tuple(blocks) + ) + if current_hint.submitted_at is not None: + self._request_hints[request_id] = replace( + current_hint, submitted_at=None + ) + kvcr._record_duration( + "hint_wait", current_hint.submitted_at, "complete" + ) op = _TargetPullOp( state=_TargetPullState.START_WRITE, @@ -257,17 +270,23 @@ def _start_target_pull( started_at=started_at, deadline=deadline, op_id=("target", op_handle), - remote_ctrl_ep=current_hint.source, + remote_ctrl_ep=source, _backend=self, ordered_keys=keys, dst_descriptors=tuple(tuple(blocks[key]) for key in keys), - request_id=request_id, + # Explicit operations have no request advice to update on completion. + request_id=request_id if current_hint is not None else None, + move=named_source is not None and named_source.move, + g2_only=named_source is not None, key_layout=( kvcr._block_record_map[keys[0]].local_dram.key_layout if local_fill else () ), ) + if named_source is not None: + # Keys sent natively need no later invalidation message. + named_source.keys.difference_update(keys) if self._nvtx is not None: op.trace = self._nvtx.operation( kvcr, @@ -275,12 +294,14 @@ def _start_target_pull( blocks, request_id, local_fill=local_fill, - hint_trace_id=current_hint.trace.trace_id - if current_hint.trace is not None - else 0, + hint_trace_id=( + current_hint.trace.trace_id + if current_hint is not None and current_hint.trace is not None + else 0 + ), range_name="target.remote", ) - if current_hint.trace is not None: + if current_hint is not None and current_hint.trace is not None: current_hint.trace.mark("hint.used", op_handle=op_handle) if op.trace is not None: op.trace.mark("target.queued") @@ -291,7 +312,7 @@ def _start_target_pull( scope, request_id, op_handle, - current_hint.source, + source, len(keys), kvcr._descriptor_bytes(chain.from_iterable(op.dst_descriptors)), ) @@ -342,6 +363,7 @@ def fetch( deadline: float, *, op_handle: OpHandle, + named_source: "_NamedSource | None" = None, ) -> bool: return self._start_target_pull( blocks, @@ -349,6 +371,7 @@ def fetch( deadline, op_handle, local_fill=True, + named_source=named_source, ) def poll_main(self, items: Collection[object]) -> None: @@ -380,6 +403,11 @@ def poll_main(self, items: Collection[object]) -> None: ) elif isinstance(item, _ProgressUpdate): self._apply_progress_update(item) + elif isinstance(item, _Invalidate): + # The destination holds these keys; drop eligible local copies. + self._kvcr._invalidation_handler.submit( + item.keys, key_layout=item.key_layout + ) elif isinstance(item, _HashResolutionOp): item.on_complete(item) else: @@ -497,6 +525,13 @@ def poll_progress( ): continue self._send_control(progress, item.endpoint, _TargetMetadata()) + elif isinstance(item, _Invalidate): + # One-way: no reply is awaited and no public result depends on it. + if not self._send_control(progress, item.source_control_endpoint, item): + logger.warning( + "KVCR invalidation could not reach %s", + item.source_control_endpoint, + ) else: raise TypeError(f"unsupported KVCR progress item: {type(item)!r}") @@ -564,6 +599,8 @@ def _process_control_messages(self, progress: _KVCRProgress) -> bool: handled_at = self._kvcr._timer() if isinstance(payload, _StartWrite): self._handle_start_write(progress, payload) + elif isinstance(payload, _Invalidate): + self._progress_outbound.append(payload) elif isinstance(payload, _TargetMetadata): self._handle_target_metadata(progress, payload) elif isinstance(payload, _TargetMetadataAck): @@ -656,7 +693,12 @@ def _resolve_hashes(self, hashes: Collection[int]) -> tuple[_ResolvedHash, ...]: record = core._block_record_map.get(key) local = record.local_dram if record is not None else None layout = local.key_layout if local is not None else None - if local is None or local.state is _LocalDramState.DISCARDING or not layout: + if ( + local is None + or record.pending_invalidation + or local.state is _LocalDramState.DISCARDING + or not layout + ): rows.append(_ResolvedHash(block_hash, reason="source_not_transferable")) continue sizes = sizes_by_layout.get(layout) @@ -691,6 +733,7 @@ def _send_control( ( _TargetMetadataAck, _WriteRefused, + _Invalidate, _ResolveHashesRequest, _ResolveHashesReply, ), @@ -925,6 +968,7 @@ def _handle_start_write( dst_descriptors=dst_descriptors, local_sources=(), _backend=self, + g2_only=payload.g2_only, route=(target_agent, self._route_generation.get(target_agent, 0)), nvtx_state=_NvtxState( requested_blocks=len(keys), @@ -933,6 +977,11 @@ def _handle_start_write( if self._nvtx is not None else None, ) + if payload.move: + # A move closes owner-side admission for its keys, so the owner + # thread admits it; the write itself still runs on progress. + self._queue_source_pin(source_write, payload.key_layout, move=True) + return prepared = self._try_local_source_write( progress, source_write, payload.key_layout ) @@ -942,7 +991,7 @@ def _handle_start_write( self._queue_source_pin(source_write, payload.key_layout) def _queue_source_pin( - self, op: _SourceWriteOp, key_layout: tuple[str, ...] + self, op: _SourceWriteOp, key_layout: tuple[str, ...], *, move: bool = False ) -> None: self._progress_outbound.append( _SourcePinOp( @@ -955,8 +1004,10 @@ def _queue_source_pin( ordered_keys=op.source_keys, dst_descriptors=op.dst_descriptors, key_layout=key_layout, + g2_only=op.g2_only, route=op.route, nvtx_state=op.nvtx_state, + move=move, ) ) @@ -1002,11 +1053,12 @@ def _try_local_source_write( ) -> bool | None: """Return None to retry contention, False when caller work is needed.""" local = self._kvcr._local_dram + framework_pinning = self._framework_pinning and not source_write.g2_only # Until prepared, the unregistered write retains the original request. ordered_keys = source_write.source_keys dst_descriptors = source_write.dst_descriptors if local is None: - if self._framework_pinning: + if framework_pinning: return False selected = () else: @@ -1014,7 +1066,7 @@ def _try_local_source_write( ordered_keys, dst_descriptors, key_layout, - allow_missing=not self._framework_pinning, + allow_missing=not framework_pinning, ) if selected is None: return None diff --git a/src/kvcr/backends/remote_ops.py b/src/kvcr/backends/remote_ops.py index 62760fa8..767d062a 100644 --- a/src/kvcr/backends/remote_ops.py +++ b/src/kvcr/backends/remote_ops.py @@ -37,6 +37,9 @@ class _StartWrite(_PeerMessage, tag="start_write"): dst_descriptors: tuple[tuple[_TransferRef, ...], ...] key_layout: tuple[str, ...] = () source_incarnation: str | None = None + move: bool = False + # Independent operations must not acquire or reuse framework memory. + g2_only: bool = False class _ResolvedHash(msgspec.Struct, frozen=True): @@ -63,6 +66,13 @@ class _ResolveHashesReply(_PeerMessage, tag="hashes_resolved"): error: str | None = None +class _Invalidate(_PeerMessage, tag="invalidate"): + """Remove eligible source G2 copies already held completely at the destination.""" + + keys: tuple[BlockKey, ...] + key_layout: tuple[str, ...] + + class _TargetMetadata(_PeerMessage, tag="target_metadata"): pass @@ -89,6 +99,7 @@ class _WriteProbeAck(_PeerMessage, tag="write_probe_ack"): _ControlMessage = ( _StartWrite + | _Invalidate | _ResolveHashesRequest | _ResolveHashesReply | _TargetMetadata @@ -216,6 +227,8 @@ class _TargetPullOp(_RemoteOp): dst_descriptors: tuple[tuple[_TransferRef, ...], ...] = () request_id: str | None = None key_layout: tuple[str, ...] = () + move: bool = False + g2_only: bool = False success: bool = False completed_keys: set[BlockKey] = field(default_factory=set) probe_sent: bool = False @@ -285,6 +298,8 @@ def _progress( keys=self.ordered_keys, dst_descriptors=self.dst_descriptors, key_layout=self.key_layout if self.local_fill else (), + move=self.move, + g2_only=self.g2_only, ), ) if not sent: @@ -473,11 +488,15 @@ class _SourceWriteOp(_RemoteOp): dst_descriptors: tuple[tuple[_TransferRef, ...], ...] _backend: "_RemoteFWDram" = field(repr=False, compare=False) src_descriptors: tuple[tuple[_TransferRef, ...], ...] = () + # Retain the source restriction across progress retries and owner admission. + g2_only: bool = False transfer_id: int | None = None success: bool = False completed_indices: tuple[int, ...] = () route: tuple[str, int] = ("", 0) nvtx_state: _NvtxState | None = None + # Keys this move marked pending at the source; reopened if the write fails. + moving: tuple[BlockKey, ...] = () # None uses caller cleanup; a tuple (even empty) is owned by progress. local_sources: tuple["_LocalDramResidency", ...] | None = None diff --git a/src/kvcr/backends/remote_sources.py b/src/kvcr/backends/remote_sources.py index 4d604cc0..1ca324c0 100644 --- a/src/kvcr/backends/remote_sources.py +++ b/src/kvcr/backends/remote_sources.py @@ -38,6 +38,10 @@ class _SourcePinOp(_Op): dst_descriptors: tuple[tuple[_TransferRef, ...], ...] route: tuple[str, int] = ("", 0) key_layout: tuple[str, ...] = () + move: bool = False + g2_only: bool = False + # Keys this move marked pending; reopened if they are not written. + moving: tuple[BlockKey, ...] = () framework_pins: set[PinHandle] = field(default_factory=set) pending_pin_ids: set[PinRequestId] = field(default_factory=set) framework_acquire_attempted: bool = False @@ -93,6 +97,8 @@ def finish_source_write(self, item: _SourceWriteOp) -> None: # KVCR-owned contents stay claimed until native quiescence. pins = self._fw_pins_by_op.pop(item.op_id, None) if item.state is _SourceWriteState.FINISHED: + if item.moving and not item.success: + self._kvcr._invalidation_handler.reopen(item.moving) self._kvcr._remove_block_dependencies(item) if item.success: self._kvcr._record_access( @@ -166,7 +172,8 @@ def _submit_prepared_source_write( framework_sources = { key: record.fw_mem.descriptors for key in source_pin.keys - if key not in local_sources + if not source_pin.g2_only + and key not in local_sources and (record := kvcr._block_record_map.get(key)) is not None and record.fw_mem is not None } @@ -177,6 +184,8 @@ def _submit_prepared_source_write( if key not in source_pin.keys: continue residency = local_sources.get(key) + if source_pin.g2_only and (residency is None or residency.pending_labels): + continue source = framework_sources.get(key) destination = source_pin.dst_descriptors[index] if residency is None and source is None: @@ -226,6 +235,14 @@ def _submit_prepared_source_write( ) retained_keys = source_pin.keys if waiting else set(completed_keys) + if source_pin.moving: + omitted = [key for key in source_pin.moving if key not in retained_keys] + # Forget rejected keys before reopening them; later retries must not + # cancel a newer invalidation for the same key. + source_pin.moving = tuple( + key for key in source_pin.moving if key in retained_keys + ) + kvcr._invalidation_handler.reopen(omitted) kvcr._release_local_dram_sources(op_id, local_sources.keys() - retained_keys) relevant_pins = { record.fw_mem.pin_handle @@ -267,9 +284,11 @@ def _submit_prepared_source_write( route=source_pin.route, nvtx_state=source_pin.nvtx_state, _backend=self._backend, + g2_only=source_pin.g2_only, source_keys=completed_keys, src_descriptors=tuple(selected_sources), completed_indices=tuple(completed_indices), + moving=source_pin.moving, ) kvcr._add_block_dependencies(source_write, new_operation=True) self._fw_pins_by_op[source_write.op_id] = framework_pins @@ -383,9 +402,44 @@ def _expire_source_pin(self, op_id: _OpId, op: _SourcePinOp) -> None: def start_source_pin(self, op: _SourcePinOp) -> None: kvcr = self._kvcr + # Decide admission once; retries retain the original request positions. + pending = kvcr._invalidation_handler._pending + if pending: + op.keys = op.keys.difference(pending) + if op.move and kvcr._local_dram is not None: + # Closing admission must not strand a producer of missing pieces. + # Reject busy keys before claiming them or marking removal pending. + busy = set() + producers = kvcr._local_dram._pending_residency_ops + for key in op.keys: + record = kvcr._block_record_map.get(key) + if record is None: + continue + residency = record.local_dram + if residency is not None: + if residency.pending_labels or residency.pending_deposits: + busy.add(key) + elif any( + producer_id[0] == "deposit" + and (producer := producers.get(producer_id)) is not None + and key not in producer.results + for producer_id in record.active_op_ids + ): + busy.add(key) + if busy: + op.keys = op.keys.difference(busy) kvcr._add_block_dependencies(op, new_operation=True) + if op.move: + # Reject new reads and deposits of these keys while they move. + op.moving = kvcr._invalidation_handler.submit( + op.keys, key_layout=op.key_layout + ) self._source_pin_ops[op.op_id] = op - self._resume_source_pin(op.op_id, op) + if op.ordered_keys and not op.keys: + # Every requested key was rejected; send failure without a transfer. + self._submit_prepared_source_write(op.op_id, op, force_failure=True) + else: + self._resume_source_pin(op.op_id, op) def _resume_source_pin(self, op_id: _OpId, op: _SourcePinOp) -> None: if self._source_pin_ops.get(op_id) is not op: @@ -397,6 +451,10 @@ def _resume_source_pin(self, op_id: _OpId, op: _SourcePinOp) -> None: ) return + if op.g2_only: + self._submit_prepared_source_write(op_id, op) + return + local_sources = kvcr._claim_local_dram_sources( op.op_id, op.keys, allow_filling=True ) @@ -427,6 +485,7 @@ def _resume_source_pin(self, op_id: _OpId, op: _SourcePinOp) -> None: def _remove_source_pin(self, op_id: _OpId, op: _SourcePinOp) -> None: self._cancel_pending_pin_for_op(op_id, op) self._source_pin_ops.pop(op_id, None) + self._kvcr._invalidation_handler.reopen(op.moving) self._kvcr._remove_block_dependencies(op) self._kvcr._release_local_dram_sources(op.op_id) framework_pins = set(op.framework_pins) diff --git a/src/kvcr/runtime/core.py b/src/kvcr/runtime/core.py index 715c1b4f..329335f1 100644 --- a/src/kvcr/runtime/core.py +++ b/src/kvcr/runtime/core.py @@ -18,7 +18,7 @@ from .. import _nvtx from ..backends.local_disk import _G3, _G3Residency -from ..backends.local_dram import _LocalDram +from ..backends.local_dram import _LocalDram, _NamedSource from ..backends.local_dram_pools import ( _layout_indices, _LocalDramResidency, @@ -49,6 +49,7 @@ ReleaseHandle, ReleaseResult, ) +from .invalidation import _InvalidationHandler from .policy_runtime import _PolicyInvoker from .progress import ( _FAILED_RESULT, @@ -139,6 +140,8 @@ class _BlockRecord(msgspec.Struct, gc=False): access_count: int = 0 last_access: float | None = None position: int = -1 + # Close G2 to new reads; retain the record until existing users drain. + pending_invalidation: bool = False def add_in_flight_op(self, op_id: _OpId) -> None: if self.in_flight_ops is None: @@ -289,6 +292,7 @@ def __init__( from ..backends.dangling_ops import _log_resilience_event from ..backends.remote_fw_dram import _RemoteFWDram + self._invalidation_handler = _InvalidationHandler(self) self._on_resilience_event_callback = ( _log_resilience_event if bindings.on_resilience_event is None @@ -435,6 +439,7 @@ def _has_complete_local(self, key: BlockKey) -> bool: residency = record.local_dram if record is not None else None return ( residency is not None + and not record.pending_invalidation and residency.state is _LocalDramState.READY and not residency.pending_labels and len(residency.slots) == len(residency.key_layout) @@ -456,6 +461,9 @@ def query( remote_keys = self._remote_delivery_keys.get(request_id, ()) for key in keys: record = self._block_record_map.get(key) + if record is not None and record.pending_invalidation: + statuses.append((QueryStatus.MISS, None)) + continue residency = record.local_dram if record is not None else None if key in remote_keys: fetching_locally = ( @@ -582,9 +590,14 @@ def deliver( local_blocks: dict[BlockKey, list[_TransferRef]] = {} g3_blocks: dict[BlockKey, list[_TransferRef]] = {} remote_blocks: dict[BlockKey, list[_TransferRef]] = {} + # Record pending-removal failures while other keys proceed. + rejected: dict[BlockKey, OpEntryResult] = {} remote_keys = self._remote_delivery_keys.get(request_id, ()) for key, destination in normalized.items(): - if key in remote_keys and ( + record = self._block_record_map.get(key) + if record is not None and record.pending_invalidation: + rejected[key] = _FAILED_RESULT + elif key in remote_keys and ( not self._has_complete_local(key) or self._remote_fw_dram.query(key, request_id) ): @@ -603,8 +616,13 @@ def deliver( remote_blocks=len(remote_blocks), g3_blocks=len(g3_blocks), ) - if sum(map(bool, (local_blocks, g3_blocks, remote_blocks))) > 1: + if ( + sum(map(bool, (local_blocks, g3_blocks, remote_blocks, rejected))) + > 1 + ): self._joined_completions[op_handle] = (set(blocks), {}) + if rejected: + self._complete(op_handle, rejected) if local_dram is not None and local_blocks: local_dram.deliver(op_handle, local_blocks, deadline=deadline) if g3_blocks and ( @@ -671,6 +689,7 @@ def deposit( ) else: self._local_dram._pools._drain_source_completions() + self._drain_invalidation() if self._g3 is not None: for key, refs in blocks.items(): if self._g3.is_ready(key): @@ -700,7 +719,15 @@ def fetch( request_id: str | None = None, expected_layout: list[str] | None = None, hints: object | None = None, + *, + source_endpoint: str | None = None, + move: bool = False, ) -> OpHandle: + if move: + if not isinstance(source_endpoint, str) or not source_endpoint: + raise ValueError("move source endpoint must be a non-empty string") + if source_endpoint == getattr(self.framework_control, "endpoint", None): + raise ValueError("move source must not be this node") declaration = self._validate_expected_layout( [""] if expected_layout is None else expected_layout ) @@ -718,6 +745,12 @@ def fetch( return op_handle local_dram._pools._drain_source_completions() + self._drain_invalidation() + named_source = ( + _NamedSource(source_endpoint, declaration, set(ordered_keys), move) + if source_endpoint is not None + else None + ) sources = {} for key in ordered_keys: record = self._block_record_map.get(key) @@ -727,7 +760,9 @@ def fetch( and residency.ready_indices(requested) is not None ): continue - if self._g3 is not None and self._g3.is_ready(key): + if named_source is not None: + sources[key] = CacheTier.REMOTE_G2 + elif self._g3 is not None and self._g3.is_ready(key): if not self._g3.is_ready(key, declaration): raise ValueError( "fetch layout must match existing G3 key layout" @@ -738,6 +773,12 @@ def fetch( ): sources[key] = CacheTier.REMOTE_G2 deadline = self._operation_deadline() + if move: + # Accept delayed completion using the existing abandonment budget. + # The source independently caps its native transfer timeout. + deadline += ( + self.config.abandon_timeout_ms - self.config.operation_timeout_ms + ) / 1000 destinations = local_dram.fetch( op_handle, ordered_keys, @@ -747,6 +788,7 @@ def fetch( hints=hints, layout=requested, key_layout=declaration, + named_source=named_source, ) for source in (CacheTier.G3, CacheTier.REMOTE_G2): self._start_local_fill( @@ -758,6 +800,7 @@ def fetch( }, request_id, deadline, + named_source=named_source, ) return op_handle @@ -765,7 +808,9 @@ def release(self, handles: Collection[ReleaseHandle]) -> list[ReleaseResult]: if self._local_dram is None: return [(handle, False) for handle in handles] with self._state_lock: - return self._local_dram._pools.release(handles) + results = self._local_dram._pools.release(handles) + self._drain_invalidation() + return results # TODO: Expose individual entry completions as they become available. def poll_completed(self) -> Iterable[OpResult]: @@ -780,6 +825,7 @@ def poll_completed(self) -> Iterable[OpResult]: if self._local_dram is not None: progress_items = self._local_dram.poll_main(progress_items) self._remote_fw_dram.poll_main(progress_items) + self._drain_invalidation() completed = self._completion_queue self._completion_queue = [] if self._nvtx_operations: @@ -789,6 +835,12 @@ def poll_completed(self) -> Iterable[OpResult]: trace.returned() return completed + def _drain_invalidation(self) -> None: + # Cheap when nothing is pending; otherwise a released or finished user + # lets its source copy drop (and free capacity) without waiting for a poll. + if self._invalidation_handler._dirty: + self._invalidation_handler.drain() + def _notify_transfer_errors(self, progress_items: list[object]) -> None: # Notify before cleanup can return a buffer to its allocator. Retain the # batch if user code raises, consuming each notification exactly once. @@ -898,6 +950,8 @@ def run(stage: Callable[[], None]) -> None: ) run(self._close_hint_handlers) + # Native work has stopped; abandon pending removals without waiting. + run(self._invalidation_handler.close) if self._g3 is not None: run(self._g3.close_main) if self._local_dram is not None: @@ -1028,10 +1082,17 @@ def _update_capacity_pressure(self, reclaimable_slots: Mapping[str, int]) -> Non # Generic operation/block dependency bookkeeping. - def _add_block_dependencies(self, op: _Op, *, new_operation: bool) -> None: + def _add_block_dependencies( + self, + op: _Op, + *, + new_operation: bool, + keys: Collection[BlockKey] | None = None, + ) -> None: if new_operation: self._outstanding_operations += 1 - for key in op.keys: + # Exclude rejected keys so work on other keys cannot delay their removal. + for key in op.keys if keys is None else keys: self._block_record(key).add_in_flight_op(op.op_id) def _remove_block_dependencies(self, op: _Op) -> None: @@ -1050,6 +1111,8 @@ def _start_local_fill( blocks: Mapping[BlockKey, list[_TransferRef]], request_id: str | None, deadline: float, + *, + named_source: _NamedSource | None = None, ) -> None: if not blocks: return @@ -1062,8 +1125,13 @@ def _start_local_fill( deadline, ) elif source is CacheTier.REMOTE_G2: + # Independent operations retain their named source across deferral. started = self._remote_fw_dram.fetch( - blocks, request_id, deadline, op_handle=fill_handle + blocks, + request_id, + deadline, + op_handle=fill_handle, + named_source=named_source, ) else: raise ValueError(f"unsupported local fill source {source.value}") @@ -1082,6 +1150,7 @@ def _claim_local_dram_sources( claimed = self._local_dram._pools.acquire_sources( tuple(key for key in keys if key not in sources), allow_filling=allow_filling, + op_id=op_id, ) if claimed: sources.update(claimed) @@ -1217,6 +1286,7 @@ def _is_local_resident(self, key: BlockKey) -> bool: record = self._block_record_map.get(key) return ( record is not None + and not record.pending_invalidation and record.local_dram is not None and record.local_dram.state is not _LocalDramState.DISCARDING ) @@ -1270,12 +1340,15 @@ def _record_access(self, keys: Collection[BlockKey]) -> None: def _prune_block_record(self, key: BlockKey) -> None: record = self._block_record_map.get(key) + if record is not None and record.pending_invalidation: + self._invalidation_handler.notify(key) if ( record is not None and record.fw_mem is None and record.local_dram is None and record.g3 is None and not record.in_flight_ops + and not record.pending_invalidation ): self._block_record_map.pop(key) diff --git a/src/kvcr/runtime/invalidation.py b/src/kvcr/runtime/invalidation.py new file mode 100644 index 00000000..d682eeeb --- /dev/null +++ b/src/kvcr/runtime/invalidation.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Drop moved source G2 copies after readers and transfers drain.""" + +from collections.abc import Collection +from typing import TYPE_CHECKING + +from ..types import BlockKey, OpEntryResult +from .progress import _FAILED_RESULT + +if TYPE_CHECKING: + from .core import _BlockRecord, _KVCRCore + + +class _InvalidationHandler: + """Remove sources under the owner lock after existing readers drain. + Submission rejects new reads and deposits; a failed move reopens its keys. + """ + + def __init__(self, core: "_KVCRCore") -> None: + self._core = core + # The layout the destination holds; None means unrestricted removal. + self._pending: dict[ + BlockKey, tuple["_BlockRecord", tuple[str, ...] | None] + ] = {} + # Recheck only on submission or release, not every idle poll. + self._dirty: dict[BlockKey, None] = {} + + def has_pending(self) -> bool: + return bool(self._pending) + + def rejected_entries( + self, keys: Collection[BlockKey] + ) -> dict[BlockKey, OpEntryResult]: + """Check new requests once; admitted operations may resume during removal.""" + if not self._pending: + return {} + return {key: _FAILED_RESULT for key in keys if key in self._pending} + + def close(self) -> None: + """Abandon pending removals after progress stops.""" + for record, _ in self._pending.values(): + record.pending_invalidation = False + self._core._outstanding_operations -= len(self._pending) + self._pending.clear() + self._dirty.clear() + + def submit( + self, + keys: Collection[BlockKey], + *, + key_layout: tuple[str, ...] | None = None, + ) -> tuple[BlockKey, ...]: + """Queue removal of G2 copies held in key_layout; return the keys marked. + + Keys protected by no_evict or held in another complete layout are left + alone. Keys without a G2 copy are skipped too unless an admitted deposit + or fetch may still create one: framework and G3 copies are never touched. + """ + core = self._core + local_dram = core._local_dram + if local_dram is None: + return () + state = local_dram._pools._source_state + marked: list[tuple[BlockKey, "_BlockRecord"]] = [] + for key in keys: + record = core._block_record_map.get(key) + if record is None or key in self._pending: + continue + local = record.local_dram + if local is None: + # A batch dependency can outlive this key's terminal result. + if not any( + (op := local_dram._pending_residency_ops.get(op_id)) is not None + and key not in op.results + for op_id in record.active_op_ids + ): + continue + elif local.retention_claims or ( + key_layout is not None and local.key_layout != key_layout + ): + continue + self._pending[key] = record, key_layout + core._outstanding_operations += 1 + self._dirty[key] = None + marked.append((key, record)) + if marked: + # Close admission atomically with progress acquiring readers, once + # per batch. No slot reclamation or callbacks run under this lock. + with state.lock: + for key, record in marked: + record.pending_invalidation = True + state.index.pop(key, None) + return tuple(key for key, _ in marked) + + def reopen(self, keys: Collection[BlockKey]) -> None: + """Cancel a failed move's own pending keys; their copies stay readable.""" + for key in keys: + pending = self._pending.get(key) + if pending is not None: + self._settle(key, pending[0], drop=False) + + def _settle(self, key: BlockKey, record: "_BlockRecord", *, drop: bool) -> bool: + """Clear the pending mark, then drop or republish the copy; True if + capacity changed.""" + core = self._core + pools = core._local_dram._pools + # Clear intent before callbacks can redeposit this key or poll again. + del self._pending[key] + self._dirty.pop(key, None) + record.pending_invalidation = False + core._outstanding_operations -= 1 + local = record.local_dram + if local is None: + core._prune_block_record(key) + return False + if drop: + pools._drop(key) + return True + pools._publish_source(key, local) + if local.retention_claims or local.source_readers: + return False + pools._make_evictable(key) + return True + + def notify(self, key: BlockKey) -> None: + if key in self._pending: + self._dirty[key] = None + + @staticmethod + def _is_ready(record: "_BlockRecord") -> bool: + # Submission withdrew the source under the reader lock; its reader count + # can now only decrease during owner-side completion processing. + local = record.local_dram + return not ( + record.in_flight_ops + or local is not None + and (local.claim_count > local.retention_claims or local.source_readers) + ) + + def drain(self) -> None: + pools = self._core._local_dram._pools + dirty, self._dirty = self._dirty, {} + capacity_changed = False + for key in dirty: + # A preceding callback may have removed this key. + pending = self._pending.get(key) + if pending is None: + continue + record, layout = pending + if not self._is_ready(record): + continue + local = record.local_dram + # An earlier producer can grant retention or create a layout the + # destination never held; keep that copy and reopen admission. + keep = local is not None and ( + local.retention_claims + or (layout is not None and local.key_layout != layout) + ) + capacity_changed |= self._settle(key, record, drop=not keep) + if capacity_changed: + pools._resume_capacity_waiters() + pools._update_capacity_pressure() diff --git a/src/kvcr/service/recovery_journal.py b/src/kvcr/service/recovery_journal.py index 0274061a..d7452393 100644 --- a/src/kvcr/service/recovery_journal.py +++ b/src/kvcr/service/recovery_journal.py @@ -394,10 +394,12 @@ def adopt(self, records: dict[BlockKey, _BlockRecord]) -> None: local_dram.claim_count = 0 local_dram.source_readers = 0 local_dram.retire_on_release = False + local_dram.retention_claims = 0 # References name the previous agent, unlike the slot layout. local_dram.refs = {} if record.g3 is not None: record.g3.claim_count = 0 + record.pending_invalidation = False record.fw_mem = None record.in_flight_ops = None record.access_count = 0 diff --git a/tests/unit/_kvcr_test_utils.py b/tests/unit/_kvcr_test_utils.py index 04fef6e2..a602a261 100644 --- a/tests/unit/_kvcr_test_utils.py +++ b/tests/unit/_kvcr_test_utils.py @@ -155,6 +155,13 @@ def create_agent(name, *_): yield +def _invalidate(kvcr: KVCR, keys: Collection[BlockKey]) -> None: + """Exercise the internal invalidation mechanism under its owner lock.""" + core = kvcr._core + with core._state_lock: + core._invalidation_handler.submit(keys) + + def _poll_until( kvcr: KVCR, predicate, diff --git a/tests/unit/test_kvcr_remote_source.py b/tests/unit/test_kvcr_remote_source.py index 9334862f..4de4a68c 100644 --- a/tests/unit/test_kvcr_remote_source.py +++ b/tests/unit/test_kvcr_remote_source.py @@ -1567,10 +1567,13 @@ def release_xfer_handle(self, handle): [("abc", (0, 1, 2)), ("ca", (2, 0)), ("x", ())], ids=["whole", "partial", "unknown"], ) +@pytest.mark.parametrize("g2_only", [False, True]) def test_a_resumed_write_holds_a_pin_another_operation_acquired( - parts, selected + parts, selected, g2_only ) -> None: - """fw_mem belongs to the block, so two writes can want the same pin.""" + """fw_mem belongs to the block; G2-only writes must not borrow its pin.""" + if g2_only: + selected = () backend = object.__new__(_RemoteFWDram) kvcr = object.__new__(_KVCRCore) kvcr._block_record_map = {} @@ -1608,6 +1611,7 @@ def test_a_resumed_write_holds_a_pin_another_operation_acquired( op_id=("source", 1), keys=frozenset((key,)), framework_pins={stale}, + g2_only=g2_only, ) backend._sources._source_pin_ops[("source", 1)] = waiting diff --git a/tests/unit/test_kvcr_remote_target.py b/tests/unit/test_kvcr_remote_target.py index a3324e6d..9c0aff4f 100644 --- a/tests/unit/test_kvcr_remote_target.py +++ b/tests/unit/test_kvcr_remote_target.py @@ -6,6 +6,7 @@ import hashlib import logging import time +from collections import Counter from concurrent.futures import ThreadPoolExecutor from unittest.mock import DEFAULT, Mock @@ -22,6 +23,7 @@ _decode_control_message, _decode_notif, _has_outstanding_operations, + _invalidate, _key_history, _mem_descriptor, _new_kvcr, @@ -1752,3 +1754,264 @@ def test_remote_framework_dram_transfers_available_keys( assert any( "target_transfer_completed" in event and completed in event for event in events ) + + +class _Peer: + def __init__(self, name): + self.name = name + self.primary = ctypes.create_string_buffer(b"a" * 16 + b"b" * 16, 32) + self.local = ctypes.create_string_buffer(32) + self.agent = _CopyingWriteAgent(metadata=name.encode()) + self.control = FakeBytesControl(f"tcp://{name}:1") + self.events, self.results = [], {} + self.kvcr = _new_kvcr( + self.agent, + FakePrimaryPinning(), + self.control, + name=name, + framework_regions=[ + _buffer_region(self.primary, 16), + _buffer_region(self.primary, 16, ":*"), + ], + local_dram=LocalDramOptions([("", ctypes.addressof(self.local), 32)]), + inventory_sink=self.events.append, + ) + self.kvcr._core._clock = lambda: 0.0 + + def ref(self, index=0, label=""): + return _mem_descriptor(index, label=label, end_point_name=self.name) + + def deposit(self, key, index=0, protected=False): + self.agent.state = "DONE" + handle = self.kvcr.deposit({key: [self.ref(index)]}, no_evict=protected) + result = dict(_poll_until(self.kvcr, bool))[handle][key] + self.agent.state = "PROC" + assert result.success + return result.release_handle + + +class _Pair: + """Use real buffers; forward control and completed fake-native notifications.""" + + def __init__(self): + self.source, self.target = _Peer("source"), _Peer("target") + self.offsets = [[0, 0], [0, 0]] + self.delivered = set() + + def pump(self): + for index, (peer, other) in enumerate( + ((self.source, self.target), (self.target, self.source)) + ): + peer.results.update(peer.kvcr.poll_completed()) + sent, notified = self.offsets[index] + other.control.incoming.extend(raw for _, raw in peer.control.sent[sent:]) + other.agent.notifs.setdefault("peer", []).extend( + raw for _, raw in peer.agent.sent_notifs[notified:] + ) + self.offsets[index] = [len(peer.control.sent), len(peer.agent.sent_notifs)] + for transfer, data in enumerate(peer.agent.xfers, 1): + token = index, transfer + if ( + token not in self.delivered + and transfer in peer.agent.released_xfers + ): + self.delivered.add(token) + if data[5] and peer.agent.state == "DONE": + other.agent.notifs.setdefault("peer", []).append(data[5]) + time.sleep(0.001) + + def until(self, predicate): + deadline = time.monotonic() + 3 + while time.monotonic() < deadline: + self.pump() + if predicate(): + return + raise AssertionError("move did not reach the expected state") + + def idle(self): + return not any( + _has_outstanding_operations(peer.kvcr) + for peer in (self.source, self.target) + ) + + +@pytest.fixture +def pair(): + pair = _Pair() + yield pair + pair.source.agent.state = pair.target.agent.state = "DONE" + pair.until(pair.idle) + + +@pytest.mark.parametrize("endpoint", [None, "", 123, "self"]) +def test_move_rejects_invalid_source_before_admission(pair, endpoint): + """A cached self-move must never enqueue removal of its own copy.""" + target = pair.target + key = BlockKey(b"cached") + target.deposit(key) + core = target.kvcr._core + before = core._next_op_handle + if endpoint == "self": + endpoint = target.control.endpoint + with pytest.raises(ValueError, match="move source"): + core.fetch((key,), source_endpoint=endpoint, move=True) + assert core._next_op_handle == before and not target.control.sent + assert target.kvcr.query((key,)) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] + assert core._block_record_map[key].local_dram.claim_count == 0 + + +@pytest.mark.parametrize( + "cached,protected,succeeds", + [ + (0, False, True), + (0, False, False), + (1, False, True), + (1, False, False), + (2, False, True), + (0, True, True), + (1, True, True), + (2, True, True), + ], +) +def test_move_bytes_removal_protection_and_cache_hits( + pair, cached, protected, succeeds +): + source, target = pair.source, pair.target + keys = (BlockKey(b"first"), BlockKey(b"second")) + retentions = [source.deposit(key, i, protected) for i, key in enumerate(keys)] + for index, key in enumerate(keys[:cached]): + target.deposit(key, index) + source.events.clear() + target.events.clear() + before = len(source.agent.xfers) + handle = target.kvcr._core.fetch( + keys, source_endpoint=source.control.endpoint, move=True + ) + pair.until(lambda: len(source.agent.xfers) > before or handle in target.results) + if not protected: + pair.until(lambda: source.kvcr.query(keys) == [(QueryStatus.MISS, None)] * 2) + else: + assert source.kvcr.query(keys) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] * 2 + source.agent.state = "DONE" if succeeds else "ERR" + pair.until(lambda: handle in target.results) + # Failed writes report first, then finish native cleanup before reopening. + source.agent.state = "DONE" + pair.until(pair.idle) + results = target.results[handle] + successful = [key for i, key in enumerate(keys) if i < cached or succeeds] + assert [key for key in keys if results[key].success] == successful + assert source.kvcr.query(keys) == [ + (QueryStatus.MISS, None) + if key in successful and not protected + else (QueryStatus.HIT, CacheTier.LOCAL_G2) + for key in keys + ] + for index, key in enumerate(keys): + result = results[key] + if result.success: + offset = result.descriptors[0].element_index * 16 + assert ( + target.local.raw[offset : offset + 16] + == (b"a" if index == 0 else b"b") * 16 + ) + assert target.kvcr.release((result.release_handle,)) == [ + (result.release_handle, True) + ] + else: + assert result.release_handle is None + for peer, removed, expected in ( + (source, True, successful if not protected else []), + (target, False, keys[cached:] if succeeds else []), + ): + assert Counter( + key for e in peer.events if e.removed is removed for key in e.keys + ) == Counter(expected) + retained = [handle for handle in retentions if handle is not None] + assert all(ok for _, ok in source.kvcr.release(retained)) + + +@pytest.mark.parametrize("protected", [False, True]) +def test_pending_removal_rejects_new_work_and_frees_after_reader(pair, protected): + source = pair.source + kvcr = source.kvcr + key = BlockKey(b"reader") + source.deposit(key) + source.deposit(BlockKey(b"other"), 1) + fetched = kvcr.fetch((key,)) + claim = dict(kvcr.poll_completed())[fetched][key].release_handle + _invalidate(kvcr, (key, key)) + assert kvcr.query((key,)) == [(QueryStatus.MISS, None)] + handles = [ + kvcr.fetch((key,)), + kvcr.deliver({key: [source.ref(1)]}), + kvcr.deposit({key: [source.ref()]}, no_evict=protected), + ] + results = dict(kvcr.poll_completed()) + assert all(not results[handle][key].success for handle in handles) + assert kvcr._core._invalidation_handler.has_pending() + assert source.local.raw == b"a" * 16 + b"b" * 16 + assert not any(event.removed for event in source.events) + assert kvcr.release((claim,)) == [(claim, True)] + assert key not in kvcr._core._block_record_map + assert not kvcr._core._invalidation_handler.has_pending() + assert Counter( + k for event in source.events if event.removed for k in event.keys + ) == Counter((key,)) + source.deposit(BlockKey(b"replacement")) + assert kvcr.query((BlockKey(b"replacement"),)) == [ + (QueryStatus.HIT, CacheTier.LOCAL_G2) + ] + + +@pytest.mark.parametrize("incomplete", [False, True]) +def test_move_rejects_busy_source_without_blocking_deposit(pair, incomplete): + source, target = pair.source, pair.target + key = BlockKey(b"busy") + layout = ["", ":tail"] if incomplete else [""] + source.agent.state = "DONE" if incomplete else "PROC" + deposit = source.kvcr.deposit({key: [source.ref()]}, expected_layout=layout) + _wait_until(lambda: bool(source.agent.xfers)) + if incomplete: + assert dict(_poll_until(source.kvcr, bool))[deposit][key].success + moved = target.kvcr._core.fetch( + (key,), + expected_layout=layout, + source_endpoint=source.control.endpoint, + move=True, + ) + pair.until(lambda: moved in target.results) + assert not target.results[moved][key].success + assert not source.kvcr._core._invalidation_handler.has_pending() + source.agent.state = "DONE" + if incomplete: + deposit = source.kvcr.deposit( + {key: [source.ref(1, ":tail")]}, expected_layout=layout + ) + pair.until(lambda: deposit in source.results) + assert source.results[deposit][key].success + assert source.kvcr.query((key,)) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] + + +def test_invalidation_preserves_an_already_admitted_remote_fetch(pair): + source, target = pair.source, pair.target + key = BlockKey(b"remote-reader") + source.deposit(key) + target.kvcr._core._remote_fw_dram._key_adapter = _ConstantHashAdapter() + target.kvcr.submit_hint(_router_hint(source.control.endpoint), "copy") + before = len(source.agent.xfers) + copied = target.kvcr.fetch((key,), "copy") + pair.until(lambda: len(source.agent.xfers) > before) + residency = source.kvcr._core._block_record_map[key].local_dram + assert residency.source_readers == 1 + _invalidate(source.kvcr, (key,)) + assert source.kvcr.query((key,)) == [(QueryStatus.MISS, None)] + assert source.kvcr._core._invalidation_handler.has_pending() + source.agent.state = "DONE" + pair.until(lambda: copied in target.results and pair.idle()) + result = target.results[copied][key] + assert result.success + assert target.local.raw[:16] == b"a" * 16 + assert key not in source.kvcr._core._block_record_map + assert target.kvcr.release((result.release_handle,)) == [ + (result.release_handle, True) + ] From c7f1f0e05a3da9018dc1537f68d122b3523aeeb7 Mon Sep 17 00:00:00 2001 From: Weiyang Wang Date: Thu, 5 Nov 2026 02:33:53 +0200 Subject: [PATCH 8/8] feat: execute independent prefetch and move hints Connect kv.prefetch and kv.move to hash resolution and named-source transfers. Actions outlive their carrier request, preserve request advice, validate local identity/layout compatibility, and aggregate per-hash results across layouts. Consume internal completions and release temporary destination claims. Extend the existing hint suite for carrier independence, copied bytes, cached moves, source disappearance, capacity waits, overlapping actions, layout groups, timeouts, and cleanup. Document the hint flow and source removal ordering, including the planned request-scoped kv.migrate hint for cold starts. Validation: full unit suite, native two-peer NIXL/UCX scenarios, Ruff, and whitespace checks. Assisted-by: OpenAI Codex Assisted-by: Anthropic Claude Signed-off-by: Weiyang Wang --- docs/design_overview.md | 16 +- docs/dev-guide.md | 2 + docs/hash_management.md | 32 ++ docs/quick-start-vllm.md | 2 + src/kvcr/hints/__init__.py | 2 +- src/kvcr/hints/dispatcher.py | 18 +- src/kvcr/hints/move.py | 26 ++ src/kvcr/hints/prefetch.py | 40 ++ src/kvcr/hints/results.py | 18 +- src/kvcr/runtime/core.py | 6 + src/kvcr/runtime/independent_hint_handler.py | 213 ++++++++++ tests/unit/test_hints_dispatcher.py | 5 +- tests/unit/test_hints_fetch.py | 389 ++++++++++++++++++- 13 files changed, 751 insertions(+), 18 deletions(-) create mode 100644 src/kvcr/hints/move.py create mode 100644 src/kvcr/hints/prefetch.py create mode 100644 src/kvcr/runtime/independent_hint_handler.py diff --git a/docs/design_overview.md b/docs/design_overview.md index 27d27076..8691788a 100644 --- a/docs/design_overview.md +++ b/docs/design_overview.md @@ -261,6 +261,20 @@ A `kv.fetch` hint identifies a source control endpoint and advertised block hash Failed or conflicting advice does not prevent local cache hits. `discard_hint` clears request advice without cancelling already-started operations. A `kv.fetch` without a request ID does not execute an independent operation. Source-less hints and `no_retain` are unsupported. +We plan to add a request-scoped `kv.migrate` hint to move KV cache for cold starts. + +### Independent hints + +`kv.prefetch` copies cached KV blocks; `kv.move` also requests removal of the source copy. Each hint names a source and block hashes and runs independently of its carrier request's keys and lifetime. Request advice and independent actions can share an envelope, and independent actions also work without a request ID. + +The destination first resolves the hashes into full keys and layouts, validates the reply, then transfers complete blocks from the source's KVCR-managed DRAM (G2). See [hash-addressed operations](hash_management.md) for detailed contracts. + +### Move and source invalidation + +For moves, the source withdraws eligible blocks from new access before transferring them. After a successful transfer, it removes their G2 copies once existing readers finish. A failed transfer keeps the source copy and reopens access. If the destination already has complete copies, it sends invalidation without copying them again. + +Removal is best effort and honors `no_evict`; framework and G3 copies remain. Destination completion reports data readiness, while source cleanup may finish later. Source removal does not wait for destination acknowledgement. + ### Peer control and liveness The router is never on the data path: KVCR instances execute transfers peer-to-peer via NIXL. Source endpoints belong to request advice or an independent operation rather than global block state, so this coordination adds no KVCR-side peer inventory. @@ -309,7 +323,7 @@ At initialization, an integration selects one built-in policy or supplies an ext Eviction scores are finite floats. LRU favors evicting later sequence keys at comparable recency. Custom policies may use ordering position or ignore it. -Policy decides whether to stage, retain, move, or drop data within KVCR-owned DRAM and downstream storage. Fetch and `no_evict` deposit claims prevent eviction while held. `no_retain` is not currently supported. +Policy decides whether to stage, retain, move, or drop data within KVCR-owned DRAM and downstream storage. Fetch and `no_evict` deposit claims prevent eviction while held. Move uses the source invalidation rules above rather than a policy eviction decision. `no_retain` is not currently supported. The policy methods are invoked as follows: diff --git a/docs/dev-guide.md b/docs/dev-guide.md index 2ca3a5d2..5f390233 100644 --- a/docs/dev-guide.md +++ b/docs/dev-guide.md @@ -160,6 +160,8 @@ For `on_resilience_event` and optional framework-memory quarantine, see the [res See the design doc's [Framework-KVCR API](design_overview.md#framework-kvcr-api) and [operating flow](design_overview.md#operating-flow) for query, transfer, completion, and release semantics, and the [Router-KVCR API](design_overview.md#router-kvcr-api) for request-scoped hints. `abort()` is currently unimplemented as it cannot be used by the frameworks. +For `kv.prefetch` and `kv.move`, configure the key adapter's reverse lookup and keep polling both peers. See [hash-addressed operations](hash_management.md) for the payload and managed-G2 contract. Hash RPC tests live in the [remote-source](../tests/unit/test_kvcr_remote_source.py) and [remote-target](../tests/unit/test_kvcr_remote_target.py) suites; the latter also covers full-key moves. The [hint tests](../tests/unit/test_hints_fetch.py) cover independent execution and cleanup. + ### Development loop Run a focused test while iterating, then the complete standalone validation before finishing. If an adapter or router contract changed, run the optional integration validation separately. diff --git a/docs/hash_management.md b/docs/hash_management.md index 08451e65..75efdc94 100644 --- a/docs/hash_management.md +++ b/docs/hash_management.md @@ -1,7 +1,39 @@ # Hash-addressed cache operations +| Hint | Behavior | +| --- | --- | +| `kv.fetch` | Stores request-scoped source advice; framework `fetch()` / `deliver()` calls initiate ordinary copying. | +| `kv.prefetch` | Resolves advertised hashes and copies their full blocks into destination G2. | +| `kv.move` | Resolves advertised hashes and invokes the internal move mechanism. | + +Independent actions work with or without a carrier request ID, can target unrelated blocks, and survive `discard_hint(carrier_id)`. A `kv.fetch` without a request ID does not become independent. + +```json +{ + "protocol_version": "0.1", + "message_id": "placement-batch", + "actions": [{ + "action_id": "placement-42", + "action_type": "kv.move", + "action_version": "1.0", + "payload": { + "source_control_endpoint": "tcp://source-worker:5600", + "block_hashes": [123, 456] + } + }] +} +``` + +Use `kv.prefetch` to preserve the source copy. One envelope can combine request advice and independent actions. + The source's `KeyAdapter.decode_from_router` maps each uint64 hash to a full `BlockKey`: `KeyError` reports `not_found`, `ValueError` reports `ambiguous_hash`, and any other exception fails only that hash with `resolution_failed`. The integration owns the reverse map and collision policy; KVCR adds no index. Lookups run on progress and must not block. Full keys remain the identities used for storage and transfers. Resolution exchanges keys, ordered piece labels, and sizes in one metadata round trip without pinning memory. Self-addressed resolution fails with `self_source`. +The destination validates identity and layout before transfer. Independent transfers require complete managed G2 at the named source; missing/incomplete copies fail without acquiring or reusing framework pins. Framework-only discovery and remote G3 staging are unsupported. Ordinary request copying retains its framework fallback. + +Keep polling both peers. A destination without managed G2 or a key adapter rejects an independent action at submission; otherwise it reports `ACCEPTED`, then `COMPLETED` or `FAILED`, with per-hash outcomes through the internal/debug result sink. Its completions are filtered from framework results and destination claims are released automatically. Closing abandons pending hint results. There is no public hint-result API. + +Moves reuse existing transfers and source-owned G2 removal, honoring readers and `no_evict`. Completion reports destination readiness, not confirmation of source cleanup; late destination failure can follow successful source removal. See [move semantics](design_overview.md#move-and-source-invalidation). + An adapter may maintain reverse mappings from inventory events. The vLLM adapter requires KV events enabled; colliding full keys reject lookup until unique again, including hashes shared across cache groups. Guard recovery does not replay inventory: reverse mappings must be restored separately. Peers need compatible key conventions and layouts; KVCR does not transform them. diff --git a/docs/quick-start-vllm.md b/docs/quick-start-vllm.md index f4505e0a..b7c940b1 100644 --- a/docs/quick-start-vllm.md +++ b/docs/quick-start-vllm.md @@ -246,4 +246,6 @@ The explicit rank selection is only for this deterministic mechanism test. In a - **No router hint or cache hit:** Use `--router-mode kv`, enable self-describing KV events, and confirm the worker advertises `router_hint` and its control endpoints. Use a fresh prefix longer than one block. - **Peer connection or transfer failure:** Ensure control and event ports are unique, with one control port per DP rank, and `control_advertise_host` is reachable by peers. Check terminal transfer status and destination data; a control acknowledgement alone does not prove delivery. +The pinned router emits request-scoped `kv.fetch` advice, and the vLLM adapter loads through copy-only `deliver`. Independent `kv.prefetch` and `kv.move` require managed destination G2 capacity, which this peer-only quick start does not configure, plus an adapter that implements reverse hash lookup. See [hash-addressed operations](hash_management.md). + For more checks, see the [developer guide](dev-guide.md#troubleshooting). diff --git a/src/kvcr/hints/__init__.py b/src/kvcr/hints/__init__.py index 053c1527..713a6667 100644 --- a/src/kvcr/hints/__init__.py +++ b/src/kvcr/hints/__init__.py @@ -3,7 +3,7 @@ """KV hint handling: envelope parsing, per-type dispatch, and the hint handlers. ``parser`` turns a router envelope into typed hints, ``dispatcher`` routes each -hint to its handler, and one module per hint type (``fetch`` today) +hint to its handler, and one module per hint type implements the parser and the handler and ties them together in a ``HintType``. Core holds one ``HintDispatcher`` and never names a handler; adding a hint type is adding a module here and listing its ``HintType`` in diff --git a/src/kvcr/hints/dispatcher.py b/src/kvcr/hints/dispatcher.py index 5b71e1eb..5825b600 100644 --- a/src/kvcr/hints/dispatcher.py +++ b/src/kvcr/hints/dispatcher.py @@ -16,14 +16,16 @@ import msgspec from .fetch import KV_FETCH +from .move import KV_MOVE from .parser import HintType, _Hint, hint_message_id, parse_hints +from .prefetch import KV_PREFETCH from .results import HintResult, HintStatus if TYPE_CHECKING: from ..runtime.core import _KVCRCore # The supported hint types. Every dispatcher builds one handler per entry. -_HINT_TYPES: tuple[HintType, ...] = (KV_FETCH,) +_HINT_TYPES: tuple[HintType, ...] = (KV_FETCH, KV_PREFETCH, KV_MOVE) logger = logging.getLogger(__name__) # Handlers run on hot paths; warn once per (phase, action type) about errors. @@ -47,8 +49,9 @@ def _warn_handler_error(action_type: str, phase: str) -> None: class _HintHandler(Protocol): """Executes one action type. - ``submit`` runs on the caller's thread without the state lock, once per - hint, and must report exactly one result for it. + ``submit`` is invoked on the caller's thread without the state lock and + reports one admission result. Independent cache operations acquire the + lock as needed and subsequently report asynchronous completion. """ action_type: str @@ -105,9 +108,9 @@ def _submit_to_handler( request_id: str | None, message_id: str | None, ) -> None: - # Contract: a handler reports the hint exactly once, or raises before - # doing so. Inline hints are terminal at submit time, so a raise - # means REJECTED, unless the raise is the fatal progress-thread + # Contract: a handler reports admission exactly once, or raises before + # doing so. ACCEPTED means admission, not asynchronous completion. + # A raise means REJECTED, unless it is the fatal progress-thread # failure, which propagates to the caller as it does without hints. try: handler.submit(hint, request_id, message_id) @@ -133,11 +136,12 @@ def report(self, result: HintResult) -> None: forward KVCR stats and may reject names they do not declare. """ logger.debug( - "KV hint result; action_id=%s action_type=%s status=%s reason=%s", + "KV hint result; action_id=%s action_type=%s status=%s reason=%s op=%s", result.action_id, result.action_type, result.status.value, result.reason, + result.op_handle, ) # Main thread, under the state lock. diff --git a/src/kvcr/hints/move.py b/src/kvcr/hints/move.py new file mode 100644 index 00000000..40f38865 --- /dev/null +++ b/src/kvcr/hints/move.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Move advertised blocks independently of the carrier request.""" + +from collections.abc import Mapping +from typing import ClassVar + +from .fetch import _KVFetchHint, _parse_kv_fetch +from .parser import HintType +from .prefetch import _PrefetchHintHandler + + +class _KVMoveHint(_KVFetchHint): + ACTION_TYPE: ClassVar[str] = "kv.move" + + +def _parse_kv_move(action_id: str | None, payload: Mapping[str, object]) -> _KVMoveHint: + hint = _parse_kv_fetch(action_id, payload) + return _KVMoveHint(action_id, hint.source_endpoint, hint.block_hashes) + + +class _MoveHintHandler(_PrefetchHintHandler): + action_type = _KVMoveHint.ACTION_TYPE + + +KV_MOVE = HintType(_KVMoveHint.ACTION_TYPE, "1.0", _parse_kv_move, _MoveHintHandler) diff --git a/src/kvcr/hints/prefetch.py b/src/kvcr/hints/prefetch.py new file mode 100644 index 00000000..7d5659ea --- /dev/null +++ b/src/kvcr/hints/prefetch.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Copy advertised blocks independently of the carrier request.""" + +from collections.abc import Mapping +from typing import ClassVar + +from .fetch import _KVFetchHint, _parse_kv_fetch +from .parser import HintType, _Hint + + +class _KVPrefetchHint(_KVFetchHint): + ACTION_TYPE: ClassVar[str] = "kv.prefetch" + + +def _parse_kv_prefetch( + action_id: str | None, payload: Mapping[str, object] +) -> _KVPrefetchHint: + hint = _parse_kv_fetch(action_id, payload) + return _KVPrefetchHint(action_id, hint.source_endpoint, hint.block_hashes) + + +class _PrefetchHintHandler: + """Dispatch hash-addressed operations to the independent runtime.""" + + action_type = _KVPrefetchHint.ACTION_TYPE + + def __init__(self, core, dispatcher) -> None: + self._core = core + + def submit(self, hint: _Hint, request_id, message_id) -> None: + self._core._independent_hint_handler.start(hint, request_id, message_id) + + def close(self) -> None: + pass + + +KV_PREFETCH = HintType( + _KVPrefetchHint.ACTION_TYPE, "1.0", _parse_kv_prefetch, _PrefetchHintHandler +) diff --git a/src/kvcr/hints/results.py b/src/kvcr/hints/results.py index 4dcd9276..cc9ae3d8 100644 --- a/src/kvcr/hints/results.py +++ b/src/kvcr/hints/results.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""Per-hint result types, internal to ``kvcr.hints``.""" +"""Admission and asynchronous cache-operation result types.""" from enum import Enum @@ -8,9 +8,19 @@ class HintStatus(Enum): - # Inline hints are terminal at submit time: ACCEPTED or REJECTED only. - ACCEPTED = "ACCEPTED" # inline work done + # Admission is distinct from an asynchronous operation's completion. + ACCEPTED = "ACCEPTED" REJECTED = "REJECTED" # invalid, unsupported, handler_error + COMPLETED = "COMPLETED" + FAILED = "FAILED" + + +class HintEntryResult(msgspec.Struct, frozen=True): + """Hint operation outcome for one hash, including its resolved full key.""" + + success: bool + key: bytes | None = None + reason: str | None = None class HintResult(msgspec.Struct, frozen=True): @@ -20,3 +30,5 @@ class HintResult(msgspec.Struct, frozen=True): reason: str | None = None request_id: str | None = None message_id: str | None = None + op_handle: int | None = None + entries: dict[int, HintEntryResult] | None = None diff --git a/src/kvcr/runtime/core.py b/src/kvcr/runtime/core.py index 329335f1..dc579c4d 100644 --- a/src/kvcr/runtime/core.py +++ b/src/kvcr/runtime/core.py @@ -49,6 +49,7 @@ ReleaseHandle, ReleaseResult, ) +from .independent_hint_handler import _IndependentHintHandler from .invalidation import _InvalidationHandler from .policy_runtime import _PolicyInvoker from .progress import ( @@ -332,6 +333,7 @@ def __init__( else None ) # Built after the backends so hint handlers can bind to them. + self._independent_hint_handler = _IndependentHintHandler(self) self._hints_dispatcher = HintDispatcher(self) kvcr_regions = ( {region.label: region for region in self._local_dram._pools.memory_regions} @@ -826,8 +828,11 @@ def poll_completed(self) -> Iterable[OpResult]: progress_items = self._local_dram.poll_main(progress_items) self._remote_fw_dram.poll_main(progress_items) self._drain_invalidation() + self._independent_hint_handler.poll_resolved() completed = self._completion_queue self._completion_queue = [] + if self._independent_hint_handler._pending_operations: + completed = self._independent_hint_handler.finish(completed) if self._nvtx_operations: for handle, _ in completed: trace = self._nvtx_operations.pop(handle, None) @@ -929,6 +934,7 @@ def close(self) -> None: def _close_hint_handlers(self) -> None: with self._state_lock: self._hints_dispatcher.close() + self._independent_hint_handler.close() def _close_main_resources(self) -> None: self._remote_delivery_keys.clear() diff --git a/src/kvcr/runtime/independent_hint_handler.py b/src/kvcr/runtime/independent_hint_handler.py new file mode 100644 index 00000000..b45d0550 --- /dev/null +++ b/src/kvcr/runtime/independent_hint_handler.py @@ -0,0 +1,213 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Execute independent hints and collect their asynchronous results.""" + +from __future__ import annotations + +from collections import deque +from typing import TYPE_CHECKING + +import msgspec + +from ..hints.dispatcher import _warn_handler_error +from ..hints.parser import reject_hint +from ..hints.results import HintEntryResult, HintResult, HintStatus + +if TYPE_CHECKING: + from ..backends.remote_ops import _HashResolutionOp + from ..hints.move import _KVMoveHint + from ..hints.prefetch import _KVPrefetchHint + from ..types import BlockKey, OpEntryResult, OpHandle + from .core import _KVCRCore + + +class _IndependentHintHandler: + """Owner-thread lifecycle for operations started by independent hints.""" + + def __init__(self, core: _KVCRCore) -> None: + self._core = core + self._resolved: deque[_HashResolutionOp] = deque() + self._resolving: dict[ + OpHandle, tuple[HintResult, _KVPrefetchHint | _KVMoveHint] + ] = {} + # Each child stores (receipt, hash count, key-to-hash map). + # Receipts share results; the last child reports the complete outcome. + self._pending_operations: dict[ + OpHandle, tuple[HintResult, int, dict[BlockKey, int]] + ] = {} + + def start( + self, + hint: _KVPrefetchHint | _KVMoveHint, + request_id: str | None, + message_id: str | None, + ) -> None: + """Admit the hint and queue resolution; resolved() starts the transfers.""" + from ..backends.remote_ops import _HashResolutionOp + + core = self._core + if core._local_dram is None or core._remote_fw_dram._key_adapter is None: + reason = ( + "destination_has_no_managed_dram" + if core._local_dram is None + else "key_adapter_unavailable" + ) + core._hints_dispatcher.report( + reject_hint( + hint.action_id, hint.ACTION_TYPE, reason, message_id, request_id + ) + ) + return + handle = core._next_op_handle + core._next_op_handle += 1 + # Queue before reporting: a failed submission must raise before any + # report, and completion only arrives through owner polling afterwards. + core._progress.submit( + _HashResolutionOp( + op_id=("resolve_hashes", handle), + keys=frozenset(), + endpoint=hint.source_endpoint, + hashes=tuple(hint.block_hashes), + _backend=core._remote_fw_dram, + on_complete=self._resolved.append, + ) + ) + receipt = HintResult( + action_id=hint.action_id, + action_type=hint.ACTION_TYPE, + status=HintStatus.ACCEPTED, + request_id=request_id, + message_id=message_id, + op_handle=handle, + ) + self._resolving[handle] = receipt, hint + core._hints_dispatcher.report(receipt) + + def poll_resolved(self) -> None: + """Start resolved transfers outside the core state lock.""" + while self._resolved: + self.resolved(self._resolved.popleft()) + + def resolved(self, resolution: _HashResolutionOp) -> None: + core = self._core + receipt, hint = self._resolving.pop(resolution.op_id[1]) + entries = {} + receipt = msgspec.structs.replace(receipt, entries=entries) + hash_count = len(resolution.hashes) + + if resolution.failure is not None: + entries.update( + (h, HintEntryResult(False, reason=resolution.failure)) + for h in resolution.hashes + ) + self._complete(receipt, hash_count, {}, {}) + return + groups = {} + # Rows share few distinct layouts; validate each one once per reply. + checked: dict[tuple[str, ...], tuple[int, ...] | None] = {} + adapter = core._remote_fw_dram._key_adapter + for row in resolution.response: + reason = row.reason + if reason is None: + if row.layout not in checked: + checked[row.layout] = self._layout_sizes(row.layout) + sizes = checked[row.layout] + try: + if sizes is None: + reason = "incompatible_key_or_layout" + elif sizes != row.sizes: + reason = "incompatible_layout" + elif adapter.encode_to_router(row.key) != row.block_hash: + reason = "hash_key_mismatch" + except Exception: + # Adapters can raise integration-specific errors for remote + # keys. Fail this hash without dropping the command's result. + reason = "incompatible_key_or_layout" + if reason is not None: + entries[row.block_hash] = HintEntryResult(False, reason=reason) + else: + groups.setdefault(row.layout, {})[row.key] = row.block_hash + if not groups: + self._complete(receipt, hash_count, {}, {}) + return + + move = hint.ACTION_TYPE == "kv.move" + for layout, hashes_by_key in groups.items(): + operation = receipt, hash_count, hashes_by_key + try: + handle = core.fetch( + tuple(hashes_by_key), + expected_layout=list(layout), + source_endpoint=resolution.endpoint, + move=move, + ) + except Exception as error: + # Fatal progress failures retain their ordinary API behavior. + core._progress.raise_if_failed() + # Fail this group without escaping the owner's poll. Layout + # conflicts are expected; anything else warns once. + if not isinstance(error, ValueError): + _warn_handler_error(hint.ACTION_TYPE, "admission") + self._complete(*operation, {}) + else: + # Synchronous hits are queued until polling, so registration + # after fetch returns is safe. + self._pending_operations[handle] = operation + + def _layout_sizes(self, layout: tuple[str, ...]) -> tuple[int, ...] | None: + """Return piece sizes for a layout this destination can store.""" + core = self._core + try: + core._validate_expected_layout(list(layout)) + except ValueError: + return None + return tuple(core._block_sizes[label.partition(":")[0]] for label in layout) + + def finish(self, completed): + """Consume independent completions and return framework operations.""" + framework = [] + for handle, entries in completed: + operation = self._pending_operations.pop(handle, None) + if operation is None: + framework.append((handle, entries)) + continue + self._complete(*operation, entries) + return framework + + def _complete( + self, + receipt: HintResult, + hash_count: int, + selected: dict[BlockKey, int], + results: dict[BlockKey, OpEntryResult], + ) -> None: + """Release private fetch claims and collect exact-key operation results.""" + handles = tuple( + result.release_handle + for result in results.values() + if result.release_handle is not None + ) + if handles: + self._core.release(handles) + entries = receipt.entries + assert entries is not None + for key, block_hash in selected.items(): + result = results.get(key) + success = result is not None and result.success + entries[block_hash] = HintEntryResult( + success, key, None if success else "operation_failed" + ) + if len(entries) == hash_count: + status = ( + HintStatus.COMPLETED + if all(result.success for result in entries.values()) + else HintStatus.FAILED + ) + self._core._hints_dispatcher.report( + msgspec.structs.replace(receipt, status=status) + ) + + def close(self) -> None: + self._resolved.clear() + self._resolving.clear() + self._pending_operations.clear() diff --git a/tests/unit/test_hints_dispatcher.py b/tests/unit/test_hints_dispatcher.py index 71e4b1d8..2bd824f1 100644 --- a/tests/unit/test_hints_dispatcher.py +++ b/tests/unit/test_hints_dispatcher.py @@ -109,7 +109,8 @@ def test_dispatcher_owns_the_hint_type_table(monkeypatch): # table, and a type added to the table appears in a fresh dispatcher # without any core change. kvcr = _new_kvcr(FakeNixlAgent(), FakePrimaryPinning(), FakeBytesControl()) - assert set(kvcr._core._hints_dispatcher._handlers) == {"kv.fetch"} + registered = {hint.action_type for hint in _HINT_TYPES} + assert set(kvcr._core._hints_dispatcher._handlers) == registered built_for: list[tuple[object, object]] = [] @@ -122,7 +123,7 @@ def factory(core, dispatcher): (*_HINT_TYPES, HintType("kv.test", "1.0", lambda i, p: _StubHint(i), factory)), ) dispatcher = HintDispatcher(kvcr._core) - assert set(dispatcher._handlers) == {"kv.fetch", "kv.test"} + assert set(dispatcher._handlers) == registered | {"kv.test"} assert built_for == [(kvcr._core, dispatcher)] diff --git a/tests/unit/test_hints_fetch.py b/tests/unit/test_hints_fetch.py index bfedbbf4..75a707d5 100644 --- a/tests/unit/test_hints_fetch.py +++ b/tests/unit/test_hints_fetch.py @@ -1,19 +1,27 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""kv.fetch hint handler tests: the engine-destination inline path.""" +"""Request advice and independent prefetch/move hint execution.""" + +from unittest.mock import Mock import pytest from _kvcr_test_utils import ( FakeBytesControl, FakeNixlAgent, FakePrimaryPinning, + _decode_control_message, _HintAction, _HintEnvelope, + _invalidate, _new_kvcr, + _poll_until, _record_hint_results, _router_hint, + _wait_until, ) +from test_kvcr_remote_target import _Pair +from kvcr.backends.remote_ops import _ResolvedHash, _TargetPullOp from kvcr.config import KVCRConfig, RemoteFWDramOptions from kvcr.hints.dispatcher import _LOGGED_HANDLER_ERRORS from kvcr.hints.parser import _LOGGED_HINT_ISSUES @@ -59,9 +67,9 @@ def _fetch_kvcr(monkeypatch): return kvcr, _record_hint_results(monkeypatch, kvcr) -def _fetch(action_id: str, source: str, hashes) -> _HintAction: +def _fetch(action_id: str, source: str, hashes, action_type="kv.fetch") -> _HintAction: return _HintAction( - action_id, "kv.fetch", "1.0", + action_id, action_type, "1.0", {"source_control_endpoint": source, "block_hashes": list(hashes)}, ) # fmt: skip @@ -151,6 +159,379 @@ def test_all_invalid_envelope_does_not_raise_from_submit_hint(monkeypatch): def test_fetch_hint_without_request_id_is_accepted(monkeypatch): kvcr, results = _fetch_kvcr(monkeypatch) - kvcr.submit_hint(_router_hint("tcp://a:1", (1,))) + envelope = _router_hint("tcp://a:1", (1,)) + envelope.actions[0].payload.update(scope="independent", mode="move") + kvcr.submit_hint(envelope) assert [(r.status, r.request_id) for r in results] == [(HintStatus.ACCEPTED, None)] assert not kvcr._core._remote_fw_dram._request_hints + assert not kvcr._core._independent_hint_handler._resolving + + +@pytest.mark.parametrize("action_type", ["kv.prefetch", "kv.move"]) +@pytest.mark.parametrize("request_id", [None, "req"]) +def test_independent_hints_require_managed_destination( + monkeypatch, action_type, request_id +): + kvcr, results = _fetch_kvcr(monkeypatch) + action = _fetch("a1", "tcp://a:1", [1], action_type) + kvcr.submit_hint(_fetch_envelope(action), request_id) + assert [(r.action_type, r.status, r.reason) for r in results] == [ + (action_type, HintStatus.REJECTED, "destination_has_no_managed_dram") + ] + assert results[0].request_id == request_id and results[0].message_id == "m1" + assert not kvcr._core._remote_fw_dram._request_hints + assert kvcr._core._outstanding_operations == 0 and not kvcr.poll_completed() + + +_KEY = BlockKey(b"full-key-with-group") +_TAIL_KEY = BlockKey(b"tail-piece-key") + + +class _Adapter: + def decode_from_router(self, block_hash): + return {123: _KEY, 124: _TAIL_KEY}[block_hash] + + def encode_to_router(self, key): + return {_KEY: 123, _TAIL_KEY: 124}[key] + + +def _pair(): + pair = _Pair() + for peer in (pair.source, pair.target): + peer.kvcr._core._remote_fw_dram._key_adapter = _Adapter() + return pair + + +def _submit(pair, action, carrier=None, hashes=(123,)): + hint = _fetch("action", pair.source.control.endpoint, hashes, action) + pair.target.kvcr.submit_hint(_fetch_envelope(hint), carrier) + + +def _idle(pair): + pair.until(pair.idle) + handler = pair.target.kvcr._core._independent_hint_handler + assert not handler._resolving and not handler._resolved + assert not handler._pending_operations and not pair.target.results + source = pair.source.kvcr._core + assert not source._framework_pin_keys and not source._local_dram_sources_by_op + for record in source._block_record_map.values(): + if record.local_dram is not None: + assert ( + record.local_dram.claim_count == record.local_dram.source_readers == 0 + ) + + +@pytest.mark.parametrize("action", ["kv.prefetch", "kv.move"]) +@pytest.mark.parametrize("carrier", [None, "unrelated-request"]) +@pytest.mark.parametrize("cached", [False, True]) +@pytest.mark.parametrize("missing", [False, True]) +def test_independent_transfer_preserves_carrier_and_releases_claims( + action, carrier, cached, missing, monkeypatch +): + """Actual bytes, cached moves, and partial failures keep their action identity.""" + pair = _pair() + pair.source.deposit(_KEY) + if cached: + pair.target.deposit(_KEY) + core = pair.target.kvcr._core + remote = core._remote_fw_dram + if carrier: + remote.submit_hint("tcp://unrelated:9", frozenset({456}), carrier) + advice = dict(remote._request_hints) + reports = _record_hint_results(monkeypatch, pair.target.kvcr) + before = len(pair.source.agent.xfers) + _submit(pair, action, carrier, (123, 999) if missing else (123,)) + assert remote._request_hints == advice + if carrier: + pair.target.kvcr.discard_hint(carrier) + advice = {} + pair.source.agent.state = "DONE" + pair.until(lambda: len(reports) == 2) + accepted, result = reports + assert accepted.status is HintStatus.ACCEPTED + assert result.status is (HintStatus.FAILED if missing else HintStatus.COMPLETED) + assert result.op_handle == accepted.op_handle and result.op_handle is not None + assert result.action_type == action and result.request_id == carrier + assert result.message_id == "m1" and result.entries[123].key == _KEY + assert result.entries[123].success + if missing: + assert result.entries[999].reason == "not_found" + assert remote._request_hints == advice + assert core._block_record_map[_KEY].local_dram.claim_count == 0 + assert pair.target.local.raw[:16] == b"a" * 16 + assert len(pair.source.agent.xfers) - before == (0 if cached else 1) + expected = ( + (QueryStatus.MISS, None) + if action == "kv.move" + else (QueryStatus.HIT, CacheTier.LOCAL_G2) + ) + pair.until(lambda: pair.source.kvcr.query((_KEY,)) == [expected]) + _idle(pair) + assert len(reports) == 2 + + +@pytest.mark.parametrize( + "action,queued", [("kv.prefetch", False), ("kv.prefetch", True), ("kv.move", False)] +) +@pytest.mark.parametrize("incomplete", [False, True]) +def test_disappearing_resolved_source_never_requests_framework_memory( + action, queued, incomplete, monkeypatch +): + """Missing or incomplete G2 fails on both source paths without pinning/waiting.""" + pair = _pair() + pair.source.deposit(_KEY) + source = pair.source.kvcr._core + pins = Mock(wraps=source._request_pin_callback) + monkeypatch.setattr(source, "_request_pin_callback", pins) + if queued: + monkeypatch.setattr( + source._remote_fw_dram, "_try_local_source_write", lambda *_: False + ) + handler = pair.target.kvcr._core._independent_hint_handler + resolved = handler.resolved + + def remove_before_fetch(operation): + _invalidate(pair.source.kvcr, (_KEY,)) + if incomplete: + pair.source.agent.state = "DONE" + handle = pair.source.kvcr.deposit( + {_KEY: [pair.source.ref()]}, expected_layout=["", ":tail"] + ) + assert dict(_poll_until(pair.source.kvcr, bool))[handle][_KEY].success + assert source._block_record_map[_KEY].local_dram.pending_labels == {":tail"} + resolved(operation) + + monkeypatch.setattr(handler, "resolved", remove_before_fetch) + reports = _record_hint_results(monkeypatch, pair.target.kvcr) + _submit(pair, action) + pair.until(lambda: len(reports) == 2) + assert reports[-1].status is HintStatus.FAILED + assert reports[-1].entries[123].reason == "operation_failed" + assert reports[-1].entries[123].key == _KEY + pins.assert_not_called() + assert pair.target.kvcr.query((_KEY,)) == [(QueryStatus.MISS, None)] + _idle(pair) + + +@pytest.mark.parametrize( + "row,reason", + [ + (_ResolvedHash(123, _KEY, ("",), (32,)), "incompatible_layout"), + (_ResolvedHash(123, _KEY, (":*",), (16,)), "incompatible_key_or_layout"), + (_ResolvedHash(123, _KEY, ("",), (16,)), "hash_key_mismatch"), + ], +) +def test_resolved_metadata_is_validated_before_fetch(row, reason, monkeypatch): + pair = _pair() + source, target = pair.source.kvcr._core, pair.target.kvcr._core + monkeypatch.setattr(source._remote_fw_dram, "_resolve_hashes", lambda _: (row,)) + if reason == "hash_key_mismatch": + monkeypatch.setattr( + target._remote_fw_dram._key_adapter, "encode_to_router", lambda _: 456 + ) + reports = _record_hint_results(monkeypatch, pair.target.kvcr) + _submit(pair, "kv.prefetch") + pair.until(lambda: len(reports) == 2) + assert reports[-1].entries[123].reason == reason + assert not pair.source.agent.xfers and not pair.target.agent.xfers + assert not target._block_record_map + _idle(pair) + + +@pytest.mark.parametrize("at_submit", [False, True]) +def test_nonfatal_error_reports_once_and_allows_later_work(at_submit, monkeypatch): + pair = _pair() + pair.source.deposit(_KEY) + core = pair.target.kvcr._core + owner, name = (core._progress, "submit") if at_submit else (core, "fetch") + original = getattr(owner, name) + monkeypatch.setattr(owner, name, Mock(side_effect=RuntimeError("nonfatal"))) + reports = _record_hint_results(monkeypatch, pair.target.kvcr) + _submit(pair, "kv.prefetch") + pair.until(lambda: len(reports) == (1 if at_submit else 2)) + if at_submit: + # The queue failure raised before ACCEPTED: admission reported once. + assert _take(reports) == [("action", HintStatus.REJECTED, "handler_error")] + else: + assert reports[-1].entries[123].reason == "operation_failed" + reports.clear() + _idle(pair) + monkeypatch.setattr(owner, name, original) + _submit(pair, "kv.prefetch") + pair.source.agent.state = "DONE" + pair.until(lambda: len(reports) == 2) + assert reports[-1].status is HintStatus.COMPLETED + _idle(pair) + + +def test_fatal_fetch_submit_error_escapes_the_same_poll(monkeypatch): + """Fail after real destination reservation, after poll's initial fatal check.""" + pair = _pair() + pair.source.deposit(_KEY) + core = pair.target.kvcr._core + handler, progress = core._independent_hint_handler, core._progress + resolved, submit = handler.resolved, progress.submit + manifests = [] + monkeypatch.setattr(handler, "resolved", manifests.append) + reports = _record_hint_results(monkeypatch, pair.target.kvcr) + _submit(pair, "kv.prefetch") + pair.until(lambda: bool(manifests)) + monkeypatch.setattr(handler, "resolved", resolved) + handler._resolved.extend(manifests) + failure = RuntimeError("fatal fetch submit") + + def fail_reserved(item): + if isinstance(item, _TargetPullOp): + assert core._block_record_map[_KEY].local_dram.remote_indices is not None + progress._failure = failure + return submit(item) + + monkeypatch.setattr(progress, "submit", fail_reserved) + with pytest.raises(RuntimeError, match="fatal fetch submit") as error: + pair.target.kvcr.poll_completed() + assert error.value is failure and len(reports) == 1 + with pytest.raises(RuntimeError, match="fatal fetch submit"): + pair.target.kvcr.close() + assert core.is_quiescent() + + +def test_resolution_timeout_does_not_revive_the_hint(monkeypatch): + """A late metadata reply cannot restart an action or report completion twice.""" + pair = _pair() + core = pair.target.kvcr._core + reports = _record_hint_results(monkeypatch, pair.target.kvcr) + _submit(pair, "kv.prefetch") + _wait_until(lambda: bool(pair.target.control.sent)) + core._clock = lambda: core.config.operation_timeout_ms / 1000 + 1 + pair.until(lambda: len(reports) == 2) + assert reports[-1].status is HintStatus.FAILED + assert reports[-1].entries[123].reason == "timeout" + pair.until(lambda: bool(pair.source.control.sent)) + pair.pump() # Forward the late reply even if timeout already made us idle. + _wait_until(lambda: not pair.target.control.incoming) + pair.target.kvcr.poll_completed() + _idle(pair) + assert len(reports) == 2 + assert not pair.source.agent.xfers and not pair.target.agent.xfers + assert not core._block_record_map + + +@pytest.mark.parametrize("vanished", [False, True]) +def test_prefetch_keeps_named_g2_source_after_capacity_wait(vanished, monkeypatch): + """A pending eviction delays admission; resumption preserves source and policy.""" + pair = _pair() + pair.source.deposit(_KEY) + occupied = (b"occupied-first", b"occupied-second") + for index, key in enumerate(occupied): + pair.target.deposit(key, index) + core = pair.target.kvcr._core + pools = core._local_dram._pools + pools.acquire_sources(occupied[:1]) + pools._capacity_eviction_key = occupied[0] + pools.retire_sources(occupied[:1]) + pins = Mock(wraps=pair.source.kvcr._core._request_pin_callback) + monkeypatch.setattr(pair.source.kvcr._core, "_request_pin_callback", pins) + reports = _record_hint_results(monkeypatch, pair.target.kvcr) + _submit(pair, "kv.prefetch") + pair.until(lambda: bool(core._local_dram._capacity_waiters)) + assert len(reports) == 1 + assert all( + _decode_control_message(raw)["type"] != "start_write" + for _, raw in pair.target.control.sent + ) + if vanished: + _invalidate(pair.source.kvcr, (_KEY,)) + pools.release_sources(occupied[:1]) + pair.source.agent.state = "DONE" + pair.until(lambda: len(reports) == 2) + assert reports[-1].entries[123].success is not vanished + requests = [_decode_control_message(raw) for _, raw in pair.target.control.sent] + writes = [request for request in requests if request["type"] == "start_write"] + assert len(writes) == 1 and writes[0]["g2_only"] is True + assert writes[0]["source_control_endpoint"] == pair.source.control.endpoint + pins.assert_not_called() + if vanished: + assert pair.target.kvcr.query((_KEY,)) == [(QueryStatus.MISS, None)] + else: + assert pair.target.local.raw[:16] == b"a" * 16 + assert core._block_record_map[_KEY].local_dram.claim_count == 0 + _idle(pair) + + +@pytest.mark.parametrize( + "actions", + [ + ("kv.prefetch", "kv.move"), + ("kv.move", "kv.prefetch"), + ("kv.prefetch", "kv.prefetch"), + ], +) +def test_overlapping_actions_share_one_transfer(actions, monkeypatch): + """The second action joins the first fill; a move still removes the source copy.""" + pair = _pair() + pair.source.deposit(_KEY) + core = pair.target.kvcr._core + handler = core._independent_hint_handler + reports = _record_hint_results(monkeypatch, pair.target.kvcr) + before = len(pair.source.agent.xfers) + for action in actions: + _submit(pair, action) + pair.until(lambda: len(handler._pending_operations) == 2) + pair.source.agent.state = "DONE" + pair.until(lambda: len(reports) == 4) + assert [report.status for report in reports[2:]] == [HintStatus.COMPLETED] * 2 + assert sorted(report.action_type for report in reports[2:]) == sorted(actions) + requests = [_decode_control_message(raw) for _, raw in pair.target.control.sent] + assert sum(request["type"] == "start_write" for request in requests) == 1 + assert len(pair.source.agent.xfers) - before == 1 + assert pair.target.local.raw[:16] == b"a" * 16 + expected = ( + (QueryStatus.MISS, None) + if "kv.move" in actions + else (QueryStatus.HIT, CacheTier.LOCAL_G2) + ) + pair.until(lambda: pair.source.kvcr.query((_KEY,)) == [expected]) + _idle(pair) + assert core._block_record_map[_KEY].local_dram.claim_count == 0 + + +@pytest.mark.parametrize("vanished", [False, True]) +def test_layout_groups_aggregate_into_one_result(vanished, monkeypatch): + """Each distinct source layout fetches separately; the action reports once.""" + pair = _pair() + pair.source.deposit(_KEY) + pair.source.agent.state = "DONE" + handle = pair.source.kvcr.deposit( + {_TAIL_KEY: [pair.source.ref(1, ":tail")]}, expected_layout=[":tail"] + ) + assert dict(_poll_until(pair.source.kvcr, bool))[handle][_TAIL_KEY].success + core = pair.target.kvcr._core + handler = core._independent_hint_handler + if vanished: + resolved = handler.resolved + + def remove_tail_then_resolve(operation): + _invalidate(pair.source.kvcr, (_TAIL_KEY,)) + resolved(operation) + + monkeypatch.setattr(handler, "resolved", remove_tail_then_resolve) + reports = _record_hint_results(monkeypatch, pair.target.kvcr) + _submit(pair, "kv.prefetch", hashes=(123, 124)) + pair.until(lambda: len(reports) == 2) + result = reports[-1] + assert result.status is (HintStatus.FAILED if vanished else HintStatus.COMPLETED) + assert result.entries[123].success and result.entries[123].key == _KEY + assert result.entries[124].key == _TAIL_KEY + assert result.entries[124].reason == ("operation_failed" if vanished else None) + requests = [_decode_control_message(raw) for _, raw in pair.target.control.sent] + writes = [request for request in requests if request["type"] == "start_write"] + assert sorted(tuple(write["key_layout"]) for write in writes) == [("",), (":tail",)] + hit = (QueryStatus.HIT, CacheTier.LOCAL_G2) + assert pair.target.kvcr.query((_KEY, _TAIL_KEY)) == [ + hit, + (QueryStatus.MISS, None) if vanished else hit, + ] + raw = pair.target.local.raw + assert b"a" * 16 in (raw[:16], raw[16:]) + assert (b"b" * 16 in (raw[:16], raw[16:])) is not vanished + _idle(pair)