@tduncan / Nvidia-Optimization-Stack-VLLM / commits / bec232a914

Secondary tier implementation for PD disaggregation (#42285)

Signed-off-by: Liran Schour <lirans@il.ibm.com> Signed-off-by: liranschour <liranschour@users.noreply.github.com> Co-authored-by: Or Ozeri <or@ozery.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>

liranschour committed Jun 30, 2026 at 07:51 UTC bec232a9146bfd61e2be83879b0b74b952818758
22 files changed +7471
docs/features/kv_offloading_usage.md
+14
@@ -120,6 +120,20 @@ To enable KV cache sharing between multiple vLLM instances using the same `root_
120 PYTHONHASHSEED=0 vllm serve ...
121 ```
122
123 +### P2P (Including P/D)
124 +
125 +The P2P tier (`type: "p2p"`) shares completed KV blocks between vLLM instances over RDMA via NIXL. Each instance binds a control socket on `host:port` and exchanges blocks directly with peers — no shared filesystem required.
126 +
127 +| Key | Required | Default | Notes |
128 +| --- | --- | --- | --- |
129 +| `type` | yes | — | Must be `p2p`. |
130 +| `host` | no | `0.0.0.0` | Address the control socket binds to. |
131 +| `port` | no | `7777` | Port for the control socket. Must be reachable from peers. |
132 +| `backends` | no | `["UCX"]` | NIXL transport backends. See [NixlConnector Usage Guide](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin) for available backends and selection guidance. |
133 +| `num_threads` | no | `4` | NIXL agent worker threads. Only used when `backends` is UCX-only; ignored when any non-UCX backend is requested. |
134 +
135 +The `backends` and `num_threads` options mirror the conditional logic used by [`NixlConnector`](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin): when any non-UCX backend is configured, NIXL is initialised with `backends=...`; otherwise it falls back to a UCX-only agent with the configured `num_threads`. This lets the P2P tier use a different transport (e.g. `MOONCAKE`, `GDS_MT`, `LIBFABRIC`) than the main `NixlConnector` running in the same process.
136 +
137 ## Tuning Tips
138
139 - `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload.
tests/v1/kv_offload/tiering/p2p/__init__.py
tests/v1/kv_offload/tiering/p2p/p2p_connector_proxy.py new
+331
@@ -0,0 +1,331 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +P2PConnector proxy server for OffloadingConnector + TieringOffloadingSpec.
5 +
6 +Unlike NixlConnector (which returns remote_host/remote_port in the prefill
7 +response), OffloadingConnector does not embed connector coordinates in its
8 +response. This proxy injects the prefiller's P2PConnector address into
9 +kv_transfer_params before forwarding the decode request so the decoder knows
10 +where to pull KV blocks from.
11 +
12 +Usage:
13 + .venv/bin/python p2p_connector_proxy.py \
14 + --port 8192 \
15 + --prefiller-host 127.0.0.1 --prefiller-port 8100 \
16 + --decoder-host 127.0.0.1 --decoder-port 8200 \
17 + --p2p-connector-host 127.0.0.1 --p2p-connector-port 7777
18 +"""
19 +
20 +import argparse
21 +import asyncio
22 +import itertools
23 +import logging
24 +import os
25 +import uuid
26 +from contextlib import asynccontextmanager
27 +
28 +import httpx
29 +from fastapi import FastAPI, Request
30 +from fastapi.responses import StreamingResponse
31 +
32 +logger = logging.getLogger(__name__)
33 +logger.setLevel(logging.DEBUG)
34 +
35 +
36 +@asynccontextmanager
37 +async def lifespan(app: FastAPI):
38 + app.state.prefill_clients = []
39 + app.state.decode_clients = []
40 +
41 + for i, (host, port) in enumerate(global_args.prefiller_instances):
42 + app.state.prefill_clients.append(
43 + {
44 + "client": httpx.AsyncClient(
45 + timeout=None,
46 + base_url=f"http://{host}:{port}/v1",
47 + limits=httpx.Limits(
48 + max_connections=None, max_keepalive_connections=None
49 + ),
50 + ),
51 + "host": host,
52 + "port": port,
53 + "id": i,
54 + }
55 + )
56 +
57 + for i, (host, port) in enumerate(global_args.decoder_instances):
58 + app.state.decode_clients.append(
59 + {
60 + "client": httpx.AsyncClient(
61 + timeout=None,
62 + base_url=f"http://{host}:{port}/v1",
63 + limits=httpx.Limits(
64 + max_connections=None, max_keepalive_connections=None
65 + ),
66 + ),
67 + "host": host,
68 + "port": port,
69 + "id": i,
70 + }
71 + )
72 +
73 + app.state.prefill_iterator = itertools.cycle(range(len(app.state.prefill_clients)))
74 + app.state.decode_iterator = itertools.cycle(range(len(app.state.decode_clients)))
75 +
76 + mode = "decoder-first" if global_args.decoder_first else "prefiller-first"
77 + pd_host = global_args.p2p_connector_host
78 + pd_port = global_args.p2p_connector_port
79 + print(
80 + f"Proxy ready [{mode}]: "
81 + f"{len(app.state.prefill_clients)} prefiller(s), "
82 + f"{len(app.state.decode_clients)} decoder(s). "
83 + f"P2PConnector at {pd_host}:{pd_port}"
84 + )
85 + yield
86 +
87 + for ci in app.state.prefill_clients:
88 + await ci["client"].aclose()
89 + for ci in app.state.decode_clients:
90 + await ci["client"].aclose()
91 +
92 +
93 +app = FastAPI(lifespan=lifespan)
94 +
95 +
96 +def parse_args():
97 + p = argparse.ArgumentParser()
98 + p.add_argument("--port", type=int, default=8192)
99 + p.add_argument("--host", type=str, default="127.0.0.1")
100 + p.add_argument("--prefiller-hosts", type=str, nargs="+", default=["127.0.0.1"])
101 + p.add_argument("--prefiller-ports", type=int, nargs="+", default=[8100])
102 + p.add_argument("--decoder-hosts", type=str, nargs="+", default=["127.0.0.1"])
103 + p.add_argument("--decoder-ports", type=int, nargs="+", default=[8200])
104 + # P2PConnector coordinates of the prefiller — injected into decode requests.
105 + p.add_argument(
106 + "--p2p-connector-host",
107 + type=str,
108 + default="127.0.0.1",
109 + help="Host of the prefiller's P2PConnector ZMQ socket",
110 + )
111 + p.add_argument(
112 + "--p2p-connector-port",
113 + type=int,
114 + default=7777,
115 + help="Port of the prefiller's P2PConnector ZMQ socket",
116 + )
117 + # P2PConnector coordinates of the decoder — injected into prefill requests
118 + # so the prefiller's submit_store can resolve the peer to push KV to.
119 + p.add_argument(
120 + "--decoder-p2p-connector-host",
121 + type=str,
122 + default="127.0.0.1",
123 + help="Host of the decoder's P2PConnector ZMQ socket",
124 + )
125 + p.add_argument(
126 + "--decoder-p2p-connector-port",
127 + type=int,
128 + default=7778,
129 + help="Port of the decoder's P2PConnector ZMQ socket",
130 + )
131 + p.add_argument(
132 + "--decoder-first",
133 + action="store_true",
134 + help="Send decode request before prefill so decoder is already "
135 + "waiting when KV blocks arrive (decoder-first mode)",
136 + )
137 + args = p.parse_args()
138 + if len(args.prefiller_hosts) != len(args.prefiller_ports):
139 + raise ValueError("Prefiller host/port count mismatch")
140 + if len(args.decoder_hosts) != len(args.decoder_ports):
141 + raise ValueError("Decoder host/port count mismatch")
142 + args.prefiller_instances = list(zip(args.prefiller_hosts, args.prefiller_ports))
143 + args.decoder_instances = list(zip(args.decoder_hosts, args.decoder_ports))
144 + return args
145 +
146 +
147 +def _get_next(app, service: str):
148 + if service == "prefill":
149 + return app.state.prefill_clients[next(app.state.prefill_iterator)]
150 + return app.state.decode_clients[next(app.state.decode_iterator)]
151 +
152 +
153 +def _auth_headers(request_id: str) -> dict:
154 + headers: dict = {"X-Request-Id": request_id}
155 + api_key = os.environ.get("OPENAI_API_KEY", "")
156 + if api_key:
157 + headers["Authorization"] = f"Bearer {api_key}"
158 + return headers
159 +
160 +
161 +async def _prefill(client_info, endpoint, req_data, request_id):
162 + """Send a prefill-only request (max_tokens=1) to the prefiller."""
163 + data = req_data.copy()
164 + data["kv_transfer_params"] = {
165 + "decode": {
166 + "kv_request_id": request_id,
167 + },
168 + }
169 + data["stream"] = False
170 + data["max_tokens"] = 1
171 + data.pop("max_completion_tokens", None)
172 + data.pop("stream_options", None)
173 + data.pop("min_tokens", None)
174 + data.pop("min_completion_tokens", None)
175 +
176 + headers = _auth_headers(request_id)
177 + resp = await client_info["client"].post(endpoint, json=data, headers=headers)
178 + resp.raise_for_status()
179 + await resp.aread()
180 + return resp
181 +
182 +
183 +async def _stream_decode(client_info, endpoint, req_data, request_id):
184 + headers = _auth_headers(request_id)
185 + async with client_info["client"].stream(
186 + "POST", endpoint, json=req_data, headers=headers
187 + ) as resp:
188 + resp.raise_for_status()
189 + async for chunk in resp.aiter_bytes():
190 + yield chunk
191 +
192 +
193 +async def _handle_completions(api: str, request: Request):
194 + try:
195 + req_data = await request.json()
196 + request_id = str(uuid.uuid4())
197 +
198 + prefill_client = _get_next(request.app, "prefill")
199 + await _prefill(prefill_client, api, req_data, request_id)
200 +
201 + # Inject the prefiller's P2PConnector address so the decoder can pull
202 + # KV blocks from it via the P2PConnector transport.
203 + req_data["kv_transfer_params"] = {
204 + "prefill": {
205 + "kv_request_id": request_id,
206 + "remote_host": global_args.p2p_connector_host,
207 + "remote_port": global_args.p2p_connector_port,
208 + },
209 + }
210 +
211 + decode_client = _get_next(request.app, "decode")
212 + logger.debug("prefill=%s decode=%s", prefill_client, decode_client)
213 +
214 + async def generate():
215 + async for chunk in _stream_decode(decode_client, api, req_data, request_id):
216 + yield chunk
217 +
218 + return StreamingResponse(generate(), media_type="application/json")
219 +
220 + except Exception as e:
221 + import sys
222 + import traceback
223 +
224 + print(f"Proxy error on {api}: {e}")
225 + print("".join(traceback.format_exception(*sys.exc_info())))
226 + raise
227 +
228 +
229 +async def _handle_completions_decoder_first(api: str, request: Request):
230 + """Decoder-first mode: send decode request before prefill.
231 +
232 + The decoder establishes its request and starts polling for KV blocks
233 + immediately. The prefill is then sent so the prefiller computes and
234 + pushes blocks to the already-waiting decoder.
235 + """
236 + try:
237 + req_data = await request.json()
238 + request_id = str(uuid.uuid4())
239 +
240 + prefill_client = _get_next(request.app, "prefill")
241 + decode_client = _get_next(request.app, "decode")
242 +
243 + decode_data = req_data.copy()
244 + decode_data["kv_transfer_params"] = {
245 + "prefill": {
246 + "kv_request_id": request_id,
247 + "remote_host": global_args.p2p_connector_host,
248 + "remote_port": global_args.p2p_connector_port,
249 + },
250 + }
251 +
252 + async def generate():
253 + queue: asyncio.Queue = asyncio.Queue()
254 +
255 + async def _run_decode():
256 + try:
257 + async for chunk in _stream_decode(
258 + decode_client, api, decode_data, request_id
259 + ):
260 + await queue.put(("data", chunk))
261 + except Exception as exc:
262 + await queue.put(("error", exc))
263 + finally:
264 + await queue.put(("done", None))
265 +
266 + # 1. Start decode request — decoder is now waiting for KV blocks
267 + asyncio.create_task(_run_decode())
268 +
269 + # 2. Send prefill — blocks are computed and pushed to the decoder
270 + try:
271 + await _prefill(prefill_client, api, req_data, request_id)
272 + except Exception as exc:
273 + logger.warning("decoder-first: prefill failed: %s", exc)
274 +
275 + logger.debug(
276 + "decoder-first: prefill done, streaming decode prefill=%s decode=%s",
277 + prefill_client,
278 + decode_client,
279 + )
280 +
281 + # 3. Stream the decode response
282 + while True:
283 + kind, value = await queue.get()
284 + if kind == "done":
285 + break
286 + if kind == "error":
287 + raise value # type: ignore[misc]
288 + yield value
289 +
290 + return StreamingResponse(generate(), media_type="application/json")
291 +
292 + except Exception as e:
293 + import sys
294 + import traceback
295 +
296 + print(f"Proxy error on {api}: {e}")
297 + print("".join(traceback.format_exception(*sys.exc_info())))
298 + raise
299 +
300 +
301 +def _route_handler(api: str):
302 + if global_args.decoder_first:
303 + return lambda req: _handle_completions_decoder_first(api, req)
304 + return lambda req: _handle_completions(api, req)
305 +
306 +
307 +@app.post("/v1/completions")
308 +async def completions(request: Request):
309 + return await _route_handler("/completions")(request)
310 +
311 +
312 +@app.post("/v1/chat/completions")
313 +async def chat_completions(request: Request):
314 + return await _route_handler("/chat/completions")(request)
315 +
316 +
317 +@app.get("/healthcheck")
318 +async def healthcheck():
319 + return {
320 + "status": "ok",
321 + "prefill_instances": len(app.state.prefill_clients),
322 + "decode_instances": len(app.state.decode_clients),
323 + }
324 +
325 +
326 +if __name__ == "__main__":
327 + global global_args
328 + global_args = parse_args()
329 + import uvicorn
330 +
331 + uvicorn.run(app, host=global_args.host, port=global_args.port)
tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh new
+312
@@ -0,0 +1,312 @@
1 +#!/bin/bash
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +#
5 +# Accuracy test driver for the p2p connector
6 +# (OffloadingConnector + TieringOffloadingSpec + p2p tier).
7 +#
8 +# Mirrors tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh:
9 +# brings up N prefillers + M decoders on the local host, fronts them
10 +# with p2p_connector_proxy.py, then runs the connector-agnostic
11 +# test_accuracy.py (gsm8k via lm_eval) against the proxy.
12 +#
13 +# Knobs (env vars unless flagged otherwise):
14 +# MODEL_NAMES space-separated model list (default: Llama-3.2-1B-Instruct)
15 +# NUM_PREFILL_INSTANCES default 1
16 +# NUM_DECODE_INSTANCES default 1
17 +# PREFILLER_TP_SIZE default 1
18 +# DECODER_TP_SIZE default 1
19 +# GPU_MEMORY_UTILIZATION default 0.45
20 +# MAX_MODEL_LEN default 512
21 +# PREFILL_BLOCK_SIZE default 128
22 +# DECODE_BLOCK_SIZE default 128
23 +# CPU_BYTES default 209715200 (200 MB)
24 +# VLLM_SERVE_EXTRA_ARGS comma-separated extra args for vllm serve
25 +# --decoder-first toggle decoder-first proxy mode
26 +#
27 +# Examples:
28 +# bash tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh
29 +# NUM_PREFILL_INSTANCES=2 NUM_DECODE_INSTANCES=2 \
30 +# bash tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh
31 +# bash tests/v1/kv_offload/tiering/p2p/run_accuracy_test.sh --decoder-first
32 +
33 +set -xe
34 +
35 +# ---------------------------------------------------------------------------
36 +# Args
37 +# ---------------------------------------------------------------------------
38 +DECODER_FIRST="false"
39 +
40 +while [[ $# -gt 0 ]]; do
41 + case $1 in
42 + --decoder-first)
43 + DECODER_FIRST="true"
44 + shift 1
45 + ;;
46 + *)
47 + echo "Unknown option $1"
48 + echo "Usage: $0 [--decoder-first]"
49 + exit 1
50 + ;;
51 + esac
52 +done
53 +
54 +# ---------------------------------------------------------------------------
55 +# Models
56 +# ---------------------------------------------------------------------------
57 +MODEL_NAMES=${MODEL_NAMES:-}
58 +if [[ -n "$MODEL_NAMES" ]]; then
59 + # shellcheck disable=SC2206
60 + MODELS=($MODEL_NAMES)
61 +else
62 + MODELS=(
63 + "meta-llama/Llama-3.2-1B-Instruct"
64 + )
65 +fi
66 +
67 +# ---------------------------------------------------------------------------
68 +# Defaults
69 +# ---------------------------------------------------------------------------
70 +NUM_PREFILL_INSTANCES=${NUM_PREFILL_INSTANCES:-1}
71 +NUM_DECODE_INSTANCES=${NUM_DECODE_INSTANCES:-1}
72 +PREFILLER_TP_SIZE=${PREFILLER_TP_SIZE:-1}
73 +DECODER_TP_SIZE=${DECODER_TP_SIZE:-1}
74 +GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.45}
75 +MAX_MODEL_LEN=${MAX_MODEL_LEN:-512}
76 +PREFILL_BLOCK_SIZE=${PREFILL_BLOCK_SIZE:-128}
77 +DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128}
78 +CPU_BYTES=${CPU_BYTES:-209715200}
79 +VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
80 +
81 +# Base ports — per-instance offsets layered on top.
82 +PREFILL_HTTP_BASE=8100
83 +DECODE_HTTP_BASE=8200
84 +PREFILL_PD_BASE=7777
85 +DECODE_PD_BASE=$((PREFILL_PD_BASE + NUM_PREFILL_INSTANCES))
86 +PROXY_PORT=8192
87 +P2P_HOST=127.0.0.1
88 +
89 +# ---------------------------------------------------------------------------
90 +# Resolve repo root + venv (works in .venv and /workspace/venv pods)
91 +# ---------------------------------------------------------------------------
92 +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
93 +GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../../.." && pwd -P)}"
94 +
95 +if [[ -z "${VLLM_BIN:-}" ]]; then
96 + if [[ -x "${GIT_ROOT}/.venv/bin/vllm" ]]; then
97 + VLLM_BIN="${GIT_ROOT}/.venv/bin/vllm"
98 + elif [[ -x "/workspace/venv/bin/vllm" ]]; then
99 + VLLM_BIN="/workspace/venv/bin/vllm"
100 + else
101 + VLLM_BIN="$(command -v vllm)"
102 + fi
103 +fi
104 +if [[ -z "${PYTHON_BIN:-}" ]]; then
105 + if [[ -x "${GIT_ROOT}/.venv/bin/python" ]]; then
106 + PYTHON_BIN="${GIT_ROOT}/.venv/bin/python"
107 + elif [[ -x "/workspace/venv/bin/python" ]]; then
108 + PYTHON_BIN="/workspace/venv/bin/python"
109 + else
110 + PYTHON_BIN="$(command -v python3 || command -v python)"
111 + fi
112 +fi
113 +echo "Using vllm: ${VLLM_BIN}"
114 +echo "Using python: ${PYTHON_BIN}"
115 +
116 +SMI_BIN=$(command -v nvidia-smi || command -v rocm-smi || echo "")
117 +
118 +# Trap SIGINT/SIGTERM/EXIT to kill background jobs.
119 +trap 'kill $(jobs -pr) 2>/dev/null || true' SIGINT SIGTERM EXIT
120 +
121 +# ---------------------------------------------------------------------------
122 +# Helpers
123 +# ---------------------------------------------------------------------------
124 +wait_for_server() {
125 + local port=$1
126 + timeout 1200 bash -c "
127 + until curl -s localhost:${port}/v1/completions > /dev/null; do
128 + sleep 1
129 + done" && return 0 || return 1
130 +}
131 +
132 +cleanup_instances() {
133 + echo "Cleaning up any running vLLM / proxy instances..."
134 + pkill -f "vllm serve" || true
135 + pkill -f "p2p_connector_proxy.py" || true
136 + sleep 2
137 +}
138 +
139 +get_num_gpus() {
140 + if [[ "$SMI_BIN" == *"nvidia"* ]]; then
141 + $SMI_BIN --query-gpu=name --format=csv,noheader | wc -l
142 + elif [[ "$SMI_BIN" == *"rocm"* ]]; then
143 + $SMI_BIN -l | grep -c GPU
144 + else
145 + echo "1"
146 + fi
147 +}
148 +
149 +# Build the OffloadingConnector kv-transfer-config for a given PD port.
150 +# Mirrors deploy_local.sh:131.
151 +build_kv_config() {
152 + local pd_port=$1
153 + printf '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"spec_name":"TieringOffloadingSpec","cpu_bytes_to_use":%s,"secondary_tiers":[{"type":"p2p","host":"%s","port":%s}]}}' \
154 + "${CPU_BYTES}" "${P2P_HOST}" "${pd_port}"
155 +}
156 +
157 +# ---------------------------------------------------------------------------
158 +# Per-model run
159 +# ---------------------------------------------------------------------------
160 +run_tests_for_model() {
161 + local model_name=$1
162 + echo "================================"
163 + echo "Testing model: $model_name"
164 + echo " prefillers=${NUM_PREFILL_INSTANCES} (tp=${PREFILLER_TP_SIZE})"
165 + echo " decoders=${NUM_DECODE_INSTANCES} (tp=${DECODER_TP_SIZE})"
166 + echo " decoder_first=${DECODER_FIRST}"
167 + echo "================================"
168 +
169 + PREFILL_HOSTS=()
170 + PREFILL_PORTS=()
171 + PREFILL_PD_PORTS=()
172 + DECODE_HOSTS=()
173 + DECODE_PORTS=()
174 + DECODE_PD_PORTS=()
175 +
176 + local num_gpus
177 + num_gpus=$(get_num_gpus)
178 + local next_gpu=0
179 +
180 + # ---- Prefillers ----
181 + for i in $(seq 0 $((NUM_PREFILL_INSTANCES-1))); do
182 + local gpu_id=$((i * PREFILLER_TP_SIZE % num_gpus))
183 + local cuda_devs="${gpu_id}"
184 + for (( j=1; j < PREFILLER_TP_SIZE; j++ )); do
185 + cuda_devs="${cuda_devs},$(((gpu_id + j) % num_gpus))"
186 + done
187 + next_gpu=$(((gpu_id + PREFILLER_TP_SIZE) % num_gpus))
188 +
189 + local http_port=$((PREFILL_HTTP_BASE + i))
190 + local pd_port=$((PREFILL_PD_BASE + i))
191 + local kv_cfg
192 + kv_cfg=$(build_kv_config "${pd_port}")
193 +
194 + echo "Prefiller $i: gpu=[${cuda_devs}] http=${http_port} pd=${pd_port}"
195 +
196 + BASE_CMD="CUDA_VISIBLE_DEVICES=${cuda_devs} \
197 + PYTHONHASHSEED=42 \
198 + ${VLLM_BIN} serve ${model_name} \
199 + --port ${http_port} \
200 + --enforce-eager \
201 + --block-size ${PREFILL_BLOCK_SIZE} \
202 + --gpu-memory-utilization ${GPU_MEMORY_UTILIZATION} \
203 + --max-model-len ${MAX_MODEL_LEN} \
204 + --tensor-parallel-size ${PREFILLER_TP_SIZE} \
205 + --kv-transfer-config '${kv_cfg}'"
206 +
207 + if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
208 + IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS"
209 + for arg in "${extra_args[@]}"; do
210 + BASE_CMD="${BASE_CMD} $arg"
211 + done
212 + fi
213 +
214 + eval "${BASE_CMD} &"
215 +
216 + PREFILL_HOSTS+=("${P2P_HOST}")
217 + PREFILL_PORTS+=("${http_port}")
218 + PREFILL_PD_PORTS+=("${pd_port}")
219 + done
220 +
221 + # ---- Decoders ----
222 + for i in $(seq 0 $((NUM_DECODE_INSTANCES-1))); do
223 + local gpu_id=$(((next_gpu + i * DECODER_TP_SIZE) % num_gpus))
224 + local cuda_devs="${gpu_id}"
225 + for (( j=1; j < DECODER_TP_SIZE; j++ )); do
226 + cuda_devs="${cuda_devs},$(((gpu_id + j) % num_gpus))"
227 + done
228 +
229 + local http_port=$((DECODE_HTTP_BASE + i))
230 + local pd_port=$((DECODE_PD_BASE + i))
231 + local kv_cfg
232 + kv_cfg=$(build_kv_config "${pd_port}")
233 +
234 + echo "Decoder $i: gpu=[${cuda_devs}] http=${http_port} pd=${pd_port}"
235 +
236 + BASE_CMD="CUDA_VISIBLE_DEVICES=${cuda_devs} \
237 + PYTHONHASHSEED=42 \
238 + ${VLLM_BIN} serve ${model_name} \
239 + --port ${http_port} \
240 + --enforce-eager \
241 + --block-size ${DECODE_BLOCK_SIZE} \
242 + --gpu-memory-utilization ${GPU_MEMORY_UTILIZATION} \
243 + --max-model-len ${MAX_MODEL_LEN} \
244 + --tensor-parallel-size ${DECODER_TP_SIZE} \
245 + --kv-transfer-config '${kv_cfg}'"
246 +
247 + if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
248 + IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS"
249 + for arg in "${extra_args[@]}"; do
250 + BASE_CMD="${BASE_CMD} $arg"
251 + done
252 + fi
253 +
254 + eval "${BASE_CMD} &"
255 +
256 + DECODE_HOSTS+=("${P2P_HOST}")
257 + DECODE_PORTS+=("${http_port}")
258 + DECODE_PD_PORTS+=("${pd_port}")
259 + done
260 +
261 + # ---- Wait for HTTP readiness ----
262 + for port in "${PREFILL_PORTS[@]}"; do
263 + echo "Waiting for prefill instance on port $port to start..."
264 + wait_for_server "$port"
265 + done
266 + for port in "${DECODE_PORTS[@]}"; do
267 + echo "Waiting for decode instance on port $port to start..."
268 + wait_for_server "$port"
269 + done
270 +
271 + # ---- Proxy ----
272 + # The proxy currently advertises a single prefiller PD address to decoders.
273 + # For the 1xM and matched NxM common cases the first prefiller's PD coords
274 + # are the right pick; multi-prefiller PD round-robin is a follow-up.
275 + PROXY_CMD="${PYTHON_BIN} ${SCRIPT_DIR}/p2p_connector_proxy.py \
276 + --port ${PROXY_PORT} \
277 + --host ${P2P_HOST} \
278 + --prefiller-hosts ${PREFILL_HOSTS[*]} \
279 + --prefiller-ports ${PREFILL_PORTS[*]} \
280 + --decoder-hosts ${DECODE_HOSTS[*]} \
281 + --decoder-ports ${DECODE_PORTS[*]} \
282 + --p2p-connector-host ${P2P_HOST} \
283 + --p2p-connector-port ${PREFILL_PD_PORTS[0]} \
284 + --decoder-p2p-connector-host ${P2P_HOST} \
285 + --decoder-p2p-connector-port ${DECODE_PD_PORTS[0]}"
286 +
287 + if [[ "${DECODER_FIRST}" == "true" ]]; then
288 + PROXY_CMD="${PROXY_CMD} --decoder-first"
289 + fi
290 +
291 + echo "Starting proxy: ${PROXY_CMD}"
292 + eval "${PROXY_CMD} &"
293 +
294 + sleep 5
295 +
296 + # ---- Run accuracy test (reused from nixl_integration) ----
297 + echo "Running tests for $model_name"
298 + TEST_MODEL=$model_name "${PYTHON_BIN}" -m pytest -s -x \
299 + "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_accuracy.py"
300 +
301 + cleanup_instances
302 + sleep 3
303 +}
304 +
305 +# ---------------------------------------------------------------------------
306 +# Drive
307 +# ---------------------------------------------------------------------------
308 +for model in "${MODELS[@]}"; do
309 + run_tests_for_model "$model"
310 +done
311 +
312 +echo "All tests completed!"
tests/v1/kv_offload/tiering/p2p/test_data_transport.py new
+330
@@ -0,0 +1,330 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for DataTransport base class and NixlTransport."""
4 +
5 +from __future__ import annotations
6 +
7 +import ctypes
8 +from unittest.mock import MagicMock, patch
9 +
10 +import numpy as np
11 +
12 +from vllm.v1.kv_offload.tiering.p2p.data.base import PollResult
13 +from vllm.v1.kv_offload.tiering.p2p.data.nixl import NixlTransport
14 +
15 +# ---------------------------------------------------------------------------
16 +# DataTransport base class tests
17 +# ---------------------------------------------------------------------------
18 +
19 +
20 +class TestDataTransportBase:
21 + """Tests for the DataTransport abstract base properties."""
22 +
23 + def _make_view(self, num_blocks: int = 8, block_len: int = 1024) -> memoryview:
24 + """Create a memoryview with the given shape."""
25 + buf = np.zeros((num_blocks, block_len), dtype=np.uint8)
26 + return memoryview(buf)
27 +
28 + def test_properties(self):
29 + """base_addr, num_blocks, block_len are set from memoryview shape."""
30 + view = self._make_view(num_blocks=4, block_len=2048)
31 +
32 + # Use NixlTransport (concrete) with NIXL mocked away
33 + with patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", None):
34 + transport = NixlTransport("test:1", view)
35 +
36 + assert transport.num_blocks == 4
37 + assert transport.block_len == 2048
38 + assert transport.base_addr == ctypes.addressof(ctypes.c_char.from_buffer(view))
39 +
40 + def test_config_fingerprint_empty_when_no_fields(self):
41 + """No config fields → empty fingerprint."""
42 + view = self._make_view()
43 + with patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", None):
44 + transport = NixlTransport("test:1", view, config_fields=None)
45 + assert transport.config_fingerprint == ""
46 +
47 + def test_config_fingerprint_deterministic(self):
48 + """Same config fields → same fingerprint."""
49 + view = self._make_view()
50 + fields = {"model": "llama", "dtype": "float16", "block_size_factor": 1}
51 + with patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", None):
52 + t1 = NixlTransport("test:1", view, config_fields=fields)
53 + t2 = NixlTransport("test:2", view, config_fields=fields)
54 + assert t1.config_fingerprint == t2.config_fingerprint
55 + assert len(t1.config_fingerprint) == 16
56 +
57 + def test_config_fingerprint_differs_for_different_fields(self):
58 + """Different config fields → different fingerprint."""
59 + view = self._make_view()
60 + with patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", None):
61 + t1 = NixlTransport("test:1", view, config_fields={"model": "a"})
62 + t2 = NixlTransport("test:2", view, config_fields={"model": "b"})
63 + assert t1.config_fingerprint != t2.config_fingerprint
64 +
65 +
66 +# ---------------------------------------------------------------------------
67 +# NixlTransport tests (with mocked NIXL agent)
68 +# ---------------------------------------------------------------------------
69 +
70 +
71 +class TestNixlTransportWithMockedAgent:
72 + """Tests for NixlTransport logic with a mocked NIXL agent."""
73 +
74 + def _make_transport(self) -> NixlTransport:
75 + """Create a NixlTransport with mocked NIXL internals."""
76 + view = memoryview(np.zeros((8, 1024), dtype=np.uint8))
77 +
78 + with patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", None):
79 + transport = NixlTransport("test:1", view)
80 +
81 + # Manually set up a mock agent after construction
82 + agent = MagicMock()
83 + agent.add_remote_agent.return_value = "nixl-peer-name"
84 + agent.get_xfer_descs.return_value = MagicMock()
85 + agent.prep_xfer_dlist.return_value = MagicMock()
86 + agent.make_prepped_xfer.return_value = MagicMock(name="handle")
87 + agent.transfer.return_value = None
88 + agent.check_xfer_state.return_value = "PROC"
89 + agent.get_agent_metadata.return_value = b"test-metadata"
90 +
91 + transport._agent = agent
92 + transport._local_dlist = MagicMock()
93 + return transport
94 +
95 + def test_available_false_without_nixl(self):
96 + """Without NIXL installed, available is False."""
97 + view = memoryview(np.zeros((4, 512), dtype=np.uint8))
98 + with patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", None):
99 + transport = NixlTransport("test:1", view)
100 + assert transport.available is False
101 +
102 + def test_available_true_with_agent(self):
103 + transport = self._make_transport()
104 + assert transport.available is True
105 +
106 + def test_get_agent_metadata(self):
107 + transport = self._make_transport()
108 + assert transport.get_agent_metadata() == b"test-metadata"
109 +
110 + def test_write_blocks_returns_none_for_unknown_peer(self):
111 + """write_blocks returns None if peer not registered."""
112 + transport = self._make_transport()
113 + result = transport.write_blocks("unknown:1", [0, 1], [2, 3])
114 + assert result is None
115 +
116 + def test_write_blocks_returns_transfer_id(self):
117 + """write_blocks returns an integer transfer_id on success."""
118 + transport = self._make_transport()
119 + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
120 +
121 + tid = transport.write_blocks("peer:1", [0, 1], [2, 3])
122 + assert tid is not None
123 + assert isinstance(tid, int)
124 +
125 + def test_write_blocks_increments_transfer_id(self):
126 + """Each write_blocks call gets a unique transfer_id."""
127 + transport = self._make_transport()
128 + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
129 +
130 + tid1 = transport.write_blocks("peer:1", [0], [1])
131 + tid2 = transport.write_blocks("peer:1", [2], [3])
132 + assert tid1 != tid2
133 +
134 + def test_poll_empty_when_no_inflight(self):
135 + """poll returns empty when nothing is inflight."""
136 + transport = self._make_transport()
137 + result = transport.poll()
138 + assert result == PollResult(done=(), failed=())
139 +
140 + def test_poll_returns_done_when_transfer_completes(self):
141 + """Completed transfer appears in poll().done."""
142 + transport = self._make_transport()
143 + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
144 +
145 + tid = transport.write_blocks("peer:1", [0], [1])
146 +
147 + # Simulate completion
148 + transport._agent.check_xfer_state.return_value = "DONE"
149 + result = transport.poll()
150 +
151 + assert tid in result.done
152 + assert result.failed == ()
153 + # Handle released
154 + transport._agent.release_xfer_handle.assert_called()
155 +
156 + def test_poll_returns_failed_for_error_state(self):
157 + """Transfer in error state appears in poll().failed."""
158 + transport = self._make_transport()
159 + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
160 +
161 + tid = transport.write_blocks("peer:1", [0], [1])
162 +
163 + transport._agent.check_xfer_state.return_value = "ERR"
164 + result = transport.poll()
165 +
166 + assert result.done == ()
167 + assert tid in result.failed
168 +
169 + def test_poll_ignores_in_progress(self):
170 + """Transfers in PROC/PEND state stay inflight."""
171 + transport = self._make_transport()
172 + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
173 +
174 + transport.write_blocks("peer:1", [0], [1])
175 +
176 + transport._agent.check_xfer_state.return_value = "PROC"
177 + result = transport.poll()
178 + assert result.done == ()
179 + assert result.failed == ()
180 +
181 + transport._agent.check_xfer_state.return_value = "PEND"
182 + result = transport.poll()
183 + assert result.done == ()
184 + assert result.failed == ()
185 +
186 + def test_cancel_removes_inflight(self):
187 + """cancel removes transfers and releases handles."""
188 + transport = self._make_transport()
189 + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
190 +
191 + tid = transport.write_blocks("peer:1", [0], [1])
192 + assert tid in transport._inflight
193 +
194 + result = transport.cancel([tid])
195 + assert result == []
196 + assert tid not in transport._inflight
197 + transport._agent.release_xfer_handle.assert_called()
198 +
199 + def test_cancel_ignores_unknown_ids(self):
200 + """cancel with unknown IDs doesn't crash."""
201 + transport = self._make_transport()
202 + assert transport.cancel([999, 1000]) == []
203 + assert transport.cancel([999, 1000], mode="wait") == []
204 +
205 + def test_cancel_wait_release_succeeds(self):
206 + """wait-mode cancel that succeeds pops the entry and returns []."""
207 + transport = self._make_transport()
208 + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
209 +
210 + tid = transport.write_blocks("peer:1", [0], [1])
211 + assert tid in transport._inflight
212 +
213 + result = transport.cancel([tid], mode="wait")
214 + assert result == []
215 + assert tid not in transport._inflight
216 + transport._agent.release_xfer_handle.assert_called_once()
217 +
218 + def test_cancel_wait_release_raises(self):
219 + """wait-mode cancel keeps the entry and returns the tid on raise."""
220 + transport = self._make_transport()
221 + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
222 +
223 + tid = transport.write_blocks("peer:1", [0], [1])
224 + transport._agent.release_xfer_handle.side_effect = RuntimeError(
225 + "NIXL_ERR_REPOST_ACTIVE"
226 + )
227 +
228 + result = transport.cancel([tid], mode="wait")
229 + assert result == [tid]
230 + assert tid in transport._inflight
231 +
232 + def test_cancel_wait_then_poll_completes(self):
233 + """A wait-cancel that left a tid pending later completes via poll."""
234 + transport = self._make_transport()
235 + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
236 +
237 + tid = transport.write_blocks("peer:1", [0], [1])
238 + transport._agent.release_xfer_handle.side_effect = RuntimeError("busy")
239 + assert transport.cancel([tid], mode="wait") == [tid]
240 + assert tid in transport._inflight
241 +
242 + transport._agent.release_xfer_handle.side_effect = None
243 + transport._agent.check_xfer_state.return_value = "DONE"
244 +
245 + result = transport.poll()
246 + assert tid in result.done
247 + assert tid not in transport._inflight
248 +
249 + def test_add_and_remove_remote_peer(self):
250 + transport = self._make_transport()
251 + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
252 + assert "peer:1" in transport._remote_dlists
253 +
254 + transport.remove_remote_peer("peer:1")
255 + assert "peer:1" not in transport._remote_dlists
256 + transport._agent.release_dlist_handle.assert_called()
257 + transport._agent.remove_remote_agent.assert_called()
258 +
259 + def test_close_releases_everything(self):
260 + """close releases all handles and clears state."""
261 + transport = self._make_transport()
262 + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
263 + transport.write_blocks("peer:1", [0], [1])
264 +
265 + transport.close()
266 + assert transport._agent is None
267 + assert transport._inflight == {}
268 + assert transport._remote_dlists == {}
269 +
270 +
271 +# ---------------------------------------------------------------------------
272 +# NIXL agent-config selection
273 +# ---------------------------------------------------------------------------
274 +
275 +
276 +class TestNixlAgentConfigSelection:
277 + """Tests that backends/num_threads pick the right nixl_agent_config call.
278 +
279 + Mirrors the conditional in
280 + vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py:325-329.
281 + """
282 +
283 + def _make_view(self) -> memoryview:
284 + return memoryview(np.zeros((4, 512), dtype=np.uint8))
285 +
286 + def test_non_ucx_backends_passes_backends_kwarg(self):
287 + """When any non-UCX backend is requested, pass backends + telemetry."""
288 + agent_cls = MagicMock()
289 + config_fn = MagicMock(return_value=MagicMock(name="cfg"))
290 + with (
291 + patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", agent_cls),
292 + patch(
293 + "vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgentConfig", config_fn
294 + ),
295 + ):
296 + NixlTransport("test:1", self._make_view(), backends=["MOONCAKE"])
297 +
298 + config_fn.assert_called_once_with(backends=["MOONCAKE"], capture_telemetry=True)
299 + # num_threads must NOT be passed on the non-UCX branch.
300 + assert "num_threads" not in config_fn.call_args.kwargs
301 +
302 + def test_ucx_only_passes_num_threads(self):
303 + """UCX-only configuration passes num_threads + telemetry, no backends."""
304 + agent_cls = MagicMock()
305 + config_fn = MagicMock(return_value=MagicMock(name="cfg"))
306 + with (
307 + patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", agent_cls),
308 + patch(
309 + "vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgentConfig", config_fn
310 + ),
311 + ):
312 + NixlTransport("test:1", self._make_view(), num_threads=8)
313 +
314 + config_fn.assert_called_once_with(num_threads=8, capture_telemetry=True)
315 + assert "backends" not in config_fn.call_args.kwargs
316 +
317 + def test_default_backends_is_ucx_only(self):
318 + """No backends arg → defaults to UCX-only branch."""
319 + agent_cls = MagicMock()
320 + config_fn = MagicMock(return_value=MagicMock(name="cfg"))
321 + with (
322 + patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", agent_cls),
323 + patch(
324 + "vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgentConfig", config_fn
325 + ),
326 + ):
327 + NixlTransport("test:1", self._make_view())
328 +
329 + # Default num_threads=4, no backends kwarg.
330 + config_fn.assert_called_once_with(num_threads=4, capture_telemetry=True)
tests/v1/kv_offload/tiering/p2p/test_manager.py new
+1370
@@ -0,0 +1,1370 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for P2PSecondaryTierManager.
4 +
5 +Tests the manager's job routing, session lifecycle, and result collection
6 +using fake transport and session objects.
7 +"""
8 +
9 +from __future__ import annotations
10 +
11 +import time
12 +from types import SimpleNamespace
13 +
14 +import numpy as np
15 +
16 +from vllm.v1.kv_offload.base import LookupResult, ReqContext
17 +from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult
18 +from vllm.v1.kv_offload.tiering.p2p import manager as manager_module
19 +from vllm.v1.kv_offload.tiering.p2p.manager import (
20 + _UNBOUND_STORE_TIMEOUT_S,
21 + P2PSecondaryTierManager,
22 +)
23 +from vllm.v1.kv_offload.tiering.p2p.session import (
24 + LoadResult,
25 + SessionPollResult,
26 + StoreResult,
27 +)
28 +
29 +# ---------------------------------------------------------------------------
30 +# Helpers
31 +# ---------------------------------------------------------------------------
32 +
33 +
34 +def _prefill_kv_params(
35 + remote_host: str = "10.0.0.1",
36 + remote_port: int = 8000,
37 + kv_request_id: str = "req-1",
38 +) -> dict:
39 + """Decoder-side kv_transfer_params: ``prefill`` sub-dict carries
40 + kv_request_id + remote_host + remote_port."""
41 + return {
42 + "prefill": {
43 + "kv_request_id": kv_request_id,
44 + "remote_host": remote_host,
45 + "remote_port": remote_port,
46 + },
47 + }
48 +
49 +
50 +def _decode_kv_params(kv_request_id: str = "req-1") -> dict:
51 + """Prefiller-side kv_transfer_params: ``decode`` sub-dict carries
52 + kv_request_id only."""
53 + return {"decode": {"kv_request_id": kv_request_id}}
54 +
55 +
56 +def _req_context(kv_params: dict | None = None) -> ReqContext:
57 + return ReqContext(req_id="test", kv_transfer_params=kv_params)
58 +
59 +
60 +def _job_metadata(
61 + job_id: int,
62 + keys: list[bytes] | None = None,
63 + block_ids: list[int] | None = None,
64 + kv_params: dict | None = None,
65 +) -> JobMetadata:
66 + if keys is None:
67 + keys = [b"key1"]
68 + if block_ids is None:
69 + block_ids = list(range(len(keys)))
70 + return JobMetadata(
71 + job_id=job_id,
72 + keys=keys,
73 + block_ids=np.array(block_ids),
74 + is_promotion=False,
75 + req_context=_req_context(kv_params),
76 + )
77 +
78 +
79 +def _make_manager() -> P2PSecondaryTierManager:
80 + """Create a manager with stubbed __init__."""
81 + mgr = P2PSecondaryTierManager.__new__(P2PSecondaryTierManager)
82 + mgr._local_id = "127.0.0.1:7777"
83 + mgr._finished_jobs = []
84 + mgr._failed_req_ids = set()
85 + mgr._sessions = {}
86 + mgr._kv_to_session = {}
87 + mgr._unbound_stores = {}
88 + return mgr
89 +
90 +
91 +# ---------------------------------------------------------------------------
92 +# Tests for _remote_id_from_params
93 +# ---------------------------------------------------------------------------
94 +
95 +
96 +class TestRemoteIdFromParams:
97 + def test_valid_params(self):
98 + result = P2PSecondaryTierManager._remote_id_from_params(
99 + {"remote_host": "10.0.0.1", "remote_port": 8000}
100 + )
101 + assert result == "10.0.0.1:8000"
102 +
103 + def test_missing_host(self):
104 + result = P2PSecondaryTierManager._remote_id_from_params({"remote_port": 8000})
105 + assert result is None
106 +
107 + def test_missing_port(self):
108 + result = P2PSecondaryTierManager._remote_id_from_params(
109 + {"remote_host": "10.0.0.1"}
110 + )
111 + assert result is None
112 +
113 + def test_empty_dict(self):
114 + result = P2PSecondaryTierManager._remote_id_from_params({})
115 + assert result is None
116 +
117 +
118 +# ---------------------------------------------------------------------------
119 +# Tests for lookup
120 +# ---------------------------------------------------------------------------
121 +
122 +
123 +class TestLookup:
124 + def test_lookup_returns_miss_without_kv_params(self):
125 + mgr = _make_manager()
126 + ctx = _req_context(kv_params=None)
127 + assert mgr.lookup(b"key", ctx) is LookupResult.MISS
128 +
129 + def test_lookup_returns_miss_without_required_fields(self):
130 + mgr = _make_manager()
131 + ctx = _req_context(kv_params={"prefill": {"remote_host": "x"}})
132 + assert mgr.lookup(b"key", ctx) is LookupResult.MISS
133 +
134 + def test_lookup_returns_hit_for_valid_request(self):
135 + mgr = _make_manager()
136 + ctx = _req_context(kv_params=_prefill_kv_params())
137 + assert mgr.lookup(b"key", ctx) is LookupResult.HIT
138 +
139 + def test_lookup_returns_miss_for_failed_request(self):
140 + mgr = _make_manager()
141 + mgr._failed_req_ids.add("req-1")
142 + ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-1"))
143 + assert mgr.lookup(b"key", ctx) is LookupResult.MISS
144 +
145 + def test_lookup_returns_hit_for_different_request_id(self):
146 + mgr = _make_manager()
147 + mgr._failed_req_ids.add("req-1")
148 + ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-2"))
149 + assert mgr.lookup(b"key", ctx) is LookupResult.HIT
150 +
151 + def test_lookup_returns_miss_without_prefill_key(self):
152 + """No ``prefill`` sub-dict means the request was not routed for
153 + remote prefill — local prefill should run instead, so lookup()
154 + returns MISS even when a stale ``decode`` block is present."""
155 + mgr = _make_manager()
156 + ctx = _req_context(kv_params=_decode_kv_params())
157 + assert mgr.lookup(b"key", ctx) is LookupResult.MISS
158 +
159 +
160 +# ---------------------------------------------------------------------------
161 +# Tests for submit_store
162 +# ---------------------------------------------------------------------------
163 +
164 +
165 +class TestSubmitStore:
166 + def test_no_decode_succeeds_immediately(self):
167 + """Without a ``decode`` block, job succeeds immediately."""
168 + mgr = _make_manager()
169 + job = _job_metadata(job_id=1, kv_params={})
170 + mgr.submit_store(job)
171 + assert mgr._finished_jobs == [JobResult(job_id=1, success=True)]
172 +
173 + def test_missing_kv_request_id_fails(self):
174 + """Missing kv_request_id inside ``decode`` fails the job."""
175 + mgr = _make_manager()
176 + params: dict = {"decode": {}}
177 + job = _job_metadata(job_id=1, kv_params=params)
178 + mgr.submit_store(job)
179 + assert mgr._finished_jobs == [JobResult(job_id=1, success=False)]
180 +
181 + def test_no_binding_yet_parks_in_unbound_stores(self):
182 + """submit_store without a bound session buffers the batch keyed
183 + by kv_request_id; no session is pre-created (peer_id is unknown
184 + to the producer at store time)."""
185 + mgr = _make_manager()
186 + job = _job_metadata(
187 + job_id=1,
188 + keys=[b"k1", b"k2"],
189 + block_ids=[3, 4],
190 + kv_params=_decode_kv_params(kv_request_id="req-1"),
191 + )
192 + mgr.submit_store(job)
193 +
194 + assert mgr._sessions == {}
195 + assert mgr._finished_jobs == []
196 + batches = mgr._unbound_stores["req-1"]
197 + assert len(batches) == 1
198 + assert (
199 + batches[0].job_id,
200 + list(batches[0].keys),
201 + list(batches[0].block_ids),
202 + ) == (
203 + 1,
204 + [b"k1", b"k2"],
205 + [3, 4],
206 + )
207 +
208 + def test_routes_to_bound_session(self):
209 + """If a session has already received FetchMsg for this
210 + kv_request_id (so _kv_to_session is populated), submit_store
211 + forwards directly to that session rather than re-buffering."""
212 + mgr = _make_manager()
213 + bound = _FakeSession(peer_id="10.0.0.1:8000", connected=True)
214 + mgr._kv_to_session["req-1"] = bound # type: ignore[assignment]
215 + # Note: _sessions is intentionally untouched — this test isolates
216 + # the kv_request_id → session fast path.
217 + job = _job_metadata(
218 + job_id=7,
219 + keys=[b"k1", b"k2"],
220 + block_ids=[3, 4],
221 + kv_params=_decode_kv_params(kv_request_id="req-1"),
222 + )
223 + mgr.submit_store(job)
224 +
225 + assert len(bound.stores_added) == 1
226 + kv_req_id, keys, _, job_id = bound.stores_added[0]
227 + assert (kv_req_id, keys, job_id) == ("req-1", [b"k1", b"k2"], 7)
228 + assert mgr._unbound_stores == {}
229 + assert mgr._finished_jobs == []
230 +
231 + def test_extra_top_level_keys_are_ignored(self):
232 + """Producer-side kv_transfer_params should not pre-create a
233 + session even when a stale caller still passes a top-level
234 + ``remote_host``/``remote_port`` next to ``decode``."""
235 + mgr = _make_manager()
236 + params = _decode_kv_params()
237 + params["remote_host"] = "stale"
238 + params["remote_port"] = 12345
239 + job = _job_metadata(job_id=1, kv_params=params)
240 + mgr.submit_store(job)
241 + # No session pre-created, no peer-keyed state.
242 + assert mgr._sessions == {}
243 + assert "req-1" in mgr._unbound_stores
244 +
245 +
246 +# ---------------------------------------------------------------------------
247 +# Tests for submit_load
248 +# ---------------------------------------------------------------------------
249 +
250 +
251 +class TestSubmitLoad:
252 + def test_missing_params_fails(self):
253 + """Missing required kv_params fields fails the job."""
254 + mgr = _make_manager()
255 + job = _job_metadata(job_id=1, kv_params={})
256 + mgr.submit_load(job)
257 + assert mgr._finished_jobs == [JobResult(job_id=1, success=False)]
258 +
259 + def test_empty_keys_succeeds_immediately(self):
260 + """Empty key list succeeds immediately."""
261 + mgr = _make_manager()
262 + job = _job_metadata(
263 + job_id=1, keys=[], block_ids=[], kv_params=_prefill_kv_params()
264 + )
265 + mgr.submit_load(job)
266 + assert mgr._finished_jobs == [JobResult(job_id=1, success=True)]
267 +
268 + def test_no_session_fails(self):
269 + """No session for peer fails and marks request failed."""
270 + mgr = _make_manager()
271 + job = _job_metadata(job_id=1, kv_params=_prefill_kv_params())
272 + mgr.submit_load(job)
273 + assert mgr._finished_jobs == [JobResult(job_id=1, success=False)]
274 + assert "req-1" in mgr._failed_req_ids
275 +
276 + def test_happy_path_with_active_session(self):
277 + """When the peer's session exists, submit_load forwards to
278 + session.request_blocks and does NOT add a finished result yet."""
279 + mgr = _make_manager()
280 + peer_id = "10.0.0.1:8000"
281 + existing = _FakeSession(peer_id=peer_id, connected=True)
282 + mgr._sessions[peer_id] = existing # type: ignore[assignment]
283 + job = _job_metadata(
284 + job_id=42,
285 + keys=[b"k1", b"k2"],
286 + block_ids=[5, 6],
287 + kv_params=_prefill_kv_params(kv_request_id="req-42"),
288 + )
289 + mgr.submit_load(job)
290 +
291 + assert existing.requests == [(42, "req-42")]
292 + assert mgr._finished_jobs == []
293 + assert "req-42" not in mgr._failed_req_ids
294 +
295 +
296 +# ---------------------------------------------------------------------------
297 +# Tests for on_request_finished
298 +# ---------------------------------------------------------------------------
299 +
300 +
301 +class TestOnRequestFinished:
302 + def _make_with_failed(self) -> P2PSecondaryTierManager:
303 + mgr = _make_manager()
304 + mgr._failed_req_ids = {"req-1"}
305 + return mgr
306 +
307 + def test_prunes_failed_req_ids(self):
308 + mgr = self._make_with_failed()
309 + ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-1"))
310 + mgr.on_request_finished(ctx)
311 + assert "req-1" not in mgr._failed_req_ids
312 +
313 + def test_no_kv_params_does_nothing(self):
314 + mgr = self._make_with_failed()
315 + ctx = _req_context(kv_params=None)
316 + mgr.on_request_finished(ctx)
317 + assert "req-1" in mgr._failed_req_ids
318 +
319 + def test_no_kv_request_id_does_nothing(self):
320 + mgr = self._make_with_failed()
321 + ctx = _req_context(kv_params={"remote_host": "x", "remote_port": 1})
322 + mgr.on_request_finished(ctx)
323 + assert "req-1" in mgr._failed_req_ids
324 +
325 + def test_decoder_side_calls_session_finish_request(self):
326 + """Decoder-side finish (``prefill`` set) still routes via peer_id
327 + because the consumer addresses the producer it loaded from. The
328 + session's finish_request cancels the client-role load."""
329 + mgr = _make_manager()
330 + peer_id = "10.0.0.1:8000"
331 + session = _FakeSession(peer_id=peer_id)
332 + mgr._sessions[peer_id] = session
333 + ctx = _req_context(kv_params=_prefill_kv_params(kv_request_id="req-1"))
334 + mgr.on_request_finished(ctx)
335 + assert session.finishes == ["req-1"]
336 +
337 + def test_prefiller_bound_id_routes_via_kv_to_session(self):
338 + """Prefiller-side finish for an id whose session is already bound
339 + (FetchMsg received) routes via _kv_to_session and pops the entry."""
340 + mgr = _make_manager()
341 + bound = _FakeSession(peer_id="some-peer:1", connected=True)
342 + mgr._kv_to_session["req-1"] = bound # type: ignore[assignment]
343 + ctx = _req_context(kv_params=_decode_kv_params(kv_request_id="req-1"))
344 + mgr.on_request_finished(ctx)
345 + assert bound.finishes == ["req-1"]
346 + assert "req-1" not in mgr._kv_to_session
347 +
348 + def test_prefiller_unbound_id_leaves_batches_parked(self):
349 + """Prefiller-side finish for an id with parked unbound batches
350 + and no session binding is a no-op on `_unbound_stores`. The
351 + parked batches survive until a peer fetches them or the
352 + `_reap_unbound_stores` timeout fires — `on_request_finished`
353 + must not evict them."""
354 + from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch
355 +
356 + mgr = _make_manager()
357 + mgr._unbound_stores["req-1"] = [
358 + _UnboundStoreBatch(job_id=10, keys=[b"k"], block_ids=[0]),
359 + _UnboundStoreBatch(job_id=11, keys=[b"k2"], block_ids=[1]),
360 + ]
361 + ctx = _req_context(kv_params=_decode_kv_params(kv_request_id="req-1"))
362 + mgr.on_request_finished(ctx)
363 + assert "req-1" in mgr._unbound_stores
364 + assert [b.job_id for b in mgr._unbound_stores["req-1"]] == [10, 11]
365 + outcomes = {(r.job_id, r.success) for r in mgr._finished_jobs}
366 + assert (10, False) not in outcomes
367 + assert (11, False) not in outcomes
368 +
369 +
370 +# ---------------------------------------------------------------------------
371 +# Tests for get_finished_jobs
372 +# ---------------------------------------------------------------------------
373 +
374 +
375 +class _FakeServerHalf:
376 + def __init__(self) -> None:
377 + self._inflight: dict[int, object] = {}
378 +
379 +
380 +class _FakeClientHalf:
381 + def __init__(self) -> None:
382 + self._inbound: dict[int, object] = {}
383 +
384 +
385 +class _FakeSession:
386 + """Fake bidirectional session that returns canned poll() results."""
387 +
388 + def __init__(
389 + self,
390 + peer_id: str = "fake:1",
391 + alive: bool = True,
392 + connected: bool = True,
393 + loads: list[LoadResult] | None = None,
394 + stores: list[StoreResult] | None = None,
395 + new_fetch_ids: list[str] | None = None,
396 + close_loads: list[tuple[int, str]] | None = None,
397 + close_stores: list[int] | None = None,
398 + ) -> None:
399 + self.peer_id = peer_id
400 + self.alive = alive
401 + self.connected = connected
402 + self.ready = True
403 + self._loads = loads or []
404 + self._stores = stores or []
405 + self._new_fetch_ids = new_fetch_ids or []
406 + self._close_loads = close_loads or []
407 + self._close_stores = close_stores or []
408 + self.requests: list[tuple[int, str]] = []
409 + self.stores_added: list[tuple[str, list, object, int]] = []
410 + self.attached: list[object] = []
411 + self.finishes: list[str] = []
412 + # Mirror P2PSession._server._inflight (transfer_id → handle) and
413 + # P2PSession._client._inbound for the shutdown-drain and drain_jobs
414 + # paths. Tests populate _server._inflight when needed.
415 + self._server = _FakeServerHalf()
416 + self._client = _FakeClientHalf()
417 +
418 + def poll(self):
419 + result = SessionPollResult(
420 + loads=self._loads,
421 + stores=self._stores,
422 + new_fetch_ids=self._new_fetch_ids,
423 + )
424 + self._loads = []
425 + self._stores = []
426 + self._new_fetch_ids = []
427 + return result
428 +
429 + def request_blocks(self, job_id, kv_request_id, keys, block_ids):
430 + self.requests.append((job_id, kv_request_id))
431 +
432 + def add_stored_blocks(self, kv_request_id, keys, block_ids, job_id):
433 + self.stores_added.append((kv_request_id, list(keys), block_ids, job_id))
434 +
435 + def attach_connection(self, conn):
436 + self.attached.append(conn)
437 + self.connected = True
438 +
439 + def finish_request(self, kv_request_id):
440 + self.finishes.append(kv_request_id)
441 +
442 + def close(self):
443 + return self._close_loads, self._close_stores
444 +
445 +
446 +class TestGetFinished:
447 + def _make(self) -> P2PSecondaryTierManager:
448 + mgr = _make_manager()
449 + mgr._finished_jobs = [
450 + JobResult(job_id=1, success=True),
451 + JobResult(job_id=2, success=False),
452 + ]
453 +
454 + class FakeControl:
455 + def poll(self):
456 + return []
457 +
458 + mgr._control = FakeControl() # type: ignore[assignment]
459 + mgr._data = None # type: ignore[assignment]
460 + return mgr
461 +
462 + def test_drains_finished_jobs(self):
463 + """get_finished_jobs returns and clears accumulated results."""
464 + mgr = self._make()
465 + results = list(mgr.get_finished_jobs())
466 + assert len(results) == 2
467 + assert JobResult(job_id=1, success=True) in results
468 + assert JobResult(job_id=2, success=False) in results
469 + # Second call returns empty.
470 + assert list(mgr.get_finished_jobs()) == []
471 +
472 + def test_reaps_dead_sessions(self):
473 + """Dead connected sessions are removed and their pending jobs failed."""
474 +
475 + class FakeData:
476 + def remove_remote_peer(self, pid):
477 + pass
478 +
479 + mgr = self._make()
480 + mgr._data = FakeData() # type: ignore[assignment]
481 + dead = _FakeSession(
482 + peer_id="dead:1234",
483 + alive=False,
484 + connected=True,
485 + close_loads=[(20, "req-load")],
486 + close_stores=[10, 11],
487 + )
488 + mgr._sessions["dead:1234"] = dead # type: ignore[assignment]
489 +
490 + results = list(mgr.get_finished_jobs())
491 + # 2 baseline + 1 failed load + 2 failed stores
492 + assert len(results) == 5
493 + assert JobResult(job_id=10, success=False) in results
494 + assert JobResult(job_id=11, success=False) in results
495 + assert JobResult(job_id=20, success=False) in results
496 + assert "dead:1234" not in mgr._sessions
497 + assert "req-load" in mgr._failed_req_ids
498 +
499 + def test_unbound_store_kept_within_timeout(self):
500 + """Recently-parked unbound stores stay across a poll."""
501 + mgr = self._make()
502 + from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch
503 +
504 + mgr._unbound_stores["req-fresh"] = [
505 + _UnboundStoreBatch(job_id=99, keys=[b"k"], block_ids=[0])
506 + ]
507 + list(mgr.get_finished_jobs())
508 + assert "req-fresh" in mgr._unbound_stores
509 +
510 + def test_unbound_store_reaped_after_timeout(self):
511 + """Unbound stores past _UNBOUND_STORE_TIMEOUT_S surface as failed
512 + and their kv_request_id lands in _failed_req_ids so a late
513 + FetchMsg/lookup doesn't try to satisfy them."""
514 + from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch
515 +
516 + mgr = self._make()
517 + stale = _UnboundStoreBatch(job_id=10, keys=[b"k"], block_ids=[0])
518 + # Backdate the submission so the head batch is past the deadline.
519 + stale.submitted_at = time.monotonic() - _UNBOUND_STORE_TIMEOUT_S - 1.0
520 + mgr._unbound_stores["req-stale"] = [
521 + stale,
522 + _UnboundStoreBatch(job_id=11, keys=[b"k2"], block_ids=[1]),
523 + ]
524 +
525 + results = list(mgr.get_finished_jobs())
526 +
527 + assert "req-stale" not in mgr._unbound_stores
528 + # 2 baseline + 2 buffered stores
529 + assert JobResult(job_id=10, success=False) in results
530 + assert JobResult(job_id=11, success=False) in results
531 + assert "req-stale" in mgr._failed_req_ids
532 +
533 + def test_submit_store_parks_unbound_batch(self):
534 + """submit_store on an unbound id appends a batch with a fresh
535 + submitted_at stamp so the unbound-store sweep can age it out."""
536 + mgr = _make_manager()
537 + job = _job_metadata(
538 + job_id=1, kv_params=_decode_kv_params(kv_request_id="req-1")
539 + )
540 + before = time.monotonic()
541 + mgr.submit_store(job)
542 + after = time.monotonic()
543 + batches = mgr._unbound_stores["req-1"]
544 + assert len(batches) == 1
545 + assert before <= batches[0].submitted_at <= after
546 +
547 +
548 +# ---------------------------------------------------------------------------
549 +# has_pending_work
550 +# ---------------------------------------------------------------------------
551 +
552 +
553 +class TestHasPendingWork:
554 + """has_pending_work() must always return True so the engine keeps
555 + ticking the offload pipeline — that's the only thread driving
556 + _control.poll() (incoming peer connects) and session.poll()
557 + (incoming fetch messages on existing sessions)."""
558 +
559 + def test_returns_true_unconditionally_to_keep_engine_ticking(self):
560 + """Even with no sessions and no jobs, has_pending_work() returns
561 + True so the engine keeps calling get_finished_jobs(), which is
562 + what drives _control.poll() for inbound peer connects."""
563 + mgr = _make_manager()
564 + assert mgr.has_pending_work() is True
565 +
566 + def test_returns_true_even_when_sessions_present(self):
567 + """The result is the same regardless of session state — there is
568 + no 'idle' branch."""
569 + mgr = _make_manager()
570 + mgr._sessions["peer:1"] = _FakeSession(peer_id="peer:1") # type: ignore[assignment]
571 + assert mgr.has_pending_work() is True
572 +
573 +
574 +# ---------------------------------------------------------------------------
575 +# Shutdown drain
576 +# ---------------------------------------------------------------------------
577 +
578 +
579 +class _ShutdownFakeData:
580 + """Fake DataTransport that records cancel/poll/close calls and
581 + drives the wait-cancel loop with a scriptable `still` queue."""
582 +
583 + def __init__(self, still_queue: list[list[int]] | None = None) -> None:
584 + # Each list in still_queue is the set of ids the next
585 + # cancel(mode="wait") should report as still inflight. The last
586 + # entry repeats once exhausted.
587 + self._still_queue = list(still_queue) if still_queue else [[]]
588 + self.cancel_calls: list[tuple[list[int], str]] = []
589 + self.poll_calls: int = 0
590 + self.close_calls: int = 0
591 +
592 + def cancel(self, transfer_ids, mode: str = "immediate") -> list[int]:
593 + ids = list(transfer_ids)
594 + self.cancel_calls.append((ids, mode))
595 + if mode == "wait":
596 + if len(self._still_queue) > 1:
597 + return list(self._still_queue.pop(0))
598 + return list(self._still_queue[0])
599 + return []
600 +
601 + def poll(self):
602 + self.poll_calls += 1
603 +
604 + class _Empty:
605 + done: list[int] = []
606 + failed: list[int] = []
607 +
608 + return _Empty()
609 +
610 + def close(self) -> None:
611 + self.close_calls += 1
612 +
613 +
614 +class _ShutdownFakeControl:
615 + def __init__(self) -> None:
616 + self.close_calls: int = 0
617 +
618 + def close(self) -> None:
619 + self.close_calls += 1
620 +
621 +
622 +class TestShutdownDrain:
623 + """shutdown() drains inflight transfers via cancel(mode='wait')
624 + before calling _data.close(), with a 3s deadline fallback to
625 + cancel(mode='immediate')."""
626 +
627 + def _prep(
628 + self,
629 + still_queue: list[list[int]] | None = None,
630 + inflight_ids: list[int] | None = None,
631 + ) -> tuple[P2PSecondaryTierManager, _ShutdownFakeData, _ShutdownFakeControl]:
632 + mgr = _make_manager()
633 + data = _ShutdownFakeData(still_queue=still_queue)
634 + control = _ShutdownFakeControl()
635 + mgr._data = data # type: ignore[assignment]
636 + mgr._control = control # type: ignore[assignment]
637 + if inflight_ids:
638 + session = _FakeSession(peer_id="peer:1", connected=True)
639 + session._server._inflight = {tid: object() for tid in inflight_ids}
640 + mgr._sessions["peer:1"] = session # type: ignore[assignment]
641 + return mgr, data, control
642 +
643 + def test_shutdown_drains_inflight_via_wait_cancel(self):
644 + # First cancel(wait) returns the input still inflight; second returns [].
645 + mgr, data, control = self._prep(
646 + still_queue=[[42, 43], []],
647 + inflight_ids=[42, 43],
648 + )
649 + mgr.shutdown()
650 +
651 + wait_calls = [c for c in data.cancel_calls if c[1] == "wait"]
652 + immediate_calls = [c for c in data.cancel_calls if c[1] == "immediate"]
653 + assert len(wait_calls) >= 1
654 + assert wait_calls[0][0] == [42, 43]
655 + assert immediate_calls == []
656 + # poll() was driven between cancel attempts.
657 + assert data.poll_calls >= 1
658 + # _data and _control were closed exactly once each, after the drain.
659 + assert data.close_calls == 1
660 + assert control.close_calls == 1
661 +
662 + def test_shutdown_force_cancels_after_timeout(self, monkeypatch):
663 + # Drain never completes — wait-cancel keeps returning the inflight set.
664 + # Use a synthetic clock so the test does not depend on real wallclock
665 + # being able to advance in <50ms on a loaded CI node:
666 + # call 1 (deadline calc): 100.0 -> deadline = 100.05
667 + # call 2 (loop predicate): 100.0 -> enters loop, one wait-cancel
668 + # call 3 (loop predicate): 100.06 -> exits, force-cancel runs
669 + monkeypatch.setattr(manager_module, "_SHUTDOWN_DRAIN_TIMEOUT_S", 0.05)
670 + monkeypatch.setattr(manager_module, "_DRAIN_SLEEP_S", 0.0)
671 + times = iter([100.0, 100.0, 100.06])
672 + # Patch via a fake module on `manager_module.time` so we do not mutate
673 + # the global `time` module — other code in the process (e.g. the
674 + # buildkite test collector's pytest_runtest_logreport hook) calls
675 + # time.monotonic() before monkeypatch teardown.
676 + fake_time = SimpleNamespace(monotonic=lambda: next(times), sleep=time.sleep)
677 + monkeypatch.setattr(manager_module, "time", fake_time)
678 +
679 + mgr, data, control = self._prep(
680 + still_queue=[[42]],
681 + inflight_ids=[42],
682 + )
683 + mgr.shutdown()
684 +
685 + wait_calls = [c for c in data.cancel_calls if c[1] == "wait"]
686 + immediate_calls = [c for c in data.cancel_calls if c[1] == "immediate"]
687 + assert len(wait_calls) == 1
688 + assert wait_calls[0][0] == [42]
689 + assert immediate_calls == [([42], "immediate")]
690 + assert data.close_calls == 1
691 + assert control.close_calls == 1
692 +
693 + def test_shutdown_no_inflight_skips_drain(self):
694 + mgr, data, control = self._prep()
695 + mgr.shutdown()
696 +
697 + assert data.cancel_calls == []
698 + assert data.poll_calls == 0
699 + assert data.close_calls == 1
700 + assert control.close_calls == 1
701 +
702 +
703 +# ---------------------------------------------------------------------------
704 +# Bidirectional regression test — both managers act as client AND server
705 +# toward each other on the same peer_id. This is the case the unification
706 +# is meant to fix.
707 +# ---------------------------------------------------------------------------
708 +
709 +
710 +class _LoopbackControl:
711 + """In-memory control transport that pairs two managers head-to-head.
712 +
713 + Each manager hands its outbound message buffer to the other's inbound
714 + queue on poll(). connect() returns a connection whose send() writes
715 + into the peer's inbound side; recv() reads what the peer's connect-or-
716 + poll path delivered for us.
717 + """
718 +
719 + def __init__(self, local_id: str) -> None:
720 + self._local_id = local_id
721 + self._peer: _LoopbackControl | None = None
722 + self._inbound_outbound: dict[str, _LoopbackConnection] = {}
723 + # Pending inbound for a peer that has not yet been registered.
724 + self._pending: list[tuple[str, dict]] = []
725 +
726 + def pair(self, peer: _LoopbackControl) -> None:
727 + self._peer = peer
728 + peer._peer = self
729 +
730 + def connect(self, peer_id: str):
731 + if peer_id in self._inbound_outbound:
732 + raise AssertionError(f"already connected to {peer_id}")
733 + conn = _LoopbackConnection(self, peer_id)
734 + self._inbound_outbound[peer_id] = conn
735 + return conn
736 +
737 + def poll(self):
738 + # Drain whatever the peer has sent toward us.
739 + new = []
740 + if self._peer is not None:
741 + for pid, msg in self._peer._drain_outbound_to(self._local_id):
742 + conn = self._inbound_outbound.get(pid)
743 + if conn is None:
744 + conn = _LoopbackConnection(self, pid)
745 + self._inbound_outbound[pid] = conn
746 + new.append(conn)
747 + conn._inbox.append(msg)
748 + return new
749 +
750 + def _drain_outbound_to(self, peer_local_id: str):
751 + # Peer is calling our poll → return all messages we've sent toward
752 + # peer_local_id (which is the peer's own local_id).
753 + out: list[tuple[str, dict]] = []
754 + for pid, conn in list(self._inbound_outbound.items()):
755 + # Each outgoing message goes to peer_local_id and is tagged
756 + # with the sender's local_id (i.e., self._local_id).
757 + for msg in conn._outbox:
758 + out.append((self._local_id, msg))
759 + conn._outbox.clear()
760 + return out
761 +
762 + def close(self):
763 + for conn in self._inbound_outbound.values():
764 + conn.close()
765 + self._inbound_outbound.clear()
766 +
767 +
768 +class _LoopbackConnection:
769 + def __init__(self, transport: _LoopbackControl, peer_id: str) -> None:
770 + self._transport = transport
771 + self.peer_id = peer_id
772 + self._inbox: list[dict] = []
773 + self._outbox: list[dict] = []
774 + self._closed = False
775 +
776 + @property
777 + def alive(self) -> bool:
778 + return not self._closed
779 +
780 + def send(self, msg: dict) -> None:
781 + if self._closed:
782 + raise RuntimeError("send on closed conn")
783 + self._outbox.append(msg)
784 +
785 + def recv(self) -> list[dict]:
786 + msgs = self._inbox
787 + self._inbox = []
788 + return msgs
789 +
790 + def mark_dead(self) -> None:
791 + self._closed = True
792 +
793 + def close(self) -> None:
794 + self._closed = True
795 +
796 +
797 +class _FakeData:
798 + """Minimal NIXL fake that lets matched transfers complete on the next poll."""
799 +
800 + def __init__(self, local_id: str) -> None:
801 + self._local_id = local_id
802 + self.block_len = 4096
803 + self.base_addr = 0x1000
804 + self.num_blocks = 16
805 + self.config_fingerprint = ""
806 + self._remote_peers: dict[str, dict] = {}
807 + self._inflight_done: list[int] = []
808 + self._next_id = 0
809 +
810 + def get_agent_metadata(self) -> bytes:
811 + return f"meta-{self._local_id}".encode()
812 +
813 + def add_remote_peer(
814 + self, peer_id, agent_metadata, base_addr, num_blocks, block_len
815 + ) -> None:
816 + self._remote_peers[peer_id] = {
817 + "agent_metadata": agent_metadata,
818 + "base_addr": base_addr,
819 + "num_blocks": num_blocks,
820 + "block_len": block_len,
821 + }
822 +
823 + def remove_remote_peer(self, peer_id: str) -> None:
824 + self._remote_peers.pop(peer_id, None)
825 +
826 + def write_blocks(self, peer_id, local_idxs, remote_idxs):
827 + if peer_id not in self._remote_peers:
828 + return None
829 + tid = self._next_id
830 + self._next_id += 1
831 + self._inflight_done.append(tid)
832 + return tid
833 +
834 + def poll(self):
835 + from vllm.v1.kv_offload.tiering.p2p.data.base import PollResult
836 +
837 + done = self._inflight_done[:]
838 + self._inflight_done.clear()
839 + return PollResult(done=done, failed=[])
840 +
841 + def cancel(self, transfer_ids) -> None:
842 + pass
843 +
844 + def close(self) -> None:
845 + pass
846 +
847 +
848 +def _build_paired_managers() -> tuple[P2PSecondaryTierManager, P2PSecondaryTierManager]:
849 + """Two managers each acting as both client and server toward the other.
850 +
851 + Wires _LoopbackControl pair + per-side _FakeData so transfers complete
852 + on the next poll. The test drives polling by calling get_finished_jobs(),
853 + which invokes _poll_once synchronously on the calling thread.
854 + """
855 + mgr_a = _make_manager()
856 + mgr_b = _make_manager()
857 + mgr_a._local_id = "A:1"
858 + mgr_b._local_id = "B:2"
859 +
860 + ctrl_a = _LoopbackControl(mgr_a._local_id)
861 + ctrl_b = _LoopbackControl(mgr_b._local_id)
862 + ctrl_a.pair(ctrl_b)
863 +
864 + mgr_a._control = ctrl_a # type: ignore[assignment]
865 + mgr_b._control = ctrl_b # type: ignore[assignment]
866 + mgr_a._data = _FakeData(mgr_a._local_id) # type: ignore[assignment]
867 + mgr_b._data = _FakeData(mgr_b._local_id) # type: ignore[assignment]
868 +
869 + return mgr_a, mgr_b
870 +
871 +
872 +class TestBidirectionalManager:
873 + """Two managers each load FROM and serve TO the other over a single peer."""
874 +
875 + def test_both_loads_succeed(self):
876 + mgr_a, mgr_b = _build_paired_managers()
877 +
878 + a_loads_kv = "req-AtoB-load" # A loads, B serves
879 + b_loads_kv = "req-BtoA-load" # B loads, A serves
880 +
881 + a_decoder_params = {
882 + "prefill": {
883 + "kv_request_id": a_loads_kv,
884 + "remote_host": "B",
885 + "remote_port": 2,
886 + },
887 + }
888 + b_decoder_params = {
889 + "prefill": {
890 + "kv_request_id": b_loads_kv,
891 + "remote_host": "A",
892 + "remote_port": 1,
893 + },
894 + }
895 + a_prefiller_params = {"decode": {"kv_request_id": b_loads_kv}}
896 + b_prefiller_params = {"decode": {"kv_request_id": a_loads_kv}}
897 +
898 + # 1. Both sides open client-role sessions toward the peer.
899 + mgr_a.on_new_request(_req_context(a_decoder_params))
900 + mgr_b.on_new_request(_req_context(b_decoder_params))
901 +
902 + # 2. Both sides store the blocks the peer will fetch.
903 + mgr_a.submit_store(
904 + _job_metadata(
905 + job_id=100,
906 + keys=[b"a-block"],
907 + block_ids=[0],
908 + kv_params=a_prefiller_params,
909 + )
910 + )
911 + mgr_b.submit_store(
912 + _job_metadata(
913 + job_id=200,
914 + keys=[b"b-block"],
915 + block_ids=[0],
916 + kv_params=b_prefiller_params,
917 + )
918 + )
919 +
920 + # 3. Both sides submit loads.
921 + mgr_a.submit_load(
922 + _job_metadata(
923 + job_id=101,
924 + keys=[b"b-block"],
925 + block_ids=[0],
926 + kv_params=a_decoder_params,
927 + )
928 + )
929 + mgr_b.submit_load(
930 + _job_metadata(
931 + job_id=201,
932 + keys=[b"a-block"],
933 + block_ids=[0],
934 + kv_params=b_decoder_params,
935 + )
936 + )
937 +
938 + # 4. Drive several poll iterations on each side. Each
939 + # get_finished_jobs() call invokes _poll_once synchronously.
940 + all_a: list[JobResult] = []
941 + all_b: list[JobResult] = []
942 + for _ in range(8):
943 + all_a.extend(list(mgr_a.get_finished_jobs()))
944 + all_b.extend(list(mgr_b.get_finished_jobs()))
945 +
946 + # Both load jobs and both store jobs must complete successfully.
947 + a_ok = {r.job_id for r in all_a if r.success}
948 + b_ok = {r.job_id for r in all_b if r.success}
949 + # A: load job 101 + store job 100
950 + assert 101 in a_ok, f"A loads succeeded: {a_ok}"
951 + assert 100 in a_ok, f"A stores succeeded: {a_ok}"
952 + # B: load job 201 + store job 200
953 + assert 201 in b_ok, f"B loads succeeded: {b_ok}"
954 + assert 200 in b_ok, f"B stores succeeded: {b_ok}"
955 +
956 +
957 +# ---------------------------------------------------------------------------
958 +# _accept_new_peers — duplicate connection rejection
959 +# ---------------------------------------------------------------------------
960 +
961 +
962 +class _RecordingConn:
963 + """Inbound connection stub that records close()/peer_id only."""
964 +
965 + def __init__(self, peer_id: str) -> None:
966 + self.peer_id = peer_id
967 + self.close_calls: int = 0
968 +
969 + def close(self) -> None:
970 + self.close_calls += 1
971 +
972 +
973 +class TestAcceptNewPeers:
974 + """A second inbound from an already-connected peer is rejected and the
975 + new conn is closed; the existing session is left untouched."""
976 +
977 + def test_duplicate_connection_is_closed_and_existing_session_untouched(self):
978 + mgr = _make_manager()
979 + peer_id = "10.0.0.1:8000"
980 + existing = _FakeSession(peer_id=peer_id, connected=True)
981 + mgr._sessions[peer_id] = existing # type: ignore[assignment]
982 +
983 + new_conn = _RecordingConn(peer_id)
984 + mgr._accept_new_peers([new_conn])
985 +
986 + # Manager swallowed the ValueError and closed the duplicate conn.
987 + assert new_conn.close_calls == 1
988 + # Existing session was NOT re-attached.
989 + assert existing.attached == []
990 + # Session map unchanged.
991 + assert mgr._sessions[peer_id] is existing
992 +
993 + def test_creates_session_for_new_peer(self):
994 + """An inbound conn from a peer with no existing session creates
995 + a fresh connected session and registers it under conn.peer_id.
996 + The prefiller has no pre-created pending session anymore — the
997 + first signal of a peer's existence is its inbound connection."""
998 +
999 + class FakeData:
1000 + block_len = 4096
1001 + base_addr = 0x1000
1002 + num_blocks = 16
1003 + config_fingerprint = ""
1004 +
1005 + def get_agent_metadata(self):
1006 + return b"meta"
1007 +
1008 + def add_remote_peer(self, *args, **kwargs):
1009 + pass
1010 +
1011 + mgr = _make_manager()
1012 + mgr._data = FakeData() # type: ignore[assignment]
1013 + peer_id = "10.0.0.1:8000"
1014 +
1015 + # Real ControlConnection-shaped fake: send/close/peer_id only.
1016 + sent: list[dict] = []
1017 +
1018 + class _Conn:
1019 + def __init__(self, pid: str) -> None:
1020 + self.peer_id = pid
1021 + self.alive = True
1022 +
1023 + def send(self, msg: dict) -> None:
1024 + sent.append(msg)
1025 +
1026 + def close(self) -> None:
1027 + self.alive = False
1028 +
1029 + mgr._accept_new_peers([_Conn(peer_id)]) # type: ignore[arg-type]
1030 +
1031 + assert peer_id in mgr._sessions
1032 + assert mgr._sessions[peer_id].connected is True
1033 + # Session sent its ConnectMsg on the new connection.
1034 + assert any(m for m in sent)
1035 +
1036 +
1037 +# ---------------------------------------------------------------------------
1038 +# _poll_once orchestration
1039 +# ---------------------------------------------------------------------------
1040 +
1041 +
1042 +class TestPollOnce:
1043 + """_poll_once must drain control, accept new peers, poll every session,
1044 + surface results, and reap dead sessions — in that order."""
1045 +
1046 + def test_orchestrates_accept_poll_and_reap(self):
1047 + mgr = _make_manager()
1048 +
1049 + # Alive session whose poll() returns one load + one store.
1050 + peer_alive = "10.0.0.2:9000"
1051 + alive = _FakeSession(
1052 + peer_id=peer_alive,
1053 + alive=True,
1054 + connected=True,
1055 + loads=[LoadResult(job_id=11, kv_request_id="req-11", success=True)],
1056 + stores=[StoreResult(job_id=22, success=True)],
1057 + )
1058 + mgr._sessions[peer_alive] = alive # type: ignore[assignment]
1059 +
1060 + # Dead session whose pending close() jobs surface as failures.
1061 + peer_dead = "10.0.0.3:9999"
1062 + dead = _FakeSession(
1063 + peer_id=peer_dead,
1064 + alive=False,
1065 + connected=True,
1066 + close_loads=[(33, "req-33")],
1067 + close_stores=[44],
1068 + )
1069 + mgr._sessions[peer_dead] = dead # type: ignore[assignment]
1070 +
1071 + class _Ctrl:
1072 + def poll(self_inner):
1073 + return []
1074 +
1075 + class _Data:
1076 + def remove_remote_peer(self_inner, pid):
1077 + pass
1078 +
1079 + mgr._control = _Ctrl() # type: ignore[assignment]
1080 + mgr._data = _Data() # type: ignore[assignment]
1081 +
1082 + mgr._poll_once()
1083 +
1084 + # Every session was polled — alive's results landed.
1085 + # Dead session was reaped — its close() failures landed.
1086 + finished = mgr._finished_jobs
1087 + ok = {(r.job_id, r.success) for r in finished}
1088 + assert (11, True) in ok # alive load result
1089 + assert (22, True) in ok # alive store result
1090 + assert (33, False) in ok # dead session's pending load
1091 + assert (44, False) in ok # dead session's pending store
1092 + assert "req-33" in mgr._failed_req_ids
1093 + assert peer_dead not in mgr._sessions
1094 + assert peer_alive in mgr._sessions
1095 +
1096 + def test_new_fetch_id_binds_and_replays_unbound_batches(self):
1097 + """When session.poll() reports a kv_request_id whose FetchMsg
1098 + arrived this tick, the manager binds it to that session and
1099 + replays every parked submit_store batch via add_stored_blocks."""
1100 + from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch
1101 +
1102 + mgr = _make_manager()
1103 + peer = "10.0.0.1:8000"
1104 + sess = _FakeSession(
1105 + peer_id=peer,
1106 + alive=True,
1107 + connected=True,
1108 + new_fetch_ids=["req-1"],
1109 + )
1110 + mgr._sessions[peer] = sess # type: ignore[assignment]
1111 + mgr._unbound_stores["req-1"] = [
1112 + _UnboundStoreBatch(job_id=5, keys=[b"k1"], block_ids=[0]),
1113 + _UnboundStoreBatch(job_id=6, keys=[b"k2"], block_ids=[1]),
1114 + ]
1115 +
1116 + class _Ctrl:
1117 + def poll(self_inner):
1118 + return []
1119 +
1120 + mgr._control = _Ctrl() # type: ignore[assignment]
1121 + mgr._poll_once()
1122 +
1123 + assert mgr._kv_to_session["req-1"] is sess
1124 + assert "req-1" not in mgr._unbound_stores
1125 + replayed = [
1126 + (kv_req_id, list(keys), job_id)
1127 + for kv_req_id, keys, _, job_id in sess.stores_added
1128 + ]
1129 + assert replayed == [("req-1", [b"k1"], 5), ("req-1", [b"k2"], 6)]
1130 +
1131 + def test_new_fetch_id_with_no_unbound_still_binds(self):
1132 + """A FetchMsg for a kv_request_id with no parked batches still
1133 + records the binding so subsequent submit_stores route fast."""
1134 + mgr = _make_manager()
1135 + peer = "10.0.0.1:8000"
1136 + sess = _FakeSession(
1137 + peer_id=peer,
1138 + alive=True,
1139 + connected=True,
1140 + new_fetch_ids=["req-fast"],
1141 + )
1142 + mgr._sessions[peer] = sess # type: ignore[assignment]
1143 +
1144 + class _Ctrl:
1145 + def poll(self_inner):
1146 + return []
1147 +
1148 + mgr._control = _Ctrl() # type: ignore[assignment]
1149 + mgr._poll_once()
1150 +
1151 + assert mgr._kv_to_session["req-fast"] is sess
1152 + assert sess.stores_added == []
1153 +
1154 + def test_failed_load_records_kv_request_id(self):
1155 + """A LoadResult(success=False) from session.poll() must add its
1156 + kv_request_id to _failed_req_ids so future lookups return MISS."""
1157 + mgr = _make_manager()
1158 + peer = "10.0.0.1:8000"
1159 + sess = _FakeSession(
1160 + peer_id=peer,
1161 + alive=True,
1162 + connected=True,
1163 + loads=[LoadResult(job_id=5, kv_request_id="req-5", success=False)],
1164 + )
1165 + mgr._sessions[peer] = sess # type: ignore[assignment]
1166 +
1167 + class _Ctrl:
1168 + def poll(self_inner):
1169 + return []
1170 +
1171 + mgr._control = _Ctrl() # type: ignore[assignment]
1172 +
1173 + mgr._poll_once()
1174 +
1175 + assert mgr._finished_jobs == [JobResult(job_id=5, success=False)]
1176 + assert "req-5" in mgr._failed_req_ids
1177 +
1178 +
1179 +# ---------------------------------------------------------------------------
1180 +# drain_jobs
1181 +# ---------------------------------------------------------------------------
1182 +
1183 +
1184 +class _DrainCtrl:
1185 + """Trivial control fake whose poll() returns an empty list."""
1186 +
1187 + def poll(self):
1188 + return []
1189 +
1190 +
1191 +class TestDrainJobs:
1192 + def test_returns_immediately_when_quiescent(self):
1193 + """No sessions and no inflight: drain_jobs returns without sleeping."""
1194 + mgr = _make_manager()
1195 + mgr._control = _DrainCtrl() # type: ignore[assignment]
1196 +
1197 + sleeps: list[float] = []
1198 + # If drain_jobs sleeps when nothing is pending, that's a regression.
1199 + import vllm.v1.kv_offload.tiering.p2p.manager as m
1200 +
1201 + original_sleep = m.time.sleep
1202 + m.time.sleep = lambda s: sleeps.append(s) # type: ignore[assignment]
1203 + try:
1204 + mgr.drain_jobs()
1205 + finally:
1206 + m.time.sleep = original_sleep # type: ignore[assignment]
1207 +
1208 + assert sleeps == []
1209 +
1210 + def test_returns_when_session_has_no_inflight_or_inbound(self):
1211 + """A session with empty _inbound and _inflight does not block drain."""
1212 + mgr = _make_manager()
1213 + mgr._control = _DrainCtrl() # type: ignore[assignment]
1214 + mgr._sessions["peer:1"] = _FakeSession(peer_id="peer:1") # type: ignore[assignment]
1215 + # Should return on the first iteration.
1216 + mgr.drain_jobs()
1217 +
1218 + def test_logs_warning_after_5s_then_completes(self, monkeypatch):
1219 + """A session that stays inflight past 5s triggers the warning, and
1220 + once it clears the loop returns."""
1221 + mgr = _make_manager()
1222 + mgr._control = _DrainCtrl() # type: ignore[assignment]
1223 + sess = _FakeSession(peer_id="peer:1")
1224 + sess._server._inflight = {1: object()} # non-empty
1225 + mgr._sessions["peer:1"] = sess # type: ignore[assignment]
1226 +
1227 + # Synthetic monotonic clock: 100.0 for the start stamp, then 106.0
1228 + # for the first elapsed-check (past the 5s warning threshold), then
1229 + # steady at 106.0 for any later checks.
1230 + clock = iter([100.0, 106.0])
1231 +
1232 + def fake_monotonic() -> float:
1233 + try:
1234 + return next(clock)
1235 + except StopIteration:
1236 + return 106.0
1237 +
1238 + monkeypatch.setattr(manager_module, "_DRAIN_SLEEP_S", 0.0)
1239 +
1240 + # Spy on the warning logger directly — vllm's logger does not
1241 + # propagate to root, so caplog can't see it.
1242 + warnings: list[str] = []
1243 +
1244 + def record_warning(msg, *args, **_kwargs):
1245 + warnings.append(msg % args if args else msg)
1246 +
1247 + monkeypatch.setattr(manager_module.logger, "warning", record_warning)
1248 +
1249 + # Clear inflight after the first sleep so drain can exit on the
1250 + # next iteration's `pending` check.
1251 + n_sleeps = 0
1252 +
1253 + def clearing_sleep(_s):
1254 + nonlocal n_sleeps
1255 + n_sleeps += 1
1256 + sess._server._inflight = {}
1257 +
1258 + # Patch via a fake module on `manager_module.time` so we do not mutate
1259 + # the global `time` module — other code in the process (e.g. the
1260 + # buildkite test collector's pytest_runtest_logreport hook) calls
1261 + # time.monotonic() before monkeypatch teardown.
1262 + fake_time = SimpleNamespace(monotonic=fake_monotonic, sleep=clearing_sleep)
1263 + monkeypatch.setattr(manager_module, "time", fake_time)
1264 +
1265 + mgr.drain_jobs()
1266 +
1267 + assert any("still draining after 5s" in w for w in warnings), warnings
1268 +
1269 +
1270 +# ---------------------------------------------------------------------------
1271 +# on_schedule_end
1272 +# ---------------------------------------------------------------------------
1273 +
1274 +
1275 +class TestOnScheduleEnd:
1276 + def test_is_noop(self):
1277 + """on_schedule_end is a documented no-op; just confirm it doesn't
1278 + raise and doesn't mutate state."""
1279 + mgr = _make_manager()
1280 + before_sessions = dict(mgr._sessions)
1281 + before_jobs = list(mgr._finished_jobs)
1282 + assert mgr.on_schedule_end() is None
1283 + assert mgr._sessions == before_sessions
1284 + assert mgr._finished_jobs == before_jobs
1285 +
1286 +
1287 +# ---------------------------------------------------------------------------
1288 +# Connection death mid-transfer (real P2PSession via paired managers)
1289 +# ---------------------------------------------------------------------------
1290 +
1291 +
1292 +class TestConnectionDeathMidTransfer:
1293 + """When a peer's control connection dies while a load is in flight,
1294 + the load surfaces as failed and its kv_request_id lands in
1295 + _failed_req_ids so future lookups route to local prefill. The
1296 + prefiller-side store no longer travels through the session at store
1297 + time (it's parked in _unbound_stores keyed by kv_request_id), so its
1298 + cleanup on connection death is via on_request_finished or the
1299 + unbound-store timeout — covered separately below."""
1300 +
1301 + def test_dead_connection_with_pending_work_surfaces_failures(self):
1302 + mgr_a, mgr_b = _build_paired_managers()
1303 +
1304 + a_decoder_params = {
1305 + "prefill": {
1306 + "kv_request_id": "req-load",
1307 + "remote_host": "B",
1308 + "remote_port": 2,
1309 + },
1310 + }
1311 + a_prefiller_params = {"decode": {"kv_request_id": "req-store"}}
1312 +
1313 + # Open the outbound session A->B and submit one load + one store.
1314 + mgr_a.on_new_request(_req_context(a_decoder_params))
1315 + mgr_a.submit_store(
1316 + _job_metadata(
1317 + job_id=900,
1318 + keys=[b"a-block"],
1319 + block_ids=[0],
1320 + kv_params=a_prefiller_params,
1321 + )
1322 + )
1323 + mgr_a.submit_load(
1324 + _job_metadata(
1325 + job_id=901,
1326 + keys=[b"b-block"],
1327 + block_ids=[0],
1328 + kv_params=a_decoder_params,
1329 + )
1330 + )
1331 +
1332 + # Drain anything the loopback can deliver synchronously, but stop
1333 + # before the remote side has had time to complete the transfers.
1334 + list(mgr_a.get_finished_jobs())
1335 +
1336 + # Sanity: store 900 is parked in unbound_stores, not in any
1337 + # session — the producer no longer learns the peer at store time.
1338 + assert "req-store" in mgr_a._unbound_stores
1339 +
1340 + # Kill the control connection out from under the session.
1341 + peer_id = "B:2"
1342 + sess = mgr_a._sessions[peer_id]
1343 + assert sess._conn is not None
1344 + sess._conn.mark_dead()
1345 +
1346 + # Reap surfaces the load (which lived inside the session) as
1347 + # failed. The store 900 is not session-scoped — it survives the
1348 + # session reap and waits for the unbound-store timeout.
1349 + results: list[JobResult] = []
1350 + for _ in range(3):
1351 + results.extend(list(mgr_a.get_finished_jobs()))
1352 +
1353 + outcomes = {(r.job_id, r.success) for r in results}
1354 + assert (901, False) in outcomes, f"load should fail: {outcomes}"
1355 + assert (900, False) not in outcomes, f"store should still be parked: {outcomes}"
1356 + assert "req-load" in mgr_a._failed_req_ids
1357 + # Session removed.
1358 + assert peer_id not in mgr_a._sessions
1359 + # Store batch is still parked.
1360 + assert "req-store" in mgr_a._unbound_stores
1361 +
1362 + # The engine signals the producer's request is done. That is a
1363 + # no-op for the parked batch — `on_request_finished` does not
1364 + # evict unbound stores; only `_reap_unbound_stores` does, after
1365 + # the unbound-store timeout. Job 900 stays unfinished here.
1366 + mgr_a.on_request_finished(_req_context(a_prefiller_params))
1367 + finishes = {(r.job_id, r.success) for r in mgr_a._finished_jobs}
1368 + assert (900, False) not in finishes
1369 + assert (900, True) not in finishes
1370 + assert "req-store" in mgr_a._unbound_stores
tests/v1/kv_offload/tiering/p2p/test_sessions.py new
+1645
@@ -0,0 +1,1645 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for the unified bidirectional P2PSession.
4 +
5 +A P2PSession owns a single ControlConnection and dispatches every message
6 +type — both client-role (FetchMsg / TransferDoneMsg / AbortAck) and
7 +server-role (FetchMsg / TransferDoneMsg / AbortAck from the peer's
8 +perspective). These tests exercise both flows independently and the
9 +bidirectional case where one session simultaneously serves a fetch and
10 +completes its own load.
11 +"""
12 +
13 +from __future__ import annotations
14 +
15 +import time
16 +
17 +import pytest
18 +
19 +from vllm.v1.kv_offload.tiering.p2p.session import (
20 + LoadResult,
21 + P2PSession,
22 + StoreResult,
23 +)
24 +from vllm.v1.kv_offload.tiering.p2p.session.client import (
25 + _ABORT_ACK_TIMEOUT_S,
26 + _LOAD_TIMEOUT_S,
27 +)
28 +from vllm.v1.kv_offload.tiering.p2p.session.protocol import (
29 + TYPE_KEY,
30 + AbortAckMsg,
31 + AbortFetchMsg,
32 + ConnectAckMsg,
33 + ConnectMsg,
34 + DisconnectMsg,
35 + FetchMsg,
36 + TransferDoneMsg,
37 +)
38 +from vllm.v1.kv_offload.tiering.p2p.session.server import (
39 + _CANCEL_DRAIN_TIMEOUT_S,
40 + _InflightXfer,
41 +)
42 +from vllm.v1.kv_offload.tiering.p2p.session.session import (
43 + _MAX_CONSECUTIVE_DISPATCH_ERRORS,
44 +)
45 +
46 +# ---------------------------------------------------------------------------
47 +# Helpers
48 +# ---------------------------------------------------------------------------
49 +
50 +
51 +class FakeDataTransport:
52 + """Minimal fake DataTransport for testing sessions."""
53 +
54 + def __init__(
55 + self,
56 + base_addr: int = 0x1000,
57 + num_blocks: int = 16,
58 + block_len: int = 4096,
59 + config_fingerprint: str = "",
60 + ) -> None:
61 + self._base_addr = base_addr
62 + self._num_blocks = num_blocks
63 + self._block_len = block_len
64 + self._config_fingerprint = config_fingerprint
65 + self._remote_peers: dict[str, dict] = {}
66 + self._transfers: dict[int, tuple] = {}
67 + self._next_id = 0
68 + self._poll_done: list[int] = []
69 + self._poll_failed: list[int] = []
70 + self._cancel_still_inflight: set[int] = set()
71 + self._cancel_calls: list[tuple[list[int], str]] = []
72 +
73 + @property
74 + def base_addr(self) -> int:
75 + return self._base_addr
76 +
77 + @property
78 + def num_blocks(self) -> int:
79 + return self._num_blocks
80 +
81 + @property
82 + def block_len(self) -> int:
83 + return self._block_len
84 +
85 + @property
86 + def config_fingerprint(self) -> str:
87 + return self._config_fingerprint
88 +
89 + def get_agent_metadata(self) -> bytes:
90 + return b"fake-metadata"
91 +
92 + def add_remote_peer(
93 + self, peer_id, agent_metadata, base_addr, num_blocks, block_len
94 + ) -> None:
95 + self._remote_peers[peer_id] = {
96 + ConnectMsg.AGENT_METADATA: agent_metadata,
97 + ConnectMsg.BASE_ADDR: base_addr,
98 + "num_blocks": num_blocks,
99 + "block_len": block_len,
100 + }
101 +
102 + def remove_remote_peer(self, peer_id: str) -> None:
103 + self._remote_peers.pop(peer_id, None)
104 +
105 + def write_blocks(self, peer_id, local_idxs, remote_idxs) -> int | None:
106 + if peer_id not in self._remote_peers:
107 + return None
108 + tid = self._next_id
109 + self._next_id += 1
110 + self._transfers[tid] = (peer_id, local_idxs, remote_idxs)
111 + return tid
112 +
113 + def poll(self):
114 + from vllm.v1.kv_offload.tiering.p2p.data.base import PollResult
115 +
116 + result = PollResult(done=list(self._poll_done), failed=list(self._poll_failed))
117 + self._poll_done.clear()
118 + self._poll_failed.clear()
119 + return result
120 +
121 + def cancel(self, transfer_ids, mode: str = "immediate") -> list[int]:
122 + ids = list(transfer_ids)
123 + self._cancel_calls.append((ids, mode))
124 + if mode == "wait":
125 + still: list[int] = []
126 + for tid in ids:
127 + if tid in self._cancel_still_inflight:
128 + still.append(tid)
129 + else:
130 + self._transfers.pop(tid, None)
131 + return still
132 + for tid in ids:
133 + self._transfers.pop(tid, None)
134 + self._cancel_still_inflight.discard(tid)
135 + return []
136 +
137 + def close(self) -> None:
138 + pass
139 +
140 +
141 +class FakeConnection:
142 + """Fake ControlConnection that captures sent messages."""
143 +
144 + def __init__(self, peer_id: str = "peer:8000") -> None:
145 + self.peer_id = peer_id
146 + self._inbox: list[dict] = []
147 + self._sent: list[dict] = []
148 + self._closed = False
149 +
150 + @property
151 + def alive(self) -> bool:
152 + return not self._closed
153 +
154 + def send(self, msg: dict) -> None:
155 + self._sent.append(msg)
156 +
157 + def recv(self) -> list[dict]:
158 + msgs = self._inbox
159 + self._inbox = []
160 + return msgs
161 +
162 + def enqueue(self, msg: dict) -> None:
163 + self._inbox.append(msg)
164 +
165 + def mark_dead(self) -> None:
166 + self._closed = True
167 +
168 + def close(self) -> None:
169 + self._closed = True
170 +
171 +
172 +def _peer_connect_msg(
173 + peer_id: str = "peer:8000",
174 + block_len: int = 4096,
175 + fingerprint: str | None = None,
176 +) -> dict:
177 + """Build a ConnectMsg as if the peer sent it."""
178 + msg = {
179 + TYPE_KEY: ConnectMsg.TYPE,
180 + ConnectMsg.PEER_ID: peer_id,
181 + ConnectMsg.AGENT_METADATA: b"peer-metadata",
182 + ConnectMsg.BASE_ADDR: 0x2000,
183 + ConnectMsg.NUM_BLOCKS: 16,
184 + ConnectMsg.BLOCK_LEN: block_len,
185 + }
186 + if fingerprint is not None:
187 + msg[ConnectMsg.CONFIG_FINGERPRINT] = fingerprint
188 + return msg
189 +
190 +
191 +def _make_session(
192 + conn: FakeConnection | None = None,
193 + transport: FakeDataTransport | None = None,
194 + peer_id: str = "peer:8000",
195 + local_id: str = "local:9000",
196 +) -> tuple[P2PSession, FakeConnection, FakeDataTransport]:
197 + if conn is None:
198 + conn = FakeConnection(peer_id=peer_id)
199 + if transport is None:
200 + transport = FakeDataTransport()
201 + session = P2PSession(
202 + peer_id=peer_id,
203 + local_id=local_id,
204 + transport=transport, # type: ignore[arg-type]
205 + local_block_len=transport.block_len,
206 + conn=conn, # type: ignore[arg-type]
207 + )
208 + return session, conn, transport
209 +
210 +
211 +def _activate(
212 + session: P2PSession, conn: FakeConnection, peer_id: str = "peer:8000"
213 +) -> None:
214 + """Drive the handshake: peer sends ConnectMsg + ConnectAckMsg."""
215 + conn.enqueue(_peer_connect_msg(peer_id=peer_id))
216 + conn.enqueue({TYPE_KEY: ConnectAckMsg.TYPE, ConnectAckMsg.PEER_ID: peer_id})
217 + session.poll()
218 +
219 +
220 +# ---------------------------------------------------------------------------
221 +# Connect / handshake
222 +# ---------------------------------------------------------------------------
223 +
224 +
225 +class TestConnectHandshake:
226 + def test_connect_msg_sent_on_creation(self):
227 + """Session sends its own ConnectMsg on connection."""
228 + session, conn, _ = _make_session()
229 + assert len(conn._sent) == 1
230 + msg = conn._sent[0]
231 + assert msg[TYPE_KEY] == ConnectMsg.TYPE
232 + assert msg[ConnectMsg.PEER_ID] == "local:9000"
233 + assert msg[ConnectMsg.NUM_BLOCKS] == 16
234 + assert msg[ConnectMsg.BLOCK_LEN] == 4096
235 + assert ConnectMsg.AGENT_METADATA in msg
236 +
237 + def test_peer_connect_triggers_add_remote_and_ack(self):
238 + """Receiving ConnectMsg registers peer and replies with ConnectAck."""
239 + session, conn, transport = _make_session()
240 + conn.enqueue(_peer_connect_msg())
241 + session.poll()
242 + assert "peer:8000" in transport._remote_peers
243 + ack = next(m for m in conn._sent if m[TYPE_KEY] == ConnectAckMsg.TYPE)
244 + assert ack[ConnectAckMsg.PEER_ID] == "local:9000"
245 +
246 + def test_connect_ack_makes_session_ready(self):
247 + """Session.ready becomes True after ConnectAckMsg."""
248 + session, conn, _ = _make_session()
249 + assert not session.ready
250 + conn.enqueue({TYPE_KEY: ConnectAckMsg.TYPE, ConnectAckMsg.PEER_ID: "peer:8000"})
251 + session.poll()
252 + assert session.ready
253 +
254 + def test_messages_queued_before_ack_flush_on_ack(self):
255 + """Outgoing messages sent before ConnectAck are flushed after."""
256 + session, conn, _ = _make_session()
257 + session.request_blocks(
258 + job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0]
259 + )
260 + # Before ack: only our ConnectMsg was sent.
261 + assert len(conn._sent) == 1
262 + assert conn._sent[0][TYPE_KEY] == ConnectMsg.TYPE
263 + # Ack arrives.
264 + conn.enqueue({TYPE_KEY: ConnectAckMsg.TYPE, ConnectAckMsg.PEER_ID: "peer:8000"})
265 + session.poll()
266 + # Queued fetch is now sent.
267 + assert any(m[TYPE_KEY] == FetchMsg.TYPE for m in conn._sent)
268 +
269 + def test_block_len_mismatch_marks_dead(self):
270 + """Mismatched block_len rejects peer and marks connection dead."""
271 + session, conn, transport = _make_session()
272 + conn.enqueue(_peer_connect_msg(block_len=8192)) # mismatch
273 + session.poll()
274 + assert "peer:8000" not in transport._remote_peers
275 + assert not session.alive
276 +
277 + def test_config_fingerprint_mismatch_marks_dead(self):
278 + """Mismatched config fingerprint rejects peer."""
279 + transport = FakeDataTransport(config_fingerprint="abc123")
280 + session, conn, _ = _make_session(transport=transport)
281 + conn.enqueue(_peer_connect_msg(fingerprint="different"))
282 + session.poll()
283 + assert "peer:8000" not in transport._remote_peers
284 + assert not session.alive
285 +
286 + def test_config_fingerprint_match_succeeds(self):
287 + """Matching fingerprints register the peer."""
288 + transport = FakeDataTransport(config_fingerprint="same_fp")
289 + session, conn, _ = _make_session(transport=transport)
290 + conn.enqueue(_peer_connect_msg(fingerprint="same_fp"))
291 + session.poll()
292 + assert "peer:8000" in transport._remote_peers
293 + assert session.alive
294 +
295 + def test_missing_fingerprint_allowed(self):
296 + """Missing fingerprint on either side is allowed."""
297 + transport = FakeDataTransport(config_fingerprint="abc123")
298 + session, conn, _ = _make_session(transport=transport)
299 + conn.enqueue(_peer_connect_msg()) # no fingerprint
300 + session.poll()
301 + assert "peer:8000" in transport._remote_peers
302 +
303 +
304 +# ---------------------------------------------------------------------------
305 +# Client-role flows
306 +# ---------------------------------------------------------------------------
307 +
308 +
309 +class TestClientFlows:
310 + def test_request_blocks_sends_fetch(self):
311 + session, conn, _ = _make_session()
312 + _activate(session, conn)
313 + session.request_blocks(
314 + job_id=1, kv_request_id="req-1", keys=[b"k1", b"k2"], block_ids=[0, 1]
315 + )
316 + lookup = conn._sent[-1]
317 + assert lookup[TYPE_KEY] == FetchMsg.TYPE
318 + assert lookup[FetchMsg.KV_REQUEST_ID] == "req-1"
319 + assert lookup[FetchMsg.BLOCK_HASHES] == [b"k1", b"k2"]
320 + assert lookup[FetchMsg.BLOCK_INDEXES] == [0, 1]
321 +
322 + def test_transfer_done_success(self):
323 + session, conn, _ = _make_session()
324 + _activate(session, conn)
325 + session.request_blocks(
326 + job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0]
327 + )
328 + conn.enqueue(
329 + {
330 + TYPE_KEY: TransferDoneMsg.TYPE,
331 + TransferDoneMsg.KV_REQUEST_ID: "req-1",
332 + TransferDoneMsg.SUCCESS: True,
333 + }
334 + )
335 + loads = session.poll().loads
336 + assert loads == [LoadResult(job_id=1, kv_request_id="req-1", success=True)]
337 +
338 + def test_transfer_done_failure(self):
339 + session, conn, _ = _make_session()
340 + _activate(session, conn)
341 + session.request_blocks(
342 + job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0]
343 + )
344 + conn.enqueue(
345 + {
346 + TYPE_KEY: TransferDoneMsg.TYPE,
347 + TransferDoneMsg.KV_REQUEST_ID: "req-1",
348 + TransferDoneMsg.SUCCESS: False,
349 + }
350 + )
351 + loads = session.poll().loads
352 + assert loads == [LoadResult(job_id=1, kv_request_id="req-1", success=False)]
353 +
354 + def test_finish_request_sends_abort(self):
355 + session, conn, _ = _make_session()
356 + _activate(session, conn)
357 + session.request_blocks(
358 + job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0]
359 + )
360 + session.finish_request("req-1")
361 + abort = conn._sent[-1]
362 + assert abort[TYPE_KEY] == AbortFetchMsg.TYPE
363 + assert abort[AbortFetchMsg.KV_REQUEST_ID] == "req-1"
364 +
365 + def test_load_timeout_sends_abort(self):
366 + session, conn, _ = _make_session()
367 + _activate(session, conn)
368 + session.request_blocks(
369 + job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0]
370 + )
371 + session._client._inbound["req-1"].submitted_at = time.monotonic() - 60.0
372 + session.poll()
373 + abort = conn._sent[-1]
374 + assert abort[TYPE_KEY] == AbortFetchMsg.TYPE
375 +
376 + def test_load_abort_ack_timeout_surfaces_failure(self):
377 + """After load timeout sends AbortFetch, if no AbortAck arrives within
378 + _ABORT_ACK_TIMEOUT_S the request is surfaced as failed and removed
379 + from _inbound — the engine cannot wait forever on a peer that won't
380 + ack.
381 + """
382 + session, conn, _ = _make_session()
383 + _activate(session, conn)
384 + session.request_blocks(
385 + job_id=7, kv_request_id="req-7", keys=[b"k"], block_ids=[0]
386 + )
387 + # 1) Trip the load timeout to send AbortFetch and stamp aborted_at.
388 + session._client._inbound["req-7"].submitted_at = (
389 + time.monotonic() - _LOAD_TIMEOUT_S - 1.0
390 + )
391 + loads = session.poll().loads
392 + assert loads == []
393 + assert any(
394 + m.get(TYPE_KEY) == AbortFetchMsg.TYPE
395 + and m[AbortFetchMsg.KV_REQUEST_ID] == "req-7"
396 + for m in conn._sent
397 + )
398 + assert session._client._inbound["req-7"].aborted_at is not None
399 +
400 + # 2) Now backdate aborted_at past the abort-ack timeout. No ack ever
401 + # arrived from the peer.
402 + session._client._inbound["req-7"].aborted_at = (
403 + time.monotonic() - _ABORT_ACK_TIMEOUT_S - 1.0
404 + )
405 + loads = session.poll().loads
406 + assert loads == [LoadResult(job_id=7, kv_request_id="req-7", success=False)]
407 + assert "req-7" not in session._client._inbound
408 +
409 + def test_load_abort_ack_clears_request(self):
410 + """After load timeout sends AbortFetch, an arriving AbortAckMsg from
411 + the peer surfaces the failure cleanly and removes the request from
412 + _inbound — covers the on_abort_ack arrival path."""
413 + session, conn, _ = _make_session()
414 + _activate(session, conn)
415 + session.request_blocks(
416 + job_id=8, kv_request_id="req-8", keys=[b"k"], block_ids=[0]
417 + )
418 + session._client._inbound["req-8"].submitted_at = (
419 + time.monotonic() - _LOAD_TIMEOUT_S - 1.0
420 + )
421 + # First poll: AbortFetch goes out.
422 + session.poll()
423 + assert session._client._inbound["req-8"].aborted_at is not None
424 +
425 + # Peer acks the abort.
426 + conn.enqueue(
427 + {
428 + TYPE_KEY: AbortAckMsg.TYPE,
429 + AbortAckMsg.KV_REQUEST_ID: "req-8",
430 + }
431 + )
432 + loads = session.poll().loads
433 + assert loads == [LoadResult(job_id=8, kv_request_id="req-8", success=False)]
434 + assert "req-8" not in session._client._inbound
435 +
436 +
437 +# ---------------------------------------------------------------------------
438 +# Server-role flows
439 +# ---------------------------------------------------------------------------
440 +
441 +
442 +class TestServerFlows:
443 + def test_store_then_fetch_matches(self):
444 + """Blocks stored before fetch demand are matched on demand arrival."""
445 + session, conn, transport = _make_session()
446 + _activate(session, conn)
447 + session.add_stored_blocks("req-1", [b"k1", b"k2"], [0, 1], job_id=1)
448 + conn.enqueue(
449 + {
450 + TYPE_KEY: FetchMsg.TYPE,
451 + FetchMsg.KV_REQUEST_ID: "req-1",
452 + FetchMsg.BLOCK_HASHES: [b"k1", b"k2"],
453 + FetchMsg.BLOCK_INDEXES: [10, 11],
454 + }
455 + )
456 + session.poll()
457 + assert len(transport._transfers) == 1
458 + _, (peer, local, remote) = next(iter(transport._transfers.items()))
459 + assert local == [0, 1]
460 + assert remote == [10, 11]
461 +
462 + def test_fetch_then_store_matches(self):
463 + """Fetch demand registered before store; store fulfills it."""
464 + session, conn, transport = _make_session()
465 + _activate(session, conn)
466 + conn.enqueue(
467 + {
468 + TYPE_KEY: FetchMsg.TYPE,
469 + FetchMsg.KV_REQUEST_ID: "req-1",
470 + FetchMsg.BLOCK_HASHES: [b"k1"],
471 + FetchMsg.BLOCK_INDEXES: [5],
472 + }
473 + )
474 + session.poll()
475 + assert len(transport._transfers) == 0
476 + session.add_stored_blocks("req-1", [b"k1"], [3], job_id=1)
477 + assert len(transport._transfers) == 1
478 +
479 + def test_transfer_completion_emits_store_result_and_done(self):
480 + """Completed transfer reports StoreResult and sends TransferDoneMsg."""
481 + session, conn, transport = _make_session()
482 + _activate(session, conn)
483 + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)
484 + conn.enqueue(
485 + {
486 + TYPE_KEY: FetchMsg.TYPE,
487 + FetchMsg.KV_REQUEST_ID: "req-1",
488 + FetchMsg.BLOCK_HASHES: [b"k1"],
489 + FetchMsg.BLOCK_INDEXES: [5],
490 + }
491 + )
492 + session.poll()
493 + tid = next(iter(transport._transfers))
494 + transport._poll_done.append(tid)
495 + stores = session.poll().stores
496 + assert StoreResult(job_id=1, success=True) in stores
497 + assert any(m[TYPE_KEY] == TransferDoneMsg.TYPE for m in conn._sent)
498 +
499 + def test_abort_fetch_replies_with_ack(self):
500 + session, conn, _ = _make_session()
501 + _activate(session, conn)
502 + conn.enqueue(
503 + {
504 + TYPE_KEY: AbortFetchMsg.TYPE,
505 + AbortFetchMsg.KV_REQUEST_ID: "req-1",
506 + }
507 + )
508 + session.poll()
509 + ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE)
510 + assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1"
511 + assert "req-1" not in session._server._pending_aborts
512 +
513 + def test_abort_fetch_defers_ack_when_cancel_pending(self):
514 + """If cancel(mode='wait') reports still-inflight tids, the ack is
515 + deferred and the abort is parked in _pending_aborts."""
516 + session, conn, transport = _make_session()
517 + _activate(session, conn)
518 + # Seed an inflight transfer for req-1 that the transport pretends
519 + # cannot be canceled yet.
520 + tid = 42
521 + session._server._inflight_add(
522 + tid,
523 + _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}),
524 + )
525 + transport._cancel_still_inflight.add(tid)
526 +
527 + conn.enqueue(
528 + {
529 + TYPE_KEY: AbortFetchMsg.TYPE,
530 + AbortFetchMsg.KV_REQUEST_ID: "req-1",
531 + }
532 + )
533 + session.poll()
534 +
535 + assert not any(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent)
536 + assert "req-1" in session._server._pending_aborts
537 + # First attempt happens inside _on_abort_fetch; the per-tick
538 + # drain runs again at the end of poll() — both are wait-mode.
539 + assert all(mode == "wait" for _, mode in transport._cancel_calls)
540 + assert tid in session._server._inflight # still tracked
541 +
542 + def test_abort_fetch_acks_after_drain(self):
543 + """Once the transport reports the tid as DONE the parked abort
544 + completes and AbortAckMsg is sent."""
545 + session, conn, transport = _make_session()
546 + _activate(session, conn)
547 + tid = 42
548 + session._server._inflight_add(
549 + tid,
550 + _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}),
551 + )
552 + transport._cancel_still_inflight.add(tid)
553 +
554 + conn.enqueue(
555 + {
556 + TYPE_KEY: AbortFetchMsg.TYPE,
557 + AbortFetchMsg.KV_REQUEST_ID: "req-1",
558 + }
559 + )
560 + session.poll()
561 + assert "req-1" in session._server._pending_aborts
562 +
563 + # Backend finishes draining: transport.poll() will return tid as
564 + # DONE, and the next cancel(mode='wait') call sees it's gone.
565 + transport._cancel_still_inflight.discard(tid)
566 + transport._poll_done.append(tid)
567 +
568 + session.poll()
569 +
570 + ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE)
571 + assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1"
572 + assert "req-1" not in session._server._pending_aborts
573 + assert tid not in session._server._inflight
574 +
575 + def test_abort_fetch_force_cancels_after_timeout(self):
576 + """If wait-mode never drains, the deadline forces immediate
577 + cancel and an ack is still sent."""
578 + session, conn, transport = _make_session()
579 + _activate(session, conn)
580 + tid = 42
581 + session._server._inflight_add(
582 + tid,
583 + _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}),
584 + )
585 + transport._cancel_still_inflight.add(tid)
586 +
587 + conn.enqueue(
588 + {
589 + TYPE_KEY: AbortFetchMsg.TYPE,
590 + AbortFetchMsg.KV_REQUEST_ID: "req-1",
591 + }
592 + )
593 + session.poll()
594 + assert "req-1" in session._server._pending_aborts
595 + # Backdate past the drain deadline.
596 + session._server._pending_aborts["req-1"] = (
597 + time.monotonic() - _CANCEL_DRAIN_TIMEOUT_S - 1.0
598 + )
599 + # Even if the transport still claims it can't cancel, the
600 + # session must force-pop and ack.
601 + transport._cancel_calls.clear()
602 +
603 + session.poll()
604 +
605 + ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE)
606 + assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1"
607 + assert "req-1" not in session._server._pending_aborts
608 + assert tid not in session._server._inflight
609 + assert ([tid], "immediate") in transport._cancel_calls
610 +
611 + def test_abort_fetch_idempotent_while_draining(self):
612 + """Receiving AbortFetchMsg twice for the same kv_request_id
613 + keeps a single pending entry and produces a single ack."""
614 + session, conn, transport = _make_session()
615 + _activate(session, conn)
616 + tid = 42
617 + session._server._inflight_add(
618 + tid,
619 + _InflightXfer(kv_request_id="req-1", block_count=1, job_ids={1}),
620 + )
621 + transport._cancel_still_inflight.add(tid)
622 +
623 + conn.enqueue(
624 + {
625 + TYPE_KEY: AbortFetchMsg.TYPE,
626 + AbortFetchMsg.KV_REQUEST_ID: "req-1",
627 + }
628 + )
629 + session.poll()
630 + first_started_at = session._server._pending_aborts["req-1"]
631 +
632 + # Second AbortFetchMsg for the same kv_request_id while still
633 + # draining must not reset the deadline.
634 + conn.enqueue(
635 + {
636 + TYPE_KEY: AbortFetchMsg.TYPE,
637 + AbortFetchMsg.KV_REQUEST_ID: "req-1",
638 + }
639 + )
640 + session.poll()
641 + assert session._server._pending_aborts["req-1"] == first_started_at
642 +
643 + # Now let the drain succeed and confirm exactly one ack ever.
644 + transport._cancel_still_inflight.discard(tid)
645 + transport._poll_done.append(tid)
646 + session.poll()
647 +
648 + acks = [m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE]
649 + assert len(acks) == 1
650 + assert acks[0][AbortAckMsg.KV_REQUEST_ID] == "req-1"
651 +
652 + def test_store_timeout(self):
653 + session, conn, _ = _make_session()
654 + _activate(session, conn)
655 + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)
656 + # Backdate.
657 + session._server._store_jobs[1] = time.monotonic() - 60.0
658 + stores = session.poll().stores
659 + assert StoreResult(job_id=1, success=False) in stores
660 +
661 + def test_store_timeout_then_late_completion_no_duplicate(self):
662 + """A job timed out by _timeout_pending_store_jobs must not also
663 + emit a contradictory StoreResult(success=True) when the transport
664 + later reports the same transfer as done."""
665 + session, conn, transport = _make_session()
666 + _activate(session, conn)
667 + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)
668 + conn.enqueue(
669 + {
670 + TYPE_KEY: FetchMsg.TYPE,
671 + FetchMsg.KV_REQUEST_ID: "req-1",
672 + FetchMsg.BLOCK_HASHES: [b"k1"],
673 + FetchMsg.BLOCK_INDEXES: [5],
674 + }
675 + )
676 + session.poll()
677 + tid = next(iter(transport._transfers))
678 +
679 + # Backdate the store job so the next poll times it out.
680 + session._server._store_jobs[1] = time.monotonic() - 60.0
681 + stores = session.poll().stores
682 + assert StoreResult(job_id=1, success=False) in stores
683 + assert StoreResult(job_id=1, success=True) not in stores
684 +
685 + # Transport later reports the same transfer as done — must not
686 + # emit a second (contradictory) StoreResult for job_id=1.
687 + transport._poll_done.append(tid)
688 + stores = session.poll().stores
689 + assert all(s.job_id != 1 for s in stores), (
690 + f"unexpected duplicate StoreResult after timeout: {stores}"
691 + )
692 +
693 + def test_store_timeout_then_late_failure_no_duplicate(self):
694 + """Symmetric guard: a timed-out job must not also emit a second
695 + StoreResult(success=False) when the transport later reports the
696 + same transfer as failed."""
697 + session, conn, transport = _make_session()
698 + _activate(session, conn)
699 + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)
700 + conn.enqueue(
701 + {
702 + TYPE_KEY: FetchMsg.TYPE,
703 + FetchMsg.KV_REQUEST_ID: "req-1",
704 + FetchMsg.BLOCK_HASHES: [b"k1"],
705 + FetchMsg.BLOCK_INDEXES: [5],
706 + }
707 + )
708 + session.poll()
709 + tid = next(iter(transport._transfers))
710 +
711 + session._server._store_jobs[1] = time.monotonic() - 60.0
712 + stores = session.poll().stores
713 + assert [s for s in stores if s.job_id == 1] == [
714 + StoreResult(job_id=1, success=False)
715 + ]
716 +
717 + transport._poll_failed.append(tid)
718 + stores = session.poll().stores
719 + assert all(s.job_id != 1 for s in stores), (
720 + f"unexpected duplicate StoreResult after timeout: {stores}"
721 + )
722 +
723 +
724 +# ---------------------------------------------------------------------------
725 +# finish_request server-role early-fail flow
726 +# ---------------------------------------------------------------------------
727 +
728 +
729 +class TestFinishRequestServerSide:
730 + def _last_transfer_done(self, conn: FakeConnection) -> dict | None:
731 + for msg in reversed(conn._sent):
732 + if msg[TYPE_KEY] == TransferDoneMsg.TYPE:
733 + return msg
734 + return None
735 +
736 + def test_no_inflight_unmatched_demand_sends_failure(self):
737 + """finish_request with unmatched demand and no inflight ->
738 + immediate TransferDoneMsg(success=False); _outbound cleared."""
739 + session, conn, _ = _make_session()
740 + _activate(session, conn)
741 + # Decoder demanded a block we never stored.
742 + conn.enqueue(
743 + {
744 + TYPE_KEY: FetchMsg.TYPE,
745 + FetchMsg.KV_REQUEST_ID: "req-1",
746 + FetchMsg.BLOCK_HASHES: [b"k1"],
747 + FetchMsg.BLOCK_INDEXES: [5],
748 + }
749 + )
750 + session.poll()
751 + assert "req-1" in session._server._outbound
752 +
753 + session.finish_request("req-1")
754 +
755 + msg = self._last_transfer_done(conn)
756 + assert msg is not None
757 + assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1"
758 + assert msg[TransferDoneMsg.SUCCESS] is False
759 + assert "req-1" not in session._server._outbound
760 +
761 + def test_with_inflight_defers_then_fires_on_last_transfer(self):
762 + """finish_request with inflight defers; last transfer fires the
763 + early-fail message and clears _outbound."""
764 + session, conn, transport = _make_session()
765 + _activate(session, conn)
766 + # Demand 2 blocks; we store 1 (kicks one inflight transfer).
767 + conn.enqueue(
768 + {
769 + TYPE_KEY: FetchMsg.TYPE,
770 + FetchMsg.KV_REQUEST_ID: "req-1",
771 + FetchMsg.BLOCK_HASHES: [b"k1", b"k2"],
772 + FetchMsg.BLOCK_INDEXES: [10, 11],
773 + }
774 + )
775 + session.poll()
776 + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)
777 + assert len(transport._transfers) == 1
778 +
779 + # finish_request while inflight: no early-fail yet.
780 + before = len(conn._sent)
781 + session.finish_request("req-1")
782 + assert len(conn._sent) == before
783 + assert "req-1" in session._server._outbound
784 + assert session._server._outbound["req-1"].finishing
785 +
786 + # Last inflight settles -> early-fail fires.
787 + tid = next(iter(transport._transfers))
788 + transport._poll_done.append(tid)
789 + session.poll()
790 +
791 + msg = self._last_transfer_done(conn)
792 + assert msg is not None
793 + assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1"
794 + assert msg[TransferDoneMsg.SUCCESS] is False
795 + assert "req-1" not in session._server._outbound
796 +
797 + def test_full_demand_satisfied_still_sends_success(self):
798 + """finish_request must not override a fully-satisfied transfer:
799 + when remaining hits 0, success=True still fires."""
800 + session, conn, transport = _make_session()
801 + _activate(session, conn)
802 + conn.enqueue(
803 + {
804 + TYPE_KEY: FetchMsg.TYPE,
805 + FetchMsg.KV_REQUEST_ID: "req-1",
806 + FetchMsg.BLOCK_HASHES: [b"k1"],
807 + FetchMsg.BLOCK_INDEXES: [10],
808 + }
809 + )
810 + session.poll()
811 + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)
812 + # Mark finishing (e.g., on_request_finished racing with the last
813 + # store) — but all demand is satisfied.
814 + session.finish_request("req-1")
815 +
816 + tid = next(iter(transport._transfers))
817 + transport._poll_done.append(tid)
818 + session.poll()
819 +
820 + msg = next(m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE)
821 + assert msg[TransferDoneMsg.SUCCESS] is True
822 +
823 + def test_prefiller_first_finish_before_fetch(self):
824 + """Prefiller-first: finish_request runs before the decoder's
825 + fetch arrives. State is held until fetch, then
826 + finalized — success=True if all demand was matched against
827 + available blocks, else success=False."""
828 + # Case A: all demand satisfied by available blocks.
829 + session, conn, transport = _make_session()
830 + _activate(session, conn)
831 + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)
832 + # finish_request first — no demand received yet -> defer.
833 + session.finish_request("req-1")
834 + assert "req-1" in session._server._outbound
835 + # Fetch arrives now: demand fully satisfied by available.
836 + conn.enqueue(
837 + {
838 + TYPE_KEY: FetchMsg.TYPE,
839 + FetchMsg.KV_REQUEST_ID: "req-1",
840 + FetchMsg.BLOCK_HASHES: [b"k1"],
841 + FetchMsg.BLOCK_INDEXES: [10],
842 + }
843 + )
844 + session.poll()
845 + # The transfer was inflight; finalize it.
846 + tid = next(iter(transport._transfers))
847 + transport._poll_done.append(tid)
848 + session.poll()
849 + msg = next(m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE)
850 + assert msg[TransferDoneMsg.SUCCESS] is True
851 +
852 + # Case B: demand exceeds available -> early-fail fires from fetch.
853 + session, conn, transport = _make_session()
854 + _activate(session, conn)
855 + session.add_stored_blocks("req-2", [b"k1"], [0], job_id=2)
856 + session.finish_request("req-2")
857 + conn.enqueue(
858 + {
859 + TYPE_KEY: FetchMsg.TYPE,
860 + FetchMsg.KV_REQUEST_ID: "req-2",
861 + FetchMsg.BLOCK_HASHES: [b"k1", b"k2"],
862 + FetchMsg.BLOCK_INDEXES: [10, 11],
863 + }
864 + )
865 + session.poll()
866 + # Inflight for k1 still in flight; nothing yet for the early-fail
867 + # — the same code path will fire from _collect_store_results.
868 + tid = next(iter(transport._transfers))
869 + transport._poll_done.append(tid)
870 + session.poll()
871 + msg = next(m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE)
872 + assert msg[TransferDoneMsg.SUCCESS] is False
873 + assert "req-2" not in session._server._outbound
874 +
875 + def test_unknown_request_is_noop(self):
876 + session, conn, _ = _make_session()
877 + _activate(session, conn)
878 + before = len(conn._sent)
879 + session.finish_request("never-existed")
880 + assert len(conn._sent) == before
881 +
882 + def test_finish_request_no_inflight_emits_store_failure(self):
883 + """finish_request with a stored-but-unmatched job and no inflight ->
884 + TransferDoneMsg(success=False) AND deferred
885 + StoreResult(success=False) for the submit_store'd job, instead of
886 + the 30s _STORE_TIMEOUT_S path."""
887 + session, conn, _ = _make_session()
888 + _activate(session, conn)
889 + # Decoder demanded b"demand"; we never stored it.
890 + conn.enqueue(
891 + {
892 + TYPE_KEY: FetchMsg.TYPE,
893 + FetchMsg.KV_REQUEST_ID: "req-1",
894 + FetchMsg.BLOCK_HASHES: [b"demand"],
895 + FetchMsg.BLOCK_INDEXES: [5],
896 + }
897 + )
898 + session.poll()
899 + # We did submit_store a different block — goes to available, never
900 + # matches demand. Without the shortcut, job 42 sits in _store_jobs
901 + # for _STORE_TIMEOUT_S.
902 + session.add_stored_blocks("req-1", [b"unrelated"], [0], job_id=42)
903 + assert session._server._outbound["req-1"].pending_job_ids == {42}
904 +
905 + session.finish_request("req-1")
906 +
907 + # Peer notified immediately with success=False (remaining > 0).
908 + msg = self._last_transfer_done(conn)
909 + assert msg is not None
910 + assert msg[TransferDoneMsg.SUCCESS] is False
911 + assert "req-1" not in session._server._outbound
912 +
913 + # Local store job surfaces on the next poll, success=False.
914 + stores = session.poll().stores
915 + assert StoreResult(job_id=42, success=False) in stores
916 + assert 42 not in session._server._store_jobs
917 +
918 + def test_finish_request_remaining_zero_emits_success_via_inflight(self):
919 + """Deferred-via-inflight path: finish_request with inflight, last
920 + transfer drains remaining to 0 -> TransferDoneMsg(success=True)
921 + AND StoreResult(success=True)."""
922 + session, conn, transport = _make_session()
923 + _activate(session, conn)
924 + conn.enqueue(
925 + {
926 + TYPE_KEY: FetchMsg.TYPE,
927 + FetchMsg.KV_REQUEST_ID: "req-1",
928 + FetchMsg.BLOCK_HASHES: [b"k1"],
929 + FetchMsg.BLOCK_INDEXES: [10],
930 + }
931 + )
932 + session.poll()
933 + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=7)
934 + # finish_request races with the inflight transfer.
935 + session.finish_request("req-1")
936 + assert "req-1" in session._server._outbound # deferred
937 +
938 + # Last inflight completes -> _finalize_outbound(success=True) fires.
939 + tid = next(iter(transport._transfers))
940 + transport._poll_done.append(tid)
941 + stores = session.poll().stores
942 +
943 + msg = self._last_transfer_done(conn)
944 + assert msg is not None
945 + assert msg[TransferDoneMsg.SUCCESS] is True
946 + assert StoreResult(job_id=7, success=True) in stores
947 + assert "req-1" not in session._server._outbound
948 +
949 + def test_write_blocks_failure_finalizes_with_failure(self):
950 + """write_blocks returning None must not leave the request hanging.
951 +
952 + The matched blocks are gone from req.demanded but no inflight
953 + will satisfy them, so remaining > 0 forever. Setting finishing
954 + and calling _finalize_outbound(success=False) immediately (no
955 + other inflight) tells the peer + emits StoreResult(success=False)
956 + without waiting on _STORE_TIMEOUT_S or _LOAD_TIMEOUT_S.
957 + """
958 + session, conn, transport = _make_session()
959 + _activate(session, conn)
960 + # Decoder demands b"k1".
961 + conn.enqueue(
962 + {
963 + TYPE_KEY: FetchMsg.TYPE,
964 + FetchMsg.KV_REQUEST_ID: "req-1",
965 + FetchMsg.BLOCK_HASHES: [b"k1"],
966 + FetchMsg.BLOCK_INDEXES: [10],
967 + }
968 + )
969 + session.poll()
970 + # Force write_blocks to fail on the next call.
971 + transport.write_blocks = lambda *a, **kw: None # type: ignore[assignment]
972 +
973 + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=42)
974 +
975 + # Outbound was finalized immediately (no other inflight).
976 + assert "req-1" not in session._server._outbound
977 + # Peer notified with success=False.
978 + msg = next(m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE)
979 + assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1"
980 + assert msg[TransferDoneMsg.SUCCESS] is False
981 + # Local store job surfaces on the next poll.
982 + stores = session.poll().stores
983 + assert StoreResult(job_id=42, success=False) in stores
984 + assert 42 not in session._server._store_jobs
985 +
986 + def test_partial_match_completes_in_two_rounds(self):
987 + """Peer demand for [k1, k2, k3]; first round only k1 is available,
988 + second round adds k2 and k3. Each round transfers what's matched
989 + and the request finalizes with success once remaining hits zero.
990 + """
991 + session, conn, transport = _make_session()
992 + _activate(session, conn)
993 +
994 + # Peer fetches three blocks.
995 + conn.enqueue(
996 + {
997 + TYPE_KEY: FetchMsg.TYPE,
998 + FetchMsg.KV_REQUEST_ID: "req-1",
999 + FetchMsg.BLOCK_HASHES: [b"k1", b"k2", b"k3"],
1000 + FetchMsg.BLOCK_INDEXES: [10, 11, 12],
1001 + }
1002 + )
1003 + session.poll()
1004 + # Demand registered, no matches yet.
1005 + assert session._server._inflight == {}
1006 + outbound = session._server._outbound["req-1"]
1007 + assert outbound.remaining == 3
1008 + assert set(outbound.demanded.keys()) == {b"k1", b"k2", b"k3"}
1009 +
1010 + # Round 1: only k1 is stored locally.
1011 + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=100)
1012 + # One inflight transfer for k1.
1013 + assert len(session._server._inflight) == 1
1014 + tid_1 = next(iter(session._server._inflight))
1015 + assert transport._transfers[tid_1][1] == [0]
1016 + assert transport._transfers[tid_1][2] == [10]
1017 + # k2 and k3 still demanded.
1018 + assert set(outbound.demanded.keys()) == {b"k2", b"k3"}
1019 +
1020 + # Transfer 1 completes.
1021 + transport._poll_done.append(tid_1)
1022 + stores = session.poll().stores
1023 + assert StoreResult(job_id=100, success=True) in stores
1024 + assert session._server._inflight == {}
1025 + assert outbound.remaining == 2
1026 + # Not yet finalized — still 2 blocks demanded.
1027 + assert "req-1" in session._server._outbound
1028 +
1029 + # Round 2: k2 and k3 arrive together.
1030 + session.add_stored_blocks("req-1", [b"k2", b"k3"], [1, 2], job_id=200)
1031 + assert len(session._server._inflight) == 1
1032 + tid_2 = next(iter(session._server._inflight))
1033 + assert tid_2 != tid_1
1034 + assert sorted(transport._transfers[tid_2][1]) == [1, 2]
1035 +
1036 + # Transfer 2 completes — request now fully satisfied.
1037 + transport._poll_done.append(tid_2)
1038 + stores = session.poll().stores
1039 + assert StoreResult(job_id=200, success=True) in stores
1040 + # _finalize_outbound fired — request gone, peer notified with success.
1041 + assert "req-1" not in session._server._outbound
1042 + done = next(m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE)
1043 + assert done[TransferDoneMsg.KV_REQUEST_ID] == "req-1"
1044 + assert done[TransferDoneMsg.SUCCESS] is True
1045 +
1046 + def test_write_blocks_failure_finalizes_after_last_inflight_completes(self):
1047 + """write_blocks returns None on a SECOND match while a first transfer
1048 + is still inflight. The request should NOT finalize until the inflight
1049 + completes, then the elif branch in collect_results
1050 + (``finishing and not _has_inflight_for(...)``) finalizes it as failure.
1051 + """
1052 + session, conn, transport = _make_session()
1053 + _activate(session, conn)
1054 +
1055 + # Peer demands two blocks.
1056 + conn.enqueue(
1057 + {
1058 + TYPE_KEY: FetchMsg.TYPE,
1059 + FetchMsg.KV_REQUEST_ID: "req-1",
1060 + FetchMsg.BLOCK_HASHES: [b"k1", b"k2"],
1061 + FetchMsg.BLOCK_INDEXES: [10, 11],
1062 + }
1063 + )
1064 + session.poll()
1065 +
1066 + # Round 1: k1 transfers cleanly.
1067 + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=100)
1068 + assert len(session._server._inflight) == 1
1069 + tid_1 = next(iter(session._server._inflight))
1070 + outbound = session._server._outbound["req-1"]
1071 + assert outbound.remaining == 2 # decrement happens on completion
1072 + assert outbound.finishing is False
1073 +
1074 + # Round 2: write_blocks fails for k2 while transfer_1 is still inflight.
1075 + transport.write_blocks = lambda *a, **kw: None # type: ignore[assignment]
1076 + session.add_stored_blocks("req-1", [b"k2"], [1], job_id=200)
1077 + # No new transfer was registered.
1078 + assert list(session._server._inflight.keys()) == [tid_1]
1079 + # Marked finishing, but NOT finalized yet (transfer_1 still inflight).
1080 + assert outbound.finishing is True
1081 + assert "req-1" in session._server._outbound
1082 + done_msgs = [m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE]
1083 + assert done_msgs == []
1084 +
1085 + # Transfer 1 completes — now ``_has_inflight_for("req-1")`` is False
1086 + # and the elif branch in collect_results fires _finalize(success=False).
1087 + # The k1 success result is direct; the k2 failure result is queued
1088 + # in _pending_store_results and surfaces on the NEXT poll.
1089 + transport._poll_done.append(tid_1)
1090 + stores_first = session.poll().stores
1091 + assert StoreResult(job_id=100, success=True) in stores_first
1092 + # Outbound state cleaned up; peer notified with success=False.
1093 + assert "req-1" not in session._server._outbound
1094 + done = next(m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE)
1095 + assert done[TransferDoneMsg.KV_REQUEST_ID] == "req-1"
1096 + assert done[TransferDoneMsg.SUCCESS] is False
1097 +
1098 + # Next poll drains the queued failure.
1099 + stores_second = session.poll().stores
1100 + assert StoreResult(job_id=200, success=False) in stores_second
1101 +
1102 +
1103 +# ---------------------------------------------------------------------------
1104 +# Bidirectional — the case the unification is meant to fix
1105 +# ---------------------------------------------------------------------------
1106 +
1107 +
1108 +class TestBidirectional:
1109 + def test_session_handles_both_roles_concurrently(self):
1110 + """Single session simultaneously serves a fetch and completes a load.
1111 +
1112 + This is the regression test for the unification: with the old
1113 + split design the inbound FetchMsg would be dispatched to a
1114 + client-only session and dropped (or a server-only session would
1115 + miss the TransferDoneMsg). One unified session handles both.
1116 + """
1117 + session, conn, transport = _make_session()
1118 + _activate(session, conn)
1119 +
1120 + # Server role: the peer fetches a block from us.
1121 + session.add_stored_blocks("req-srv", [b"served"], [0], job_id=100)
1122 + conn.enqueue(
1123 + {
1124 + TYPE_KEY: FetchMsg.TYPE,
1125 + FetchMsg.KV_REQUEST_ID: "req-srv",
1126 + FetchMsg.BLOCK_HASHES: [b"served"],
1127 + FetchMsg.BLOCK_INDEXES: [7],
1128 + }
1129 + )
1130 +
1131 + # Client role: we ask the peer for a different block.
1132 + session.request_blocks(
1133 + job_id=200, kv_request_id="req-cli", keys=[b"loaded"], block_ids=[3]
1134 + )
1135 +
1136 + # Both flows progress in the same poll.
1137 + session.poll()
1138 +
1139 + # Server side: write_blocks was submitted.
1140 + assert len(transport._transfers) == 1
1141 + # Client side: the lookup was sent.
1142 + assert any(
1143 + m[TYPE_KEY] == FetchMsg.TYPE and m[FetchMsg.KV_REQUEST_ID] == "req-cli"
1144 + for m in conn._sent
1145 + )
1146 +
1147 + # Peer now signals: server-side transfer completes AND a
1148 + # TransferDoneMsg arrives for our client-side request, all in
1149 + # one batch on the same connection.
1150 + tid = next(iter(transport._transfers))
1151 + transport._poll_done.append(tid)
1152 + conn.enqueue(
1153 + {
1154 + TYPE_KEY: TransferDoneMsg.TYPE,
1155 + TransferDoneMsg.KV_REQUEST_ID: "req-cli",
1156 + TransferDoneMsg.SUCCESS: True,
1157 + }
1158 + )
1159 + result_ = session.poll()
1160 + loads = result_.loads
1161 + stores = result_.stores
1162 +
1163 + assert LoadResult(job_id=200, kv_request_id="req-cli", success=True) in loads
1164 + assert StoreResult(job_id=100, success=True) in stores
1165 +
1166 +
1167 +# ---------------------------------------------------------------------------
1168 +# Pending sessions (no connection yet)
1169 +# ---------------------------------------------------------------------------
1170 +
1171 +
1172 +class TestPendingSession:
1173 + def test_pending_session_buffers_stored_blocks(self):
1174 + """Pending session accepts add_stored_blocks but cannot send."""
1175 + transport = FakeDataTransport()
1176 + session = P2PSession(
1177 + peer_id="peer:8000",
1178 + local_id="local:9000",
1179 + transport=transport, # type: ignore[arg-type]
1180 + local_block_len=4096,
1181 + conn=None,
1182 + )
1183 + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)
1184 + result_ = session.poll()
1185 + loads = result_.loads
1186 + stores = result_.stores
1187 + assert loads == []
1188 + assert stores == []
1189 + assert not session.connected
1190 + assert session.alive
1191 +
1192 + def test_attach_connection_sends_connect(self):
1193 + """attach_connection triggers our ConnectMsg send."""
1194 + transport = FakeDataTransport()
1195 + session = P2PSession(
1196 + peer_id="peer:8000",
1197 + local_id="local:9000",
1198 + transport=transport, # type: ignore[arg-type]
1199 + local_block_len=4096,
1200 + conn=None,
1201 + )
1202 + conn = FakeConnection(peer_id="peer:8000")
1203 + session.attach_connection(conn) # type: ignore[arg-type]
1204 + assert conn._sent
1205 + assert conn._sent[0][TYPE_KEY] == ConnectMsg.TYPE
1206 +
1207 + def test_attach_connection_twice_raises(self):
1208 + """attach_connection on an already-connected session raises."""
1209 + session, conn, _ = _make_session()
1210 + with pytest.raises(ValueError, match="already connected"):
1211 + session.attach_connection(FakeConnection()) # type: ignore[arg-type]
1212 +
1213 + def test_pending_close_returns_pending_stores(self):
1214 + """Closing a pending session reports buffered stores as failed."""
1215 + transport = FakeDataTransport()
1216 + session = P2PSession(
1217 + peer_id="peer:8000",
1218 + local_id="local:9000",
1219 + transport=transport, # type: ignore[arg-type]
1220 + local_block_len=4096,
1221 + conn=None,
1222 + )
1223 + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)
1224 + session.add_stored_blocks("req-2", [b"k2"], [1], job_id=2)
1225 + failed_loads, failed_stores = session.close()
1226 + assert failed_loads == []
1227 + assert set(failed_stores) == {1, 2}
1228 +
1229 +
1230 +# ---------------------------------------------------------------------------
1231 +# Disconnect
1232 +# ---------------------------------------------------------------------------
1233 +
1234 +
1235 +class TestDisconnect:
1236 + def test_disconnect_marks_session_dead(self):
1237 + session, conn, _ = _make_session()
1238 + _activate(session, conn)
1239 + conn.enqueue({TYPE_KEY: DisconnectMsg.TYPE})
1240 + session.poll()
1241 + assert not session.alive
1242 +
1243 + def test_close_returns_pending_loads_and_stores(self):
1244 + session, conn, _ = _make_session()
1245 + _activate(session, conn)
1246 + session.request_blocks(1, "req-1", [b"k"], [0])
1247 + session.request_blocks(2, "req-2", [b"k"], [0])
1248 + session.add_stored_blocks("req-srv", [b"k"], [0], job_id=10)
1249 + failed_loads, failed_stores = session.close()
1250 + assert set(failed_loads) == {(1, "req-1"), (2, "req-2")}
1251 + assert set(failed_stores) == {10}
1252 +
1253 +
1254 +# ---------------------------------------------------------------------------
1255 +# Adversarial / malformed messages
1256 +# ---------------------------------------------------------------------------
1257 +
1258 +
1259 +class TestAdversarial:
1260 + def test_unknown_message_type_logged(self):
1261 + session, conn, _ = _make_session()
1262 + _activate(session, conn)
1263 + conn.enqueue({TYPE_KEY: "evil_command"})
1264 + result_ = session.poll()
1265 + loads = result_.loads
1266 + stores = result_.stores
1267 + assert loads == []
1268 + assert stores == []
1269 +
1270 + def test_empty_message(self):
1271 + session, conn, _ = _make_session()
1272 + _activate(session, conn)
1273 + conn.enqueue({})
1274 + result_ = session.poll()
1275 + loads = result_.loads
1276 + stores = result_.stores
1277 + assert loads == []
1278 + assert stores == []
1279 +
1280 + def test_non_dict_message(self):
1281 + session, conn, _ = _make_session()
1282 + _activate(session, conn)
1283 + conn._inbox.append(42) # type: ignore[arg-type]
1284 + result_ = session.poll()
1285 + loads = result_.loads
1286 + stores = result_.stores
1287 + assert loads == []
1288 + assert stores == []
1289 +
1290 + def test_fetch_mismatched_lengths(self):
1291 + session, conn, transport = _make_session()
1292 + _activate(session, conn)
1293 + conn.enqueue(
1294 + {
1295 + TYPE_KEY: FetchMsg.TYPE,
1296 + FetchMsg.KV_REQUEST_ID: "req-bad",
1297 + FetchMsg.BLOCK_HASHES: [b"k1", b"k2"],
1298 + FetchMsg.BLOCK_INDEXES: [1],
1299 + }
1300 + )
1301 + session.poll()
1302 + assert len(transport._transfers) == 0
1303 + # Protocol violation: session disconnects immediately so the peer
1304 + # can't keep wedging us with malformed traffic.
1305 + assert not session.alive
1306 + assert any(m[TYPE_KEY] == DisconnectMsg.TYPE for m in conn._sent)
1307 +
1308 + def test_transfer_done_missing_kv_request_id(self):
1309 + session, conn, _ = _make_session()
1310 + _activate(session, conn)
1311 + conn.enqueue({TYPE_KEY: TransferDoneMsg.TYPE, TransferDoneMsg.SUCCESS: True})
1312 + loads = session.poll().loads
1313 + assert loads == []
1314 + assert not session.alive
1315 + assert any(m[TYPE_KEY] == DisconnectMsg.TYPE for m in conn._sent)
1316 +
1317 + def test_duplicate_connect_ack(self):
1318 + session, conn, _ = _make_session()
1319 + _activate(session, conn)
1320 + conn.enqueue({TYPE_KEY: ConnectAckMsg.TYPE, ConnectAckMsg.PEER_ID: "peer:8000"})
1321 + session.poll()
1322 + assert session.ready
1323 +
1324 +
1325 +class TestDispatchErrorHandling:
1326 + """Errors raised by message handlers split into two classes:
1327 +
1328 + - Protocol-contract violations from the peer (ValueError) → disconnect
1329 + on the first occurrence; retrying won't help and may corrupt state.
1330 + - Anything else is treated as an internal bug: log loudly, count, and
1331 + only disconnect once errors arrive in a tight burst. A successful
1332 + dispatch in between resets the counter.
1333 + """
1334 +
1335 + def test_value_error_disconnects_on_first_occurrence(self):
1336 + """A FetchMsg that fails validate() raises ValueError and must
1337 + terminate the session immediately, with a DisconnectMsg sent."""
1338 + session, conn, _ = _make_session()
1339 + _activate(session, conn)
1340 + # length mismatch → FetchMsg.validate raises ValueError
1341 + conn.enqueue(
1342 + {
1343 + TYPE_KEY: FetchMsg.TYPE,
1344 + FetchMsg.KV_REQUEST_ID: "req-bad",
1345 + FetchMsg.BLOCK_HASHES: [b"k1", b"k2"],
1346 + FetchMsg.BLOCK_INDEXES: [1],
1347 + }
1348 + )
1349 + session.poll()
1350 + assert not session.alive
1351 + assert any(m[TYPE_KEY] == DisconnectMsg.TYPE for m in conn._sent)
1352 +
1353 + def test_transfer_done_missing_field_disconnects(self):
1354 + """Same contract for a malformed TransferDoneMsg from the peer."""
1355 + session, conn, _ = _make_session()
1356 + _activate(session, conn)
1357 + conn.enqueue({TYPE_KEY: TransferDoneMsg.TYPE, TransferDoneMsg.SUCCESS: True})
1358 + session.poll()
1359 + assert not session.alive
1360 +
1361 + def test_internal_error_does_not_disconnect_once(self):
1362 + """A non-ValueError raised by a handler is treated as an internal
1363 + bug: counter increments, session stays alive on a single hit."""
1364 + session, conn, _ = _make_session()
1365 + _activate(session, conn)
1366 +
1367 + def _boom(*args, **kwargs):
1368 + raise RuntimeError("simulated internal bug")
1369 +
1370 + session._server.on_fetch = _boom # type: ignore[assignment]
1371 + conn.enqueue(
1372 + {
1373 + TYPE_KEY: FetchMsg.TYPE,
1374 + FetchMsg.KV_REQUEST_ID: "req-1",
1375 + FetchMsg.BLOCK_HASHES: [b"k1"],
1376 + FetchMsg.BLOCK_INDEXES: [0],
1377 + }
1378 + )
1379 + session.poll()
1380 + assert session.alive
1381 + assert session._dispatch_error_count == 1
1382 +
1383 + def test_internal_error_threshold_disconnects(self):
1384 + """Once consecutive non-protocol errors hit the threshold, the
1385 + session tears down via _protocol_error."""
1386 + session, conn, _ = _make_session()
1387 + _activate(session, conn)
1388 +
1389 + def _boom(*args, **kwargs):
1390 + raise RuntimeError("simulated internal bug")
1391 +
1392 + session._server.on_fetch = _boom # type: ignore[assignment]
1393 + for _ in range(_MAX_CONSECUTIVE_DISPATCH_ERRORS):
1394 + conn.enqueue(
1395 + {
1396 + TYPE_KEY: FetchMsg.TYPE,
1397 + FetchMsg.KV_REQUEST_ID: "req-1",
1398 + FetchMsg.BLOCK_HASHES: [b"k1"],
1399 + FetchMsg.BLOCK_INDEXES: [0],
1400 + }
1401 + )
1402 + session.poll()
1403 + assert not session.alive
1404 + assert any(m[TYPE_KEY] == DisconnectMsg.TYPE for m in conn._sent)
1405 +
1406 + def test_internal_error_counter_resets_on_success(self):
1407 + """A successful dispatch between errors prevents the threshold
1408 + from being reached."""
1409 + session, conn, _ = _make_session()
1410 + _activate(session, conn)
1411 +
1412 + original_on_fetch = session._server.on_fetch
1413 +
1414 + def _boom(*args, **kwargs):
1415 + raise RuntimeError("simulated internal bug")
1416 +
1417 + # Alternate (boom, success) (_MAX-1) times: counter rises to 1
1418 + # then resets to 0 each cycle, never reaching the threshold.
1419 + for _ in range(_MAX_CONSECUTIVE_DISPATCH_ERRORS - 1):
1420 + session._server.on_fetch = _boom # type: ignore[assignment]
1421 + conn.enqueue(
1422 + {
1423 + TYPE_KEY: FetchMsg.TYPE,
1424 + FetchMsg.KV_REQUEST_ID: "req-1",
1425 + FetchMsg.BLOCK_HASHES: [b"k1"],
1426 + FetchMsg.BLOCK_INDEXES: [0],
1427 + }
1428 + )
1429 + session.poll()
1430 + session._server.on_fetch = original_on_fetch # type: ignore[assignment]
1431 + # A benign no-op message (unknown type) dispatches cleanly
1432 + # and resets the consecutive-error counter.
1433 + conn.enqueue({TYPE_KEY: "unknown_for_test"})
1434 + session.poll()
1435 +
1436 + assert session.alive
1437 + assert session._dispatch_error_count == 0
1438 +
1439 +
1440 +class TestInflightPerReqInvariant:
1441 + """`_inflight_per_req` is the O(1) replacement for the previous
1442 + O(N) scan in `_has_inflight_for`. These tests check that every
1443 + mutation site keeps the counter in sync with `_inflight` and that
1444 + the lookup is correct under high fan-out.
1445 + """
1446 +
1447 + def test_invariant_holds_through_lifecycle(self):
1448 + """Run a full submit→complete sequence for two concurrent
1449 + kv_request_ids and assert the counter matches `_inflight` at
1450 + every observable step, including the empty-after-finish case."""
1451 + session, conn, transport = _make_session()
1452 + _activate(session, conn)
1453 +
1454 + def _invariant_holds() -> bool:
1455 + counted = sum(session._server._inflight_per_req.values())
1456 + return counted == len(session._server._inflight) and all(
1457 + v > 0 for v in session._server._inflight_per_req.values()
1458 + )
1459 +
1460 + assert _invariant_holds()
1461 +
1462 + # Two requests, two blocks each, all dispatched in one batch.
1463 + session.add_stored_blocks("req-A", [b"a1", b"a2"], [0, 1], job_id=10)
1464 + session.add_stored_blocks("req-B", [b"b1", b"b2"], [2, 3], job_id=11)
1465 + for kv_id, hashes, indexes in (
1466 + ("req-A", [b"a1", b"a2"], [100, 101]),
1467 + ("req-B", [b"b1", b"b2"], [102, 103]),
1468 + ):
1469 + conn.enqueue(
1470 + {
1471 + TYPE_KEY: FetchMsg.TYPE,
1472 + FetchMsg.KV_REQUEST_ID: kv_id,
1473 + FetchMsg.BLOCK_HASHES: hashes,
1474 + FetchMsg.BLOCK_INDEXES: indexes,
1475 + }
1476 + )
1477 + session.poll()
1478 +
1479 + assert _invariant_holds()
1480 + assert session._server._has_inflight_for("req-A")
1481 + assert session._server._has_inflight_for("req-B")
1482 + assert not session._server._has_inflight_for("req-C")
1483 +
1484 + # Complete req-A's transfer first; req-B should still be inflight.
1485 + a_tids = [
1486 + tid
1487 + for tid, x in session._server._inflight.items()
1488 + if x.kv_request_id == "req-A"
1489 + ]
1490 + for tid in a_tids:
1491 + transport._poll_done.append(tid)
1492 + session.poll()
1493 +
1494 + assert _invariant_holds()
1495 + assert not session._server._has_inflight_for("req-A")
1496 + assert "req-A" not in session._server._inflight_per_req # entry was removed
1497 + assert session._server._has_inflight_for("req-B")
1498 +
1499 + # Complete req-B; counter must drain to empty.
1500 + b_tids = [
1501 + tid
1502 + for tid, x in session._server._inflight.items()
1503 + if x.kv_request_id == "req-B"
1504 + ]
1505 + for tid in b_tids:
1506 + transport._poll_done.append(tid)
1507 + session.poll()
1508 +
1509 + assert _invariant_holds()
1510 + assert session._server._inflight == {}
1511 + assert session._server._inflight_per_req == {}
1512 +
1513 + def test_has_inflight_for_correct_with_many_requests(self):
1514 + """Populate many inflight xfers across many ids; lookup must
1515 + match the actual presence in `_inflight` for both hits and
1516 + misses. The whole point of the counter is that this lookup is
1517 + constant-time, but we assert correctness, not timing."""
1518 + session, _, _ = _make_session()
1519 + for kv_id_idx in range(100):
1520 + kv_id = f"req-{kv_id_idx}"
1521 + for j in range(10):
1522 + tid = kv_id_idx * 10 + j
1523 + session._server._inflight_add(
1524 + tid,
1525 + _InflightXfer(kv_request_id=kv_id, block_count=1, job_ids={tid}),
1526 + )
1527 + assert sum(session._server._inflight_per_req.values()) == len(
1528 + session._server._inflight
1529 + )
1530 + assert session._server._has_inflight_for("req-0")
1531 + assert session._server._has_inflight_for("req-99")
1532 + assert not session._server._has_inflight_for("req-missing")
1533 +
1534 + # Drain all entries for req-50 and confirm the entry disappears.
1535 + tids_50 = [
1536 + tid
1537 + for tid, x in session._server._inflight.items()
1538 + if x.kv_request_id == "req-50"
1539 + ]
1540 + for tid in tids_50:
1541 + session._server._inflight_pop(tid)
1542 + assert "req-50" not in session._server._inflight_per_req
1543 + assert not session._server._has_inflight_for("req-50")
1544 + # Other ids unaffected.
1545 + assert session._server._has_inflight_for("req-49")
1546 +
1547 +
1548 +# ---------------------------------------------------------------------------
1549 +# Protocol validation tests (unchanged from the prior file)
1550 +# ---------------------------------------------------------------------------
1551 +
1552 +
1553 +class TestConnectMsgValidation:
1554 + def _valid_msg(self) -> dict:
1555 + return {
1556 + TYPE_KEY: ConnectMsg.TYPE,
1557 + ConnectMsg.PEER_ID: "peer:1",
1558 + ConnectMsg.AGENT_METADATA: b"meta",
1559 + ConnectMsg.BASE_ADDR: 0x1000,
1560 + ConnectMsg.NUM_BLOCKS: 8,
1561 + ConnectMsg.BLOCK_LEN: 4096,
1562 + }
1563 +
1564 + def test_valid_message_passes(self):
1565 + ConnectMsg.validate(self._valid_msg())
1566 +
1567 + def test_missing_peer_id(self):
1568 + msg = self._valid_msg()
1569 + del msg[ConnectMsg.PEER_ID]
1570 + with pytest.raises(ValueError, match="peer_id"):
1571 + ConnectMsg.validate(msg)
1572 +
1573 + def test_missing_agent_metadata(self):
1574 + msg = self._valid_msg()
1575 + del msg[ConnectMsg.AGENT_METADATA]
1576 + with pytest.raises(ValueError, match="agent_metadata"):
1577 + ConnectMsg.validate(msg)
1578 +
1579 + def test_agent_metadata_wrong_type(self):
1580 + msg = self._valid_msg()
1581 + msg[ConnectMsg.AGENT_METADATA] = "not bytes"
1582 + with pytest.raises(ValueError, match="agent_metadata"):
1583 + ConnectMsg.validate(msg)
1584 +
1585 + def test_base_addr_negative(self):
1586 + msg = self._valid_msg()
1587 + msg[ConnectMsg.BASE_ADDR] = -1
1588 + with pytest.raises(ValueError, match="base_addr"):
1589 + ConnectMsg.validate(msg)
1590 +
1591 + def test_num_blocks_zero(self):
1592 + msg = self._valid_msg()
1593 + msg[ConnectMsg.NUM_BLOCKS] = 0
1594 + with pytest.raises(ValueError, match="num_blocks"):
1595 + ConnectMsg.validate(msg)
1596 +
1597 + def test_block_len_zero(self):
1598 + msg = self._valid_msg()
1599 + msg[ConnectMsg.BLOCK_LEN] = 0
1600 + with pytest.raises(ValueError, match="block_len"):
1601 + ConnectMsg.validate(msg)
1602 +
1603 +
1604 +class TestFetchMsgValidation:
1605 + def _valid_msg(self) -> dict:
1606 + return {
1607 + TYPE_KEY: FetchMsg.TYPE,
1608 + FetchMsg.KV_REQUEST_ID: "req-1",
1609 + FetchMsg.BLOCK_HASHES: [b"k1", b"k2"],
1610 + FetchMsg.BLOCK_INDEXES: [0, 1],
1611 + }
1612 +
1613 + def test_valid_message_passes(self):
1614 + FetchMsg.validate(self._valid_msg())
1615 +
1616 + def test_length_mismatch(self):
1617 + msg = self._valid_msg()
1618 + msg[FetchMsg.BLOCK_INDEXES] = [0]
1619 + with pytest.raises(ValueError, match="length mismatch"):
1620 + FetchMsg.validate(msg)
1621 +
1622 + def test_negative_index(self):
1623 + msg = self._valid_msg()
1624 + msg[FetchMsg.BLOCK_INDEXES] = [0, -1]
1625 + with pytest.raises(ValueError, match="invalid index"):
1626 + FetchMsg.validate(msg)
1627 +
1628 +
1629 +class TestTransferDoneMsgValidation:
1630 + def test_valid_message_passes(self):
1631 + msg = {
1632 + TYPE_KEY: TransferDoneMsg.TYPE,
1633 + TransferDoneMsg.KV_REQUEST_ID: "req-1",
1634 + TransferDoneMsg.SUCCESS: True,
1635 + }
1636 + TransferDoneMsg.validate(msg)
1637 +
1638 + def test_success_wrong_type(self):
1639 + msg = {
1640 + TYPE_KEY: TransferDoneMsg.TYPE,
1641 + TransferDoneMsg.KV_REQUEST_ID: "req-1",
1642 + TransferDoneMsg.SUCCESS: 1,
1643 + }
1644 + with pytest.raises(ValueError, match="success"):
1645 + TransferDoneMsg.validate(msg)
tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py new
+230
@@ -0,0 +1,230 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for vllm.v1.kv_offload.tiering.p2p.control.zmq."""
4 +
5 +from __future__ import annotations
6 +
7 +import socket
8 +import time
9 +
10 +import pytest
11 +import zmq
12 +
13 +from vllm.v1.kv_offload.tiering.p2p.control.zmq import (
14 + ZmqConnection,
15 + ZmqTransport,
16 + _Sockets,
17 +)
18 +
19 +
20 +def _free_port() -> int:
21 + """Find a free TCP port."""
22 + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
23 + s.bind(("127.0.0.1", 0))
24 + return s.getsockname()[1]
25 +
26 +
27 +def _make_transport(host: str = "127.0.0.1", attempts: int = 8):
28 + """Construct a ZmqTransport on a fresh port, retrying on bind collisions.
29 +
30 + Why: _free_port() releases the probe socket before ZmqTransport binds the
31 + same port — a parallel test run can steal it in between. Retrying on
32 + ZMQError/OSError closes that race without a production change.
33 + """
34 + last_err: Exception | None = None
35 + for _ in range(attempts):
36 + port = _free_port()
37 + try:
38 + return ZmqTransport(f"{host}:{port}", host, port), port
39 + except (zmq.ZMQError, OSError) as e:
40 + last_err = e
41 + assert last_err is not None
42 + raise last_err
43 +
44 +
45 +def _wait_for_inbound(transport: ZmqTransport, deadline: float = 2.0):
46 + """Poll until at least one new inbound connection is accepted, or fail."""
47 + end = time.monotonic() + deadline
48 + while time.monotonic() < end:
49 + new = transport.poll()
50 + if new:
51 + return new
52 + time.sleep(0.005)
53 + raise AssertionError(f"no inbound connection within {deadline}s")
54 +
55 +
56 +def _wait_for_messages(
57 + transport: ZmqTransport,
58 + conn: ZmqConnection,
59 + n: int,
60 + deadline: float = 2.0,
61 +) -> list[dict]:
62 + """Poll until `conn` has received at least `n` messages, then return them."""
63 + end = time.monotonic() + deadline
64 + msgs: list[dict] = []
65 + while time.monotonic() < end:
66 + transport.poll()
67 + msgs.extend(conn.recv())
68 + if len(msgs) >= n:
69 + return msgs
70 + time.sleep(0.005)
71 + raise AssertionError(f"got {len(msgs)}/{n} messages within {deadline}s")
72 +
73 +
74 +def _make_mock_connection(peer_id: str = "test:1234") -> ZmqConnection:
75 + """Create a ZmqConnection with mock sockets for unit testing."""
76 + from unittest.mock import MagicMock
77 +
78 + sockets = _Sockets(dealer=MagicMock(), monitor=MagicMock())
79 + return ZmqConnection(peer_id, sockets)
80 +
81 +
82 +class TestZmqConnection:
83 + """Tests for ZmqConnection in isolation (no real sockets)."""
84 +
85 + def test_enqueue_and_recv(self):
86 + """Messages enqueued are returned by recv() in order."""
87 + conn = _make_mock_connection()
88 +
89 + conn.enqueue({"type": "a"})
90 + conn.enqueue({"type": "b"})
91 +
92 + msgs = conn.recv()
93 + assert list(msgs) == [{"type": "a"}, {"type": "b"}]
94 + # Second recv is empty
95 + assert not conn.recv()
96 +
97 + def test_recv_returns_empty_initially(self):
98 + conn = _make_mock_connection()
99 + assert not conn.recv()
100 +
101 + def test_alive_initially_true(self):
102 + conn = _make_mock_connection()
103 + assert conn.alive is True
104 +
105 + def test_mark_dead(self):
106 + conn = _make_mock_connection()
107 + conn.mark_dead()
108 + assert conn.alive is False
109 +
110 + def test_send_raises_when_closed(self):
111 + conn = _make_mock_connection()
112 + conn.mark_dead()
113 +
114 + with pytest.raises(RuntimeError, match="closed connection"):
115 + conn.send({"type": "test"})
116 +
117 +
118 +class TestZmqTransportConnectivity:
119 + """Integration tests for ZmqTransport with real ZMQ sockets."""
120 +
121 + def test_connect_and_send_message(self):
122 + """Two transports can connect and exchange messages."""
123 + transport_a, port_a = _make_transport()
124 + transport_b, port_b = _make_transport()
125 +
126 + try:
127 + peer_a_id = f"127.0.0.1:{port_a}"
128 + conn_b_to_a = transport_b.connect(peer_a_id)
129 + conn_b_to_a.send({"type": "hello", "data": 42})
130 +
131 + new_conns = _wait_for_inbound(transport_a)
132 + assert len(new_conns) == 1
133 +
134 + conn_a_from_b = new_conns[0]
135 + assert conn_a_from_b.peer_id == f"127.0.0.1:{port_b}"
136 +
137 + msgs = _wait_for_messages(transport_a, conn_a_from_b, 1)
138 + assert msgs == [{"type": "hello", "data": 42}]
139 + finally:
140 + transport_a.close()
141 + transport_b.close()
142 +
143 + def test_bidirectional_messaging(self):
144 + """Both sides can send and receive after connection."""
145 + transport_a, port_a = _make_transport()
146 + transport_b, _ = _make_transport()
147 +
148 + try:
149 + conn_b = transport_b.connect(f"127.0.0.1:{port_a}")
150 + conn_b.send({"type": "connect", "from": "b"})
151 +
152 + new_conns = _wait_for_inbound(transport_a)
153 + assert len(new_conns) == 1
154 + conn_a = new_conns[0]
155 +
156 + conn_a.send({"type": "reply", "from": "a"})
157 +
158 + msgs = _wait_for_messages(transport_b, conn_b, 1)
159 + assert msgs == [{"type": "reply", "from": "a"}]
160 + finally:
161 + transport_a.close()
162 + transport_b.close()
163 +
164 + def test_poll_returns_empty_when_no_connections(self):
165 + transport, _ = _make_transport()
166 + try:
167 + assert not transport.poll()
168 + finally:
169 + transport.close()
170 +
171 + def test_multiple_messages(self):
172 + """Multiple messages are buffered and returned together."""
173 + transport_a, port_a = _make_transport()
174 + transport_b, _ = _make_transport()
175 +
176 + try:
177 + conn_b = transport_b.connect(f"127.0.0.1:{port_a}")
178 + conn_b.send({"seq": 1})
179 + conn_b.send({"seq": 2})
180 + conn_b.send({"seq": 3})
181 +
182 + new_conns = _wait_for_inbound(transport_a)
183 + assert len(new_conns) == 1
184 + conn_a = new_conns[0]
185 +
186 + msgs = _wait_for_messages(transport_a, conn_a, 3)
187 + assert [m["seq"] for m in msgs] == [1, 2, 3]
188 + finally:
189 + transport_a.close()
190 + transport_b.close()
191 +
192 + def test_duplicate_connect_asserts(self):
193 + """Connecting to the same peer twice raises AssertionError."""
194 + # port_a is never bound — we just need a syntactically-valid peer id.
195 + port_a = _free_port()
196 + transport_b, _ = _make_transport()
197 + try:
198 + transport_b.connect(f"127.0.0.1:{port_a}")
199 + with pytest.raises(AssertionError, match="already exists"):
200 + transport_b.connect(f"127.0.0.1:{port_a}")
201 + finally:
202 + transport_b.close()
203 +
204 + def test_dead_connection_removed_on_poll(self):
205 + """Dead connections are cleaned up during poll."""
206 + transport_a, port_a = _make_transport()
207 + transport_b, _ = _make_transport()
208 +
209 + try:
210 + conn_b = transport_b.connect(f"127.0.0.1:{port_a}")
211 + conn_b.send({"type": "hello"})
212 +
213 + new_conns = _wait_for_inbound(transport_a)
214 + assert len(new_conns) == 1
215 +
216 + # Mark the inbound connection dead manually.
217 + new_conns[0].mark_dead()
218 +
219 + # Pruning is synchronous within poll().
220 + transport_a.poll()
221 + assert len(transport_a._connections) == 0
222 + finally:
223 + transport_a.close()
224 + transport_b.close()
225 +
226 + def test_close_is_idempotent(self):
227 + """Calling close() twice doesn't raise."""
228 + transport, _ = _make_transport()
229 + transport.close()
230 + transport.close() # should not raise
vllm/v1/kv_offload/tiering/factory.py
+6
@@ -66,6 +66,12 @@ SecondaryTierFactory.register_tier(
66 "FileSystemTierManager",
67 )
68
69 +SecondaryTierFactory.register_tier(
70 + "p2p",
71 + "vllm.v1.kv_offload.tiering.p2p.manager",
72 + "P2PSecondaryTierManager",
73 +)
74 +
75 SecondaryTierFactory.register_tier(
76 "obj",
77 "vllm.v1.kv_offload.tiering.obj.manager",
vllm/v1/kv_offload/tiering/p2p/__init__.py
vllm/v1/kv_offload/tiering/p2p/control/__init__.py new
+17
@@ -0,0 +1,17 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from vllm.v1.kv_offload.tiering.p2p.control.base import (
4 + ControlConnection,
5 + ControlTransport,
6 +)
7 +from vllm.v1.kv_offload.tiering.p2p.control.zmq import (
8 + ZmqConnection,
9 + ZmqTransport,
10 +)
11 +
12 +__all__ = [
13 + "ControlConnection",
14 + "ControlTransport",
15 + "ZmqConnection",
16 + "ZmqTransport",
17 +]
vllm/v1/kv_offload/tiering/p2p/control/base.py new
+166
@@ -0,0 +1,166 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Abstract base classes for the P2P control-plane transport.
4 +
5 +The control plane handles peer discovery, connection lifecycle, and
6 +message routing. It is message-content agnostic — it moves opaque
7 +dicts between peers without interpreting them.
8 +
9 +Architecture
10 +------------
11 +
12 + ControlTransport (one per node)
13 + ├── listen for inbound connections
14 + ├── connect() to outbound peers
15 + └── poll() → new ControlConnections
16 +
17 + ControlConnection (one per peer)
18 + ├── send(msg) — enqueue a message to the peer
19 + ├── recv() — drain buffered inbound messages
20 + ├── mark_dead() — signal that the peer is gone
21 + └── close() — tear down the connection
22 +
23 +Threading model: all I/O is driven by the caller invoking poll().
24 +No background threads. poll() must be called periodically to:
25 + - receive messages (buffered per-connection)
26 + - accept new inbound peers
27 + - detect disconnections
28 +
29 +Implementor contracts
30 +---------------------
31 +
32 +- ControlConnection.send() must not block. Messages are serialized
33 + and queued for the next I/O pass.
34 +- ControlConnection.recv() returns all messages received since the
35 + last call (may be empty). Messages are dicts (already deserialized).
36 +- ControlConnection.alive returns False after mark_dead() or close().
37 +- ControlTransport.poll() returns newly accepted connections only —
38 + not previously returned ones. Each connection appears exactly once.
39 +- ControlTransport.connect() creates an outbound connection to a peer
40 + identified by peer_id (format: "host:port"). Raises on failure.
41 +- Messages may arrive from unknown peers (new inbound connections).
42 + The transport creates a ControlConnection and returns it from poll()
43 + with the first message(s) already in its recv() buffer.
44 +"""
45 +
46 +from __future__ import annotations
47 +
48 +from abc import ABC, abstractmethod
49 +from collections.abc import Sequence
50 +
51 +
52 +class ControlConnection(ABC):
53 + """Bidirectional message channel to a single remote peer.
54 +
55 + Lifecycle:
56 + 1. Created by ControlTransport (connect() or poll())
57 + 2. Used for send/recv by sessions
58 + 3. Marked dead on disconnect (mark_dead())
59 + 4. Cleaned up with close()
60 +
61 + Once mark_dead() is called, alive becomes False and the owning
62 + session should stop using this connection. close() releases
63 + underlying resources (sockets, monitors).
64 + """
65 +
66 + def __init__(self, peer_id: str) -> None:
67 + self.peer_id = peer_id
68 +
69 + @property
70 + @abstractmethod
71 + def alive(self) -> bool:
72 + """True if the connection is usable. False after mark_dead/close."""
73 + ...
74 +
75 + @abstractmethod
76 + def send(self, msg: dict) -> None:
77 + """Enqueue a message for delivery to the peer.
78 +
79 + Must not block. Serialization happens internally.
80 + Raises on closed connection.
81 + """
82 + ...
83 +
84 + @abstractmethod
85 + def recv(self) -> Sequence[dict]:
86 + """Drain and return all messages received since the last call.
87 +
88 + Returns an empty sequence if no messages are pending.
89 + The returned sequence is read-only — callers must not mutate it.
90 + Messages are dicts deserialized from the wire format.
91 + """
92 + ...
93 +
94 + @abstractmethod
95 + def mark_dead(self) -> None:
96 + """Mark this connection as dead (peer disconnected).
97 +
98 + After this call, alive returns False. The session should
99 + stop using this connection and the transport will clean it up.
100 + """
101 + ...
102 +
103 + @abstractmethod
104 + def close(self) -> None:
105 + """Release all resources (sockets, monitors).
106 +
107 + Idempotent. After close(), alive returns False.
108 + """
109 + ...
110 +
111 +
112 +class ControlTransport(ABC):
113 + """Manages peer connections and drives all control-plane I/O.
114 +
115 + Owns the listening socket and all active connections.
116 + The caller must invoke poll() periodically to process I/O.
117 +
118 + Lifecycle:
119 + 1. Constructed with a local identity and listen address
120 + 2. connect() to reach remote peers
121 + 3. poll() to accept inbound peers and process messages
122 + 4. close() to shut down
123 + """
124 +
125 + @abstractmethod
126 + def connect(self, peer_id: str) -> ControlConnection:
127 + """Create an outbound connection to a remote peer.
128 +
129 + Args:
130 + peer_id: Remote peer identity (format: "host:port").
131 +
132 + Returns:
133 + A new ControlConnection ready for send/recv.
134 +
135 + The connection's send queue is live immediately — messages
136 + sent before the remote peer's poll() will be buffered.
137 + """
138 + ...
139 +
140 + @abstractmethod
141 + def poll(self) -> Sequence[ControlConnection]:
142 + """Process all pending I/O and return newly accepted connections.
143 +
144 + This is the main I/O driver. Each call:
145 + - Receives messages from all connected peers (buffered in
146 + each connection's recv() queue)
147 + - Accepts new inbound peers and creates connections for them
148 + (first message already in recv() buffer)
149 + - Detects disconnections and marks connections dead
150 +
151 + Returns:
152 + Newly accepted inbound connections (not previously returned).
153 + The returned sequence is read-only — callers must not mutate
154 + it. The caller is responsible for creating sessions for
155 + these connections.
156 + """
157 + ...
158 +
159 + @abstractmethod
160 + def close(self) -> None:
161 + """Shut down the transport and all connections.
162 +
163 + Closes the listening socket and all active connections.
164 + Idempotent.
165 + """
166 + ...
vllm/v1/kv_offload/tiering/p2p/control/zmq.py new
+307
@@ -0,0 +1,307 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +ZMQ-based transport layer for P2P KV cache sharing.
5 +
6 +Provides ZmqConnection (per-peer messaging) and ZmqTransport (connection
7 +management). Message-content agnostic.
8 +"""
9 +
10 +from __future__ import annotations
11 +
12 +from collections.abc import Sequence
13 +from dataclasses import dataclass
14 +
15 +import msgspec
16 +import zmq
17 +import zmq.utils.monitor
18 +
19 +from vllm.logger import init_logger
20 +from vllm.v1.kv_offload.tiering.p2p.control.base import (
21 + ControlConnection,
22 + ControlTransport,
23 +)
24 +
25 +logger = init_logger(__name__)
26 +
27 +_HEARTBEAT_IVL_MS = 2000
28 +_HEARTBEAT_TIMEOUT_MS = 10000
29 +_HEARTBEAT_TTL_MS = 10000
30 +
31 +# Shared sentinels returned when there is nothing to report.
32 +_EMPTY_INBOX: tuple[dict, ...] = ()
33 +_EMPTY_NEW_CONNECTIONS: tuple[ControlConnection, ...] = ()
34 +
35 +
36 +def _tcp_addr(host: str, port: int | str) -> str:
37 + return f"tcp://{host}:{port}"
38 +
39 +
40 +def _apply_heartbeat(sock: zmq.Socket) -> None:
41 + sock.setsockopt(zmq.HEARTBEAT_IVL, _HEARTBEAT_IVL_MS)
42 + sock.setsockopt(zmq.HEARTBEAT_TIMEOUT, _HEARTBEAT_TIMEOUT_MS)
43 + sock.setsockopt(zmq.HEARTBEAT_TTL, _HEARTBEAT_TTL_MS)
44 +
45 +
46 +@dataclass
47 +class _Sockets:
48 + dealer: zmq.Socket
49 + monitor: zmq.Socket
50 +
51 +
52 +class ZmqConnection(ControlConnection):
53 + """Bidirectional message channel to a single remote peer."""
54 +
55 + def __init__(self, peer_id: str, sockets: _Sockets) -> None:
56 + super().__init__(peer_id)
57 + self._sockets = sockets
58 + self._closed = False
59 + self._inbox: list[dict] = []
60 +
61 + def send(self, msg: dict) -> None:
62 + """Send a msgpack-encoded message to this peer."""
63 + if self._closed:
64 + raise RuntimeError(
65 + f"ZmqConnection: send on closed connection to {self.peer_id}"
66 + )
67 + data = msgspec.msgpack.encode(msg)
68 + self._sockets.dealer.send(data)
69 +
70 + def recv(self) -> Sequence[dict]:
71 + """Drain and return all buffered incoming messages."""
72 + if not self._inbox:
73 + return _EMPTY_INBOX
74 + msgs = self._inbox
75 + self._inbox = []
76 + return msgs
77 +
78 + @property
79 + def alive(self) -> bool:
80 + return not self._closed
81 +
82 + def close(self) -> None:
83 + if self._closed:
84 + return
85 + self._closed = True
86 + logger.info("ZmqConnection: closing connection to %s", self.peer_id)
87 + self._sockets.monitor.close()
88 + self._sockets.dealer.close()
89 +
90 + def enqueue(self, msg: dict) -> None:
91 + """Buffer an incoming message."""
92 + self._inbox.append(msg)
93 +
94 + def mark_dead(self) -> None:
95 + """Mark connection as disconnected."""
96 + self._closed = True
97 +
98 + @property
99 + def monitor_socket(self) -> zmq.Socket:
100 + """Monitor socket for disconnect detection (used by ZmqTransport)."""
101 + return self._sockets.monitor
102 +
103 +
104 +class ZmqTransport(ControlTransport):
105 + """ZMQ implementation of ControlTransport.
106 +
107 + Manages a ROUTER socket for accepting connections and DEALER sockets
108 + for outbound connections. Message-content agnostic.
109 + """
110 +
111 + def __init__(self, local_id: str, host: str, port: int) -> None:
112 + self._local_id = local_id
113 + self._closed = False
114 +
115 + self._connections: dict[str, ZmqConnection] = {}
116 + self._pending_inbound: list[tuple[str, dict]] = []
117 +
118 + self._zmq_ctx = zmq.Context()
119 + self._router: zmq.Socket = self._zmq_ctx.socket(zmq.ROUTER)
120 + _apply_heartbeat(self._router)
121 + bind_addr = _tcp_addr(host, port)
122 + self._router.bind(bind_addr)
123 + logger.info("ZmqTransport %s: ROUTER bound on %s", self._local_id, bind_addr)
124 +
125 + # ------------------------------------------------------------------
126 + # ZmqConnection lifecycle
127 + # ------------------------------------------------------------------
128 +
129 + def connect(self, peer_id: str) -> ZmqConnection:
130 + """Open an outbound connection to a remote peer."""
131 + assert peer_id not in self._connections, (
132 + f"ZmqConnection to {peer_id} already exists"
133 + )
134 + logger.info(
135 + "ZmqTransport %s: opening OUTBOUND connection to %s",
136 + self._local_id,
137 + peer_id,
138 + )
139 + return self._open_connection(peer_id, direction="outbound")
140 +
141 + def poll(self) -> Sequence[ControlConnection]:
142 + """Process all pending I/O. Returns newly accepted connections.
143 +
144 + - Receives messages (buffered in each connection's inbox)
145 + - Creates connections for new inbound peers (connect msg in inbox)
146 + - Checks monitors for disconnections
147 + - Removes and closes dead connections
148 + """
149 + self._recv_router()
150 + self._check_monitors()
151 +
152 + # Create connections for new inbound peers
153 + new_connections: list[ControlConnection] | None = None
154 + for sender_id, msg in self._pending_inbound:
155 + conn = self._connections.get(sender_id)
156 + if conn is None:
157 + logger.info(
158 + "ZmqTransport %s: accepting INBOUND connection from %s",
159 + self._local_id,
160 + sender_id,
161 + )
162 + conn = self._open_connection(sender_id, direction="inbound")
163 + if new_connections is None:
164 + new_connections = []
165 + new_connections.append(conn)
166 + conn.enqueue(msg)
167 + self._pending_inbound.clear()
168 +
169 + # Remove dead connections
170 + for pid in [p for p, c in self._connections.items() if not c.alive]:
171 + self._connections.pop(pid).close()
172 +
173 + return (
174 + new_connections if new_connections is not None else _EMPTY_NEW_CONNECTIONS
175 + )
176 +
177 + def close(self) -> None:
178 + if self._closed:
179 + return
180 + self._closed = True
181 +
182 + for conn in self._connections.values():
183 + conn.close()
184 + self._connections.clear()
185 +
186 + self._router.setsockopt(zmq.LINGER, 0)
187 + self._router.close()
188 + self._zmq_ctx.destroy(linger=0)
189 +
190 + # ------------------------------------------------------------------
191 + # Internal
192 + # ------------------------------------------------------------------
193 +
194 + def _open_connection(
195 + self, peer_id: str, direction: str = "outbound"
196 + ) -> ZmqConnection:
197 + """Create a DEALER socket + monitor and register the connection."""
198 + host, port_str = peer_id.rsplit(":", 1)
199 + dealer_addr = _tcp_addr(host, port_str)
200 +
201 + logger.debug(
202 + "ZmqTransport %s: creating DEALER for %s peer %s -> %s",
203 + self._local_id,
204 + direction,
205 + peer_id,
206 + dealer_addr,
207 + )
208 +
209 + dealer = self._zmq_ctx.socket(zmq.DEALER)
210 + _apply_heartbeat(dealer)
211 + dealer.identity = self._local_id.encode()
212 +
213 + safe_id = peer_id.replace(":", "-").replace("/", "-")
214 + monitor_addr = f"inproc://p2p-monitor-{safe_id}"
215 + dealer.monitor(monitor_addr, zmq.EVENT_DISCONNECTED)
216 +
217 + monitor_sock = self._zmq_ctx.socket(zmq.PAIR)
218 + monitor_sock.connect(monitor_addr)
219 +
220 + dealer.connect(dealer_addr)
221 +
222 + sockets = _Sockets(dealer=dealer, monitor=monitor_sock)
223 + conn = ZmqConnection(peer_id, sockets)
224 + self._connections[peer_id] = conn
225 + logger.info(
226 + "ZmqTransport %s: %s connection established to %s (active connections: %d)",
227 + self._local_id,
228 + direction,
229 + peer_id,
230 + len(self._connections),
231 + )
232 + return conn
233 +
234 + def _recv_router(self) -> None:
235 + """Non-blocking: receive all pending messages from ROUTER."""
236 + while True:
237 + try:
238 + frames = self._router.recv_multipart(zmq.NOBLOCK)
239 + except zmq.Again:
240 + break
241 + except zmq.ZMQError as exc:
242 + logger.warning("ZmqTransport %s: recv error: %s", self._local_id, exc)
243 + break
244 +
245 + if len(frames) != 2:
246 + logger.warning(
247 + "ZmqTransport %s: dropping message with %d frames (expected 2)",
248 + self._local_id,
249 + len(frames),
250 + )
251 + continue
252 +
253 + identity, data = frames
254 + sender_id = identity.decode()
255 +
256 + logger.debug(
257 + "ZmqTransport %s: ROUTER recv from %s (%d bytes)",
258 + self._local_id,
259 + sender_id,
260 + len(data),
261 + )
262 +
263 + try:
264 + msg = msgspec.msgpack.decode(data)
265 + except Exception as exc:
266 + logger.warning(
267 + "ZmqTransport %s: failed to decode message from %s: %s",
268 + self._local_id,
269 + sender_id,
270 + exc,
271 + )
272 + continue
273 +
274 + conn = self._connections.get(sender_id)
275 +
276 + if conn is not None:
277 + conn.enqueue(msg)
278 + else:
279 + self._pending_inbound.append((sender_id, msg))
280 +
281 + def _check_monitors(self) -> None:
282 + """Non-blocking: check all monitor sockets for disconnection."""
283 + for conn in self._connections.values():
284 + if not conn.alive:
285 + continue
286 + try:
287 + event = zmq.utils.monitor.recv_monitor_message(
288 + conn.monitor_socket, zmq.NOBLOCK
289 + )
290 + except zmq.Again:
291 + continue
292 + except zmq.ZMQError as exc:
293 + logger.warning(
294 + "ZmqTransport %s: monitor error for peer %s: %s",
295 + self._local_id,
296 + conn.peer_id,
297 + exc,
298 + )
299 + continue
300 +
301 + if event["event"] == zmq.EVENT_DISCONNECTED:
302 + logger.debug(
303 + "ZmqTransport %s: peer %s disconnected",
304 + self._local_id,
305 + conn.peer_id,
306 + )
307 + conn.mark_dead()
vllm/v1/kv_offload/tiering/p2p/data/__init__.py new
+10
@@ -0,0 +1,10 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from vllm.v1.kv_offload.tiering.p2p.data.base import DataTransport, PollResult
4 +from vllm.v1.kv_offload.tiering.p2p.data.nixl import NixlTransport
5 +
6 +__all__ = [
7 + "DataTransport",
8 + "NixlTransport",
9 + "PollResult",
10 +]
vllm/v1/kv_offload/tiering/p2p/data/base.py new
+259
@@ -0,0 +1,259 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Abstract base class for the P2P data-plane transport.
4 +
5 +The data plane handles RDMA (or similar) block transfers between peers.
6 +It is independent of the control plane — the control plane establishes
7 +who can talk to whom, then the data plane moves blocks at wire speed.
8 +
9 +Architecture
10 +------------
11 +
12 + DataTransport (one per node)
13 + ├── owns the local KV block memory region
14 + ├── registers remote peers (add_remote_peer / remove_remote_peer)
15 + ├── submits block writes (write_blocks → transfer_id)
16 + ├── polls for completion (poll → done/failed IDs)
17 + └── cancels inflight transfers (cancel)
18 +
19 +Memory model
20 +------------
21 +
22 +The local node exposes a contiguous block region:
23 +
24 + base_addr ──► ┌─────────────┐ block 0
25 + ├─────────────┤ block 1
26 + ├─────────────┤ ...
27 + └─────────────┘ block (num_blocks - 1)
28 +
29 + Each block is block_len bytes.
30 +
31 +Remote peers expose the same layout. write_blocks() copies local
32 +blocks to a remote peer's block region by index:
33 +
34 + write_blocks("peer:1", local_idxs=[0, 3], remote_idxs=[5, 7])
35 + → writes local block 0 → remote block 5
36 + local block 3 → remote block 7
37 +
38 +Transfer lifecycle
39 +------------------
40 +
41 +1. Register remote peer: add_remote_peer(peer_id, metadata, ...)
42 + - Provides the remote memory layout so transfers can target it
43 +2. Submit: write_blocks(peer_id, local_idxs, remote_idxs) → int
44 + - Returns a transfer_id (opaque int) for tracking
45 + - Returns None if peer not registered or submission fails
46 +3. Poll: poll() → PollResult(done=[...], failed=[...])
47 + - Returns transfer_ids that completed or failed since last poll
48 + - Completed transfers are automatically cleaned up
49 +4. Cancel: cancel(transfer_ids, mode="immediate" | "wait")
50 + - Best-effort cancellation of inflight transfers
51 + - mode="wait" returns ids still in PROC/PEND so the caller can poll
52 + them to completion
53 +
54 +Implementor contracts
55 +---------------------
56 +
57 +- write_blocks() must not block. The transfer runs asynchronously.
58 +- poll() must be called periodically. It drives completion checking.
59 +- transfer_ids are unique across the lifetime of the transport.
60 +- add_remote_peer() must be called before write_blocks() to that peer.
61 +- get_agent_metadata() returns opaque bytes that the remote peer
62 + needs to call add_remote_peer() (e.g., RDMA connection info).
63 +- config_fingerprint is a content hash of the model configuration.
64 + Peers with different fingerprints are incompatible and must not
65 + exchange blocks (validated during the control-plane handshake).
66 +- close() releases all resources (memory registrations, handles).
67 + After close(), no other methods may be called.
68 +
69 +Threading model: no background threads. All I/O driven by poll().
70 +"""
71 +
72 +from __future__ import annotations
73 +
74 +import ctypes
75 +import hashlib
76 +import json
77 +from abc import ABC, abstractmethod
78 +from collections.abc import Iterable, Sequence
79 +from typing import Literal, NamedTuple
80 +
81 +CancelMode = Literal["immediate", "wait"]
82 +
83 +
84 +class PollResult(NamedTuple):
85 + """Result of polling inflight transfers.
86 +
87 + Attributes:
88 + done: Transfer IDs that completed successfully.
89 + failed: Transfer IDs that failed (error, timeout, etc.).
90 + """
91 +
92 + done: Sequence[int]
93 + failed: Sequence[int]
94 +
95 +
96 +class DataTransport(ABC):
97 + """Abstract data-plane transport for RDMA-style block transfers.
98 +
99 + Owns the local KV block memory region and manages transfers to/from
100 + registered remote peers.
101 +
102 + Construction:
103 + view: A 2D memoryview (num_blocks × block_len bytes) over the
104 + local KV cache block storage.
105 + config_fields: Dict of model config values used to compute the
106 + compatibility fingerprint. None → empty fingerprint
107 + (compatible with any peer).
108 + """
109 +
110 + def __init__(self, view: memoryview, config_fields: dict | None = None) -> None:
111 + assert view.shape is not None
112 + self._view = view
113 + self._base_addr = ctypes.addressof(ctypes.c_char.from_buffer(view))
114 + self._num_blocks = view.shape[0]
115 + self._block_len = view.shape[1]
116 + self._config_fingerprint = self._compute_fingerprint(config_fields)
117 +
118 + @property
119 + def base_addr(self) -> int:
120 + """Base address of the local block memory region."""
121 + return self._base_addr
122 +
123 + @property
124 + def num_blocks(self) -> int:
125 + """Number of blocks in the local region."""
126 + return self._num_blocks
127 +
128 + @property
129 + def block_len(self) -> int:
130 + """Size of each block in bytes."""
131 + return self._block_len
132 +
133 + @property
134 + def config_fingerprint(self) -> str:
135 + """Content-hash of the model configuration (hex string).
136 +
137 + Peers must have matching fingerprints to exchange blocks.
138 + Empty string means no fingerprint (always compatible).
139 + """
140 + return self._config_fingerprint
141 +
142 + @staticmethod
143 + def _compute_fingerprint(config_fields: dict | None) -> str:
144 + if not config_fields:
145 + return ""
146 + canonical = json.dumps(config_fields, sort_keys=True, separators=(",", ":"))
147 + return hashlib.sha256(canonical.encode()).hexdigest()[:16]
148 +
149 + @abstractmethod
150 + def get_agent_metadata(self) -> bytes:
151 + """Return opaque metadata needed by remote peers to connect.
152 +
153 + The returned bytes are sent during the control-plane handshake
154 + and passed to the remote peer's add_remote_peer().
155 + """
156 + ...
157 +
158 + @abstractmethod
159 + def add_remote_peer(
160 + self,
161 + peer_id: str,
162 + agent_metadata: bytes,
163 + base_addr: int,
164 + num_blocks: int,
165 + block_len: int,
166 + ) -> None:
167 + """Register a remote peer for block transfers.
168 +
169 + Must be called before write_blocks() to this peer.
170 +
171 + Args:
172 + peer_id: Unique identifier for the remote peer.
173 + agent_metadata: Opaque bytes from the peer's get_agent_metadata().
174 + base_addr: Base address of the peer's block memory region.
175 + num_blocks: Number of blocks in the peer's region.
176 + block_len: Size of each block (must match local block_len).
177 + """
178 + ...
179 +
180 + @abstractmethod
181 + def remove_remote_peer(self, peer_id: str) -> None:
182 + """Unregister a remote peer and release associated resources.
183 +
184 + Inflight transfers to this peer should be cancelled first.
185 + """
186 + ...
187 +
188 + @abstractmethod
189 + def write_blocks(
190 + self,
191 + peer_id: str,
192 + local_idxs: list[int],
193 + remote_idxs: list[int],
194 + ) -> int | None:
195 + """Submit a WRITE transfer: local blocks → remote peer's blocks.
196 +
197 + Args:
198 + peer_id: Target peer (must be registered via add_remote_peer).
199 + local_idxs: Indexes of local blocks to read from.
200 + remote_idxs: Indexes of remote blocks to write to.
201 + Must be same length as local_idxs.
202 +
203 + Returns:
204 + A unique transfer_id (int) to track this transfer, or
205 + None if the peer is not registered or submission failed.
206 + """
207 + ...
208 +
209 + @abstractmethod
210 + def poll(self) -> PollResult:
211 + """Poll all inflight transfers for completion.
212 +
213 + Returns:
214 + PollResult with lists of completed and failed transfer_ids.
215 + Completed/failed transfers are removed from the inflight set.
216 +
217 + Must be called periodically to drive progress checking.
218 + """
219 + ...
220 +
221 + @abstractmethod
222 + def cancel(
223 + self,
224 + transfer_ids: Iterable[int],
225 + mode: CancelMode = "immediate",
226 + ) -> list[int]:
227 + """Cancel inflight transfers by their IDs.
228 +
229 + Best-effort: transfers that already completed are ignored.
230 +
231 + Args:
232 + transfer_ids: IDs to cancel. Unknown IDs are ignored.
233 + mode:
234 + "immediate" (default): pop and release each handle and
235 + return []. Matches the legacy fire-and-forget
236 + behavior — the caller does not wait for the
237 + underlying transfer to drain.
238 + "wait": attempt to release each handle. If the release
239 + cannot complete because the transfer is still
240 + PROC/PEND, the entry stays in the inflight set and
241 + its id is included in the returned list. The
242 + caller is expected to keep calling poll() until
243 + every returned id surfaces in done/failed.
244 +
245 + Returns:
246 + For mode="wait", the subset of *transfer_ids* still
247 + tracked as inflight after the cancel attempt. For
248 + mode="immediate", always [].
249 + """
250 + ...
251 +
252 + @abstractmethod
253 + def close(self) -> None:
254 + """Release all resources (registrations, handles, memory).
255 +
256 + Cancels any remaining inflight transfers. Idempotent.
257 + After close(), no other methods may be called.
258 + """
259 + ...
vllm/v1/kv_offload/tiering/p2p/data/nixl.py new
+297
@@ -0,0 +1,297 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +NixlTransport: Data-plane transport for RDMA-based KV block transfers via NIXL.
5 +"""
6 +
7 +from __future__ import annotations
8 +
9 +import itertools
10 +from collections.abc import Iterable
11 +from typing import Any
12 +
13 +from vllm.distributed.nixl_utils import NixlWrapper as _NixlAgent
14 +from vllm.distributed.nixl_utils import nixl_agent_config as _NixlAgentConfig
15 +from vllm.logger import init_logger
16 +from vllm.v1.kv_offload.tiering.p2p.data.base import (
17 + CancelMode,
18 + DataTransport,
19 + PollResult,
20 +)
21 +
22 +logger = init_logger(__name__)
23 +
24 +# Shared sentinel returned by poll() in the steady state (no inflight, or
25 +# no transfer changed state since the last poll). Tuples make it immutable;
26 +# callers only iterate / membership-test / equality-check.
27 +_EMPTY_POLL_RESULT: PollResult = PollResult(done=(), failed=())
28 +
29 +
30 +class NixlTransport(DataTransport):
31 + """Manages a NIXL agent, memory registration, and block transfers.
32 +
33 + Wraps the NIXL C library behind a Python interface so the rest of the
34 + P2P tier code never touches NIXL types directly. Tracks inflight
35 + handles internally and returns completed/failed tags on poll.
36 + """
37 +
38 + def __init__(
39 + self,
40 + local_id: str,
41 + view: memoryview,
42 + config_fields: dict | None = None,
43 + backends: list[str] | None = None,
44 + num_threads: int = 4,
45 + ) -> None:
46 + super().__init__(view, config_fields=config_fields)
47 + self._local_id = local_id
48 + self._backends = list(backends) if backends else ["UCX"]
49 + self._num_threads = num_threads
50 + self._agent: Any = None
51 + self._reg: Any = None
52 + self._local_dlist: Any = None
53 + self._remote_dlists: dict[str, object] = {}
54 + self._peer_nixl_names: dict[str, str] = {}
55 + self._inflight: dict[int, object] = {} # transfer_id → handle
56 + self._next_id = itertools.count()
57 +
58 + self._init(view)
59 +
60 + @property
61 + def available(self) -> bool:
62 + return self._agent is not None
63 +
64 + def _init(self, view: memoryview) -> None:
65 + if _NixlAgent is None:
66 + return
67 +
68 + non_ucx_backends = [b for b in self._backends if b != "UCX"]
69 + if non_ucx_backends:
70 + cfg = _NixlAgentConfig(backends=self._backends, capture_telemetry=True)
71 + logger.info(
72 + "NixlTransport %s: NIXL backends=%s",
73 + self._local_id,
74 + self._backends,
75 + )
76 + else:
77 + cfg = _NixlAgentConfig(
78 + num_threads=self._num_threads, capture_telemetry=True
79 + )
80 + logger.info(
81 + "NixlTransport %s: NIXL backends=[UCX] num_threads=%d",
82 + self._local_id,
83 + self._num_threads,
84 + )
85 + self._agent = _NixlAgent(self._local_id, cfg)
86 +
87 + total_size = self._num_blocks * self._block_len
88 + reg_descs = [(self._base_addr, total_size, 0, "")]
89 + self._reg = self._agent.register_memory(reg_descs, mem_type="DRAM")
90 +
91 + block_tuples = [
92 + (self._base_addr + i * self._block_len, self._block_len, 0)
93 + for i in range(self._num_blocks)
94 + ]
95 + xfer_dlist = self._agent.get_xfer_descs(block_tuples, mem_type="DRAM")
96 + self._local_dlist = self._agent.prep_xfer_dlist("NIXL_INIT_AGENT", xfer_dlist)
97 + logger.info(
98 + "NixlTransport %s: registered %d blocks", self._local_id, self._num_blocks
99 + )
100 +
101 + def get_agent_metadata(self) -> bytes:
102 + assert self._agent is not None
103 + return self._agent.get_agent_metadata()
104 +
105 + # ------------------------------------------------------------------
106 + # Peer management
107 + # ------------------------------------------------------------------
108 +
109 + def add_remote_peer(
110 + self,
111 + peer_id: str,
112 + agent_metadata: bytes,
113 + base_addr: int,
114 + num_blocks: int,
115 + block_len: int,
116 + ) -> None:
117 + nixl_name = self._agent.add_remote_agent(agent_metadata)
118 + block_descs = [
119 + (base_addr + i * block_len, block_len, 0) for i in range(num_blocks)
120 + ]
121 + xfer_dlist = self._agent.get_xfer_descs(block_descs, mem_type="DRAM")
122 + remote_dlist = self._agent.prep_xfer_dlist(nixl_name, xfer_dlist)
123 + self._peer_nixl_names[peer_id] = nixl_name
124 + self._remote_dlists[peer_id] = remote_dlist
125 +
126 + def remove_remote_peer(self, peer_id: str) -> None:
127 + nixl_name = self._peer_nixl_names.pop(peer_id, None)
128 + dlist = self._remote_dlists.pop(peer_id, None)
129 + if self._agent is not None:
130 + if dlist is not None:
131 + self._agent.release_dlist_handle(dlist)
132 + if nixl_name:
133 + self._agent.remove_remote_agent(nixl_name)
134 +
135 + # ------------------------------------------------------------------
136 + # Transfer submission and polling
137 + # ------------------------------------------------------------------
138 +
139 + def write_blocks(
140 + self,
141 + peer_id: str,
142 + local_idxs: list[int],
143 + remote_idxs: list[int],
144 + ) -> int | None:
145 + """Submit a WRITE transfer to *peer_id*.
146 +
147 + Returns a transfer ID, or None if the peer is not registered.
148 + The ID is returned via poll() when the transfer completes or fails.
149 + """
150 + remote_dlist = self._remote_dlists.get(peer_id)
151 + if remote_dlist is None:
152 + logger.warning(
153 + "NixlTransport %s: write_blocks NO REMOTE DLIST for peer=%s "
154 + "(known peers=%s)",
155 + self._local_id,
156 + peer_id,
157 + list(self._remote_dlists.keys()),
158 + )
159 + return None
160 + logger.debug(
161 + "NixlTransport %s: write_blocks NIXL.transfer peer=%s blocks=%d",
162 + self._local_id,
163 + peer_id,
164 + len(local_idxs),
165 + )
166 + handle = self._agent.make_prepped_xfer(
167 + "WRITE",
168 + self._local_dlist,
169 + local_idxs,
170 + remote_dlist,
171 + remote_idxs,
172 + )
173 + self._agent.transfer(handle)
174 + transfer_id = next(self._next_id)
175 + self._inflight[transfer_id] = handle
176 + return transfer_id
177 +
178 + def poll(self) -> PollResult:
179 + """Poll all inflight transfers.
180 +
181 + Returns PollResult(done=..., failed=...) with transfer IDs.
182 + Completed handles are released automatically.
183 + """
184 + if not self._inflight:
185 + return _EMPTY_POLL_RESULT
186 +
187 + done_ids: list[int] | None = None
188 + failed_ids: list[int] | None = None
189 +
190 + for transfer_id, handle in self._inflight.items():
191 + try:
192 + state = self._agent.check_xfer_state(handle)
193 + except Exception as exc:
194 + logger.warning(
195 + "NixlTransport %s: check_xfer_state failed for transfer_id=%d: %s",
196 + self._local_id,
197 + transfer_id,
198 + exc,
199 + )
200 + continue
201 + if state == "DONE":
202 + if done_ids is None:
203 + done_ids = []
204 + done_ids.append(transfer_id)
205 + elif state not in ("PROC", "PEND"):
206 + if failed_ids is None:
207 + failed_ids = []
208 + failed_ids.append(transfer_id)
209 +
210 + if done_ids is None and failed_ids is None:
211 + return _EMPTY_POLL_RESULT
212 +
213 + handles_to_release = []
214 + for tid in done_ids or ():
215 + handles_to_release.append(self._inflight.pop(tid))
216 + for tid in failed_ids or ():
217 + handles_to_release.append(self._inflight.pop(tid))
218 + self._release_handles(handles_to_release)
219 +
220 + return PollResult(
221 + done=done_ids if done_ids is not None else _EMPTY_POLL_RESULT.done,
222 + failed=failed_ids if failed_ids is not None else _EMPTY_POLL_RESULT.failed,
223 + )
224 +
225 + def cancel(
226 + self,
227 + transfer_ids: Iterable[int],
228 + mode: CancelMode = "immediate",
229 + ) -> list[int]:
230 + """Cancel inflight transfers by their IDs.
231 +
232 + See ``DataTransport.cancel`` for the contract. In "wait" mode,
233 + transfers whose ``release_xfer_handle`` raises (NIXL could not
234 + complete the abort because the backend is still draining) stay
235 + in ``self._inflight`` so a later ``poll()`` will observe them.
236 + """
237 + if mode == "immediate":
238 + handles = [
239 + self._inflight.pop(tid) for tid in transfer_ids if tid in self._inflight
240 + ]
241 + self._release_handles(handles)
242 + return []
243 +
244 + still_inflight: list[int] = []
245 + for tid in transfer_ids:
246 + handle = self._inflight.get(tid)
247 + if handle is None:
248 + continue
249 + try:
250 + self._agent.release_xfer_handle(handle)
251 + except Exception as exc:
252 + logger.debug(
253 + "NixlTransport %s: cancel pending for transfer_id=%d: %s",
254 + self._local_id,
255 + tid,
256 + exc,
257 + )
258 + still_inflight.append(tid)
259 + continue
260 + del self._inflight[tid]
261 + return still_inflight
262 +
263 + # ------------------------------------------------------------------
264 + # Lifecycle
265 + # ------------------------------------------------------------------
266 +
267 + def close(self) -> None:
268 + if self._agent is None:
269 + return
270 + self._release_handles(list(self._inflight.values()))
271 + self._inflight.clear()
272 + for peer_id in list(self._remote_dlists):
273 + self.remove_remote_peer(peer_id)
274 + if self._local_dlist is not None:
275 + self._agent.release_dlist_handle(self._local_dlist)
276 + self._local_dlist = None
277 + if self._reg is not None:
278 + self._agent.deregister_memory(self._reg)
279 + self._reg = None
280 + self._agent = None
281 +
282 + # ------------------------------------------------------------------
283 + # Internal
284 + # ------------------------------------------------------------------
285 +
286 + def _release_handles(self, handles: list[object]) -> None:
287 + if self._agent is None:
288 + return
289 + for handle in handles:
290 + try:
291 + self._agent.release_xfer_handle(handle)
292 + except Exception as exc:
293 + logger.warning(
294 + "NixlTransport %s: release_xfer_handle failed: %s",
295 + self._local_id,
296 + exc,
297 + )
vllm/v1/kv_offload/tiering/p2p/manager.py new
+664
@@ -0,0 +1,664 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +P2PSecondaryTierManager: Secondary tier for P2P KV cache sharing.
5 +
6 +Owns transports and a single bidirectional P2PSession per remote peer.
7 +"""
8 +
9 +from __future__ import annotations
10 +
11 +import time
12 +from collections.abc import Iterable, Sequence
13 +from dataclasses import dataclass, field
14 +from typing import TYPE_CHECKING
15 +
16 +from typing_extensions import override
17 +
18 +from vllm.logger import init_logger
19 +from vllm.v1.kv_offload.base import (
20 + LookupResult,
21 + OffloadKey,
22 + ReqContext,
23 + RequestOffloadingContext,
24 +)
25 +from vllm.v1.kv_offload.file_mapper import FileMapper
26 +from vllm.v1.kv_offload.tiering.base import (
27 + JobMetadata,
28 + JobResult,
29 + SecondaryTierManager,
30 +)
31 +from vllm.v1.kv_offload.tiering.p2p.control import ControlTransport, ZmqTransport
32 +from vllm.v1.kv_offload.tiering.p2p.data import DataTransport, NixlTransport
33 +from vllm.v1.kv_offload.tiering.p2p.session import P2PSession
34 +
35 +if TYPE_CHECKING:
36 + from vllm.v1.kv_offload.base import OffloadingSpec
37 + from vllm.v1.kv_offload.tiering.p2p.control.base import ControlConnection
38 +
39 +logger = init_logger(__name__)
40 +
41 +# Reap unbound store batches that have been parked without a FetchMsg
42 +# binding them to a session for longer than this. Protects against the
43 +# prefiller buffering blocks for a decoder that never asks (decoder died,
44 +# network partition, lost kv_request_id). Must be longer than the per-store
45 +# deadline so the store-timeout path fires first for individual jobs.
46 +_UNBOUND_STORE_TIMEOUT_S = 60.0
47 +
48 +# Time we wait during shutdown for inflight transfers to drain via
49 +# cancel(mode="wait") before falling back to mode="immediate". Bounded
50 +# so a wedged peer can't hang shutdown.
51 +_SHUTDOWN_DRAIN_TIMEOUT_S = 3.0
52 +
53 +# Sleep between iterations of the bounded drain loops in drain_jobs() and
54 +# _drain_inflight_for_shutdown(). Short enough to keep latency low, long
55 +# enough to avoid busy-spinning the scheduler thread.
56 +_DRAIN_SLEEP_S = 0.001
57 +
58 +
59 +def _prefill_params(kv_params: dict | None) -> dict | None:
60 + """Return the ``prefill`` sub-dict, or None if absent.
61 +
62 + Set on decoder requests; carries kv_request_id, remote_host, remote_port.
63 + """
64 + if not kv_params:
65 + return None
66 + return kv_params.get("prefill")
67 +
68 +
69 +def _decode_params(kv_params: dict | None) -> dict | None:
70 + """Return the ``decode`` sub-dict, or None if absent.
71 +
72 + Set on prefiller requests; carries kv_request_id.
73 + """
74 + if not kv_params:
75 + return None
76 + return kv_params.get("decode")
77 +
78 +
79 +@dataclass
80 +class _UnboundStoreBatch:
81 + """A submit_store batch parked at the manager before any peer has fetched.
82 +
83 + Indexed by kv_request_id only — the prefiller no longer learns the peer
84 + identity at store time. When a FetchMsg(kv_request_id) arrives on some
85 + session, the manager binds the kv_request_id to that session and replays
86 + every parked batch into ServerRole via session.add_stored_blocks.
87 + """
88 +
89 + job_id: int
90 + keys: list[OffloadKey]
91 + block_ids: Sequence[int]
92 + submitted_at: float = field(default_factory=time.monotonic)
93 +
94 +
95 +class P2PSecondaryTierManager(SecondaryTierManager):
96 + """Secondary tier for P2P KV cache sharing.
97 +
98 + A single P2PSession per remote peer handles both client-role (loading
99 + blocks from the peer) and server-role (serving blocks to the peer)
100 + over the same control connection.
101 +
102 + Single-threaded: every public method runs on the scheduler thread, and
103 + the engine drives polling via ``get_finished_jobs()`` once per step.
104 + ``has_pending_work()`` keeps the engine ticking so the control transport
105 + and existing sessions are polled even when no requests are scheduled.
106 + """
107 +
108 + def __init__(
109 + self,
110 + offloading_spec: OffloadingSpec,
111 + primary_kv_view: memoryview,
112 + tier_type: str = "p2p",
113 + host: str = "0.0.0.0",
114 + port: int = 7777,
115 + backends: list[str] | None = None,
116 + num_threads: int = 4,
117 + **kwargs,
118 + ) -> None:
119 + """Initialize the P2P secondary tier manager.
120 +
121 + All keyword arguments after ``primary_kv_view`` come from the
122 + ``secondary_tiers`` entry in ``kv_connector_extra_config``. See
123 + ``docs/features/kv_offloading_usage.md`` for the user-facing
124 + configuration reference.
125 +
126 + Args:
127 + offloading_spec: Owning ``OffloadingSpec`` (provides
128 + ``vllm_config`` and the offloaded block layout).
129 + primary_kv_view: Memoryview over the CPU primary tier; the
130 + NIXL agent registers this region for RDMA transfers.
131 + tier_type: Tier identifier (defaults to ``"p2p"``).
132 + host: Address the ZMQ control socket binds to.
133 + port: Port for the ZMQ control socket. Must be reachable
134 + from peers.
135 + backends: NIXL transport backends (e.g. ``["UCX"]``,
136 + ``["MOONCAKE"]``, ``["LIBFABRIC"]``). Defaults to
137 + ``["UCX"]``. When any non-UCX backend is requested, the
138 + NIXL agent is initialized with ``backends=...``;
139 + otherwise it falls back to a UCX-only agent with
140 + ``num_threads`` threads.
141 + num_threads: NIXL agent worker threads for the UCX-only
142 + branch. Ignored when ``backends`` contains a non-UCX
143 + entry.
144 + **kwargs: Reserved for future tier-specific options.
145 + """
146 + super().__init__(offloading_spec, primary_kv_view, tier_type)
147 + port = int(port)
148 + self._local_id = f"{host}:{port}"
149 +
150 + config_fields = FileMapper.from_offloading_spec(
151 + root_dir="",
152 + offloading_spec=offloading_spec,
153 + gpu_blocks_per_file=offloading_spec.block_size_factor,
154 + parallel_agnostic=True,
155 + ).get_run_config()
156 + self._data: DataTransport = NixlTransport(
157 + self._local_id,
158 + primary_kv_view,
159 + config_fields=config_fields,
160 + backends=backends,
161 + num_threads=int(num_threads),
162 + )
163 + self._control: ControlTransport = ZmqTransport(self._local_id, host, port)
164 +
165 + self._sessions: dict[str, P2PSession] = {}
166 + # kv_request_id → session, set when the bound session has received
167 + # FetchMsg for that id. submit_store after binding routes directly
168 + # to the session; before binding, batches are parked in
169 + # _unbound_stores below. Stays in sync with _sessions: entries
170 + # pointing to a reaped session are purged in _reap_dead_sessions.
171 + self._kv_to_session: dict[str, P2PSession] = {}
172 + # kv_request_id → list of batches submit_store'd before any peer
173 + # asked for that id. Drained into a session by _on_session_fetch
174 + # when the corresponding FetchMsg arrives, or surfaced as failures
175 + # by _reap_unbound_stores after _UNBOUND_STORE_TIMEOUT_S.
176 + self._unbound_stores: dict[str, list[_UnboundStoreBatch]] = {}
177 +
178 + self._finished_jobs: list[JobResult] = []
179 + # kv_request_ids that hit a transport/session failure; On load lookup()
180 + # rejects them so the request falls back to local prefill.
181 + self._failed_req_ids: set[str] = set()
182 +
183 + # ------------------------------------------------------------------
184 + # SecondaryTierManager interface
185 + # ------------------------------------------------------------------
186 +
187 + @override
188 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
189 + prefill = _prefill_params(req_context.kv_transfer_params)
190 + if (
191 + not prefill
192 + or not prefill.get("remote_host")
193 + or not prefill.get("remote_port")
194 + or not prefill.get("kv_request_id")
195 + ):
196 + return LookupResult.MISS
197 +
198 + kv_request_id = prefill["kv_request_id"]
199 + if kv_request_id in self._failed_req_ids:
200 + return LookupResult.MISS
201 + return LookupResult.HIT
202 +
203 + @override
204 + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
205 + """Open the outbound session toward the producer if needed.
206 +
207 + On the decoder side (``prefill`` set), open a session toward the
208 + producer at remote_host:remote_port so submit_load can issue
209 + FetchMsg as soon as it fires. On the prefiller side, sessions
210 + are created when the consumer's inbound connection arrives in
211 + _accept_new_peers — submit_store no longer pre-creates anything.
212 + """
213 + prefill = _prefill_params(req_context.kv_transfer_params)
214 + if prefill:
215 + peer_id = self._remote_id_from_params(prefill)
216 + if peer_id:
217 + self._get_or_create_session(peer_id)
218 + return RequestOffloadingContext()
219 +
220 + @override
221 + def on_request_finished(self, req_context: ReqContext) -> None:
222 + """Cancels pending loads and prunes session-scoped state.
223 +
224 + Decoder side (``prefill`` set): looks up the session by peer_id
225 + because the producer's address is what addresses the client-role
226 + load to cancel. Prefiller side (``decode`` set): looks up via
227 + kv_request_id because peer_id is no longer carried on store-time
228 + kv_transfer_params; if a session has bound the id, finish it. If
229 + no session has bound the id yet, this is a no-op: parked batches
230 + in `_unbound_stores` are left in place and cleaned up only by
231 + `_reap_unbound_stores` after `_UNBOUND_STORE_TIMEOUT_S`.
232 + """
233 + kv_params = req_context.kv_transfer_params
234 + if not kv_params:
235 + return
236 + prefill = _prefill_params(kv_params)
237 + decode = _decode_params(kv_params)
238 + kv_request_id = (prefill or decode or {}).get("kv_request_id")
239 + if not kv_request_id:
240 + return
241 + self._failed_req_ids.discard(kv_request_id)
242 +
243 + if prefill:
244 + peer_id = self._remote_id_from_params(prefill)
245 + if peer_id:
246 + session = self._sessions.get(peer_id)
247 + if session is not None:
248 + session.finish_request(kv_request_id)
249 + return
250 +
251 + # Prefiller-side finish: identify the session via kv_request_id.
252 + session = self._kv_to_session.pop(kv_request_id, None)
253 + if session is not None:
254 + session.finish_request(kv_request_id)
255 + return
256 +
257 + @override
258 + def submit_store(self, job_metadata: JobMetadata) -> None:
259 + job_id = job_metadata.job_id
260 + keys = list(job_metadata.keys)
261 + block_ids = job_metadata.block_ids
262 +
263 + assert len(keys) == len(block_ids)
264 +
265 + kv_params = job_metadata.req_context.kv_transfer_params
266 + decode = _decode_params(kv_params)
267 + logger.debug(
268 + "P2P %s: submit_store ENTRY job_id=%d blocks=%d decode=%s kv_request_id=%s",
269 + self._local_id,
270 + job_id,
271 + len(block_ids),
272 + decode is not None,
273 + (decode or {}).get("kv_request_id"),
274 + )
275 + # Absent ``decode`` block => not a remote-decode request: succeed
276 + # locally without parking. An empty/malformed dict is still a
277 + # remote-decode signal and must fail the missing-id check below.
278 + if decode is None:
279 + self._finished_jobs.append(JobResult(job_id=job_id, success=True))
280 + return
281 +
282 + kv_request_id = decode.get("kv_request_id")
283 + if not kv_request_id:
284 + logger.warning(
285 + "P2P %s: submit_store missing kv_request_id",
286 + self._local_id,
287 + )
288 + self._finished_jobs.append(JobResult(job_id=job_id, success=False))
289 + return
290 +
291 + # Fast path: a session has already received FetchMsg for this id,
292 + # so we can route the batch straight into its ServerRole.
293 + session = self._kv_to_session.get(kv_request_id)
294 + if session is not None:
295 + session.add_stored_blocks(kv_request_id, keys, block_ids, job_id)
296 + return
297 +
298 + # No session bound yet — park the batch keyed by kv_request_id.
299 + # _on_session_fetch drains it on the first FetchMsg; if no peer
300 + # ever asks, _reap_unbound_stores surfaces the job as failed.
301 + self._unbound_stores.setdefault(kv_request_id, []).append(
302 + _UnboundStoreBatch(
303 + job_id=job_id,
304 + keys=keys,
305 + block_ids=block_ids,
306 + )
307 + )
308 + logger.debug(
309 + "P2P %s: parked submit_store kv_request_id=%s job_id=%d blocks=%d",
310 + self._local_id,
311 + kv_request_id,
312 + job_id,
313 + len(block_ids),
314 + )
315 +
316 + @override
317 + def submit_load(self, job_metadata: JobMetadata) -> None:
318 + job_id = job_metadata.job_id
319 + keys = list(job_metadata.keys)
320 + block_ids = job_metadata.block_ids
321 +
322 + prefill = _prefill_params(job_metadata.req_context.kv_transfer_params)
323 + logger.debug(
324 + "P2P %s: submit_load ENTRY job_id=%d blocks=%d kv_request_id=%s peer=%s",
325 + self._local_id,
326 + job_id,
327 + len(block_ids),
328 + (prefill or {}).get("kv_request_id"),
329 + self._remote_id_from_params(prefill or {}),
330 + )
331 + if (
332 + not prefill
333 + or not prefill.get("remote_host")
334 + or not prefill.get("remote_port")
335 + or not prefill.get("kv_request_id")
336 + ):
337 + logger.debug(
338 + "P2P %s: submit_load job_id=%d FAILED missing prefill params",
339 + self._local_id,
340 + job_id,
341 + )
342 + self._finished_jobs.append(JobResult(job_id=job_id, success=False))
343 + return
344 +
345 + kv_request_id = prefill["kv_request_id"]
346 + peer_id = self._remote_id_from_params(prefill)
347 + assert peer_id is not None # guaranteed by prefill checks above
348 +
349 + if not keys:
350 + logger.debug(
351 + "P2P %s: submit_load job_id=%d short-circuit success (no keys)",
352 + self._local_id,
353 + job_id,
354 + )
355 + self._finished_jobs.append(JobResult(job_id=job_id, success=True))
356 + return
357 +
358 + session = self._sessions.get(peer_id)
359 + if session is None:
360 + logger.warning(
361 + "P2P %s: submit_load job_id=%d NO SESSION for peer=%s",
362 + self._local_id,
363 + job_id,
364 + peer_id,
365 + )
366 + self._finished_jobs.append(JobResult(job_id=job_id, success=False))
367 + self._failed_req_ids.add(kv_request_id)
368 + return
369 + logger.debug(
370 + "P2P %s: submit_load job_id=%d -> request_blocks peer=%s "
371 + "kv_request_id=%s blocks=%d session_ready=%s",
372 + self._local_id,
373 + job_id,
374 + peer_id,
375 + kv_request_id,
376 + len(block_ids),
377 + session.ready,
378 + )
379 + session.request_blocks(job_id, kv_request_id, keys, block_ids)
380 +
381 + @override
382 + def get_finished_jobs(self) -> Iterable[JobResult]:
383 + # Drive one polling sweep on the scheduler thread, then hand off
384 + # whatever has accumulated. The engine calls this once per step
385 + # (and keeps stepping while has_pending_work() is True).
386 + self._poll_once()
387 + result = self._finished_jobs
388 + self._finished_jobs = []
389 + return result
390 +
391 + @override
392 + def has_pending_work(self) -> bool:
393 + # The engine tick is the only driver of _control.poll() and
394 + # session.poll(); without it we miss new peer connects and
395 + # inbound fetch messages on existing sessions. Keep the engine
396 + # ticking for the lifetime of this manager.
397 + return True
398 +
399 + @override
400 + def drain_jobs(self) -> None:
401 + """Block until every submitted load/store job has completed or failed.
402 +
403 + Loops calling ``_poll_once()`` until no session has outstanding
404 + inbound loads or in-flight outbound stores. Mid-flight transfers
405 + are NOT cancelled — the caller (``TieringOffloadingManager.reset_cache``)
406 + needs the primary memoryview to be quiescent, not aborted. Results
407 + accumulate in ``_finished_jobs`` and are surfaced by the next
408 + ``get_finished_jobs()`` call.
409 + """
410 + start = time.monotonic()
411 + warned = False
412 + while True:
413 + self._poll_once()
414 + pending = any(
415 + s._client._inbound or s._server._inflight
416 + for s in self._sessions.values()
417 + )
418 + if not pending:
419 + return
420 + if not warned and time.monotonic() - start > 5.0:
421 + logger.warning(
422 + "P2PSecondaryTierManager.drain_jobs: still draining "
423 + "after 5s; a stuck transfer will block the engine.",
424 + )
425 + warned = True
426 + time.sleep(_DRAIN_SLEEP_S)
427 +
428 + @override
429 + def on_schedule_end(self) -> None:
430 + return
431 +
432 + # ------------------------------------------------------------------
433 + # Internal
434 + # ------------------------------------------------------------------
435 +
436 + @staticmethod
437 + def _remote_id_from_params(role_params: dict) -> str | None:
438 + """Build peer_id from a role-scoped sub-dict (``prefill``/``p2p``)."""
439 + host = role_params.get("remote_host")
440 + port = role_params.get("remote_port")
441 + if host and port:
442 + return f"{host}:{port}"
443 + return None
444 +
445 + def _get_or_create_session(self, peer_id: str) -> P2PSession:
446 + """Return the existing session for peer_id, or open one outbound.
447 +
448 + Decoder-side helper for on_new_request: when ``prefill`` is set,
449 + the consumer must reach the producer at peer_id. If we already
450 + have a session toward that peer (from a prior load or a
451 + peer-initiated inbound), reuse it; otherwise open an outbound
452 + ControlConnection and build a connected session.
453 + """
454 + session = self._sessions.get(peer_id)
455 + if session is not None:
456 + return session
457 + conn = self._control.connect(peer_id)
458 + session = P2PSession(
459 + peer_id=peer_id,
460 + local_id=self._local_id,
461 + transport=self._data,
462 + local_block_len=self._data.block_len,
463 + conn=conn,
464 + )
465 + self._sessions[peer_id] = session
466 + return session
467 +
468 + def _accept_new_peers(self, new_connections: Sequence[ControlConnection]) -> None:
469 + for conn in new_connections:
470 + logger.info(
471 + "P2P %s: accepting incoming connection from %s",
472 + self._local_id,
473 + conn.peer_id,
474 + )
475 + try:
476 + existing = self._sessions.get(conn.peer_id)
477 + if existing is not None:
478 + raise ValueError(f"duplicate connection from {conn.peer_id}")
479 + self._sessions[conn.peer_id] = P2PSession(
480 + peer_id=conn.peer_id,
481 + local_id=self._local_id,
482 + transport=self._data,
483 + local_block_len=self._data.block_len,
484 + conn=conn,
485 + )
486 + logger.info(
487 + "P2P %s: created connected session for %s",
488 + self._local_id,
489 + conn.peer_id,
490 + )
491 + except (ValueError, KeyError, TypeError, AssertionError) as exc:
492 + logger.error("P2P %s: rejecting peer: %s", self._local_id, exc)
493 + conn.close()
494 +
495 + def _reap_dead_sessions(self) -> None:
496 + # Reap connected sessions whose connection died — peer is gone.
497 + # Stranded prefiller-side stores are no longer tracked through a
498 + # session (they live in _unbound_stores keyed by kv_request_id);
499 + # _reap_unbound_stores handles their timeout independently.
500 + dead: list[str] | None = None
501 + for pid, s in self._sessions.items():
502 + if s.connected and not s.alive:
503 + if dead is None:
504 + dead = []
505 + dead.append(pid)
506 + if dead is None:
507 + return
508 + for pid in dead:
509 + session = self._sessions.pop(pid)
510 + # Purge any kv_request_id → session entries pointing at this
511 + # session so subsequent submit_stores fall back to the unbound
512 + # path (which will time out into failure if no peer rebinds).
513 + stale_kv_ids = [
514 + kid for kid, s in self._kv_to_session.items() if s is session
515 + ]
516 + for kid in stale_kv_ids:
517 + del self._kv_to_session[kid]
518 + failed_loads, failed_stores = session.close()
519 + for job_id, kv_request_id in failed_loads:
520 + self._finished_jobs.append(JobResult(job_id=job_id, success=False))
521 + self._failed_req_ids.add(kv_request_id)
522 + for job_id in failed_stores:
523 + self._finished_jobs.append(JobResult(job_id=job_id, success=False))
524 + self._data.remove_remote_peer(pid)
525 + logger.warning("P2P %s: peer %s down", self._local_id, pid)
526 +
527 + def _reap_unbound_stores(self) -> None:
528 + """Time out submit_store batches that no peer has ever fetched.
529 +
530 + Walks `_unbound_stores` for entries whose oldest batch is older
531 + than `_UNBOUND_STORE_TIMEOUT_S`. Drops the kv_request_id, surfaces
532 + every batched job as failed, and adds the id to `_failed_req_ids`
533 + so a late inbound FetchMsg short-circuits to a clean rejection.
534 + """
535 + if not self._unbound_stores:
536 + return
537 + deadline = time.monotonic() - _UNBOUND_STORE_TIMEOUT_S
538 + expired: list[str] | None = None
539 + for kid, batches in self._unbound_stores.items():
540 + # Batches are appended in arrival order, so the head is oldest.
541 + if batches and batches[0].submitted_at <= deadline:
542 + if expired is None:
543 + expired = []
544 + expired.append(kid)
545 + if expired is None:
546 + return
547 + for kid in expired:
548 + batches = self._unbound_stores.pop(kid)
549 + self._failed_req_ids.add(kid)
550 + for batch in batches:
551 + self._finished_jobs.append(
552 + JobResult(job_id=batch.job_id, success=False)
553 + )
554 + logger.warning(
555 + "P2P %s: unbound store kv_request_id=%s timed out after %.0fs "
556 + "without a fetch — failing %d job(s)",
557 + self._local_id,
558 + kid,
559 + _UNBOUND_STORE_TIMEOUT_S,
560 + len(batches),
561 + )
562 +
563 + # ------------------------------------------------------------------
564 + # Polling
565 + # ------------------------------------------------------------------
566 +
567 + def _poll_once(self) -> None:
568 + """One sweep of the polling work.
569 +
570 + Drains the control transport, polls every session, accumulates
571 + their results into ``_finished_jobs``, and reaps any dead sessions.
572 + Runs on the scheduler thread.
573 + """
574 + new_connections = self._control.poll()
575 + if new_connections:
576 + logger.info(
577 + "P2P %s: _poll_once got %d new connection(s): %s",
578 + self._local_id,
579 + len(new_connections),
580 + [c.peer_id for c in new_connections],
581 + )
582 +
583 + self._accept_new_peers(new_connections)
584 +
585 + for session in self._sessions.values():
586 + result = session.poll()
587 + for lr in result.loads:
588 + self._finished_jobs.append(
589 + JobResult(job_id=lr.job_id, success=lr.success)
590 + )
591 + if not lr.success:
592 + self._failed_req_ids.add(lr.kv_request_id)
593 + for sr in result.stores:
594 + self._finished_jobs.append(
595 + JobResult(job_id=sr.job_id, success=sr.success)
596 + )
597 + # Bind kv_request_id → session for any FetchMsg this tick and
598 + # replay any submit_store batches parked while no peer was
599 + # asking. ServerRole.on_fetch already recorded the demand
600 + # inline in dispatch, so the replayed add_stored_blocks calls
601 + # match that demand and submit transfers immediately.
602 + for kv_request_id in result.new_fetch_ids:
603 + self._kv_to_session[kv_request_id] = session
604 + for batch in self._unbound_stores.pop(kv_request_id, ()):
605 + session.add_stored_blocks(
606 + kv_request_id, batch.keys, batch.block_ids, batch.job_id
607 + )
608 +
609 + self._reap_dead_sessions()
610 + self._reap_unbound_stores()
611 +
612 + # ------------------------------------------------------------------
613 + # Lifecycle
614 + # ------------------------------------------------------------------
615 +
616 + @override
617 + def shutdown(self) -> None:
618 + self._drain_inflight_for_shutdown()
619 + for session in self._sessions.values():
620 + session.close()
621 + self._sessions.clear()
622 + self._kv_to_session.clear()
623 + # Surface buffered store jobs as failed so the engine doesn't
624 + # leak them; the manager is going away after this call.
625 + for batches in self._unbound_stores.values():
626 + for batch in batches:
627 + self._finished_jobs.append(
628 + JobResult(job_id=batch.job_id, success=False)
629 + )
630 + self._unbound_stores.clear()
631 + self._control.close()
632 + self._data.close()
633 +
634 + def _drain_inflight_for_shutdown(self) -> None:
635 + """Best-effort drain of inflight transfers before closing _data.
636 +
637 + Mirrors session._drain_abort but as a single bounded loop. Collects
638 + inflight transfer_ids from each session, repeatedly calls
639 + _data.cancel(..., mode="wait") and _data.poll() so handles can
640 + surface as done/failed, and falls back to mode="immediate" once
641 + _SHUTDOWN_DRAIN_TIMEOUT_S elapses so a wedged peer can't hang us.
642 + """
643 + ids = [tid for s in self._sessions.values() for tid in s._server._inflight]
644 + if not ids:
645 + return
646 + deadline = time.monotonic() + _SHUTDOWN_DRAIN_TIMEOUT_S
647 + still: list[int] = ids
648 + while still and time.monotonic() < deadline:
649 + still = list(self._data.cancel(still, mode="wait"))
650 + if not still:
651 + break
652 + # poll() advances NIXL handle state so the next wait-cancel
653 + # has a chance to release the handles.
654 + self._data.poll()
655 + time.sleep(_DRAIN_SLEEP_S)
656 + if still:
657 + logger.warning(
658 + "P2P %s: shutdown drain timed out after %.1fs with %d "
659 + "transfers still inflight — force-cancelling",
660 + self._local_id,
661 + _SHUTDOWN_DRAIN_TIMEOUT_S,
662 + len(still),
663 + )
664 + self._data.cancel(still, mode="immediate")
vllm/v1/kv_offload/tiering/p2p/session/__init__.py new
+15
@@ -0,0 +1,15 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from vllm.v1.kv_offload.tiering.p2p.session.session import (
4 + LoadResult,
5 + P2PSession,
6 + SessionPollResult,
7 + StoreResult,
8 +)
9 +
10 +__all__ = [
11 + "LoadResult",
12 + "P2PSession",
13 + "SessionPollResult",
14 + "StoreResult",
15 +]
vllm/v1/kv_offload/tiering/p2p/session/client.py new
+210
@@ -0,0 +1,210 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Client-role state machine for a single peer session.
4 +
5 +Handles outgoing fetch requests, abort-on-timeout, abort-ack timeout,
6 +and produces ``LoadResult`` for completed loads. The session coordinator
7 +parses wire messages and dispatches typed arguments here; this module
8 +never touches ``ControlConnection`` directly — it emits via the ``send``
9 +callback injected by the coordinator (which gates on ConnectAck).
10 +"""
11 +
12 +from __future__ import annotations
13 +
14 +import time
15 +from collections.abc import Callable, Sequence
16 +from dataclasses import dataclass
17 +from typing import TYPE_CHECKING, NamedTuple
18 +
19 +from vllm.logger import init_logger
20 +from vllm.v1.kv_offload.tiering.p2p.session.protocol import (
21 + TYPE_KEY,
22 + AbortFetchMsg,
23 + FetchMsg,
24 +)
25 +
26 +if TYPE_CHECKING:
27 + from vllm.v1.kv_offload.tiering.base import JobId
28 +
29 +logger = init_logger(__name__)
30 +
31 +_LOAD_TIMEOUT_S = 30.0
32 +_ABORT_ACK_TIMEOUT_S = 10.0
33 +
34 +
35 +@dataclass
36 +class _InboundRequestState:
37 + """Client-role state for a single load request."""
38 +
39 + job_id: int # opaque ID assigned by the manager to this load request
40 + kv_request_id: str
41 + submitted_at: float
42 + aborted_at: float | None = None
43 +
44 +
45 +class LoadResult(NamedTuple):
46 + """Result from a session poll, client side."""
47 +
48 + job_id: int
49 + kv_request_id: str
50 + success: bool
51 +
52 +
53 +class ClientRole:
54 + """Client-side load state machine for one peer session.
55 +
56 + The coordinator owns the connection and the send-gating; this role
57 + is given a ``send`` callback and a ``peer_id`` for log messages and
58 + is otherwise self-contained.
59 + """
60 +
61 + def __init__(self, peer_id: str, send: Callable[[dict], None]) -> None:
62 + self._peer_id = peer_id
63 + self._send = send
64 + self._inbound: dict[str, _InboundRequestState] = {}
65 + self._completed_loads: list[LoadResult] = []
66 +
67 + # ------------------------------------------------------------------
68 + # Public API
69 + # ------------------------------------------------------------------
70 +
71 + def request_blocks(
72 + self,
73 + job_id: JobId,
74 + kv_request_id: str,
75 + keys: Sequence[bytes],
76 + block_ids: Sequence[int],
77 + send_ready: bool,
78 + ) -> None:
79 + """Register a load request and send the FetchMsg."""
80 + logger.debug(
81 + "P2PSession %s: request_blocks job_id=%d kv_request_id=%s "
82 + "blocks=%d ready=%s",
83 + self._peer_id,
84 + job_id,
85 + kv_request_id,
86 + len(block_ids),
87 + send_ready,
88 + )
89 + self._inbound[kv_request_id] = _InboundRequestState(
90 + job_id=job_id,
91 + kv_request_id=kv_request_id,
92 + submitted_at=time.monotonic(),
93 + )
94 + self._send(
95 + {
96 + TYPE_KEY: FetchMsg.TYPE,
97 + FetchMsg.KV_REQUEST_ID: kv_request_id,
98 + FetchMsg.BLOCK_HASHES: list(keys),
99 + FetchMsg.BLOCK_INDEXES: [int(idx) for idx in block_ids],
100 + }
101 + )
102 +
103 + def cancel(self, kv_request_id: str) -> None:
104 + """Cancel a pending load. Sends AbortFetchMsg if still active."""
105 + req = self._inbound.pop(kv_request_id, None)
106 + if req is not None and req.aborted_at is None:
107 + self._send(
108 + {
109 + TYPE_KEY: AbortFetchMsg.TYPE,
110 + AbortFetchMsg.KV_REQUEST_ID: kv_request_id,
111 + }
112 + )
113 +
114 + def on_transfer_done(self, kv_request_id: str, success: bool) -> None:
115 + """Handle a TransferDoneMsg from the peer."""
116 + req = self._inbound.pop(kv_request_id, None)
117 + if req is not None:
118 + self._completed_loads.append(
119 + LoadResult(
120 + job_id=req.job_id,
121 + kv_request_id=kv_request_id,
122 + success=success,
123 + )
124 + )
125 + else:
126 + # No matching _inbound entry: either a duplicate
127 + # transfer_done from the peer (protocol violation) or a
128 + # benign race with a local cancel/abort/timeout that
129 + # already popped the entry. We don't track terminated ids,
130 + # so we can't tell — log so it's findable.
131 + logger.warning(
132 + "P2PSession %s: transfer_done for unknown kv_request_id=%s "
133 + "(duplicate from peer, or raced with local cancel/timeout)",
134 + self._peer_id,
135 + kv_request_id,
136 + )
137 +
138 + def on_abort_ack(self, kv_request_id: str) -> None:
139 + """Handle an AbortAckMsg from the peer."""
140 + req = self._inbound.pop(kv_request_id, None)
141 + if req is not None:
142 + self._completed_loads.append(
143 + LoadResult(
144 + job_id=req.job_id,
145 + kv_request_id=kv_request_id,
146 + success=False,
147 + )
148 + )
149 + else:
150 + # See on_transfer_done: same ambiguity (duplicate ack
151 + # vs. raced with local cancel/timeout that already popped).
152 + logger.warning(
153 + "P2PSession %s: abort_ack for unknown kv_request_id=%s "
154 + "(duplicate from peer, or raced with local cancel/timeout)",
155 + self._peer_id,
156 + kv_request_id,
157 + )
158 +
159 + def collect_results(self) -> list[LoadResult]:
160 + """Walk timeouts and drain completed loads.
161 +
162 + Active requests past ``_LOAD_TIMEOUT_S`` get an AbortFetchMsg
163 + sent and enter the aborting phase. Aborting requests past
164 + ``_ABORT_ACK_TIMEOUT_S`` are surfaced as failed loads.
165 + """
166 + now = time.monotonic()
167 + to_remove: list[str] = []
168 + for req_id, req in self._inbound.items():
169 + if req.aborted_at is None:
170 + if now - req.submitted_at >= _LOAD_TIMEOUT_S:
171 + req.aborted_at = now
172 + logger.warning(
173 + "P2PSession %s: %s timed out, sending abort",
174 + self._peer_id,
175 + req_id,
176 + )
177 + self._send(
178 + {
179 + TYPE_KEY: AbortFetchMsg.TYPE,
180 + AbortFetchMsg.KV_REQUEST_ID: req_id,
181 + }
182 + )
183 + else:
184 + if now - req.aborted_at >= _ABORT_ACK_TIMEOUT_S:
185 + to_remove.append(req_id)
186 + self._completed_loads.append(
187 + LoadResult(
188 + job_id=req.job_id,
189 + kv_request_id=req_id,
190 + success=False,
191 + )
192 + )
193 + logger.warning(
194 + "P2PSession %s: abort_ack timed out for kv_request_id=%s",
195 + self._peer_id,
196 + req_id,
197 + )
198 + for req_id in to_remove:
199 + self._inbound.pop(req_id)
200 +
201 + results = self._completed_loads
202 + self._completed_loads = []
203 + return results
204 +
205 + def close(self) -> list[tuple[int, str]]:
206 + """Tear down. Returns ``(job_id, kv_request_id)`` for pending loads."""
207 + failed = [(req.job_id, req.kv_request_id) for req in self._inbound.values()]
208 + self._inbound.clear()
209 + self._completed_loads.clear()
210 + return failed
vllm/v1/kv_offload/tiering/p2p/session/protocol.py new
+247
@@ -0,0 +1,247 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +P2P KV cache sharing protocol constants and documentation.
5 +
6 +Protocol Overview
7 +=================
8 +
9 +Two session types communicate over a bidirectional message channel:
10 +- P2PClientSession (requests blocks from a server)
11 +- P2PServerSession (serves blocks to a client)
12 +
13 +Connection Lifecycle
14 +--------------------
15 +
16 +1. Client opens a connection and sends ConnectMsg with its identity,
17 + RDMA metadata, memory layout, and config fingerprint.
18 +2. Server validates block_len and config_fingerprint, registers the
19 + RDMA peer, and replies with ConnectAckMsg.
20 +3. Client receives ConnectAckMsg and transitions to ready state
21 + (flushes any queued messages).
22 +4. Either side may send DisconnectMsg to gracefully close.
23 +
24 +Block Transfer Flow (happy path)
25 +---------------------------------
26 +
27 +1. Client sends FetchMsg with a kv_request_id and lists of
28 + block keys + remote indexes where it wants the data written.
29 +2. Server matches requested blocks against locally stored blocks:
30 + - Blocks already available are transferred immediately via RDMA.
31 + - Blocks not yet available are recorded as "demanded" and
32 + transferred when the server later stores them.
33 +3. When all blocks for a kv_request_id are transferred, the server
34 + sends TransferDoneMsg (success=True) to the client.
35 +4. Client reports the load job as complete.
36 +
37 +Abort Flow (timeout path)
38 +--------------------------
39 +
40 +1. If the client times out waiting for TransferDoneMsg, it sends
41 + AbortFetchMsg to cancel the request.
42 +2. Server cancels inflight transfers for that kv_request_id and
43 + replies with AbortAckMsg.
44 +3. Client receives AbortAckMsg and reports the load job as failed.
45 +4. If AbortAckMsg itself times out, the client fails the job anyway.
46 +
47 +Message Format
48 +--------------
49 +
50 +All messages are dicts serialized with msgpack. Every message has a
51 +TYPE_KEY key identifying its type. Additional fields depend on the
52 +message type (see per-message class docstrings below).
53 +
54 +Security
55 +--------
56 +
57 +Sessions wrap all incoming message handling in try/except to guard
58 +against malformed messages from adversarial or buggy peers. Invalid
59 +messages are logged and dropped without crashing the session.
60 +
61 +The config_fingerprint field in ConnectMsg ensures peers have
62 +compatible model configurations (model, dtype, block sizes). Mismatches
63 +are rejected during the handshake.
64 +"""
65 +
66 +TYPE_KEY = "type"
67 +
68 +
69 +# ---------------------------------------------------------------------------
70 +# Validation helpers
71 +# ---------------------------------------------------------------------------
72 +
73 +
74 +def _require(msg: dict, key: str, typ: type, *, name: str = "") -> None:
75 + """Raise ValueError if msg[key] is missing or not isinstance(typ)."""
76 + val = msg.get(key)
77 + if not isinstance(val, typ):
78 + label = name or key
79 + raise ValueError(f"{label}: expected {typ.__name__}, got {type(val).__name__}")
80 +
81 +
82 +def _require_pos_int(msg: dict, key: str, *, name: str = "") -> None:
83 + """Raise ValueError if msg[key] is not a positive int."""
84 + val = msg.get(key)
85 + if not isinstance(val, int) or val <= 0:
86 + label = name or key
87 + raise ValueError(f"{label}: expected positive int, got {val!r}")
88 +
89 +
90 +def _require_non_neg_int(msg: dict, key: str, *, name: str = "") -> None:
91 + """Raise ValueError if msg[key] is not a non-negative int."""
92 + val = msg.get(key)
93 + if not isinstance(val, int) or val < 0:
94 + label = name or key
95 + raise ValueError(f"{label}: expected non-negative int, got {val!r}")
96 +
97 +
98 +def _require_list(msg: dict, key: str, *, name: str = "") -> None:
99 + """Raise ValueError if msg[key] is not a list."""
100 + val = msg.get(key)
101 + if not isinstance(val, list):
102 + label = name or key
103 + raise ValueError(f"{label}: expected list, got {type(val).__name__}")
104 +
105 +
106 +# ---------------------------------------------------------------------------
107 +# Message classes
108 +# ---------------------------------------------------------------------------
109 +
110 +
111 +class ConnectMsg:
112 + """Client → Server: initial handshake request.
113 +
114 + Fields:
115 + PEER_ID: Local peer identity string.
116 + AGENT_METADATA: RDMA agent metadata (opaque bytes).
117 + BASE_ADDR: Base memory address of the KV block region.
118 + NUM_BLOCKS: Number of blocks in the KV block region.
119 + BLOCK_LEN: Size in bytes of each block (must match between peers).
120 + CONFIG_FINGERPRINT: SHA-256 prefix of the model configuration.
121 + Peers with different fingerprints are incompatible.
122 + """
123 +
124 + TYPE = "connect"
125 + PEER_ID = "peer_id"
126 + AGENT_METADATA = "agent_metadata"
127 + BASE_ADDR = "base_addr"
128 + NUM_BLOCKS = "num_blocks"
129 + BLOCK_LEN = "block_len"
130 + CONFIG_FINGERPRINT = "config_fingerprint"
131 +
132 + @staticmethod
133 + def validate(msg: dict) -> None:
134 + """Raise ValueError if any field has an invalid type or value."""
135 + _require(msg, ConnectMsg.PEER_ID, str)
136 + _require(msg, ConnectMsg.AGENT_METADATA, bytes)
137 + _require_non_neg_int(msg, ConnectMsg.BASE_ADDR)
138 + _require_pos_int(msg, ConnectMsg.NUM_BLOCKS)
139 + _require_pos_int(msg, ConnectMsg.BLOCK_LEN)
140 +
141 +
142 +class ConnectAckMsg:
143 + """Server → Client: handshake acknowledgement.
144 +
145 + Fields:
146 + PEER_ID: Server's peer identity string.
147 + """
148 +
149 + TYPE = "connect_ack"
150 + PEER_ID = "peer_id"
151 +
152 + @staticmethod
153 + def validate(msg: dict) -> None:
154 + """Raise ValueError if any field has an invalid type or value."""
155 + _require(msg, ConnectAckMsg.PEER_ID, str)
156 +
157 +
158 +class DisconnectMsg:
159 + """Either → Either: graceful connection close.
160 +
161 + No additional fields beyond TYPE_KEY.
162 + """
163 +
164 + TYPE = "disconnect"
165 +
166 +
167 +class FetchMsg:
168 + """Client → Server: request blocks by key.
169 +
170 + Fields:
171 + KV_REQUEST_ID: Identifies this block transfer request.
172 + BLOCK_HASHES: List of block keys (OffloadKey bytes).
173 + BLOCK_INDEXES: List of remote block indexes (same length as BLOCK_HASHES).
174 + """
175 +
176 + TYPE = "fetch"
177 + KV_REQUEST_ID = "kv_request_id"
178 + BLOCK_HASHES = "block_hashes"
179 + BLOCK_INDEXES = "block_indexes"
180 +
181 + @staticmethod
182 + def validate(msg: dict) -> None:
183 + """Raise ValueError if any field has an invalid type or value."""
184 + _require(msg, FetchMsg.KV_REQUEST_ID, str)
185 + _require_list(msg, FetchMsg.BLOCK_HASHES)
186 + _require_list(msg, FetchMsg.BLOCK_INDEXES)
187 + hashes = msg[FetchMsg.BLOCK_HASHES]
188 + indexes = msg[FetchMsg.BLOCK_INDEXES]
189 + if len(hashes) != len(indexes):
190 + raise ValueError(
191 + f"block_hashes/block_indexes length mismatch: "
192 + f"{len(hashes)} vs {len(indexes)}"
193 + )
194 + for idx in indexes:
195 + if not isinstance(idx, int) or idx < 0:
196 + raise ValueError(f"block_indexes: invalid index {idx!r}")
197 +
198 +
199 +class TransferDoneMsg:
200 + """Server → Client: all blocks transferred for a request.
201 +
202 + Fields:
203 + KV_REQUEST_ID: The request that completed.
204 + SUCCESS: Whether the transfer completed successfully.
205 + """
206 +
207 + TYPE = "transfer_done"
208 + KV_REQUEST_ID = "kv_request_id"
209 + SUCCESS = "success"
210 +
211 + @staticmethod
212 + def validate(msg: dict) -> None:
213 + """Raise ValueError if any field has an invalid type or value."""
214 + _require(msg, TransferDoneMsg.KV_REQUEST_ID, str)
215 + _require(msg, TransferDoneMsg.SUCCESS, bool)
216 +
217 +
218 +class AbortFetchMsg:
219 + """Client → Server: cancel a pending request.
220 +
221 + Fields:
222 + KV_REQUEST_ID: The request to cancel.
223 + """
224 +
225 + TYPE = "abort_fetch"
226 + KV_REQUEST_ID = "kv_request_id"
227 +
228 + @staticmethod
229 + def validate(msg: dict) -> None:
230 + """Raise ValueError if any field has an invalid type or value."""
231 + _require(msg, AbortFetchMsg.KV_REQUEST_ID, str)
232 +
233 +
234 +class AbortAckMsg:
235 + """Server → Client: acknowledge cancellation.
236 +
237 + Fields:
238 + KV_REQUEST_ID: The request that was cancelled.
239 + """
240 +
241 + TYPE = "abort_ack"
242 + KV_REQUEST_ID = "kv_request_id"
243 +
244 + @staticmethod
245 + def validate(msg: dict) -> None:
246 + """Raise ValueError if any field has an invalid type or value."""
247 + _require(msg, AbortAckMsg.KV_REQUEST_ID, str)
vllm/v1/kv_offload/tiering/p2p/session/server.py new
+601
@@ -0,0 +1,601 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Server-role state machine for a single peer session.
4 +
5 +Owns block matching (supply vs. demand), inflight RDMA transfers, store-job
6 +timeouts, abort-drain, and produces ``StoreResult`` for completed stores.
7 +The session coordinator parses wire messages and dispatches typed
8 +arguments here; this module never touches ``ControlConnection`` directly
9 +— it emits via the ``send`` callback injected by the coordinator (which
10 +gates on ConnectAck).
11 +
12 +Protocol violations the role can detect (today: duplicate ``FetchMsg``
13 +for the same ``kv_request_id``) are surfaced as ``ValueError`` so the
14 +coordinator's ``_dispatch_message`` can reuse its existing
15 +``_protocol_error`` path.
16 +"""
17 +
18 +from __future__ import annotations
19 +
20 +import time
21 +from collections.abc import Callable, Sequence
22 +from dataclasses import dataclass, field
23 +from typing import TYPE_CHECKING, NamedTuple
24 +
25 +from vllm.logger import init_logger
26 +from vllm.v1.kv_offload.base import OffloadKey
27 +from vllm.v1.kv_offload.tiering.p2p.session.protocol import (
28 + TYPE_KEY,
29 + AbortAckMsg,
30 + TransferDoneMsg,
31 +)
32 +
33 +if TYPE_CHECKING:
34 + from vllm.v1.kv_offload.tiering.base import JobId
35 + from vllm.v1.kv_offload.tiering.p2p.data import DataTransport
36 +
37 +logger = init_logger(__name__)
38 +
39 +_STORE_TIMEOUT_S = 30.0
40 +_CANCEL_DRAIN_TIMEOUT_S = 10.0
41 +
42 +
43 +class StoreResult(NamedTuple):
44 + """Result from a session poll, server side."""
45 +
46 + job_id: int
47 + success: bool
48 +
49 +
50 +class _InflightXfer(NamedTuple):
51 + """Metadata for a single inflight RDMA transfer, keyed by transfer_id."""
52 +
53 + kv_request_id: str
54 + block_count: int
55 + # The set of store job IDs that contributed blocks to this transfer.
56 + job_ids: set[int]
57 +
58 +
59 +class _MatchResult(NamedTuple):
60 + """Result of block matching: pairs ready for transfer."""
61 +
62 + local_idxs: list[int]
63 + remote_idxs: list[int]
64 + # The set of store job IDs that contributed blocks
65 + job_ids: set[int]
66 +
67 +
68 +@dataclass
69 +class _OutboundRequestState:
70 + """Server-role state for a single peer fetch request.
71 +
72 + The owning ``kv_request_id`` is the dict key in
73 + ``ServerRole._outbound`` and is not duplicated on the value.
74 + """
75 +
76 + demand_received: bool = False
77 + available: dict[OffloadKey, tuple[int, int]] = field(
78 + default_factory=dict
79 + ) # key → (job_id, local_block_idx): blocks we have, awaiting demand
80 + demanded: dict[OffloadKey, int] = field(
81 + default_factory=dict
82 + ) # key → remote_block_idx: blocks peer wants, awaiting supply
83 + remaining: int = 0 # blocks that need to be transferred to client
84 + finishing: bool = False # Signal finish request ASAP
85 + # Job IDs that submit_store'd blocks for this request and have not
86 + # yet emitted a StoreResult. The terminal-finalize helper drains
87 + # this set; poll-done and poll-failed discard entries as their
88 + # StoreResults fire.
89 + pending_job_ids: set[int] = field(default_factory=set)
90 +
91 + def add_stored_blocks(
92 + self,
93 + block_hashes: Sequence[OffloadKey],
94 + block_ids: Sequence[int],
95 + job_id: int,
96 + ) -> _MatchResult:
97 + """Add locally-stored blocks. Returns matched pairs."""
98 + self.pending_job_ids.add(job_id)
99 + local_idxs: list[int] = []
100 + remote_idxs: list[int] = []
101 + for block_hash, local_idx in zip(block_hashes, block_ids):
102 + remote_idx = self.demanded.pop(block_hash, None)
103 + if remote_idx is not None:
104 + local_idxs.append(local_idx)
105 + remote_idxs.append(remote_idx)
106 + else:
107 + self.available[block_hash] = (job_id, local_idx)
108 + return _MatchResult(
109 + local_idxs=local_idxs,
110 + remote_idxs=remote_idxs,
111 + job_ids={job_id} if local_idxs else set(),
112 + )
113 +
114 + def add_fetch_demand(
115 + self,
116 + block_hashes: Sequence[OffloadKey],
117 + block_indexes: Sequence[int],
118 + ) -> _MatchResult:
119 + """Register the peer's fetch demand. Returns matched pairs."""
120 + self.demand_received = True
121 + self.remaining = len(block_hashes)
122 +
123 + local_idxs: list[int] = []
124 + remote_idxs: list[int] = []
125 + job_ids: set[int] = set()
126 + for block_hash, remote_idx in zip(block_hashes, block_indexes):
127 + stored_entry = self.available.pop(block_hash, None)
128 + if stored_entry is not None:
129 + stored_job_id, local_idx = stored_entry
130 + local_idxs.append(local_idx)
131 + remote_idxs.append(remote_idx)
132 + job_ids.add(stored_job_id)
133 + else:
134 + self.demanded[block_hash] = remote_idx
135 + return _MatchResult(
136 + local_idxs=local_idxs,
137 + remote_idxs=remote_idxs,
138 + job_ids=job_ids,
139 + )
140 +
141 +
142 +class ServerRole:
143 + """Server-side store/serve state machine for one peer session.
144 +
145 + The coordinator owns the connection and the send-gating; this role
146 + is given a ``send`` callback, the ``DataTransport``, and the
147 + ``peer_id`` for transport calls and log messages.
148 + """
149 +
150 + def __init__(
151 + self,
152 + peer_id: str,
153 + transport: DataTransport,
154 + send: Callable[[dict], None],
155 + ) -> None:
156 + self._peer_id = peer_id
157 + self._transport = transport
158 + self._send = send
159 +
160 + self._outbound: dict[str, _OutboundRequestState] = {}
161 + # transfer_id → xfer. Mutate ONLY via _inflight_add / _inflight_pop
162 + # so the per-request count below stays in sync.
163 + self._inflight: dict[int, _InflightXfer] = {}
164 + # kv_request_id → number of entries in _inflight for that id.
165 + # Kept in sync with _inflight; entries that hit zero are removed
166 + # so `kv_request_id in self._inflight_per_req` is an exact
167 + # "has any inflight transfer" predicate (O(1) replacement for
168 + # the previous O(N) scan).
169 + self._inflight_per_req: dict[str, int] = {}
170 + self._store_jobs: dict[int, float] = {} # job_id → submitted_at
171 + self._pending_aborts: dict[str, float] = {} # kv_request_id → start
172 + # StoreResults queued by _finalize_outbound for the next poll
173 + # tick to surface. Mirrors the deferred-result pattern used for
174 + # load timeouts.
175 + self._pending_store_results: list[StoreResult] = []
176 +
177 + # ------------------------------------------------------------------
178 + # Public API
179 + # ------------------------------------------------------------------
180 +
181 + def add_stored_blocks(
182 + self,
183 + kv_request_id: str,
184 + keys: Sequence[OffloadKey],
185 + block_ids: Sequence[int],
186 + job_id: JobId,
187 + ) -> None:
188 + """New blocks stored locally — match against pending fetch demand."""
189 + self._store_jobs[job_id] = time.monotonic()
190 + req = self._outbound.setdefault(kv_request_id, _OutboundRequestState())
191 + result = req.add_stored_blocks(keys, block_ids, job_id)
192 + if result.local_idxs and req.demand_received:
193 + self._submit_transfer(kv_request_id, result)
194 +
195 + def on_fetch(
196 + self,
197 + kv_request_id: str,
198 + block_hashes: Sequence[OffloadKey],
199 + block_indexes: Sequence[int],
200 + ) -> None:
201 + """Handle a FetchMsg from the peer.
202 +
203 + Raises ``ValueError`` on a duplicate fetch for the same
204 + ``kv_request_id``; the coordinator's dispatch loop turns that
205 + into a protocol-error disconnect.
206 + """
207 + logger.debug(
208 + "P2PSession %s: fetch RECEIVED kv_request_id=%s blocks=%d",
209 + self._peer_id,
210 + kv_request_id,
211 + len(block_hashes),
212 + )
213 + existing = self._outbound.get(kv_request_id)
214 + if existing is not None and existing.demand_received:
215 + # A second fetch for the same kv_request_id would overwrite
216 + # `remaining` and leak inflight bookkeeping. Treat as a
217 + # protocol violation.
218 + raise ValueError(f"duplicate fetch for kv_request_id={kv_request_id}")
219 + req = self._outbound.setdefault(kv_request_id, _OutboundRequestState())
220 + result = req.add_fetch_demand(block_hashes, block_indexes)
221 + if result.local_idxs:
222 + self._submit_transfer(kv_request_id, result)
223 + # Prefiller-first mode: finish_request may have run before
224 + # fetch arrived. If so, finalize once we know what was
225 + # demanded — fully satisfied → success, else early-fail.
226 + if req.finishing and not self._has_inflight_for(kv_request_id):
227 + self._finalize_outbound(kv_request_id)
228 +
229 + def on_abort_fetch(self, kv_request_id: str) -> None:
230 + """Handle an AbortFetchMsg from the peer."""
231 + # Abort for an unknown id may be a benign race/duplicate or a
232 + # real protocol violation; we don't track completed ids, so warn.
233 + if kv_request_id not in self._outbound and not self._has_inflight_for(
234 + kv_request_id
235 + ):
236 + logger.warning(
237 + "P2PSession %s: abort_fetch for unknown kv_request_id=%s "
238 + "(no outbound or inflight state); benign race or stale",
239 + self._peer_id,
240 + kv_request_id,
241 + )
242 + # Idempotent: receiving AbortFetchMsg again before we've sent the
243 + # ack just triggers another drain attempt without resetting the
244 + # deadline.
245 + self._pending_aborts.setdefault(kv_request_id, time.monotonic())
246 + self._drain_abort(kv_request_id)
247 +
248 + def finish(self, kv_request_id: str) -> None:
249 + """Mark an outbound request finishing.
250 +
251 + No more submit_store calls will arrive for this id. Any blocks
252 + the peer demanded but we never stored will never come; tell the
253 + peer to stop waiting (TransferDoneMsg success=False) instead of
254 + letting it hit _LOAD_TIMEOUT_S.
255 +
256 + If the decoder hasn't sent fetch yet (no demand received),
257 + defer — on_fetch will finalize once demand arrives.
258 +
259 + If inflight transfers exist for this id, defer — the last
260 + completing transfer in collect_results will fire the message.
261 + """
262 + req = self._outbound.get(kv_request_id)
263 + if req is None:
264 + return
265 + req.finishing = True
266 + if not req.demand_received:
267 + return
268 + if self._has_inflight_for(kv_request_id):
269 + return
270 + # Remaining > 0 here: if it had hit 0, the poll-done success
271 + # branch would have already popped _outbound and we'd have
272 + # returned at `req is None` above. Helper derives success from
273 + # remaining and emits StoreResult(success=False) for any
274 + # leftover pending jobs.
275 + self._finalize_outbound(kv_request_id)
276 +
277 + def collect_results(self) -> list[StoreResult]:
278 + """Drain timeouts, deferred results, and transport completions."""
279 + results: list[StoreResult] = self._timeout_pending_store_jobs()
280 +
281 + if self._pending_store_results:
282 + results.extend(self._pending_store_results)
283 + self._pending_store_results.clear()
284 +
285 + poll_result = self._transport.poll()
286 +
287 + for tid in poll_result.done:
288 + xfer = self._inflight_pop(tid)
289 + if xfer is None:
290 + # Bug signal: transport reported a transfer we have no
291 + # bookkeeping for. Likely a double-completion in the
292 + # transport or a stale removal in the session. The
293 + # attached job(s) still live in _store_jobs and will be
294 + # surfaced as failures by _timeout_pending_store_jobs
295 + # after _STORE_TIMEOUT_S, but log loudly so the
296 + # underlying bug is findable.
297 + logger.error(
298 + "P2PSession %s: transport reported done for unknown "
299 + "transfer_id=%d; attached job(s) will fail via "
300 + "store-timeout instead of completing now",
301 + self._peer_id,
302 + tid,
303 + )
304 + continue
305 + req = self._outbound.get(xfer.kv_request_id)
306 + for job_id in xfer.job_ids:
307 + if self._store_jobs.pop(job_id, None) is None:
308 + # Already reported (timeout, cancellation, etc.) —
309 + # don't double-emit a contradictory success result.
310 + continue
311 + results.append(StoreResult(job_id=job_id, success=True))
312 + if req is not None:
313 + req.pending_job_ids.discard(job_id)
314 + if req is not None and req.demand_received:
315 + req.remaining -= xfer.block_count
316 + assert req.remaining >= 0, (
317 + f"remaining went negative for kv_request_id={xfer.kv_request_id}"
318 + )
319 + if req.remaining == 0:
320 + self._finalize_outbound(xfer.kv_request_id, success=True)
321 + elif req.finishing and not self._has_inflight_for(xfer.kv_request_id):
322 + self._finalize_outbound(xfer.kv_request_id, success=False)
323 +
324 + failed_kv_request_ids: set[str] | None = None
325 + for tid in poll_result.failed:
326 + xfer = self._inflight_pop(tid)
327 + if xfer is None:
328 + # See the matching error log in the done branch above.
329 + logger.error(
330 + "P2PSession %s: transport reported failed for unknown "
331 + "transfer_id=%d; attached job(s) will fail via "
332 + "store-timeout instead of completing now",
333 + self._peer_id,
334 + tid,
335 + )
336 + continue
337 + if failed_kv_request_ids is None:
338 + failed_kv_request_ids = set()
339 + failed_kv_request_ids.add(xfer.kv_request_id)
340 + req_for_xfer = self._outbound.get(xfer.kv_request_id)
341 + for job_id in xfer.job_ids:
342 + if self._store_jobs.pop(job_id, None) is None:
343 + # Already reported (timeout, cancellation, etc.) —
344 + # don't double-emit.
345 + continue
346 + results.append(StoreResult(job_id=job_id, success=False))
347 + if req_for_xfer is not None:
348 + req_for_xfer.pending_job_ids.discard(job_id)
349 + req = self._outbound.pop(xfer.kv_request_id, None)
350 + if req is not None and req.demand_received:
351 + self._send(
352 + {
353 + TYPE_KEY: TransferDoneMsg.TYPE,
354 + TransferDoneMsg.KV_REQUEST_ID: xfer.kv_request_id,
355 + TransferDoneMsg.SUCCESS: False,
356 + }
357 + )
358 +
359 + # Cancel other inflight for the same failed kv_request_ids
360 + if failed_kv_request_ids:
361 + ids_to_cancel = [
362 + tid
363 + for tid, xfer in self._inflight.items()
364 + if xfer.kv_request_id in failed_kv_request_ids
365 + ]
366 + for tid in ids_to_cancel:
367 + self._inflight_pop(tid)
368 + self._transport.cancel(ids_to_cancel)
369 +
370 + return results
371 +
372 + def collect_idle_timeouts(self) -> list[StoreResult]:
373 + """Run only the store-job timeout sweep.
374 +
375 + Used by the coordinator's no-conn poll path: a pending session
376 + cannot have inflight transfers (no peer registered yet), so we
377 + skip the transport poll and the deferred-result drain.
378 + """
379 + return self._timeout_pending_store_jobs()
380 +
381 + def drain_pending_aborts(self) -> None:
382 + """Re-attempt every parked abort once per poll tick."""
383 + if not self._pending_aborts:
384 + return
385 + for kv_request_id in list(self._pending_aborts):
386 + self._drain_abort(kv_request_id)
387 +
388 + def close(self) -> list[int]:
389 + """Tear down. Cancels inflight, returns failed store job ids."""
390 + failed_stores = list(self._store_jobs.keys())
391 + self._store_jobs.clear()
392 + if self._inflight:
393 + self._transport.cancel(list(self._inflight.keys()))
394 + self._inflight.clear()
395 + self._inflight_per_req.clear()
396 + self._outbound.clear()
397 + self._pending_aborts.clear()
398 + self._pending_store_results.clear()
399 + return failed_stores
400 +
401 + # ------------------------------------------------------------------
402 + # Internal — inflight bookkeeping
403 + # ------------------------------------------------------------------
404 +
405 + def _has_inflight_for(self, kv_request_id: str) -> bool:
406 + return kv_request_id in self._inflight_per_req
407 +
408 + def _inflight_add(self, tid: int, xfer: _InflightXfer) -> None:
409 + """Insert an inflight transfer and bump the per-request count."""
410 + self._inflight[tid] = xfer
411 + self._inflight_per_req[xfer.kv_request_id] = (
412 + self._inflight_per_req.get(xfer.kv_request_id, 0) + 1
413 + )
414 +
415 + def _inflight_pop(self, tid: int) -> _InflightXfer | None:
416 + """Pop an inflight transfer and decrement the per-request count.
417 +
418 + Removes the per-request entry once the count hits zero so the
419 + dict stays bounded and `_has_inflight_for` remains exact.
420 + """
421 + xfer = self._inflight.pop(tid, None)
422 + if xfer is None:
423 + return None
424 + new_count = self._inflight_per_req.get(xfer.kv_request_id, 0) - 1
425 + if new_count > 0:
426 + self._inflight_per_req[xfer.kv_request_id] = new_count
427 + else:
428 + self._inflight_per_req.pop(xfer.kv_request_id, None)
429 + return xfer
430 +
431 + # ------------------------------------------------------------------
432 + # Internal — finalize / abort drain
433 + # ------------------------------------------------------------------
434 +
435 + def _finalize_outbound(
436 + self,
437 + kv_request_id: str,
438 + success: bool | None = None,
439 + ) -> None:
440 + """Pop the outbound state and emit terminal results.
441 +
442 + Called when no further work will happen for this kv_request_id
443 + on the server side: either request_finish has fired and there
444 + are no inflight transfers, or the last inflight just completed
445 + while finishing.
446 +
447 + If ``success`` is None, derive it from ``req.remaining == 0``.
448 + The same flag is used for both the peer's TransferDoneMsg and
449 + the StoreResult(s) emitted for any leftover pending job_ids.
450 + """
451 + req = self._outbound.pop(kv_request_id, None)
452 + if req is None:
453 + return
454 + if success is None:
455 + success = req.remaining == 0
456 + for job_id in req.pending_job_ids:
457 + self._store_jobs.pop(job_id, None)
458 + self._pending_store_results.append(
459 + StoreResult(job_id=job_id, success=success)
460 + )
461 + self._send(
462 + {
463 + TYPE_KEY: TransferDoneMsg.TYPE,
464 + TransferDoneMsg.KV_REQUEST_ID: kv_request_id,
465 + TransferDoneMsg.SUCCESS: success,
466 + }
467 + )
468 +
469 + def _drain_abort(self, kv_request_id: str) -> None:
470 + """One drain attempt for a pending abort.
471 +
472 + Stops accepting more blocks for ``kv_request_id``, then asks the
473 + transport to cancel any matching inflight transfers in
474 + ``mode="wait"``. Sends ``AbortAckMsg`` once nothing remains
475 + inflight, or after ``_CANCEL_DRAIN_TIMEOUT_S`` falls back to
476 + ``mode="immediate"`` and acks anyway.
477 + """
478 + self._outbound.pop(kv_request_id, None)
479 + ids = [
480 + tid
481 + for tid, xfer in self._inflight.items()
482 + if xfer.kv_request_id == kv_request_id
483 + ]
484 + if not ids:
485 + self._finalize_abort(kv_request_id)
486 + return
487 +
488 + started_at = self._pending_aborts.get(kv_request_id)
489 + expired = (
490 + started_at is not None
491 + and time.monotonic() - started_at >= _CANCEL_DRAIN_TIMEOUT_S
492 + )
493 + if expired:
494 + for tid in ids:
495 + self._inflight_pop(tid)
496 + self._transport.cancel(ids, mode="immediate")
497 + logger.warning(
498 + "P2PSession %s: cancel drain timed out for kv_request_id=%s,"
499 + " force-canceled %d transfers",
500 + self._peer_id,
501 + kv_request_id,
502 + len(ids),
503 + )
504 + self._finalize_abort(kv_request_id)
505 + return
506 +
507 + still = self._transport.cancel(ids, mode="wait")
508 + # Tids the transport successfully released are gone from its
509 + # _inflight; mirror that in session bookkeeping so they don't
510 + # block the drain forever waiting for a poll() event that will
511 + # never come.
512 + still_set = set(still)
513 + for tid in ids:
514 + if tid not in still_set:
515 + self._inflight_pop(tid)
516 + if not still:
517 + self._finalize_abort(kv_request_id)
518 +
519 + def _finalize_abort(self, kv_request_id: str) -> None:
520 + self._pending_aborts.pop(kv_request_id, None)
521 + self._send(
522 + {
523 + TYPE_KEY: AbortAckMsg.TYPE,
524 + AbortAckMsg.KV_REQUEST_ID: kv_request_id,
525 + }
526 + )
527 +
528 + # ------------------------------------------------------------------
529 + # Internal — transfers and store-job timeouts
530 + # ------------------------------------------------------------------
531 +
532 + def _submit_transfer(self, kv_request_id: str, result: _MatchResult) -> None:
533 + logger.debug(
534 + "P2PSession %s: NIXL write_blocks CALL kv_request_id=%s "
535 + "local_idxs=%d remote_idxs=%d",
536 + self._peer_id,
537 + kv_request_id,
538 + len(result.local_idxs),
539 + len(result.remote_idxs),
540 + )
541 + transfer_id = self._transport.write_blocks(
542 + self._peer_id, result.local_idxs, result.remote_idxs
543 + )
544 + if transfer_id is not None:
545 + logger.debug(
546 + "P2PSession %s: NIXL write_blocks SUBMITTED kv_request_id=%s "
547 + "transfer_id=%d blocks=%d",
548 + self._peer_id,
549 + kv_request_id,
550 + transfer_id,
551 + len(result.local_idxs),
552 + )
553 + self._inflight_add(
554 + transfer_id,
555 + _InflightXfer(
556 + kv_request_id=kv_request_id,
557 + block_count=len(result.local_idxs),
558 + job_ids=result.job_ids,
559 + ),
560 + )
561 + else:
562 + logger.warning(
563 + "P2PSession %s: write_blocks failed for %s (%d blocks)",
564 + self._peer_id,
565 + kv_request_id,
566 + len(result.local_idxs),
567 + )
568 + # The matched blocks were popped from req.demanded /
569 + # req.available, but no inflight will satisfy them, so
570 + # remaining will never reach 0 on its own. Mark the
571 + # request as finishing so the existing terminal paths
572 + # clean up: if other inflight is in flight, the last one
573 + # to drain will fire _finalize_outbound(success=False)
574 + # via the elif branch in collect_results. If
575 + # nothing else is in flight, finalize now so the peer
576 + # and the local store jobs don't wait for finish_request
577 + # or for _STORE_TIMEOUT_S / _LOAD_TIMEOUT_S.
578 + req = self._outbound.get(kv_request_id)
579 + if req is not None:
580 + req.finishing = True
581 + if not self._has_inflight_for(kv_request_id):
582 + self._finalize_outbound(kv_request_id, success=False)
583 +
584 + def _timeout_pending_store_jobs(self) -> list[StoreResult]:
585 + if not self._store_jobs:
586 + return []
587 + deadline = time.monotonic() - _STORE_TIMEOUT_S
588 + timed_out: list[int] | None = None
589 + for jid, submitted_at in self._store_jobs.items():
590 + if submitted_at <= deadline:
591 + if timed_out is None:
592 + timed_out = []
593 + timed_out.append(jid)
594 + if timed_out is None:
595 + return []
596 + results: list[StoreResult] = []
597 + for jid in timed_out:
598 + del self._store_jobs[jid]
599 + results.append(StoreResult(job_id=jid, success=False))
600 + logger.warning("P2PSession %s: store job %d timed out", self._peer_id, jid)
601 + return results
vllm/v1/kv_offload/tiering/p2p/session/session.py new
+440
@@ -0,0 +1,440 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +P2PSession — bidirectional session combining client + server roles.
5 +
6 +A single P2PSession per remote peer handles BOTH directions of the P2P
7 +protocol on one ControlConnection: it can request blocks from the peer
8 +("client" role, in :mod:`.client`) AND serve blocks to the peer
9 +("server" role, in :mod:`.server`). This module is the thin coordinator
10 +that owns the connection, the handshake, send-gating, and the message
11 +dispatch — each parsed message is forwarded to the corresponding role.
12 +
13 +Wire protocol is unchanged. Both sides advertise their NIXL metadata
14 +via ConnectMsg when their session is connected; the peer's ConnectMsg
15 +triggers transport.add_remote_peer; ConnectAckMsg confirms the peer
16 +received our ConnectMsg, after which queued outgoing messages are flushed.
17 +"""
18 +
19 +from __future__ import annotations
20 +
21 +import contextlib
22 +from collections.abc import Sequence
23 +from typing import TYPE_CHECKING, NamedTuple
24 +
25 +from vllm.logger import init_logger
26 +from vllm.v1.kv_offload.base import OffloadKey
27 +from vllm.v1.kv_offload.tiering.p2p.control.base import ControlConnection
28 +from vllm.v1.kv_offload.tiering.p2p.session.client import ClientRole, LoadResult
29 +from vllm.v1.kv_offload.tiering.p2p.session.protocol import (
30 + TYPE_KEY,
31 + AbortAckMsg,
32 + AbortFetchMsg,
33 + ConnectAckMsg,
34 + ConnectMsg,
35 + DisconnectMsg,
36 + FetchMsg,
37 + TransferDoneMsg,
38 +)
39 +from vllm.v1.kv_offload.tiering.p2p.session.server import (
40 + ServerRole,
41 + StoreResult,
42 +)
43 +
44 +if TYPE_CHECKING:
45 + from vllm.v1.kv_offload.tiering.base import JobId
46 + from vllm.v1.kv_offload.tiering.p2p.data import DataTransport
47 +
48 +logger = init_logger(__name__)
49 +
50 +# Cap on consecutive non-protocol dispatch exceptions before we tear
51 +# down the session. Protocol violations (ValueError) disconnect on the
52 +# first occurrence; this threshold protects against repeated internal
53 +# bugs that may indicate a peer-induced bad state. Reset on any
54 +# successful dispatch.
55 +_MAX_CONSECUTIVE_DISPATCH_ERRORS = 5
56 +
57 +
58 +class SessionPollResult(NamedTuple):
59 + """Result of one P2PSession.poll() tick.
60 +
61 + `loads`/`stores` are the same per-role results the manager has always
62 + consumed. `new_fetch_ids` reports kv_request_ids whose FetchMsg
63 + arrived this tick — the manager uses them to bind kv_request_id →
64 + session and replay any submit_store batches parked while no peer had
65 + asked yet. Reporting (rather than calling back into the manager
66 + mid-dispatch) keeps the dependency strictly top-down.
67 + """
68 +
69 + loads: list[LoadResult]
70 + stores: list[StoreResult]
71 + new_fetch_ids: list[str]
72 +
73 +
74 +class P2PSession:
75 + """Bidirectional session — coordinator over ClientRole + ServerRole.
76 +
77 + Lifecycle:
78 + - Constructor with conn=None ⇒ pending. Accepts add_stored_blocks
79 + but cannot send (used by the prefiller to buffer blocks before
80 + the decoder connects).
81 + - Constructor with conn != None ⇒ connected. Sends our own ConnectMsg
82 + immediately; the peer's ConnectMsg arrives in poll() and is
83 + dispatched to _on_connect (which calls transport.add_remote_peer
84 + and replies with ConnectAckMsg). Outgoing sends are queued until
85 + ConnectAckMsg confirms our metadata reached the peer.
86 + - attach_connection(conn) on a pending session ⇒ same as above,
87 + starting from pending.
88 + """
89 +
90 + def __init__(
91 + self,
92 + peer_id: str,
93 + local_id: str,
94 + transport: DataTransport,
95 + local_block_len: int,
96 + conn: ControlConnection | None = None,
97 + ) -> None:
98 + self.peer_id = peer_id
99 + self._local_id = local_id
100 + self._transport = transport
101 + self._local_block_len = local_block_len
102 + self._conn: ControlConnection | None = None
103 +
104 + self._send_ready = False # True after the peer acked our ConnectMsg
105 + # Msgs waiting to be sent on connection establishment
106 + self._queued: list[dict] = []
107 +
108 + # Consecutive non-protocol dispatch errors. Reset on success.
109 + self._dispatch_error_count: int = 0
110 +
111 + # kv_request_ids whose FetchMsg arrived during the current poll
112 + # tick. Drained and returned in the next poll() result so the
113 + # manager can bind kv_request_id → session and replay any
114 + # submit_store batches parked before the binding existed.
115 + self._new_fetch_ids: list[str] = []
116 +
117 + self._client = ClientRole(peer_id=peer_id, send=self._send)
118 + self._server = ServerRole(peer_id=peer_id, transport=transport, send=self._send)
119 +
120 + if conn is not None:
121 + self.attach_connection(conn)
122 +
123 + # ------------------------------------------------------------------
124 + # Properties
125 + # ------------------------------------------------------------------
126 +
127 + @property
128 + def alive(self) -> bool:
129 + # Pending sessions (awaiting connection) are alive — only a
130 + # closed real connection counts as dead.
131 + return self._conn is None or self._conn.alive
132 +
133 + @property
134 + def connected(self) -> bool:
135 + return self._conn is not None
136 +
137 + @property
138 + def ready(self) -> bool:
139 + """True after the peer acked our ConnectMsg (we may send freely)."""
140 + return self._send_ready
141 +
142 + # ------------------------------------------------------------------
143 + # Connection lifecycle
144 + # ------------------------------------------------------------------
145 +
146 + def attach_connection(self, conn: ControlConnection) -> None:
147 + """Attach a connection to a pending session and announce ourselves.
148 +
149 + Symmetric: every side advertises its NIXL metadata on connect, so
150 + whichever peer receives a session first can register the other.
151 + """
152 + if self._conn is not None:
153 + raise ValueError(f"P2PSession {self.peer_id}: already connected")
154 + self._conn = conn
155 + self._send_connect()
156 +
157 + # ------------------------------------------------------------------
158 + # Public API
159 + # ------------------------------------------------------------------
160 +
161 + def request_blocks(
162 + self,
163 + job_id: JobId,
164 + kv_request_id: str,
165 + keys: Sequence[bytes],
166 + block_ids: Sequence[int],
167 + ) -> None:
168 + """Send fetch to the peer."""
169 + self._client.request_blocks(
170 + job_id, kv_request_id, keys, block_ids, send_ready=self._send_ready
171 + )
172 +
173 + def add_stored_blocks(
174 + self,
175 + kv_request_id: str,
176 + keys: Sequence[OffloadKey],
177 + block_ids: Sequence[int],
178 + job_id: JobId,
179 + ) -> None:
180 + """New blocks stored locally — match against pending fetch demand."""
181 + self._server.add_stored_blocks(kv_request_id, keys, block_ids, job_id)
182 +
183 + def finish_request(self, kv_request_id: str) -> None:
184 + """Called when the request is finishing locally.
185 +
186 + Cancels any inbound load (client role) and finalizes any
187 + outbound serving (server role) for this id. Roles that aren't
188 + active for this id are silent no-ops.
189 + """
190 + self._client.cancel(kv_request_id)
191 + self._server.finish(kv_request_id)
192 +
193 + def poll(self) -> SessionPollResult:
194 + """Process incoming messages, drive transfers, apply timeouts."""
195 + if self._conn is None:
196 + # Pending session — store-job timeouts still apply so buffered
197 + # jobs that never get picked up are surfaced as failures.
198 + return SessionPollResult(
199 + loads=[],
200 + stores=self._server.collect_idle_timeouts(),
201 + new_fetch_ids=[],
202 + )
203 +
204 + for msg in self._conn.recv():
205 + self._on_message(msg)
206 +
207 + loads = self._client.collect_results()
208 + stores = self._server.collect_results()
209 + self._server.drain_pending_aborts()
210 +
211 + new_fetch_ids = self._new_fetch_ids
212 + self._new_fetch_ids = []
213 + return SessionPollResult(
214 + loads=loads, stores=stores, new_fetch_ids=new_fetch_ids
215 + )
216 +
217 + def close(self) -> tuple[list[tuple[int, str]], list[int]]:
218 + """Shut down. Returns (failed_loads, failed_stores).
219 +
220 + failed_loads: list of (job_id, kv_request_id) pairs.
221 + failed_stores: list of job_ids.
222 + """
223 + failed_loads = self._client.close()
224 + failed_stores = self._server.close()
225 +
226 + if self._conn is not None:
227 + with contextlib.suppress(Exception):
228 + self._conn.send({TYPE_KEY: DisconnectMsg.TYPE})
229 + self._conn.close()
230 + self._conn = None
231 +
232 + return failed_loads, failed_stores
233 +
234 + # ------------------------------------------------------------------
235 + # Message dispatch
236 + # ------------------------------------------------------------------
237 +
238 + def _on_message(self, msg: dict) -> None:
239 + msg_type = msg.get(TYPE_KEY) if isinstance(msg, dict) else msg
240 + try:
241 + self._dispatch_message(msg)
242 + except ValueError as exc:
243 + # Protocol contract violation from the peer — *Msg.validate()
244 + # and handler-level checks raise ValueError. Retrying won't
245 + # help and may corrupt session state, so disconnect now.
246 + self._protocol_error(f"malformed {msg_type!r}: {exc}")
247 + return
248 + except Exception as exc:
249 + # Anything else is most likely an internal bug rather than a
250 + # peer fault. Log loudly with a traceback so it doesn't
251 + # disappear, but don't kill the session on a single hiccup.
252 + # Disconnect only if errors keep arriving — that pattern is
253 + # consistent with a peer wedging us into a broken state.
254 + self._dispatch_error_count += 1
255 + logger.exception(
256 + "P2PSession %s: error handling message %r (count=%d): %s",
257 + self.peer_id,
258 + msg_type,
259 + self._dispatch_error_count,
260 + exc,
261 + )
262 + if self._dispatch_error_count >= _MAX_CONSECUTIVE_DISPATCH_ERRORS:
263 + self._protocol_error(
264 + f"too many consecutive dispatch errors "
265 + f"({self._dispatch_error_count})"
266 + )
267 + return
268 + self._dispatch_error_count = 0
269 +
270 + def _protocol_error(self, reason: str) -> None:
271 + """Log a protocol violation and disconnect.
272 +
273 + Best-effort sends ``DisconnectMsg`` so the peer learns why we're
274 + going away, then marks the connection dead. The manager reaps
275 + the session on the next poll via ``alive``.
276 + """
277 + logger.error(
278 + "P2PSession %s: protocol error: %s — disconnecting",
279 + self.peer_id,
280 + reason,
281 + )
282 + if self._conn is not None:
283 + with contextlib.suppress(Exception):
284 + self._conn.send({TYPE_KEY: DisconnectMsg.TYPE})
285 + self._conn.mark_dead()
286 +
287 + def _dispatch_message(self, msg: dict) -> None:
288 + # Drop messages buffered before disconnect: a poll batch can
289 + # contain msg-after-DisconnectMsg, and dispatching them would
290 + # mutate state on a dead session.
291 + if self._conn is not None and not self._conn.alive:
292 + return
293 + msg_type = msg.get(TYPE_KEY) if isinstance(msg, dict) else None
294 + if msg_type == ConnectMsg.TYPE:
295 + self._on_connect(msg)
296 + elif msg_type == ConnectAckMsg.TYPE:
297 + ConnectAckMsg.validate(msg)
298 + self._on_connect_ack()
299 + elif msg_type == FetchMsg.TYPE:
300 + FetchMsg.validate(msg)
301 + kv_request_id = msg[FetchMsg.KV_REQUEST_ID]
302 + block_hashes = [
303 + OffloadKey(bh if isinstance(bh, bytes) else bytes(bh))
304 + for bh in msg[FetchMsg.BLOCK_HASHES]
305 + ]
306 + block_indexes = msg[FetchMsg.BLOCK_INDEXES]
307 + # Run the server-role state machine inline as today —
308 + # add_fetch_demand records demand against any blocks we've
309 + # already seen in `available`. Report the kv_request_id so
310 + # the manager (after poll() returns) can replay any parked
311 + # submit_store batches; their add_stored_blocks calls hit
312 + # the demand recorded here and submit transfers immediately.
313 + self._server.on_fetch(kv_request_id, block_hashes, block_indexes)
314 + self._new_fetch_ids.append(kv_request_id)
315 + elif msg_type == AbortFetchMsg.TYPE:
316 + AbortFetchMsg.validate(msg)
317 + self._server.on_abort_fetch(msg[AbortFetchMsg.KV_REQUEST_ID])
318 + elif msg_type == TransferDoneMsg.TYPE:
319 + TransferDoneMsg.validate(msg)
320 + self._client.on_transfer_done(
321 + msg[TransferDoneMsg.KV_REQUEST_ID],
322 + msg[TransferDoneMsg.SUCCESS],
323 + )
324 + elif msg_type == AbortAckMsg.TYPE:
325 + AbortAckMsg.validate(msg)
326 + self._client.on_abort_ack(msg[AbortAckMsg.KV_REQUEST_ID])
327 + elif msg_type == DisconnectMsg.TYPE:
328 + if self._conn is not None:
329 + self._conn.mark_dead()
330 + else:
331 + logger.warning(
332 + "P2PSession %s: unknown message type %r", self.peer_id, msg_type
333 + )
334 +
335 + # ------------------------------------------------------------------
336 + # Handshake
337 + # ------------------------------------------------------------------
338 +
339 + def _on_connect(self, msg: dict) -> None:
340 + # Validation failures here mean an incompatible or malicious peer.
341 + # Mark the connection dead so the manager reaps the session;
342 + # don't call add_remote_peer or send connect_ack.
343 + if self._send_ready:
344 + # We've already received connect_ack, so the handshake is
345 + # complete. A second connect from the peer is a protocol
346 + # violation — re-registering would corrupt transport state.
347 + self._protocol_error("duplicate connect after handshake")
348 + return
349 + try:
350 + ConnectMsg.validate(msg)
351 + if msg[ConnectMsg.BLOCK_LEN] != self._local_block_len:
352 + raise ValueError(
353 + f"block_len mismatch from {self.peer_id}: "
354 + f"remote={msg[ConnectMsg.BLOCK_LEN]}, "
355 + f"local={self._local_block_len}"
356 + )
357 + remote_fp = msg.get(ConnectMsg.CONFIG_FINGERPRINT, "")
358 + local_fp = self._transport.config_fingerprint
359 + if local_fp and remote_fp and remote_fp != local_fp:
360 + raise ValueError(
361 + f"config fingerprint mismatch from {self.peer_id}: "
362 + f"remote={remote_fp!r}, local={local_fp!r}"
363 + )
364 + self._transport.add_remote_peer(
365 + self.peer_id,
366 + agent_metadata=msg[ConnectMsg.AGENT_METADATA],
367 + base_addr=msg[ConnectMsg.BASE_ADDR],
368 + num_blocks=msg[ConnectMsg.NUM_BLOCKS],
369 + block_len=msg[ConnectMsg.BLOCK_LEN],
370 + )
371 + except ValueError as exc:
372 + logger.error("P2PSession %s: rejecting peer connect: %s", self.peer_id, exc)
373 + if self._conn is not None:
374 + self._conn.mark_dead()
375 + return
376 +
377 + if self._conn is not None:
378 + self._conn.send(
379 + {
380 + TYPE_KEY: ConnectAckMsg.TYPE,
381 + ConnectAckMsg.PEER_ID: self._local_id,
382 + }
383 + )
384 +
385 + def _on_connect_ack(self) -> None:
386 + if self._queued:
387 + logger.debug(
388 + "P2PSession %s: connect_ack received, flushing %d queued msg(s)",
389 + self.peer_id,
390 + len(self._queued),
391 + )
392 + self._send_ready = True
393 + for queued in self._queued:
394 + self._do_send(queued)
395 + self._queued.clear()
396 +
397 + # ------------------------------------------------------------------
398 + # Send helpers
399 + # ------------------------------------------------------------------
400 +
401 + def _send_connect(self) -> None:
402 + """Send our ConnectMsg announcing local NIXL metadata."""
403 + assert self._conn is not None
404 + self._conn.send(
405 + {
406 + TYPE_KEY: ConnectMsg.TYPE,
407 + ConnectMsg.PEER_ID: self._local_id,
408 + ConnectMsg.AGENT_METADATA: self._transport.get_agent_metadata(),
409 + ConnectMsg.BASE_ADDR: self._transport.base_addr,
410 + ConnectMsg.NUM_BLOCKS: self._transport.num_blocks,
411 + ConnectMsg.BLOCK_LEN: self._transport.block_len,
412 + ConnectMsg.CONFIG_FINGERPRINT: self._transport.config_fingerprint,
413 + }
414 + )
415 +
416 + def _send(self, msg: dict) -> None:
417 + if self._conn is None or not self._send_ready:
418 + logger.debug(
419 + "P2PSession %s: queueing %s (ready=%s queue_depth=%d)",
420 + self.peer_id,
421 + msg.get(TYPE_KEY),
422 + self._send_ready,
423 + len(self._queued) + 1,
424 + )
425 + self._queued.append(msg)
426 + return
427 + self._do_send(msg)
428 +
429 + def _do_send(self, msg: dict) -> None:
430 + if self._conn is None:
431 + return
432 + try:
433 + self._conn.send(msg)
434 + logger.debug("P2PSession %s: sent %s", self.peer_id, msg.get(TYPE_KEY))
435 + except Exception:
436 + logger.warning(
437 + "P2PSession %s: failed to send %s",
438 + self.peer_id,
439 + msg.get(TYPE_KEY),
440 + )